diff --git a/.buildkite/amd/scripts/bootstrap-amd-omni.sh b/.buildkite/amd/scripts/bootstrap-amd-omni.sh index 32097cb8c..4eb7f6eb4 100755 --- a/.buildkite/amd/scripts/bootstrap-amd-omni.sh +++ b/.buildkite/amd/scripts/bootstrap-amd-omni.sh @@ -1,4 +1,7 @@ #!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # vllm-omni customized version # Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh # Last synced: 2025-12-15 @@ -6,6 +9,9 @@ set -euo pipefail +# The bootstrap runs from the repository root; ShellCheck does not follow +# sourced files unless invoked with -x. +# shellcheck disable=SC1091 source .buildkite/common/scripts/resolve_skip_ci.sh if [[ -z "${RUN_ALL:-}" ]]; then @@ -63,6 +69,8 @@ upload_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 + # Installed by the minijinja bootstrap above and only present on the CI agent. + # shellcheck disable=SC1091 source /var/lib/buildkite-agent/.cargo/env if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then @@ -90,22 +98,82 @@ upload_pipeline() { cd .buildkite/amd # Select test definition file: merge suite for main, ready suite for PRs. - # For debugging, DEBUG_TEST_YAML can override the selection — accepts - # "merge" or "ready" (case-insensitive). + # For debugging, DEBUG_TEST_YAML accepts a comma-separated list containing + # "merge" and/or "ready" (case-insensitive). Multiple suites are combined + # before rendering so they share one amd-build step. if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then - case "${DEBUG_TEST_YAML,,}" in - merge) - TEST_YAML="test-amd-merge.yml" - ;; - ready) - TEST_YAML="test-amd-ready.yml" - ;; - *) - echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2 + declare -a DEBUG_TEST_SPECS=() + declare -A SEEN_DEBUG_TESTS=() + IFS=',' read -ra REQUESTED_DEBUG_TESTS <<< "${DEBUG_TEST_YAML,,}" + + for requested_test in "${REQUESTED_DEBUG_TESTS[@]}"; do + # Permit readable values such as "ready, merge," and ignore the + # empty item produced by a trailing comma. + requested_test="${requested_test//[[:space:]]/}" + [[ -z "$requested_test" ]] && continue + + if [[ -n "${SEEN_DEBUG_TESTS[$requested_test]:-}" ]]; then + echo "ERROR: duplicate DEBUG_TEST_YAML suite '$requested_test'" >&2 exit 1 - ;; - esac - echo "DEBUG_TEST_YAML override: using $TEST_YAML" + fi + SEEN_DEBUG_TESTS[$requested_test]=1 + + case "$requested_test" in + ready) + DEBUG_TEST_SPECS+=("READY_TESTS:test-amd-ready.yml") + ;; + merge) + DEBUG_TEST_SPECS+=("MERGE_TESTS:test-amd-merge.yml") + ;; + *) + echo "ERROR: DEBUG_TEST_YAML entries must be 'merge' or 'ready', got '$requested_test'" >&2 + exit 1 + ;; + esac + done + + if [[ ${#DEBUG_TEST_SPECS[@]} -eq 0 ]]; then + echo "ERROR: DEBUG_TEST_YAML did not contain a test suite" >&2 + exit 1 + elif [[ ${#DEBUG_TEST_SPECS[@]} -eq 1 ]]; then + TEST_YAML="${DEBUG_TEST_SPECS[0]#*:}" + else + TEST_YAML=$(mktemp "${TMPDIR:-/tmp}/amd-debug-tests.XXXXXX.yml") + python - "$TEST_YAML" "${DEBUG_TEST_SPECS[@]}" <<'PY' +import sys + +import yaml + + +output_path, *suite_specs = sys.argv[1:] +combined = {"env": {}, "steps": []} + +for suite_spec in suite_specs: + group_name, input_path = suite_spec.split(":", 1) + with open(input_path, encoding="utf-8") as test_file: + suite = yaml.safe_load(test_file) + + for name, value in (suite.get("env") or {}).items(): + previous = combined["env"].get(name, value) + if previous != value: + raise ValueError( + f"Conflicting environment value for {name}: {previous!r} != {value!r}" + ) + combined["env"][name] = value + + suite_steps = [] + for entry in suite.get("steps") or []: + if "group" in entry: + suite_steps.extend(entry.get("steps") or []) + else: + suite_steps.append(entry) + combined["steps"].append({"group": group_name, "steps": suite_steps}) + +with open(output_path, "w", encoding="utf-8") as output_file: + yaml.safe_dump(combined, output_file, sort_keys=False) +PY + fi + echo "DEBUG_TEST_YAML override: using ${DEBUG_TEST_SPECS[*]}" elif [[ $BUILDKITE_BRANCH == "main" ]]; then TEST_YAML="test-amd-merge.yml" else @@ -130,19 +198,20 @@ upload_pipeline() { > pipeline.yaml ) cat pipeline.yaml + if [[ "$TEST_YAML" == "${TMPDIR:-/tmp}/amd-debug-tests."*.yml ]]; then + rm -f -- "$TEST_YAML" + fi 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)) + 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) + git diff --name-only --diff-filter=ACMDR HEAD~1 } file_diff=$(get_diff) @@ -181,7 +250,7 @@ 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 + if [[ $file == "$pattern"* ]] || [[ $file == "$pattern" ]]; then matches_pattern=1 break fi @@ -191,7 +260,7 @@ for file in $file_diff; do if [[ $matches_pattern -eq 1 ]]; then matches_ignore=0 for ignore in "${ignore_patterns[@]}"; do - if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then + if [[ $file == "$ignore"* ]] || [[ $file == "$ignore" ]]; then matches_ignore=1 break fi diff --git a/.buildkite/amd/scripts/run-amd-test.sh b/.buildkite/amd/scripts/run-amd-test.sh index c49614e70..c1c80c495 100755 --- a/.buildkite/amd/scripts/run-amd-test.sh +++ b/.buildkite/amd/scripts/run-amd-test.sh @@ -1,153 +1,103 @@ #!/bin/bash -# vllm-omni customized version -# Based on: vllm/.buildkite/scripts/hardware_ci/run-amd-test.sh -# Last synced: 2025-12-15 -# Modifications: docker image name for vllm-omni -# This script runs test inside the corresponding ROCm docker container. -set -o pipefail +# Run vLLM-Omni ROCm tests directly in the MI300 Kubernetes pod. The pod's +# container image is selected by test-template-amd-omni.j2; MI300 has no DinD. +set -euo pipefail -# Export Python path -export PYTHONPATH=".." +: "${PYTHONFAULTHANDLER:=1}" +: "${HF_HOME:=/home/buildkite-agent/huggingface}" +: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" +: "${HF_HUB_ETAG_TIMEOUT:=60}" +: "${MIOPEN_DEBUG_CONV_DIRECT:=0}" +: "${MIOPEN_DEBUG_CONV_GEMM:=0}" +: "${VLLM_ROCM_USE_AITER:=0}" +export PYTHONFAULTHANDLER HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT +export MIOPEN_DEBUG_CONV_DIRECT MIOPEN_DEBUG_CONV_GEMM VLLM_ROCM_USE_AITER +export PYTORCH_ROCM_ARCH="" +export PYTHONPATH="${PYTHONPATH:-..}" + +if [[ "${VLLM_CI_DOCKER_DISABLED:-0}" != "1" ]]; then + echo "Error: MI300 CI must run natively with Docker disabled." >&2 + exit 1 +fi + +if [[ -n "${TEST_COMMANDS:-}" ]]; then + commands="${TEST_COMMANDS}" +else + commands="$*" +fi +if [[ -z "${commands}" ]]; then + echo "Error: No test commands provided." >&2 + exit 1 +fi + +echo "--- Native in-pod ROCm CI" + +job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}" +job_id="${job_id//[^A-Za-z0-9_.-]/_}" +native_cache_root="/tmp/vllm-omni-native-${job_id}" +TMPDIR="${native_cache_root}/tmp" +TORCHINDUCTOR_CACHE_DIR="${native_cache_root}/torchinductor" +TRITON_CACHE_DIR="${native_cache_root}/triton" +VLLM_CACHE_ROOT="${native_cache_root}/vllm" +XDG_CACHE_HOME="${native_cache_root}/xdg" +HF_DATASETS_CACHE="${native_cache_root}/huggingface/datasets" +export TMPDIR TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT +export XDG_CACHE_HOME HF_DATASETS_CACHE +mkdir -p "${TMPDIR}" "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" \ + "${VLLM_CACHE_ROOT}" "${XDG_CACHE_HOME}" "${HF_HOME}" \ + "${HF_DATASETS_CACHE}" + +if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then + if ! command -v findmnt >/dev/null 2>&1; then + echo "Error: findmnt is required to verify the Hugging Face cache mount." >&2 + exit 1 + fi + hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true) + if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then + echo "Error: MI300 CI requires a persistent volume at or above ${HF_HOME}." >&2 + exit 1 + fi +fi -# Print ROCm version -echo "--- ROCm info" rocminfo -# cleanup older docker images -cleanup_docker() { - # Get Docker's root directory - docker_root=$(docker info -f '{{.DockerRootDir}}') - if [ -z "$docker_root" ]; then - echo "Failed to determine Docker root directory." - exit 1 - fi - echo "Docker root directory: $docker_root" - # Check disk usage of the filesystem where Docker's root directory is located - disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//') - # Define the threshold - threshold=70 - if [ "$disk_usage" -gt "$threshold" ]; then - echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..." - # Remove dangling images (those that are not tagged and not used by any container) - docker image prune -f - # Remove unused volumes / force the system prune for old images as well. - docker volume prune -f && docker system prune --force --filter "until=72h" --all - echo "Docker images and volumes cleanup completed." - else - echo "Disk usage is below $threshold%. No cleanup needed." - fi -} +expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}" +python3 - <<'PY' +import os -# Call the cleanup docker function -cleanup_docker +import torch -echo "--- Pulling container" -## Temporary change to use AMD Docker Hub to store the vllm-omni image -# to bypass the rate limit issue with ECR Public Gallery. -# Images are now stored in a separate repository for vllm-omni, instead of vllm-ci. -# TODO: @tjtanaa point back to ECR Public Gallery -# once the amd agents are configured to use ECR Public Gallery. -# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni" -image_name="rocm/vllm-omni:${BUILDKITE_COMMIT}" -container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)" +expected = int(os.environ.get("VLLM_CI_EXPECTED_GPU_COUNT", "1")) +assert torch.version.hip, "PyTorch is not a ROCm build" +assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch" +actual = torch.cuda.device_count() +assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}" +PY -# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery. -# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images -# sudo apt-get update && sudo apt-get install -y awscli -## Use safe docker login helper to prevent race conditions -# source "$(dirname "${BASH_SOURCE[0]}")/../../common/scripts/docker_login_ecr_public.sh" -# safe_docker_login_ecr_public - -## Pull the container from AMD Docker Hub - -docker pull "${image_name}" - -remove_docker_container() { - docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true - docker image rm -f "${image_name}" || true -} -trap remove_docker_container EXIT - -echo "--- Running container" - -HF_CACHE="$(realpath ~)/huggingface" -mkdir -p "${HF_CACHE}" -HF_MOUNT="/root/.cache/huggingface" - -if [[ -n "${TEST_COMMAND:-}" ]]; then - commands="$TEST_COMMAND" -else - commands="$@" -fi -echo "Commands:$commands" - -PARALLEL_JOB_COUNT=8 -MYPYTHONPATH=".." - -# Test that we're launching on the machine that has -# proper access to GPUs -render_gid=$(getent group render | cut -d: -f3) -if [[ -z "$render_gid" ]]; then - echo "Error: 'render' group not found. This is required for GPU access." >&2 - exit 1 -fi - -# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs. -# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels. -if [[ $commands == *"--shard-id="* ]]; then - # assign job count as the number of shards used - commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g') - for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do - # assign shard-id for each shard - commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g') - echo "Shard ${GPU} commands:$commands_gpu" - echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES" - docker run \ - --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ - --network=host \ - --shm-size=16gb \ - --group-add "$render_gid" \ - -e MIOPEN_DEBUG_CONV_DIRECT=0 \ - -e MIOPEN_DEBUG_CONV_GEMM=0 \ - -e VLLM_ROCM_USE_AITER=0 \ - -e HIP_VISIBLE_DEVICES="${GPU}" \ - -e HF_TOKEN \ - -v "${HF_CACHE}:${HF_MOUNT}" \ - -e "HF_HOME=${HF_MOUNT}" \ - -e "PYTHONPATH=${MYPYTHONPATH}" \ - --name "${container_name}_${GPU}" \ - "${image_name}" \ - /bin/bash -c "${commands_gpu}" \ - |& while read -r line; do echo ">>Shard $GPU: $line"; done & - PIDS+=($!) - done - #wait for all processes to finish and collect exit codes - for pid in "${PIDS[@]}"; do - wait "${pid}" - STATUS+=($?) - done - for st in "${STATUS[@]}"; do - if [[ ${st} -ne 0 ]]; then - echo "One of the processes failed with $st" - exit "${st}" +echo "Commands:${commands}" +if /bin/bash -o pipefail -c ' +set -E +test_status=0 +trap '\'' + command_status=$? + # Keep running subsequent commands, but retain a failure for the final exit. + # Prefer a test failure over pytest "no tests collected" when both occur. + if (( test_status == 0 || (test_status == 5 && command_status != 5) )); then + test_status=${command_status} fi - done +'\'' ERR +eval "$1" +exit "${test_status}" +' _ "${commands}"; then + exit 0 else - echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES" - docker run \ - --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \ - --network=host \ - --shm-size=16gb \ - --group-add "$render_gid" \ - -e MIOPEN_DEBUG_CONV_DIRECT=0 \ - -e MIOPEN_DEBUG_CONV_GEMM=0 \ - -e VLLM_ROCM_USE_AITER=0 \ - -e HF_TOKEN \ - -v "${HF_CACHE}:${HF_MOUNT}" \ - -e "HF_HOME=${HF_MOUNT}" \ - -e "PYTHONPATH=${MYPYTHONPATH}" \ - --name "${container_name}" \ - "${image_name}" \ - /bin/bash -c "${commands}" + exit_code=$? fi + +if [[ ${exit_code} -eq 5 && "${VLLM_CI_ALLOW_NO_TESTS:-0}" == "1" ]]; then + echo "Pytest collected no tests; VLLM_CI_ALLOW_NO_TESTS=1, treating exit code 5 as success." + exit 0 +fi + +exit "${exit_code}" diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml index 4e80d1968..56083b774 100644 --- a/.buildkite/amd/test-amd-merge.yml +++ b/.buildkite/amd/test-amd-merge.yml @@ -7,84 +7,113 @@ steps: - group: ":card_index_dividers: Simple Test" steps: - - label: "Simple · Diffusion & Model Executor Test (1/2)" - agent_pool: mi325_1 + - label: "Simple · Model Executor Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 2 depends_on: amd-build - timeout_in_minutes: 70 mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - # CPU tests on ROCm gpu machines are very slow, so we give them more time. - - "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" - - label: "Simple · Diffusion & Model Executor Test (2/2)" - agent_pool: mi325_1 + - label: "Simple · Diffusion Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 4 depends_on: amd-build - timeout_in_minutes: 70 mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_ROCM_USE_AITER=0 - # CPU tests on ROCm gpu machines are very slow, so we give them more time. - - "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" + + - label: "Simple · Engine&Entrypoints Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'" - label: "Simple · Other Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build - timeout_in_minutes: 50 mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py" - group: ":card_index_dividers: Diffusion Test" steps: - label: "Diffusion · LoRA Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking + artifact_paths: + - lora-server.log commands: - - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py + - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model" 2>&1 | tee lora-server.log - label: "Diffusion Model CPU offloading Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - | - timeout 20m bash -c ' - set +e - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py - EXIT1=\$? - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py - EXIT2=\$? - exit \$((EXIT1 | EXIT2)) - ' + - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py + - >- + pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py + --deselect='tests/diffusion/offloader/test_diffusion_layerwise_offload.py::test_layerwise_offload_diffusion_model[stabilityai/stable-audio-open-1.0]' - - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)" - agent_pool: mi325_4 + # TODO: Diagnose unstable Stable Audio layerwise-offload memory savings on ROCm. + - label: "Diffusion Stable Audio Layerwise Offload Test (quarantined)" + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] + grade: NonBlocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - >- + pytest -s -v + 'tests/diffusion/offloader/test_diffusion_layerwise_offload.py::test_layerwise_offload_diffusion_model[stabilityai/stable-audio-open-1.0]' + + - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 1/2" + agent_pool: mi300_4 + depends_on: amd-build + mirror_hardwares: [amdproduction] + timeout_in_minutes: 40 grade: Blocking commands: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py - - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py + - pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py + + - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs) · Shard 2/2" + agent_pool: mi300_4 + depends_on: amd-build + mirror_hardwares: [amdproduction] + timeout_in_minutes: 40 + grade: Blocking + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py # merge-only tests - label: "Diffusion Tensor Parallelism Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -92,7 +121,7 @@ steps: - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py - label: "Diffusion Model Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -100,7 +129,7 @@ steps: - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py - label: "Diffusion GPU Worker Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -108,7 +137,7 @@ steps: - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py - label: "Diffusion · Wan22 Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -121,7 +150,7 @@ steps: ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 # - label: "Audio Generation Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -132,7 +161,7 @@ steps: # - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model - label: "Engine Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -147,7 +176,7 @@ steps: - group: ":card_index_dividers: E2E Test" steps: - label: "Omni Model Test Qwen2-5-Omni" - agent_pool: mi325_4 + agent_pool: mi300_4 depends_on: amd-build mirror_hardwares: [amdproduction] commands: @@ -155,30 +184,55 @@ steps: - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model" - label: "Omni Model Test Qwen3-Omni" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' + - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model' - - label: "Qwen3-TTS CustomVoice E2E Test" - agent_pool: mi325_1 + - label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 3 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking + timeout_in_minutes: 80 commands: - | - timeout 50m bash -c ' + timeout --signal=TERM --kill-after=2m 80m bash -c ' export VLLM_LOGGING_LEVEL=DEBUG export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" + TESTS=( + "tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_default_cuda_graph_startup[async_chunk]" + "tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_text_to_audio_002[async_chunk]" + "tests/e2e/offline_inference/test_qwen3_tts_customvoice.py::test_text_to_audio_001[no_cuda_graph]" + ) + TEST="$${TESTS[$$BUILDKITE_PARALLEL_JOB]}" + pytest -s -v "$$TEST" -m "advanced_model" --run-level "advanced_model" + ' + + # TODO: Restore this case to the blocking sharded step once its long-running + # Whisper validation on AMD CI is resolved. + - label: "Qwen3-TTS CustomVoice async_chunk Test (quarantined)" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: NonBlocking + timeout_in_minutes: 80 + commands: + - | + timeout --signal=TERM --kill-after=2m 75m bash -c ' + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + pytest -s -v "tests/e2e/online_serving/test_qwen3_tts_customvoice.py::test_text_to_audio_001[async_chunk]" -m "advanced_model" --run-level "advanced_model" ' - label: "Qwen3-TTS Base E2E Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -196,7 +250,7 @@ steps: # Need to debug before reneable numerical changes across large PRs # # split Bagel Model Test with H100 (Real Weights) into three tests # - label: "Bagel Text2Img Model Test (1/3)" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -208,7 +262,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm" # - label: "Bagel Img2Img Model Test (2/3)" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -220,7 +274,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm" # - label: "Bagel Online Serving Test (3/3)" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 8f99ce909..695f47dc1 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -7,8 +7,9 @@ steps: - group: ":card_index_dividers: Simple Test" steps: - - label: "Simple · Diffusion & Model Executor Test" - agent_pool: mi325_1 + - label: "Simple · Model Executor Test · Shard %N/%t" + agent_pool: mi300_1 + parallelism: 2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -16,10 +17,30 @@ steps: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" + + - label: "Simple · Diffusion Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM + # that is not supported on CPU + - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py" + + - label: "Simple · Engine&Entrypoints Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'" - label: "Simple · Other Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -27,19 +48,19 @@ steps: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --ignore=tests/diffusion/cache/test_teacache_extractors.py" - label: "Custom Pipeline Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" + - pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" - label: "Voxtral TTS CUDA Unit Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -49,33 +70,25 @@ steps: - group: ":card_index_dividers: Diffusion Test" steps: - label: "Diffusion · Batch Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model" + - pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model" - label: "Diffusion · Offloader Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - | - timeout 20m bash -c ' - set +e - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py - EXIT1=\$? - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py - EXIT2=\$? - exit \$((EXIT1 | EXIT2)) - ' + - pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model" - label: "Diffusion · Cache Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -83,10 +96,10 @@ steps: - export GPU_ARCHS=gfx942 - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4" + - pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model" - label: "Diffusion Sequence Parallelism Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -96,7 +109,7 @@ steps: - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model - label: "Diffusion GPU Worker Test" - agent_pool: mi325_2 + agent_pool: mi300_2 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -104,7 +117,7 @@ steps: - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py - label: "Diffusion · Model Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -112,7 +125,7 @@ steps: - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model" - label: "Diffusion · Wan22 Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -126,7 +139,7 @@ steps: ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274 # - label: "Audio Generation Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -137,7 +150,7 @@ steps: # - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model - label: "Engine Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -149,7 +162,7 @@ steps: ' # - label: "Omni Model Test Qwen3-Omni" -# agent_pool: mi325_2 +# agent_pool: mi300_2 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -162,7 +175,7 @@ steps: - group: ":card_index_dividers: E2E Test" steps: - label: "Qwen3-TTS E2E Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking @@ -173,20 +186,17 @@ steps: - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model" - label: "CosyVoice3-TTS E2E Test" - agent_pool: mi325_1 + agent_pool: mi300_1 depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking commands: - - | - timeout 20m bash -c ' - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model" - ' + - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model" # TODO: Bagel test on ROCm is very unstable. @tjtanaa # Need to debug before reneable numerical changes across large PRs # - label: "Bagel Text2Img Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -198,7 +208,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm" # - label: "Bagel Img2Img Model Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking @@ -210,7 +220,7 @@ steps: # - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm" # - label: "Bagel Online Serving Test" -# agent_pool: mi325_1 +# agent_pool: mi300_1 # depends_on: amd-build # mirror_hardwares: [amdproduction] # grade: Blocking diff --git a/.buildkite/amd/test-template-amd-omni.j2 b/.buildkite/amd/test-template-amd-omni.j2 index 2f1a9aed3..e2042b96a 100644 --- a/.buildkite/amd/test-template-amd-omni.j2 +++ b/.buildkite/amd/test-template-amd-omni.j2 @@ -1,11 +1,65 @@ {# vllm-omni customized version Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/test-template-amd.j2 - Last synced: 2025-12-15 - Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests + Last synced: 2026-08-04 + Modifications: Use the shared MI300 queues and Omni command-passing contract #} {% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %} {% set default_working_dir = "/app/vllm-omni" %} +{# MI300 runs tests natively in Kubernetes; Docker-in-Docker is unavailable. #} +{% macro amd_native_k8s_plugin(indent, step) -%} +{{ indent }}plugins: +{{ indent }} - kubernetes: +{{ indent }} podSpecPatch: +{{ indent }} automountServiceAccountToken: false +{{ indent }} securityContext: +{{ indent }} seccompProfile: +{{ indent }} type: RuntimeDefault +{{ indent }} imagePullSecrets: +{{ indent }} - name: docker-config +{{ indent }} containers: +{{ indent }} - name: container-0 +{{ indent }} image: {{ docker_image_amd }} +{{ indent }} imagePullPolicy: Always +{{ indent }} securityContext: +{{ indent }} allowPrivilegeEscalation: false +{{ indent }} capabilities: +{{ indent }} add: +{{ indent }} - IPC_LOCK +{{ indent }} resources: +{{ indent }} limits: +{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}" +{{ indent }} requests: +{{ indent }} amd.com/gpu: "{{ step.agent_pool[-1:] }}" +{{ indent }} volumeMounts: +{{ indent }} - name: devshm +{{ indent }} mountPath: /dev/shm +{{ indent }} env: +{{ indent }} - name: AMD_CI_RUNTIME +{{ indent }} value: native +{{ indent }} - name: NATIVE_CI +{{ indent }} value: "true" +{{ indent }} - name: VLLM_CI_DOCKER_DISABLED +{{ indent }} value: "1" +{{ indent }} - name: VLLM_CI_EXPECTED_GPU_COUNT +{{ indent }} value: "{{ step.agent_pool[-1:] }}" +{{ indent }} - name: HF_HOME +{{ indent }} value: /home/buildkite-agent/huggingface +{{ indent }} - name: MIOPEN_DEBUG_CONV_DIRECT +{{ indent }} value: "0" +{{ indent }} - name: MIOPEN_DEBUG_CONV_GEMM +{{ indent }} value: "0" +{{ indent }} - name: VLLM_ROCM_USE_AITER +{{ indent }} value: "0" +{{ indent }} - name: PYTORCH_ROCM_ARCH +{{ indent }} value: "" +{{ indent }} volumes: +{{ indent }} - name: devshm +{{ indent }} emptyDir: +{{ indent }} medium: Memory +{{ indent }} sizeLimit: 16Gi +{%- endmacro %} + {% if env %} env: {% for name in env %} @@ -53,19 +107,31 @@ steps: {% if step.agent_pool %} queue: amd_{{ step.agent_pool }} {% else %} - queue: amd_mi325_1 + queue: amd_mi300_1 {% endif %} +{{ amd_native_k8s_plugin(" ", step) }} {% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %} {% set indented_cmd = cmd_body | replace("\n", "\n ") %} command: bash .buildkite/amd/scripts/run-amd-test.sh env: - DOCKER_BUILDKIT: "1" + AMD_CI_RUNTIME: "native" + NATIVE_CI: "true" + VLLM_CI_DOCKER_DISABLED: "1" + VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}" + VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1" + HF_HOME: "/home/buildkite-agent/huggingface" + HF_HUB_DOWNLOAD_TIMEOUT: "300" + HF_HUB_ETAG_TIMEOUT: "60" + MIOPEN_DEBUG_CONV_DIRECT: "0" + MIOPEN_DEBUG_CONV_GEMM: "0" + VLLM_ROCM_USE_AITER: "0" + PYTORCH_ROCM_ARCH: "" {% if step.env %} {% for name in step.env %} {{ name }}: "{{ step.env[name] }}" {% endfor %} {% endif %} - TEST_COMMAND: |- + TEST_COMMANDS: |- (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} {% if env %} {% for name in env %} @@ -82,6 +148,15 @@ steps: {% endif %} {{ indented_cmd | safe }} priority: 100 +{% if step.parallelism %} + parallelism: {{ step.parallelism }} +{% endif %} +{% if step.artifact_paths %} + artifact_paths: +{% for path in step.artifact_paths %} + - "{{ path }}" +{% endfor %} +{% endif %} {% if step.timeout_in_minutes %} timeout_in_minutes: {{ step.timeout_in_minutes }} {% endif %} @@ -100,19 +175,31 @@ steps: {% if step.agent_pool %} queue: amd_{{ step.agent_pool }} {% else %} - queue: amd_mi325_1 + queue: amd_mi300_1 {% endif %} +{{ amd_native_k8s_plugin(" ", step) }} {% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %} {% set indented_cmd = cmd_body | replace("\n", "\n ") %} command: bash .buildkite/amd/scripts/run-amd-test.sh env: - DOCKER_BUILDKIT: "1" + AMD_CI_RUNTIME: "native" + NATIVE_CI: "true" + VLLM_CI_DOCKER_DISABLED: "1" + VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}" + VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1" + HF_HOME: "/home/buildkite-agent/huggingface" + HF_HUB_DOWNLOAD_TIMEOUT: "300" + HF_HUB_ETAG_TIMEOUT: "60" + MIOPEN_DEBUG_CONV_DIRECT: "0" + MIOPEN_DEBUG_CONV_GEMM: "0" + VLLM_ROCM_USE_AITER: "0" + PYTORCH_ROCM_ARCH: "" {% if step.env %} {% for name in step.env %} {{ name }}: "{{ step.env[name] }}" {% endfor %} {% endif %} - TEST_COMMAND: |- + TEST_COMMANDS: |- (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} {% if env %} {% for name in env %} @@ -129,6 +216,15 @@ steps: {% endif %} {{ indented_cmd | safe }} priority: 100 +{% if step.parallelism %} + parallelism: {{ step.parallelism }} +{% endif %} +{% if step.artifact_paths %} + artifact_paths: +{% for path in step.artifact_paths %} + - "{{ path }}" +{% endfor %} +{% endif %} {% if step.timeout_in_minutes %} timeout_in_minutes: {{ step.timeout_in_minutes }} {% endif %} diff --git a/.buildkite/common/ci_mirror_hardwares.yml b/.buildkite/common/ci_mirror_hardwares.yml index 00c6e4f82..e68cbd382 100644 --- a/.buildkite/common/ci_mirror_hardwares.yml +++ b/.buildkite/common/ci_mirror_hardwares.yml @@ -9,6 +9,11 @@ # a2b3_npu_1 — ascend-a2b3, resource_class npu-1 (A2/B3 CI image) # a2b3_npu_4 — ascend-a2b3, resource_class npu-4 # a3_npu_2 — ascend-a3, resource_class npu-2 +# +# The ``l4_*`` presets run under the ``docker`` plugin, which does not put +# ``buildkite-agent`` on PATH inside the container the way the ``kubernetes`` +# presets do, so they set ``mount-buildkite-agent: true``; without it any step +# calling ``buildkite-agent`` (artifact upload, annotate, meta-data) fails there. x-npu-k8s-plugins: &npu_k8s_plugins - kubernetes: @@ -26,6 +31,7 @@ mirror_hardwares: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT always-pull: true propagate-environment: true + mount-buildkite-agent: true shm-size: "8gb" environment: - HF_HOME=/fsx/hf_cache @@ -41,6 +47,7 @@ mirror_hardwares: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT always-pull: true propagate-environment: true + mount-buildkite-agent: true shm-size: "8gb" environment: - HF_HOME=/fsx/hf_cache diff --git a/.buildkite/common/scripts/run_cov_split.sh b/.buildkite/common/scripts/run_cov_split.sh new file mode 100755 index 000000000..0ed61ec17 --- /dev/null +++ b/.buildkite/common/scripts/run_cov_split.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# Run a model's E2E tests once per entry mode as separate coverage runs and upload +# the reports, so per-mode attribution survives (see vllm-project/vllm-omni#5332). +# +# Usage: run_cov_split.sh --model-id [--offline ] [--online ] \ +# -- +# +# --model-id Names the artifacts: coverage--{offline,online}-.xml.gz. +# Use the model's directory name under vllm_omni/*/models/. +# --offline Pytest path(s) for the offline half. Repeatable. +# --online Pytest path(s) for the online half. Repeatable. +# -- Everything after this is forwarded verbatim to every pytest run, +# quoting preserved: the job's own -m expression, --run-level, and +# anything else it needs (e.g. --test-config-file). +# +# At least one of --offline/--online is required, so a job with only one mode runs +# just that half. Total runtime is bounded by the step's timeout_in_minutes. +# +# Coverage + artifact upload run only on scheduled weekly E2E sweeps +# (WEEKLY=1 and BUILDKITE_BRANCH=main — same gate as bootstrap --e2e). +# All other invocations (PR merge-test, regular main merge, local) run the +# same combined pytest as the job used before coverage was added. +# +# Exits non-zero if any half or the upload failed. +set -uo pipefail + +MODEL_ID="" +OFFLINE=() +ONLINE=() +PYTEST_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --model-id) MODEL_ID="$2"; shift 2 ;; + --offline) OFFLINE+=("$2"); shift 2 ;; + --online) ONLINE+=("$2"); shift 2 ;; + --) shift; PYTEST_ARGS=("$@"); break ;; + *) echo "run_cov_split.sh: unknown argument: $1" >&2; exit 2 ;; + esac +done + +if [[ -z "${MODEL_ID}" ]]; then + echo "run_cov_split.sh: missing --model-id" >&2 + exit 2 +fi + +if ((${#OFFLINE[@]} == 0 && ${#ONLINE[@]} == 0)); then + echo "run_cov_split.sh: need at least one of --offline / --online" >&2 + exit 2 +fi + +_should_run_coverage() { + [[ "${WEEKLY:-}" == "1" && "${BUILDKITE_BRANCH:-}" == "main" ]] +} + +if ! _should_run_coverage; then + PATHS=("${OFFLINE[@]}" "${ONLINE[@]}") + echo "--- pytest (no coverage): ${MODEL_ID}" + pytest -s -v "${PATHS[@]}" "${PYTEST_ARGS[@]}" + exit $? +fi + +# Two steps can test the same model, and artifacts are scoped to the build (a +# nightly one carries the ready, merge and nightly tiers at once), so unsuffixed +# reports would share a path -- a glob download then races to rename onto one file. +# Not BUILDKITE_JOB_ID: the step id survives retries, so attempts collapse under +# Buildkite's latest-attempt lookup instead of accumulating. A parallel step would +# need BUILDKITE_PARALLEL_JOB as well; no pilot is one. +REPORT_SUFFIX="${BUILDKITE_STEP_ID:-local}" +OFFLINE_REPORT="coverage-${MODEL_ID}-offline-${REPORT_SUFFIX}.xml" +ONLINE_REPORT="coverage-${MODEL_ID}-online-${REPORT_SUFFIX}.xml" + +run_half() { + local mode="$1" + shift + local report="coverage-${MODEL_ID}-${mode}-${REPORT_SUFFIX}.xml" + echo "--- coverage: ${MODEL_ID} ${mode}" + pytest -s -v "$@" "${PYTEST_ARGS[@]}" --cov=vllm_omni --cov-report="xml:${report}" + local status=$? + # Cobertura lists every statement in vllm_omni whether or not it ran, so the + # report is ~7MB no matter how little the job covers. The repetition is what + # dominates it, and gzip takes that to ~400KB. + gzip -9 -f "${report}" || status=1 + return "${status}" +} + +# Clear both modes, not just the ones being run: the upload glob matches both, so +# a report left by an earlier run would otherwise be published as if the mode it +# belongs to had run this time. +rm -f "${OFFLINE_REPORT}" "${ONLINE_REPORT}" "${OFFLINE_REPORT}.gz" "${ONLINE_REPORT}.gz" + +EXIT=0 +((${#OFFLINE[@]})) && { run_half offline "${OFFLINE[@]}" || EXIT=1; } +((${#ONLINE[@]})) && { run_half online "${ONLINE[@]}" || EXIT=1; } + +buildkite-agent artifact upload "coverage-${MODEL_ID}-*-${REPORT_SUFFIX}.xml.gz" || EXIT=1 + +exit "${EXIT}" diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index b27f78c3a..a9c0b0561 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Render and optionally upload Buildkite pipeline YAML with diff-aware logic. Bootstrap mode (``bootstrap-upload-steps.yml``): @@ -58,11 +59,12 @@ BOOTSTRAP_UPLOAD_IF_KEYS = { E2E_GROUP_MARKER = "E2E Test" CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml" -CUDA_NIGHTLY_ONLY = ( - '(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")' -) -NPU_NIGHTLY_ONLY = ( - '(build.branch == "main" && build.env("NIGHTLY") == "1") || ' +# Bootstrap Buildkite ``if`` expressions. +# ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN). +# ``*_UPLOAD_IF``: full gate for uploading that child pipeline. +NIGHTLY_MAIN_IF = 'build.branch == "main" && build.env("NIGHTLY") == "1"' +NIGHTLY_LABEL_IF = ( + f"({NIGHTLY_MAIN_IF}) || " '(build.branch != "main" && (' 'build.pull_request.labels includes "nightly-test" || ' 'build.pull_request.labels includes "omni-test" || ' @@ -71,6 +73,19 @@ NPU_NIGHTLY_ONLY = ( 'build.pull_request.labels includes "diffusion-x2v-test"' "))" ) +WEEKLY_E2E_IF = 'build.branch == "main" && build.env("WEEKLY") == "1"' +WEEKLY_MAIN_IF = 'build.branch == "main" && (build.env("WEEKLY") == "1" || build.env("NON_CRITICAL") == "1")' +WEEKLY_LABEL_IF = f'({WEEKLY_MAIN_IF}) || (build.branch != "main" && build.pull_request.labels includes "weekly-test")' +READY_LABEL_IF = 'build.branch != "main" && build.pull_request.labels includes "ready"' +MERGE_LABEL_IF = 'build.branch != "main" && build.pull_request.labels includes "merge-test"' +MERGE_MAIN_IF = ( + 'build.branch == "main" && build.env("NIGHTLY") != "1" && ' + 'build.env("WEEKLY") != "1" && build.env("NON_CRITICAL") != "1"' +) +READY_UPLOAD_IF = f"({WEEKLY_E2E_IF}) || ({READY_LABEL_IF})" +MERGE_UPLOAD_IF = f"({WEEKLY_E2E_IF}) || (({MERGE_MAIN_IF}) || ({MERGE_LABEL_IF}))" +BOOTSTRAP_DISABLED_IF = "false" +BOOTSTRAP_ENABLED_IF = "true" # --- Logging --- @@ -102,74 +117,47 @@ def _format_bootstrap_if(expr: str) -> str: def _compute_bootstrap_if_exprs(*, decision, platform: str) -> dict[str, str]: - disabled = "false" - nightly_main = 'build.branch == "main" && build.env("NIGHTLY") == "1"' - nightly_only = NPU_NIGHTLY_ONLY if platform == "npu" else CUDA_NIGHTLY_ONLY - if platform == "npu": - ready_pr = ( - 'build.branch != "main" && (' - 'build.pull_request.labels includes "npu-test" || ' - 'build.pull_request.labels includes "ready"' - ")" - ) - nightly_label_if = nightly_only - weekly_label_if = disabled - merge_base = disabled + ready_upload = READY_LABEL_IF + merge_upload = BOOTSTRAP_DISABLED_IF + weekly_label_if = BOOTSTRAP_DISABLED_IF else: - ready_pr = 'build.branch != "main" && build.pull_request.labels includes "ready"' - merge_main = 'build.branch == "main" && build.env("NIGHTLY") != "1" && build.env("WEEKLY") != "1"' - merge_pr = 'build.branch != "main" && build.pull_request.labels includes "merge-test"' - nightly_label_if = ( - '(build.branch == "main" && build.env("NIGHTLY") == "1") || ' - '(build.branch != "main" && (' - 'build.pull_request.labels includes "nightly-test" || ' - 'build.pull_request.labels includes "omni-test" || ' - 'build.pull_request.labels includes "tts-test" || ' - 'build.pull_request.labels includes "diffusion-x2iat-test" || ' - 'build.pull_request.labels includes "diffusion-x2v-test"' - "))" - ) - weekly_label_if = ( - '(build.branch == "main" && build.env("WEEKLY") == "1") || ' - '(build.branch != "main" && build.pull_request.labels includes "weekly-test")' - ) - merge_base = f"({nightly_main}) || (({merge_main}) || ({merge_pr}))" - - ready_base = f"({nightly_main}) || ({ready_pr})" - weekly_main = 'build.branch == "main" && build.env("WEEKLY") == "1"' + ready_upload = READY_UPLOAD_IF + merge_upload = MERGE_UPLOAD_IF + weekly_label_if = WEEKLY_LABEL_IF if decision.skip_all: # Docs / skip-mark only: no PR-label escape hatch. Main scheduled - # NIGHTLY=1 still runs L4 + L2/L3 (--e2e); WEEKLY=1 still runs L5. - image_expr = f"({nightly_main}) || ({weekly_main})" if platform == "cuda" else nightly_main - ready_expr = nightly_main - merge_expr = nightly_main if platform == "cuda" else disabled - nightly_expr = nightly_main - weekly_expr = weekly_main if platform == "cuda" else disabled + # NIGHTLY=1 still runs L4; WEEKLY=1 / NON_CRITICAL=1 still run L5. + # main+WEEKLY=1 also uploads L2/L3 (those steps then pass --e2e). + image_expr = f"({NIGHTLY_MAIN_IF}) || ({WEEKLY_MAIN_IF})" if platform == "cuda" else NIGHTLY_MAIN_IF + ready_expr = WEEKLY_E2E_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF + merge_expr = WEEKLY_E2E_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF + nightly_expr = NIGHTLY_MAIN_IF + weekly_expr = WEEKLY_MAIN_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF elif decision.skip_l2_l3: l2_enabled = decision.is_run("npu", "l2") if platform == "npu" else decision.is_run("cuda", "l2") l3_enabled = platform == "cuda" and decision.is_run("cuda", "l3") - ready_expr = ready_base if l2_enabled else disabled - merge_expr = merge_base if l3_enabled else disabled - nightly_expr = nightly_label_if - weekly_expr = weekly_label_if if platform == "cuda" else disabled + ready_expr = ready_upload if l2_enabled else BOOTSTRAP_DISABLED_IF + merge_expr = merge_upload if l3_enabled else BOOTSTRAP_DISABLED_IF + nightly_expr = NIGHTLY_LABEL_IF + weekly_expr = weekly_label_if if platform == "cuda" else BOOTSTRAP_DISABLED_IF - image_parts = [f"({nightly_label_if})"] + image_parts = [f"({NIGHTLY_LABEL_IF})"] if platform == "cuda": image_parts.append(f"({weekly_label_if})") if l2_enabled: - image_parts.insert(0, f"({ready_base})") + image_parts.insert(0, f"({ready_upload})") if l3_enabled: - image_parts.insert(1 if l2_enabled else 0, f"({merge_base})") + image_parts.insert(1 if l2_enabled else 0, f"({merge_upload})") image_expr = " || ".join(image_parts) else: - image_expr = "true" - ready_expr = ready_base - merge_expr = merge_base if platform == "cuda" else disabled - nightly_expr = nightly_label_if - weekly_expr = weekly_label_if if platform == "cuda" else disabled + image_expr = BOOTSTRAP_ENABLED_IF + ready_expr = ready_upload + merge_expr = merge_upload if platform == "cuda" else BOOTSTRAP_DISABLED_IF + nightly_expr = NIGHTLY_LABEL_IF + weekly_expr = weekly_label_if if platform == "cuda" else BOOTSTRAP_DISABLED_IF return { "image": _format_bootstrap_if(image_expr), diff --git a/.buildkite/cuda/bootstrap-upload-steps.yml b/.buildkite/cuda/bootstrap-upload-steps.yml index f2467bf42..ab4254feb 100644 --- a/.buildkite/cuda/bootstrap-upload-steps.yml +++ b/.buildkite/cuda/bootstrap-upload-steps.yml @@ -8,28 +8,28 @@ steps: agents: queue: "cpu_queue_premerge" - # L2 Test — PR ready label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group) + # L2 Test — PR ready label (diff-aware) or main+WEEKLY=1 (--e2e, full E2E group) - label: "Upload Ready Pipeline" depends_on: image-build key: upload-ready-pipeline commands: - | EXTRA="" - if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then EXTRA="--e2e" fi python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-ready.yml agents: queue: "cpu_queue_premerge" - # L3 Test — PR merge-test label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group) + # L3 Test — PR merge-test label (diff-aware) or main+WEEKLY=1 (--e2e, full E2E group) - label: "Upload Merge Pipeline" depends_on: image-build key: upload-merge-pipeline commands: - | EXTRA="" - if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then EXTRA="--e2e" fi python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-merge.yml @@ -45,7 +45,7 @@ steps: agents: queue: "cpu_queue_premerge" - # L5 Test — main+WEEKLY=1 (scheduled), or PR with weekly-test label + # L5 Test — main+WEEKLY=1 / NON_CRITICAL=1 (scheduled), or PR with weekly-test label - label: "Upload Weekly Pipeline" depends_on: image-build key: upload-weekly-pipeline diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 363e8fbdc..53c70154a 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -9,23 +9,27 @@ steps: depends_on: upload-merge-pipeline steps: - label: "Simple · Model Executor Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/model_executor -m 'core_model and cpu'" mirror_hardwares: l4_1 - label: "Simple · Diffusion Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/diffusion -m 'core_model and cpu'" mirror_hardwares: l4_1 - label: "Simple · Engine&Entrypoints Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'" mirror_hardwares: l4_1 - label: "Simple · Other Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine" mirror_hardwares: l4_1 - group: ":card_index_dividers: Diffusion Test" @@ -76,8 +80,14 @@ steps: - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/common/alias_free_activation.py - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - vllm_omni/deploy/qwen3_tts.yaml + - vllm_omni/entrypoints/openai/serving_speech.py + - vllm_omni/entrypoints/openai/tts_adapters/ + - vllm_omni/engine/stage_pool.py commands: - | timeout 40m bash -c " @@ -92,49 +102,58 @@ steps: - tests/e2e/online_serving/test_qwen3_tts_base.py - tests/e2e/offline_inference/test_qwen3_tts_base.py - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/common/alias_free_activation.py - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - vllm_omni/deploy/qwen3_tts.yaml + - vllm_omni/entrypoints/openai/serving_speech.py + - vllm_omni/entrypoints/openai/tts_adapters/ + - vllm_omni/engine/stage_pool.py + - .buildkite/common/scripts/run_cov_split.sh + - pyproject.toml + # 40 -> 60: the coverage split runs the model twice. + timeout_in_minutes: 60 commands: - | - timeout 40m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model' - " - mirror_hardwares: l4_1 - - - label: "TTS · Higgs-Audio-V3 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_higgs_audio_v3.py - - vllm_omni/model_executor/models/higgs_audio_v3/ - - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py - - vllm_omni/deploy/higgs_multimodal_qwen3.yaml - commands: - - | - timeout 20m bash -c " - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_USE_DEEP_GEMM=0 - export VLLM_MOE_USE_DEEP_GEMM=0 - pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'advanced_model' --run-level 'advanced_model' - " + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 + .buildkite/common/scripts/run_cov_split.sh \ + --model-id qwen3_tts \ + --offline tests/e2e/offline_inference/test_qwen3_tts_base.py \ + --online tests/e2e/online_serving/test_qwen3_tts_base.py \ + -- -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: l4_1 - label: "Omni · Qwen3-Omni Test" source_file_dependencies: - tests/e2e/offline_inference/test_qwen3_omni.py + - tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py - tests/e2e/online_serving/test_qwen3_omni.py + - tests/helpers/runtime.py + - vllm_omni/entrypoints/openai/serving_chat.py - vllm_omni/model_executor/models/qwen3_omni/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py + - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - vllm_omni/deploy/qwen3_omni_moe.yaml + - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml timeout_in_minutes: 50 commands: - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' + - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Test" source_file_dependencies: - tests/e2e/offline_inference/test_minicpmo_4_5.py - tests/e2e/online_serving/test_minicpmo_4_5.py + - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/models/step_audio2/ + - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ + - vllm_omni/model_executor/models/cosyvoice3/utils.py + - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - vllm_omni/deploy/minicpmo_4_5.yaml timeout_in_minutes: 50 commands: @@ -155,63 +174,36 @@ steps: - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py + - vllm_omni/deploy/minicpmo_4_5.yaml - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - - vllm_omni/deploy/minicpmo_4_5_duplex.yaml - vllm_omni/experimental/fullduplex/ - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/models/step_audio2/ + - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ + - vllm_omni/model_executor/models/cosyvoice3/utils.py - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - vllm_omni/core/sched/ - vllm_omni/outputs/multimodal_accumulation.py - vllm_omni/engine/stage_init_utils.py - vllm_omni/engine/async_omni_engine.py + - vllm_omni/engine/orchestrator.py - vllm_omni/entrypoints/cli/serve.py + - vllm_omni/entrypoints/openai/api_server.py + - vllm_omni/entrypoints/openai/duplex_capability.py + - vllm_omni/config/stage_config.py timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_1 - - label: "TTS · VoxCPM2 Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_voxcpm2_tts.py - - tests/e2e/online_serving/test_voxcpm2_tts.py - - vllm_omni/model_executor/models/voxcpm2/ - - vllm_omni/deploy/voxcpm2.yaml - timeout_in_minutes: 30 - commands: - - | - timeout 30m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model' - " - mirror_hardwares: l4_1 - - - label: "Diffusion · Qwen Image Edit Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image_edit.py - - vllm_omni/diffusion/models/qwen_image/ - timeout_in_minutes: 20 - commands: - - pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model' - mirror_hardwares: h100_1 - - label: "Diffusion · Bagel Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_bagel.py - - tests/e2e/online_serving/test_bagel.py - - vllm_omni/model_executor/models/bagel/ - - vllm_omni/diffusion/models/bagel/ - - vllm_omni/model_executor/stage_input_processors/bagel.py - - vllm_omni/deploy/bagel.yaml - timeout_in_minutes: 20 - commands: - - pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model' - mirror_hardwares: h100_1 - - label: "Diffusion · Qwen Image Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen_image.py - vllm_omni/diffusion/models/qwen_image/ + - vllm_omni/diffusion/models/dmd2/ + - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py commands: - | timeout 40m bash -c ' @@ -220,23 +212,12 @@ steps: ' mirror_hardwares: h100_1 - - label: "Diffusion · Qwen Image Layered Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image_layered.py - - vllm_omni/diffusion/models/qwen_image/ - commands: - - | - timeout 40m bash -c " - export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_model' - " - mirror_hardwares: h100_1 - - label: "Diffusion · Wan22 Test" source_file_dependencies: - tests/e2e/offline_inference/test_wan22_t2v.py - tests/e2e/online_serving/test_wan22_t2v.py - vllm_omni/diffusion/models/wan2_2/ + - vllm_omni/diffusion/models/dmd2/ commands: - | timeout 40m bash -c " @@ -245,20 +226,14 @@ steps: " mirror_hardwares: h100_1 - - label: "Diffusion · Z Image Test" + - label: "Diffusion · HunyuanVideo-1.5 Test" source_file_dependencies: - - tests/e2e/offline_inference/test_z_image.py - - tests/e2e/online_serving/test_z_image.py - - vllm_omni/diffusion/models/z_image/ - timeout_in_minutes: 30 + - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py + - vllm_omni/diffusion/models/hunyuan_video/ commands: - - pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' - mirror_hardwares: l4_1 - - label: "Diffusion · FLUX.2 Klein Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_flux2_klein.py - - vllm_omni/diffusion/models/flux2_klein/ - timeout_in_minutes: 30 - commands: - - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' - mirror_hardwares: l4_1 + - | + timeout 40m bash -c " + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' + " + mirror_hardwares: h100_1 diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index d3246a98f..9f44ba2b0 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -12,16 +12,22 @@ steps: if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test" steps: - - label: ":full_moon: Omni · Function Test with H100" - timeout_in_minutes: 120 + - label: ":full_moon: Omni · Function Test with H100 · Single-GPU" + timeout_in_minutes: 90 commands: - - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py + - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py + mirror_hardwares: h100_1 + + - label: ":full_moon: Omni · Function Test with H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py mirror_hardwares: h100_2 - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test" timeout_in_minutes: 50 commands: - - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni" --run-level "full_model" + - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Omni · Doc Test with L4" @@ -31,10 +37,10 @@ steps: - pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model" mirror_hardwares: l4_4 - - label: ":full_moon: Omni · Doc Test with H100" + - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU" timeout_in_minutes: 90 commands: - - pytest -s -v tests/examples/ -m "full_model and omni and H100" --run-level "full_model" + - pytest -s -v tests/examples/ -m "full_model and omni and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Omni · Accuracy Test" @@ -83,38 +89,12 @@ steps: - export BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and full_model" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT mirror_hardwares: h100_2 - - label: ":full_moon: Omni · Perf Test · vLLM Text" - key: nightly-omni-performance-vllm-text - timeout_in_minutes: 180 - commands: - - export BENCHMARK_DIR=tests/dfx/perf/results - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" - exit $$EXIT - mirror_hardwares: h100_2 - - - label: ":full_moon: Omni · Perf Test · Multi-Replica" - key: nightly-omni-performance-multi-replicas - timeout_in_minutes: 180 - commands: - - export BENCHMARK_DIR=tests/dfx/perf/results - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" - exit $$EXIT - mirror_hardwares: h100_3 - - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" key: nightly-omni-performance-minicpmo-4-5 timeout_in_minutes: 180 @@ -159,64 +139,37 @@ steps: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - export VLLM_USE_DEEP_GEMM="0" - export VLLM_MOE_USE_DEEP_GEMM="0" - - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py + - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy mirror_hardwares: l4_1 - - label: ":full_moon: TTS · Accuracy Test · VoxCPM2" - timeout_in_minutes: 30 - commands: - - | - timeout 30m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 - pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model' - " - mirror_hardwares: l4_1 - - - label: ":full_moon: TTS · Perf Test" - key: nightly-tts-performance + - label: ":full_moon: TTS · Perf Test · Single-GPU" + key: nightly-tts-performance-single timeout_in_minutes: 180 commands: - export BENCHMARK_DIR=tests/dfx/perf/results - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_1" + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + mirror_hardwares: h100_1 + + - label: ":full_moon: TTS · Perf Test · 2-GPU" + key: nightly-tts-performance-2gpu + timeout_in_minutes: 180 + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json -m "H100 and cards_2" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT mirror_hardwares: h100_2 - - label: ":full_moon: TTS · VoxCPM2 · Perf Test" - key: nightly-tts-performance-voxcpm2 - timeout_in_minutes: 180 - commands: - - export BENCHMARK_DIR=tests/dfx/perf/results - - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" - exit $$EXIT - mirror_hardwares: h100_1 - - - label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test" - key: nightly-tts-performance-higgs-audio-v3 - timeout_in_minutes: 180 - commands: - - export BENCHMARK_DIR=tests/dfx/perf/results - - export VLLM_USE_DEEP_GEMM="0" - - export VLLM_MOE_USE_DEEP_GEMM="0" - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" - exit $$EXIT - mirror_hardwares: h100_1 - - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]" key: nightly-tiny-model-test-group depends_on: upload-nightly-pipeline @@ -246,8 +199,6 @@ steps: build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test" steps: - # Use explicit file shards so single-card cases do not reserve 2+ GPUs, - # and Bagel feature coverage is not duplicated via a broad tests/e2e sweep. - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU" timeout_in_minutes: 120 @@ -255,66 +206,24 @@ steps: - >- pytest -sv tests/e2e/online_serving/test_qwen_image_expansion.py - tests/e2e/online_serving/test_qwen_image_edit_expansion.py - tests/e2e/online_serving/test_qwen_image_layered_expansion.py - tests/e2e/online_serving/test_boogu_image.py - tests/e2e/online_serving/test_boogu_image_edit.py - tests/e2e/online_serving/test_boogu_image_expansion.py - tests/e2e/online_serving/test_boogu_image_edit_expansion.py - tests/e2e/online_serving/test_bagel_expansion.py - tests/e2e/online_serving/test_krea2_expansion.py - tests/e2e/offline_inference/test_krea2_expansion.py - -m "full_model and diffusion and H100 and not distributed_cuda" + -m "full_model and diffusion and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Qwen-Image" + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU" timeout_in_minutes: 120 commands: - >- pytest -sv tests/e2e/online_serving/test_qwen_image_expansion.py - tests/e2e/online_serving/test_qwen_image_edit_expansion.py - tests/e2e/offline_inference/test_krea2_expansion.py - -m "full_model and diffusion and H100 and distributed_cuda" + -m "full_model and diffusion and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX" - timeout_in_minutes: 90 - commands: - - >- - pytest -sv - tests/e2e/online_serving/test_qwen_image_layered_expansion.py - -m "full_model and diffusion and H100 and distributed_cuda" - --run-level "full_model" - mirror_hardwares: h100_2 - - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel" + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · Multi-GPU" timeout_in_minutes: 120 commands: - - >- - pytest -sv - tests/e2e/online_serving/test_bagel_expansion.py - -m "full_model and diffusion and H100 and distributed_cuda" - --run-level "full_model" - mirror_hardwares: h100_2 - - - label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100" - timeout_in_minutes: 120 - commands: - - >- - pytest -sv - tests/e2e/online_serving/test_glm_image_expansion.py - tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py - -m "full_model and diffusion and H100" - --run-level "full_model" - mirror_hardwares: h100_2 - - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4" - timeout_in_minutes: 120 - commands: - - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy + - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and not cards_1" --run-level "full_model" --ignore=tests/e2e/accuracy mirror_hardwares: l4_4 - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test" @@ -327,7 +236,7 @@ steps: timeout_in_minutes: 180 commands: - export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120 - - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model + - pytest -s -v tests/e2e/accuracy/test_qwen_image.py --run-level full_model - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model mirror_hardwares: h100_1 @@ -339,8 +248,14 @@ steps: - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model mirror_hardwares: h100_4 - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image" - key: nightly-diffusion-x2iat-performance-qwen-image + - label: ":full_moon: HiDream-O1-Image · Accuracy Test" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m full_model --run-level full_model + mirror_hardwares: h100_2 + + - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU" + key: nightly-diffusion-x2iat-performance-qwen-image-single timeout_in_minutes: 180 commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results @@ -348,63 +263,15 @@ steps: - export CACHE_DIT_VERSION=1.5.0 - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" - exit $$EXIT - mirror_hardwares: h100_4 - - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit" - key: nightly-diffusion-x2iat-performance-qwen-image-edit - timeout_in_minutes: 180 - commands: - - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.5.0 - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" - exit $$EXIT - mirror_hardwares: h100_4 - - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511" - key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511 - timeout_in_minutes: 180 - commands: - - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.5.0 - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" - exit $$EXIT - mirror_hardwares: h100_4 - - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered" - key: nightly-diffusion-x2iat-performance-qwen-image-layered - timeout_in_minutes: 120 - commands: - - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.5.0 - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT mirror_hardwares: h100_1 - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL" - key: nightly-diffusion-x2iat-performance-bagel + - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-GPU" + key: nightly-diffusion-x2iat-performance-qwen-image-4gpu timeout_in_minutes: 180 commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results @@ -412,12 +279,12 @@ steps: - export CACHE_DIT_VERSION=1.5.0 - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json -m "H100 and cards_4" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - mirror_hardwares: h100_1 + mirror_hardwares: h100_4 - group: ":card_index_dividers: Diffusion X2V Model Test" key: nightly-diffusion-x2v-group @@ -428,24 +295,36 @@ steps: build.pull_request.labels includes "diffusion-x2v-test" steps: - - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test" + - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · Single-GPU" timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda" --run-level "full_model" - mirror_hardwares: h100_2 - - - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test" - timeout_in_minutes: 90 - commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model" - mirror_hardwares: h100_2 - - - label: ":full_moon: Diffusion X2V · LingBot Video MoE Function Test" - timeout_in_minutes: 90 - commands: - - pytest -s -v tests/e2e/online_serving/test_lingbot_video_moe.py -k "video_generation" -m "full_model and cuda" --run-level "full_model" + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · 2-GPU" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" + mirror_hardwares: h100_2 + + - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · Single-GPU" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" + mirror_hardwares: h100_1 + + - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · 2-GPU" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" + mirror_hardwares: h100_2 + + - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Function Test" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model" + mirror_hardwares: h100_2 + - label: ":full_moon: Diffusion X2V · Doc Test" timeout_in_minutes: 60 commands: @@ -456,7 +335,6 @@ steps: timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" - - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" @@ -477,21 +355,51 @@ steps: commands: - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model - - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model - - pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2V · Perf Test" - key: nightly-diffusion-x2v-performance + - label: ":full_moon: Diffusion X2V · Perf Test · Single-GPU" + key: nightly-diffusion-x2v-performance-single timeout_in_minutes: 180 commands: - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - | set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_1" + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" + buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" + exit $$EXIT + mirror_hardwares: h100_1 + + - label: ":full_moon: Diffusion X2V · Perf Test · 2-GPU" + key: nightly-diffusion-x2v-performance-2gpu + timeout_in_minutes: 180 + commands: + - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results + - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_2" EXIT1=$$? - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json -m "H100 and cards_2" + EXIT2=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" + buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" + if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi + exit $$EXIT2 + mirror_hardwares: h100_2 + + - label: ":full_moon: Diffusion X2V · Perf Test · HunyuanVideo-1.5" + key: nightly-diffusion-x2v-performance-hunyuanvideo15 + timeout_in_minutes: 180 + commands: + - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json + EXIT1=$$? + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json EXIT2=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" @@ -525,10 +433,16 @@ steps: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model" mirror_hardwares: l4_4 - - label: ":full_moon: Diffusion · Distributed Test with H100" + - label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU" timeout_in_minutes: 60 commands: - - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model" + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_2' --run-level "full_model" + mirror_hardwares: h100_2 + + - label: ":full_moon: Diffusion · Distributed Test with H100 · 4-GPU" + timeout_in_minutes: 60 + commands: + - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_4' --run-level "full_model" mirror_hardwares: h100_4 - label: ":full_moon: Diffusion · Offloader Test" @@ -552,20 +466,10 @@ steps: - label: ":full_moon: Diffusion · Quantization Test with L4" timeout_in_minutes: 60 commands: - - pip install "vllm-gguf-plugin>=0.0.3" + - pip install "vllm-gguf-plugin==0.0.4" - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model" mirror_hardwares: l4_1 - - label: ":bar_chart: Testcase Statistics" - key: nightly-testcase-statistics - timeout_in_minutes: 120 - depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" - commands: - - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html - - buildkite-agent artifact upload "tests/dfx/perf/results/*.html" - mirror_hardwares: h100_1 - # NOTE: ":email: Nightly Collection & Email" is deprecated and will be removed soon. # Do not add new depends_on entries or extend this step. - label: ":email: Nightly Collection & Email" @@ -573,35 +477,25 @@ steps: depends_on: - nightly-omni-performance-no-async-chunk - nightly-omni-performance-async-chunk - - nightly-omni-performance-vllm-text - - nightly-omni-performance-multi-replicas - - nightly-tts-performance - - nightly-tts-performance-higgs-audio-v3 - - nightly-diffusion-x2iat-performance-qwen-image - - nightly-diffusion-x2iat-performance-qwen-image-edit - - nightly-diffusion-x2iat-performance-qwen-image-edit-2511 - - nightly-diffusion-x2iat-performance-qwen-image-layered - - nightly-diffusion-x2iat-performance-bagel - - nightly-diffusion-x2v-performance - - nightly-testcase-statistics + - nightly-tts-performance-single + - nightly-tts-performance-2gpu + - nightly-diffusion-x2iat-performance-qwen-image-single + - nightly-diffusion-x2iat-performance-qwen-image-4gpu + - nightly-diffusion-x2v-performance-single + - nightly-diffusion-x2v-performance-2gpu if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1" commands: - pip install openpyxl - export DEFAULT_INPUT_DIR=tests/dfx/perf/results - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3 + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-single + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-2gpu - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511 - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance - - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-single + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-4gpu + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance-single + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance-2gpu - python tools/nightly/generate_nightly_perf_excel.py - python tools/nightly/generate_nightly_perf_html.py - python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html" diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index f416631ff..7fc46501b 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -9,23 +9,27 @@ steps: depends_on: upload-ready-pipeline steps: - label: "Simple · Model Executor Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/model_executor -m 'core_model and cpu'" mirror_hardwares: l4_1 - label: "Simple · Diffusion Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/diffusion -m 'core_model and cpu'" mirror_hardwares: l4_1 - label: "Simple · Engine&Entrypoints Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu'" mirror_hardwares: l4_1 - label: "Simple · Other Test" + timeout_in_minutes: 20 commands: - - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine" mirror_hardwares: l4_1 - group: ":card_index_dividers: Diffusion Test" @@ -76,7 +80,7 @@ steps: - label: "Tiny Model Tests · Diffusion (base)" commands: - - timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model" + - timeout 15m pytest -sv -n 4 tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model" mirror_hardwares: l4_1 - label: "Engine&Model Executor Test" @@ -121,32 +125,8 @@ steps: - vllm_omni/outputs/ commands: - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" - agents: - queue: "gpu_1_queue" - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - # vLLM 0.27 added a pre-flight capacity check in - # distributed/device_communicators/shm_broadcast.py. Without this - # the container gets Docker's 64 MiB default and every test in - # this step dies at engine start with - # RuntimeError: Insufficient space in /dev/shm: - # 240 MiB required, 64 MiB available - # (9 failed in 8.7s on build #2949). vLLM 0.26 had no such check - # and allocated regardless, so this step passed on main while - # being under-provisioned the whole time — 0.27 turned a silent - # risk into an explicit error. - # - # This is the only inline docker plugin block in .buildkite/cuda/; - # every other step inherits common/ci_mirror_hardwares.yml, which - # already sets 8gb. Matching that rather than inventing a value. - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" + mirror_hardwares: l4_1 + - label: "RLHF · VeRL-Omni E2E Test" source_file_dependencies: @@ -168,7 +148,13 @@ steps: - label: "Omni · Qwen3-Omni Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen3_omni.py + - tests/helpers/runtime.py + - vllm_omni/entrypoints/openai/serving_chat.py - vllm_omni/model_executor/models/qwen3_omni/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py + - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - vllm_omni/deploy/qwen3_omni_moe.yaml commands: @@ -181,6 +167,11 @@ steps: - label: "Omni · MiniCPM-o 4.5 Offline Test" source_file_dependencies: - tests/e2e/offline_inference/test_minicpmo_4_5.py + - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/models/step_audio2/ + - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ + - vllm_omni/model_executor/models/cosyvoice3/utils.py + - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - vllm_omni/deploy/minicpmo_4_5.yaml commands: - | @@ -192,6 +183,11 @@ steps: - label: "Omni · MiniCPM-o 4.5 Online Test" source_file_dependencies: - tests/e2e/online_serving/test_minicpmo_4_5.py + - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/models/step_audio2/ + - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ + - vllm_omni/model_executor/models/cosyvoice3/utils.py + - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - vllm_omni/deploy/minicpmo_4_5.yaml commands: - | @@ -209,17 +205,25 @@ steps: - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py + - vllm_omni/deploy/minicpmo_4_5.yaml - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - - vllm_omni/deploy/minicpmo_4_5_duplex.yaml - vllm_omni/experimental/fullduplex/ - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/models/step_audio2/ + - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ + - vllm_omni/model_executor/models/cosyvoice3/utils.py - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - vllm_omni/core/sched/ - vllm_omni/outputs/multimodal_accumulation.py - vllm_omni/engine/stage_init_utils.py - vllm_omni/engine/async_omni_engine.py + - vllm_omni/engine/orchestrator.py - vllm_omni/entrypoints/cli/serve.py + - vllm_omni/entrypoints/openai/api_server.py + - vllm_omni/entrypoints/openai/duplex_capability.py + - vllm_omni/config/stage_config.py commands: - | timeout 20m bash -c " @@ -239,27 +243,18 @@ steps: - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda' mirror_hardwares: h100_1 - - label: "TTS · VoxCPM2 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_voxcpm2_tts.py - - vllm_omni/model_executor/models/voxcpm2/ - - vllm_omni/deploy/voxcpm2.yaml - timeout_in_minutes: 30 - commands: - - | - timeout 30m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model' - " - mirror_hardwares: l4_1 - - label: "TTS · Qwen3-TTS CustomVoice Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/common/alias_free_activation.py - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - vllm_omni/deploy/qwen3_tts.yaml + - vllm_omni/entrypoints/openai/serving_speech.py + - vllm_omni/entrypoints/openai/tts_adapters/ + - vllm_omni/engine/stage_pool.py commands: - | timeout 20m bash -c " @@ -268,68 +263,11 @@ steps: " mirror_hardwares: l4_1 - - label: "TTS · Qwen3-TTS Base Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen3_tts_base.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - - vllm_omni/deploy/qwen3_tts.yaml - commands: - - | - timeout 20m bash -c " - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 - export VLLM_OMNI_USE_V2_RUNNER=1 - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model' - " - mirror_hardwares: l4_1 - - - label: "TTS · Higgs-Audio-V3 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_higgs_audio_v3.py - - vllm_omni/model_executor/models/higgs_audio_v3/ - - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py - - vllm_omni/deploy/higgs_multimodal_qwen3.yaml - commands: - - | - timeout 20m bash -c " - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_USE_DEEP_GEMM=0 - export VLLM_MOE_USE_DEEP_GEMM=0 - pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model' - " - mirror_hardwares: l4_1 - - - label: "Diffusion · Qwen Image Edit Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image_edit.py - - vllm_omni/diffusion/models/qwen_image/ - commands: - - | - timeout 20m bash -c " - pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'core_model' --run-level 'core_model' - " - mirror_hardwares: h100_1 - - - label: "Diffusion · Bagel Test" - source_file_dependencies: - - tests/e2e/online_serving/test_bagel.py - - vllm_omni/model_executor/models/bagel/ - - vllm_omni/diffusion/models/bagel/ - - vllm_omni/model_executor/stage_input_processors/bagel.py - - vllm_omni/deploy/bagel.yaml - commands: - - | - timeout 40m bash -c " - export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model' - " - mirror_hardwares: h100_1 - - label: "Diffusion · Wan22 Test" source_file_dependencies: - tests/e2e/online_serving/test_wan22_t2v.py - vllm_omni/diffusion/models/wan2_2/ + - vllm_omni/diffusion/models/dmd2/ commands: - | timeout 40m bash -c " @@ -342,6 +280,9 @@ steps: source_file_dependencies: - tests/e2e/online_serving/test_cosmos3.py - vllm_omni/diffusion/models/cosmos3/ + - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py + - vllm_omni/experimental/world_models/ + - vllm_omni/model_extras/cosmos3.py commands: - | timeout 40m bash -c " @@ -350,10 +291,30 @@ steps: " mirror_hardwares: h100_1 + - label: "Diffusion · MiniMax-H3 DLO DP2 T2VA Test" + timeout_in_minutes: 90 + source_file_dependencies: + - tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py + - tests/helpers/assertions.py + - tests/helpers/runtime.py + - vllm_omni/diffusion/models/minimax_h3/ + - vllm_omni/diffusion/offloader/ + - vllm_omni/diffusion/distributed/ + - vllm_omni/entrypoints/openai/ + commands: + - | + timeout 80m bash -c " + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py -m 'core_model and cuda and diffusion' --run-level 'core_model' + " + mirror_hardwares: h100_2 + - label: "Diffusion · Qwen Image Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen_image.py - vllm_omni/diffusion/models/qwen_image/ + - vllm_omni/diffusion/models/dmd2/ + - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py commands: - | timeout 40m bash -c " @@ -362,14 +323,14 @@ steps: " mirror_hardwares: h100_1 - - label: "Diffusion · Qwen Image Layered Test" + - label: "Diffusion · HunyuanVideo-1.5 Test" source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image_layered.py - - vllm_omni/diffusion/models/qwen_image/ + - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py + - vllm_omni/diffusion/models/hunyuan_video/ commands: - | timeout 40m bash -c " export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model' + pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'core_model and diffusion' --run-level 'core_model' " mirror_hardwares: h100_1 diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml index cc7f99647..85b5e4c3f 100644 --- a/.buildkite/cuda/test-weekly.yml +++ b/.buildkite/cuda/test-weekly.yml @@ -5,85 +5,146 @@ env: HF_HUB_ETAG_TIMEOUT: 60 steps: - - - label: "Reliability Test - qwen3-omni" - timeout_in_minutes: 180 + - label: "Simple · CPU Coverage Test" + key: weekly-simple-cpu-coverage depends_on: upload-weekly-pipeline - if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + if: build.env("WEEKLY") == "1" + timeout_in_minutes: 60 commands: - - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" - mirror_hardwares: h100_2 - - - label: "Reliability Test - wan22" - timeout_in_minutes: 180 - depends_on: upload-weekly-pipeline - if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" - commands: - - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" - mirror_hardwares: h100_2 - - - label: "Reliability Test - hunyuan-image" - timeout_in_minutes: 180 - depends_on: upload-weekly-pipeline - if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" - commands: - - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow" - mirror_hardwares: h100_4 - - - label: "Reliability Test - voxcpm2" - timeout_in_minutes: 180 - depends_on: upload-weekly-pipeline - if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" - commands: - - pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow" + - | + set +e + REPORT="coverage-core-model-cpu-$${BUILDKITE_STEP_ID:-local}.xml" + pytest -sv tests/ -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:$${REPORT} + EXIT=$$? + gzip -9 -f "$${REPORT}" + buildkite-agent artifact upload "$${REPORT}.gz" + exit $$EXIT mirror_hardwares: l4_1 - - group: "Reliability Test - Invalid parameters Test" + - group: "Reliability Test" depends_on: upload-weekly-pipeline if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" steps: - - label: "Invalid parameters Test · H100" - timeout_in_minutes: 120 + - label: "Reliability Test · qwen3-omni" + timeout_in_minutes: 180 commands: - - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100" + - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" mirror_hardwares: h100_2 - - label: "Invalid parameters Test · L4" + - label: "Reliability Test · wan22" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" + mirror_hardwares: h100_2 + + - label: "Reliability Test · hunyuan-image" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow" + mirror_hardwares: h100_4 + + - label: "Reliability Test · Invalid parameters · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_1" + mirror_hardwares: h100_1 + + - label: "Reliability Test · Invalid parameters · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_2" + mirror_hardwares: h100_2 + + - label: "Reliability Test · Invalid parameters · L4" timeout_in_minutes: 120 commands: - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4" mirror_hardwares: l4_1 - - group: ":card_index_dividers: E2E Tests" + - group: "Perf Test" depends_on: upload-weekly-pipeline if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" steps: - - label: "Omni · L4" + - label: ":full_moon: Omni · Perf Test · vLLM Text" + key: weekly-omni-performance-vllm-text + timeout_in_minutes: 180 + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + mirror_hardwares: h100_2 + + - label: ":full_moon: Omni · Perf Test · Async Chunk · Random" + key: weekly-omni-performance-async-chunk-random + timeout_in_minutes: 180 + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json -m "H100 and slow" + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + mirror_hardwares: h100_2 + + - label: ":full_moon: Omni · Perf Test · Multi-Replica" + key: weekly-omni-performance-multi-replicas + timeout_in_minutes: 180 + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + mirror_hardwares: h100_3 + + - group: ":card_index_dividers: E2E Tests" + depends_on: upload-weekly-pipeline + if: build.env("NON_CRITICAL") == "1" + steps: + + # Qwen2.5 expansion (CUDA 4-card / online 2-card) stays on this l4_4 job. + - label: "Omni · L4 · Multi-GPU" timeout_in_minutes: 90 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model" + - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and not cards_1" --run-level "core_model" - >- pytest -s -v tests/e2e/ --ignore=tests/e2e/online_serving/test_qwen2_5_omni_expansion.py --ignore=tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py - -m "slow and L4 and omni" + -m "slow and L4 and omni and not cards_1" --run-level "full_model" mirror_hardwares: l4_4 - - label: "Omni · H100" + - label: "Omni · H100 · Single-GPU" timeout_in_minutes: 90 commands: - >- pytest -sv - tests/e2e/ -m "slow and H100 and omni" + tests/e2e/ -m "slow and H100 and omni and cards_1" --run-level "full_model" - mirror_hardwares: h100_2 + mirror_hardwares: h100_1 + + - label: "Omni · H100 · 4-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and H100 and omni and cards_4" + --run-level "full_model" + mirror_hardwares: h100_4 - label: "TTS · L4" - timeout_in_minutes: 90 + timeout_in_minutes: 120 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - export VLLM_USE_DEEP_GEMM="0" @@ -108,29 +169,63 @@ steps: --run-level "full_model" mirror_hardwares: h100_1 - - label: "Diffusion · L4" + # Own step: the image has no NanoCodec, so this file import-skips out of + # the mark-selected TTS steps above. The overrides keep nemo_toolkit[tts] + # from downgrading transformers; see + # examples/offline_inference/text_to_speech/README.md. + - label: "TTS · Gepard-1.0 · L4" + timeout_in_minutes: 60 + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt + - uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt + - >- + pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py + -m "tts and cuda" + --run-level "full_model" + mirror_hardwares: l4_1 + + - label: "Diffusion · L4 · Single-GPU" timeout_in_minutes: 60 commands: - >- pytest -sv - tests/e2e/ -m "slow and diffusion and L4" + tests/e2e/ -m "slow and diffusion and L4 and cards_1" + --run-level "full_model" + mirror_hardwares: l4_1 + + - label: "Diffusion · L4 · Multi-GPU" + timeout_in_minutes: 60 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and diffusion and L4 and not cards_1" --run-level "full_model" mirror_hardwares: l4_4 - label: "Diffusion · H100 · Single-GPU" - timeout_in_minutes: 120 + timeout_in_minutes: 150 commands: - >- pytest -sv - tests/e2e/ -m "slow and diffusion and H100 and not distributed_cuda" + tests/e2e/ -m "slow and diffusion and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - - label: "Diffusion · H100 · Multi-GPU" + - label: "Diffusion · H100 · 2-GPU" timeout_in_minutes: 90 commands: - >- pytest -sv - tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda" + tests/e2e/ -m "slow and diffusion and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 + + - label: "Diffusion · H100 · 4-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and diffusion and H100 and cards_4" + --run-level "full_model" + mirror_hardwares: h100_4 diff --git a/.buildkite/intel/pipeline-intel.yml b/.buildkite/intel/pipeline-intel.yml index 2a1507456..d6134b028 100644 --- a/.buildkite/intel/pipeline-intel.yml +++ b/.buildkite/intel/pipeline-intel.yml @@ -10,7 +10,7 @@ steps: DOCKER_BUILDKIT: "1" # Buildkite will automatically replace this with the actual commit hash VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}" - VLLM_VERSION: "v0.27.0" + VLLM_VERSION: "v0.28.0" priority: 100 timeout_in_minutes: 60 soft_fail: false diff --git a/.buildkite/intel/scripts/run-xpu-test.sh b/.buildkite/intel/scripts/run-xpu-test.sh index 742e39b78..20a3a5d62 100755 --- a/.buildkite/intel/scripts/run-xpu-test.sh +++ b/.buildkite/intel/scripts/run-xpu-test.sh @@ -5,7 +5,12 @@ set -ex omni_source_dir=$(git rev-parse --show-toplevel) -base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}" +: "${VLLM_VERSION:?VLLM_VERSION must be set}" +# Official upstream XPU image. It is published from vllm's own +# docker/Dockerfile.xpu, which is exactly what build_vllm_base() clones and +# builds below, so pulling it is equivalent to building the fallback base. +upstream_base_image="vllm/vllm-openai-xpu:${VLLM_VERSION}" +local_base_image="xpu/vllm-omni-ci-base:${VLLM_VERSION}" image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}" container_name="xpu_${BUILDKITE_COMMIT}_$( tr -dc A-Za-z0-9 &2 + base_image_name="${upstream_base_image}" + return + fi + + echo "WARNING: ${upstream_base_image} unavailable remotely and locally" >&2 + base_image_name="${VLLM_BASE:-${local_base_image}}" + if [ -z "$(docker images -q "${base_image_name}")" ]; then + echo "WARNING: building ${base_image_name} from vLLM ${VLLM_VERSION} source" >&2 + build_vllm_base + fi +} + +resolve_vllm_base # Try building the docker image docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" @@ -87,8 +119,9 @@ time timeout -k 30 30m docker run \ echo $ZE_AFFINITY_MASK pip install tblib==3.1.0 cd /workspace/vllm-omni - pytest -v -s -m "core_model and xpu and B60" + XPU_TEST_PATHS="tests/diffusion tests/dfx tests/e2e" + pytest -v -s $XPU_TEST_PATHS -m "core_model and xpu and B60" pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py - pytest -v -s -m "advanced_model and xpu and B60" - pytest -v -s -m "omni and xpu and B60" + pytest -v -s $XPU_TEST_PATHS -m "advanced_model and xpu and B60" + pytest -v -s $XPU_TEST_PATHS -m "omni and xpu and B60" ' diff --git a/.buildkite/npu/test-npu-nightly.yml b/.buildkite/npu/test-npu-nightly.yml index 050cef4f7..e049783ce 100644 --- a/.buildkite/npu/test-npu-nightly.yml +++ b/.buildkite/npu/test-npu-nightly.yml @@ -69,12 +69,67 @@ steps: VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" HF_TOKEN: "${HF_TOKEN}" commands: + # Pin the Seed-TTS scorer to the second card: _get_eval_device() otherwise + # auto-picks npu:0, which the omni server under test already fills with its + # KV cache, so Whisper OOMs on `.to(device)` and every clip reports + # "ASR / WER failed". + - export SEED_TTS_WER_EVAL=1 + - export SEED_TTS_EVAL_DEVICE=npu:1 - | set +e pytest -s -v tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py -m "full_model and npu and A3" --run-level full_model EXIT=$$? buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json" exit $$EXIT + - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" + key: nightly-omni-performance-minicpmo-4-5-npu + timeout_in_minutes: 180 + mirror_hardwares: a3_npu_2 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + HF_TOKEN: "${HF_TOKEN}" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test (A2)" + key: nightly-omni-performance-minicpmo-4-5-npu-a2 + timeout_in_minutes: 180 + mirror_hardwares: a2b3_npu_1 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + HF_TOKEN: "${HF_TOKEN}" + HF_HUB_OFFLINE: "1" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A2" --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test" + key: nightly-omni-performance-minicpmo-4-5-duplex-seed-tts-npu + timeout_in_minutes: 180 + mirror_hardwares: a3_npu_2 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + HF_TOKEN: "${HF_TOKEN}" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group depends_on: upload-nightly-pipeline @@ -96,33 +151,6 @@ steps: EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" - key: nightly-diffusion-x2iat-group - depends_on: upload-nightly-pipeline - if: >- - build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2iat-test" - steps: - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511" - key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511 - timeout_in_minutes: 180 - mirror_hardwares: a3_npu_8 - env: - VLLM_WORKER_MULTIPROC_METHOD: spawn - VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" - HF_TOKEN: "${HF_TOKEN}" - DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA - commands: - - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export CACHE_DIT_VERSION=1.5.0 - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json - EXIT=$$? - buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" - exit $$EXIT - group: ":card_index_dividers: Diffusion I2I Model Test" key: nightly-HunyuanImage3-x2i-test-group depends_on: upload-nightly-pipeline diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index f23ba7120..9aa55e5d0 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -1,4 +1,27 @@ steps: + - label: "Platforms Test" + key: ready-platforms-test + depends_on: upload-ready-pipeline + mirror_hardwares: a3_npu_2 + commands: + - >- + pytest -s -v + tests/platforms/npu/test_minicpmo_code2wav_npugraph.py + -m 'core_model and npu and A3 and omni' + --run-level 'core_model' + + - label: "Diffusion NPU layer tests" + key: ready-diffusion-npu-layer-tests + depends_on: upload-ready-pipeline + mirror_hardwares: a3_npu_2 + commands: + - >- + pytest -s -v + tests/diffusion/layers/test_activation.py + tests/diffusion/layers/test_fused_qk_norm_rope_npu.py + -m 'core_model and npu and A3' + --run-level 'core_model' + - group: ":card_index_dividers: Omni Model Test" key: ready-omni-group depends_on: upload-ready-pipeline @@ -32,7 +55,7 @@ steps: depends_on: upload-ready-pipeline steps: - label: ":sound: TTS · Qwen3-TTS CustomVoice Test" - mirror_hardwares: a2b3_npu_1 + mirror_hardwares: a3_npu_2 env: VLLM_WORKER_MULTIPROC_METHOD: spawn VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index bae8cd1b5..16b68b06e 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -105,7 +105,7 @@ For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` ## Path B: Custom/Private Repo Model -For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns. +For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns. ### B1. Understand the reference repo @@ -126,9 +126,9 @@ This is the key design decision for custom models. Follow these placement rules: | **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models//pipeline_.py` | Always required | | **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models//_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` | | **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models//` as separate files | `autoencoder.py`, `fusion.py` | -| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone | +| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | original-repo package via git clone | | **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts | -| **Download/setup script** | `examples/offline_inference//download_.py` | `download_dreamid_omni.py` | +| **Download/setup script** | `examples/offline_inference//download_.py` | `download_.py` | | **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` | ### B3. Handle external dependencies @@ -157,7 +157,7 @@ Copy the essential model files into `vllm_omni/diffusion/models//` and ada Custom models have two common patterns for weight loading: -**Pattern 1: Bypass standard loader** (DreamID-Omni style) +**Pattern 1: Bypass standard loader** (eager custom init) When the original model has complex custom init functions that load weights in `__init__`: @@ -338,7 +338,7 @@ Classify the model's **CI priority** first: **L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`). -**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (1-GPU vs `distributed_cuda`). +**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (`cards_1` vs `not cards_1`). **L4 design (medium / low):** One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority. @@ -610,6 +610,24 @@ See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling. --- +## Pre-commit conventions + +New library files must pass the local gates in +[docs/contributing/README.md](../../../docs/contributing/README.md#linting). +That page is the full hook list (SPDX, forbidden imports including Hugging Face +Hub / Triton / pickle, `torch.cuda`, mypy, test marks, markdownlint, Buildkite, +shellcheck). In particular: + +- SPDX copyright is `vLLM-Omni project` (stale `vLLM project` is rewritten). +- Use `import regex as re` and `pybase64` in `vllm_omni/`; do not import stdlib + `re` or `base64`. Hugging Face Hub downloads go through + `vllm.transformers_utils.repo_utils`. +- Do not add `torch.cuda.*` call sites; use `current_omni_platform`. +- New `tests/**/test_*.py` files need a CI level mark and a hardware mark. +- Do not expand `CHECK_IMPORTS[*].allowed_files` or `ALLOWED_FILES` without review. +- GitHub Actions skips SPDX/shellcheck/mypy-3.10/test-marks/markdownlint; run + `pre-commit` locally. + ## Iterative Development Tips 1. **Start minimal**: Basic generation first, no parallelism/caching diff --git a/.claude/skills/add-diffusion-model/references/custom-model-patterns.md b/.claude/skills/add-diffusion-model/references/custom-model-patterns.md index 2434e0b5d..4e25f9762 100644 --- a/.claude/skills/add-diffusion-model/references/custom-model-patterns.md +++ b/.claude/skills/add-diffusion-model/references/custom-model-patterns.md @@ -16,20 +16,19 @@ vllm_omni/diffusion/models/wan2_2/ The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`. -### Custom model with external deps (e.g., DreamID-Omni) +### Custom model with external deps ``` -vllm_omni/diffusion/models/dreamid_omni/ +vllm_omni/diffusion/models// ├── __init__.py # Exports pipeline only -├── pipeline_dreamid_omni.py # Pipeline: loads ALL weights in __init__ via custom helpers -├── fusion.py # Custom fusion architecture (video + audio cross-attention) -└── wan2_2.py # Re-implemented Wan backbone with split API +├── pipeline_.py # Pipeline: loads ALL weights in __init__ via custom helpers +└── .py # Custom transformer / fusion architecture -examples/offline_inference/x_to_video_audio/ -└── download_dreamid_omni.py # Downloads weights from 3 HF repos + clones code repo +examples/offline_inference// +└── download_.py # Downloads weights and/or clones code repo ``` -All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except. +All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency imported with try/except. ### Custom model with ported code (e.g., BAGEL) @@ -76,7 +75,7 @@ self.weights_sources = [ ] ``` -### Pattern 3: Fully custom loading (DreamID-Omni) +### Pattern 3: Fully custom loading ``` init → load ALL weights eagerly via custom helpers → load_weights() = no-op @@ -108,8 +107,8 @@ Standard diffusers `model_index.json`: Custom model `model_index.json` (minimal): ```json { - "_class_name": "DreamIDOmniPipeline", - "fusion": "DreamID-Omni/dreamid_omni.safetensors" + "_class_name": "YourCustomPipeline", + "fusion": "weights/model.safetensors" } ``` @@ -117,7 +116,7 @@ The only **required** field is `_class_name` — it must match a key in `_DIFFUS ## External Dependency Management -### Git clone + .pth injection (DreamID-Omni pattern) +### Git clone + .pth injection ```python def download_dependency(): @@ -163,7 +162,7 @@ The engine checks `isinstance(pipeline, SupportImageInput)` at startup to config Diffusers models use `config.json` in each subfolder. Custom models often use: -**Module-level config dicts** (DreamID-Omni): +**Module-level config dicts**: ```python VIDEO_CONFIG = { "patch_size": [1, 2, 2], "model_type": "ti2v", @@ -181,7 +180,7 @@ vae_cfg = bagel_cfg.get("vae_config", {}) ## Custom Architecture Patterns -### Split forward API (DreamID-Omni) +### Split forward API When a fusion model needs to interleave blocks from two backbones: diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md index ff86ad457..8c11d04ba 100644 --- a/.claude/skills/add-tts-model/SKILL.md +++ b/.claude/skills/add-tts-model/SKILL.md @@ -138,9 +138,12 @@ A shared buffer silently corrupts audio across concurrent requests — the sympt - Load codec weights (may need lazy loading from separate checkpoint) - Implement `forward()`: codec codes -> audio waveform - Return `OmniOutput` with `multimodal_outputs` -5. **Create stage config YAML** defining both stages, memory allocation, and model paths -6. **Create stage input processor** for prompt building -7. **Write end2end.py** test script +5. **Define and register the pipeline topology** in + `vllm_omni/model_executor/models//pipeline.py` +6. **Create the deploy YAML** under `vllm_omni/deploy/` for placement, + memory sizing, connectors, and runtime overrides +7. **Create stage input processor** for prompt building and inter-stage handoff +8. **Write end2end.py** test script ### Critical Parameters to Get Right @@ -174,7 +177,7 @@ These bugs appear in almost every new TTS PR. Check all before the first push. S - **All return paths must emit `model_outputs`** — if any early-return branch skips setting `model_outputs`, the serving layer silently drops that step's audio (fish speech: `fix: ensure ALL return paths emit model_outputs`) - **Per-request state isolation** — for batched concurrent requests, key all state by request ID; a shared buffer corrupts audio across requests (fish speech: `fix: per-request vocode + delta emission`) - **Codec tensor device** — move codec codes to the codec decoder's device before calling decode; mismatches cause silent CPU fallback or crashes (fish speech: `fix: use model device for CUDA stream`) -- **AR stage `max_num_seqs`** — set to at least 4 in production configs; for single-stage models this is the only stage. For two-stage models, Stage 0 (AR) needs `max_num_seqs ≥ 4` to pipeline concurrent requests; Stage 1 (codec decoder) typically uses `max_num_seqs: 1` intentionally. Default of 1 everywhere causes audio gaps under concurrency because the codec window round-robins across requests (RFC #2568) +- **AR stage `max_num_seqs`** — set to at least 4 in production deploy configs; for single-stage models this is the only stage. For two-stage models, Stage 0 (AR) needs `max_num_seqs ≥ 4` to pipeline concurrent requests; Stage 1 (codec decoder) is model-specific and may intentionally use `max_num_seqs: 1`. Defaulting the AR stage to 1 causes audio gaps under concurrency because the codec window round-robins across requests (RFC #2568) ### Optional Dependency Handling @@ -190,18 +193,18 @@ pattern, signature constraints, and MOSS-TTS-Nano reference. When the upstream model cannot be cleanly split into an AR stage and a separate decoder, run the full pipeline inside a single AR worker and stream audio through a per-request `inference_stream()` generator keyed by -`_omni_req_id`. Stage config must set `worker_type: ar`, -`engine_output_type: audio`, `final_output: true`, `is_comprehension: true`, -and `async_chunk: false` at the top level. Only extract params from +`_omni_req_id`. Define one `StagePipelineConfig` with +`execution_type=StageExecutionType.LLM_AR`, `engine_output_type="audio"`, +`final_output=True`, and `owns_tokenizer=True`. Set `async_chunk: false` in +the deploy YAML. Only extract params from `additional_information` that you actually forward, or pre-commit fails `ruff F841`. Full walkthrough with the complete `forward()` / `_create_stream_gen()` -skeleton and stage-config fields: +skeleton plus pipeline/deploy definitions: [references/single-stage-ar.md](references/single-stage-ar.md). For an in-tree reference, look for any single-stage AR model under -`vllm_omni/model_executor/models/` — e.g. the MOSS-TTS-Nano integration when -it lands. +`vllm_omni/model_executor/models/`, such as MOSS-TTS-Nano. **VoxCPM2 is a different pattern** and should not reuse this skeleton — it runs the base LM under vLLM PagedAttention with external side-computation. @@ -210,7 +213,8 @@ See `plan/voxcpm2_native_ar_design.md`. ### Deliverables - Model files in `vllm_omni/model_executor/models//` -- Stage config YAML +- Registered `pipeline.py` topology +- Deploy YAML under `vllm_omni/deploy/` - Working `end2end.py` at `examples/offline_inference/text_to_speech//end2end.py` - New section in `examples/offline_inference/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/offline_inference//` dir or a per-model `README.md` inside `text_to_speech//` — the hub README is the documented surface and the mkdocs `generate_examples` hook only descends one level into `examples//`. @@ -341,8 +345,8 @@ Heavier scenarios in the same file use **`advanced_model` only** (streaming, ext @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None: - openai_client.send_audio_speech_request({...}) +def test_voice_clone_en_non_streaming_001(omni_server, online_client) -> None: + online_client.send_audio_speech_request({...}) ``` **L4 consolidation:** Prefer parametrized `OmniServerParams` rows (`default`, `async_chunk`, feature flags) in one expansion module rather than many merge-only files ([#1832](https://github.com/vllm-project/vllm-omni/issues/1832)). @@ -359,8 +363,8 @@ def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None: | Path | Fixture | Call | |------|---------|------| -| Online `/v1/*` | `openai_client` | `openai_client.send_*_request(request_config)` | -| Offline inference | `omni_runner_handler` | `omni_runner_handler.send_*_request(request_config)` | +| Online `/v1/*` | `online_client` | `online_client.send_*_request(request_config)` | +| Offline inference | `offline_client` | `offline_client.send_*_request(request_config)` | 1. **Grep `runtime.py` first** — reuse `send_omni_request`, `send_diffusion_request`, `send_audio_speech_request` (online + offline Qwen-style TTS), `send_single_stage_tts_request` (Coqui XTTS / MOSS-TTS-Nano offline), etc. 2. **No matching helper** → add `send__request` (or `send__http_request` for negative/dfx) in **`runtime.py`** with general `assert_*` bundled inside, **then** call it from the test. @@ -392,7 +396,7 @@ See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers `/health` return non-200 until you're actually ready. - **Mark tests per the CI Levels table** — baseline smoke: `core_model` + `advanced_model`; heavier cases: `advanced_model` only; L4 expansion: `full_model` - **No per-model helper modules** — do not create `tests/helpers/{slug}.py`; keep constants and `request_config` payloads in the test file -- **Online and offline e2e go through `runtime.py`** — `openai_client.send_audio_speech_request` (online); `omni_runner_handler.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests +- **Online and offline e2e go through `runtime.py`** — `online_client.send_audio_speech_request` (online); `offline_client.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests ## Phase 4: Async Chunk (Streaming) @@ -400,21 +404,29 @@ See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers ### Steps -1. **Update stage config YAML**: +1. **Update the pipeline topology** so the producing stage declares its async + handoff processor with `async_chunk_process_next_stage_input_func`. +2. **Update the deploy YAML** to enable async chunk and configure the connector: ```yaml async_chunk: true - codec_chunk_frames: 25 # frames per chunk - codec_left_context_frames: 25 # overlap for smooth boundaries + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + codec_chunk_frames: 25 + codec_left_context_frames: 25 ``` -2. **Implement chunk handling in Stage 1**: +3. **Implement chunk handling in Stage 1**: - Accept partial input (chunk of codec codes) - Handle left context for smooth audio boundaries - Return partial audio in `OmniOutput` -3. **Test streaming**: +4. **Test streaming**: - Verify audio quality matches non-streaming output - Check for artifacts at chunk boundaries - Measure TTFA (time to first audio) -4. **Update online serving** to support `stream=true` with PCM output +5. **Update online serving** to support `stream=true` with PCM output ### Streaming Architecture @@ -435,7 +447,8 @@ Stage 0 (AR) Stage 1 (Decoder) ### Deliverables -- Updated stage config with async_chunk enabled +- Updated pipeline async handoff processor and deploy config with + `async_chunk: true` - Smooth streaming audio without boundary artifacts - TTFA metrics @@ -507,14 +520,15 @@ Use this checklist when integrating a new TTS model: - [ ] Config classes created with `model_type` registration - [ ] Stage 0 (AR) implemented and generates correct tokens - [ ] Stage 1 (Decoder) produces correct audio from tokens — dtype float32 for codec decoder -- [ ] Stage 1 `max_num_seqs` ≥ 4 in production config (default 1 causes gaps under concurrency) +- [ ] AR stage `max_num_seqs` ≥ 4 in the production deploy config unless the model has a tested lower limit - [ ] Optional dependency fallbacks handled at `load_weights()` time (torchaudio/soundfile/etc.) - [ ] Streaming: codec codes accumulated across AR steps (not reset per step) - [ ] Streaming: delta audio emitted per chunk, not full re-decoded waveform - [ ] Streaming: all `forward()` return paths emit `model_outputs` - [ ] Streaming: per-request state keyed by request ID (not shared across requests) - [ ] Streaming: codec tensors moved to codec decoder device before decode -- [ ] Stage config YAML created +- [ ] Pipeline topology defined in `pipeline.py` and registered in `OMNI_PIPELINES` +- [ ] Deploy YAML created under `vllm_omni/deploy/` - [ ] `end2end.py` produces audio matching reference quality - [ ] README.md written @@ -531,7 +545,8 @@ Use this checklist when integrating a new TTS model: - [ ] Documentation added (offline + online docs, nav, supported models) ### Phase 4: Async Chunk -- [ ] Stage config updated with `async_chunk: true` +- [ ] Pipeline declares the async handoff processor +- [ ] Deploy config sets `async_chunk: true` and connector chunk parameters - [ ] Stage 1 handles partial chunks correctly - [ ] No audio artifacts at chunk boundaries - [ ] Streaming via API (`stream=true`) works @@ -548,7 +563,7 @@ Use this checklist when integrating a new TTS model: - [ ] `pre-commit run --files ` passes before every push - [ ] Every commit has `Signed-off-by` matching the author email (`git commit -s`) - [ ] `git config user.email` matches the email registered on your GitHub account -- [ ] Details and failure-recovery commands: [references/precommit-dco.md](references/precommit-dco.md) +- [ ] Hook list, failure table, and DCO recovery: [references/precommit-dco.md](references/precommit-dco.md) ## References diff --git a/.claude/skills/add-tts-model/references/precommit-dco.md b/.claude/skills/add-tts-model/references/precommit-dco.md index 86a1f42ce..bb67db740 100644 --- a/.claude/skills/add-tts-model/references/precommit-dco.md +++ b/.claude/skills/add-tts-model/references/precommit-dco.md @@ -30,6 +30,18 @@ but the changes are correct — stage the modified files and re-commit. | `ruff F841` | Variable extracted but never forwarded to model call | Remove the extraction or wire it through | | `ruff E402` | Import added below function definitions | Move to top-level import block | | `ruff format` | Line length, spacing, quote style | Accept auto-fix, stage, re-commit | +| `check-spdx-header` | Missing header, or copyright still says `vLLM project` | Two-line Omni header (`vLLM-Omni project`); restage the rewrite | +| `check-forbidden-imports` | Stdlib `re`/`base64`, pickle, Hugging Face Hub API, or direct Triton/TileLang | `import regex as re` and `pybase64`; use `vllm.transformers_utils.repo_utils`. Do **not** add the file to `allowed_files` without review | +| `check-torch-cuda-call` | New `torch.cuda.*` call site | Use `current_omni_platform` / `OmniPlatform`; do not grow `ALLOWED_FILES` without review | +| `check-tts-adapter-migration` | New `self._tts_model_type` branch in `serving_speech.py` | Put per-model logic in `tts_adapters/`. Lower `MAX_MODEL_TYPE_BRANCHES` when removing branches; do not raise it without review | +| `check-mark` | New `tests/**/test_*.py` missing level/hardware mark, or direct `pytest.mark.H100` | Add `core_model`/`advanced_model`/… plus `cpu`/`cuda`/`hardware_test(`; SKU marks only via helpers | +| `mypy-3.10` | Type error on changed `vllm_omni/` files | Fix the types; model trees are excluded. Extra versions: `pre-commit run --hook-stage manual mypy-3.12` | +| `markdownlint-cli2` | Docs/README markdown lint | Accept auto-fix under `docs/` / `recipes/` / root README | +| `check-buildkite` | Invalid `.buildkite/*.yml` | Fix the YAML; do not grow `SKIP_FILES` without review | +| `shellcheck` | No `shellcheck` on PATH, or script warning | `apt-get`/`dnf`/`brew install shellcheck`, or `shellcheck.exe` on Windows PATH | + +Canonical policy, CI `SKIP` list, and allowlist locations: +[docs/contributing/README.md](../../../docs/contributing/README.md#linting). ## DCO sign-off diff --git a/.claude/skills/add-tts-model/references/single-stage-ar.md b/.claude/skills/add-tts-model/references/single-stage-ar.md index ed53d3026..7a48197af 100644 --- a/.claude/skills/add-tts-model/references/single-stage-ar.md +++ b/.claude/skills/add-tts-model/references/single-stage-ar.md @@ -86,12 +86,51 @@ class YourModelForCausalLM(nn.Module): - Handle both `"audio"` (incremental) and `"result"` (final combined) event types from upstream models. -## Stage config +## Pipeline and deploy config -Single stage with `worker_type: ar`, `engine_output_type: audio`, -`final_output: true`, `is_comprehension: true`, and `async_chunk: false` at -the top level. Omitting any of these causes silent misclassification in the -serving layer. +Define the fixed single-stage topology in `pipeline.py`: + +```python +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +YOUR_TTS_PIPELINE = PipelineConfig( + model_type="your_model_name", + default_deploy_config_name="your_model_name.yaml", + model_arch="YourModelForCausalLM", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="your_model_name", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="audio", + final_output=True, + final_output_type="audio", + ), + ), +) +``` + +Register it in `vllm_omni/config/pipeline_registry.py`. Keep placement and +runtime sizing in `vllm_omni/deploy/your_model_name.yaml`: + +```yaml +async_chunk: false + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 4 + gpu_memory_utilization: 0.8 +``` + +Do not move `execution_type`, output ownership, tokenizer ownership, or model +architecture into the deploy YAML; those are pipeline topology. ## Lint discipline @@ -101,8 +140,5 @@ pre-commit. ## Reference implementation -Look for any single-stage AR model under -`vllm_omni/model_executor/models/` — e.g. `moss_tts_nano/` when its -integration lands. If none is in tree yet, follow the skeleton above and -cross-check against the `is_comprehension: true` / `async_chunk: false` -dispatch in `vllm_omni/entrypoints/openai/serving_speech.py`. +Use `vllm_omni/model_executor/models/moss_tts_nano/pipeline.py` and +`vllm_omni/deploy/moss_tts_nano.yaml` as the in-tree reference. diff --git a/.claude/skills/find-simplifications/SKILL.md b/.claude/skills/find-simplifications/SKILL.md new file mode 100644 index 000000000..a89981d1e --- /dev/null +++ b/.claude/skills/find-simplifications/SKILL.md @@ -0,0 +1,135 @@ +--- +name: find-simplifications +description: Find evidence-backed simplification candidates in vLLM-Omni and, when requested, turn them into focused proposals or code changes. Use for audits of dead, duplicated, speculative, over-generalized, unnecessarily defensive, or hand-rolled code. Use review-pr for ordinary correctness review and diffusion-perf-opt for performance-first optimization. +--- + +# Find vLLM-Omni Simplifications + +Find a small number of well-proven ways to reduce code, state, APIs, or maintenance cost without weakening supported behavior. Treat this as an architecture and ownership audit, not a dead-code search or style cleanup. + +The default deliverable is a read-only report. Do not add TODOs, design documents, issues, or code unless the user asks for those changes. + +## Establish the Contract + +Before judging a candidate: + +1. Read the repository instructions supplied by the host. +2. Read [`docs/design/architecture_overview.md`](../../../docs/design/architecture_overview.md). +3. Select only the relevant active module and feature documents from [`docs/design/index.md`](../../../docs/design/index.md). Archived design pages are historical context, not current contracts. +4. Identify the comparison base with `git merge-base HEAD origin/main`; do not assume the current branch represents the intended baseline. +5. If tests or CI coverage may change, use [`vllm-omni-test`](../vllm-omni-test/SKILL.md) for test placement, markers, and commands. + +Separate an implementation accident from an intentional boundary. In particular: + +- vLLM-Omni extends and sometimes overrides upstream vLLM. A similar upstream implementation is evidence for reuse only after comparing API, lifecycle, and multi-stage semantics. +- Registries, import strings, model configuration, deployment YAML, plugins, and optional dependencies create dynamic consumers that a static reference count can miss. +- Hardware backends, model families, execution modes, and connector implementations may be intentionally parallel. Do not propose collapsing one only because it is inactive in the local environment. +- Stage, replica, request, cache, and transport state can look redundant while representing different owners. Prove that two fields or barriers encode the same fact before merging them. +- A shorter hot path is not a simplification if it adds synchronization, tensor copies, host materialization, recompilation, or a measurable quality/performance regression. + +## Strong Candidates + +A strong candidate removes a meaningful surface and has a clear replacement or deletion boundary. Examples include: + +- A public method, configuration option, registry entry, route, message field, or event has no production consumer and no compatibility obligation. +- Tests or documentation are the only consumers of behavior that is not part of a supported contract. +- Multiple layers mirror the same request, lifecycle, cache-validity, readiness, or topology state. +- AR, diffusion, platform, or model-specific implementations duplicate a stable common operation and can share it without hiding performance-critical differences. +- An inherited upstream vLLM primitive is reimplemented locally without an Omni-specific semantic difference. +- Compatibility or fallback branches protect environments the project no longer supports. +- Serialization, IPC, output formatting, or connector code materializes the same payload more than once or converts between equivalent representations. +- A helper, abstraction, package split, or extension point has only one real implementation and adds more indirection than policy. +- Hand-written infrastructure can be replaced by an existing project dependency, Python standard-library feature, PyTorch primitive, or upstream vLLM facility with net deletion and equal observability. +- Tests maintain large mocks, snapshots, or expected data solely for an unused API or representation. + +Do not elevate typo fixes, formatting, isolated renames, or vague complexity complaints into simplification candidates. Bundle small cleanups only when they support one substantive removal. + +## Survey by Ownership Boundary + +For broad audits, cover the largest or riskiest production deltas first. Useful domains are: + +- `vllm_omni/engine/`: orchestration, stage pools, request state, output convergence, cancellation, and replica lifecycle. +- `vllm_omni/worker/` and AR model execution: upstream vLLM overlays, scheduler/cache state, collective RPC, adapters, and weight loading. +- `vllm_omni/diffusion/`: request/step execution, batching, output materialization, IPC, offload, parallelism, and model pipelines. +- `vllm_omni/entrypoints/`: duplicated validation/defaulting, offline/online parity, streaming, and route compatibility. +- Connectors and distributed code: ownership, copies, barriers, handles, cleanup, and failure propagation across process/device/node boundaries. +- Platform backends: genuine vendor differences versus copied control flow that has drifted. +- Configuration, registries, recipes, and examples: duplicate sources of truth and options that never reach a runtime consumer. +- Tests and CI: redundant fixtures, mocks of obsolete contracts, and expensive coverage that can move to a cheaper level without losing the guarded behavior. + +When the user explicitly requests a repository-wide or many-candidate audit and parallel agent work is available and authorized, divide these domains among agents and require the same evidence format from each. Otherwise survey them sequentially. + +## Prove or Reject Each Candidate + +Use `rg` first, then read every relevant caller and implementation. Search exact Python symbols plus serialized names, route paths, configuration keys, registry strings, YAML values, and CLI spellings. + +Classify evidence: + +- **Production:** `vllm_omni/`, runtime configuration, registries, serving entrypoints, connectors, recipes used for supported deployments, and executable loader paths. +- **Validation/support:** `tests/`, benchmarks, examples, docs, CI, and developer tooling. +- **External/dynamic:** upstream vLLM contracts, third-party plugins, model repositories, optional backends, and string-loaded objects. State what was checked and what remains unknown. + +For each candidate, answer: + +1. What exact surface would disappear or merge? +2. Who owns the state or behavior today? +3. Which production consumers exist, including dynamic consumers? +4. What supported behavior, compatibility, performance, memory, or observability could change? +5. What is the smallest validation that would catch a bad simplification? +6. Is the result net deletion after tests, adapters, migration glue, and documentation are counted? + +Reject or downgrade the candidate when a live consumer exists and removal is actually a feature decision; when an active design contract explains the separation; when the change only moves complexity; or when hardware/model evidence is unavailable for a performance-sensitive claim. + +## Lifecycle and Data-Movement Audits + +For asynchronous or distributed code, map the ownership graph before recommending deletion: + +```text +request admission -> orchestrator -> stage/replica -> worker/device + -> connector or IPC boundary -> output aggregation -> terminal cleanup +``` + +Associate every lock, event, sentinel, queue, readiness flag, timeout, abort path, and cache marker with an owner and transition. Coalesce mechanisms only when they represent the same state and have the same failure boundary. Preserve separate mechanisms when they protect publication versus rollback, local versus remote ownership, first-terminal-outcome arbitration, or cleanup after partial failure. + +For tensors and media, also trace device, process, and ownership transitions. Count GPU-to-CPU copies, shared-memory/object-store handles, serialization, decode/encode steps, and unlink/release responsibility. Prefer eliminating a transfer or representation over merely hiding it behind a helper. + +## Dependency Replacements + +Prefer facilities already available in the repository before adding a dependency. For a proposed replacement: + +- Name the exact local behavior it covers and any semantics that remain in glue code. +- Check maintenance, license, optional-dependency boundaries, platform support, and transitive footprint. +- Compare net deletion, test burden, import/startup cost, and runtime behavior. +- Benchmark when the replaced code is on a model, scheduler, connector, serialization, or serving hot path. + +A wrapper that retains the same state machine or conversion logic is not a simplification. + +## Deliver the Result + +Report only candidates supported by concrete evidence. For each one include: + +- **Surface:** files and symbols involved. +- **Evidence:** production consumers found, dynamic paths checked, and relevant design contract. +- **Simplification:** exactly what to remove, merge, inline, or reuse. +- **Benefit:** deleted API/state/code, reduced data movement, or reduced maintenance/test burden. +- **Risk and validation:** behavior that could change and the narrowest useful tests or measurements. +- **Confidence:** high or medium, with the unresolved evidence named for medium-confidence items. + +Also list representative ideas that were rejected when that prevents repeated investigation. Prefer three high-confidence candidates over a long speculative inventory. + +If the user asks for a durable design proposal, place it under `docs/design/feature/` and register it in `docs/design/index.md`; use current design-document conventions and keep implementation details proportional to the decision. If the user asks for implementation, make the smallest coherent change, update affected tests/docs, and preserve public compatibility unless the user explicitly accepts a breaking change. + +## Folding Work from Another Branch + +Diff the source branch against its own merge base with `origin/main`, not against the current feature branch. Port only independently justified candidates. Consolidate overlaps into the candidate with stronger evidence, and do not preserve a candidate count by carrying weaker duplicates. + +## Validation + +Match validation to the outgoing diff: + +- Documentation or skill changes: run the repository skill validator when available, check links/paths, search for stale source-specific terminology, and run `git diff --check`. +- Python changes: run the narrowest relevant pytest coverage through the repository environment, then lint/format the changed files. +- Distributed, model, hardware, performance, or accuracy changes: run available CPU/static checks and state the exact GPU/NPU/model gap; never present simulated evidence as device validation. +- Before submission, use [`precheck-pr`](../precheck-pr/SKILL.md) for the branch-level self-review. + +When summarizing, distinguish checks that passed from checks that were not runnable, and state the snapshot or branch against which the audit was performed. diff --git a/.claude/skills/precheck-pr/SKILL.md b/.claude/skills/precheck-pr/SKILL.md index 4c8086029..b7d6d270a 100644 --- a/.claude/skills/precheck-pr/SKILL.md +++ b/.claude/skills/precheck-pr/SKILL.md @@ -67,6 +67,10 @@ Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type **Also run the Examples-Policy check on every PR** using [references/examples-policy.md](references/examples-policy.md). Inspect only Python paths introduced by the diff. Treat a new model-, checkpoint-, vendor-, or family-specific Python example as blocking; do not report pre-existing example debt. +**Also run the Simplification pass on every PR** using [find-simplifications](../find-simplifications/SKILL.md). Keep it diff-scoped: inspect added or modified public APIs, state, abstractions, fallback and compatibility paths, data movement, and only the adjacent ownership needed to prove a candidate. Do not expand the pass into a repository-wide audit. Zero candidates is a valid pass; report only evidence-backed opportunities. Mark a candidate as blocking only when it proves a correctness or merge-readiness problem, such as newly unreachable code. Otherwise report it as a warning. + +**Also confirm local pre-commit gates** from [docs/contributing/README.md](../../../docs/contributing/README.md#linting). GitHub Actions `SKIP` does not prove they passed. Full new-hook list: Omni SPDX (`vLLM-Omni project`); forbidden imports (stdlib `re`/`base64`, pickle, Hugging Face Hub API, Triton/TileLang); no new `torch.cuda` call sites; shellcheck (macOS/Windows must install the binary); mypy-3.10; test files have CI level + hardware marks; TTS `_tts_model_type` branches in `serving_speech.py` must not increase; markdownlint on `docs/`; Buildkite YAML schema. ✗ for growing `allowed_files` / `ALLOWED_FILES` or raising `MAX_MODEL_TYPE_BRANCHES` without review. + ### Step 5: Print Report ``` @@ -80,6 +84,7 @@ Pre-check report for PR title format ✓ Code quality ⚠ 1 broad except, 2 Any hints Examples policy ✗ new model-specific Python example + Simplification ⚠ helper duplicates an existing owner PR desc integrity ✓ Registry/config ✓ Dead code ⚠ 2 warnings diff --git a/.claude/skills/precheck-pr/references/checklists.md b/.claude/skills/precheck-pr/references/checklists.md index b65a66713..dc7f85ea9 100644 --- a/.claude/skills/precheck-pr/references/checklists.md +++ b/.claude/skills/precheck-pr/references/checklists.md @@ -164,6 +164,7 @@ Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model - [ ] **PR body matches diff:** description doesn't claim changes not in the diff - [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts - [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists +- [ ] **Local pre-commit gates:** full list in [Linting](../../../../docs/contributing/README.md#linting). Blocking when the diff introduces: stale `vLLM project` SPDX; stdlib `re`/`base64` or pickle in library code; new `torch.cuda` call sites; `tests/**/test_*.py` without level+hardware marks; new `_tts_model_type` branches in `serving_speech.py`; invalid `.buildkite` YAML. macOS/Windows need a native `shellcheck`. ✗ for growing `CHECK_IMPORTS[*].allowed_files`, `ALLOWED_FILES`, or `MAX_MODEL_TYPE_BRANCHES` without review. ### Full (adds) diff --git a/.claude/skills/readme.md b/.claude/skills/readme.md index b5aaf28b2..95edc5021 100644 --- a/.claude/skills/readme.md +++ b/.claude/skills/readme.md @@ -31,9 +31,12 @@ self-check. - [`diffusion-perf-opt`](diffusion-perf-opt/SKILL.md): guides diffusion model performance optimization, including profiling traces, parallel strategies, stage timing analysis, and benchmark-driven tuning +- [`find-simplifications`](find-simplifications/SKILL.md): finds evidence-backed + opportunities to remove or merge dead, duplicated, speculative, + over-generalized, or unnecessarily defensive vLLM-Omni code - [`precheck-pr`](precheck-pr/SKILL.md): self-checks a branch before creating a - PR by validating title format, dead code, accuracy and performance claims, - and merge readiness + PR by validating title format, dead code, simplification opportunities, + accuracy and performance claims, and merge readiness - [`quantization`](quantization/SKILL.md): guides quantization method selection, model integration, checkpoint loading, and quality/performance validation for vLLM-Omni diff --git a/.claude/skills/review-pr/SKILL.md b/.claude/skills/review-pr/SKILL.md index d8b1b43c5..0b19b2b37 100644 --- a/.claude/skills/review-pr/SKILL.md +++ b/.claude/skills/review-pr/SKILL.md @@ -21,7 +21,10 @@ Make every finding: - **Calibrated:** match severity to user and maintainer impact. Do not report unrelated backlog, style already enforced by pre-commit, or a -missing test that would not protect changed behavior. +missing test that would not protect changed behavior. Do report new allowlist +or budget entries in `check_forbidden_imports.py` / `check_torch_cuda.py` / +`check_tts_adapter.py` / `check_buildkite.py` unless the PR justifies them: +those are policy changes, not lint noise. ## Select the input and depth @@ -92,12 +95,13 @@ If docs and live code disagree, verify the code/tests and report the drift. | Reference | Read when | | --- | --- | -| [model-addition-checklist.md](references/checks/model-addition-checklist.md) | A model, architecture, loader, processor, registry, or stage config is added. | +| [model-addition-checklist.md](references/checks/model-addition-checklist.md) | A model, architecture, loader, processor, registry, pipeline config, or deploy config is added. | | [perf-verification.md](references/checks/perf-verification.md) | The PR makes a latency, throughput, memory, or quality claim. | | [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. | | [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. | | [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. | | [examples-policy.md](../precheck-pr/references/examples-policy.md) | The PR adds, copies, or renames Python under `examples/`; apply the canonical policy shared with `precheck-pr`. | +| [find-simplifications](../find-simplifications/SKILL.md) | Every review; run a diff-scoped subtraction and simplification pass after correctness blockers. | ### Delivery and reviewer coordination @@ -180,9 +184,13 @@ promotion gate. Candidate or draft rules are questions, not blockers, unless current code, tests, or policy enforce them. Inspect both sides of any config, registry, serialization, connector, cache, or stage boundary. -When a diff adds or expands a helper, class, fallback, compatibility branch, or -public behavior, run a subtraction pass: remove out-of-scope behavior and check -whether each new abstraction can be deleted, merged, moved, or inlined. +Read and apply [find-simplifications](../find-simplifications/SKILL.md) on every +review. Constrain it to the diff and the adjacent ownership, callers, or +consumers needed to prove a candidate. Check whether added or expanded helpers, +classes, state, fallback and compatibility branches, data movement, or public +behavior can be deleted, merged, moved, or inlined. Zero candidates is a valid +result. Do not widen the review into repository backlog or report speculative +style preferences as simplification findings. ### 6. Verify the changed path diff --git a/.claude/skills/review-pr/references/checks/model-addition-checklist.md b/.claude/skills/review-pr/references/checks/model-addition-checklist.md index b60fc936e..bde60d7ab 100644 --- a/.claude/skills/review-pr/references/checks/model-addition-checklist.md +++ b/.claude/skills/review-pr/references/checks/model-addition-checklist.md @@ -1,7 +1,7 @@ # Model Addition Checklist Read this reference when a PR adds or expands a model architecture, pipeline, -loader, processor, registry entry, or stage configuration. +loader, processor, registry entry, pipeline config, or deploy config. Official docs: [model contribution guides](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/), [adding an omni model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_omni_model/), diff --git a/.claude/skills/review-pr/references/modules/configuration.md b/.claude/skills/review-pr/references/modules/configuration.md index bdb3a5159..6ecf08e99 100644 --- a/.claude/skills/review-pr/references/modules/configuration.md +++ b/.claude/skills/review-pr/references/modules/configuration.md @@ -3,7 +3,7 @@ Primary design: [vLLM-Omni Configuration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/vllm_omni_config/). Use for config construction, deploy YAML, registries, schema, defaults, -aliases, CLI projection, stage config, and topology. Inspect the reviewed +aliases, CLI projection, pipeline config, and topology. Inspect the reviewed page's status: while its contract remains deferred or draft, current code and tests—not proposed precedence rules—are authoritative. diff --git a/.claude/skills/review-pr/references/process/review-execution.md b/.claude/skills/review-pr/references/process/review-execution.md index f06f27933..7c0305254 100644 --- a/.claude/skills/review-pr/references/process/review-execution.md +++ b/.claude/skills/review-pr/references/process/review-execution.md @@ -111,6 +111,17 @@ Mark early findings as preliminary and continue. This is not a GitHub comment. but do not publish review events without separate authorization. - Record DCO, pre-commit, required CI, and mergeability. Pending or unknown gates do not block source review. +- GitHub Actions `SKIP` omits SPDX, shellcheck, markdownlint, mypy-3.10, and + test-mark coverage. A green GHA pre-commit job does not prove those local + gates passed. New files still need the full Linting list: Omni SPDX + (`vLLM-Omni project`), no stdlib `re`/`base64` in `vllm_omni/`, no new + pickle / Hugging Face Hub API / `torch.cuda` call sites, test marks, TTS + adapter ratchet, Buildkite schema, and native shellcheck on macOS/Windows. + See [Linting](https://vllm-omni.readthedocs.io/en/latest/contributing/#linting). +- Expanding `CHECK_IMPORTS[*].allowed_files`, `ALLOWED_FILES`, + `MAX_MODEL_TYPE_BRANCHES`, or Buildkite `SKIP_FILES` is a policy change. + Do not rubber-stamp allowlist/budget growth; require justification and + prefer fixing the call site. - Treat a failed gate as its own evidence. Do not restate its formatting/lint output as a new code-review finding. - Open CI logs only when the first failing step overlaps the frozen diff or diff --git a/.claude/skills/review-pr/references/process/review-routing.md b/.claude/skills/review-pr/references/process/review-routing.md index 10768d329..f6455d492 100644 --- a/.claude/skills/review-pr/references/process/review-routing.md +++ b/.claude/skills/review-pr/references/process/review-routing.md @@ -48,7 +48,7 @@ compatibility changes. | Signal | Read | Optional repo-local skill | | --- | --- | --- | -| New or expanded model, loader, processor, registry, or stage config | [model-addition-checklist.md](../checks/model-addition-checklist.md) | [`add-tts-model`](../../../add-tts-model/SKILL.md) or [`add-diffusion-model`](../../../add-diffusion-model/SKILL.md) | +| New or expanded model, loader, processor, registry, pipeline config, or deploy config | [model-addition-checklist.md](../checks/model-addition-checklist.md) | [`add-tts-model`](../../../add-tts-model/SKILL.md) or [`add-diffusion-model`](../../../add-diffusion-model/SKILL.md) | | Latency, throughput, memory, scaling, precision, or quality claim | [perf-verification.md](../checks/perf-verification.md) | [`diffusion-perf-opt`](../../../diffusion-perf-opt/SKILL.md) for diffusion | | Tests changed, absent for risky behavior, or test-only | [test-quality-evaluation.md](../checks/test-quality-evaluation.md) | [`vllm-omni-test`](../../../vllm-omni-test/SKILL.md) | | CI, examples, docs, public behavior, or contributor evidence | [tests-docs-checklist.md](../checks/tests-docs-checklist.md) | None | diff --git a/.claude/skills/vllm-omni-test/SKILL.md b/.claude/skills/vllm-omni-test/SKILL.md index d5102b3f2..2d4226788 100644 --- a/.claude/skills/vllm-omni-test/SKILL.md +++ b/.claude/skills/vllm-omni-test/SKILL.md @@ -52,8 +52,8 @@ Always attach markers deliberately: - `tts` — speech synthesis / TTS-only models (`/v1/audio/speech`, voice clone, etc.) - `diffusion` — generative diffusion models (image / audio / text / video from noise) - **Cross-cutting area** (when relevant): `parallel`, `cache`, `example`, `benchmark` -- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, … -- Optional: `slow`, distributed markers when multi-card is required +- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, …; SKU (`H100`, `L4`, …) and `cards_{n}` via `@hardware_test(...)` / `hardware_marks(...)` only — never `pytest.mark.H100` (`check-mark` rejects it) +- Optional: `slow`. Multi-card: `num_cards=N` on the helper (filter with `cards_N` / `not cards_1`) **Baseline smoke (L2 + L3):** The simplest e2e case per model — default deploy, minimal request — should usually carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** test function so `test-ready.yml` and `test-merge.yml` share one test. `send_*_request` picks validation depth from `--run-level`. References: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. Heavier scenarios use **`advanced_model` only**; L4 expansion uses **`full_model`**. @@ -94,8 +94,8 @@ Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-styl | Scenario | Typical location | Fixtures / runner pattern | Baseline markers & level | |----------|------------------|---------------------------|---------------------------| -| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required | -| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` | +| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `offline_client`. **Function (isolation only):** `omni_runner_function` + `offline_client_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required | +| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `online_client`. **Function (isolation only):** `omni_server_function` + `online_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` | | **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/cuda/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). | | **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: JSON `mark` with `full_model` + `omni`/`tts`/`diffusion`; wire `test-nightly.yml` Perf steps | | **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) | @@ -110,10 +110,10 @@ After choosing offline/online/docs/perf, classify the **product under test** and |-----------|------------------------------|----------------------------|----------------------------------------| | **What it is** | Multimodal LLM pipeline (thinker/talker/stages; text + vision + audio I/O) | Speech synthesis / voice models | Generative diffusion (noise → image, audio, text, or video) | | **Examples** | Qwen2.5-Omni, Qwen3-Omni | Qwen3-TTS, VoxCPM2, Higgs-Audio, Step-Audio2 | Qwen-Image, BAGEL, Wan2.2, HunyuanVideo | -| **Offline runner** | `omni_runner` + `omni_runner_handler`, `generate_multimodal` | `Omni(...).generate` or stage YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` | -| **Online client** | `openai_client.send_omni_request` (chat completions, modalities) | `openai_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) | -| **Typical assertions** | Stage outputs, text/audio keywords via `OmniRunnerHandler` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` | -| **Stage / deploy YAML** | Per-model omni stage configs (`ci/qwen3_omni_moe.yaml`, …) | `qwen3_tts.yaml`, `voxcpm2.yaml`, … | Often default serve; parallel/offload YAML for heavy DiT | +| **Offline runner** | `omni_runner` + `offline_client`, `generate_multimodal` | `Omni(...).generate` with deploy YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` | +| **Online client** | `online_client.send_omni_request` (chat completions, modalities) | `online_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) | +| **Typical assertions** | Stage outputs, text/audio keywords via `OfflineOmniClient` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` | +| **Deploy YAML** | Per-model deploy configs (`ci/qwen3_omni_moe.yaml`, …) | `qwen3_tts.yaml`, `voxcpm2.yaml`, … | Often default serve; parallel/offload YAML for heavy DiT | | **Nightly group** (`test-nightly.yml`) | **Omni Model Test** — `-m "full_model and omni"` | **TTS Model Test** — `-m "full_model and tts"` | **Diffusion X2I(&A&T)** *or* **Diffusion X2V** (see Step 3 table; same `diffusion` marker, different YAML group / file shard) | | **L4 pressure** | Expansion per modality/model as needed | Expansion + accuracy/perf in TTS group | X2I: merge feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); X2V: separate nightly group | @@ -175,9 +175,9 @@ See **L1 unit test constraints (mocking)** below for the full do/don't list. @pytest.mark.omni @hardware_test(...) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_(omni_runner, omni_runner_handler) -> None: +def test_(omni_runner, offline_client) -> None: request_config = {"prompts": ..., "modalities": [...]} # optional: images, videos, audios - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) ``` *Offline generative e2e — **Diffusion** (representative):* @@ -187,8 +187,8 @@ def test_(omni_runner, omni_runner_handler) -> None: @pytest.mark.diffusion @hardware_test(...) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_text_to_image_001(omni_runner_handler) -> None: - omni_runner_handler.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}}) +def test_text_to_image_001(offline_client) -> None: + offline_client.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}}) ``` *Offline TTS e2e — **Qwen3-TTS** (two-stage; representative):* @@ -198,8 +198,8 @@ def test_text_to_image_001(omni_runner_handler) -> None: @pytest.mark.tts @hardware_test(...) @pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: - omni_runner_handler.send_audio_speech_request({ +def test_text_to_audio_001(omni_runner, offline_client) -> None: + offline_client.send_audio_speech_request({ "input": "...", "task_type": "Base", "ref_audio": REF_AUDIO_URL, @@ -214,8 +214,8 @@ def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: @pytest.mark.tts @hardware_test(...) @pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) -def test_voice_clone_001(omni_runner_handler) -> None: - omni_runner_handler.send_single_stage_tts_request({ +def test_voice_clone_001(offline_client) -> None: + offline_client.send_single_stage_tts_request({ "input": "...", "language": "en", "prompt_audio_path": REF_AUDIO_PATH, @@ -231,9 +231,9 @@ def test_voice_clone_001(omni_runner_handler) -> None: @pytest.mark.omni @hardware_test(...) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_001(omni_server, openai_client) -> None: +def test_text_to_text_001(omni_server, online_client) -> None: request_config = {"model": omni_server.model, "messages": ..., "modalities": ["text"]} - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) ``` *Online serving e2e — **TTS** (representative):* @@ -244,9 +244,9 @@ def test_text_to_text_001(omni_server, openai_client) -> None: @pytest.mark.tts @hardware_test(...) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_server, openai_client) -> None: +def test_text_to_audio_001(omni_server, online_client) -> None: request_config = {"model": omni_server.model, "input": "...", "response_format": "wav", ...} - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) ``` *Online serving e2e — **Diffusion** X2I (representative):* @@ -255,8 +255,8 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: @pytest.mark.core_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True) -def test_text_to_image_001(omni_server, openai_client) -> None: - openai_client.send_diffusion_request({...}) # chat completions + extra_body +def test_text_to_image_001(omni_server, online_client) -> None: + online_client.send_diffusion_request({...}) # chat completions + extra_body ``` *Online serving e2e — **Diffusion** X2V (representative):* @@ -264,8 +264,8 @@ def test_text_to_image_001(omni_server, openai_client) -> None: ```python @pytest.mark.core_model @pytest.mark.diffusion -def test_text_to_video_001(omni_server, openai_client) -> None: - openai_client.send_video_diffusion_request({"model": ..., "form_data": {...}}) # /v1/videos +def test_text_to_video_001(omni_server, online_client) -> None: + online_client.send_video_diffusion_request({"model": ..., "form_data": {...}}) # /v1/videos ``` *Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts). @@ -274,6 +274,7 @@ def test_text_to_video_001(omni_server, openai_client) -> None: **3. Cross-cutting rules** +- New `.py` files need the Omni SPDX header (`Copyright contributors to the vLLM-Omni project`, not `vLLM project`). `tests/` may keep stdlib `re`/`base64`; pickle is still banned unless the file is already on the pickle allowlist. Every collected `tests/**/test_*.py` needs a CI **level** mark and a **hardware** platform mark/helper (`check-mark` runs locally; GHA skips it). Do not write `pytest.mark.H100` / `pytest.mark.L4`; use the helpers so `cards_{n}` is attached. Do not grow `allowed_files` to land a test. Policy: [docs/contributing/README.md](../../../docs/contributing/README.md#linting). - Reuse existing fixtures for the chosen scenario; do not mix “online client” assumptions into offline `OmniRunner` tests without a clear reason. - Avoid external network dependency in assertions unless the scenario is explicitly “online serving” or doc examples that require a model hub (then align with CI secrets/cache). - Keep **one test function = one intent** (one modality combo, one endpoint contract, or one acceleration combo). @@ -323,21 +324,21 @@ vLLM-Omni e2e tests start a real **OmniServer** (online) or **OmniRunner** (offl | Scope | Online fixtures | Offline fixtures | When to use | |-------|-----------------|------------------|-------------| -| **Module** (default) | `omni_server` → `openai_client` | `omni_runner` → `omni_runner_handler` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. | -| **Function** | `omni_server_function` → `openai_client_function` | `omni_runner_function` → `omni_runner_handler_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. | +| **Module** (default) | `omni_server` → `online_client` | `omni_runner` → `offline_client` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. | +| **Function** | `omni_server_function` → `online_client_function` | `omni_runner_function` → `offline_client_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. | **Rules:** 1. **Default to module scope** (`omni_server` / `omni_runner`) unless the scenario **explicitly** needs a fresh instance per test function. -2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `openai_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `openai_client_function`. Do not mix module fixture with function client (or vice versa). +2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `online_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `online_client_function`. Do not mix module fixture with function client (or vice versa). 3. Different `OmniServerParams` per `test_*` (e.g. default vs `--enable-cpu-offload`) is still OK with **module** `omni_server` — pytest parametrizes per test node; only switch to `_function` when isolation between tests matters, not merely because `server_args` differ. 4. One `test_*` with **many** server configs (expansion matrix) → single function + `@pytest.mark.parametrize("omni_server", [...], indirect=True)` + module `omni_server` (see in-tree `test_qwen_image_expansion.py`). ```python # Default — module-scoped server (L4 expansion) @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=H100)], indirect=True) -def test_foo_images_generations_default_1024(omni_server, openai_client) -> None: - openai_client.send_images_generations_request({...}) +def test_foo_images_generations_default_1024(omni_server, online_client) -> None: + online_client.send_images_generations_request({...}) # Function-scoped — reliability / per-test clean state only @pytest.mark.parametrize( @@ -345,9 +346,9 @@ def test_foo_images_generations_default_1024(omni_server, openai_client) -> None [pytest.param(OmniServerParams(model=MODEL), marks=H100)], indirect=True, ) -def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> None: - openai_client_function.send_omni_sleep_http_request({...}) - openai_client_function.send_omni_wakeup_http_request({...}) +def test_foo_sleep_wakeup_cycle(omni_server_function, online_client_function) -> None: + online_client_function.send_omni_sleep_http_request({...}) + online_client_function.send_omni_wakeup_http_request({...}) ``` #### Runtime send helpers — online **and** offline (`tests/helpers/runtime.py`) @@ -356,19 +357,19 @@ def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> | Principle | Action | |-----------|--------| -| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OpenAIClientHandler` (online) or `OmniRunnerHandler` (offline) that matches the endpoint / pipeline shape. | +| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OnlineOmniClient` (online) or `OfflineOmniClient` (offline) that matches the endpoint / pipeline shape. | | **Extend when close** | If an existing `send_*_request` almost fits (e.g. missing one optional field), extend it in `runtime.py` — do not fork logic in the test file. | | **Add when missing** | No suitable helper → add **`send__request`** (high-level: call + general `assert_*`) or **`send___http_request`** (low-level HTTP for negative/dfx) in **`runtime.py` first**, then call it from tests. | | **Test module owns payload only** | In `test_*.py`: `MODEL`, deploy path, vendored media, `get_prompt()`, and inline **`request_config` dicts** only. **No** `omni.generate(...)`, raw `requests.post`, OpenAI SDK calls, or `_collect_audio()` / `_process_output()` in e2e tests. | -**Online** (`openai_client` from `omni_server`): `OpenAIClientHandler.send_*_request`. +**Online** (`online_client` from `omni_server`): `OnlineOmniClient.send_*_request`. -**Offline** (`omni_runner_handler` from `omni_runner`): `OmniRunnerHandler.send_*_request`. +**Offline** (`offline_client` from `omni_runner`): `OfflineOmniClient.send_*_request`. | Do | Don't | |----|-------| -| Online: `openai_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test | -| Offline: `omni_runner_handler.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` | +| Online: `online_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test | +| Offline: `offline_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` | | Add missing `send_*` to **`runtime.py`** first; bundle general `assert_*` inside the send helper | A one-off `def _post_*` / `def _collect_audio` at the bottom of a test module | | Mirror naming/style of neighboring `send_*` (docstring, `request_config` dict, optional `run_level`, `err_code` / `err_message` for negative cases) | Different parameter shapes per test file for the same endpoint | @@ -391,7 +392,7 @@ def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** special validation — the general `assert_*` should read those from `request_config` inside `send_*`. -**`send_*` ↔ `assert_*` pairing (online `OpenAIClientHandler`):** +**`send_*` ↔ `assert_*` pairing (online `OnlineOmniClient`):** | High-level `send_*` (prefer in L2+ e2e) | Assert already invoked inside `runtime.py` | Low-level HTTP-only sibling | |----------------------------------------|--------------------------------------------|-----------------------------| @@ -404,7 +405,7 @@ Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** **Rule:** Tests call high-level `send_*_request` for the general contract. **Never** call the same bundled `assert_*` again. Call an extra `assert_*` in the test **only** for special, case-specific validation. -**Common `OpenAIClientHandler` entry points** (non-exhaustive — grep `runtime.py` before adding): +**Common `OnlineOmniClient` entry points** (non-exhaustive — grep `runtime.py` before adding): | Area | High-level (SDK + assert) | Low-level HTTP (`*_http_request`) | |------|---------------------------|-----------------------------------| @@ -415,7 +416,7 @@ Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** | TTS (online) | `send_audio_speech_request` | `send_audio_speech_http_request`, `send_audio_generate_http_request`, … | | Ops / meta | — | `send_health_http_request`, `send_models_http_request`, `send_omni_sleep_http_request`, … | -**Common `OmniRunnerHandler` entry points** (offline — grep `runtime.py` before adding): +**Common `OfflineOmniClient` entry points** (offline — grep `runtime.py` before adding): | Area | High-level (`send_*_request` + assert) | Notes | |------|----------------------------------------|-------| @@ -425,7 +426,7 @@ Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** | Single-stage TTS (Coqui XTTS, MOSS-TTS-Nano, …) | `send_single_stage_tts_request`, `send_single_stage_tts_batch_request` | `prompt` + `additional_information` + `omni.generate` — **do not** duplicate in tests | | New model family | **Add** `send__request` here first | Then call from `tests/e2e/offline_inference/test_*.py` | -L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OpenAIClientHandler` / `OmniRunnerHandler`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers. +L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OnlineOmniClient` / `OfflineOmniClient`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers. #### Invalid parameter validation (`tests/dfx/reliability/invalid_param_test/`) @@ -451,7 +452,7 @@ When the user asks for **invalid parameter validation**, **invalid request bodie | **Module markers** | `pytestmark = [pytest.mark.slow, pytest.mark.]` | | **Server fixture** | `_PARAMS` / `_QWEN3_TTS_SPEECH`-style list of `pytest.param(OmniServerParams(...), id="...", marks=hardware_marks(...))`; `@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)` | | **Hardware** | `hardware_marks(res={"cuda": "H100"})` for heavy diffusion/omni/video; `hardware_marks(res={"cuda": "L4"})` for smaller TTS models (must match weekly `-m "slow and L4"` step) | -| **HTTP client** | **Low-level** `openai_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) | +| **HTTP client** | **Low-level** `online_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) | | **Case shape** | Prefer **one `test_*` per route family** + `@pytest.mark.parametrize("body_spec, err_message", [...])` with stable `id=` per case; or dedicated `test__malformed_json` when not parametrized | | **Body helpers** | `_minimal__json(omni_server)` / `_minimal_*_form_data()`; merge overrides with `body.update(body_spec)` | | **Known gaps** | `pytest.mark.skip(reason="…#3649")` as `_SKIP_ISSUE_3649` when server validation is not yet strict (mirror neighboring cases) | @@ -484,13 +485,13 @@ When the user asks for **invalid parameter validation**, **invalid request bodie @pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) def test_images_generations_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, body_spec: dict[str, object], err_message: str | tuple[str, ...], ) -> None: body = _minimal_images_gen_json(omni_server) body.update(body_spec) - openai_client.send_images_generations_http_request( + online_client.send_images_generations_http_request( {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message} ) ``` @@ -541,18 +542,18 @@ def send_images_generations_request(self, request_config: dict[str, Any], ...) - # tests/e2e/.../test_foo_expansion.py — one case per test_* (no case_id branching) @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) -def test_foo_images_generations_default_1024(omni_server, openai_client) -> None: - openai_client.send_images_generations_request({"json": body, "timeout": 300}) +def test_foo_images_generations_default_1024(omni_server, online_client) -> None: + online_client.send_images_generations_request({"json": body, "timeout": 300}) @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) -def test_foo_images_generations_sizes_256_512_1024(omni_server, openai_client) -> None: +def test_foo_images_generations_sizes_256_512_1024(omni_server, online_client) -> None: for size in ("256x256", "512x512", "1024x1024"): - openai_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300}) + online_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300}) # chat diffusion — separate test; send_diffusion_request bundles assert_diffusion_response @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) -def test_foo_chat_completions_t2i_fallback(omni_server, openai_client) -> None: - openai_client.send_diffusion_request({...}) +def test_foo_chat_completions_t2i_fallback(omni_server, online_client) -> None: + online_client.send_diffusion_request({...}) ``` L1 tests under `tests/entrypoints/` may keep **minimal** asserts next to the handler under test when validating a single branch; still prefer `assertions.py` when the same JSON/media check appears in more than one file or level (e.g. L1 protocol test + L4 e2e). diff --git a/.claude/skills/vllm-omni-test/references/test-routing.md b/.claude/skills/vllm-omni-test/references/test-routing.md index 394046b91..5b90957d6 100644 --- a/.claude/skills/vllm-omni-test/references/test-routing.md +++ b/.claude/skills/vllm-omni-test/references/test-routing.md @@ -79,7 +79,7 @@ E2E (L2+) should not rely on mocks unless documenting a rare exception; prefer r - `full_model` — L4 nightly (`test-nightly.yml`); some expansion tests still carry both `advanced_model` and `full_model` during migration 2. **Model type** (pick one): `omni`, `tts`, or `diffusion` 3. **Cross-cutting** when relevant: `parallel`, `cache`, `example`, `benchmark` -4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, … +4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `cards_1`, … 5. Multi-card: `@hardware_test(...)` in `tests/helpers/mark.py` ## Command Templates @@ -244,7 +244,7 @@ pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py ```bash cd tests -pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model +pytest -s -v -m "core_model and L4 and not cards_1" --run-level=core_model ``` ### Concrete e2e paths (common in-tree) @@ -279,7 +279,7 @@ When adding or modifying tests, do not stop at “where the file lives” — al 4. **API client + assert placement**: **General** validation → implement in `assertions.py`, call **inside** `send_*_request` in `runtime.py`; tests call **`send_*_request` only**. **Special** case validation → `assert_*` in `assertions.py`, called **in the test** after `send_*_request`. Low-level `send_*_http_request` is for negative/dfx tests (`err_code`), not ordinary L2+ success e2e. 5. **Shared assertions**: logic in **`tests/helpers/assertions.py`**; general checks inside `send_*_request`; special checks only in the test. No `_assert_*` in `test_*.py`. 6. **One case → one `test_*`**: function name reflects what is validated; no `if case_id == ...` mega-test merging multiple scenarios. -7. **Fixture scope**: default **`omni_server` + `openai_client`** / **`omni_runner` + `omni_runner_handler`** (module). Use **`omni_server_function` + `openai_client_function`** / **`omni_runner_function` + `omni_runner_handler_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`). +7. **Fixture scope**: default **`omni_server` + `online_client`** / **`omni_runner` + `offline_client`** (module). Use **`omni_server_function` + `online_client_function`** / **`omni_runner_function` + `offline_client_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`). 8. **Type marker**: `omni`, `tts`, or `diffusion` on every model e2e module. 9. **Diffusion L4 Function**: wire `*_expansion.py` into **X2I(&A&T)** or **X2V** **Function Test** in **`test-nightly.yml` only** — do not add `test-merge.yml` unless the user also requested L3. 10. **Diffusion L4 Perf** (only when requested): add `tests/dfx/perf/tests/test__vllm_omni.json` + **Perf Test · <Model>** step (artifact upload); not part of “L4 functional cases” by default. diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 2f7f80673..c92c7b22f 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -79,8 +79,6 @@ /vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian /vllm_omni/model_extras/ @tzhouam @gcanlin @princepride /vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi -# Config override within model_executor — docs/design/module/vllm_omni_config.md -/vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks # Execution platforms (hardware backends) — docs/design/module/execution_platforms.md /vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi diff --git a/.github/ISSUE_TEMPLATE/400-bug-report.yml b/.github/ISSUE_TEMPLATE/400-bug-report.yml index 7cbf0ac0e..6e04ae9e7 100644 --- a/.github/ISSUE_TEMPLATE/400-bug-report.yml +++ b/.github/ISSUE_TEMPLATE/400-bug-report.yml @@ -83,7 +83,7 @@ body: omni = Omni( model="Qwen/Qwen-Image", - stage_configs_path="/path/to/stage_configs.yaml", + deploy_config="/path/to/deploy_config.yaml", ) prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}] diff --git a/.github/mergify.yml b/.github/mergify.yml index 4e4c68425..fdba32b17 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -45,7 +45,7 @@ pull_request_rules: > # For mypy (substitute "3.10" with the failing version if needed) > pre-commit run --hook-stage manual mypy-3.10 > # For markdownlint - > pre-commit run --hook-stage manual markdownlint + > pre-commit run markdownlint-cli2 > ``` > diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 2bd0cc133..853609d28 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -24,9 +24,7 @@ jobs: python-version: "3.12" cache: 'pip' - uses: pre-commit/action@v3.0.1 - # Don't run the pytest mark check against everything in the CI - # for now since a lot of files are still missing them. This check - # is used to help contributors remember to add the marks when they - # commit. + # Don't run historical-debt hooks against the whole tree in CI yet. + # They still run locally on changed files. env: - SKIP: check-test-ci-coverage + SKIP: check-mark,markdownlint-cli2,shellcheck,check-spdx-header,mypy-3.10 diff --git a/.gitignore b/.gitignore index 3c24e40ae..542bc5035 100644 --- a/.gitignore +++ b/.gitignore @@ -46,6 +46,8 @@ htmlcov/ .cache nosetests.xml coverage.xml +coverage-*.xml +coverage-*.xml.gz *.cover *.py,cover .hypothesis/ @@ -240,9 +242,9 @@ data/ datasets/ *.csv *.json +!recipes/cosmos3/negative_prompt.json !tests/dfx/perf/tests/*.json !apps/ComfyUI-vLLM-Omni/example_workflows/*.json -!tests/assets/soulxsinger/*.json *.jsonl *.parquet @@ -252,6 +254,7 @@ outputs/ results/ generated/ output_*/ +/tests/e2e/accuracy/artifacts/ # Configuration overrides configs/local.yaml @@ -275,3 +278,6 @@ tests/.ci_generated/ # Local Lance precision-matching dump scripts (not part of the repo) lance_compare/ + +# Linting binaries downloaded by local pre-commit hooks +shellcheck*/ diff --git a/.markdownlint.yaml b/.markdownlint.yaml new file mode 100644 index 000000000..329597772 --- /dev/null +++ b/.markdownlint.yaml @@ -0,0 +1,17 @@ +# Keep these rules aligned with vLLM so MkDocs Material syntax is tolerated. +MD007: + indent: 4 +MD013: false +# GitHub READMEs start with a centered logo, not an H1. +MD041: false +MD024: + siblings_only: true +MD025: + # Allow front matter title to be different from the first heading in the document. + front_matter_title: "" +MD031: + list_items: false +MD033: false +MD046: false +MD052: false +MD059: false diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 491123225..21cf0b9a4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -24,7 +24,7 @@ repos: - id: ruff-format - repo: https://github.com/crate-ci/typos - rev: typos-dict-v0.13.13 + rev: v1.49.0 hooks: - id: typos # only for staged files @@ -38,7 +38,19 @@ repos: - id: actionlint files: ^\.github/workflows/.*\.ya?ml$ + - repo: https://github.com/DavidAnson/markdownlint-cli2 + rev: v0.23.2 + hooks: + - id: markdownlint-cli2 + language_version: lts + args: [--fix] + files: (^docs/|^recipes/|(^|/)README\.md$|^CONTRIBUTING\.md$) + # Agent skill docs and CLAUDE.md are not user-facing docs/READMEs. + exclude: (^|/)CLAUDE\.md$|^\.claude/|^\.cursor/ + # Keep `local` last so `suggestion` stays the final hook. Autofix repos + # (ruff, typos, markdownlint, trailing-whitespace, etc.) must run first; + # otherwise their "files were modified" output displaces the SKIP tip. - repo: local hooks: - id: signoff-commit @@ -54,8 +66,47 @@ repos: verbose: true stages: [commit-msg] - - id: check-test-ci-coverage - name: Ensure test files have pytestmarks + # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup + # is less awkward. Local wrapper groups files so follow_imports=skip on + # tests does not leak into library code. + # mypy-3.10 runs on git commit (changed files). CI SKIPs it because + # --all-files is not clean yet. Extra versions are manual-stage only + # (not in default `pre-commit run --all-files`): + # pre-commit run --hook-stage manual mypy-3.12 + - id: mypy-3.10 + name: Run mypy for Python 3.10 + entry: python tools/pre_commit/mypy.py "3.10" + <<: &mypy_common + language: python + types_or: [python, pyi] + require_serial: true + additional_dependencies: + - mypy==1.11.1 + - regex + - types-cachetools + - types-setuptools + - types-PyYAML + - types-requests + - types-torch + - pydantic + - id: mypy-3.11 + name: Run mypy for Python 3.11 + entry: python tools/pre_commit/mypy.py "3.11" + <<: *mypy_common + stages: [manual] + - id: mypy-3.12 + name: Run mypy for Python 3.12 + entry: python tools/pre_commit/mypy.py "3.12" + <<: *mypy_common + stages: [manual] + - id: mypy-3.13 + name: Run mypy for Python 3.13 + entry: python tools/pre_commit/mypy.py "3.13" + <<: *mypy_common + stages: [manual] + + - id: check-mark + name: Ensure test files have CI marks (no direct SKU) entry: python tools/pre_commit/check_test_marks.py language: python types: [python] @@ -66,12 +117,50 @@ repos: entry: python tools/pre_commit/check_tts_adapter.py language: python pass_filenames: false - files: ^(vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)|tools/pre_commit/check_tts_adapter\.py)$ + files: ^(vllm_omni/entrypoints/openai/serving_speech\.py|tools/pre_commit/check_tts_adapter\.py)$ # The budget-is-now-slack reminder is printed on a PASSING run, and # pre-commit hides output from passing hooks unless the hook is verbose. verbose: true - # Keep `suggestion` last + - id: shellcheck + name: Lint shell scripts + # language: script breaks on Windows (shebang path loses slashes). + # Invoke the vLLM-style wrapper with bash so Linux CI and Win/Git Bash work. + entry: bash + args: [tools/pre_commit/shellcheck.sh] + language: system + types: [shell] + + - id: check-spdx-header + name: Check SPDX headers + entry: python tools/pre_commit/check_spdx_header.py + language: python + types_or: [python, pyi, rust, proto, shell] + files: \.(py|pyi|rs|proto|sh)$ + + - id: check-forbidden-imports + name: Check forbidden imports + entry: python tools/pre_commit/check_forbidden_imports.py + language: python + types: [python] + additional_dependencies: [regex] + + - id: check-torch-cuda-call + name: Prevent new torch.cuda API calls + entry: python tools/pre_commit/check_torch_cuda.py + language: python + types: [python] + additional_dependencies: [regex] + + - id: check-buildkite + name: Validate Buildkite Pipelines + entry: python tools/pre_commit/check_buildkite.py + language: python + files: ^\.buildkite/.*\.ya?ml$ + additional_dependencies: [pyyaml, check-jsonschema==0.38.0] + + # Keep `suggestion` as the final hook. Run after autofix hooks so the + # bypass/SKIP tip is not displaced. - 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=."' @@ -79,10 +168,3 @@ repos: 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] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 4ace89fd6..36cefcf0d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,5 +1,11 @@ # Contributing to vLLM-Omni -You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/) +You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/). + +Local `pre-commit` gates (SPDX, forbidden imports, `torch.cuda`, shellcheck, +mypy, test marks, TTS adapter ratchet, markdownlint, Buildkite schema, and +how to extend allowlists) are documented in the +[Linting](https://vllm-omni.readthedocs.io/en/latest/contributing/#linting) +section of that guide. Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions. diff --git a/README.md b/README.md index 418d0b38d..e5147e857 100644 --- a/README.md +++ b/README.md @@ -12,10 +12,11 @@ Easy, fast, and cheap omni-modality model serving for everyone | Documentation | DeepWiki | User Forum | Developer Slack | WeChat | Paper | Slides |

- --- *Latest News* 🔥 + +- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0). - [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support. - [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support. - [2026/06] Starting with [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0), vLLM-Omni publishes a stable release aligned with every even-numbered upstream vLLM minor version. [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0), [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0), [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0), and [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) continued this cadence, expanding omni and world-model support with [NVIDIA Cosmos3](recipes/cosmos3/Cosmos3-Nano.md) and DreamZero, adding models such as MiniCPM-o 4.5, MOSS-TTS, and Lance, and advancing TTS, diffusion, distributed execution, quantization, RL integration through [VeRL-Omni](https://github.com/verl-project/verl-omni), and CUDA/ROCm/MUSA/NPU/XPU coverage. @@ -88,11 +89,18 @@ If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.or ``` ## 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) + + + + + Star History Chart + + ## License diff --git a/benchmarks/tts/README.md b/benchmarks/tts/README.md index d3a0f9676..c6b6615d5 100644 --- a/benchmarks/tts/README.md +++ b/benchmarks/tts/README.md @@ -18,8 +18,8 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --port 8000 ``` The server auto-loads its Deploy YAML from `vllm_omni/deploy/qwen3_tts.yaml` -(Pipeline + Deploy schema introduced in #2383). No `--stage-configs-path` or -`--deploy-config` flag is needed for any registered model. +(Pipeline + Deploy schema introduced in #2383). No explicit `--deploy-config` +flag is needed for a registered model. ### 2. Run the benchmark (`vllm bench serve --omni`) diff --git a/benchmarks/tts/benchmark_cosyvoice3_trt_streams.py b/benchmarks/tts/benchmark_cosyvoice3_trt_streams.py new file mode 100644 index 000000000..40b4ce135 --- /dev/null +++ b/benchmarks/tts/benchmark_cosyvoice3_trt_streams.py @@ -0,0 +1,151 @@ +"""Benchmark CosyVoice3-style CUDA stream handoff strategies.""" + +import argparse +import json +import statistics +import time + +import torch + +from vllm_omni.platforms import current_omni_platform + + +def run_case( + *, + device: torch.device, + steps: int, + warmup_steps: int, + producer_cycles: int, + estimator_cycles: int, + use_stream_dependencies: bool, +) -> dict[str, float | int | str]: + caller_stream = torch.cuda.Stream(device=device) + estimator_stream = torch.cuda.Stream(device=device) + input_tensor = torch.zeros(1024, device=device) + output_tensor = torch.empty_like(input_tensor) + submission_ms = [] + + def step(measure: bool) -> None: + with torch.cuda.stream(caller_stream): + torch.cuda._sleep(producer_cycles) + input_tensor.add_(1) + + start = time.perf_counter_ns() + if use_stream_dependencies: + estimator_stream.wait_stream(caller_stream) + else: + caller_stream.synchronize() + + with torch.cuda.stream(estimator_stream): + torch.cuda._sleep(estimator_cycles) + output_tensor.copy_(input_tensor) + + if use_stream_dependencies: + caller_stream.wait_stream(estimator_stream) + else: + estimator_stream.synchronize() + + if measure: + submission_ms.append((time.perf_counter_ns() - start) / 1e6) + + for _ in range(warmup_steps): + step(measure=False) + caller_stream.synchronize() + estimator_stream.synchronize() + input_tensor.zero_() + torch.accelerator.reset_peak_memory_stats(device) + initial_allocated = torch.accelerator.memory_allocated(device) + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + with torch.cuda.stream(caller_stream): + start_event.record() + for _ in range(steps): + step(measure=True) + with torch.cuda.stream(caller_stream): + end_event.record() + end_event.synchronize() + wall_ms = (time.perf_counter() - wall_start) * 1e3 + peak_delta_bytes = torch.accelerator.max_memory_allocated(device) - initial_allocated + + expected = torch.full_like(output_tensor, steps) + torch.testing.assert_close(output_tensor, expected, rtol=0, atol=0) + return { + "mode": "wait_stream" if use_stream_dependencies else "synchronize", + "steps": steps, + "submission_median_ms": statistics.median(submission_ms), + "submission_p95_ms": statistics.quantiles(submission_ms, n=20)[18], + "submission_mean_ms": statistics.mean(submission_ms), + "gpu_elapsed_ms": start_event.elapsed_time(end_event), + "wall_elapsed_ms": wall_ms, + "peak_allocation_delta_bytes": peak_delta_bytes, + "final_value": int(output_tensor[0].item()), + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--steps", type=int, default=100) + parser.add_argument("--warmup-steps", type=int, default=10) + parser.add_argument("--repeats", type=int, default=7) + parser.add_argument("--producer-cycles", type=int, default=1_000_000) + parser.add_argument("--estimator-cycles", type=int, default=2_000_000) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + device = torch.device(args.device) + current_omni_platform.set_device(device) + torch.manual_seed(0) + + old_results = [] + new_results = [] + for repeat in range(args.repeats): + common_args = { + "device": device, + "steps": args.steps, + "warmup_steps": args.warmup_steps, + "producer_cycles": args.producer_cycles, + "estimator_cycles": args.estimator_cycles, + } + if repeat % 2 == 0: + old_results.append(run_case(**common_args, use_stream_dependencies=False)) + new_results.append(run_case(**common_args, use_stream_dependencies=True)) + else: + new_results.append(run_case(**common_args, use_stream_dependencies=True)) + old_results.append(run_case(**common_args, use_stream_dependencies=False)) + + old_submission_ms = statistics.median(result["submission_median_ms"] for result in old_results) + new_submission_ms = statistics.median(result["submission_median_ms"] for result in new_results) + old_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in old_results) + new_gpu_ms = statistics.median(result["gpu_elapsed_ms"] for result in new_results) + print( + json.dumps( + { + "device": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "repeats": args.repeats, + "warmup_steps": args.warmup_steps, + "producer_cycles": args.producer_cycles, + "estimator_cycles": args.estimator_cycles, + "old_submission_median_ms": old_submission_ms, + "new_submission_median_ms": new_submission_ms, + "submission_median_speedup": old_submission_ms / new_submission_ms, + "old_gpu_elapsed_median_ms": old_gpu_ms, + "new_gpu_elapsed_median_ms": new_gpu_ms, + "gpu_elapsed_change_percent": (new_gpu_ms - old_gpu_ms) / old_gpu_ms * 100, + "parity": all(result["final_value"] == args.steps for result in old_results + new_results), + "old_runs": old_results, + "new_runs": new_results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/collect_env.py b/collect_env.py index 8b09379e1..6a172b0f0 100644 --- a/collect_env.py +++ b/collect_env.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # ruff: noqa # code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py @@ -494,6 +494,16 @@ def is_xnnpack_available(): def get_env_vars(): + try: + from vllm_omni.config.environment_variable_inventory import ENVIRONMENT_VARIABLE_INVENTORY + except Exception: + # Keep environment collection usable when an incomplete/broken Omni + # installation is the reason the user is running this diagnostic. + # Importing an Omni submodule executes package initialization, which + # may fail with RuntimeError, AssertionError, or OSError as well as + # ImportError when vLLM or the accelerator stack is broken. + ENVIRONMENT_VARIABLE_INVENTORY = {} + env_vars = "" secret_terms = ("secret", "token", "api", "access", "password") report_prefix = ( @@ -507,12 +517,17 @@ def get_env_vars(): "MKL_", "NVIDIA", ) - for k, v in os.environ.items(): - if any(term in k.lower() for term in secret_terms): + for k, v in sorted(os.environ.items()): + omni_classification = ENVIRONMENT_VARIABLE_INVENTORY.get(k) + if (omni_classification is not None and omni_classification.redact_value) or 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): + if ( + k in environment_variables + or (omni_classification is not None and omni_classification.is_public_omni) + or k.startswith(report_prefix) + ): env_vars = env_vars + "{}={}".format(k, v) + "\n" return env_vars diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index a7e6ccad0..b84b344a8 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,4 +1,4 @@ -# Target vllm version: v0.27.0 (releases/v0.27.0 branch). +# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5). # Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase. # # Release-boundary form: there is no precompiled-wheel pin here, because the @@ -14,11 +14,14 @@ # core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy # under numba's ceiling. All of it was compensation for a base image and a vLLM # that disagreed. At a real release boundary they agree, so none of it is -# needed. If you pin an untagged commit again, restore the reinstall block — -# agent/tasks/20_update_wheel_url.sh fails the run rather than letting the -# Dockerfile silently ship the wrong vLLM. +# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels +# rode on the v0.27.0 base image with the full reinstall block until the +# official v0.28.0 image shipped, at which point this file dropped back to the +# form you are reading. If you pin an untagged commit again, restore the +# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather +# than letting the Dockerfile silently ship the wrong vLLM. ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.27.0 +ARG VLLM_BASE_TAG=v0.28.0 # We use registry-based layer caching for the dependencies, # where the tag is determined based on the hash of the contents diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index f2e80eaf3..3881ca2e5 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0 +ARG BASE_IMAGE=vllm/vllm-openai:v0.28.0 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index eaeef693c..3ebacdc98 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.27.0 +ARG VLLM_ASCEND_TAG=v0.28.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm @@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \ "stepaudio2-minicpmo>=0.1.1" \ "step-audio2>=1.0.0" +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index 2dc6334d9..d341d456d 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.27.0-a3 +ARG VLLM_ASCEND_TAG=v0.28.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm @@ -30,6 +30,7 @@ RUN python3 -m pip install -v --no-deps \ "stepaudio2-minicpmo>=0.1.1" \ "step-audio2>=1.0.0" +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index f5af5fb31..029a905fd 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.27.0 +ARG VLLM_ASCEND_TAG=v0.28.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni @@ -65,6 +65,7 @@ RUN python3 -m pip install -v --no-deps \ --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ --default-timeout=60 +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 index 0529cb2da..fe3741d04 100644 --- a/docker/Dockerfile.npu.ci.a3 +++ b/docker/Dockerfile.npu.ci.a3 @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.27.0-a3 +ARG VLLM_ASCEND_TAG=v0.28.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni @@ -66,6 +66,7 @@ RUN python3 -m pip install -v --no-deps \ --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ --default-timeout=60 +ENV VLLM_TARGET_DEVICE=npu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn ENTRYPOINT [] diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 6b5a12309..597f5d134 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0 FROM ${BASE_IMAGE} AS base # Declare a variable to know if we want to use the nightly build or the stable build. @@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base # we should swap over to use stable release ASAP. # We should avoid relying on custom commits. ARG USE_NIGHTLY_BUILD=0 -ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0 +ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0 ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} @@ -79,9 +79,16 @@ RUN printf '%s\n' \ ' os._exit(_exit_code)' \ > ${COMMON_WORKDIR}/vllm-omni/conftest.py +# Some optional dependencies pull in the CPU ONNX Runtime distribution. The +# CPU and ROCm wheels own the same import path and cannot coexist, so remove +# the CPU wheel and reinstall the ROCm wheel after dependency resolution. RUN cd ${COMMON_WORKDIR}/vllm-omni && \ uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \ --no-cache-dir ".[dev]" --no-build-isolation && \ + python3 -m pip uninstall -y onnxruntime && \ + python3 -m pip install --no-cache-dir --force-reinstall --no-deps \ + "onnxruntime-rocm>=1.22.2" && \ + python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \ python3 -c "from torchcodec.decoders import VideoDecoder" RUN ln -sf /usr/bin/python3 /usr/bin/python diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 65ae05dd3..acfa71588 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -1,140 +1,33 @@ -# Argument to configure vllm base image if pre-built -ARG VLLM_BASE=vllm-base - -FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base - -WORKDIR /workspace/ - -ARG PYTHON_VERSION=3.12 -ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu" - -RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ - echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list - -RUN apt clean && apt-get update -y && \ - apt-get install -y --no-install-recommends --fix-missing \ - curl \ - ffmpeg \ - git \ - libsndfile1 \ - libsm6 \ - libxext6 \ - libgl1 \ - lsb-release \ - libaio-dev \ - numactl \ - wget \ - vim \ - python3.12 \ - python3.12-dev \ - python3-pip - -# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer. -RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \ - echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \ - printf '%s\n' \ - 'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \ - 'Pin: version 2025.3*' \ - 'Pin-Priority: 1001' \ - > /etc/apt/preferences.d/oneapi-2025.3.pref && \ - apt-get update -y && \ - apt-get install -y --no-install-recommends \ - intel-oneapi-compiler-dpcpp-cpp-2025.3 \ - intel-oneapi-mkl-devel-2025.3 \ - intel-oneapi-dnnl-devel-2025.3 && \ - rm -rf /var/lib/apt/lists/* - -# Install UMD with fixed version -RUN mkdir neo && \ - cd neo && \ - wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \ - wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \ - wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \ - wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \ - dpkg -i *.deb && \ - cd .. && \ - rm -rf neo - -ENV PATH="/root/.local/bin:$PATH" -ENV VIRTUAL_ENV="/opt/venv" -ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python -RUN curl -LsSf https://astral.sh/uv/install.sh | sh -RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV} -ENV PATH="$VIRTUAL_ENV/bin:$PATH" - -# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3. -ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh" -RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \ - bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \ - rm "${ONECCL_INSTALLER}" && \ - echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \ - echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc -RUN rm -f /opt/intel/oneapi/ccl/latest && \ - ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest - -SHELL ["bash", "-c"] -CMD ["bash", "-c", "source /root/.bashrc && exec bash"] - -WORKDIR /workspace/ - -ENV UV_HTTP_TIMEOUT=500 - -# Configure package index for XPU -ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} -ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} -ENV UV_INDEX_STRATEGY="unsafe-best-match" -ENV UV_LINK_MODE="copy" - -ARG VLLM_VERSION=v0.27.0 -RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm -WORKDIR /workspace/vllm - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --upgrade pip && \ - uv pip install -r requirements/xpu.txt && \ - uv pip install setuptools_scm &&\ - uv pip install grpcio-tools protobuf nanobind && \ - source /opt/intel/oneapi/setvars.sh --force && \ - source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \ - export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" - -ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/" - -ENV VLLM_TARGET_DEVICE=xpu -ENV VLLM_WORKER_MULTIPROC_METHOD=spawn - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip uninstall triton triton-xpu && \ - uv pip install triton-xpu==3.7.2 && \ - uv pip install --no-build-isolation --no-deps . - -CMD ["/bin/bash"] - -FROM vllm-base AS vllm-openai - -# install nixl from source code -RUN uv pip install nixl - -# install additional dependencies for openai api server -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope - -# install development dependencies (for testing) -RUN uv pip install -e tests/vllm_test_utils - -# ensure vllm is properly installed -RUN python -c "import vllm, inspect; print(vllm.__file__)" -RUN uv pip show vllm - -CMD ["/bin/bash"] - -ENTRYPOINT [] +# Target vLLM version: v0.28.0. Keep VLLM_VERSION aligned with the target vLLM +# release on every rebase. +# +# This file layers vLLM-Omni on top of the published upstream XPU image, which +# already ships a self-consistent XPU runtime (torch-xpu, oneCCL/UCX+NIXL, the +# pinned UMD/compute-runtime stack, and vLLM itself) built from +# https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu. Docker +# pulls that base automatically, so there is nothing to build here beyond the +# Omni layers. Mirrors how docker/Dockerfile.ci consumes vllm/vllm-openai. +# +# If the tag does not exist yet (targeting an unreleased vLLM commit), build the +# base from upstream's own Dockerfile and point VLLM_BASE at the result: +# +# git clone --depth 1 --branch $VLLM_VERSION https://github.com/vllm-project/vllm /tmp/vllm +# docker build -t vllm-openai-xpu:local --target vllm-openai \ +# -f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm +# docker build -f docker/Dockerfile.xpu --build-arg VLLM_BASE=vllm-openai-xpu:local . +# +# .buildkite/intel/scripts/run-xpu-test.sh automates that pull-then-build fallback. +ARG VLLM_BASE_IMAGE=vllm/vllm-openai-xpu +ARG VLLM_VERSION=v0.28.0 +ARG VLLM_BASE=${VLLM_BASE_IMAGE}:${VLLM_VERSION} FROM ${VLLM_BASE} AS vllm-omni +# The base image ships ENTRYPOINT ["vllm", "serve"]; clear it so this stage stays +# usable as a plain shell/test container regardless of which VLLM_BASE was used. +ENTRYPOINT [] +CMD ["/bin/bash"] + RUN apt-get update && \ apt-get install -y espeak-ng ffmpeg git jq && \ apt-get clean && \ diff --git a/docs/.nav.yml b/docs/.nav.yml index f349e6fe6..94c4de213 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -32,7 +32,6 @@ nav: - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md - X-To-Text: user_guide/examples/offline_inference/x_to_text.md - - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md - Online Serving: - Image-To-Image: user_guide/examples/online_serving/image_to_image.md - Image-To-Video: user_guide/examples/online_serving/image_to_video.md @@ -51,6 +50,7 @@ nav: - Overview: features/README.md - Runtime and Stage Execution: - Execution Modes and Streaming: user_guide/diffusion/execution_modes.md + - Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md - Sleep Mode: features/sleep_mode.md - Quantization: - Overview: user_guide/quantization/overview.md @@ -61,6 +61,7 @@ nav: - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md - MXFP8 W8A8: user_guide/quantization/mxfp8.md - MXFP4 W4A4: user_guide/quantization/mxfp4.md + - SVDQuant W4A4: user_guide/quantization/svdquant.md - ModelOpt: user_guide/quantization/modelopt.md - AutoRound: user_guide/quantization/autoround.md - msModelSlim: user_guide/quantization/msmodelslim.md @@ -68,7 +69,11 @@ nav: - Overview and Compatibility: - Overview: user_guide/diffusion_features.md - Feature Compatibility: user_guide/feature_compatibility.md - - CPU Offloading: user_guide/diffusion/cpu_offload.md + - CPU Offloading: + - Overview: user_guide/diffusion/cpu_offload.md + - Model-Level: user_guide/diffusion/offloader/cpu_offload.md + - Layerwise: user_guide/diffusion/offloader/layerwise_offload.md + - Distributed Layerwise: user_guide/diffusion/offloader/distributed_layerwise_offload.md - Cache Acceleration: - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md @@ -80,7 +85,13 @@ nav: - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md - - Attention Backends: user_guide/diffusion/attention_backends.md + - Attention Backends: + - Overview: user_guide/diffusion/attention_backends.md + - Dense Backends: user_guide/diffusion/attention_backends/dense_backends.md + - TRTLLM: user_guide/diffusion/attention_backends/trtllm.md + - SageAttention: user_guide/diffusion/attention_backends/sage.md + - Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md + - RainFusion: user_guide/diffusion/attention_backends/rainfusion.md - Regional Compilation: user_guide/diffusion/regional_compilation.md - Frame Interpolation: user_guide/diffusion/frame_interpolation.md - Startup and Loading: user_guide/diffusion/startup_and_loading.md @@ -110,10 +121,12 @@ nav: - Feature Design: - Runtime and Stage Execution: - Disaggregated Inference: design/feature/disaggregated_inference.md + - Host Weight Runtime: design/feature/host_weight_runtime.md - Async Chunk: design/feature/async_chunk.md - Async Diffusion Output: design/feature/async_diffusion_output.md - Async Omni Output Materialization: design/feature/omni_async_output_materialization.md - Automatic Prefix Caching: design/feature/prefix_caching.md + - Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md - Communication: - OmniConnector Implementations: - Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md @@ -133,11 +146,16 @@ nav: - Tensor Parallel: design/feature/tensor_parallel.md - VAE Patch Parallelism: design/feature/vae_parallel.md - Attention Optimization: + - Backend Selection: design/feature/attention_backend_selection.md - Skip-Softmax: design/feature/skip_softmax.md + - CPU Offloading: + - Overview: design/feature/offloader/README.md + - Model-Level: design/feature/offloader/cpu_offload.md + - Layerwise: design/feature/offloader/layerwise_offload.md + - Distributed Layerwise: design/feature/offloader/distributed_layerwise_offload.md - Cache-DiT: design/feature/cache_dit.md - TeaCache: design/feature/teacache.md - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md - - Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md - Infrastructure and Performance: - Prometheus Metrics: design/metrics.md - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md @@ -160,6 +178,7 @@ nav: - Offloader: design/module/diffusion/offloader.md - Execution Platforms: design/module/execution_platforms.md - Cache Management: design/module/cache_management.md + - Host Weight Runtime: design/module/host_weight_runtime.md - Quantization: design/module/quantization.md - Observability: design/module/observability.md - Profiling: design/module/profiling.md diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index 04818fe04..e5a763192 100644 Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ diff --git a/docs/cli/README.md b/docs/cli/README.md index 40858ca1a..88620acd8 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -2,6 +2,11 @@ The CLI for vLLM-Omni inherits from vllm with some additional arguments. +Environment variables are process configuration rather than CLI-only options. +See the [Environment Variables](../configuration/environment_variables.md) +reference for ownership, precedence, evaluation time, and per-stage `env` +scoping. + ## serve Starts the vLLM-Omni OpenAI Compatible API server. @@ -39,4 +44,5 @@ vllm bench serve --omni \ --num-prompts 5 ``` -See [vllm bench serve](./bench/serve.md) for the full reference of all available arguments. +See [vllm bench serve](./bench/serve.md) for serving benchmark arguments and +dataset-specific examples, including OmniInteract native-duplex sessions. diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md index 2a1acd1f0..3e16ae554 100644 --- a/docs/cli/bench/serve.md +++ b/docs/cli/bench/serve.md @@ -2,7 +2,7 @@ The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models. ## Notes -vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends. +vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` serving benchmark backends. It also adds the `omniinteract` dataset. ## Basic Parameter Description You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below: @@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio. - `--backend` - Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` to the upstream vLLM backend choices. + Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, `daily-omni`, and `openai-realtime-duplex` to the upstream vLLM backend choices. - `--model` The model identifier to load, filled according to the models supported by vLLM-Omni. @@ -19,7 +19,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para The API endpoint exposed externally, to which clients send their requests. - `--dataset-name` - The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio). + The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio), while `omniinteract` replays official OmniInteract videos. - `--num-prompts` The total number of requests to send, an integer. @@ -269,6 +269,44 @@ We use audio generation time / audio duration to calculate RTF. +### OmniInteract Realtime Benchmark + +OmniInteract runs each video as one native-duplex WebSocket sample in the standard serving-benchmark lifecycle: + +```bash +vllm bench serve --omni \ + --backend openai-realtime-duplex \ + --dataset-name omniinteract \ + --dataset-path /path/to/OmniInteract \ + --model openbmb/MiniCPM-o-4_5 \ + --base-url http://127.0.0.1:8000 \ + --endpoint /v1/realtime \ + --omniinteract-ref-audio /path/to/reference.wav \ + --omniinteract-output-dir ./omniinteract-artifacts \ + --num-prompts 3 \ + --num-warmups 0 +``` + +`--dataset-path` accepts an extracted directory, `data.tar[.gz]`, or a Hugging Face dataset ID; omitting it uses +`lucky-lance/OmniInteract`. `--num-prompts` is the total across subsets and defaults to 3 for OmniInteract; explicit `0` +selects all and oversize values use all available cases. Reference audio is required, and OmniInteract uses the +`/v1/realtime` endpoint. + +Audio is replayed as 16 kHz PCM16 in 200 ms chunks and video at 1 FPS with real-time pacing. All selected media is decoded +before timing and remains in client memory for the run, so `--max-concurrency` does not limit media preparation memory; use +explicit `--num-prompts 0` only when the client has enough RAM for the full dataset. Media commands are bounded by +`--omniinteract-media-timeout-s`, and concurrency defaults to 1. Standard request-rate, warmup, result-saving, and summary +options apply. Use `--omniinteract-require-response` only for functional E2E cases; LISTEN is a valid benchmark result. + +Each completed case writes `output.wav`, `wav_transcript.json`, `events.json`, `result.json`, and a final `.done` marker under +`--omniinteract-output-dir`. The root also contains `batch_summary.json` and `official_eval_manifest.jsonl`; failed cases write +`.failed.json`. Runs sharing one output root are serialized. Completion validates transport, response lifecycle, and artifacts, +not answer accuracy. Transcript timestamps are serialized playback-queue times. Clipped or cancelled outputs are ineligible and +omitted from the official manifest; `audio_clipped_bytes` records output beyond the rounded video horizon. + +TTFT, TTFP, and RTF start at client receipt of `response.created`. TPOT/ITL use engine stage-0 timing; ITL is emitted only when +every token interval is present. + ### Multi-Modal Benchmark
diff --git a/docs/configuration/README.md b/docs/configuration/README.md index 09772ccd1..e68668e06 100644 --- a/docs/configuration/README.md +++ b/docs/configuration/README.md @@ -4,20 +4,20 @@ This section lists the most common options for running vLLM-Omni. For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration guide](https://docs.vllm.ai/en/v0.26.0/configuration/index.html). -Currently, the main options are maintained by stage configs for each model. +Each model defines fixed topology in a registered `PipelineConfig` and runtime +overrides in a deploy YAML. + +For process-level settings shared by the CLI, deploy YAML, workers, and model +integrations, see [Environment Variables](./environment_variables.md). For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/deploy/qwen2_5_omni.yaml). The matching frozen pipeline topology lives at [vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py](gh-file:vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py). -For introduction, please check [Introduction for stage config](./stage_configs.md) +For an introduction, see [Pipeline and deploy configurations](./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 -## Multi-Stage Recipes - -- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML - ## Optimization Features - **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models diff --git a/docs/configuration/composable_parallel.md b/docs/configuration/composable_parallel.md index 473605c69..00af4cad4 100644 --- a/docs/configuration/composable_parallel.md +++ b/docs/configuration/composable_parallel.md @@ -1,6 +1,6 @@ # Composable parallel strategies -In vLLM-Omni, the `composable_parallel` system layers a declarative *parallel strategy* on top of the existing stage configs. A strategy YAML maps each pipeline stage (by its `model_stage` name) to a stack of axis declarations — one per parallelism dimension — that the engine translates into per-stage parallel sizing (`tensor_parallel_size`, `data_parallel_size`, `pipeline_parallel_size`, expert-parallel toggle, replica count) and, when relevant, an orchestrator-wide load-balancer policy. +In vLLM-Omni, the `composable_parallel` system layers a declarative *parallel strategy* on top of the registered pipeline and deploy config. A strategy YAML maps each pipeline stage (by its `model_stage` name) to a stack of axis declarations — one per parallelism dimension — that the engine translates into per-stage parallel sizing (`tensor_parallel_size`, `data_parallel_size`, `pipeline_parallel_size`, expert-parallel toggle, replica count) and, when relevant, an orchestrator-wide load-balancer policy. Strategies are intentionally narrow: each one only writes the axes it declares, leaving every other engine arg untouched. This makes a strategy file a small, reusable overlay you can swap independently of the deploy YAML (see `configuration/stage_configs.md`). @@ -195,7 +195,7 @@ Because `omni_lb_policy` is read once at `AsyncOmniEngine` construction — not Default: no default — required for every entry under `strategies:`. -Strategy keys MUST be `model_stage` *names* (strings), the same labels stage configs use for `engine_args.model_stage` (see [`model_stage` in `stage_configs.md`](stage_configs.md#engine_argsmodel_stage)). Integer `stage_id` values are not accepted; the resolver matches on `model_stage` exactly (`_resolve_stage` in `vllm_omni/config/composable_parallel/apply.py`). +Strategy keys MUST be `model_stage` *names* (strings), the same labels defined by [`StagePipelineConfig.model_stage`](stage_configs.md#pipeline-configuration). Integer `stage_id` values are not accepted; the resolver matches on `model_stage` exactly (`_resolve_stage` in `vllm_omni/config/composable_parallel/apply.py`). ## Common pitfalls diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md new file mode 100644 index 000000000..ace897048 --- /dev/null +++ b/docs/configuration/environment_variables.md @@ -0,0 +1,238 @@ +# Environment Variables + +Environment variables are process-level configuration shared by the CLI, +deploy YAML, Python APIs, worker processes, and model integrations. They are +therefore documented under **User Guide → Configuration**, rather than under +CLI Reference. + +This page describes the public vLLM-Omni environment-variable contract on the +`latest` documentation branch. Use the documentation for your installed +vLLM-Omni release when behavior differs. + +!!! important "Public contract" + Only names in [Public vLLM-Omni variables](#public-vllm-omni-variables) are + public Omni configuration. A name found in source code is not automatically + supported: model-specific, benchmark, platform, and internal variables have + separate ownership and lifecycle rules. + +## Naming and implementation boundary + +New public variables owned by vLLM-Omni must use the `VLLM_OMNI_` prefix. +Existing public names with older prefixes remain supported for compatibility, +but they do not establish naming precedent for new settings. Inherited vLLM +and third-party names retain the prefix chosen by their owning project. + +The source [environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py) +records ownership, migration disposition, and redaction metadata. Unlike +`vllm.envs`, it is not an executable value resolver: existing consumers still +own parsing, defaults, precedence, and evaluation time. The public tables below +describe those live consumer contracts. Stable model or stage behavior should +move to typed configuration, and request-varying behavior should move to a +request schema instead of adding another environment variable. + +## Precedence and evaluation time + +There is no repository-wide rule that makes environment variables override (or +lose to) every CLI, YAML, or request option. The tables below state precedence +for each public setting. In general: + +- request fields should control behavior that can vary per request; +- typed deploy or stage configuration should control stable model and stage + behavior; +- environment variables are appropriate for process launch, platform + integration, emergency diagnostics, and compatibility fallbacks. + +An environment variable can be read while a module is imported, when server or +model configuration is constructed, when a worker is spawned, or whenever a +property is accessed. Set import- and startup-time variables before invoking +`vllm serve` or importing vLLM-Omni. Changing the parent shell after workers +start does not update existing worker environments. + +The lifecycle labels used below are: + +- **Stable** — reviewed public configuration; changes require compatibility + consideration. +- **Experimental** — public for evaluation, but its contract may still evolve. +- **Diagnostic** — supported for investigation or temporary mitigation, not as + a normal tuning interface. +- **Deprecated** — compatibility alias; migrate to the named replacement. + +## Public vLLM-Omni variables + +### Diffusion backend and cache selection + +| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | +| --- | --- | --- | --- | --- | +| `DIFFUSION_ATTENTION_BACKEND` | Backend name or `auto`; default `auto` (platform selection) | Diffusion stages; read when `OmniDiffusionConfig` is constructed | `diffusion_attention_config.default` wins. An unknown backend fails during backend resolution. | Stable fallback | +| `DIFFUSION_CACHE_BACKEND` | `none`, `cache_dit`, `tea_cache`, `mag_cache`, `step_cache`, `stepcache`, or `step_cache_dit`; default `none` | Diffusion runner startup | Explicit `cache_backend` in config wins. Otherwise this name wins over the deprecated alias. Unsupported values raise `ValueError` during runner setup. | Stable fallback | +| `DIFFUSION_CACHE_ADAPTER` | Same values as `DIFFUSION_CACHE_BACKEND`; default `none` | Diffusion runner startup | Used only when neither explicit `cache_backend` nor `DIFFUSION_CACHE_BACKEND` is set. Unsupported values raise `ValueError`. | Deprecated; use `DIFFUSION_CACHE_BACKEND` | +| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE` | Boolean: `1`, `true`, `yes`, `on`, `0`, `false`, `no`, or `off`; default disabled | Each diffusion runner; resolved during model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental | +| `OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE` | Positive integer; default `32` entries | Each diffusion runner; resolved during model setup | A valid environment value overrides the explicit cache size. A non-integer logs a warning; a non-positive value is ignored. | Experimental | +| `OMNI_DIFFUSION_SESSION_STATE_MANAGER` | Boolean with the same accepted spellings as the prompt cache; default disabled | Experimental diffusion session manager; model setup | A recognized environment value overrides the explicit enable setting. An unrecognized value is ignored. | Experimental | +| `OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS` | Positive integer; default `64` | Experimental diffusion session manager; model setup | A valid environment value overrides the explicit maximum. A non-integer logs a warning; a non-positive value is ignored. | Experimental | + +Backend names for `DIFFUSION_ATTENTION_BACKEND` are the members of +`DiffusionAttentionBackendEnum`, such as `FLASH_ATTN`, `TORCH_SDPA`, +`SAGE_ATTN`, `FLASHINFER_ATTN`, and `TRTLLM_ATTN`. Platform support still +depends on the installed kernels and model path. + +### Serving and runtime + +| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | +| --- | --- | --- | --- | --- | +| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable | +| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable | +| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental | +| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental | +| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control | +| `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `/vllm_omni_orch_monitor_.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic | +| `VLLM_OMNI_VIDEO_SYNC_TIMEOUT` | Float seconds; default `600` | Synchronous Videos API; read when the API server module imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | +| `VLLM_VIDEO_ASYNC_CHUNK` | `on` or `off`; default `on` | Streaming video output; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `on`. | Experimental | +| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast` or `slow`; default `fast` | Streaming video audio deltas; read on attribute access | Environment-only setting. Values are trimmed and case-normalized; an invalid value warns once and uses `fast`. | Experimental | + +### Server storage + +Storage names use Pydantic's nested-settings delimiter (`__`). They are read +when `ServerSettings` is initialized, normally during server import/startup. +Invalid integer values raise a Pydantic validation error. + +| Name | Type and default | Meaning | Lifecycle | +| --- | --- | --- | --- | +| `VLLM_OMNI_SERVER_STORAGE__PATH` | Path; default `/tmp/storage` | Directory for completed server files. | Stable | +| `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY` | Integer; default `4` | Maximum concurrent file operations. No additional range validation is currently applied. | Stable | +| `VLLM_OMNI_SERVER_STORAGE__FILE_TTL` | Integer seconds or unset; default unset | Optional lifetime for locally stored files. | Stable | +| `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL` | Integer seconds or unset | Sweep frequency. If file TTL is set and this is unset, the effective default is `300`. | Stable | +| `VLLM_OMNI_STORAGE_PATH` | Path | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__PATH`. The new name wins when both are set. | Deprecated | +| `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` | Integer | Deprecated alias for `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`. The new name wins when both are set. | Deprecated | + +The deprecated aliases emit `DeprecationWarning`. They do not affect file TTL +settings. + +### Quantization diagnostics and performance + +| Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | +| --- | --- | --- | --- | --- | +| `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | +| `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control | + +`QUACK_CACHE_DIR` is owned by the external quack library and is not an +Omni-owned variable, even though vLLM-Omni supplies a persistent default when +it is unset. + +## Per-stage environment + +Deploy configurations can set arbitrary environment keys for one stage: + +```yaml +stages: + - stage_id: 0 + env: + HCCL_IF_BASE_PORT: 23000 + MY_MODEL_CACHE: /mnt/model-cache +``` + +The legacy nested form under `runtime.env` is also accepted. vLLM-Omni converts +keys and values to strings, applies them while the stage engine or child process +is launched, logs keys but not values, and restores the parent process's prior +environment afterward. A child process inherits the scoped values. This keeps +sibling stages from permanently changing one another's parent environment. + +This scoping has two consequences: + +1. A setting read inside stage/model initialization can see the stage value. +2. A setting already cached by an import that occurred before stage launch + cannot be changed retroactively with stage `env`. + +Stage `env` is an escape hatch for process and third-party settings, not a +replacement for typed stage configuration. Do not put request-varying behavior +there. Secret values may reach child processes; logs and diagnostics must emit +their keys only. + +## Inherited vLLM variables + +vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to +the [vLLM 0.26 environment-variable reference](https://docs.vllm.ai/en/v0.26.0/configuration/env_vars.html) +for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm, +XPU, ModelScope, and FlashInfer workspace settings. + +An inherited name does not become Omni-owned because Omni reads or forwards it. +Where behavior differs, the Omni-specific page should document only that +difference. For example, device-visibility values can be further narrowed by a +stage's `devices` setting. + +## Platform, external, and secret variables + +Hardware runtimes, distributed launchers, and third-party libraries retain +ownership of their variables. Common groups include: + +- device visibility (`ASCEND_RT_VISIBLE_DEVICES`, `HIP_VISIBLE_DEVICES`, and + `MUSA_VISIBLE_DEVICES`); +- distributed launch (`MASTER_ADDR`, `MASTER_PORT`, `RANK`, and `WORLD_SIZE`); +- Hugging Face cache and authentication (`HF_HOME`, `HF_TOKEN`, and related + names); +- accelerator libraries such as NCCL, FlashInfer, MindIE, MORI, and quack. + +Consult the owning project's documentation for accepted values and support. +`HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, and `OPENAI_API_KEY` are explicitly marked +for redaction in the reviewed inventory. Never print their values in environment +collection, examples, bug reports, or logs. + +## Model-specific variables + +The audit found 56 variables read by a single model or pipeline family. They are +not listed as public usage options here because doing so would turn implementation +escape hatches into an accidental compatibility contract. + +Every audited model-specific name has a migration disposition in the +[environment-variable inventory](gh-file:vllm_omni/config/environment_variable_inventory.py): + +| Disposition | Count | Required outcome | +| --- | ---: | --- | +| Promote | 33 | Move a stable setting into typed stage or model configuration. | +| Request scope | 6 | Move request-varying behavior into a declared request-option schema. | +| External | 0 | Retain only when a supported third-party library owns the contract. | +| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. | +| Deprecate/remove | 6 | Remove a compatibility escape hatch that has no continuing contract. | + +The disposition is a migration target, not a statement that the existing +environment switch is stable. Promote or request-scope work should land in +model-owner-reviewed follow-up changes with validation and model-page updates. +The family-sized migration packages and their ownership gate are tracked in +[#6232](https://github.com/vllm-project/vllm-omni/issues/6232). + +## Environment collection and troubleshooting + +Run the repository collector when preparing a bug report: + +```bash +python collect_env.py +``` + +The environment section includes the safe public Omni allowlist, registered +vLLM variables, and selected platform prefixes. It omits names marked for +redaction and names containing common secret terms. Review the output before +sharing it. A key supplied through stage `env` is not collected merely because +it appears in deploy configuration; the same allowlist and prefix rules apply. + +If a setting appears ineffective, first check its read time in the tables +above. Restart processes after changing import- or startup-time values, verify +that the variable reaches the intended worker, and check whether an explicit +typed configuration field takes precedence. + +## Benchmark and internal variables + +The reviewed inventory marks 20 benchmark variables as transitional while +benchmark configuration moves to CLI/request schemas in +[issue #5376](https://github.com/vllm-project/vllm-omni/issues/5376). They are not +part of the serving/runtime contract. + +Internal process-coordination and debug variables are deliberately excluded +from this page. The CI scanner resolves string literals and module-level string +constants used by direct `os` access, local wrapper helpers, aliased `os` +imports, and environment membership checks. Server-storage names generated by +Pydantic settings are checked separately against their model fields. + +Names assembled from cross-module metadata and arbitrary keys supplied through +stage `env` cannot be inferred reliably. Contributors must classify those names +explicitly. CI fails when a statically resolved name has no classification or +when a reviewed public name is absent from this page. diff --git a/docs/configuration/gpu_memory_utilization.md b/docs/configuration/gpu_memory_utilization.md index 741066035..9f0e7439a 100644 --- a/docs/configuration/gpu_memory_utilization.md +++ b/docs/configuration/gpu_memory_utilization.md @@ -80,19 +80,15 @@ When multiple stages share the same GPU, you must ensure the sum of their `gpu_m **Example: Two stages on GPU 0** ```yaml -stage_args: +stages: - stage_id: 0 - runtime: - devices: "0" - engine_args: - gpu_memory_utilization: 0.6 # Uses 60% of GPU 0 + devices: "0" + 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) + devices: "0" + 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. @@ -103,14 +99,12 @@ When using `tensor_parallel_size > 1`, the model is split across multiple GPUs, **Example: 2-way tensor parallelism** ```yaml -stage_args: +stages: - 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 + devices: "0,1" # Uses both GPUs + tensor_parallel_size: 2 + gpu_memory_utilization: 0.6 # 60% per GPU + # Model is split, so each GPU uses ~30% of model memory ``` ## Examples @@ -118,27 +112,25 @@ stage_args: ### Qwen3-Omni-MoE on 2x H100-80GB ```yaml -stage_args: +base_config: /path/to/vllm_omni/deploy/qwen3_omni_moe.yaml + +stages: - 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 + devices: "0,1" + 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 + devices: "1" + 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 + devices: "0" + 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. +**Note:** Stage 0 uses GPUs 0 and 1, so the combined utilization is `0.7` on +GPU 0 and `0.9` on GPU 1. Keep the sum of all resident stages below `1.0` on +each device. ## Troubleshooting diff --git a/docs/configuration/pd_disaggregation.md b/docs/configuration/pd_disaggregation.md deleted file mode 100644 index 93c623ebd..000000000 --- a/docs/configuration/pd_disaggregation.md +++ /dev/null @@ -1,171 +0,0 @@ -# Prefill-Decode (PD) Disaggregation - -PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode -stages so prompt processing and token generation can run on different workers. - -This is documented as a stage-config recipe instead of a bundled YAML because the -deployment-specific values usually change per environment: - -- GPU placement -- `tensor_parallel_size` -- connector backend and connector ports -- connector IPs or bootstrap addresses - -Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml) -and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the -changes below. - -## Requirements - -- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav -- A KV connector supported by vLLM, such as `MooncakeConnector` -- Matching `tensor_parallel_size` on the prefill and decode thinker stages - -## 1. Split the thinker into prefill and decode stages - -Replace the original thinker stage with two stages: - -```yaml -stage_args: - - stage_id: 0 - stage_type: llm - is_prefill_only: true - runtime: - devices: "0" - engine_args: - max_num_seqs: 16 - model_stage: thinker - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.9 - enforce_eager: true - trust_remote_code: true - engine_output_type: latent - distributed_executor_backend: "mp" - enable_prefix_caching: false - max_num_batched_tokens: 32768 - hf_config_name: thinker_config - tensor_parallel_size: 1 - kv_transfer_config: - kv_connector: "MooncakeConnector" - kv_role: "kv_producer" - kv_rank: 0 - kv_parallel_size: 2 - kv_connector_extra_config: - mooncake_bootstrap_port: 25201 - final_output: false - is_comprehension: true - default_sampling_params: - temperature: 0.4 - top_p: 0.9 - top_k: 1 - max_tokens: 2048 - seed: 42 - detokenize: True - repetition_penalty: 1.05 - - - stage_id: 1 - stage_type: llm - is_decode_only: true - runtime: - devices: "1" - engine_args: - max_num_seqs: 64 - model_stage: thinker - model_arch: Qwen3OmniMoeForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.9 - enforce_eager: true - trust_remote_code: true - engine_output_type: latent - distributed_executor_backend: "mp" - enable_prefix_caching: false - max_num_batched_tokens: 32768 - hf_config_name: thinker_config - tensor_parallel_size: 1 - kv_transfer_config: - kv_connector: "MooncakeConnector" - kv_role: "kv_consumer" - kv_rank: 1 - kv_parallel_size: 2 - kv_connector_extra_config: - mooncake_bootstrap_port: 25202 - engine_input_source: [0] - final_output: true - final_output_type: text - is_comprehension: true - default_sampling_params: - temperature: 0.4 - top_p: 0.9 - top_k: 1 - max_tokens: 2048 - seed: 42 - detokenize: True - repetition_penalty: 1.05 -``` - -Notes: - -- `is_prefill_only: true` marks the thinker stage that only saves KV. -- `is_decode_only: true` marks the thinker stage that resumes from remote KV. -- `kv_transfer_config` is required on both stages. -- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the - prefill side only processes the prompt and exports KV. - -## 2. Shift the downstream stages by one index - -After inserting the extra thinker stage, renumber the remaining stages: - -```yaml - - stage_id: 2 - runtime: - devices: "2" - engine_input_source: [1] - sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only - - - stage_id: 3 - runtime: - devices: "2" - engine_args: - max_num_seqs: 1 - engine_input_source: [2] - sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only -``` - -Compared with the default Qwen3-Omni config: - -- the talker becomes stage `2` instead of stage `1` -- the code2wav stage becomes stage `3` instead of stage `2` -- the talker now reads from decode stage `1` - -## 3. Add runtime edges for the four-stage pipeline - -```yaml -runtime: - enabled: true - edges: - - from: 0 - to: 1 - - from: 1 - to: 2 - - from: 2 - to: 3 -``` - -## 4. Launch with your custom config - -```bash -vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ - --deploy-config /path/to/qwen3_omni_pd.yaml -``` - -## Operational Notes - -- `MooncakeConnector` does not support heterogeneous TP sizes across the PD - pair. Keep prefill and decode at the same `tensor_parallel_size`. -- If the thinker requires TP=2, both thinker stages must use TP=2 and be given - separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode. -- Choose connector ports and addresses that match your deployment. The values - shown above are examples only. diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 290c15b94..855390739 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -3,52 +3,104 @@ In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. !!! note - Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when `--deploy-config` is omitted. The model registry resolves the right pipeline and deploy YAML by `model_type`. + Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when `--deploy-config` is omitted. The resolved pipeline selects its default through `default_deploy_config_name`. + +## Pipeline configuration + +`PipelineConfig` and its `StagePipelineConfig` entries are Python definitions +owned by the model implementation and registered in `OMNI_PIPELINES`. They +describe fixed model topology and are not accepted in deploy YAMLs. + +Common `PipelineConfig` fields include: + +| Field | Description | +| ------- | ------------- | +| `model_type` | Pipeline identifier used during model and config resolution. | +| `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. | +| `model_arch` | Default Hugging Face architecture for the pipeline. | +| `hf_architectures` | Architecture names used to identify checkpoints whose `model_type` is shared. | +| `hf_config_predicate` | Optional predicate used to select between pipelines with otherwise identical HF metadata. | +| `diffusers_class_name` | Diffusers `_class_name` used to identify Diffusers-style repositories. | +| `stages` | Ordered tuple of fixed `StagePipelineConfig` definitions. | + +Common `StagePipelineConfig` fields include: + +| Field | Description | +| ------- | ------------- | +| `stage_id` | Stable stage identifier. | +| `model_stage` | Logical stage name used by runtime and strategy resolution. | +| `execution_type` | `LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`. | +| `input_sources` | Upstream stage IDs that provide this stage's inputs. | +| `final_output` / `final_output_type` | Whether the stage produces a user-visible output and its modality. | +| `owns_tokenizer` | Whether this stage owns the pipeline tokenizer. | +| `model_arch`, `hf_config_name` | Stage-specific model architecture and nested HF config selector. | +| `engine_output_type` | Runtime output representation such as `text`, `latent`, or `audio`. | +| `custom_process_input_func` | Processor applied to this stage's incoming payload. | +| `custom_process_next_stage_input_func` | Processor used for full-payload handoff to the next stage. | +| `async_chunk_process_next_stage_input_func` | Processor used for async chunk handoff. | +| `sampling_constraints` | Model-owned sampling constraints that deploy defaults cannot override. | + +To add or change topology, define and register a new pipeline variant. Use +deploy YAML only for runtime placement, resource sizing, connectors, and other +deployment overrides. ## Deploy configuration schema The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields: | Field | Type | Required | Default | Description | -|-------|------|----------|---------|-------------| +| ------- | ------ | ---------- | --------- | ------------- | | `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. | | `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. | +| `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. | +| `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. | +| `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. | | `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). | | `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. | -| `stages` | list | required | — | Per-stage engine args + wiring (see [Stage fields](#stage-fields)). | +| `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. | | `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. | -| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). | -| `trust_remote_code` | bool | optional | `true` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage. | +| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only` / `qwen3_omni_moe_thinker_only`). | +| `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. | | `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. | | `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. | | `quantization` | str \| null | optional | `null` | **Pipeline-wide.** Quantization method for every stage. | -| `enable_prefix_caching` | bool | optional | `false` | **Pipeline-wide.** Prefix cache toggle applied to every stage. | +| `enable_prefix_caching` | bool \| null | optional | `null` | **Pipeline-wide.** Prefix cache toggle applied to every stage when specified. | | `enable_chunked_prefill` | bool \| null | optional | `null` | **Pipeline-wide.** Chunked prefill toggle applied to every stage. | -| `data_parallel_size` | int | optional | `1` | **Pipeline-wide.** DP degree for every stage. | -| `pipeline_parallel_size` | int | optional | `1` | **Pipeline-wide.** PP degree for every stage. | +| `data_parallel_size` | int \| null | optional | `null` | **Pipeline-wide.** DP degree for every stage. | +| `pipeline_parallel_size` | int \| null | optional | `null` | **Pipeline-wide.** PP degree for every stage. | +| `custom_voice_dir` | str \| null | optional | `null` | **Pipeline-wide.** Directory containing custom voice profiles for supported TTS models. | -Note: for diffusion path, `distributed_executor_backend` currently defaults to -`mp`, and `ray` / `external_launcher` are not fully supported yet. +For fields whose deploy default is `null`, the deploy layer contributes no +override. The effective value may still come from a platform section, an +explicit CLI or stage override, or the downstream vLLM engine default. + +Note: for the diffusion path, an omitted `distributed_executor_backend` selects +`uni` on a single GPU (in-process worker, no MessageQueue / `/dev/shm` output +segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker +subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet. ### Stage fields -Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. +Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`. | Field | Type | Required | Default | Description | -|-------|------|----------|---------|-------------| +| ------- | ------ | ---------- | --------- | ------------- | | `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. | -| `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. | -| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. | -| `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. | -| `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. | -| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. | -| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). | +| `max_num_seqs` | int \| null | optional | `null` | Max concurrent sequences per stage. | +| `gpu_memory_utilization` | float \| null | optional | `null` | Per-stage total memory target; used for automatic KV-cache sizing. | +| `kv_cache_memory_bytes` | int \| null | optional | `null` | Explicit per-rank KV-cache budget in bytes; overrides automatic sizing. | +| `tensor_parallel_size` | int \| null | optional | `null` | TP degree for this stage. | +| `enforce_eager` | bool \| null | optional | `null` | Disable CUDA graphs. | +| `max_num_batched_tokens` | int \| null | optional | `null` | Per-stage prefill/token budget; also contributes to the native maximum in-flight token limit. | +| `max_model_len` | int \| null | optional | `null` | Per-sequence context or KV length; `-1` enables native cache-capacity auto-fitting, while values above the HF default auto-set `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1`. | | `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. | -| `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. | +| `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. | | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_`; values are names registered under top-level `connectors:`. | | `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_`; values are names registered under top-level `connectors:`. | | `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). | -| `engine_extras` | dict | optional | `{}` | Catch-all for keys not listed above; deep-merged across overlays. Also carries per-stage overrides of pipeline-wide settings (e.g. stage-specific `dtype`). | +| `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. | + +**Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage. ### Connector schema @@ -60,12 +112,12 @@ connectors: name: # required — class registered in vllm_omni.distributed extra: # optional — forwarded to the connector's __init__ : - ... + # Additional connector-specific options ``` -| Connector class | Use case | `extra` keys | -|-----------------|----------|--------------| -| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. | +| Connector class | Use case | `extra` keys | +|--------------------------|-----------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------| +| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | None. All payloads use shared memory. | | `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). | A stage references a connector by name in its `input_connectors` / `output_connectors`: @@ -85,7 +137,7 @@ stages: ### CLI flags | Flag | Description | -|------|-------------| +| ------ | ------------- | | `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/.yaml` is auto-loaded by the model registry. | | `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. | | `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. | @@ -199,7 +251,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ Effective config per stage after the merge: | Stage | Field | Final value | Source | -|-------|-------|-------------|--------| +| ------- | ------- | ------------- | -------- | | 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) | | 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` | | 0 | `max_model_len` | `16384` | global CLI | @@ -223,300 +275,17 @@ when loading a deploy configuration. Examples: For offline inference (assuming the necessary dependencies have been imported): + ```python model_name = "Qwen/Qwen2.5-Omni-7B" omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml") ``` For online serving: + ```bash vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml ``` !!! important We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers. - -## Legacy stage configuration reference - -The remainder of this page documents the deprecated `stage_args` schema for models that have not migrated to deploy configurations. The following historical Qwen2.5-Omni configuration illustrates that schema; current Qwen2.5-Omni deployments should use `vllm_omni/deploy/qwen2_5_omni.yaml` instead. - -```python -# Historical stage configuration for Qwen2.5-Omni. -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" # Logical device index for this stage (mapped through CUDA_VISIBLE_DEVICES / ASCEND_RT_VISIBLE_DEVICES if set) - engine_args: # Engine arguments for a certain engine - model_stage: thinker - max_num_seqs: 1 - model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py - worker_type: ar # 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" - engine_args: - model_stage: talker - max_num_seqs: 3 - model_arch: Qwen2_5OmniForConditionalGeneration - worker_type: ar - 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] - 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 - engine_args: - model_stage: code2wav - max_num_seqs: 1 - model_arch: Qwen2_5OmniForConditionalGeneration - worker_type: generation - 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 - - edges: - - from: 0 # thinker → talker: trigger only after receiving full input (-1) - to: 1 - - from: 1 # talker → code2wav: trigger only after receiving full input (-1) - to: 2 - -``` - -## Legacy 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`). - -### `prompt_expand_func` (Optional) - -A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously. - -### `cfg_kv_collect_func` (Optional) - -A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers. - -### `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` - -Logical device indices for this stage, specified as a string. Values are **logical indices** (`0`, `1`, `2`, ...) — not physical GPU IDs — and are mapped through the platform's visibility env var (`CUDA_VISIBLE_DEVICES` on CUDA, `ASCEND_RT_VISIBLE_DEVICES` on NPU) before being applied via `torch.cuda.set_device()` (or the equivalent). - -Example: if `CUDA_VISIBLE_DEVICES=0,2,4` is set in the environment, then `devices: "0"` selects physical GPU 0 (the first visible), `devices: "1"` selects physical GPU 2, and `devices: "0,1"` makes physical GPUs 0 and 2 available to the stage. If no visibility env var is set, logical and physical IDs coincide. - -Default: `"0"` - -#### `engine_args.max_num_seqs` - -The maximum number of sequences for concurrent processing in this stage. For -LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. -For diffusion stages, this controls scheduler wave capacity in both request -batch mode and step batch mode. `max_num_seqs: 1` is the serial/conservative -path; values above `1` allow compatible requests to batch when the selected -diffusion pipeline supports that batching mode. - -Default: `1` - -#### `engine_args.request_batch_max_wait_ms` - -The maximum time, in milliseconds, that a diffusion request-mode stage may wait -before the first `schedule()` of a new scheduler wave so compatible requests can -accumulate for request-level batching. This only applies to diffusion pipelines -that support request-level batching with `step_execution` disabled. - -Use this together with `max_num_seqs > 1` for bursty serving traffic. `0` -disables admission waiting and preserves the lowest first-request latency. -For diffusion execution and batching tuning, see -[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md). - -Default: `0.0` - -### `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` - -### `tts_args` (TTS stages only) - -Configuration for Text-to-Speech specific parameters. This section is only applicable to TTS model stages (e.g., `qwen3_tts`). - -#### `tts_args.max_instructions_length` - -Maximum character length for voice style/emotion instructions. Instructions exceeding this limit will be rejected with a validation error. - -Default: `500` - -This value can be overridden at runtime using the `--tts-max-instructions-length` CLI parameter when starting the server. diff --git a/docs/contributing/DOCS_GUIDE.md b/docs/contributing/DOCS_GUIDE.md index 916f5e60f..57c2fdc54 100644 --- a/docs/contributing/DOCS_GUIDE.md +++ b/docs/contributing/DOCS_GUIDE.md @@ -46,7 +46,7 @@ class Omni: Args: model: Model name or path - stage_configs: Optional stage configurations + deploy_config: Optional path to a deploy configuration **kwargs: Additional arguments passed to the engine Example: diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 2224c667e..dbfc917bf 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -44,6 +44,151 @@ vLLM-Omni's `pre-commit` hooks will now run automatically every time you commit. pre-commit run --show-diff-on-failure --color=always --all-files # runs on all files (short for --all-files) ``` +!!! warning + GitHub Actions `pre-commit` **skips** several local gates so a whole-tree + `--all-files` run stays green while historical debt is cleaned up. The + current `SKIP` list in + [`.github/workflows/pre-commit.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.github/workflows/pre-commit.yml) + is `check-mark`, `markdownlint-cli2`, `shellcheck`, + `check-spdx-header`, and `mypy-3.10`. Those hooks still run on **your + commit** for changed files. A passing GitHub pre-commit check does not mean + they passed locally. Hooks **not** on that list (forbidden imports, + `torch.cuda`, TTS adapter ratchet, Buildkite schema, Ruff, typos, …) run in + both places. + +The hooks below are new relative to the previous Omni config (Ruff, typos, +actionlint, YAML/whitespace, DCO sign-off, and the pickle-only checker were +already there). `check-pickle-imports` is gone: pickle is now one rule inside +`check-forbidden-imports`. + +| Hook id | Enforces | GitHub Actions | +| ------- | -------- | -------------- | +| `markdownlint-cli2` | Markdown in `docs/`, `recipes/`, `README.md`, `CONTRIBUTING.md` (not `.claude/` / `.cursor/` / `CLAUDE.md`) | skipped | +| `mypy-3.10` | Type-check changed `vllm_omni/` files (model trees excluded). `tests/` uses `--follow-imports skip` | skipped | +| `mypy-3.11` / `3.12` / `3.13` | Same checker, extra Python versions | not installed; `pre-commit run --hook-stage manual mypy-3.12` | +| `check-mark` | Every `tests/**/test_*.py` has a CI level mark and a hardware platform mark or helper; no direct `pytest.mark.H100` / SKU | skipped | +| `check-tts-adapter-migration` | `self._tts_model_type` branches in `serving_speech.py` must not increase | runs | +| `shellcheck` | `*.sh` quoting / undefined vars | skipped | +| `check-spdx-header` | Omni SPDX header on `.py` / `.pyi` / `.sh` / `.rs` / `.proto` | skipped | +| `check-forbidden-imports` | pickle, stdlib `re`/`base64`, Triton, TileLang, Hugging Face Hub download/API | runs | +| `check-torch-cuda-call` | Direct `torch.cuda.*` helpers outside platform adapters | runs | +| `check-buildkite` | `.buildkite/*.yml` against the Buildkite schema | runs | + +Related config (not a new hook id): Ruff `TID251` also bans `librosa` and several +`torch.cuda.*` names; typos is pinned and ignores git-describe hashes (`+g` + +hex) and NumPy `writeable`. Keep the `suggestion` hook last in +`.pre-commit-config.yaml` so autofix output does not bury the SKIP tip. + +#### SPDX headers + +Source files must start with: + +```text +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +``` + +Rust and proto files use `//` instead of `#`. Empty files are exempt. A stale +`Copyright contributors to the vLLM project` line is **rewritten in place** to +the Omni copyright; the hook then exits non-zero so you restage the rewrite. + +#### Forbidden imports + +Under `vllm_omni/`, do not use: + +- stdlib `re` → `import regex as re` +- stdlib `base64` → `import pybase64` (or `import pybase64 as base64`) +- `pickle` / `cloudpickle` (also banned in most tests) +- `from huggingface_hub import snapshot_download` / `hf_hub_download` / `HfApi` + / similar → `vllm.transformers_utils.repo_utils` +- direct `triton` / `tilelang` → `vllm.triton_utils` / `vllm.tilelang_utils` + +`examples/`, `tests/`, `benchmarks/`, `apps/`, `docs/`, `tools/`, `scripts/`, +`.buildkite/`, and `.github/` may keep stdlib `re` and `base64` +(`_NON_LIBRARY_DIRS`). Pickle is stricter: tests are not exempt unless the +file is on the pickle allowlist. + +#### `torch.cuda` call sites + +Prefer `current_omni_platform` / `OmniPlatform`. Platform adapters under +`vllm_omni/platforms/` are allowed; other files must not add new +`torch.cuda.*` helpers (`device_count`, `synchronize`, `empty_cache`, +`device()`, `manual_seed`, and similar). `import torch` is fine. + +#### shellcheck + +The hook does **not** download a binary. Install `shellcheck` with a signed +package manager, then re-run: + +- Debian/Ubuntu/WSL: `sudo apt-get install shellcheck` +- Fedora: `sudo dnf install ShellCheck` +- macOS: `brew install shellcheck` +- Git Bash (MINGW): `scoop install shellcheck` so `shellcheck.exe` is on PATH. + WSL is Linux: use `apt-get`, not a Windows `.exe`. + +If no native binary is found, the wrapper prints those commands (and +[the upstream install page](https://github.com/koalaman/shellcheck?tab=readme-ov-file#installing)) +and exits 1. AMD and other vendor scripts are **not** excluded. + +#### mypy + +`mypy-3.10` runs on commit for changed files. `[tool.mypy]` uses Python 3.12, +`ignore_missing_imports = true`, and `follow_imports = "silent"`. Wrapper +`tools/pre_commit/mypy.py` skips `vllm_omni/model_executor/models/` and +`vllm_omni/diffusion/models/`, and type-checks `tests/` with +`--follow-imports skip`. Extra versions: + +```bash +pre-commit run --hook-stage manual mypy-3.12 +``` + +#### Test CI marks + +`check-mark` requires each collected test module under `tests/` to +have at least one CI **level** mark (`core_model`, `advanced_model`, +`full_model`, `local_model`, or `slow`) and a **hardware** platform mark +(`cpu`, `cuda`, …) or helper (`hardware_test(` / `hardware_marks(`). SKU marks +(`H100`, `L4`, …) must come from those helpers so `cards_{n}` is attached; +direct `pytest.mark.H100` is rejected. GitHub Actions skips this hook; local +commit does not. See the +[test writing guide](./ci/test_writing_guide.md). + +#### TTS adapter ratchet + +`check-tts-adapter-migration` counts `self._tts_model_type` comparisons in +`vllm_omni/entrypoints/openai/serving_speech.py`. The count must stay at or +below `MAX_MODEL_TYPE_BRANCHES` in +[`tools/pre_commit/check_tts_adapter.py`](https://github.com/vllm-project/vllm-omni/blob/main/tools/pre_commit/check_tts_adapter.py). +New per-model behavior belongs in +`vllm_omni/entrypoints/openai/tts_adapters/`. Removing branches **must** lower +the budget in the same change; raising it is a reviewable policy edit. + +#### Markdownlint and Buildkite + +`markdownlint-cli2` auto-fixes `docs/`, `recipes/`, and the root README / +CONTRIBUTING (rules in `.markdownlint.yaml`). Skill files under `.claude/` and +`.cursor/` are excluded. + +`check-buildkite` validates `.buildkite/*.yml` against the official schema +after the same expansion `upload_pipeline.py` uses. Omni-only keys such as +`mirror_hardwares` are stripped or rendered first. A few files are skipped in +`SKIP_FILES` inside `tools/pre_commit/check_buildkite.py`. + +#### Extending allowlists and budgets + +Prefer fixing the call site. Growing an allowlist or raising a budget is a +policy change and needs review — do not edit these just to make the hook pass. + +| Gate | File | What to edit | +| ---- | ---- | ------------ | +| pickle, `re`, `base64`, Hugging Face Hub, Triton/TileLang | [`tools/pre_commit/check_forbidden_imports.py`](https://github.com/vllm-project/vllm-omni/blob/main/tools/pre_commit/check_forbidden_imports.py) | `CHECK_IMPORTS[].allowed_files`. Use `allowed_dirs` only when a whole tree should stay exempt. | +| `torch.cuda` | [`tools/pre_commit/check_torch_cuda.py`](https://github.com/vllm-project/vllm-omni/blob/main/tools/pre_commit/check_torch_cuda.py) | `ALLOWED_FILES`, or `ALLOWED_PREFIXES` for a platform adapter tree. | +| TTS `_tts_model_type` branches | [`tools/pre_commit/check_tts_adapter.py`](https://github.com/vllm-project/vllm-omni/blob/main/tools/pre_commit/check_tts_adapter.py) | `MAX_MODEL_TYPE_BRANCHES` (down only, unless the PR justifies a raise). | +| Buildkite skip list | [`tools/pre_commit/check_buildkite.py`](https://github.com/vllm-project/vllm-omni/blob/main/tools/pre_commit/check_buildkite.py) | `SKIP_FILES` for pipelines that cannot be expanded. | + +To skip one hook for a single commit (discouraged): `SKIP= git commit`. +To bypass every hook: `git commit --no-verify` (also discouraged). + ### 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`. @@ -162,12 +307,14 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-review against project conventions: The skill offers two modes: + - **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status - **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) for the full checklist. ### Local Test + Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_execution_guide.md) for running the test cases. ### Automatic skip-ci (docs and pytest skip marks) @@ -189,7 +336,7 @@ These PR shapes all trigger skip-ci: Skip-ci does **not** apply when the diff also touches product code (for example `vllm_omni/`), or when test files change assertions, imports, fixtures, or other non-skip logic. If the diff cannot be resolved (non-PR branches outside `main`), CI runs as usual. !!! note - Skipping L2/L3 does **not** disable the Docker image build step. Nightly (L4) upload can still run when the PR has a `nightly-test` label or on scheduled `main` builds with `NIGHTLY=1`. Bootstrap child steps live in `bootstrap-upload-steps.yml`; the hook entry step runs `upload_pipeline.py --upload /bootstrap-upload-steps.yml`, which injects `if` by step `key` from skip-ci before upload. See [CI Settings — Diff-aware CI](./ci/ci_settings.md#diff-aware-ci) and [Test System Overview](./ci/test_system_overview.md). + Skipping L2/L3 does **not** disable the Docker image build step. Nightly (L4) upload can still run when the PR has a `nightly-test` label or on scheduled `main` builds with `NIGHTLY=1`. Weekly (L5) upload can still run with `weekly-test`, `WEEKLY=1`, or `NON_CRITICAL=1`. Bootstrap child steps live in `bootstrap-upload-steps.yml`; the hook entry step runs `upload_pipeline.py --upload /bootstrap-upload-steps.yml`, which injects `if` by step `key` from skip-ci before upload. See [CI Settings — Diff-aware CI](./ci/ci_settings.md#diff-aware-ci) and [Test System Overview](./ci/test_system_overview.md). ### Code Quality diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index ff4b8e107..a3ffed765 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -8,7 +8,7 @@ For CI levels (L1–L5) and triggers, see [Test System Overview](./test_system_o Canonical layout (prefer these paths for new changes): -``` +```text .buildkite/ ├── common/ # Shared across platforms │ ├── scripts/ @@ -48,7 +48,7 @@ Canonical layout (prefer these paths for new changes): └── scripts/ ``` -**Placement rules** +### Placement rules | Rule | Detail | | ---- | ------ | @@ -81,7 +81,14 @@ There are still **legacy copies** at `.buildkite/*.yaml` (without the `cuda/` pr | L2 | `cuda/test-ready.yml` | `ready` label | | L3 | `cuda/test-merge.yml` | `merge-test` label / main merge | | L4 | `cuda/test-nightly.yml` | `nightly-test` label or `NIGHTLY=1` | - | L5 | `cuda/test-weekly.yml` | `weekly-test` label or `WEEKLY=1` | + | L5 | `cuda/test-weekly.yml` | `weekly-test` label, `WEEKLY=1`, or `NON_CRITICAL=1` (see step `if`s below) | + + On scheduled `main` builds, `test-weekly.yml` is uploaded when **`WEEKLY=1` or `NON_CRITICAL=1`**. Inside that file: + + - **Reliability** / **Perf Test** / **Simple · CPU Coverage Test** → `WEEKLY=1` (or PR label `weekly-test` for Reliability/Perf). Pipeline upload for scheduled env vars still requires `main`. + - **E2E Tests** group (slow Omni/TTS/Diffusion sweeps) → `NON_CRITICAL=1` + + **`--e2e`:** when uploading ready/merge on `main` with **`WEEKLY=1`**, bootstrap passes `--e2e` so only the **E2E Test** group is kept (same flag `run_cov_split.sh` uses to enable per-model coverage). **Upload:** `upload_pipeline.py --upload` expands uploader-only keys before Buildkite upload. @@ -112,6 +119,8 @@ There are still **legacy copies** at `.buildkite/*.yaml` (without the `cuda/` pr **Test YAML:** `npu/test-npu-ready.yml` (L2), `test-npu-nightly.yml` (L4). + **L2 trigger:** PR label `ready` (CUDA/AMD L2). PR labels `ready` + `npu-test` trigger NPU CI only. + **Upload:** same as CUDA—`upload_pipeline.py --upload`. **Hardware in YAML:** `mirror_hardwares` preset (string), expanded to `agents`, top-level `image`, and `plugins`. Presets: `a2b3_npu_1`, `a2b3_npu_4`, `a3_npu_2` in `common/ci_mirror_hardwares.yml`. @@ -204,15 +213,15 @@ Label triggers (`ready`, `merge-test`) are unchanged—diff-aware logic only red | Diff shape | Path | Default L2/L3 | | --- | --- | --- | | Product code / non-whitelisted paths | normal CI | all on | -| Docs and/or qualifying skip-mark only | `skip_all` | all off (scheduled `main` NIGHTLY/WEEKLY exceptions below) | +| Docs and/or qualifying skip-mark only | `skip_all` | all off (scheduled `main` NIGHTLY/WEEKLY/`NON_CRITICAL` exceptions below) | | Whitelisted CI YAML only | yaml-gated (`skip_l2_l3`) | per platform/level matrix | | Docs/skip-mark **+** whitelisted CI YAML | yaml-gated (`skip_l2_l3`) | same as CI-YAML-only (does **not** widen to normal CI) | | Non-qualifying skip-mark and no CI YAML | normal CI | all on | | Diff unavailable | normal CI | all on | -**`skip_all` exceptions (CUDA/NPU bootstrap only):** PR labels (`nightly-test`, `merge-test`, `npu-test`, `weekly-test`, …) do **not** revive jobs. On `main`, scheduled `NIGHTLY=1` still builds the image and uploads L4 plus L2/L3 (with `--e2e`); `WEEKLY=1` still builds the image and uploads L5 (CUDA). +**`skip_all` exceptions (CUDA/NPU bootstrap only):** PR labels (`nightly-test`, `merge-test`, `weekly-test`, `ready`, …) do **not** revive jobs under `skip_all`. On `main`, scheduled `NIGHTLY=1` still builds the image and uploads L4; `WEEKLY=1` or `NON_CRITICAL=1` still builds the image and uploads L5 (CUDA). `WEEKLY=1` on `main` also uploads L2/L3 with `--e2e`. -**Yaml-gated nightly/weekly:** L4/L5 upload steps keep their normal label / `NIGHTLY` / `WEEKLY` conditions (for example PR `nightly-test` still uploads nightly). Only L2/L3 upload steps are matrix-gated. +**Yaml-gated nightly/weekly:** L4/L5 upload steps keep their normal label / `NIGHTLY` / `WEEKLY` / `NON_CRITICAL` conditions (for example PR `nightly-test` still uploads nightly). Only L2/L3 upload steps are matrix-gated. ##### Whitelisted CI YAML → platform @@ -285,12 +294,12 @@ CUDA **L2** (`.buildkite/cuda/test-ready.yml`) and **L3** (`.buildkite/cuda/test | `source_file_dependencies` | List of path **prefixes**. If any changed file equals a prefix or starts with `prefix/`, keep the step (or group); otherwise omit it. | | `mirror_hardwares` | Expand to `agents` + `plugins` (+ optional `image`) from `ci_mirror_hardwares.yml`. | -**Policy** +### Policy - **Always uploaded** (no key): groups outside **E2E Test**—Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3). - **Diff-gated**: every **E2E Test** leaf job. List the smallest prefix set per step—pytest file(s), model code under `vllm_omni/model_executor/models/` or `vllm_omni/diffusion/models/`, plus `stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable. A **group** may define the key instead; the whole group drops if no prefix matches. -**YAML examples** +### YAML examples ```yaml - label: "Diffusion · Qwen Image Test" @@ -312,7 +321,7 @@ CUDA **L2** (`.buildkite/cuda/test-ready.yml`) and **L3** (`.buildkite/cuda/test mirror_hardwares: l4_4 ``` -**Local dry-run** +### Local dry-run ```bash python3 .buildkite/common/scripts/upload_pipeline.py .buildkite/cuda/test-ready.yml @@ -324,15 +333,109 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for ### Labels and grouping -- Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` nightly runs on main. +- Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` weekly runs on main (`WEEKLY=1`). - Prefix labels by model domain: **Omni ·**, **TTS ·**, **Diffusion ·**, **Simple ·**, etc., matching existing steps. +### Per-model coverage + +Pilot (v1): one Merge-tier (L3) job in `.buildkite/cuda/test-merge.yml` uploads +per-model, per-entry-mode coverage as Buildkite artifacts — "TTS · Qwen3-TTS Base +Test". Coverage + artifact upload run only when bootstrap uploaded that job with +`--e2e` (`WEEKLY=1` and `BUILDKITE_BRANCH=main`, via `run_cov_split.sh`); otherwise +the same job runs plain pytest without `--cov`. It runs on a single GPU so the +pilot is cheap to reproduce. + +**L1 package coverage (CPU):** ready/merge Simple Test jobs run without `--cov`. Scheduled weekly on `main` with `WEEKLY=1` runs **Simple · CPU Coverage Test** in `.buildkite/cuda/test-weekly.yml` (`pytest -m 'core_model and cpu'` + Cobertura XML artifact). + +**Naming convention**: `coverage---.xml.gz`, where +`` is the model's directory name under +`vllm_omni/diffusion/models//` or +`vllm_omni/model_executor/models//` (e.g. `z_image`, `qwen3_omni`), +`` is `online` or `offline`, and `` is `BUILDKITE_STEP_ID` +(`local` outside Buildkite). `gunzip` before feeding a report to a Cobertura +consumer. + +Reports are gzipped because `--cov=vllm_omni` sets coverage's `source`, which +makes it walk the package and emit one `` element per statement for every +file — including the ones a single model's tests never import. That inventory, +not the coverage, is the size: a Bagel report is 7.1 MB, of which 96.5% is +`` elements and 77% belongs to the 840 of 1169 files that run never +imported. A run covering zero lines produces the same 7 MB. gzip takes it to +roughly 410 KB and changes nothing about the data. + +`` is what keeps two steps that cover the same model apart. Artifacts are +scoped to the build, and a weekly build can carry ready/merge (with `--e2e`) and +weekly tiers at once, so several same-model steps land in one namespace — `qwen3_tts` already has +both a Base and a CustomVoice job, and `minicpmo_4_5` a base and a duplex one. Two +uploads on one path do not fail loudly: an exact-name download can report the path +as ambiguous, and a glob download fetches every match concurrently and renames each +onto the same destination, so whichever finishes last wins. The step id remains the +logical step identifier across retries, while Buildkite's default artifact lookup +selects only the latest attempt — `BUILDKITE_JOB_ID` would instead give every +attempt its own filename, leaving a failed attempt's partial report to be picked up +alongside the real one. A step using `parallelism`/matrix would need +`BUILDKITE_PARALLEL_JOB` added, since its jobs share a step id; neither pilot is +one. + +**Opting in a new model**: replace the job's combined `pytest` command with +[`run_cov_split.sh`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/run_cov_split.sh), +which runs each mode, names the reports, uploads them, and fails the step if +either half or the upload failed: + +```yaml +commands: + - | + .buildkite/common/scripts/run_cov_split.sh \ + --model-id \ + --offline \ + --online \ + -- +``` + +Everything after `--` is forwarded verbatim to both runs with quoting preserved, +so an expression like `-m 'advanced_model and cuda'` survives intact and a job can +pass anything else it needs. `--offline`/`--online` are repeatable and only one is +required, so a job with tests in a single mode runs just that half. Runtime is +bounded solely by the step's `timeout_in_minutes` — raise it, since the split +loads the model once per mode. + +Before splitting, confirm each half actually collects at least one test under the job's forwarded filters — pytest exits with code 5 ("no tests collected") if a half is empty, which the script reports as a failure, red-ing a job that used to pass as one combined invocation. + +Also check the tests' `num_cards` against the job's `mirror_hardwares` preset. +`_cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a +test asking for more GPUs than the preset provides is silently skipped and its +coverage artifact comes back empty. Splitting such a job produces a file that +looks fine and measures nothing — "Diffusion · Z Image Test" is the current +example (online tests want 4 cards, the job runs on `l4_1`), which is why it is +not a coverage pilot. + +Online-mode coverage depends on the `[tool.coverage.run]` `patch`/`sigterm` +settings in `pyproject.toml`: online tests launch the server with +`subprocess.Popen` and stop it with SIGTERM, and without those settings the XML +reflects only the pytest parent process, not the server's code paths. + +The upload depends on `buildkite-agent` being callable inside the container. The +`kubernetes` presets (`h100_*`, `*_npu_*`) provide it; the `docker` ones (`l4_*`) +only do because they set `mount-buildkite-agent: true`. A new docker preset that +runs a coverage job needs the same. + +List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's +`source_file_dependencies` — both change what the job measures, so without them a +change there is filtered out of normal PR builds and only surfaces in a later +nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with +a synthetic job (it does not pin real merge labels). +Editing only the surrounding CI YAML still does not schedule them, so a PR that +touches just the wiring needs a full E2E run (or the commands run on a GPU host) +to produce artifacts. When checking a new model's +artifacts, compare `lines-covered` between the online and offline XML rather than +just confirming both files exist. + ### Validation checklist | Check | Command / location | | ----- | ------------------ | | CUDA render | `python3 .buildkite/common/scripts/upload_pipeline.py .buildkite/cuda/test-.yml` | -| No leaked uploader keys | `… \| grep -E 'mirror_hardwares|source_file_dependencies'` → empty | +| No leaked uploader keys | Pipe the render into `grep -E 'mirror_hardwares\|source_file_dependencies'` — expect empty | | Skip-ci / upload unit tests | `pytest tests/buildkite/` | | Local job replay (CUDA L2+) | `tools/run_ready_jobs.sh`, `tools/run_merge_jobs.sh`, `tools/nightly/run_nightly_jobs.sh` (read YAML from `cuda/`) | diff --git a/docs/contributing/ci/failures.md b/docs/contributing/ci/failures.md index 9743097cf..e4b846cca 100644 --- a/docs/contributing/ci/failures.md +++ b/docs/contributing/ci/failures.md @@ -6,7 +6,7 @@ When you open a PR against vLLM-Omni, several CI checks run automatically: | Check | Platform | What it does | | ----- | -------- | ------------ | -| **pre-commit** | GitHub Actions | Runs linting (Ruff), formatting, spell-checking (typos), and YAML validation. | +| **pre-commit** | GitHub Actions | Runs Ruff, formatting, typos, YAML validation, forbidden imports, `torch.cuda`, TTS adapter ratchet, and Buildkite schema. Local-only gates (SPDX, shellcheck, markdownlint, mypy-3.10, test marks) are in the workflow `SKIP` env — see [Linting](../README.md#linting). | | **Build Wheel** | GitHub Actions | Builds Python wheels for Python 3.11 and 3.12 on Ubuntu. Skipped for docs-only or Markdown-only changes (controlled by `paths-ignore` in the workflow). | | **DCO** | GitHub | Verifies every commit has a `Signed-off-by` line. | | **docs/readthedocs.org:vllm-omni** | Read the Docs | Builds the MkDocs documentation site. | @@ -29,6 +29,13 @@ pre-commit run --all-files Then commit the fixes and push. +A green GitHub pre-commit job can still miss SPDX, shellcheck, mypy-3.10, +markdownlint, and test-mark coverage because CI skips those IDs. If your +**local** commit hook fails, follow [Linting](../README.md#linting) for the +full hook list (forbidden imports, `torch.cuda`, TTS adapter ratchet, and +Buildkite schema still run in GHA). Do not grow allowlists or raise +`MAX_MODEL_TYPE_BRANCHES` without review. + ### DCO failures Every commit must include a `Signed-off-by` line. If you forgot, amend your commits: diff --git a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md index 5e65d2f48..3411b4377 100644 --- a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md @@ -1,8 +1,8 @@ -**Scope** +#### Scope For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/). -**Test Case Design** +#### Test Case Design For a *high priority* model (currently listed in [issue #1832](https://github.com/vllm-project/vllm-omni/issues/1832)), we use several test cases, each with multiple features turned on, so that each supported feature is tested in at least one test case. This is to relieve the GPU workload on the CI machine. The suggested test case combination is as follows: @@ -34,12 +34,12 @@ For a *normal priority* model, further reduce the number of test cases. Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. -**Code Style** +#### Code Style -- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this. +- Validation: test that the multimodal output files of your model have the correct shapes. `OnlineOmniClient.send_diffusion_request` should have taken care of this. - Test marks: always add `full_model` and `diffusion` for L4 nightly `test_*_expansion.py` cases. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_writing_guide/#markers-for-tests). - To maximize code reuse, you may refer to - - `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper. + - `tests/conftest.py` for `omni_server` (running server in subprocess) and `online_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper. - `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`. - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases. - Doc: add a concise docstring for each test function. diff --git a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md index 8ff27ee08..feb4f38ed 100644 --- a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md @@ -1,14 +1,15 @@ When you want to add L4-level ***performance test*** cases, add entries to JSON files under `tests/dfx/perf/tests/` and run them via `tests/dfx/perf/scripts/run_benchmark.py` (omni / TTS) or `run_diffusion_benchmark.py` (diffusion). -**Config file layout (in-tree examples)** +#### Config file layout (in-tree examples) | Model type | Runner | Example JSON files | | ---------- | ------ | ------------------ | -| Omni | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` | +| Omni (nightly) | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json` (`full_model` without `slow` in `mark`) | +| Omni (weekly) | `run_benchmark.py` | `test_qwen3_omni_async_chunk.json` (CUDA only), `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` (`slow` in `mark`; **Perf Test** in `test-weekly.yml`) | | TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` | | Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … | -**How runners pick cases** +#### How runners pick cases Without **`--test-config-file`**, each runner scans all `*.json` under `tests/dfx/perf/tests/` but only keeps its own model type: @@ -17,9 +18,9 @@ Without **`--test-config-file`**, each runner scans all `*.json` under `tests/df Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-omni"`) or `"diffusion"` in the `mark` array. Omni / TTS JSON has neither. -**Running perf cases** +#### Running perf cases -Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"`). +Pass **`--test-config-file`** to load one JSON file, or omit it for the bulk scan above. In either mode, pytest **`-m`** filters each case using its JSON `mark`. A single file can contain nightly, weekly, and platform-specific case objects. For example, the nightly object in `test_qwen3_omni_async_chunk.json` declares both H100 and A3, while its `slow` weekly object is CUDA-only; there is no NPU weekly counterpart. Use `-m "H100 and full_model and not slow"`, `-m "H100 and full_model and slow"`, or `-m "npu"` to select the intended schedule and platform. ```JSON { @@ -31,7 +32,7 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul ], "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_config_name": "qwen3_omni.yaml" + "stage_config_name": "qwen3_omni_moe.yaml" }, "benchmark_params": [ { @@ -55,31 +56,33 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul } ``` -**Parameter Explanation** +#### Parameter Explanation -*Overview* +##### Overview -| Field | Required | Description | -| ---------------- | -------- | --------------------------------------------------------------- | -| test_name | Yes | Unique identifier for the test case | -| mark | No | Pytest marks for this case (see **`mark` field** below). Omit only for configs not meant to be filtered by `-m`. | -| server_params | Yes | Server-side configuration parameters | -| benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) | -| server_type | Diffusion only | When set (e.g. `"vllm-omni"`), routes the case to `run_diffusion_benchmark.py`. | -| benchmark_endpoint | Diffusion only | API path for the benchmark client (e.g. `/v1/videos`, `/v1/images/generations`). | +| Field | Required | Description | +| ------------------ | -------------- | -------------------------------------------- | +| test_name | Yes | Unique identifier for the test case | +| mark | No | Pytest marks; see **`mark` field** below | +| server_params | Yes | Server-side configuration parameters | +| benchmark_params | Yes | Benchmark running parameters | +| server_type | Diffusion only | Routes case to run_diffusion_benchmark.py | +| benchmark_endpoint | Diffusion only | Benchmark API path | -**`mark` field** +Omit `mark` only for configs not meant to be filtered by `-m`. `server_type` is typically `"vllm-omni"`. `benchmark_endpoint` examples: `/v1/videos`, `/v1/images/generations`. + +#### `mark` field Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`. -When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, `diffusion`, or `local_model`. +When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `slow`, `omni`, `tts`, `diffusion`, or `local_model`. Supported form: | Form | Example | Effect | | ---- | ------- | ------ | -| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` > 1 adds `distributed_*` (+ CUDA `skipif_cuda` when applicable). String entries become extra pytest marks. | -| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325", "npu": "A2"}, "num_cards": {"cuda": 2, "rocm": 2, "npu": 1}}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** in one case. Filter examples: `-m "full_model and H100 and cuda"`, `-m "full_model and MI325 and rocm"`. | +| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` adds `cards_{n}` (including `cards_1` when omitted). String entries become extra pytest marks. | +| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325"}, "num_cards": 2}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** on one item. All platforms must share the same `num_cards` (int or identical dict values). Different counts need `@hardware_test` or one `hardware_marks()` per `pytest.param`. Filter examples: `-m "full_model and H100 and cards_2 and cuda"`, `-m "full_model and MI325 and cards_2 and rocm"`. | Recommended for L4 perf cases: @@ -119,11 +122,11 @@ Multi-GPU diffusion (example: Cosmos3 with `cfg-parallel-size=2`): - Diffusion perf: include `"diffusion"` in the `mark` array - HunyuanImage local-weight cases: add `"local_model"` to the `mark` array -**Parametrization IDs** +#### Parametrization IDs Each `(server, benchmark index)` pair becomes one `pytest.param` with id `{test_name}-{suffix}`. The suffix comes from `benchmark_params[].name` when set (for example `test_tts-p0`, `test_omni-p1`); otherwise it is derived from fields like `task` / `eval_phase`. -**Benchmark result filenames** +#### Benchmark result filenames Result files use the **runtime** hardware label from `get_runtime_resource_label()` (detected GPU/NPU name on the machine that ran the job), **not** `mark.hardware_marks.res`. On the default H100 CI pool, `H100` is **omitted** from filenames (`resource_label_for_filename`). @@ -132,32 +135,40 @@ Examples: - Omni/TTS: `result_{test_name}_{optional_hw}_{dataset}_....json` under `BENCHMARK_DIR` - Diffusion: one aggregate `diffusion_result_{config_stem}_{optional_hw}_{timestamp}.json` per source JSON file (array of all runs from that file) -**Local commands** +#### Local commands ```bash # Bulk load + filter by JSON mark pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and H100 and diffusion" pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100" -# Single file (same as nightly CI Perf steps) +# Single file (same selectors as the CI Perf steps) pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \ --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \ - --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json + --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \ + -m "H100 and full_model and not slow" +pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \ + --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \ + -m "H100 and full_model and slow" +pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \ + --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json \ + -m "npu" ``` -See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …). + +See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `cards_1`, …). -**`server_params` Configuration** +#### `server_params` Configuration -*Basic Parameters* +##### Basic Parameters -| Parameter | Required | Example | Description | -| ----------------- | -------- | ---------------------------------- | ----------------------------- | -| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path | -| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name | +| Parameter | Required | Example | Description | +| ----------------- | -------- | ------------------------- | ------------------------------ | +| model | Yes | Qwen/Qwen3-Omni-30B-A3B… | Model name or path | +| stage_config_name | Yes | qwen3_omni_moe.yaml | Deploy configuration file name | -*Dynamic Configuration (update/delete)* +##### Dynamic Configuration (update/delete) Supports incremental modifications based on the basic configuration: @@ -166,38 +177,40 @@ Supports incremental modifications based on the basic configuration: | update | Update or add configuration items | | delete | Delete specified configuration items | -**Example**: +##### Example -``` +```json "update": { "async_chunk": true, // Enable asynchronous chunk processing - "stage_args": { + "stages": { "0": { - "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk" + "max_num_seqs": 64 } } }, "delete": { - "stage_args": { - "2": ["custom_process_input_func"] // Delete this configuration for stage 2 + "stages": { + "2": ["max_num_batched_tokens"] // Delete this deploy override for stage 2 } } ``` -**`benchmark_params` Configuration** +#### `benchmark_params` Configuration You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows: -1. Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file. -2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file. -3. Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field. -4. Set `"name"` on each `benchmark_params` entry for stable pytest ids and readable result keys. -5. The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below: +1. Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file. +2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file. +3. Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field. +4. Set `"name"` on each `benchmark_params` entry for stable pytest ids and readable result keys. +5. The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below: -| Parameter | Type | Required | Example/Values | Description | -| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| name | string | No | `"1024x1024_steps4"` | Stable suffix for pytest param id and logs. | -| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. | -| request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. | -| max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). | -| baseline | object | No | see below | Optional thresholds. **Hardware-nested only**: top-level keys must be `[hardware-resource]` pytest markers from `pyproject.toml` (`get_hardware_mark_list()`). Under each hardware key, metric names are free-form; values are scalars or sweep-aligned lists. Example: `{"H100": {"throughput_qps": […], "custom_stage_ms": 1.0}, "A3": {…}}`. Runtime filename aliases (`_RUNTIME_DEVICE_ALIASES`) are separate and do not expand this allowlist. | +| Parameter | Type | Required | Example/Values | Description | +| --------------- | ------------- | -------- | -------------------- | ------------------------------------ | +| name | string | No | `"1024x1024_steps4"` | Stable suffix for pytest param id | +| num_prompts | int / array | Yes | 10, [10, 20, 30] | Number of requests (scalar or sweep) | +| request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second (QPS mode) | +| max_concurrency | int / array | No | 1, [1, 2, 3] | Max in-flight requests | +| baseline | object | No | see below | Optional hardware-nested thresholds | + +`num_prompts`, `request_rate`, and `max_concurrency` accept scalars or arrays; single values expand to match sweep length. QPS and concurrency modes are mutually exclusive. `baseline` top-level keys must be hardware pytest markers from `pyproject.toml`; metric names and values are free-form (scalars or sweep-aligned lists). Omit `baseline` to run without threshold assertions. diff --git a/docs/contributing/ci/test_execution_guide.md b/docs/contributing/ci/test_execution_guide.md index d1c2c560f..7597f72e4 100644 --- a/docs/contributing/ci/test_execution_guide.md +++ b/docs/contributing/ci/test_execution_guide.md @@ -1,13 +1,20 @@ # Test Execution Guide + ## Setting Up the Test Environment + ### Creating a Container + vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on. For a local test environment, you can follow the steps below to create a container: + ## Installing Dependencies + ### vLLM & vLLM-Omni + vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment. ### Test Case Dependencies + When running test cases, you may need to install the following dependencies: ```bash @@ -16,6 +23,7 @@ apt-get install -y espeak-ng jq ``` ## Running Tests + Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases. ### CI job runners (L2–L4): logs, timing, and timeouts @@ -25,7 +33,7 @@ Our test scripts use the pytest framework. First, please use `git clone https:// **Log layout** (default under `logs/ready_jobs`, `logs/merge_jobs`, or `logs/nightly_jobs`; override with `--log-dir`): | Path | Description | -|------|-------------| +| ------ | ------------- | | `jobs/*.sh` | Generated per-step bash wrappers | | `jobs/.job_timeouts` | Job key → `timeout_in_minutes` from Buildkite YAML (when set) | | `.log` | Tee output for each job | @@ -33,7 +41,7 @@ Our test scripts use the pytest framework. First, please use `git clone https:// **Timing:** while each job runs, stderr prints `finished in …` or `failed after …`. When all jobs finish, a summary is printed and written to `timing_summary.log`, for example: -``` +```text === Job timing summary === TTS_VoxCPM2_Test 25m 40s OK Omni_Qwen3-Omni_Test 8m 05s TIMED OUT @@ -80,7 +88,7 @@ Failed jobs: 1/2 ```bash cd tests pytest -s -v test_xxxx.py --run-level=core_model - pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model + pytest -s -v -m "core_model and L4 and not cards_1" --run-level=core_model ``` === "L3 level" @@ -111,10 +119,9 @@ Failed jobs: 1/2 ```bash pytest -s -v test_xxxx.py --run-level=advanced_model - pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model + pytest -s -v -m "advanced_model and L4 and not cards_1" --run-level=advanced_model ``` - === "L4 level" **Recommended:** run CI-aligned nightly jobs from the repo root via [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh). The script reads [`.buildkite/cuda/test-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-nightly.yml), generates per-step bash wrappers, runs pytest (perf jobs first, then aggregates perf JSON into Excel), and tees logs under `logs/nightly_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts); `generate_nightly_perf_excel` is included in the timing summary). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`). @@ -146,12 +153,12 @@ Failed jobs: 1/2 ```bash cd tests pytest -s -v test_xxxx.py --run-level=full_model - pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model + pytest -s -v -m "full_model and L4 and not cards_1" --run-level=full_model pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model ``` If you only want to run specific test cases on a particular platform, you can use: ```bash - pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model + pytest -s -v -m "full_model and L4 and not cards_1" --run-level=full_model ``` Note: ``run_benchmark.py`` and ``run_diffusion_benchmark.py`` accept an optional ``--test-config-file``. If omitted, each loads every ``*.json`` under ``tests/dfx/perf/tests/`` (omni/tts vs diffusion split by ``is_diffusion_perf_config``) and pytest ``-m`` filters by each case's JSON ``mark``: ```bash @@ -160,7 +167,7 @@ Failed jobs: 1/2 pytest -sv tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json ``` - Nightly **Perf Test** jobs in ``test-nightly.yml`` use ``--test-config-file`` only (no ``-m``). E2e L4 function tests still use ``full_model`` + ``--run-level full_model`` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml)). Example: + Nightly **Perf Test** jobs in [``test-nightly.yml``](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-nightly.yml) use ``--test-config-file`` only (no ``-m``). Weekly **Perf Test** in [``test-weekly.yml``](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml) runs ``test_qwen3_omni_vllm_text.json`` and ``test_qwen3_omni_multi_replicas.json`` (JSON ``mark`` includes ``slow``). E2e L4 function tests use ``full_model`` + ``--run-level full_model``. Example: === "L5 level" @@ -184,14 +191,15 @@ Failed jobs: 1/2 # Reliability: HunyuanImage DiT (H100 × 4) pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow - # Reliability: VoxCPM2 (L4) - pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow - ``` - The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml). + The latest L5 CI jobs are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml): + + - **Reliability** (+ invalid-parameter), **Perf Test** (Omni vLLM Text / Multi-Replica), and **Simple · CPU Coverage Test** → `WEEKLY=1` (or PR label `weekly-test` for Reliability/Perf). Scheduled upload of those env-gated steps still requires `main`. + - **E2E Tests** (slow Omni / TTS / Diffusion sweeps) → `NON_CRITICAL=1` (pipeline upload also accepts `WEEKLY=1` / `weekly-test`; the group `if` still needs `NON_CRITICAL`) You can find more information about markers in [Test Writing Guide — Markers for Tests](./test_writing_guide.md#markers-for-tests). ## Adding New Test Cases + Please refer to the [Test System Overview](./test_system_overview.md) and [Test Writing Guide](./test_writing_guide.md). diff --git a/docs/contributing/ci/test_system_overview.md b/docs/contributing/ci/test_system_overview.md index 335ed5c53..b4d3ce121 100644 --- a/docs/contributing/ci/test_system_overview.md +++ b/docs/contributing/ci/test_system_overview.md @@ -12,8 +12,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Level Scope & Focus Model Coverage Strategy - Feature Coverage Strategy - Interface Coverage Strategy + Module Strategy Tags Time Cost Test Dir @@ -31,7 +30,6 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari / / / - / PR Checklist / / @@ -43,7 +41,6 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari / / / - / CI Failures / / @@ -53,7 +50,6 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Unit tests for components like entrypoints, models / / - / core_model and cpu <15min @@ -71,9 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L2
(E2E across models & GPU-required UT) Online (basic deployment scenarios):
dummy, normal inference function (output format, stream), some instance startup UT - High-priority models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy) - High-priority features and component tests that require launching instances - High-priority interfaces (using random lightweight models) + Key models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy) + Key features and modules that require launching instances (prefer random weights) core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion Model E2E:
@@ -95,9 +90,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L3
(Important Perf & Integration & Accuracy) Online & Offline (multiple deployment scenarios):
real model, normal inference function, normal accuracy - High/medium-priority models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy - Medium-priority features and component tests that require launching instances - Medium-priority interfaces (using random lightweight models) + Key models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy + Key features and modules that require launching instances (using real weights) advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <30min @@ -120,9 +114,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L4
(Perf & Integration & Accuracy) Online: full functional scenarios + performance test + doc test + accuracy test - High-priority models: function, performance, accuracy, and doc testing
Medium-priority models: function and doc testing - Low-priority features and component tests that require launching instances - Low-priority interfaces (using real weights) + Key models: function, performance, accuracy, and doc testing + Other features and modules that require launching instances (using real weights) full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <3 hour @@ -147,11 +140,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari GPU - L5
(Stability & Reliability) - Online: long-term stability test + reliability test - Long-term stability and reliability testing for high-priority models
Low-priority models: function and doc testing + L5
(Stability & Reliability & selected Perf) + Online: long-term stability + reliability; coverage test + Long-term stability and reliability testing for key models
Non-critical scenario performance and accuracy testing for key models
Other models: function and doc testing / - Invalid-parameter validation for high-priority interfaces slow and hardware_test(H100, L4, etc.) and omni/tts/diffusion Depends on reality @@ -178,8 +170,8 @@ For per-level test authoring (markers, examples), see [Test Writing Guide](./tes Before entering specific testing levels, the project establishes two common specifications aimed at standardizing the development process and quickly locating issues. -1. ***PR Checklist ([`.github/PULL_REQUEST_TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/.github/PULL_REQUEST_TEMPLATE.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance. -2. ***CI Failure Explanation ([CI Failures](./failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency. +1. ***PR Checklist ([`.github/PULL_REQUEST_TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/.github/PULL_REQUEST_TEMPLATE.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance. +2. ***CI Failure Explanation ([CI Failures](./failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency. ## Notes diff --git a/docs/contributing/ci/test_writing_guide.md b/docs/contributing/ci/test_writing_guide.md index 397f5e83c..469e69cac 100644 --- a/docs/contributing/ci/test_writing_guide.md +++ b/docs/contributing/ci/test_writing_guide.md @@ -17,47 +17,43 @@ Summary: - **Reliability:** `tests/dfx/reliability/` - Do **not** add new top-level `tests/` directories unrelated to a component (or to `e2e` / `dfx` / `helpers` / `examples` / `buildkite`) -Deploy YAMLs: `vllm_omni/deploy/` via `tests.helpers.stage_config.get_deploy_config_path` (see [stage configs](../../configuration/stage_configs.md)). +Deploy YAMLs: `vllm_omni/deploy/` via `tests.helpers.stage_config.get_deploy_config_path` (see [pipeline and deploy configs](../../configuration/stage_configs.md)). ## 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` | L1&L2 tests (run in each PR) | -| `advanced_model` | L3 tests (run on each merge to main) | -| `full_model` | L4 tests (run nightly) | -| `diffusion` | Diffusion model tests | -| `omni` | Omni multimodal model tests | -| `tts` | TTS model tests | -| `cache` | Cache backend tests | -| `parallel` | Parallelism/distributed tests | -| `cpu` | Tests that run on CPU | -| `gpu` | Tests that run on GPU * | -| `cuda` | Tests that run on CUDA * | -| `rocm` | Tests that run on AMD/ROCm * | -| `xpu` | Tests that run on Intel XPU * | -| `npu` | Tests that run on NPU/Ascend * | -| `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) | +| Marker | Description | +| ------------------ | ------------------------------------------------------------------------------------------------------- | +| `core_model` | L1&L2 tests (run in each PR) | +| `advanced_model` | L3 tests (run on each merge to main) | +| `full_model` | L4 tests (run nightly) | +| `diffusion` | Diffusion model tests | +| `omni` | Omni multimodal model tests | +| `tts` | TTS model tests | +| `cache` | Cache backend tests | +| `parallel` | Parallelism/distributed tests | +| `cpu` | Tests that run on CPU | +| `gpu` | Tests that run on GPU * | +| `cuda` | Tests that run on CUDA * | +| `rocm` | Tests that run on AMD/ROCm * | +| `xpu` | Tests that run on Intel XPU * | +| `npu` | Tests that run on NPU/Ascend * | +| `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 * | +| `cards_{n}` | Tests that require *n* accelerator cards (`cards_1` … `cards_8`); auto-added from `num_cards` * | +| `slow` | Slow tests (may skip in quick CI) | | `benchmark` | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) | -| `local_model` | Tests requiring local / non-HF-hub model weights | +| `local_model` | Tests requiring local / non-HF-hub model weights | -\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility). +\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility). Do **not** write `pytest.mark.H100` / `pytest.mark.L4` (or other SKU names) by hand: `check-mark` rejects that so `cards_{n}` cannot be skipped. #### Example usage for markers @@ -80,28 +76,26 @@ def test_video_to_audio() 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/helpers/mark.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`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`). - ``` + Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`). + ```python @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) - ``` +2. **Handles card-count filtering** + Adds a filterable marker `cards_{n}` (including `cards_1` when `num_cards` is omitted / 1). Use this with the SKU and platform, e.g. `-m "H100 and cards_4 and cuda"` vs `-m "H100 and cards_2 and cuda"`. Any multi-card CUDA test is `-m "cuda and not cards_1"`. + + When `res` lists more than one platform **and** `num_cards` differs by platform, `@hardware_test` collects **one variant per platform** so those `cards_{n}` marks do not stack. `-m "H100 and cards_4 and cuda"` then selects only the CUDA/H100 4-card variant, not a ROCm 2-card sibling. The same count on every platform stays a single item. 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. - + On CUDA and XPU, when `num_cards > 1`, attaches pytest's built-in `skipif` if the machine has fewer devices than requested. This is not a named `skipif_*` marker and is not used for `-m` filtering. 5. **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"`). + Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "H100 and cards_2"`, `-m "L4 and not cards_1"`). ##### Example usage for decorator + - Single call for multiple platforms: ```python @hardware_test( @@ -116,13 +110,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring num_cards=2, ) ``` -- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), XPU (B60), MUSA (S5000), NPU (A2/A3) +- `res` must be a dict. SKU values come from `pyproject.toml` markers tagged `[hardware-resource]` plus the platform key (`[cuda]`, `[rocm]`, `[xpu]`, `[npu]`, `[musa]`). - `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing -- 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"` + - CUDA H100 4-card: `pytest -m "H100 and cards_4 and cuda"` + - CUDA H100 2-card: `pytest -m "H100 and cards_2 and cuda"` + - CUDA multi-card (any count): `pytest -m "L4 and not cards_1"` + - ROCm multi-card: `pytest -m "MI325 and not cards_1"` + - NPU multi-card: `pytest -m "npu and not cards_1"` #### Function: `hardware_marks` @@ -178,37 +173,30 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator [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", + "tpu: [hardware-platform] Tests that require TPU device", + # Add [gpu] next to [hardware-platform] if this platform should also get pytest.mark.gpu. + "TPU_V3: [hardware-resource] [tpu] Tests that require TPU v3 hardware", + "cards_1: [hardware-cards] Tests that require 1 accelerator card", + "cards_2: [hardware-cards] Tests that require 2 accelerator cards", ] ``` + `check-mark` treats `[hardware-platform]` names (including `cpu`/`gpu`) as allowed `pytest.mark.` values. `hardware_test(res=...)` still rejects `cpu`/`gpu` as `res` keys. 2. **Implement a marker construction function for your platform** in `vllm-omni/tests/helpers/mark.py`: ```python # In vllm-omni/tests/helpers/mark.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] + def _tpu_marks(*, res: str, num_cards: int): + _require_sku("tpu", res) + return [pytest.mark.tpu, getattr(pytest.mark, res), _cards_mark(num_cards)] ``` -3. **Update `hardware_marks` to recognize your new platform**: - In the relevant place (see the `hardware_marks` implementation), add: +3. **Register the builder** in `_marks_for_platform`: ```python - if platform == "tpu": - marks = tpu_marks(res=resource, num_cards=cards) + builders = { + # ... existing platforms ... + "tpu": _tpu_marks, + } ``` - (`hardware_test` calls `hardware_marks` internally, so both will pick up the change.) + (`hardware_test` / `hardware_marks` both go through `_marks_for_platform`.) 4. **(Recommended) Add a test using your new markers**: ```python @hardware_test( @@ -220,12 +208,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator ``` **Summary**: -- Add pytest markers for your new platform/resources -- Implement a marker function (`xxx_marks`) -- Plug into `hardware_marks` -- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation! -See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`). +- Add pytest markers for your new platform/resources +- Implement a private marker function (`_xxx_marks`) +- Register it in `_marks_for_platform` +- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers and `cards_{n}` filtering. + +See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`_cuda_marks`, `_rocm_marks`, `_npu_marks`). ## Test case style @@ -237,47 +226,46 @@ L1 and L2 level testing form the foundation of the quality assurance system. L1 L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle. - - #### 1.2 Testing Content and Scope -- ***L1 (Unit & Logic Testing)***: -- - ***Scope***: Tests internal functions and methods of core components such as `entrypoints`, `models`. - - ***Focus***: Branch coverage, exception handling, algorithm logic correctness. Does not involve external dependencies or the complete service stack. - - ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback. -- ***L2 (Basic End-to-End Testing)***: -- - ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference). - - ***Focus***: Uses `dummy` weights (via deploy YAML patching at `core_model`) or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, and **cheap payload checks** at `--run-level core_model` (see below). Also includes some unit tests that require launching independent service instances. - - ***L2 response validation (`core_model`)***: Implemented in `tests/helpers/assertions.py` and invoked by `OpenAIClientHandler` based on `--run-level`. At L2 we require **request success** plus minimal output sanity—not full accuracy: - - **Speech / TTS** (`assert_audio_speech_response`): decoded audio must be present and non-empty (or exceed `min_audio_bytes` when set in `request_config`); `response_format` must match the returned content-type (e.g. `wav`, `pcm`). Whisper transcript similarity, PCM HNR, and preset-voice gender checks run only at L3+. - - **Diffusion** (`assert_diffusion_response`): at least one non-empty image, video, or audio artifact. Resolution/frame-count parity and other deep checks run only at L3+. - - **Omni multimodal** (`assert_omni_response`): L2 asserts successful completion; keyword, transcript, and cross-modal similarity checks run only at L3+. - - ***Characteristic***: Requires ***GPU*** resources to perform model computations. +- ***L1 (Unit & Logic Testing)***: + - ***Scope***: Tests internal functions and methods of core components such as `entrypoints`, `models`. + - ***Focus***: Branch coverage, exception handling, algorithm logic correctness. Does not involve external dependencies or the complete service stack. + - ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback. +- ***L2 (Basic End-to-End Testing)***: + - ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference). + - ***Focus***: Uses `dummy` weights (via deploy YAML patching at `core_model`) or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, and **cheap payload checks** at `--run-level core_model` (see below). Also includes some unit tests that require launching independent service instances. + - ***L2 response validation (`core_model`)***: Implemented in `tests/helpers/assertions.py` and invoked by `OpenAIClientHandler` based on `--run-level`. At L2 we require **request success** plus minimal output sanity—not full accuracy: + - **Speech / TTS** (`assert_audio_speech_response`): decoded audio must be present and non-empty (or exceed `min_audio_bytes` when set in `request_config`); `response_format` must match the returned content-type (e.g. `wav`, `pcm`). Whisper transcript similarity, PCM HNR, and preset-voice gender checks run only at L3+. + - **Diffusion** (`assert_diffusion_response`): at least one non-empty image, video, or audio artifact. Resolution/frame-count parity and other deep checks run only at L3+. + - **Omni multimodal** (`assert_omni_response`): L2 asserts successful completion; keyword, transcript, and cross-modal similarity checks run only at L3+. + - ***Characteristic***: Requires ***GPU*** resources to perform model computations. #### 1.3 Test Directory and Execution Files A clear directory structure is key to managing test cases efficiently. See the **Test Dir** column in [Test System Overview](./test_system_overview.md) for the canonical placement rules. -- ***L1 (component / unit)***: `tests/{component_name}/…` mirroring `vllm_omni/{component_name}/` (e.g. `tests/diffusion/`, `tests/entrypoints/`). Do **not** invent new top-level `tests/` folders unrelated to a component. -- ***L2 (basic model E2E)***: - - Online serving: `tests/e2e/online_serving/test_{model_name}.py` - - Offline inference: `tests/e2e/offline_inference/test_{model_name}.py` -- ***Feature-level integration*** (cross-cutting, not a single-model smoke): `tests/e2e/features//` (e.g. `fullduplex/`, `custom_pipeline/`, `rlhf_test/`). +- ***L1 (component / unit)***: `tests/{component_name}/…` mirroring `vllm_omni/{component_name}/` (e.g. `tests/diffusion/`, `tests/entrypoints/`). Do **not** invent new top-level `tests/` folders unrelated to a component. +- ***L2 (basic model E2E)***: + - Online serving: `tests/e2e/online_serving/test_{model_name}.py` + - Offline inference: `tests/e2e/offline_inference/test_{model_name}.py` +- ***Feature-level integration*** (cross-cutting, not a single-model smoke): `tests/e2e/features//` (e.g. `fullduplex/`, `custom_pipeline/`, `rlhf_test/`). #### 1.4 Execution Method and Example -- ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered. -- ***Diff-aware step skipping***: On L2, **E2E Test** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](ci_settings.md#step-filtering) prefixes; non-E2E groups still always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci). -- ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment. -- ***Script Example***: +- ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered. +- ***Diff-aware step skipping***: On L2, **E2E Test** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](ci_settings.md#step-filtering) prefixes; non-E2E groups still always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci). +- ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment. +- ***Script Example***:
L1 Test Examples Examples from `tests/model_executor/models/qwen2_5_omni/test_audio_length.py` + ```python # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest @@ -331,6 +319,7 @@ def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames) if max_mel_frames is not None and int(max_mel_frames) > 0 and int(max_mel_frames) >= repeats: assert target_mel_len <= int(max_mel_frames) ``` +
@@ -338,7 +327,7 @@ def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames) You can refer to Test Examples in the L3 section to see example test cases that incorporate both L2 and L3 testing logic.
-- - ***Run Command***: +- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py` `pytest -s -v -m 'core_model and cpu' --run-level=core_model` @@ -351,36 +340,34 @@ L3 level testing executes after code is merged into the main branch. Its core pu . It acts as the "quality gatekeeper" for the main branch, ensuring that no merge breaks the core capabilities of the model service. Testing needs to provide clear conclusions within a relatively short time (<30min), balancing test depth with feedback speed. - - #### 2.2 Testing Content and Scope -- ***Deployment Scenarios***: Covers richer `online` and `offline` deployment configurations, which may include different hardware configurations, batch sizes, concurrency levels, etc. -- ***Core Verification***: -- 1. ***Inference Functionality***: Ensures real models can perform forward computation normally and return results. - 2. ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues. - 3. ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds. -- ***L3 response validation (`advanced_model`)***: At `--run-level advanced_model`, `tests/helpers/assertions.py` adds semantic checks on top of L2 payload gates (see L1 & L2 above): - - **Speech / TTS** (`assert_audio_speech_response`): Whisper transcript of returned audio vs `request_config["input"]` (cosine similarity > 0.9 when input text is set); optional `min_audio_bytes` floor; PCM harmonic-to-noise ratio when `response_format` is `pcm`; **preset voice gender** via `_assert_preset_voice_gender_from_audio` when `voice` matches a known preset in `_PRESET_VOICE_GENDER_MAP` (pitch/F0-based classifier; skipped for unknown voices or `pcm` output). - - **Omni multimodal** (`assert_omni_response`): non-empty text/audio outputs; `key_words` in transcript or text; text–audio similarity / containment; preset **`speaker`** gender check (same helper as TTS). - - **Diffusion** (`assert_*_diffusion_response`): image/video dimension and frame-count parity with request parameters where configured. - - **Accuracy suites**: pixel/video similarity and metric baselines under `/tests/e2e/accuracy/` (separate from inline helper assertions). +- ***Deployment Scenarios***: Covers richer `online` and `offline` deployment configurations, which may include different hardware configurations, batch sizes, concurrency levels, etc. +- ***Core Verification***: + 1. ***Inference Functionality***: Ensures real models can perform forward computation normally and return results. + 2. ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues. + 3. ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds. +- ***L3 response validation (`advanced_model`)***: At `--run-level advanced_model`, `tests/helpers/assertions.py` adds semantic checks on top of L2 payload gates (see L1 & L2 above): + - **Speech / TTS** (`assert_audio_speech_response`): Whisper transcript of returned audio vs `request_config["input"]` (cosine similarity > 0.9 when input text is set); optional `min_audio_bytes` floor; PCM harmonic-to-noise ratio when `response_format` is `pcm`; **preset voice gender** via `_assert_preset_voice_gender_from_audio` when `voice` matches a known preset in `_PRESET_VOICE_GENDER_MAP` (pitch/F0-based classifier; skipped for unknown voices or `pcm` output). + - **Omni multimodal** (`assert_omni_response`): non-empty text/audio outputs; `key_words` in transcript or text; text–audio similarity / containment; preset **`speaker`** gender check (same helper as TTS). + - **Diffusion** (`assert_*_diffusion_response`): image/video dimension and frame-count parity with request parameters where configured. + - **Accuracy suites**: pixel/video similarity and metric baselines under `/tests/e2e/accuracy/` (separate from inline helper assertions). #### 2.3 Test Directory and Execution Files -- ***Functional / model E2E***: - - Online serving: `tests/e2e/online_serving/test_{model_name}_expansion.py` - - Offline inference: `tests/e2e/offline_inference/test_{model_name}_expansion.py` - - Accuracy / similarity: `tests/e2e/accuracy/` - - (`_expansion.py` holds broader scenario coverage than the L2 smoke file.) -- ***Feature-level integration***: `tests/e2e/features//` (not under a new top-level `tests/` directory). +- ***Functional / model E2E***: + - Online serving: `tests/e2e/online_serving/test_{model_name}_expansion.py` + - Offline inference: `tests/e2e/offline_inference/test_{model_name}_expansion.py` + - Accuracy / similarity: `tests/e2e/accuracy/` + - (`_expansion.py` holds broader scenario coverage than the L2 smoke file.) +- ***Feature-level integration***: `tests/e2e/features//` (not under a new top-level `tests/` directory). #### 2.4 Execution Method and Example -- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge). -- ***Diff-aware step skipping***: Same [`source_file_dependencies`](ci_settings.md#step-filtering) mechanism as L2—**all E2E Test** leaf steps are diff-gated; other L3 groups always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci). -- ***Execution Environment***: ***GPU*** servers. -- ***Script Example***: +- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge). +- ***Diff-aware step skipping***: Same [`source_file_dependencies`](ci_settings.md#step-filtering) mechanism as L2—**all E2E Test** leaf steps are diff-gated; other L3 groups always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci). +- ***Execution Environment***: ***GPU*** servers. +- ***Script Example***: ???+ example "Test Examples" @@ -496,9 +483,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu **Audio output debugging**: Deep validation may transcribe returned audio via `convert_audio_bytes_to_text` (Whisper). If an audio keyword or text–audio similarity assertion fails, set `VLLM_OMNI_KEEP_REQUEST_MEDIA=1` before running pytest to keep the intermediate WAV files for inspection (see [Test helper environment variables](test_system_overview.md#test-helper-environment-variables)). -- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model` +- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model` -- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model` +- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model` ### L4 Level Testing - Full Functionality, Performance, and Documentation Testing @@ -506,27 +493,25 @@ L3 level testing executes after code is merged into the main branch. Its core pu L4 level testing is a comprehensive quality audit before a version release. It expands upon L3, executing ***full*** functional scenarios, conducting systematic ***performance stress tests***, and simultaneously verifying the correctness of accompanying ***example documentation***. Its purpose is to perform deep validation of the system during off-peak nighttime hours, providing quality trend reports for daytime development and data support for release decisions. - - #### 3.2 Testing Content and Scope -- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling. -- ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen3_omni_*.json` (Omni), `test_tts.json` / `test_voxcpm2.json` / `test_higgs_audio_v3.json` (TTS), and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` or `run_diffusion_benchmark.py` via `--test-config-file` in nightly **Perf Test** steps) to drive throughput, latency, and memory benchmarks. Each JSON **case** may declare an optional top-level **`mark`** array: exactly one ``hardware_marks`` object plus pytest marker name strings (`full_model`, `omni` / `tts` / `diffusion`, …). Runners attach those marks to each parametrized `(server, benchmark index)` pair so **local** bulk runs can filter with `-m` (for example `-m "full_model and H100 and diffusion"`). Nightly CI perf jobs select workloads by **`--test-config-file`**, not `-m`. Details are in the Performance Tests example below. -- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description. +- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling. +- ***Performance Testing***: JSON configs under `tests/dfx/perf/tests/` drive Omni, TTS, and diffusion serving benchmarks via `run_benchmark.py` or `run_diffusion_benchmark.py`, measuring throughput, latency, and related metrics under configured workloads. +- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description. #### 3.3 Test Directory and Execution Files -- ***Functional Testing***: Same directories as L3. -- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen3_omni_*.json`, `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`). Optional per-case `mark` for local `-m` filtering is documented in Section 3.4 Performance Tests. -- ***Documentation Example Tests***: - - `tests/examples/online_serving/test_{model_name}.py` - - `tests/examples/offline_inference/test_{model_name}.py` +- ***Functional Testing***: Same directories as L3. +- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen3_omni_*.json`, `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`). Optional per-case `mark` for local `-m` filtering is documented in Section 3.4 Performance Tests. +- ***Documentation Example Tests***: + - `tests/examples/online_serving/test_{model_name}.py` + - `tests/examples/offline_inference/test_{model_name}.py` #### 3.4 Execution Method and Example -- ***Trigger Timing***: **`Nightly`**, automatically executed every night. -- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing. -- ***Script Example***: +- ***Trigger Timing***: **`Nightly`**, automatically executed every night. +- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing. +- ***Script Example***: ??? example "Test Examples: Documentation Example Tests" @@ -540,7 +525,7 @@ L4 level testing is a comprehensive quality audit before a version release. It e --8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md" -- ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`). +- ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`). ### L5 Level Testing - Stability and Reliability Testing @@ -548,41 +533,43 @@ L4 level testing is a comprehensive quality audit before a version release. It e L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness. - - #### 4.2 Testing Content and Scope -- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable. -- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `
` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill). +- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable. +- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See Section 4.3 for file-level responsibilities and Section 4.4 Test Examples (Reliability) for CI markers (`slow`, `hardware_test`, POSIX-only kill). #### 4.3 Test Directory and Execution Files -- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family) -- ***Reliability Test Suite*** (`tests/dfx/reliability/`): - - `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default). - - `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery). - - `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4). - - `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1). - - `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`). - - `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready. - - `README.md` — Short local run commands for this directory. +- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family) +- ***Reliability Test Suite*** (`tests/dfx/reliability/`): + - `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default). + - `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery). + - `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4). + - `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`). + - `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready. + - `README.md` — Short local run commands for this directory. #### 4.4 Execution Method and Example -- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower. -- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment. -- ***Script Example***: +- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower. +- ***Run Command***: + - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_.py` alongside a matching JSON config) + - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`). +- ***Script Example***: +
Test Examples -When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). The following illustrates the Qwen3-Omni shape: +##### Stability + +When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). Pair the JSON with `tests/dfx/stability/scripts/test_stability_.py`. The following illustrates the Qwen3-Omni shape: ```json { "test_name": "test_qwen3_omni_stability", "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_config_name": "qwen3_omni.yaml" + "stage_config_name": "qwen3_omni_moe.yaml" }, "benchmark_params": [ { @@ -601,112 +588,61 @@ When you want to add L5-level stability test cases, add or extend the appropriat } ``` -##### Parameter Explanation +###### Parameter Explanation ***Overview*** | Field | Required | Description | | ---------------- | -------- | --------------------------------------------------------------------------- | | test_name | Yes | Unique identifier for the stability test case | -| server_params | Yes | Server-side configuration parameters (model, stage configuration, etc.) | +| server_params | Yes | Server-side configuration parameters (model, deploy configuration, etc.) | | benchmark_params | Yes | Stability benchmark running parameters (supports multiple configurations) | -##### server_params Configuration +###### server_params Configuration -###### Basic Parameters +Basic parameters: | Parameter | Required | Example | Description | | ----------------- | -------- | ---------------------------------- | ----------------------------------- | | model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path | -| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name | +| stage_config_name | Yes | "qwen3_omni_moe.yaml" | Deploy configuration file name | -###### Dynamic Configuration (update/delete) - -Supports incremental modifications based on the basic configuration: +Dynamic configuration (`update` / `delete`) supports incremental modifications based on the basic configuration: | Operation | Description | | --------- | ------------------------------------ | | update | Update or add configuration items | | delete | Delete specified configuration items | -***Example***: -You can refer to Test Examples in L4 §3.4 +***Example***: You can refer to Test Examples in L4 §3.4. -##### benchmark_params Configuration +###### benchmark_params Configuration For stability testing, the key parameters are: -- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached. -- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch: - - `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate. - - `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally. -- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`. +- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached. +- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch: + - `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate. + - `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally. +- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`. All other optional parameters follow the same rules as in L4 §3.4. -
+##### Reliability -
- Reliability test suite (tests/dfx/reliability) +Add pytest modules under `tests/dfx/reliability/` (for example `test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`). Reuse `helpers.py` and `conftest.py` rather than duplicating OOM / process-kill / raw HTTP plumbing. -##### Purpose and relationship to stability - -Reliability tests are **short fault-injection** integration runs (L5 **(b)** in `tests/dfx/reliability/README.md`). They complement **stability** JSON-driven long runs: instead of hours of steady traffic, they **perturb** the server (GPU OOM sidecar, fatal signals on selected processes) and check **failure mode** and **latency bounds** (e.g. chat or `/v1/videos` must not hang under concurrent fault-time load). - -##### Directory layout - -| Path | Responsibility | -| ---- | -------------- | -| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). | -| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). | -| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). | -| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). | -| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. | -| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. | -| `README.md` | Minimal run / collect examples. | - -##### Parametrization and markers +###### Parametrization and markers - Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`. -- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4). +- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100). - **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**. -##### CI trigger - -Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model suite (trigger: `WEEKLY=1` or PR label `weekly-test`), for example: - -| Buildkite step | Test file | CI hardware | -| -------------- | --------- | ----------- | -| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) | -| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) | -| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) | -| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) | - -```bash -pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" -pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" -pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow" -pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow" -``` - -##### Local commands - -```bash -pytest --collect-only tests/dfx/reliability -pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow -pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow -pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow -pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow -``` - -##### Adding a new model suite +###### Adding a new model suite 1. Add `test_reliability_.py` under `tests/dfx/reliability/`. -2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e stage-config directory (same pattern as existing files). -3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs stage YAML. +2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e deploy-config directory (same pattern as existing files). +3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs deploy YAML. 4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/cuda/test-weekly.yml`** when the suite should run in weekly L5.
- -- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_.py` alongside a matching JSON config) - - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models) diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 4698caf9c..42b9a53d7 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -9,7 +9,7 @@ This guide walks through the process of adding a new multi-stage model to vLLM-O 3. [Step-by-Step Implementation](#step-by-step-implementation) 4. [Key Components](#key-components) 5. [Model Registration](#model-registration) -6. [Stage Configuration](#stage-configuration) +6. [Pipeline and Deploy Configuration](#pipeline-and-deploy-configuration) 7. [Stage Input Processors](#stage-input-processors) 8. [Testing](#testing) 9. [Adding a Model Recipe](#adding-a-model-recipe) @@ -316,18 +316,28 @@ _OMNI_MODELS = { The registry uses lazy loading, so the model class is imported only when needed. -## Stage Configuration +## Pipeline and Deploy Configuration -Create a YAML configuration file in `vllm_omni/deploy/`. For a complete example, see the [Qwen3-Omni configuration file](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml). +Define fixed topology in `vllm_omni/model_executor/models//pipeline.py` +with `PipelineConfig` and `StagePipelineConfig`, then register it in +`vllm_omni/config/pipeline_registry.py`. For a complete example, see the +[Qwen3-Omni pipeline](gh-file:vllm_omni/model_executor/models/qwen3_omni/pipeline.py). -### Key Configuration Fields +### Pipeline topology 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.) +- **`model_stage`**: Logical stage name ("thinker", "talker", "code2wav", etc.) +- **`execution_type`**: Runtime family (`LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`) +- **`model_arch`**: Stage architecture name, when different from the pipeline default +- **`input_sources`**: Upstream stage IDs that provide input to this stage +- **`custom_process_input_func`**: Function that processes inputs from previous stages +- **`final_output`**: Whether this stage produces user-facing output +- **`final_output_type`**: Output modality ("text", "audio", "image", etc.) +- **`owns_tokenizer`**: Whether this stage owns the pipeline tokenizer + +Create the matching runtime defaults in `vllm_omni/deploy/.yaml` and set +`PipelineConfig.default_deploy_config_name` to that file. The deploy YAML owns +GPU placement, memory sizing, parallelism, connectors, and sampling defaults; +it must not redefine topology fields. ## Stage Input Processors @@ -374,7 +384,7 @@ The stage transition process follows these steps: 2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it -3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition: +3. **Input Processing**: The stage input processor registered by the pipeline (implemented under `vllm_omni/model_executor/stage_input_processors/`) handles the transition: ```python def process_engine_inputs( self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None @@ -577,7 +587,7 @@ 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 +4. **Define pipeline topology and deploy defaults** in `pipeline.py` and `vllm_omni/deploy/` 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) @@ -590,7 +600,8 @@ For a complete reference implementation, see: - **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/deploy/qwen3_omni_moe.yaml` +- **Pipeline topology**: `vllm_omni/model_executor/models/qwen3_omni/pipeline.py` +- **Deploy config**: `vllm_omni/deploy/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` @@ -598,4 +609,4 @@ For a complete reference implementation, see: For more information, see: - [Architecture Overview](../../design/architecture_overview.md) - [Supported Models](../../models/supported_models.md) -- [Stage Configuration Guide](../../configuration/stage_configs.md) +- [Pipeline and Deploy Configuration Guide](../../configuration/stage_configs.md) diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 45680ee5d..264287306 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -20,7 +20,7 @@ pattern, refer to MOSS-TTS-Nano. 4. [Step-by-Step Implementation](#step-by-step-implementation) 5. [Key Components](#key-components) 6. [Model Registration](#model-registration) -7. [Stage Configuration](#stage-configuration) +7. [Pipeline and Deploy Configuration](#pipeline-and-deploy-configuration) 8. [Stage Input Processors](#stage-input-processors) 9. [Online Serving Integration](#online-serving-integration) 10. [Single-Stage Models](#single-stage-models) @@ -78,7 +78,7 @@ Profile before optimizing. new TTS model must pass all three levels: | Layer | Catches | Tool | -|-------|---------|------| +| ------- | --------- | ------ | | Offline RTF / duration | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e | | Browser streaming playback | Delta-vs-cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` | | Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke with 4+ parallel prompts | @@ -94,10 +94,10 @@ the symptom is crosstalk or truncation under load, nothing in single-request tes vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently and can be placed on different devices. Qwen3-TTS has two stages: -| Stage | Name | Input | Output | -|-------|------|-------|--------| -| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes | -| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform | +| Stage | Name | Input | Output | +| ----- | ------------------- | --------------- | ------------------------ | +| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes | +| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform | Each stage is a separate model class configured independently via YAML. The two stages are connected by the `async_chunk` framework, which enables inter-stage streaming for @@ -170,7 +170,7 @@ and Qwen3-Omni). When adding a new TTS model, create the following structure: -``` +```text vllm_omni/model_executor/models/ your_model_name/ __init__.py @@ -201,7 +201,7 @@ it would not be picked up — the hub README is the documented surface. **Qwen3-TTS reference files:** | File | Purpose | -|------|---------| +| ------ | --------- | | `models/qwen3_tts/qwen3_tts.py` | Unified model class | | `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR | | `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder | @@ -427,112 +427,93 @@ _OMNI_MODELS = { The registry uses lazy loading - model classes are only imported when needed. -## Stage Configuration +## Pipeline and Deploy Configuration -Each stage has a `worker_type` that determines how it is scheduled: +Define fixed topology in `models//pipeline.py`. The execution type +selects the scheduler and worker family: -- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention -- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler` +- `StageExecutionType.LLM_AR` uses the autoregressive runtime with PagedAttention. +- `StageExecutionType.LLM_GENERATION` uses the generation runtime for stages such as audio decoders. -Key configuration fields: +Key topology fields belong to `StagePipelineConfig`: | Field | Description | -|-------|-------------| -| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) | -| `model_arch` | Architecture name, must match `registry.py` | -| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) | -| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final | -| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) | -| `final_output` | Whether this stage produces the final user-facing output | -| `final_output_type` | Type of final output (`audio`, `text`, etc.) | +| ------- | ------------- | +| `model_stage` | Logical stage name (`ar_stage`, `decoder`, etc.). | +| `execution_type` | Runtime family for this stage. | +| `input_sources` | Upstream stage IDs that provide input. | +| `model_arch` | Stage-specific architecture name, when it differs from the pipeline default. | +| `engine_output_type` | Output representation such as `latent` or `audio`. | +| `custom_process_next_stage_input_func` | Processor for full-payload stage handoff. | +| `async_chunk_process_next_stage_input_func` | Processor for streaming chunk handoff. | +| `final_output` / `final_output_type` | Whether this stage produces user-facing output and its modality. | +| `owns_tokenizer` | Whether this stage supplies the pipeline tokenizer. | -!!! note - New in-tree models should define frozen topology in `models//pipeline.py`, - register it in `vllm_omni/config/pipeline_registry.py`, and put deployment - knobs in `vllm_omni/deploy/.yaml`. Use `--deploy-config` to load a - custom deployment YAML. +For example, a two-stage TTS topology can be defined as: -### Batch mode +```python +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) -```yaml -# stage_configs/your_model_name.yaml +_PROC = "vllm_omni.model_executor.stage_input_processors.your_model_name" -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - devices: "0" - engine_args: - model_stage: ar_stage - max_num_seqs: 64 - model_arch: YourTTSModelForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - engine_output_type: latent - default_sampling_params: - temperature: 0.9 - top_k: 50 - max_tokens: 2048 - - - stage_id: 1 - stage_type: llm - runtime: - devices: "0" - engine_args: - model_stage: decoder - model_arch: YourTTSModelForConditionalGeneration - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - engine_output_type: audio - engine_input_source: [0] - final_output: true - final_output_type: audio +YOUR_TTS_PIPELINE = PipelineConfig( + model_type="your_model_name", + default_deploy_config_name="your_model_name.yaml", + model_arch="YourTTSModelForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="ar_stage", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + custom_process_next_stage_input_func=f"{_PROC}.ar2decoder", + async_chunk_process_next_stage_input_func=f"{_PROC}.ar2decoder_async_chunk", + ), + StagePipelineConfig( + stage_id=1, + model_stage="decoder", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + model_arch="YourTTSDecoder", + engine_output_type="audio", + final_output=True, + final_output_type="audio", + ), + ), +) ``` -### Streaming mode (async_chunk) - -Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func` -on Stage 0 to define how intermediate outputs are chunked and forwarded: +Register the pipeline in `vllm_omni/config/pipeline_registry.py`. Runtime +placement and sizing belong in `vllm_omni/deploy/your_model_name.yaml`: ```yaml -# stage_configs/your_model_name_async_chunk.yaml - async_chunk: true -stage_args: +stages: - stage_id: 0 - stage_type: llm - runtime: - devices: "0" - engine_args: - model_stage: ar_stage - max_num_seqs: 64 - model_arch: YourTTSModelForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - engine_output_type: latent - custom_process_next_stage_input_func: > - vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk + devices: "0" + max_num_seqs: 64 default_sampling_params: temperature: 0.9 top_k: 50 max_tokens: 2048 - stage_id: 1 - stage_type: llm - runtime: - devices: "0" - engine_args: - model_stage: decoder - model_arch: YourTTSModelForConditionalGeneration - worker_type: generation - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - engine_output_type: audio - engine_input_source: [0] - final_output: true - final_output_type: audio + devices: "0" + max_num_seqs: 1 ``` +Set `async_chunk: false` for end-to-end handoff. The registered pipeline +selects the full-payload or async chunk processor; processor paths and topology +must not be moved into the deploy YAML. Use `--deploy-config` to load a custom +deployment overlay. + ## Stage Input Processors Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in @@ -547,9 +528,9 @@ Understanding what's available in stage outputs: - `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects - Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects - Each `RequestOutput` has: - - `token_ids` - generated token IDs - - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs` - - `prompt_token_ids` - original prompt token IDs + - `token_ids` - generated token IDs + - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs` + - `prompt_token_ids` - original prompt token IDs ### Batch mode (non-streaming) @@ -629,6 +610,7 @@ def ar2decoder_async_chunk( ``` Key points: + - `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle - Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps - The final (possibly partial) chunk is flushed when `is_finished` is true @@ -683,6 +665,7 @@ Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`: ```python # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """YourModel serving adapter.""" from typing import TYPE_CHECKING @@ -821,7 +804,7 @@ directly from the model's own generator. ### Directory structure -``` +```text vllm_omni/model_executor/models/your_model_name/ __init__.py modeling_your_model_name.py # unified class: load_weights + forward + streaming @@ -832,32 +815,40 @@ vllm_omni/deploy/your_model_name.yaml No stage input processor is needed. -### Stage config +### Pipeline and deploy config -Use a single stage with `worker_type: ar`. The `is_comprehension: true` field and the -top-level `async_chunk: false` are required — omitting them causes silent -misclassification in the serving layer. Set `max_num_seqs` to at least 4 for -concurrent production use. +Use a single `LLM_AR` stage that owns the tokenizer and produces final audio: + +```python +YOUR_TTS_PIPELINE = PipelineConfig( + model_type="your_model_name", + default_deploy_config_name="your_model_name.yaml", + model_arch="YourModelForCausalLM", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="your_model_stage_key", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="audio", + final_output=True, + final_output_type="audio", + ), + ), +) +``` + +The deploy YAML controls runtime behavior. Set `max_num_seqs` to at least 4 +for concurrent production use: ```yaml -# stage_configs/your_model_name.yaml async_chunk: false -stage_args: +stages: - stage_id: 0 - stage_type: llm - is_comprehension: true # required for serving_speech.py dispatch - runtime: - devices: "0" - engine_args: - model_stage: your_model_stage_key - model_arch: YourModelForCausalLM - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - engine_output_type: audio - max_num_seqs: 4 # min 4 for concurrent requests; default 1 causes gaps - final_output: true - final_output_type: audio + devices: "0" + max_num_seqs: 4 ``` ### Generator-based streaming pattern @@ -961,10 +952,17 @@ is correct. Stage the modified files and commit again — do not revert the chan Common failures and fixes: | Check | Cause | Fix | -|-------|-------|-----| +| ------- | ------- | ----- | | `ruff F841` | Local variable assigned but never used | Remove the extraction or forward it to the model call | | `ruff E402` | Module-level import not at top of file | Move import to the top-level import block | | `ruff format` | Line length, spacing, or quote style | Accept the auto-fix, stage, and re-commit | +| `check-spdx-header` | Missing header, or copyright still says `vLLM project` | Use `Copyright contributors to the vLLM-Omni project`; restage the rewrite | +| `check-forbidden-imports` | Stdlib `re`/`base64`, pickle, Hugging Face Hub API, or direct Triton/TileLang | `import regex as re` and `pybase64`; do not grow the allowlist without review | +| `check-torch-cuda-call` | New `torch.cuda.*` call site | Use `current_omni_platform`; do not grow `ALLOWED_FILES` without review | +| `check-tts-adapter-migration` | New `_tts_model_type` branch in `serving_speech.py` | Put logic in `tts_adapters/`; lower `MAX_MODEL_TYPE_BRANCHES` when removing branches | +| `check-mark` | Test file missing level/hardware mark, or direct `pytest.mark.H100` | Add `core_model`/… plus `cpu`/`cuda`/`hardware_test(`; SKU marks only via helpers | +| `shellcheck` | Native `shellcheck` missing, or script warning | Install via apt/dnf/brew (or `shellcheck.exe` on Windows). See [Linting](../README.md#linting) | +| `mypy-3.10` / `markdownlint-cli2` / `check-buildkite` | Types, docs markdown, or Buildkite YAML | See [Linting](../README.md#linting) | ### DCO sign-off @@ -1019,7 +1017,7 @@ Adding a TTS model to vLLM-Omni involves: ### Qwen3-TTS Reference Files | File | Purpose | -|------|---------| +| ------ | --------- | | `models/qwen3_tts/qwen3_tts.py` | Unified model class | | `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops | | `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` | @@ -1031,4 +1029,4 @@ For more information, see: - [Architecture Overview](../../design/architecture_overview.md) - [Async Chunk Design](../../design/feature/async_chunk.md) -- [Stage Configuration Guide](../../configuration/stage_configs.md) +- [Pipeline and Deploy Configuration Guide](../../configuration/stage_configs.md) diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index b57de441e..62fff29c5 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -322,6 +322,15 @@ Each 1-second window records: On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima). +`loop_idle` and `loop_active` count loop iterations, so their absolute scale is +tied to which orchestration loop is running. Under the default poll loop an idle +orchestrator records roughly one iteration per millisecond. Under the +event-driven loop (`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, see +[Speech API](../serving/speech_api.md#orchestration-loop-experimental)) an idle +orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still +records one iteration per routed output. Window counts and the `loop_active_pct` +summary are therefore not comparable across the two modes. + ### Relationship to other diagnostics This monitor is intentionally separate from the existing profiling tools: diff --git a/docs/design/ar_diffusion_pipeline_capability.md b/docs/design/ar_diffusion_pipeline_capability.md index c9e05c8e8..3c1b4b9cd 100644 --- a/docs/design/ar_diffusion_pipeline_capability.md +++ b/docs/design/ar_diffusion_pipeline_capability.md @@ -4,6 +4,9 @@ > `experimental/ar_diffusion` contract. It follows the ownership and lifecycle > direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137), > but it is not yet part of the stable module-design hierarchy. +> +> The end-to-end request, identity, failure, and memory contracts are documented +> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md). `ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep autoregressive attention KV across requests. The engine only selects @@ -51,14 +54,13 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting those paths requires a batch-aware state-binding contract rather than silently bypassing `bind_ar_diffusion_state()`. -The current implementation retains one session per model instance. A request -for a different `session_id` closes the resident session through the same LRU -lifecycle used by explicit close and failure cleanup before allocating the new -one. Self-attention and cross-attention memory are therefore sized for one -resident session even if the pipeline declares a larger `session_capacity`. -This prioritizes rollout throughput and prevents block-pool exhaustion before a -count-based eviction can run. Multi-session residency can be added later with a -memory-backed capacity calculation or allocation-pressure eviction. +The implementation may retain multiple sessions per model instance. It selects +the largest feasible capacity up to the pipeline's `session_capacity`, based on +self-attention pages, cross-attention allocations, scratch, and declared +model-owned state. One session is admitted whenever it fits actual free device +memory; `gpu_memory_fraction` controls additional residency. Requests beyond +that capacity evict the least-recently-used session through the same lifecycle +hook used by explicit close and failure cleanup. Scratch storage is also capability-driven. `frames_per_block` reserves space for an uncommitted current video block, while `max_scratch_tokens_per_branch` diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index aa6f72770..d7938b62c 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -335,8 +335,8 @@ The single resolution boundary is implemented by `StageConfigFactory.create_from_model()` and `VllmOmniConfig.from_pipeline_config()` in [`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config). -The legacy `stage_args` YAML path remains only for models that have not yet -migrated to `PipelineConfig` and `DeployConfig`. +The legacy `stage_args` YAML path has been removed. Model topology now resolves +through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`. In the typed path, each stage configuration derives from `BaseVllmOmniStageConfig` and is specialized as @@ -406,7 +406,7 @@ the source for configuration, compatibility, and implementation details. stage contract. * **Quantization and memory efficiency:** [Quantization](feature/quantization.md) resolves per-pipeline or per-component quantization configurations, while - [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md) + [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md) streams diffusion blocks from host memory within the existing parallel topology. diff --git a/docs/design/feature/async_chunk.md b/docs/design/feature/async_chunk.md index 57b4209b8..068a5ac1b 100644 --- a/docs/design/feature/async_chunk.md +++ b/docs/design/feature/async_chunk.md @@ -211,38 +211,39 @@ In sequential mode, each stage must wait for the previous stage to complete enti ## Configuration -Enable async_chunk in stage configuration YAML: +Enable async chunk mode in the deploy YAML: ```yaml async_chunk: true -stage_args: - - stage_id: 0 - engine_args: - custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk - - stage_id: 1 - engine_args: - custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk ``` -### Stage Configuration +The registered `PipelineConfig` owns the async handoff processor functions. +For Qwen3-Omni these are `thinker2talker_async_chunk` and +`talker2code2wav_async_chunk`; deploy YAML only selects the mode and runtime +sizing. + +### Configuration ownership - `async_chunk: bool`: Enable/disable async chunk mode -- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk` -- `stage_connector_config: dict`: Connector configuration -- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling) +- `async_chunk_process_next_stage_input_func: str`: Pipeline-owned chunk processor path +- `execution_type: StageExecutionType`: Pipeline-owned runtime family +- `connectors: dict`: Deploy-owned connector definitions - `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage - ### Connector Configuration ```yaml connectors: - - from_stage: 0 - to_stage: 1 - spec: - name: SharedMemoryConnector - extra: - stage_id: 0 + connector_of_shared_memory: + name: SharedMemoryConnector + +stages: + - stage_id: 0 + output_connectors: + to_stage_1: connector_of_shared_memory + - stage_id: 1 + input_connectors: + from_stage_0: connector_of_shared_memory ``` ### Code2Wav Batch Configuration @@ -250,13 +251,10 @@ connectors: For optimal performance with async_chunk, the code2wav stage should be configured with batching: ```yaml -stage_args: +stages: - stage_id: 2 # code2wav stage - runtime: - devices: "1" - engine_args: - model_stage: code2wav - max_num_seqs: 64 # Enables batched audio generation + devices: "1" + max_num_seqs: 64 # Enables batched audio generation ``` ## Related Files diff --git a/docs/design/feature/async_diffusion_output.md b/docs/design/feature/async_diffusion_output.md index a81ed6811..733478883 100644 --- a/docs/design/feature/async_diffusion_output.md +++ b/docs/design/feature/async_diffusion_output.md @@ -117,7 +117,22 @@ After (async D2H): ### Batch Split -When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. +When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output. + +### Reliability, Lifecycle & Timeout Behavior + +1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**: + When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation. + +2. **Timeout Diagnostics (`_async_output_timeout()`)**: + The engine waits for `wait_output_ready()` with a timeout of `_async_output_timeout()` — 600.0s by default, overridable per run with the `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` environment variable (a non-numeric or non-positive value is ignored with a warning and the default applies, since this resolves on the request path). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter. The packing itself takes milliseconds, but the wait is queued behind the step's GPU work, so its wall-clock duration tracks step time; the bound is generous because a hung engine is surfaced by the worker monitor and `check_health()`, not by this timeout. + +3. **Queue Serialization & Atomic Future Resolution**: + - `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`). + - `ResultPumpThread` unpacks shared-memory payloads first, and then atomically resolves waiting futures or populates `_completed_outputs` under `_futures_lock`, avoiding race conditions with concurrent `wait_output_ready()` invocations. + +4. **Storage-Aware IPC Packing**: + `pack_diffusion_output_shm()` evaluates `max(view_bytes, storage_bytes)` against `_SHM_TENSOR_THRESHOLD` (1 MB) to ensure that per-request slices sharing a large batch storage are moved to POSIX shared memory rather than serialized through pickle over the MessageQueue inline buffer. ## Model Coverage @@ -131,6 +146,7 @@ All models running in request-mode (`step_execution=False`, the default) automat |---|---|---| | **HunyuanImage-3.0** | Image | `False` | | **Qwen-Image** | Image | `True` | +| **LTX-2.3** | Video / Audio | `False` | Other models with `step_execution=False` are also supported but not yet verified. diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md new file mode 100644 index 000000000..41cbaf3ff --- /dev/null +++ b/docs/design/feature/attention_backend_selection.md @@ -0,0 +1,106 @@ +# Diffusion Attention Backend Selection + +This document defines the selection and extension contract for diffusion +attention backends. User-facing backend choices, installation, and tuning are +in the [attention backend guides](../../user_guide/diffusion/attention_backends.md). + +## Scope + +The contract applies to DiT and other diffusion attention layers. It is +separate from vLLM's autoregressive attention selector. + +The selection path has four responsibilities: + +1. normalize user configuration into `AttentionConfig` and `AttentionSpec`; +2. resolve a spec for an attention role; +3. ask the active platform to validate an explicit backend or choose a + hardware default; and +4. load the selected `AttentionBackend` class from the registry path. + +## Resolution contract + +`get_attn_backend_for_role()` resolves configuration in this order: + +1. exact `per_role[role]`; +2. category `per_role[role_category]`; +3. `default`; +4. platform default. + +An explicit resolution returns both the backend class and its +`AttentionSpec`. A platform-default resolution returns the class and `None`. +Layers must therefore treat the spec as optional and must not infer that a +missing spec means SDPA. + +Class resolution is cached by backend name, head size, and whether TRTLLM may +be selected automatically. Logging is cached separately by role and source so +the startup record identifies why each role received its backend. + +## Registry and platform boundary + +`DiffusionAttentionBackendEnum` maps stable configuration names to default +qualified class paths. `register_diffusion_backend()` may replace a path at +runtime without changing the public enum value. + +The active `OmniPlatform` owns compatibility policy through +`get_diffusion_attn_backend_cls()`. It must: + +- validate explicit selections and fail with an actionable error when the + requested kernel cannot run; +- choose only an available, compatible default when no backend is explicit; +- account for head size and other platform-visible constraints; and +- return a qualified class path implementing `AttentionBackend`. + +The selector must not duplicate device capability or package-availability +policy that belongs to the platform. + +## Typed backend options + +Backend-specific settings remain typed fields on `AttentionSpec` rather than +unstructured keyword dictionaries: + +- `quant` is consumed by FlashInfer and TRTLLM with backend-specific value + validation; +- `skip_softmax` is consumed by TRTLLM; and +- `block_sparse` is consumed by block-sparse backends such as RainFusion; and +- `fastvideo_vsa_topk` is consumed by FastVideo VSA. + +A backend reads only the fields it owns and rejects incompatible values. New +options should be added to a shared typed spec only when more than one backend +shares their semantics; otherwise add a dedicated typed block. + +## Model integration contract + +Each diffusion attention call declares a stable role and, when appropriate, a +broader category. Model-specific roles permit precise overrides; categories +preserve common `self` or `cross` policies. A model also decides whether its +path is compatible with automatic TRTLLM selection, because only the model +knows whether masking and packing satisfy the kernel contract. + +## Adding or changing a backend + +An implementation change is complete only when it: + +1. implements the `AttentionBackend` interface and exposes a stable enum name; +2. adds platform validation and default routing where appropriate; +3. defines and validates any typed configuration; +4. tests explicit selection, default selection, and incompatible paths; +5. documents installation, fallback behavior, and quality implications in the + corresponding user guide; and +6. updates the overview matrix without moving algorithm details back into the + landing page. + +## FastVideo VSA model metadata + +FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in +the original latent video grid. Wan integrations therefore attach the +post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The +backend validates that `T * H * W` equals the sequence length, derives the +runtime block count, and selects the configured top-k key/value blocks for +each query block. + +Checkpoint-specific compensation remains model-owned. Wan layers pass an +optional learned `gate_compress` projection through attention metadata when +the checkpoint contains those weights. Checkpoints without the projection do +not allocate or execute it. When top-k selects every block, native checkpoints +route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path +to retain checkpoint semantics. diff --git a/docs/design/feature/cfg_parallel.md b/docs/design/feature/cfg_parallel.md index 053360494..db8f738c1 100644 --- a/docs/design/feature/cfg_parallel.md +++ b/docs/design/feature/cfg_parallel.md @@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. ### N-Branch CFG (3+ branches) -Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches. +Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches. `predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`): diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md index c31086c0d..ba8c8f9fc 100644 --- a/docs/design/feature/disaggregated_inference.md +++ b/docs/design/feature/disaggregated_inference.md @@ -59,19 +59,18 @@ This `metadata` must be passed through the control plane so `get()` can locate t ## Configuration Model -Define connectors in runtime: +Define connectors at the top level of the deploy YAML: ```yaml -runtime: - connectors: - connector_of_shared_memory: - name: SharedMemoryConnector +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector ``` Wire stages to connectors: ```yaml -stage_args: +stages: - stage_id: 0 output_connectors: to_stage_1: connector_of_shared_memory diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md deleted file mode 100644 index da735d60b..000000000 --- a/docs/design/feature/distributed_layerwise_offload.md +++ /dev/null @@ -1,156 +0,0 @@ -# Distributed Layerwise Offload - -This document describes distributed layerwise offload (DLO) for diffusion -models. DLO keeps only a small number of DiT blocks on the accelerator and -streams the remaining blocks from host memory. The distributed backend can -either shard those host-side weights across an existing parallel group or keep -the standard loader's rank-local weights and avoid an additional collective. - -For user-facing commands, see the -[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md) -and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md). - -## Status - -DLO is implemented for multi-device diffusion execution. The default -AllGather path is the primary path for DP and SP deployments. The -`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful -for standard-loader sharding, workstation bring-up, and systems where an -additional DLO collective is undesirable, but it does not reduce host weight -storage across ranks. - -The compatibility matrix below describes the current implementation. The -unit-level guards are covered, but not every parallelism combination has a -full model-and-hardware end-to-end test. - -## Design - -### DLO consumes the existing parallel topology - -DLO does not create a new DP, TP, or SP topology. It reads the configured -`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the -standard distributed groups have been initialized. - -The DLO weight-sharding group is selected as follows: - -1. Use the existing DP group when `data_parallel_size > 1`. -2. When DP is one and SP is greater than one, use the SP group. -3. Otherwise, run rank-locally without a DLO process group. - -TP is deliberately not used as DLO's AllGather group. HSDP has its own -parameter-sharding lifecycle and is not allowed to be sharded a second time by -DLO's AllGather path. - -### AllGather path - -With the default `dlo_use_allgather=True`, each rank stores approximately -`1 / group_size` of each streamable block in pinned host memory. The next -block's shard is copied to a device buffer and reconstructed with -`all_gather_into_tensor` on a communication stream while the current block is -executing. - -```text -Compute: [Block N] [Block N+1] [Block N+2] -H2D: [shard N+1] [shard N+2] -AllGather: [full N+1] [full N+2] -Buffers: [current slot] [prefetch slot] [current slot] -``` - -![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif) - -The backend uses two shared device buffers, so accelerator weight residency is -bounded by the largest streamed blocks rather than the complete model. - -When DP is greater than one, the engine can process one request per DP rank in -the same denoising wave. Because AllGather is a collective, all participating -requests must take the same execution path at every denoising step. - -### Rank-local path without DLO AllGather - -With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to -one. The regular model loader remains responsible for preparing each rank's -weights, including TP-local tensors or HSDP-managed parameters. DLO then -streams those rank-local tensors block by block using H2D copies only. - -This mode means: - -- DP still provides independent replicas, but DLO does not shard weights - across DP ranks. -- SP still performs its normal activation/attention collectives, but DLO does - not shard weights across SP ranks. -- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only - DLO's additional weight AllGather is disabled. -- Pure DP deployments keep a full host-side model copy per rank, subject to - shared OS page-cache behavior. -- The scheduler does not require a synchronized DP request wave for DLO. - -This path is intentionally not implemented by reusing the AllGather mmap -loader. It relies on the standard loader so model-specific TP/HSDP transforms -remain intact. - -## Parallelism compatibility - -| Parallelism | DLO + AllGather | DLO without AllGather | -|---|---|---| -| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. | -| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. | -| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. | -| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. | - -### Combined dimensions - -- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater - than one; SP continues to use its own sequence-parallel group. If DP is one, - the SP group becomes DLO's sharding group in AllGather mode. -- **DP + TP/SP without AllGather:** standard model loading defines the - rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight - collective. -- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but - DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight - materialization and synchronization. -- **HSDP + DP or TP:** rejected independently by the diffusion parallel - configuration. - -## Request and loading constraints - -AllGather DP multi-concurrency requires: - -- explicit `num_inference_steps`; -- the same `num_inference_steps` for all requests in a wave; and -- identical request arguments that affect the collective execution path. - -The no-AllGather path does not impose these DLO-specific synchronized-wave -requirements. - -The mmap loader is used only by the supported DLO+AllGather path when the -model has a compatible checkpoint layout. Online quantization is incompatible -with that sharded mmap path; use `--dlo-no-use-allgather` or disable online -quantization. Models that do not meet the mmap requirements use the regular -loader path. - -## Validation coverage - -Current source-level validation includes: - -- HSDP + DLO + AllGather rejection; -- HSDP + DLO without AllGather acceptance at configuration level; -- TP rejection in the DLO+AllGather mmap path; -- resident-layer requests requiring no-AllGather; -- DP request-wave validation for denoising-step compatibility; -- sharding, double-buffer, AllGather-size, and heterogeneous-block regression - tests. - -The highest-value missing coverage is end-to-end numerical comparison against -ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather -on the target CUDA/NCCL or CANN/HCCL hardware. - -## Recommendations - -- Use **DP + DLO AllGather** for the supported throughput and host-memory - scaling path. -- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is - not the goal. -- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with - the expectation of higher host-memory use and lower validation confidence. -- Prefer **HSDP alone** for production HSDP deployments until the combined - HSDP + DLO no-AllGather path has broader end-to-end coverage. diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md new file mode 100644 index 000000000..a9bb1b0b8 --- /dev/null +++ b/docs/design/feature/host_weight_runtime.md @@ -0,0 +1,386 @@ +# Host Weight Runtime + +The Host Weight Runtime is a loader-facing feature for reusing immutable, +runtime-ready model weights across workers on one host. It lets a model +integration replace repeated checkpoint loading and transformation with an +exact lookup of a previously published host representation, while keeping the +canonical model source authoritative. + +This feature is shared infrastructure. Diffusion, autoregressive, speech, and +future model integrations may depend on the top-level +`vllm_omni.host_weight_runtime` package without depending on one another's +loader or execution stack. The internal component and filesystem contracts are +specified in the [Host Weight Runtime module design](../module/host_weight_runtime.md). + +## Status + +The implementation provides contracts and a CPU local-filesystem store. The +diffusion consumer additionally defines typed, +representation-independent final-layout identity/restoration mechanics plus a +concrete BF16-with-preserved-FP32 policy for MiniMax H3 and +`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration +selects, publishes, restores, and transfers these artifacts. + +V1 includes: + +- exact runtime-weight identity and immutable manifests; +- coordinated local lookup and one-producer publication; +- descriptor-backed safetensors mmap leases; +- preferred and required resolution policy; +- explicit, separately reported post-load publication; +- validation, deny, quarantine, cleanup, and capacity controls; and +- typed reports for every terminal resolution outcome. + +V1 does not include: + +- default-on activation or consumers outside no-AllGather DLO; +- online FP8, quantized, merged-adaptation, or additional model producers; +- HWR interaction with DLO AllGather; +- a remote artifact provider or cross-node coordination; +- automatic eviction; or +- a change to DLO collective or execution behavior. + +## Motivation and use cases + +Model loading can create the same final host representation repeatedly. This +is especially expensive when loading performs checkpoint decoding, tensor +renaming, TP slicing, quantization, packing, or scale construction before GPU +transfer. Independent workers may then retain private CPU copies of identical +runtime weights. + +The Host Weight Runtime is useful when: + +- multiple same-node workers request semantically identical final weights; +- constructing those weights is materially more expensive than validating and + mapping a local artifact; +- the runtime layout can be identified exactly and reproduced deterministically; +- the workers can share immutable file-backed pages through the OS page cache; + and +- canonical loading remains available when policy permits fallback. + +Typical consumers include same-node diffusion replicas, DLO host-weight +sources, and future transformed or quantized model loaders. Request routing and +replica orchestration remain outside this feature. + +This feature is not zero-copy GPU execution. A lease exposes stable CPU tensor +views and mapped ranges. A separate transport still decides whether to use +registered mmap, private pinned staging, synchronous copies, or asynchronous +H2D transfer. + +## Resolution behavior + +The loader resolves the immutable canonical source and computes the exact +requested identity before asking the runtime for a representation. + +```mermaid +flowchart TD + L["Loader resolves canonical source"] --> I["Compute exact WeightArtifactIdentity"] + I --> M{"Runtime mode"} + M -->|"disabled"| C["Canonical loader"] + M -->|"preferred or required"| A{"Validated local artifact?"} + A -->|"yes"| R["Return HostWeightLease"] + A -->|"no"| X["Remote exact artifact (future)"] + X --> P{"Registered producer available?"} + P -->|"yes"| B["Build, validate, and publish atomically"] + B --> R + P -->|"no or recoverable failure"| F{"Policy permits fallback?"} + F -->|"preferred"| C + F -->|"required"| E["Fail startup"] + P -->|"nonretryable failure"| E + C --> W{"Post-load publication enabled?"} + W -->|"yes"| B2["Publish final model through POST_LOAD_ONLY producer"] + B2 --> CL["Close validated publication lease"] + CL --> R2["Return separate publication report"] + R2 --> D + W -->|"no"| D["Keep canonical model"] +``` + +The modes are: + +| Mode | Observable behavior | +| --- | --- | +| `disabled` | Use canonical loading directly without probing storage, identity, credentials, or topology. | +| `preferred` | Return an exact lease when available; otherwise use canonical fallback only for a miss, invalid cache entry, or typed retryable failure. | +| `required` | Fail when an exact lease cannot be acquired. | + +During pre-load resolution, an unsupported capability, semantic identity +collision, producer failure, or publication failure is nonretryable and remains +visible even in preferred mode. A post-load publication failure is likewise +visible in its own report, but cannot revise the canonical-fallback outcome. +Storage policy must never disguise a semantic or configuration error as a cache +miss. + +## Loader and restoration sequence + +The model integration owns the end-to-end loading transaction. The runtime +does not construct a model or invoke the canonical loader. + +```mermaid +sequenceDiagram + participant L as Loader adapter + participant R as HostWeightRuntime + participant S as HostWeightStore + participant P as WeightProducer + participant X as WeightRestorer + participant T as GPU transport + + L->>L: Resolve canonical revision and exact identity + L->>R: resolve(identity, producer) + R->>S: lookup(exact identity) + alt validated hit + S-->>R: HostWeightLease + R-->>L: lease and LOCAL_HIT report + L->>X: plan_restore(model, lease) + X-->>L: validation-only restore plan + L->>X: commit() once + else miss with allowed producer + R->>S: get_or_build(identity, producer) + S->>P: produce(store-scoped writer) + P-->>S: final-layout tensors and metadata + S-->>R: validated HostWeightLease + R-->>L: lease and LOCAL_PRODUCTION report + L->>X: plan_restore(model, lease) + X-->>L: validation-only restore plan + L->>X: commit() once + else policy permits canonical fallback + R-->>L: CANONICAL_FALLBACK + L->>L: Run canonical loader + opt explicit post-load publication enabled + L->>R: publish_after_load(identity, POST_LOAD_ONLY producer) + R->>S: get_or_build(identity, producer) + S-->>R: validated HostWeightLease or typed failure + R->>R: close publication lease + R-->>L: separate publication report + end + else required or nonretryable failure + R-->>L: FAILED report + end + L->>T: consume final CPU tensors and any lease + T-->>L: transfer teardown complete + L->>L: close lease when one was acquired +``` + +`plan_restore()` must not mutate the model or lease. `commit() -> None` is the +sole one-shot model mutation. If planning fails, canonical fallback may reuse +the untouched model. If commit begins and fails, the partially hydrated model +must be discarded and canonical fallback must construct a fresh model. + +`publish_after_load()` is synchronous in V1 and accepts only a +`POST_LOAD_ONLY` producer. Its report is separate from the terminal resolution, +so a publication failure cannot rewrite a successful canonical fallback. A +successful publication closes the store-returned lease inside the runtime and +warms only future startups. It does not restore, rebind, or otherwise mutate the +canonically loaded model serving the current startup. +`allow_local_build` gates producers during pre-load resolution, while +`allow_post_load_publish` independently gates this explicit post-load path. + +## Exact representation identity + +The store performs exact lookup. It never silently converts one representation +into another or substitutes the first backing that responds. + +Identity includes: + +- immutable model revision and source fingerprint; +- model component and ownership boundary; +- representation name, dtype, and format metadata; +- final tensor layout and semantic parallel coordinates; +- static adaptation identity; and +- producer implementation, manifest, and restorer schema versions. + +The requested representation is selected by the loader. The store chooses +where to obtain that exact representation: validated local artifact, future +remote materialization, or a registered producer. + +### Parallelism + +Parallel coordinates are included only when they change bytes, shape, layout, +or component ownership. + +| Dimension | Identity and sharing rule | +| --- | --- | +| DP | Exclude DP rank for replicated weights so same-node replicas can acquire the same artifact. Include a coordinate only if DP changes weight ownership. | +| TP | Include TP size and rank when each rank owns different tensor slices or packed layouts. Each distinct TP shard is a separate exact artifact. | +| SP | Include SP size and backend when they change the weight layout. Exclude SP rank when all SP ranks consume identical bytes. | +| PP | Encode component ownership or stage-local layout when PP changes which weights belong to the consumer. | +| EP | Encode expert ownership and layout whenever ranks own different expert weights. | + +The runtime does not infer these rules from process-group topology. The model +adapter is responsible for constructing a correct identity. + +### Quantization and adaptations + +Runtime BF16, FP8, packed quantized weights, and future formats are distinct +representations. Conversion is performed by an explicitly versioned producer, +never by cache coercion. A producer must publish the final layout consumed by +its matching restorer. + +Dynamic LoRA overlays are not part of a reusable base-weight artifact. A +statically merged adapter is cacheable only as a separate identity containing +the adapter fingerprint and merge semantics. + +### Initial diffusion final-layout contract + +The shared diffusion contract covers complete final-layout DiT parameters and +persistent buffers. Text encoders, VAEs, non-persistent derived state, and other +pipeline components remain outside the artifact. One explicit representation +policy selects allowed dtypes, tensor roles, physical layout identity, producer +ABI, manifest schema, and restoration schema. + +The contract is intentionally separate from loader activation: + +- `FinalLayoutRequest` contains typed loader identity/configuration fingerprints, + TP coordinate, and conservative SP semantics. It has no open metadata bag, + DP coordinate, SP rank, device identity, DLO transfer mode, registration + policy, or store path. +- `FinalLayoutArtifactSpec` binds one `WeightRepresentation` and runtime-layout + name to explicit producer/restorer schemas and a canonical, versioned + implementation ABI descriptor. Compatibility never depends on reflective + source inspection. +- `PreparedWeightSource` snapshots immutable revisions or exact local file + content plus a typed checkpoint-adapter identity before ordinary + materialization. Source replacement before or during production fails + publication. A hash-looking symlink basename is trusted only for an explicit + Hugging Face Hub source whose repository ID and + `models--.../snapshots/ -> blobs/` topology validate; every + local or otherwise unverified symlink target is content-hashed. +- the tensor ownership digest records exact runtime names, kinds, shapes, + semantic roles, dtypes, and strides from a CPU or meta model skeleton; +- `FinalLayoutTensorRestorer` accepts only an exact lease identity, validates + complete policy-defined coverage without mutation, and returns a one-shot + commit plan; +- each model declares one dtype-neutral `FinalLayoutModelContract` with an + explicit implementation version and a post-commit validator; and +- `FinalLayoutBF16Producer` accepts only the matching identity context and a + finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP + coordinate. Its BF16 policy preserves model-declared FP32 parameters and + buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates + FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout. + +Other representations reuse source identity, typed parallel identity, tensor +ownership, and exact restoration only when their policy proves those semantics. +For example, runtime FP8 needs a separate policy/producer for generated scales, +quantization metadata, and Cutlass physical layouts; it is not enabled by +changing a dtype string on the BF16 producer. + +This stage makes no startup, sharing, or DLO performance claim. A following +consumer PR owns disabled/preferred/required precedence, mixed-component loader +transactions, warm-hit restoration, and transactional lease handoff. A TP2 +prewarm deployment will require a matching TP2 producer cohort to populate both +TP-coordinate identities even though the store coordinates each artifact +independently. + +## Host sharing and GPU transport + +Every process receives its own virtual mappings and `HostWeightLease`, but +workers mapping the same local artifact can share physical file-backed pages +through the OS page cache. This avoids independent private tensor allocations; +it does not guarantee that process PSS reports exactly one model copy. + +The ownership boundary is: + +- the store owns immutable artifact files, publication, and lifecycle; +- the kernel owns page-cache residency and placement; +- the lease owns process-local mappings and the shared artifact lock; and +- transport owns page registration, page locking, private staging, H2D copies, + streams, device buffers, and lease release ordering. + +DLO AllGather and no-AllGather remain transport and execution choices. A DLO +integration may consume a lease, but the Host Weight Runtime does not choose a +parallel collective or orchestrate DP requests. See the +[DLO feature design](offloader/distributed_layerwise_offload.md). + +## Locality and NUMA policy + +The V1 backend is node-local. It accepts an allowlist of known local disk +filesystems plus tmpfs/ramfs, records the detected kernel filesystem type, and +rejects known remote or unknown filesystems. NFS, CIFS, Lustre, Ceph, and other +cross-node mounts cannot silently satisfy the local backend because their +page-cache and advisory-lock behavior does not provide the required node-local +contract. + +A NUMA-specific root creates a separate storage domain but does not by itself +place page-cache pages on that NUMA node. A topology-aware integration must +also pin workers, establish local first touch or prefault, perform transport +registration in the same domain, and collect residency evidence. + +Tmpfs artifacts consume host memory and may consume swap. They must be +accounted as memory-backed storage rather than ordinary disk capacity. + +## Concurrency and timeouts + +One process per exact identity owns a build; other workers wait and then acquire +leases for the published artifact. Publication is invisible until all payloads +and metadata are validated, hashed, fsynced, and atomically renamed. + +`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not +cancel synchronous validation, a producer that has already started, or atomic +publication. A hung in-process producer therefore blocks its owning process and +must be handled by external process supervision. Enforceable producer +cancellation requires a future process-isolated producer contract. + +Leases keep a shared artifact lock for their lifetime. A forked child may close +its inherited descriptors but cannot unlock the parent process's lease. Cleanup +uses a nonblocking exclusive lock and reports an active lease instead of +removing live mappings. + +## Failure and fallback contract + +| Condition | Preferred mode | Required mode | +| --- | --- | --- | +| Exact local hit | Return lease | Return lease | +| Local miss | Try later backing or canonical fallback | Try later backing or fail | +| Invalid or denied artifact | Quarantine/rebuild when possible; otherwise canonical fallback | Quarantine/rebuild or fail | +| Retryable lock, domain, or capacity failure | Canonical fallback | Fail | +| Unsupported producer or backend | Fail | Fail | +| Identity collision | Fail | Fail | +| Pre-load producer or atomic publication failure | Fail | Fail | +| Post-load publication failure after canonical fallback | Keep canonical model; report publication failure | Not reached through required resolution | +| Restore planning failure | Canonical fallback may reuse untouched model | Fail | +| Restore commit failure | Discard model; fallback requires a fresh instance | Fail | + +Cache failures should normally reduce performance rather than availability, +but only when the typed failure explicitly permits fallback. Canonical source, +authentication, configuration, and semantic errors remain visible. + +## Consumer integration requirements + +A consumer PR must: + +1. Resolve an immutable canonical model source before cache lookup. +2. Construct the exact final representation identity, including relevant + parallel, quantization, and adaptation semantics. +3. Register a deterministic producer for that identity or support lookup-only + operation. +4. Provide a validation-only restorer with a one-shot commit. +5. Keep GPU transport outside the producer, store, and restorer contracts. +6. Handle preferred fallback and required failure without reusing a model after + a failed restore commit. +7. Retain the lease until every transport operation that may access mapped + memory has completed. +8. Emit and test the terminal resolution report and any separate post-load + publication report. + +Consumer validation must prove: + +- output parity with canonical loading; +- exact warm-hit work avoidance; +- correct TP/SP/PP/EP ownership and layout identity; +- correct BF16, FP8, quantized, or adaptation semantics; +- shared-backing evidence when host-memory savings are claimed; +- clean unmapping, unregistration, and artifact-lock release; and +- startup, host-memory, and transport effects for the claimed deployment. + +## Rollout + +The rollout is intentionally staged: + +1. Land the neutral contracts and local filesystem store. +2. Add one model-specific producer/restorer integration with parity evidence. +3. Connect eligible DLO or other transport paths without moving transport + ownership into the runtime. +4. Add additional final-layout and quantized producers independently. +5. Introduce remote materialization only through the same local lease contract. + +The architectural decisions and deferred work are tracked in +[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414). diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md new file mode 100644 index 000000000..f071196b1 --- /dev/null +++ b/docs/design/feature/offloader/README.md @@ -0,0 +1,80 @@ +# CPU Offloading + +This document defines the shared architecture for diffusion CPU offloading. +The three strategies have separate user and design pages: + +| Strategy | User guide | Design contract | +| --- | --- | --- | +| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) | +| Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) | +| Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) | + +## Strategy selection + +`OffloadConfig.from_od_config()` converts diffusion configuration into one +`OffloadStrategy`. If callers set multiple flags, selection is deterministic: + +1. `DISTRIBUTED_LAYERWISE`; +2. `LAYERWISE`; +3. `MODEL_LEVEL`; and +4. `NONE`. + +The factory logs when a higher-priority strategy disables another. It derives +parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide +a separate offload group size. + +`get_offload_backend()` then validates platform offload support, resolves the +device, and creates exactly one backend. Returning `None` means offloading is +disabled or unsupported and must not leave partially installed hooks. + +## Shared lifecycle + +Every backend implements `OffloadBackend`: + +- `enable(pipeline)` discovers modules, establishes initial residency, and + installs hooks; +- `disable()` removes owned hooks and resources; and +- `is_enabled()` reports lifecycle state. + +`disable()` does not promise to restore the pipeline's original device +placement. The caller owns any subsequent rematerialization. + +Hooks are registered through `HookRegistry` and `ModelHook`; offload backends +must use distinct hook names and remove only hooks they own. + +## Discovery and topology + +Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: + +- `_dit_modules`; +- `_encoder_modules`; +- `_vae_modules`; and +- `_resident_modules`. + +Dotted paths are supported. Legacy pipelines may use the fallback scan of +well-known attribute names, but new integrations should declare components +explicitly. + +Ordinary layerwise offload obtains topology from declared +`_layerwise_offload_blocks_attrs` (with a deprecated singular-name +compatibility path). Distributed layerwise offload may instead use a pipeline +`OffloadPlan`, then falls back to declared block attributes and its documented +compatibility discovery. Discovery metadata describes structure only; the +backend remains responsible for transfer, synchronization, and storage +ownership. + +## Cross-strategy invariants + +- At most one strategy owns offload hooks for a pipeline. +- A parameter has one authoritative host representation while offloaded. +- Device storage is not freed until dependent compute or transfer work is + complete. +- Non-persistent and model-specific buffers remain correct after movement or + checkpoint rematerialization. +- Unsupported parallel or loading combinations fail before hooks mutate the + model. +- Platform streams, events, synchronization, and cache management go through + the vLLM-Omni platform abstraction. + +The diffusion [Offloader module design](../../module/diffusion/offloader.md) +describes how these feature contracts fit into the larger diffusion runtime. diff --git a/docs/design/feature/offloader/cpu_offload.md b/docs/design/feature/offloader/cpu_offload.md new file mode 100644 index 000000000..2318611cc --- /dev/null +++ b/docs/design/feature/offloader/cpu_offload.md @@ -0,0 +1,56 @@ +# Model-Level Offload + +Model-level offload enforces mutual exclusion between pipeline component +groups. For commands and model support, see the +[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md). + +## Generic hook path + +`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It +makes encoders and VAEs device resident for initialization, then installs +`SequentialOffloadHook` instances: + +- each DiT hook offloads all encoders and other DiTs before loading itself; +- each encoder hook offloads every DiT before loading itself; and +- VAEs are not part of the mutual-exclusion hook set. + +The hook moves parameters and buffers without recursively calling +`module.to()`. This avoids recursion through modules that retain references to +other transformer blocks. Host tensors are pinned when configured; DTensors +are not pinned through the ordinary tensor path. + +Before forward returns control to the model, the hook synchronizes the active +platform so compute cannot race the transfer. XPU uses blocking device-to-host +copies because its allocator cannot safely evict storage while an asynchronous +copy is pending. + +## Split-model delegation + +If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend +delegates instead of installing generic encoder/DiT hooks. This supports models +whose mutually exclusive components are nested inside one transformer, such +as a reasoner and generator. + +The delegated pipeline owns its internal contexts, but must preserve the same +contract: + +- only the executing component is device resident; +- enable and disable are idempotent; +- transfers use the supplied device, pinning, and HSDP settings; and +- disabling removes all model-local offload hooks or contexts. + +The delegation is currently duck typed. If another independent implementation +is added, formalize the enable/disable pair as a protocol rather than expanding +attribute-name heuristics. + +## Extension invariants + +- New ordinary pipelines declare both DiT and encoder groups through + `SupportsComponentDiscovery`. +- Models with multiple DiTs offload the inactive DiTs as well as encoders. +- Empty or parameterless modules do not trigger transfers. +- Hook removal does not move modules; callers must not assume original + residency is restored. + +Shared selection and lifecycle behavior is defined in the +[CPU Offloading design](README.md). diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md new file mode 100644 index 000000000..2b5f929ed --- /dev/null +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -0,0 +1,555 @@ +# Distributed Layerwise Offload + +This document describes distributed layerwise offload (DLO) for diffusion +models. DLO keeps only a small number of DiT blocks on the accelerator and +streams the remaining blocks from host memory. The distributed backend can +either shard those host-side weights across an existing parallel group or keep +complete rank-local block sources and avoid an additional collective. + +For user-facing commands, see the +[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) +and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md). + +## Feature compatibility + +Host-storage optimization and runtime compatibility are separate decisions. +When direct checkpoint mmap is unavailable, DLO can still use tensors produced +by the ordinary loader. "Compatibility path" below means that fallback is +implemented but has less end-to-end coverage than the primary path. + +Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported. + +| Feature | DLO + AllGather | DLO without AllGather | +| --- | --- | --- | +| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. | +| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. | +| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | +| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. | +| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. | +| **Online INT8 linears** | ✅ Ordinary loader finalizes the int8 weight (contiguous or transposed view) and fp32 scale before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. | +| **Online MXFP8 linears** | ✅ Ordinary loader finalizes the fp8 weight (contiguous) and e8m0 block scale (contiguous or transposed view) before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. | +| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. | +| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. | +| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. | + +See [Parallelism compatibility](#parallelism-compatibility) and +[Request and loading constraints](#request-and-loading-constraints) for the +detailed contracts and validation boundaries. + +## Status + +DLO is implemented for multi-device diffusion execution. The default +AllGather path is the primary path for DP and SP deployments. The +`--dlo-no-use-allgather` path streams complete blocks independently and adds no +DLO weight collective. + +Host storage is selected separately from the transfer protocol. The loader can +produce a direct-checkpoint mmap plan for a proven-compatible runtime layout; +otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on +the same node can share immutable checkpoint pages when direct mmap is +selected, while the ordinary-loader fallback still keeps a private runtime +copy per process. + +The Phase A shared-mmap support boundary is TP1. TP greater than one is an +ordinary-loader compatibility path: DLO can consume the resulting TP-local +tensors, but those configurations do not use checkpoint mmap and must not be +used to claim shared-mmap host-memory savings. + +The compatibility matrix below describes the current implementation. The +unit-level guards are covered, but not every parallelism combination has a +full model-and-hardware end-to-end test. + +## Design + +### DLO consumes the existing parallel topology + +DLO does not create a new DP, TP, or SP topology. It reads the configured +`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the +standard distributed groups have been initialized. + +The DLO weight-sharding group is selected as follows: + +1. Use the existing DP group when `data_parallel_size > 1`. +2. When DP is one and SP is greater than one, use the SP group. +3. Otherwise, run rank-locally without a DLO process group. + +TP is deliberately not used as DLO's AllGather group. HSDP has its own +parameter-sharding lifecycle and is not allowed to be sharded a second time by +DLO's AllGather path. + +### The loader owns host-weight planning + +Before it decides whether ordinary weight materialization can be skipped, the +diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is +accepted only when preflight proves all of the following: + +- every required DiT parameter and persistent buffer has exactly one source; +- runtime names, checkpoint keys, shapes, and dtypes match; +- the runtime topology is TP1 without HSDP or online quantization; and +- every custom loader operation is represented by a loader-owned checkpoint + adapter. + +The exact plan object is handed to DLO. The backend does not rescan checkpoint +files, repeat the capability decision, or reconstruct names from its block +topology. If preflight fails, the loader materializes weights normally and DLO +consumes those runtime tensors. + +The plan owns only dedicated DiT component sources. If a pipeline also exposes +ordinary sources for a text encoder or another non-DiT component, the loader +still consumes those sources and includes their loaded names in its strict +coverage check. Only the source prefixes covered by the plan skip ordinary +materialization. A source that mixes DiT and non-DiT weights fails closed to +the complete ordinary-loader path because it cannot be skipped safely as a +unit. + +This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline +flags such as `_supports_mmap_loading` or parameter attributes for mmap-only +transforms. Model-specific direct-layout knowledge, when required, lives in a +checkpoint adapter beside the ordinary loader. + +### AllGather path + +With the default `dlo_use_allgather=True`, each rank stores approximately +`1 / group_size` of each streamable block in pinned host memory. The next +block's shard is copied to a device buffer and reconstructed with +`all_gather_into_tensor` on a communication stream while the current block is +executing. + +```text +Compute: [Block N] [Block N+1] [Block N+2] +H2D: [shard N+1] [shard N+2] +AllGather: [full N+1] [full N+2] +Buffers: [current slot] [prefetch slot] [current slot] +``` + +![DLO double-buffer prefetch pipeline](../../figures/dlo/dlo_pipeline.gif) + +The backend uses two shared device buffers, so accelerator weight residency is +bounded by the largest streamed blocks rather than the complete model. + +When direct checkpoint mmap is selected, the checkpoint mappings are only the +source used to prepare each rank's persistent shard. They can be closed after +shard preparation. Across the AllGather group, those private shards total +approximately one runtime model copy. + +An effective DLO group size of one performs no collective, even when +`dlo_use_allgather=True`; it follows the rank-local transfer path described +below. + +When DP is greater than one, the engine can process one request per DP rank in +the same denoising wave. Because AllGather is a collective, all participating +requests must take the same execution path at every denoising step. + +### Rank-local path without DLO AllGather + +With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to +one and streams complete blocks using H2D copies only. The host backing may be +either a loader-approved checkpoint mapping or ordinary runtime tensors. + +For direct mmap, each process retains immutable safetensors views and uses two +bounded pinned host staging slots. Processes on the same node that map the same +files share physical checkpoint pages through the OS page cache. This removes +the persistent private full-model copy per pure-DP process, but each process +still packs and transfers every complete block. Sharing is node-local; each +node has its own page cache. + +When direct mmap preflight fails, the regular model loader remains responsible +for preparing each rank's weights, including TP-local tensors or HSDP-managed +parameters. In that fallback, each pure-DP process keeps a private full runtime +copy. + +This mode means: + +- DP still provides independent replicas, but DLO does not shard weights + across DP ranks. +- SP still performs its normal activation/attention collectives, but DLO does + not shard weights across SP ranks. +- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only + DLO's additional weight AllGather is disabled. +- Pure DP deployments share one checkpoint-backed copy per node when direct + mmap is selected; the ordinary-loader fallback keeps one private runtime + copy per rank. +- The scheduler does not require a synchronized DP request wave for DLO. + +### Component allocator-cache retention + +MiniMax-H3 DLO keeps one bounded PyTorch allocator cache across its staged +encoder, DiT, and VAE component boundaries so allocations released by one +component can be reused by the next. The DiT prefetch path and its two shared +device buffers are unchanged. + +After a component is offloaded, cached-but-unallocated memory is retained only +while it is at most 25% of device capacity and at least 5% of device capacity +remains physically free. Crossing either bound releases the allocator cache. +Missing allocator telemetry also releases it conservatively. Component or +staging failure forces a release, and an out-of-memory allocation gets one +retry after release. This is a component-local policy rather than a global +`empty_cache` override, so the executor's unconditional shutdown cleanup is +preserved. + +Retaining allocator blocks trades higher reserved and peak device memory for +lower lifecycle latency. The benefit is device- and workload-dependent, and a +smaller or externally loaded device may cross a bound and return to the +ordinary per-component cache-release behavior. Current performance validation +covers MiniMax-H3 no-AllGather DLO on CUDA. Other MiniMax-H3 DLO transfer +topologies and platforms receive the same bounded policy but are not claimed +performance targets; other DLO models remain unchanged unless they explicitly +adopt it. + +### Final-layout Host Weight Runtime consumer + +> **Status:** PR2 landed in +> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds +> registered direct H2D plus generic HWR publication/source-digest +> optimizations without changing AllGather behavior. + +Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies +only to no-AllGather DLO. Enable it with +`--host-weight-runtime-mode preferred` (or `required`) and +`--host-weight-runtime-root `. + +The modes express operator fallback policy, not different artifact formats: + +- `preferred` consumes an exact hit, but on a miss it allows canonical loading + followed by post-load publication. It is the normal population path. +- `required` consumes the same exact artifact but fails startup on a miss or + unusable artifact. It never invokes canonical DiT fallback or post-load + publication and therefore cannot populate an empty store. + +PR2 has no separate prewarm command. Operators populate one matching producer +cohort per node-local storage domain in `preferred` mode, then restart the same +model revision and parallel layout with `required`. TP coordinates have distinct +identities, while equivalent DP replicas share them. + +The current producer/consumer boundary is the model-declared final-layout BF16 +contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein +contract covers both transformer block stacks, constructor-stable packed QKV +mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps` +buffers. The shared HWR machinery supports ordinary-loader final layouts for +TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence +in this change is TP1-only; its TP2/SP layouts rely on the existing +per-coordinate identity mechanics and remain unmeasured. Online quantization, +HSDP, LoRA/adapted weights, and non-default load formats remain ineligible. +HWR mode +`disabled`, DLO-disabled, and DLO AllGather configurations stop before source +identity or store construction and retain the existing checkpoint-mmap or +ordinary-loader path. + +Eligibility is decided before constructing HWR, resolving canonical sources, +hashing identity inputs, probing a filesystem, or emitting an HWR observer +event: + +```text +DLO disabled / HWR disabled / DLO AllGather + -> zero final-layout HWR interaction + -> preserve current checkpoint-mmap or ordinary-loader behavior +``` + +When final-layout HWR is explicitly enabled for no-AllGather DLO, runtime +outcomes are authoritative: + +- `LOCAL_HIT`: plan and commit an exact restoration, then transfer the lease + transactionally to DLO. +- `CANONICAL_FALLBACK`: bypass checkpoint mmap, canonically materialize and + finalize the DiT, then call `publish_after_load()` for a future startup. The + current startup retains its canonical tensors. +- `FAILED`: fail startup. Preferred mode does not reinterpret nonretryable + identity, configuration, or compatibility failures as misses. +- Required mode cannot bootstrap an empty store through a `POST_LOAD_ONLY` + producer. + +Checkpoint mmap remains an unchanged control path whenever final-layout HWR is +not selected. + +Local source identity and publication avoid repeated full-payload passes while +preserving the same semantic contract: + +- Canonical files without a trusted immutable Hub blob identity are hashed in + parallel. Their content IDs are cached per storage domain under the HWR root + and reused only when the path, inode, size, timestamps, symlink target, and + cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or + cache-I/O failure falls back to hashing the canonical file directly. +- Producers that emit payload bytes in canonical storage-key order may declare + `ordered=True`. The filesystem writer then computes file and tensor SHA256 + values during the write and overlaps payload `fsync` with later producer + work. Producers without that guarantee retain unordered writes and use a + parallel readback checksum fallback. +- Manifest and `READY` publication still wait for every payload `fsync` and + checksum to complete. These optimizations do not weaken artifact validation + or change lease ownership. + +#### Registered mmap transport + +After DLO takes an HWR lease, the no-AllGather backend may register the lease's +complete immutable mapped ranges under the existing `pin_cpu_memory` policy. +Registration is transport state: it does not change artifact identity, the +store, tensor ownership, or H2D payload. On success, each tensor view copies +directly into the existing rotating HBM block buffers and the two private host +staging slots are not allocated. + +`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling +over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned +memory policy, unsupported platform/capability, over-budget mapping, or fully +rolled-back registration error selects the existing two-slot staging path. A +partial registration that cannot be rolled back aborts startup because closing +the lease would unmap memory still owned by the platform. + +Direct checkpoint mmap remains unchanged and continues to use staging. It may +require loader-owned per-block transforms, while the HWR artifact already +contains final runtime bytes. DLO AllGather never receives an HWR final-layout +lease and therefore never enters this registration path. + +#### Pre-service transaction boundary + +The startup transaction does not end at restore commit. A lease-backed model is +disposable until backend setup and initial prefetch complete: + +```text +UNRESOLVED + -> LEASE_OWNED_BY_LOADER + -> RESTORE_PLANNED + -> COMMIT_STARTED + -> CARRIER_OWNED + -> BACKEND_OWNED + -> BACKEND_READY + -> IN_SERVICE +``` + +Failure before `COMMIT_STARTED` closes the loader-owned lease. Preferred mode +may canonically load using the untouched skeleton; required mode fails. + +Failure from `COMMIT_STARTED` through `BACKEND_READY` must: + +1. synchronize or otherwise quiesce partial backend and initial-prefetch work; +2. release hook, staging, restore-plan, and model references; +3. close the lease through its current owner; +4. discard the restored model; and +5. in preferred mode, construct a fresh canonical model while bypassing HWR + lookup, HWR publication, and checkpoint mmap for that recovery attempt. + +Required mode fails instead of constructing the fresh fallback. Once the model +enters service, failures follow normal runtime handling rather than startup +fallback. + +#### Finalization phases + +Cold canonical loading separates byte-changing work from shared runtime-state +finalization: + +```text +Cold-only byte-changing: + casting, reordering, packing, quantization, generated scales, + model-specific weight transforms, and calibration that mutates + parameters or persistent buffers + +Shared cold/warm non-byte: + validation, hook installation, eval state, bookkeeping, + and non-persistent runtime state +``` + +The warm path must snapshot restored tensor bytes and backing pointers before +shared finalization and prove both are unchanged afterward. + +#### Lease ownership + +Lease transfer is single-owner and single-take: + +```text +loader owns + -> restore plan borrows + -> commit succeeds + -> runner carrier takes once + -> backend takes once before asynchronous work + -> backend drains pending H2D work + -> backend releases hooks, staging, and model references + -> backend unregisters every mapped range + -> backend closes lease +``` + +The runner carrier is process-local, rejects serialization, and rejects a +duplicate `take()`. If the carrier still owns the lease, runner cleanup closes +it. Once the backend takes ownership, only backend abort or teardown may drain +work and close it. + +The implementation keeps the final-layout lease through backend setup and +initial prefetch. Successful registration copies directly from immutable HWR +views; otherwise the backend uses the existing two bounded rank-local staging +slots. Transactional cleanup drains device work, removes source references, +unregisters mappings, and only then closes the lease. An unregistration failure +retains both registration and lease for retry/process teardown. Preferred mode +then uses the runner's fresh canonical retry; required mode propagates the +failure. + +#### PR2 and PR3 promotion gates + +- Warm hit performs zero ordinary DiT materialization and zero producer calls. +- Shared warm finalization changes neither restored bytes nor backing pointers. +- Preferred `FAILED` outcomes do not silently fall back. +- Planning failure leaves the skeleton untouched. +- Any pre-service failure after commit begins discards the restored model; + preferred mode uses a fresh canonical model and required mode fails. +- Mixed components load normally and no required tensor remains on `meta`. +- Disabled mode and AllGather emit zero final-layout HWR interaction. +- The lease carrier rejects duplicate take and serialization. +- Backend setup or prefetch failure drains asynchronous work before lease close. +- Registration is all-or-nothing; rollback/unregistration failure never closes + a mapping still owned by the platform. +- Registered HWR transport bypasses host staging while preserving output and + H2D payload; unsupported registration retains the bounded staging path. +- Compatible checkpoint mmap remains an unchanged benchmark and control path. +- Prewarm uses one matching producer cohort per TP/SP topology and storage + domain. +- Benchmarks compare ordinary loading, cold publication, and warm staging with + output parity, startup latency, aggregate PSS, page-cache state, HBM, and H2D + payload. + +## Parallelism compatibility + +| Parallelism | DLO + AllGather | DLO without AllGather | +| --- | --- | --- | +| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. | +| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. | +| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. | +| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. | + +### Combined dimensions + +- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater + than one; SP continues to use its own sequence-parallel group. If DP is one, + the SP group becomes DLO's sharding group in AllGather mode. +- **DP + TP/SP without AllGather:** standard model loading defines the + rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight + collective. When the model declares the final-layout contract, HWR keys the + reusable artifact by TP rank/size and SP layout. +- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but + DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight + materialization and synchronization. +- **HSDP + DP or TP:** rejected independently by the diffusion parallel + configuration. + +## Request and loading constraints + +AllGather DP multi-concurrency requires: + +- explicit `num_inference_steps`; +- the same `num_inference_steps` for all requests in a wave; and +- identical request arguments that affect the collective execution path. + +The no-AllGather path does not impose these DLO-specific synchronized-wave +requirements. + +Direct checkpoint mmap can back either transfer path. It is currently limited +to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the +ordinary loader. Per-tensor online FP8, online INT8, and online MXFP8 linears +can use DLO AllGather after the ordinary loader finalizes their runtime +weights and scales; DLO then shards and reconstructs those tensors with their +recorded layouts (they all keep plain 1-byte dtypes over ordinary strided +views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an +fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view +finalized tensor, all covered by the same physical-order packing as online +FP8). Other online methods must use `--dlo-no-use-allgather` or disable +online quantization until their runtime layouts are validated. + +The Host Weight Runtime representation, publication, and no-AllGather consumer +contracts are merged; see +[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414), +[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and +[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered +direct H2D is an optional transport layer over that merged lease contract. + +## Validation coverage + +Current source-level validation includes: + +- HSDP + DLO + AllGather rejection; +- HSDP + DLO without AllGather acceptance at configuration level; +- loader preflight fallback for TP, HSDP, online quantization, unknown custom + loaders, missing keys, and shape/dtype mismatches; +- ordinary-loader fallback for per-tensor online FP8, online INT8, and online + MXFP8 linears followed by DLO sharding of finalized weights and scales; +- exact loader-to-backend plan transfer and ordinary-loader fallback; +- ordered publication hashing, overlapped durability, unordered parallel + checksum fallback, and node-local source-digest reuse/invalidation; +- rank-local mmap source retention, bounded two-slot staging, and adapter + transforms without parameter-side flags; +- read-only registration preflight, direct source-to-device copies, safe + fallback, partial rollback, retryable unregistration, and lease ordering; +- resident-layer requests requiring no-AllGather; +- DP request-wave validation for denoising-step compatibility; +- bounded component allocator-cache retention, conservative/forced release, + OOM retry, and immutable encoder non-block staging; +- sharding, double-buffer, AllGather-size, and heterogeneous-block regression + tests. + +### B300 parallel-topology smoke matrix + +A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed, +CUDNN attention backend, 256x256 output, two denoising steps, and +`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and +VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap +or shared-mmap host-memory savings. + +| Configuration | Result | Warm E2E | Peak device memory | Host PSS | +| --- | ---: | ---: | ---: | ---: | +| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB | +| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB | +| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB | +| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB | + +Within each topology, the AllGather and no-AllGather video and audio outputs +were byte-identical. All four runs completed without an `ERROR` or traceback +and released their device allocations. For DP4xTP1, no-AllGather direct mmap +reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to +125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the +ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above +AllGather, because DP replicas did not share checkpoint-backed runtime +weights. This is a functional and memory smoke test, not a production-quality +performance or output-quality benchmark. + +### Host-memory measurement + +A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the +ordinary-loader fallback with direct mmap. Both runs used +DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a +256x256 four-second request. The ordinary-loader workers were sampled after +initialization. The mmap workers were sampled after one completed request, so +the checkpoint working set had been faulted into the page cache; this is the +more conservative point for mmap. + +The values below come from `/proc//smaps_rollup` and include the whole +worker, not only the DiT. The stable rank-to-rank difference comes from other +pipeline components, so each worker should be compared with the same worker in +the other storage mode. + +| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction | +| --- | ---: | ---: | ---: | ---: | ---: | +| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB | +| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB | +| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** | + +The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only +31.20 GiB `Pss_File`, which is the proportional charge expected when the same +resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell +from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker +1, a reduction of about 97–98 GiB per worker. RSS understates this benefit +because it counts a shared physical page in every process that maps it; summed +PSS is the appropriate node-memory comparison. + +The highest-value missing coverage is broader end-to-end numerical and +lifecycle comparison against ordinary layerwise offload for DP+SP, +HSDP+SP+no-AllGather, and TP greater than one across additional models and +target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not +change the Phase A direct-mmap TP1 support boundary. + +## Recommendations + +- Use **DP + DLO AllGather** for the supported throughput and host-memory + scaling path. +- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is + not the goal. +- Use **no-AllGather** when independent replica execution is required. TP1 + direct-mmap deployments can share checkpoint pages per node; other layouts + retain the ordinary loader's private host memory behavior and are outside + the Phase A shared-mmap support scope. +- Prefer **HSDP alone** for production HSDP deployments until the combined + HSDP + DLO no-AllGather path has broader end-to-end coverage. diff --git a/docs/design/feature/offloader/layerwise_offload.md b/docs/design/feature/offloader/layerwise_offload.md new file mode 100644 index 000000000..3fff501da --- /dev/null +++ b/docs/design/feature/offloader/layerwise_offload.md @@ -0,0 +1,57 @@ +# Layerwise Offload + +Layerwise offload streams transformer blocks through one accelerator. For +commands and model support, see the +[layerwise user guide](../../../user_guide/diffusion/offloader/layerwise_offload.md). + +## Hook ring + +`LayerWiseOffloadBackend` discovers every streamable block and registers one +`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned +host representation of its registered block and a reference to the next block +it must prefetch; the final hook prefetches the first block for the next +denoising iteration. + +During initialization, parameters and buffers are grouped by dtype and copied +into contiguous pinned CPU tensors. Original tensor storage is replaced with +an empty or meta placeholder while identity and module structure are retained. +Metadata records names, offsets, shapes, and dtypes for rematerialization. + +## Forward lifecycle + +For each block: + +1. the previous hook's completion event ensures this block is materialized; +2. the copy stream asynchronously allocates and fills the next block's device + tensors; +3. parameter and buffer storage is rebound to views of those tensors; +4. the current block computes on the normal stream; and +5. its post-forward hook replaces device storage with placeholders after all + dependent work is safe. + +The copy stream waits on the compute stream before reusing storage, and the +compute stream waits on the per-block prefetch event before forward. These +dependencies, not incidental global synchronization, define correctness. + +## Topology contract + +A DiT declares ordered block-container attributes using +`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must +be an executable `nn.Module`; ordering must match forward execution. Multiple +containers are concatenated in execution order. + +Non-block DiT modules, encoders, VAEs, and declared resident modules are not +streamed by this backend. + +## Invariants and limitations + +- Each offloaded tensor has one pinned host backing store. +- Parameter identity is preserved when storage is rebound. +- Blocks with heterogeneous shapes and dtypes use their recorded metadata; + implementations must not assume identical blocks. +- The ring is initialized before the first forward so block zero is available. +- This backend is single-device; distributed weight reconstruction belongs to + [Distributed Layerwise Offload](distributed_layerwise_offload.md). + +Shared selection and lifecycle behavior is defined in the +[CPU Offloading design](README.md). diff --git a/docs/design/feature/omni_connectors/mooncake_store_connector.md b/docs/design/feature/omni_connectors/mooncake_store_connector.md index 2ff116859..ed8e1ef49 100644 --- a/docs/design/feature/omni_connectors/mooncake_store_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_store_connector.md @@ -32,26 +32,25 @@ mooncake_master \ ## Configuration -Define the connector in runtime: +Define the connector at the top level of the deploy YAML: ```yaml -runtime: - connectors: - connector_of_mooncake: - name: MooncakeStoreConnector - extra: - host: "127.0.0.1" - metadata_server: "http://:8080/metadata" - master: ":50051" - segment: 512000000 - localbuf: 64000000 - proto: "tcp" +connectors: + connector_of_mooncake: + name: MooncakeStoreConnector + extra: + host: "127.0.0.1" + metadata_server: "http://:8080/metadata" + master: ":50051" + segment: 512000000 + localbuf: 64000000 + proto: "tcp" ``` Wire stages to the connector: ```yaml -stage_args: +stages: - stage_id: 0 output_connectors: to_stage_1: connector_of_mooncake diff --git a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md index 306a0620b..e23ea897b 100644 --- a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md @@ -21,26 +21,25 @@ InfiniBand/RoCE NICs). ## Configuration -Define the connector in runtime: +Define the connector at the top level of the deploy YAML: ```yaml -runtime: - connectors: - rdma_connector: - name: MooncakeTransferEngineConnector - extra: - host: "auto" # Auto-detect local RDMA IP - zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below) - protocol: "rdma" # "rdma" or "tcp" - device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect - memory_pool_size: 4294967296 # 4 GB (CPU); use 2147483648 (2 GB) for GPU - memory_pool_device: "cpu" # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA +connectors: + rdma_connector: + name: MooncakeTransferEngineConnector + extra: + host: "auto" # Auto-detect local RDMA IP + zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below) + protocol: "rdma" # "rdma" or "tcp" + device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect + memory_pool_size: 4294967296 # 4 GB (CPU); use 2147483648 (2 GB) for GPU + memory_pool_device: "cpu" # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA ``` Wire stages to the connector: ```yaml -stage_args: +stages: - stage_id: 0 output_connectors: to_stage_1: rdma_connector @@ -415,7 +414,7 @@ A receiver connector: - allocates receive buffers from its own pool - requests metadata or transfer service from the sender -The role is not inferred dynamically. It is injected by the stage configuration layer: +The role is not inferred dynamically. It is injected from the deploy connector wiring: - incoming edge for a stage -> `role="receiver"` - outgoing edge for a stage -> `role="sender"` diff --git a/docs/design/feature/omni_connectors/shared_memory_connector.md b/docs/design/feature/omni_connectors/shared_memory_connector.md index 62d77cd6d..1c0272340 100644 --- a/docs/design/feature/omni_connectors/shared_memory_connector.md +++ b/docs/design/feature/omni_connectors/shared_memory_connector.md @@ -13,10 +13,9 @@ segment name is returned in metadata. ## Configuration ```yaml -runtime: - connectors: - connector_of_shared_memory: - name: SharedMemoryConnector +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector ``` ## Notes diff --git a/docs/design/feature/omni_connectors/yuanrong_connector.md b/docs/design/feature/omni_connectors/yuanrong_connector.md index 12127ba3e..828feb99f 100644 --- a/docs/design/feature/omni_connectors/yuanrong_connector.md +++ b/docs/design/feature/omni_connectors/yuanrong_connector.md @@ -65,23 +65,22 @@ dscli stop --worker_address "${WORKER_IP}:31501" ## Configuration -Define the connector in runtime: +Define the connector at the top level of the deploy YAML: ```yaml -runtime: - connectors: - connector_of_yuanrong: - name: YuanrongConnector - extra: - host: "127.0.0.1" - port: 31501 - get_sub_timeout_ms: 1000 +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: +stages: - stage_id: 0 output_connectors: to_stage_1: connector_of_yuanrong diff --git a/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md index 21d5f1841..c98a34896 100644 --- a/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md +++ b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md @@ -94,26 +94,23 @@ For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU memory pool: ```yaml -runtime: - enabled: true - defaults: +connectors: + yuanrong_te_connector: + name: YuanrongTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + rpc_port: "auto" + protocol: "ascend" + device_name: "auto" + memory_pool_size: 1073741824 + memory_pool_device: "npu" + +edges: + - from: 0 + to: 1 window_size: -1 max_inflight: 1 - connectors: - yuanrong_te_connector: - name: YuanrongTransferEngineConnector - extra: - host: "auto" - zmq_port: 50051 - rpc_port: "auto" - protocol: "ascend" - device_name: "auto" - memory_pool_size: 1073741824 - memory_pool_device: "npu" - edges: - - from: 0 - to: 1 - window_size: -1 ``` Important fields: @@ -141,31 +138,27 @@ For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in both stage `omni_kv_config` blocks: ```yaml -stage_args: +stages: - stage_id: 0 - runtime: - devices: "0,1,2,3" - engine_args: - tensor_parallel_size: 4 - omni_kv_config: - need_send_cache: true - rank_mapping: - from_tp: 4 - to_tp: 4 + devices: "0,1,2,3" + tensor_parallel_size: 4 + omni_kv_config: + need_send_cache: true + rank_mapping: + from_tp: 4 + to_tp: 4 output_connectors: to_stage_1: yuanrong_te_connector - stage_id: 1 - runtime: - devices: "4,5,6,7" - engine_args: - parallel_config: - tensor_parallel_size: 4 - omni_kv_config: - need_recv_cache: true - rank_mapping: - from_tp: 4 - to_tp: 4 + devices: "4,5,6,7" + parallel_config: + tensor_parallel_size: 4 + omni_kv_config: + need_recv_cache: true + rank_mapping: + from_tp: 4 + to_tp: 4 input_connectors: from_stage_0: yuanrong_te_connector ``` diff --git a/docs/design/feature/prefix_caching.md b/docs/design/feature/prefix_caching.md index ebad8b691..a8fe4fc50 100644 --- a/docs/design/feature/prefix_caching.md +++ b/docs/design/feature/prefix_caching.md @@ -41,7 +41,7 @@ With this in mind, consider the set of blocks in a 2D layout, where the row repr ### Example !!! note "Note 3" - Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. It can be enabled/disabled per-stage via the `enable_prefix_caching` parameter in the model's stage config. + Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config. The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following: diff --git a/docs/design/feature/quantization.md b/docs/design/feature/quantization.md index 37dd5e26b..84c9641b3 100644 --- a/docs/design/feature/quantization.md +++ b/docs/design/feature/quantization.md @@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter. ModelOpt checkpoint metadata is detected from its method, producer metadata, or `quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration. +Unless the input is a per-component dictionary, the resolved configuration is +passed unchanged to every quantization-aware component constructed by the +pipeline. It affects only layers supported by that component's quantization +method; it does not rewrite arbitrary `torch.nn` modules. A per-component +dictionary narrows this global scope through runtime layer-prefix routing. + The factory uses Omni-specific overrides for methods whose diffusion behavior or platform dispatch is not provided by the vLLM registry. The current override set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/ @@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for: Model code must keep precision-sensitive or unsupported modules unquantized. The shared `safe_quant_config()` helper is used by model integrations for normalization and modulation layers where a generic FP8 configuration would be -unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers -are not implicitly quantized just because the main transformer is quantized; -they require explicit component routing and method support. +unsafe. Encoders exposed through the quantization factory receive a global +configuration just like the main transformer, but only their supported vLLM +quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs, +schedulers, and tokenizers are not rewritten automatically. A component map can +select a narrower scope, but it cannot add quantization support to a layer. Weight quantization and quantized KV cache are separate paths. The former is resolved through `quantization_config`; runtime attention/KV-cache quantization @@ -182,12 +190,15 @@ configured parallelism strategy: - `WeightsMapper` transformations must be applied before matching ignored-layer lists or component prefixes. -Distributed layerwise offload has an additional boundary: its sharded mmap -AllGather path is not compatible with online quantization because the path -expects already materialized checkpoint weights. Use the rank-local -`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when -combining these features. See the [DLO design document](distributed_layerwise_offload.md) -for the full compatibility matrix. +Distributed layerwise offload has an additional boundary: direct checkpoint +mmap cannot prepare online-quantized tensors because it bypasses the ordinary +loader that creates their runtime weights and scales. Per-tensor online FP8 +linears can still use DLO AllGather through the ordinary loader: each rank first +finalizes the FP8 tensors, then DLO shards those runtime tensors for +reconstruction during prefetch. Other online methods remain limited to +`--dlo-no-use-allgather` until their packing and scale layouts are validated. +See the [DLO design document](offloader/distributed_layerwise_offload.md) for +the full compatibility matrix. ## Adding a quantization backend diff --git a/docs/design/feature/realtime_ar_diffusion.md b/docs/design/feature/realtime_ar_diffusion.md new file mode 100644 index 000000000..5963befe7 --- /dev/null +++ b/docs/design/feature/realtime_ar_diffusion.md @@ -0,0 +1,213 @@ +# Realtime AR-Diffusion sessions + +> **Status:** experimental. This document describes the internal execution +> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not +> define a public HTTP or WebSocket schema. + +## Scope + +Realtime AR-Diffusion serves world models that generate one autoregressive +latent block per request while retaining model state across requests. The +generic runtime owns identity, event ordering, session lifecycle, routing, and +KV allocation. A model adapter owns model-specific controls and conditioning. + +LingBot World v2 is the first single-KV-branch integration. DreamZero uses the +same runner capability with two logical CFG branches. + +## Identities + +Four identifiers have distinct lifetimes: + +| Identifier | Lifetime | Owner | Contract | +| --- | --- | --- | --- | +| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. | +| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. | +| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. | +| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. | + +`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick +`request_id` remains inside the immutable AR-Diffusion snapshot and the +standard output metadata, while caller-visible `OmniRequestOutput.request_id` +continues to use AsyncOmni's existing external-ID contract. Protocol identity +therefore does not require a private entrypoint override. + +## Request lifecycle + +1. `ARDiffusionSession.accept_event()` validates and queues an immutable + `ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in + the same event. +2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls + `_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot + of the current prompt, reduced controls, ordered event IDs, and chunk index. +3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling + parameters, stores the typed tick in + `sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni + prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`. +4. `AsyncOmni` allocates its normal unique engine routing ID and submits one + `OmniDiffusionRequest` through the orchestrator. +5. `ARDiffusionModelRunner._request_session()` parses the typed tick. + `_get_or_create_session()` obtains runner-owned paged KV, and + `bind_ar_diffusion_state()` exposes it to the pipeline only for the duration + of `execute_model()`. +6. The model pipeline validates the tick and generates exactly one latent + block. LingBot interprets camera controls in + `LingBotCameraControlReducer` and + `LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse + camera or action schemas. +7. The pipeline returns one `DiffusionOutput` with the generated payload and an + `ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`. + `AsyncOmni` exposes it as + `OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`. +8. The consumer parses the standard envelope. The session commits its reducer + state, prompt, controls, event IDs, and next chunk index only when returned + metadata exactly equals the submitted tick snapshot. + +The metadata envelope is: + +```json +{ + "metadata": { + "ar_diffusion": { + "session_id": "world-7", + "request_id": "ar-world-7-3-...", + "chunk_index": 3, + "applied_event_ids": [10, 11] + } + } +} +``` + +## Transaction and failure semantics + +Accepted events are not acknowledged as applied until the model output and +metadata are validated. Reducer `prepare()` is speculative; reducer `commit()` +and the session snapshot form one logical commit. A model, metadata, or reducer +failure leaves the chunk index and queued events uncommitted and transitions the +session to `FAILED`. + +After a tick failure, worker state is closed because a forward may have +partially modified or already released KV. The same chunk is not retried in +place. A caller must explicitly reset the session, starting again at chunk +zero, or close it. + +Explicit close follows this state machine: + +```text +ACTIVE or FAILED + | + v + CLOSING ---- worker cleanup succeeds ----> CLOSED + | + +------- worker cleanup fails --------> CLEANUP_FAILED + | + +-- close retry --> CLOSING +``` + +`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects +creation of the same `session_id`, and rejects reset. Only an idempotent close +retry that succeeds removes the session. Transport disconnect uses the same +path. + +## KV ownership and capacity + +`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It +declares: + +- TP-local self-attention geometry, AR block size, sink, and recent window; +- logical KV branches and their worker-local slots; +- fixed-length cross-attention caches; +- scratch tokens required by an uncommitted forward; +- persistent model-owned CUDA bytes per resident session; and +- the requested resident-session capacity. + +The runner computes: + +```text +required(capacity) = + scratch + + managed_self_attention_pages(capacity) + + capacity * cross_attention_bytes_per_session + + capacity * model_owned_state_bytes_per_session +``` + +`gpu_memory_fraction` is a soft expansion budget: + +```text +configured_budget = available_device_bytes * gpu_memory_fraction +effective_budget = max(configured_budget, required(1)) +``` + +If `required(1)` exceeds actual available device memory, initialization fails. +Otherwise the runner always admits one viable session and uses the configured +fraction to select additional resident sessions up to the model-declared cap. +All model-owned reservations are deducted before the paged self-attention pool +is allocated. + +LingBot reports its persistent full-horizon image-condition tensor. DreamZero +reports the measured 603 MiB per-session Wan VAE causal-convolution state. + +## Concurrency and routing limits + +The current runner executes one request at a time: + +- `max_num_seqs` must be one; +- request batching and diffusion step execution are rejected; +- one AR block is generated per engine request; and +- an AR-Diffusion stage must have exactly one replica. + +The single-replica restriction is intentional. Multi-replica support requires +session-affine routing so every tick, reset, and close reaches the worker that +owns the session. Random or round-robin routing is not correct. + +Multiple resident sessions may share one runner. Capacity selection accounts +for all persistent pools, and the runner performs LRU eviction through the same +model lifecycle hook used by explicit close. + +## Model boundary + +The generic runtime transports controls as opaque +`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate +and reduce those values. This keeps target/velocity, key state, pose, +intrinsics, and other model-specific conventions out of the session protocol. + +For LingBot World v2: + +- `lingbot_world/actions.py` parses key-state and trajectory controls; +- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings; +- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and +- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and + produces one block plus the standard metadata envelope. + +## Non-goals + +This contract does not currently provide: + +- a public HTTP or WebSocket request schema (the structured-interaction + frontend is tracked separately in + [#5527](https://github.com/vllm-project/vllm-omni/pull/5527)); +- camera/action semantics shared by every world model; +- session migration or replication across workers; +- retry of an ambiguous partially executed chunk; +- cross-stage KV transfer; or +- stateful streaming VAE decode. + +Those features can be layered on top without adding LingBot-specific identity +or lifecycle fields to the generic runtime. + +## Required regression coverage + +CPU contract tests cover: + +- separation of tick correlation IDs from internal engine routing IDs; +- event ordering, metadata equality, and reducer fail-closed commit; +- close/disconnect failure tombstones and cleanup retry; +- capacity one under the default fraction for shipped LingBot geometry; +- capacity expansion under a tuned fraction; +- DreamZero model-owned state accounting; +- runner session reuse, reset, close, LRU eviction, and forward cleanup; and +- LingBot model registration, imports, camera/action reduction, and one-block + output metadata. + +GPU validation must additionally exercise real model loading, action input, +prompt switching, rolling-window replacement, reset, close, finite output +latents, and exact metadata/request identity. diff --git a/docs/design/feature/skip_softmax.md b/docs/design/feature/skip_softmax.md index 8e6375ee8..0527c7eeb 100644 --- a/docs/design/feature/skip_softmax.md +++ b/docs/design/feature/skip_softmax.md @@ -2,8 +2,10 @@ Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys and how to pick an operating point — is in -[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains -the algorithm. +[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax). +The shared selector contract is documented in +[Diffusion Attention Backend Selection](attention_backend_selection.md). This +page explains the algorithm. ## Motivation diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md index c465449f8..c6cd19987 100644 --- a/docs/design/feature/vae_parallel.md +++ b/docs/design/feature/vae_parallel.md @@ -22,59 +22,60 @@ We use **Qwen-Image** as the reference implementation for decode parallel, and * **VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is: -+ Split into multiple spatial tiles +- Split into multiple spatial tiles -+ Distributed across multiple ranks +- Distributed across multiple ranks -+ Encoded/Decoded in parallel +- Encoded/Decoded in parallel -+ Merged to reconstruct the final output +- Merged to reconstruct the final output This approach: -+ Distributes computation across multiple devices +- Distributes computation across multiple devices -+ Reduces peak memory usage per device +- Reduces peak memory usage per device -+ Accelerates encoding/decoding latency +- Accelerates encoding/decoding latency ### When to Use Encode vs Decode Parallel -| Operation | Use Case | Example | -|-----------|----------|---------| -| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video | -| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) | +| Operation | Use Case | Example | +| ------------------- | ---------------------------- | --------------------------------- | +| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video | +| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) | ### Architecture + We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding. The executor is model-agnostic and accepts three function parameters: -+ split – Partition the latent into tiles +- split – Partition the latent into tiles -+ exec – Decode a single tile +- exec – Decode a single tile -+ merge – Combine decoded tiles into the final output +- merge – Combine decoded tiles into the final output #### Execution Flow -+ Call split(z) to generate a list of TileTask and a GridSpec +- Call split(z) to generate a list of TileTask and a GridSpec -+ Dispatch tasks across ranks using workload-based balancing +- Dispatch tasks across ranks using workload-based balancing -+ Each rank executes exec(task) on its assigned tiles +- Each rank executes exec(task) on its assigned tiles -+ Gather decoded tile results to rank 0 +- Gather decoded tile results to rank 0 -+ Rank 0 performs merge(...) +- Rank 0 performs merge(...) -+ (Optional) Broadcast final result to all ranks +- (Optional) Broadcast final result to all ranks This design separates: -+ Distributed execution logic +- Distributed execution logic -+ Model-specific tiling and merging logic +- Model-specific tiling and merging logic #### Why split / exec / merge is necessary? @@ -82,23 +83,23 @@ The latent tensor cannot be arbitrarily partitioned. During decoding: -+ Each output pixel may depend on neighboring pixels +- Each output pixel may depend on neighboring pixels -+ The receptive field is model-dependent +- The receptive field is model-dependent Therefore: -+ Tiles must include overlap +- Tiles must include overlap -+ Merge must perform blending to avoid seams +- Merge must perform blending to avoid seams ## Step-by-Step Implementation (Decode) ### Step 1: Implement DistributedAutoencoderKLQwenImage + `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version: - -``` +```python class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin): @classmethod def from_pretrained(cls, *args: Any, **kwargs: Any): @@ -106,28 +107,35 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi model.init_distributed() return model ``` + **Key points**: -+ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin -+ Call init_distributed() after loading weights + +- Inherit both AutoencoderKLQwenImage and DistributedVaeMixin +- Call init_distributed() after loading weights ### Step 2: Implement split/exec/merge + Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`: -``` + +```python class GridSpec: split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W)) grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols) tile_spec: dict = field(default_factory=dict) # Metadata required for merging output_dtype: torch.dtype | None = None # Final output dtype ``` -``` + +```python class TileTask: tile_id: int # task id grid_coord: tuple[int, ...] # Tile position in grid tensor: torch.Tensor | list[torch.Tensor] # The tile tensor workload: int | float = 1 # Used for load balancing (e.g., tile area) ``` + And tiled base split/exec/merge as follow: -``` + +```python def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: # mostly copy from AutoencoderKL _, _, num_frames, height, width = z.shape @@ -210,11 +218,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid ``` ### Step 3: Override tiled_decode + We need to override tiled_decode, the main logic is: -+ check distributed is enabled -+ select split/exec/merge -+ Invoke self.distributed_executor.execute to decode -``` + +- check distributed is enabled +- select split/exec/merge +- Invoke self.distributed_executor.execute to decode + +```python def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): if not self.is_distributed_enabled(): return super().tiled_decode(z, return_dict=return_dict) @@ -230,11 +241,14 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): return DecoderOutput(sample=result) ``` + `broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0. ### Step 4: Modify Pipeline + Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage -``` + +```diff class YourModelPipeline(nn.Module): def __init__( self, @@ -259,8 +273,8 @@ For models that require VAE encoding (e.g., Image-to-Video), you can also parall Similar to decode, split the input tensor into tiles. Key considerations: -+ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly -+ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x) +- **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly +- **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x) ```python def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]: @@ -406,23 +420,28 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor: **Key differences from decode parallel:** | Aspect | Decode Parallel | Encode Parallel | -|--------|-----------------|-----------------| +| -------- | ----------------- | ----------------- | | `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) | | Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` | | Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) | ## Testing -Verify numerical consistency between: -+ vae_patch_parallel_size = 1 -+ vae_patch_parallel_size = N +Verify numerical consistency between: + +- vae_patch_parallel_size = 1 +- vae_patch_parallel_size = N Example: + +```python torch.allclose(output_1, output_n, atol=1e-5) +``` Testing requirements: -+ Fix random seed -+ Use identical tiling strategy + +- Fix random seed +- Use identical tiling strategy ```python m = Omni( @@ -434,6 +453,7 @@ m = Omni( ), ) ``` + When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead. ## Reference Implementations @@ -441,7 +461,7 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati Complete examples in the codebase: | Model | Path | Decode Parallel | Encode Parallel | -|-------|------|-----------------|-----------------| +| ------- | ------ | ----------------- | ----------------- | | **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ | | **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ | | **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ | @@ -456,7 +476,7 @@ The tile-parallel executor above assigns independent spatial **tiles** to ranks. ### How it differs from tile parallel | Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) | -|--------|--------------------------|-----------------------------------------------| +| -------- | -------------------------- | ----------------------------------------------- | | Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W | | Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) | | Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) | @@ -481,7 +501,7 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs ### Notes - The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance. -- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`). +- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `cards_2`). --- diff --git a/docs/design/index.md b/docs/design/index.md index d81ccb806..066d257d1 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -18,10 +18,12 @@ implementation contract; it is not, by itself, a general support claim. ### Runtime and stage execution - [Disaggregated Inference](feature/disaggregated_inference.md) +- [Host Weight Runtime](feature/host_weight_runtime.md) - [Async Chunk](feature/async_chunk.md) - [Async Diffusion Output](feature/async_diffusion_output.md) - [Async Omni Output Materialization](feature/omni_async_output_materialization.md) - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md) +- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md) ### Communication @@ -53,15 +55,22 @@ implementation contract; it is not, by itself, a general support claim. #### Attention optimization The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md) -guide lists all selectable backends and their platform defaults. Skip-Softmax -is the backend-related optimization with a standalone feature-design contract: +guides list selectable backends, platform defaults, installation, and tuning. +The design contracts separate selection mechanics from backend algorithms: +- [Attention Backend Selection](feature/attention_backend_selection.md) - [Skip-Softmax](feature/skip_softmax.md) +#### CPU offloading + +- [Overview and Shared Contracts](feature/offloader/README.md) +- [Model-Level Offload](feature/offloader/cpu_offload.md) +- [Layerwise Offload](feature/offloader/layerwise_offload.md) +- [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md) + - [Cache-DiT](feature/cache_dit.md) - [TeaCache](feature/teacache.md) - [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md) -- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md) ## Infrastructure and Performance @@ -80,14 +89,15 @@ is the backend-related optimization with a standalone feature-design contract: - [Model Integration](module/model_integration.md) - [Autoregressive Runtime](module/ar_runtime.md) - Diffusion - - [Overview](module/diffusion/index.md) - - [Runtime](module/diffusion/diffusion_runtime.md) - - [Model Integration](module/diffusion/diffusion_model_integration.md) - - [Continuous Batching](module/diffusion/continuous_batching.md) - - [Parallelism](module/diffusion/parallelism.md) - - [Offloader](module/diffusion/offloader.md) + - [Overview](module/diffusion/index.md) + - [Runtime](module/diffusion/diffusion_runtime.md) + - [Model Integration](module/diffusion/diffusion_model_integration.md) + - [Continuous Batching](module/diffusion/continuous_batching.md) + - [Parallelism](module/diffusion/parallelism.md) + - [Offloader](module/diffusion/offloader.md) - [Execution Platforms](module/execution_platforms.md) - [Cache Management](module/cache_management.md) +- [Host Weight Runtime](module/host_weight_runtime.md) - [Quantization](module/quantization.md) - [Observability](module/observability.md) - [Profiling](module/profiling.md) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 7ee35813f..85eb7ac38 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -6,7 +6,7 @@ This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage - Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope). - Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically. -- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss. +- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss, plus image and diffusion workload, queueing, execution, and memory signals. - Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable. - Keep the metrics collection overhead low enough that it does not regress TTFA or throughput. @@ -20,7 +20,7 @@ vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` col ### The Problem -vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`. +vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm_omni:` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm_omni:` / `vllm_omni`. Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below). @@ -38,7 +38,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by | +--------+--------+--------+--------+--------+ | | - vllm:omni_* vllm:* + vllm_omni:* vllm:* collectors collectors | | +----+--------+ +-----------+ +----------+ +-----------+ @@ -58,17 +58,30 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by ### Data Flow -There are four independent paths for metric collection. +There are five independent paths for metric collection. -**Path 1: Pipeline-level metrics (`vllm:omni_*`)** +**Path 1: Pipeline-level metrics (`vllm_omni:*`)** `OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress: - `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`. -- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`. -- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion. +- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm_omni:requests_success_total{finished_reason}`. +- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter. +- `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`. -**Path 2: Audio modality metrics (`vllm:omni_audio_*`)** +**Path 2: Image and diffusion metrics** + +Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice. + +- `stage_gen_time_s`, `stage_in_queue_s`, image workload, inference-step, and per-stage peak-memory values are emitted from `StageRequestStats` when a stage finishes. `image_pixels` and `image_count_total` are restricted to image outputs; `num_inference_steps` is restricted to diffusion stages. +- Diffusion engine timings are accumulated per request by `OrchestratorAggregator.accumulate_diffusion_metrics`. Engine-side millisecond values are converted to seconds before `observe_modality_at_finalize(...)` dispatches the per-replica Histogram observations. +- `request_queue_wait_s` is emitted once at final request completion from `pipeline_timings["queue_wait_ms"]`. A present value of `0` is a valid observation; a missing key produces no sample. +- `stage_waiting_requests` is the sum of the latest scheduler snapshots for all live replicas in a stage. AR snapshots arrive through `SchedulerStats`; diffusion snapshots ride normal diffusion output metrics. A dead replica's snapshot is removed. Because this deliberately reuses existing message paths, the gauge can retain its last value while a replica produces no output. +- `kv_wait_s` measures the AR scheduler interval from entering the KV-transfer-free wait state until extraction acknowledgement. The sample is carried to the Orchestrator through the existing engine output and labeled by connector type. Abort/error/removal paths discard incomplete timers rather than emitting partial waits. + +Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active. + +**Path 3: Audio modality metrics (`vllm_omni:audio_*`)** `OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites: @@ -76,7 +89,7 @@ There are four independent paths for metric collection. - `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry. - `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`. -**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)** +**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)** `OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution. @@ -87,7 +100,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected). -**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)** +**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)** The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics. @@ -97,9 +110,9 @@ The Orchestrator runs in a background thread. The API server (OmniBase) runs in ### Metric Registration and Lifecycle -All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data. +All `vllm_omni:*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data. -The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime. +The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm_omni:*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime. ## OmniPrometheusStatLogger Wrap @@ -130,11 +143,11 @@ A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestr ## Gating: `--log-stats` -All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`. +All metrics — both the omni-specific `vllm_omni:*` families and the upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`. Behavior with `--log-stats=off` (default): -- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them. +- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm_omni:*` stay materialized but never have data written to them. - `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all. - The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate. @@ -162,10 +175,43 @@ second. | Metric | Type | Labels | Description | |--------|------|--------|-------------| -| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages | -| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled | -| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | -| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds | +| `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages | +| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled | +| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | +| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds | + +### Image and diffusion service-level metrics + +| Metric | Type | Labels | Description | +|--------|------|--------|-------------| +| `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing | +| `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded | +| `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage | +| `vllm_omni:stage_in_queue_s` | Histogram | `model_name`, `stage` | Diffusion scheduler wait from enqueue to initial execution admission | +| `vllm_omni:num_inference_steps` | Histogram | `model_name` | Denoising-step distribution for diffusion stages | +| `vllm_omni:image_count_total` | Counter | `model_name` | Cumulative number of generated image outputs | +| `vllm_omni:image_pixels` | Histogram | `model_name` | Pixel-count distribution for image outputs | +| `vllm_omni:peak_memory_mb` | Gauge | `model_name`, `stage` | Positive per-stage peak-memory samples reported by engine outputs | +| `vllm_omni:requests_failed_total` | Counter | `model_name`, `reason` | Request failures with bounded reason values (`client_abort`, `client_disconnect`, `stage_error`, `unknown`) | +| `vllm_omni:kv_wait_s` | Histogram | `model_name`, `connector_type` | Completed AR KV-transfer wait intervals, grouped by connector backend | + +### Diffusion execution breakdown + +Labels: `{model_name, stage, replica}`. + +| Metric | Type | Description | +|--------|------|-------------| +| `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request | +| `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count | +| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time | +| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing time | +| `vllm_omni:vae_decode_s` | Histogram | VAE decode time when the engine reports it | +| `vllm_omni:diffusion_forward_s` | Histogram | Forward-only denoise-loop time when pipeline profiling is enabled | +| `vllm_omni:diffusion_kv_load_s` | Histogram | Diffusion-side KV receive/load time when upstream KV is used | +| `vllm_omni:image_ttfp_s` | Histogram | Image-stage submit to first materialized image output | +| `vllm_omni:denoise_step_latency_s` | Histogram | Mean diffusion forward time per denoise step for image output stages | + +The optional breakdown and memory families are sparse by design: when an engine does not produce the source measurement, no sample is emitted. This differs from queue timings, where an explicitly measured zero is meaningful and is retained. ### Audio (7) @@ -173,13 +219,13 @@ Labels: `{model_name, stage, replica}` plus the listed extra label. | Metric | Type | Extra label | Description | |--------|------|-------------|-------------| -| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | -| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | -| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` | -| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated | -| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | -| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | -| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | +| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | +| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | +| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` | +| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frames generated | +| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps | +| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | +| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | ### Cross-stage transfer (4) @@ -187,10 +233,10 @@ Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`. | Metric | Type | Description | |--------|------|-------------| -| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | -| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | -| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | -| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | +| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | +| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | +| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | +| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) | ### LLM stage-level (wrapped `vllm:*`) @@ -203,7 +249,7 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution. - Counters use the `_total` suffix (auto-appended by `prometheus_client`). - Sizes use the `_bytes` suffix. -- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`. +- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`. ## Logging vs. Prometheus diff --git a/docs/design/module/diffusion/offloader.md b/docs/design/module/diffusion/offloader.md index 919b1f39a..7ac5d2cc1 100644 --- a/docs/design/module/diffusion/offloader.md +++ b/docs/design/module/diffusion/offloader.md @@ -18,7 +18,7 @@ validation_paths: - tests/diffusion/offloader/** upstream_refs: - torch.nn.Module.to -last_reviewed: 2026-07-16 +last_reviewed: 2026-08-25 --- # Diffusion offloader @@ -48,6 +48,11 @@ and correctness required by the selected execution mode. **Rule:** Retained host and device copies MUST respect configured limits and provide deterministic teardown. +Allocator-cache retention MUST remain local to an explicit component owner, +declare both cache and physical-free-memory bounds, release conservatively +when telemetry is unavailable, and force release on failure or memory +pressure. It MUST NOT replace unconditional executor shutdown cleanup. + ## Safe-change guide Test enabled and disabled paths, repeated execution, asynchronous transfer, diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index 5bb0ea0df..f3f07a43f 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -37,6 +37,7 @@ validation_paths: - tests/engine/test_async_omni_engine_input.py - tests/engine/test_async_omni_engine_outputs.py - tests/engine/test_async_omni_engine_abort.py + - tests/e2e/offline_inference/test_qwen3_omni.py - tests/engine/test_async_omni_engine_stage_init.py - tests/engine/test_orchestrator.py - tests/engine/test_orchestrator_error_handling.py @@ -66,6 +67,12 @@ the boundary affected by the in-flight stage client/process refactor in [#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and responsibilities proposed only by that PR are not current contracts. +The orchestration loop also has an opt-in event-driven mode +(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in +[#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll +cadence only: the routing, ordering, and terminal-state contracts below hold +identically on both loops. + ## Ownership boundary This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation, @@ -107,6 +114,15 @@ Numbers become append-only after normative promotion. Test routing, output ordering, queue correlation, cancellation, failure propagation, shutdown ordering, and representative multi-stage execution. +AR abort of a final-stage LLM request must deliver a terminal output with +`finish_reason="abort"` and the generated prefix so collocated training can +resume; diffusion abort remains whole-sample retry. Aborting a parallel-sampling +child must drop `parent_requests` once no children remain. EngineCore control +RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must +propagate worker exceptions rather than returning `{"supported": False}`. +When a request has multiple final output stages, only the last abort message +is request-terminal. Output-processor abort state is committed only after the +physical EngineCore abort succeeds. Changes that cross into stage runtime or public error behavior require review from that contract's owners. diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md index 160f12514..fcdfe23ea 100644 --- a/docs/design/module/entrypoints.md +++ b/docs/design/module/entrypoints.md @@ -43,6 +43,7 @@ depends_on: validation_paths: - tests/entrypoints/test_omni_entrypoints.py - tests/entrypoints/test_async_omni.py + - tests/entrypoints/test_async_omni_pause_sleep_routing.py - tests/entrypoints/test_async_omni_duplex.py - tests/entrypoints/test_serve.py - tests/entrypoints/test_stream_finish_reason.py @@ -119,7 +120,13 @@ promotion. Test request validation, protocol conversion, model-adapter routing, streaming, session identity, disconnect/cancellation, and error mapping for -each affected offline or serving entrypoint. +each affected offline or serving entrypoint. Sleep must wait for in-flight +`generate()` admission before EngineCore offload; `wake_up` does not resume +admission — callers must `resume_generation()`. Sleeping tags are tracked per +stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`. +Streaming input pumps take an admission slot immediately before each EngineCore +ADD or update, not while waiting for the next client chunk. Frontend abort +keeps `request_states` until `generate()` consumes the terminal output. ## Promotion gate diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md new file mode 100644 index 000000000..ac92711fc --- /dev/null +++ b/docs/design/module/host_weight_runtime.md @@ -0,0 +1,341 @@ +--- +title: Host Weight Runtime +kind: module +status: draft +owners: + - "@hsliuustc0106" +primary_code_paths: + - vllm_omni/host_weight_runtime/** +related_code_paths: + - vllm_omni/diffusion/model_loader/** + - vllm_omni/diffusion/offloader/** +depends_on: + - model_integration.md + - quantization.md +validation_paths: + - tests/host_weight_runtime/** +upstream_refs: + - https://github.com/vllm-project/vllm-omni/issues/6414 +last_reviewed: 2026-08-25 +--- + +# Host Weight Runtime + +The Host Weight Runtime is a loader-adjacent, transport-independent foundation +for reusing immutable, runtime-ready model weights. It separates the semantic +question—"are these exactly the weights this loader requested?"—from the +physical question—"where is the validated host representation stored?" + +The runtime can save repeated checkpoint loading, transformation, and peak host +memory when several workers need the same final representation. It also creates +a second, versioned path for reconstructing a valid model, so a cache miss or +recoverable store failure normally falls back to the canonical loader. + +This module does not enable host-weight caching for any model by itself. A model +integration must still provide an exact identity, a producer, and a restorer. +For feature behavior and consumer integration requirements, see the +[Host Weight Runtime feature design](../feature/host_weight_runtime.md). + +## Component boundary + +```mermaid +flowchart LR + LA["Model loader adapter"] -->|"exact identity + policy"| RT["HostWeightRuntime"] + RT -->|"lookup / get_or_build"| ST["HostWeightStore"] + PR["WeightProducer"] -->|"runtime-ready tensors"| ST + ST -->|"validated mmap views"| LE["HostWeightLease"] + LE --> RE["WeightRestorer"] + RE --> MO["Restored host model"] + MO --> TR["DLO or another GPU transport"] + + CS["Canonical model source"] --> PR + LA -. "fallback" .-> CS +``` + +The contracts have deliberately narrow ownership: + +- `HostWeightRuntime` applies resolution and fallback policy and emits one + terminal resolution report. Explicit post-load publication emits a separate + report and closes its store-returned lease before returning; it does not + revise the completed resolution. + The runtime does not own a model or call the canonical loader. +- `HostWeightStore` owns artifact lookup, coordinated construction, + publication, validation, quarantine, and lifecycle. +- `HostWeightLease` owns process-local tensor views, mapped-file resources, and + the shared artifact lock that stabilizes their lifetime. +- `WeightProducer` creates one declared final-layout representation. It writes + only through a store-scoped writer and cannot publish paths directly. +- `WeightRestorer` validates a lease-to-model plan without mutation. Its + one-shot `commit() -> None` is the sole model-mutating phase. Concrete + restorers remain with the model or quantization integration. +- GPU transport owns page registration, private staging, H2D scheduling, and + lease release. DLO must not parse manifests or manage store artifacts. + +## Exact semantic identity + +Storage lookup is exact; the store never coerces one representation into +another. `WeightArtifactIdentity` includes: + +- immutable model revision and source-file fingerprint; +- component and ownership boundary; +- runtime representation, dtype, and representation metadata; +- final tensor layout and semantic parallel coordinates; +- static adaptation identity; and +- producer version, implementation fingerprint, manifest schema, and restorer + schema. + +Loader integrations may persist content digests for canonical source shards in +the selected storage domain. Reuse requires the same path, inode, size, +timestamps, symlink target, and cache-record checksum. Corrupt records are +rebuilt, and coordination or cache-I/O failure falls back to hashing the +canonical source directly. This is an identity-computation optimization, not a +second source of truth or an artifact substitution rule. + +The identity excludes transfer policy such as registered mmap versus private +pinned staging. Those paths move the same host representation and must not +produce duplicate artifacts. + +Parallel coordinates follow semantics rather than process topology: + +- TP size and rank are included when they change tensor contents or shape. +- DP rank is excluded for replicated weights. +- SP size and backend are included when they affect layout. SP rank is excluded + only when all SP ranks consume an identical representation. +- PP, EP, or future parallel modes must be added by the model adapter whenever + they change component ownership or tensor layout. + +Dynamic LoRA overlays are not cached as base weights. A statically merged +adapter is a distinct artifact identity with its own fingerprint. + +## Deterministic resolution + +The loader first resolves the canonical source and computes the exact requested +identity. The runtime then follows this order: + +```text +validated local artifact + -> exact remote artifact materialized locally (future) + -> registered local producer + -> canonical loader fallback, if policy permits +``` + +An already published local artifact wins because it avoids network transfer and +is already in node-accessible storage. "Wins" means the highest-priority valid +representation, not whichever backing responds first. + +Configuration modes are: + +- `disabled`: return the canonical-direct outcome without probing a root, + credentials, topology, or identity; +- `preferred`: recoverable store failures may return canonical fallback; and +- `required`: fail when no exact lease can be acquired. + +Failure-to-action policy is centralized in `HostWeightRuntime`. A clean miss or +invalid cache artifact may fall back in preferred mode. A typed failure may do +so only when it is marked retryable; unsupported capabilities, identity +collisions, producer failures, and publication failures remain visible as +startup failures. + +`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it +is not a hard wall-clock deadline for in-process validation, production, or +atomic publication. Once this V1 implementation becomes the producer, the +synchronous producer and publication run to completion. A hung producer blocks +its owning process and requires external process supervision; enforceable +producer cancellation requires a future process-isolated producer contract. A +waiter timing out never cancels another process's valid build. + +Remote providers are represented by protocols but are explicitly unsupported +in the first implementation. Pre-load `resolve()` invokes only +`PRE_LOAD_SAFE` producers. A loader that completed canonical materialization +may explicitly call `publish_after_load()` with a `POST_LOAD_ONLY` producer +when `allow_post_load_publish` is enabled. + +`allow_local_build` controls producer use during pre-load resolution; +`allow_post_load_publish` independently controls the explicit post-load entry +point. This permits lookup-only warm starts that populate a missing artifact +only after the canonical loader has established the current model. + +Post-load publication is synchronous in V1, but its outcome is independent of +the already completed resolution. A store failure is returned in a separate +publication report and cannot turn a valid canonical model into a failed +resolution. On success, the runtime closes the validated store lease before +returning the report. Post-load publication only warms future startups; it does +not restore, rebind, or otherwise mutate the canonically loaded model serving +the current startup. + +## Restoration transaction boundary + +`WeightRestorer.plan_restore()` is validation-only: it must not modify the model +or lease. The returned plan may be committed exactly once, `commit()` returns +`None`, and commit is the sole phase allowed to mutate the model. Registration, +staging, and GPU transfer remain transport responsibilities rather than restore +side effects. + +If planning fails, preferred-mode canonical fallback may reuse the untouched +model. If commit begins and then fails, the model is considered partially +hydrated and must be discarded; canonical fallback must construct a fresh model +instance instead of reusing it. + +## Filesystem domain and artifact lifecycle + +A storage domain describes physical locality and policy; it is not part of +semantic identity. The first implementation supports node- or NUMA-scoped +roots on verified local disk filesystems or tmpfs. The detected kernel +filesystem type is preserved in `domain.json` and inspection output. Known +remote filesystems and unknown filesystem types are rejected rather than being +silently treated as local disk, because cross-node page cache and advisory-lock +semantics do not satisfy this backend's contract. A NUMA root alone does not +guarantee page placement; CPU pinning, local first-touch/prefault, and residency +evidence belong to a later topology-aware integration. + +```text +/ +├── domain.json +├── domain-policy.json +├── locks/ +│ ├── .build.lock +│ └── .artifact.lock +├── artifacts// +│ ├── READY.json +│ ├── manifest.json +│ └── *.safetensors +├── tmp/ +├── quarantine/ +└── deny/ +``` + +Lock files live outside replaceable artifact directories. Every cooperative +mutation follows one lock order: build lock, then exclusive artifact lock. +Lookup acquires a shared artifact lock, validates `READY` and the manifest, +opens payloads with non-following descriptor-based access, creates mmap-backed +views, and transfers that shared lock to the lease. + +Lease acquisition rechecks the producer and restorer schemas against the exact +identity. Malformed safetensors payloads are reported as typed invalid +artifacts and denied rather than escaping the resolution policy as raw parser +exceptions. + +Each `FileLock` records its creator PID. A forked child closes an inherited lock +descriptor without issuing `LOCK_UN`, because `flock` state belongs to the +shared open-file description and unlocking in the child would also release the +parent's active lease lock. + +Publication follows this lifecycle: + +```text +ABSENT -> invisible temporary build -> fully hashed and fsynced + -> atomic rename -> READY + +READY -> denied -> quarantined -> explicit removal +``` + +`READY.json` is written last in the temporary directory. Readers therefore see +either a complete validated artifact or a miss, never an in-progress build. +Failed process-local builds are left invisible and a later build owner removes +their stale temporary directories. Explicit cleanup is nonblocking and reports +an active builder or lease instead of unlinking live mappings. + +## Validation and capacity + +Every publication computes full file and tensor hashes. Local lookup supports: + +- `manifest_and_metadata`, which validates identity, immutable publication + metadata, file descriptors, sizes, tensor keys, shapes, dtypes, and strides; + and +- `full_checksum`, which additionally reads and hashes every payload file. + +An ordered producer may compute file and tensor digests while payload bytes are +written and overlap each closed payload's `fsync` with later producer work. An +unordered producer retains parallel readback hashing. In both cases manifest +and `READY` publication wait for every payload digest and `fsync`; the atomic +publication and validation contract is unchanged. + +Filesystem verity is modeled but unsupported. Invalid artifacts receive an +external deny marker before their shared lookup lock is released; a build owner +then quarantines the immutable directory under the exclusive lock. +Lookup checks the marker both before validation and after constructing the +lease so a weaker overlapping lookup cannot return a lease after a stronger +validation has denied the artifact. + +If publishing the external deny marker fails, lookup returns a retryable +storage failure and marks the artifact directory with a persistent invalidation +mode before releasing the shared lock. Cooperating readers check that mode both +while opening and immediately before returning a lease. Lookup then makes a +bounded attempt to take the exclusive artifact lock and quarantine the entry; +the quarantine transition removes the internal mode marker. If a cooperating +builder replaces the marked inode while the locks are exchanged, containment +preserves that new publication instead of quarantining it. If neither the +external marker nor the directory-mode marker can be persisted, cross-process +exclusion cannot be guaranteed on an unavailable filesystem and the domain +must remain failed rather than treating the artifact as an ordinary invalid +cache entry. + +New publications carry the internal invalidation mode through the staging +rename. Hardening clears it and is the logical authority transition for a +process crash; the temporary- and artifact-parent syncs then make the rename +durable, with an artifact-inode sync first making the hardened mode durable. A +process that dies before hardening leaves a persistently denied +entry, while one that dies after successful hardening leaves an authoritative +ready entry. Any synchronous failure after rename is marked invalid and +contained in quarantine before `PUBLICATION_FAILED` is returned, so that +nonretryable result cannot later resolve as a ready hit. If the filesystem +refuses both the inode marker and the quarantine move, the store instead returns +a retryable storage failure with `outcome_uncertain` details; as with deny-marker +double failure, cross-process exclusion cannot be guaranteed until storage +recovers. Once the artifact-parent sync succeeds, a later artifact- or +build-lock release error is logged and lookup resolves the committed entry +instead of reporting a contradictory publication failure. The same authority +rule applies when a competing publisher already installed byte-identical +semantic content; staging-cleanup errors are retried without downgrading that +validated ready entry to publication failure. + +Before widening a hardened artifact for a lifecycle move, the store durably +sets the same invalidation mode. A process crash before rename therefore leaves +the writable source denied. After rename, the store re-hardens the destination +but retains the marker until both parent-directory syncs make the move durable; +only then does it durably clear the marker. An interrupted move therefore leaves +a marked transition that the next exact-key lifecycle operation re-hardens. +Explicit cleanup uses the locked move and retries exact-key +`.cleanup.*` tombstones left by an interrupted removal; a later build performs +the same tombstone reconciliation before producing replacement content. + +An operational failure while inspecting a noncooperative competing publication +is a retryable storage failure. The competing entry remains authoritative and +is not mislabeled as corrupt or quarantined merely because its manifest could +not be read transiently. + +`inspect_domain()` and `inspect_artifact()` return typed, read-only snapshots. +Inspection validation does not create a deny marker. Reported lock activity is +point-in-time evidence only; kernel locks are not a persistent owner registry. + +Capacity policy covers all store-owned bytes, including ready, temporary, and +quarantined data. The local writer preflights `max_artifact_bytes`, +`max_store_bytes`, and `min_free_bytes`, and preallocates payloads where the +filesystem supports it. `ENOSPC` remains a normal typed store failure because +concurrent preflight is inherently racy. Automatic eviction and strict +concurrent reservations are not implemented. + +For tmpfs, artifact bytes are host-memory consumption and may consume swap; +they must not be reported operationally as ordinary disk capacity. The store +owns files and reports capacity. The kernel controls page-cache residency, and +transport controls registration or page locking. + +## First implementation scope + +The foundation implements a CPU local filesystem store and +`SINGLE_PROCESS` production for one exact artifact. It intentionally adds no: + +- public CLI or DLO configuration; +- BF16, FP8, or model-specific producer/restorer; +- change to DLO AllGather or no-AllGather behavior; +- CUDA registration, pinned staging, or H2D scheduling; +- remote provider, group-cohort producer, artifact bundle, eviction, or + cross-node coordination. + +Consumer PRs must prove ordinary-loader parity, warm-hit work avoidance, +parallel-layout correctness, clean teardown, shared-backing evidence, and +representation-specific startup and memory behavior. Generic foundation tests +remain CPU-only and make no production performance claim. + +The complete rationale and staged rollout are tracked in +[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414). diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 1d8d7ab73..2e848a49c 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -42,7 +42,7 @@ and execution behavior to stable runtime contracts. ### MODEL-INV-001: Registration is explicit **Rule:** A model integration MUST declare how its model class, loader, input -processor, and stage configuration are selected. +processor, and pipeline configuration are selected. ### MODEL-INV-002: Model code does not route stages diff --git a/docs/design/module/vllm_omni_config.md b/docs/design/module/vllm_omni_config.md index 085e4cee6..1ee303b8c 100644 --- a/docs/design/module/vllm_omni_config.md +++ b/docs/design/module/vllm_omni_config.md @@ -10,16 +10,15 @@ owners: primary_code_paths: - vllm_omni/config/** - vllm_omni/deploy/** - - vllm_omni/model_executor/stage_configs/** -related_code_paths: - - vllm_omni/platforms/*/stage_configs/** + - vllm_omni/model_executor/models/*/pipeline.py +related_code_paths: [] depends_on: [] validation_paths: - tests/config/** upstream_refs: - vllm.config -last_reviewed: 2026-08-07 -last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc +last_reviewed: 2026-08-11 +last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f --- # vLLM-Omni configuration @@ -39,11 +38,66 @@ The PR review proposal to capture environment-derived values in configuration objects at initialization is an open design question, not a current invariant. +## Current structured-stage reuse behavior + +The current RFC #4021 implementation uses upstream vLLM config classes as the +schema source for AR and generation stages: + +- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`. +- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`. +- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`. +- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`. + +This inheritance preserves the existing load, cache, scheduler, and parallel +concern boundaries while reducing duplicate downstream declarations. It does +not, by itself, make every inherited field an effective vLLM-Omni runtime +option. Construction and runtime consumption are separate surfaces: + +| Surface | Current behavior | +| --- | --- | +| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. | +| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. | +| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. | +| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. | + +The generated projection maps include known upstream naming differences such +as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and +`data_parallel_master_ip` to `data_parallel_address`. A field added to both an +upstream concern config and `EngineArgs` therefore enters the AR/generation +projection without requiring a second downstream allowlist. Explicit inherited +fields that have no projection raise an error rather than being silently +dropped. + +Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`, +cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns +`distributed_executor_backend` and `worker_cls`; vLLM's private API-process +fields are terminal internals. These exclusions prevent one input from being +constructed or projected through two config concerns. Effective-engine-argument +tests cover both the dynamic mapping and this exclusion boundary. + +`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields: +`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML +inputs and complete upstream value objects. Mapping inputs are validated and +materialized as their concrete upstream types when the structured config is +constructed, so downstream consumers see one resolved representation. +Prebuilt upstream objects retain their type across the structured projection +boundary. + +Quantization remains on its existing Omni-owned transport contract. This +reuse step does not adopt upstream `QuantizationConfigArgs`, because the +current engine-specific quantization split needs to be resolved separately. + +Diffusion keeps its existing engine-field ownership and projection surface. +Shared Python inheritance must not be treated as evidence that an LLM-only +parallel input is effective for a diffusion stage. + ## Safe-change guide while deferred Changes should trace every affected producer to its runtime consumer and test -the structured, legacy, CLI, and deployment paths that are actually supported. -Do not infer a stable contract from this placeholder page. +the structured pipeline, CLI, and deployment paths that are actually supported. +For inherited vLLM fields, tests should separately cover constructor behavior +and effective `EngineArgs` projection. Do not infer a stable contract from +this placeholder page. ## Promotion gate diff --git a/docs/design/qwen3_omni_tts_performance_optimization.md b/docs/design/qwen3_omni_tts_performance_optimization.md index bc1e29a7f..ff783a02a 100644 --- a/docs/design/qwen3_omni_tts_performance_optimization.md +++ b/docs/design/qwen3_omni_tts_performance_optimization.md @@ -111,7 +111,7 @@ This post walks through each optimization in the same order they are typically e For both Qwen3-Omni and Qwen3-TTS, batching is a pipeline-level optimization: -- Requests are grouped per stage using `runtime.max_batch_size` +- Requests are grouped per stage using `stages[].max_num_seqs` - Each stage executes batch inference with its own scheduler/worker - Stage outputs are routed to downstream stages with per-request mapping preserved @@ -171,7 +171,7 @@ At concurrency 10, batching alone brings Qwen3-TTS RTF from 2.19 (slower than re In decode-heavy serving, repeatedly launching many small kernels from CPU can become a visible overhead. CUDA Graph reduces this overhead by capturing and replaying stable execution graphs. -In stage configs, this is represented by `enforce_eager: false` for stages where graph capture is desired (Thinker/Talker), while Code2Wav keeps eager mode depending on stage behavior. +In deploy configs, this is represented by `enforce_eager: false` for stages where graph capture is desired (Thinker/Talker), while Code2Wav keeps eager mode depending on stage behavior. ### CUDA Graph results on top of batching @@ -320,7 +320,7 @@ In the async chunk pipeline, the standard `codec_chunk_frames` is 25 (each chunk **Dynamic initial chunk sizing (default behavior):** -Rather than using a fixed initial chunk size, vLLM-Omni dynamically selects it based on current server load. The initial chunk size is chosen from power-of-2 steps [2, 4, 8, 16] based on load factor (`active_requests / max_batch_size`): +Rather than using a fixed initial chunk size, vLLM-Omni dynamically selects it based on current server load. The initial chunk size is chosen from power-of-2 steps [2, 4, 8, 16] based on load factor (`active_requests / max_num_seqs`): | Server load | Initial chunk frames | Rationale | | --- | --- | --- | @@ -349,16 +349,15 @@ This overrides the dynamic calculation for that request. **Config (server-side):** ```yaml -runtime: - connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - codec_streaming: true - codec_chunk_frames: 25 # standard chunk size (~2s of audio) - codec_left_context_frames: 25 - # initial chunk is computed dynamically by default - # set initial_codec_chunk_frames: 2 to force a fixed value +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + codec_chunk_frames: 25 # standard chunk size (~2s of audio) + codec_left_context_frames: 25 + # initial chunk is computed dynamically by default + # set initial_codec_chunk_frames: 2 to force a fixed value ``` The 64 ms TTFP result reported above for Qwen3-TTS at concurrency 1 uses the dynamic initial chunk, which picks `initial_codec_chunk_frames=2` at low load. At higher concurrency the dynamic sizing increases the initial chunk to maintain decode efficiency. @@ -408,7 +407,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \ Notes: -- `runtime.max_batch_size` controls stage-level batching. +- `stages[].max_num_seqs` controls stage-level batching. - Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths. - Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior. - Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster @@ -426,33 +425,26 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \ ```yaml async_chunk: true -stage_args: +stages: - stage_id: 0 # thinker - runtime: - max_batch_size: 64 - engine_args: - enforce_eager: false - max_num_batched_tokens: 32768 - custom_process_next_stage_input_func: >- - vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk + max_num_seqs: 64 + enforce_eager: false + max_num_batched_tokens: 32768 - stage_id: 1 # talker - runtime: - max_batch_size: 64 - engine_args: - enforce_eager: false - max_num_batched_tokens: 32768 - custom_process_next_stage_input_func: >- - vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk + max_num_seqs: 64 + enforce_eager: false + max_num_batched_tokens: 32768 - stage_id: 2 # code2wav - runtime: - max_batch_size: 64 - engine_args: - enforce_eager: true - max_num_batched_tokens: 51200 + max_num_seqs: 64 + enforce_eager: true + max_num_batched_tokens: 51200 ``` +The async handoff processors are fixed by the registered Qwen3-Omni +`PipelineConfig`; deploy YAML only selects `async_chunk` and runtime sizing. + #### Reproduce Qwen3-Omni benchmarks ```bash @@ -483,7 +475,7 @@ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ The default config (`qwen3_tts.yaml`) enables the full optimization stack: -- Batching with `max_batch_size: 10` on the Talker stage +- Batching with `max_num_seqs: 10` on the Talker stage - CUDA Graph on the Talker (`enforce_eager: false`) - Async chunk with streaming transport @@ -500,33 +492,29 @@ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ ```yaml async_chunk: true -stage_args: +stages: - stage_id: 0 # Talker (AR decoder) - runtime: - max_batch_size: 10 - engine_args: - enforce_eager: false - max_num_batched_tokens: 512 - custom_process_next_stage_input_func: >- - vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk + max_num_seqs: 10 + enforce_eager: false + max_num_batched_tokens: 512 - stage_id: 1 # Code2Wav (vocoder) - runtime: - max_batch_size: 1 - engine_args: - enforce_eager: true - max_num_batched_tokens: 8192 + max_num_seqs: 1 + enforce_eager: true + max_num_batched_tokens: 8192 -runtime: - connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - codec_streaming: true - codec_chunk_frames: 25 - codec_left_context_frames: 25 +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + codec_chunk_frames: 25 + codec_left_context_frames: 25 ``` +The async Talker-to-Code2Wav processor is defined by the registered Qwen3-TTS +pipeline rather than the deploy YAML. + #### Reproduce Qwen3-TTS benchmarks ```bash diff --git a/docs/features/README.md b/docs/features/README.md index 0dee558a4..7c3e89676 100644 --- a/docs/features/README.md +++ b/docs/features/README.md @@ -20,13 +20,19 @@ an implementation contract has a user-facing workflow. | Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) | | Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide | +The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md) +describes the experimental Qwen3-Omni runtime topology and its current +integration boundaries. It is not yet a supported launch recipe. + Some runtime designs are deliberately not promoted as standalone User Guide features yet: -- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and +- Generic, OmniConnector-based + [Disaggregated Inference](../design/feature/disaggregated_inference.md) is not + yet a standalone User Guide feature. Its [OmniConnector implementations](../design/index.md#communication) describe - topology and transport contracts. Practical deployment configuration remains - under [Pipeline and deploy configurations](../configuration/stage_configs.md). + transport contracts, while practical connector configuration remains under + [Pipeline and deploy configurations](../configuration/stage_configs.md). - [Async Chunk](../design/feature/async_chunk.md) and [Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md) are model- and pipeline-dependent. Use the selected model's deploy @@ -48,10 +54,10 @@ defines the shared configuration and backend extension points. | Goal | User guide | Related design contract | | --- | --- | --- | | Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) | -| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) | +| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [CPU Offloading](../design/feature/offloader/README.md) | | Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) | | Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) | -| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) | +| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) | | Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide | | Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide | | Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide | diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md index 78d680f26..887dfa769 100644 --- a/docs/features/custom_pipeline.md +++ b/docs/features/custom_pipeline.md @@ -101,7 +101,7 @@ outputs = omni.generate( ) # Access custom trajectory data -output = outputs[0].request_output +output = outputs[0] print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}") print(f"Trajectory latents shape: {output.latents.shape}") ``` diff --git a/docs/features/pd_disaggregation.md b/docs/features/pd_disaggregation.md new file mode 100644 index 000000000..ae8497e81 --- /dev/null +++ b/docs/features/pd_disaggregation.md @@ -0,0 +1,140 @@ +# Prefill-Decode (PD) Disaggregation (experimental) + +!!! warning "Experimental design reference" + PD disaggregation is only partially integrated on the current `main` branch. + The runtime contains PD detection and routing scaffolding, but there is no + supported deploy overlay or end-to-end validated launch recipe yet. Treat this + page as a description of the current design, not as production guidance. + +Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical +stages: a prefill worker that processes the prompt and produces KV cache, and a +decode worker that imports that cache and generates tokens. The Talker and +Code2Wav stages remain downstream; PD does not disaggregate the entire +Qwen3-Omni pipeline. + +## Current design + +### Runtime topology + +```mermaid +flowchart LR + client["Client request"] --> orchestrator["Omni Orchestrator"] + orchestrator -->|"original prompt
max_tokens = 1"| prefill["Thinker prefill
KV producer"] + prefill -->|"KV blocks"| connector[("vLLM KV connector")] + prefill -->|"KV routing metadata"| orchestrator + orchestrator -->|"original prompt
remote KV parameters"| decode["Thinker decode
KV consumer"] + connector -->|"remote KV load"| decode + decode -->|"generated text"| client +``` + +The orchestrator owns the request lifecycle, while the vLLM KV connector owns +the bulk KV transfer. The original prompt is submitted to both Thinker stages: +the decode worker still needs the request tokens and multimodal features to +construct its request, while the imported KV cache prevents it from repeating +the completed prefill computation. + +For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage: + +| Stage | Role | Output or handoff | +| --- | --- | --- | +| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output | +| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning | +| 2 | Talker | Converts combined Thinker conditioning into codec codes | +| 3 | Code2Wav | Converts codec codes into audio | + +Text remains a final output of the decode stage, and audio remains a final +output of Code2Wav. + +### Request lifecycle + +1. At startup, PD detection looks for one stage marked `is_prefill_only` and + one downstream stage marked `is_decode_only`. The decode stage must list the + prefill stage as an input source. +2. The orchestrator clones the Thinker sampling parameters for prefill, forces + `max_tokens=1`, clears stop conditions, and sets producer-side + `kv_transfer_params`. This makes prefill finish by length after exporting KV. +3. When prefill finishes, the orchestrator captures its KV routing metadata and + Qwen multimodal output under the original request ID. +4. The orchestrator submits the original prompt to decode with consumer-side + `kv_transfer_params`. Those parameters identify the transfer, prefill engine, + bootstrap endpoint, and, when required by the connector, remote request. +5. The decode engine imports the KV cache and continues generation. The normal + Omni stage path then carries Thinker conditioning through Talker and + Code2Wav. + +### Control and data paths + +PD uses two different connector layers. They solve different handoffs and are +not interchangeable. + +| Path | Payload | Owner | +| --- | --- | --- | +| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator | +| PD data path | Thinker attention KV blocks | vLLM KV connector | +| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors | + +### Detection and validation + +The current runtime accepts one logical prefill/decode pair per pipeline. Each +stage may use tensor parallelism, but both sides of the pair must satisfy these +invariants: + +- both stages define `kv_transfer_config` +- prefill uses `kv_producer` or `kv_both` +- decode uses `kv_consumer` or `kv_both` +- both stages use the same KV connector +- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages +- both stages use the same `tensor_parallel_size` + +These checks happen before request processing so an inconsistent PD topology +fails during engine initialization. + +## What is implemented today + +| Area | Current state | +| --- | --- | +| Pair detection and topology validation | Implemented in the shared PD helper | +| Prefill sampling preparation | Implemented for sync and async entrypoints | +| Prefill-to-decode routing | Implemented in the current Orchestrator | +| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology | +| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor | +| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata | +| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration | +| Performance | No non-regression or throughput/latency result is established for the current path | + +Because these pieces are still converging, the previous legacy `stage_args` YAML +and `vllm serve` command have been removed. The serve CLI now accepts the +deploy-based `stages:` schema, so retaining the legacy recipe would imply a +launch path that current `main` rejects. + +## Configuration contract under development + +The intended deploy integration starts from the +[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml) +and expands its frozen three-stage pipeline into the four stages shown above. +A complete integration must: + +- split the Thinker definition without duplicating the Talker or Code2Wav model + contracts +- remap downstream stage IDs, input sources, and stage connectors +- inject producer and consumer `kv_transfer_config` values into the two Thinker + stages +- allocate separate device sets for prefill and decode +- preserve both final outputs: decode text and Code2Wav audio +- provide connector-specific routing metadata without hard-coding one connector + contract + +The required number of GPUs depends on tensor-parallel sizes and whether Talker +and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one +for decode, and one for the two downstream audio stages. + +## Implementation map + +- PD detection, validation, and sampling-parameter preparation: + [`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py) +- startup detection and connector metadata: + [`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py) +- request-scoped prefill-to-decode routing: + [`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py) +- Qwen3 downstream conditioning helpers: + [Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py) diff --git a/docs/getting_started/installation/README.md b/docs/getting_started/installation/README.md index b49cee794..448ca81f2 100644 --- a/docs/getting_started/installation/README.md +++ b/docs/getting_started/installation/README.md @@ -1,7 +1,7 @@ # Installation !!! important - vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0. + vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0. vLLM-Omni supports the following hardware platforms: diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 59d2aa482..196b6a638 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr ### Pre-built wheels -Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). +Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). === "NVIDIA CUDA" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index a68bfdb58..a8f795861 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.26.0 --torch-backend=auto +uv pip install vllm==0.28.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]' # --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.26.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library ```bash -uv pip install vllm==0.26.0 --torch-backend=auto +uv pip install vllm==0.28.0 --torch-backend=auto ``` -The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. +The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.26.0 +git checkout v0.28.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 -#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl +#For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ --ipc=host \ - vllm/vllm-omni:v0.26.0 \ + vllm/vllm-omni:v0.28.0 \ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -124,7 +124,7 @@ If you want to specify the base vLLM version: ```bash DOCKER_BUILDKIT=1 docker build \ -f docker/Dockerfile.cuda \ - --build-arg BASE_IMAGE=vllm/vllm-openai:v0.26.0 \ + --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \ -t vllm-omni-cuda . ``` diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index c44407299..e0f1614ae 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 +uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 ``` #### Installation of vLLM-Omni @@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm # --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.26.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library ```bash -uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 +uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 ``` -The pre-built 0.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. +The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.26.0 +git checkout v0.28.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` @@ -130,7 +130,7 @@ docker run --rm \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ - vllm/vllm-omni-rocm:v0.26.0 \ + vllm/vllm-omni-rocm:v0.28.0 \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -149,7 +149,7 @@ docker run --rm -it \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ --entrypoint bash \ - vllm/vllm-omni-rocm:v0.26.0 + vllm/vllm-omni-rocm:v0.28.0 ``` # --8<-- [end:pre-built-images] diff --git a/docs/getting_started/installation/gpu/xpu.inc.md b/docs/getting_started/installation/gpu/xpu.inc.md index 8822e3a7c..b9a6e117c 100644 --- a/docs/getting_started/installation/gpu/xpu.inc.md +++ b/docs/getting_started/installation/gpu/xpu.inc.md @@ -23,6 +23,19 @@ vLLM-Omni currently recommends using the Docker image setup steps below. DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g . ``` +This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:` +base image, which Docker pulls automatically. To target a different vLLM release, +pass `--build-arg VLLM_VERSION=`. If that tag has not been published yet, build +the base from upstream's own Dockerfile first and point the build at it: + +```bash +git clone --depth 1 --branch "$VLLM_VERSION" https://github.com/vllm-project/vllm /tmp/vllm +DOCKER_BUILDKIT=1 docker build -t vllm-openai-xpu:local --target vllm-openai \ + -f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm +DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \ + --build-arg VLLM_BASE=vllm-openai-xpu:local . +``` + #### Launch the docker image ##### Launch with OpenAI API Server diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 0d8323760..036c2da17 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,10 +19,10 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.26.0 --torch-backend=auto +uv pip install vllm==0.28.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723 +uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst !!! note It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. - If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. + If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. ## Offline Inference @@ -47,7 +47,7 @@ if __name__ == "__main__": omni = Omni(model="Tongyi-MAI/Z-Image-Turbo") prompt = "a cup of coffee on the table" outputs = omni.generate(prompt) - images = outputs[0].request_output.images + images = outputs[0].images images[0].save("coffee.png") ``` @@ -66,7 +66,7 @@ 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 + # deploy_config="./deploy-config.yaml", # Optional deploy override ) prompts = [ "a cup of coffee on a table", @@ -75,8 +75,7 @@ if __name__ == "__main__": ] omni_outputs = omni.generate(prompts) for i_prompt, prompt_output in enumerate(omni_outputs): - this_request_output = prompt_output.request_output - this_images = this_request_output.images + this_images = prompt_output.images for i_image, image in enumerate(this_images): image.save(f"p{i_prompt}-img{i_image}.jpg") print("saved to", f"p{i_prompt}-img{i_image}.jpg") diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index 8ddb63561..4b72f9a95 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/hooks/generate_examples.py @@ -37,7 +37,6 @@ GENERAL_EXAMPLE_SLUGS = frozenset( "text_to_speech", "text_to_video", "x_to_text", - "x_to_video_audio", } ) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b9686047a..4333c9934 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -21,43 +21,54 @@ th { |--------------|--------|-------------------|------------|---------|-----|-----------|--------| | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) | | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | -| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen1) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — | +| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen1) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — | | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — | | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — | | `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — | | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | -| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) | -| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) | +| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) | +| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) | | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — | | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | -| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — | +| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — | | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | +| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — | | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | +| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) | | `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | -| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — | +| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — | | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — | -| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — | +| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — | +| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — | +| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — | +| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) | +| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) | +| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) | +| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) | | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — | +| `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) | +| `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) | | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | +| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — | -| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — | | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) | -| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — | +| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) | | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — | | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — | -| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — | -| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — | +| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | +| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | +| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — | | `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — | | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | @@ -66,13 +77,11 @@ th { | `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — | | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — | | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — | -| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — | -| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — | -| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — | -| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) | +| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — | +| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | -| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — | +| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | ✅︎ | — | | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — | | `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — | | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — | @@ -86,17 +95,18 @@ th { | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — | | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — | | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — | -| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — | | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — | | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — | | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — | -|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — | +|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — | | `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — | | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | +| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — | |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — | +| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — | The `Recipe` column links the corresponding deployment recipe. `Published` is preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the @@ -109,6 +119,10 @@ support metadata until their recipe is audited. ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend. +See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md) +for offline and experimental realtime usage, supported controls, validation, +and current limitations. + H3 MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend. Other AMD SKUs (e.g. MI325X) are not listed until their diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md index fef28142e..a38512342 100644 --- a/docs/serving/image_edit_api.md +++ b/docs/serving/image_edit_api.md @@ -123,7 +123,14 @@ When `stream=false` or omitted, the endpoint returns the standard image edit res } ], "output_format": null, - "size": null + "size": null, + "metrics": { + "stage_durations": { + "queue_wait_ms": 1.0, + "stage_0_gen_ms": 2.0 + }, + "peak_memory_mb": 1024.0 + } } ``` diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 0deea7f8c..98fb7f380 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -131,7 +131,7 @@ Content-Type: application/json | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) | +| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`). Local files fold `mtime_ns` and `size` into cache keys to automatically reload on-disk edits; HTTP URLs and base64 URIs remain cached by string locator. | | `ref_text` | string | null | Transcript of reference audio | | `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) | @@ -141,7 +141,20 @@ The response shape depends on the streaming parameters: **Non-streaming (default).** With `stream=false` and no `stream_format`, returns the complete clip as binary audio data with an appropriate `Content-Type` header (e.g. -`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported. +`audio/wav`). Because the raw-bytes body has no JSON carrier, successful +non-streaming responses from non-diffusion speech servers report usage through +response headers: + +| Header | Description | +| --- | --- | +| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). | +| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. | +| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. | +| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). | +| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. | + +Diffusion-mode speech servers route through a separate response path and do not +emit these headers. **Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or WAV) as they are decoded. @@ -675,7 +688,7 @@ for result in response.json()["results"]: |-----------|--------|---------|-------------| | `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request | -All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed. +All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_num_seqs` and queues the rest — no client-side throttling needed. For best throughput, set both stages' `max_num_seqs` above 1 via `--stage-overrides`. On the current Qwen3-TTS CustomVoice benchmark, stage 1 performed best at `max_num_seqs: 10`: @@ -798,6 +811,51 @@ If you encounter OOM errors: Use `/v1/audio/voices` to list available voices for the loaded model. +## Orchestration Loop (experimental) + +Multi-stage omni deployments route stage outputs through a single orchestrator +loop. By default that loop polls every stage replica on a 1 ms cadence. An +opt-in event-driven mode replaces the poll with one reader task per live stage +replica awaiting its client directly, and switches the serving-side +final-output drain to a condition-variable wakeup at the same time. + +**Configuration (environment variables):** + +| Variable | Default | Description | +|----------|---------|-------------| +| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. | + +Set it on the process that runs the orchestrator (stage 0 of an omni +deployment) before starting the server: + +```bash +export VLLM_OMNI_EVENT_DRIVEN_ORCH=1 +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \ + --omni \ + --port 8091 +``` + +The server logs the selected loop mode and its reader/poller counts once at +startup, so you can confirm which loop is live. + +Routing, output ordering, and terminal-state behavior are identical on both +loops; only the poll cadence changes. Leaving the variable unset keeps the +legacy poll loop, which is the supported default. + +**Known limitations:** + +- The measured serving A/B (idle CPU 2.43% to 0.07%; TTFP p99 -32% at + concurrency 8) predates the rebuild on the per-replica fault-isolation work + in [#4285](https://github.com/vllm-project/vllm-omni/pull/4285). That work + changed dead-replica handling and reader/poller lifecycle rather than the + steady-state output path, and the parity suite covers it, but the serving + A/B has not been re-run on the current head. +- The diffusion-poller branch is covered by unit tests only. Deployments whose + stages all run as standard engine cores never exercise it, including GLM-TTS, + which deploys its DiT without `stage_type: diffusion`. +- Concurrency 1 and 32 measured at parity with the legacy loop. At 32 the + latency is admission-bound, which this mode does not address. + ## Development Enable debug logging: diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md index 88f74affc..688f5fe9c 100644 --- a/docs/serving/video_stream_api.md +++ b/docs/serving/video_stream_api.md @@ -10,7 +10,7 @@ Each server instance runs a single model specified at startup with `vllm serve < ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \ - --deploy-config vllm_omni/deploy/qwen3_omni.yaml \ + --deploy-config vllm_omni/deploy/qwen3_omni_moe.yaml \ --omni \ --port 8000 \ --trust-remote-code diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 6a94f9dab..7fc7d0ba8 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -161,11 +161,29 @@ are not implemented yet. Models may expose additional V2V controls through `extra_params`. For example, Cosmos3 supports `condition_frame_indexes_vision` and `condition_video_keep` to select which decoded reference frames are used as clean conditioning. Cosmos3 transfer mode -also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus -transfer options such as `control_path`, `control_guidance`, -`control_guidance_interval`, `num_video_frames_per_chunk`, -`num_conditional_frames`, `show_control_condition`, and `show_input`; see the -Cosmos3 recipe for complete examples. +also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints. Each hint +may specify its own `control_path` and `control_weight`. Request-level transfer +options include +`control_guidance`, `control_guidance_interval`, +`emphasize_control_in_prompt`, `num_video_frames_per_chunk`, +`num_conditional_frames`, `show_control_condition`, and `show_input`. Transfer +uses its transfer-specific system prompt and, by default, appends a +control-adherence directive to the positive prompt. It adds duration/FPS and +resolution metadata to both CFG branches but does not add a negative prompt +automatically. The Cosmos3 recipe includes an optional reference negative +prompt and shows how to pass it. Set `emphasize_control_in_prompt`, +`use_duration_template`, or `use_resolution_template` to `false` to disable the +corresponding addition. `negative_metadata_mode` accepts `same`, `inverse`, or +`none` and defaults to `same` for transfer. See the Cosmos3 recipe for complete +examples. + +HTTP redirects for `image_reference.image_url` follow vLLM's +`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to +`1` (the default) to allow redirects or `0` to reject them. A redirect target +can differ from the original host, and vLLM-Omni does not yet fully implement +upstream vLLM's media URL allowlist protection. Deployments should therefore +treat remote media URLs as untrusted and choose this setting as part of their +URL access policy. ### Speech-to-Video diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index b6fd13b12..e2a3b9bfc 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -1,494 +1,195 @@ # Diffusion Attention Backends -This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention. +Use this page to select and configure a diffusion attention backend. Backend- +specific installation, tuning, and compatibility details live in separate +guides so that the selection contract stays easy to scan. -## Overview +Diffusion attention backend selection applies to DiT attention in image and +video generation models. It does **not** change autoregressive LLM attention, +which uses vLLM's attention selector. -Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured. +For the internal selector, registry, and platform contract, see +[Attention Backend Selection](../../design/feature/attention_backend_selection.md). -This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector. +## Choose a guide -The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. +| Need | Guide | +| --- | --- | +| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) | +| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) | +| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) | +| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | +| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) | +| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) | -## Backend Options - -| Value | Notes | -|---|---| -| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. | -| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. | -| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). | -| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. | -| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. | -| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. | -| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. | -| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. | -| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. | -| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. | +## Backend options +| Value | Family | Primary use | Detail | +| --- | --- | --- | --- | +| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) | +| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) | +| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) | +| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) | +| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) | +| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) | +| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) | +| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | +| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | +| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) | +| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) | ## Configuration -Diffusion attention backends can be configured three ways, in priority order: +Backend selection is resolved in this order: -1. **`--diffusion-attention-config`** — structured per-role config (highest priority). -2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend. -3. **Platform default** — used when nothing is configured. +1. `--diffusion-attention-config` per-role configuration. +2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a + global default. +3. The current platform's default. -`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`. +`--diffusion-attention-backend` is shorthand for +`--diffusion-attention-config.default.backend`. Do not pass it together with an +explicit `default.backend` in the structured configuration. ### Global default -Set the default backend for every diffusion attention layer: - ```bash -# CLI flag -vllm-omni serve --diffusion-attention-backend SAGE_ATTN +vllm-omni serve --diffusion-attention-backend FLASH_ATTN -# Environment variable (also recognized for backwards compatibility) -export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN +# Backwards-compatible environment variable +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN ``` ### Per-role configuration -Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order: +Roles are declared by each diffusion model. Common categories are `self` and +`cross`; a model may also use a more specific role such as +`ltx2.audio_to_video`. Resolution order is: -1. Exact `per_role[role]` match -2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`) -3. `default` -4. Platform default - -Use vLLM-style dotted flags or one JSON blob: +1. Exact `per_role[role]` match. +2. Category `per_role[role_category]` match. +3. `default`. +4. Platform default. ```bash # Dotted flags vllm-omni serve \ - --diffusion-attention-config.default.backend FLASH_ATTN \ - --diffusion-attention-config.per_role.cross.backend TORCH_SDPA + --diffusion-attention-config.default.backend FLASH_ATTN \ + --diffusion-attention-config.per_role.cross.backend TORCH_SDPA -# JSON +# Equivalent JSON vllm-omni serve \ - --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}' + --diffusion-attention-config \ + '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}' ``` -A backend that needs configuration exposes it as a typed field on the spec: -`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and -`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)): - -```bash ---diffusion-attention-config.default.backend TRTLLM_ATTN \ ---diffusion-attention-config.default.skip_softmax.target_sparsity 0.5 -``` - -### Programmatic API - -When constructing `OmniDiffusionConfig` directly: +### Python API ```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig +from vllm_omni.diffusion.data import ( + AttentionConfig, + AttentionSpec, + OmniDiffusionConfig, +) config = OmniDiffusionConfig( diffusion_attention_config=AttentionConfig( default=AttentionSpec(backend="FLASH_ATTN"), - per_role={ - "cross": AttentionSpec(backend="TORCH_SDPA"), - }, + per_role={"cross": AttentionSpec(backend="TORCH_SDPA")}, ), ..., ) ``` -A plain dict is also accepted and normalized to `AttentionConfig`. +Backend-specific typed blocks are documented with their consumers: -#### FlashInfer QK16/V8 quantized attention +- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention) + and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization). +- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax). +- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration). +- `fastvideo_vsa_topk`: [FastVideo VSA](attention_backends/fastvideo_vsa.md#choose-top-k). -`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8: - -```python -config = OmniDiffusionConfig( - diffusion_attention_config=AttentionConfig( - default=AttentionSpec( - backend="FLASHINFER_ATTN", - quant=AttnQuantSpec( - dtype_qk="bfloat16", - dtype_vo="fp8_e4m3", - ), - ), - ), - ..., -) -``` - -`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype. - -## Platform Defaults +## Platform defaults ### Blackwell (sm_100 / sm_103 / sm_120 / sm_121) -Auto-route preference, in order: +The CUDA auto-route preference is: -1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path** -2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels) -3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5 -4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel -5. `TORCH_SDPA` — last resort +1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is + available, `head_dim=128`, and the model declares a compatible packed or + mask-free path. +2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available. +3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old. +4. `FLASH_ATTN` when a compatible package is installed. +5. `TORCH_SDPA`. -`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below. +`TRTLLM_ATTN` is not auto-selected on workstation Blackwell +(sm_120/sm_121), for other head dimensions, or for paths that require an +unsupported mask. -`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback. +### Hopper, Ada, and Ampere -The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route. +The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back +to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options. -**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models. +Other platforms validate an explicit backend and choose their own default +through the platform implementation. Check the startup log to confirm the +resolved backend. -### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86) +## Choosing a backend manually -Auto-route preference: +Override the platform default when you need: -1. `FLASH_ATTN` — when `flash-attn` is installed -2. `TORCH_SDPA` — fallback +- a correctness reference (`TORCH_SDPA`); +- a backend-specific workaround; +- training/rollout kernel alignment (Hub backends); or +- an explicitly validated sparse or quantized speedup. -`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware. +The startup log prints the resolved backend and whether it came from explicit +configuration or platform defaulting. If no resolution message appears, check +earlier logs for diffusion-stage initialization failures. -## TRTLLM_ATTN Backend and Skip-Softmax +## Reference benchmark -`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of -that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed -(algorithm: [design doc](../../design/feature/skip_softmax.md)). +The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell +with the same prompt and seed across runs. Treat them as reference results, not +portable guarantees. -Enable it through the typed `skip_softmax` block on the attention spec: - -| Key | Valid values | Meaning | -|---|---|---| -| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. | -| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. | -| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. | - -```bash -vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN", - "skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}' -``` - -Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction): - -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec - -AttentionConfig( - default=AttentionSpec( - backend="TRTLLM_ATTN", - skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86), - ), -) -``` - -**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only -as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with -sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no -`skip_softmax` block. - -Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting -`TRTLLM_ATTN` raises rather than silently degrading. - -## TRTLLM_ATTN SAGE Quantization - -On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen -FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run -in low precision. Enable it through the typed `quant` block on the attention spec: - -| Key | Valid values | Meaning | -|---|---|---| -| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). | -| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). | -| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). | - -V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so -neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels. - -The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the -fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`), -`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`), -`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises. - -```bash -vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN", - "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}' -``` - -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec - -AttentionConfig( - default=AttentionSpec( - backend="TRTLLM_ATTN", - quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16), - ), -) -``` - -**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel -availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103** -(B300); `int8` QK kernels are compiled for **SM100 only**. - -## RAINFUSION_ATTN Backend and Block-Sparse Video Attention - -`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each -128-token block of keys, ranks them per query block, and computes attention against only the -top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are -spatiotemporal neighbours rather than arbitrary rows of the packed sequence. - -Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the -first-frame blocks are always kept dense, which is why the realized sparsity is lower than the -nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an -exempt layer, a sequence with no published video geometry, a video segment under 32 blocks — -delegates to `FLASH_ATTN`, so a model can select this backend unconditionally. - -Enable it through the typed `block_sparse` block on the attention spec, shared by every -block-sparse backend: - -| Key | Valid values | Meaning | -|---|---|---| -| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. | -| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. | -| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. | - -```bash -vllm-omni serve MiniMaxAI/MiniMax-H3 \ - --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN", - "block_sparse": {"sparsity": 0.8, "start_step": 0}}}' -``` - -Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction): - -```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec - -AttentionConfig( - default=AttentionSpec( - backend="RAINFUSION_ATTN", - block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"), - ), -) -``` - -**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the -cheapest way to recover structure, because the early high-noise steps decide global layout while the -later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for -`skip_layers` last, once an A/B against dense at the same seed points at specific blocks. - -Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also -incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank -blocks — use Ulysses SP (`ring_degree=1`). - -### Which geometries run sparse - -MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when -necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2 -does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the -promoted rows remain visible to every query and key. - -For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row -128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix. -MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask -and kernel tiling aligned, at the cost of retaining more blocks. - -8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8 -creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent -grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release -that implements protected video tails for rf_v2. - -## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell) - -Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`. - -| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN | -|---|---|---|---|---| +| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` | +| --- | --- | ---: | ---: | ---: | | HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s | | Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** | | Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s | | FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s | -Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise. +Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads +were close enough that users should benchmark their exact model and shape. -## Known Limitations +## Compatibility anchors -### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile +The following headings preserve links to sections that moved into dedicated +guides. -LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121). +## TRTLLM_ATTN Backend and Skip-Softmax -**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0: +See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax). -```bash -DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model Lightricks/LTX-2 ... -``` +## TRTLLM_ATTN SAGE Quantization -### FlashAttention-4 on Blackwell +See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization). -Install the optional CUDA 13 extra to use the CuTe-based FA4 path: +## RAINFUSION_ATTN Backend and Block-Sparse Video Attention -```bash -pip install 'vllm-omni[fa4]' -``` - -`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2 -when it is unavailable. Version `4.0.0b18` is required; earlier beta -wheels had known JIT failures on Blackwell. - -## Choosing a Backend Manually - -### When to override the default - -- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested. -- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate. -- **Workaround for known issues**: see Known Limitations above. - -### Verifying which backend is in use - -The startup log prints one of: - -``` -Using diffusion attention backend 'CUDNN_ATTN' # explicit override -Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route -Defaulting to diffusion attention backend SDPA # nothing else available -``` - -If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures. +See [RainFusion](attention_backends/rainfusion.md). ## SageAttention Installation -vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni. - -Build from source: - -```bash -git clone https://github.com/thu-ml/SageAttention.git -cd SageAttention - -export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 -pip install . --no-build-isolation -``` - -Quick check: - -```bash -python -c "import sageattention; print(sageattention.__file__)" -``` +See [SageAttention](attention_backends/sage.md#installation). ## SageAttention3 Installation -vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni. - -Build from source: - -```bash -git clone https://github.com/thu-ml/SageAttention.git -cd SageAttention/sageattention3_blackwell -python setup.py install -``` - -Quick check: - -```bash -python -c "import sageattn3; print(sageattn3.__file__)" -``` - -Notes: - -- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class. -- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness. +See [SageAttention](attention_backends/sage.md#sageattention3-installation). ## HuggingFace Kernels Hub Backends -To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout. - -The following backend options are supported: -- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`) -- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only) - -### Installation - -To use these backends, you must install the `kernels` library: - -```bash -pip install kernels==0.14.1 -``` - -If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`. - -### Usage - -Select a Hub backend using the global CLI flag or environment variables: - -```bash -# Environment variable -export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB - -# CLI flag -vllm-omni serve --diffusion-attention-backend FLASH_ATTN_HUB -``` - -## Usage Examples - -### Default (auto-route) - -```bash -python examples/offline_inference/text_to_video/text_to_video.py \ - --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 50 --seed 42 --guidance-scale 6.0 \ - --output hv15.mp4 -``` - -On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line. - -### Explicit backend selection - -```bash -DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model Lightricks/LTX-2 \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 40 --seed 42 --guidance-scale 4.0 \ - --output ltx2.mp4 -``` - -### SageAttention (lossy) - -```bash -DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \ - --tensor-parallel-size 2 \ - --output hv15_sage.mp4 -``` - -Example: Wan2.2 TI2V 5B - -```bash -DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ - --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ - --prompt "A dog running across a field of golden wheat." \ - --height 704 --width 1280 --num-frames 49 \ - --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \ - --tensor-parallel-size 2 \ - --output outputs/wan22_sage.mp4 -``` - -### Enable SageAttention3 - -Example: - -```bash -DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \ - --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ - --prompt "A dog running across a field of golden wheat." \ - --height 480 --width 832 --num-frames 33 \ - --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \ - --tensor-parallel-size 2 \ - --output outputs/hv15_sage3.mp4 -``` +See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md). diff --git a/docs/user_guide/diffusion/attention_backends/dense_backends.md b/docs/user_guide/diffusion/attention_backends/dense_backends.md new file mode 100644 index 000000000..030550dac --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/dense_backends.md @@ -0,0 +1,110 @@ +# Dense Attention Backends + +Use these backends when you want dense attention without backend-specific +sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark +the fastest compatible kernel for your model, shape, and hardware. + +For selection precedence and per-role configuration, see the +[attention backend overview](../attention_backends.md). + +## `TORCH_SDPA` + +`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's +dispatcher choose the implementation. It is always available and is the most +conservative reference when validating another backend. + +```bash +vllm-omni serve --diffusion-attention-backend TORCH_SDPA +``` + +## `FLASH_ATTN` + +`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it +prefers the CuTe-based FlashAttention 4 path and falls back to +FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and +Ampere it is the preferred automatic route when a compatible package is +installed. + +```bash +vllm-omni serve --diffusion-attention-backend FLASH_ATTN +``` + +### FlashAttention 4 on Blackwell + +Install the optional CUDA 13 extra: + +```bash +pip install 'vllm-omni[fa4]' +``` + +Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on +Blackwell. If the CuTe path is unavailable, the backend falls back to the +compatible FlashAttention 3 or 2 path. + +## `CUDNN_ATTN` + +`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful +for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN +9.5 or newer is available and the higher-priority TRTLLM route is not +compatible. + +```bash +vllm-omni serve --diffusion-attention-backend CUDNN_ATTN +``` + +### LTX-2.0 limitation + +LTX-2 audio attention has a symbolic head dimension during `torch.compile` +tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo +aborts compilation. This is tracked in +[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121). + +Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround: + +```bash +DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \ + python examples/offline_inference/text_to_video/text_to_video.py \ + --model Lightricks/LTX-2 ... +``` + +## `FLASHINFER_ATTN` + +`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit +option on CUDA platforms and an automatic Blackwell fallback when FlashInfer +is installed but cuDNN is too old for `CUDNN_ATTN`. + +```bash +vllm-omni serve --diffusion-attention-backend FLASHINFER_ATTN +``` + +### FlashInfer quantized attention + +The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K +in FP16 or BF16 and use FP8 E4M3 for V: + +```python +from vllm_omni.diffusion.data import ( + AttentionConfig, + AttentionSpec, + AttnQuantSpec, + OmniDiffusionConfig, +) + +config = OmniDiffusionConfig( + diffusion_attention_config=AttentionConfig( + default=AttentionSpec( + backend="FLASHINFER_ATTN", + quant=AttnQuantSpec( + dtype_qk="bfloat16", + dtype_vo="fp8_e4m3", + ), + ), + ), + ..., +) +``` + +`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype +configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization +schema is also consumed by TRTLLM, but each backend validates its own allowed +fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization). diff --git a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md new file mode 100644 index 000000000..2439633d4 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md @@ -0,0 +1,125 @@ +# FastVideo VSA + +FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's +self-attention by partitioning the post-patch latent grid into spatiotemporal +blocks. For every query block, VSA scores the key/value blocks and computes +attention only against the selected top-k blocks. + +The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope. + +VSA is a CUDA-only, explicitly selected backend. It requires the +`fastvideo-kernel` package and currently supports non-causal self-attention +with equal query and key/value sequence lengths. Unsupported shapes, masks, +dtypes, sequence-parallel execution, or kernel failures fall back to +`TORCH_SDPA` and emit a warning with the reason. + +## Enable the backend + +For online serving, select the backend with the existing attention backend +flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained +for every query block: + +```bash +vllm-omni serve \ + --diffusion-attention-backend FASTVIDEO_VSA \ + --fastvideo-vsa-topk 64 +``` + +The backwards-compatible environment variable selects the backend with the +default `topk=64`: + +```bash +export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA +vllm-omni serve +``` + +To tune top-k, pass the CLI backend and top-k flags together as shown above, +or use the equivalent structured configuration: + +```bash +vllm-omni serve \ + --diffusion-attention-config \ + '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}' +``` + +Do not combine `--diffusion-attention-backend` with an explicit +`diffusion_attention_config.default.backend`. The top-k value must be positive +and is valid only when the default backend is `FASTVIDEO_VSA`. + +For a deploy YAML stage, use either the shorthand fields: + +```yaml +stages: + - stage_id: 0 + diffusion_attention_backend: FASTVIDEO_VSA + fastvideo_vsa_topk: 64 +``` + +or the structured configuration: + +```yaml +stages: + - stage_id: 0 + diffusion_attention_config: + default: + backend: FASTVIDEO_VSA + fastvideo_vsa_topk: 64 +``` + +## Choose top-k + +At runtime the backend logs the sequence shape and derived block count: + +```text +FASTVIDEO_VSA routing: seq_len=27280, dit_seq_shape=(31, 22, 40), +block_size=(4, 8, 8), num_blocks=120, topk=64, keep_ratio=53.3%, +checkpoint_mode=native, route=VSA +``` + +Use `num_blocks` as the upper bound when tuning top-k. Top-k is the number of +key/value blocks retained for each query block, not the number of tokens or +the total number of blocks processed by the layer. + +- A smaller top-k increases sparsity and may reduce attention computation, but + it can remove relevant blocks and reduce visual quality or temporal + consistency. Routing and padding overhead can also make a smaller value + slower for some shapes. +- A larger top-k retains more context and generally approaches dense-attention + quality, but reduces the potential speedup and increases memory traffic. +- `topk > num_blocks` is invalid for the runtime shape and falls back to SDPA. +- For a native checkpoint, `topk == num_blocks` routes to SDPA because scoring + every block provides no sparsity benefit. +- For a FastVideo DMD checkpoint, `topk == num_blocks` stays on the VSA + all-block path so the checkpoint keeps its trained compensation semantics. + +There is no universally optimal value. Resolution, frame count, GPU, kernel +version, and checkpoint all affect both quality and latency. Start from the +logged `num_blocks`, test several keep ratios on the target workload, and +compare output quality against the same checkpoint running dense attention. + +## Checkpoint behavior + +The backend does not expose a user-selectable gate mode. Wan checkpoints that +contain learned `to_gate_compress` weights use that projection automatically. +When those weights are absent, the unused projection is removed and VSA uses +the sparse branch without learned compensation. + +FastVideo DMD checkpoints use their fixed distilled timestep schedule. Native +Wan checkpoints keep their normal scheduler and inference-step configuration; +selecting VSA does not turn a native checkpoint into a distilled model. + +## Verify routing and fallback + +Check the startup and first-forward logs instead of assuming that selecting +the backend guarantees sparse execution: + +- `route=VSA` means top-k block selection is active. +- `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all + blocks through the VSA kernel. +- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA + executed; the warning includes the reason. + +The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks, +standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask, +and no active sequence-parallel context. NPU and XPU paths do not execute the +FastVideo VSA CUDA kernel. diff --git a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md new file mode 100644 index 000000000..946786b97 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md @@ -0,0 +1,32 @@ +# Hugging Face Hub Attention Backends + +Hub backends load the same published kernels commonly used with Hugging Face +Diffusers. They are useful when training and serving must use matching kernel +implementations to reduce numerical drift and sampling divergence. + +| Backend | Kernel | Platform | +| --- | --- | --- | +| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs | +| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer | + +## Installation + +```bash +pip install kernels==0.14.1 +``` + +If `kernels` is unavailable, vLLM-Omni warns and falls back to the +corresponding local FlashAttention implementation. On CUDA devices older than +Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`. + +## Usage + +```bash +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB + +# Equivalent CLI selection +vllm-omni serve --diffusion-attention-backend FLASH_ATTN_HUB +``` + +For per-role kernel alignment, configure the Hub backend on only the relevant +role; see the [attention backend overview](../attention_backends.md#per-role-configuration). diff --git a/docs/user_guide/diffusion/attention_backends/rainfusion.md b/docs/user_guide/diffusion/attention_backends/rainfusion.md new file mode 100644 index 000000000..00f4e37a1 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/rainfusion.md @@ -0,0 +1,80 @@ +# RainFusion Attention + +`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video +attention on Ascend NPU. It pools 128-token key blocks, ranks them per query +block, and attends to the highest-scoring blocks after arranging video tokens +in `(t, h, w)` order. + +Only the video segment is sparse. Prefix rows such as text, visual conditions, +and audio, plus first-frame blocks, remain dense. Unsupported calls—including +warmup steps, skipped layers, missing video geometry, and video segments below +32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion +globally. + +## Configuration + +| Key | Valid values | Meaning | +| --- | --- | --- | +| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity | +| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense | +| `end_step` | integer, `>= 0` | Number of final denoise steps kept dense (tail fallback) | +| `precision` | `"bf16"`, `"fp8"`, `"mix"` | Kernel precision mode; default `"bf16"`. Requires MindIE-SD with `sparse_attention(precision=...)` | +| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense | + +```bash +vllm-omni serve MiniMaxAI/MiniMax-H3 \ + --diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\ + "block_sparse":{"sparsity":0.8,"start_step":0}}}' +``` + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec + +config = AttentionConfig( + default=AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse=BlockSparseSpec( + sparsity=0.8, + start_step=0, + skip_layers="0-1", + ), + ), +) +``` + +Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense +early-step window is usually the cheapest way to recover global structure; +use layer exclusions only after a same-seed dense comparison identifies +sensitive blocks. + +## Requirements and compatibility + +RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another +platform raises. It is incompatible with ring sequence parallelism because +the kernel needs the complete key sequence for ranking; use Ulysses sequence +parallelism with `ring_degree=1`. + +The `precision` knob requires a MindIE-SD release whose `sparse_attention` +accepts `precision=`; older releases accept it through `**kwargs` but +silently ignore it and run the BF16 path. `RAINFUSION_ATTN` raises at runtime +when a non-`bf16` precision is requested without that support, so pin a +compatible MindIE-SD release when using `precision="mix"` or `"fp8"`. + +## Geometry handling + +RainFusion handles arbitrary video grids by rearranging video spatially and, +when necessary, promoting a real-video suffix to the always-kept prefix. It +does not pad because `rf_v2` does not consume a padding attention mask. + +For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows. +The implementation promotes 2,976 rows and leaves 59,520 sparse rows +(`465 x 128`) so the mask and kernel tiling remain aligned. + +Spatial grids aligned to 8x8 are still preferable. A latent height or width +not divisible by 8 creates an always-kept suffix and reduces realized +sparsity. At the image level, that means width and height divisible by 256. +Protected video tails require a compatible MindIE-SD release. + +For common configuration and selector behavior, see the +[attention backend overview](../attention_backends.md) and the +[backend selection design](../../../design/feature/attention_backend_selection.md). diff --git a/docs/user_guide/diffusion/attention_backends/sage.md b/docs/user_guide/diffusion/attention_backends/sage.md new file mode 100644 index 000000000..1e1fa5b60 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/sage.md @@ -0,0 +1,63 @@ +# SageAttention + +SageAttention backends provide lossy low-precision attention for diffusion +models. Validate output quality against `TORCH_SDPA` at the same seed before +using either backend in production. + +## `SAGE_ATTN` + +`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16 +accumulation. + +### Installation + +Install SageAttention into the same environment as vLLM-Omni: + +```bash +git clone https://github.com/thu-ml/SageAttention.git +cd SageAttention +export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 +pip install . --no-build-isolation +``` + +Verify the installation: + +```bash +python -c "import sageattention; print(sageattention.__file__)" +``` + +Select it globally: + +```bash +vllm-omni serve --diffusion-attention-backend SAGE_ATTN +``` + +## `SAGE_ATTN_3` + +`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation. + +### SageAttention3 installation + +```bash +git clone https://github.com/thu-ml/SageAttention.git +cd SageAttention/sageattention3_blackwell +python setup.py install +``` + +Verify the installation: + +```bash +python -c "import sageattn3; print(sageattn3.__file__)" +``` + +```bash +vllm-omni serve --diffusion-attention-backend SAGE_ATTN_3 +``` + +`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class +GPU. Its kernel assumes the query-head count equals the key/value-head count. +GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for +correctness. + +For common configuration and platform routing, see the +[attention backend overview](../attention_backends.md). diff --git a/docs/user_guide/diffusion/attention_backends/trtllm.md b/docs/user_guide/diffusion/attention_backends/trtllm.md new file mode 100644 index 000000000..dd36ba61e --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends/trtllm.md @@ -0,0 +1,103 @@ +# TRTLLM Attention + +`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense +BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization. + +For common selection and per-role configuration, see the +[attention backend overview](../attention_backends.md). + +## Requirements + +Selecting `TRTLLM_ATTN` requires: + +- datacenter Blackwell, sm_100 or sm_103; +- `head_dim=128`; +- FlashInfer; and +- a model-declared packed or mask-free compatible path. + +It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible +explicit selection raises instead of silently changing the backend. + +## Dense mode + +Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16: + +```bash +vllm-omni serve --diffusion-attention-backend TRTLLM_ATTN +``` + +## Skip-Softmax + +Skip-Softmax trades some fidelity for speed by skipping low-contribution +attention work. See the [feature design](../../../design/feature/skip_softmax.md) +for the algorithm and implementation contract. + +| Key | Valid values | Meaning | +| --- | --- | --- | +| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve | +| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` | +| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` | + +```bash +vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\ + "skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}' +``` + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec + +config = AttentionConfig( + default=AttentionSpec( + backend="TRTLLM_ATTN", + skip_softmax=SkipSoftmaxSpec( + target_sparsity=0.65, + disabled_until_timestep=0.86, + ), + ), +) +``` + +Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase +sparsity only after comparing quality against dense output at the same seed. +The benefit grows with sequence length because only attention is accelerated. + +## SAGE quantization + +TRTLLM SAGE quantizes Q/K per block and V per channel so both attention +matrix multiplications use low precision. + +| Key | Valid values | Meaning | +| --- | --- | --- | +| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense | +| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` | +| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` | + +V is always quantized per channel to FP8 E4M3 and K-smoothing is internal. +Only block sizes with compiled kernels are accepted. + +```bash +vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\ + "quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}' +``` + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec + +config = AttentionConfig( + default=AttentionSpec( + backend="TRTLLM_ATTN", + quant=AttnQuantSpec( + dtype_qk="fp8_e4m3", + q_block_size=1, + k_block_size=16, + ), + ), +) +``` + +This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are +available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only. +The shared `AttnQuantSpec` is backend typed: values intended for +`FLASHINFER_ATTN` are rejected here and vice versa. diff --git a/docs/user_guide/diffusion/cpu_offload.md b/docs/user_guide/diffusion/cpu_offload.md index 02ab9f074..93ac5e48e 100644 --- a/docs/user_guide/diffusion/cpu_offload.md +++ b/docs/user_guide/diffusion/cpu_offload.md @@ -1,389 +1,74 @@ # CPU Offloading for Diffusion Models -## Overview +vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose +the coarsest strategy that meets your memory target: finer-grained streaming +saves more device memory but adds more transfer and setup overhead. -vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models: +For the shared factory, discovery, and lifecycle contract, see the +[CPU Offloading design](../../design/feature/offloader/README.md). -1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time. -2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap. +## Choose a strategy -Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority. +| Strategy | Device residency | Parallel scope | Guide | +| --- | --- | --- | --- | +| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) | +| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) | +| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) | +All strategies use pinned host memory for faster transfers where applicable. +Configuration priority is: + +1. Distributed layerwise offloading. +2. Layerwise offloading. +3. Model-level offloading. + +Treat the flags as mutually exclusive. If more than one is enabled, the +higher-priority strategy is selected. + +## Quick selection + +- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole + encoders and DiTs is enough to fit and phase-boundary transfers are + acceptable. +- Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video + DiTs where block transfers can overlap computation. +- Use [distributed layerwise offloading](offloader/distributed_layerwise_offload.md) + when a multi-rank deployment also needs bounded device residency and, + optionally, sharded host weights. + +## Supported models + +| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes | +| --- | --- | --- | :---: | :---: | :---: | --- | +| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` | +| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` | +| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` | +| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` | +| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` | +| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` | +| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` | +| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` | +| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` | +| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules | +| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` | + +Model-level support requires discoverable DiT and encoder components. +Layerwise support requires transformer block topology. Distributed support +reuses that topology but still requires validation for the model, checkpoint, +and parallel configuration. + +## Compatibility anchors + +These headings preserve links to sections that moved into dedicated guides. ## Model-level (Sequential) Offloading -### How It Works - -Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks: - -- **When encoders run**: DiT transformer is offloaded to CPU -- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU. -- **VAE**: Stays resident on GPU - -Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed. - -### Usage - -**Python API:** -```python -from vllm_omni import Omni - -m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True) -``` - -**CLI:** -```bash -vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload -``` - -### To Support a Model - -Implement the `SupportsComponentDiscovery` protocol to declare which -submodules serve as pipeline components (used by offloading, HSDP -sharding, and other framework features): - -```python -from typing import ClassVar -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery - -class MyPipeline(nn.Module, SupportsComponentDiscovery): - _dit_modules: ClassVar[list[str]] = ["transformer"] - _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"] - _vae_modules: ClassVar[list[str]] = ["vae"] - _resident_modules: ClassVar[list[str]] = [] # optional - - def __init__(self): - super().__init__() - self.transformer = ... # DiT — stays on GPU during denoising - self.text_encoder = ... # Encoder — offloaded to CPU during denoising - self.vision_model = ... # Encoder — offloaded to CPU during denoising - self.vae = ... # VAE — always on GPU -``` - -- `_dit_modules`: attribute names of denoising submodules (kept on GPU - during the diffusion loop). -- `_encoder_modules`: attribute names of encoder/vision submodules - (offloaded to CPU during the diffusion loop). -- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not - part of the mutual exclusion hooks). -- `_resident_modules`: attribute names of small submodules that must - stay on GPU during layerwise offloading (e.g. embedders, connectors). - Optional — defaults to `[]`. - -All attribute names support dotted paths for nested submodules -(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`). - -Both DiT and encoder lists are needed because the offload hooks use -mutual exclusion: when one group runs, the other moves to CPU. - -### Limitations -- Cold start latency increases -- Adds overhead from CPU-GPU transfers between encoder and denoising phases -- Support single GPU only for now - - -### Component offloading for split models (e.g. Cosmos3) - -Some models split their transformer into mutually-exclusive *components* that run -in different phases of a single forward pass rather than as separate pipeline -components -- e.g. Cosmos3's understanding (reasoner) component runs once per -generation while the generation (generator) component runs every denoising step. -Such models have no separate text encoder to swap against, so the transformer -owns a small model-local offload path and wraps each phase with -`with self._offload_context(name):` - -```python -class Cosmos3VFMTransformer(nn.Module): - def forward(self, ...): - with self._offload_context("reasoner"): - ... # understanding pass, runs once - with self._offload_context("generator"): - ... # denoising pass, runs every step -``` - -Model-level offloading then keeps exactly one component GPU-resident at a time -(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The -pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the -transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading -works for these models too -- each component declares its own block container via -`_layerwise_offload_blocks_attrs`. - +See [Model-Level Offloading](offloader/cpu_offload.md). ## Layerwise (Blockwise) Offloading -### How It Works - -Layerwise offloading keeps only one transformer block on GPU at a time. - -As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU. - -**Execution Flow:** - -| Block | Pre-forward Hook | Forward | Post-forward Hook | -|-------|------------------|---------|-------------------| -| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 | -| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 | -| ... | ... | ... | ... | -| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) | - -Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward. - -Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines. - -### Usage - -**Python API:** -```python -from vllm_omni import Omni - -# Text-to-video -m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True) - -# Or image-to-video -m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True) -``` - -**CLI:** -```bash -# Text-to-video -vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload - -# Or image-to-video -vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload -``` - -### To Support a Model - -Models must define the blocks attribute name for layerwise offloading: - -```python -class WanTransformer3DModel(nn.Module): - _layerwise_offload_blocks_attrs = ["blocks"] # Attribute names containing transformer blocks - - def __init__(self): - self.blocks = nn.ModuleList([...]) # Transformer blocks -``` - -For models with multiple block types: - -```python -class Flux2Transformer2DModel(nn.Module): - _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] -``` - -### Limitations -- Cold start latency increases because offloaded components must be moved to CPU - during setup; layerwise offload may add extra weight consolidation and pinning - work. -- Performance depends on compute cost and H2D bandwidth as well -- Support single GPU only for now - +See [Layerwise Offloading](offloader/layerwise_offload.md). ## Distributed Layerwise Offloading -### How It Works - -Distributed layerwise offloading extends single-GPU layerwise offloading to -multi-device deployments. In the default mode, each rank stores only -**1/dp_size** of the model weights in host memory; full layer weights are -reconstructed at runtime via **AllGather** on a dedicated communication stream, -overlapped with computation via a fixed double-buffer scheme. The -`--dlo-no-use-allgather` mode instead streams standard-loader rank-local -weights independently and does not shard weights across ranks. - -For the implementation design and compatibility matrix, see -[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md). - -**Key features:** -- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights, - reconstructed per-layer via `all_gather_into_tensor` -- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any - time, regardless of model size -- **DP multi-concurrency**: N concurrent requests processed in parallel - (AllGather only gathers weight shards, which are request-independent) -- **mmap weight loading**: weights loaded as mmap views pointing to shared OS - page cache, eliminating O(dp_size × model_size) RSS during model creation -- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU - (CANN/HCCL) via vLLM-Omni's platform abstraction - -**Execution Flow (per device):** - -``` -Compute Stream: [Layer N] [Layer N+1] [Layer N+2] -H2D Stream: [H2D Shard N+1] [H2D Shard N+2] -AllGather: [AG N+1] [AG N+2] -Slot Usage: Slot0: Layer N Slot1: Layer N+1 -``` - -### Usage - -**CLI:** -```bash -# 4× GPU/NPU with AllGather (recommended) -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --data-parallel-size 4 - -# Without DLO AllGather (each rank streams standard-loader rank-local weights) -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --data-parallel-size 4 \ - --dlo-no-use-allgather - -# With SP instead of DP (long sequences) -vllm serve /path/to/model --omni \ - --enable-distributed-layerwise-offload \ - --usp 4 -``` - -**Python API:** -```python -from vllm_omni import Omni - -m = Omni( - model="/path/to/model", - enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, # default -) -``` - -### CLI Flags - -| Flag | Description | Default | -|------|-------------|---------| -| `--enable-distributed-layerwise-offload` | Enable DLO | `false` | -| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` | -| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` | -| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` | - -### How model weights are loaded (mmap path) - -When DLO + AllGather is active, the offloader: - -1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`) - from the normally-created transformer -2. **Converts to meta device** via `to_empty(device="meta")`, releasing random - initialization weights -3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`, - which return views into the OS page cache (shared across all ranks, 0 RSS) -4. **Calls `post_load_weights()`** to apply model-specific dtype conversions - (e.g. Cosmos3's `time_embedder` → FP32) -5. **Restores non-persistent buffers** from saved copies - -This approach requires **zero model-specific code changes** — no pipeline or -transformer modifications are needed. - -### OffloadPlan (declarative topology metadata) - -Models can optionally declare an `OffloadPlan` class variable to provide -topology metadata (block attribute names, submodules to offload) without -any offload-specific logic: - -```python -from vllm_omni.diffusion.offloader import OffloadPlan - -class MyPipeline(nn.Module): - _dit_modules = ["transformer"] - _offload_plan = OffloadPlan( - block_attrs={"transformer": ("blocks",)}, - offload_submodules={"context_encoder": "layers"}, - ) -``` - -When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs` -and heuristic attribute search. This is backward-compatible — existing models -work without any changes. - -### DP Multi-concurrency - -When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches -up to `dp_size` requests per denoise step. Each DP rank processes a different -request while AllGather synchronizes only weight shards (request-independent). - -**Requirements for concurrent requests:** -- `num_inference_steps` must be specified explicitly (None is not allowed) -- All concurrent requests must have the same `num_inference_steps` value -- These constraints exist because AllGather is a collective that requires - every rank to participate at each denoise step - -### Limitations -- Online quantization (FP8) is incompatible with mmap loading — falls back to - regular loading only when the model is outside the DLO+AllGather mmap path; - online quantization with DLO+AllGather is rejected. Use - `--dlo-no-use-allgather` or disable online quantization. -- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware - weight-loader callbacks are bypassed. The no-AllGather path retains the - standard TP loader and is experimental. -- HSDP + AllGather is rejected because it would double-shard parameters. HSDP - with no-AllGather is accepted at configuration level, but has limited - end-to-end validation. -- `num_inference_steps=None` is not allowed in DP multi-concurrency mode - -**Module Discovery** - -The offloader discovers pipeline components in two ways: - -1. **Protocol-based** (preferred): If the pipeline implements - `SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`, - `_vae_modules`, and `_resident_modules` class variables are used - directly. All attribute names support dotted paths (e.g. - `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules. - -2. **Fallback attribute scan**: Otherwise, the offloader scans for - well-known attribute names: - - **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model` - - **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` - - **VAE**: `vae`, `audio_vae` - -**Hook System** - -Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code. - -**Backend Architecture** - -``` -OffloadBackend (base class) -├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap) -│ (delegates to a pipeline's enable_omni_model_cpu_offload -│ for split models like Cosmos3) -├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook -│ (single-GPU, full weights on host) -└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook - (multi-GPU, sharded AllGather or rank-local mode) -``` - -Factory function `get_offload_backend()` selects the appropriate backend based on -configuration. - -For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's -`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its -reasoner/generator components inside the model forward pass. - - -## Supported Models - -| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) | -|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------| -| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` | -| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` | -| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` | -| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` | -| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` | -| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` | -| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` | -| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | -| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | -| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` | -| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` | -| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` | - -**Notes:** -- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally -- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks -- Distributed Layerwise Offloading reuses the layerwise block topology, but the - model and parallelism combination still needs validation. AllGather mode - additionally requires a compatible checkpoint/mmap layout; no-AllGather mode - retains the standard loader's rank-local tensors. See the - [DLO design](../../design/feature/distributed_layerwise_offload.md) and - [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) - for usage examples. +See [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md). diff --git a/docs/user_guide/diffusion/execution_modes.md b/docs/user_guide/diffusion/execution_modes.md index 9611a8210..f9fbc1be0 100644 --- a/docs/user_guide/diffusion/execution_modes.md +++ b/docs/user_guide/diffusion/execution_modes.md @@ -101,8 +101,16 @@ set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure `--max-num-seqs >1`. See the [HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md) for its validated configuration. Helios supports single-request step execution only: use -`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's -documentation and source for the latest support status. +`--step-execution --max-num-seqs 1` for Helios. MiniMax H3 supports step-wise +continuous batching by packing co-batched requests into one sequence that keeps +a separate attention document per request; that layout needs a backend which +honors the packed `cu_seqlens` metadata, so run it with +`--diffusion-attention-backend FLASH_ATTN` (other backends stay correct but +fall back to one transformer forward per request). Batching H3 does not improve +its throughput — see the measured numbers in the +[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md) +— so keep `--max-num-seqs 1` unless you need step-level scheduling. Consult the +selected pipeline's documentation and source for the latest support status. ## Streaming Output @@ -207,18 +215,16 @@ outputs = omni.generate( ) ``` -## Stage Configuration +## Deploy Configuration Use the equivalent engine arguments in a deployment YAML: ```yaml -stage_args: +stages: - stage_id: 0 - stage_type: diffusion - engine_args: - step_execution: false - max_num_seqs: 4 - request_batch_max_wait_ms: 20 + step_execution: false + max_num_seqs: 4 + request_batch_max_wait_ms: 20 ``` For step execution, set `step_execution: true` and remove diff --git a/docs/user_guide/diffusion/lora.md b/docs/user_guide/diffusion/lora.md index 44705129f..3c8582af0 100644 --- a/docs/user_guide/diffusion/lora.md +++ b/docs/user_guide/diffusion/lora.md @@ -12,15 +12,17 @@ LoRA adapters are lightweight, model-specific fine-tuning weights that can be ap ## LoRA Adapter Format ### PEFT (Parameter-Efficient Fine-Tuning) format (default) + A typical PEFT-format LoRA adapter directory structure: -``` +```text lora_adapter/ ├── adapter_config.json └── adapter_model.safetensors ``` The `adapter_config.json` file contains metadata about the LoRA adapter, including: + - `r`: LoRA rank - `lora_alpha`: LoRA alpha scaling factor - `target_modules`: List of module names to apply LoRA to @@ -67,7 +69,7 @@ For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization. | Pipeline | Supported distilled LoRA repo | Notes | -|----------|--------------------------------|-------| +| ---------- | -------------------------------- | ------- | | `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. | | `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. | | `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. | @@ -240,6 +242,37 @@ Notes: - Output quality and speed depend on the replacement checkpoints and sampling params you choose. - If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path .safetensors .safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above. +## MiniMax-H3 adapter-declared schedules + +MiniMax-H3 supports two few-step mechanisms that must not be conflated: + +- **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into + `model_index.json`. Requests must pass `num_inference_steps` as the interval + count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`). +- **Runtime Turbo LoRA**: LightX2V Turbo artifacts request five sigma points and + enforce `flow_shift=6`, `audio_flow_shift=3`. +- **Runtime native LoRA**: FlashGen-style artifacts declare + `key_format=minimax-h3-native` and embed `base_schedule` in safetensors + metadata. When active, the adapter schedule overrides the base checkpoint and + requests must use the interval-count contract (`num_inference_steps=4`). + Request-mode generation may omit the field and take the count from the adapter + schedule; step execution requires it explicitly, because the step scheduler + derives the total step count from the request before the adapter schedule is + known. + +Native artifacts also declare `qkv_layout=grouped`. The H3 loader reorders fused +`qkv_proj` LoRA rows with the same `_reorder_grouped_qkv_to_qkv` helper used for +base-weight loading, then binds the packed Q/K/V slices through the legacy PEFT +manager without modifying `DiffusionLoRAManager`. + +Both runtime adapters run with distributed layerwise offload in request-mode +generation, where the manager keeps the LoRA A/B buffers resident on the compute +device while DLO streams the base blocks. MiniMax-H3 step execution still +rejects DLO, so `--step-execution` cannot be combined with +`--enable-distributed-layerwise-offload`. Model-level CPU offload and standard +layerwise offload are rejected in every mode, because the dynamic LoRA tensors +are neither parameters nor registered buffers and therefore do not participate +in those weight lifecycles. ## See Also diff --git a/docs/user_guide/diffusion/offloader/cpu_offload.md b/docs/user_guide/diffusion/offloader/cpu_offload.md new file mode 100644 index 000000000..558985e66 --- /dev/null +++ b/docs/user_guide/diffusion/offloader/cpu_offload.md @@ -0,0 +1,81 @@ +# Model-Level Offloading + +Model-level, or sequential, offloading keeps only the pipeline component group +currently executing on the accelerator. It is the simplest offload strategy +and is selected with `--enable-cpu-offload`. + +## How it works + +Pre-forward hooks enforce mutual exclusion between DiT and encoder modules: + +- before an encoder runs, the DiT moves to CPU; +- before a DiT runs, encoders and other DiTs move to CPU; and +- VAE modules remain on the accelerator. + +Pinned host memory reduces transfer overhead. Transfers occur at phase +boundaries, so cold-start and encoder-to-denoiser transitions become slower. + +## Usage + +```python +from vllm_omni import Omni + +omni = Omni( + model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + enable_cpu_offload=True, +) +``` + +```bash +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --omni --enable-cpu-offload +``` + +## Model integration + +Pipelines should implement `SupportsComponentDiscovery`: + +```python +from typing import ClassVar + +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery + + +class MyPipeline(nn.Module, SupportsComponentDiscovery): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] +``` + +All entries may be dotted paths. DiT and encoder lists are both required for +mutual exclusion. VAE modules are pinned but not swapped; resident modules are +small modules that must stay on the accelerator for layerwise paths. + +## Split-model components + +Some models, such as Cosmos3, split one transformer into mutually exclusive +components that run in different phases. The pipeline exposes +`enable_omni_model_cpu_offload`, and the backend delegates to the model-local +contexts: + +```python +class Cosmos3VFMTransformer(nn.Module): + def forward(self, ...): + with self._offload_context("reasoner"): + ... + with self._offload_context("generator"): + ... +``` + +This preserves the same invariant—exactly one component is device resident— +while reusing sequential `.to()` movers. + +## Limitations + +- Single device only. +- Higher cold-start latency. +- Transfers between encoder and denoising phases add latency. + +See the [model-level design](../../../design/feature/offloader/cpu_offload.md) +for lifecycle and extension invariants. diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md new file mode 100644 index 000000000..45ff8df2d --- /dev/null +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -0,0 +1,285 @@ +# Distributed Layerwise Offloading + +Distributed layerwise offloading (DLO) extends block streaming to multi-device +deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of +the host weights and reconstructs each layer at runtime. Without AllGather, +each rank streams a complete block independently. Compatible TP1 deployments +can share checkpoint-backed host pages among processes on the same node; +otherwise DLO streams the ordinary loader's rank-local tensors. + +See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md) +for the implementation contract and compatibility matrix. + +## Execution model + +DLO overlaps three operations with a fixed two-block device buffer: + +```text +Compute stream: [Layer N] [Layer N+1] [Layer N+2] +H2D stream: [H2D shard N+1] [H2D shard N+2] +AllGather: [AG N+1] [AG N+2] +Slots: slot 0: Layer N slot 1: Layer N+1 +``` + +AllGather communicates only request-independent weight shards, so data- +parallel ranks may process different requests concurrently. + +## Usage + +```bash +# Four ranks with sharded host weights and AllGather +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 + +# Standard-loader rank-local weights, without DLO AllGather +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 \ + --dlo-no-use-allgather + +# Sequence parallel deployment +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --usp 4 +``` + +```python +from vllm_omni import Omni + +omni = Omni( + model="/path/to/model", + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, +) +``` + +## Flags + +| Flag | Meaning | Default | +| --- | --- | --- | +| `--enable-distributed-layerwise-offload` | Enable DLO | `false` | +| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` | +| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` | +| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` | +| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | +| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` | +| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset | +| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` | + +## Host-weight loading + +The diffusion loader chooses host storage before DLO is enabled. It first +attempts to build a complete, validated direct-checkpoint mmap plan. If names, +coverage, shape, dtype, topology, or loader-callback compatibility cannot be +proven, it runs the ordinary model loader instead. DLO consumes that result and +does not make a second checkpoint-compatibility decision. + +The shared-mmap optimization in this phase is supported only with TP1. TP +greater than one falls back before model mutation to ordinary TP-aware loading. +DLO may still consume those TP-local tensors, but this is a compatibility path: +it does not share checkpoint-backed runtime weights across DP replicas and +provides no shared-mmap host-memory guarantee. + +The mmap plan skips only dedicated DiT weight sources. Other component sources, +such as a text encoder loaded through the shared diffusion loader, continue to +use their ordinary component loader. A checkpoint source that mixes DiT and +non-DiT weights falls back completely rather than leaving an unplanned +component uninitialized. + +With direct checkpoint mmap, the loader: + +1. saves non-persistent buffers such as RoPE frequencies; +2. moves the normally created transformer to the meta device; +3. loads checkpoint tensors as mmap views backed by the shared OS page cache; +4. applies any loader-owned bounded layout adapters while packing blocks; +5. restores saved non-persistent buffers; and +6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle + hooks. + +For AllGather with a group larger than one, each process copies only its +persistent shard and then releases the source mapping. For no-AllGather, each +process keeps the mapping open and packs complete blocks through two bounded +pinned staging slots. Processes mapping the same files on one node share the +immutable pages; no-AllGather still performs a complete-block H2D copy in each +process. + +When the effective DLO group size is one, `dlo_use_allgather=True` does not +perform a collective and uses the same rank-local transfer behavior. + +### Final-layout Host Weight Runtime + +HWR is an opt-in startup optimization for models that declare the final-layout +BF16 restore contract. Use it only with no-AllGather DLO: + +The validated BF16 model contracts currently cover MiniMax H3 and +`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model +class but has not been validated against this contract. FLUX.2-dev, online +FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside +the validated scope. + +```bash +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --host-weight-runtime-mode preferred \ + --host-weight-runtime-root /var/cache/vllm-omni/hwr +``` + +#### Choosing a mode + +| Mode | Exact local artifact hit | Miss or recoverable artifact/store problem | Intended use | +| --- | --- | --- | --- | +| `disabled` | HWR is not consulted | Use the existing checkpoint-mmap or ordinary-loader path | Default compatibility path | +| `preferred` | Restore the final-layout artifact | Load canonically, serve with those tensors, and attempt to publish an artifact for the next startup | Normal deployment and store population | +| `required` | Restore the final-layout artifact | Fail startup without canonical DiT fallback or publication | Enforce a pre-populated store in controlled rollouts or CI | + +Both enabled modes still fail on non-retryable configuration, identity, or +compatibility errors. `preferred` is a fallback policy for a cache miss or a +recoverable cache problem; it does not hide an invalid deployment. + +#### Populating a store for `required` mode + +`required` is deliberately consume-only. PR2 does not include a separate +prewarm command, so populate each node-local storage domain with one matching +`preferred` producer cohort: + +1. Choose a persistent, writable root visible to every diffusion worker in the + storage domain. Do not use a process-private temporary directory. +2. Start the deployment with `preferred`, using the exact model revision, + dtype, TP size, SP configuration, and other weight-layout settings intended + for serving. No inference request is needed; publication happens during + model startup. Wait for the engine to become healthy, then shut it down + cleanly. +3. Restart with the same arguments and root, changing only the mode to + `required`. A successful startup proves that every worker acquired a valid + artifact; a miss, corrupt artifact, or incompatible identity fails startup. +4. Repeat the population start on every node or storage domain because this + store is node-local. + +For example: + +```bash +# First startup: canonically load and populate the exact artifacts. +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --host-weight-runtime-mode preferred \ + --host-weight-runtime-root /var/cache/vllm-omni/hwr + +# After a healthy startup and clean shutdown, enforce cache hits. +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --host-weight-runtime-mode required \ + --host-weight-runtime-root /var/cache/vllm-omni/hwr +``` + +Include the same TP and SP layout flags in both commands. DP rank and DP size +are excluded from artifact identity, so the population and serving DP sizes may +differ and equivalent DP replicas share artifacts. TP rank is included, so a +TP-N deployment normally needs N rank-specific artifacts in each storage +domain; launching the matching TP cohort creates that set. If the model +revision or layout changes, run `preferred` again for the new identity before +returning to `required`. + +On a cold start, the canonical loader remains authoritative and publishes a +validated final-layout artifact for later workers. A warm start restores the +DiT final tensors without ordinary DiT materialization. DLO then attempts to +register the immutable mapping for direct H2D; if registration is unavailable, +it streams through the same two bounded host staging slots. The artifact +identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards, +and distinct SP layouts do not alias one another. Publication failure remains +separate from a `preferred` serving startup; the next `required` startup +provides the explicit artifact-availability check. + +For local canonical checkpoints, the first eligible worker may hash source +shards to establish immutable identity. HWR caches those digests in the same +node-local storage domain and validates file metadata before reuse, so later +workers normally avoid repeating that read. Cold BF16 publication also hashes +ordered payloads as they are written and overlaps payload durability work with +later shards; this changes startup work only, not artifact contents or runtime +H2D behavior. + +#### Registered direct H2D + +Registration is attempted automatically only for an eligible warm HWR hit when +`pin_cpu_memory` is enabled. A successful path is: + +```text +shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel +``` + +It removes the recurrent mmap-to-private-staging CPU copy and does not allocate +the two host staging slots. It does not reduce the H2D payload or make GPU +kernels access host memory directly. + +CUDA requires read-only host-registration capability for the immutable HWR +mapping. Unsupported capability, a positive registration limit smaller than +the complete page-aligned mapping, or a safely rolled-back registration error +falls back to bounded staging. Programmatic configurations can set +`pin_cpu_memory=False` to disable registration explicitly. A successful +registration locks the complete mapped range in host memory for that worker's +lifetime, so use +`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is +required. + +Shutdown drains H2D work, releases hook/source references, unregisters every +range, and only then closes the HWR lease. AllGather and direct checkpoint mmap +retain their existing transfer paths. + +When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does +not resolve HWR sources or construct its store. + +## Declarative topology + +Models may declare an `OffloadPlan` instead of embedding offload logic: + +```python +from vllm_omni.diffusion.offloader import OffloadPlan + + +class MyPipeline(nn.Module): + _dit_modules = ["transformer"] + _offload_plan = OffloadPlan( + block_attrs={"transformer": ("blocks",)}, + offload_submodules={"context_encoder": "layers"}, + ) +``` + +When no plan exists, discovery falls back to +`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup. + +## Data-parallel concurrency + +With `data_parallel_size > 1` and AllGather enabled, the scheduler can process +up to `dp_size` requests per denoising step. Every concurrent request must set +the same explicit `num_inference_steps`; `None` is rejected because every rank +must enter each collective. + +## Limitations + +- Direct checkpoint mmap currently requires TP1. TP greater than one is + outside the Phase A shared-mmap support scope and falls back before model + mutation to the ordinary TP-aware loader. DLO can stream that runtime layout, + and eligible no-AllGather configurations may use HWR final-layout artifacts, + but direct checkpoint mmap provides no shared-mmap host-memory guarantee for + TP greater than one. +- HSDP plus AllGather is rejected to avoid double sharding. HSDP without + AllGather has limited end-to-end validation. +- Per-tensor online FP8 linears use the ordinary loader and can run with either + DLO transfer path. With AllGather, every rank temporarily materializes the + complete FP8 model in host memory before DLO retains only its shard. Other + online quantization methods require no-AllGather until their runtime layouts + are validated. +- Resident leading layers require `--dlo-no-use-allgather` and a model + `OffloadPlan` that declares eligible `resident_dit_paths`. +- DP concurrency requires an explicit, identical inference-step count. + +Sharing quantized or otherwise unvalidated transformed layouts through a +normalized HWR producer is a follow-up design in +[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of +the current BF16 final-layout path. + +See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) +for an end-to-end example. diff --git a/docs/user_guide/diffusion/offloader/layerwise_offload.md b/docs/user_guide/diffusion/offloader/layerwise_offload.md new file mode 100644 index 000000000..c88b67fcc --- /dev/null +++ b/docs/user_guide/diffusion/offloader/layerwise_offload.md @@ -0,0 +1,65 @@ +# Layerwise Offloading + +Layerwise, or blockwise, offloading keeps one transformer block on the +accelerator and prefetches the next block while the current block computes. +It is best suited to compute-heavy video DiTs whose block execution can hide +host-to-device transfers. + +## Execution flow + +Each block has a pre-forward and post-forward hook. Parameters are consolidated +in pinned host tensors and rematerialized for execution on a dedicated copy +stream. + +| Block | Pre-forward hook | Forward | Post-forward hook | +| --- | --- | --- | --- | +| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 | +| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 | +| ... | ... | ... | ... | +| last block | Prefetch block 0 | Compute last block | Free last block | + +Encoders, VAE modules, and non-block DiT modules such as embeddings and norms +remain device resident. + +## Usage + +```python +from vllm_omni import Omni + +omni = Omni( + model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + enable_layerwise_offload=True, +) +``` + +```bash +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --omni --enable-layerwise-offload +``` + +## Model integration + +Transformer classes declare containers of executable blocks: + +```python +class WanTransformer3DModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + +class Flux2Transformer2DModel(nn.Module): + _layerwise_offload_blocks_attrs = [ + "transformer_blocks", + "single_transformer_blocks", + ] +``` + +See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for +the discovery and hook invariants. `OffloadPlan` is a separate declarative +topology path for distributed layerwise offload. + +## Limitations + +- Single device only. +- Setup consolidates and pins block parameters, increasing cold-start time. +- Performance depends on block compute time and host-to-device bandwidth; + lightweight blocks may not hide transfers. diff --git a/docs/user_guide/diffusion/parallelism/hsdp.md b/docs/user_guide/diffusion/parallelism/hsdp.md index 96a357c86..9491ed9f5 100644 --- a/docs/user_guide/diffusion/parallelism/hsdp.md +++ b/docs/user_guide/diffusion/parallelism/hsdp.md @@ -64,6 +64,41 @@ omni = Omni( ) ``` +### Combined with Online FP8 Quantization + +HSDP can be combined with online FP8 quantization for diffusion models that +support both features. HSDP shards the quantized weights across GPUs to reduce +the per-GPU weight memory footprint. + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="", + quantization="fp8", + parallel_config=DiffusionParallelConfig( + use_hsdp=True, + hsdp_shard_size=2, + ), +) +``` + +For online serving: + +```bash +vllm serve \ + --omni \ + --quantization fp8 \ + --use-hsdp \ + --hsdp-shard-size 2 +``` + +The model must support both HSDP and online FP8 quantization. HSDP remains +incompatible with Tensor Parallelism. See the +[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and +hardware. + --- ## Example Script @@ -117,7 +152,8 @@ In `DiffusionParallelConfig`: **Constraints:** - `hsdp_replicate_size × hsdp_shard_size == world_size` -- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1) +- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism, + Pipeline Parallelism, or Expert Parallelism. --- @@ -147,3 +183,4 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co 1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models 2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism 3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled +4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint diff --git a/docs/user_guide/diffusion/regional_compilation.md b/docs/user_guide/diffusion/regional_compilation.md index 469a62b2c..448f19518 100644 --- a/docs/user_guide/diffusion/regional_compilation.md +++ b/docs/user_guide/diffusion/regional_compilation.md @@ -24,7 +24,7 @@ stages: For an experimental whole-transformer compile scope, set `--diffusion-compile-granularity full` or use -`diffusion_compile_granularity: full` in the stage configuration. Full scope may +`diffusion_compile_granularity: full` in the deploy configuration. Full scope may still contain graph breaks; it does not force one graph. It is rejected when HSDP, sequence parallelism, CPU offload, or layerwise offload is enabled. Use regional scope with those features. diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 7c664e88b..c0b29cf6e 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -106,7 +106,10 @@ The following tables show which models support each feature: > Notes: -> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md). +> 1. CPU Offload has three strategies: model-level, layerwise, and distributed +> layerwise. The tables below show **layerwise support** only. Split models +> like Cosmos3 swap their reasoner/generator components for model-level +> offload; see the [CPU Offload Guide](diffusion/cpu_offload.md). > 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details. ### ImageGen @@ -121,11 +124,11 @@ The following tables show which models support each feature: | **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | | **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **HiDream-O1-Image** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅* | | **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | @@ -144,7 +147,7 @@ The following tables show which models support each feature: > Notes: > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. -> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary. +> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary. > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload. > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). @@ -158,20 +161,29 @@ The following tables show which models support each feature: | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ | +| **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | -| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ | +| **SANA-WM** | ❌ | ❌ | ❌5 | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | + +> Notes: +> 5. SANA-WM cannot support sequence parallelism: its bidirectional gated delta +> recurrence carries state across frames, so a rank cannot denoise a slice of +> the token sequence in isolation. Doing so would need a distributed scan or +> an all-gather before every GDN block. The remaining ❌ columns are simply +> unvalidated on this model, not known-broken. > **Step execution note:** Helios supports single-request step execution only; > use `max_num_seqs=1`. **Frame Interpolation Support** -- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines -- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni +- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines; SANA-Video-2B native text-to-video and image-to-video pipelines +- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM; SANA-Video-2B Diffusers-adapter text-to-video and image-to-video pipelines ### AudioGen @@ -214,7 +226,7 @@ The following tables show which models support each feature: 3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible. 4. The CPU Offloading (Layerwise) row describes local layerwise offload. Multi-device Distributed Layerwise Offload has a separate topology and - compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading). + compatibility matrix in the [Distributed Layerwise Offloading guide](diffusion/offloader/distributed_layerwise_offload.md). 5. The compatibility matrix uses FP8 as the representative quantization method. 6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches). diff --git a/docs/user_guide/examples/offline_inference/glm_image.md b/docs/user_guide/examples/offline_inference/glm_image.md index 09d927ff7..267819cc5 100644 --- a/docs/user_guide/examples/offline_inference/glm_image.md +++ b/docs/user_guide/examples/offline_inference/glm_image.md @@ -43,7 +43,7 @@ if __name__ == "__main__": ), ], ) - outputs[0].request_output.images[0].save("output.png") + outputs[0].images[0].save("output.png") ``` ## Image-to-Image (Image Editing) @@ -72,7 +72,7 @@ if __name__ == "__main__": ), ], ) - outputs[0].request_output.images[0].save("output.png") + outputs[0].images[0].save("output.png") ``` ## Generation Parameters diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md index db7a31ce5..ff57063b8 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -4,8 +4,8 @@ Source **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md index f631cd962..24cab3a3c 100644 --- a/docs/user_guide/examples/offline_inference/text_to_image.md +++ b/docs/user_guide/examples/offline_inference/text_to_image.md @@ -55,7 +55,7 @@ if __name__ == "__main__": omni = Omni(model="Qwen/Qwen-Image") prompt = "a cup of coffee on the table" outputs = omni.generate(prompt) - images = outputs[0].request_output.images + images = outputs[0].images images[0].save("coffee.png") ``` @@ -179,7 +179,7 @@ if __name__ == "__main__": ] outputs = omni.generate(prompts) for i, output in enumerate(outputs): - output.request_output.images[0].save(f"{i}.jpg") + output.images[0].save(f"{i}.jpg") ``` !!! info @@ -207,7 +207,7 @@ if __name__ == "__main__": } ]) for i, output in enumerate(outputs): - output.request_output.images[0].save(f"{i}.jpg") + output.images[0].save(f"{i}.jpg") ``` You can also pass a negative prompt via the CLI argument `--negative-prompt`: diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index ad7777014..a6544c1e5 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -69,7 +69,7 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \ ### Notes - Output: 48 kHz mono WAV. -- Stage config: `vllm_omni/model_executor/stage_configs/voxcpm2.yaml` (default). +- Deploy config: `vllm_omni/deploy/voxcpm2.yaml` (default). --- @@ -152,7 +152,7 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a streaming test." \ --streaming ``` -Streaming requires `async_chunk: true` in the stage config. +Streaming requires `async_chunk: true` in the deploy config. ### Notes - Output: 44.1 kHz mono WAV. @@ -274,7 +274,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --streaming \ --output-dir /tmp/out_stream ``` -Streaming requires `async_chunk: true` in the stage config. +Streaming requires `async_chunk: true` in the deploy config. ### Batched decoding The Code2Wav stage supports batched decoding through the SpeechTokenizer. Configure both stages with `max_num_seqs > 1` via `--stage-overrides` and pass multiple prompts via `--txt-prompts`: diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index 1a2082f82..ec61bc749 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -16,6 +16,8 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | +| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE | +| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | ~23.6 GiB | ## Local CLI Usage @@ -122,6 +124,31 @@ python text_to_video.py \ --output quick_test.mp4 ``` +### SANA-Video-2B + +```bash +python text_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaVideoPipeline \ + --prompt "A cinematic tracking shot of a sailboat crossing the ocean at sunset." \ + --height 480 \ + --width 832 \ + --num-frames 81 \ + --num-inference-steps 50 \ + --guidance-scale 6.0 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 \ + --output sana_video_480p.mp4 +``` + +For the 720p checkpoint, switch the model to +`Efficient-Large-Model/SANA-Video_2B_720p_diffusers`, use +`--height 704 --width 1280`, and write to a different output path. SANA-Video +accepts `motion_score`, `clean_caption`, and `use_resolution_binning` through +`--extra-body`. For image-to-video, use the shared +`examples/offline_inference/image_to_video/image_to_video.py` example with +`--model-class-name SanaImageToVideoPipeline` and `--image `. + ## Key Arguments ### Common diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md deleted file mode 100644 index d5cd0ea46..000000000 --- a/docs/user_guide/examples/offline_inference/x_to_video_audio.md +++ /dev/null @@ -1,115 +0,0 @@ -# X-To-Video-Audio - -Source . - - -The `DreamID-Omni` pipeline generates short videos from text, image and video. - -## Local CLI Usage -### Download the Model locally -Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally. - -```bash -python download_dreamid_omni.py --output-dir ./dreamid_omni -``` -After download, the model directory will look like this: - -``` -dreamid_omni/ -├── DreamID-Omni/ -│ ├── dreamid_omni.safetensors -├── MMAudio/ -│ ├── ext_weights/ -│ │ ├── best_netG.pt -│ │ ├── v1-16.pth -├── Wan2.2-TI2V-5B/ -│ ├── google/* -│ ├── models_t5_umt5-xxl-enc-bf16.pth -│ ├── Wan2.2_VAE.pth -│ -├── model_index.json # create by download_dreamid_omni.py -``` - -### Run the Inference -```python -python x_to_video_audio.py \ - --model /path/to/dreamid_omni \ - --prompt "Two people walking together and singing happily" \ - --image-path ./example0.png ./example1.png \ - --audio-path ./example0.wav ./example1.wav \ - --video-negative-prompt "jitter, bad hands, blur, distortion" \ - --audio-negative-prompt "robotic, muffled, echo, distorted" \ - --cfg-parallel-size 2 \ - --num-inference-steps 45 \ - --height 704 \ - --width 1280 \ - --output out_dreamid_omni_twoip.mp4 -``` -In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled. -The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload. -For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP: - -```python -python x_to_video_audio.py \ - --model /path/to/dreamid_omni \ - --prompt "Two people walking together and singing happily" \ - --image-path ./example0.png ./example1.png \ - --audio-path ./example0.wav ./example1.wav \ - --video-negative-prompt "jitter, bad hands, blur, distortion" \ - --audio-negative-prompt "robotic, muffled, echo, distorted" \ - --cfg-parallel-size 2 \ - --num-inference-steps 45 \ - --height 704 \ - --width 1280 \ - --use-hsdp \ - --hsdp-shard-size 2 \ - --output out_dreamid_omni_twoip.mp4 -``` - - -You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni - -For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing. - -```python -# Example usage for oneip, ref media from the official repo DreamID-Omni -python x_to_video_audio.py \ - --model /path/to/dreamid_omni \ - --prompt ": In the frame, a woman with black long hair is identified as .\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ - --image-path 9.png \ - --audio-path 9.wav \ - --video-negative-prompt "jitter, bad hands, blur, distortion" \ - --audio-negative-prompt "robotic, muffled, echo, distorted" \ - --cfg-parallel-size 2 \ - --num-inference-steps 45 \ - --height 704 \ - --width 1280 \ - --output out_dreamid_omni_oneip.mp4 -``` - - -Key arguments: -- `--prompt`: text description (string). -- `--model`: path to the model local directory. -- `--height/--width`: output resolution (defaults 704 * 1024). -- `--image-path`: path to the input image list. -- `--audio-path`: path to the input audio list, indicate the timbre of the output video. -- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1). -- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks. -- `--hsdp-shard-size`: number of GPUs used for HSDP sharding. -- `--hsdp-replicate-size`: number of HSDP replica groups. -- `--num-inference-steps`: number of denoising steps (defaults 45). -- `--video-negative-prompt`: negative prompt for video generation. -- `--audio-negative-prompt`: negative prompt for audio generation. -- `--enable-cpu-offload`: enable CPU offload (defaults False). - -## Example materials - -??? abstract "download_dreamid_omni.py" - ``````py - --8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py" - `````` -??? abstract "x_to_video_audio.py" - ``````py - --8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py" - `````` diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index b4c05bdbe..f92463f84 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -28,6 +28,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --deploy-config /path/to/deploy_config_file ``` +To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct +weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key +`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ + --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml +``` + +Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select +the same single-stage pipeline without `--deploy-config`. + ### Launch individual stages (stage-based CLI) Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage. @@ -271,7 +283,7 @@ The script supports the following arguments: vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` -If you have custom stage configs file: +If you have a custom deploy config file: ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file ``` diff --git a/docs/user_guide/examples/online_serving/text_to_audio.md b/docs/user_guide/examples/online_serving/text_to_audio.md deleted file mode 100644 index 1b6dbc463..000000000 --- a/docs/user_guide/examples/online_serving/text_to_audio.md +++ /dev/null @@ -1,193 +0,0 @@ -# Text-To-Audio - -Source . - -This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni. - -## Supported Models - -| Model | Description | -|-------|-------------| -| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo | - -## Start Server - -### Basic Start - -```bash -vllm serve stabilityai/stable-audio-open-1.0 \ - --host 0.0.0.0 \ - --port 8091 \ - --gpu-memory-utilization 0.9 \ - --trust-remote-code \ - --enforce-eager \ - --omni -``` - -## API Calls - -### Method 1: Using curl - -```bash -# Run all curl examples -bash curl_examples.sh - -# Or execute directly -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "The sound of a cat purring", - "audio_length": 10.0 - }' --output cat.wav -``` - -### Method 2: Using Python Client - -```bash -cd examples/online_serving/stable_audio - -# Simple generation -python stable_audio_client.py \ - --text "The sound of a cat purring" - -# With custom duration -python stable_audio_client.py \ - --text "A dog barking" \ - --audio_length 5.0 - -# With all parameters -python stable_audio_client.py \ - --text "Thunder and rain" \ - --audio_length 15.0 \ - --negative_prompt "Low quality" \ - --guidance_scale 7.0 \ - --num_inference_steps 100 \ - --seed 42 \ - --output thunder.wav -``` - -The Python client supports the following command-line arguments: - -- `--api_url`: API endpoint URL (default: `http://localhost:8091/v1/audio/generate`) -- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`) -- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`) -- `--audio_start`: Audio start time in seconds (default: `0.0`) -- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`) -- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`) -- `--num_inference_steps`: Number of inference steps (default: `100`) -- `--seed`: Random seed for reproducibility (default: `None`) -- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm` -- `--output`: Output file path (default: `stable_audio_output.wav`) - -### Method 3: Using Python httpx - -```python -import httpx - -response = httpx.post( - "http://localhost:8091/v1/audio/generate", - json={ - "input": "The sound of ocean waves crashing on a beach", - "audio_length": 10.0, - "negative_prompt": "Low quality, distorted", - "guidance_scale": 7.0, - "num_inference_steps": 100, - }, - timeout=300.0, -) - -with open("ocean.wav", "wb") as f: - f.write(response.content) -``` - -## Request Format - -### Simple Generation - -```json -{ - "input": "The sound of ocean waves" -} -``` - -### Generation with Parameters - -```json -{ - "input": "A piano playing a gentle melody", - "audio_length": 10.0, - "negative_prompt": "Low quality, distorted, noisy", - "guidance_scale": 8.0, - "num_inference_steps": 150, - "seed": 42, - "response_format": "wav" -} -``` - -## API Reference - -### Endpoint - -``` -POST /v1/audio/generate -Content-Type: application/json -``` - -### Generation Parameters - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `input` | string | **required** | Text prompt describing the audio to generate | -| `model` | string | server's model | Model to use (optional, should match server if specified) | -| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus | -| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) | -| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) | -| `audio_start` | float | 0.0 | Audio start time in seconds | -| `negative_prompt` | string | null | Text describing what to avoid in generation | -| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) | -| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) | -| `seed` | int | null | Random seed for reproducible generation | - -### Response Format - -Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`). - -## Tuning Tips - -1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`. -2. **Quality vs Speed**: - - 50 steps: Fast, decent quality (quick previews) - - 100 steps: Good balance (general purpose) - - 150+ steps: High quality, slower (final / critical audio) -3. **Guidance Scale**: - - Lower (3 - 5): More creative / varied output - - Default (7): Good balance - - Higher (10+): Strict adherence to the prompt -4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`. -5. **Seeds**: Set a fixed seed to get deterministic, reproducible results. - -## File Description - -| File | Description | -|------|-------------| -| `curl_examples.sh` | Curl examples covering common use cases | -| `stable_audio_client.py` | Python client with full CLI argument support | - -## Troubleshooting - -1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors. -2. **Connection refused**: Make sure the server is running on the correct port. -3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`. -4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`. -5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`. - -## Example materials - -??? abstract "stable_audio_client.py" - ``````py - --8<-- "examples/online_serving/stable_audio/stable_audio_client.py" - `````` -??? abstract "curl_examples.sh" - ``````sh - --8<-- "examples/online_serving/stable_audio/curl_examples.sh" - `````` diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index c6a613723..5d4ccbf3b 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -419,7 +419,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "response_format": "pcm" }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` -Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. +Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. ### Streaming WebSocket The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: diff --git a/docs/user_guide/fault_injection_reliability_matrix.md b/docs/user_guide/fault_injection_reliability_matrix.md index 28271affa..4a2fa0a16 100644 --- a/docs/user_guide/fault_injection_reliability_matrix.md +++ b/docs/user_guide/fault_injection_reliability_matrix.md @@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b ## Source of Conclusions -The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`. +The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, and **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`. ### Example: `SIGTERM` fault injection log (Qwen3-Omni) diff --git a/docs/user_guide/feature_compatibility.md b/docs/user_guide/feature_compatibility.md index ea05fe81e..eeb087510 100644 --- a/docs/user_guide/feature_compatibility.md +++ b/docs/user_guide/feature_compatibility.md @@ -165,7 +165,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \ - **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771). -- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored. +- **CPU Offloading — Strategies Are Mutually Exclusive**: Distributed layerwise, layerwise, and model-level offload should not be enabled together. Selection priority is distributed layerwise, then layerwise, then model-level. - **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism. diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index ddca3ea19..010ef4ba7 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -80,20 +80,22 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) ## Model Type Support -### Diffusion Model (Qwen-Image, Wan2.2) +### Diffusion Models | Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization | |-------|-----------|:-------:|:------:|----------------|------------------|------------------| | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | | | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | | +| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | +| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ | | FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | | | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ | | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | | -| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | | +| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan deploy config for multi-stage runs | None | | | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | | | Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | | -| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | | +| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) @@ -136,6 +138,39 @@ outputs = omni.generate( ) ``` +### Global and per-component scope + +A plain, global quantization configuration is passed unchanged to every +quantization-aware component constructed by a pipeline. This includes an +eligible encoder when that encoder is implemented with vLLM quantizable +layers; it does not rewrite arbitrary `torch.nn` modules. A structured +component map is the only way to narrow that scope. Pipeline integrations that +do not yet expose an encoder through the quantization factory remain DiT-only. + +For a pipeline that exposes both a transformer and a quantization-aware text +encoder, the scope is: + +| Configuration | Transformer | Text encoder | Components without supported quantizable layers | +|---------------|-------------|--------------|-------------------------------------------------| +| `quantization="fp8"` | FP8 | FP8 | checkpoint precision | +| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision | +| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision | + +Use `quantization_config` for component-selective Python configuration. For +example, quantize only the text decoder: + +```python +omni = Omni( + model="", + quantization_config={"text_encoder": {"method": "fp8"}}, +) +``` + +Component keys are runtime prefixes exposed by the pipeline integration; common +keys include `transformer` and `text_encoder`. Entries may be combined, and +`ignored_layers` can keep named eligible layers in checkpoint precision. Check +the model recipe for supported components and their runtime prefixes. + CLI: ```bash @@ -144,6 +179,24 @@ python text_to_image.py --model --quantization fp8 --ignored-layers vllm serve --omni --quantization fp8 ``` +### Distributed layerwise offload + +Per-tensor online FP8 linears can be combined with DLO's default AllGather path: + +```bash +vllm serve \ + --omni \ + --quantization fp8 \ + --enable-distributed-layerwise-offload +``` + +Direct checkpoint mmap is not used for this combination. Each rank runs the +ordinary loader to create the final FP8 weights and scales before DLO shards +them, so startup temporarily materializes the complete FP8 model in host memory +on every rank. Runtime DLO host residency is sharded normally. Use +`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract +does not fit the workload. + ## Parameters | Parameter | Type | Default | Description | diff --git a/docs/user_guide/quantization/msmodelslim.md b/docs/user_guide/quantization/msmodelslim.md index c03628b9d..ea42f6250 100644 --- a/docs/user_guide/quantization/msmodelslim.md +++ b/docs/user_guide/quantization/msmodelslim.md @@ -92,4 +92,4 @@ For HunyuanImage-3.0, use the Hunyuan-specific adaptation linked above. 2. The `ascend` quantization method expects weights produced by the Ascend tooling; it is not a load-time CUDA quantizer. 3. Keep the quantized checkpoint aligned with the same model architecture and - stage config used for BF16 inference. + pipeline and deploy config used for BF16 inference. diff --git a/docs/user_guide/quantization/quantized_kvcache.md b/docs/user_guide/quantization/quantized_kvcache.md index e560032d3..99a90e08a 100644 --- a/docs/user_guide/quantization/quantized_kvcache.md +++ b/docs/user_guide/quantization/quantized_kvcache.md @@ -83,23 +83,23 @@ Online serving: vllm serve --omni --diffusion-kv-cache-dtype fp8 ``` -Stage config: +Deploy config: ```yaml -stage_args: +stages: - stage_id: 0 - stage_type: diffusion - engine_args: - model_stage: dit - diffusion_kv_cache_dtype: "fp8" - diffusion_kv_cache_skip_steps: "0,1" - diffusion_kv_cache_skip_layers: "0-2" + diffusion_kv_cache_dtype: "fp8" + diffusion_kv_cache_skip_steps: "0,1" + diffusion_kv_cache_skip_layers: "0-2" ``` -Legacy YAML keys `kv_cache_dtype`, `kv_cache_skip_steps`, and -`kv_cache_skip_layers` are still accepted when constructing -`OmniDiffusionConfig` (for example via `from_kwargs`); prefer the `diffusion_*` -names for new configs. +The `model_stage` and diffusion execution type belong to the registered +`PipelineConfig`; the deploy YAML only carries runtime overrides. + +The legacy keyword aliases `kv_cache_dtype`, `kv_cache_skip_steps`, and +`kv_cache_skip_layers` remain accepted when constructing +`OmniDiffusionConfig` directly. They are not deploy YAML fields; prefer the +`diffusion_*` names for new code. ## Parameters diff --git a/docs/user_guide/quantization/svdquant.md b/docs/user_guide/quantization/svdquant.md new file mode 100644 index 000000000..4fbd8040d --- /dev/null +++ b/docs/user_guide/quantization/svdquant.md @@ -0,0 +1,50 @@ +# SVDQuant W4A4 + +## Overview + +[SVDQuant](https://arxiv.org/abs/2411.05007) combines four-bit weights +and activations with a small low-rank branch that corrects part of the +quantization error. vLLM-Omni consumes an offline-quantized checkpoint; +it does not calibrate the model while loading. + +## Checkpoint contract + +Place the following entry in the diffusion transformer's `config.json`: + +```json +{ + "quantization_config": { + "quant_method": "svdquant", + "rank": 32, + "precision": "nvfp4", + "act_unsigned": false, + "modules_to_not_convert": [] + } +} +``` + +For a quantized linear with input size `K`, output size `N`, and correction +rank `R`, the checkpoint stores: + +| Suffix | Shape | dtype | +| --- | --- | --- | +| `qweight` | `(N, K / 2)` | `int8` (two packed FP4 values per byte) | +| `wscales` | `(K / 16, N)` | `float8_e4m3fn` | +| `proj_down` | `(K, R)` | `bfloat16` | +| `proj_up` | `(N, R)` | `bfloat16` | +| `smooth_factor` | `(K,)` | `bfloat16` | +| `wcscales` | `(N,)` | `bfloat16` | +| `wtscale` | `(1,)` | `bfloat16` | + +`K` must be divisible by 16 on every tensor-parallel rank. Set `wcscales` to +ones when no per-output correction is needed. Modules listed in +`modules_to_not_convert` keep their checkpoint precision. + +## Runtime support + +The compatibility path accepts BF16 inputs and executes an NVFP4 GEMM followed +by the BF16 rank correction. It supports vLLM's FlashInfer, CUTLASS, and FBGEMM +NVFP4 tensor layouts; incompatible forced backends fail during model loading. +SM103 is the currently validated and enabled hardware target. Native fusion of +the NVFP4 GEMM and rank correction is separate from this checkpoint-loading +contract. diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 000000000..78a85387e --- /dev/null +++ b/examples/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""vLLM-Omni example scripts, importable as a package. + +Keep this ``__init__.py``: with the upstream vLLM installed editable (which also +places an ``examples`` package on ``sys.path``), a regular package at this repo +root wins over the upstream ``examples`` namespace during ``import examples``. +""" diff --git a/examples/offline_inference/audex/README.md b/examples/offline_inference/audex/README.md index 793b69eb9..4b570b06a 100644 --- a/examples/offline_inference/audex/README.md +++ b/examples/offline_inference/audex/README.md @@ -17,7 +17,7 @@ Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond pair contract) lives in `common.py`; each task script keeps only its own prompt recipe and flow. Note `text_to_audio.py` here drives the Audex two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio -Open, AudioX) use the generic +Open) use the generic `examples/offline_inference/text_to_audio/text_to_audio.py` instead. ## text_to_speech.py — text → speech diff --git a/examples/offline_inference/audex/text_to_audio.py b/examples/offline_inference/audex/text_to_audio.py index c357819c3..e81a53bb3 100644 --- a/examples/offline_inference/audex/text_to_audio.py +++ b/examples/offline_inference/audex/text_to_audio.py @@ -7,7 +7,7 @@ thinker generates interleaved 4-codebook RVQ tokens under an RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz mono WAV. -(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use +(For DIFFUSION text-to-audio models — Stable Audio Open — use ``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex TTA is a two-stage autoregressive pipeline with a different interface.) diff --git a/examples/offline_inference/aura_omni/end2end.py b/examples/offline_inference/aura_omni/end2end.py index bac385d8e..29ccb3539 100644 --- a/examples/offline_inference/aura_omni/end2end.py +++ b/examples/offline_inference/aura_omni/end2end.py @@ -24,7 +24,6 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays from vllm.multimodal.media.audio import load_audio from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.omni import Omni from vllm_omni.model_executor.stage_input_processors.aura_omni import ( DEFAULT_QWEN3_TTS_REF_TEXT, @@ -228,7 +227,6 @@ def parse_args(): parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.") parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.") parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.") - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/covo_audio/README.md b/examples/offline_inference/covo_audio/README.md index 80a822fa2..083666ebd 100644 --- a/examples/offline_inference/covo_audio/README.md +++ b/examples/offline_inference/covo_audio/README.md @@ -2,7 +2,8 @@ ## Setup -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. +Use `--deploy-config` for deployment overrides such as stage memory allocation. +See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/). > **Note** > Covo-Audio code2wav requires `torchdiffeq`. Install it with: `pip install torchdiffeq` @@ -41,7 +42,7 @@ python end2end.py -m /path/to/Covo-Audio-Chat --output-dir ./my_output | `--sampling-rate` | | `16000` | Sampling rate for audio loading (Hz) | | `--output-dir` | | `./output_audio` | Output directory for generated files | | `--num-prompts` | | `1` | Number of prompts to generate | -| `--stage-configs-path` | | (auto) | Path to stage configs YAML file | +| `--deploy-config` | | (auto) | Optional path to a deploy YAML override | | `--log-stats` | | `false` | Enable detailed statistics logging | | `--stage-init-timeout` | | `300` | Stage initialization timeout (seconds) | | `--batch-timeout` | | `5` | Batching timeout (seconds) | diff --git a/examples/offline_inference/covo_audio/end2end.py b/examples/offline_inference/covo_audio/end2end.py index 5b7ac5fef..b5050f62c 100644 --- a/examples/offline_inference/covo_audio/end2end.py +++ b/examples/offline_inference/covo_audio/end2end.py @@ -59,7 +59,7 @@ def main(args): omni = Omni( model=args.model_name, - stage_configs_path=args.stage_configs_path, + deploy_config=args.deploy_config, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, batch_timeout=args.batch_timeout, @@ -167,10 +167,10 @@ def parse_args(): help="Sampling rate for audio loading (default: 16000).", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, default=None, - help="Path to stage configs YAML file.", + help="Path to a deploy config YAML file.", ) parser.add_argument( "--log-stats", diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py index 5f4311050..851d18e18 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -53,7 +53,6 @@ from typing import Any import torch from PIL import Image -from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -139,12 +138,6 @@ async def main(): # ---- Torch setup ---- generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, - ring_degree=args.ring_degree, - cfg_parallel_size=args.cfg_parallel_size, - tensor_parallel_size=args.tensor_parallel_size, - ) # ---- Cache Config ---- if args.cache_backend == "cache_dit": @@ -171,7 +164,10 @@ async def main(): vae_use_tiling=args.vae_use_tiling, cache_backend=args.cache_backend, cache_config=cache_config, - parallel_config=parallel_config, + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, enforce_eager=args.enforce_eager, enable_cpu_offload=args.enable_cpu_offload, diffusion_load_format="dummy", diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py new file mode 100644 index 000000000..b3a4ae0d1 --- /dev/null +++ b/examples/offline_inference/diffusion/lingbot_world_v2.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Generate a LingBot-World v2 video from an image and camera trajectory. + +The official checkpoint is licensed separately under CC BY-NC-SA and is +restricted to non-commercial use. This vLLM-Omni integration remains +Apache-2.0 licensed. + +Example: + python examples/offline_inference/diffusion/lingbot_world_v2.py \ + --prompt "The camera moves slowly forward through the scene." \ + --image /path/to/first_frame.png \ + --action-dir /path/to/actions/forward \ + --output lingbot_world_v2.mp4 + +``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The +example configures its parent as the trusted action root and sends the +contained directory through ``extra_args["action_path"]``. +""" + +from __future__ import annotations + +import argparse +import math +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np + +_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" +_NUM_INFERENCE_STEPS = 4 +_MAX_SEQUENCE_LENGTH = 512 +_MAX_PIXEL_AREA = 480 * 832 +_MAX_RAW_FRAMES = 117 +_MAX_SOURCE_CAMERA_FRAMES = 4096 +_TEMPORAL_COMPRESSION = 4 +_LATENT_FRAMES_PER_BLOCK = 3 +_SPATIAL_ALIGNMENT = 16 + + +@dataclass(frozen=True) +class LingBotPaths: + image: Path + action_dir: Path + action_root: Path + action_relative: Path + camera_frames: int + output: Path + + +def _validate_camera_arrays(action_dir: Path) -> int: + """Validate the bounded camera contract before starting the engine.""" + + try: + poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False) + intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False) + except (OSError, TypeError, ValueError): + raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None + if poses.ndim != 3 or poses.shape[1:] != (4, 4): + raise ValueError("poses.npy must have shape [frames, 4, 4].") + if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,): + raise ValueError("intrinsics.npy must have shape [frames, 4].") + if poses.shape[0] != intrinsics.shape[0]: + raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.") + if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES: + raise ValueError("Camera arrays must contain between 1 and 4096 source frames.") + if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu": + raise ValueError("Camera arrays must contain real numeric values.") + if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all(): + raise ValueError("Camera arrays must contain only finite values.") + return int(poses.shape[0]) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory." + ) + parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.") + parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.") + parser.add_argument("--image", required=True, help="Path to the first-frame image.") + parser.add_argument( + "--action-dir", + required=True, + help="Camera action directory containing poses.npy and intrinsics.npy.", + ) + parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.") + parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.") + parser.add_argument( + "--num-frames", + type=int, + default=9, + help="Raw output frame count. Nine frames are one three-latent-frame causal block.", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed.") + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Number of GPUs used for tensor parallelism inside the DiT.", + ) + parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.") + parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.") + parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.") + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable compilation and force eager execution.", + ) + return parser.parse_args(argv) + + +def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths: + """Canonicalize local inputs and construct a safe trusted-root mapping.""" + + image = Path(args.image).expanduser().resolve() + if not image.is_file(): + raise ValueError("--image must point to an existing file.") + + action_dir = Path(args.action_dir).expanduser().resolve() + if not action_dir.is_dir(): + raise ValueError("--action-dir must point to an existing directory.") + if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")): + raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.") + camera_frames = _validate_camera_arrays(action_dir) + + action_root = action_dir.parent.resolve() + try: + action_relative = action_dir.relative_to(action_root) + except ValueError: + raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None + if not action_relative.parts: + raise ValueError("--action-dir must name a directory contained by a trusted action root.") + + output = Path(args.output).expanduser().resolve() + return LingBotPaths( + image=image, + action_dir=action_dir, + action_root=action_root, + action_relative=action_relative, + camera_frames=camera_frames, + output=output, + ) + + +def _positive_finite(value: float, flag: str) -> float: + result = float(value) + if not math.isfinite(result) or result <= 0: + raise ValueError(f"{flag} must be a positive finite number.") + return result + + +def build_omni_kwargs( + args: argparse.Namespace, + paths: LingBotPaths, +) -> dict[str, Any]: + """Build Omni configuration without overriding checkpoint class discovery.""" + + if args.tensor_parallel_size <= 0: + raise ValueError("--tensor-parallel-size must be a positive integer.") + flow_shift = _positive_finite(args.flow_shift, "--flow-shift") + model_path = Path(args.model).expanduser() + model = str(model_path.resolve()) if model_path.exists() else args.model + return { + "model": model, + "flow_shift": flow_shift, + "tensor_parallel_size": args.tensor_parallel_size, + "enforce_eager": args.enforce_eager, + "model_config": {"lingbot_action_root": str(paths.action_root)}, + } + + +def build_request( + args: argparse.Namespace, + paths: LingBotPaths, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Build the prompt and fixed v1 causal-DMD sampling contract.""" + + prompt_text = args.prompt.strip() + if not prompt_text: + raise ValueError("--prompt must contain non-whitespace text.") + if args.height <= 0 or args.width <= 0: + raise ValueError("--height and --width must be positive integers.") + if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT: + raise ValueError("--height and --width must each be divisible by 16.") + if args.height * args.width > _MAX_PIXEL_AREA: + raise ValueError("--height * --width must not exceed 480 * 832 pixels.") + if args.num_frames <= 0: + raise ValueError("--num-frames must be a positive integer.") + if args.num_frames > _MAX_RAW_FRAMES: + raise ValueError("--num-frames must not exceed 117 raw frames.") + if args.num_frames > paths.camera_frames: + raise ValueError("Camera arrays must contain at least --num-frames frames.") + if (args.num_frames - 1) % _TEMPORAL_COMPRESSION: + raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.") + latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1 + if latent_frames % _LATENT_FRAMES_PER_BLOCK: + raise ValueError("--num-frames must map to whole three-frame latent blocks.") + if args.fps <= 0: + raise ValueError("--fps must be a positive integer.") + flow_shift = _positive_finite(args.flow_shift, "--flow-shift") + + prompt = { + "prompt": prompt_text, + "multi_modal_data": {"image": str(paths.image)}, + } + sampling_kwargs = { + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "num_inference_steps": _NUM_INFERENCE_STEPS, + "max_sequence_length": _MAX_SEQUENCE_LENGTH, + "seed": args.seed, + "fps": args.fps, + "extra_args": { + "action_path": paths.action_relative.as_posix(), + "flow_shift": flow_shift, + }, + } + return prompt, sampling_kwargs + + +def extract_video_array(outputs: Any) -> np.ndarray: + """Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``.""" + + if not isinstance(outputs, list) or len(outputs) != 1: + raise ValueError("Expected exactly one Omni output for the single LingBot request.") + output = outputs[0] + request_output = getattr(output, "request_output", None) + if request_output is not None and getattr(request_output, "images", None): + output = request_output + + images = getattr(output, "images", None) + if not images: + multimodal_output = getattr(output, "multimodal_output", None) or {} + images = multimodal_output.get("video") + if images is None: + raise ValueError("No video frames found in Omni output.") + if isinstance(images, list) and len(images) == 1: + images = images[0] + + video: np.ndarray = np.asarray(images) + if video.ndim == 5: + if video.shape[0] != 1: + raise ValueError(f"Expected one generated video, got batch shape {video.shape}.") + video = video[0] + if video.ndim != 4 or video.shape[-1] not in (3, 4): + raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.") + return video + + +def main(argv: Sequence[str] | None = None) -> Path: + args = parse_args(argv) + paths = resolve_cli_paths(args) + prompt, sampling_kwargs = build_request(args, paths) + paths.output.parent.mkdir(parents=True, exist_ok=True) + + # Keep vLLM imports below pure CLI/path validation so ``--help`` and the + # helper tests work on machines without the native vLLM runtime. + from diffusers.utils import export_to_video + + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + omni_kwargs = build_omni_kwargs(args, paths) + omni = Omni(**omni_kwargs) + try: + outputs = omni.generate( + prompt, + OmniDiffusionSamplingParams(**sampling_kwargs), + use_tqdm=False, + ) + finally: + omni.close() + + frames = extract_video_array(outputs) + export_to_video(frames, str(paths.output), fps=args.fps) + print(f"Saved {frames.shape[0]} frames to {paths.output}") + return paths.output + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py new file mode 100644 index 000000000..86461f9f8 --- /dev/null +++ b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session. + +This example exercises the internal session API. It is not a public HTTP or +WebSocket protocol. Each JSONL event advances the world by one three-latent +frame AR block and writes the latent plus its identity metadata. The current +scoped implementation supports at most ten ticks per generation epoch because +its image condition is bounded to 117 pixel frames; reset or create a session +to start a new world. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import math +import time +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" +_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1" +_FRAMES_PER_BLOCK = 3 +# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames. +_MAX_REALTIME_TICKS = 10 + + +def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]: + """Build the event-side script consumed by LingBotCameraControlReducer.""" + return {"mode": "script", "frames": frames} + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.") + parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.") + parser.add_argument("--image", required=True, help="Initial RGB image.") + parser.add_argument("--prompt", required=True, help="Initial scene prompt.") + parser.add_argument( + "--events", + required=True, + help="JSONL file with one prompt/action event per AR block; at most 10 events.", + ) + parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.") + parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--gpu-memory-fraction", type=float, default=0.1) + parser.add_argument("--tensor-parallel-size", type=int, default=1) + parser.add_argument("--enforce-eager", action="store_true") + return parser.parse_args(argv) + + +def _load_events(path: Path) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1): + line = raw_line.strip() + if not line: + continue + try: + value = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None + if not isinstance(value, dict): + raise ValueError(f"events line {line_number} must be a JSON object.") + event_id = value.get("event_id") + if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0: + raise ValueError(f"events line {line_number} requires a non-negative integer event_id.") + prompt = value.get("prompt") + if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()): + raise ValueError(f"events line {line_number} prompt must be a non-empty string.") + frames = value.get("frames") + if frames is not None: + if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK: + raise ValueError(f"events line {line_number} frames must contain exactly three lists.") + for frame in frames: + if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame): + raise ValueError(f"events line {line_number} frames must contain only action strings.") + if prompt is None and frames is None: + raise ValueError(f"events line {line_number} must update prompt and/or frames.") + events.append({"event_id": event_id, "prompt": prompt, "frames": frames}) + if not events: + raise ValueError("events file must contain at least one event.") + if len(events) > _MAX_REALTIME_TICKS: + raise ValueError( + "events file must contain at most " + f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime " + "image-condition horizon is 117 pixel frames." + ) + event_ids = [event["event_id"] for event in events] + if event_ids != sorted(set(event_ids)): + raise ValueError("event_id values must be unique and strictly increasing.") + return events + + +def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]: + image = Path(args.image).expanduser().resolve() + events = Path(args.events).expanduser().resolve() + output_dir = Path(args.output_dir).expanduser().resolve() + if not image.is_file(): + raise ValueError("--image must point to an existing file.") + if not events.is_file(): + raise ValueError("--events must point to an existing JSONL file.") + if not args.prompt.strip(): + raise ValueError("--prompt must contain non-whitespace text.") + if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16: + raise ValueError("--height and --width must be positive multiples of 16.") + if args.tensor_parallel_size <= 0: + raise ValueError("--tensor-parallel-size must be positive.") + if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1: + raise ValueError("--gpu-memory-fraction must be in (0, 1].") + return image, events, output_dir + + +async def run(argv: Sequence[str] | None = None) -> Path: + args = parse_args(argv) + image, events_path, output_dir = _validate_args(args) + events = _load_events(events_path) + output_dir.mkdir(parents=True, exist_ok=True) + + # Imports stay below pure input validation so --help and helper tests do not + # require a CUDA-enabled vLLM installation. + import torch + + from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer + from vllm_omni.entrypoints.async_omni import AsyncOmni + from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer + from vllm_omni.experimental.ar_diffusion.session import ( + ARDiffusionSessionEvent, + ARDiffusionSessionManager, + ARDiffusionWorkerLifecycle, + ) + from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + engine = AsyncOmni( + model=args.model, + engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine", + enforce_eager=args.enforce_eager, + tensor_parallel_size=args.tensor_parallel_size, + max_num_seqs=1, + model_config={ + "ar_diffusion_height": args.height, + "ar_diffusion_width": args.width, + "ar_diffusion_kv_config": { + "gpu_memory_fraction": args.gpu_memory_fraction, + "warmup_cudagraph": True, + }, + }, + ) + sampling = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_frames=9, + num_inference_steps=4, + max_sequence_length=512, + seed=args.seed, + output_type="latent", + extra_args={"flow_shift": 5.0}, + ) + consumer = ARDiffusionOmniTickConsumer( + engine, + prompt_provider=lambda tick: { + "prompt": tick.prompt, + "multi_modal_data": {"image": str(image)}, + }, + sampling_params_list=[sampling], + diffusion_stage_id=0, + ) + manager = ARDiffusionSessionManager( + tick_consumer=consumer, + lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0), + max_pending_events=32, + control_reducer_factory=LingBotCameraControlReducer, + ) + session = await manager.create_session(args.session_id) + measurements: list[dict[str, Any]] = [] + try: + for chunk_index, event in enumerate(events): + controls = () + if event["frames"] is not None: + controls = ( + ARDiffusionControlInput( + track="camera", + schema=_CAMERA_ACTION_SCHEMA, + data=_camera_event_data(event["frames"]), + ), + ) + await session.accept_event( + ARDiffusionSessionEvent( + event_id=event["event_id"], + prompt=event["prompt"] + if event["prompt"] is not None + else (args.prompt if chunk_index == 0 else None), + controls=controls, + ) + ) + started = time.perf_counter() + output = await session.next_chunk() + elapsed = time.perf_counter() - started + if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor): + raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.") + latent = output.images[0].detach().float().cpu() + metadata = consumer.chunk_metadata(output).to_dict() + torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt") + (output_dir / f"chunk_{chunk_index:03d}.json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n" + ) + measurements.append( + { + "chunk_index": chunk_index, + "latency_seconds": elapsed, + "shape": list(latent.shape), + "finite": bool(torch.isfinite(latent).all()), + "metadata": metadata, + } + ) + print(json.dumps(measurements[-1], sort_keys=True), flush=True) + finally: + try: + await manager.close_session(args.session_id) + finally: + engine.shutdown() + + summary_path = output_dir / "summary.json" + summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n") + return summary_path + + +def main(argv: Sequence[str] | None = None) -> Path: + return asyncio.run(run(argv)) + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/dynin_omni/README.md b/examples/offline_inference/dynin_omni/README.md index 19fa8c69e..6cf834142 100644 --- a/examples/offline_inference/dynin_omni/README.md +++ b/examples/offline_inference/dynin_omni/README.md @@ -107,4 +107,4 @@ python /examples/offline_inference/dynin_omni/end2end.py \ - If you want to force local config resolution, pass `--dynin-config-path `. - If you see the warning `max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`, - reduce `max_num_batched_tokens` in stage config (for example, `4096` in CI config). + reduce `max_num_batched_tokens` in the deploy config (for example, `4096` in the CI config). diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py index 0cab46fee..74d33fe01 100644 --- a/examples/offline_inference/dynin_omni/end2end.py +++ b/examples/offline_inference/dynin_omni/end2end.py @@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace: parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES) parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.") parser.add_argument( - "--stage-config-path", + "--deploy-config", type=str, default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"), help="Path to deploy config yaml.", @@ -1394,8 +1394,8 @@ def main() -> None: from vllm_omni.entrypoints.omni import Omni - stage_config_path = str(Path(args.stage_config_path).expanduser()) - omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype) + deploy_config_path = str(Path(args.deploy_config).expanduser()) + omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype) sampling_params_list = [ SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False) for _ in range(omni.num_stages) diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md index 8b90e6b7f..b4d5af2c7 100644 --- a/examples/offline_inference/hunyuan_image3/README.md +++ b/examples/offline_inference/hunyuan_image3/README.md @@ -7,17 +7,17 @@ YAMLs under `vllm_omni/deploy/`. | File | Topology | Default use | | :--- | :--- | :--- | -| `vllm_omni/deploy/hunyuan_image3.yaml` | AR + DiT | Default for `text2img` and `img2img`. | +| `vllm_omni/deploy/hunyuan_image_3_moe.yaml` | AR + DiT | Default for `text2img` and `img2img`. | | `vllm_omni/deploy/hunyuan_image3_ar.yaml` | AR only | Default for `img2text` and `text2text`. | | `vllm_omni/deploy/hunyuan_image3_dit.yaml` | DiT only | Standalone diffusion stage. Pass it explicitly with `--deploy-config`. | -The example chooses a deploy config automatically when `--deploy-config` and -`--stage-configs-path` are both omitted: +The example chooses a deploy config automatically when `--deploy-config` is +omitted: | `--modality` | `mode` passed to Omni | Default deploy | | :--- | :--- | :--- | -| `text2img` | `text-to-image` | `hunyuan_image3.yaml` | -| `img2img` | `image-editing` | `hunyuan_image3.yaml` | +| `text2img` | `text-to-image` | `hunyuan_image_3_moe.yaml` | +| `img2img` | `image-editing` | `hunyuan_image_3_moe.yaml` | | `img2text` | `image-to-text` | `hunyuan_image3_ar.yaml` | | `text2text` | `text-to-text` | `hunyuan_image3_ar.yaml` | @@ -35,8 +35,8 @@ path; the image endpoints infer the image task from the endpoint and payload. | Online scenario | Server deploy | Request | | :--- | :--- | :--- | -| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image3.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. | -| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image3.yaml` | `POST /v1/images/edits`. | +| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. | +| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/edits`. | | Image/text to text | `--deploy-config vllm_omni/deploy/hunyuan_image3_ar.yaml` | `POST /v1/chat/completions` for text output, for example with `"modalities": ["text"]`. | | DiT-only image generation | `--deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml` | `POST /v1/images/generations`. | @@ -96,7 +96,7 @@ Override the default full AR + DiT deploy explicitly: python examples/offline_inference/hunyuan_image3/end2end.py \ --model tencent/HunyuanImage-3.0-Instruct \ --modality text2img \ - --deploy-config vllm_omni/deploy/hunyuan_image3.yaml \ + --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ --prompts "A cute cat" ``` @@ -117,8 +117,7 @@ python end2end.py --modality text2img \ | Argument | Description | | :--- | :--- | -| `--deploy-config` | Preferred config path for unified deploy YAMLs. | -| `--stage-configs-path` | Legacy stage config path, kept only for compatibility. Prefer `--deploy-config`. | +| `--deploy-config` | Optional path for unified deploy YAML overrides. | | `--additional-config` | JSON object forwarded to diffusion worker `additional_config`. | | `--modality` | Offline-only convenience flag. One of `text2img`, `img2img`, `img2text`, `text2text`. It selects prompt formatting, internal `mode`, and default deploy config for this script. Online serving uses `--deploy-config` plus the endpoint and, for chat completions, request `modalities` instead. | | `--steps` | Number of diffusion inference steps for image generation. | @@ -132,7 +131,7 @@ python end2end.py --modality text2img \ - `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy. - `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`. -- The old HunyuanImage3 YAMLs under `model_executor/stage_configs/` and `platforms/*/stage_configs/` have been folded into the deploy YAMLs. +- Platform-specific overrides are included in the unified deploy YAMLs. ## Prompt Format diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 9a2350156..fbe4289be 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -73,7 +73,6 @@ def parse_args(): ) parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.") parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.") - parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.") parser.add_argument( "--stream", action="store_true", @@ -146,13 +145,7 @@ def main(): else: bot_task = args.bot_task - if args.deploy_config is not None and args.stage_configs_path is not None: - raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.") - - deploy_config = args.deploy_config - stage_configs_path = args.stage_configs_path - if deploy_config is None and stage_configs_path is None: - deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality] + deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality] omni_kwargs = { "model": args.model, @@ -168,10 +161,7 @@ def main(): if additional_config is not None: omni_kwargs["additional_config"] = additional_config - if deploy_config is not None: - omni_kwargs["deploy_config"] = deploy_config - else: - omni_kwargs["stage_configs_path"] = stage_configs_path + omni_kwargs["deploy_config"] = deploy_config omni = Omni(**omni_kwargs) @@ -272,10 +262,7 @@ def main(): print(f" Modality: {args.modality}") print(f" Prompt task: {task}") print(f" Bot task: {bot_task}") - if deploy_config is not None: - print(f" Deploy config: {deploy_config}") - else: - print(f" Stage config: {stage_configs_path}") + print(f" Deploy config: {deploy_config}") print(f" Num stages: {omni.num_stages}") if args.modality in ("text2img", "img2img"): print(f" Inference steps: {args.steps}") diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index ae84353f7..136d99857 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -96,7 +96,6 @@ from typing import Any import torch from PIL import Image -from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni @@ -452,14 +451,6 @@ def main(): generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, - ring_degree=args.ring_degree, - cfg_parallel_size=args.cfg_parallel_size, - tensor_parallel_size=args.tensor_parallel_size, - enable_expert_parallel=args.enable_expert_parallel, - ) - # Configure cache based on backend type cache_config = None if args.cache_backend == "cache_dit": @@ -490,7 +481,11 @@ def main(): vae_use_tiling=args.vae_use_tiling, cache_backend=args.cache_backend, cache_config=cache_config, - parallel_config=parallel_config, + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, + enable_expert_parallel=args.enable_expert_parallel, enable_cpu_offload=args.enable_cpu_offload, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, profiler_config=args.profiler_config, diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index c8f7fe617..02f233642 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -1,7 +1,7 @@ # Image-To-Video This shared example generates videos from images with VACE, Wan2.2, LTX-2, -HunyuanVideo-1.5, Cosmos3, and other compatible pipelines. +HunyuanVideo-1.5, SANA-Video, Cosmos3, and other compatible pipelines. - `image_to_video.py`: command-line script for single video generation with advanced options. @@ -27,6 +27,8 @@ This folder provides a unified CLI script for image-to-video generation using vL | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card | | `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism | +| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480 x 832 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; Wan VAE | +| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704 x 1280 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; LTX-2 Video VAE | !!! info Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading. @@ -96,7 +98,7 @@ if __name__ == "__main__": ), ) from diffusers.utils import export_to_video - frames = outputs[0].request_output.images + frames = outputs[0].images export_to_video(frames, "quick_test_i2v.mp4", fps=16) ``` @@ -227,6 +229,48 @@ python image_to_video.py \ See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints, pipeline selection, T2V, defaults, and advanced options. +### SANA-Video-2B + +The checkpoint metadata selects the upstream T2V pipeline, so pass +`SanaImageToVideoPipeline` explicitly for I2V: + +```bash +python image_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaImageToVideoPipeline \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms sway in the breeze as petals drift past the camera." \ + --negative-prompt "blurry, low quality, temporal artifacts" \ + --height 480 \ + --width 832 \ + --num-frames 81 \ + --num-inference-steps 50 \ + --guidance-scale 6.0 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 \ + --seed 42 \ + --output sana_video_i2v_480p.mp4 +``` + +Use the 720p checkpoint with `--height 704 --width 1280`. In this example, +81 frames at 16 FPS means the standard approximately five-second SANA-Video +checkpoint request. It does not enable the separate LongSANA/LongLive +minute-scale autoregressive workflow. + +The native `SanaImageToVideoPipeline` supports both checkpoint variants. The +480p path uses vLLM-Omni's `DistributedAutoencoderKLWan` wrapper and the 720p +path uses `DistributedAutoencoderKLLTX2Video`; both wrap their corresponding +Diffusers VAE implementations. The native denoising loop intentionally keeps +Diffusers' checkpoint-compatible `DPMSolverMultistepScheduler`. + +The Diffusers adapter has completed end-to-end I2V serving validation at both +480p and 720p. For online adapter serving, select the matching checkpoint and +use `WIDTH=1280 HEIGHT=704` for the 720p request. Adapter T2V is also validated +at both resolutions; see the recipe for the exact backend matrix. + +See the [SANA-Video recipe](../../../recipes/NVIDIA/SANA-Video-2B.md) for +online serving, native/adapter validation boundaries, and hardware notes. + ## Advanced Features ### CFG Parallel diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 222502e8a..65e7389f0 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -1,15 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3, -HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE. +HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE. Supports: - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B) - LTX2 image-to-video pipeline - HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning +- SANA-Video 2B: first-frame latent conditioning at 480p or 720p - Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning Usage: @@ -38,6 +39,13 @@ Usage: --image input.jpg --prompt "A cat playing with yarn" \ --flow-shift 5.0 --guidance-scale 6.0 + # SANA-Video 2B I2V (480p) + python image_to_video.py --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaImageToVideoPipeline \ + --image input.jpg --prompt "A cat turns toward the camera." \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 50 \ + --guidance-scale 6.0 + # Cosmos3 I2V (image conditioning) python image_to_video.py --model nvidia/Cosmos3-Nano \ --image input.jpg --prompt "The scene comes to life with smooth, natural motion." \ @@ -56,7 +64,6 @@ import numpy as np import PIL.Image import torch -from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -68,6 +75,7 @@ from vllm_omni.model_extras import ( from vllm_omni.model_extras import ( get_extra_body_params, get_model_class_name, + should_preserve_reference_image_size, ) from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -103,17 +111,27 @@ def build_image_to_video_prompt( return result +def _validate_video_output_type(output_type: str) -> None: + if output_type not in {"image", "video"}: + raise ValueError( + f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation." + ) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "Generate a video from one or more images " - "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE)." + "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)." ) ) parser.add_argument( "--model", default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", - help="Diffusers I2V model ID or local path (Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE).", + help=( + "Diffusers I2V model ID or local path " + "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)." + ), ) parser.add_argument( "--model-class-name", @@ -377,6 +395,7 @@ def main(): is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name + is_sana = model_class_name == "SanaImageToVideoPipeline" or "sana-video" in model_name is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower()) is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower) @@ -434,6 +453,17 @@ def main(): 512 * 768, 32, ) + elif is_sana: + is_720p = "720p" in model_name + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( + 16, + 6.0, + 81, + 50, + 5.0, + (704 * 1280) if is_720p else (480 * 832), + 32 if is_720p else 16, + ) else: # Wan2.2 / HunyuanVideo-1.5 d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16 @@ -454,7 +484,13 @@ def main(): media_inputs: dict[str, Any] = {} if image is not None: - media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + preserve_image_size = should_preserve_reference_image_size( + model_class_name, + model=args.model, + ) + media_inputs["image"] = ( + image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS) + ) if last_image is not None: media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS) if mask_image is not None: @@ -496,17 +532,6 @@ def main(): } profiler_enabled = args.profiler_config is not None - parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, - ring_degree=args.ring_degree, - cfg_parallel_size=args.cfg_parallel_size, - tensor_parallel_size=args.tensor_parallel_size, - vae_patch_parallel_size=args.vae_patch_parallel_size, - use_hsdp=args.use_hsdp, - hsdp_shard_size=args.hsdp_shard_size, - hsdp_replicate_size=args.hsdp_replicate_size, - pipeline_parallel_size=args.pipeline_parallel_size, - ) omni_kwargs = dict( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, @@ -517,7 +542,15 @@ def main(): diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps, diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers, enable_cpu_offload=args.enable_cpu_offload, - parallel_config=parallel_config, + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, + use_hsdp=args.use_hsdp, + hsdp_shard_size=args.hsdp_shard_size, + hsdp_replicate_size=args.hsdp_replicate_size, + pipeline_parallel_size=args.pipeline_parallel_size, enforce_eager=args.enforce_eager, model_class_name=model_class_name, cache_backend=args.cache_backend, @@ -639,10 +672,7 @@ def main(): frames = frames[0] if frames else None if isinstance(frames, OmniRequestOutput): - if frames.final_output_type != "image": - raise ValueError( - f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation." - ) + _validate_video_output_type(frames.final_output_type) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py index d26910161..2b70bd749 100644 --- a/examples/offline_inference/lance/end2end.py +++ b/examples/offline_inference/lance/end2end.py @@ -82,9 +82,6 @@ def parse_args(): # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below # and ``create_default_diffusion`` materializes the stage config. - from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults - - nullify_stage_engine_defaults(p) return p.parse_args() diff --git a/examples/offline_inference/longcat_video/end2end.py b/examples/offline_inference/longcat_video/end2end.py index 8071dfc25..6e0c5593e 100644 --- a/examples/offline_inference/longcat_video/end2end.py +++ b/examples/offline_inference/longcat_video/end2end.py @@ -15,7 +15,6 @@ import imageio import numpy as np import torch -from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import ( prepare_longcat_video_avatar_model_for_omni, ) @@ -236,14 +235,6 @@ def main(): dtype="bfloat16", enforce_eager=args.enforce_eager, additional_config=additional_config, - parallel_config=DiffusionParallelConfig( - ulysses_degree=1, - ring_degree=1, - cfg_parallel_size=1, - tensor_parallel_size=1, - vae_patch_parallel_size=1, - pipeline_parallel_size=1, - ), ) audio_input: str | dict[str, str] | None = args.audio diff --git a/examples/offline_inference/mimo_audio/README.md b/examples/offline_inference/mimo_audio/README.md index 60f433a4f..126896259 100644 --- a/examples/offline_inference/mimo_audio/README.md +++ b/examples/offline_inference/mimo_audio/README.md @@ -19,7 +19,7 @@ MiMo-Audio provides multiple task variants for audio understanding and generatio ## Setup -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. +See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation for your hardware. ### Environment Variables @@ -189,7 +189,7 @@ Note: This task uses hardcoded message lists in the script. ### Other -- If the model or stage config fails to load, check [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings. +- If the model or deploy config fails to load, check the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings. - For errors when reading/writing WAV (e.g. unsupported format), ensure input files are standard WAV/MP3 and that `soundfile` is linked to a working libsndfile (see above). ## Notes diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md index ee7e93fd0..29fd59a8d 100644 --- a/examples/offline_inference/ming_flash_omni/README.md +++ b/examples/offline_inference/ming_flash_omni/README.md @@ -14,7 +14,7 @@ For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Te ## Setup -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. +See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation for your hardware. When no `--deploy-config` is passed, the model registry auto-loads the full thinker+talker `vllm_omni/deploy/ming_flash_omni.yaml` (See [Omni-Speech](#omni-speech-thinker--talker)). diff --git a/examples/offline_inference/minicpmo/README.md b/examples/offline_inference/minicpmo/README.md index b80fc0aa6..b827a4e98 100644 --- a/examples/offline_inference/minicpmo/README.md +++ b/examples/offline_inference/minicpmo/README.md @@ -8,7 +8,7 @@ Two-stage pipeline: **thinker** (multimodal understanding) → **talker + Token2 - Install talker deps: `pip install stepaudio2-minicpmo` - `--trust-remote-code` is always passed by `end2end.py` (required for MiniCPMO) -- See [stage configuration docs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory tuning +- See [pipeline and deploy configuration docs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory tuning ## Run examples diff --git a/examples/offline_inference/minicpmo/end2end.py b/examples/offline_inference/minicpmo/end2end.py index e9c9019c0..8e536735d 100644 --- a/examples/offline_inference/minicpmo/end2end.py +++ b/examples/offline_inference/minicpmo/end2end.py @@ -335,7 +335,7 @@ def main(args): print(f"use_tts (tts_bos): {use_tts}") for stage_outputs in omni_generator: - output = stage_outputs.request_output + output = stage_outputs if stage_outputs.final_output_type == "text": request_id = output.request_id text_output = output.outputs[0].text @@ -452,12 +452,6 @@ def parse_args(): default=None, help="Path to a .txt file with one prompt per line (preferred).", ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="Path to a stage configs file (deprecated; prefer --deploy-config).", - ) parser.add_argument( "--deploy-config", type=str, diff --git a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py b/examples/offline_inference/minimax_h3/dlo_lifecycle.py similarity index 72% rename from examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py rename to examples/offline_inference/minimax_h3/dlo_lifecycle.py index f8117b0f9..5630777cb 100644 --- a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py +++ b/examples/offline_inference/minimax_h3/dlo_lifecycle.py @@ -1,21 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior. +"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior. + +The DLO modes accept any positive data-parallel size. With AllGather enabled, +the script submits one complete DP wave and verifies recovery after an invalid +wave. It also checks that every mode shuts down all worker processes cleanly. Examples: VLLM_WORKER_MULTIPROC_METHOD=spawn \ VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \ + CUDA_VISIBLE_DEVICES=0,1,2,3 \ + python examples/offline_inference/minimax_h3/dlo_lifecycle.py \ + --model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4 + + # Rank-local mmap storage, with no DLO collective or DP wave scheduling. + VLLM_WORKER_MULTIPROC_METHOD=spawn \ + VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ CUDA_VISIBLE_DEVICES=0,1 \ - python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \ - --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2 + python examples/offline_inference/minimax_h3/dlo_lifecycle.py \ + --model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2 VLLM_WORKER_MULTIPROC_METHOD=spawn \ VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ CUDA_VISIBLE_DEVICES=0,1 \ - python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \ - --model /path/to/MiniMax-H3/FL2VA --mode request + python examples/offline_inference/minimax_h3/dlo_lifecycle.py \ + --model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2 """ from __future__ import annotations @@ -40,14 +51,33 @@ DEFAULT_PROMPTS = ( ) +def positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA") parser.add_argument( "--mode", - choices=("dlo-dp2", "request"), + choices=("dlo", "dlo-no-allgather", "request"), required=True, - help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode", + help="DLO (with or without AllGather) or ordinary non-DLO request mode", + ) + parser.add_argument( + "--dp-size", + type=positive_int, + default=2, + help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)", + ) + parser.add_argument( + "--tp-size", + type=positive_int, + default=2, + help="Tensor-parallel size used only by request mode (default: 2)", ) parser.add_argument("--steps", type=int, default=2) parser.add_argument("--duration", type=float, default=5.0) @@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace: def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]: + is_dlo = args.mode in {"dlo", "dlo-no-allgather"} common: dict[str, Any] = { "model": args.model, "trust_remote_code": True, - "num_gpus": 2, + "num_gpus": args.dp_size if is_dlo else args.tp_size, "usp": 1, "ring": 1, "vae_parallel_mode": "tile", @@ -74,22 +105,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]: "stage_init_timeout": args.init_timeout, "init_timeout": args.init_timeout, } - if args.mode == "dlo-dp2": + if is_dlo: common.update( tensor_parallel_size=1, - data_parallel_size=2, + data_parallel_size=args.dp_size, text_encoder_tp_size=1, vae_patch_parallel_size=1, enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, + dlo_use_allgather=args.mode == "dlo", dlo_resident_layers=0, ) else: common.update( - tensor_parallel_size=2, + tensor_parallel_size=args.tp_size, data_parallel_size=1, - text_encoder_tp_size=2, - vae_patch_parallel_size=2, + text_encoder_tp_size=args.tp_size, + vae_patch_parallel_size=args.tp_size, enable_distributed_layerwise_offload=False, ) return common @@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]: "engine_kwargs": engine_kwargs(args), } try: - if args.mode == "dlo-dp2": + uses_dp_wave = args.mode == "dlo" and args.dp_size > 1 + if uses_dp_wave: started = time.perf_counter() asymmetric = await asyncio.gather( generate_one( @@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]: prompt="", seed=1001, ), - generate_one( - engine, - args, - request_id="invalid-peer", - prompt=DEFAULT_PROMPTS[0], - seed=1002, + *( + generate_one( + engine, + args, + request_id=f"invalid-peer-{index}", + prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)], + seed=1001 + index, + ) + for index in range(1, args.dp_size) ), return_exceptions=True, ) @@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]: summary["asymmetric_errors"] = [ f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException) ] - if len(summary["asymmetric_errors"]) != 2: - raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch") + if len(summary["asymmetric_errors"]) != args.dp_size: + raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch") started = time.perf_counter() - request_count = 2 if args.mode == "dlo-dp2" else 1 + request_count = args.dp_size if uses_dp_wave else 1 outputs = await asyncio.gather( *( generate_one( engine, args, request_id=f"recovery-{index}", - prompt=DEFAULT_PROMPTS[index], + prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)], seed=2000 + index, ) for index in range(request_count) diff --git a/examples/offline_inference/nemotron_voicechat/end2end.py b/examples/offline_inference/nemotron_voicechat/end2end.py index b020e5bad..ac5fac376 100644 --- a/examples/offline_inference/nemotron_voicechat/end2end.py +++ b/examples/offline_inference/nemotron_voicechat/end2end.py @@ -67,7 +67,7 @@ def main() -> None: parser.add_argument("--output-dir", default="results/nemotron_voicechat") parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT) parser.add_argument( - "--stage-configs-path", + "--deploy-config", default=None, help=( "Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml " @@ -134,8 +134,8 @@ def main() -> None: ] omni_kwargs: dict = {} - if args.stage_configs_path: - omni_kwargs["stage_configs_path"] = args.stage_configs_path + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs) outputs = omni.generate(inputs, sampling_params_list) diff --git a/examples/offline_inference/qwen2_5_omni/README.md b/examples/offline_inference/qwen2_5_omni/README.md index 059c10226..490711d34 100644 --- a/examples/offline_inference/qwen2_5_omni/README.md +++ b/examples/offline_inference/qwen2_5_omni/README.md @@ -1,7 +1,7 @@ # Qwen2.5-Omni: Offline inference ## Setup -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. +See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation for your hardware. ## Run examples diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md index 2f4ee11cc..d216042b4 100644 --- a/examples/offline_inference/qwen3_omni/README.md +++ b/examples/offline_inference/qwen3_omni/README.md @@ -1,7 +1,8 @@ # Qwen3-Omni: Offline inference ## Setup -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. +Use `--deploy-config` for deployment overrides such as stage memory allocation. +See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/). ## Run examples @@ -70,8 +71,8 @@ For true stage-level concurrency -- where downstream stages (Talker, Code2Wav) start **before** the upstream stage (Thinker) finishes -- use the async_chunk example. This requires: -1. A deploy config YAML with ``async_chunk: true`` (e.g. - ``qwen3_omni_moe.yaml``). +1. A deploy YAML with ``async_chunk: true`` (for example, an overlay based on + ``vllm_omni/deploy/qwen3_omni_moe.yaml``). 2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage config). @@ -97,11 +98,11 @@ bash run_multiple_prompts_async_chunk.sh --max-in-flight 4 python end2end_async_chunk.py --query-type text --modalities text ``` -#### Custom stage config +#### Custom deploy config ```bash python end2end_async_chunk.py \ --query-type use_audio \ - --deploy-config /path/to/your_deploy_config.yaml + --deploy-config /path/to/your_async_chunk.yaml ``` > **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 9434b3058..faa200bc9 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -484,10 +484,10 @@ def parse_args(): help="Path to a .txt file with one prompt per line (preferred).", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, default=None, - help="Path to a stage configs file.", + help="Path to a deploy config YAML.", ) parser.add_argument( "--video-path", diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py index c224586c4..c4570abae 100644 --- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py +++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py @@ -386,8 +386,8 @@ async def run_all(args): # argparse defaults do not silently override deploy YAML values. async_omni = AsyncOmni.from_cli_args(args, model=args.model) - # Use default sampling params from stage config (they are pre-configured - # in the YAML for each stage). + # Use default sampling params from the resolved pipeline and deploy + # config. # # NOTE: Since we do not set the sampling params directly, .generate in # will automatically set the output kind to delta, since this is what diff --git a/examples/offline_inference/speech_to_video/speech_to_video.py b/examples/offline_inference/speech_to_video/speech_to_video.py index 964cd4874..62dd29720 100644 --- a/examples/offline_inference/speech_to_video/speech_to_video.py +++ b/examples/offline_inference/speech_to_video/speech_to_video.py @@ -23,7 +23,6 @@ from typing import Any import torch -from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -250,17 +249,6 @@ def main(): "scm_steps_policy": "dynamic", } - parallel_config = DiffusionParallelConfig( - tensor_parallel_size=args.tensor_parallel_size, - cfg_parallel_size=args.cfg_parallel_size, - vae_patch_parallel_size=args.vae_patch_parallel_size, - ring_degree=args.ring_degree, - ulysses_degree=args.ulysses_degree, - use_hsdp=args.use_hsdp, - hsdp_shard_size=args.hsdp_shard_size, - hsdp_replicate_size=args.hsdp_replicate_size, - ) - profiler_enabled = args.profiler_config is not None omni_kwargs = dict( @@ -271,7 +259,14 @@ def main(): vae_use_tiling=args.vae_use_tiling, enable_cpu_offload=args.enable_cpu_offload, enable_layerwise_offload=args.enable_layerwise_offload, - parallel_config=parallel_config, + tensor_parallel_size=args.tensor_parallel_size, + cfg_parallel_size=args.cfg_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, + ring_degree=args.ring_degree, + ulysses_degree=args.ulysses_degree, + use_hsdp=args.use_hsdp, + hsdp_shard_size=args.hsdp_shard_size, + hsdp_replicate_size=args.hsdp_replicate_size, enforce_eager=args.enforce_eager, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, cache_backend=args.cache_backend, diff --git a/examples/offline_inference/step_audio2/README.md b/examples/offline_inference/step_audio2/README.md index abd21f371..7457863f5 100644 --- a/examples/offline_inference/step_audio2/README.md +++ b/examples/offline_inference/step_audio2/README.md @@ -29,7 +29,7 @@ Step-Audio2 is a two-stage audio model: **Notes:** - Single GPU mode requires high VRAM due to both stages sharing memory - Multi GPU mode separates Stage 0 (Thinker) and Stage 1 (Token2Wav) across GPUs -- VRAM usage can be adjusted via `gpu_memory_utilization` in stage config +- VRAM usage can be adjusted via `gpu_memory_utilization` in the deploy config ## Performance Benchmark @@ -197,11 +197,6 @@ This mode: ### Advanced Options ```bash -# Use custom stage configuration -python end2end.py --query-type audio_to_text \ - --model stepfun-ai/Step-Audio-2-mini \ - --stage-configs-path /path/to/custom_config.yaml - # Use custom deploy configuration python end2end.py --query-type audio_to_text \ --model stepfun-ai/Step-Audio-2-mini \ diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py index f175db750..418a7f129 100644 --- a/examples/offline_inference/step_audio2/end2end.py +++ b/examples/offline_inference/step_audio2/end2end.py @@ -232,15 +232,8 @@ def main(args): raise ValueError(f"Unknown query type: {args.query_type}") # Initialize vLLM-Omni with Step-Audio2 - # Resolve deploy config path. Keep explicit stage_configs_path for legacy - # custom configs, but use bundled deploy configs by default. - deploy_config_path = None - stage_config_path = None - if args.stage_configs_path and args.deploy_config: - raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive") - if args.stage_configs_path: - stage_config_path = args.stage_configs_path - elif args.deploy_config: + # Resolve an explicit deploy config or use the bundled task default. + if args.deploy_config: deploy_config_path = args.deploy_config else: configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy" @@ -262,10 +255,7 @@ def main(args): "ray_address": args.ray_address, "trust_remote_code": True, } - if stage_config_path is not None: - omni_kwargs["stage_configs_path"] = stage_config_path - else: - omni_kwargs["deploy_config"] = deploy_config_path + omni_kwargs["deploy_config"] = deploy_config_path omni_llm = Omni(**omni_kwargs) @@ -418,12 +408,6 @@ def parse_args(): " - HuggingFace ID: stepfun-ai/Step-Audio-2-mini" ), ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="Path to a legacy stage config YAML file (default: use bundled deploy config)", - ) parser.add_argument( "--deploy-config", type=str, diff --git a/examples/offline_inference/text_to_audio/README.md b/examples/offline_inference/text_to_audio/README.md index 89ffc8bd8..f096d96c8 100644 --- a/examples/offline_inference/text_to_audio/README.md +++ b/examples/offline_inference/text_to_audio/README.md @@ -1,11 +1,10 @@ # Text-To-Audio -A unified script for text/video-to-audio generation. Supported models: +A unified script for text-to-audio generation. Supported models: | Model | Tasks | Notes | |-------|-------|-------| | `stabilityai/stable-audio-open-1.0` | text-to-audio | gated; uses `--audio-length` | -| `zhangj1an/AudioX` | `t2a` / `t2m` / `v2a` / `v2m` / `tv2a` / `tv2m` | pass `--task`; video tasks need `--video` | The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts. @@ -50,42 +49,12 @@ python text_to_audio.py \ --output stable_audio_output.wav ``` -### AudioX - -AudioX supports six tasks. Sampler and reference knobs (declared in -`vllm_omni/model_extras/audiox.py`) are passed via the generic `--extra-body` -JSON flag, routed into sampling `extra_args`. - -Text tasks (`t2a` / `t2m`): - -```bash -python text_to_audio.py \ - --model zhangj1an/AudioX --task t2a \ - --prompt "Fireworks burst twice, followed by a clock ticking." \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 --seed 42 \ - --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \ - --output t2a.wav -``` - -Video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) require `--video`: - -```bash -python text_to_audio.py \ - --model zhangj1an/AudioX --task tv2a \ - --prompt "drum beating sound and human talking" \ - --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ - --output tv2a.wav -``` - Key arguments: - `--prompt`: text description (string). -- `--task`: [AudioX] one of `t2a`/`t2m`/`v2a`/`v2m`/`tv2a`/`tv2m`. -- `--video`: [AudioX `v2*`/`tv2*`] video file/URL for conditioning (→ `video_path`). -- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s` for Stable Audio, `seconds_start` for AudioX). -- `--audio-length`: audio duration in seconds (audio length for Stable Audio, `seconds_total` for AudioX). -- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. For AudioX, sampler/reference knobs go here, e.g. `'{"sigma_min": 0.03, "sigma_max": 1000.0, "cfg_rescale": 0.0, "audio_path": "ref.wav"}'`. +- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s`). +- `--audio-length`: audio duration in seconds (audio length for Stable Audio). +- `--extra-body`: JSON dict of model-specific knobs, merged into sampling `extra_args`. - `--negative-prompt`: negative prompt for classifier-free guidance. - `--seed`: integer seed for deterministic generation. - `--guidance-scale`: classifier-free guidance scale. diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index eff9c9d1a..094a1c66e 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -2,15 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Example script for text/video-to-audio generation with diffusion audio models. +Example script for text-to-audio generation with diffusion audio models. This script supports: * Stable Audio Open (`stabilityai/stable-audio-open-1.0`) — text-to-audio. - * AudioX (`zhangj1an/AudioX`) — six tasks: t2a / t2m / v2a / v2m / tv2a / tv2m. -Model-specific generation knobs are declared in `vllm_omni/model_extras` and -routed into `sampling_params.extra_args` via `apply_declared_extra_args`, so -the same script serves both models without per-model branching in the engine. +Model-specific generation knobs can be passed via `--extra-body` and are +merged into `sampling_params.extra_args`. Usage: # Stable Audio Open @@ -18,15 +16,6 @@ Usage: python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0 python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality" python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache - - # AudioX (text-to-audio / text-to-music) - python text_to_audio.py --model zhangj1an/AudioX --task t2a \ - --prompt "Fireworks burst twice, then a clock ticking." \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 - # AudioX video-conditioned (v2*/tv2*) — pass a video file/URL - python text_to_audio.py --model zhangj1an/AudioX --task tv2a \ - --prompt "drum beating sound and human talking" --video clip.mp4 \ - --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' """ import argparse @@ -38,17 +27,11 @@ from typing import Any import numpy as np import torch -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.model_extras import get_extra_body_params, get_model_class_name from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform -AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m") -AUDIOX_VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"}) - def parse_extra_body(value: str) -> dict[str, Any]: """Parse a JSON object string for --extra-body.""" @@ -66,7 +49,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--model", default="stabilityai/stable-audio-open-1.0", - help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0, zhangj1an/AudioX.", + help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0.", ) parser.add_argument( "--prompt", @@ -97,19 +80,13 @@ def parse_args() -> argparse.Namespace: "--audio-start", type=float, default=0.0, - help=( - "Audio start offset in seconds. Applies to both models " - "(audio_start_in_s for Stable Audio, seconds_start for AudioX)." - ), + help="Audio start offset in seconds (maps to audio_start_in_s).", ) parser.add_argument( "--audio-length", type=float, default=10.0, - help=( - "Audio duration in seconds. Maps to audio length for Stable Audio " - "(max ~47s for stable-audio-open-1.0) and seconds_total for AudioX." - ), + help=("Audio duration in seconds. Maps to audio length for Stable Audio (max ~47s for stable-audio-open-1.0)."), ) parser.add_argument( "--num-inference-steps", @@ -123,41 +100,20 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of audio waveforms to generate for the given prompt.", ) - parser.add_argument( - "--task", - type=str, - default=None, - choices=list(AUDIOX_TASKS), - help=( - "[AudioX only] Generation task: t2a/t2m/v2a/v2m/tv2a/tv2m. " - "v2*/tv2* require --video. Ignored for Stable Audio Open." - ), - ) - parser.add_argument( - "--video", - type=str, - default=None, - help="[AudioX v2*/tv2*] Video file path or URL used for conditioning.", - ) parser.add_argument( "--extra-body", type=parse_extra_body, default=None, help=( - "Model-specific generation params as a JSON object, e.g. " - '\'{"sigma_min": 0.03, "sigma_max": 1000.0}\'. Each key is filtered ' - "against the model's declared extra_body_params (see vllm_omni/model_extras), " - "so unknown keys are silently dropped. Values here override the equivalent flags." + "Model-specific generation params as a JSON object, merged into " + "sampling extra_args. Values here override the equivalent flags." ), ) parser.add_argument( "--output", type=str, default=None, - help=( - "Path to save the generated audio (WAV format). Defaults to " - "audiox_output.wav for AudioX and stable_audio_output.wav otherwise." - ), + help=("Path to save the generated audio (WAV format). Defaults to stable_audio_output.wav."), ) parser.add_argument( "--sample-rate", @@ -276,8 +232,6 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410 def main(): args = parse_args() - model_name = str(args.model).lower() if args.model else "" - is_audiox = "audiox" in model_name generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) cache_config = None if args.cache_backend == "tea_cache": @@ -286,11 +240,9 @@ def main(): } print(f"\n{'=' * 60}") - print("Text/Video-to-Audio Generation") + print("Text-to-Audio Generation") print(f"{'=' * 60}") print(f" Model: {args.model}") - if args.task is not None: - print(f" Task: {args.task}") print(f" Prompt: {args.prompt}") print(f" Negative prompt: {args.negative_prompt}") print(f" Duration: {args.audio_length}s") @@ -306,28 +258,18 @@ def main(): print(f" Seed: {args.seed}") print(f"{'=' * 60}\n") - parallel_config = DiffusionParallelConfig( + omni_kwargs: dict[str, Any] = dict( + model=args.model, use_hsdp=args.use_hsdp, hsdp_shard_size=args.hsdp_shard_size, hsdp_replicate_size=args.hsdp_replicate_size, - ) - - # Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio - # is resolved from its config. - omni_kwargs: dict[str, Any] = dict( - model=args.model, - parallel_config=parallel_config, cache_backend=args.cache_backend, cache_config=cache_config, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, enable_cpu_offload=args.enable_cpu_offload, enable_layerwise_offload=args.enable_layerwise_offload, ) - if is_audiox: - omni_kwargs["model_class_name"] = "AudioXPipeline" omni = Omni(**omni_kwargs) - model_class_name = get_model_class_name(omni) - declared_extra_body_params = get_extra_body_params(model_class_name) diffusion_params = OmniDiffusionSamplingParams( generator=generator, @@ -337,38 +279,19 @@ def main(): seed=args.seed, ) - # Negative prompt is opt-in for both models (off by default). Including a + # Negative prompt is opt-in (off by default). Including a # non-empty negative prompt enables the separate CFG conditioning branch. prompt: dict[str, Any] = {"prompt": args.prompt} if args.negative_prompt: prompt["negative_prompt"] = args.negative_prompt - if model_class_name == "AudioXPipeline": - if args.task is None: - raise ValueError("AudioX requires --task (one of t2a/t2m/v2a/v2m/tv2a/tv2m).") - if args.task in AUDIOX_VIDEO_TASKS and not args.video: - raise ValueError(f"AudioX task {args.task!r} requires --video.") - # AudioX core knobs are surfaced as shared flags; sampler/reference knobs - # (sigma_min/sigma_max/cfg_rescale/audio_path) arrive via --extra-body and - # are filtered against the declared extra_body contract. - user_extra: dict[str, Any] = { - "audiox_task": args.task, - "seconds_start": args.audio_start, - "seconds_total": args.audio_length, - "video_path": args.video, - } - if args.extra_body: - user_extra.update(args.extra_body) - apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra) - else: - # Stable Audio Open path (backward-compatible). - audio_end_in_s = args.audio_start + args.audio_length - diffusion_params.extra_args = { - "audio_start_in_s": args.audio_start, - "audio_end_in_s": audio_end_in_s, - } - if args.extra_body: - diffusion_params.extra_args.update(args.extra_body) + audio_end_in_s = args.audio_start + args.audio_length + diffusion_params.extra_args = { + "audio_start_in_s": args.audio_start, + "audio_end_in_s": audio_end_in_s, + } + if args.extra_body: + diffusion_params.extra_args.update(args.extra_body) generation_start = time.perf_counter() @@ -379,7 +302,7 @@ def main(): print(f"Total generation time: {generation_time:.2f} seconds") - output = args.output or ("audiox_output.wav" if model_class_name == "AudioXPipeline" else "stable_audio_output.wav") + output = args.output or "stable_audio_output.wav" output_path = Path(output) output_path.parent.mkdir(parents=True, exist_ok=True) suffix = output_path.suffix or ".wav" diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index b71c998b3..cfe05255d 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -74,7 +74,7 @@ if __name__ == "__main__": omni = Omni(model="Qwen/Qwen-Image") prompt = "a cup of coffee on the table" outputs = omni.generate(prompt) - images = outputs[0].request_output.images + images = outputs[0].images images[0].save("coffee.png") ``` @@ -265,7 +265,7 @@ if __name__ == "__main__": ] outputs = omni.generate(prompts) for i, output in enumerate(outputs): - output.request_output.images[0].save(f"{i}.jpg") + output.images[0].save(f"{i}.jpg") ``` !!! info @@ -301,7 +301,7 @@ if __name__ == "__main__": } ]) for i, output in enumerate(outputs): - output.request_output.images[0].save(f"{i}.jpg") + output.images[0].save(f"{i}.jpg") ``` You can also pass a negative prompt via the CLI argument `--negative-prompt`: diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 360c726fd..2ee09e178 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -89,15 +89,9 @@ def parse_args() -> argparse.Namespace: "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, " "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, " "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, " - "meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, " + "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, " "stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.", ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.", - ) parser.add_argument( "--deploy-config", type=str, @@ -475,8 +469,6 @@ def main(): omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager - if args.stage_configs_path: - omni_kwargs["stage_configs_path"] = args.stage_configs_path if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config if use_nextstep: @@ -515,8 +507,8 @@ def main(): print(f" Image size: {args.width}x{args.height}") if args.lora_path: print(f" LoRA: scale={args.lora_scale}") - if args.stage_configs_path: - print(f" stage-configs-path: {args.stage_configs_path}") + if args.deploy_config: + print(f" deploy-config: {args.deploy_config}") print(f"{'=' * 60}\n") # Build LoRA request when --lora-path is set diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index d246f9e45..92e08fab5 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -16,6 +16,7 @@ list of supported architectures across all modalities, see |---|---|---|---|---|---|---| | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | +| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz | | GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz | | Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz | | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz | @@ -23,6 +24,7 @@ list of supported architectures across all modalities, see | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz | | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz | +| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz | | IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz | | IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz | @@ -154,7 +156,7 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \ --text "Hello, this is a streaming test." \ --streaming ``` -Streaming requires `async_chunk: true` in the stage config. +Streaming requires `async_chunk: true` in the deploy config. ### Notes - Output: 44.1 kHz mono WAV. @@ -369,7 +371,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --streaming \ --output-dir /tmp/out_stream ``` -Streaming requires `async_chunk: true` in the stage config. +Streaming requires `async_chunk: true` in the deploy config. ### Word Timestamps Word-level timestamps are currently a serving-path feature: launch @@ -378,13 +380,13 @@ Word-level timestamps are currently a serving-path feature: launch example is not provided in this release. ### Batched decoding -The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates): +The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, create an overlay based on `vllm_omni/deploy/qwen3_tts.yaml` and pass it with `--deploy-config`: ```bash python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --query-type CustomVoice \ --txt-prompts examples/offline_inference/text_to_speech/qwen3_tts/benchmark_prompts.txt \ --batch-size 4 \ - --stage-configs-path /path/to/qwen3_tts_batched.yaml + --deploy-config /path/to/qwen3_tts_batched.yaml ``` `--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…). @@ -394,6 +396,57 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ --- +## Gepard-1.0 + +Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up). + +### Prerequisites +The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately: + +```bash +uv pip install "nemo_toolkit[tts]==2.7.3" \ + --overrides <(printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0") +``` + +The overrides are what make that command work. `nemo_toolkit[tts]` declares +`transformers~=4.57.0`, so a plain install downgrades the version +`requirements/common.txt` pins and takes `huggingface-hub` back to 0.x with it — after +which `import vllm` fails. Both packages are named here rather than left to resolve on +their own: transformers 5.x happens to require `huggingface-hub` 1.x, but relying on +that would make the second repair an accident of the first. `--constraint` cannot be +used instead, since it narrows a version range rather than overriding a dependency a +package declares, and the two sets never intersect. NeMo 2.7.3 decodes correctly +against transformers 5.x, so nothing is lost by holding it there. + +(The same clash is why this is documented here rather than declared in +`pyproject.toml`: an extra resolves together with the base dependencies, so +`vllm-omni[gepard]` would be unsatisfiable rather than opt-in.) + +On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a +consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's +sampler JIT-compiles through FlashInfer on first use, and where that build fails the engine +dies inside `profile_run` long after every import has succeeded, which reads as a model +failure rather than an environment one. The PyTorch sampler it falls back to costs nothing +here. + +### Quick start (zero-shot, default voice) +```bash +python examples/offline_inference/text_to_speech/gepard/end2end.py \ + --text "Hello, this is a Gepard demo." +``` + +### Voice cloning +Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR. + +### Notes +- Output: 22.05 kHz mono WAV. +- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower. +- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR. +- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it). +- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`. +- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced. +- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed. + ## VoxCPM2 Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`. @@ -430,6 +483,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [ --- +## dots.tts + +Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ + --model rednote-hilab/dots.tts-soar \ + --text "Hello, this is a test of dots TTS running on vLLM Omni." +``` + +### Voice cloning +Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet. + +### Streaming +The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example. + +### Notes +- Output: 48 kHz mono WAV. +- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`). +- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet. +- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap. + +--- + ## IndexTTS-2 and IndexTTS-2.5 Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning. @@ -538,100 +616,3 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T - `--num-prompts N` replicates the prompt for performance measurement. - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`. - Run `--help` for the full argument surface. - ---- - -## SoulX-Singer - -Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`. - -### Prerequisites - -Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it. - -```bash -# 1. DiT weights -export BASE=path/to/SoulX-Singer -export PREPROCESS=path/to/SoulX-Singer-Preprocess -export SVC_DIR=path/to/SoulX-Singer-svc - -huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE" - -# 2. Preprocess weights (required) -huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS" -export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS" - -# 3. SVS / SVC view directories -mkdir -p "$SVC_DIR" -cp $BASE/{config.yaml,README.md,assets} $SVC_DIR -mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt - -cat > "$BASE/config.json" <<'EOF' -{ - "model_type": "soulxsinger", - "architectures": ["SoulXSingerPipeline"], - "max_num_seqs": 1 -} -EOF - -cat > "$SVC_DIR/config.json" <<'EOF' -{ - "model_type": "soulxsinger", - "architectures": ["SoulXSingerSVCPipeline"], - "max_num_seqs": 1 -} -EOF -``` - -`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches. - -**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs: - -```bash -pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC -``` - -Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`: - -```bash -# install optional dependencies: -pip install -e ".[soulx-svs]" -``` - -English SVS adds NeMo ASR and NLTK data; pass `--language English`: - -```bash -pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2 -python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')" -``` - -**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni. - -### Quick start - -```bash -# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3 -python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \ - --model "$BASE" \ - --preprocess-weights-dir "$PREPROCESS" \ - --control score \ - --num-inference-steps 32 \ - -o output.wav - -python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \ - --model "$SVC_DIR" \ - --preprocess-weights-dir "$PREPROCESS" \ - --svc \ - --num-inference-steps 32 \ - -o output_svc.wav -``` - -`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options. - -### Notes - -- Output: 24 kHz mono WAV; batch only. -- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on. -- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only. - ---- diff --git a/examples/offline_inference/text_to_speech/dots_tts/end2end.py b/examples/offline_inference/text_to_speech/dots_tts/end2end.py new file mode 100644 index 000000000..8705f03f3 --- /dev/null +++ b/examples/offline_inference/text_to_speech/dots_tts/end2end.py @@ -0,0 +1,125 @@ +"""Offline dots.tts inference example (single-stage native AR pipeline). + +Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml). +Reference checkpoint: rednote-hilab/dots.tts-soar. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import soundfile as sf +import torch + +from vllm_omni import Omni +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +REPO_ROOT = Path(__file__).resolve().parents[4] +SAMPLE_RATE = 48_000 + + +def parse_args(): + parser = TrackingArgumentParser(description="Offline dots.tts native AR inference") + parser.add_argument( + "--model", + type=str, + default="rednote-hilab/dots.tts-soar", + help="dots.tts model path or HuggingFace repo ID.", + ) + parser.add_argument( + "--text", + type=str, + default="Hello, this is a test of dots TTS running on vLLM Omni.", + help="Text to synthesize.", + ) + parser.add_argument( + "--output-dir", + type=str, + default="output_audio", + help="Directory for output WAV files.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path. If unset, auto-loads " + "vllm_omni/deploy/dots_tts.yaml based on the HF model_type.", + ) + return parser.parse_args() + + +def extract_audio(multimodal_output: dict) -> torch.Tensor: + """Extract the final complete audio tensor from multimodal output. + + The output processor concatenates per-step delta tensors under + ``model_outputs``. Falls back to ``audio`` for backwards compat. + """ + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}") + + if isinstance(audio, list): + valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] + if not valid: + raise ValueError("Audio list is empty or all elements are None.") + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] + + return torch.as_tensor(audio).float().cpu().reshape(-1) + + +def main(): + args = parse_args() + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + engine = Omni( + model=args.model, + deploy_config=args.deploy_config, + ) + + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import ( + build_dots_tts_prompt, + ) + + tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + + prompt = build_dots_tts_prompt( + tokenizer=tokenizer, + text=args.text, + ) + + print(f"Model : {args.model}") + print(f"Text : {args.text}") + print(f"Output dir : {output_dir}") + print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens") + + t_start = time.perf_counter() + outputs = engine.generate([prompt]) + elapsed = time.perf_counter() - t_start + + request_output = outputs[0] + mm = request_output.outputs[0].multimodal_output + audio = extract_audio(mm) + + duration = audio.numel() / SAMPLE_RATE + rtf = elapsed / duration if duration > 0 else float("inf") + + output_path = output_dir / "output.wav" + sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV") + + print(f"Saved : {output_path}") + print(f"Duration : {duration:.2f}s") + print(f"Inference : {elapsed:.2f}s") + print(f"RTF : {rtf:.3f}") + + +if __name__ == "__main__": + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + main() diff --git a/examples/offline_inference/text_to_speech/fish_speech/end2end.py b/examples/offline_inference/text_to_speech/fish_speech/end2end.py index 18b4d79f8..e2e5738fd 100644 --- a/examples/offline_inference/text_to_speech/fish_speech/end2end.py +++ b/examples/offline_inference/text_to_speech/fish_speech/end2end.py @@ -38,17 +38,6 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser logger = logging.getLogger(__name__) DEFAULT_MODEL = "fishaudio/s2-pro" -DEFAULT_STAGE_CONFIG = os.path.join( - os.path.dirname(__file__), - "..", - "..", - "..", - "..", - "vllm_omni", - "model_executor", - "stage_configs", - "fish_speech_s2_pro.yaml", -) def build_prompt( @@ -154,8 +143,6 @@ def _build_omni_kwargs(args, model_name: str) -> dict: ) if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config - else: - omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG return omni_kwargs @@ -278,12 +265,6 @@ def parse_args(): default=None, help="Deploy config YAML path (under vllm_omni/deploy/).", ) - parser.add_argument( - "--stage-configs-path", - type=str, - default=None, - help="Path to stage configs YAML (legacy, prefer --deploy-config).", - ) parser.add_argument( "--output-dir", default="output_audio", diff --git a/examples/offline_inference/text_to_speech/gepard/end2end.py b/examples/offline_inference/text_to_speech/gepard/end2end.py new file mode 100644 index 000000000..2f6b61d89 --- /dev/null +++ b/examples/offline_inference/text_to_speech/gepard/end2end.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni. + +Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ +frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono +waveform. Zero-shot: the default learned voice, no reference audio. + +Generation length and seed are stage settings rather than flags here. One +output token is one audio frame, so ``max_tokens`` in the deploy YAML is the +frame budget; pass a copy of the YAML via ``--deploy-config`` to change it. + +Usage: + python end2end.py --text "Hello, this is Gepard speaking." +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import soundfile as sf +import torch + +import vllm_omni +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + +from vllm_omni import Omni # noqa: E402 + +MODEL = "nineninesix/gepard-1.0" +SAMPLE_RATE = 22050 +# Absolute: the loader hands the path straight to OmegaConf, which resolves a +# bare name against the current working directory, not the package. +DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml") + + +def build_request(text: str) -> dict: + """Build an Omni request payload for Gepard (zero-shot, no ref audio). + + The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so + a bare text prompt would have no speaker slots and no SOS. + """ + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig + from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids + + # From the checkpoint's own sidecar, which is what the worker builds its + # config from too. A layout that disagrees with the trained one does not + # fail, it just degrades the speech. + cfg = GepardConfig.from_checkpoint(MODEL) + tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) + prompt_token_ids = build_gepard_prompt_ids( + tokenizer(text, add_special_tokens=False)["input_ids"], + config=cfg, + ) + return { + "prompt_token_ids": prompt_token_ids, + "additional_information": {"text": [text]}, + } + + +def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None: + sf.write(path, waveform.float().numpy(), sample_rate) + seconds = waveform.numel() / sample_rate + print( + f" Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, " + f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})" + ) + + +def main(args) -> None: + omni = Omni( + model=MODEL, + deploy_config=args.deploy_config, + stage_init_timeout=args.stage_init_timeout, + init_timeout=args.init_timeout, + ) + + # No explicit SamplingParams: one supplied here would replace the stage + # defaults rather than merge over them, dropping the pipeline's stop token. + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"Synthesizing: {args.text!r}") + inputs = build_request(text=args.text) + + for i, stage_outputs in enumerate(omni.generate(inputs)): + for j, out in enumerate(stage_outputs.outputs): + mm = out.multimodal_output + if mm is None: + print(f" [req {i}] No audio output.") + continue + # The consolidation path renames "model_outputs" to "audio". + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if isinstance(audio, list): + audio = audio[0] if len(audio) else None + if audio is None: + print(f" [req {i}] No waveform in multimodal_output (keys: {list(mm)}).") + continue + sr_tensor = mm.get("sr") + sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE + save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr) + + print("Done.") + + +def parse_args(): + parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)") + parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.") + parser.add_argument( + "--output-dir", + default=os.path.join( + os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")), + "gepard_output", + ), + help="Directory for WAV outputs (default: ~/.cache/gepard_output).", + ) + parser.add_argument( + "--deploy-config", + default=DEFAULT_DEPLOY_CONFIG, + help=( + "Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must " + "stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes " + "Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed." + ), + ) + # The upstream defaults are too tight for a cold start. + parser.add_argument("--stage-init-timeout", type=int, default=900) + parser.add_argument("--init-timeout", type=int, default=1800) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/offline_inference/text_to_speech/glm_tts/end2end.py b/examples/offline_inference/text_to_speech/glm_tts/end2end.py index 6624c2b35..3f66ed8c3 100644 --- a/examples/offline_inference/text_to_speech/glm_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/glm_tts/end2end.py @@ -145,7 +145,7 @@ def main(args): omni = Omni( model=args.model, - stage_configs_path=deploy_config_path, + deploy_config=deploy_config_path, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, ) diff --git a/examples/offline_inference/text_to_speech/ming_tts/README.md b/examples/offline_inference/text_to_speech/ming_tts/README.md index 3fdfaf4d5..87f2df1f5 100644 --- a/examples/offline_inference/text_to_speech/ming_tts/README.md +++ b/examples/offline_inference/text_to_speech/ming_tts/README.md @@ -131,7 +131,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \ --enforce-eager ``` -`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently +`--streaming` uses `AsyncOmni` and the async chunk deploy config. It currently supports one prompt per process invocation; use blocking mode for `--num-prompts > 1`. diff --git a/examples/offline_inference/text_to_speech/ming_tts/end2end.py b/examples/offline_inference/text_to_speech/ming_tts/end2end.py index 1a6991e40..89d5fda73 100644 --- a/examples/offline_inference/text_to_speech/ming_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/ming_tts/end2end.py @@ -13,6 +13,7 @@ import yaml from transformers import AutoTokenizer from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( KEY_CFG, KEY_MAX_DECODE_STEPS, @@ -21,7 +22,6 @@ from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( SAMPLE_RATE, ) from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt -from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor try: from .runner import ( @@ -142,7 +142,8 @@ def _resolve_speaker_embedding(args, case: dict, paths: list[str]): return _load_speaker_embedding(args.speaker_embedding) if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths: return None - embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths) + # Runs before the engine is built, so nothing has pre-fetched campplus.onnx. + embs = SpeakerEmbeddingExtractor(args.model, allow_download=True).extract_many(paths) if not embs: raise RuntimeError("Speaker extraction produced no embeddings") return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0) diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py index 7c0ce125e..75f562af6 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py @@ -490,10 +490,10 @@ def parse_args(): help="Path to a .txt file with one prompt per line (preferred).", ) parser.add_argument( - "--stage-configs-path", + "--deploy-config", type=str, default=None, - help="Path to a stage configs file.", + help="Path to a deploy config YAML.", ) parser.add_argument( "--audio-path", diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py index 62f5e2caa..7b282c560 100644 --- a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py +++ b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py @@ -32,11 +32,10 @@ from end2end import _estimate_prompt_len from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni import Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config -def _default_stage_config() -> str: +def _default_deploy_config() -> str: repo_root = Path(__file__).resolve().parents[4] return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml") @@ -109,7 +108,7 @@ def main(args: Any) -> None: "output_dir", ): omni_kwargs.pop(key, None) - omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config() + omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config() omni_kwargs["log_stats"] = args.log_stats omni = Omni(**omni_kwargs) @@ -156,7 +155,7 @@ def parse_args() -> Any: default=None, help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)", ) - parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML") + parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML") parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align") parser.add_argument("--language", default="English", help="Qwen3-TTS language field") parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name") @@ -164,7 +163,6 @@ def parse_args() -> Any: parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens") parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar") parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging") - nullify_stage_engine_defaults(parser) return parser.parse_args() diff --git a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py deleted file mode 100644 index 14dd5e145..000000000 --- a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings). - -Usage: - python benchmark.py --model /path/to/SoulX-Singer --svs \\ - --prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\ - --target-audio ../SoulX-Singer/example/audio/music.mp3 \\ - --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ - -o benchmark.wav - - python benchmark.py ... --enable-diffusion-pipeline-profiler -""" - -import argparse -import statistics -import time -from pathlib import Path - -import soundfile as sf -from end2end import ( - SVC_DEPLOY_CONFIG, - SVS_DEPLOY_CONFIG, - add_inference_args, - build_sampling, - extract_audio, - resolve_preprocess_weights_dir, -) - -from vllm_omni.entrypoints.omni import Omni - - -def _audio_duration_sec(outputs) -> float: - audio_np, sr = extract_audio(outputs) - return audio_np.size / sr - - -def _stage_durations_from_output(outputs) -> dict[str, float]: - omni_out = outputs[0] - durations = getattr(omni_out, "stage_durations", None) or {} - if durations: - return dict(durations) - ro = omni_out - if ro is not None: - return dict(getattr(ro, "stage_durations", None) or {}) - return {} - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark") - add_inference_args(parser) - parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)") - parser.add_argument("--runs", type=int, default=3, help="Measured runs") - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable low-overhead stage timing in DiffusionOutput.stage_durations", - ) - parser.add_argument( - "--enforce-eager", - action="store_true", - help="Disable torch.compile on diff_estimator (eager baseline for A/B)", - ) - parser.add_argument( - "--dtype", - type=str, - default=None, - choices=["bfloat16", "bf16", "float16", "fp16", "half"], - help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.", - ) - parser.add_argument( - "--output", - "-o", - type=str, - default=None, - help="Save WAV from the last measured run (after warmup)", - ) - parser.add_argument( - "--save-each-run", - action="store_true", - help="With --output, also write measured runs as _run1.wav, _run2.wav, ...", - ) - return parser.parse_args() - - -def main() -> None: - args = parse_args() - preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir) - deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG) - mode = "SVS" if args.svs else "SVC" - - omni_kwargs: dict = { - "model": args.model, - "deploy_config": deploy_config, - "async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed) - } - if args.enable_diffusion_pipeline_profiler: - omni_kwargs["enable_diffusion_pipeline_profiler"] = True - if args.enforce_eager: - omni_kwargs["enforce_eager"] = True - if args.dtype is not None: - omni_kwargs["dtype"] = args.dtype - - compile_mode = "eager" if args.enforce_eager else "torch.compile" - dtype_label = args.dtype or "from deploy YAML" - print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]") - - omni = Omni(**omni_kwargs) - kind = "svs" if args.svs else "svc" - sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind) - prompt = {"prompt_token_ids": [0]} - - latencies_ms: list[float] = [] - rtfs: list[float] = [] - last_stages: dict[str, float] = {} - last_measured_outputs = None - - total_runs = args.warmup + args.runs - for run_idx in range(total_runs): - is_warmup = run_idx < args.warmup - label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}" - t0 = time.perf_counter() - outputs = list(omni.generate([prompt], sampling)) - elapsed_ms = (time.perf_counter() - t0) * 1000.0 - audio_sec = _audio_duration_sec(outputs) - rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf") - print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}") - if not is_warmup: - latencies_ms.append(elapsed_ms) - rtfs.append(rtf) - last_stages = _stage_durations_from_output(outputs) - last_measured_outputs = outputs - if args.output and args.save_each_run: - measured_idx = run_idx - args.warmup + 1 - out_path = Path(args.output) - run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}") - audio_np, sr = extract_audio(outputs) - sf.write(str(run_path), audio_np, sr) - print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)") - - omni.close() - - if args.output and last_measured_outputs is not None: - audio_np, sr = extract_audio(last_measured_outputs) - out_path = Path(args.output) - out_path.parent.mkdir(parents=True, exist_ok=True) - sf.write(str(out_path), audio_np, sr) - print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)") - - if latencies_ms: - print("\n=== Summary (measured runs) ===") - print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}") - print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}") - if last_stages: - print("\nStage durations (last measured run):") - for name, value in sorted(last_stages.items()): - if name.endswith("_ms"): - print(f" {name}: {value:.1f} ms") - else: - print(f" {name}: {value * 1000.0:.1f} ms") - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py deleted file mode 100644 index a1bb112af..000000000 --- a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py +++ /dev/null @@ -1,285 +0,0 @@ -"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline). - -Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference. - -Usage: - # SVS (default) - python end2end.py --model /path/to/SoulX-Singer \\ - --prompt-audio path/to/prompt/audio.mp3 \\ - --target-audio path/to/target/audio.mp3 \\ - --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ - -o output.wav - - # SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view) - python end2end.py --model /path/to/SoulX-Singer-svc \\ - --prompt-audio path/to/prompt/audio.mp3 \\ - --target-audio path/to/target/audio.mp3 \\ - --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ - -o output_svc.wav -""" - -import argparse -import os -from pathlib import Path - -import numpy as np -import soundfile as sf - -from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -REPO_ROOT = Path(__file__).resolve().parents[4] -SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml" -SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml" -DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3" -DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3" -_SAMPLE_RATE = 24000 - - -def _require_paths(paths: dict[str, str | None]) -> None: - missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)] - if missing: - raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing)) - - -def resolve_preprocess_weights_dir(explicit: str | None) -> str: - if explicit: - path = Path(explicit).expanduser().resolve() - if not (path / "rmvpe" / "rmvpe.pt").is_file(): - raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}") - return str(path) - env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR") - if env: - path = Path(env).expanduser().resolve() - if (path / "rmvpe" / "rmvpe.pt").is_file(): - return str(path) - raise FileNotFoundError( - "SoulX preprocess weights not found. Pass --preprocess-weights-dir or set " - "SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt " - "(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)." - ) - - -def build_sampling( - args: argparse.Namespace, - *, - preprocess_weights_dir: str, - kind: str, -) -> OmniDiffusionSamplingParams: - extra: dict = { - "prompt_audio": os.path.abspath(args.prompt_audio), - "target_audio": os.path.abspath(args.target_audio), - "preprocess_weights_dir": preprocess_weights_dir, - "vocal_sep": args.vocal_sep, - "auto_shift": args.auto_shift, - "pitch_shift": args.pitch_shift, - } - if kind == "svs": - extra["language"] = args.language - extra["control"] = args.control - if args.prompt_metadata_path: - extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path) - if args.target_metadata_path: - extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path) - if args.prompt_metadata_path or args.target_metadata_path: - extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio) - - extra = validate_soulx_extra_args(kind, extra) - - return OmniDiffusionSamplingParams( - num_inference_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - seed=args.seed, - extra_args=extra, - ) - - -def extract_audio(outputs) -> tuple[np.ndarray, int]: - out = outputs[0] - mm = getattr(out, "multimodal_output", None) - if mm is None: - ro = out - if ro is not None: - mm = getattr(ro, "multimodal_output", None) - if mm is None and ro.outputs: - mm = getattr(ro.outputs[0], "multimodal_output", None) - if not mm or "audio" not in mm: - raise RuntimeError("No audio in multimodal_output") - - audio_val = mm["audio"] - if isinstance(audio_val, list): - chunks = [ - np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1) - for chunk in audio_val - if chunk is not None - ] - audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32) - elif hasattr(audio_val, "cpu"): - audio_np = audio_val.detach().cpu().numpy().squeeze() - else: - audio_np = np.asarray(audio_val, dtype=np.float32).squeeze() - - sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE)) - if isinstance(sr_val, list) and sr_val: - sr_val = sr_val[-1] - if hasattr(sr_val, "item"): - sr_val = sr_val.item() - return audio_np.astype(np.float32, copy=False), int(sr_val) - - -def add_inference_args(parser: argparse.ArgumentParser) -> None: - mode = parser.add_mutually_exclusive_group() - mode.add_argument( - "--svc", - action="store_true", - help="Assert SVC mode (must match config.json architectures)", - ) - parser.add_argument( - "--model", - type=str, - required=True, - help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)", - ) - parser.add_argument( - "--deploy-config", - type=str, - default=None, - help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)", - ) - parser.add_argument( - "--prompt-audio", - type=str, - default=str(DEFAULT_PROMPT_AUDIO), - help="Prompt / reference audio (mp3 or wav)", - ) - parser.add_argument( - "--target-audio", - type=str, - default=str(DEFAULT_TARGET_AUDIO), - help="Target content audio (mp3 or wav)", - ) - parser.add_argument( - "--prompt-metadata-path", - type=str, - default=None, - help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)", - ) - parser.add_argument( - "--target-metadata-path", - type=str, - default=None, - help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)", - ) - parser.add_argument( - "--prompt-wav-path", - type=str, - default=None, - help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)", - ) - parser.add_argument( - "--preprocess-weights-dir", - type=str, - default=None, - help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)", - ) - parser.add_argument( - "--language", - type=str, - default="Mandarin", - help="SVS lyrics language (Mandarin, Cantonese, English)", - ) - parser.add_argument( - "--control", - type=str, - default="melody", - choices=["score", "melody"], - help="SVS control mode", - ) - parser.add_argument( - "--vocal-sep", - action=argparse.BooleanOptionalAction, - default=False, - help="Run vocal separation during preprocess", - ) - parser.add_argument( - "--auto-shift", - action=argparse.BooleanOptionalAction, - default=True, - help="Auto pitch shift to align target pitch with prompt", - ) - parser.add_argument("--pitch-shift", type=int, default=0) - parser.add_argument("--num-inference-steps", type=int, default=32) - parser.add_argument("--guidance-scale", type=float, default=3.0) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Optional CFM RNG seed. Omit for non-deterministic sampling.", - ) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)") - add_inference_args(parser) - parser.add_argument("--output", "-o", type=str, default="output.wav") - args = parser.parse_args() - return args - - -def main() -> None: - args = parse_args() - kind = resolve_soulx_kind(args.model) - - if not (args.svc ^ (kind == "svs")): - raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}") - - required = { - "prompt_audio": args.prompt_audio, - "target_audio": args.target_audio, - } - if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path): - raise ValueError( - "SVS precomputed metadata requires both --prompt-metadata-path and " - "--target-metadata-path (or neither, for online preprocess)." - ) - if args.prompt_metadata_path: - required["prompt_metadata_path"] = args.prompt_metadata_path - required["target_metadata_path"] = args.target_metadata_path - if args.prompt_wav_path: - required["prompt_wav_path"] = args.prompt_wav_path - _require_paths(required) - preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir) - deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG) - mode = kind.upper() - - print(f"Loading SoulX-Singer {mode} from {args.model}") - print(f" deploy: {deploy_config}") - print(f" prompt: {args.prompt_audio}") - print(f" target: {args.target_audio}") - - omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only - sampling = build_sampling( - args, - preprocess_weights_dir=preprocess_dir, - kind=kind, - ) - prompt = {"prompt_token_ids": [0], "multi_modal_data": {}} - - print("Running generate (inline preprocess + DiT)...") - outputs = list(omni.generate([prompt], sampling)) - if not outputs: - raise RuntimeError("No output from omni.generate") - - err = getattr(outputs[0], "error", None) - if err is not None: - raise RuntimeError(f"Pipeline error: {err}") - - audio_np, sr = extract_audio(outputs) - sf.write(args.output, audio_np, sr) - print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)") - omni.close() - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py index c9225c212..e67e9c350 100644 --- a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py @@ -304,7 +304,7 @@ def parse_args() -> Namespace: "--cfg-alpha", type=float, default=None, - help="CFG alpha for flow-matching guidance (default: use value from stage config, typically 1.2).", + help="CFG alpha for flow-matching guidance (default: use the deploy config value, typically 1.2).", ) parser.add_argument( "--quantization", diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 339c7dac9..0cb75a8d4 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -14,6 +14,8 @@ A unified script for text-to-video generation. Supports multiple models with mod | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) | | `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — | +| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE wrapper | +| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | BF16 DiT + LTX-2 Video VAE wrapper | ## Local CLI Usage @@ -151,6 +153,40 @@ python text_to_video.py \ --output quick_test.mp4 ``` +### SANA-Video-2B + +Select the native T2V pipeline explicitly: + +```bash +python text_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaVideoPipeline \ + --prompt "A cinematic tracking shot of a sailboat crossing the ocean at sunset." \ + --height 480 \ + --width 832 \ + --num-frames 81 \ + --num-inference-steps 50 \ + --guidance-scale 6.0 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 \ + --output sana_video_480p.mp4 +``` + +For the 720p checkpoint, switch the model to +`Efficient-Large-Model/SANA-Video_2B_720p_diffusers`, use +`--height 704 --width 1280`, and write to a different output path. Both native +checkpoint variants are supported. The Diffusers adapter T2V compatibility +path is also validated at both resolutions through online serving. + +The native 480p path uses vLLM-Omni's `DistributedAutoencoderKLWan` wrapper; +the native 720p path uses `DistributedAutoencoderKLLTX2Video`. Both wrappers +retain their underlying Diffusers autoencoders. The pipeline intentionally +uses Diffusers' checkpoint-compatible `DPMSolverMultistepScheduler`. + +For image-to-video, use the shared +[`image_to_video.py`](../image_to_video/README.md#sana-video-2b) example with +`--model-class-name SanaImageToVideoPipeline` and `--image `. + ### Cosmos3 ```bash diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 129ca860a..02383c6e4 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import argparse import json @@ -10,7 +10,6 @@ from typing import Any import numpy as np import torch -from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -115,12 +114,32 @@ _MODEL_PRESETS = { "fps": 24, "output": "ltx23_output.mp4", }, + "sana_480p": { + "height": 480, + "width": 832, + "num_frames": 81, + "num_inference_steps": 50, + "guidance_scale": 6.0, + "fps": 16, + "output": "sana_video_480p.mp4", + }, + "sana_720p": { + "height": 704, + "width": 1280, + "num_frames": 81, + "num_inference_steps": 50, + "guidance_scale": 6.0, + "fps": 16, + "output": "sana_video_720p.mp4", + }, } def _detect_preset(model: str, model_class_name: str | None = None) -> dict: model_lower = model.lower() class_lower = (model_class_name or "").lower() + if "sana-video" in model_lower or "sana_video" in model_lower or "sanavideo" in class_lower: + return _MODEL_PRESETS["sana_720p" if "720p" in model_lower else "sana_480p"] if "lingbot" in model_lower or "lingbotvideo" in class_lower: return _MODEL_PRESETS["lingbot"] if "ltx" in class_lower or "ltx" in model_lower: @@ -155,6 +174,13 @@ def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict return result +def _validate_video_output_type(output_type: str) -> None: + if output_type not in {"image", "video"}: + raise ValueError( + f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation." + ) + + def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor: """Normalize decoded frames according to the pipeline's output contract.""" if not tensor.is_floating_point(): @@ -462,17 +488,6 @@ def main(): "scm_steps_policy": "dynamic", } - # Configure parallel settings - parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, - ring_degree=args.ring_degree, - cfg_parallel_size=args.cfg_parallel_size, - tensor_parallel_size=args.tensor_parallel_size, - vae_patch_parallel_size=args.vae_patch_parallel_size, - pipeline_parallel_size=args.pipeline_parallel_size, - enable_expert_parallel=args.enable_expert_parallel, - ) - profiler_enabled = args.profiler_config is not None omni_kwargs = dict( @@ -481,7 +496,13 @@ def main(): vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, enable_cpu_offload=args.enable_cpu_offload, - parallel_config=parallel_config, + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, + pipeline_parallel_size=args.pipeline_parallel_size, + enable_expert_parallel=args.enable_expert_parallel, enforce_eager=args.enforce_eager, model_class_name=model_class_name, cache_backend=args.cache_backend, @@ -570,6 +591,8 @@ def main(): num_frames=args.num_frames, extra_args=extra_args, ) + if args.frame_rate is not None: + sampling_kwargs["frame_rate"] = args.frame_rate if args.guidance_scale_high is not None: sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high @@ -604,10 +627,7 @@ def main(): frames = frames[0] if frames else None if isinstance(frames, OmniRequestOutput): - if frames.final_output_type != "image": - raise ValueError( - f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation." - ) + _validate_video_output_type(frames.final_output_type) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) diff --git a/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py b/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py deleted file mode 100644 index 2f66d5f77..000000000 --- a/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py +++ /dev/null @@ -1,108 +0,0 @@ -import argparse -import fcntl -import os -import site -import subprocess -import tempfile -import time -from pathlib import Path - -from huggingface_hub import snapshot_download - -DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git" -DEPENDENCY_BRANCH = "omni" -CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency" -LOCK_FILE = CACHE_DIR / ".install.lock" -DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni" - - -def download_dependency(): - CACHE_DIR.mkdir(parents=True, exist_ok=True) - - with open(LOCK_FILE, "w") as f: - fcntl.flock(f, fcntl.LOCK_EX) - if not DEPENDENCY_DIR.exists(): - print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...") - subprocess.run( - ["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)], - check=True, - ) - print("Download finished.") - fcntl.flock(f, fcntl.LOCK_UN) - - # write .pth to site-packages - site_packages = Path(site.getsitepackages()[0]) - pth_file = site_packages / "vllm_omni_dependency.pth" - pth_file.write_text(str(DEPENDENCY_DIR)) - print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}") - - -def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None): - """Download files from HF repo and log time + destination.""" - if os.path.exists(local_dir): - print(f"Directory {local_dir} already exists. Skipping download.") - return - print(f"Starting download from {repo_id} into {local_dir}") - start_time = time.time() - - snapshot_download( - repo_id=repo_id, - local_dir=local_dir, - local_dir_use_symlinks=False, - allow_patterns=allow_patterns, - ) - - elapsed = time.time() - start_time - print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}") - - -def main(output_dir: str): - # Wan2.2 - wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B") - timed_download( - repo_id="Wan-AI/Wan2.2-TI2V-5B", - local_dir=wan_dir, - allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"], - ) - - # MMAudio - mm_audio_dir = os.path.join(output_dir, "MMAudio") - timed_download( - repo_id="hkchengrex/MMAudio", - local_dir=mm_audio_dir, - allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"], - ) - - dreamid_dir = os.path.join(output_dir, "DreamID-Omni") - - timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir) - - # Now we construct the config file - import json - - data = { - "_class_name": "DreamIDOmniPipeline", - } - - with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f: - json.dump(data, f, indent=2) - - print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}") - - transformer_dir = os.path.join(output_dir, "transformer") - os.makedirs(transformer_dir, exist_ok=True) - with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f: - json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f) - print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}") - - # now we download the dependency code - download_dependency() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Download models from Hugging Face") - parser.add_argument( - "--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models" - ) - args = parser.parse_args() - main(args.output_dir) diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md deleted file mode 100644 index 8305c785f..000000000 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ /dev/null @@ -1,163 +0,0 @@ -# X-To-Video-Audio - -The `DreamID-Omni` pipeline generates short videos from text, image and video. - -## Local CLI Usage -### Download the Model locally -Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally. - -```bash -python download_dreamid_omni.py --output-dir ./dreamid_omni -``` -After download, the model directory will look like this: - -``` -dreamid_omni/ -├── DreamID-Omni/ -│ ├── dreamid_omni.safetensors -├── MMAudio/ -│ ├── ext_weights/ -│ │ ├── best_netG.pt -│ │ ├── v1-16.pth -├── Wan2.2-TI2V-5B/ -│ ├── google/* -│ ├── models_t5_umt5-xxl-enc-bf16.pth -│ ├── Wan2.2_VAE.pth -│ -├── model_index.json -└── transformer/ - └── config.json # create by download_dreamid_omni.py -``` - -### Run the Inference -```python -python x_to_video_audio.py \ - --model /path/to/dreamid_omni \ - --prompt "Two people walking together and singing happily" \ - --image-path ./example0.png ./example1.png \ - --audio-path ./example0.wav ./example1.wav \ - --video-negative-prompt "jitter, bad hands, blur, distortion" \ - --audio-negative-prompt "robotic, muffled, echo, distorted" \ - --cfg-parallel-size 2 \ - --num-inference-steps 45 \ - --height 704 \ - --width 1280 \ - --output out_dreamid_omni_twoip.mp4 -``` -In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled. -The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload. -For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP: - -```python -python x_to_video_audio.py \ - --model /path/to/dreamid_omni \ - --prompt "Two people walking together and singing happily" \ - --image-path ./example0.png ./example1.png \ - --audio-path ./example0.wav ./example1.wav \ - --video-negative-prompt "jitter, bad hands, blur, distortion" \ - --audio-negative-prompt "robotic, muffled, echo, distorted" \ - --cfg-parallel-size 2 \ - --num-inference-steps 45 \ - --height 704 \ - --width 1280 \ - --use-hsdp \ - --hsdp-shard-size 2 \ - --output out_dreamid_omni_twoip.mp4 -``` - - -You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni - -For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing. - -```python -# Example usage for oneip, ref media from the official repo DreamID-Omni -python x_to_video_audio.py \ - --model /path/to/dreamid_omni \ - --prompt ": In the frame, a woman with black long hair is identified as .\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ - --image-path 9.png \ - --audio-path 9.wav \ - --video-negative-prompt "jitter, bad hands, blur, distortion" \ - --audio-negative-prompt "robotic, muffled, echo, distorted" \ - --cfg-parallel-size 2 \ - --num-inference-steps 45 \ - --height 704 \ - --width 1280 \ - --output out_dreamid_omni_oneip.mp4 -``` - - -## MagiHuman (`--model-type magi-human`) - -MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param -T5Gemma text encoder. A detailed text prompt is the only required input; an optional -image and/or audio file may be supplied for conditioning. Natively supports Tensor -Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard -the MoE weights and the text encoder. - -> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct -> attention-kernel behaviour (the pipeline otherwise falls back to stubs). - -Text-only generation: - -```bash -python x_to_video_audio.py \ - --model-type magi-human \ - --model /path/to/daVinci-MagiHuman \ - --prompt "A young woman with long, wavy golden blonde hair... " \ - --tensor-parallel-size 4 \ - --height 256 --width 448 \ - --num-inference-steps 8 \ - --seed 52 \ - --extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \ - --output output_magihuman.mp4 -``` - -With optional image and audio conditioning: - -```bash -python x_to_video_audio.py \ - --model-type magi-human \ - --model /path/to/daVinci-MagiHuman \ - --prompt "A young woman..." \ - --tensor-parallel-size 4 \ - --height 256 --width 448 \ - --num-inference-steps 8 \ - --seed 52 \ - --extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \ - --output output_magihuman.mp4 -``` - -MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared -in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`): - -- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set, - because the audio length then determines the number of frames). -- `audio_path`: path to an audio file for audio-to-video conditioning. When provided, - the audio drives the frame count and is encoded as a latent condition into the DiT - (`is_a2v` mode). Omit for pure text-to-video+audio generation. -- `image_path`: path to an image file for visual conditioning. Applied at both the - base-resolution (BR) and super-resolution (SR) stages when SR is enabled. -- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped - when these are omitted. -- `sr_num_inference_steps`: denoising steps for the SR stage. - -## DreamID-Omni arguments - -Key arguments: -- `--prompt`: text description (string). -- `--model`: path to the model local directory. -- `--height/--width`: output resolution (defaults 704 * 1024). -- `--image-path`: path to the input image list. -- `--audio-path`: path to the input audio list, indicate the timbre of the output video. -- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1). -- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks. -- `--hsdp-shard-size`: number of GPUs used for HSDP sharding. -- `--hsdp-replicate-size`: number of HSDP replica groups. -- `--num-inference-steps`: number of denoising steps (defaults 45). -- `--video-negative-prompt`: negative prompt for video generation. -- `--audio-negative-prompt`: negative prompt for audio generation. -- `--enable-cpu-offload`: enable CPU offload (defaults False). -- `--cache-backend`: enable `cache_dit` for acceleration. -- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`. - (VAEs, the T5 text encoder, norms and modulation stay bf16.) diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py deleted file mode 100644 index 9512ac65c..000000000 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ /dev/null @@ -1,284 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import argparse -import json -import re -import time -from typing import Any - -import numpy as np -from PIL import Image -from vllm.multimodal.media.audio import load_audio - -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"} -AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"} - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.") - parser.add_argument("--model", required=True, help="Model ckpt root directory.") - parser.add_argument( - "--model-type", - default="dreamid-omni", - choices=["dreamid-omni", "magi-human"], - help="Model type.", - ) - parser.add_argument("--prompt", default=None, help="Text prompt.") - - parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path") - parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path") - parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.") - parser.add_argument( - "--extra-body", - type=str, - default=None, - help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), " - 'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.', - ) - - parser.add_argument("--height", type=int, default=704, help="Video height.") - parser.add_argument("--width", type=int, default=1280, help="Video width.") - parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.") - parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.") - parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.") - parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.") - parser.add_argument( - "--cfg-parallel-size", - type=int, - default=1, - choices=[1, 2, 3, 4], - help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).", - ) - parser.add_argument( - "--tensor-parallel-size", - type=int, - default=1, - help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."), - ) - parser.add_argument( - "--video-negative-prompt", - default="jitter, bad hands, blur, distortion", - help="Negative prompt for video.", - ) - parser.add_argument( - "--audio-negative-prompt", - default="robotic, muffled, echo, distorted", - help="Negative prompt for audio.", - ) - parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.") - parser.add_argument( - "--quantization", - type=str, - default=None, - choices=["fp8", "int8"], - help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "), - ) - parser.add_argument( - "--enable-cpu-offload", - action="store_true", - default=False, - help="Enable CPU offloading for diffusion models.", - ) - parser.add_argument( - "--enable-layerwise-offload", - action="store_true", - help="Enable layerwise (blockwise) offloading on DiT modules.", - ) - parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.") - parser.add_argument( - "--use-hsdp", - action="store_true", - help="Enable HSDP for DreamID-Omni fused transformer blocks.", - ) - parser.add_argument( - "--hsdp-shard-size", - type=int, - default=1, - help="Number of GPUs used for HSDP sharding when HSDP is enabled.", - ) - parser.add_argument( - "--hsdp-replicate-size", - type=int, - default=1, - help="Number of HSDP replica groups. Default 1 means pure sharding.", - ) - return parser.parse_args() - - -def load_image_and_audio(image_paths, audio_paths): - image = [] - audio = [] - - for path in image_paths: - with Image.open(path) as img: - img = img.convert("RGB") - image.append(img) - - for path in audio_paths: - audio_array, sr = load_audio(path, sr=16000) - audio_array = audio_array[int(sr * 1) : int(sr * 3)] - audio.append(audio_array) - return image, audio - - -def _extract_peak_memory_mb(result: Any) -> float: - """Pull worker-reported peak VRAM (MiB) generation result. - - Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb. - """ - if isinstance(result, list): - result = result[0] if result else None - if result is None: - return 0.0 - val = getattr(result, "peak_memory_mb", 0.0) - if not val: - inner = result - if isinstance(inner, list): - inner = inner[0] if inner else None - val = getattr(inner, "peak_memory_mb", 0.0) - try: - return float(val or 0.0) - except (TypeError, ValueError): - return 0.0 - - -def main() -> None: - args = parse_args() - if args.prompt is None and args.prompt_file is None: - raise ValueError("Either --prompt or --prompt-file must be provided.") - - text_prompt = args.prompt - if args.prompt_file: - with open(args.prompt_file) as f: - text_prompt = json.load(f) - text_prompt = re.sub( - r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL - ).strip() - text_prompt = re.sub( - r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL - ).strip() - text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt) - text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip() - - if args.model_type == "magi-human": - prompt = text_prompt - extra_args = json.loads(args.extra_body) if args.extra_body else {} - sampling_params = OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - num_inference_steps=args.num_inference_steps, - seed=args.seed, - extra_args=extra_args, - ) - else: - image, audio = load_image_and_audio(args.image_path, args.audio_path) - - prompt = { - "prompt": text_prompt, - "video_negative_prompt": args.video_negative_prompt, - "audio_negative_prompt": args.audio_negative_prompt, - "multi_modal_data": {"image": image, "audio": audio}, - } - - sampling_params = OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - num_inference_steps=args.num_inference_steps, - seed=args.seed, - extra_args={ - "solver_name": args.solver_name, - "shift": args.shift, - }, - ) - - parallel_config = DiffusionParallelConfig( - cfg_parallel_size=args.cfg_parallel_size, - tensor_parallel_size=args.tensor_parallel_size, - use_hsdp=args.use_hsdp, - hsdp_shard_size=args.hsdp_shard_size, - hsdp_replicate_size=args.hsdp_replicate_size, - ) - - cache_config = None - if args.cache_backend == "cache_dit": - cache_config = { - "Fn_compute_blocks": 1, - "Bn_compute_blocks": 0, - "max_warmup_steps": 4, - "max_cached_steps": 20, - "residual_diff_threshold": 0.24, - "max_continuous_cached_steps": 3, - "enable_taylorseer": False, - "taylorseer_order": 1, - "scm_steps_mask_policy": None, - "scm_steps_policy": "dynamic", - } - - omni_kwargs = dict( - model=args.model, - parallel_config=parallel_config, - model_type=args.model_type, - enable_cpu_offload=args.enable_cpu_offload, - enable_layerwise_offload=args.enable_layerwise_offload, - cache_backend=args.cache_backend, - cache_config=cache_config, - ) - if args.quantization is not None: - omni_kwargs["quantization"] = args.quantization - omni = Omni(**omni_kwargs) - start = time.perf_counter() - outputs = omni.generate(prompt, sampling_params) - elapsed = time.perf_counter() - start - - if not outputs: - raise RuntimeError("No output returned from the model.") - result = outputs[0] - if not result.images: - raise RuntimeError("No video frames found in model output.") - generated_video = result.images[0] - mm = result.multimodal_output or {} - generated_audio = mm.get("audio") - - if args.model_type == "magi-human": - # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux. - fps = float(mm.get("fps", 25)) - sample_rate = int(mm.get("audio_sample_rate", 24000)) - frames = generated_video - else: - fps = float(mm.get("fps", 24)) - sample_rate = int(mm.get("audio_sample_rate", 16000)) - # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1]. - # mux_video_audio_bytes expects (F, H, W, C) uint8. - if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4: - raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}") - frames = generated_video.transpose(1, 2, 3, 0) - frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8) - - audio_np = None - if generated_audio is not None: - audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32) - - output_path = args.output - video_bytes = mux_video_audio_bytes( - frames, - audio_np, - fps=float(fps), - audio_sample_rate=sample_rate, - ) - with open(output_path, "wb") as f: - f.write(video_bytes) - print(f"Saved generated video to {output_path}") - print(f"Total time: {elapsed:.2f}s") - peak_mb = _extract_peak_memory_mb(outputs) - if peak_mb: - print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/audiox/README.md b/examples/online_serving/audiox/README.md deleted file mode 100644 index bf068b3eb..000000000 --- a/examples/online_serving/audiox/README.md +++ /dev/null @@ -1,64 +0,0 @@ -# AudioX online serving - -Launches the `AudioXPipeline` behind vLLM-Omni's OpenAI-compatible chat endpoint and provides a -minimal Python client that covers all six tasks (`t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`). - -## Start the server - -```bash -cd examples/online_serving/audiox -bash run_server.sh # defaults: MODEL=zhangj1an/AudioX, PORT=8099 -``` - -Environment overrides: `MODEL`, `PORT`, `DIFFUSION_ATTENTION_BACKEND`. - -## Call from Python - -```bash -# text-to-audio -python openai_chat_client.py --task t2a \ - --prompt "Fireworks burst twice, followed by a period of silence before a clock begins ticking." \ - --output t2a.wav - -# text-to-music -python openai_chat_client.py --task t2m \ - --prompt "Uplifting ukulele tune for a travel vlog" \ - --output t2m.wav - -# video-to-audio (no text) -python openai_chat_client.py --task v2a --video path/to/clip.mp4 --output v2a.wav - -# text+video-to-audio -python openai_chat_client.py --task tv2a \ - --prompt "drum beating sound and human talking" \ - --video path/to/clip.mp4 \ - --output tv2a.wav -``` - -The client sends: - -- `num_inference_steps`, `guidance_scale`, `seed`, `audiox_task`, `seconds_start`, - `seconds_total`, `sigma_min`, `sigma_max` as flattened keys under `extra_body`. The server - routes the declared knobs (see `vllm_omni/model_extras/audiox.py`) into the pipeline's - `sampling_params.extra_args`. -- For `v2*` / `tv2*` tasks, the video as a `video_url` content item (data URI for local files) - -## curl - -```bash -curl -sS -X POST http://localhost:8099/v1/chat/completions \ - -H 'Content-Type: application/json' \ - -d '{ - "model": "zhangj1an/AudioX", - "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}], - "extra_body": { - "num_inference_steps": 250, - "guidance_scale": 7.0, - "seed": 42, - "audiox_task": "t2m", - "seconds_total": 10.0, - "sigma_min": 0.3, - "sigma_max": 500.0 - } - }' > t2m.json -``` diff --git a/examples/online_serving/audiox/openai_chat_client.py b/examples/online_serving/audiox/openai_chat_client.py deleted file mode 100755 index fd9aa8cc5..000000000 --- a/examples/online_serving/audiox/openai_chat_client.py +++ /dev/null @@ -1,119 +0,0 @@ -#!/usr/bin/env python3 -"""AudioX OpenAI-compatible chat client. - -AudioX supports 6 tasks (t2a, t2m, v2a, v2m, tv2a, tv2m). Text-only tasks send the prompt as the -chat message; video-conditioned tasks additionally attach the video via a ``video_url`` content -item (data URI for local files). Task + generation knobs (steps, cfg, sigma range, seconds, seed) -are sent as flattened keys under ``extra_body``; the server routes the declared knobs into -``sampling_params.extra_args`` (see ``vllm_omni/model_extras/audiox.py``). - -Usage: - python openai_chat_client.py --task t2a --prompt "Fireworks burst twice..." --output t2a.wav - python openai_chat_client.py --task tv2a --prompt "drum beating" --video clip.mp4 -o tv2a.wav -""" - -from __future__ import annotations - -import argparse -import base64 -import io -import mimetypes -import sys -from pathlib import Path - -import requests -import soundfile -import torch - -VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"}) -TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"}) - - -def _to_data_url(path: str) -> str: - mime, _ = mimetypes.guess_type(path) - mime = mime or "video/mp4" - with open(path, "rb") as f: - data = base64.b64encode(f.read()).decode("ascii") - return f"data:{mime};base64,{data}" - - -def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None: - audio = audio.to(torch.float32) - audio = audio / audio.abs().max().clamp(min=1e-8) - path.parent.mkdir(parents=True, exist_ok=True) - # soundfile expects channels-last (T, C); project convention is (C, T). - soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16") - - -def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]: - for choice in body.get("choices", []): - audio_obj = choice.get("message", {}).get("audio") - if not (isinstance(audio_obj, dict) and audio_obj.get("data")): - continue - data, sr = soundfile.read(io.BytesIO(base64.b64decode(audio_obj["data"])), dtype="float32", always_2d=True) - return torch.from_numpy(data).transpose(0, 1), sr - brief = {k: v for k, v in body.items() if k != "choices"} - raise RuntimeError(f"no audio in response message.audio: {brief}") - - -def main() -> int: - p = argparse.ArgumentParser(description="AudioX OpenAI chat client") - p.add_argument("--task", required=True, choices=["t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m"]) - p.add_argument("--prompt", "-p", default="", help="Text prompt (required for t2*/tv2*).") - p.add_argument("--video", help="Video path or URL (required for v2*/tv2*).") - p.add_argument("--output", "-o", default="audiox_out.wav") - p.add_argument("--server", "-s", default="http://localhost:8099") - p.add_argument("--model", default="zhangj1an/AudioX") - p.add_argument("--steps", type=int, default=250) - p.add_argument("--guidance-scale", type=float, default=7.0) - p.add_argument("--seed", type=int, default=42) - p.add_argument("--seconds-total", type=float, default=10.0) - p.add_argument("--seconds-start", type=float, default=0.0) - p.add_argument("--sigma-min", type=float, default=0.03) - p.add_argument("--sigma-max", type=float, default=1000.0) - args = p.parse_args() - - if args.task in VIDEO_TASKS and not args.video: - print(f"ERROR: task {args.task!r} requires --video", file=sys.stderr) - return 2 - if args.task in TEXT_TASKS and not args.prompt.strip() and args.task not in {"v2a", "v2m"}: - print(f"ERROR: task {args.task!r} requires --prompt", file=sys.stderr) - return 2 - - content: list[dict] = [{"type": "text", "text": args.prompt}] - if args.task in VIDEO_TASKS: - vurl = args.video if args.video.startswith(("http://", "https://")) else _to_data_url(args.video) - content.append({"type": "video_url", "video_url": {"url": vurl}}) - - payload = { - "model": args.model, - "messages": [{"role": "user", "content": content}], - "extra_body": { - "num_inference_steps": args.steps, - "guidance_scale": args.guidance_scale, - "seed": args.seed, - "audiox_task": args.task, - "seconds_start": args.seconds_start, - "seconds_total": args.seconds_total, - "sigma_min": args.sigma_min, - "sigma_max": args.sigma_max, - }, - } - - print(f"POST {args.server}/v1/chat/completions task={args.task} steps={args.steps}") - r = requests.post( - f"{args.server}/v1/chat/completions", - headers={"Content-Type": "application/json"}, - json=payload, - timeout=600, - ) - r.raise_for_status() - audio, sr = _decode_audio_from_response(r.json()) - _save_wav(audio, Path(args.output), sr) - dur = audio.shape[-1] / sr - print(f"saved {args.output} sr={sr}Hz duration={dur:.2f}s channels={audio.shape[0]}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/examples/online_serving/audiox/run_server.sh b/examples/online_serving/audiox/run_server.sh deleted file mode 100755 index d1a259bfd..000000000 --- a/examples/online_serving/audiox/run_server.sh +++ /dev/null @@ -1,16 +0,0 @@ -#!/bin/bash -# AudioX online serving startup script. - -MODEL="${MODEL:-zhangj1an/AudioX}" -PORT="${PORT:-8099}" -DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-FLASH_ATTN}" - -echo "Starting AudioX server..." -echo "Model: $MODEL" -echo "Port: $PORT" -echo "Diffusion attention backend: $DIFFUSION_ATTENTION_BACKEND" - -DIFFUSION_ATTENTION_BACKEND="$DIFFUSION_ATTENTION_BACKEND" \ - vllm serve "$MODEL" --omni \ - --model-class-name AudioXPipeline \ - --port "$PORT" diff --git a/examples/online_serving/image_to_video/README.md b/examples/online_serving/image_to_video/README.md index e6d9b070f..7fac613df 100644 --- a/examples/online_serving/image_to_video/README.md +++ b/examples/online_serving/image_to_video/README.md @@ -1,5 +1,19 @@ # Image-To-Video +This example demonstrates how to deploy image-to-video models for online +video generation using vLLM-Omni. + +## Supported Models + +| Model | Model ID | +|-------|----------| +| Wan2.2 I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | +| SANA-Video 2B (480p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | +| SANA-Video 2B (720p) | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | + +For SANA-Video 2B image-to-video online serving, see the +[SANA-Video 2B recipe](../../../recipes/NVIDIA/SANA-Video-2B.md). + This example demonstrates online image-to-video generation with vLLM-Omni. The startup script defaults to Wan2.2, but `MODEL` can select any supported image-to-video model. The existing curl helper remains a concrete Wan2.2 diff --git a/examples/online_serving/image_to_video/run_curl_sana_video.sh b/examples/online_serving/image_to_video/run_curl_sana_video.sh new file mode 100755 index 000000000..35806335f --- /dev/null +++ b/examples/online_serving/image_to_video/run_curl_sana_video.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# Submit and download a standard SANA-Video-2B image-to-video request. + +set -euo pipefail + +INPUT_IMAGE="${INPUT_IMAGE:-}" +BASE_URL="${BASE_URL:-http://localhost:8099}" +OUTPUT_PATH="${OUTPUT_PATH:-sana_video_i2v.mp4}" +WIDTH="${WIDTH:-832}" +HEIGHT="${HEIGHT:-480}" +POLL_INTERVAL="${POLL_INTERVAL:-2}" + +if [ -z "$INPUT_IMAGE" ] || [ ! -f "$INPUT_IMAGE" ]; then + echo "Set INPUT_IMAGE to an existing image file." + exit 1 +fi + +create_response="$( + curl -sS -X POST "${BASE_URL}/v1/videos" \ + -H "Accept: application/json" \ + -F "prompt=A cat turns toward the camera with smooth, natural motion." \ + -F "negative_prompt=blurry, low quality, temporal artifacts" \ + -F "input_reference=@${INPUT_IMAGE}" \ + -F "width=${WIDTH}" \ + -F "height=${HEIGHT}" \ + -F "num_frames=81" \ + -F "fps=16" \ + -F "num_inference_steps=50" \ + -F "guidance_scale=6.0" \ + -F "seed=42" \ + -F 'extra_params={"motion_score":30}' +)" + +video_id="$(echo "${create_response}" | jq -r '.id')" +if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then + echo "Failed to create SANA-Video job:" + echo "${create_response}" | jq . + exit 1 +fi + +while true; do + status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")" + status="$(echo "${status_response}" | jq -r '.status')" + + case "${status}" in + queued|in_progress) + echo "Video job ${video_id} status: ${status}" + sleep "${POLL_INTERVAL}" + ;; + completed) + break + ;; + failed) + echo "SANA-Video generation failed:" + echo "${status_response}" | jq . + exit 1 + ;; + *) + echo "Unexpected status response:" + echo "${status_response}" | jq . + exit 1 + ;; + esac +done + +curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}" +echo "Saved SANA-Video to ${OUTPUT_PATH}" diff --git a/examples/online_serving/image_to_video/run_server_sana_video.sh b/examples/online_serving/image_to_video/run_server_sana_video.sh new file mode 100755 index 000000000..257cb2b00 --- /dev/null +++ b/examples/online_serving/image_to_video/run_server_sana_video.sh @@ -0,0 +1,13 @@ +#!/bin/bash +# SANA-Video-2B native image-to-video serving. + +set -euo pipefail + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8099}" + +vllm serve "$MODEL" \ + --omni \ + --model-class-name SanaImageToVideoPipeline \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh b/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh new file mode 100755 index 000000000..bf0573529 --- /dev/null +++ b/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# SANA-Video-2B image-to-video serving through the Diffusers adapter. +# Both 480p and 720p checkpoints are validated. For 720p: +# MODEL=Efficient-Large-Model/SANA-Video_2B_720p_diffusers bash run_server_sana_video_diffusers.sh +# INPUT_IMAGE=input.jpg WIDTH=1280 HEIGHT=704 bash run_curl_sana_video.sh + +set -euo pipefail + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8099}" + +vllm serve "$MODEL" \ + --omni \ + --model-class-name SanaImageToVideoPipeline \ + --diffusion-load-format diffusers \ + --diffusion-attention-backend TORCH_SDPA \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py index 3030fa3e3..f2aa159cf 100644 --- a/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py @@ -328,7 +328,7 @@ def run_multimodal_generation(args) -> None: extra_body = { "sampling_params_list": sampling_params_list - # Optional, it has a default setting in stage_configs of the corresponding model. + # Optional; defaults come from the resolved pipeline and deploy config. } # Build messages list (filter None so concurrent tasks get valid structure) diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index 6ed7d8c32..96c021a44 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -32,12 +32,14 @@ GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That profile admits at most four concurrent sequences per stage. | deploy config | GPUs | Notes | -|---|---|---| +| --- | --- | --- | | `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. | | `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. | | `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. | | `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. | -| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. | + +Every profile sets `session_mode: duplex`, so native full-duplex is served +from the same process as `/v1/chat/completions`. The split pipeline preserves native-duplex epoch/turn identity, segment text, turn completion, reference voice, and terminal-audio metadata through @@ -55,8 +57,8 @@ vllm serve openbmb/MiniCPM-o-4_5 \ ``` For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the -checkpoint path. To start the experimental native duplex backend, use -`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`. +checkpoint path. For native full-duplex, connect `/v1/realtime?duplex=1` on +this same server. ### Per-stage overrides @@ -145,7 +147,7 @@ assistant template, so it is not an apples-to-apples accuracy run. ## Run the Realtime duplex CLI demo -After the duplex backend is running, stream one WAV through the Realtime +After the server is running, stream one WAV through the Realtime WebSocket endpoint: ```bash @@ -157,6 +159,33 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \ --output-dir /tmp/minicpmo_realtime_duplex_demo ``` +Video input uses the same session. PyAV demuxes JPEG frames at `--video-fps` +(default 1.0) and vLLM `load_audio` extracts a 16 kHz mono WAV unless +`--input-wav` overrides the soundtrack. Frames keep their capture size +(`--frame-max-side 0`); the server `process_image` normalizes at 448. + +`--stack-frames N` raises the visual refresh rate the way official duplex does: +each 1 s unit also samples the `N-1` sub-frames captured inside it and tiles +them into a single composite image sent next to that unit's base frame. The +audio cadence never changes — a unit is always one second — and the wire always +carries 2 images per unit however large `N` is, because the sub-frames share one +composite. Official uses 5 for high refresh rate mode: + +```bash +python examples/online_serving/minicpmo/realtime_duplex_demo.py \ + --url ws://localhost:8099/v1/realtime?duplex=1 \ + --model openbmb/MiniCPM-o-4_5 \ + --input-video /path/to/clip.mp4 \ + --stack-frames 5 \ + --ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \ + --output-dir /tmp/minicpmo_realtime_duplex_video_demo +``` + +Detail inside a composite is capped by `scale_resolution=448` at +`max_slice_nums=1`: official suggests HD slicing (`max_slice_nums=[2, 1]`) for +stacked frames, which the duplex adapter does not implement yet. Reading small +text or digits out of a wide scene is limited by that, not by frame timing. + ## Open the experimental browser client The browser UI serves the page and proxies the same-origin Realtime WebSocket to @@ -190,6 +219,64 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built two-response WAV, its `--input-sha256`, and an `--expect-second-response-substring` value. +## Run Omni-DuplexEval + +Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex +endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge +served from the same vLLM-Omni environment. This validates the vLLM-Omni +duplex client and local judge integration; it does not reproduce the paper's +original MiniCPM-o inference implementation. + +Start the duplex generation server: + +```bash +vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \ + --deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \ + --served-model-name openbmb/MiniCPM-o-4_5 \ + --host 0.0.0.0 --port 8099 +``` + +Start a separate multimodal judge. The allowed path must contain any local +videos passed with `--judge-video-mode video_url`: + +```bash +vllm serve Qwen/Qwen2.5-VL-7B-Instruct \ + --served-model-name Qwen/Qwen2.5-VL-7B-Instruct \ + --allowed-local-media-path /data/omni-duplex-eval \ + --host 0.0.0.0 --port 8000 +``` + +Generate, evaluate, and summarize: + +```bash +vllm bench omni-duplex-eval --omni generate \ + --url ws://127.0.0.1:8099/v1/realtime?duplex=1 \ + --model openbmb/MiniCPM-o-4_5 \ + --ref-audio /data/ref.wav \ + --dataset Hothan/Omni-DuplexEval \ + --concurrency 2 \ + --response-root /data/omni-duplex-eval/responses + +vllm bench omni-duplex-eval --omni evaluate \ + --dataset Hothan/Omni-DuplexEval \ + --response-root /data/omni-duplex-eval/responses \ + --score-root /data/omni-duplex-eval/scores \ + --judge-base-url http://127.0.0.1:8000 \ + --judge-model Qwen/Qwen2.5-VL-7B-Instruct \ + --judge-video-mode video_url \ + --eval-workers 4 + +vllm bench omni-duplex-eval --omni summarize \ + --score-root /data/omni-duplex-eval/scores +``` + +The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face +splits, not dataset configurations. Pass one with `--split`, or omit it to +run all nine splits. `--limit 1` is useful for a smoke test. Realtime +generation records `clock=media`; artifacts generated with +`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects +them unless `--allow-invalid-clock` is explicit. + ## Related examples - [Offline MiniCPM-o inference](../../offline_inference/minicpmo/) diff --git a/examples/online_serving/minicpmo/realtime_duplex_demo.py b/examples/online_serving/minicpmo/realtime_duplex_demo.py index b6c861747..b12c20a56 100644 --- a/examples/online_serving/minicpmo/realtime_duplex_demo.py +++ b/examples/online_serving/minicpmo/realtime_duplex_demo.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Minimal single-input MiniCPM-o 4.5 Realtime duplex demo. Run this after starting the duplex server. Strict lifecycle, overlap, and @@ -9,7 +12,9 @@ from __future__ import annotations import argparse import asyncio import base64 +import io import json +import math import sys import time from pathlib import Path @@ -28,6 +33,17 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 wait_for, write_pcm16_wav, ) +from vllm_omni.experimental.fullduplex.client import chunk_period_ms as _chunk_period_ms # noqa: E402 +from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 + has_residual_model_unit as _has_residual_model_unit, +) +from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 + reference_audio_data_url as _ref_audio_data_url, +) +from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402 + concat_frames_b64, + unit_subframe_offsets, +) class _StreamingOutputWriter: @@ -133,31 +149,29 @@ def _latest_model_decision( return decision -def _chunk_period_ms(events: list[dict[str, object]]) -> int: - for event in reversed(events): - session = event.get("session") - if not isinstance(session, dict): - continue - capabilities = session.get("capabilities") - if not isinstance(capabilities, dict): - continue - chunk_period_ms = capabilities.get("chunk_period_ms") - if isinstance(chunk_period_ms, int) and chunk_period_ms > 0: - return chunk_period_ms - return 1000 - - -def _has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool: - unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000 - return bool(unit_bytes > 0 and len(pcm16) % unit_bytes) - - def _response_in_progress(events: list[dict[str, object]]) -> bool: return sum(event.get("type") == "response.created" for event in events) > sum( event.get("type") == "response.done" for event in events ) +async def _commit_input_after_playback_checkpoint(client: RealtimeDuplexClient) -> float: + """Commit input only after preserving already-played assistant history. + + Native duplex can finish several responses while the user is still + streaming. Under the ``ack_only`` playback policy, acknowledging those + responses after the final input commit is intentionally rejected: it could + append an old assistant message after the newly committed user message. + Checkpoint playback first so the server commits each assistant item at its + original history position. A second acknowledgement after the commit can + then safely advance response-local playback that drained in the meantime. + """ + await client.acknowledge_playback() + commit_sent_at_s = time.monotonic() + await client.commit() + return commit_sent_at_s + + def _event_count_after( events: list[dict[str, object]], event_type: str, @@ -168,21 +182,258 @@ def _event_count_after( return sum(event.get("type") == event_type for event in events[index + 1 :]) -def _ref_audio_data_url(path: str | None) -> str | None: - if path is None: - return None - ref_path = Path(path).expanduser() - return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii") +def _probe_video_frames_and_fps(video_path: Path) -> tuple[int, float]: + """Frame count and average FPS via PyAV (same approach as Daily-Omni).""" + import av + + with av.open(str(video_path)) as container: + stream = container.streams.video[0] + avg_fps = float(stream.average_rate) if stream.average_rate else 0.0 + num_frames = stream.frames or 0 + if num_frames <= 0: + duration: float | None = None + if stream.duration is not None and stream.time_base is not None: + duration = float(stream.duration * stream.time_base) + elif container.duration is not None: + duration = container.duration / av.time_base + if duration and avg_fps: + num_frames = int(math.floor(duration * avg_fps)) + if num_frames <= 0: + stream.thread_type = "AUTO" + num_frames = sum(1 for _ in container.decode(stream)) + if num_frames <= 0: + raise ValueError(f"No decodable video frames in {video_path}") + return num_frames, avg_fps or 1.0 + + +def _upright(image, rotation: int): + """Apply a container display-matrix rotation, in degrees, to ``image``. + + PyAV hands back the stored raster and leaves the display matrix for the + caller, unlike the ffmpeg CLI which autorotates. Phone clips routinely + carry +/-90, so skipping this feeds the model a sideways frame: arrows and + digits stop being readable and the model mirrors the motion it reports. + """ + from PIL import Image + + quarter_turns = { + 90: Image.Transpose.ROTATE_90, + 180: Image.Transpose.ROTATE_180, + 270: Image.Transpose.ROTATE_270, + } + turn = quarter_turns.get(int(rotation) % 360) + return image.transpose(turn) if turn is not None else image + + +def _decode_video_frames_rgb(video_path: Path, frame_idx: list[int]) -> list: + """Decode presentation-order ``frame_idx`` into upright RGB ``PIL.Image``s.""" + import av + + wanted = sorted({int(i) for i in frame_idx}) + if not wanted: + return [] + + decoded: dict[int, object] = {} + with av.open(str(video_path)) as container: + stream = container.streams.video[0] + stream.thread_type = "AUTO" + cursor = 0 + for position, frame in enumerate(container.decode(stream)): + if wanted[cursor] > position: + continue + image = _upright(frame.to_image(), getattr(frame, "rotation", 0) or 0) + while cursor < len(wanted) and wanted[cursor] <= position: + decoded[wanted[cursor]] = image + cursor += 1 + if cursor >= len(wanted): + break + + if not decoded: + raise ValueError(f"Decoded no frames from {video_path} for indices {wanted[:8]}") + last_image = decoded[max(decoded)] + return [decoded.get(int(i), last_image) for i in frame_idx] + + +def _resize_max_side(image, max_side: int): + width, height = image.size + longest = max(width, height) + if longest <= max_side or max_side <= 0: + return image + scale = max_side / float(longest) + new_size = (max(1, int(round(width * scale))), max(1, int(round(height * scale)))) + return image.resize(new_size) + + +def _jpeg_bytes(image, *, quality: int = 95) -> bytes: + buf = io.BytesIO() + image.convert("RGB").save(buf, format="JPEG", quality=quality) + return buf.getvalue() + + +def _to_pcm16_bytes(audio_np) -> bytes: + import numpy as np + + arr = np.asarray(audio_np).reshape(-1) + if arr.dtype != np.int16: + arr = (np.clip(arr, -1.0, 1.0) * 32767.0).astype(np.int16) + return arr.tobytes() + + +def _pcm16_wav_duration_s(path: Path) -> float: + pcm16 = read_pcm16_wav(path) + return len(pcm16) / float(PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE) + + +def _extract_video_frames( + video_path: Path, + work_dir: Path, + *, + fps: float, + max_side: int, + duration_s: float, + stack_frames: int = 1, +) -> tuple[list[str], list[str | None]]: + """Demux JPEG frames only. Frame count follows ``duration_s``, not soundtrack. + + Returns the base frames (one per model unit) and the parallel stacked track. + With ``stack_frames=N`` each unit also samples its ``N-1`` interior + sub-frames and tiles them into one composite, which is how official duplex + raises the visual refresh rate while keeping one second of audio per unit. + A unit whose interior falls past the end of the clip stacks nothing. + """ + num_frames, avg_fps = _probe_video_frames_and_fps(video_path) + sample_fps = max(float(fps), 1e-6) + unit_s = 1.0 / sample_fps + timeline_s = max(float(duration_s), unit_s) + num_samples = max(1, int(math.ceil(timeline_s * sample_fps - 1e-9))) + offsets = unit_subframe_offsets(stack_frames, unit_s=unit_s) + + def frame_at(seconds: float) -> int: + # Official get_video_frame_audio_segments uses truncating int(), not + # round(), so 29.97 fps lands on the same presentation index. + return min(int(seconds * avg_fps), num_frames - 1) + + base_idx = [frame_at(unit * unit_s) for unit in range(num_samples)] + sub_times = [ + [unit * unit_s + offset for offset in offsets if unit * unit_s + offset < timeline_s] + for unit in range(num_samples) + ] + sub_idx = [frame_at(seconds) for times in sub_times for seconds in times] + images = _decode_video_frames_rgb(video_path, base_idx + sub_idx) + sub_images = iter(images[len(base_idx) :]) + + frame_dir = work_dir / "frames" + frame_dir.mkdir(parents=True, exist_ok=True) + for stale in frame_dir.glob("*.jpg"): + stale.unlink() + frames_b64: list[str] = [] + stacked_b64: list[str | None] = [] + for index, (image, times) in enumerate(zip(images[: len(base_idx)], sub_times), start=1): + jpeg = _jpeg_bytes(_resize_max_side(image, max_side)) + (frame_dir / f"frame_{index:05d}.jpg").write_bytes(jpeg) + frames_b64.append(base64.b64encode(jpeg).decode("ascii")) + interior = [_resize_max_side(next(sub_images), max_side) for _ in times] + if not interior: + stacked_b64.append(None) + continue + composite = concat_frames_b64(interior) + (frame_dir / f"stack_{index:05d}.jpg").write_bytes(base64.b64decode(composite)) + stacked_b64.append(composite) + return frames_b64, stacked_b64 + + +def _extract_video_soundtrack(video_path: Path, work_dir: Path) -> Path: + """Write a 16 kHz mono WAV from the video soundtrack.""" + try: + from vllm.multimodal.media.audio import load_audio + except ImportError as exc: # pragma: no cover + raise SystemExit("--input-video needs vLLM audio helpers (load_audio)") from exc + + audio_np, sr = load_audio(str(video_path), sr=PCM16_SAMPLE_RATE, mono=True) + work_dir.mkdir(parents=True, exist_ok=True) + audio_path = work_dir / "input_16k.wav" + write_pcm16_wav(audio_path, _to_pcm16_bytes(audio_np), sample_rate_hz=int(sr)) + return audio_path + + +def _extract_video_input( + video_path: Path, + work_dir: Path, + *, + fps: float, + max_side: int, + stack_frames: int, +) -> tuple[Path, list[str], list[str | None]]: + """Demux a video into duplex inputs: 16 kHz mono WAV plus JPEG frames. + + Use this only when the video soundtrack is the audio source. When the + caller already has ``--input-wav``, extract frames against that WAV + duration instead so a silent clip cannot fail ``load_audio``. + """ + audio_path = _extract_video_soundtrack(video_path, work_dir) + frames_b64, stacked_b64 = _extract_video_frames( + video_path, + work_dir, + fps=fps, + max_side=max_side, + duration_s=_pcm16_wav_duration_s(audio_path), + stack_frames=stack_frames, + ) + return audio_path, frames_b64, stacked_b64 + + +def _resolve_duplex_av_inputs( + *, + input_wav: str | None, + input_video: str | None, + work_dir: Path, + fps: float, + max_side: int, + stack_frames: int = 1, +) -> tuple[str, list[str], list[str | None]]: + """Return ``(wav_path, frames_b64, stacked_b64)``. External WAV wins over video audio.""" + if input_video: + video_path = Path(input_video) + if input_wav: + frames, stacked = _extract_video_frames( + video_path, + work_dir, + fps=fps, + max_side=max_side, + duration_s=_pcm16_wav_duration_s(Path(input_wav)), + stack_frames=stack_frames, + ) + return input_wav, frames, stacked + extracted_wav, frames, stacked = _extract_video_input( + video_path, + work_dir, + fps=fps, + max_side=max_side, + stack_frames=stack_frames, + ) + return str(extracted_wav), frames, stacked + if not input_wav: + raise SystemExit("provide --input-wav, --input-video, or both") + return input_wav, [], [] async def run_demo(args: argparse.Namespace) -> dict[str, object]: - input_pcm16 = read_pcm16_wav(Path(args.input_wav)) - if not input_pcm16: - raise ValueError("input WAV has no audio") - output_dir = Path(args.output_dir) stream_writer = _StreamingOutputWriter(output_dir) + input_wav, video_frames, stacked_video_frames = _resolve_duplex_av_inputs( + input_wav=args.input_wav, + input_video=args.input_video, + work_dir=output_dir / "video_input", + fps=args.video_fps, + max_side=args.frame_max_side, + stack_frames=args.stack_frames, + ) + + input_pcm16 = read_pcm16_wav(Path(input_wav)) + if not input_pcm16: + raise ValueError("input WAV has no audio") + url = build_realtime_url( args.url, args.model, @@ -200,10 +451,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: timeout_s=args.timeout_s, ) stream_event_cursor = len(client.events.events) - await client.stream_pcm16( + frames_sent = await client.stream_pcm16( input_pcm16, chunk_ms=args.chunk_ms, realtime=not args.no_realtime_pacing, + video_frames=video_frames, + stacked_video_frames=stacked_video_frames, ) commit_event_cursor = len(client.events.events) stream_decision = _latest_model_decision(client.events.events, stream_event_cursor) @@ -212,8 +465,7 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: chunk_period_ms=_chunk_period_ms(client.events.events), ) wait_for_post_commit_decision = False - commit_sent_at_s = time.monotonic() - await client.commit() + commit_sent_at_s = await _commit_input_after_playback_checkpoint(client) wait_error: str | None = None committed_index: int | None = None post_commit_decision: str | None = None @@ -326,6 +578,10 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: committed_index, ), }, + "input_video": args.input_video, + "video_frames_extracted": len(video_frames), + "video_frames_sent": frames_sent, + "video_stacked_frames_extracted": sum(frame is not None for frame in stacked_video_frames), "audio_bytes": len(audio), "audio_chunk_count": len(stream_writer.audio_chunk_paths), "audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths], @@ -372,7 +628,48 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1") parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5") parser.add_argument("--session-id") - parser.add_argument("--input-wav", required=True) + parser.add_argument( + "--input-wav", + help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.", + ) + parser.add_argument( + "--input-video", + help=( + "Video to stream as omni duplex input. PyAV demuxes frames and " + "vLLM load_audio extracts a 16 kHz mono WAV (unless --input-wav overrides the audio)." + ), + ) + parser.add_argument( + "--video-fps", + type=float, + default=1.0, + help="Frame extraction rate. The wire contract carries ~1 frame per second of audio.", + ) + parser.add_argument( + "--frame-max-side", + type=int, + default=0, + help=( + "Optional client-side downscale of the longest side, in pixels. " + "0 (default) keeps the capture resolution so the server " + "process_image can normalize at scale_resolution=448 from a " + "higher-quality source. The web demo never resizes on the client." + ), + ) + parser.add_argument( + "--stack-frames", + type=int, + default=1, + help=( + "Frames sampled per 1 s audio unit. 1 (default) sends the unit's " + "frame alone. N > 1 also samples the N-1 sub-frames inside the unit " + "and tiles them into one composite sent next to the base frame, so " + "motion (arrows, doors, digits changing) is visible at N fps while " + "the audio stays 1 s per unit. Official uses 5 for high refresh " + "rate mode; the wire carries 2 images per unit at any N because the " + "sub-frames share one composite." + ), + ) parser.add_argument( "--ref-audio", required=True, @@ -388,7 +685,7 @@ def parse_args() -> argparse.Namespace: "--temperature", type=float, default=None, - help="Stage0 sampling temperature; response-required soft-interrupt uses 0.0.", + help="Stage 0 sampling temperature; omit to preserve the model default.", ) parser.add_argument("--no-realtime-pacing", action="store_true") parser.add_argument("--require-audio", action="store_true") diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py index 63a3f9ec4..fd74ebe71 100644 --- a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py @@ -415,7 +415,7 @@ def run_multimodal_generation(args, client: OpenAI) -> None: audio_path=audio_path, ) extra_body = { - # Optional, it has default settings in stage configs. you can override them here. + # Optional; defaults come from the resolved pipeline and deploy config. } if args.query_type == "use_audio_in_video": extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True} diff --git a/examples/online_serving/personaplex/README.md b/examples/online_serving/personaplex/README.md index cf105dd6c..3b737d275 100644 --- a/examples/online_serving/personaplex/README.md +++ b/examples/online_serving/personaplex/README.md @@ -38,7 +38,7 @@ full-duplex control plane: CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \ /path/to/personaplex-7b-v1 \ --omni \ - --stage-configs-path vllm_omni/deploy/personaplex.yaml + --deploy-config vllm_omni/deploy/personaplex.yaml ``` Validate the actual `/v1/realtime?duplex=1` scheduler path with paced 24 kHz diff --git a/examples/online_serving/pi0/README.md b/examples/online_serving/pi0/README.md new file mode 100644 index 000000000..1f424a88b --- /dev/null +++ b/examples/online_serving/pi0/README.md @@ -0,0 +1,104 @@ +# π0 (Pi-Zero) VLA — OpenPI realtime serving + +[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action +model from Physical Intelligence: multi-camera images + a language instruction + +robot proprioceptive state → a continuous action chunk via flow-matching denoising +(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime +websocket protocol at `/v1/realtime/robot/openpi`. + +## Install extras + +The core `pip install -e .` does not include the OpenPI client used here: + +- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi) + repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy` + +## Weights + +π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server +download it (`MODEL=lerobot/pi0_base`) or point at a local copy +(`MODEL=/path/to/pi0_base`). + +## Run the server + +```bash +vllm serve lerobot/pi0_base --omni --port 8000 \ + --served-model-name pi0 \ + --deploy-config vllm_omni/deploy/pi0.yaml \ + --enforce-eager --disable-log-stats +``` + +The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage +(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake +metadata (3 cameras, `joint_position`, action horizon 50, action dim 32). + +### Commonly adjusted `model_config` keys + +These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`: + +| Key | Default | Meaning | +|---|---|---| +| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. | +| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. | +| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). | +| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). | +| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). | +| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. | +| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). | +| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. | +| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). | + +The `policy_server_config` block below them is the OpenPI handshake metadata +advertised to the client; keep its `action_horizon` / `action_dim` / +`image_resolution` in sync with the `model_config` values above. The e2e test +`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a +live server and asserts the advertised metadata matches these `pi0.yaml` values. + +## Run the client + +```bash +python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \ + --prompt "pick up the red block and place it in the bin" +``` + +It connects, prints the server metadata, sends robot observations, and prints the +returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank +cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and +proprioceptive state to drive a robot. + +### Observation format + +The client sends a flat dict per inference: + +```python +{ + "observation.images.base_0_rgb": np.uint8[H, W, 3], + "observation.images.left_wrist_0_rgb": np.uint8[H, W, 3], + "observation.images.right_wrist_0_rgb":np.uint8[H, W, 3], + "state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side + "prompt": "pick up the red block", + "session_id": "", # accepted but ignored (π0 is stateless) +} +``` + +Camera keys must match the server's `image_feature_keys` (the checkpoint's +`input_features` order). If your robot uses different camera names, set +`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys. + +## Correctness + +π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference +(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see +`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is +version-stable across transformers releases (the version-stability checks in +`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e +lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`. + +## Limitations + +- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is + fully supported (state passes through, actions are returned in the model's space). + Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but + stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are + not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge + added before its actions are in real-world units. diff --git a/examples/online_serving/pi0/openpi_client.py b/examples/online_serving/pi0/openpi_client.py new file mode 100644 index 000000000..f2154b33d --- /dev/null +++ b/examples/online_serving/pi0/openpi_client.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server. + +Connects to a running π0 server (see README.md), sends a single robot +observation (3 cameras + state + language instruction), and prints the returned +continuous action chunk ``[action_horizon, action_dim]``. + +The server speaks the OpenPI websocket protocol: + connect -> server sends msgpack(policy_server_config metadata) + infer -> client sends msgpack(obs), server replies msgpack(actions ndarray) + +Usage: + python openpi_client.py --host 127.0.0.1 --port 8000 \ + --prompt "pick up the red block" +""" + +from __future__ import annotations + +import argparse +import logging +import sys +import uuid +from pathlib import Path +from typing import Any + +import numpy as np + +try: + # Avoid shadowing the installed ``openpi_client`` package with this dir. + example_dir = str(Path(__file__).resolve().parent) + removed = sys.path and sys.path[0] == example_dir + if removed: + sys.path.pop(0) + try: + from openpi_client.websocket_client_policy import WebsocketClientPolicy + finally: + if removed: + sys.path.insert(0, example_dir) +except ImportError as exc: # pragma: no cover - runtime dependency guard + raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc + +logger = logging.getLogger("pi0_openpi_client") + +DEFAULT_HOST = "127.0.0.1" +DEFAULT_PORT = 8000 +DEFAULT_PATH = "/v1/realtime/robot/openpi" +DEFAULT_PROMPT = "pick up the red block and place it in the bin" + +# π0 / pi0_base camera identities (must match the server's image_feature_keys). +CAMERA_KEYS = ( + "observation.images.base_0_rgb", + "observation.images.left_wrist_0_rgb", + "observation.images.right_wrist_0_rgb", +) +STATE_DIM = 32 + + +def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]: + """Build a single π0 observation with blank cameras + zero state. + + Replace the zeros with real camera frames (HWC uint8) and proprioceptive + state to drive an actual robot. + """ + obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS} + obs["state"] = np.zeros(STATE_DIM, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + return obs + + +def _ws_uri(host: str, port: int, path: str) -> str: + """Build the full websocket URI. WebsocketClientPolicy takes the whole URI as + ``host`` (it has no separate path arg).""" + if host.startswith(("ws://", "wss://")): + return host + return f"ws://{host}:{port}{path}" + + +def _as_action_array(response: Any) -> np.ndarray: + if isinstance(response, dict) and response.get("type") == "error": + raise RuntimeError(f"Inference failed: {response.get('message', response)}") + return np.asarray(response, dtype=np.float32) + + +def main() -> int: + parser = argparse.ArgumentParser(description="π0 OpenPI client") + parser.add_argument("--host", default=DEFAULT_HOST) + parser.add_argument("--port", type=int, default=DEFAULT_PORT) + parser.add_argument("--path", default=DEFAULT_PATH) + parser.add_argument("--prompt", default=DEFAULT_PROMPT) + parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.") + args = parser.parse_args() + logging.basicConfig(level=logging.INFO, format="%(message)s") + + policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path)) + metadata = dict(policy.get_server_metadata()) + logger.info("Server metadata: %s", metadata) + + session_id = str(uuid.uuid4()) + for step in range(args.num_steps): + obs = _make_dummy_obs(args.prompt, session_id) + actions = _as_action_array(policy.infer(obs)) + logger.info( + "[step %d] actions shape=%s mean=%.4f std=%.4f", + step, + actions.shape, + float(actions.mean()), + float(actions.std()), + ) + if not np.isfinite(actions).all(): + raise RuntimeError("Server returned non-finite actions.") + logger.info("Done — π0 OpenPI round-trip OK.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh index ed09bac14..d0ed84b85 100644 --- a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh +++ b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh @@ -51,7 +51,7 @@ code2wav_sampling_params='{ "detokenize": true, "repetition_penalty": 1.1 }' -# Above is optional, it has a default setting in stage_configs of the corresponding model. +# The block above is optional; defaults come from the resolved pipeline and deploy config. # Define URLs for assets MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index 74a18af32..926a58682 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -24,6 +24,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --deploy-config /path/to/your_deploy_config.yaml ``` +To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct +weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key +`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ + --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml +``` + +Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select +the same single-stage pipeline without `--deploy-config`. + For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2), use `--stage-overrides` on top of the default config: @@ -513,7 +525,7 @@ The script supports the following arguments: vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` -If you have custom stage configs file: +If you have a custom deploy config file: ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file ``` diff --git a/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh index 399c65481..0a059a03e 100644 --- a/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh +++ b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh @@ -50,7 +50,7 @@ code2wav_sampling_params='{ "detokenize": true, "repetition_penalty": 1.1 }' -# Above is optional, it has a default setting in stage_configs of the corresponding model. +# The block above is optional; defaults come from the resolved pipeline and deploy config. # Define URLs for assets MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" diff --git a/examples/online_serving/replica_data_parallel/bench_replica_dp.py b/examples/online_serving/replica_data_parallel/bench_replica_dp.py index 977f475ac..d17899b78 100755 --- a/examples/online_serving/replica_data_parallel/bench_replica_dp.py +++ b/examples/online_serving/replica_data_parallel/bench_replica_dp.py @@ -15,7 +15,7 @@ so an output is a deterministic function of its own input and nothing else. Measuring replica scaling ------------------------- This client is replica-agnostic. Vary the *server's* replica count -(`runtime.num_replicas` in the stage config; see run_server.sh) and keep the +(`stages[].num_replicas` in the deploy config; see run_server.sh) and keep the client workload fixed, then compare throughput: N=1 -> baseline diff --git a/examples/online_serving/stable_audio/README.md b/examples/online_serving/stable_audio/README.md deleted file mode 100644 index 954ac4c85..000000000 --- a/examples/online_serving/stable_audio/README.md +++ /dev/null @@ -1,234 +0,0 @@ -# Stable Audio Online Serving - -Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint. - -## Features - -- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint -- **Flexible control**: Adjust audio length, guidance scale, inference steps -- **Quality control**: Use negative prompts to avoid unwanted characteristics -- **Reproducible**: Set random seed for deterministic generation - -## Quick Start - -### 1. Start the Server - -```bash -vllm serve stabilityai/stable-audio-open-1.0 \ - --host 0.0.0.0 \ - --port 8091 \ - --gpu-memory-utilization 0.9 \ - --trust-remote-code \ - --enforce-eager \ - --omni -``` - -### 2. Generate Audio - -#### Using curl - -```bash -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "The sound of a cat purring", - "audio_length": 10.0 - }' --output cat.wav -``` - -#### Using Python Client - -```bash -python stable_audio_client.py \ - --text "The sound of a cat purring" \ - --audio_length 10.0 \ - --output cat.wav -``` - -#### Using Bash Script - -```bash -bash curl_examples.sh -``` - -## API Reference - -### Endpoint - -``` -POST /v1/audio/generate -``` - -### Request Body - -```json -{ - "input": "Text description of the audio", - "audio_length": 10.0, - "audio_start": 0.0, - "negative_prompt": "Low quality", - "guidance_scale": 7.0, - "num_inference_steps": 100, - "seed": 42, - "response_format": "wav" -} -``` - -### Parameters - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `input` | string | **required** | Text prompt describing the audio to generate | -| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) | -| `audio_start` | float | 0.0 | Audio start time in seconds | -| `negative_prompt` | string | null | Text describing what to avoid in generation | -| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) | -| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) | -| `seed` | int | null | Random seed for reproducibility | -| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm | - -### Response - -Returns audio data in the requested format (default: WAV). - -## Usage Examples - -### Basic Generation - -```bash -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "The sound of ocean waves" - }' --output ocean.wav -``` - -### Custom Duration - -```bash -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "A dog barking", - "audio_length": 5.0 - }' --output dog_5s.wav -``` - -### High Quality with Negative Prompt - -```bash -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "A piano playing a gentle melody", - "audio_length": 10.0, - "negative_prompt": "Low quality, distorted, noisy", - "guidance_scale": 8.0, - "num_inference_steps": 150 - }' --output piano_hq.wav -``` - -### Reproducible Generation - -```bash -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Thunder and rain sounds", - "audio_length": 15.0, - "seed": 42 - }' --output thunder.wav -``` - -### Quick Generation (Fewer Steps) - -For faster generation with slightly lower quality: - -```bash -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Birds chirping in a forest", - "audio_length": 8.0, - "num_inference_steps": 50 - }' --output birds_quick.wav -``` - -## Python Client Examples - -### Simple Generation - -```bash -python stable_audio_client.py \ - --text "The sound of a cat purring" -``` - -### Custom Parameters - -```bash -python stable_audio_client.py \ - --text "Thunder and rain" \ - --audio_length 15.0 \ - --negative_prompt "Low quality" \ - --guidance_scale 7.0 \ - --num_inference_steps 100 \ - --seed 42 \ - --output thunder.wav -``` - -### Different Output Format - -```bash -python stable_audio_client.py \ - --text "Guitar playing" \ - --response_format mp3 \ - --output guitar.mp3 -``` - -## Tips - -1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0` -2. **Quality vs Speed**: - - 50 steps: Fast, decent quality - - 100 steps: Good balance (default) - - 150+ steps: High quality, slower -3. **Guidance Scale**: - - Lower (3-5): More creative/varied - - Default (7): Good balance - - Higher (10+): More literal to prompt -4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc. -5. **Seeds**: Use same seed for reproducible results - -## Performance - -| Inference Steps | Quality | Speed | Use Case | -|----------------|---------|-------|----------| -| 50 | Good | Fast | Quick previews | -| 100 (default) | Very Good | Medium | Production | -| 150+ | Excellent | Slow | Final/critical audio | - -## Troubleshooting - -### Server not responding -- Check if server is running: `curl http://localhost:8091/health` -- Check server logs for errors - -### Audio quality issues -- Increase `num_inference_steps` (e.g., 150) -- Add negative prompts: `"Low quality, distorted, noisy"` -- Increase `guidance_scale` for more prompt adherence - -### Generation timeout -- Reduce `num_inference_steps` -- Reduce `audio_length` -- Check GPU memory with `nvidia-smi` - -### Wrong audio length -- Ensure `audio_length` is within model limits (~47s max) -- Adjust `audio_start` if trimming is needed - -## See Also - -- [Offline Inference Example](../../offline_inference/text_to_audio/README.md) -- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0) -- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni) diff --git a/examples/online_serving/stable_audio/curl_examples.sh b/examples/online_serving/stable_audio/curl_examples.sh deleted file mode 100755 index 16a12ef1c..000000000 --- a/examples/online_serving/stable_audio/curl_examples.sh +++ /dev/null @@ -1,54 +0,0 @@ -#!/bin/bash -# Examples for using Stable Audio with curl via /v1/audio/generate endpoint - -# Example 1: Simple request with default parameters -echo "Example 1: Simple request with default parameters" -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "The sound audience clapping and cheering in a stadium" - }' --output stadium.wav - -# Example 2: Request with custom audio_length -echo "Example 2: Custom audio length (5 seconds)" -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "The sound of a dog barking", - "audio_length": 5.0 - }' --output dog_5s.wav - -# Example 3: Request with negative prompt for quality control -echo "Example 3: With negative prompt" -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "A piano playing a gentle melody", - "audio_length": 10.0, - "negative_prompt": "Low quality, distorted, noisy" - }' --output piano.wav - -# Example 4: Full control with all parameters -echo "Example 4: Full control (custom length, guidance, steps, seed)" -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Thunder and rain sounds", - "audio_length": 15.0, - "negative_prompt": "Low quality", - "guidance_scale": 7.0, - "num_inference_steps": 100, - "seed": 42 - }' --output thunder_rain.wav - -# Example 5: Quick generation with fewer steps (faster but lower quality) -echo "Example 5: Quick generation (fewer steps)" -curl -X POST http://localhost:8091/v1/audio/generate \ - -H "Content-Type: application/json" \ - -d '{ - "input": "Ocean waves crashing on a beach", - "audio_length": 8.0, - "num_inference_steps": 50 - }' --output ocean.wav - -echo "All examples completed!" diff --git a/examples/online_serving/stable_audio/stable_audio_client.py b/examples/online_serving/stable_audio/stable_audio_client.py deleted file mode 100755 index 7f28a5245..000000000 --- a/examples/online_serving/stable_audio/stable_audio_client.py +++ /dev/null @@ -1,170 +0,0 @@ -#!/usr/bin/env python3 -""" -OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint. - -This script demonstrates how to use the OpenAI-compatible speech API -to generate audio from text using Stable Audio models. - -Examples: - # Simple generation - python stable_audio_client.py --text "The sound of a cat purring" - - # With custom duration - python stable_audio_client.py --text "A dog barking" --audio_length 5.0 - - # With all parameters - python stable_audio_client.py --text "Thunder and rain" \ - --audio_length 15.0 \ - --negative_prompt "Low quality" \ - --guidance_scale 7.0 \ - --num_inference_steps 100 \ - --seed 42 \ - --output thunder.wav -""" - -import argparse -import sys - -import requests - - -def parse_args(): - parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API") - parser.add_argument( - "--api_url", - default="http://localhost:8091/v1/audio/generate", - help="API endpoint URL", - ) - parser.add_argument( - "--text", - default="The sound of a cat purring", - help="Text prompt for audio generation", - ) - parser.add_argument( - "--audio_length", - type=float, - default=10.0, - help="Audio length in seconds (max ~47s for stable-audio-open-1.0)", - ) - parser.add_argument( - "--audio_start", - type=float, - default=0.0, - help="Audio start time in seconds", - ) - parser.add_argument( - "--negative_prompt", - default="Low quality", - help="Negative prompt for classifier-free guidance", - ) - parser.add_argument( - "--guidance_scale", - type=float, - default=7.0, - help="Guidance scale for diffusion (higher = more adherence to prompt)", - ) - parser.add_argument( - "--num_inference_steps", - type=int, - default=100, - help="Number of inference steps (higher = better quality, slower)", - ) - parser.add_argument( - "--seed", - type=int, - default=None, - help="Random seed for reproducibility", - ) - parser.add_argument( - "--output", - default="stable_audio_output.wav", - help="Output file path", - ) - parser.add_argument( - "--response_format", - default="wav", - choices=["wav", "mp3", "flac", "pcm"], - help="Audio output format", - ) - return parser.parse_args() - - -def generate_audio(args): - """Generate audio using the API.""" - - # Build request payload - payload = { - "input": args.text, - "audio_length": args.audio_length, - "audio_start": args.audio_start, - "response_format": args.response_format, - } - - # Add optional parameters - if args.negative_prompt: - payload["negative_prompt"] = args.negative_prompt - if args.guidance_scale: - payload["guidance_scale"] = args.guidance_scale - if args.num_inference_steps: - payload["num_inference_steps"] = args.num_inference_steps - if args.seed is not None: - payload["seed"] = args.seed - - print(f"\n{'=' * 60}") - print("Stable Audio - Text-to-Audio Generation") - print(f"{'=' * 60}") - print(f"API URL: {args.api_url}") - print(f"Prompt: {args.text}") - print(f"Audio length: {args.audio_length}s") - print(f"Negative prompt: {args.negative_prompt}") - print(f"Guidance scale: {args.guidance_scale}") - print(f"Inference steps: {args.num_inference_steps}") - if args.seed is not None: - print(f"Seed: {args.seed}") - print(f"Output: {args.output}") - print(f"{'=' * 60}\n") - - try: - # Make the API request - print("Generating audio...") - response = requests.post( - args.api_url, - json=payload, - headers={"Content-Type": "application/json"}, - timeout=300, # 5 minute timeout for long generations - ) - - # Check for errors - if response.status_code != 200: - print(f"Error: API returned status code {response.status_code}") - print(f"Response: {response.text}") - return False - - # Save the audio - with open(args.output, "wb") as f: - f.write(response.content) - - print(f"✓ Audio saved to {args.output}") - print(f" File size: {len(response.content) / 1024:.1f} KB") - return True - - except requests.exceptions.Timeout: - print("Error: Request timed out. Try reducing inference steps or audio length.") - return False - except requests.exceptions.ConnectionError: - print(f"Error: Could not connect to {args.api_url}") - print("Make sure the server is running.") - return False - except Exception as e: - print(f"Error: {e}") - return False - - -def main(): - args = parse_args() - success = generate_audio(args) - sys.exit(0 if success else 1) - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/step_audio2/README.md b/examples/online_serving/step_audio2/README.md index 25e2a6785..ba2797dae 100644 --- a/examples/online_serving/step_audio2/README.md +++ b/examples/online_serving/step_audio2/README.md @@ -189,5 +189,5 @@ Async chunk reduces TTFP by **67%** by streaming audio token chunks from Thinker - Check server logs for errors ### Out of memory -- Reduce `gpu_memory_utilization` in stage configs +- Reduce `gpu_memory_utilization` in the deploy config - Use a smaller batch size diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 1ea191a37..69a50540b 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -25,7 +25,6 @@ For the full list of supported architectures across all modalities, see | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ | -| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) | CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead. @@ -570,7 +569,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "response_format": "pcm" }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` -Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. +Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. ### Streaming WebSocket The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: @@ -715,65 +714,3 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). - Voice cloning is gated upstream and may require a recent `mistral_common`. - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now. - ---- - -## SoulX-Singer - -Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`). - -### Prerequisites - -Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory. - -### Launch - -```bash -# SVS (default) -export MODEL=/path/to/SoulX-Singer -export PREPROCESS=/path/to/SoulX-Singer-Preprocess -bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh - -# SVC -export MODE=svc -export MODEL=/path/to/SoulX-Singer-svc -bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh -``` - -Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly: -```bash -export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS -vllm serve $MODEL --omni \ - --deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \ - --port 8192 --trust-remote-code --enforce-eager -``` - -### Sending requests - -Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`. - -```bash -# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3 -python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \ - --prompt-audio /path/on/server/zh_prompt.mp3 \ - --target-audio /path/on/server/music.mp3 \ - --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \ - -o output.wav -``` - -Use precomputed metadata to skip online preprocess with following command: -```bash -python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \ - --prompt-metadata-path /path/on/server/zh_prompt.json \ - --target-metadata-path /path/on/server/music.json \ - --audio-path /path/on/server/zh_prompt.mp3 \ - -o output.wav -``` - -`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`. - -### Notes - -- Output: 24 kHz mono WAV; batch only. -- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on. -- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only. diff --git a/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py b/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py deleted file mode 100644 index 91035e625..000000000 --- a/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py +++ /dev/null @@ -1,186 +0,0 @@ -#!/usr/bin/env python3 -"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC). - -Sends prompt audio via ``input_audio`` and target accompaniment via -``extra_args['target_audio']`` (server-local path). For integrated preprocess, -also pass ``preprocess_weights_dir`` in ``extra_args``. - -Usage: - python openai_chat_client.py \\ - --prompt-audio /path/on/server/zh_prompt.mp3 \\ - --target-audio /path/on/server/music.mp3 \\ - --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\ - -o output.wav -""" - -from __future__ import annotations - -import argparse -import base64 -import io -import sys -from pathlib import Path - -import requests -import soundfile -import torch - - -def _audio_to_data_url(path: Path) -> str: - with path.open("rb") as handle: - data = base64.b64encode(handle.read()).decode("ascii") - return f"data:audio/mpeg;base64,{data}" - - -def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None: - audio = audio.to(torch.float32) - peak = audio.abs().max().clamp(min=1e-8) - audio = audio / peak - path.parent.mkdir(parents=True, exist_ok=True) - soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16") - - -def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]: - for choice in body.get("choices", []): - audio_obj = choice.get("message", {}).get("audio") - if isinstance(audio_obj, dict) and audio_obj.get("data"): - data, sr = soundfile.read( - io.BytesIO(base64.b64decode(audio_obj["data"])), - dtype="float32", - always_2d=True, - ) - return torch.from_numpy(data).transpose(0, 1), sr - brief = {k: v for k, v in body.items() if k != "choices"} - raise RuntimeError(f"no audio in response message.audio: {brief}") - - -def main() -> int: - repo_root = Path(__file__).resolve().parents[4] - default_assets = repo_root / "tests" / "assets" / "soulxsinger" - - parser = argparse.ArgumentParser(description="SoulX-Singer online chat client") - parser.add_argument("--port", type=int, default=8192) - parser.add_argument("--model", default="Soul-AILab/SoulX-Singer") - parser.add_argument( - "--prompt-audio", - default=str(default_assets / "zh_prompt.mp3"), - help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)", - ) - parser.add_argument( - "--target-audio", - default=str(default_assets / "music.mp3"), - help="Target accompaniment path on the server (extra_args['target_audio'])", - ) - parser.add_argument( - "--prompt-metadata-path", - default=None, - help="SVS precomputed prompt metadata.json", - ) - parser.add_argument( - "--target-metadata-path", - default=None, - help="SVS precomputed target metadata.json", - ) - parser.add_argument( - "--audio-path", - default=None, - help="SVS prompt vocal wav for precomputed metadata", - ) - parser.add_argument("--preprocess-weights-dir", default=None) - parser.add_argument("--output", "-o", default="soulxsinger_out.wav") - parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs") - parser.add_argument("--language", default="Mandarin") - parser.add_argument("--num-inference-steps", type=int, default=32) - parser.add_argument("--guidance-scale", type=float, default=3.0) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Optional CFM seed. Omit for non-deterministic sampling.", - ) - parser.add_argument( - "--auto-shift", - action=argparse.BooleanOptionalAction, - default=True, - help="Auto pitch shift (default: on, original upstream infer.sh)", - ) - parser.add_argument( - "--control", - default="melody", - choices=["melody", "score"], - help="SVS control mode", - ) - parser.add_argument("--vocal-sep", action="store_true") - args = parser.parse_args() - - meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path) - if any(meta_paths) and not all(meta_paths): - print( - "ERROR: precomputed metadata requires --prompt-metadata-path, " - "--target-metadata-path, and --audio-path together.", - file=sys.stderr, - ) - return 2 - - extra_args: dict = { - "vocal_sep": args.vocal_sep, - "auto_shift": args.auto_shift, - "pitch_shift": 0, - } - if all(meta_paths): - extra_args.update( - { - "prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()), - "target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()), - "audio_path": str(Path(args.audio_path).expanduser().resolve()), - } - ) - content = [{"type": "text", "text": "soulx-singer"}] - else: - prompt_path = Path(args.prompt_audio).expanduser().resolve() - if not prompt_path.is_file(): - print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr) - return 2 - extra_args["prompt_audio"] = str(prompt_path) - extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve()) - if args.preprocess_weights_dir: - extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve()) - content = [ - {"type": "text", "text": "soulx-singer"}, - { - "type": "input_audio", - "input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"}, - }, - ] - if not args.svc: - extra_args["language"] = args.language - extra_args["control"] = args.control - - payload = { - "model": args.model, - "modalities": ["audio"], - "messages": [{"role": "user", "content": content}], - "num_inference_steps": args.num_inference_steps, - "guidance_scale": args.guidance_scale, - "extra_args": extra_args, - } - if args.seed is not None: - payload["seed"] = args.seed - - print(f"POST http://localhost:{args.port}/v1/chat/completions") - response = requests.post( - f"http://localhost:{args.port}/v1/chat/completions", - headers={"Content-Type": "application/json"}, - json=payload, - timeout=1800, - ) - response.raise_for_status() - audio, sample_rate = _decode_audio_from_response(response.json()) - _save_wav(audio, Path(args.output), sample_rate) - duration = audio.shape[-1] / sample_rate - print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/examples/online_serving/text_to_speech/soulxsinger/run_server.sh b/examples/online_serving/text_to_speech/soulxsinger/run_server.sh deleted file mode 100755 index 3e20de39b..000000000 --- a/examples/online_serving/text_to_speech/soulxsinger/run_server.sh +++ /dev/null @@ -1,39 +0,0 @@ -#!/bin/bash -# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline). -# -# Usage: -# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \ -# -# Audio paths in client extra_args must be readable on the server host. - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" - -MODEL="${MODEL:-Soul-AILab/SoulX-Singer}" -MODE="${MODE:-svs}" -PORT="${PORT:-8192}" -GPUS="${GPUS:-0}" - -if [[ "$MODE" == "svc" ]]; then - DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}" -else - DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}" -fi - -echo "Starting SoulX-Singer server" -echo " MODEL=$MODEL" -echo " MODE=$MODE" -echo " PORT=$PORT" -echo " DEPLOY_CONFIG=$DEPLOY_CONFIG" -echo " CUDA_VISIBLE_DEVICES=$GPUS" - -CUDA_VISIBLE_DEVICES="$GPUS" \ -vllm serve "$MODEL" \ - --omni \ - --deploy-config "$DEPLOY_CONFIG" \ - --host 0.0.0.0 \ - --port "$PORT" \ - --trust-remote-code \ - --enforce-eager diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md index 51512fe0c..cf303ccb1 100644 --- a/examples/online_serving/text_to_video/README.md +++ b/examples/online_serving/text_to_video/README.md @@ -11,6 +11,11 @@ This example demonstrates how to deploy text-to-video models for online video ge | Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | | LTX-2 | `Lightricks/LTX-2` | | Helios (Base / Mid / Distilled) | `BestWishYsh/Helios-Base`, `Helios-Mid`, `Helios-Distilled` | +| SANA-Video 2B (480p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | +| SANA-Video 2B (720p) | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | + +For SANA-Video 2B text-to-video online serving, see the +[SANA-Video 2B recipe](../../../recipes/NVIDIA/SANA-Video-2B.md). ## Wan2.2 T2V diff --git a/examples/online_serving/text_to_video/run_curl_sana_video.sh b/examples/online_serving/text_to_video/run_curl_sana_video.sh new file mode 100755 index 000000000..490a743e5 --- /dev/null +++ b/examples/online_serving/text_to_video/run_curl_sana_video.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +BASE_URL="${BASE_URL:-http://localhost:8091}" +OUTPUT_PATH="${OUTPUT_PATH:-sana_video.mp4}" + +curl -sS -X POST "$BASE_URL/v1/videos/sync" \ + -F "prompt=A cinematic tracking shot of a sailboat crossing the ocean at sunset" \ + -F "height=480" \ + -F "width=832" \ + -F "num_frames=81" \ + -F "fps=16" \ + -F "num_inference_steps=50" \ + -F "guidance_scale=6.0" \ + -F "seed=42" \ + -F 'extra_params={"motion_score":30}' \ + --output "$OUTPUT_PATH" diff --git a/examples/online_serving/text_to_video/run_server_sana_video.sh b/examples/online_serving/text_to_video/run_server_sana_video.sh new file mode 100755 index 000000000..446a75239 --- /dev/null +++ b/examples/online_serving/text_to_video/run_server_sana_video.sh @@ -0,0 +1,11 @@ +#!/bin/bash +# SANA-Video-2B text-to-video serving. + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8091}" + +vllm serve "$MODEL" \ + --omni \ + --model-class-name SanaVideoPipeline \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh b/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh new file mode 100755 index 000000000..adfb125c0 --- /dev/null +++ b/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# SANA-Video through the black-box Diffusers backend adapter. + +MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}" +PORT="${PORT:-8091}" + +vllm serve "$MODEL" \ + --omni \ + --diffusion-load-format diffusers \ + --diffusion-attention-backend TORCH_SDPA \ + --dtype bfloat16 \ + --port "$PORT" diff --git a/mkdocs.yml b/mkdocs.yml index 1a3b5109f..8a0e322e3 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -79,8 +79,13 @@ plugins: user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md user_guide/quantization/index.md: user_guide/quantization/overview.md + configuration/pd_disaggregation.md: features/pd_disaggregation.md design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md + design/feature/cpu_offloading.md: design/feature/offloader/README.md + design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md + design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md + design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md - glightbox - git-revision-date-localized: # exclude files diff --git a/pyproject.toml b/pyproject.toml index 0e91bfae6..92dcb1a05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,8 +40,11 @@ dev = [ "pytest-asyncio==1.4.0", "pytest-cov==7.1.0", "pytest-mock==3.15.1", + "pytest-shard==0.1.2", + "pytest-xdist==3.8.0", "websockets==16.1", "pytest-subtests==0.15.0", + "filelock>=3.16.1,<4", "datasets==3.6.0", "mypy==1.11.1", "pre-commit==4.0.1", @@ -68,9 +71,11 @@ dev = [ "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests "torchdiffeq==0.2.5", # for covo_audio Test "FlagEmbedding==1.4.0", # for hunyuan_image3 tests - "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0) "s3tokenizer==0.3.0", # for tts test + # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo, + # whose top-level stepaudio2 package conflicts with this distribution. "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml) + "silero-vad==6.1", # Realtime server-VAD duplex E2E ] demo = [ @@ -80,6 +85,10 @@ demo = [ "cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次") ] +server-vad = [ + "silero-vad==6.1", +] + forced-aligner = [ "qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner ] @@ -132,13 +141,6 @@ docs = [ "pydantic", ] -soulx-svs = [ - "funasr>=1.0.0", - "g2pM>=0.1.2.5", - "g2p-en>=2.1.0", - "ToJyutping>=3.2.0", -] - [project.urls] Homepage = "https://github.com/vllm-project/vllm-omni" Repository = "https://github.com/vllm-project/vllm-omni" @@ -189,6 +191,7 @@ select = [ "N", # pep8-naming "UP", # pyupgrade "TID251", # flake8-tidy-imports.banned-api + "TC004", # flake8-type-checking: runtime-used imports ] ignore = [ "E203", # whitespace before ':' (conflicts with black) @@ -224,19 +227,14 @@ ignore = [ "torch.cuda.synchronize".msg = "Use torch.accelerator.synchronize" [tool.mypy] -python_version = "3.12, 3.13" -warn_return_any = true -warn_unused_configs = true -disallow_untyped_defs = true -disallow_incomplete_defs = true +python_version = "3.12" +# Align with vLLM: silent follow for library code; missing third-party stubs +# are ignored everywhere (including tests). tests/ still get +# --follow-imports skip from tools/pre_commit/mypy.py. +plugins = ["pydantic.mypy"] +ignore_missing_imports = true check_untyped_defs = true -disallow_untyped_decorators = true -no_implicit_optional = true -warn_redundant_casts = true -warn_unused_ignores = true -warn_no_return = true -warn_unreachable = true -strict_equality = true +follow_imports = "silent" [tool.pytest.ini_options] asyncio_mode = "auto" @@ -246,14 +244,15 @@ python_classes = ["Test*"] python_functions = ["test_*"] addopts = [ "--strict-markers", - "--strict-config" + "--strict-config", + "--dist=loadgroup" ] markers = [ - # ci/cd required - "core_model: L1&L2 tests (run in each PR)", - "advanced_model: L3 level tests (run on each merge)", - "full_model: L4 level tests (run nightly)", - "local_model: tests (run locally)", + # ci/cd required — tag with [ci-level] for tests.helpers.mark.get_level_markers(). + "core_model: [ci-level] L1&L2 tests (run in each PR)", + "advanced_model: [ci-level] L3 level tests (run on each merge)", + "full_model: [ci-level] L4 level tests (run nightly)", + "local_model: [ci-level] tests (run locally)", # function module markers "diffusion: Diffusion model tests", "omni: Omni model tests", @@ -263,44 +262,64 @@ markers = [ "sp: Sequence parallelism tests", "example: Doc example code tests", # platform markers - "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)", - "xpu: Tests that run on XPU (auto-added)", - "npu: Tests that run on NPU/Ascend (auto-added)", - "musa: Tests that run on MUSA/Moore Threads (auto-added)", + # Tag with [hardware-platform] for tests.helpers.mark.get_supported_platforms(). + # Add [gpu] on accelerator platforms that should also get pytest.mark.gpu. + # cpu/gpu are pytest marks; hardware_test(res=...) rejects them as keys. + "cpu: [hardware-platform] Tests that run on CPU", + "gpu: [hardware-platform] Tests that run on GPU (auto-added)", + "cuda: [hardware-platform] [gpu] Tests that run on CUDA (auto-added)", + "rocm: [hardware-platform] [gpu] Tests that run on AMD/ROCm (auto-added)", + "xpu: [hardware-platform] [gpu] Tests that run on XPU (auto-added)", + "npu: [hardware-platform] Tests that run on NPU/Ascend (auto-added)", + "musa: [hardware-platform] [gpu] Tests that run on MUSA/Moore Threads (auto-added)", # specified computation resources marks (auto-added) - # Tag with [hardware-resource] for tests.helpers.mark.get_hardware_mark_list(). - "H100: [hardware-resource] Tests that require H100 GPU", - "H200: [hardware-resource] Tests that require H200 GPU", - "L4: [hardware-resource] Tests that require L4 GPU", - "B200: [hardware-resource] Tests that require B200 GPU", - "MI325: [hardware-resource] Tests that require MI325 GPU (AMD/ROCm)", - "B60: [hardware-resource] Tests that require Intel Arc Pro B60 XPU", - "S5000: [hardware-resource] Tests that require S5000 GPU (Moore Threads/MUSA)", - "A2: [hardware-resource] Tests that require A2 NPU", - "A3: [hardware-resource] 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_xpu: Tests that require multi cards on XPU platform", - "distributed_npu: Tests that require multi cards on NPU platform", - "distributed_musa: Tests that require multi cards on MUSA 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_xpu: Skip if the num of XPU cards is less than the required", - "skipif_npu: Skip if the num of NPU cards is less than the required", - "skipif_musa: Skip if the num of MUSA cards is less than the required", + # [hardware-resource] + [cuda]/[rocm]/[xpu]/[npu]/[musa] for hardware_test(res=...). + # Do not put [gpu] on SKU lines: that would make gpu a res key. + "H100: [hardware-resource] [cuda] Tests that require H100 GPU", + "H800: [hardware-resource] [cuda] Tests that require H800 GPU", + "H200: [hardware-resource] [cuda] Tests that require H200 GPU", + "L4: [hardware-resource] [cuda] Tests that require L4 GPU", + "B200: [hardware-resource] [cuda] Tests that require B200 GPU", + "MI325: [hardware-resource] [rocm] Tests that require MI325 GPU (AMD/ROCm)", + "B60: [hardware-resource] [xpu] Tests that require Intel Arc Pro B60 XPU", + "S5000: [hardware-resource] [musa] Tests that require S5000 GPU (Moore Threads/MUSA)", + "A2: [hardware-resource] [npu] Tests that require A2 NPU", + "A3: [hardware-resource] [npu] Tests that require A3 NPU", + # Tag with [hardware-cards] for tests.helpers.mark.get_supported_card_counts(). + "cards_1: [hardware-cards] Tests that require 1 accelerator card", + "cards_2: [hardware-cards] Tests that require 2 accelerator cards", + "cards_3: [hardware-cards] Tests that require 3 accelerator cards", + "cards_4: [hardware-cards] Tests that require 4 accelerator cards", + "cards_5: [hardware-cards] Tests that require 5 accelerator cards", + "cards_6: [hardware-cards] Tests that require 6 accelerator cards", + "cards_7: [hardware-cards] Tests that require 7 accelerator cards", + "cards_8: [hardware-cards] Tests that require 8 accelerator cards", # more detailed markers "sp: Sequence parallelism tests (multi-GPU)", - "slow: Slow tests (may skip in quick CI)", + "slow: [ci-level] Slow tests (may skip in quick CI)", "benchmark: Benchmark tests", + "xdist: xdist-parallelizable tests", ] filterwarnings = [ "ignore:.*does not have '__test__' attribute.*:UserWarning", "ignore:.*does not have '__bases__' attribute.*:UserWarning", ] +[tool.coverage.run] +core = "sysmon" +source = ["vllm_omni"] +# Online-serving tests launch the server via subprocess.Popen and stop it with +# SIGTERM. pytest-cov no longer instruments subprocesses implicitly, and a +# SIGTERM-ed process discards its data, so without patch/sigterm a coverage run +# over an online test measures only the pytest parent. +patch = ["subprocess"] +concurrency = ["multiprocessing", "thread"] +sigterm = true + +[tool.coverage.report] +show_missing = true +skip_covered = true + [tool.typos.default] extend-ignore-identifiers-re = [ ".*_thw", @@ -327,6 +346,8 @@ extend-ignore-identifiers-re = [ ] extend-ignore-re = [ "\\.\\.\\.n,d->\\.\\.\\.nd", + # setuptools-scm / git-describe local versions, e.g. 0.26.1.dev14+gae6577ea + "\\+g[0-9a-f]{7,}", ] [tool.typos.default.extend-words] ue = "ue" @@ -351,3 +372,5 @@ Leary = "Leary" fo = "fo" ba = "ba" SER = "SER" +# NumPy ndarray.flags.writeable (canonical flag name, not "writable"). +writeable = "writeable" diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md index 4c2d75aca..3b4b6c594 100644 --- a/recipes/Bagel/BAGEL-7B-MoT.md +++ b/recipes/Bagel/BAGEL-7B-MoT.md @@ -57,14 +57,14 @@ move the diffusion stage to a second GPU in a custom deploy config. #### Offline Commands Run text-to-image with the shared offline example from the repository root. -Pick the topology with `--stage-configs-path` and forward BAGEL-specific +Pick the topology with `--deploy-config` and forward BAGEL-specific generation parameters as a JSON object through `--extra-body`: ```bash # Two-stage (Thinker + DiT), shares one GPU by default python examples/offline_inference/text_to_image/text_to_image.py \ --model ByteDance-Seed/BAGEL-7B-MoT \ - --stage-configs-path vllm_omni/deploy/bagel.yaml \ + --deploy-config vllm_omni/deploy/bagel.yaml \ --prompt "A beautiful sunset over mountains" \ --height 1024 \ --width 1024 \ @@ -79,7 +79,7 @@ python examples/offline_inference/text_to_image/text_to_image.py \ # Single-stage (DiT only, with internal LLM/ViT/VAE) python examples/offline_inference/text_to_image/text_to_image.py \ --model ByteDance-Seed/BAGEL-7B-MoT \ - --stage-configs-path vllm_omni/deploy/bagel_single_stage.yaml \ + --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \ --prompt "A beautiful sunset over mountains" \ --height 1024 \ --width 1024 \ diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index b9bff789f..b1edf0b4e 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -1,11 +1,14 @@ # Boogu-Image -> Text-to-image and image-editing online serving (Boogu-Image-0.1-Base / -Edit) +> Text-to-image and image-editing online serving +(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo) ## Summary - Vendor: Boogu -- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing) +- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), + `Boogu/Boogu-Image-0.1-Edit` (image editing), and + `Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing) - Task: Text-to-image generation and text-guided image editing (TI2I) - Mode: Online serving with the OpenAI-compatible API - Maintainer: Community @@ -13,19 +16,23 @@ ## When to use this recipe Use this recipe when you want a known-good starting point for serving -`Boogu/Boogu-Image-0.1-Base` with vLLM-Omni's native pipeline (no +`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or +`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no `--diffusion-load-format diffusers`, and the upstream `boogu` package is not required). Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model -family. This recipe covers the Base text-to-image checkpoint, which pairs a -Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match -Euler scheduler with time-shift. It handles photorealistic generation and -Chinese/English text rendering. +family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder +with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with +time-shift. It handles photorealistic generation and Chinese/English text +rendering. The Edit and Edit-Turbo checkpoints use the same native +`BooguImagePipeline` for image editing. ## References - Upstream model card: +- Edit model card: +- Edit-Turbo 1K hotfix: - Project page: - GitHub: - Related example: [`examples/online_serving/text_to_image/`](../../examples/online_serving/text_to_image/README.md) @@ -93,7 +100,7 @@ curl -s http://localhost:8091/v1/chat/completions \ - **Memory usage:** ~34.6 GiB on GPU. Use `--vae-use-slicing --vae-use-tiling` to trim peak VRAM if needed. - **Key flags:** - - `--omni` — enables vLLM-Omni diffusion serving. + - `--omni` — enables vLLM-Omni diffusion serving. - **Guidance:** Boogu-Image uses `guidance_scale` (mapped to the upstream `text_guidance_scale`); the default is `4.0`. Classifier-free guidance is active whenever `guidance_scale > 1.0`. @@ -111,13 +118,13 @@ the image-editing (TI2I) path activates automatically when a request carries a reference image. The Base text-to-image path is unaffected (no reference image is sent). -#### Command +### Command ```bash vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091 ``` -#### Verification +### Verification Edit an image with `/v1/images/edits` (the model-card example — change a photo to a colored-pencil drawing). Diffusion parameters are plain multipart form @@ -157,14 +164,14 @@ curl -s http://localhost:8091/v1/chat/completions \ }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited.png ``` -#### Notes +### Notes - **Single reference image:** only one input image is supported for now (the upstream "Only support 1 reference image for now" limit). - **Guidance semantics:** - - `guidance_scale` = text guidance (upstream `text_guidance_scale`, default + - `guidance_scale` = text guidance (upstream `text_guidance_scale`, default `4.0`); `> 1.0` enables text CFG. Editing typically uses `5.0`. - - `guidance_scale_2` = image guidance (upstream `image_guidance_scale`, + - `guidance_scale_2` = image guidance (upstream `image_guidance_scale`, default `1.0` = off). Setting it `> 1.0` enables the double-guidance path (3 model predictions per step), steering more strongly toward the reference image. @@ -178,3 +185,69 @@ curl -s http://localhost:8091/v1/chat/completions \ supports Boogu-Image-Edit directly (`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional `--guidance-scale-2`). + +## Fast image editing (Boogu-Image-0.1-Edit-Turbo) + +Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json` +declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit. +The upstream 1K hotfix is recommended for stable results and is pinned below +to avoid silently loading the older checkpoint from the repository's default +revision. + +### Command + +```bash +vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \ + --omni \ + --revision hotfix-1k-20260708 \ + --port 8091 +``` + +### Verification + +Edit an image with the distilled four-step settings: + +```bash +curl -s http://localhost:8091/v1/images/edits \ + -F model="Boogu/Boogu-Image-0.1-Edit-Turbo" \ + -F image="@input.png" \ + -F prompt="Change the style to a colored pencil drawing." \ + -F num_inference_steps=4 \ + -F guidance_scale=1.0 \ + -F seed=42 \ + | jq -r '.data[0].b64_json' | base64 -d > edited-turbo.png +``` + +Or via chat completions: + +```bash +curl -s http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + {"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,"}}, + {"type": "text", "text": "Change the style to a colored pencil drawing."} + ]} + ], + "extra_body": { + "num_inference_steps": 4, + "guidance_scale": 1.0, + "seed": 42 + } + }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited-turbo.png +``` + +### Notes + +- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K + hotfix over the 1.5K variant for more stable results. +- **Recommended settings:** `num_inference_steps=4` and + `guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance + behavior is already baked into the distilled checkpoint. A scale of `1.0` + avoids applying additional classifier-free guidance; unlike the regular Edit + checkpoint, it should not be increased to `5.0`. +- **Reference images:** the same single-reference and `align_res` behavior as + the regular Edit checkpoint applies. +- **Known limitations:** the same single-GPU limitations as Base and Edit apply; + CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated. diff --git a/recipes/BosonAI/Higgs-Audio-V3-TTS.md b/recipes/BosonAI/Higgs-Audio-V3-TTS.md index b924ac472..0391fe03c 100644 --- a/recipes/BosonAI/Higgs-Audio-V3-TTS.md +++ b/recipes/BosonAI/Higgs-Audio-V3-TTS.md @@ -45,7 +45,7 @@ MusicGen-style delay pattern). **Online serving:** ```bash -vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ +vllm serve bosonai/higgs-audio-v3-tts-4b \ --host 0.0.0.0 --port 8095 \ --trust-remote-code --omni ``` diff --git a/recipes/GLM/GLM-Image.md b/recipes/GLM/GLM-Image.md index bd11558c9..df6b512cd 100644 --- a/recipes/GLM/GLM-Image.md +++ b/recipes/GLM/GLM-Image.md @@ -52,7 +52,7 @@ Start the server from the repository root: vllm serve zai-org/GLM-Image --omni --port 8091 ``` -To use the bundled stage config explicitly (same default as above): +To use the bundled deploy config explicitly (same default as above): ```bash vllm serve zai-org/GLM-Image \ @@ -106,7 +106,7 @@ Overall summary from the run’s metrics. Rough wall-time split: **Stage 0 (AR)* #### Notes - Memory usage: Roughly **~38 GiB + KV** on Stage 0 (AR) and **~20 GiB** on Stage 1 (DiT+VAE) per the user guide; two 80 GB cards match the default split. -- Key flags: `--omni` is required; `--stage-configs-path` is optional unless you use a custom YAML (for example single-GPU). +- Key flags: `--omni` is required; `--deploy-config` is optional unless you use a custom YAML (for example single-GPU). - Keep **Transformers ≥ 5.5.1** (this recipe used **5.5.4**) so `glm_image` configs resolve; otherwise Stage 0 can fail at `ModelConfig` validation. - Known limitations: This starter recipe follows the dual-GPU online path documented under `examples/online_serving/glm_image`. The first request may be slower due to warmup. - Generation time: about **61 s** wall time end-to-end for the sample above (50 inference steps, 1024×1024). diff --git a/recipes/Helios/Helios.md b/recipes/Helios/Helios.md index fdc0660bb..f1c11ed92 100644 --- a/recipes/Helios/Helios.md +++ b/recipes/Helios/Helios.md @@ -79,7 +79,7 @@ Total generation time: seconds ( ms) Saved generated video to helios_t2v_base.mp4 ``` -#### Important flags and stage configs +#### Important flags and deploy config - `--model BestWishYsh/Helios-Base` selects the base Helios checkpoint used by this recipe. @@ -88,7 +88,7 @@ Saved generated video to helios_t2v_base.mp4 - `--cache-backend cache_dit` enables the cache-dit acceleration path. - `--enable-cache-dit-summary` prints cache-dit summary information after diffusion forward passes. -- No separate stage config is required for this offline recipe; the shared +- No separate deploy config is required for this offline recipe; the shared `text_to_video.py` example configures the pipeline through its arguments. - Helios-specific knobs (declared in `vllm_omni/model_extras/helios.py`) are passed via the generic `--extra-body` JSON flag: diff --git a/recipes/HiDream-ai/HiDream-O1-Image.md b/recipes/HiDream-ai/HiDream-O1-Image.md new file mode 100644 index 000000000..1864324b2 --- /dev/null +++ b/recipes/HiDream-ai/HiDream-O1-Image.md @@ -0,0 +1,113 @@ +# HiDream-O1-Image + +> HiDream-O1 text-to-image generation through the shared offline example + +## Summary + +- Vendor: HiDream.ai +- Model: `HiDream-ai/HiDream-O1-Image` +- Task: Text-to-image +- Mode: Offline inference +- Maintainer: Community + +## When to use this recipe + +Use this recipe to generate native-resolution images with HiDream-O1 on one +or two CUDA GPUs. The two-GPU command shards the Qwen3-VL attention and MLP +weights, as well as the timestep and pixel projections. + +## References + +- Upstream model: + [`HiDream-ai/HiDream-O1-Image`](https://huggingface.co/HiDream-ai/HiDream-O1-Image) +- Related offline example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Related issue: + [`vllm-project/vllm-omni#3479`](https://github.com/vllm-project/vllm-omni/issues/3479) + +## Hardware Support + +## GPU + +### 1x H100 80GB + +#### Environment + +- OS: Linux +- Python: Match the repository requirements for your checkout +- Driver / runtime: NVIDIA CUDA +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model HiDream-ai/HiDream-O1-Image \ + --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \ + --height 2048 \ + --width 2048 \ + --num-inference-steps 50 \ + --guidance-scale 5.0 \ + --seed 42 \ + --output hidream_o1_output.png +``` + +#### Verification + +```bash +python -c "from PIL import Image; im = Image.open('hidream_o1_output.png'); print(im.mode, im.size)" +``` + +Expected output: + +```text +RGB (2048, 2048) +``` + +#### Notes + +- The pipeline supports BF16 text-to-image generation. +- Requested sizes are matched to the closest resolution supported by the checkpoint. + +### 2x H100 80GB + +#### Environment + +- OS: Linux +- Python: Match the repository requirements for your checkout +- Driver / runtime: NVIDIA CUDA +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model HiDream-ai/HiDream-O1-Image \ + --tensor-parallel-size 2 \ + --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \ + --height 2048 \ + --width 2048 \ + --num-inference-steps 50 \ + --guidance-scale 5.0 \ + --seed 42 \ + --output hidream_o1_tp2_output.png +``` + +#### Verification + +```bash +python -c "from PIL import Image; im = Image.open('hidream_o1_tp2_output.png'); print(im.mode, im.size)" +``` + +Expected output: + +```text +RGB (2048, 2048) +``` + +#### Notes + +- Tensor parallelism shards the attention, MLP, timestep, and pixel projection weights across both GPUs. +- Image editing and image-conditioned generation are not supported. diff --git a/recipes/LTX/LTX-2.5.md b/recipes/LTX/LTX-2.5.md new file mode 100644 index 000000000..f0be6948b --- /dev/null +++ b/recipes/LTX/LTX-2.5.md @@ -0,0 +1,217 @@ +# LTX-2.5 + +> Text-to-video and first-frame image-to-video generation with synchronized audio + +## Overview + +vLLM-Omni supports the gated +[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers) +checkpoint for T2V and first-frame I2V. Every output includes synchronized +48 kHz stereo audio. + +The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5) +repository is not directly loadable with `--model`. It supplies the canonical +Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by +the Full/SFT two-stage pipeline, so accept both model licenses and authenticate +before first use. + +The vLLM-Omni integration code is provided under the project's +[Apache-2.0 license](../../LICENSE). The LTX-2.5 model weights are separate +artifacts covered by the +[LTX-2.x Community License](https://github.com/Lightricks/LTX-2/blob/main/LICENSE.md), +not the vLLM-Omni license. Users are responsible for reviewing and complying +with the model license before downloading or using the weights. + +## Choose a pipeline + +| Pipeline | Mode | Default size | Schedule | +| --- | --- | ---: | ---: | +| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps | +| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps | +| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps | +| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps | + +Both two-stage pipelines generate at half resolution, apply the official x2 +latent upsampler, and run a three-step refinement stage. Select the class with +`--model-class-name`; no `--task-type` flag is required. Supplying one initial +image selects I2V, while omitting it selects T2V. + +## Video decoder + +LTX-2.5 uses the canonical Native Diffusion VAE decoder by default. To opt in +to the legacy convolutional VAE, set the startup-only stage-0 model extra +`ltx2_use_conv_vae: true`. The choice applies to all requests and all four +pipeline classes above. DiffVAE loads its NATTEN kernel from Hugging Face Hub +during startup. `decode_timestep` and `decode_noise_scale` condition only the +legacy ConvVAE and have no effect when DiffVAE is selected. + +For offline Python usage: + +```python +omni = Omni( + model="Lightricks/LTX-2.5-Diffusers", + model_class_name="LTX2Pipeline", + stage_overrides='{"0":{"extras":{"ltx2_use_conv_vae":true}}}', +) +``` + +For online serving: + +```bash +vllm serve Lightricks/LTX-2.5-Diffusers \ + --omni \ + --model-class-name LTX2Pipeline \ + --stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}' +``` + +Both decoders are untiled by default. Set `vae_use_tiling` for memory-saving +serial tiling; DiffVAE tiles only above 80 frames or 768 pixels in either spatial +dimension. DiffVAE also supports distributed VAE decode; see the +[VAE Parallelism Guide](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md). +DiffVAE is decoder-only, so I2V still uses the convolutional VAE encoder. + +## Prerequisites + +```bash +hf auth login +export MODEL=Lightricks/LTX-2.5-Diffusers +``` + +Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and +`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the +`libx264` encoder. + +## Hardware + +| GPU | Status | Recommended scope | +| --- | --- | --- | +| NVIDIA B300 | Verified | All four canonical pipelines | +| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines | +| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines | +| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 | + +The 1920x1088 two-stage examples require about 114 GB of peak GPU memory. +80 GB GPUs do not have enough safety margin for the canonical two-stage +configuration; use a larger GPU or reduce the output configuration. + +### H100 80 GB FP8 + +Use the distilled one-stage pipeline with FP8 and cuDNN attention: + +```bash +vllm serve Lightricks/LTX-2.5-Diffusers \ + --omni \ + --model-class-name LTX2DistilledOneStagePipeline \ + --quantization fp8 \ + --diffusion-attention-backend CUDNN_ATTN \ + --host 0.0.0.0 \ + --port 8000 +``` + +Use the T2V or I2V request below after the server is ready. + +## Offline inference + +Choose values from the pipeline table. For example, the distilled two-stage +path uses: + +```bash +export PIPELINE=LTX2DistilledTwoStagePipeline +export WIDTH=1920 +export HEIGHT=1088 +export STEPS=8 +``` + +### Text-to-video + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model "${MODEL}" \ + --model-class-name "${PIPELINE}" \ + --prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \ + --width "${WIDTH}" \ + --height "${HEIGHT}" \ + --num-frames 121 \ + --num-inference-steps "${STEPS}" \ + --frame-rate 24 \ + --fps 24 \ + --seed 42 \ + --output ltx25-t2v.mp4 +``` + +### First-frame image-to-video + +```bash +python examples/offline_inference/image_to_video/image_to_video.py \ + --model "${MODEL}" \ + --model-class-name "${PIPELINE}" \ + --image /absolute/path/to/first-frame.png \ + --prompt "The red fox walks forward while the camera tracks alongside." \ + --width "${WIDTH}" \ + --height "${HEIGHT}" \ + --num-frames 121 \ + --num-inference-steps "${STEPS}" \ + --frame-rate 24 \ + --fps 24 \ + --seed 42 \ + --output ltx25-i2v.mp4 +``` + +LTX-2.5 uses the official CRF-18 first-frame conditioning path by default. + +## Online serving + +Start one server for the selected pipeline: + +```bash +vllm serve "${MODEL}" \ + --omni \ + --model-class-name "${PIPELINE}" \ + --host 0.0.0.0 \ + --port 8000 \ + --stage-init-timeout 900 +``` + +T2V request: + +```bash +curl -sS --fail-with-body \ + -X POST http://127.0.0.1:8000/v1/videos/sync \ + -F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \ + -o ltx25-online-t2v.mp4 +``` + +For I2V, add exactly one first frame: + +```bash +export FIRST_FRAME=/absolute/path/to/first-frame.png + +curl -sS --fail-with-body \ + -X POST http://127.0.0.1:8000/v1/videos/sync \ + -F "input_reference=@${FIRST_FRAME};type=image/png" \ + -F 'prompt=The red fox walks forward while the camera tracks alongside.' \ + -o ltx25-online-i2v.mp4 +``` + +Restart the server after changing `PIPELINE` because each class selects a +different weight layout and execution topology. + +## Constraints + +- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps` + are optional and use the selected pipeline defaults when omitted. `seed` is + optional and random when omitted. +- `num_frames` must be `8k+1` when overridden. +- One-stage width and height must be divisible by 32; two-stage final + dimensions must be divisible by 64. +- Full/SFT pipelines accept negative prompts. Distilled pipelines are + positive-only and reject negative prompts. +- I2V accepts exactly one initial image per prompt. +- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`. + +## References + +- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers) +- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2) +- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py) +- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py) diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index 3142f9fd1..ec4e8fb3b 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -4,11 +4,22 @@ ## Pipelines -| Model | Checkpoint | Task | `--model-class-name` | Batching | -|---|---|---|---|---| -| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes | -| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No | -| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes | +| `--model-class-name` | Task | Required checkpoint repositories | +|---|---|---| +| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` | +| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` | +| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` | +| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` | +| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` | +| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`
`Lightricks/LTX-2.3` | +| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` | +| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`
`Lightricks/LTX-2.3` | + +Repositories in the table are download units. A full pipeline repository +contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler, +and tokenizer; an additional repository supplies LoRA or upsampler sidecars. +The runtime first searches the model root for sidecars under their official +filenames, then downloads them from the matching Lightricks Hub repository. `LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial @@ -16,10 +27,13 @@ image selects I2V. Both one-stage repositories declare this class, so `--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint; the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable. -`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V -or provide one initial image for I2V. As in the official LTX pipelines, -distilled mode is selected explicitly rather than inferred from the checkpoint -name or metadata; always pass `--model-class-name LTX2DistilledPipeline`. +`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples, +then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a +merged distilled Transformer without upsampling, while +`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support +T2V and I2V; select their class explicitly. The deprecated +`LTX2DistilledPipeline` name remains an alias for +`LTX2DistilledTwoStagePipeline`. ## API Migration @@ -43,8 +57,8 @@ The consolidation also removes these registry names without aliases: | `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 | | `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` | | `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 | -| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` | -| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` | +| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` | +| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` | Passing any second positional argument now raises `TypeError`. These changes affect direct Python callers and explicit `--model-class-name` overrides; @@ -68,24 +82,28 @@ The model recipe also supplies the default negative prompt. Top-level API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts default to `121`. -## Distilled Two-Stage Defaults +## Two-Stage Defaults -| Parameter | Value | -|---|---:| -| Final width × height | 1536 × 1024 | -| Stage 1 width × height | 768 × 512 | -| Frames / frame rate | 121 / 24 | -| Stage 1 / Stage 2 denoise steps | 8 / 3 | -| Guidance | Fixed positive-only | +| Parameter | Ordinary | Full-distilled | +|---|---:|---:| +| Final width × height | 1536 × 1024 | 1536 × 1024 | +| Stage 1 width × height | 768 × 512 | 768 × 512 | +| Frames / frame rate | 121 / 24 | 121 / 24 | +| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 | +| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only | -The dimensions passed to the API are the final output dimensions and must be -divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`. -The distilled recipe fixes both sigma schedules; `num_inference_steps` must be -`8` when supplied. +API dimensions are final dimensions and must be divisible by 64. All LTX +requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or +LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so +`num_inference_steps`, when supplied, must be `8`. Both entries reject custom +sigmas and input latents. + +Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer +and automatically switches to dynamic LoRA when quantization is enabled. ## Serving -Start either checkpoint: +Start either one-stage checkpoint: ```bash vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600 @@ -95,11 +113,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600 vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600 ``` -Start the distilled checkpoint with its two-stage entry: +Start a two-stage checkpoint with its explicit class: ```bash vllm serve rootonchair/LTX-2-19b-distilled --omni \ - --model-class-name LTX2DistilledPipeline --stage-init-timeout 600 + --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600 +# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent +vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \ + --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600 +# Ordinary LTX-2.3 two-stage +vllm serve diffusers/LTX-2.3-Diffusers --omni \ + --model-class-name LTX2TwoStagePipeline \ + --enable-layerwise-offload \ + --stage-init-timeout 600 ``` The same server handles T2V and I2V. A T2V request using the selected recipe's @@ -135,8 +161,9 @@ together with `input_reference`. ## Guidance -LTX supports independent video/audio CFG, spatio-temporal guidance (STG), -cross-modality guidance, and rescaling: +One-stage and ordinary Stage 1 support independent video/audio CFG, +spatio-temporal guidance (STG), cross-modality guidance, and rescaling. +Distilled stages and ordinary Stage 2 are fixed positive-only. | Parameter | Default | Effect | Alias | |---|---:|---|---| @@ -172,9 +199,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers text CFG, STG, and cross-modality guidance passes; guidance rescaling is applied after all pass predictions are gathered. -The default one-stage recipe has four Transformer passes per denoise step: -`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced -configurations are therefore: +The default one-stage recipe and ordinary Stage 1 have four Transformer passes +per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod` +(cross-modality). The useful balanced configurations are therefore: | `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes | |---:|---:|---:|---| @@ -195,9 +222,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \ ``` The total device count is the product of `cfg_parallel_size` and the other -configured parallel dimensions. Distilled pipelines use positive-only -guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for -them. +configured parallel dimensions. Positive-only distilled phases do not benefit +from `cfg_parallel_size > 1`. ### Python API @@ -242,22 +268,22 @@ noted below. | `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. | | `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. | | `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. | -| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. | +| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. | | `width` | `int`, `None` | Same precedence and alignment as `height`. | | `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. | | `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. | -| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. | +| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. | | `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. | | `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. | -| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. | +| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. | | `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. | | `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. | | `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. | | `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. | -| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. | -| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. | +| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. | +| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. | | `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. | -| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. | +| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. | | `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. | | `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. | | `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. | @@ -273,13 +299,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`. ### Recipe-Specific Request Capabilities -| Override | One-stage | Distilled two-stage | -|---|---|---| -| Guidance | Supported | Fixed positive-only | -| Negative prompt/embeddings | Supported | Rejected | -| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 | -| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules | -| Video/audio latents | Supported | Rejected | +| Override | One-stage | Ordinary two-stage | Distilled two-stage | +|---|---|---|---| +| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only | +| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected | +| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 | +| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules | +| Video/audio latents | Supported | Rejected | Rejected | These capability checks apply equally to direct `forward` keywords and values in `OmniDiffusionSamplingParams`; unsupported values fail instead of @@ -309,6 +335,10 @@ bundled offline CLI do not currently expose `sigmas`. `--max-num-seqs 1`. - `--cfg-parallel-size` shards the complete LTX guidance plan, including STG, modality guidance, and rescale-compatible prediction gathering. +- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding; + Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a + divisible request shape. +- Cache-DiT is one-stage only; multi-stage configurations are rejected. ## Operational Notes @@ -323,14 +353,15 @@ bundled offline CLI do not currently expose `sigmas`. request parameter. - For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with `tests/dfx/perf/scripts/run_diffusion_benchmark.py`. -- Current two-stage support is limited to the distilled checkpoint; official - LTX two-stage/HQ execution remains out of scope. +- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution + remains out of scope. ## References - - - +- - [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md) - [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md) - [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md) diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md index 7056a6f8f..6dc5f12e9 100644 --- a/recipes/MammothModa2/MammothModa2.md +++ b/recipes/MammothModa2/MammothModa2.md @@ -81,7 +81,7 @@ MammothModa2 generation parameters as a JSON object through `--extra-body`: ```bash python examples/offline_inference/text_to_image/text_to_image.py \ --model ./MammothModa2-Preview \ - --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \ + --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \ --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ --height 1024 \ --width 1024 \ @@ -132,6 +132,44 @@ ls -lh mammoth_t2i.png python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)" ``` +### 1x AMD MI300X, MammothModa2 Preview + +#### Environment + +- OS: Linux 6.8.0-134-generic, x86_64 +- Container: official ROCm image built from `docker/Dockerfile.rocm` +- Python: 3.12.13 +- PyTorch: 2.11.0+gitd0c8b1f +- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211 +- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM +- vLLM version: 0.27.0+rocm723 +- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0` +- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm` + +#### Offline Commands + +The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT: + +```bash +python3 examples/offline_inference/text_to_image/text_to_image.py \ + --model bytedance-research/MammothModa2-Preview \ + --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \ + --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ + --height 1024 \ + --width 1024 \ + --seed 42 \ + --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \ + --enable-diffusion-pipeline-profiler \ + --log-stats \ + --output mammoth_t2i.png +``` + +#### Verification + +The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting. + +The output was a valid 1024 by 1024 RGB PNG. + ## MammothModa2-Dev unified inference MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses diff --git a/recipes/MiniMaxAI/MiniMax-H3-4090.md b/recipes/MiniMaxAI/MiniMax-H3-4090.md index 1e22f5d2e..782c1483b 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-4090.md +++ b/recipes/MiniMaxAI/MiniMax-H3-4090.md @@ -110,9 +110,10 @@ outputs of each pair agreed to within 115 bytes. Four-GPU T2VA repeats: 273.8 s and 269.5 s, peaking at 15,560 MiB and 15,612 MiB. Four-GPU Ref2VA is a single successful end-to-end run (545.2 s, 16,448 MiB). Follow-up Ref2VA requests on this build completed diffusion but -then hit the engine's hardcoded 30 s async output wait -(`_ASYNC_OUTPUT_TIMEOUT` in `diffusion_engine.py`) during post-compute D2H, -so they are not reported as repeats. +then hit the engine's async output wait during post-compute D2H, so they are +not reported as repeats. That wait was a hardcoded 30 s when these numbers were +taken; it is now `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` (default 600 s), so a rerun on +current main should not lose these repeats to the bound. Relative to the two-GPU baseline on the same shape, four-GPU TP2×USP2 was about 1.6× faster for both tasks, while per-GPU peak HBM stayed in the mid-15 GiB diff --git a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md new file mode 100644 index 000000000..debc89d89 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md @@ -0,0 +1,82 @@ +# MiniMax H3 disaggregated text encoder + +This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends +its hidden states and token-role metadata to an encoder-free diffusion stage. +The standard MiniMax H3 recipes remain single-stage and continue to load the +text encoder inside the diffusion pipeline. + +## Start the server + +Choose the topology explicitly and load its deployment defaults: + +```bash +vllm-omni serve MiniMaxAI/MiniMax-H3 \ + --omni \ + --deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml +``` + +The default deployment assigns stage 0 to GPUs 0-1 with tensor parallel size +2 and `max_num_seqs: 1`. Stage 1 uses GPUs 2-5 with diffusion tensor parallel +size 1, Ulysses degree 4, and VAE patch parallel size 4. Adjust the +`devices`, `tensor_parallel_size`, and stage 1 `parallel_config` values in a +deployment override for the available hardware. Diffusion quantization, +layerwise offload, distributed layerwise offload, VAE parallelism, and USP +settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md). + +For example, this five-GPU topology assigns one GPU to the encoder and four +to the diffusion stage. `--stage-overrides` keeps placement and parallelism +scoped to the owning stage rather than broadcasting an override to both: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4 \ +vllm-omni serve MiniMaxAI/MiniMax-H3 \ + --omni \ + --deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \ + --stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}' +``` + +For memory-constrained deployments, start from the CPU-offload or distributed +layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and +quantization options only to Stage 1 with `--stage-overrides`; retain the +encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select +one offload strategy per deployment: + +```bash +# Stage 1 model-level CPU offload. +--stage-overrides '{"1":{"enable_cpu_offload":true}}' + +# Stage 1 distributed layerwise offload. Tune resident layers for available RAM. +--stage-overrides '{"1":{"enable_distributed_layerwise_offload":true,"dlo_use_allgather":false,"dlo_resident_layers":20}}' + +# Stage 1 online FP8 quantization of the DiT only. +--stage-overrides '{"1":{"diffusion_quantization_config":"{\"transformer\":{\"method\":\"fp8\"}}"}}' +``` + +The Stage 1 VAE patch-parallel options remain independent of offload and +quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and +hardware-qualified profiles before combining these options. + +Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download +text-encoder weights. This H3 topology explicitly keeps its single-replica +diffusion stage inline, avoiding serialization of decoded video through a +subprocess. It expects both stages to run in one deployment; cross-node payload +transport is outside this configuration. + +The `/v1/videos` request schema and `extra_params.task` values (`t2va`, +`fl2va`, and `ref2va`) are unchanged from the single-stage recipe. + +## Turbo LoRA + +MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and +`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so +requests that omit sampling controls do not inherit the 50-step base schedule: + +```bash +vllm-omni serve MiniMaxAI/MiniMax-H3 \ + --omni \ + --lora-path /path/to/MiniMax-H3-Turbo \ + --deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml +``` + +The base deployment intentionally retains 50 inference steps for non-LoRA +quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA. diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md b/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md new file mode 100644 index 000000000..36eb78a5d --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md @@ -0,0 +1,195 @@ +# MiniMax-H3 on a single NPU 950PR + +This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR +NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at +1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see +[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md). + +## Capacity requirements + +| Resource | Requirement | +| --- | ---: | +| NPU | 1x NPU 950PR | +| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) | +| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) | +| Checkpoint storage | 135 GiB per partition, local disk strongly preferred | +| Container shared memory | 8 GiB, or a writable local filesystem for spill | + +The observed high-water mark leaves about 12.3 GiB of headroom, so this card is +sized for exactly one resident task partition at 1024x576. Do not expect a +larger output shape or concurrency greater than one to fit. + +Because a single 950PR holds the whole partition in HBM, this route does not +need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB +GPU recipes does not apply. The validated run completed inside a container +limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is +actively harmful: it stages weights in non-reclaimable host memory and the +container's OOM killer terminates the server with exit code -9. + + +## Environment + +- Host architecture: x86_64 +- driver: 25.7.rc1.6 (hal 7.35.23) +- firmware: 9.0.0.105.229 +- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`) +- npu-smi: 25.7.rc1.6 +- Python: 3.12.13 +- PyTorch: 2.10.0+cpu +- torch_npu: 2.10.0.post2 +- vLLM: 0.26.0 +- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`) + +Install vLLM-Omni from a checkout with MiniMax-H3 support: + +```bash +uv venv +source .venv/bin/activate +uv pip install -e . +``` + +## Start a server + +```bash +export MODEL=/path/to/MiniMax-H3/FL2VA +export PORT=8000 +export VLLM_WORKER_MULTIPROC_METHOD=spawn +export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 + +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --num-gpus 1 \ + --tensor-parallel-size 1 \ + --usp 1 \ + --ring 1 \ + --text-encoder-tp-size 1 \ + --vae-patch-parallel-size 1 \ + --vae-parallel-mode tile \ + --vae-use-tiling \ + --init-timeout 14400 \ + --stage-init-timeout 14400 \ + --quantization int8 +``` + +Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no +NPU implementation and fails at the first denoise step. Leave the backend +unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is +installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)). + +H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. + +## Validated evidence + +Measured on one NPU 950PR with the server command above plus +`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from +the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3 +requested 60 denoise steps and executed 59 denoise updates, so per-step latency +is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps, +124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples. + +Five requests were issued back to back with no idle time between them. The card +is measurably faster on the first request than in steady state, so both regimes +are reported. + +| Stage | First request, cold card | Thermal steady state | +| --- | ---: | ---: | +| End-to-end request | 462.13 s | 506.38 s | +| Text encode | 0.55 s | 0.05 s | +| Denoise (59 updates) | 429.29 s | 471.74 s | +| Per denoise update | 7,276 ms | 7,996 ms | +| VAE decode | 8.43 s | 8.14 s | +| Worker-to-server handoff | 24 s | 27 s | +| MP4 muxing | 0.57 s | 0.50 s | +| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms | +| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB | +| Average NPU power | 566.9 W | 556.9 W | +| Peak NPU temperature | 104 C | 104 C | + +Sample counts differ per row. The cold column is a single request. In the steady +column, text encode, denoise and VAE decode are means over the four subsequent +requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while +end-to-end, handoff and muxing are means over the two of those four that +survived the handoff timeout described above. + +**Report the steady-state column.** The cold card completes denoise 9.9% faster +while drawing only 1.8% more average power, and both regimes reach the same +104 C ceiling, so the first request is buying a short window of higher clocks +before the die saturates. Any benchmark that issues a single request against an +idle card will overstate throughput by about 10%. + +**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides +the entire stage wall time by the 60 *requested* steps, so it absorbs text +encoding, VAE decode and the 837 MiB handoff. All three successful requests +satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The +per-update figure in the table is `diffuse / 59` from the pipeline profiler. + +Denoise accounts for 93% of the steady-state request and the handoff accounts +for 5.3%, so keeping the transfer in shared memory should bring end-to-end down +to roughly 480 s. + +Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a +reserved high-water mark rather than live allocation: it reads 111,272 MiB after +weight loading and before the first request, and does not fall back after +requests complete. Re-measure for longer outputs, a different output shape, or +concurrency greater than one. + +### Startup + +Weight loading took **31 min 50 s** on the validated host (first log line +08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a +shared GlusterFS network mount. This is why the server command sets +`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and +1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143. +Reduce both timeouts when the partition is staged on local disk. + +## T2VA request example + +```bash +export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync" + +curl -sS --max-time 1800 -X POST "${API_URL}" \ + -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \ + -F 'width=1024' \ + -F 'height=576' \ + -F 'aspect_ratio=16:9' \ + -F 'fps=24' \ + -F 'num_inference_steps=60' \ + -F 'flow_shift=12' \ + -F 'seed=1101' \ + -F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \ + -o t2va.mp4 +``` + +## Known limitations + +- INT8 online quantization is validated for T2VA on this card. Use BF16 for + FL2VA and Ref2VA, or re-measure before relying on it. +- Single-card serving loads one task partition at a time. For Ref2VA, stop the + server and restart it against the `Ref2VA` directory. +- Sustained load drives the die to 104 C and costs about 10% of denoise + throughput relative to a cold card. Warm the server with at least one full + request before measuring, and treat steady state as the reportable number. +- Requests can fail during result handoff on hosts where the IPC spill path is + slow. See [§ Result handoff and shared + memory](#result-handoff-and-shared-memory). +- Loading the checkpoint from a network mount takes over 30 minutes and forces + very large init timeouts. Stage the partition on local disk when possible. +- The configuration measured here is 1024x576 at 60 steps, which differs from + the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise + cost scales superlinearly with token count, so the numbers are not directly + comparable across recipes. +- The image ships a `triton` package whose backend is not built + (`No module named 'triton._C.libtriton.'`). vLLM logs this as an error + at startup and disables Triton; the diffusion path does not need it and the + run completes normally. + +## Additional resources + +- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide +- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide +- [Int8 quantization](../../docs/user_guide/quantization/int8.md) +- [Supported models](../../docs/models/supported_models.md) +- [Video API](../../docs/serving/videos_api.md) diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU.md b/recipes/MiniMaxAI/MiniMax-H3-NPU.md index ffe8eee09..b4563c1f7 100755 --- a/recipes/MiniMaxAI/MiniMax-H3-NPU.md +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU.md @@ -118,6 +118,21 @@ On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at above) or HSDP — see [§ Memory and attention optimizations](#memory-and-attention-optimizations-a3). +### CPU MP4 response encoding (Atlas A2) + +Non-streaming MP4 responses use one public automatic encoder. It checks the +runtime frame shape, common dtype, and RGB channel-plane contiguity for every +request; compatible inputs use direct planar PyAV frames, while unsupported +inputs fall back to the legacy muxer before opening the PyAV container. No CLI +flag, model declaration, or user configuration is required. Streaming fMP4 +output remains on its existing incremental path. + +Current performance and correctness validation is limited to one Atlas A2 host +with 8x Ascend 910B4-1 NPUs, the `FL2VA`/`t2va` partition, one request at a +time, 1344x768 at 24 fps, and 5, 8.7, and 15 second requests. This optimization +only changes CPU MP4 response encoding; it does not change DiT execution or +stage 0. + ### Optional optimizations Two independent optimizations may be enabled on top of the configuration diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index 114352343..71de893fd 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -65,6 +65,32 @@ On AMD ROCm, install without the `[fa4]` extra (FA4 is CUDA-only) and use `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for reference-video preparation and MP4 output. +## CPU MP4 response encoding + +For CUDA and ROCm deployments, non-streaming MP4 responses are encoded on the +host CPU through PyAV/libx264 after generation. The response encoder selects the +path automatically at runtime. The server-owned parallel converter accepts +supported frame shapes and dtypes with either per-channel-contiguous or strided +RGB planes, including interleaved arrays materialized by output transport. +Standalone callers without a parallel converter retain the legacy fallback for +strided planes. No CLI flag, model declaration, or user configuration is +required. Streaming fMP4 output is unchanged. + +A community benchmark on 2x Xeon 8480C reported the following comparison +between the legacy and direct planar paths +([full result](https://github.com/vllm-project/vllm-omni/pull/6288#issuecomment-5337546499)): + +| Metric | Legacy | Direct planar | Change | +| --- | ---: | ---: | ---: | +| Median wall time | 1.805 s | 1.394 s | -22.8% | +| Median process CPU time | 3.613 s | 3.207 s | -11.2% | +| Peak RSS | 3182 MiB | 2794 MiB | -387 MiB (-12.2%) | + +Across the 1.0-8.7 s sweep, wall-time improvement was approximately 21.8-22.6%; +outputs were byte-identical and full decode passed. This is evidence for host +CPU response encoding. Actual gains depend on the CPU and runtime, and should +not be interpreted as GPU, DiT, or stage 0 speedups. + ## Start a server Pass the repository ID directly. The pipeline uses `FL2VA` for model discovery @@ -158,7 +184,7 @@ vllm serve "${MODEL}" \ Use the profile that matches the per-GPU memory capacity: | Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status | -|---|---:|---:|---:|---|---|---| +| --- | ---: | ---: | ---: | --- | --- | --- | | `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated | | `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point | @@ -177,8 +203,8 @@ increasing it on a different request shape. At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on 2 x RTX 5090 without OOM: -| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation | -|---:|---:|---:|---:|---| +| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation | +| -------: | ------------: | ---------: | ---------------------: | -----------------------------------------------------------: | | 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed | This is a single end-to-end validation run, not a warmed multi-run latency @@ -247,7 +273,7 @@ vllm serve "${MODEL}" \ Do not add `--enforce-eager` to this performance configuration. The first request includes regional compilation; warm the server once before measuring steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain -1. The H3 VAE supports its native `tile` mode, not +`1`. The H3 VAE supports its native `tile` mode, not `spatial_shard_height` or `spatial_shard_width`. ### Attention Backends @@ -355,14 +381,78 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while `task=ref2va` routes to `Ref2VA/transformer`. T2VA uses the FL2VA DiT. +### Step execution and continuous batching + +H3 implements the step-wise execution contract, so the scheduler can admit and +retire requests between denoise steps instead of running one request end to +end. Add the feature gate, then raise `--max-num-seqs` to co-batch: + +```bash +--step-execution --max-num-seqs 4 +``` + +Co-batched requests are packed into a single sequence that keeps one attention +document per request, so a batch costs one DiT forward. That packing requires +`--diffusion-attention-backend FLASH_ATTN`; other backends fall back to one +forward per request. `--max-num-seqs 1` keeps the conservative single-request +step path. Cache acceleration (`--cache-backend`) is not available in step mode. + +!!! warning "Co-batching does not improve H3 throughput for large simultaneous requests" + Keep `--max-num-seqs 1` unless you specifically need scheduler-level control + (admitting and retiring requests between denoise steps). Measured on two + H100s (TP2, BF16, 672x384, 209 frames, 30 steps, 4 requests at concurrency + 4 with `--request-rate inf`, i.e. all four submitted at time 0; one packed + request is 16384 rows): + + | Configuration | Wall time | Mean latency | Peak memory | + |---------------|-----------|--------------|-------------| + | request mode | 174.8 s | 111.5 s | 72.4 GB | + | `--step-execution --max-num-seqs 1` | 179.0 s | 113.8 s | 72.4 GB | + | `--step-execution --max-num-seqs 4` | 182.1 s | 175.7 s | 78.3 GB | + + A single H3 denoise step is a compute-bound dense GEMM over an already long + packed sequence, so fusing N requests costs N times the FLOPs and buys almost + no amortization — unlike LLM decoding, which is memory-bandwidth bound. + Going from one request per step to four cuts the per-request denoise cost + only from 1.323 s to 1.291 s (2.4%), which the step-mode bookkeeping then + spends. Mean latency degrades further because co-batched requests finish + together instead of staggered. Quantization moves the absolute numbers + without changing this: with online `int8` the same workload runs in 153.3 s + at 56.9 GB (request mode), and `--max-num-seqs 4` is still 5.0% slower than + request mode. + + Two workloads outside this table are unmeasured and may behave differently: + + - **Staggered arrivals (admission latency).** A single 16384-row H3 request + already saturates the GPUs, so submitting all requests at time 0 is the + one arrival pattern where co-batching cannot win: it can only bunch + completions. In request mode a new request queues behind the whole + in-flight generation (~45 s at this size); step mode admits at the next + denoise-step boundary (~1.3 s). Whether that shows up as a wall-clock + benefit under Poisson arrivals is not yet measured for H3. To reproduce, + run request mode vs `--step-execution --max-num-seqs 4` with + `diffusion_benchmark_serving.py --request-rate 0.05` (roughly one request + every 20 s) and report mean / p95 latency, which then includes queueing. + - **Small requests.** The numbers above are drawn from 672x384 / 209-frame + requests. A short clip at lower resolution packs a few thousand rows and + may not saturate the hardware; co-batching may amortize better there. + This is also unmeasured. + ### Online FP8 quantization -MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint -remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and -uses dynamic activation scaling during inference. By default, attention and -MLP linears in the token refiner and main DiT blocks, the condition -projection, and all AdaLN projections use FP8. Patch, timestep, and final -projections remain FP32; the text encoder and VAEs are unchanged. +MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL +text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8 +weights at runtime and uses dynamic activation scaling during inference. By +default, `--quantization fp8` quantizes eligible attention and MLP linears in +the text decoder, token refiner, and main DiT blocks, as well as the DiT +condition and AdaLN projections. The Qwen vision tower, embeddings, norms, +RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections +keep checkpoint precision. + +To select a component, use `--diffusion-quantization-config` with +`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or +`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two +entries can be combined. The shorthand below enables both components. Add this option to an existing H3 server command: @@ -370,6 +460,28 @@ Add this option to an existing H3 server command: --quantization fp8 ``` +#### Single 96 GB GPU, no-offload capacity check + +Use the FL2VA-only partition for this capacity test. Loading the combined +service would also load the Ref2VA DiT and would test a different memory +budget. A no-offload capacity check should contain none of +`--enable-cpu-offload`, `--enable-layerwise-offload`, or +`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement +but does not offload model weights to the CPU. + +The run passes the capacity check when the server initializes, the request +finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the +card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus +32 kHz stereo AAC audio. Report the measured headroom rather than assuming +that every nominal 96 GB SKU exposes the same MiB total. + +As a capacity proxy only, the same five-second case on one B300 measured a +92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its +encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s, +6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but +they are not an RTX PRO 6000 validation: kernels, allocator behavior, and +reported device capacity differ across GPUs. + Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3 resolves the `transformer` component before constructing the DiT, so names do not start with `transformer.`. Entries are exact runtime linear prefixes; a @@ -386,10 +498,11 @@ For example, keep the first main block's attention projections in BF16 with: '{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}' ``` -The structured option replaces `--quantization fp8`. Online FP8 is currently -incompatible with H3 layerwise offload because the offload path produces a -weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor -parallelism and VAE tiling instead. +The structured option replaces `--quantization fp8`. Online FP8 can be used +with H3 layerwise offload and with both DLO transfer paths. The default +AllGather path uses the ordinary loader to finalize FP8 weights and scales +before sharding them across ranks. `--dlo-no-use-allgather` instead retains +complete rank-local tensors and avoids the synchronized request-wave contract. ## AMD ROCm (gfx942 / gfx950) @@ -475,7 +588,7 @@ vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the `0.26.0+rocm723` wheel (HIP 7.2). | Workload | Configuration | Observed result | -|----------|---------------|-----------------| +| ---------- | --------------- | ----------------- | | T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC | | FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC | | T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup | @@ -657,7 +770,7 @@ at most 15 seconds combined. ## Official input matrix and limits | Task | Supported references | Limits | -|------|----------------------|--------| +| ------ | ---------------------- | -------- | | T2VA | text only | prompt must be non-empty | | FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` | | Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference | @@ -688,13 +801,17 @@ Omitting `quality` preserves the startup default: it uses the reference path normally, or the server-configured profile when the server was started with `--cache-backend cache_dit`. +Turbo is independent of this quality switch: `quality` selects a Cache-DiT +policy, while Turbo changes the active LoRA weights and sampling schedule. See +[LoRA](#lora) below. + The following result was measured on 4× NVIDIA H200 with SP4, text-encoder TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full `lossless` warmup was excluded, followed by three fixed prompt/seed pairs in balanced switch order. | `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off | -|---|---:|---:|---:|---:|---| +| --- | ---: | ---: | ---: | ---: | --- | | `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path | | `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation | @@ -704,10 +821,173 @@ balanced switch order. > workload. The values above apply to this deployment and are not universal > guarantees. `lossless` remains the exact reference path. +## LoRA + +### Turbo LoRA + +Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported: + +```text +lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors +``` + +Download only that file: + +```bash +export TURBO_DIR=/path/to/minimax-h3-turbo +export TURBO_FILE=minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors +hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}" +export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}" +``` + +Start from a non-offloaded or DLO FL2VA server command and add +`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`. +`--lora-path` preloads the adapter; each request still activates it and uses +the published sampling settings: + +```bash +-F 'num_inference_steps=5' \ +-F 'flow_shift=6' \ +-F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \ +-F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}" +``` + +For FL2VA, change `task` and add `input_reference` as shown above. The 8-step, +ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is +dynamic-only and does not support prefusion or LoRA composition. DLO is +supported by keeping the request-switchable LoRA A/B buffers resident on the +accelerator while DLO streams only the base blocks; budget for this additional +fixed HBM usage. Model-level and standard layerwise offload remain unsupported. +The five requested sigma points produce the four denoiser evaluations expected +by the Turbo artifact. + +### FlashGen native LoRA + +The FlashGen 4-step T2VA artifact uses the native MiniMax-H3 module layout and +declares its distilled sigma schedule in safetensors metadata. It is published on +[ModelScope](https://modelscope.cn/models/FlashGen/Minimax-H3-4step-lora-flashgen): + +```text +FlashGen/Minimax-H3-4step-lora-flashgen/minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors +``` + +Download only that file: + +```bash +python -m pip install modelscope +export FLASHGEN_DIR=/path/to/minimax-h3-flashgen-lora +export FLASHGEN_FILE=minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors +modelscope download FlashGen/Minimax-H3-4step-lora-flashgen \ + --local_dir "${FLASHGEN_DIR}" \ + --include "${FLASHGEN_FILE}" +export FLASHGEN_LORA="${FLASHGEN_DIR}/${FLASHGEN_FILE}" +``` + +Start from a non-offloaded or DLO FL2VA server command and add +`--task-type fl2va --lora-backend peft --lora-path "${FLASHGEN_LORA}"`. +Each request must use T2VA and the distilled interval-count contract: + +```bash +-F 'num_inference_steps=4' \ +-F 'extra_params={"task":"t2va","duration":5.2}' \ +-F "lora={\"name\":\"h3-flashgen-v1.0\",\"path\":\"${FLASHGEN_LORA}\",\"scale\":1.0}" +``` + +This path rejects Ref2VA and checkpoints that already pin `base_schedule` in +`model_index.json`. The adapter metadata carries +`base_schedule=1.0,0.7,0.4,0.15,0.0`, so `num_inference_steps=4` means four +denoiser evaluations, not five sigma points. Request-mode generation may omit +the field and take the count from the adapter schedule; `--step-execution` +requires it explicitly, because the step scheduler reads the total step count +off the request at admission, before the adapter schedule is known. + +DLO is supported in request-mode generation on the same terms as the Turbo +adapter: the request-switchable LoRA A/B buffers stay resident on the +accelerator while DLO streams only the base blocks, so budget for that +additional fixed HBM usage. The native artifact is rank 64 over 259 target +modules, and its packed `qkv_proj` and `fc1` layers reuse the full-input A +tensor per slice while B carries slice-local output rows, so the resident +footprint exceeds the on-disk payload; measure it for your parallel layout +rather than assuming the checkpoint size. Pure Ulysses replicates the adapter +on every rank, while DiT tensor parallelism shards the B buffers. Model-level +and standard layerwise offload remain unsupported, and `--step-execution` +cannot be combined with `--enable-distributed-layerwise-offload`. + +To validate a deployment, post the same fixed-seed T2VA request twice with the +adapter and twice without it, then compare the four output digests. The adapter +is bound and deterministic when each pair matches internally and the two pairs +differ from each other. + +### FastH3 adapter + +[FastH3](https://haoailab.com/blogs/fasth3-preview/) is FastVideo's four-step +DMD2 student of H3-Base. It reuses H3's text encoder, VAEs, tokenizers, and +schedulers unchanged, replacing 49 denoiser evaluations with four. It is fused +into the checkpoint at load time rather than switched per request, because it +carries full-rank deltas that no LoRA layer can express. + +The bundle publishes four variants, so download one and point `--lora-path` at +it; the repository root is refused rather than guessed at: + +```bash +export FASTH3_DIR=/path/to/fasth3 +hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA \ + dense-datafree/adapter_model.safetensors --local-dir "${FASTH3_DIR}" +export FASTH3_LORA="${FASTH3_DIR}/dense-datafree/adapter_model.safetensors" +``` + +Add `--task-type fl2va --lora-path "${FASTH3_LORA}"` to a non-offloaded server +command. T2VA is served by the FL2VA partition, so `--task-type fl2va` is +correct even though FastH3 preview v1 distills T2VA only. Because the adapter is +fused, `--lora-backend` does not apply and a request carrying a `lora=` field is +rejected rather than served without the adapter it asked for. + +```bash +-F 'num_inference_steps=4' \ +-F 'extra_params={"task":"t2va","duration":4.4}' +``` + +Requests must ask for `num_inference_steps=4` and `task=t2va`: the release's five +sigma points bound four denoiser evaluations, and that count is what the step +scheduler admits a request on. The server denoises on the release's own ladder, +keeping H3's per-modality shifts at the checkpoint values, so a request that +overrides `flow_shift` or `audio_flow_shift` is rejected - it would sample the +student at noise levels it was never distilled at. + +Only a release that identifies itself as FastH3 is fused; any other +`fastvideo-lora-v2` adapter stays on the dynamic LoRA route. A claimed artifact +is then held to its own metadata: one that misdeclares its tensor counts or +leaves a transformer block unedited is refused at startup instead of serving +mostly base H3 weights on a four-step schedule. Offload is refused for the same +reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and +`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast. + +> [!NOTE] +> Only the **Dense / Data-Free** variant is supported today. The three VSA +> variants carry `.set_weight` compression gates for modules that do not exist +> in vLLM-Omni's dense H3 attention, and the release states dense attention is +> not a supported substitute for them, so they are refused at startup rather +> than silently run wrong. Video Sparse Attention for H3's packed +> `[text | cond | audio | video]` sequence is follow-up work. + +Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at +1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded: + +| Adapter | Steps | End-to-end | Diffusion engine | +| --- | ---: | ---: | ---: | +| none (base H3) | 50 | 25.8 / 26.4 s | 16.22 / 16.28 s | +| FastH3 Dense | 4 (5 sigma points) | 11.7 / 11.8 s | 2.37 / 2.36 s | + +The denoising speedup is 6.9x. End-to-end is 2.2x because text encoding, VAE +decoding and muxing are a fixed cost that dominates a clip this short; longer +generations move the end-to-end figure toward the denoising one. Fusing the +adapter does not measurably change startup: weight loading took 77.3 s with it +against 85.8 s without. + ## Key parameters | Parameter | Recommended value | Notes | -|-----------|-------------------|-------| +| ----------- | ------------------- | ------- | | `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile | | `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` | | `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT | @@ -732,9 +1012,9 @@ Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 servi The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one excluded warmup followed by three requests. -| Workload | Configuration | Observed result | -|----------|---------------|-----------------| -| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency | +| Workload | Configuration | Observed result | +| -------------------------------------- | ------------------------------------------- | ------------------------------------ | +| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency | | Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean | These measurements describe the validated shapes rather than a general @@ -750,6 +1030,13 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22% reduction. +For direct human inspection, see the external +[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/), +which includes matched five-second T2VA, I2VA, and Ref2VA videos. The +[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16) +also record per-task fidelity metrics and provenance without storing generated +media in this repository. + ## TeaCache acceleration TeaCache reuses DiT block residuals across denoising steps when consecutive @@ -819,7 +1106,12 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ - TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached. - Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads both task-specific DiTs, and loads shared components once from `FL2VA`. -- H3 currently executes one generation request per diffusion batch. +- Request mode executes one generation request per diffusion batch. Use + `--step-execution` with `--max-num-seqs N` to admit several requests at once + (see + [Execution modes](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/execution_modes.md)); + step + mode does not support `cache_backend`. - The first regional-compile request is a warmup and should not be included in steady-state performance measurements. - The serving path accepts fewer references than the model supports. H3 documents up @@ -834,7 +1126,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \ H3 native `tile` mode only. - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use pure Ulysses. -- Online FP8 is not compatible with layerwise offload. +- Online FP8 with DLO AllGather temporarily materializes the complete FP8 model + in host memory on every rank during startup before retaining only each rank's + shard. Size startup host memory for that transient peak. - TeaCache and Cache-DiT cannot be enabled on the same server. - Image+audio Ref2VA accepts exactly one image and one audio reference. - Video Ref2VA accepts one or more video files, but not an additional standalone diff --git a/recipes/MiniMaxAI/MiniMax-Music3.md b/recipes/MiniMaxAI/MiniMax-Music3.md new file mode 100644 index 000000000..3d078f645 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-Music3.md @@ -0,0 +1,173 @@ +# MiniMax Music 3 + +> Text-to-music: lyrics plus a style caption in, a 32 kHz stereo song out + +## Summary + +- Vendor: MiniMax +- Model: `MiniMaxAI/MiniMax-Music3` +- Task: Text-to-music generation (vocals and instrumental together) +- Mode: Online serving through the OpenAI-compatible `/v1/audio/speech` API +- Maintainer: @linyueqian + +## When to use this recipe + +Use it to serve `MiniMaxAI/MiniMax-Music3` for full-song generation up to six +minutes. This is not a TTS model with a musical voice: there is no reference +speaker, no `voice`, and no `temperature`. What you control is the lyrics, the +caption, the seed, and the length. + +## References + +- Model card: [MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3) +- Demo samples: [MiniMax Music 3 demo](https://minimax-ai.github.io/music3-demo/) +- Deploy config: `vllm_omni/deploy/minimax_music3.yaml` (single GPU), + `vllm_omni/deploy/minimax_music3_2gpu.yaml` (split across two) + +## Architecture + +| Component | Spec | +|---|---| +| Backbone | Qwen3 decoder, 36 layers, hidden 4096, GQA 32/8, vocab 200k | +| Depth decoder | 4 layers, hidden 4096, 16 heads; 7 residual codebooks of 1024 | +| Frame | 8 codebooks; `c0` in the backbone vocabulary, `c1..c7` in the depth decoder | +| AR guidance | Classifier-free, scale 1.5, `c0` masked to the conditioned branch's top 50 | +| Acoustic stage | Flow-matching transformer (36 layers, dim 2048) plus a DAC-style decoder, 512x upsampling | +| Solver | Euler, 30 steps, acoustic guidance scale 1.7 | +| Window | 200 frames per acoustic chunk, 100-frame hop | +| Frame rate | 25 frames per second | +| Context length | 10,240 tokens | +| Output | 32 kHz stereo WAV | + +Generation is two-stage. The backbone predicts one audio frame per decode step, +each frame eight RVQ codebooks deep. Every 200 frames are handed to the +flow-matching transformer, which solves a latent that the decoder turns into +waveform. + +**Guidance doubles the rows.** Every request decodes twice: one row sees the +real prompt, the other sees the same prompt with the caption-and-lyrics span +replaced by `<|audio_cfg|>`. Both rows hold their own KV cache for the whole +song, so a request costs twice the KV its length suggests. Size the pool for +rows, not requests. + +## Hardware Support + +## GPU + +### 1x H200 141GB + +#### Environment + +- OS: Linux +- Python: 3.12 +- CUDA: 12.9 +- vLLM version: 0.27.0 +- vLLM-Omni version or commit: branch `feat/minimax-music3` + +#### Command + +```bash +vllm serve MiniMaxAI/MiniMax-Music3 \ + --host 0.0.0.0 --port 8000 \ + --trust-remote-code --omni +``` + +The default deploy config `vllm_omni/deploy/minimax_music3.yaml` is discovered +automatically from the checkpoint's `model_type=minimax_music3`. It colocates +both stages on device 0. + +#### Verification + +```bash +curl -X POST http://localhost:8000/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "MiniMaxAI/MiniMax-Music3", + "input": "[Verse]\nWalking down the empty street at midnight\nStreetlights flicker like a broken dream\n[Chorus]\nAnd I keep on walking\nTill the morning finds me", + "instructions": "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. Intimate bedroom production, gentle tape saturation, no bright cymbals.", + "seed": 1, + "max_new_tokens": 750 + }' \ + --output song.wav +``` + +Expected: a 30 second 32 kHz stereo WAV. `max_new_tokens` counts audio frames +at 25 per second, so 750 frames is at most 30 seconds. It is a cap, not a +target: the model ends the song itself when it emits the audio-end token, and a +response shorter than the cap is the model finishing rather than a truncation. + +#### Notes + +- **Memory:** stage 0 takes `gpu_memory_utilization: 0.6` for the bf16 backbone + and its KV pool; stage 1 takes 0.2 for the transformer and decoder weights, + which sit outside any KV budget. +- **Key flags:** none required. Guidance, the solver step count and the window + geometry are fixed by the checkpoint and are not request parameters. +- **Known limitations:** + - The external API is non-streaming; `stream: true` is rejected. + - `max_model_len` is 10,240, so a maximum-length prompt (5,000 tokens) and a + maximum-length song (9,000 frames) cannot both fit. Long captions shorten + the maximum song. This is a property of the checkpoint. + - The acoustic stage runs in float32. bfloat16 is accepted but measurably + degrades the solver. + +### 2x H200 141GB + +Use `minimax_music3_2gpu.yaml` to put the acoustic stage on the second device: + +```bash +vllm serve MiniMaxAI/MiniMax-Music3 \ + --host 0.0.0.0 --port 8000 \ + --trust-remote-code --omni \ + --deploy-config minimax_music3_2gpu.yaml +``` + +Only the placement differs; both layouts run the acoustic stage in float32. + +## Writing the request + +Two fields carry the request, and both are required. + +- `input` is the **lyrics**. Structure tags (`[Verse]`, `[Chorus]`, `[Bridge]`, + `[Outro]`) steer the arrangement and are part of the prompt contract. + **Put a tag on its own line.** Normalization keeps only the tags on a line + that starts with one and drops the rest of that line, so lyrics written next + to a tag are silently lost: + + ```text + "[Verse]\nWalking down the street" -> [start] [verse] Walking down the street + "[Verse] Walking down the street" -> [start] [verse] + ``` + +- `instructions` is the **caption** describing genre, instrumentation, tempo, + mood and production. It is the strongest control available: vague captions + give generic arrangements. Long structured captions work well, since the + model was trained on descriptions covering global attributes, emotional + progression, vocal detail and arrangement. + +A request is deterministic in its seed: the same lyrics, caption, seed and +length return the same audio. Changing only the seed gives another take of the +same song. + +## Request parameters + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `input` | string | required | Lyrics, non-empty. Tags on their own lines | +| `instructions` | string | required | Caption, non-empty | +| `seed` | int | `0` | Non-negative. Fixes the output for a given request | +| `max_new_tokens` | int | `9000` | Cap on audio frames, 25 per second. Maximum 9,000 | +| `response_format` | string | `"wav"` | Output container | + +## Parameters the model rejects + +The request is refused rather than silently ignored, so a mistake is visible +immediately. + +| Parameter | Reason | +|---|---| +| `temperature`, `top_p`, `top_k`, `repetition_penalty` | Sampling is fixed: guidance at 1.5, then a seeded top-k 50 draw | +| `voice` | There is no speaker to select; the vocal comes from the caption | +| `ref_audio`, `ref_text`, `language`, `task_type` | No reference-audio conditioning and no language tag | +| `speed` | Only `1.0`. Tempo belongs in the caption, e.g. "at 92 BPM" | +| `stream: true` | The external API is non-streaming | diff --git a/recipes/NVIDIA/GR00T-N1.7.md b/recipes/NVIDIA/GR00T-N1.7.md index 0a5270222..826309835 100644 --- a/recipes/NVIDIA/GR00T-N1.7.md +++ b/recipes/NVIDIA/GR00T-N1.7.md @@ -44,7 +44,7 @@ vllm serve nvidia/GR00T-N1.7-3B \ --host 127.0.0.1 \ --port 8000 \ --served-model-name gr00t-n1d7 \ - --stage-configs-path vllm_omni/deploy/Gr00tN1d7.yaml + --deploy-config vllm_omni/deploy/Gr00tN1d7.yaml ``` Notes: diff --git a/recipes/NVIDIA/Nemotron-Labs-Audex.md b/recipes/NVIDIA/Nemotron-Labs-Audex.md index 21ee7e655..13be59f4c 100644 --- a/recipes/NVIDIA/Nemotron-Labs-Audex.md +++ b/recipes/NVIDIA/Nemotron-Labs-Audex.md @@ -51,10 +51,10 @@ NemotronH (~3B active parameters). Pick the pipeline by task: `model_type` auto-resolves `audex_tts.yaml`): ```bash -vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \ +vllm serve nvidia/Nemotron-Labs-Audex-2B \ --host 0.0.0.0 --port 8097 \ --trust-remote-code --omni -# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml +# other modes: add --deploy-config vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml # or use the launcher: MODE=s2s PORT=8098 examples/online_serving/audex/run_server.sh ``` @@ -109,10 +109,10 @@ The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type` auto-resolves to the 2B-tuned config otherwise): ```bash -vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \ +vllm serve nvidia/Nemotron-Labs-Audex-30B-A3B \ --host 0.0.0.0 --port 8097 \ --trust-remote-code --omni \ - --stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml + --deploy-config vllm_omni/deploy/audex_tts_30b.yaml # or: SIZE=30b MODE=tts examples/online_serving/audex/run_server.sh ``` diff --git a/recipes/NVIDIA/NemotronLabs-VoiceChat.md b/recipes/NVIDIA/NemotronLabs-VoiceChat.md index 8faa268d2..dbe7d5431 100644 --- a/recipes/NVIDIA/NemotronLabs-VoiceChat.md +++ b/recipes/NVIDIA/NemotronLabs-VoiceChat.md @@ -1,7 +1,7 @@ # NemotronLabs VoiceChat 11B -> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply -> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline. +> Speech-to-speech on a 3-stage vLLM-Omni pipeline: offline single-turn +> inference and experimental model-native, frame-locked Realtime serving. ## Summary @@ -11,8 +11,8 @@ timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into 31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio. -- Mode: Offline single-turn inference (batch=1). Online/duplex serving, - function calling, and batch>1 are documented follow-ups. +- Mode: Offline single-turn inference (batch=1), plus experimental native + duplex Realtime serving at one 80 ms microphone frame per scheduler wake. - Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang) ## When to use this recipe @@ -33,6 +33,11 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed. [`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/) - Staged pipeline config: [`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml) +- Native duplex config and E2E probe: + [`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml), + [`tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py`](../../tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py) +- Nightly model-level gate: + [`tests/e2e/online_serving/test_nemotron_voicechat_duplex.py`](../../tests/e2e/online_serving/test_nemotron_voicechat_duplex.py) - Upstream model card: [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B) - Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat` @@ -60,7 +65,7 @@ output stayed within one word of the reference in testing. - OS: Linux - Python: 3.12 -- vLLM version: 0.26.0 +- vLLM version: 0.27.0 - vLLM-Omni version or commit: this PR / current `main` #### Command @@ -88,6 +93,55 @@ The reply text should read as a coherent spoken-style answer to the question in the input WAV, and the WAV should transcribe to (approximately) the same text with any ASR model. +#### Native duplex serving + +The checkpoint does not contain the underlying Nemotron text tokenizer. Set +`NEMOTRON_VOICECHAT_LLM_PATH` to a local snapshot of +`nvidia/NVIDIA-Nemotron-Nano-9B-v2` before starting an offline deployment. + +```bash +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +export NEMOTRON_VOICECHAT_LLM_PATH=/path/to/NVIDIA-Nemotron-Nano-9B-v2 + +vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \ + --omni \ + --served-model-name nemotron-voicechat \ + --deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml +``` + +The single-GPU profile runs three eager fp32 engine processes on one device. + +In another shell, stream only the user channel of one of NVIDIA's bundled +stereo conversation fixtures: + +```bash +python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \ + --model nemotron-voicechat \ + --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/turn_taking.wav \ + --input-channel 0 \ + --output-dir results/nemotron_voicechat_duplex \ + --timeout-s 600 +``` + +The probe requires a completed response, non-silent 22.05 kHz output, and the +advertised native 80 ms append capabilities. Add `--expect-function-call` and +use `tool_call.wav` to validate the function-call channel. To validate the +Realtime tool round trip, also provide the expected arguments and return a +tool result: + +```bash +python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \ + --model nemotron-voicechat \ + --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/tool_call.wav \ + --output-dir results/nemotron_voicechat_tool_call \ + --expect-function-call \ + --expected-function-name generate_random_number \ + --expected-function-arguments '{"min":1,"max":50}' \ + --function-output 20 \ + --expected-post-tool-text "random number" +``` + #### Notes - Memory usage: the shipped yaml runs all three stages on one GPU @@ -109,5 +163,14 @@ with any ASR model. instead of pausing. - The talker's `max_tokens` is 16383 (its stage prompt is one placeholder token, and the stage context is 16384). -- Known limitations: batch=1 only; offline single turn (no online/duplex - serving yet); the function-call channel is decoded but not yet acted on. +- The 80 ms frame period is a model/protocol contract, not a throughput + guarantee. In an eager fp32 H200 profile, Stage 2 emitted 80 ms packets at + roughly 0.87--0.88 s intervals (RTF about 10.8--10.9); placing the three + stages on separate GPUs did not remove the serial 8-step EAR-TTS bottleneck. + The current native duplex path is functionally streaming but not wall-clock + realtime. +- Known limitations: batch=1 only. The native duplex deployment allows one + active session and does not support barge-in. Tool execution remains + client-owned: the server emits function-call events, accepts a validated + `function_call_output`, and resumes the live model with the returned result; + it does not execute arbitrary tools itself. diff --git a/recipes/NVIDIA/SANA-Video-2B.md b/recipes/NVIDIA/SANA-Video-2B.md new file mode 100644 index 000000000..c85d8ca40 --- /dev/null +++ b/recipes/NVIDIA/SANA-Video-2B.md @@ -0,0 +1,245 @@ +# SANA-Video 2B + +> Native text-to-video and image-to-video generation at 480p and 720p, +> plus a validated Diffusers-adapter compatibility path + +## Summary + +- Vendor: NVIDIA +- Model: `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, + `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` +- Task: Text-to-video and image-to-video +- Mode: Offline inference and OpenAI-compatible online serving +- Maintainer: Community + +## When to use this recipe + +Use the native `SanaVideoPipeline` for T2V and +`SanaImageToVideoPipeline` for I2V. Both native pipelines support the 480p +and 720p checkpoints. Use `--diffusion-load-format diffusers` when you need +the black-box Diffusers compatibility baseline; adapter T2V and I2V are +validated at both resolutions. + +The native pipeline loads the 480p checkpoint through +`DistributedAutoencoderKLWan` and the 720p checkpoint through +`DistributedAutoencoderKLLTX2Video`. These are vLLM-Omni distributed wrappers +around the corresponding Diffusers autoencoders, not independent VAE +implementations. The denoising loop also intentionally loads Diffusers' +`DPMSolverMultistepScheduler` from the checkpoint to preserve its scheduler +configuration and numerical behavior. + +## References + +- Upstream project: +- Model cards: + [480p](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_480p_diffusers), + [720p](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_720p_diffusers) +- Diffusers documentation: +- Online serving guides: + [Text-to-Video](../../docs/user_guide/examples/online_serving/text_to_video.md), + [Image-to-Video](../../docs/user_guide/examples/online_serving/image_to_video.md) +- Offline T2V example: + [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py) +- Offline I2V example: + [`examples/offline_inference/image_to_video/image_to_video.py`](../../examples/offline_inference/image_to_video/image_to_video.py) +- Support discussion: [vLLM-Omni issue #5432](https://github.com/vllm-project/vllm-omni/issues/5432) + +## Hardware Support + +## GPU + +### 1x RTX 5090 32GB + +#### Environment + +- OS: Ubuntu 22.04.5 LTS +- Python: 3.12.3 +- Driver / runtime: NVIDIA driver 580.95.05; PyTorch 2.11.0+cu130 +- Diffusers: 0.38.0 +- vLLM version: 0.26.0 +- vLLM-Omni version or commit: PR #5508, commit `22037901` + +#### Command + +##### Native text-to-video inference + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_720p_diffusers \ + --model-class-name SanaVideoPipeline \ + --prompt "A cat walking on the grass, facing the camera." \ + --negative-prompt "blurry, low quality, temporal artifacts" \ + --height 704 --width 1280 --num-frames 81 \ + --num-inference-steps 50 --guidance-scale 6 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 --seed 42 --output sana_video_720p.mp4 +``` + +For 480p, select `SANA-Video_2B_480p_diffusers` and use +`--height 480 --width 832`. + +##### Native image-to-video inference + +SANA checkpoints declare `SanaVideoPipeline` in `model_index.json`, so I2V +must be selected explicitly with `--model-class-name +SanaImageToVideoPipeline`. + +```bash +python examples/offline_inference/image_to_video/image_to_video.py \ + --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + --model-class-name SanaImageToVideoPipeline \ + --image input.jpg \ + --prompt "A cat turns toward the camera with smooth, natural motion." \ + --negative-prompt "blurry, low quality, temporal artifacts" \ + --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 50 --guidance-scale 6 \ + --extra-body '{"motion_score": 30}' \ + --fps 16 --seed 42 --output sana_video_i2v_480p.mp4 +``` + +The same pipeline class supports the 720p checkpoint through vLLM-Omni's +distributed LTX-2 VAE wrapper; use `--height 704 --width 1280`. + +For online I2V serving: + +```bash +MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + bash examples/online_serving/image_to_video/run_server_sana_video.sh + +INPUT_IMAGE=input.jpg OUTPUT_PATH=sana_video_i2v.mp4 \ + bash examples/online_serving/image_to_video/run_curl_sana_video.sh +``` + +##### Native online serving + +```bash +MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + bash examples/online_serving/text_to_video/run_server_sana_video.sh + +bash examples/online_serving/text_to_video/run_curl_sana_video.sh +``` + +To run the black-box compatibility backend for T2V, replace the server script +with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request +works; `num_frames` is adapted to Diffusers' `frames` argument. The script +selects `TORCH_SDPA` because SANA-Video uses an attention mask that the +AITER-backed Diffusers attention path does not accept. + +##### Diffusers-adapter image-to-video serving + +The validated I2V adapter commands are: + +```bash +# 480p +MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \ + bash examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh + +INPUT_IMAGE=input.jpg OUTPUT_PATH=sana_video_i2v_adapter.mp4 \ + bash examples/online_serving/image_to_video/run_curl_sana_video.sh + +# 720p +MODEL=Efficient-Large-Model/SANA-Video_2B_720p_diffusers \ + bash examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh + +INPUT_IMAGE=input.jpg WIDTH=1280 HEIGHT=704 \ + OUTPUT_PATH=sana_video_i2v_adapter_720p.mp4 \ + bash examples/online_serving/image_to_video/run_curl_sana_video.sh +``` + +#### Verification + +Check the encoded 720p output metadata after running a generation command: + +```bash +ffprobe -v error -select_streams v:0 \ + -show_entries stream=width,height,r_frame_rate,nb_frames \ + -of default=noprint_wrappers=1 sana_video_720p.mp4 +``` + +The standard 720p request above should report: + +```text +width=1280 +height=704 +r_frame_rate=16/1 +nb_frames=81 +``` + +For 480p, expect `width=832` and `height=480`. + +The automated serving matrix covers both checkpoint variants: + +| Backend | 480p T2V | 720p T2V | 480p I2V | 720p I2V | +| --- | --- | --- | --- | --- | +| Native vLLM-Omni | Validated | Validated | Validated | Validated | +| Diffusers adapter | Validated | Validated | Validated | Validated | + +Use the native `SanaVideoPipeline` and `SanaImageToVideoPipeline` for the +primary SANA execution paths. The Diffusers adapter is retained as a +validated compatibility/reference backend. + +#### Notes + +- Key flags: select I2V explicitly with `--model-class-name + SanaImageToVideoPipeline`; SANA checkpoint `model_index.json` files declare + the T2V class. Pass `motion_score` through `--extra-body`. The supplied + Diffusers-adapter server scripts select `TORCH_SDPA` because the AITER-backed + path does not accept SANA's attention mask. +- Output profile: 81 frames at 16 FPS is the standard checkpoint profile and + produces approximately five seconds of video. Minute-scale generation + requires the separate LongSANA/LongLive block-autoregressive workflow. +- Backend boundary: the native pipelines and Transformer are owned by + vLLM-Omni. The 480p Wan VAE and 720p LTX-2 VAE run through vLLM-Omni + distributed wrappers derived from the corresponding Diffusers VAE classes. + The denoising loop intentionally retains the checkpoint-compatible + Diffusers `DPMSolverMultistepScheduler`. +- Known limitations: + - Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution + are not validated for the native pipeline. + - The Diffusers backend is a compatibility path and does not provide native + vLLM-Omni parallelism or continuous batching. + - Native describes pipeline and Transformer ownership, not a zero-Diffusers + dependency guarantee. + +### 1x NVIDIA L20X + +#### Environment + +- Driver: 570.133.20 +- PyTorch: 2.13.0+cu130 +- Diffusers: 0.38.0 +- vLLM: 0.27.0 +- Nsight Systems: 2026.1 +- vLLM-Omni commit: `54fcc6d` + +#### Same-workload T2V comparison + +Both backends used +`Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, the same T2V prompt +and seeds, 832×480 output, 81 frames, and four inference steps. One warmup +request was excluded, followed by three measured requests. + +| Metric | Native vLLM-Omni | Diffusers adapter | Native delta | +| --- | ---: | ---: | ---: | +| E2E latency | 7.1600 ± 0.0043 s | 6.9761 ± 0.0517 s | +2.636% | +| Throughput | 0.1397 req/s | 0.1433 req/s | -2.568% | +| Peak reserved VRAM | 23,936 MiB | 19,060 MiB | +25.582% | +| Initialization | 11.669 s | 20.236 s | -42.336% | + +The peak-memory result is not a Transformer-only comparison. For the 480p +checkpoint, the native pipeline deliberately loads and decodes with the Wan +VAE in FP32. This matches the upstream SANA-Video reference and the recorded +accuracy-golden setup for this integration. The adapter launch uses +`--dtype bfloat16` as the Diffusers pipeline-wide dtype, so its Wan VAE also +runs in BF16. + +Peak VRAM is the CUDA allocator's maximum reserved memory over the complete +request. It includes VAE decode activations and allocator-reserved blocks, not +only live Transformer tensors. The different VAE precision policies are a +material code-level difference and a likely contributor to the higher native +high-water mark; the available full-request measurements do not isolate how +much of the delta they account for. The 25.582% result should therefore not be +interpreted as an isolated native Transformer memory regression. Exact +attribution requires both paths to use the same VAE dtype. Changing the native +480p VAE to BF16 would instead change the recorded golden/reference setup and +requires a new accuracy study. diff --git a/recipes/NVIDIA/SANA-WM.md b/recipes/NVIDIA/SANA-WM.md new file mode 100644 index 000000000..2926cc46f --- /dev/null +++ b/recipes/NVIDIA/SANA-WM.md @@ -0,0 +1,193 @@ +# SANA-WM + +> Camera-controllable first-frame image-to-video world model. + +## Summary + +- Vendor: Efficient-Large-Model / NVlabs SANA +- Model: `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` (standard diffusers layout, converted offline from the NVlabs release; Stage-1 transformer + VAE only) +- Task: First-frame image-to-video generation with camera control +- Mode: Online serving with the OpenAI-compatible video API +- Model weights: about 13 GB for the Stage-1 transformer (10 GB) and VAE (2.3 GB) +- Local disk: reserve about 40 GB for the Hugging Face cache and runtime artifacts +- Recommended GPU: 24 GB or larger CUDA GPU +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to serve SANA-WM through `/v1/videos` or +`/v1/videos/sync`. The model takes a text prompt, a first-frame image, and +either an action DSL string or explicit camera poses. vLLM-Omni serves the +SANA-WM Stage-1 DiT, decoded through the SANA VAE. The optional LTX-2 refiner +stage is not supported by this integration; it is a planned follow-up. + +## References + +- Upstream model card: +- Video API: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md) + +## Hardware Support + +## GPU + +### 1x NVIDIA RTX PRO 6000 Blackwell 96GB + +#### Capacity + +- Model storage: the Stage-1 transformer is about 10 GB and the VAE about + 2.3 GB. The Gemma text encoder is a separate 4.9 GB repo. The engine prefetches + the whole model repo at startup (`allow_patterns=["*"]`), which is why the + Stage-1 weights live in their own repo — the two-stage one carries an + additional 84 GB `refiner/` that this path never loads. +- Text encoder: the pipeline tries `google/gemma-2-2b-it` first, then falls back + to the ungated mirror `Efficient-Large-Model/gemma-2-2b-it`. The first repo is + gated, so without an accepted licence and `hf auth login` you will see one + failed load before the fallback succeeds — that is expected, not an error to + chase. Set `VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER` to pin a specific repo or a + local path and skip both. +- Disk sizing: provision about 40 GB of local disk or Hugging Face cache volume + so the model, temporary downloads, and generated artifacts fit without cache + eviction. +- GPU sizing: the default 1280x704, 161-frame, 60-step serving profile peaks at + 22.6 GB of device memory and takes about 133 s to generate on one RTX PRO + 6000 Blackwell. The peak lands in the VAE decode, which is why the pipeline + forces VAE tiling on regardless of `vae_use_tiling` — without it the same + request costs about 9 GB more, and a 321-frame one OOMs outright. On smaller + GPUs, lower `width`, `height`, or `num_frames` before serving production + requests. + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: NVIDIA driver with CUDA runtime supported by your PyTorch + build +- Recommended operator library: Triton, installed through the vLLM/vLLM-Omni + Python environment +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +The repo ships the standard Diffusers layout (`model_index.json` + +`transformer/`, `vae/`), and its `model_index.json` names `SanaWmPipeline`, so +the pipeline class resolves on its own. + +```bash +CUDA_VISIBLE_DEVICES=0 \ +vllm serve BBBBruce/SANA-WM_bidirectional-stage1-diffusers \ + --omni \ + --host 0.0.0.0 \ + --port 8091 +``` + +No deploy config: single-stage diffusion models are deliberately absent from +`OMNI_PIPELINES` (`vllm_omni/config/pipeline_registry.py`), so stage resolution +falls back to the default stage config and a YAML's stage settings — including +`default_sampling_params` — would not be applied. The production generation +settings therefore live in the model (`num_inference_steps=60`, +`guidance_scale=5.0`); a request that omits a field gets them. The examples +below still pass every field explicitly so the numbers are visible. + +If you point this at the older two-stage repo +(`BBBBruce/SANA-WM_bidirectional-diffusers`), startup fails with `Model class +SanaWmTwoStagesPipeline not found in diffusion model registry`, because that +repo's `model_index.json` names a class this build does not register. Add +`--model-class-name SanaWmPipeline` to override it. + +#### Verification + +Use a short smoke request first: + +```bash +curl -sS -X POST http://localhost:8091/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "prompt=A slow forward camera move through a quiet city street." \ + -F "negative_prompt=blurry, low quality, distorted, watermark" \ + -F "input_reference=@/path/to/first_frame.png;type=image/png" \ + -F "width=1280" \ + -F "height=704" \ + -F "num_frames=9" \ + -F "fps=16" \ + -F "num_inference_steps=2" \ + -F "guidance_scale=5.0" \ + -F "seed=42" \ + --form-string 'extra_params={"sana_wm":{"action":"w-8","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \ + -o sana_wm_smoke.mp4 +``` + +For a production-length request, note that the action durations must sum to +`num_frames - 1` — the rollout includes the identity start pose — and a +mismatch is rejected rather than padded or truncated: + +```bash +curl -sS -X POST http://localhost:8091/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "prompt=A slow forward camera move through a quiet city street." \ + -F "negative_prompt=blurry, low quality, distorted, watermark" \ + -F "input_reference=@/path/to/first_frame.png;type=image/png" \ + -F "width=1280" \ + -F "height=704" \ + -F "num_frames=161" \ + -F "fps=16" \ + -F "num_inference_steps=60" \ + -F "guidance_scale=5.0" \ + -F "seed=42" \ + --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \ + -o sana_wm_output.mp4 +``` + +Use `POST /v1/videos` instead when you want job storage and polling rather than +inline MP4 bytes. It accepts the same form fields as `/v1/videos/sync`. + +```bash +create_response=$(curl -sS -X POST http://localhost:8091/v1/videos \ + -H "Accept: application/json" \ + -F "prompt=A slow forward camera move through a quiet city street." \ + -F "negative_prompt=blurry, low quality, distorted, watermark" \ + -F "input_reference=@/path/to/first_frame.png;type=image/png" \ + -F "width=1280" \ + -F "height=704" \ + -F "num_frames=161" \ + -F "fps=16" \ + -F "num_inference_steps=60" \ + -F "guidance_scale=5.0" \ + -F "seed=42" \ + --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}') + +video_id=$(echo "$create_response" | jq -r '.id') +curl -sS "http://localhost:8091/v1/videos/${video_id}" | jq . +curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o sana_wm_output.mp4 +``` + +#### Notes + +- Sequence parallelism is not supported. The bidirectional gated delta + recurrence carries state across frames, so a rank cannot denoise a slice + of the token sequence in isolation; supporting it needs a distributed scan + or an all-gather before the GDN blocks. + +- `input_reference` is required for the first frame. Use `image_reference` only + when you need a JSON-safe image URL or data URL instead of a multipart upload. +- `sana_wm` must provide exactly one of `action` or `camera`. +- Action strings use comma-separated `-` segments. Supported + keys are `w`, `a`, `s`, `d` for translation and `i`, `j`, `k`, `l` for + pitch/yaw rotation. The durations must sum to `num_frames - 1`. +- Explicit camera control (alternative to `action`): pass + `"camera": {"poses": [...]}` where `poses` is a list of `num_frames` + camera-to-world 4x4 matrices (row-major, OpenCV `+X right, +Y down, +Z forward` + convention), e.g. + `extra_params={"sana_wm":{"camera":{"poses":[[[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], ...]},"intrinsics":{...}}}`. + Most callers should prefer `action`; explicit poses exist for callers that + already have a per-frame trajectory. +- Explicit `intrinsics` are recommended and take the mapping form + `{"fx":640,"fy":640,"cx":640,"cy":352}` (for 1280x704). This `{fx,fy,cx,cy}` + mapping is the only accepted intrinsics form; omit `intrinsics` to derive them + from the output resolution. All four values must be finite, and `fx`/`fy` + must be positive — the ray map divides by them. +- The video API returns decoded MP4 bytes and has no `output_type` field, so + raw Stage-1 latents are reachable only from the offline API + (`OmniDiffusionSamplingParams(output_type="latent")`); see + [`tests/e2e/offline_inference/test_sana_wm.py`](../../tests/e2e/offline_inference/test_sana_wm.py). + Putting `output_type` in `extra_params` does not work: it lands in + `sampling_params.extra_args`, while the pipeline reads the top-level field. diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index 99b064310..86f0883f7 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -71,6 +71,35 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended two-GPU profile. The removed fused two-stage implementation is not retained as a fallback because it would duplicate state machines and correctness paths. +### Graph execution + +Graph boundaries follow each stage's state model: + +| Stage | Graph mode on Ascend | Eager boundary | +| --- | --- | --- | +| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing | +| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates | +| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit | + +Stage 2 keeps `enforce_eager: true` for the generation runner because its +Python request metadata and per-request cache dictionaries are not valid +outer-graph inputs. The backend captures only the deterministic CFM DiT +estimator, with an exact graph key for each tensor shape. Timestep embedding +stays eager because the upstream implementation creates a CPU frequency tensor; +HiFT stays eager because it generates random phase. The graph cache stores up +to 32 entries by default, after which unseen shapes run eagerly. + +An ACL graph capture failure is fatal to that Stage-2 process because older +torch-npu releases can leave allocator and RNG capture state invalid. Restart +the service after a capture failure. To run without capture, set +`code2wav_enable_npu_graph: false` under the Stage-2 +`platforms.npu.stages[].additional_config` block before startup. Tune the cache +limit there with `code2wav_max_npu_graphs`. Graph mode also requires +`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`. + +The inner NPUGraph is independent of the outer runner setting, so do not use a +global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired. + ## GPU ### 1 x GPU (default — single command) @@ -211,9 +240,9 @@ speech output (TTS)"** checkbox on / off. model processes still share one CUDA device. - `--trust-remote-code` is required — the HF repo ships a custom `MiniCPMO` config / model class. -- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains - eager because its request-owned Flow/HiFT caches and variable chunk/cache - shapes are not yet exposed through a static exact-shape graph wrapper. +- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its + request orchestration eager while using inner CFM and HiFT CUDA Graphs. On + Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay. - All default stages use `max_num_seqs: 4` to reduce cross-process GPU contention. Talker AR state and Code2Wav caches are request-owned; Code2Wav batches only @@ -224,8 +253,8 @@ speech output (TTS)"** checkbox on / off. - `StageRequestStats.batch_size` is a request-scoped placeholder, not the scheduler's execution batch. - Single-GPU co-location trades throughput for hardware density: Stage 0/1 - CUDA Graph replay and eager Stage 2 vocoder kernels compete across three - CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for + CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA + contexts. Use the 8x4090 config or a custom multi-GPU mapping for throughput-sensitive serving. ### 8 x RTX 4090 24GB (consumer-GPU layout) diff --git a/recipes/Qwen/Qwen3-TTS.md b/recipes/Qwen/Qwen3-TTS.md index b6d2c9293..ede83d47d 100644 --- a/recipes/Qwen/Qwen3-TTS.md +++ b/recipes/Qwen/Qwen3-TTS.md @@ -18,11 +18,11 @@ examples in this repository. Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint: -| Task Type | Model | Description | -| ------------- | ---------------------------------------- | ------------------------------------------------------------ | -| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control | -| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description | -| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript | +| Task Type | Model | Description | +| ------------- | ---------------------------------------- | ------------------------------------------------------------- | +| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control | +| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description | +| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript | Smaller 0.6B variants are also available for `CustomVoice` and `Base`. @@ -166,6 +166,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty - Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config. - Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`. - Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail. +- Base codec termination: Base requests without an explicit `max_new_tokens` + use a text-scaled safety ceiling (at least 192 codec frames and no more than + the configured model limit). If the Talker reaches that ceiling without + codec EOS, non-streaming serving discards the incomplete audio and retries + once with a fresh seed; an explicit `seed` or `max_new_tokens` disables the + retry. SSE and WebSocket clients receive structured errors and must discard + previously emitted audio when the error contains `"action":"discard"`; + raw PCM streams terminate with a connection error after any partial bytes + already sent. - Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model. ## Hardware Support @@ -252,3 +261,53 @@ curl -X POST http://localhost:8091/v1/audio/speech \ - stage_id: 1 gpu_memory_utilization: 0.15 ``` + +### 1x AMD MI300X, 1.7B checkpoints + +#### Environment + +- OS: Linux 6.8.0-134-generic, x86_64 +- Container: official ROCm image built from `docker/Dockerfile.rocm` +- Python: 3.12.13 +- PyTorch: 2.11.0+gitd0c8b1f +- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211 +- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM +- vLLM version: 0.27.0+rocm723 +- vLLM Omni version or commit: `3fecb6953ca8dc51210cc0421ef24552267a41ef` +- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm` +- ONNX Runtime: onnxruntime-rocm 1.22.2.post3 with `ROCMExecutionProvider` +- transformers: 5.15.0 + +#### Command + +The MI300X tests also worked without `enforce_eager`. To use that setup, remove `enforce_eager: true` from stage 1 under `platforms.rocm` in `vllm_omni/deploy/qwen3_tts.yaml`. + +Set this environment variable before you start Qwen3-TTS: + +```bash +export MIOPEN_FIND_MODE=FAST +``` + +#### Verification + +Each row records one offline request with one prompt and batch size one. Request time is the logged `e2e_wall_time_ms`, and stage 1 time is the logged `e2e_stage_1_wall_time_ms`. + +| Checkpoint | Code2Wav configuration | `MIOPEN_FIND_MODE` | Request time | Stage 1 time | Output duration | Real time factor | Maximum sampled device memory | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | +| CustomVoice | Eager | Unset | 1.683 s | 1.613 s | 5.68 s | 0.30 | 60.91 GiB | +| CustomVoice | Graphs allowed | `FAST` | 1.326 s | 1.259 s | 5.68 s | 0.23 | 62.19 GiB | +| VoiceDesign | Eager | Unset | 1.195 s | 1.180 s | 4.40 s | 0.27 | 60.91 GiB | +| VoiceDesign | Graphs allowed | `FAST` | 1.127 s | 1.054 s | 4.40 s | 0.26 | 62.21 GiB | +| Base | Eager | Unset | 13.330 s | 13.263 s | 4.32 s | 3.09 | 61.74 GiB | +| Base | Graphs allowed | `FAST` | 14.314 s | 14.245 s | 4.80 s | 2.98 | 62.84 GiB | + +All six processes finished successfully, and each output passed the 24 kHz mono WAV checks. + +#### Notes + +- The current ROCm config keeps Code2Wav in eager mode. +- The tested alternative allows graphs and sets [`MIOPEN_FIND_MODE=FAST`](https://rocm.docs.amd.com/projects/MIOpen/en/develop/reference/env_variables.html). MIOpen uses a saved kernel choice when available and uses its immediate fallback otherwise. +- CustomVoice used a graph for one of four logged Code2Wav batches. VoiceDesign used a graph for two of four batches. Base used no graphs for its four logged batches. +- The recorded request time with graphs allowed was lower for CustomVoice and VoiceDesign and higher for Base. The Base runs generated different amounts of audio, so their request times are not a direct comparison. +- Each setup was tested with one request, so the timing is not a performance benchmark. +- GPU memory was sampled once per second with `rocm-smi`. diff --git a/recipes/README.md b/recipes/README.md index e010ade67..f82a5955d 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -39,7 +39,6 @@ recipes/ | Recipe | Task | Hardware | |--------|------|----------| -| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB | | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB | | [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs | | [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ | @@ -47,18 +46,23 @@ recipes/ | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB | | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB | | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 | +| [`HiDream-ai/HiDream-O1-Image.md`](./HiDream-ai/HiDream-O1-Image.md) | Text-to-image with shared offline example | 1x or 2x H100 80GB | | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB | | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) | +| [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) | | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | | [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB | | [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 | | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) | | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU | +| [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified | | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom | +| [`NVIDIA/SANA-Video-2B.md`](./NVIDIA/SANA-Video-2B.md) | Native and Diffusers-adapter 480p/720p T2V+I2V | 1x RTX 5090 32GB | | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) | +| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 | | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | | [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) | | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 | @@ -67,21 +71,23 @@ recipes/ | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB | | [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB | | [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB | +| [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk | | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 | | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 | | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB | | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB | -| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB | -| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) | -| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB | +| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X | +| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X | +| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X | | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB | | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 | | [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) | | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) | | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB | | [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) | +| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB | | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB | -| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB | +| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB / 2x non-standard RTX 4090 48GB (49,140 MiB each) | | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB | | [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ | diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md new file mode 100644 index 000000000..851aeb930 --- /dev/null +++ b/recipes/Robbyant/LingBot-World-2.0.md @@ -0,0 +1,122 @@ +# LingBot-World 2.0 + +> Offline and experimental realtime interactive world generation + +## Summary + +- Vendor: Robbyant +- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` +- Task: image-conditioned interactive world generation +- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks +- Hardware validated: NVIDIA H200 and B200 +- Maintainer: Community + +The checkpoint is separately licensed under CC BY-NC-SA and restricted to +non-commercial use. The vLLM-Omni integration code remains Apache-2.0. + +## Offline generation + +The offline path consumes one source image and an action directory containing +`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape +`[frames, 4]`. + +```bash +python examples/offline_inference/diffusion/lingbot_world_v2.py \ + --prompt "The camera moves slowly forward through the scene." \ + --image /path/to/first_frame.png \ + --action-dir /path/to/actions/forward \ + --num-frames 81 \ + --output lingbot_world_v2.mp4 +``` + +The current causal-fast checkpoint uses four DMD sampling steps. Raw frame +counts must be `9 + 12k`, up to 117 frames. Camera files may contain a +longer bounded source trajectory (the official example contains 269 frames); +the runtime consumes only the prefix needed by the request. + +## Realtime in-process generation + +The realtime example keeps AR-Diffusion KV and model-owned state across +requests. Each JSONL line describes the prompt and/or three latent-frame camera +actions applied at the next chunk boundary: + +```json +{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]} +{"event_id":2,"frames":[["w"],["w"],["w"]]} +{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]} +``` + +Run: + +```bash +python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \ + --image /path/to/first_frame.png \ + --events /path/to/events.jsonl \ + --output-dir /tmp/lingbot-realtime \ + --gpu-memory-fraction 0.6 +``` + +The example writes one latent tensor and one metadata JSON file per chunk. It +exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer +-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket +transport. + +This PR intentionally does not define a public realtime HTTP/WebSocket schema. +An online serving client should be added together with the public transport +API rather than exposing LingBot-specific event fields from the generic model +runtime. The structured camera-interaction frontend and transport work is +tracked separately in +[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527). + +## Realtime identity and controls + +- `session_id` identifies the persistent world and its worker-owned state. +- `event_id` identifies a prompt/control update and remains monotonic across reset. +- `chunk_index` is contiguous from zero and restarts from zero after reset. +- `request_id` correlates one chunk snapshot with its output metadata. +- AsyncOmni uses a separate UUID-suffixed internal engine routing ID. + +The generic runtime transports controls as opaque snapshots. LingBot's adapter +accepts: + +- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`, + `s`, `d`, `i`, `j`, `k`, and `l`; +- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories. + +## Validation + +Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42. +The exercised matrix includes: + +- TP=1 and TP=2 execution; +- two interleaved resident sessions; +- action input and prompt switching; +- seven contiguous chunks crossing the sink plus recent rolling window; +- direct versus paged replay; +- CUDA Graph execution; +- reset, close, failure cleanup, and exact metadata matching; and +- nine-frame VAE decode. + +Official `generate.py` accuracy and performance numbers, together with +generation artifacts, are recorded in the PR validation section for the exact +tested commit. + +## Current limitations + +- Only the 14B causal-fast checkpoint is supported. +- The realtime control plane is internal; there is no public server transport yet. +- AR-Diffusion stages currently require one replica because session-affine + routing across replicas is not implemented. +- One AR block is generated per request and `max_num_seqs` must be one. +- Stateful streaming VAE decode is not implemented; the realtime example emits + latent chunks. +- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization, + Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed. +- No AMD GPU, Ascend NPU, or Intel GPU support is claimed. + +## References + +- Checkpoint: +- Official implementation: +- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py) +- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md) diff --git a/recipes/SenseNova/SenseNova-U1.md b/recipes/SenseNova/SenseNova-U1.md index be25df570..8e9091d63 100644 --- a/recipes/SenseNova/SenseNova-U1.md +++ b/recipes/SenseNova/SenseNova-U1.md @@ -178,6 +178,48 @@ with the expected 1536×2720 resolution. - The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear` for fused QKV and gate/up projections with TP support. +### 1x AMD MI300X 192GB + +#### Environment + +- OS: Linux 6.8.0-134-generic, x86_64 +- Container: official ROCm image built from `docker/Dockerfile.rocm` +- Python: 3.12.13 +- PyTorch: 2.11.0+gitd0c8b1f +- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211 +- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM +- vLLM version: 0.27.0+rocm723 +- vLLM Omni version or commit: `a704c8759c96e123c0d7c89b11f120b1c0f120cf` +- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm` + +#### Command + +```bash +python3 examples/offline_inference/text_to_image/text_to_image.py \ + --model SenseNova/SenseNova-U1-8B-MoT \ + --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look." \ + --width 1536 \ + --height 2720 \ + --seed 42 \ + --num-inference-steps 50 \ + --cfg-scale 4.0 \ + --extra-body '{"think": true, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \ + --enable-diffusion-pipeline-profiler \ + --log-stats \ + --output sensenova_u1_mi300x.png +``` + +#### Verification + +The command completed and wrote a valid 1536 by 2720 RGB PNG. + +#### Notes + +- Generation with 50 inference steps took 34.011 seconds. +- Model loading used 32.774 GiB and took 13.987 seconds. +- The internal profiler recorded 35.67 GB reserved and 35.10 GB allocated for the request. +- The highest one second whole device memory sample was 37.93 GiB. + ## Online Serving SenseNova-U1 supports all four modalities via the OpenAI-compatible diff --git a/recipes/StabilityAI/Stable-Audio-Open.md b/recipes/StabilityAI/Stable-Audio-Open.md index a7f18128d..d24d516ca 100644 --- a/recipes/StabilityAI/Stable-Audio-Open.md +++ b/recipes/StabilityAI/Stable-Audio-Open.md @@ -23,8 +23,6 @@ offline validation sample with TeaCache and online serving through the - Model: - Offline example: [`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio) -- Online example: - [`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio) - Audio generation API: [`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md) - Related issue: @@ -78,7 +76,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \ Start the online serving endpoint: ```bash -vllm-omni serve /path/to/stable-audio-open-1.0 \ +vllm serve /path/to/stable-audio-open-1.0 \ --host 0.0.0.0 \ --port 8091 \ --gpu-memory-utilization 0.9 \ @@ -103,7 +101,7 @@ curl -X POST http://localhost:8091/v1/audio/generate \ "seed": 42, "response_format": "wav" }' \ - --output examples/online_serving/stable_audio/piano_10s.wav + --output piano_10s.wav ``` ##### Verification @@ -137,12 +135,12 @@ PY Validate the online output: ```bash -ls -lh examples/online_serving/stable_audio/piano_10s.wav +ls -lh piano_10s.wav python - <<'PY' import soundfile as sf -path = "examples/online_serving/stable_audio/piano_10s.wav" +path = "piano_10s.wav" audio, sample_rate = sf.read(path) print("sample_rate:", sample_rate) print("shape:", audio.shape) @@ -164,7 +162,50 @@ PY `torchaudio==2.11.0+cu130`. - The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary warnings observed during validation did not prevent successful generation. -- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts, - ROCm, XPU, and NPU setups are not covered here. +- The RTX 4090 entry was validated on one 24 GB GPU. The MI300X entry below covers one 192 GB GPU. - Long generations, higher inference-step counts, and non-WAV response formats were not benchmarked in this recipe. + +#### 1x AMD MI300X 192GB + +##### Environment + +- OS: Linux 6.8.0-134-generic, x86_64 +- Container: official ROCm image built from `docker/Dockerfile.rocm` +- Python: 3.12.13 +- PyTorch: 2.11.0+gitd0c8b1f +- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211 +- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM +- vLLM version: 0.27.0+rocm723 +- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0` +- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm` + +##### Commands + +```bash +python3 examples/offline_inference/text_to_audio/text_to_audio.py \ + --model stabilityai/stable-audio-open-1.0 \ + --prompt "A gentle piano melody with soft room ambience" \ + --negative-prompt "Low quality, distorted, noisy" \ + --seed 42 \ + --guidance-scale 7.0 \ + --audio-length 10.0 \ + --num-inference-steps 50 \ + --cache-backend tea_cache \ + --enable-diffusion-pipeline-profiler \ + --output stable_audio_10s.wav +``` + +##### Verification + +The command completed and wrote a valid 44.1 kHz stereo WAV with 10.00 seconds of audio. + +##### Notes + +- TeaCache ran with `rel_l1_thresh=0.2`. +- Model loading used 2.7891 GiB and took 3.706 seconds. +- Generation took 4.750 seconds, which gives a real time factor of 0.475 for the 10.00 second output. +- The internal profiler recorded 15.65 GB reserved and 9.68 GB allocated for the request. +- The highest one second whole device memory sample was 19.61 GiB. +- The output RMS was 0.0887, and its peak absolute amplitude was 0.5761. +- The full process took 384 seconds, including startup and compilation. diff --git a/recipes/TEMPLATE.md b/recipes/TEMPLATE.md index 6af66f721..f2a145660 100644 --- a/recipes/TEMPLATE.md +++ b/recipes/TEMPLATE.md @@ -39,6 +39,9 @@ platform section, document one or more tested hardware configurations. #### Command +Serve commands use `vllm serve --omni`. The `--omni` flag is what +selects the Omni pipeline, so it is required and cannot be dropped. + ```bash # Add the exact command(s) here ``` diff --git a/recipes/XiaomiMiMo/MiMo-Audio.md b/recipes/XiaomiMiMo/MiMo-Audio.md new file mode 100644 index 000000000..2850cad2a --- /dev/null +++ b/recipes/XiaomiMiMo/MiMo-Audio.md @@ -0,0 +1,161 @@ +# MiMo-Audio for omni audio generation and understanding + +> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue + +## Summary + +- Vendor: Xiaomi MiMo +- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct` +- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue +- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a known-good starting point for running +MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D). +The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml) +splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids +FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce +Stage-1 GPU pressure. + +## References + +- Upstream model card: +- Audio tokenizer: +- Offline example: + [`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py) +- Online example: + [`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md) +- Default deploy config (generic single GPU): + [`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml) +- 5090D-tuned deploy config: + [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml) +- Pipeline / deploy docs: + [Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) +- Related issue or discussion: + [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645) +- Tokenizer device alignment fix (required for CPU tokenizer on stage 1): + [PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539) + +## Hardware Support + +## GPU + +### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB + +Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested +offline with `tts_sft`. + +#### Environment + +- OS: Ubuntu (Linux x86_64) +- Python: 3.12 +- PyTorch: 2.11.0+cu129 +- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x) +- GPU: NVIDIA GeForce RTX 5090 D, 32 GB +- vLLM: 0.26.0+cu129 +- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support) + +Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack. +On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using +PyTorch 2.11+cu129. + +Download or cache: + +- `XiaomiMiMo/MiMo-Audio-7B-Instruct` +- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo) + +Export the runtime knobs below before offline or online runs: + +```bash +export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}" +export MIMO_AUDIO_TOKENIZER_DEVICE=cpu +export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0 +export VLLM_USE_FLASHINFER_SAMPLER=0 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +``` + +Optional (mirror / download stability): + +```bash +export HF_HUB_DISABLE_XET=1 +export HF_ENDPOINT=https://hf-mirror.com # if you use a Hugging Face mirror +``` + +Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure +`attention_backend: TRITON_ATTN` in the deploy YAML instead. + +#### Command + +Offline TTS validation from the repository root: + +```bash +cd examples/offline_inference/mimo_audio + +python3 -u end2end.py \ + --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \ + --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \ + --query-type tts_sft \ + --text "The weather is so nice today." +``` + +Online serving (requires MiMo chat template): + +```bash +export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer" + +vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \ + --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \ + --served-model-name "MiMo-Audio-7B-Instruct" \ + --port 18091 \ + --chat-template ./examples/online_serving/mimo_audio/chat_template.jinja +``` + +Pass the same environment exports to the server process when using a CPU +tokenizer or disabling FlashInfer sampling. + +#### Verification + +Offline: + +```bash +ls -lh output_audio/tts_sft/*.wav +``` + +Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt` +with the decoded text. A successful run completes both stages without a +Stage-1 `decode_vq` device mismatch. + +Online: + +```bash +curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health +``` + +Then use the client under +[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/). + +#### Notes + +- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB + with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1. +- Key flags: + - `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory + and `TRITON_ATTN`. + - `MIMO_AUDIO_TOKENIZER_PATH` is mandatory. + - `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires + [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent + fix) so code2wav decode uses the tokenizer device. + - `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler + issues on consumer GPUs. +- Known limitations: + - Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional + tooling; the tested path uses TRITON attention and cu129 wheels. + - Without a working `flash-attn` build, audio quality may be slightly + metallic; install FlashAttention when a compatible wheel is available. + - First request triggers Triton JIT for rotary / attention kernels (latency + spike); subsequent requests are faster. + - For longer outputs or batching, you may need to raise `max_model_len` in + the deploy YAML and match `max_position_embeddings` in the model config + (see the offline example README). diff --git a/recipes/audiox/AudioX.md b/recipes/audiox/AudioX.md deleted file mode 100644 index a2e260736..000000000 --- a/recipes/audiox/AudioX.md +++ /dev/null @@ -1,130 +0,0 @@ -# AudioX - -> AudioX MMDiT for unified audio + music generation: t2a / t2m / v2a / v2m / tv2a / tv2m - -## Summary - -- Vendor: HKUSTAudio (project), `zhangj1an/AudioX` weight bundle -- Model: `zhangj1an/AudioX` -- Task: Text/video → audio or music. Six tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. -- Mode: Offline inference + online serving (pure diffusion) -- Maintainer: Community - -## When to use this recipe - -Use this recipe to run AudioX for sound-effect (`*2a`) or music (`*2m`) generation -from a text prompt and/or video clip. AudioX is a unified diffusion transformer -that produces stereo 44.1 kHz audio up to ~10 s per call. - -## References - -- Project page: -- vLLM-Omni weight bundle: -- Pipeline: `vllm_omni.diffusion.models.audiox.pipeline_audiox.AudioXPipeline` -- Input transforms: `vllm_omni.transformers_utils.processors.audiox` -- Param contract: `vllm_omni/model_extras/audiox.py` (declared `extra_body` knobs) -- Offline example: [`examples/offline_inference/text_to_audio/text_to_audio.py`](../../examples/offline_inference/text_to_audio/text_to_audio.py) -- Online: standard OpenAI chat-completions endpoint (see commands below) - -## Hardware Support - -## GPU - -### 1x L4 24GB - -#### Environment - -- OS: Ubuntu 22.04 -- Python: 3.10+ -- Driver / runtime: CUDA 12.4 -- vLLM version: 0.20.0 -- vLLM-Omni version: 0.1.x - -#### Command - -AudioX uses the **standard** `text_to_audio` example offline and the standard -OpenAI chat-completions endpoint online. All model-specific knobs -(`audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max`, -`cfg_rescale`, `video_path`, `audio_path`) are declared in -`vllm_omni/model_extras/audiox.py` and routed through `extra_body` (online) / -`--extra-body` (offline). The offline example also exposes `--task`, `--video`, -`--audio-start`, and `--audio-length` shortcuts for the most common knobs. - -Offline — text tasks (`t2a` / `t2m`): - -```bash -huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights -python examples/offline_inference/text_to_audio/text_to_audio.py \ - --model ./audiox_weights --task t2a \ - --prompt "Fireworks burst twice, followed by a clock ticking." \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ - --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \ - --output t2a.wav -``` - -Offline — video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) need `--video`: - -```bash -python examples/offline_inference/text_to_audio/text_to_audio.py \ - --model ./audiox_weights --task tv2a \ - --prompt "drum beating sound and human talking" \ - --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \ - --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ - --output tv2a.wav -``` - -Online: - -```bash -DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN \ - vllm serve zhangj1an/AudioX --omni --model-class-name AudioXPipeline --port 8099 - -curl -sS -X POST http://localhost:8099/v1/chat/completions \ - -H 'Content-Type: application/json' \ - -d '{ - "model": "zhangj1an/AudioX", - "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}], - "extra_body": { - "num_inference_steps": 250, - "guidance_scale": 7.0, - "seed": 42, - "audiox_task": "t2m", - "seconds_total": 10.0, - "sigma_min": 0.3, - "sigma_max": 500.0 - } - }' > t2m.json -``` - -For `v2*` / `tv2*` online, attach the video as a `video_url` content item -(data URI for local files) in the message `content`. - -#### Verification - -```bash -# Health check -curl http://localhost:8099/health - -# Listen to the saved file (stereo, 44.1 kHz, sigma_min=0.03, sigma_max=1000 — upstream defaults) -ffprobe t2a.wav -``` - -#### Notes - -- Memory usage: ~10 GB peak with `num_inference_steps=250`, 10 s of audio. -- Output rate: 44.1 kHz stereo regardless of requested rate. -- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via the - `audiox_task` knob (in `extra_body` online, or `--task` offline). -- Video conditioning: `v2*` and `tv2*` require a video; offline via `--video` - (→ `video_path` knob), online via a `video_url` content item. -- Cache acceleration is **not** supported (AudioXPipeline is in `_NO_CACHE_ACCELERATION`). -- Tensor parallelism is supported via `--tensor-parallel-size` (DiT QKV is sharded with - `QKVParallelLinear`); cross-attention K/V is also TP-sharded. - -### Known limitations - -- Inference uses an inlined DPM-Solver++(3M) SDE sampler (k-diffusion port). Replacing it with - diffusers' `EDMDPMSolverMultistepScheduler` introduces a fixed ~861 Hz resonance and is not - recommended. -- Generation is fixed at 10 s (configured by the bundle's `sample_size`); longer outputs require - a different bundle. diff --git a/recipes/cosmos3/Cosmos3-DistOffload.md b/recipes/cosmos3/Cosmos3-DistOffload.md index 8eeea5cc2..2866e7a0c 100644 --- a/recipes/cosmos3/Cosmos3-DistOffload.md +++ b/recipes/cosmos3/Cosmos3-DistOffload.md @@ -189,8 +189,10 @@ wait - `--dlo-use-allgather` (default) requires all concurrent requests to have the same `num_inference_steps` (AllGather is a collective that needs all ranks synchronized) -- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path; - use `--dlo-no-use-allgather` or disable online quantization +- Online per-tensor FP8 linears with DLO AllGather use the ordinary loader + instead of direct checkpoint mmap, then shard the finalized FP8 weights and + scales. Size host memory for a transient complete FP8 model on every rank at + startup - Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with `--dlo-no-use-allgather` retains the standard TP loader and should be treated as an experimental compatibility path diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index b27d1b473..4aaebf0be 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -195,11 +195,13 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ # Transfer V2V with a precomputed depth control video. `control_path` can point # to a local image/video; edge and blur can also be computed from `input_reference` -# by passing `"edge":true` or `"blur":true`. +# by passing `"edge":true` or `"blur":true`. The reference negative prompt is +# optional; omit the `negative_prompt` form field to use an empty negative branch. curl -sS -X POST http://localhost:8000/v1/videos/sync \ -H "Accept: video/mp4" \ -F "model=nvidia/Cosmos3-Nano" \ -F "prompt=Generate a realistic scene following the provided control video." \ + --form-string "negative_prompt=$(jq -c . recipes/cosmos3/negative_prompt.json)" \ -F "size=1280x720" -F "num_frames=121" \ -F "num_inference_steps=50" -F "seed=125" \ -F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \ @@ -322,15 +324,34 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ `extra_params={"guardrails":false}` (per request) toggles safety. The per-request flag only takes effect when the server was launched **with** guardrails enabled (it cannot re-enable them on a `--no-guardrails` server). - `use_resolution_template` / `use_duration_template` are off by default and only - needed when not using upsampled prompts that already encode resolution/duration. + Outside transfer mode, `use_resolution_template` / `use_duration_template` + are off by default and only needed when not using upsampled prompts that + already encode resolution/duration. For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the API decodes the first or last needed reference frames (`"first"` by default). - **Transfer controls:** `extra_params` may include `edge`, `blur`, `depth`, `seg`, or `wsm`. Each hint accepts `true`, a path string, or an object such as `{"control_path": "/path/to/control.mp4"}`; `edge` also accepts - `preset_edge_threshold` and `blur` accepts `preset_blur_strength`. + `preset_edge_threshold` and `blur` accepts `preset_blur_strength`. Every hint + accepts a non-negative `control_weight`; weights are normalized across active + controls and therefore only set their relative influence. A single positive + weight always normalizes to `1.0`; use `control_guidance` to change the + absolute strength of a single control. With two or more active controls, the + per-control attention passes run replicated on every sequence-parallel + (Ulysses) rank, so Ulysses does not reduce per-rank memory or latency for + multi-control transfer requests. + Transfer always uses Cosmos3's transfer-specific system prompt. By default it + also appends a directive naming every active hint and asking the model to + follow its shape, position, and motion precisely; set the request-level + `emphasize_control_in_prompt` option to `false` for prompt ablations. Transfer + does not add a negative prompt automatically. An optional reference prompt is + provided in [`negative_prompt.json`](negative_prompt.json); compact the JSON + with `jq -c` and pass it through `negative_prompt` as shown in the transfer + example. Transfer enables duration/FPS and resolution metadata on both CFG + branches. Set `use_duration_template` or `use_resolution_template` to + `false` to disable either template. `negative_metadata_mode` accepts `same` + (the transfer default), `inverse`, or `none`. Transfer-level options include `control_guidance`, `control_guidance_interval`, `num_video_frames_per_chunk` (default `93`, `101` for WSM), `num_conditional_frames` (default `1`), diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index 54d4d7e32..c8bfb6bbd 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -125,6 +125,23 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ under the top-level `action` field. Verified on the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`) and the rollout video reliably. +- **Transfer controls:** same semantics as Nano — `extra_params` may include + `edge`, `blur`, `depth`, `seg`, or `wsm` hints; see the Cosmos3-Nano recipe + for the request shapes and full option list. Every hint accepts a + non-negative `control_weight`; weights are normalized across active controls + and therefore only set their relative influence. A single positive weight + always normalizes to `1.0`; use `control_guidance` to change the absolute + strength of a single control. With two or more active controls, the + per-control attention passes run replicated on every sequence-parallel + (Ulysses) rank, so Ulysses does not reduce per-rank memory or latency for + multi-control transfer requests. Transfer always uses Cosmos3's + transfer-specific system prompt, appends a control-adherence directive by + default (`emphasize_control_in_prompt: false` disables it), and enables + duration/FPS and resolution metadata on both CFG branches + (`use_duration_template` / `use_resolution_template` to disable; + `negative_metadata_mode` accepts `same`, `inverse`, or `none`, defaulting to + `same`). Transfer does not add a negative prompt automatically; an optional + reference prompt is provided in [`negative_prompt.json`](negative_prompt.json). ## NPU diff --git a/recipes/cosmos3/negative_prompt.json b/recipes/cosmos3/negative_prompt.json new file mode 100644 index 000000000..9cefcb4d5 --- /dev/null +++ b/recipes/cosmos3/negative_prompt.json @@ -0,0 +1,108 @@ +{ + "subjects": [ + { + "description": "Blurry, poorly defined subjects with inconsistent shapes and unrealistic proportions.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "Extremely low-quality subjects with visible rendering artifacts, broken mesh geometry, and completely unrealistic proportions throughout.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "Poorly generated subjects exhibiting all hallmarks of failed neural rendering \u2014 flickering edges, inconsistent depth, and uncanny spatial relationships.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + } + ], + "background_setting": "A poorly rendered, flat background with visible seams, repeated textures, and inconsistent depth cues. The environment lacks volumetric depth and appears as a painted backdrop rather than a three-dimensional space. Vegetation looks like flat cutouts with no volumetric depth. The background appears to have been composited from multiple source materials at different resolutions, creating visible seams and edge artifacts where elements meet. Textures swim and shift across surfaces in a way that breaks the illusion of solidity \u2014 patterns drift laterally rather than staying anchored to the geometry they belong to. Background elements flicker in and out of existence between frames, particularly at the edges of the field of view. The rendering resolution is visibly lower for distant elements, creating a jarring transition between near and far objects. Cloud textures repeat obviously in the sky with visible tiling. Water surfaces lack proper reflection and refraction, appearing as flat animated textures. Fog and atmospheric effects pop in and out rather than smoothly transitioning. Trees and vegetation exhibit obvious LOAD (level-of-detail) switching. Building facades have inconsistent window spacing and pattern repetition. The overall scene feels like a poorly assembled collage of individually rendered elements rather than a coherent whole.", + "lighting": { + "conditions": "Harsh, flat lighting with no natural variation. The scene appears uniformly lit as if by a single overhead fluorescent light, removing all sense of depth and atmosphere.", + "direction": "Inconsistent light sources \u2014 shadows point in multiple contradictory directions, breaking physical plausibility.", + "shadows": "Hard-edged, unrealistic shadows that pop in and out of existence between frames. Some objects cast no shadows while others have impossibly dark ones that don't animate smoothly with the object's motion. Shadow edges exhibit visible staircase aliasing artifacts. Shadow maps appear to have been rendered at extremely low resolution, creating blocky patterns. Self-shadowing on characters shows visible peter-panning artifacts where shadows detach from their source. Contact shadows between objects and the ground appear and disappear as objects move slightly. Shadow color is pure black with no ambient contribution, creating an unnaturally harsh contrast that flattens the image. Multiple shadow cascades have visible boundaries where resolution changes. The shadow rendering appears to be temporally unstable \u2014 even static objects have shadows that shimmer and crawl frame to frame, breaking the illusion of a stable light source.", + "illumination_effect": "No bounce light, no ambient occlusion, no subtle color interactions between surfaces. The scene looks like a poorly lit 3D render from the early 2000s." + }, + "aesthetics": { + "composition": "Cluttered, poorly framed composition with no clear focal point. Important elements are cut off by the frame edges. The rule of thirds is completely ignored, leading to an unbalanced and visually unpleasant arrangement.", + "color_scheme": "Oversaturated, garish colors that clash violently. Color banding is visible in gradient areas. The overall palette feels artificial and digitally processed rather than natural.", + "mood_atmosphere": "Unsettling, uncanny atmosphere that fails to evoke any intended emotional response. The scene feels lifeless and sterile despite attempting to portray dynamic action.", + "patterns": "Visible tiling artifacts in textures, moir\u00e9 patterns, and aliasing on edges." + }, + "cinematography": { + "camera_motion": "Extremely shaky, unstable camera with visible rolling shutter artifacts. The motion is jerky and discontinuous, causing motion sickness and making the scene impossible to follow.", + "framing": "Poorly framed shots that cut off important elements and include unnecessary empty space.", + "camera_angle": "Awkward, disorienting camera angles that provide no useful spatial information about the scene. The camera path exhibits visible mathematical artifacts suggesting simple interpolation between keyframes rather than natural camera operation. Camera motion is completely disconnected from the scene content \u2014 panning away from action, dollying during dialogue, and shaking during still moments. The camera appears to pass through solid objects occasionally. Zoom is applied digitally rather than optically, revealing progressively worse resolution. Camera motion exhibits non-physical acceleration profiles \u2014 instant starts and stops rather than smooth ease-in/ease-out. Rolling shutter simulation is applied inconsistently, present in some frames but not others. The camera occasionally exhibits impossible motion like teleporting between positions. Virtual camera stabilization creates an uncanny floating sensation disconnected from any physical camera rig.", + "depth_of_field": "Uniform focus throughout, creating a flat, documentary-like appearance with no cinematic depth separation.", + "focus": "Soft, out-of-focus imagery with visible chromatic aberration and lens distortion that was not corrected in post-processing.", + "lens_focal_length": "Inappropriate focal length causing barrel distortion and unnatural perspective compression." + }, + "style_medium": "Low quality compressed digital video with visible encoding artifacts", + "artistic_style": "Amateur, unpolished with inconsistent visual style", + "context": "A poorly produced video with numerous technical and artistic flaws that detract from any intended narrative or visual impact.", + "actions": [ + { + "time": "0:00-0:08", + "description": "Subjects attempt to move but their motion is jerky, temporally inconsistent, and physically implausible. Background elements flicker and shift between frames." + } + ], + "text_and_signage_elements": [], + "segments": [ + { + "segment_index": 0, + "time_range": "0:00-0:08", + "description": "A single continuous shot suffering from severe temporal inconsistencies \u2014 subjects that morph and deform between frames, backgrounds that shift and wobble, and rendering quality that fluctuates visibly over time. Motion blur is applied incorrectly, smearing in directions that don't match actual movement. Frame-to-frame coherence breaks down with individual pixels changing color randomly in flat areas. Texture detail level fluctuates between frames as if the rendering budget varied shot to shot. Color grading drifts over the duration with no creative motivation. Noise patterns change between frames in ways that draw attention rather than being invisible. Overall visual quality degrades progressively from start to finish.", + "key_changes": "No meaningful progression or narrative development. Visual quality degrades over time.", + "camera": "Unstable, poorly controlled camera work with visible mathematical interpolation artifacts." + } + ], + "transitions": [], + "temporal_caption": "The scene opens at 0.0 seconds with a poorly rendered establishing shot that immediately reveals low production quality. At 1.0 seconds, subjects begin to move but their motion is jerky and inconsistent, with limbs bending at unnatural angles and objects clipping through each other. From 2.0 to 4.0 seconds, the camera shakes violently while the scene exhibits visible compression artifacts, color banding in the sky, and flickering in the shadows. Between 4.0 and 6.0 seconds, temporal coherence breaks down as elements appear and disappear between frames, textures swim and morph unnaturally, and the lighting shifts abruptly without physical cause. In the final 2 seconds, the overall visual quality deteriorates further with increasing noise, blur, and a general loss of spatial coherence that makes the scene nearly unwatchable. Additionally, the frame rate appears inconsistent with visible judder and stuttering throughout. Color temperature shifts randomly between warm and cool tones with no motivation. The encode quality degrades in complex regions showing macro-blocking and mosquito noise around moving edges. Temporal noise patterns are spatially correlated, creating swimming artifacts on flat surfaces.", + "audio_description": "", + "physical_realism": "No adherence to physical laws. Objects defy gravity, pass through solid surfaces, and change mass and momentum without cause. Fluid dynamics, cloth simulation, and rigid body physics are all fundamentally broken. Furthermore, conservation of energy is violated as objects gain or lose kinetic energy spontaneously. Elastic collisions produce inelastic results and vice versa. Surface friction is inconsistent \u2014 objects slide on rough surfaces while sticking to smooth ones. Air resistance appears to affect only some objects while others move through the atmosphere unimpeded." + } diff --git a/recipes/inclusionAI/Ming-omni-tts.md b/recipes/inclusionAI/Ming-omni-tts.md index e42959e8a..3666dce0f 100644 --- a/recipes/inclusionAI/Ming-omni-tts.md +++ b/recipes/inclusionAI/Ming-omni-tts.md @@ -5,20 +5,21 @@ ## Summary - Vendor: inclusionAI -- Model: `inclusionAI/Ming-omni-tts-0.5B` +- Model: `inclusionAI/Ming-omni-tts-16.8B-A3B`, `inclusionAI/Ming-omni-tts-0.5B` - Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls - Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API - Maintainer: Community ## References -- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B) +- [Huggingface: inclusionAI/Ming-omni-tts-16.8B-A3B](https://huggingface.co/inclusionAI/Ming-omni-tts-16.8B-A3B) +- [Huggingface: inclusionAI/Ming-omni-tts-0.5B](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B) - Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts) ## Hardware Support -This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment. +This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B and the MoE 16.8B two-stage TTS pipeline deployment. Other hardware is welcome as community validation lands. ## CUDA @@ -39,11 +40,7 @@ Other hardware is welcome as community validation lands. Launch the two-stage talker: ```bash -vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \ - --deploy-config vllm_omni/deploy/ming_tts.yaml \ - --omni \ - --port 8091 \ - --enforce-eager +vllm serve inclusionAI/Ming-omni-tts-16.8B-A3B --omni --port 8091 ``` #### Verification @@ -54,7 +51,7 @@ Basic synthesis (save the WAV bytes): curl -X POST http://localhost:8091/v1/audio/speech \ -H "Content-Type: application/json" \ -d '{ - "model": "inclusionAI/Ming-omni-tts-0.5B", + "model": "inclusionAI/Ming-omni-tts-16.8B-A3B", "input": "你好,这是 Ming 在线语音合成测试。", "response_format": "wav" }' --output ming_tts_basic.wav @@ -74,7 +71,7 @@ Reference-audio zero-shot cloning: REF_AUDIO="$BASE/10002287-00000094.wav" jq -n --arg ref_audio "$REF_AUDIO" \ '{ - model: "inclusionAI/Ming-omni-tts-0.5B", + model: "inclusionAI/Ming-omni-tts-16.8B-A3B", input: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。", ref_audio: $ref_audio, ref_text: "在此奉劝大家别乱打美白针。", @@ -107,7 +104,7 @@ jq -n \ --arg a "$REF_A" \ --arg b "$REF_B" \ '{ - model: "inclusionAI/Ming-omni-tts-0.5B", + model: "inclusionAI/Ming-omni-tts-16.8B-A3B", input: $input, ref_audio: [$a, $b], ref_text: $ref_text, @@ -170,6 +167,35 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ `--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text` for the dialect case. +**Offline smoke result on ROCm 7.2:** + +A later offline smoke run used the same MI300X model with the following software: + +- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942 +- Kernel: Linux 6.8.0-134-generic, x86_64 +- Python: 3.12.13 +- PyTorch: 2.11.0+gitd0c8b1f +- ROCm or HIP: 7.2.53211 +- vLLM: 0.27.0+rocm723 +- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0` + +```bash +python3 examples/offline_inference/text_to_speech/ming_tts/end2end.py \ + --model inclusionAI/Ming-omni-tts-0.5B \ + --deploy-config vllm_omni/deploy/ming_tts.yaml \ + --case basic \ + --text "你好,这是 AMD MI300X ROCm 配方验证。" \ + --output-dir ./ming_tts_output \ + --output-name ming_omni_tts_mi300x.wav \ + --enforce-eager \ + --log-stats \ + --metadata-json ./ming_tts_manifest.json +``` + +The checked run produced a valid 44.1 kHz mono WAV with 4.80 seconds of audio, RMS 0.1449, and peak absolute amplitude 0.8621. The first request took 20.406 seconds, which gives a real time factor of 4.25. The stage weight loads used 1.31 GiB and 1.47 GiB and took 1.756 seconds and 1.264 seconds. The largest one second whole device memory sample was 91.05 GiB, including the reserved KV cache. Both stages used eager mode, and the AR stage selected `TRITON_ATTN`. + +The run used the official ROCm image built from `docker/Dockerfile.rocm`. + ## Notes - The official ROCm image includes the platform dependencies. diff --git a/recipes/k2-fsa/OmniVoice.md b/recipes/k2-fsa/OmniVoice.md new file mode 100644 index 000000000..74e1b12f5 --- /dev/null +++ b/recipes/k2-fsa/OmniVoice.md @@ -0,0 +1,98 @@ +# OmniVoice + +> Offline and online text to speech with automatic voice, voice design, and reference voice cloning + +## Summary + +- Vendor: k2-fsa +- Model: `k2-fsa/OmniVoice` +- Task: Multilingual text to speech +- Mode: Offline inference for all modes and online serving for automatic voice +- Maintainer: Community + +## References + +- [Hugging Face model](https://huggingface.co/k2-fsa/OmniVoice) +- [Offline text to speech examples](../../docs/user_guide/examples/offline_inference/text_to_speech.md#omnivoice) +- [Online text to speech examples](../../docs/user_guide/examples/online_serving/text_to_speech.md#omnivoice) +- [Deploy config](../../vllm_omni/deploy/omnivoice.yaml) + +## Requirements + +Download the checkpoint before the first run: + +```bash +huggingface-cli download k2-fsa/OmniVoice +``` + +Automatic voice and voice design require `transformers>=4.57.0`. Reference voice cloning requires `transformers>=5.3.0`. + +The vLLM Omni deploy config runs the OmniVoice generator and decoder inside one float32 diffusion stage. + +## ROCm + +### 1x AMD MI300X + +#### Environment + +- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942 +- Kernel: Linux 6.8.0-134-generic, x86_64 +- Python: 3.12.13 +- PyTorch: 2.11.0+gitd0c8b1f +- ROCm or HIP: 7.2.53211 +- vLLM: 0.27.0+rocm723 +- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0` +- transformers: 5.15.0 + +The run used the official ROCm image built from `docker/Dockerfile.rocm`. + +#### Automatic voice + +Run the offline example from the repository root: + +```bash +python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \ + --model k2-fsa/OmniVoice \ + --deploy-config vllm_omni/deploy/omnivoice.yaml \ + --text "Hello, this is OmniVoice running on one AMD MI300X." \ + --seed 42 \ + --output omnivoice_mi300x.wav +``` + +The checked run produced a valid 24 kHz mono WAV with 3.60 seconds of audio, RMS 0.1194, and peak absolute amplitude 0.9129. The first request took 0.685 seconds, which gives a real time factor of 0.19. Model loading used 3.5739 GiB and took 9.564 seconds. The request used 4.56 GB reserved and 3.64 GB allocated according to the internal profiler, while the largest one second whole device sample was 8.33 GiB. The log states that the OmniVoice Triton kernels loaded. + +#### Reference voice cloning + +The offline interface also accepts a reference WAV and its transcript: + +```bash +python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \ + --model k2-fsa/OmniVoice \ + --text "Hello, this is a cloned voice." \ + --ref-audio ref.wav \ + --ref-text "This is the reference transcription." \ + --output cloned.wav +``` + +#### Voice design + +```bash +python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \ + --model k2-fsa/OmniVoice \ + --text "Hello, this is a designed voice." \ + --instruct "female, low pitch, British accent" \ + --output designed.wav +``` + +## Online serving + +Online serving currently supports automatic voice: + +```bash +vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code +``` + +```bash +cd examples/online_serving/text_to_speech/omnivoice +python speech_client.py --text "Hello, how are you?" +``` diff --git a/recipes/lerobot/Pi0.md b/recipes/lerobot/Pi0.md new file mode 100644 index 000000000..6341e57a0 --- /dev/null +++ b/recipes/lerobot/Pi0.md @@ -0,0 +1,108 @@ +# Pi0 (π0, Pi-Zero) VLA + +> Vision-Language-Action policy serving over the OpenPI realtime websocket + +## Summary + +- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot) +- Model: `lerobot/pi0_base` +- Task: Vision-Language-Action — multi-camera images + a language instruction + + robot proprioceptive state → a continuous action chunk + `[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does + **not** emit text tokens. +- Mode: Online serving via the OpenPI realtime robot API + (`/v1/realtime/robot/openpi` websocket) +- Maintainer: Community + +## When to use this recipe + +Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot +client connects over the OpenPI websocket and, per control step, sends an +observation (camera frames + state + instruction) and receives an action chunk. + +A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across +calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but +ignored. + +## Weights + +π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the +server download it (`MODEL=lerobot/pi0_base`) or point at a local copy. + +## Hardware Support + +### 1x H200 141GB (Online serving) + +#### Environment + +- OS: Ubuntu 22.04+ +- Python: 3.12+ +- Driver / runtime: NVIDIA CUDA environment +- vLLM / vLLM-Omni: use the commit you are deploying from +- OpenPI client extras (for the example client): `openpi-client` + (`pip install -e packages/openpi-client` from the + [openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`, + `msgpack`, `msgpack-numpy` + +#### Command + +```bash +vllm serve lerobot/pi0_base \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --served-model-name pi0 \ + --deploy-config vllm_omni/deploy/pi0.yaml \ + --enforce-eager --disable-log-stats +``` + +The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage +(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config` +handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim +32) that enables the OpenPI realtime serving layer. + +#### Verification + +Run the example client (sends a synthetic observation and prints the returned +action chunk): + +```bash +python examples/online_serving/pi0/openpi_client.py \ + --host 127.0.0.1 --port 8000 \ + --prompt "pick up the red block and place it in the bin" +``` + +It connects, prints the server metadata, sends robot observations, and prints the +returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation +is a flat dict per inference: + +```python +{ + "observation.images.base_0_rgb": np.uint8[H, W, 3], + "observation.images.left_wrist_0_rgb": np.uint8[H, W, 3], + "observation.images.right_wrist_0_rgb": np.uint8[H, W, 3], + "state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side + "prompt": "pick up the red block", + "session_id": "", # accepted but ignored (π0 is stateless) +} +``` + +Camera keys must match the server's `image_feature_keys` (the checkpoint's +`input_features` order). If your robot uses different camera names, set +`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys. + +## Correctness + +π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` +reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see +`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`). + +## References + +- Model blog: +- Checkpoint: +- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py) +- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml) +- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/) +- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/), + OpenPI websocket e2e in + [`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py) diff --git a/recipes/rednote-hilab/dots.tts.md b/recipes/rednote-hilab/dots.tts.md new file mode 100644 index 000000000..d5c7868cb --- /dev/null +++ b/recipes/rednote-hilab/dots.tts.md @@ -0,0 +1,152 @@ +# dots.tts + +> Offline continuous-AR TTS at 48 kHz (rednote-hilab) + +## Summary + +- Vendor: rednote-hilab +- Model: `rednote-hilab/dots.tts-soar` +- Task: Text-to-speech, zero-shot synthesis only +- Mode: Offline end-to-end example only (not yet wired for online serving — + see [Known limitations](#known-limitations)) +- Maintainer: Community + +## When to use this recipe + +Use this recipe as a known-good starting point for running +`rednote-hilab/dots.tts-soar` offline on vLLM-Omni on consumer-class GPUs. +dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM ++ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio. +It follows the same "vLLM-native base LM + side-path computation" pattern as +VoxCPM2 — single-stage pipeline +`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR +loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone +instead of MiniCPM4, and no FSQ / residual-LM stage. + +This is an early integration (see [Known limitations](#known-limitations) +below before depending on it for anything beyond offline experimentation). + +## References + +- Offline end-to-end script: + [`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py) +- Example guide: + [`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts) +- Default deploy config: + [`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) +- Talker / pipeline source: + [`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/) +- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts) + +## Hardware Support + +This recipe documents one tested 16 GB consumer-GPU configuration. Other +vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as +community validation lands. + +## GPU + +### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended) + +dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++ +speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The +bundled default config at +[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) +(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`, +`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model +weights and ~0.3 GiB peak activation; the remainder of the configured +budget is available for KV cache. Total resident footprint at idle is +roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph +roadmap would be step 8's graph capture (not implemented yet; this +release runs `enforce_eager: true`, so it doesn't apply today). + +#### Environment + +- OS: Linux (WSL2) +- Python: 3.12 +- Driver / runtime: NVIDIA driver 595.95 +- torch: 2.11.0+cu130 +- vLLM: 0.26.0 +- vLLM-Omni: 0.22.1.dev (current `main`) + +#### Command + +```bash +python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ + --model rednote-hilab/dots.tts-soar \ + --text "Hello, this is a test of dots TTS running on vLLM Omni." +``` + +The deploy config at +[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml) +is loaded automatically by the model registry (HF `model_type=dots_tts`). +Pass `--deploy-config ` to override. + +#### Verification + +**T1 — offline zero-shot synthesis**: + +```bash +python examples/offline_inference/text_to_speech/dots_tts/end2end.py \ + --model rednote-hilab/dots.tts-soar \ + --text "Hello, this is a test of dots TTS running on vLLM Omni." +``` + +Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single +one-shot process (init → one `generate()` → exit), so the reported +numbers include engine init and first-request warmup, not just steady-state +per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`, +`RTF: 1.569`. + +Weight-loading breakdown from the same run (all tensors matched, no +missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270 +patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder. + +Whisper transcription of the output matched the input text with no +dropped leading word (confirms the streaming-vocoder patch-boundary fix +described in [Known limitations](#known-limitations)). + +#### Notes + +- Output: 48 kHz mono WAV. +- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default + used throughout this recipe. `rednote-hilab/dots.tts-base` shares the + same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf` + (MeanFlow, 2-4 step) is not supported — see below. +- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy + config are load-bearing, not just conservative defaults: with prefix + caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path + does not preserve the `sparse_audio` marker this model relies on to + route audio output correctly, and generation silently truncates to a + single ~160 ms patch. Do not override `enable_prefix_caching` for this + model until that framework-level gap is fixed. + +## Known limitations + +- **No online serving yet.** dots.tts is not registered in + `vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS + models in this repo use for `/v1/audio/speech`), so `vllm serve + rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible + speech endpoint does not know how to build request params for it yet. + Offline `Omni()` (as used in `end2end.py`) is the only supported path + today. +- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder + weights load and are exercised by `load_weights()`, but nothing in the + prompt builder consumes a reference audio yet — generation is zero-shot + only. `end2end.py` has no `--ref-audio`/`--ref-text` flags. +- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only + the fixed 10-step Euler DiT sampler used by `dots.tts-soar` / + `dots.tts-base` is implemented. +- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three + captured graphs (base LM decode, CFM solver, VAE decode) have no + dots.tts equivalent yet. +- **Concurrent requests do not scale.** Each request's 10-step DiT Euler + integration runs serially in the side path (no cross-request batching, + unlike voxcpm2's `enable_batched_cfm`). A community review of this + integration measured no throughput gain at `c=4` concurrent requests + versus `c=1`. +- **`SamplingParams.seed` does not control audio-generation randomness.** + The DiT's flow-matching noise is deterministically derived per-request + from a fixed internal seed (reproducible run-to-run, matching + voxcpm2's `deterministic_cfm_seed` convention) rather than from the + caller-supplied `seed` field — the same limitation voxcpm2 has today. diff --git a/recipes/zai-org/GLM-TTS.md b/recipes/zai-org/GLM-TTS.md index 2bf1a202a..79bb85c47 100644 --- a/recipes/zai-org/GLM-TTS.md +++ b/recipes/zai-org/GLM-TTS.md @@ -1,4 +1,4 @@ -# GLM-TTS for Chinese/English TTS on 1x GPU +# GLM-TTS for Chinese/English TTS ## Summary @@ -44,20 +44,20 @@ conditioned on reference audio and its transcript. Start the server from the repository root: ```bash -vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 +vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 ``` Async chunking is enabled by default in the bundled deployment config. For the sync (non-streaming) path: ```bash -vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk +vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk ``` Use a custom deploy config for advanced cases: ```bash -vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \ +vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \ --deploy-config /path/to/your_glm_tts_overrides.yaml ``` @@ -103,3 +103,137 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ - Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback. - Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side. - Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40. + +### 2x non-standard RTX 4090 48GB + +This configuration was personally validated on two non-standard RTX 4090 cards +reporting 49,140 MiB each in `nvidia-smi`. Stock 24GB RTX 4090 cards were not +validated. + +#### Environment + +- OS: Ubuntu 22.04.1 LTS +- NVIDIA driver: 580.76.05 +- Python: 3.12.13 +- PyTorch: 2.11.0+cu130 +- CUDA runtime and JIT toolkit: 13.0 +- vLLM: 0.26.0 +- vLLM-Omni: commit `66a9b0c8`, version + `0.26.0rc2.dev5+g66a9b0c84` +- FlashInfer: 0.6.14 + +The host's default `nvcc` was 11.8. The working CUDA toolkit came from the +Python environment, so `CUDA_HOME` was set to +`$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13`. + +#### Command + +From the repository root, start the async-chunk server with: + +```bash +export CUDA_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13" +export PATH="$CUDA_HOME/bin:$PATH" +export LD_LIBRARY_PATH="$CUDA_HOME/lib:${LD_LIBRARY_PATH:-}" + +CUDA_VISIBLE_DEVICES=0,1 vllm-omni serve zai-org/GLM-TTS \ + --omni \ + --trust-remote-code \ + --port 8091 \ + --deploy-config vllm_omni/deploy/glm_tts.yaml \ + --stage-overrides '{"1":{"devices":"1"}}' \ + --allowed-local-media-path "${PWD}/tests/assets" +``` + +Stage 0 runs on GPU 0, while the override places stage 1 on GPU 1. For the sync +variant, add `--no-async-chunk` to the same command. + +#### Verification + +Run the Chinese voice-cloning check: + +```bash +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --model zai-org/GLM-TTS \ + --text "今天天气很好,我们一起学习语音模型推理。" \ + --ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \ + --ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \ + --response-format wav \ + --output glm_tts_zh.wav +``` + +Run the English voice-cloning check with the same Chinese reference: + +```bash +python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ + --model zai-org/GLM-TTS \ + --text "Hello, this is a voice cloning test with GLM-TTS." \ + --ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \ + --ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \ + --response-format wav \ + --output glm_tts_en.wav +``` + +Inspect either output: + +```bash +ffprobe -v error -select_streams a:0 \ + -show_entries stream=codec_name,sample_rate,channels \ + -of default=noprint_wrappers=1 glm_tts_zh.wav +``` + +Expected output: + +```text +codec_name=pcm_s16le +sample_rate=24000 +channels=1 +``` + +The English smoke test produced a valid WAV with an actual duration of 5.020 +seconds. + +#### Observed warm-cache results + +These observations are from one controlled run, not a hardware or performance +guarantee. Cold-start and JIT costs were excluded. Values are mean +/- standard +deviation. + +##### Measurement procedure + +- For the single-GPU server, use the documented command with + `CUDA_VISIBLE_DEVICES=0` and omit `--stage-overrides`; both bundled stages + remain on GPU 0. For the dual-GPU server, use the command as shown. +- Async chunking is the default. Add `--no-async-chunk` for sync mode. +- For non-streaming measurements, use the exact Chinese client command above. + For streaming, add `--stream` to that command and change the output filename + to `glm_tts_zh.pcm`; streaming output is raw PCM, not WAV. +- The fixed Chinese prompt and reference were exactly those shown in + Verification and yielded 3.900 seconds of audio. Each case used one warm-up + followed by three measured requests at concurrency 1. +- Wall time starts immediately before the HTTP POST and stops after the complete + response body. TTFA runs from immediately before the POST to the first + non-empty audio bytes. RTF is wall time divided by audio duration. +- Peak total device memory used was sampled every 0.2 seconds and reported by + NVML via `nvmlDeviceGetMemoryInfo.used`. The readings include background and + baseline use and are not baseline-subtracted. + +| GPUs | Mode | Response | Wall time (s) | TTFA (s) | RTF | +|---:|---|---|---:|---:|---:| +| 1 | Async | Non-streaming | 2.144 +/- 0.076 | -- | 0.550 +/- 0.019 | +| 2 | Async | Non-streaming | 1.969 +/- 0.130 | -- | 0.505 +/- 0.033 | +| 1 | Async | Streaming | 2.135 +/- 0.093 | 1.245 +/- 0.178 | 0.548 +/- 0.024 | +| 2 | Async | Streaming | 1.891 +/- 0.006 | 1.170 +/- 0.042 | 0.485 +/- 0.002 | +| 1 | Sync | Non-streaming | 1.822 +/- 0.016 | -- | 0.467 +/- 0.004 | +| 2 | Sync | Non-streaming | 1.792 +/- 0.023 | -- | 0.459 +/- 0.006 | + +Moving stage 1 to GPU 1 reduced async non-streaming wall time by 8.2%. It +reduced async streaming wall time by 11.4% and TTFA by 6.0%, and reduced sync +wall time by 1.6%. Single- and dual-GPU outputs were bitwise identical only +within matching execution and response modes. + +| Configuration | GPU 0 | GPU 1 | +|---|---:|---:| +| Single-GPU async | 31.52 GiB | -- | +| Dual-GPU async | 29.23 GiB | 2.91 GiB | +| Single-GPU sync | 31.47 GiB | -- | +| Dual-GPU sync | 29.23 GiB | 2.91 GiB | diff --git a/requirements/common.txt b/requirements/common.txt index e0233bd8d..99d7412d2 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -1,18 +1,20 @@ # Common dependencies for all platforms -# transformers: kept exactly aligned with the upstream vLLM constraint. -# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0". -# The previous extra "< 5.9.0" exclusion was removed when re-aligning with -# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from +# transformers: required unified multimodal model classes are unavailable +# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction +# regression. transformers 5.9 also changed +# _adjust_missing_and_unexpected_keys from # list+list to set|set, breaking models that declare # _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano) # with "TypeError: unsupported operand type(s) for |: 'list' and 'set'". # That breakage is now handled in code instead of by pinning: # MossTTSNanoForGeneration wraps weight loading in # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py). -transformers >= 5.5.3 +transformers >= 5.10.1, < 5.15 av>=14.0.0 omegaconf>=2.3.0 -diffusers==0.38.0 +diffusers==0.40.0 +# Hub kernels, including LTX-2.5 NATTEN +kernels==0.14.1 safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 diff --git a/tests/assets/hidream_o1/hidream_o1_baseline.png b/tests/assets/hidream_o1/hidream_o1_baseline.png new file mode 100644 index 000000000..38920f301 Binary files /dev/null and b/tests/assets/hidream_o1/hidream_o1_baseline.png differ diff --git a/tests/assets/soulxsinger/music.json b/tests/assets/soulxsinger/music.json deleted file mode 100644 index d9e0d8882..000000000 --- a/tests/assets/soulxsinger/music.json +++ /dev/null @@ -1,16 +0,0 @@ -[ - { - "index": "vocal_240_51240", - "language": "Mandarin", - "time": [ - 240, - 51240 - ], - "duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14", - "text": " 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 再 也 没 能 忘 掉 你 容 颜 梦 想 着 着 偶 偶 然 然 有 一 一 天 再 相 见 从 此 我 开 始 始 孤 孤 单 思 念 念 想 想 你 时 你 你 在 天 边 想 你 时 你 在 眼 前 前 想 你 时 你 你 在 脑 海 ", - "phoneme": " zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 ", - "note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0", - "note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1", - "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 309.0 320.1 329.6 334.4 334.4 334.5 338.2 346.0 355.1 361.4 366.1 369.5 372.7 373.3 372.9 375.2 375.8 374.7 367.6 360.1 0.0 0.0 0.0 0.0 0.0 0.0 283.6 280.4 277.2 274.3 271.8 271.0 273.3 277.2 278.8 278.1 277.2 277.4 278.8 279.3 280.2 281.4 279.2 274.9 268.3 254.1 239.9 0.0 0.0 267.7 267.7 270.1 271.8 273.1 276.3 278.3 269.1 256.1 0.0 0.0 0.0 0.0 286.0 281.1 275.7 272.7 271.7 271.7 272.5 274.0 275.5 277.9 281.1 287.1 314.0 342.2 363.0 376.8 375.4 365.8 351.8 323.3 0.0 0.0 0.0 0.0 299.6 322.7 336.5 336.9 335.4 332.9 330.0 326.6 323.4 322.3 324.4 326.7 328.5 328.5 326.4 324.0 322.5 323.2 325.5 327.7 328.6 328.1 325.4 320.9 316.5 316.1 320.8 327.7 333.7 333.2 325.6 315.5 306.9 305.5 315.1 331.0 343.3 341.6 331.4 319.5 308.3 307.6 318.1 329.5 338.2 342.0 337.4 329.4 321.0 316.6 319.8 330.8 340.6 344.5 342.3 331.4 319.3 315.7 322.9 329.3 336.5 344.0 341.1 328.3 318.7 320.0 328.2 333.5 337.9 339.7 338.5 327.3 323.3 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 285.3 284.6 284.3 288.2 295.1 305.6 312.2 315.5 287.8 0.0 0.0 0.0 0.0 0.0 309.8 312.4 313.8 312.8 310.8 310.7 312.5 313.2 310.7 305.7 305.9 308.1 309.9 311.3 311.0 309.8 309.7 310.6 311.1 311.1 311.0 311.9 313.7 316.4 319.5 311.6 281.3 283.4 306.6 0.0 0.0 0.0 334.4 326.7 322.6 321.6 321.7 324.5 332.0 336.7 333.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 332.5 331.7 332.7 334.8 334.7 335.9 339.7 345.6 353.9 362.3 370.2 374.0 373.8 371.4 370.0 368.2 367.0 368.0 370.4 371.2 371.3 370.0 360.3 0.0 0.0 0.0 319.1 326.8 331.5 332.0 331.6 327.2 328.2 334.1 343.1 350.7 358.0 363.5 368.4 370.8 371.8 371.1 370.2 368.5 368.0 371.1 375.5 376.3 369.5 0.0 0.0 0.0 280.4 267.2 263.6 265.8 271.2 275.0 276.4 276.9 277.1 280.9 287.5 284.7 280.7 0.0 0.0 0.0 0.0 0.0 0.0 231.6 226.0 222.9 224.1 226.2 228.5 234.3 240.9 246.6 249.8 249.6 246.3 245.2 245.8 247.2 248.3 249.3 249.0 245.9 242.7 235.8 226.3 217.8 213.1 209.3 208.0 207.0 205.9 205.8 203.4 201.5 205.3 207.5 209.1 210.4 208.9 203.2 199.2 199.3 201.2 205.6 206.3 204.0 202.7 202.2 203.3 206.1 205.6 201.9 198.8 195.5 195.6 198.7 204.2 214.5 218.3 214.3 207.0 199.1 192.4 189.9 193.6 201.0 210.6 212.4 209.2 202.0 195.2 190.0 190.4 196.0 204.2 209.2 208.9 203.0 195.8 190.7 189.8 194.6 200.8 208.7 213.5 210.5 202.0 193.9 187.7 187.8 193.1 199.8 202.8 204.1 203.8 200.4 197.0 193.9 191.6 192.5 198.4 204.8 203.9 203.5 201.3 200.2 198.4 198.5 201.0 204.0 204.6 205.4 202.0 199.2 194.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 426.6 432.4 436.6 437.2 434.6 433.0 434.0 453.1 473.8 492.8 503.2 501.8 491.3 479.6 478.1 489.7 503.6 508.3 503.2 489.5 482.3 484.2 495.7 505.4 506.6 501.6 495.4 492.7 497.1 500.4 496.1 490.1 480.4 453.8 412.6 373.6 357.7 354.5 356.8 361.0 365.0 367.8 369.5 370.1 371.1 371.3 370.8 370.1 369.7 369.0 369.5 370.4 371.1 370.9 370.2 370.8 371.3 372.7 380.2 386.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.0 427.8 417.0 411.6 412.6 411.2 410.7 412.5 413.7 414.9 414.8 415.7 415.4 414.0 413.8 414.6 415.2 416.3 415.2 413.6 413.4 414.9 416.8 416.4 414.7 413.8 413.2 414.7 417.5 420.7 428.2 443.1 459.8 481.7 502.3 510.8 509.1 497.8 490.6 493.5 500.8 508.3 511.8 504.6 493.5 487.3 492.3 507.2 518.8 514.8 496.4 489.6 490.4 491.8 493.3 0.0 0.0 0.0 0.0 0.0 391.9 368.8 360.2 358.3 360.0 362.0 361.8 360.9 360.3 361.8 365.7 368.7 369.3 367.7 364.7 365.6 368.4 370.4 371.7 369.0 366.8 367.5 367.9 369.4 374.6 384.6 396.8 388.2 382.8 0.0 0.0 0.0 0.0 0.0 352.6 339.5 328.2 326.0 326.5 326.9 328.2 329.4 329.4 329.3 337.3 349.1 350.9 343.6 333.4 326.5 327.1 331.5 333.4 326.8 318.4 0.0 0.0 0.0 282.0 283.6 281.6 279.5 273.2 267.9 268.8 273.2 277.5 279.8 278.7 276.8 273.3 268.7 265.0 266.1 272.4 282.3 284.9 279.4 268.5 255.3 253.0 259.9 270.2 281.5 284.3 279.7 270.1 260.2 256.4 260.1 266.6 273.0 279.6 283.6 282.3 272.3 262.1 255.8 259.8 269.6 274.9 278.7 278.4 271.5 261.7 255.8 256.8 263.9 272.1 279.7 278.4 269.9 259.0 255.9 262.1 269.2 274.5 279.6 282.1 280.7 274.3 269.9 270.9 270.7 273.8 276.9 274.9 271.1 266.5 266.7 269.4 276.9 282.9 282.0 279.5 274.9 271.2 267.9 263.1 270.2 281.3 285.5 283.9 280.6 271.1 262.8 263.9 263.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 362.9 364.0 366.8 366.6 363.6 359.8 357.7 359.9 367.6 371.6 365.4 359.3 358.6 359.5 365.5 375.7 378.1 371.5 359.5 350.8 358.6 374.7 381.7 376.8 363.6 353.3 356.9 367.8 378.3 380.8 375.7 372.3 372.6 374.7 372.9 365.2 346.4 314.8 285.6 277.3 275.4 275.3 277.5 279.2 279.2 278.8 277.6 276.4 276.5 277.6 278.8 279.9 281.8 282.7 281.0 279.2 278.6 279.4 279.9 279.2 278.7 281.7 285.8 289.9 288.7 284.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 428.0 417.9 410.0 409.5 411.2 414.0 417.5 417.9 418.1 419.3 419.0 417.8 417.1 417.2 416.6 415.8 415.0 415.1 414.5 410.9 401.5 388.2 371.5 359.2 363.0 371.7 375.7 370.6 359.6 353.4 358.7 367.0 374.5 378.2 373.7 362.8 358.0 361.6 368.5 373.7 377.7 375.8 368.8 361.7 360.5 366.7 372.5 375.6 380.1 384.9 383.5 376.6 0.0 0.0 0.0 0.0 0.0 331.7 320.0 316.8 320.7 325.1 327.0 327.4 326.6 326.4 325.5 325.1 326.2 327.6 328.5 328.4 327.7 328.7 329.9 329.8 329.7 327.7 326.5 327.8 328.9 329.3 329.8 329.2 328.8 329.4 330.4 331.0 331.0 330.5 329.2 328.2 328.3 328.7 329.4 327.5 322.9 314.8 299.3 290.8 292.3 299.0 305.4 309.2 312.1 313.9 316.8 320.4 327.6 333.1 338.0 341.2 342.1 338.4 331.0 0.0 0.0 0.0 0.0 0.0 0.0 311.7 292.5 281.8 276.1 275.1 276.6 277.6 277.4 274.0 264.8 249.6 239.1 238.8 243.0 245.7 245.0 241.9 237.9 234.6 237.1 242.7 250.4 252.5 248.3 240.7 231.6 228.3 234.0 240.6 247.0 250.6 248.5 242.1 233.6 226.8 227.2 234.0 240.7 246.9 248.9 244.4 237.6 230.8 228.8 236.4 246.9 250.2 249.2 245.4 239.9 232.5 224.7 226.5 238.5 252.0 257.4 255.9 248.4 237.4 230.5 231.6 238.8 247.6 252.6 253.9 253.0 248.9 241.8 236.4 234.1 236.5 246.2 257.7 255.4 248.2 238.5 234.8 237.8 244.4 250.5 253.8 251.7 246.7 240.0 238.3 244.8 251.6 256.4 256.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 484.0 484.5 487.5 490.6 487.2 479.7 478.9 483.8 494.1 505.7 510.3 501.1 485.4 472.4 472.3 495.6 515.8 517.1 503.8 486.5 475.1 479.9 491.4 507.9 509.9 505.0 496.2 485.7 484.5 490.8 495.4 497.7 496.6 489.9 469.1 424.3 387.2 373.0 367.3 364.4 364.6 366.2 366.6 368.8 371.9 373.1 371.4 368.9 367.1 367.0 367.7 368.8 370.1 369.8 369.2 369.8 370.0 369.0 369.0 371.0 375.2 385.3 391.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 425.2 418.5 409.4 409.5 408.1 406.3 408.4 410.3 412.8 413.8 412.6 411.2 410.8 410.4 410.2 410.5 412.2 413.6 413.4 412.6 412.6 413.4 415.6 417.9 418.5 417.0 415.0 415.0 416.8 421.5 437.9 455.0 472.6 493.0 502.2 496.7 484.6 480.6 487.2 498.6 506.3 506.2 499.4 489.7 489.5 496.2 503.8 507.2 501.8 491.1 489.9 501.1 520.5 533.7 525.8 0.0 0.0 0.0 0.0 0.0 421.8 365.2 351.0 352.0 358.0 362.6 365.6 365.4 363.2 362.5 362.2 362.4 364.1 365.2 364.9 363.5 363.4 366.3 368.5 370.0 369.7 368.5 366.9 365.7 366.3 368.6 370.8 372.2 370.2 373.1 377.0 376.5 371.0 346.8 326.1 325.2 328.4 331.4 334.0 332.2 327.2 324.4 331.3 349.7 363.2 362.0 346.1 327.4 321.8 329.8 340.7 348.4 349.9 346.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 269.5 275.7 282.6 282.9 278.8 275.2 274.3 273.5 273.9 275.3 276.4 277.3 276.1 272.3 268.0 268.3 272.5 279.1 281.7 280.1 271.3 261.0 257.3 262.7 273.2 280.8 283.3 279.3 269.1 258.3 258.0 266.3 278.2 286.7 287.8 283.4 274.9 264.7 257.7 260.0 272.0 286.3 294.7 289.8 274.5 263.8 263.4 268.8 277.3 284.1 286.5 285.2 281.7 272.7 264.3 260.3 267.7 281.7 289.2 289.0 281.2 266.2 256.0 254.4 261.5 276.4 286.8 288.8 286.7 273.8 260.9 260.8 270.2 283.4 291.3 292.8 286.1 273.8 265.2 264.0 271.3 283.8 290.3 289.7 277.7 266.9 260.5 263.0 267.8 281.5 286.7 286.1 285.5 280.1 279.0 283.0 284.1 284.5 285.5 282.2 0.0 0.0 280.0 276.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" - } -] diff --git a/tests/assets/soulxsinger/music.mp3 b/tests/assets/soulxsinger/music.mp3 deleted file mode 100644 index f3126e79e..000000000 Binary files a/tests/assets/soulxsinger/music.mp3 and /dev/null differ diff --git a/tests/assets/soulxsinger/zh_prompt.json b/tests/assets/soulxsinger/zh_prompt.json deleted file mode 100644 index 86ff3614e..000000000 --- a/tests/assets/soulxsinger/zh_prompt.json +++ /dev/null @@ -1,16 +0,0 @@ -[ - { - "index": "vocal_320_10687", - "language": "Mandarin", - "time": [ - 320, - 10687 - ], - "duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45", - "text": " 除 了 想 你 你 除 了 了 爱 你 你 我 什 么 什 么 都 愿 愿 意", - "phoneme": " zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4", - "note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0", - "note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2", - "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" - } -] diff --git a/tests/assets/soulxsinger/zh_prompt.mp3 b/tests/assets/soulxsinger/zh_prompt.mp3 deleted file mode 100644 index bc306410a..000000000 Binary files a/tests/assets/soulxsinger/zh_prompt.mp3 and /dev/null differ diff --git a/tests/benchmarks/benchmark_video_response_encoding.py b/tests/benchmarks/benchmark_video_response_encoding.py new file mode 100644 index 000000000..ec589e734 --- /dev/null +++ b/tests/benchmarks/benchmark_video_response_encoding.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Benchmark legacy and automatic MP4 response encoding on synthetic input.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time +from collections.abc import Callable +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import cast + +import numpy as np + +from vllm_omni.entrypoints.openai.video_api_utils import ( + _encode_video_bytes, + _encode_video_bytes_legacy, + _PlanarFrameConverter, +) + + +@dataclass(frozen=True) +class BenchmarkVariant: + label: str + encoder: Callable[..., bytes] + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("value must be at least 1") + return parsed + + +def _build_inputs( + *, + frames: int, + height: int, + width: int, + fps: int, + audio_sample_rate: int, + seed: int, + input_layout: str = "planar-backed", +) -> tuple[np.ndarray, np.ndarray]: + rng = np.random.default_rng(seed) + planar = rng.random((3, frames, height, width), dtype=np.float32) + video = planar.transpose(1, 2, 3, 0) + if input_layout == "interleaved": + video = np.ascontiguousarray(video) + if not video.flags.c_contiguous: + raise RuntimeError("synthetic video is not C-contiguous") + elif input_layout == "planar-backed" and not all(video[..., channel].flags.c_contiguous for channel in range(3)): + raise RuntimeError("synthetic video does not expose contiguous channel planes") + elif input_layout != "planar-backed": + raise ValueError(f"unsupported input layout: {input_layout}") + + sample_count = max(1, round(frames / fps * audio_sample_rate)) + timeline = np.arange(sample_count, dtype=np.float32) / audio_sample_rate + left = np.sin(2 * np.pi * 220 * timeline).astype(np.float32, copy=False) + right = np.sin(2 * np.pi * 330 * timeline).astype(np.float32, copy=False) + return video, np.stack((left, right)) + + +def _measure( + variant: BenchmarkVariant, + *, + video: np.ndarray, + audio: np.ndarray, + fps: int, + audio_sample_rate: int, +) -> tuple[bytes, dict[str, object]]: + cpu_start = time.process_time_ns() + wall_start = time.perf_counter_ns() + output = variant.encoder( + video, + fps=fps, + audio=audio, + audio_sample_rate=audio_sample_rate, + video_codec_options={"preset": "ultrafast", "threads": "0"}, + ) + wall_ms = (time.perf_counter_ns() - wall_start) / 1_000_000 + process_cpu_ms = (time.process_time_ns() - cpu_start) / 1_000_000 + return output, { + "label": variant.label, + "wall_ms": wall_ms, + "process_cpu_ms": process_cpu_ms, + "output_bytes": len(output), + "output_sha256": hashlib.sha256(output).hexdigest(), + } + + +def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]: + selected = [record for record in records if record["label"] == label] + wall_values = [cast(float, record["wall_ms"]) for record in selected] + cpu_values = [cast(float, record["process_cpu_ms"]) for record in selected] + return { + "runs": len(selected), + "wall_ms": { + "median": statistics.median(wall_values), + "min": min(wall_values), + "max": max(wall_values), + }, + "process_cpu_ms": { + "median": statistics.median(cpu_values), + "min": min(cpu_values), + "max": max(cpu_values), + }, + } + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--frames", type=_positive_int, default=12) + parser.add_argument("--height", type=_positive_int, default=96) + parser.add_argument("--width", type=_positive_int, default=160) + parser.add_argument("--fps", type=_positive_int, default=24) + parser.add_argument("--audio-sample-rate", type=_positive_int, default=32000) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--rounds", type=_positive_int, default=3) + parser.add_argument("--seed", type=int, default=20260817) + parser.add_argument( + "--input-layout", + choices=("planar-backed", "interleaved"), + default="planar-backed", + ) + parser.add_argument( + "--frame-conversion-workers", + type=_positive_int, + default=8, + ) + parser.add_argument("--output-json", type=Path) + args = parser.parse_args() + if args.warmup < 0: + parser.error("--warmup must be non-negative") + return args + + +def main() -> None: + args = _parse_args() + video, audio = _build_inputs( + frames=args.frames, + height=args.height, + width=args.width, + fps=args.fps, + audio_sample_rate=args.audio_sample_rate, + seed=args.seed, + input_layout=args.input_layout, + ) + baseline = BenchmarkVariant("legacy", _encode_video_bytes_legacy) + converter = _PlanarFrameConverter(max_workers=args.frame_conversion_workers) + candidate = BenchmarkVariant( + "automatic", + partial(_encode_video_bytes, frame_converter=converter), + ) + + try: + for _ in range(args.warmup): + for variant in (baseline, candidate): + _measure( + variant, + video=video, + audio=audio, + fps=args.fps, + audio_sample_rate=args.audio_sample_rate, + ) + + records: list[dict[str, object]] = [] + for round_index in range(args.rounds): + order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline) + for variant in order: + _, record = _measure( + variant, + video=video, + audio=audio, + fps=args.fps, + audio_sample_rate=args.audio_sample_rate, + ) + record["round"] = round_index + 1 + records.append(record) + finally: + converter.shutdown() + + output_hashes = {str(record["output_sha256"]) for record in records} + if len(output_hashes) != 1: + raise RuntimeError(f"benchmark variants produced different outputs: {sorted(output_hashes)}") + + paired_savings_ms = [] + paired_savings_percent = [] + for round_index in range(args.rounds): + pair = [record for record in records if record["round"] == round_index + 1] + baseline_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == baseline.label)) + candidate_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == candidate.label)) + paired_savings_ms.append(baseline_ms - candidate_ms) + paired_savings_percent.append((baseline_ms - candidate_ms) / baseline_ms * 100) + + result = { + "config": { + "frames": args.frames, + "height": args.height, + "width": args.width, + "fps": args.fps, + "audio_sample_rate": args.audio_sample_rate, + "warmup": args.warmup, + "rounds": args.rounds, + "seed": args.seed, + "input_layout": args.input_layout, + "frame_conversion_workers": args.frame_conversion_workers, + "video_shape": list(video.shape), + "video_strides": list(video.strides), + "codec_options": {"preset": "ultrafast", "threads": "0"}, + }, + "records": records, + "summary": { + baseline.label: _summarize(records, baseline.label), + candidate.label: _summarize(records, candidate.label), + "paired_wall_saving_ms": { + "median": statistics.median(paired_savings_ms), + "min": min(paired_savings_ms), + "max": max(paired_savings_ms), + }, + "paired_wall_saving_percent": { + "median": statistics.median(paired_savings_percent), + "min": min(paired_savings_percent), + "max": max(paired_savings_percent), + }, + "output_sha256": output_hashes.pop(), + }, + } + output_json = json.dumps(result, indent=2, sort_keys=True) + if args.output_json is not None: + args.output_json.write_text(output_json + "\n", encoding="utf-8") + print(output_json) + + +if __name__ == "__main__": + main() diff --git a/tests/benchmarks/duplex/test_omni_duplex_eval.py b/tests/benchmarks/duplex/test_omni_duplex_eval.py new file mode 100644 index 000000000..5f77aacc0 --- /dev/null +++ b/tests/benchmarks/duplex/test_omni_duplex_eval.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import hashlib +import json +import sys +from types import SimpleNamespace + +import pytest + +from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module +from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import ( + extract_timed_sentences, + normalize_response_items, + split_text, + validate_clock, +) +from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import ( + DuplexSample, + canonical_task_type, + family_for_split, + load_samples, + task_type_for_split, +) +from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import ( + PROTOCOL_PIN, + build_content_prompt, + build_reminder_prompt, + build_temporal_prompt, + parse_judge_json, + reminder_window, + summarize_pr_results, + temporal_window, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_dataset_mapping_and_manifest(tmp_path): + manifest = tmp_path / "samples.json" + manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8") + sample = load_samples(manifest, media_root=tmp_path)[0] + assert sample.family == "rtd" + assert sample.video == str(tmp_path / "clip.mp4") + assert family_for_split("PR_event_reminder") == "pr" + assert canonical_task_type("post-event-reminder") == "post_event_reminder" + assert task_type_for_split("PR_correction") == "correction" + assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction" + + +def test_hugging_face_subset_names_are_splits(monkeypatch): + calls = [] + + def fake_load_dataset(name, **kwargs): + calls.append((name, kwargs)) + return [{"id": "pr", "question_text": "Correct this."}] + + monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset)) + sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0] + assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})] + assert sample.split == "PR_correction" + assert sample.task_type == "correction" + + +def test_response_aliases_and_clock_guard(): + assert split_text("One. Two!") == ["One.", "Two!"] + assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [ + {"sentence": "x", "start": 800.0, "end": 800.0} + ] + with pytest.raises(ValueError, match="clock=invalid"): + validate_clock({"clock": "invalid"}) + validate_clock({"clock": "invalid"}, allow_invalid=True) + timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}]) + assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}] + with pytest.raises(ValueError, match="only clock=media"): + extract_timed_sentences([], clock="wall") + + +def test_protocol_windows_and_parsing(): + assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48" + assert temporal_window(4, 5, 10) == (2.0, 3.0) + assert temporal_window(1, 1.2, 10) == (0.0, 0.5) + assert reminder_window(5) == (5.0, 15.0) + assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == { + "content_score": 2.35, + "is_relevant": 1, + } + assert parse_judge_json('noise "success_score": 1') == {"success_score": 1} + assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0 + + +def test_protocol_prompt_wording_is_pinned(): + prompts = [ + build_temporal_prompt(1, 2, "response", "question"), + build_content_prompt("response", "question", ["one", "two"]), + build_reminder_prompt("instruction", "response", "proactive_reminder"), + build_reminder_prompt("instruction", "response", "post_event_reminder"), + build_reminder_prompt("instruction", "response", "correction", "answer"), + ] + assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [ + "5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5", + "1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457", + "d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c", + "08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a", + "3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73", + ] + + +def test_frame_sample_content_passes_frames(tmp_path, monkeypatch): + response = tmp_path / "response.json" + response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8") + response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8") + video = tmp_path / "video.mp4" + video.write_bytes(b"video") + monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg") + + class Judge: + def temporal(self, *args, **kwargs): + return '{"temporal_score": 3, "is_relevant": 1}' + + def content(self, prompt, video, frames=None, *, mode="video_url"): + assert mode == "frame-sample" + assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"] + return '{"content_score": 3}' + + sample = DuplexSample( + id="rtd", + split="RTD_OCR", + family="rtd", + task_type=None, + video=video, + video_duration=4, + ) + score = eval_module.evaluate_sample( + sample, + response, + tmp_path / "score.json", + Judge(), + judge_video_mode="frame-sample", + ) + assert score["content"]["frame_count"] == 2 diff --git a/tests/benchmarks/metrics/test_metrics.py b/tests/benchmarks/metrics/test_metrics.py index 32d9d7958..315685258 100644 --- a/tests/benchmarks/metrics/test_metrics.py +++ b/tests/benchmarks/metrics/test_metrics.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for metrics.py """ +import math + import pytest from vllm.benchmarks.serve import TaskType @@ -66,6 +68,22 @@ def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutp return output +def _calculate_test_metrics(outputs, goodput=None): + return calculate_metrics( + input_requests=[], + outputs=outputs, + dur_s=1.0, + tokenizer=None, + selected_percentiles=[50.0], + goodput_config_dict=goodput or {}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=[], + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=1.0, + )[0] + + # ============================================================================ # total_input Tests # ============================================================================ @@ -126,6 +144,228 @@ def test_audio_continuity_aggregation(): assert p99 is not None and p99 > 0.4 +def test_unmeasured_duplex_latency_does_not_add_zero_samples(): + measured = _make_output(100, output_tokens=2) + measured.ttft = 0.1 + measured.audio_ttfp = 0.2 + measured.audio_rtf = 0.5 + measured.duplex_session_metrics = { + "mean_ttft_ms": 100.0, + "mean_ttfp_ms": 200.0, + "mean_rtf": 0.5, + } + listen_only = _make_output(100, output_tokens=0) + listen_only.ttft = listen_only.audio_ttfp = listen_only.audio_rtf = 0.0 + listen_only.duplex_session_metrics = { + "mean_ttft_ms": None, + "mean_ttfp_ms": None, + "mean_rtf": None, + } + + metrics = _calculate_test_metrics([measured, listen_only], {"ttft": 150.0}) + + assert metrics.mean_ttft_ms == 100.0 + assert metrics.mean_audio_ttfp_ms == 200.0 + assert metrics.mean_audio_rtf == 0.5 + assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (1, 1, 1) + assert metrics.request_goodput == 1.0 + + +def test_all_unmeasured_duplex_latency_is_not_reported_as_zero(): + listen_only = _make_output(100, output_tokens=0) + listen_only.duplex_session_metrics = { + "mean_ttft_ms": None, + "mean_ttfp_ms": None, + "mean_rtf": None, + } + + metrics = _calculate_test_metrics( + [listen_only], + {"ttft": float("inf"), "audio_ttft": float("inf")}, + ) + + assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (0, 0, 0) + assert math.isnan(metrics.mean_ttft_ms) + assert math.isnan(metrics.mean_audio_ttfp_ms) + assert math.isnan(metrics.mean_audio_rtf) + assert metrics.request_goodput == 0.0 + + +def test_unmeasured_duplex_tpot_does_not_add_zero_or_misalign_goodput(): + missing_tpot = _make_output(100, output_tokens=5) + missing_tpot.itl = [] + missing_tpot.text_latency = missing_tpot.ttft = 0.1 + missing_tpot.tpot_measured = False + missing_tpot.duplex_session_metrics = {"mean_ttft_ms": 100.0} + + slow_ttft = _make_output(100, output_tokens=5) + slow_ttft.ttft = 1.0 + slow_ttft.text_latency = 1.4 + slow_ttft.itl = [0.1] * 4 + slow_ttft.duplex_session_metrics = {"mean_ttft_ms": 1000.0} + + metrics = _calculate_test_metrics( + [missing_tpot, slow_ttft], + {"ttft": 500.0, "tpot": 200.0}, + ) + + assert metrics.num_tpot_samples == 1 + assert metrics.mean_tpot_ms == pytest.approx(100.0) + assert metrics.request_goodput == 0.0 + + +def test_all_unmeasured_duplex_token_timing_is_not_reported_as_zero(): + output = _make_output(100, output_tokens=5) + output.itl = [] + output.text_latency = output.ttft + output.tpot_measured = False + output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + + metrics = _calculate_test_metrics([output], {"tpot": float("inf")}) + + assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 0) + assert math.isnan(metrics.mean_tpot_ms) + assert math.isnan(metrics.mean_itl_ms) + assert metrics.request_goodput == 0.0 + + +def test_duplex_response_timings_do_not_build_a_session_token_timeline(): + output = _make_output(100, output_tokens=5) + output.latency = 101.0 + output.duplex_request_metrics = [ + {"response_id": "r1", "stage0_tokens": {"itls_ms": [100.0, 100.0]}}, + {"response_id": "r2", "stage0_tokens": {"itls_ms": [100.0, 100.0]}}, + ] + output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + + metrics = _calculate_test_metrics([output]) + + assert math.isnan(metrics.max_output_tokens_per_s) + assert metrics.max_concurrent_requests == 1 + assert metrics.mean_tpot_ms == pytest.approx(100.0) + assert metrics.mean_itl_ms == pytest.approx(100.0) + + +def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback(): + output = _make_output(100, output_tokens=0) + output.generated_text = "timing metadata missing" + output.itl = [] + output.text_latency = output.ttft + output.tpot_measured = False + output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + + def tokenizer(text, *, add_special_tokens): + assert text == output.generated_text + assert add_special_tokens is False + return type("Tokenized", (), {"input_ids": [1, 2, 3]})() + + metrics, _ = calculate_metrics( + input_requests=[], + outputs=[output], + dur_s=1.0, + tokenizer=tokenizer, + selected_percentiles=[50.0], + goodput_config_dict={"tpot": float("inf")}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=[], + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=1.0, + ) + + assert metrics.total_output == 3 + assert metrics.num_tpot_samples == 0 + assert math.isnan(metrics.mean_tpot_ms) + assert metrics.request_goodput == 0.0 + + +def test_zero_itl_does_not_create_zero_tpot(): + output = _make_output(100, output_tokens=3) + output.itl = [0.0, 0.0] + output.duplex_session_metrics = {"mean_ttft_ms": 100.0} + + metrics = _calculate_test_metrics([output], {"tpot": 1.0}) + + assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 2) + assert math.isnan(metrics.mean_tpot_ms) + assert metrics.mean_itl_ms == 0.0 + assert metrics.request_goodput == 0.0 + + +def test_stage_output_tokens_without_client_timing_omit_tpot(capsys): + output = _make_output(100, output_tokens=5) + output.itl = [] + output.text_latency = 0.0 + output.ttft = 0.1 + + metrics, _ = calculate_metrics( + input_requests=[], + outputs=[output], + dur_s=1.0, + tokenizer=None, + selected_percentiles=[50.0, 99.0], + goodput_config_dict={}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=["tpot"], + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=1.0, + ) + + printed = capsys.readouterr().out + assert metrics.num_tpot_samples == 0 + assert math.isnan(metrics.mean_tpot_ms) + assert "Time per Output Token" not in printed + assert "Mean TPOT" not in printed + + +def test_consistent_client_latency_can_supply_tpot_fallback(): + output = _make_output(100, output_tokens=10) + output.itl = [] + output.ttft = 0.1 + output.text_latency = 1.0 + + metrics = _calculate_test_metrics([output]) + + assert metrics.num_tpot_samples == 1 + assert metrics.mean_tpot_ms == pytest.approx(100.0) + + +def test_single_token_responses_do_not_report_zero_tpot(capsys): + output = _make_output(100, output_tokens=1) + + metrics, _ = calculate_metrics( + input_requests=[], + outputs=[output], + dur_s=1.0, + tokenizer=None, + selected_percentiles=[50.0], + goodput_config_dict={}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=["tpot"], + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=1.0, + ) + + printed = capsys.readouterr().out + assert metrics.num_tpot_samples == 0 + assert math.isnan(metrics.mean_tpot_ms) + assert "Time per Output Token" not in printed + + +def test_duplex_goodput_does_not_pair_measurements_from_different_requests(): + text_only, audio_only = _make_output(100), _make_output(100) + text_only.ttft, text_only.audio_ttfp = 0.1, 0.0 + text_only.duplex_session_metrics = {"mean_ttft_ms": 100.0, "mean_ttfp_ms": None, "mean_rtf": None} + audio_only.ttft, audio_only.audio_ttfp = 0.0, 0.2 + audio_only.duplex_session_metrics = {"mean_ttft_ms": None, "mean_ttfp_ms": 200.0, "mean_rtf": None} + + metrics = _calculate_test_metrics([text_only, audio_only], {"ttft": 500.0, "audio_ttft": 500.0}) + + assert metrics.request_goodput == 0.0 + + # ============================================================================ # TTFT suppression for pure-audio (TTS) benchmarks # ============================================================================ diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index e0a9ff49b..d5c714be8 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for patch.py @@ -9,22 +9,98 @@ import asyncio import base64 import json import time +from argparse import Namespace from types import SimpleNamespace import pytest from pytest_mock import MockerFixture from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput +from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( + SeedTTSSampleRequest, + SeedTTSTextSampleRequest, +) from vllm_omni.benchmarks.patch.patch import ( MixRequestFuncOutput, + _apply_stage0_token_timings, + _attach_seed_tts_to_request_func_input, async_request_openai_chat_omni_completions, async_request_openai_realtime_duplex, + should_request_stage_metrics, ) from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu] +def _seed_tts_request_func_input() -> RequestFuncInput: + return RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="target text", + api_url="http://test.com/v1/chat/completions", + prompt_len=2, + output_len=20, + extra_body={"modalities": ["text", "audio"]}, + ) + + +def test_seed_tts_chat_request_carries_reference_audio_once() -> None: + reference_audio_url = "data:audio/wav;base64,AAAA" + sample = SeedTTSSampleRequest( + prompt="target text", + prompt_len=2, + expected_output_len=20, + multi_modal_data=None, + seed_tts_system_prompt="Clone the supplied voice.", + seed_tts_speech_extra={ + "ref_audio": reference_audio_url, + "ref_text": "reference text", + "task_type": "Base", + }, + ) + request_func_input = _seed_tts_request_func_input() + + _attach_seed_tts_to_request_func_input(sample, request_func_input) + + assert request_func_input.omni_chat_messages == [ + { + "role": "system", + "content": [{"type": "text", "text": "Clone the supplied voice."}], + }, + {"role": "user", "content": [{"type": "text", "text": "target text"}]}, + ] + assert request_func_input.extra_body["ref_audio"] == reference_audio_url + serialized_request = json.dumps( + {"messages": request_func_input.omni_chat_messages, **request_func_input.extra_body}, + ) + assert serialized_request.count(reference_audio_url) == 1 + + +def test_seed_tts_text_chat_messages_do_not_add_reference_audio() -> None: + sample = SeedTTSTextSampleRequest( + prompt="target text", + prompt_len=2, + expected_output_len=20, + multi_modal_data=None, + seed_tts_system_prompt="Use the configured voice.", + seed_tts_speech_extra=None, + ) + request_func_input = _seed_tts_request_func_input() + + _attach_seed_tts_to_request_func_input(sample, request_func_input) + + assert getattr(request_func_input, "seed_tts_row", False) is True + assert request_func_input.omni_chat_messages == [ + { + "role": "system", + "content": [{"type": "text", "text": "Use the configured voice."}], + }, + {"role": "user", "content": [{"type": "text", "text": "target text"}]}, + ] + assert "ref_audio" not in request_func_input.extra_body + + class MockResponse: """Mock aiohttp response for testing""" @@ -102,7 +178,24 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) received_at_s=now + 0.03, ) self.events.add( - {"type": "response.done", "response": {"id": response_id}}, + { + "type": "response.done", + "response": { + "id": response_id, + "metadata": { + "vllm_omni": { + "stage_metrics": { + "0": { + "num_tokens_out": 3, + "vllm_ttft_ms": 20.0, + "vllm_tpot_ms": 10.0, + "vllm_itls_ms": [10.0, 10.0], + } + } + } + }, + }, + }, received_at_s=now + 0.04, ) @@ -143,6 +236,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) assert client.configure_kwargs["native_duplex"] is False assert client.configure_kwargs["extra_body"] == { "ref_audio": "data:audio/wav;base64,AAAA", + "return_stage_metrics": True, } assert client.sent == [ event @@ -167,6 +261,10 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) assert output.audio_ttfp > output.ttft assert output.audio_rtf > 0 assert output.latency > 0 + assert output.output_tokens == 12 + assert output.itl == [0.01, 0.01] * 4 + assert output.text_latency == pytest.approx(output.ttft + 0.08) + assert output.tpot_measured is True assert output.tts_turn_pcm_bytes == [b"\x00\x00" * 2400] * 4 assert output.tts_output_pcm_bytes == b"\x00\x00" * 9600 session_id = output.duplex_request_metrics[0]["session_id"] @@ -198,11 +296,160 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch) } +def test_stage0_token_timings_use_weighted_tpot_when_itls_are_incomplete(): + output = MixRequestFuncOutput(ttft=0.1) + + measured = _apply_stage0_token_timings( + output, + [ + {"output_token_count": 3, "itls_ms": [], "tpot_ms": 10.0}, + {"output_token_count": 2, "itls_ms": [], "tpot_ms": 40.0}, + ], + expected_output_tokens=5, + ) + + assert measured is True + assert output.itl == [] + assert output.tpot_measured is True + assert (output.text_latency - output.ttft) / 4 == pytest.approx(0.02) + + +def test_stage0_zero_itls_fall_back_to_positive_tpot(): + output = MixRequestFuncOutput(ttft=0.1) + + measured = _apply_stage0_token_timings( + output, + [{"output_token_count": 3, "itls_ms": [0.0, 0.0], "tpot_ms": 25.0}], + expected_output_tokens=3, + ) + + assert measured is True + assert output.itl == [] + assert output.tpot_measured is True + assert (output.text_latency - output.ttft) / 2 == pytest.approx(0.025) + + +def test_stage0_token_count_without_timing_is_not_measured(): + output = MixRequestFuncOutput(ttft=0.1) + + measured = _apply_stage0_token_timings( + output, + [{"output_token_count": 3, "itls_ms": [], "tpot_ms": None}], + expected_output_tokens=3, + ) + + assert measured is False + assert output.itl == [] + assert output.text_latency == output.ttft + assert output.tpot_measured is False + + def create_sse_chunk(data_dict): """Helper to create SSE formatted chunk""" return f"data: {json.dumps(data_dict)}\n\n".encode() +def test_chat_text_timing_metrics_request_stage_metrics(): + args = Namespace( + backend="openai-chat-omni", + percentile_metrics="ttft,tpot,itl,e2el", + print_stage=False, + extra_body={}, + ) + + assert should_request_stage_metrics(args) is True + + +@pytest.mark.asyncio +async def test_bundled_first_text_chunk_uses_stage0_token_timings(mocker: MockerFixture): + """Engine timings recover TPOT when every text token arrives together.""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "ABCD"}}], + "modality": "text", + "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, + "metrics": { + "stage_metrics": { + "0": { + "num_tokens_out": 4, + "vllm_itls_ms": [10.0, 11.0, 12.0], + "vllm_tpot_ms": 11.0, + } + } + }, + } + ), + b"data: [DONE]\n\n", + ] + mock_response = MockResponse(200, chunks) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + assert output.success is True + assert output.output_tokens == 4 + assert output.itl == pytest.approx([0.010, 0.011, 0.012]) + assert output.text_latency - output.ttft == pytest.approx(0.033) + assert output.tpot_measured is True + + +@pytest.mark.asyncio +async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture): + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "A"}}], + "modality": "text", + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": "B"}}], + "modality": "text", + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + "metrics": { + "stage_metrics": { + "0": { + "num_tokens_out": 2, + "vllm_itls_ms": [1000.0], + "vllm_tpot_ms": 1000.0, + } + } + }, + } + ), + b"data: [DONE]\n\n", + ] + mock_response = MockResponse(200, chunks, delay_between_chunks=0.01) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + assert output.success is True + assert len(output.itl) == 1 + assert 0.0 < output.itl[0] < 0.1 + + # ============================================================================ # output_tokens Tests # ============================================================================ diff --git a/tests/benchmarks/test_daily_omni_local_source.py b/tests/benchmarks/test_daily_omni_local_source.py index 35045a04b..13ae0637a 100644 --- a/tests/benchmarks/test_daily_omni_local_source.py +++ b/tests/benchmarks/test_daily_omni_local_source.py @@ -9,13 +9,17 @@ from pathlib import Path import pytest from vllm_omni.benchmarks.data_modules.daily_omni_dataset import ( + DailyOmniDataset, daily_omni_local_qa_json, daily_omni_local_videos_dir, + ensure_daily_omni_hub_root, resolve_daily_omni_local_root, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark] +_HUB_ID = "liarliar/Daily-Omni" + def _write_qa_json(root: Path) -> Path: qa = root / "qa.json" @@ -23,6 +27,28 @@ def _write_qa_json(root: Path) -> Path: return qa +def _stage_offline_hub_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, with_qa: bool = True) -> Path: + """Build a hub cache holding ``liarliar/Daily-Omni`` and cut off all network access. + + Reproduces the air-gapped NPU runner: the dataset is pre-staged with ``hf download``, so it + only exists in the *hub* cache -- never in the ``datasets`` cache that ``load_dataset`` reads. + """ + constants = pytest.importorskip("huggingface_hub").constants + + cache = tmp_path / "hub" + repo_dir = cache / "datasets--liarliar--Daily-Omni" + snapshot = repo_dir / "snapshots" / "deadbeef" + snapshot.mkdir(parents=True) + if with_qa: + _write_qa_json(snapshot) + (repo_dir / "refs").mkdir() + (repo_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8") + + monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache)) + monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True) + return snapshot + + def test_resolve_local_root_returns_none_for_hub_id() -> None: assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None assert resolve_daily_omni_local_root("") is None @@ -90,3 +116,37 @@ def test_missing_qa_json_reports_none(tmp_path: Path) -> None: root = resolve_daily_omni_local_root(str(tmp_path)) assert daily_omni_local_qa_json(root) is None assert daily_omni_local_videos_dir(root) is None + + +def test_hub_id_resolves_to_offline_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch) + + root = ensure_daily_omni_hub_root(_HUB_ID) + + assert root == snapshot.resolve() + assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve() + + +def test_dataset_loads_hub_id_without_network(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The regression this file exists for: a Hub id must not require ``load_dataset``.""" + snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch) + # Any use of `datasets` here means we went back to the path that cannot read the hub cache. + monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None) + + dataset = DailyOmniDataset(dataset_path=_HUB_ID, random_seed=0) + + assert dataset.qa_json_path == (snapshot / "qa.json").resolve() + assert [row["video_id"] for row in dataset.data] == ["vid0"] + + +def test_uncached_hub_id_reports_both_attempts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _stage_offline_hub_cache(tmp_path, monkeypatch, with_qa=False) + monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None) + + with pytest.raises(RuntimeError) as excinfo: + DailyOmniDataset(dataset_path="someone-else/Daily-Omni", random_seed=0) + + message = str(excinfo.value) + assert "someone-else/Daily-Omni" in message + assert "snapshot_download failed" in message + assert "--daily-omni-qa-json" in message diff --git a/tests/benchmarks/test_daily_omni_pack_modes.py b/tests/benchmarks/test_daily_omni_pack_modes.py index 0fd231857..34d2a5637 100644 --- a/tests/benchmarks/test_daily_omni_pack_modes.py +++ b/tests/benchmarks/test_daily_omni_pack_modes.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes. The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches @@ -27,6 +30,8 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset" if _MODULE_NAME not in sys.modules: _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH) + if _spec is None or _spec.loader is None: + raise ImportError(f"Cannot load {_MODULE_NAME} from {_MODULE_PATH}") _mod = importlib.util.module_from_spec(_spec) sys.modules[_MODULE_NAME] = _mod _spec.loader.exec_module(_mod) @@ -87,7 +92,9 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) -> def fake_extract(video_path, *, audio_path, include_audio, **_kw): counter.append(1) frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)] - segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else [] + segments: list[np.ndarray] = ( + [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else [] + ) return frames, segments monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract) diff --git a/tests/benchmarks/test_omniinteract.py b/tests/benchmarks/test_omniinteract.py new file mode 100644 index 000000000..c38e001bf --- /dev/null +++ b/tests/benchmarks/test_omniinteract.py @@ -0,0 +1,782 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import argparse +import asyncio +import base64 +import io +import json +import subprocess +import tarfile +import threading +import time +import wave +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest +from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput +from vllm.benchmarks.serve import TaskType + +from vllm_omni.benchmarks import omniinteract as oi +from vllm_omni.benchmarks import serve as benchmark_serve +from vllm_omni.benchmarks.data_modules import omniinteract_dataset as data +from vllm_omni.benchmarks.metrics.metrics import calculate_metrics +from vllm_omni.benchmarks.patch import patch as benchmark_patch +from vllm_omni.entrypoints.cli.benchmark.cli_args import preprocess_serve_args +from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand +from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark] +MODEL = "openbmb/MiniCPM-o-4_5" + + +def _case(tmp_path: Path, *, subset: str = "1q1a", name: str = "video.mp4") -> data.OmniInteractCase: + video, annotation = tmp_path / name, tmp_path / f"{Path(name).stem}.json" + video.parent.mkdir(parents=True, exist_ok=True) + video.touch() + annotation.write_text("{}") + return data.OmniInteractCase(subset, name, video, annotation, "multi_turn") + + +def _collector(*events: tuple[dict[str, object], float]) -> RealtimeEventCollector: + collector = RealtimeEventCollector() + for event, received_at in events: + collector.add(event, received_at_s=received_at) + return collector + + +def _created(response_id: str = "r1") -> dict[str, object]: + return {"type": "response.created", "response": {"id": response_id}} + + +def _done(response_id: str = "r1", status: str = "completed") -> dict[str, object]: + return {"type": "response.done", "response": {"id": response_id, "status": status}} + + +def _audio( + response_id: str = "r1", + *, + samples: int = 2400, + rate: int | None = 24_000, + value: int = 1, +) -> dict[str, object]: + event: dict[str, object] = { + "type": "response.audio.delta", + "response_id": response_id, + "format": "pcm16", + "delta": base64.b64encode(bytes((value, 0)) * samples).decode(), + } + if rate is not None: + event["sample_rate_hz"] = rate + return event + + +def _text(response_id: str = "r1", value: str = "hello") -> dict[str, object]: + return {"type": "response.audio_transcript.delta", "response_id": response_id, "delta": value} + + +def _listen(*, buffering: bool = False) -> dict[str, object]: + return { + "type": "response.listen", + "response": {"metadata": {"model_listen": True, "buffering": buffering}}, + } + + +def _session_options(tmp_path: Path) -> data.OmniInteractSessionOptions: + return data.OmniInteractSessionOptions(tmp_path / "artifacts", 10.0, 5.0, "reference.wav", True) + + +def _sample(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions): + return data.OmniInteractSampleRequest( + prompt="", + prompt_len=0, + expected_output_len=0, + multi_modal_data=None, + request_id="measured-0", + omniinteract_case=case, + omniinteract_options=options, + ) + + +def _write_dataset(root: Path) -> Path: + data_root = root / "data" + for subset in ("1q1a", "1q1a_math"): + subset_root = data_root / subset + (subset_root / "videos").mkdir(parents=True) + (subset_root / "annotations").mkdir() + (subset_root / "videos" / f"{subset}.mp4").touch() + (subset_root / "annotations" / f"{subset}.json").write_text("{}") + (subset_root / "video_json_map.json").write_text( + json.dumps( + { + "entries": [ + { + "video": f"videos/{subset}.mp4", + "annotation": f"annotations/{subset}.json", + "scene_type": "multi_turn", + } + ] + } + ) + ) + one_to_many = data_root / "1qna" + (one_to_many / "videos_bench" / "nested").mkdir(parents=True) + (one_to_many / "annotations" / "nested").mkdir(parents=True) + (one_to_many / "videos_bench" / "nested" / "guide.mp4").touch() + (one_to_many / "annotations" / "nested" / "guide.json").write_text("{}") + return data_root + + +def _successful_output( + tmp_path: Path, + *, + audio_time: float = 10.0, + status: str = "completed", +) -> tuple[data.OmniInteractCase, RealtimeEventCollector, oi.OmniInteractCaseResult]: + case = _case(tmp_path) + collector = _collector( + (_created(), 9.9), + (_audio(), audio_time), + (_text(), audio_time + 0.01), + (_done(status=status), audio_time + 0.1), + ) + result = oi.OmniInteractCaseResult(case.subset, str(case.video_path), "") + return case, collector, result + + +def _write_success( + root: Path, + case: data.OmniInteractCase, + collector: RealtimeEventCollector, + result: oi.OmniInteractCaseResult, + **kwargs, +): + require_response = kwargs.pop("require_response", True) + return oi.prepare_success_artifacts( + root, + case, + collector, + stream_start=10.0, + video_duration_s=1.0, + require_response=require_response, + result=result, + **kwargs, + ) + + +def test_dataset_discovers_official_layouts_and_total_selection(tmp_path: Path): + root = _write_dataset(tmp_path) + all_cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True) + selected = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=2, disable_shuffle=True) + assert [case.subset for case in all_cases] == ["1q1a", "1q1a_math", "1qna"] + assert all_cases[-1].video_rel == "videos_bench/nested/guide.mp4" + assert [case.subset for case in selected] == ["1q1a", "1q1a_math"] + + +def _archive(path: Path, member: str) -> None: + with tarfile.open(path, "w") as archive: + content = b"{}" + item = tarfile.TarInfo(member) + item.size = len(content) + archive.addfile(item, io.BytesIO(content)) + + +def test_archive_is_safe_and_atomically_shared(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + bad = tmp_path / "bad.tar" + _archive(bad, "../escape") + with pytest.raises(ValueError, match="Unsafe path"): + data._extract_archive(bad, tmp_path / "bad-cache") + + good = tmp_path / "good.tar" + _archive(good, "data/1q1a/video_json_map.json") + target = tmp_path / "cache" + barrier, safe_extract = threading.Barrier(2), data._safe_extract + monkeypatch.setattr(data, "_safe_extract", lambda *args: (barrier.wait(), safe_extract(*args))) + with ThreadPoolExecutor(2) as executor: + roots = list(executor.map(lambda _: data._extract_archive(good, target), range(2))) + assert roots[0] == roots[1] + assert (roots[0] / "1q1a" / "video_json_map.json").is_file() + assert not list(target.glob(".tmp-*")) + + +def test_hub_archive_uses_vllm_filesystem(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + source = tmp_path / "source.tar" + _archive(source, "data/1q1a/video_json_map.json") + + class FS: + def get_file(self, remote: str, local: str) -> None: + assert remote == "datasets/org/repo/data.tar.gz" + Path(local).write_bytes(source.read_bytes()) + + monkeypatch.setenv("HF_HOME", str(tmp_path / "hf")) + monkeypatch.setattr(data, "hf_fs", lambda: FS()) + assert (data.resolve_omniinteract_root(None, "org/repo") / "1q1a").is_dir() + + +def test_media_commands_are_bounded(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + calls: list[float] = [] + commands: list[list[str]] = [] + bounded_run = oi._run_media_command + + def run(command, *, timeout_s, text=False): + calls.append(timeout_s) + commands.append(command) + if command[0] == "ffprobe": + return subprocess.CompletedProcess(command, 0, "1.0", "") + return subprocess.CompletedProcess(command, 0, b"", b"") + + monkeypatch.setattr(oi, "_run_media_command", run) + duration, pcm, frames = oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=1.0) + assert (duration, len(pcm), frames, calls) == (1.0, 32_000, [None], [3.0, 3.0, 3.0]) + assert commands[1][commands[1].index("-t") + 1] == "1.0" + assert commands[1][commands[1].index("-fs") + 1] == "32000" + + monkeypatch.setattr(subprocess, "run", lambda *a, **k: (_ for _ in ()).throw(subprocess.TimeoutExpired("x", 3))) + with pytest.raises(TimeoutError, match="timed out"): + bounded_run(["ffmpeg"], timeout_s=3.0) + + +def test_media_duration_is_bounded_before_decode(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + calls = 0 + + def run(command, *, timeout_s, text=False): + nonlocal calls + calls += 1 + return subprocess.CompletedProcess(command, 0, "601.0", "") + + monkeypatch.setattr(oi, "_run_media_command", run) + with pytest.raises(ValueError, match="Invalid video duration"): + oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=600.0) + assert calls == 1 + + +def test_final_commit_keeps_a_partial_model_unit(): + """Exact units retain a sample for the server's correlated final flush.""" + + events = [{"session": {"capabilities": {"chunk_period_ms": 1000}}}] + exact_unit = b"\0\0" * 16_000 + assert oi._ensure_final_commit_tail(exact_unit, events) == exact_unit[:-2] + assert oi._ensure_final_commit_tail(exact_unit + b"\0\0", events) == exact_unit + b"\0\0" + + +@pytest.mark.parametrize( + ("events", "match"), + [ + (((_done(), 0.0),), "without response.created"), + (((_created(), 0.0), (_created(), 0.1)), "duplicate response.created"), + (((_created(), 0.0), (_done(), 0.1), (_done(), 0.2)), "duplicate response.done"), + ( + ( + (_created(), 0.0), + ({"type": "response.done", "response": {"id": "r1", "status_details": {"type": "failed"}}}, 0.1), + ), + "reports failure", + ), + ], +) +def test_response_ledger_rejects_identity_errors(events, match: str): + with pytest.raises(ValueError, match=match): + oi.response_ledger(_collector(*events)) + + +class _CompletionClient: + def __init__(self, collector: RealtimeEventCollector): + self.events, self.acks = collector, [] + + def raise_if_reader_stopped(self) -> None: + return None + + async def send_playback_ack(self, response_id: str, played_ms: int) -> None: + self.acks.append((response_id, played_ms)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tail", "succeeds"), + [ + (((_created(), 1.1), (_done(), 1.2)), True), + (((_listen(), 1.1),), True), + (((_listen(buffering=True), 1.1),), False), + ((({"type": "session.closed", "reason": "timeout"}, 1.1),), False), + ], +) +async def test_completion_requires_a_final_model_decision(tail, succeeds: bool): + events = _collector(({"type": "input_audio_buffer.committed"}, 1.0), *tail) + coroutine = oi.wait_for_session_completion( + _CompletionClient(events), oi._Playback(), commit_from=0, timeout_s=0.08, settle_s=0.0 + ) + if succeeds: + assert await coroutine == 0 + else: + with pytest.raises((TimeoutError, RuntimeError)): + await coroutine + + +@pytest.mark.asyncio +async def test_deferred_commit_ignores_the_prior_response_terminal(): + events = _collector( + (_created("old"), 0.0), + ({"type": "input_audio_buffer.committed", "event": {"overlap_deferred": True}}, 1.0), + (_done("old"), 1.1), + (_listen(), 1.2), + ) + assert ( + await oi.wait_for_session_completion( + _CompletionClient(events), oi._Playback(), commit_from=1, timeout_s=0.1, settle_s=0.0 + ) + == 1 + ) + + +@pytest.mark.asyncio +async def test_playback_acks_only_after_serial_audio_drain(): + collector = _collector((_audio(samples=2400), 10.0), (_done(), 10.01)) + client, playback = _CompletionClient(collector), oi._Playback() + await playback.acknowledge(client, now=10.05) + assert client.acks == [] + await playback.acknowledge(client, now=10.1) + assert client.acks == [("r1", 100)] + + +@pytest.mark.parametrize( + ("event", "match"), + [ + (_audio(rate=16_000), "24000 Hz"), + ({**_audio(), "format": "opus"}, "pcm16"), + ({**_audio(), "delta": "!"}, "base64"), + ], +) +def test_playback_rejects_non_official_audio(event: dict[str, object], match: str): + playback = oi._Playback() + with pytest.raises(ValueError, match=match): + playback.ingest(_collector((event, 1.0))) + + +def test_artifacts_publish_official_bundle_and_sparse_deferred_state(tmp_path: Path): + case, collector, result = _successful_output(tmp_path) + summary = _write_success( + tmp_path / "out", + case, + collector, + result, + ) + context = result._artifact_context + assert result.success and result.eligible_for_official_eval + assert summary["scene_type"] == "multi_turn" + assert context and context.spans + assert all("delta" not in event for event in context.events if event["type"] == "response.audio.delta") + oi.publish_deferred_case_artifacts(tmp_path / "out", case, result) + directory = oi._output_dir(tmp_path / "out", case) + assert all((directory / name).is_file() for name in oi.SUCCESS_ARTIFACTS) + with wave.open(str(directory / "output.wav")) as output: + assert (output.getframerate(), output.getnframes()) == (24_000, 24_000) + transcript = json.loads((directory / "wav_transcript.json").read_text()) + assert transcript["text"] == "hello" and transcript["chunks"][0]["timestamp"] == [0.0, 0.1] + assert "serialized playback queue" in transcript["timestamp_semantics"] + assert result._artifact_context is None + + +@pytest.mark.parametrize( + ("audio_time", "status", "reason"), + [(10.95, "completed", "audio_clipped"), (10.0, "cancelled", "cancelled_response")], +) +def test_ineligible_outputs_are_excluded_from_manifest(tmp_path: Path, audio_time: float, status: str, reason: str): + case, collector, result = _successful_output(tmp_path, audio_time=audio_time, status=status) + _write_success(tmp_path / "out", case, collector, result, require_response=False) + assert reason in result.official_eval_ineligible_reasons + oi.write_batch_artifacts(tmp_path / "out", [case], [result]) + assert (tmp_path / "out" / "official_eval_manifest.jsonl").read_text() == "" + + +def test_artifacts_require_complete_audio_and_transcript_response(tmp_path: Path): + case, _, result = _successful_output(tmp_path) + collector = _collector((_created(), 10.0), (_audio(), 10.0), (_done(), 10.1)) + with pytest.raises(ValueError, match="audio and transcript"): + _write_success(tmp_path, case, collector, result) + + +def test_publication_failure_removes_partial_success_bundle(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + case, collector, result = _successful_output(tmp_path) + _write_success(tmp_path, case, collector, result) + original = oi._atomic_write_json + + def fail(path: Path, value: object) -> None: + if path.name == "events.json": + raise OSError("disk full") + original(path, value) + + monkeypatch.setattr(oi, "_atomic_write_json", fail) + with pytest.raises(OSError, match="disk full"): + oi.publish_deferred_case_artifacts(tmp_path, case, result) + directory = oi._output_dir(tmp_path, case) + assert not any((directory / name).exists() for name in oi.SUCCESS_ARTIFACTS) + assert not result.success and "artifact_write_failed" in result.official_eval_ineligible_reasons + + +def test_atomic_write_preserves_destination_on_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + destination = tmp_path / "output.wav" + destination.write_bytes(b"old") + monkeypatch.setattr(oi, "write_pcm16_wav", lambda *a, **k: (_ for _ in ()).throw(OSError("write"))) + with pytest.raises(OSError): + oi._atomic_write_wav(destination, b"\0\0", 24_000) + assert destination.read_bytes() == b"old" + assert not list(tmp_path.glob(".*.tmp")) + + +class _RealtimeClient: + instances: list[_RealtimeClient] = [] + + def __init__(self, url: str, **kwargs): + self.url, self.events, self.acks, self.configure_kwargs = url, RealtimeEventCollector(), [], {} + self.instances.append(self) + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return None + + async def configure(self, model: str, **kwargs) -> None: + self.configure_kwargs = kwargs + self.events.add({"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 1000}}}) + + async def send(self, event: dict[str, object]) -> None: + return None + + async def commit(self) -> None: + now = time.monotonic() + for event, offset in ( + ({"type": "input_audio_buffer.committed"}, 0.0), + (_created(), 0.001), + (_audio(samples=24), 0.002), + (_text(value="answer"), 0.003), + (_done(), 0.004), + ): + self.events.add(event, received_at_s=now + offset) + + async def send_playback_ack(self, response_id: str, played_ms: int) -> None: + self.acks.append((response_id, played_ms)) + + async def close_session(self, **kwargs) -> None: + self.events.add({"type": "session.closed"}) + + def raise_if_reader_stopped(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + case = _case(tmp_path) + config = oi.OmniInteractBenchmarkConfig( + model=MODEL, + output_root=tmp_path / "out", + ref_audio="ref.wav", + require_response=True, + extra_body={"custom": "value"}, + ) + prepared = data.OmniInteractPreparedInput(1.0, b"\0\0\0\0", ("frame",), "data:audio/wav;base64,ref") + monkeypatch.setattr(oi, "RealtimeDuplexClient", _RealtimeClient) + monkeypatch.setattr(oi, "_COMPLETION_SETTLE_S", 0.0) + result = await oi.run_omniinteract_case( + case, + config, + request_index=0, + prepared_input=prepared, + ) + assert result.success and result.transcript == "answer" and result._artifact_context is not None + assert not config.output_root.exists() + assert "autostart=0" in _RealtimeClient.instances[-1].url + assert _RealtimeClient.instances[-1].configure_kwargs["extra_body"] == {"custom": "value"} + assert _RealtimeClient.instances[-1].acks == [("r1", 1)] + + +def test_standard_sample_loading_prepares_media_before_timing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + _write_dataset(tmp_path) + ref = tmp_path / "ref.wav" + ref.touch() + parser = TrackingArgumentParser() + OmniBenchmarkServingSubcommand.add_cli_args(parser) + args = parser.parse_args( + [ + "--backend", + "openai-realtime-duplex", + "--dataset-name", + "omniinteract", + "--dataset-path", + str(tmp_path), + "--model", + MODEL, + "--endpoint", + "/v1/realtime", + "--num-prompts", + "0", + "--omniinteract-ref-audio", + str(ref), + ] + ) + preprocess_serve_args(args) + assert args.endpoint == "/v1/realtime" + monkeypatch.setattr(benchmark_patch, "reference_audio_data_url", lambda _: "data:audio/wav;base64,ref") + monkeypatch.setattr(benchmark_patch, "prepare_media", lambda *a, **k: (1.0, b"pcm", ["frame"])) + samples = benchmark_patch.get_samples(args, None) + assert len(samples) == args.num_prompts == 3 + assert all(sample.omniinteract_prepared_input.video_frames == ("frame",) for sample in samples) + + +def _request(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions) -> RequestFuncInput: + request = RequestFuncInput( + model=MODEL, + model_name=MODEL, + prompt="", + api_url="http://server:8000/v1/realtime", + prompt_len=0, + output_len=0, + logprobs=None, + multi_modal_content=None, + ignore_eos=False, + request_id="measured-0", + ) + request.omniinteract_case, request.omniinteract_options = case, options + return request + + +@pytest.mark.asyncio +async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + case, options = _case(tmp_path), _session_options(tmp_path) + request = _request(case, options) + output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None) + assert not output.success and "prepared before benchmark request timing" in output.error + + prepared = data.OmniInteractPreparedInput(1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref") + request.omniinteract_prepared_input = prepared + request.extra_headers = {"X-Proxy": "test"} + request.extra_body = {"custom": "value"} + monkeypatch.setenv("OPENAI_API_KEY", "secret") + seen = {} + + async def run(_case_arg, config, **kwargs): + seen["headers"] = config.extra_headers + seen["extra_body"] = config.extra_body + seen.update(kwargs) + return oi.OmniInteractCaseResult( + case.subset, + str(case.video_path), + str(options.output_root), + success=True, + transcript="timing metadata missing", + output_tokens=0, + duplex_request_metrics=[{"request_metrics": {"ttft_ms": 1.0}}], + duplex_session_metrics={"mean_ttft_ms": 10.0}, + ) + + monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run) + output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None) + assert output.success and seen["prepared_input"] is prepared + assert seen["headers"] == {"Authorization": "Bearer secret", "X-Proxy": "test", "x-request-id": "measured-0"} + assert seen["extra_body"] == {"custom": "value"} + assert output.itl == [] + assert output.tpot_measured is False + assert seen["capture_artifacts"] is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_metrics", "expected_itl", "expected_text_latency", "expected_tpot_measured"), + [ + ( + [ + {"stage0_tokens": {"output_token_count": 3, "itls_ms": [10.0, 20.0]}}, + {"stage0_tokens": {"output_token_count": 2, "itls_ms": [30.0]}}, + ], + [0.01, 0.02, 0.03], + 0.16, + True, + ), + ( + [ + {"stage0_tokens": {"output_token_count": 3, "tpot_ms": 10.0}}, + {"stage0_tokens": {"output_token_count": 2, "tpot_ms": 40.0}}, + ], + [], + 0.18, + True, + ), + ( + [ + {"stage0_tokens": {"output_token_count": 3, "tpot_ms": 0.0}}, + {"stage0_tokens": {"output_token_count": 2, "tpot_ms": 0.0}}, + ], + [], + 0.1, + False, + ), + ], +) +async def test_adapter_reports_exact_or_weighted_token_timing( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + request_metrics: list[dict[str, object]], + expected_itl: list[float], + expected_text_latency: float, + expected_tpot_measured: bool, +): + case, options = _case(tmp_path), _session_options(tmp_path) + request = _request(case, options) + request.omniinteract_prepared_input = data.OmniInteractPreparedInput( + 1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref" + ) + + async def run(*args, **kwargs): + return oi.OmniInteractCaseResult( + case.subset, + str(case.video_path), + str(options.output_root), + success=True, + output_tokens=5, + duplex_request_metrics=request_metrics, + duplex_session_metrics={"mean_ttft_ms": 100.0}, + ) + + monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run) + output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None) + + assert output.success + assert output.itl == expected_itl + assert output.text_latency == pytest.approx(expected_text_latency) + assert output.tpot_measured is expected_tpot_measured + + +def test_batch_finalization_publishes_only_measured_results(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + case, options = _case(tmp_path), _session_options(tmp_path) + options.output_root.mkdir() + sample = _sample(case, options) + result = oi.OmniInteractCaseResult( + case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True + ) + output = benchmark_patch.MixRequestFuncOutput(success=True) + output.omniinteract_case_result = result + published = [] + monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: published.append(args)) + summary = benchmark_patch._finalize_omniinteract_batch([sample], [output]) + assert summary and summary["total"] == summary["success"] == summary["eligible_for_official_eval"] == 1 + assert summary["failed"] == summary["successful_but_ineligible"] == summary["audio_clipped_bytes"] == 0 + assert published == [(options.output_root, case, result)] + + +def test_case_artifact_failures_do_not_discard_serving_metrics(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + case, options = _case(tmp_path), _session_options(tmp_path) + sample = _sample(case, options) + result = oi.OmniInteractCaseResult( + case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True + ) + output = benchmark_patch.MixRequestFuncOutput(success=True, error="inference warning") + output.omniinteract_case_result = result + monkeypatch.setattr( + benchmark_patch, + "publish_deferred_case_artifacts", + lambda *args: (_ for _ in ()).throw(OSError("case disk full")), + ) + monkeypatch.setattr( + benchmark_patch, + "write_failure_artifacts", + lambda *args: (_ for _ in ()).throw(OSError("failure disk full")), + ) + monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", lambda *args: None) + + summary = benchmark_patch._finalize_omniinteract_batch([sample], [output]) + metrics, _ = calculate_metrics( + [], + [output], + 1.0, + None, + [50.0], + {}, + TaskType.GENERATION, + [], + None, + float("inf"), + 1.0, + ) + + assert summary and summary["failed"] == 1 + assert summary["artifacts_complete"] is False + assert len(summary["artifact_errors"]) == 2 + assert output.success and metrics.completed == 1 + assert output.error.startswith("inference warning\n") + assert "Artifact publication failed: case disk full" in output.error + + +def test_batch_artifact_failure_removes_partial_batch_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + case, options = _case(tmp_path), _session_options(tmp_path) + sample = _sample(case, options) + result = oi.OmniInteractCaseResult( + case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True + ) + output = benchmark_patch.MixRequestFuncOutput(success=True) + output.omniinteract_case_result = result + monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: None) + + def fail_batch(root, *args): + root.mkdir(parents=True, exist_ok=True) + (root / "batch_summary.json").write_text("partial") + raise OSError("batch disk full") + + monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", fail_batch) + + summary = benchmark_patch._finalize_omniinteract_batch([sample], [output]) + + assert summary and summary["success"] == 1 and summary["artifacts_complete"] is False + assert summary["artifact_errors"] == ["Batch artifact publication failed: batch disk full"] + assert output.success + assert not (options.output_root / "batch_summary.json").exists() + + +@pytest.mark.parametrize( + ("field", "value", "match"), + [ + ("omniinteract_ref_audio", None, "ref-audio"), + ("endpoint", "/v1/completions", "v1/realtime"), + ("max_concurrency", 0, "positive"), + ("skip_tokenizer_init", True, "skip-tokenizer-init"), + ], +) +def test_cli_rejects_incompatible_omniinteract_options(field: str, value, match: str): + args = argparse.Namespace( + dataset_name="omniinteract", + backend="openai-realtime-duplex", + endpoint="/v1/realtime", + omniinteract_ref_audio="ref.wav", + max_concurrency=1, + skip_tokenizer_init=False, + ) + setattr(args, field, value) + with pytest.raises(ValueError, match=match): + preprocess_serve_args(args) + + +def test_whole_benchmark_lock_serializes_a_shared_output_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + entered, release = threading.Event(), threading.Event() + calls: list[int] = [] + + async def fake_main_async(args): + calls.append(len(calls)) + if len(calls) == 1: + entered.set() + await asyncio.to_thread(release.wait) + return {"total": 1} + + monkeypatch.setattr(benchmark_serve, "main_async", fake_main_async) + args = argparse.Namespace(dataset_name="omniinteract", omniinteract_output_dir=tmp_path, extra_body=None) + with ThreadPoolExecutor(2) as executor: + first = executor.submit(benchmark_serve.main, args) + assert entered.wait(1) + second = executor.submit(benchmark_serve.main, args) + time.sleep(0.05) + assert calls == [0] + release.set() + assert first.result() == second.result() == {"total": 1} + assert calls == [0, 1] diff --git a/tests/benchmarks/test_seed_tts_dataset_variants.py b/tests/benchmarks/test_seed_tts_dataset_variants.py index fe1e6b62f..31f66bd7d 100644 --- a/tests/benchmarks/test_seed_tts_dataset_variants.py +++ b/tests/benchmarks/test_seed_tts_dataset_variants.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for SeedTTSTextDataset, SeedTTSTextSampleRequest, SeedTTSDesignDataset, and SeedTTSDesignSampleRequest. @@ -234,34 +237,6 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo assert req.seed_tts_utterance_id == "ok" -def test_attach_sets_seed_tts_row_even_without_extra_body(): - """seed_tts_row=True must be set for SeedTTSTextSampleRequest (no extra body).""" - from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSTextSampleRequest - - req = SeedTTSTextSampleRequest( - prompt="hello world", - prompt_len=2, - expected_output_len=100, - multi_modal_data=None, - request_id="test-0", - seed_tts_speech_extra=None, - seed_tts_ref_wav_path="", - ) - assert req.seed_tts_speech_extra is None - assert req.seed_tts_ref_wav_path == "" - # The fix ensures that even with speech_extra=None, the function - # sets seed_tts_row=True. We verify the source code has the fix. - import inspect - - import vllm_omni.benchmarks.patch.patch as patch_mod - - src = inspect.getsource(patch_mod._attach_seed_tts_to_request_func_input) - # seed_tts_row must be set BEFORE the 'if not ex: return' check - row_pos = src.index("seed_tts_row") - not_ex_pos = src.index("if not ex:") - assert row_pos < not_ex_pos, "seed_tts_row must be set before 'if not ex: return'" - - def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch): from vllm_omni.benchmarks.data_modules import seed_tts_eval diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py index 7dd12d000..5365d56ac 100644 --- a/tests/benchmarks/test_serve_cli.py +++ b/tests/benchmarks/test_serve_cli.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import argparse import json import os @@ -22,6 +25,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark] "dataset_name", [ "daily-omni", + "omniinteract", "seed-tts", "seed-tts-text", "seed-tts-design", @@ -55,6 +59,7 @@ def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) - [ (["--print-stage"], "print_stage", True), (["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"), + (["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]), (["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"), ], ) @@ -80,12 +85,14 @@ def test_add_omni_args_preserves_implicit_defaults() -> None: args = parser.parse_args([]) assert args.print_stage is False assert args.daily_omni_input_mode == "all" + assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"] assert args.seed_tts_locale == "en" assert args.explicit_keys == set() def test_update_omni_help_updates_upstream_actions() -> None: parser = TrackingArgumentParser() + parser.add_argument("--num-prompts", default=1000, help="Number of prompts.") parser.add_argument("--percentile-metrics", help="Upstream percentile help.") parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit help.") parser.add_argument("--random-mm-bucket-config", help="Upstream bucket help.") @@ -93,6 +100,7 @@ def test_update_omni_help_updates_upstream_actions() -> None: update_omni_help(parser) actions = {action.dest: action for action in parser._actions} + assert "OmniInteract uses 3" in actions["num_prompts"].help assert all(metric in actions["percentile_metrics"].help for metric in ("ttfc", "tpop", "audio_rtf")) assert "probabilities are renormalized" in actions["random_mm_limit_mm_per_prompt"].help assert "Currently allows for 3 modalities" in actions["random_mm_bucket_config"].help @@ -120,6 +128,31 @@ def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body( assert args.extra_body == expected +@pytest.mark.parametrize( + ("argv", "expected"), + [ + ([], 3), + (["--num-prompts", "0"], 0), + (["--num-prompts", "12"], 12), + ], +) +def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default( + argv: list[str], + expected: int, +) -> None: + parser = TrackingArgumentParser() + parser.add_argument("--dataset-name", default="omniinteract") + parser.add_argument("--backend", default="openai-realtime-duplex") + parser.add_argument("--endpoint", default="/v1/realtime") + parser.add_argument("--omniinteract-ref-audio", default="reference.wav") + parser.add_argument("--num-prompts", type=int, default=1000) + args = parser.parse_args(argv) + + preprocess_serve_args(args) + + assert args.num_prompts == expected + + @pytest.mark.parametrize( ("argv", "expected_extra_body", "expected_explicit"), [ diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 24a2a731e..c994f8dd2 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import sys from pathlib import Path import pytest +import yaml sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "common" / "scripts")) @@ -38,8 +40,6 @@ def _render(changed_files: list[str]) -> str: def test_bootstrap_if_injected_by_step_key() -> None: - import yaml - rendered = _render_bootstrap_pipeline( BOOTSTRAP_STEPS_TEMPLATE, decision=resolve_ci_decision([]), @@ -62,16 +62,23 @@ def test_bootstrap_steps_loaded_from_file() -> None: def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None: - """skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 still gates scheduled CI.""" + """skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 / NON_CRITICAL=1 still gates scheduled CI.""" rendered = _render(["docs/foo.md"]) assert "key: image-build" in rendered assert "key: upload-nightly-pipeline" in rendered assert "key: upload-weekly-pipeline" in rendered - # Scheduled nightly also uploads L2/L3 with --e2e - assert "key: upload-ready-pipeline" in rendered - assert "key: upload-merge-pipeline" in rendered - assert 'build.env("NIGHTLY") == "1"' in rendered + # Scheduled main+WEEKLY=1 uploads L2/L3 with --e2e; NIGHTLY still gates L4 only. + doc = yaml.safe_load(rendered) + by_key = {step["key"]: step for step in doc["steps"]} + assert "NIGHTLY" not in by_key["upload-ready-pipeline"]["if"] + assert 'build.branch == "main"' in by_key["upload-ready-pipeline"]["if"] + assert 'build.env("WEEKLY") == "1"' in by_key["upload-ready-pipeline"]["if"] + assert "NIGHTLY" not in by_key["upload-merge-pipeline"]["if"] + assert 'build.branch == "main"' in by_key["upload-merge-pipeline"]["if"] + assert 'build.env("WEEKLY") == "1"' in by_key["upload-merge-pipeline"]["if"] + assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"] assert 'build.env("WEEKLY") == "1"' in rendered + assert 'build.env("NON_CRITICAL") == "1"' in rendered assert "nightly-test" not in rendered assert "weekly-test" not in rendered assert "merge-test" not in rendered @@ -79,6 +86,19 @@ def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None: assert "if: false" not in rendered +def test_npu_docs_only_does_not_upload_ready_on_nightly() -> None: + """NPU skip_all: scheduled NIGHTLY still uploads L4, not L2 ready.""" + rendered = _render_bootstrap_pipeline( + BOOTSTRAP_STEPS_TEMPLATE, + decision=resolve_ci_decision(["docs/foo.md"]), + path=Path(".buildkite/npu/bootstrap-upload-steps.yml"), + ) + doc = yaml.safe_load(rendered) + by_key = {step["key"]: step for step in doc["steps"]} + assert "upload-ready-pipeline" not in by_key + assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"] + + def test_yaml_gated_l45_only_does_not_unconditionally_build_image() -> None: rendered = _render([".buildkite/cuda/test-nightly.yml"]) assert "if: true" not in rendered @@ -139,3 +159,65 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None: assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [ {"name": "swr-secret"}, ] + + +def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]: + rendered = _render_test_pipeline(doc, changed_files=changed_files) + labels: set[str] = set() + + def walk(steps: list | None) -> None: + for step in steps or []: + if not isinstance(step, dict): + continue + if "label" in step: + labels.add(step["label"]) + walk(step.get("steps")) + + walk(rendered.get("steps")) + return labels + + +# Synthetic coverage-style job: shared inputs that change what the split measures. +_COVERAGE_SHARED_INPUTS_DOC = { + "steps": [ + { + "label": "Coverage Pilot", + "source_file_dependencies": [ + "tests/e2e/online_serving/test_example.py", + ".buildkite/common/scripts/run_cov_split.sh", + "pyproject.toml", + ], + "commands": [".buildkite/common/scripts/run_cov_split.sh --model-id example"], + }, + { + "label": "Unrelated Model Test", + "source_file_dependencies": [ + "tests/e2e/online_serving/test_other.py", + ], + "commands": ["pytest -sv tests/e2e/online_serving/test_other.py"], + }, + ], +} + + +@pytest.mark.parametrize( + "changed_file", + [ + ".buildkite/common/scripts/run_cov_split.sh", + "pyproject.toml", + ], +) +def test_coverage_shared_inputs_select_dependent_job(changed_file: str) -> None: + """Jobs that list coverage shared inputs must stay selected when those files change.""" + labels = _surviving_labels(_COVERAGE_SHARED_INPUTS_DOC, [changed_file]) + assert "Coverage Pilot" in labels + assert "Unrelated Model Test" not in labels + + +def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None: + labels = _surviving_labels( + _COVERAGE_SHARED_INPUTS_DOC, + ["vllm_omni/entrypoints/openai/serving_chat.py"], + ) + assert "Coverage Pilot" not in labels + assert "Unrelated Model Test" not in labels diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index c3c3edbd1..e3c6c585f 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for StageConfigFactory and related classes. """ import importlib import warnings -from dataclasses import dataclass, fields +from dataclasses import dataclass from pathlib import Path from unittest.mock import patch @@ -14,7 +14,8 @@ import pytest from transformers import PretrainedConfig, Qwen3OmniMoeConfig from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage -from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config import config_factory as config_factory_module +from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config @@ -33,13 +34,18 @@ from vllm_omni.config.stage_config import ( build_stage_runtime_overrides, load_deploy_config, merge_pipeline_deploy, + normalize_pipeline_cli_overrides, pipeline_cfg_resolver, - strip_parent_engine_args, ) -from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys +from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +SINGLE_STAGE_PIPE_CFG = (StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),) +# Validation strings to match against in post init +NO_STAGES_MATCH_STR = "no stages" +NO_TERMINAL_STAGE_MATCH_STR = "No terminal stage" + @pytest.fixture(autouse=True) def _stable_test_platform(monkeypatch): @@ -56,6 +62,7 @@ def clear_config_factory_caches(): yield StageConfigFactory.get_hf_config.cache_clear() StageConfigFactory.try_infer_model_type.cache_clear() + _materialize_object_storage_configs.cache_clear() Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True) @@ -428,6 +435,20 @@ class TestStageConfigFactory: assert cfg["final_output"] is True assert cfg["final_output_type"] == "image" + @pytest.mark.parametrize( + "model_class_name", + [ + "MiniMaxH3Pipeline", + "WanPipeline", + "LTX2Pipeline", + "Cosmos3OmniDiffusersPipeline", + ], + ) + def test_default_video_diffusion_output_type(self, model_class_name): + configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name}) + + assert configs[0]["final_output_type"] == "video" + def test_default_diffusion_with_parallel_config(self): """Test diffusion config calculates devices from parallel_config.""" @@ -475,7 +496,7 @@ class TestStageConfigFactory: cli_overrides = { "gpu_memory_utilization": 0.9, "model": "some_model", # Internal - "stage_configs_path": "/path", # Internal + "deploy_config": "/path", # Internal "batch_timeout": 10, # Internal } @@ -483,7 +504,7 @@ class TestStageConfigFactory: assert overrides["gpu_memory_utilization"] == 0.9 assert "model" not in overrides - assert "stage_configs_path" not in overrides + assert "deploy_config" not in overrides assert "batch_timeout" not in overrides def test_per_stage_override_excludes_internal_keys(self): @@ -525,49 +546,6 @@ class TestStageResolutionHelpers: assert "model" not in overrides assert "parallel_config" not in overrides - def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self): - parent_fields = {f.name: f for f in fields(EngineArgs)} - filtered, overridden = strip_parent_engine_args( - { - "model": "some/model", - "stage_configs_path": "/tmp/stages.yaml", - "tensor_parallel_size": 4, - "worker_extension_cls": "some.Extension", - "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"}, - }, - parent_fields=parent_fields, - keep_keys={"worker_extension_cls"}, - strip_keys={"stage_configs_path"}, - no_warn_keys={"model"}, - ) - - assert filtered == { - "worker_extension_cls": "some.Extension", - "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"}, - } - assert overridden == ["tensor_parallel_size"] - - def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self): - parent_fields = {f.name: f for f in fields(EngineArgs)} - filtered, overridden = strip_parent_engine_args( - { - "model": "some/model", - "stage_configs_path": "/tmp/stages.yaml", - "allowed_local_media_path": "/data/qwentts", - "allowed_media_domains": ["example.com"], - }, - parent_fields=parent_fields, - keep_keys={"allowed_local_media_path", "allowed_media_domains"}, - strip_keys={"stage_configs_path"}, - no_warn_keys={"model"}, - ) - - assert filtered == { - "allowed_local_media_path": "/data/qwentts", - "allowed_media_domains": ["example.com"], - } - assert overridden == [] - class TestPipelineDiscovery: """Tests for the central pipeline registry (``OMNI_PIPELINES``).""" @@ -576,6 +554,7 @@ class TestPipelineDiscovery: """Check that specific models are in OMNI_PIPELINES.""" assert "qwen2_5_omni" in OMNI_PIPELINES assert "qwen3_omni_moe" in OMNI_PIPELINES + assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES assert "qwen3_tts" in OMNI_PIPELINES def test_registry_resolver_qwen3_omni_all_stages(self): @@ -600,6 +579,18 @@ class TestPipelineDiscovery: assert pipeline.default_deploy_config_name is None assert len(pipeline.stages) == 1 # thinker only + def test_registry_qwen3_omni_thinker_only_key_is_static(self): + """Explicit registry key must not re-enter the HF-config resolver.""" + pipeline = resolve_pipeline_config( + "qwen3_omni_moe_thinker_only", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"] + assert pipeline.model_type == "qwen3_omni_moe_thinker_only" + assert len(pipeline.stages) == 1 + assert pipeline.stages[0].engine_output_type == "text" + assert pipeline.stages[0].final_output_type == "text" + @pytest.mark.parametrize( "pipeline", [registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)], @@ -632,6 +623,7 @@ class TestPipelineDiscovery: p = PipelineConfig( model_type="custom_collide", hf_architectures=("SomeCollidingArch",), + stages=SINGLE_STAGE_PIPE_CFG, ) assert p.hf_architectures == ("SomeCollidingArch",) @@ -653,33 +645,63 @@ class TestStagePipelineConfig: class TestPipelineConfigNew: - def test_frozen(self): - p = PipelineConfig(model_type="t", model_arch="A") - with pytest.raises(AttributeError): - p.model_type = "changed" + def test_simple_init(self): + # Ensure pipeline config validates in post init. + PipelineConfig( + model_type="t", + model_arch="A", + stages=SINGLE_STAGE_PIPE_CFG, + ) - def test_validate_valid(self): + def test_frozen(self): p = PipelineConfig( + model_type="t", + model_arch="A", + stages=SINGLE_STAGE_PIPE_CFG, + ) + with pytest.raises(AttributeError): + p.model_type = "changed" + + def test_validate_no_stages(self): + with pytest.raises(ValueError, match=NO_STAGES_MATCH_STR): + PipelineConfig(model_type="t", model_arch="A") + + def test_validate_no_terminal_stage(self): + """A pipeline with no ``final_output`` stage can never emit a result.""" + with pytest.raises(ValueError, match=NO_TERMINAL_STAGE_MATCH_STR): + PipelineConfig( + model_type="t", + model_arch="A", + stages=( + StagePipelineConfig(stage_id=0, model_stage="a"), + StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)), + ), + ) + + def test_pipeline_can_have_final_output_in_any_stage(self): + """Ensure any stage may carry ``final_output``.""" + PipelineConfig( model_type="t", model_arch="A", stages=( - StagePipelineConfig(stage_id=0, model_stage="a"), + StagePipelineConfig(stage_id=0, model_stage="a", final_output=True), StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)), ), ) - assert p.validate() == [] - - def test_validate_no_stages(self): - p = PipelineConfig(model_type="t", model_arch="A") - assert any("no stages" in e.lower() for e in p.validate()) class TestPipelineRegistration: def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path): deploy_key = "deploy_selected_pipeline" model_type_key = "hf_model_type_pipeline" - deploy_pipe = PipelineConfig(model_type=deploy_key) - model_type_pipe = PipelineConfig(model_type=model_type_key) + deploy_pipe = PipelineConfig( + model_type=deploy_key, + stages=SINGLE_STAGE_PIPE_CFG, + ) + model_type_pipe = PipelineConfig( + model_type=model_type_key, + stages=SINGLE_STAGE_PIPE_CFG, + ) register_pipeline(deploy_pipe) register_pipeline(model_type_pipe) @@ -705,7 +727,11 @@ class TestPipelineRegistration: tmp_path, ): model_type_key = "registered_model_type_pipeline" - register_pipeline(PipelineConfig(model_type=model_type_key)) + pipeline_cfg = PipelineConfig( + model_type=model_type_key, + stages=SINGLE_STAGE_PIPE_CFG, + ) + register_pipeline(pipeline_cfg) deploy_path = tmp_path / "deploy.yaml" deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8") @@ -740,6 +766,7 @@ class TestPipelineRegistration: pipe_cfg = PipelineConfig( model_type=pipeline_key, hf_architectures=("ArchitectureFallbackForTest",), + stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipe_cfg) @@ -770,15 +797,18 @@ class TestPipelineRegistration: model_type="rejected_by_predicate", hf_architectures=(shared_arch,), hf_config_predicate=rejecting_predicate, + stages=SINGLE_STAGE_PIPE_CFG, ) raises_cfg = PipelineConfig( model_type="raises_in_predicate", hf_architectures=(shared_arch,), hf_config_predicate=raising_predicate, + stages=SINGLE_STAGE_PIPE_CFG, ) accept_cfg = PipelineConfig( model_type="accepted_by_predicate", hf_architectures=(shared_arch,), + stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(reject_cfg) register_pipeline(raises_cfg) @@ -812,6 +842,7 @@ class TestPipelineRegistration: model_type="predicate_without_arch_match", hf_architectures=("DifferentArchitectureForTest",), hf_config_predicate=predicate, + stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipe_cfg) @@ -835,6 +866,7 @@ class TestPipelineRegistration: resolved_cfg = PipelineConfig( model_type="callable_resolved_pipeline", hf_architectures=("CallableArchitectureForTest",), + stages=SINGLE_STAGE_PIPE_CFG, ) seen_hf_configs = [] @@ -908,7 +940,10 @@ class TestPipelineRegistration: def test_pipeline_registration(self, clean_pipeline_registry): """Ensure that we can register and create a custom pipeline config.""" new_model_type = "new_model_type" - pipe_cfg = PipelineConfig(model_type=new_model_type) + pipe_cfg = PipelineConfig( + model_type=new_model_type, + stages=SINGLE_STAGE_PIPE_CFG, + ) # Register the new PipelineConfig assert new_model_type not in OMNI_PIPELINES @@ -951,7 +986,10 @@ class TestPipelineRegistration: def custom_resolver( hf_config: FakeConfig, ) -> PipelineConfig: - return PipelineConfig(model_type=resolved_type) + return PipelineConfig( + model_type=resolved_type, + stages=SINGLE_STAGE_PIPE_CFG, + ) # Register the new PipelineConfig assert new_model_type not in OMNI_PIPELINES @@ -1008,7 +1046,7 @@ class TestPipelineRegistration: pipe_cfg = PipelineConfig( model_type=pipeline_key, model_arch="DeployOnlyArch", - stages=(StagePipelineConfig(stage_id=0, model_stage="diffusion"),), + stages=SINGLE_STAGE_PIPE_CFG, ) register_pipeline(pipe_cfg) @@ -1041,7 +1079,11 @@ class TestPipelineRegistration: def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path): pipeline_key = "single_load_pipeline" - register_pipeline(PipelineConfig(model_type=pipeline_key)) + pipeline_cfg = PipelineConfig( + model_type=pipeline_key, + stages=SINGLE_STAGE_PIPE_CFG, + ) + register_pipeline(pipeline_cfg) deploy_path = tmp_path / "single_load.yaml" deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8") @@ -1071,6 +1113,7 @@ class TestPipelineRegistration: pipeline = PipelineConfig( model_type="pipeline_with_default", default_deploy_config_name=default_name, + stages=SINGLE_STAGE_PIPE_CFG, ) with patch( @@ -1090,8 +1133,16 @@ class TestPipelineRegistration: OmniServingCapability.COMPLETIONS, "pipeline_a blocks completions", ) - pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,)) - pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=()) + pipe_a = PipelineConfig( + model_type="detect_type", + endpoint_restrictions=(restriction,), + stages=SINGLE_STAGE_PIPE_CFG, + ) + pipe_b = PipelineConfig( + model_type="override_type", + endpoint_restrictions=(), + stages=SINGLE_STAGE_PIPE_CFG, + ) register_pipeline(pipe_a) register_pipeline(pipe_b) @@ -1144,8 +1195,42 @@ class TestResolveScheduler: class TestDeployConfigLoading: - def test_load_minicpmo_duplex_deploy_config(self): - deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml")) + def test_rejects_legacy_stage_args_schema(self, tmp_path): + deploy_path = tmp_path / "legacy.yaml" + deploy_path.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8") + + with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"): + load_deploy_config(deploy_path) + + @pytest.mark.parametrize( + ("filename", "max_sessions"), + [ + ("minicpmo_4_5.yaml", 4), + ("minicpmo_4_5_2gpu.yaml", 4), + ("minicpmo_4_5_3gpu.yaml", 4), + ("minicpmo_4_5_8x4090.yaml", 1), + ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4), + ("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4), + ("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1), + ], + ) + def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int): + """Every MiniCPM-o profile serves /v1/realtime, so the retired + minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on. + """ + deploy = load_deploy_config(Path(get_deploy_config_path(filename))) + pipeline = resolve_pipeline_config("minicpmo_4_5") + assert isinstance(pipeline, PipelineConfig) + stages = merge_pipeline_deploy(pipeline, deploy) + + assert deploy.session_mode == "duplex" + assert deploy.active_stream_window == max_sessions + assert deploy.duplex_session.max_sessions == max_sessions + assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] + assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] + + def test_load_minicpmo_default_deploy_config(self): + deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml")) deploy = load_deploy_config(deploy_path) pipeline = resolve_pipeline_config("minicpmo_4_5") @@ -1153,18 +1238,18 @@ class TestDeployConfigLoading: stages = merge_pipeline_deploy(pipeline, deploy) - assert deploy.session_mode == "duplex" assert deploy.async_chunk is True - assert deploy.active_stream_window == 1 - assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] - assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] - assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False] - assert all("Async" not in (stage.scheduler_cls or "") for stage in stages) + assert stages[0].yaml_engine_args["async_scheduling"] is True + assert stages[1].yaml_engine_args["async_scheduling"] is True + assert "Async" in (stages[0].scheduler_cls or "") + assert "Async" in (stages[1].scheduler_cls or "") assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"] assert deploy.stages[1].enforce_eager is False - assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048 - assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0 - assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1] + assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096 + assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50 + assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8 + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] + assert "codec_sampling_params" not in stages[1].yaml_engine_args @pytest.mark.parametrize( ("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"), @@ -1233,14 +1318,22 @@ stages: assert deploy.connectors is not None assert deploy.platforms is not None - def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self): + @pytest.mark.parametrize( + ("hf_config", "model"), + [ + (Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"), + (Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"), + ], + ids=["thinker_checkpoint", "instruct_forced_thinker_only"], + ) + def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model): deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml")) with patch( "vllm_omni.config.config_factory.get_config", - return_value=Q3_OMNI_THINKER_HF_CONFIG, + return_value=hf_config, ): pipeline = StageConfigFactory.get_pipeline_config( - "Qwen/Qwen3-Omni-30B-A3B-Thinking", + model, trust_remote_code=False, deploy_config_path=str(deploy_path), ) @@ -1249,6 +1342,7 @@ stages: assert pipeline.model_type == "qwen3_omni_moe_thinker_only" deploy = load_deploy_config(deploy_path) + assert deploy.pipeline == "qwen3_omni_moe_thinker_only" stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 @@ -1310,6 +1404,13 @@ stages: ("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"), ("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None), ("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None), + ( + "minimax_h3_disaggregated.yaml", + "minimax_h3_disaggregated", + 2, + "video", + "minimax_h3_disaggregated", + ), ], ) def test_load_new_registry_backed_deploy_configs( @@ -1330,15 +1431,98 @@ stages: assert len(stages) == stage_count assert stages[-1].final_output is True assert stages[-1].final_output_type == final_output_type + if pipeline_name == "minimax_h3_disaggregated": + assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root") + assert ( + stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path") + ) + assert "model_path_resolver" not in stages[0].yaml_extras + assert "model_path_resolver" not in stages[1].yaml_extras if deploy.trust_remote_code is not None: assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code} + def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch): + from vllm_omni.engine import stage_init_utils + + model_root = tmp_path / "MiniMax-H3" + (model_root / "FL2VA" / "text_encoder").mkdir(parents=True) + (model_root / "Ref2VA" / "text_encoder").mkdir(parents=True) + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None) + + deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml")) + assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False} + stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy) + resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages] + + assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder") + assert resolved[1]["model"] == str(model_root / "FL2VA") + assert all("model_path_resolver" not in engine_args for engine_args in resolved) + + def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self): + pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] + stages, _ = StageConfigFactory._create_legacy_from_registry( + pipeline, + {"text_encoder_tp_size": 4}, + ) + + assert stages[0].runtime_overrides["tensor_parallel_size"] == 4 + assert "text_encoder_tp_size" not in stages[0].runtime_overrides + assert "text_encoder_tp_size" not in stages[1].runtime_overrides + assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1 + + def test_minimax_h3_disaggregated_defaults_match_validated_topology(self): + pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] + deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml"))) + stages = merge_pipeline_deploy(pipeline, deploy) + + assert stages[0].yaml_engine_args["max_num_seqs"] == 1 + assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root") + assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path") + parallel = stages[1].yaml_engine_args["parallel_config"] + assert parallel["tensor_parallel_size"] == 1 + assert parallel["ulysses_degree"] == 4 + assert parallel["vae_patch_parallel_size"] == 4 + assert stages[1].yaml_engine_args["inline_diffusion"] is True + assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50 + + turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml"))) + turbo_stages = merge_pipeline_deploy(pipeline, turbo) + turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"] + assert turbo_sampling["num_inference_steps"] == 5 + assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0} + + def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self): + pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] + + assert normalize_pipeline_cli_overrides(pipeline, {}) == {} + assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == { + "stage_0_tensor_parallel_size": 4 + } + + def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self): + pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] + + with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"): + normalized = normalize_pipeline_cli_overrides( + pipeline, + {"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2}, + ) + + assert normalized == {"stage_0_tensor_parallel_size": 2} + + def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self): + pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] + + with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"): + normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4}) + @pytest.mark.parametrize( ("config_json", "model_index", "expected_pipeline"), [ ({"model_type": "step_audio_2"}, None, "step_audio_2"), (None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"), (None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"), + (None, {"_class_name": "WanDMDPipeline"}, "wan2_2_ti2v"), ], ) def test_migrated_models_are_discovered_without_explicit_deploy( @@ -1367,6 +1551,29 @@ stages: assert pipeline is not None assert pipeline.model_type == expected_pipeline + def test_minimax_h3_disaggregation_is_not_auto_discovered(self): + def get_model_file(filename, _model, revision=None): + del revision + if filename == "config.json": + return None + if filename == "model_index.json": + return {"_class_name": "MiniMaxH3Pipeline"} + return None + + with ( + patch.object(StageConfigFactory, "get_hf_config", return_value=None), + patch( + "vllm_omni.config.config_factory.get_hf_file_to_dict", + side_effect=get_model_file, + ), + ): + pipeline = StageConfigFactory.get_pipeline_config( + model="/models/unrelated-checkpoint-name", + trust_remote_code=False, + ) + + assert pipeline is None + @pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"]) def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name): deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) @@ -1419,7 +1626,8 @@ stages: assert async_stages[1].custom_process_input_func is None def test_no_bundled_legacy_stage_config_yamls(self): - stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs" + repo_root = Path(__file__).resolve().parents[2] + stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs" assert not list(stage_config_dir.glob("*.yaml")) def test_merge_pipeline_deploy(self): @@ -1474,6 +1682,7 @@ stages: model_stage="ar", execution_type=StageExecutionType.LLM_AR, requires_multimodal_data=True, + final_output=True, ), ), ) @@ -1494,6 +1703,7 @@ stages: model_stage="ar", execution_type=StageExecutionType.LLM_AR, scheduler_cls=scheduler_cls, + final_output=True, ), ), ) @@ -1693,6 +1903,7 @@ stages: ] } original_foo = fake_config["stages"][0]["foo"] + assert isinstance(original_foo, dict) original_b = original_foo["b"] with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): @@ -1738,7 +1949,6 @@ class TestQwen3OmniPipeline: assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration" assert len(p.stages) == 3 - assert p.validate() == [] def test_thinker(self): p = resolve_pipeline_config( @@ -1792,7 +2002,6 @@ class TestQwen2_5OmniPipeline: assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen2_5OmniForConditionalGeneration" assert len(p.stages) == 3 - assert p.validate() == [] def test_thinker(self): p = resolve_pipeline_config("qwen2_5_omni") @@ -1841,7 +2050,6 @@ class TestQwen3TTSPipeline: assert p is not None assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration" assert len(p.stages) == 2 - assert p.validate() == [] def test_talker_stage(self): p = resolve_pipeline_config("qwen3_tts") @@ -1921,7 +2129,6 @@ class TestMingFlashOmniPipeline: assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 - assert p.validate() == [] def test_thinker_stage(self): p = resolve_pipeline_config("ming_flash_omni") @@ -1978,7 +2185,6 @@ class TestMingFlashOmniPipeline: assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration" assert len(p.stages) == 1 - assert p.validate() == [] def test_tts_stage(self): p = resolve_pipeline_config("ming_flash_omni_tts") @@ -2035,7 +2241,6 @@ class TestMingFlashOmniPipeline: assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 1 - assert p.validate() == [] def test_thinker_only_stage(self): p = resolve_pipeline_config("ming_flash_omni_thinker_only") @@ -2074,7 +2279,6 @@ class TestMingFlashOmniPipeline: assert p is not None assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 - assert p.validate() == [] def test_image_thinker_stage(self): s = resolve_pipeline_config("ming_flash_omni_image").get_stage(0) @@ -2145,6 +2349,26 @@ class TestMingFlashOmniPipeline: class TestBaseConfigInheritance: """Test deploy YAML base_config inheritance.""" + def test_minicpmo_overlays_inherit_talker_sampling_params(self): + """Overlays must keep the base Talker codec Sampler knobs.""" + for filename in ( + "minicpmo_4_5_2gpu.yaml", + "minicpmo_4_5_3gpu.yaml", + "minicpmo_4_5_3gpu_stage1_replicas.yaml", + "minicpmo_4_5_4gpu_stage1_replicas.yaml", + "minicpmo_4_5_8x4090_stage1_replicas.yaml", + ): + path = Path(get_deploy_config_path(filename)) + if not path.exists(): + pytest.skip(f"{filename} not found") + deploy = load_deploy_config(path) + sampling = deploy.stages[1].default_sampling_params + assert sampling is not None, f"{filename} stage 1 lost default_sampling_params" + assert sampling["temperature"] == 0.8, filename + assert sampling["top_k"] == 25, filename + assert sampling["top_p"] == 0.85, filename + assert sampling["repetition_penalty"] == 1.05, filename + def test_ci_inherits_from_main(self): ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml")) if not ci_path.exists(): @@ -2235,6 +2459,42 @@ class TestPlatformOverrides: assert rocm.stages[0].enforce_eager is None assert rocm.stages[1].enforce_eager is True + def test_higgs_audio_v3_rocm_uses_triton_attention(self): + deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml")) + + base = load_deploy_config(deploy_path) + assert base.stages[0].engine_extras["attention_backend"] == "FLASHINFER" + + rocm = _apply_platform_overrides(base, platform="rocm") + assert rocm.stages[0].engine_extras["attention_backend"] == "TRITON_ATTN" + + pipeline = resolve_pipeline_config("higgs_multimodal_qwen3") + assert isinstance(pipeline, PipelineConfig) + stages = merge_pipeline_deploy(pipeline, rocm) + assert stages[0].yaml_engine_args["attention_backend"] == "TRITON_ATTN" + + def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self): + deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) + + cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda") + cuda_stages = merge_pipeline_deploy(pipeline, cuda) + expected = {"custom_ops": ["+rotary_embedding"]} + assert cuda_stages[0].yaml_engine_args["compilation_config"] == expected + assert all("compilation_config" not in stage.yaml_engine_args for stage in cuda_stages[1:]) + + for platform in ("cpu", "musa", "npu", "rocm", "xpu"): + deploy = _apply_platform_overrides( + load_deploy_config(deploy_path), + platform=platform, + ) + config = deploy.stages[0].compilation_config or {} + assert "+rotary_embedding" not in config.get("custom_ops", []) + def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self): pipeline = resolve_pipeline_config("minicpmo_4_5") assert isinstance(pipeline, PipelineConfig) @@ -2741,22 +3001,23 @@ class TestSentinelDefaultPrecedence: ) def test_ming_flash_omni_topology(self): - """Guard ming_flash_omni's SIP cleanup: stage 0 has no full-payload - producer hook (arch is not in ``_FULL_PAYLOAD_INPUT_STAGES``), and - stage 1 uses only ``thinker2talker_token_only`` (sync_process_input_func). + """Guard ming_flash_omni's SIP cleanup: stage 0 has no + ``custom_process_next_stage_input_func``, and stage 1 does not set + ``requires_full_payload_input=True``. Stage 1 uses only + ``thinker2talker_token_only`` (sync_process_input_func). The dead ``thinker2talker`` (custom_process_input_func) was removed because ming_flash_omni has no async_chunk support and both functions called ``_build_talker_inputs`` identically. - Merge under either async_chunk mode must not re-introduce a - stage-0 full-payload hook.""" + Merge under either async_chunk mode must not re-introduce either + full-payload transport gate.""" pipeline = resolve_pipeline_config("ming_flash_omni") assert isinstance(pipeline, PipelineConfig) stage0, stage1 = pipeline.stages assert stage0.custom_process_next_stage_input_func is None, ( - "ming_flash_omni stage 0 must not declare a full-payload producer " - "(connector path is not active for this arch)." + "ming_flash_omni stage 0 must not declare a full-payload producer hook." ) + assert stage1.requires_full_payload_input is False assert stage1.custom_process_input_func is None assert stage1.sync_process_input_func is not None assert stage1.sync_process_input_func.endswith("thinker2talker_token_only") @@ -2809,3 +3070,90 @@ class TestPipelineConfigResolvers: pass assert resolver(NotTheRightHfConfig()) is None + + +class TestObjectStorageConfigResolution: + """Object-storage URIs must be materialized locally before any HF-style reads. + + Regression coverage for review on the Run:AI PR: ``get_config`` and the + config.json/model_index.json fallbacks crashed with ``HFValidationError`` + for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so + a bucket named after another pipeline could hijack pipeline selection. + """ + + @pytest.fixture + def fake_object_storage(self, monkeypatch, tmp_path): + """Replace ObjectStorageModel with a fake that "pulls" into a tmp dir. + + Returns a recorder: calls to the returned callable write files into + the materialized directory, and ``recorder.pulls`` records every + ``pull_files`` invocation. + """ + pulls: list[tuple[str, list[str] | None, list[str] | None]] = [] + materialized_files: list[tuple[str, str]] = [] + + class FakeObjectStorageModel: + def __init__(self, url: str): + self.dir = str(tmp_path) + + def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None): + pulls.append((model_path, allow_pattern, ignore_pattern)) + for name, content in materialized_files: + (tmp_path / name).write_text(content) + + monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel) + + class Recorder: + def __init__(self) -> None: + self.pulls = pulls + + def set_files(self, files: list[tuple[str, str]]) -> None: + materialized_files.clear() + materialized_files.extend(files) + + return Recorder() + + def test_passthrough_for_non_uri(self, fake_object_storage): + assert _materialize_object_storage_configs("org/model") == "org/model" + assert _materialize_object_storage_configs("/local/model") == "/local/model" + assert fake_object_storage.pulls == [] + + def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage): + uri = "s3://bucket/model" + + first = _materialize_object_storage_configs(uri) + second = _materialize_object_storage_configs(uri) + + assert first == second + assert len(fake_object_storage.pulls) == 1 + pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0] + assert pulled_uri == uri + assert "*.json" in allow_pattern + assert ignore_pattern is None + + def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage): + """A misleading bucket name must not beat the real config.json content.""" + fake_object_storage.set_files( + [ + ( + "config.json", + '{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}', + ) + ] + ) + uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct" + + model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False) + + assert model == "qwen3_omni_moe" + assert fake_object_storage.pulls # materialization happened + + def test_name_match_fallback_scans_basename_only(self, fake_object_storage): + # Empty config.json (CosyVoice3 style) forces the name-match fallback. + fake_object_storage.set_files([("config.json", "{}")]) + + deceptive = "s3://qwen3-tts-models/plain-checkpoint" + assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None + + matching = "s3://any-bucket/my-cosyvoice3-model" + assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3" diff --git a/tests/config/test_config_import_cycle.py b/tests/config/test_config_import_cycle.py new file mode 100644 index 000000000..e149b63ee --- /dev/null +++ b/tests/config/test_config_import_cycle.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Guard the NPU pytest circular import through DiffusionOutput.""" + +from __future__ import annotations + +import ast +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _top_level_import_modules(path: Path) -> set[str]: + tree = ast.parse(path.read_text(encoding="utf-8")) + modules: set[str] = set() + for node in tree.body: + if isinstance(node, ast.ImportFrom) and node.module: + modules.add(node.module) + elif isinstance(node, ast.Import): + modules.update(alias.name for alias in node.names) + return modules + + +def _assert_isolated_import_succeeds(script: str) -> None: + env = os.environ.copy() + env.pop("CUDA_VISIBLE_DEVICES", None) + env.pop("HIP_VISIBLE_DEVICES", None) + result = subprocess.run( + [sys.executable, "-c", script], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=120, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_config_package_does_not_eagerly_import_pipeline_registry() -> None: + imported = _top_level_import_modules(REPO_ROOT / "vllm_omni" / "config" / "__init__.py") + assert "vllm_omni.config.config_factory" not in imported + assert "vllm_omni.config.pipeline_registry" not in imported + + +def test_config_package_lazy_exports_still_resolve() -> None: + from vllm_omni.config import StageConfigFactory, register_pipeline + from vllm_omni.config.config_factory import StageConfigFactory as DirectStageConfigFactory + from vllm_omni.config.pipeline_registry import register_pipeline as direct_register_pipeline + + assert StageConfigFactory is DirectStageConfigFactory + assert register_pipeline is direct_register_pipeline + + +def test_npu_platform_defers_minimax_h3_encoder_patch() -> None: + src = (REPO_ROOT / "vllm_omni" / "platforms" / "npu" / "platform.py").read_text(encoding="utf-8") + init_start = src.index("def __init__(self) -> None:") + runtime_start = src.index("def init_diffusion_model_runner_runtime") + init_src = src[init_start:runtime_start] + runtime_src = src[runtime_start:] + assert "apply_minimax_h3_qwen3vl_patch" not in init_src + assert "apply_minimax_h3_qwen3vl_patch" in runtime_src + assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src + assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src + assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src + + +def test_lora_config_import_does_not_load_pipeline_registry() -> None: + """Import LoRAConfig and diffusion.data in a fresh interpreter. + + The NPU failure was ``data.py`` importing LoRAConfig while still loading, + then ``vllm_omni.config`` re-entering ``pipeline_registry`` → PI0 → + DiffusionOutput. This child process must not preload those exports. + """ + _assert_isolated_import_succeeds( + """ +import sys + +from vllm_omni.config.lora import LoRAConfig + +if "vllm_omni.config.pipeline_registry" in sys.modules: + raise SystemExit("LoRAConfig import loaded pipeline_registry") +if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules: + raise SystemExit("LoRAConfig import loaded PI0_PIPELINE") + +from vllm_omni.diffusion.data import DiffusionOutput + +if DiffusionOutput is None or LoRAConfig is None: + raise SystemExit("expected LoRAConfig and DiffusionOutput") +""" + ) + + +def test_pipeline_registry_does_not_import_pi0_runtime() -> None: + """Resolving the pi0 topology must not import its runtime pipeline. + + ``pipeline_pi0`` imports ``diffusion.data``. Importing it while + ``diffusion.data`` is only partially initialized closes the Ascend plugin's + circular-import loop, so the registry must depend only on a lightweight + pipeline-config module. + """ + _assert_isolated_import_succeeds( + """ +import sys + +from vllm_omni.config.pipeline_registry import resolve_pipeline_config + +if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules: + raise SystemExit("pipeline_registry imported the pi0 runtime pipeline") + +pipeline = resolve_pipeline_config("pi0") +if pipeline is None or pipeline.model_type != "pi0": + raise SystemExit("expected the registered pi0 pipeline config") +""" + ) diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py new file mode 100644 index 000000000..00e0dffdb --- /dev/null +++ b/tests/config/test_environment_variables.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Drift checks for the reviewed environment-variable inventory.""" + +import ast +import builtins +from collections import Counter +from pathlib import Path + +import pytest + +from vllm_omni.config.environment_variable_inventory import ( + ENVIRONMENT_VARIABLE_INVENTORY, + EnvironmentVariableCategory, + ModelEnvironmentVariableDisposition, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_REPO_ROOT = Path(__file__).parents[2] +_PACKAGE_ROOT = _REPO_ROOT / "vllm_omni" +_INVENTORY_MODULE = _PACKAGE_ROOT / "config" / "environment_variable_inventory.py" +_REFERENCE_PAGE = _REPO_ROOT / "docs" / "configuration" / "environment_variables.md" + + +def _module_string_constants(tree: ast.Module) -> dict[str, str]: + constants: dict[str, str] = {} + for statement in tree.body: + targets: list[ast.expr] = [] + value: ast.expr | None = None + if isinstance(statement, ast.Assign): + targets = statement.targets + value = statement.value + elif isinstance(statement, ast.AnnAssign): + targets = [statement.target] + value = statement.value + if isinstance(value, ast.Constant) and isinstance(value.value, str): + for target in targets: + if isinstance(target, ast.Name): + constants[target.id] = value.value + return constants + + +def _os_import_aliases(tree: ast.Module) -> tuple[set[str], set[str], set[str]]: + os_names: set[str] = set() + getenv_names: set[str] = set() + environ_names: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for imported in node.names: + if imported.name == "os": + os_names.add(imported.asname or "os") + elif isinstance(node, ast.ImportFrom) and node.module == "os": + for imported in node.names: + if imported.name == "getenv": + getenv_names.add(imported.asname or imported.name) + elif imported.name == "environ": + environ_names.add(imported.asname or imported.name) + return os_names, getenv_names, environ_names + + +def _is_environ(node: ast.expr, os_names: set[str], environ_names: set[str]) -> bool: + return ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id in os_names + and node.attr == "environ" + ) or (isinstance(node, ast.Name) and node.id in environ_names) + + +def _environment_key_expressions( + node: ast.AST, + os_names: set[str], + getenv_names: set[str], + environ_names: set[str], +) -> list[ast.expr]: + """Return key expressions used by one direct environment access.""" + if isinstance(node, ast.Call) and node.args: + func = node.func + if ( + isinstance(func, ast.Attribute) + and isinstance(func.value, ast.Name) + and func.value.id in os_names + and func.attr == "getenv" + ) or (isinstance(func, ast.Name) and func.id in getenv_names): + return [node.args[0]] + if ( + isinstance(func, ast.Attribute) + and _is_environ(func.value, os_names, environ_names) + and func.attr in {"get", "pop", "setdefault"} + ): + return [node.args[0]] + elif isinstance(node, ast.Subscript) and _is_environ(node.value, os_names, environ_names): + return [node.slice] + elif isinstance(node, ast.Compare): + expressions: list[ast.expr] = [] + left_values = [node.left, *node.comparators[:-1]] + for left, operator, right in zip(left_values, node.ops, node.comparators): + if isinstance(operator, (ast.In, ast.NotIn)) and _is_environ(right, os_names, environ_names): + expressions.append(left) + return expressions + return [] + + +def _resolve_environment_name(expression: ast.expr, constants: dict[str, str]) -> str | None: + value: str | None = None + if isinstance(expression, ast.Constant) and isinstance(expression.value, str): + value = expression.value + elif isinstance(expression, ast.Name): + value = constants.get(expression.id) + return value or None + + +def _environment_wrapper_parameters( + tree: ast.Module, + os_names: set[str], + getenv_names: set[str], + environ_names: set[str], +) -> dict[str, set[tuple[str, int | None]]]: + """Find local helpers that forward one parameter as an environment key.""" + wrappers: dict[str, set[tuple[str, int | None]]] = {} + for function in (node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))): + positional = [*function.args.posonlyargs, *function.args.args] + parameter_indexes: dict[str, int | None] = {parameter.arg: index for index, parameter in enumerate(positional)} + parameter_indexes.update({parameter.arg: None for parameter in function.args.kwonlyargs}) + forwarded: set[tuple[str, int | None]] = set() + for node in ast.walk(function): + for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names): + if isinstance(expression, ast.Name) and expression.id in parameter_indexes: + forwarded.add((expression.id, parameter_indexes[expression.id])) + if forwarded: + wrappers[function.name] = forwarded + return wrappers + + +def _environment_accesses(path: Path) -> set[str]: + """Return statically resolvable names accessed through ``os``.""" + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + constants = _module_string_constants(tree) + os_names, getenv_names, environ_names = _os_import_aliases(tree) + wrappers = _environment_wrapper_parameters(tree, os_names, getenv_names, environ_names) + names: set[str] = set() + + for node in ast.walk(tree): + for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names): + if name := _resolve_environment_name(expression, constants): + names.add(name) + + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + for parameter_name, position in wrappers.get(node.func.id, set()): + argument = node.args[position] if position is not None and position < len(node.args) else None + if argument is None: + argument = next( + (keyword.value for keyword in node.keywords if keyword.arg == parameter_name), + None, + ) + if argument is not None and (name := _resolve_environment_name(argument, constants)): + names.add(name) + + return names + + +def test_inventory_matches_reviewed_snapshot_counts(): + """Make an inventory expansion an explicit review decision.""" + category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) + assert category_counts == { + EnvironmentVariableCategory.PUBLIC_OMNI: 24, + EnvironmentVariableCategory.INHERITED_VLLM: 20, + EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, + EnvironmentVariableCategory.MODEL_SPECIFIC: 56, + EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, + EnvironmentVariableCategory.INTERNAL: 2, + } + + disposition_counts = Counter( + item.model_disposition + for item in ENVIRONMENT_VARIABLE_INVENTORY.values() + if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC + ) + assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { + ModelEnvironmentVariableDisposition.PROMOTE: 33, + ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, + ModelEnvironmentVariableDisposition.EXTERNAL: 0, + ModelEnvironmentVariableDisposition.INTERNALIZE: 11, + ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6, + } + + +def test_new_public_omni_names_use_project_prefix(): + """Grandfather legacy names without allowing more prefix exceptions.""" + legacy_public_names = { + "DIFFUSION_ATTENTION_BACKEND", + "DIFFUSION_CACHE_ADAPTER", + "DIFFUSION_CACHE_BACKEND", + "OMNI_DIFFUSION_PROMPT_EMBED_CACHE", + "OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", + "OMNI_DIFFUSION_SESSION_STATE_MANAGER", + "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS", + "SPEAKER_MAX_UPLOADED", + "SPEAKER_SAMPLES_DIR", + "VLLM_VIDEO_ASYNC_CHUNK", + "VLLM_VIDEO_AUDIO_DELTA_MODE", + } + public_names = { + name + for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() + if item.category is EnvironmentVariableCategory.PUBLIC_OMNI + } + + assert {name for name in public_names if not name.startswith("VLLM_OMNI_")} == legacy_public_names + + +def test_statically_resolvable_environment_accesses_are_classified(): + discovered: set[str] = set() + for path in _PACKAGE_ROOT.rglob("*.py"): + discovered.update(_environment_accesses(path)) + + assert discovered - ENVIRONMENT_VARIABLE_INVENTORY.keys() == set() + + +def test_model_and_benchmark_inventory_entries_are_still_referenced(): + """Reject transitional rows after their implementation has disappeared.""" + package_source = "\n".join( + path.read_text(encoding="utf-8") for path in _PACKAGE_ROOT.rglob("*.py") if path != _INVENTORY_MODULE + ) + transitional_categories = { + EnvironmentVariableCategory.MODEL_SPECIFIC, + EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL, + } + + stale_names = { + name + for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() + if item.category in transitional_categories and name not in package_source + } + + assert stale_names == set() + + +def test_environment_scanner_covers_indirection_aliases_membership_and_casing(): + expected_by_path = { + "metrics/definitions.py": { + "VLLM_OMNI_BENCH_AUDIO_CHANNELS", + "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE", + }, + "entrypoints/openai/video_stream_envs.py": { + "VLLM_VIDEO_ASYNC_CHUNK", + "VLLM_VIDEO_AUDIO_DELTA_MODE", + }, + "model_executor/models/mimo_audio/mimo_audio.py": {"model_stage"}, + "model_executor/models/moss_tts/modeling_moss_tts_local_depth.py": {"MOSS_TTS_DEBUG_STOP"}, + "distributed/ray_utils/utils.py": {"RAY_RAYLET_PID"}, + } + for relative_path, expected in expected_by_path.items(): + assert expected <= _environment_accesses(_PACKAGE_ROOT / relative_path) + + +def test_generated_server_storage_environment_names_are_classified(): + from vllm_omni.config.server_settings import FileBackend + + generated_names = { + f"VLLM_OMNI_SERVER_STORAGE__{field_name.upper()}" + for field_name in FileBackend.model_fields + if field_name != "type" + } + assert generated_names <= ENVIRONMENT_VARIABLE_INVENTORY.keys() + + +def test_public_omni_variables_are_in_the_reference_page(): + reference = _REFERENCE_PAGE.read_text(encoding="utf-8") + missing = { + name + for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() + if item.is_public_omni and f"`{name}`" not in reference + } + assert missing == set() + + +def test_secret_values_are_marked_for_redaction(): + assert {name for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() if item.redact_value} == { + "HF_TOKEN", + "HUGGINGFACE_HUB_TOKEN", + "OPENAI_API_KEY", + } + + +def test_collect_env_reports_only_safe_public_omni_values(monkeypatch): + from collect_env import get_env_vars + + monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache") + monkeypatch.setenv("HF_TOKEN", "hf-secret") + monkeypatch.setenv("OPENAI_API_KEY", "openai-secret") + monkeypatch.setenv("VLLM_OMNI_REPLICA_ID", "internal-replica") + + report = get_env_vars() + + assert "DIFFUSION_CACHE_BACKEND=tea_cache" in report.splitlines() + assert "hf-secret" not in report + assert "openai-secret" not in report + assert "VLLM_OMNI_REPLICA_ID" not in report + + +@pytest.mark.parametrize("error_type", [ImportError, OSError, AssertionError, RuntimeError]) +def test_collect_env_survives_inventory_import_error(monkeypatch, error_type): + from collect_env import get_env_vars + + real_import = builtins.__import__ + + def fail_inventory_import(name, *args, **kwargs): + if name == "vllm_omni.config.environment_variable_inventory": + raise error_type("broken vllm-omni installation") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fail_inventory_import) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + + assert "CUDA_VISIBLE_DEVICES=0" in get_env_vars().splitlines() diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1fb32239d..126b51c18 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1,15 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the additive structured Omni config.""" from __future__ import annotations from dataclasses import fields +from inspect import Parameter, signature from pathlib import Path +import msgspec import pytest from pydantic import ValidationError +from pydantic.fields import FieldInfo from transformers import Qwen3OmniMoeConfig +from vllm.config import CacheConfig as VllmCacheConfig +from vllm.config import CompilationConfig as VllmCompilationConfig +from vllm.config import LoadConfig as VllmLoadConfig +from vllm.config import ParallelConfig as VllmParallelConfig +from vllm.config import ProfilerConfig as VllmProfilerConfig +from vllm.config import SchedulerConfig as VllmSchedulerConfig from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.config import omni_config as omni_config_module @@ -36,10 +45,12 @@ from vllm_omni.config.stage_config import ( PipelineConfig, StageDeployConfig, StageExecutionType, + StagePipelineConfig, load_deploy_config, merge_pipeline_deploy, ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.engine.stage_engine_startup import _serialize_stage_config from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -82,10 +93,23 @@ def _from_pipeline_key( ) +def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None: + config = _from_pipeline_key( + "minimax_h3_disaggregated", + cli_overrides={"text_encoder_tp_size": 4}, + ) + + assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4 + diffusion = config.stage_by_id(1) + assert diffusion.parallel_config.tensor_parallel_size == 1 + assert diffusion.parallel_config.text_encoder_tp_size == 1 + + def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None: omni_config = _from_pipeline_key("personaplex") assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2] + assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["duplex", "duplex"] def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None: @@ -95,6 +119,7 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) - omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path)) assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1] + assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"] @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) @@ -122,6 +147,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod engine_args = legacy_stage.yaml_engine_args assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1) + assert omni_stage.model_config.session_mode == engine_args.get("session_mode", "turn") assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False) assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto") assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto") @@ -234,6 +260,36 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) +@pytest.mark.parametrize( + ("engine_extras", "cli_overrides"), + [ + ({"requires_full_payload_input": False}, {}), + ({}, {"stage_1_requires_full_payload_input": False}), + ], + ids=["deploy", "stage-cli"], +) +def test_from_pipeline_config_rejects_full_payload_input_capability_overrides( + engine_extras, + cli_overrides, +): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + deploy = DeployConfig( + stages=[ + StageDeployConfig( + stage_id=1, + engine_extras=engine_extras, + ) + ] + ) + + with pytest.raises(ValueError, match=r"no structured config owner: requires_full_payload_input"): + VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=deploy, + cli_overrides=cli_overrides, + ) + + @pytest.mark.parametrize( ("cli_overrides", "stage_id"), [ @@ -243,9 +299,29 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras ], ids=["global", "ar-stage", "generation-stage"], ) -def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id): - with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"): - _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides) +def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id): + stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id) + + assert stage.cache_config.cache_dtype == "fp8" + assert "cache_dtype" in stage.cache_config._omni_explicit_fields + + +@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"]) +def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + deploy = DeployConfig( + stages=[ + StageDeployConfig( + stage_id=0, + engine_extras={"parallel_config": {field_name: "10.0.0.1"}}, + ) + ] + ) + + stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0) + + assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1" + assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage(): @@ -413,10 +489,12 @@ def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating async_config = _from_pipeline_key("qwen3_tts") assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk") assert async_config.stage_by_id(1).custom_process_input_func is None + assert async_config.stage_by_id(1).model_config.requires_full_payload_input is True sync_config = _from_pipeline_key("qwen3_tts", cli_overrides={"async_chunk": False}) assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload") assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only") + assert sync_config.stage_by_id(1).model_config.requires_full_payload_input is True assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload") assert pipeline.get_stage(1).custom_process_input_func is None @@ -437,6 +515,8 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): "connector_config", "runtime_config", "parallel_config", + "compilation_config", + "profiler_config", "quantization_config", } assert "diffusion_config" not in public_fields @@ -455,10 +535,47 @@ def test_runtime_config_fields_match_structured_runtime_scope(): "num_gpus", "log_level", "log_stats", - "profiler_config", } +def test_upstream_config_compatible_fields_materialize_mapping_inputs(): + compilation_config = {"backend": "eager"} + profiler_config = {"profiler": "cuda"} + + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0], + compilation_config=compilation_config, + profiler_config=profiler_config, + ) + + assert isinstance(stage_config.compilation_config, VllmCompilationConfig) + assert stage_config.compilation_config.backend == "eager" + assert isinstance(stage_config.profiler_config, VllmProfilerConfig) + assert stage_config.profiler_config.profiler == "cuda" + + +def test_upstream_config_compatible_fields_validate_mapping_inputs(): + with pytest.raises(ValidationError): + VllmOmniARStageConfig( + stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0], + profiler_config={"delay_iterations": -1}, + ) + + +def test_upstream_config_compatible_fields_keep_prebuilt_config_types(): + compilation_config = VllmCompilationConfig(backend="eager") + profiler_config = VllmProfilerConfig(profiler="cuda") + + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0], + compilation_config=compilation_config, + profiler_config=profiler_config, + ) + + assert stage_config.compilation_config is compilation_config + assert stage_config.profiler_config is profiler_config + + def test_sub_config_fields_match_structured_scopes(): assert {f.name for f in fields(OmniStageModelConfig)} == { "model", @@ -471,16 +588,19 @@ def test_sub_config_fields_match_structured_scopes(): "trust_remote_code", "dtype", "attention_backend", + "attention_config", "moe_backend", "hf_overrides", "limit_mm_per_prompt", "interleave_mm_strings", "media_io_kwargs", "active_stream_window", + "session_mode", "duplex_max_sessions", "enable_sleep_mode", "default_sampling_params", "subtalker_sampling_params", + "silence_ban_frames", "has_sampling_extra_args", "custom_voice_dir", "task_type", @@ -488,7 +608,6 @@ def test_sub_config_fields_match_structured_scopes(): "enforce_eager", "max_cudagraph_capture_size", "enable_flashinfer_autotune", - "compilation_config", "enable_multithread_weight_load", "num_weight_load_threads", "disable_autocast", @@ -497,30 +616,27 @@ def test_sub_config_fields_match_structured_scopes(): # legacy engine-args path. "model_subdir", "tokenizer_subdir", + "requires_full_payload_input", } - assert {f.name for f in fields(OmniStageLoadConfig)} == { + vllm_load_fields = {f.name for f in fields(VllmLoadConfig)} + assert issubclass(OmniStageLoadConfig, VllmLoadConfig) + assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | { "tokenizer", - "download_dir", "skip_tokenizer_init", - "load_format", "tokenizer_mode", "config_format", "skip_mm_profiling", } - assert {f.name for f in fields(OmniStageCacheConfig)} == { - "kv_cache_memory_bytes", - "gpu_memory_utilization", - "enable_prefix_caching", + assert OmniStageLoadConfig(load_format="PT").load_format == "pt" + assert issubclass(OmniStageCacheConfig, VllmCacheConfig) + assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | { "disable_hybrid_kv_cache_manager", "mm_processor_cache_gb", "mamba_ssm_cache_dtype", } - assert {f.name for f in fields(OmniStageSchedulerConfig)} == { - "max_num_seqs", - "max_num_batched_tokens", + assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig) + assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | { "max_model_len", - "enable_chunked_prefill", - "async_scheduling", } assert {f.name for f in fields(OmniStageConnectorConfig)} == { "async_chunk", @@ -529,44 +645,104 @@ def test_sub_config_fields_match_structured_scopes(): "output_connectors", "input_connectors", } - assert {f.name for f in fields(OmniStageParallelConfig)} == { - "pipeline_parallel_size", - "data_parallel_size", - "tensor_parallel_size", - "enable_expert_parallel", - "world_size", - } - assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == { - "pipeline_parallel_size", - "data_parallel_size", - "tensor_parallel_size", + vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)} + assert issubclass(OmniStageParallelConfig, VllmParallelConfig) + assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields + assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | { "sequence_parallel_size", "ulysses_degree", "ring_degree", "allgather_degree", "ulysses_mode", + "ulysses_a2a_permute", "cfg_parallel_size", "vae_patch_parallel_size", - "text_encoder_tp_size", "vae_parallel_mode", + "text_encoder_tp_size", "use_hsdp", "mask_sp_padding", "hsdp_shard_size", "hsdp_replicate_size", - "enable_expert_parallel", - "world_size", } +def test_inherited_sub_configs_initialize_transport_safe_derived_fields(): + scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096) + assert scheduler_config.max_num_encoder_input_tokens == 4096 + assert scheduler_config.encoder_cache_size == 4096 + + unresolved_scheduler_config = OmniStageSchedulerConfig() + assert unresolved_scheduler_config.max_num_encoder_input_tokens is None + assert unresolved_scheduler_config.encoder_cache_size is None + + parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2) + assert parallel_config.data_parallel_index == 2 + + diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2) + assert diffusion_parallel_config.data_parallel_index == 2 + + +@pytest.mark.parametrize( + "config_cls", + [ + OmniStageLoadConfig, + OmniStageCacheConfig, + OmniStageSchedulerConfig, + OmniStageParallelConfig, + OmniStageDiffusionParallelConfig, + ], +) +def test_inherited_sub_configs_are_keyword_only(config_cls): + assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values()) + + with pytest.raises(TypeError): + config_cls(1) + + +@pytest.mark.parametrize( + "config_cls", + [ + OmniStageLoadConfig, + OmniStageCacheConfig, + OmniStageSchedulerConfig, + OmniStageParallelConfig, + OmniStageDiffusionParallelConfig, + ], +) +def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls): + config = config_cls() + + uninitialized_fields = [ + config_field.name + for config_field in fields(config) + if isinstance(getattr(config, config_field.name), FieldInfo) + ] + assert uninitialized_fields == [] + + msgspec.msgpack.encode(_serialize_stage_config(config)) + + +def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe(): + omni_config = _from_pipeline_key("qwen3_tts") + + for stage_config in omni_config.stage_configs: + payload = { + "stage_id": stage_config.stage_id, + "stage_config": _serialize_stage_config(stage_config), + } + msgspec.msgpack.encode(payload) + + def test_diffusion_parallel_config_fields_cover_legacy_surface(): from vllm_omni.diffusion.data import DiffusionParallelConfig legacy_fields = {f.name for f in fields(DiffusionParallelConfig)} structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)} expected_upstream_fields = {"mask_sp_padding"} + vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)} assert legacy_fields | expected_upstream_fields <= structured_fields - assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"} + assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set() def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface(): @@ -609,8 +785,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding(): cfg = OmniStageDiffusionParallelConfig( - pipeline_parallel_size=2, - ulysses_degree=2, + ulysses_degree=4, use_hsdp=True, hsdp_shard_size=-1, hsdp_replicate_size=2, @@ -621,10 +796,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding(): def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp(): - with pytest.raises(ValueError, match="cannot be used with TP or DP"): + with pytest.raises(ValueError, match="not compatible with TP"): OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2) - with pytest.raises(ValueError, match="cannot be used with TP or DP"): + with pytest.raises(ValueError, match="not compatible with DP"): OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2) @@ -685,6 +860,23 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr assert stage.parallel_config.world_size == 2 +def test_from_pipeline_config_preserves_deployed_ulysses_a2a_permute() -> None: + pipeline = _resolve_pipeline_or_skip("dreamzero") + deploy = DeployConfig( + stages=[ + StageDeployConfig( + stage_id=0, + ulysses_a2a_permute=True, + ) + ] + ) + + stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.parallel_config.ulysses_a2a_permute is True + + def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism(): cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4) @@ -779,6 +971,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic(): omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false") +def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path): + deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " extras:", + " ltx2_use_conv_vae: true", + ] + ) + ) + + stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0) + + assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True + + +def test_stage_override_routes_ltx2_conv_vae_extra(): + stage = _from_pipeline_key( + "dreamzero", + cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}}, + ).stage_by_id(0) + + assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True + + def test_structured_diffusion_config_rejects_invalid_compile_granularity(): with pytest.raises(ValidationError, match="diffusion_compile_granularity"): omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block") @@ -819,7 +1040,10 @@ def test_from_pipeline_config_uses_hf_config_for_callable_resolver(): def test_from_pipeline_config_accepts_pre_resolved_pipeline(): - resolved_pipeline = PipelineConfig(model_type="callable_resolved_variant") + resolved_pipeline = PipelineConfig( + model_type="callable_resolved_variant", + stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),), + ) omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline) @@ -852,6 +1076,7 @@ def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_p pipeline = PipelineConfig( model_type="pipeline_with_default", default_deploy_config_name=default_name, + stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),), ) loaded_paths = [] @@ -966,7 +1191,10 @@ def test_diffusion_config_preserves_existing_coercion_hooks(): def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): + from vllm_omni.platforms import current_omni_platform + monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE") + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( diffusion_attention_backend="flash_attn", @@ -976,6 +1204,7 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): kv_cache_skip_layers=[2], static_lora_scale=0.25, diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=2, diffusers_load_kwargs=None, diffusers_call_kwargs=None, ) @@ -992,7 +1221,10 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): assert cfg.diffusers_call_kwargs == {} -def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path): +def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml" deploy_path.write_text( "\n".join( @@ -1004,6 +1236,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar " diffusion_attention_backend: flash_attn", " fa_deterministic: true", " diffusion_kv_mode: paged_scheduler", + " diffusion_kv_max_rows_per_request: 2", ] ) ) @@ -1017,6 +1250,36 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn" assert stage.diffusion_config.fa_deterministic is True assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2 + + +def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + deploy_path = tmp_path / "dreamzero_fastvideo_vsa.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " diffusion_attention_backend: FASTVIDEO_VSA", + " fastvideo_vsa_topk: 96", + ] + ) + ) + + stage = _from_pipeline_key( + "dreamzero", + deploy_config_path=str(deploy_path), + ).stage_by_id(0) + + attention_config = stage.diffusion_config.diffusion_attention_config + assert attention_config.default is not None + assert attention_config.default.backend == "FASTVIDEO_VSA" + assert attention_config.default.fastvideo_vsa_topk == 96 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path): @@ -1059,6 +1322,7 @@ def test_diffusion_config_field_classification_covers_current_fields(): "prompt_embed_cache_size", "diffusion_kv_cache_dtype", "diffusion_kv_mode", + "diffusion_kv_max_rows_per_request", } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS assert { "revision", @@ -1077,3 +1341,21 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config) assert cfg.quantization_config == quantization_config + + +def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + quantization_config = {"method": "int8", "activation_scheme": "dynamic"} + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( + quantization_config=quantization_config, + ) + + # Exercise the terminal construction performed by the future typed startup + # path without probing ports or loading remote model metadata. + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) + monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None) + cfg.enrich_config() + + assert cfg.quantization_config is not None + assert cfg.quantization_config.get_name() == "int8" diff --git a/tests/config/test_pipeline_registry.py b/tests/config/test_pipeline_registry.py index 4bc19b116..88109e052 100644 --- a/tests/config/test_pipeline_registry.py +++ b/tests/config/test_pipeline_registry.py @@ -1,16 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for out of tree registration to OMNI_PIPELINES.""" import pytest from transformers import PretrainedConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline -from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver +from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig, pipeline_cfg_resolver pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def build_fake_pipeline_config(model_type: str) -> PipelineConfig: + return PipelineConfig( + model_type=model_type, stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),) + ) + + @pytest.fixture def custom_resolver(): """Build a reusable custom resolver for PipelineConfigs.""" @@ -22,7 +28,7 @@ def custom_resolver(): def custom_resolver( hf_config: CustomConfigType, ) -> PipelineConfig: - return PipelineConfig(model_type="resolved_type") + return build_fake_pipeline_config("resolved_type") return custom_resolver @@ -30,7 +36,7 @@ def custom_resolver(): def test_register_pipeline_config(clean_pipeline_registry): """Ensure that we can register a custom pipeline config to OMNI_PIPELINES.""" new_model_type = "new_model_type" - pipe_cfg = PipelineConfig(model_type=new_model_type) + pipe_cfg = build_fake_pipeline_config(new_model_type) assert new_model_type not in OMNI_PIPELINES register_pipeline(pipe_cfg) assert new_model_type in OMNI_PIPELINES @@ -41,7 +47,7 @@ def test_register_pipeline_config_with_model_type(clean_pipeline_registry): """Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES.""" new_model_type = "new_model_type" unused_model_type = "foo" - pipe_cfg = PipelineConfig(model_type=unused_model_type) + pipe_cfg = build_fake_pipeline_config(unused_model_type) assert new_model_type not in OMNI_PIPELINES assert unused_model_type not in OMNI_PIPELINES @@ -66,3 +72,10 @@ def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_r """Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type.""" with pytest.raises(ValueError): register_pipeline(custom_resolver) + + +def test_minimax_h3_disaggregation_is_explicit_opt_in(): + assert "minimax_h3" not in OMNI_PIPELINES + pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"] + assert isinstance(pipeline, PipelineConfig) + assert pipeline.model_type == "minimax_h3_disaggregated" diff --git a/tests/conftest.py b/tests/conftest.py index b453bbcbf..fa55a23da 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Root pytest entrypoint for the vLLM-Omni test suite. -- `tests/conftest.py` stays thin: plugin registration + compatibility re-exports. +- `tests/conftest.py` stays thin: plugin registration only. - Importable utilities live under `tests/helpers/`. - Helper unit tests live under `tests/helpers/tests/`. - Fixtures live under `tests/helpers/fixtures/` and are loaded via `pytest_plugins`. @@ -19,57 +22,15 @@ bootstrap_vllm_layer_custom_op_modules() pytest_plugins = ( "tests.helpers.fixtures.config", - "tests.helpers.fixtures.env", + "tests.helpers.fixtures.clean", "tests.helpers.fixtures.log", - "tests.helpers.fixtures.run_args", + "tests.helpers.fixtures.media", + "tests.helpers.fixtures.pytest_collection", + "tests.helpers.fixtures.pytest_run_args", "tests.helpers.fixtures.runtime", - "tests.helpers.fixtures.speaker_cache", ) def pytest_terminal_summary(terminalreporter, exitstatus, config): # Marker for Buildkite log folding before pytest summary lines. terminalreporter.write_line("--- Running Summary") - - -# Backward-compatible lazy re-exports. -# (Many tests still import from `tests.conftest`; migrate these imports to `tests.helpers.*` over time.) -# Keep these lazy so conftest import does not trigger heavy helper dependencies. -_ASSERTION_EXPORT_NAMES = ( - "assert_audio_speech_response", - "assert_diffusion_response", - "assert_http_error", - "assert_image_diffusion_response", - "assert_image_valid", - "assert_omni_response", - "assert_video_diffusion_response", - "assert_video_valid", -) -_MEDIA_EXPORT_NAMES = ( - "convert_audio_bytes_to_text", - "convert_audio_file_to_text", - "cosine_similarity_text", - "decode_b64_image", - "generate_synthetic_audio", - "generate_synthetic_image", - "generate_synthetic_video", -) -_STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",) -_RUNTIME_EXPORT_NAMES = ( - "DiffusionResponse", - "HttpResponse", - "OmniResponse", - "OmniRunner", - "OmniRunnerHandler", - "OmniServer", - "OmniServerParams", - "OmniServerStageCli", - "OpenAIClientHandler", - "dummy_messages_from_mix_data", -) -_LAZY_EXPORT_MODULES = { - **{name: "tests.helpers.assertions" for name in _ASSERTION_EXPORT_NAMES}, - **{name: "tests.helpers.media" for name in _MEDIA_EXPORT_NAMES}, - **{name: "tests.helpers.stage_config" for name in _STAGE_CONFIG_EXPORT_NAMES}, - **{name: "tests.helpers.runtime" for name in _RUNTIME_EXPORT_NAMES}, -} diff --git a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py new file mode 100644 index 000000000..d9fbc452b --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py @@ -0,0 +1,89 @@ +"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``. + +``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before +handing control to upstream, because upstream ``Scheduler.schedule()`` raises +``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it +admits in a finished state -- which kills the stage's engine core, not just +the request. + +Regression for the sweep missing ``skipped_waiting``, the third queue +upstream admits from. An aborted duplex session parked there was re-selected +by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the +stage. +""" + +from __future__ import annotations + +import pytest +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_POLICIES = [ + pytest.param(SchedulingPolicy.FCFS, id="fcfs"), + pytest.param(SchedulingPolicy.PRIORITY, id="priority"), +] + + +class _StubRequest: + """Minimal ``Request`` stub with the surface the sweep exercises. + + Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue`` + relies on so the same stub works under both scheduling policies. + """ + + def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None: + self.request_id = request_id + self.status = status + self.priority = 0 + self.arrival_time = arrival_time + + def __lt__(self, other: _StubRequest) -> bool: + return (self.priority, self.arrival_time) < (other.priority, other.arrival_time) + + +def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()): + """Bare scheduler with just the surface the sweep reads.""" + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(policy) + for req in waiting: + scheduler.waiting.add_request(req) + scheduler.skipped_waiting = create_request_queue(policy) + for req in skipped_waiting: + scheduler.skipped_waiting.add_request(req) + scheduler.running = list(running) + return scheduler + + +@pytest.mark.parametrize("policy", _POLICIES) +def test_schedule_sweeps_skipped_waiting_before_upstream_selection( + monkeypatch: pytest.MonkeyPatch, + policy: SchedulingPolicy, +) -> None: + """The public schedule path must remove aborted skipped requests first.""" + aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED) + scheduler = _make_scheduler(policy, skipped_waiting=[aborted]) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.num_waiting_for_streaming_input = 0 + scheduler._process_pending_omni_inputs = lambda model_mode: None + scheduler._resync_streaming_input_counter = lambda: None + scheduler._should_defer_waiting_admission = lambda: False + scheduler._restore_omni_wait_queues = lambda: None + scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None + scheduler.get_finished_requests_needing_kv_transfer = lambda: {} + scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output + + sentinel = object() + + def upstream_schedule(_self, _throttle_prefills: bool = False): + assert list(scheduler.skipped_waiting) == [] + return sentinel + + monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule) + + assert OmniARScheduler.schedule(scheduler) is sentinel diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index fb92414ef..5c1290fc4 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py @@ -100,6 +100,7 @@ def _make_drain_sched(session: Request) -> MagicMock: sched.pending_stop_after_extraction = set() sched.connector = None sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None return sched @@ -141,13 +142,16 @@ def test_exact_drain_delivers_new_segment_frame() -> None: session.num_computed_tokens = 6 session.num_output_placeholders = 1 session.num_in_flight_tokens = 1 + session.async_tokens_to_discard = 1 _replace_streaming_session(session) assert session.num_stale_output_tokens == 1 + assert session.async_tokens_to_discard == 1 sched = _make_drain_sched(session) assert _run_step(sched, session, num_scheduled=1, token=42) is False # late frame dropped assert session.num_stale_output_tokens == 0 + assert session.async_tokens_to_discard == 0 assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index ca80323a7..a1aa7351f 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for Omni AR streaming-session async placeholder handling.""" @@ -18,21 +18,26 @@ from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate -from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler # isort: on pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler: +def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler: sched = OmniARScheduler.__new__(OmniARScheduler) sched._new_prompt_len_snapshot = {} - sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id)) + sched.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode), + ) sched.num_waiting_for_streaming_input = 0 sched.log_stats = False sched.chunk_transfer_adapter = None sched.skipped_waiting = set() + sched._free_request_blocks = MagicMock() + sched.encoder_cache_manager = MagicMock() + sched._inflight_prefills = set() return sched @@ -61,6 +66,9 @@ def _run_resumable_segment_stop( session: Request, *, session_finished: bool = False, + handle_stopped=None, + chunk_transfer_adapter=None, + inter_stage_output=None, ): sched = MagicMock() sched.requests = {session.request_id: session} @@ -72,11 +80,14 @@ def _run_resumable_segment_stop( return [42], True sched._update_request_with_output.side_effect = stop_request - sched._handle_stopped_request.return_value = session_finished # vLLM 0.26 returns (kv_xfer_params, ec_xfer_params); an unconfigured # MagicMock iterates empty and fails to unpack at the call site. sched._free_request.return_value = (None, None) - sched.chunk_transfer_adapter = None + if handle_stopped is None: + sched._handle_stopped_request.return_value = session_finished + else: + sched._handle_stopped_request.side_effect = handle_stopped + sched.chunk_transfer_adapter = chunk_transfer_adapter sched.running = [session] sched.waiting_for_transfer_free = set() sched.transfer_triggered_requests = set() @@ -84,6 +95,7 @@ def _run_resumable_segment_stop( sched.pending_stop_after_extraction = set() sched.connector = None sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None @@ -102,6 +114,7 @@ def _run_resumable_segment_stop( model_runner_output.cudagraph_stats = None model_runner_output.req_id_to_index = {session.request_id: 0} model_runner_output.routed_experts = None + model_runner_output.inter_stage_outputs = [inter_stage_output] if inter_stage_output is not None else None return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) @@ -157,6 +170,44 @@ def test_update_from_output_settles_in_flight_tokens() -> None: assert session.num_in_flight_tokens == 0 +def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> None: + """The old segment's transfer must not observe the next segment's reset. + + A queued streaming update can replace the same mutable Request while the + stop output is still being handled. The connector needs the confirmed + token count from before that replacement to avoid dropping the boundary. + """ + session = _make_request() + session.status = RequestStatus.RUNNING + session.resumable = True + session.num_computed_tokens = 26 + adapter = MagicMock() + adapter._confirmed_num_computed_tokens.return_value = 26 + inter_stage_output = {"codes": {"audio": [7]}} + + def replace_with_next_segment(request: Request) -> bool: + request.num_computed_tokens = 0 + request._omni_segment_generation = 1 + request.status = RequestStatus.WAITING + return False + + _run_resumable_segment_stop( + session, + handle_stopped=replace_with_next_segment, + chunk_transfer_adapter=adapter, + inter_stage_output=inter_stage_output, + ) + + adapter.save_async.assert_called_once_with( + inter_stage_output, + session, + True, + new_token_ids=[42], + confirmed_num_computed_tokens=26, + segment_generation=0, + ) + + def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None: """A decode step that neither stops nor carries an inter-stage payload. @@ -181,6 +232,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non sched.pending_stop_after_extraction = set() sched.connector = None sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None @@ -207,6 +259,171 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non sched.chunk_transfer_adapter.save_async.assert_not_called() +def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None: + """Native duplex + async: the next unit is usually already queued when the + current one stops, so _handle_stopped_request applies the update and the + stop site below it fences the same in-flight decode a second time. An + accumulated residue outlives the drain and swallows the next unit's + listen/speak, which the client sees as silence until its timeout. + """ + session = _make_request() + session.status = RequestStatus.RUNNING + session.resumable = True + session.append_output_token_ids([7]) + session.num_computed_tokens = 4 + session.num_output_placeholders = 1 + # This step's frame (settled to 0 by update_from_output) plus one extra + # async decode still in flight. + session.num_in_flight_tokens = 2 + + sched = _make_scheduler(stage_id=0) + sched._enqueue_waiting_request = MagicMock() + + def handle_stopped(request: Request) -> bool: + sched._update_request_as_session(request, _make_update([10, 20])) + return False + + _run_resumable_segment_stop(session, handle_stopped=handle_stopped) + + # Exactly the one unreported decode, so the drain reaches zero before the + # new segment's first frame arrives. + assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1 + + +def test_stale_async_frame_is_dropped_before_output_processing() -> None: + session = _make_request() + session.status = RequestStatus.RUNNING + session.num_in_flight_tokens = 2 + session.num_computed_tokens = session.num_tokens + 1 + session.num_output_placeholders = 1 + session.async_tokens_to_discard = 1 + session.sampling_params = SimpleNamespace(num_logprobs=1) + num_computed_tokens = session.num_computed_tokens + num_output_placeholders = session.num_output_placeholders + + sched = MagicMock() + sched.requests = {session.request_id: session} + sched.perf_metrics = None + sched.structured_output_manager.should_advance.return_value = False + + def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]: + request.async_tokens_to_discard = 0 + return token_ids, False + + sched._update_request_with_output.side_effect = discard_stale_output + sched._process_kv_transfer_trigger.return_value = False + sched.chunk_transfer_adapter = MagicMock() + sched.running = [session] + sched.waiting_for_transfer_free = set() + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.pending_stop_after_extraction = set() + sched.connector = None + sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 + sched.finished_req_ids_dict = {} + sched.make_stats.return_value = None + + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {session.request_id: 1} + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.multimodal_outputs = None + model_runner_output.inter_stage_outputs = [{"hidden": object()}] + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {session.request_id: 0} + model_runner_output.routed_experts = None + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + assert session.async_tokens_to_discard == 0 + assert session.status == RequestStatus.RUNNING + assert session.num_computed_tokens == num_computed_tokens + assert session.num_output_placeholders == num_output_placeholders + sched.chunk_transfer_adapter.save_async.assert_not_called() + + session.sampling_params = SamplingParams(max_tokens=8) + next_payload = {"hidden": object()} + sched._update_request_with_output.side_effect = None + sched._update_request_with_output.return_value = ([43], False) + model_runner_output.sampled_token_ids = [[43]] + model_runner_output.inter_stage_outputs = [next_payload] + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + sched.chunk_transfer_adapter.save_async.assert_called_once_with( + next_payload, + session, + False, + new_token_ids=[43], + confirmed_num_computed_tokens=None, + ) + + +def test_legacy_stale_async_marker_bypasses_real_placeholder_accounting() -> None: + request = _make_request() + request.status = RequestStatus.RUNNING + request.num_in_flight_tokens = 1 + request.num_output_placeholders = 0 + request.async_tokens_to_discard = 1 + request.num_stale_output_tokens = 0 + + sched = OmniARAsyncScheduler.__new__(OmniARAsyncScheduler) + sched.requests = {request.request_id: request} + sched.perf_metrics = None + sched.chunk_transfer_adapter = None + sched.connector = None + sched.ec_connector = None + sched.waiting_for_transfer_free = set() + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.pending_stop_after_extraction = set() + sched.finished_req_ids_dict = {} + sched._new_prompt_len_snapshot = {} + sched.kv_cache_manager = MagicMock() + sched.kv_cache_manager.take_events.return_value = None + sched._remove_stopped_requests_from_queues = MagicMock() + sched._handle_failed_kv_load_outputs = MagicMock(return_value=[]) + sched._cleanup_kv_tracking = MagicMock() + sched._aggregate_kv_connector_stats = MagicMock(return_value=None) + sched._publish_kv_cache_events = MagicMock() + sched._attach_finished_request_sets = MagicMock() + sched._attach_scheduler_stats = MagicMock() + sched._capture_omni_connector_output = MagicMock() + + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {request.request_id: 1} + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.multimodal_outputs = None + model_runner_output.inter_stage_outputs = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {request.request_id: 0} + model_runner_output.routed_experts = None + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + assert request.async_tokens_to_discard == 0 + assert request.num_output_placeholders == 0 + assert request._output_token_ids == [] + + def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None: sched = _make_scheduler(stage_id=0) session = _make_request() @@ -276,29 +493,33 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None: "meta": {"turn_eos_token_id": 99}, } assert session.status == RequestStatus.WAITING + sched._free_request_blocks.assert_called_once_with(session) + sched.encoder_cache_manager.free.assert_called_once_with(session) -def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None: +def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_watermark() -> None: sched = _make_scheduler(stage_id=1) + session = _make_request() sched.chunk_transfer_adapter = SimpleNamespace( receives_chunks=False, segment_finished_requests=set(), + requests_num_chunks_sent={session.external_req_id: 59}, ) - session = _make_request() session.status = RequestStatus.WAITING_FOR_STREAMING_REQ session.prompt_token_ids = [0] * 59 session._all_token_ids.clear() session._all_token_ids.extend(session.prompt_token_ids) session.num_prompt_tokens = 59 session.num_computed_tokens = 59 + session.num_in_flight_tokens = 2 update = _make_update([0] * 10) update.additional_information = None update.model_intermediate_buffer = { "ids": {"tts": list(range(8))}, "hidden_states": {"tts": [[0.0]] * 8}, "meta": { - "replace_streaming_prompt": True, "next_stage_prompt_len": 10, + "replace_streaming_prompt": True, }, } @@ -308,6 +529,110 @@ def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None assert list(session._all_token_ids) == [0] * 10 assert session.num_prompt_tokens == 10 assert session.num_computed_tokens == 0 + assert session.num_stale_output_tokens == 2 assert session.additional_information is None assert session.model_intermediate_buffer == update.model_intermediate_buffer assert session.status == RequestStatus.WAITING + assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {} + sched._free_request_blocks.assert_called_once_with(session) + sched.encoder_cache_manager.free.assert_called_once_with(session) + + +def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None: + sched = _make_scheduler(stage_id=1) + session = _make_request() + session.external_req_id = "external-ar-streaming-test" + session.num_in_flight_tokens = 2 + # _update_request_as_session() may have already fenced this frame before + # the connector marks the explicit replacement ready. + session.num_stale_output_tokens = 2 + session.num_output_placeholders = 2 + session.spec_token_ids = [-1, -1] + sched.requests = {session.request_id: session} + sched._inflight_prefills.add(session) + sched.chunk_transfer_adapter = SimpleNamespace( + replaced_streaming_prompt_ids={session.request_id}, + requests_with_ready_chunks={session.request_id}, + requests_num_chunks_sent={session.external_req_id: 4090}, + ) + + sched._reset_ready_async_chunk_replacements() + sched._reset_ready_async_chunk_replacements() + + sched._free_request_blocks.assert_called_once_with(session) + sched.encoder_cache_manager.free.assert_called_once_with(session) + assert session not in sched._inflight_prefills + assert session.num_stale_output_tokens == 2 + assert session.num_output_placeholders == 0 + assert session.spec_token_ids == [] + assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set() + assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id} + assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {} + + +def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None: + """A WAITING_FOR_CHUNK stop must not delete its newly parked session.""" + session = _make_request() + session.resumable = True + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + + def queue(*requests): + result = MagicMock() + result.requests = set(requests) + result.add_request.side_effect = result.requests.add + result.remove_requests.side_effect = result.requests.difference_update + return result + + sched = OmniARScheduler.__new__(OmniARScheduler) + sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex")) + sched.running = [] + sched.waiting = queue() + sched.skipped_waiting = queue(session) + sched.num_waiting_for_streaming_input = 1 + sched.chunk_transfer_adapter = SimpleNamespace( + receives_chunks=True, + segment_finished_requests={session.request_id}, + ) + + sched._resume_downstream_chunk_receiver(session) + sched._remove_stopped_requests_from_queues(set(), {session}) + + assert session.status == RequestStatus.WAITING + assert session in sched.waiting.requests + assert session not in sched.skipped_waiting.requests + assert sched.num_waiting_for_streaming_input == 0 + assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests + + +@pytest.mark.parametrize( + ("receives_chunks", "session_mode"), + [ + (False, "duplex"), + (True, "turn"), + ], +) +def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked( + receives_chunks: bool, + session_mode: str, +) -> None: + """Only duplex connector-driven receivers resume without an update.""" + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + + sched = OmniARScheduler.__new__(OmniARScheduler) + sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode)) + sched.num_waiting_for_streaming_input = 1 + sched.chunk_transfer_adapter = SimpleNamespace( + receives_chunks=receives_chunks, + segment_finished_requests={session.request_id}, + ) + sched.skipped_waiting = MagicMock() + sched._enqueue_waiting_request = MagicMock() + + sched._resume_downstream_chunk_receiver(session) + + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + assert sched.num_waiting_for_streaming_input == 1 + assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests + sched.skipped_waiting.remove_requests.assert_not_called() + sched._enqueue_waiting_request.assert_not_called() diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py index 27087d253..0e0ed04ad 100644 --- a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py +++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py @@ -76,6 +76,9 @@ def _make_scheduler( else: hf_config.voxcpm2_runtime_config = runtime_config sched.vllm_config = SimpleNamespace(model_config=model_config) + # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the + # abort sweep at the top of ``schedule()`` reads it. + sched.skipped_waiting = _MockQueue() return sched diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index 0f654715a..3e1b53afe 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -138,6 +138,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None: sched.recompute_kv_load_failures = False sched.connector = None sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None @@ -410,12 +411,13 @@ class TestRealignRequestStatusToQueues: assert stale.status == RequestStatus.RUNNING assert clean.status == RequestStatus.RUNNING - def test_finished_request_is_skipped(self) -> None: + def test_non_resumable_finished_request_is_skipped(self) -> None: """Already-finished requests must not be touched -- they may have legitimate finished statuses (FINISHED_STOPPED etc.) that a status flip would corrupt. """ req = _make_request(request_id="req-finished") + req.resumable = False req.status = RequestStatus.FINISHED_STOPPED sched = _RealignSchedulerStub( diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py new file mode 100644 index 000000000..36b084f7c --- /dev/null +++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review). + +Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every +non-empty scheduled step and, at the top of ``update_from_output()``, +advances ``processed_step_seq`` and drains the deferred frees. The Omni +schedulers each carried only half of that contract: + +- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence + advances) but its reimplemented ``update_from_output()`` never advanced + ``processed_step_seq`` nor drained -> deferred blocks accumulate without + bound under async/PP KV-consumer configurations. +- ``OmniGenerationScheduler`` reimplements both: the fast-path + ``schedule()`` never advanced ``sched_step_seq`` (so finished requests + freed blocks an in-flight GPU step could still be writing) and its + ``update_from_output()`` never drained. + +These tests pin the restored contract at the call boundary; the drain and +free logic itself is upstream-tested. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +# Imports must run in this order: vllm_omni applies patches to vllm.v1.request +# before Request / RequestStatus are bound in this module. +# isort: off +import vllm_omni # noqa: F401 - import for side effects (patch vLLM) +from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.interface import PauseState +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.request_queue import SchedulingPolicy +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +# isort: on + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_update_sched(scheduler_cls) -> MagicMock: + """MagicMock scheduler driving the real ``update_from_output`` with an + empty batch; only the fence state carries real values.""" + sched = MagicMock() + sched.requests = {} + sched.perf_metrics = None + sched.chunk_transfer_adapter = None + sched.defer_block_free = True + sched.processed_step_seq = 0 + return sched + + +def _empty_frame(*, total_num_scheduled_tokens: int): + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {} + scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + return scheduler_output, model_runner_output + + +@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler]) +def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None: + sched = _make_update_sched(scheduler_cls) + scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3) + + scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output) + + assert sched.processed_step_seq == 1 + sched._drain_deferred_frees.assert_called_once_with() + + +@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler]) +def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None: + """0-token steps do not advance the fence upstream (their seq was never + advanced in schedule()); draining on them would run ahead of the GPU.""" + sched = _make_update_sched(scheduler_cls) + scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0) + + scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output) + + assert sched.processed_step_seq == 0 + sched._drain_deferred_frees.assert_not_called() + + +@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler]) +def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None: + sched = _make_update_sched(scheduler_cls) + sched.defer_block_free = False + scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3) + + scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output) + + assert sched.processed_step_seq == 0 + sched._drain_deferred_frees.assert_not_called() + + +def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]: + """MagicMock scheduler that drives the real generation fast-path + ``schedule()`` over one running one-shot request.""" + request = Request( + request_id="req-fence", + prompt_token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=4), + pooling_params=None, + arrival_time=100.0, + block_hasher=None, + ) + request.status = RequestStatus.RUNNING + + sched = MagicMock() + sched.max_num_scheduled_tokens = 8 + sched._pause_state = PauseState.UNPAUSED + sched.requests = {request.request_id: request} + sched.running = [request] + sched.waiting = [] # falsy: the waiting loop never runs + sched.policy = SchedulingPolicy.FCFS + sched.chunk_transfer_adapter = None + sched.scheduler_config = MagicMock(enable_chunked_prefill=True) + sched.num_lookahead_tokens = 0 + sched.log_stats = False + sched.kv_cache_config = MagicMock(kv_cache_groups=[]) + sched.use_v2_model_runner = False + sched.prev_step_scheduled_req_ids = set() + sched.needs_kv_cache_zeroing = False + sched.connector = None + sched.ec_connector = None + sched.defer_block_free = defer_block_free + sched.sched_step_seq = 0 + # last_sched_seq stamping order: _update_after_schedule must observe the + # already-advanced fence value, as upstream schedule() guarantees. + sched.seq_at_update_after_schedule = None + + def _capture(_scheduler_output): + sched.seq_at_update_after_schedule = sched.sched_step_seq + + sched._update_after_schedule.side_effect = _capture + return sched, request + + +def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None: + sched, request = _make_fast_path_sched(defer_block_free=True) + + OmniGenerationScheduler.schedule(sched) + + assert sched.sched_step_seq == 1 + assert sched.seq_at_update_after_schedule == 1 + + +def test_fast_path_schedule_defer_disabled_leaves_fence() -> None: + sched, request = _make_fast_path_sched(defer_block_free=False) + + OmniGenerationScheduler.schedule(sched) + + assert sched.sched_step_seq == 0 + + +def test_generation_schedule_sweeps_and_resyncs_around_pending_inputs() -> None: + sched, request = _make_fast_path_sched(defer_block_free=False) + calls: list[str] = [] + sched._drop_aborted_queued_requests.side_effect = lambda: calls.append("sweep") + sched._process_pending_omni_inputs.side_effect = lambda model_mode: calls.append(model_mode) + sched._resync_streaming_input_counter.side_effect = lambda: calls.append("resync") + + OmniGenerationScheduler.schedule(sched) + + assert calls == ["sweep", "generation", "sweep", "resync"] diff --git a/tests/core/sched/test_omni_sched_ec_request_finish.py b/tests/core/sched/test_omni_sched_ec_request_finish.py new file mode 100644 index 000000000..81b6577fd --- /dev/null +++ b/tests/core/sched/test_omni_sched_ec_request_finish.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review). + +Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()`` +BEFORE the encoder cache is freed (so the connector can inspect per-request +state), honors its delayed-free result, and returns the generated +``ec_transfer_params``; ``update_from_output()`` then carries them into the +``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and +always returned ``None``; both Omni schedulers discarded the second tuple +item, so a remote encoder-cache handle could never reach the frontend. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +# Imports must run in this order: vllm_omni applies patches to vllm.v1.request +# before Request / RequestStatus are bound in this module. +# isort: off +import vllm_omni # noqa: F401 - import for side effects (patch vLLM) +from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +# isort: on + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"} + + +def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]: + """Minimal real-`_free_request` harness; `order` records the EC hook + firing relative to the encoder-cache free.""" + order = MagicMock() + sched = OmniARScheduler.__new__(OmniARScheduler) + sched._omits_kv_transfer_cache = {} + sched._connector_finished = lambda request: (False, None) + sched.ec_connector = ec_connector + if ec_connector is not None: + order.attach_mock(ec_connector.request_finished, "request_finished") + sched.encoder_cache_manager = MagicMock() + order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free") + sched.finished_req_ids = set() + sched._new_prompt_len_snapshot = {} + sched.finished_req_ids_dict = None + sched._should_transfer_kv_for_request = lambda req_id: False + sched._free_blocks = MagicMock() + sched._free_input_coordinator_request = MagicMock() + sched.chunk_transfer_adapter = None + return sched, order + + +class _FakeFinishedRequest: + def __init__(self, request_id: str) -> None: + self.request_id = request_id + + def is_finished(self) -> bool: + return True + + +def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None: + ec_connector = MagicMock() + ec_connector.request_finished.return_value = (False, EC_PARAMS) + sched, order = _make_free_request_sched(ec_connector=ec_connector) + request = _FakeFinishedRequest("req-ec") + + kv_params, ec_params = sched._free_request(request) + + assert ec_params == EC_PARAMS + ec_connector.request_finished.assert_called_once_with(request) + hook_calls = [name for name, *_ in order.mock_calls] + assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free") + sched._free_blocks.assert_called_once_with(request) + + +def test_free_request_honors_ec_delayed_free() -> None: + ec_connector = MagicMock() + ec_connector.request_finished.return_value = (True, EC_PARAMS) + sched, _ = _make_free_request_sched(ec_connector=ec_connector) + request = _FakeFinishedRequest("req-ec-delay") + + kv_params, ec_params = sched._free_request(request) + + assert ec_params == EC_PARAMS + sched._free_blocks.assert_not_called() + + +def test_free_request_without_ec_connector_returns_none() -> None: + sched, _ = _make_free_request_sched(ec_connector=None) + request = _FakeFinishedRequest("req-no-ec") + + kv_params, ec_params = sched._free_request(request) + + assert ec_params is None + sched._free_blocks.assert_called_once_with(request) + + +def _make_finishing_request() -> Request: + request = Request( + request_id="req-ec-finish", + prompt_token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=4), + pooling_params=None, + arrival_time=100.0, + block_hasher=None, + ) + request.status = RequestStatus.RUNNING + return request + + +def _make_finish_sched(request: Request) -> MagicMock: + """MagicMock scheduler driving the real ``update_from_output`` through a + finishing request; ``_free_request`` reports EC transfer params.""" + sched = MagicMock() + sched.requests = {request.request_id: request} + sched.perf_metrics = None + sched.defer_block_free = False + sched.structured_output_manager.should_advance.return_value = False + sched._process_kv_transfer_trigger.return_value = False + sched._maybe_decode_pooling_output.return_value = None + sched._handle_stopped_request.return_value = True + sched._free_request.return_value = (None, EC_PARAMS) + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 + sched.chunk_transfer_adapter = None + sched.finished_req_ids_dict = None + sched._new_prompt_len_snapshot = {} + return sched + + +def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None: + request = _make_finishing_request() + sched = _make_finish_sched(request) + sched._update_request_with_output.return_value = ([42], True) + + engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request)) + + (eco,) = engine_core_outputs[request.client_index].outputs + assert eco.ec_transfer_params == EC_PARAMS + sched._free_request.assert_called_once_with(request) + + +def _finish_frame_parts(request: Request): + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {request.request_id: 1} + scheduler_output.total_num_scheduled_tokens = 1 + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {request.request_id: 0} + model_runner_output.routed_experts = None + return scheduler_output, model_runner_output + + +def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None: + request = _make_finishing_request() + # One-shot generation finish: the whole prompt is computed. + request.num_computed_tokens = len(request.prompt_token_ids) + sched = _make_finish_sched(request) + + engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request)) + + (eco,) = engine_core_outputs[request.client_index].outputs + assert eco.ec_transfer_params == EC_PARAMS + sched._free_request.assert_called_once_with(request) diff --git a/tests/core/sched/test_omni_sched_prefill_stats_finalize.py b/tests/core/sched/test_omni_sched_prefill_stats_finalize.py new file mode 100644 index 000000000..d436367db --- /dev/null +++ b/tests/core/sched/test_omni_sched_prefill_stats_finalize.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review). + +Upstream v0.28 finalizes prefill statistics with +``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))`` +before emitting them, which computes ``num_cache_creation_tokens`` (tokens +computed and written to the prefix cache). Both Omni schedulers previously +emitted ``take_prefill_stats()`` raw, so creation usage was always 0. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +# Imports must run in this order: vllm_omni applies patches to vllm.v1.request +# before Request / RequestStatus are bound in this module. +# isort: off +import vllm_omni # noqa: F401 - import for side effects (patch vLLM) +from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +# isort: on + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_request() -> Request: + request = Request( + request_id="req-prefill-stats", + prompt_token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=4), + pooling_params=None, + arrival_time=100.0, + block_hasher=None, + ) + request.status = RequestStatus.RUNNING + # A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache. + request.prefill_stats.set( + num_prompt_tokens=3, + num_local_cached_tokens=1, + num_external_cached_tokens=0, + ) + return request + + +def _make_sched(request: Request) -> MagicMock: + sched = MagicMock() + sched.requests = {request.request_id: request} + sched.perf_metrics = None + sched.defer_block_free = False + sched.structured_output_manager.should_advance.return_value = False + sched._process_kv_transfer_trigger.return_value = False + sched._maybe_decode_pooling_output.return_value = None + sched._handle_stopped_request.return_value = True + # The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created. + sched.kv_cache_manager.estimate_cached_tokens.return_value = 3 + + def _free_and_invalidate(req): + # Lifecycle guard: freeing releases the KV blocks, after which the + # manager reports 0 cached tokens. If finalize ran after + # _free_request (the reviewed bug), creation count would compute + # against 0 and the assertions below would see 0, not 2. + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 + return (None, None) + + sched._free_request.side_effect = _free_and_invalidate + sched.chunk_transfer_adapter = None + sched.finished_req_ids_dict = None + sched._new_prompt_len_snapshot = {} + return sched + + +def _frame(request: Request): + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {request.request_id: 1} + scheduler_output.total_num_scheduled_tokens = 1 + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {request.request_id: 0} + model_runner_output.routed_experts = None + return scheduler_output, model_runner_output + + +def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None: + request = _make_request() + sched = _make_sched(request) + sched._update_request_with_output.return_value = ([42], True) + + engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request)) + + (eco,) = engine_core_outputs[request.client_index].outputs + assert eco.prefill_stats is not None + assert eco.prefill_stats.num_cached_tokens == 1 + assert eco.prefill_stats.num_cache_creation_tokens == 2 + sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request) + + +def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None: + request = _make_request() + request.num_computed_tokens = len(request.prompt_token_ids) + sched = _make_sched(request) + + engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request)) + + (eco,) = engine_core_outputs[request.client_index].outputs + assert eco.prefill_stats is not None + assert eco.prefill_stats.num_cached_tokens == 1 + assert eco.prefill_stats.num_cache_creation_tokens == 2 + sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request) diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py index fbfeb384f..770c85578 100644 --- a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py +++ b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py @@ -157,3 +157,71 @@ def test_finish_requests_leaves_healthy_running_intact( ) assert scheduler.running == [alive] + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_preserves_unrelated_resumable_segment( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + paused = _StubRequest("req-paused", RequestStatus.FINISHED_STOPPED) + paused.resumable = True + leaving = _StubRequest("req-leaving", RequestStatus.FINISHED_STOPPED) + scheduler = _make_scheduler( + scheduler_cls, + requests={paused.request_id: paused}, + running=[paused], + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + return [leaving] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + scheduler_cls.finish_requests(scheduler, [leaving.request_id], RequestStatus.FINISHED_ABORTED) + + assert scheduler.running == [paused] + assert scheduler.requests == {paused.request_id: paused} + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_materializes_single_pass_request_ids( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + request = _StubRequest("req-generator", RequestStatus.RUNNING) + scheduler = _make_scheduler( + scheduler_cls, + requests={request.request_id: request}, + running=[request], + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + assert tuple(request_ids) == (request.request_id,) + return [request] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + result = scheduler_cls.finish_requests( + scheduler, + (request_id for request_id in [request.request_id]), + RequestStatus.FINISHED_ABORTED, + ) + + assert result == [request] + + +@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler]) +def test_finish_requests_does_not_reopen_off_queue_deferred_free_terminal(scheduler_cls) -> None: + terminal = _StubRequest("req-deferred-free", RequestStatus.FINISHED_STOPPED) + terminal.resumable = True + scheduler = _make_scheduler(scheduler_cls, requests={terminal.request_id: terminal}, running=[], waiting=[]) + scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed") + + assert scheduler_cls.finish_requests(scheduler, [terminal.request_id], RequestStatus.FINISHED_ABORTED) == [] + assert scheduler.requests == {terminal.request_id: terminal} + assert terminal.status == RequestStatus.FINISHED_STOPPED diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py index 9aaa4c656..bb9fd8859 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_shared.py +++ b/tests/core/sched/test_omni_scheduler_mixin_shared.py @@ -4,6 +4,7 @@ from types import SimpleNamespace import pytest from vllm.v1.engine import FinishReason +from vllm_omni.core.sched import omni_scheduler_mixin from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.output import OmniChunkRecvHandle @@ -14,6 +15,77 @@ class _Scheduler(OmniSchedulerMixin): pass +def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch): + created_adapters = [] + + def adapter_factory(config): + adapter = SimpleNamespace(vllm_config=config) + created_adapters.append(adapter) + return adapter + + monkeypatch.setattr(omni_scheduler_mixin, "OmniChunkTransferAdapter", adapter_factory) + + def init_scheduler(**model_config): + scheduler = _Scheduler() + scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(**model_config)) + scheduler._init_omni_io_scheduling_state() + return scheduler + + producer = init_scheduler( + stage_id=0, + async_chunk=True, + requires_full_payload_input=False, + custom_process_next_stage_input_func="test.pipeline.produce_async_chunk", + ) + receiver = init_scheduler( + stage_id=1, + async_chunk=True, + requires_full_payload_input=True, + custom_process_next_stage_input_func=None, + ) + terminal_stage = init_scheduler( + stage_id=2, + async_chunk=True, + requires_full_payload_input=False, + custom_process_next_stage_input_func=None, + ) + + assert producer.chunk_transfer_adapter is not None + assert receiver.chunk_transfer_adapter is not None + assert terminal_stage.chunk_transfer_adapter is not None + assert receiver.input_coordinator is None + assert len(created_adapters) == 3 + + +@pytest.mark.parametrize( + ("stage_id", "async_chunk", "required", "enabled"), + [ + (0, False, True, False), + (1, True, True, False), + (1, False, False, False), + (1, False, True, True), + ], +) +def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, async_chunk, required, enabled): + monkeypatch.setattr( + omni_scheduler_mixin, + "OmniChunkTransferAdapter", + lambda _config: SimpleNamespace(), + ) + scheduler = _Scheduler() + scheduler.vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + stage_id=stage_id, + async_chunk=async_chunk, + requires_full_payload_input=required, + ) + ) + + scheduler._init_omni_io_scheduling_state() + + assert (scheduler.input_coordinator is not None) is enabled + + def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): calls = [] scheduler = _Scheduler() @@ -22,13 +94,25 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): scheduler.requests = {"request": object()} scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode)) scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",)) + + def _collect_timed_out(timeout_s): + calls.append(("chunk-timeouts", timeout_s)) + return set() + + def _collect_failed_sends(): + calls.append(("failed-sends",)) + return {} + scheduler.chunk_transfer_adapter = SimpleNamespace( + receives_chunks=True, process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append( ("process", waiting, running, scheduler_requests) ), restore_queues=lambda waiting, running, scheduler_requests: calls.append( ("restore-chunks", waiting, running, scheduler_requests) ), + collect_timed_out_request_ids=_collect_timed_out, + collect_failed_send_request_ids=_collect_failed_sends, ) scheduler.input_coordinator = SimpleNamespace( restore_queues=lambda waiting: calls.append(("restore-full", waiting)) @@ -41,6 +125,10 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): ("consume", "ar"), ("timeouts",), ("process", scheduler.waiting, scheduler.running, scheduler.requests), + # The chunk deadline runs after chunks are applied, so a chunk that + # arrived this cycle resets the clock before it is measured (R1.1). + ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S), + ("failed-sends",), ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests), ("restore-full", scheduler.waiting), ] diff --git a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py index 5fad4340c..e7d126cd9 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py +++ b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit coverage for _process_pending_input_timeouts. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit coverage for the two stage-input timeout hooks. + +``_process_pending_input_timeouts`` covers the full-payload transport; +``_process_pending_chunk_timeouts`` covers async-chunk. Both share +``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``. Verifies that the mixin correctly *delegates* timed-out requests to the base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``. @@ -12,6 +16,10 @@ this file only asserts the wiring from the mixin to that API. from __future__ import annotations +import importlib +import logging +import math +import os from types import SimpleNamespace import pytest @@ -102,5 +110,280 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch): assert scheduler.finish_calls == [] +class _FakeChunkAdapter: + """Stands in for OmniChunkTransferAdapter's timeout-collection surface.""" + + def __init__(self, timed_out_ids, *, receives_chunks=True): + self._timed_out_ids = set(timed_out_ids) + self.receives_chunks = receives_chunks + self.calls = [] + + def collect_timed_out_request_ids(self, timeout_s): + self.calls.append(timeout_s) + return set(self._timed_out_ids) + + +class _FakeChunkScheduler(OmniSchedulerMixin): + def __init__(self, requests, adapter): + self.requests = requests + self.chunk_transfer_adapter = adapter + self.finish_calls = [] + + def finish_requests(self, req_ids, status): + self.finish_calls.append((set(req_ids), status)) + + +def test_process_pending_chunk_timeouts_delegates_to_finish_requests(): + """A request stalled in WAITING_FOR_CHUNK is failed, not left parked. + + This is the async-chunk half of the safety net (RFC #4855 R1.1); before it + existed, the chunk path had no deadline at all and issue #3833 could park a + request forever. + """ + req_id = "stalled-stream" + adapter = _FakeChunkAdapter(timed_out_ids={req_id}) + scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter) + + scheduler._process_pending_chunk_timeouts() + + assert len(adapter.calls) == 1 + assert adapter.calls[0] > 0 + assert len(scheduler.finish_calls) == 1 + finished_ids, status = scheduler.finish_calls[0] + assert finished_ids == {req_id} + assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR") + + +def test_process_pending_chunk_timeouts_skips_already_freed_request(): + adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"}) + scheduler = _FakeChunkScheduler(requests={}, adapter=adapter) + + scheduler._process_pending_chunk_timeouts() + + assert scheduler.finish_calls == [] + + +def test_process_pending_chunk_timeouts_noop_on_a_producer_stage(): + """Stage 0 does not receive chunks, so it has nothing to expire.""" + adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False) + scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter) + + scheduler._process_pending_chunk_timeouts() + + assert adapter.calls == [] + assert scheduler.finish_calls == [] + + +def test_process_pending_chunk_timeouts_noop_without_adapter(): + class _NoAdapter(OmniSchedulerMixin): + def __init__(self): + self.requests = {} + self.chunk_transfer_adapter = None + self.finish_calls = [] + + def finish_requests(self, req_ids, status): + self.finish_calls.append((set(req_ids), status)) + + scheduler = _NoAdapter() + scheduler._process_pending_chunk_timeouts() + assert scheduler.finish_calls == [] + + +def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch): + """The shared knob disables both transports' nets together.""" + from vllm_omni.core.sched import omni_scheduler_mixin + + monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0) + + adapter = _FakeChunkAdapter(timed_out_ids={"r1"}) + scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter) + scheduler._process_pending_chunk_timeouts() + + assert adapter.calls == [] + assert scheduler.finish_calls == [] + + +class _FakeSendFailureAdapter: + """Stands in for the adapter's drained send-failure map (R1.2).""" + + def __init__(self, failures): + self._failures = dict(failures) + self.collect_calls = 0 + + def collect_failed_send_request_ids(self): + self.collect_calls += 1 + failures, self._failures = self._failures, {} + return failures + + +class _FakeSendFailureScheduler(OmniSchedulerMixin): + def __init__(self, requests, adapter): + self.requests = requests + self.chunk_transfer_adapter = adapter + + +def test_log_failed_chunk_sends_reports_a_live_request(caplog): + adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"}) + scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter) + + with caplog.at_level(logging.ERROR): + scheduler._log_failed_chunk_sends() + + assert "req-live" in caplog.text + assert "No space left on device" in caplog.text + + +def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog): + """``collect_*`` drains the map, so a failure for a request that already left + this scheduler is the one place a drop could still go unrecorded.""" + adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"}) + scheduler = _FakeSendFailureScheduler({}, adapter) + + with caplog.at_level(logging.WARNING): + scheduler._log_failed_chunk_sends() + + assert adapter.collect_calls == 1 + assert "req-gone" in caplog.text + assert "No space left on device" in caplog.text + assert "finished or aborted" in caplog.text + + +def test_log_failed_chunk_sends_noop_without_adapter(): + scheduler = _FakeSendFailureScheduler({}, None) + scheduler._log_failed_chunk_sends() + + +class _FakeReceiveFailureAdapter: + def __init__(self, failures): + self._failures = dict(failures) + + def collect_failed_receive_request_ids(self): + failures, self._failures = self._failures, {} + return failures + + +class _FakeReceiveFailureScheduler(OmniSchedulerMixin): + def __init__(self, requests, adapter): + self.requests = requests + self.chunk_transfer_adapter = adapter + self.finish_calls = [] + + def finish_requests(self, req_ids, status): + self.finish_calls.append((set(req_ids), status)) + + +def test_process_chunk_receive_failures_delegates_to_finish_requests(): + request_id = "invalid-chunk" + adapter = _FakeReceiveFailureAdapter({request_id: "prompt exceeds limit"}) + scheduler = _FakeReceiveFailureScheduler({request_id: SimpleNamespace(request_id=request_id)}, adapter) + + scheduler._process_chunk_receive_failures() + + assert len(scheduler.finish_calls) == 1 + finished_ids, status = scheduler.finish_calls[0] + assert finished_ids == {request_id} + assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR") + + +# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of +# these reloads the module under a different environment. One knob now gates the +# deadline on both transports, which is why a silent 0 matters more than it did +# when only the full-payload path used it. + + +@pytest.fixture +def reload_with(monkeypatch): + """Reload the mixin module under a given env value, then put it back. + + ``importlib.reload`` re-executes into the live module dict, so whatever the + last test parsed is what every later test in the session sees -- ``-k + zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and + silently disable both deadlines. The rejection cases are worse: the + ValueError fires *after* the rebind, so ``nan`` lingers in the module. + ``monkeypatch`` restores the environment but not the module, so reload once + more on the way out. + """ + import vllm_omni.core.sched.omni_scheduler_mixin as mod + + def _reload(raw): + if raw is None: + monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False) + else: + monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw) + return importlib.reload(mod) + + yield _reload + + monkeypatch.undo() + importlib.reload(mod) + + +@pytest.mark.parametrize( + "raw,expected", + [ + (None, 600.0), + ("30", 30.0), + ("0.5", 0.5), + ("0", 0.0), # explicit opt-out is still honoured + ], +) +def test_timeout_accepts_valid_values(reload_with, raw, expected): + mod = reload_with(raw) + assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected + + +@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"]) +def test_negative_timeout_is_rejected(reload_with, raw): + """`-1` is a common "no limit" idiom; substituting 600 would give an operator + who meant "never time out" mysterious failures ten minutes in. Fail loudly.""" + with pytest.raises(ValueError, match="negative"): + reload_with(raw) + + +@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"]) +def test_non_finite_timeout_is_rejected(reload_with, raw): + """nan compares false against every deadline check and inf is unreachable; + both disable the net while looking like a number.""" + with pytest.raises(ValueError, match="not finite"): + reload_with(raw) + + +def test_zero_warns_and_names_both_transports(reload_with, caplog): + with caplog.at_level(logging.WARNING): + mod = reload_with("0") + assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0 + text = caplog.text + assert "DISABLED" in text + assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text + + +def test_unparsable_timeout_falls_back(reload_with): + mod = reload_with("not-a-number") + assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0 + + +def test_a_rejected_value_does_not_outlive_its_test(monkeypatch): + """The reason ``reload_with`` reloads on the way out, done here by hand. + + ``nan`` is bound before the ValueError fires, so the reload that rejects it + still leaves it in the module -- and nan compares false against every + deadline check, which is exactly the silent-disable this item exists to + stop. Restoring the environment is not enough; the module has to be reloaded + too. + """ + import vllm_omni.core.sched.omni_scheduler_mixin as mod + + ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600")) + + monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan") + with pytest.raises(ValueError, match="not finite"): + importlib.reload(mod) + assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S) + + monkeypatch.undo() + importlib.reload(mod) + assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py new file mode 100644 index 000000000..6f44aadac --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py @@ -0,0 +1,168 @@ +"""Tests for the ``num_waiting_for_streaming_input`` resync placed in +``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests`` +and in ``OmniARScheduler.schedule``. + +Upstream keeps that counter incrementally (``+1`` when a resumable request +parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the +next update arrives or the request is finished) and +``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting +queues. Omni rewrites ``request.status`` outside both hooks -- the +chunk-transfer adapter's park/restore and +``_realign_request_status_to_queues`` -- so a counted request can leave +``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the +counter for the lifetime of the process. + +An inflated counter does not degrade throughput, it hangs the stage: +``EngineCore.has_work()`` reads false while a live request sits in +``waiting``, so the engine blocks in ``input_queue.get()`` and never calls +``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the +server answered the first conversation after a boot and was silent from the +second on while still reporting healthy. +""" + +from __future__ import annotations + +import pytest +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod +import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _StubRequest: + def __init__(self, request_id: str, status: RequestStatus) -> None: + self.request_id = request_id + self.status = status + + def is_finished(self) -> bool: + return RequestStatus.is_finished(self.status) + + +_SCHEDULER_PARAMS = [ + pytest.param(OmniARScheduler, ar_sched_mod, id="ar"), + pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"), +] + + +def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None): + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.requests = requests + scheduler.running = running + scheduler.waiting = waiting + scheduler.skipped_waiting = skipped if skipped is not None else [] + scheduler.num_waiting_for_streaming_input = counter + return scheduler + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_clears_counter_left_by_a_stomped_status( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """The abort case that wedged the talker. + + The request was counted while parked for streaming input, then had its + status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's + ``finish_requests`` took the running branch and never decremented. With + no parked requests left, the counter must land at 0. + """ + aborted = _StubRequest("req-aborted", RequestStatus.RUNNING) + scheduler = _make_scheduler( + scheduler_cls, + requests={}, # upstream already dropped it + running=[], + waiting=[], + counter=1, + ) + + monkeypatch.setattr( + scheduler_mod.VLLMScheduler, + "finish_requests", + lambda self, request_ids, finished_status: [aborted], + ) + + scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED) + + assert scheduler.num_waiting_for_streaming_input == 0 + + +def test_resync_counts_parked_requests_in_skipped_waiting() -> None: + """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``, + so the resync has to look at both or it under-counts and reintroduces the + spin the counter exists to prevent.""" + parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ) + scheduler = _make_scheduler( + OmniARScheduler, + requests={parked.request_id: parked}, + running=[], + waiting=[], + counter=0, + skipped=[parked], + ) + + scheduler._resync_streaming_input_counter() + + assert scheduler.num_waiting_for_streaming_input == 1 + + +def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None: + """Pins placement in ``schedule``: the chunk-transfer adapter rewrites + status during ``process_pending_chunks``, so the resync has to run after + that and before upstream schedules, while ``self.waiting`` is still the + real queue.""" + parked = _StubRequest("req-resumed", RequestStatus.RUNNING) + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.requests = {parked.request_id: parked} + scheduler.running = [parked] + scheduler.waiting = [] + scheduler.skipped_waiting = [] + scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite + scheduler.policy = "fcfs" + + seen: dict[str, int] = {} + + monkeypatch.setattr( + ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None + ) + monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None) + monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False) + + def fake_schedule(self, throttle_prefills=False): + seen["counter"] = self.num_waiting_for_streaming_input + raise _StopScheduleError + + class _StopScheduleError(Exception): + pass + + monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule) + + with pytest.raises(_StopScheduleError): + OmniARScheduler.schedule(scheduler) + + assert seen["counter"] == 0 + + +def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -> None: + aborted = _StubRequest("req-generation-aborted", RequestStatus.FINISHED_ABORTED) + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.waiting.add_request(aborted) + scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.skipped_waiting.add_request(aborted) + scheduler.running = [aborted] + + scheduler._drop_aborted_queued_requests() + + assert list(scheduler.waiting) == [] + assert list(scheduler.skipped_waiting) == [] + assert scheduler.running == [] diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index 444d21c34..ae1b3924e 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -20,7 +20,6 @@ import torch import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod from vllm_omni.core.sched.omni_scheduling_coordinator import ( OmniSchedulingCoordinator, - uses_full_payload_input_coordinator, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -96,109 +95,6 @@ class MockQueue: # ------------------------------------------------------------------ # -class TestFullPayloadCoordinatorSelection(unittest.TestCase): - """Tests for the (model_arch, model_stage) whitelist gate. - - The init_omni_connectors arch allowlist is keyed by ``model_arch`` and - is a superset of the stages registered here -- consumer-wait stages - must be registered explicitly in ``_FULL_PAYLOAD_INPUT_STAGES``, while - the init allowlist covers both producer- and consumer-side runners. - These tests pin which ``(arch, stage)`` pairs the gate fires for today. - """ - - # Expected whitelist (model_arch, model_stage). Hardcoded to avoid the - # tautology of importing _FULL_PAYLOAD_INPUT_STAGES and asserting it - # against itself; any drift between this matrix and the whitelist will - # fail loudly here. - EXPECTED_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset( - { - ("Qwen3OmniMoeForConditionalGeneration", "talker"), - ("Qwen3OmniMoeForConditionalGeneration", "code2wav"), - ("Qwen2_5OmniForConditionalGeneration", "talker"), - ("Qwen2_5OmniForConditionalGeneration", "code2wav"), - ("CovoAudioForConditionalGeneration", "code2wav"), - ("MiMoAudioModel", "code2wav"), - ("Qwen3TTSCode2Wav", "code2wav"), - ("MiniCPMO45Code2Wav", "code2wav"), - # nemotron_voicechat: sync-path RVQ-VAE decode waits for the - # talker's full [frames, 31] code-stack payload. - ("NemotronVoiceChatCode2Wav", "code2wav"), - ("CosyVoice3Model", "cosyvoice3_code2wav"), - # audex: sync-path TTS decoder (async_chunk: false deploys) and - # the TTA XCodec1 stage (always sync full-payload). - ("AudexCode2Wav", "audex_code2wav"), - ("AudexXCodec1", "audex_xcodec"), - ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"), - ("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"), - ("DyninOmniForConditionalGeneration", "token2image"), - ("DyninOmniForConditionalGeneration", "token2audio"), - } - ) - - def test_whitelist_matches_expected_matrix(self): - """_FULL_PAYLOAD_INPUT_STAGES must equal the hardcoded expected matrix. - - Catches both accidental additions (which would silently enable the - consumer-wait gate for a new arch) and accidental removals (which - would silently disable an enabled arch). - """ - from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES - - self.assertEqual( - frozenset(_FULL_PAYLOAD_INPUT_STAGES), - self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES, - msg="_FULL_PAYLOAD_INPUT_STAGES drifted from the expected matrix; " - "update EXPECTED_FULL_PAYLOAD_INPUT_STAGES if intentional.", - ) - - def test_all_whitelisted_arch_stage_pairs_fire_gate(self): - """Every (arch, stage) pair in the expected matrix must fire - the gate when stage_id > 0 and async_chunk=False. - """ - for arch, stage in self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES: - model_config = SimpleNamespace( - stage_id=1, - async_chunk=False, - model_arch=arch, - model_stage=stage, - ) - self.assertTrue( - uses_full_payload_input_coordinator(model_config), - msg=f"expected gate to fire for {arch}/{stage}", - ) - - def test_other_arch_or_stage_or_mode_does_not_fire(self): - cases = [ - SimpleNamespace( - stage_id=1, async_chunk=True, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="talker" - ), - SimpleNamespace( - stage_id=0, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="thinker" - ), - SimpleNamespace( - stage_id=1, - async_chunk=False, - model_arch="Qwen3OmniMoeForConditionalGeneration", - model_stage="some_future_stage", - ), - SimpleNamespace( - stage_id=1, async_chunk=False, model_arch="Qwen3TTSForConditionalGeneration", model_stage="code2wav" - ), - SimpleNamespace( - stage_id=1, async_chunk=False, model_arch="MingFlashOmniForConditionalGeneration", model_stage="talker" - ), - SimpleNamespace(stage_id=1, async_chunk=False, model_arch=None, model_stage="talker"), - SimpleNamespace( - stage_id=1, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage=None - ), - ] - for model_config in cases: - self.assertFalse( - uses_full_payload_input_coordinator(model_config), - msg=f"expected gate OFF for {model_config}", - ) - - class TestCoordinatorUpdateRequestMetadata(unittest.TestCase): """Test update_request_metadata applies scheduling metadata to requests.""" diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index 7e3ab77c7..864fb7ab3 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -105,7 +105,7 @@ def create_unique_server_pytest_params( configs: list[dict[str, Any]], stage_configs_dir: Path, ) -> list[Any]: - """Like :func:`create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks.""" + """Like :func:`_create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks.""" marks_by_name = _marks_by_test_name(configs) return [ pytest.param( @@ -113,7 +113,7 @@ def create_unique_server_pytest_params( marks=marks_by_name.get(row[0], []), id=row[0], ) - for row in create_unique_server_params(configs, stage_configs_dir) + for row in _create_unique_server_params(configs, stage_configs_dir) ] @@ -150,7 +150,7 @@ def create_paired_omni_benchmark_pytest_params( """Paired params for ``run_benchmark.py`` (omni/tts).""" mapping = create_test_parameter_mapping(configs) marks_by_name = _marks_by_test_name(configs) - server_entries = [(row, row[0]) for row in create_unique_server_params(configs, stage_configs_dir)] + server_entries = [(row, row[0]) for row in _create_unique_server_params(configs, stage_configs_dir)] params_by_test_name = { test_name: get_benchmark_params_for_server(test_name, mapping) for _, test_name in server_entries } @@ -232,7 +232,7 @@ def _build_serve_args(serve_args: Any) -> list[str]: return args -def create_unique_server_params( +def _create_unique_server_params( configs: list[dict[str, Any]], stage_configs_dir: Path, ) -> list[tuple[str, str, str | None, str | None, tuple[str, ...], bool]]: @@ -311,7 +311,7 @@ def create_reliability_omni_server_params( configs: list[dict[str, Any]], stage_configs_dir: Path ) -> list[OmniServerParams]: adjusted_configs = configs_with_platform_stage_configs(configs) - unique_params = create_unique_server_params(adjusted_configs, stage_configs_dir) + unique_params = _create_unique_server_params(adjusted_configs, stage_configs_dir) server_args_by_name = extract_server_args_by_test_name(adjusted_configs) return [ OmniServerParams( @@ -653,6 +653,7 @@ def run_benchmark( random_input_len: Any | None = None, random_output_len: Any | None = None, resource_label: str | None = None, + num_warmups: int = 2, ) -> dict[str, Any]: """Run one ``vllm bench serve --omni`` iteration and return parsed metrics. @@ -677,7 +678,7 @@ def run_benchmark( + args + [ "--num-warmups", - "2", + str(num_warmups), "--save-result", "--result-dir", os.environ.get("BENCHMARK_DIR", "tests"), diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 03dfac1a0..64285dca3 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -1,4 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import json +import math import os import threading from collections.abc import Callable @@ -159,6 +163,18 @@ def benchmark_params(request): def assert_result(result, params, num_prompt) -> None: assert result["completed"] == num_prompt, "Request failures exist" + baseline = params.get("baseline") + hardware = result.get("Hardware") + hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None + if isinstance(hardware_baseline, dict) and "mean_tpot_ms" in hardware_baseline: + num_tpot_samples = result.get("num_tpot_samples") + mean_tpot_ms = result.get("mean_tpot_ms") + assert isinstance(num_tpot_samples, int) and not isinstance(num_tpot_samples, bool) and num_tpot_samples > 0, ( + "TPOT baseline is configured, but no measurable TPOT samples were produced" + ) + assert ( + isinstance(mean_tpot_ms, int | float) and not isinstance(mean_tpot_ms, bool) and math.isfinite(mean_tpot_ms) + ), "TPOT baseline is configured, but mean_tpot_ms is not finite" expected_audio_turns = params.get("expected_duplex_audio_turns_per_session") if expected_audio_turns is not None: session_metrics = result.get("duplex_session_metrics") @@ -279,5 +295,6 @@ def test_performance_benchmark(omni_server, benchmark_params): random_input_len=params.get("random_input_len"), random_output_len=params.get("random_output_len"), resource_label=resource_label, + num_warmups=max(2, int(concurrency)), ) assert_result(result, params, num_prompt) diff --git a/tests/dfx/perf/tests/test_bagel_vllm_omni.json b/tests/dfx/perf/tests/test_bagel_vllm_omni.json index 9f6ff0241..1189ddd9d 100644 --- a/tests/dfx/perf/tests/test_bagel_vllm_omni.json +++ b/tests/dfx/perf/tests/test_bagel_vllm_omni.json @@ -10,7 +10,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task", @@ -43,8 +43,8 @@ "H100": { "throughput_qps": 0.495, "latency_mean": 2.0238, - "peak_memory_mb_max": 29393.4286, - "peak_memory_mb_mean": 29392.7143 + "peak_memory_mb_max": 29392.0, + "peak_memory_mb_mean": 29392.0 } } } @@ -61,7 +61,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task", @@ -93,10 +93,10 @@ "num-input-images": 1, "baseline": { "H100": { - "throughput_qps": 0.3451, - "latency_mean": 2.9031, - "peak_memory_mb_max": 30480.2857, - "peak_memory_mb_mean": 30479.2857 + "throughput_qps": 0.3499, + "latency_mean": 2.8701, + "peak_memory_mb_max": 30466.0, + "peak_memory_mb_mean": 30465.0 } } } @@ -113,7 +113,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task", @@ -164,7 +164,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task", @@ -196,8 +196,8 @@ "num-input-images": 1, "baseline": { "H100": { - "throughput_qps": 0.149, - "latency_mean": 6.7157, + "throughput_qps": 0.1533, + "latency_mean": 6.5243, "peak_memory_mb_max": 30574.0, "peak_memory_mb_mean": 30574.0 } diff --git a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json index c4d2305e6..bcd649de1 100644 --- a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json +++ b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json @@ -49,8 +49,8 @@ }, "baseline": { "H100": { - "throughput_qps": 1.107, - "latency_mean": 0.8929, + "throughput_qps": 1.1958, + "latency_mean": 0.843, "peak_memory_mb_mean": 80000.0 } } @@ -118,9 +118,9 @@ }, "baseline": { "H100": { - "throughput_qps": 0.027, - "latency_mean": 36.6824, - "peak_memory_mb_mean": 30155.4 + "throughput_qps": 0.0307, + "latency_mean": 32.7665, + "peak_memory_mb_mean": 28762.4762 } } } @@ -186,9 +186,9 @@ }, "baseline": { "H100": { - "throughput_qps": 0.027, - "latency_mean": 36.968, - "peak_memory_mb_mean": 30443.6 + "throughput_qps": 0.0305, + "latency_mean": 33.1009, + "peak_memory_mb_mean": 29441.1429 } } } @@ -258,9 +258,9 @@ }, "baseline": { "H100": { - "throughput_qps": 0.027, - "latency_mean": 37.2117, - "peak_memory_mb_mean": 30135.6 + "throughput_qps": 0.0299, + "latency_mean": 33.925, + "peak_memory_mb_mean": 29443.1429 } } } diff --git a/tests/dfx/perf/tests/test_higgs_audio_v3.json b/tests/dfx/perf/tests/test_higgs_audio_v3.json index 75e5983b0..1197c88f1 100644 --- a/tests/dfx/perf/tests/test_higgs_audio_v3.json +++ b/tests/dfx/perf/tests/test_higgs_audio_v3.json @@ -10,7 +10,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "tts" ], "server_params": { diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json index c4ef34b96..098dcd5a0 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json +++ b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json @@ -137,9 +137,7 @@ "H200": { "throughput_qps": 0.1874, "latency_p99": 101.852, - "latency_mean": 79.0385, - "peak_memory_mb_max": 70000, - "peak_memory_mb_mean": 70000 + "latency_mean": 79.0385 } } } diff --git a/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json new file mode 100644 index 000000000..8186edf2c --- /dev/null +++ b/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json @@ -0,0 +1,106 @@ +[ + { + "test_name": "test_hunyuanvideo15_i2v_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v", + "serve_args": { + "flow-shift": 5.0, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "832x480_frames33_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 33, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "guidance_scale": 6.0, + "flow_shift": 5.0 + } + } + ] + }, + { + "test_name": "test_hunyuanvideo15_i2v_cachedit_tp2_vae_patch2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling", + "server_type": "vllm-omni", + "server_params": { + "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v", + "serve_args": { + "flow-shift": 5.0, + "cache-backend": "cache_dit", + "tensor-parallel-size": 2, + "vae-patch-parallel-size": 2, + "vae-use-tiling": true, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "832x480_frames33_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 33, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "guidance_scale": 6.0, + "flow_shift": 5.0 + } + } + ] + } +] diff --git a/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json b/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json new file mode 100644 index 000000000..3771c1c81 --- /dev/null +++ b/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json @@ -0,0 +1,104 @@ +[ + { + "test_name": "test_hunyuanvideo15_t2v_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", + "serve_args": { + "flow-shift": 5.0, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "832x480_frames33_steps4", + "dataset": "random", + "task": "t2v", + "num-prompts": 10, + "max-concurrency": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 33, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "guidance_scale": 6.0, + "flow_shift": 5.0 + } + } + ] + }, + { + "test_name": "test_hunyuanvideo15_t2v_cachedit_tp2_vae_patch2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling", + "server_type": "vllm-omni", + "server_params": { + "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", + "serve_args": { + "flow-shift": 5.0, + "cache-backend": "cache_dit", + "tensor-parallel-size": 2, + "vae-patch-parallel-size": 2, + "vae-use-tiling": true, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "832x480_frames33_steps4", + "dataset": "random", + "task": "t2v", + "num-prompts": 10, + "max-concurrency": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 33, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "guidance_scale": 6.0, + "flow_shift": 5.0 + } + } + ] + } +] diff --git a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json index 1a0f31e98..9b0634a20 100644 --- a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json +++ b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json @@ -10,7 +10,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Single-device baseline with enforce-eager (no torch.compile)", @@ -64,7 +64,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Single-device with torch.compile (default, no enforce-eager)", @@ -117,7 +117,7 @@ "num_cards": 2 } }, - "full_model", + "local_model", "diffusion" ], "description": "CFG-parallel=2 with enforce-eager", @@ -172,7 +172,7 @@ "num_cards": 2 } }, - "full_model", + "local_model", "diffusion" ], "description": "CFG-parallel=2 with torch.compile", @@ -226,7 +226,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "CacheDiT with enforce-eager", diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5.json b/tests/dfx/perf/tests/test_minicpmo_4_5.json index cfb57169b..cac1880c5 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5.json @@ -39,6 +39,7 @@ 8 ], "no_oversample": true, + "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", "extra_body": { @@ -55,33 +56,116 @@ "baseline": { "H100": { "request_throughput": [ - 0.7803, - 1.1008, - 1.2027 + 0.9731, + 1.1986, + 1.686 ], "mean_ttft_ms": [ - 166.9764, - 307.7133, - 323.584 + 152.5665, + 270.2559, + 301.455 ], "mean_e2el_ms": [ - 1361.1682, - 3800.6332, - 6961.8636 + 1027.2953, + 3306.3735, + 4686.5366 ], "mean_audio_ttfp_ms": [ - 635.8099, - 1615.6591, - 2888.9831 + 458.1553, + 2030.2335, + 2068.6233 ], "mean_audio_rtf": [ - 0.3151, - 0.8568, - 1.5425 + 0.2357, + 0.791, + 1.0324 + ] + }, + "A3": { + "request_throughput": [ + 0.5383, + 0.6042, + 0.7547 + ], + "mean_ttft_ms": [ + 333.2633, + 533.8656, + 547.3388 + ], + "mean_e2el_ms": [ + 1857.2154, + 6600.5101, + 10450.6811 + ], + "mean_audio_ttfp_ms": [ + 986.4666, + 3411.0754, + 3352.7515 + ], + "mean_audio_rtf": [ + 0.4423, + 1.5734, + 2.3024 ] } } } ] + }, + { + "test_name": "test_minicpmo_4_5_a2", + "mark": [ + { + "hardware_marks": { + "res": { + "npu": "A2" + }, + "num_cards": 1 + } + }, + "full_model", + "omni" + ], + "server_params": { + "model": "openbmb/MiniCPM-o-4_5", + "extra_cli_args": [ + "--deploy-config", + "vllm_omni/deploy/minicpmo_4_5.yaml", + "--trust-remote-code" + ] + }, + "benchmark_params": [ + { + "dataset_name": "seed-tts", + "dataset_path": "zhaochenyang20/seed-tts-eval", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8 + ], + "no_oversample": true, + "disable_shuffle": true, + "trust_remote_code": true, + "seed_tts_locale": "en", + "extra_body": { + "modalities": [ + "text", + "audio" + ], + "chat_template_kwargs": { + "enable_thinking": false, + "use_tts_template": true + } + }, + "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" + } + ] } ] diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json index cfbcc3abb..801743134 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json @@ -18,7 +18,7 @@ "model": "openbmb/MiniCPM-o-4_5", "extra_cli_args": [ "--deploy-config", - "vllm_omni/deploy/minicpmo_4_5_duplex.yaml", + "vllm_omni/deploy/minicpmo_4_5.yaml", "--trust-remote-code" ] }, @@ -35,6 +35,7 @@ 1 ], "no_oversample": true, + "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", "seed_tts_turns_per_session": 4, @@ -42,9 +43,26 @@ "extra_body": { "save_duplex_request_metrics": true }, - "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration", + "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration", "baseline": { "H100": { + "request_throughput": [ + 0.3038 + ], + "mean_ttft_ms": [ + 139.997 + ], + "mean_e2el_ms": [ + 3269.8807 + ], + "mean_audio_ttfp_ms": [ + 456.073 + ], + "mean_audio_rtf": [ + 0.2254 + ] + }, + "A3": { "request_throughput": [ 0.1594 ], diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json index e327be510..26dc54aab 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json @@ -37,6 +37,7 @@ 4 ], "no_oversample": true, + "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", "extra_body": { @@ -53,24 +54,24 @@ "baseline": { "H100": { "request_throughput": [ - 0.7803, - 1.1008 + 0.9731, + 1.1986 ], "mean_ttft_ms": [ - 166.9764, - 307.7133 + 152.5665, + 270.2559 ], "mean_e2el_ms": [ - 1361.1682, - 3800.6332 + 1027.2953, + 3306.3735 ], "mean_audio_ttfp_ms": [ - 635.8099, - 1615.6591 + 458.1553, + 2030.2335 ], "mean_audio_rtf": [ - 0.3151, - 0.8568 + 0.2357, + 0.791 ] } } diff --git a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json index a73abc5fb..6b0247aaf 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json @@ -21,54 +21,6 @@ ] }, "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 4, - 16, - 32, - 64, - 128 - ], - "max_concurrency": [ - 1, - 4, - 8, - 16, - 32 - ], - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "H100": { - "mean_ttft_ms": [ - 96.4061, - 140.8147, - 271.9025, - 362.2916, - 507.776 - ], - "mean_e2el_ms": [ - 18507.4582, - 28365.4956, - 31907.9613, - 48161.5829, - 72630.8601 - ], - "mean_audio_rtf": [ - 0.1175, - 0.18, - 0.2359, - 0.3278, - 0.6296 - ] - } - } - }, { "dataset_name": "random-mm", "backend": "openai-chat-omni", @@ -94,14 +46,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [ - 354.0951 - ], - "mean_e2el_ms": [ - 3952.5615 - ], "mean_audio_rtf": [ - 0.117 + 0.1098 ] } } @@ -133,14 +79,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [ - 259.6317 - ], - "mean_e2el_ms": [ - 4253.3009 - ], "mean_audio_rtf": [ - 0.1448 + 0.1335 ] } } @@ -174,60 +114,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [ - 457.4453 - ], - "mean_e2el_ms": [ - 5912.5623 - ], "mean_audio_rtf": [ - 0.1967 - ] - } - } - }, - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 4, - 16, - 32, - 64, - 128 - ], - "max_concurrency": [ - 1, - 4, - 8, - 16, - 32 - ], - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "H100": { - "mean_ttft_ms": [ - 90.8033, - 202.1779, - 326.5864, - 544.8914, - 961.3929 - ], - "mean_e2el_ms": [ - 4517.3411, - 7857.9711, - 9861.1687, - 13283.3363, - 19541.0736 + 0.1798 ] } } @@ -262,11 +150,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { "H100": { - "mean_ttft_ms": [ - 318.048 - ], "mean_e2el_ms": [ - 807.6114 + 756.8415 ] } } @@ -303,11 +188,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { "H100": { - "mean_ttft_ms": [ - 186.168 - ], "mean_e2el_ms": [ - 726.1137 + 694.9777 ] } } @@ -346,9 +228,6 @@ "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { "H100": { - "mean_ttft_ms": [ - 375.5282 - ], "mean_e2el_ms": [ 1006.1663 ] @@ -356,5 +235,61 @@ } } ] + }, + { + "test_name": "test_qwen3_omni_chunk_cuda", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "omni", + "slow" + ], + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "extra_cli_args": [ + "--async-chunk" + ] + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 4, + 32, + 128 + ], + "max_concurrency": [ + 1, + 8, + 32 + ], + "random_input_len": 100, + "random_output_len": 100, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "H100": { + "mean_e2el_ms": [ + 4354.8945, + 9861.1687, + 18716.709 + ] + } + } + } + ] } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json index a3d8aaea8..5453d6a90 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json @@ -10,7 +10,7 @@ "num_cards": 3 } }, - "full_model", + "slow", "omni" ], "server_params": { @@ -36,10 +36,9 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085], - "mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485], - "mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275], - "mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322] + "mean_tpot_ms": [7.6955, 10.6457, 12.7363, 15.1608], + "mean_audio_ttfp_ms": [325.1624, 602.5192, 699.3891, 1024.088], + "mean_audio_rtf": [0.1689, 0.2088, 0.2397, 0.2791] } } }, @@ -64,9 +63,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [340.5122], - "mean_audio_ttfp_ms": [463.0752], - "mean_audio_rtf": [0.1153] + "mean_audio_ttfp_ms": [436.4663], + "mean_audio_rtf": [0.1069] } } }, @@ -93,9 +91,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [236.926], - "mean_audio_ttfp_ms": [395.2182], - "mean_audio_rtf": [0.1259] + "mean_audio_ttfp_ms": [376.4911], + "mean_audio_rtf": [0.1145] } } }, @@ -124,9 +121,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [470.6126], - "mean_audio_ttfp_ms": [665.5474], - "mean_audio_rtf": [0.166] + "mean_audio_ttfp_ms": [638.2227], + "mean_audio_rtf": [0.1438] } } } diff --git a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json index 38dbfb2e3..d73629666 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json @@ -45,26 +45,12 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [ - 115.8348, - 173.1823, - 446.672, - 464.5985, - 3420.0634 - ], - "mean_e2el_ms": [ - 25070.9523, - 31992.0646, - 37619.7669, - 50393.8257, - 136619.2743 - ], "mean_audio_rtf": [ 0.169, - 0.2111, - 0.2475, - 0.3551, - 0.9266 + 0.2041, + 0.2382, + 0.3177, + 0.4595 ] } } @@ -94,14 +80,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [ - 343.1377 - ], - "mean_e2el_ms": [ - 4811.1123 - ], "mean_audio_rtf": [ - 0.1463 + 0.138 ] } } @@ -133,14 +113,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [ - 245.5446 - ], - "mean_e2el_ms": [ - 5019.1256 - ], "mean_audio_rtf": [ - 0.1713 + 0.1609 ] } } @@ -174,14 +148,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "H100": { - "mean_ttft_ms": [ - 506.1407 - ], - "mean_e2el_ms": [ - 6814.9656 - ], "mean_audio_rtf": [ - 0.2213 + 0.2048 ] } } diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json index 8775ee381..ef8d97a94 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json @@ -10,7 +10,7 @@ "num_cards": 2 } }, - "full_model", + "slow", "omni" ], "server_params": { @@ -29,16 +29,12 @@ "endpoint": "/v1/completions", "num_prompts": [ 4, - 16, 32, - 64, 128 ], "max_concurrency": [ 1, - 4, 8, - 16, 32 ], "random_input_len": 2500, @@ -52,19 +48,10 @@ "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { "H100": { - "mean_ttft_ms": [ - 69.8081, - 173.0651, - 271.8108, - 413.4079, - 497.2546 - ], "mean_e2el_ms": [ - 4720.8225, - 7937.2532, - 9962.3566, - 13724.5479, - 18324.091 + 4619.0068, + 9760.6175, + 18007.6918 ] } } @@ -100,53 +87,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { "H100": { - "mean_ttft_ms": [ - 25.9477 - ], "mean_e2el_ms": [ - 517.3155 - ] - } - } - }, - { - "dataset_name": "random-mm", - "backend": "vllm", - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "endpoint": "/v1/completions", - "num_prompts": [ - 40 - ], - "request_rate": [ - 0.5 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "random_mm_base_items_per_request": 2, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.5, - "(720, 1280, 2)": 0.5 - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "H100": { - "mean_ttft_ms": [ - 26.7056 - ], - "mean_e2el_ms": [ - 551.308 + 491.2534 ] } } @@ -186,11 +128,8 @@ "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { "H100": { - "mean_ttft_ms": [ - 27.6906 - ], "mean_e2el_ms": [ - 596.5284 + 559.6674 ] } } diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json index dcd2188c5..34d275ac9 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json @@ -10,7 +10,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Single-device baseline (two input images)", @@ -35,10 +35,10 @@ "enable-negative-prompt": true, "baseline": { "H100": { - "throughput_qps": 0.0721, - "latency_mean": 13.8398, - "peak_memory_mb_max": 57610.0, - "peak_memory_mb_mean": 57610.0 + "throughput_qps": 0.0684, + "latency_mean": 14.6249, + "peak_memory_mb_max": 57632.0, + "peak_memory_mb_mean": 57632.0 } } }, @@ -55,10 +55,10 @@ "enable-negative-prompt": true, "baseline": { "H100": { - "throughput_qps": 0.0186, - "latency_mean": 53.6942, - "peak_memory_mb_max": 67790.0, - "peak_memory_mb_mean": 67790.0 + "throughput_qps": 0.0177, + "latency_mean": 56.5285, + "peak_memory_mb_max": 67840.0, + "peak_memory_mb_mean": 67840.0 } } } @@ -76,7 +76,7 @@ "num_cards": 4 } }, - "full_model", + "local_model", "diffusion" ], "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4", @@ -105,10 +105,10 @@ "enable-negative-prompt": true, "baseline": { "H100": { - "throughput_qps": 0.0538, - "latency_mean": 18.5591, - "peak_memory_mb_max": 56449.4286, - "peak_memory_mb_mean": 56449.4286 + "throughput_qps": 0.0515, + "latency_mean": 19.4345, + "peak_memory_mb_max": 56348.0, + "peak_memory_mb_mean": 56348.0 } } } @@ -125,7 +125,7 @@ "num_cards": 4 } }, - "full_model", + "local_model", "diffusion" ], "description": "Ulysses SP=2 + CFG=2 + CacheDiT", @@ -164,10 +164,10 @@ "enable-negative-prompt": true, "baseline": { "H100": { - "throughput_qps": 0.2776, - "latency_mean": 3.5938, - "peak_memory_mb_max": 58113.7143, - "peak_memory_mb_mean": 58075.1429 + "throughput_qps": 0.2479, + "latency_mean": 4.0389, + "peak_memory_mb_max": 57750.0, + "peak_memory_mb_mean": 57750.0 } } }, @@ -184,10 +184,10 @@ "enable-negative-prompt": true, "baseline": { "H100": { - "throughput_qps": 0.1051, - "latency_mean": 9.4918, - "peak_memory_mb_max": 68107.4286, - "peak_memory_mb_mean": 68107.4286 + "throughput_qps": 0.0948, + "latency_mean": 10.5462, + "peak_memory_mb_max": 67764.0, + "peak_memory_mb_mean": 67764.0 } } } diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json deleted file mode 100644 index 686bd6160..000000000 --- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json +++ /dev/null @@ -1,190 +0,0 @@ -[ - { - "test_name": "test_qwen_image_edit_single_device", - "mark": [ - { - "hardware_marks": { - "res": { - "cuda": "H100" - }, - "num_cards": 1 - } - }, - "full_model", - "diffusion" - ], - "description": "Single-device baseline", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit", - "serve_args": { - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20_i2i", - "dataset": "random", - "task": "i2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "H100": { - "throughput_qps": 0.1129, - "latency_mean": 8.8352, - "peak_memory_mb_max": 58572.0, - "peak_memory_mb_mean": 58572.0 - } - } - }, - { - "name": "1536x1536_steps35_i2i", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "H100": { - "throughput_qps": 0.024, - "latency_mean": 41.5686, - "peak_memory_mb_max": 66368.0, - "peak_memory_mb_mean": 66368.0 - } - } - } - ] - }, - { - "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4", - "mark": [ - { - "hardware_marks": { - "res": { - "cuda": "H100" - }, - "num_cards": 4 - } - }, - "full_model", - "diffusion" - ], - "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "vae-patch-parallel-size": 4, - "vae-use-tiling": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "1536x1536_steps35_i2i", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "H100": { - "throughput_qps": 0.0672, - "latency_mean": 14.8517, - "peak_memory_mb_max": 57076.0, - "peak_memory_mb_mean": 57076.0 - } - } - } - ] - }, - { - "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit", - "mark": [ - { - "hardware_marks": { - "res": { - "cuda": "H100" - }, - "num_cards": 4 - } - }, - "full_model", - "diffusion" - ], - "description": "Ulysses SP=2 + CFG=2 + CacheDiT", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "cache-backend": "cache_dit", - "cache-config": { - "Fn_compute_blocks": 1, - "Bn_compute_blocks": 0, - "max_warmup_steps": 4, - "residual_diff_threshold": 0.24, - "max_continuous_cached_steps": 3, - "enable_taylorseer": false, - "taylorseer_order": 1, - "scm_steps_mask_policy": null, - "scm_steps_policy": "dynamic" - }, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20_i2i", - "dataset": "random", - "task": "i2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "H100": { - "throughput_qps": 0.2879, - "latency_mean": 3.4701, - "peak_memory_mb_max": 58758.0, - "peak_memory_mb_mean": 58758.0 - } - } - }, - { - "name": "1536x1536_steps35_i2i", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "H100": { - "throughput_qps": 0.1113, - "latency_mean": 8.9656, - "peak_memory_mb_max": 67424.0, - "peak_memory_mb_mean": 67424.0 - } - } - } - ] - } -] diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json index bc4d93a24..bc0f5c923 100644 --- a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -10,7 +10,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "diffusion" ], "description": "Single-device baseline", @@ -55,8 +55,8 @@ "H100": { "throughput_qps": 0.0417, "latency_mean": 23.914, - "peak_memory_mb_max": 63225.1429, - "peak_memory_mb_mean": 63225.1429 + "peak_memory_mb_max": 62912.0, + "peak_memory_mb_mean": 62912.0 } } } diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json index ebf9ccc29..4f64d6ec1 100644 --- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -36,7 +36,7 @@ "H100": { "throughput_qps": 0.4153, "latency_mean": 2.4065, - "peak_memory_mb_mean": 55121.4286 + "peak_memory_mb_mean": 55062.0 } } }, @@ -54,7 +54,7 @@ "H100": { "throughput_qps": 0.0423, "latency_mean": 23.5991, - "peak_memory_mb_mean": 65118.0 + "peak_memory_mb_mean": 64546.0 } } } @@ -74,13 +74,14 @@ "full_model", "diffusion" ], - "description": "Single-device baseline (no parallelism) with step execution", + "description": "Single-device step execution: sequential 512/1536 plus 512 concurrency sweep. One server (max-num-seqs 8).", "server_type": "vllm-omni", "server_params": { "model": "Qwen/Qwen-Image", "serve_args": { "enable-diffusion-pipeline-profiler": true, - "step-execution": true + "step-execution": true, + "max-num-seqs": 8 } }, "benchmark_params": [ @@ -98,7 +99,7 @@ "H100": { "throughput_qps": 0.4175, "latency_mean": 2.3953, - "peak_memory_mb_mean": 55088.5714 + "peak_memory_mb_mean": 54952.0 } } }, @@ -116,7 +117,28 @@ "H100": { "throughput_qps": 0.0423, "latency_mean": 23.5768, - "peak_memory_mb_mean": 66012.2857 + "peak_memory_mb_mean": 64500.0 + } + } + }, + { + "name": "512x512_steps20_high_concurrency", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": [20, 40, 80, 160], + "max-concurrency": [1, 2, 4, 8], + "warmup-requests": 8, + "warmup-concurrency": 8, + "warmup-num-inference-steps": 20, + "enable-negative-prompt": true, + "baseline": { + "H100": { + "throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752], + "latency_mean": [2.5466, 3.0235, 6.3948, 10.1944], + "peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325] } } } @@ -163,7 +185,7 @@ "H100": { "throughput_qps": 0.1214, "latency_mean": 8.2211, - "peak_memory_mb_mean": 54501.1429 + "peak_memory_mb_mean": 54474.0 } } } @@ -220,7 +242,7 @@ "H100": { "throughput_qps": 0.559, "latency_mean": 1.794, - "peak_memory_mb_mean": 55290.5714 + "peak_memory_mb_mean": 55174.0 } } }, @@ -238,58 +260,10 @@ "H100": { "throughput_qps": 0.1918, "latency_mean": 5.2022, - "peak_memory_mb_mean": 64838.8571 + "peak_memory_mb_mean": 64750.0 } } } ] - }, - { - "test_name": "test_qwen_image_single_device_step_execution_high_concurrency", - "mark": [ - { - "hardware_marks": { - "res": { - "cuda": "H100" - }, - "num_cards": 1 - } - }, - "full_model", - "diffusion" - ], - "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image", - "serve_args": { - "enable-diffusion-pipeline-profiler": true, - "step-execution": true, - "max-num-seqs": 8 - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20_high_concurrency", - "dataset": "random", - "task": "t2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": [20, 40, 80, 160], - "max-concurrency": [1, 2, 4, 8], - "warmup-requests": 8, - "warmup-concurrency": 8, - "warmup-num-inference-steps": 20, - "enable-negative-prompt": true, - "baseline": { - "H100": { - "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159], - "latency_mean": [2.4215, 2.992, 5.9293, 9.656], - "peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821] - } - } - } - ] } ] diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py index 74ad29cf6..2e98dcaf7 100644 --- a/tests/dfx/perf/tests/test_runner_metadata.py +++ b/tests/dfx/perf/tests/test_runner_metadata.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for DFX runner metadata field exclusion.""" import json @@ -105,7 +108,7 @@ def test_resolve_pytest_marks_hardware_dict_with_extra(): assert "H100" in names assert "cuda" in names assert "gpu" in names - assert "distributed_cuda" in names + assert "cards_2" in names assert "full_model" in names assert "diffusion" in names assert "local_model" in names @@ -512,3 +515,41 @@ def test_omni_duplex_expected_audio_turns_rejects_incomplete_session(): {"expected_duplex_audio_turns_per_session": 4}, 1, ) + + +def test_omni_tpot_baseline_accepts_measured_finite_sample(): + from tests.dfx.perf.scripts.run_benchmark import assert_result + + assert_result( + { + "completed": 1, + "Hardware": "H100", + "num_tpot_samples": 1, + "mean_tpot_ms": 10.0, + }, + {"baseline": {"H100": {"mean_tpot_ms": 20.0}}}, + 1, + ) + + +@pytest.mark.parametrize( + ("num_tpot_samples", "mean_tpot_ms", "match"), + [ + (0, float("nan"), "no measurable TPOT samples"), + (1, float("nan"), "mean_tpot_ms is not finite"), + ], +) +def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples, mean_tpot_ms, match): + from tests.dfx.perf.scripts.run_benchmark import assert_result + + with pytest.raises(AssertionError, match=match): + assert_result( + { + "completed": 1, + "Hardware": "H100", + "num_tpot_samples": num_tpot_samples, + "mean_tpot_ms": mean_tpot_ms, + }, + {"baseline": {"H100": {"mean_tpot_ms": 20.0}}}, + 1, + ) diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json index f398e3271..46aa5e83f 100644 --- a/tests/dfx/perf/tests/test_tts.json +++ b/tests/dfx/perf/tests/test_tts.json @@ -424,10 +424,10 @@ "baseline": { "H100": { "median_audio_ttfp_ms": [ - 82.1886 + 82.6344 ], "median_audio_rtf": [ - 0.1241 + 0.1277 ] } } @@ -453,13 +453,13 @@ "baseline": { "H100": { "median_audio_ttfp_ms": [ - 160.9168 + 170.3737 ], "median_audio_rtf": [ - 0.1907 + 0.2076 ], "audio_throughput": [ - 39.2332 + 37.0834 ] } } @@ -485,13 +485,13 @@ "baseline": { "H100": { "median_audio_ttfp_ms": [ - 100.1265 + 101.0293 ], "median_audio_rtf": [ - 0.1437 + 0.1498 ], "audio_throughput": [ - 11.7605 + 10.1186 ] } } diff --git a/tests/dfx/perf/tests/test_voxcpm2.json b/tests/dfx/perf/tests/test_voxcpm2.json index e5163e5bb..1756bc206 100644 --- a/tests/dfx/perf/tests/test_voxcpm2.json +++ b/tests/dfx/perf/tests/test_voxcpm2.json @@ -10,7 +10,7 @@ "num_cards": 1 } }, - "full_model", + "local_model", "tts" ], "server_params": { diff --git a/tests/dfx/reliability/README.md b/tests/dfx/reliability/README.md index 8dc69dc55..79d862025 100644 --- a/tests/dfx/reliability/README.md +++ b/tests/dfx/reliability/README.md @@ -7,12 +7,10 @@ The detailed RFC document is maintained outside this repository (local/internal - Qwen3 tests: `test_reliability_qwen3_omni.py` - Wan2.2 tests: `test_reliability_wan22.py` - HunyuanImage tests: `test_reliability_hunyuan_image.py` -- VoxCPM2 tests: `test_reliability_voxcpm2.py` ```bash pytest --collect-only tests/dfx/reliability pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow -pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow ``` diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py index ee1515fa6..f14dcf8c1 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py @@ -1,13 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """``POST /v1/audio/generate`` validation (live Stable Audio Open server).""" from __future__ import annotations +from typing import Any + import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient pytestmark = [pytest.mark.slow, pytest.mark.diffusion] @@ -68,10 +70,14 @@ _PARAMS = [ pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649), pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649), pytest.param( - {"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649 + {"num_inference_steps": 0}, + ("num_inference_steps", "greater than or equal to 1"), + id="num_inference_steps_zero", ), pytest.param( - {"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative" + {"num_inference_steps": -1}, + ("num_inference_steps", "greater than or equal to 1"), + id="num_inference_steps_negative", ), pytest.param( {"num_inference_steps": 6000}, @@ -95,20 +101,20 @@ _PARAMS = [ @pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) def test_audio_generate_invalid_field_values( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, overrides: dict[str, object], err_message: str | tuple[str, ...], ) -> None: - body = {"model": omni_server.model, "input": "ambient electronic pad"} + body: dict[str, Any] = {"model": omni_server.model, "input": "ambient electronic pad"} body.update(overrides) - openai_client.send_audio_generate_http_request( + online_client.send_audio_generate_http_request( {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} ) @pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) -def test_audio_generate_missing_input(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_audio_generate_http_request( +def test_audio_generate_missing_input(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_audio_generate_http_request( { "json": {"model": omni_server.model}, "timeout": 120, diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py index 61cd30ec8..9861b609e 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """HTTP + WebSocket validation for Qwen3-TTS Base, Higgs Audio V3 and Audex TTA ``/v1/audio/speech`` (and related routes).""" from __future__ import annotations @@ -13,8 +13,8 @@ from typing import Any import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.media import load_test_audio_data_url -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient from tests.helpers.stage_config import get_deploy_config_path # Speech / voice-upload numeric caps for these tests only (not exported from vllm_omni). @@ -26,7 +26,7 @@ _SPEECH_API_MAX_NEW_TOKENS = 4096 _VOICE_UPLOAD_MAX_CONSENT_LEN = 1024 _VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192 _VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS = 2048 -REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." # Qwen3-TTS 0.6B expects 1024-dim speaker embeddings (see serving_speech / talker hidden_size). _QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM = 1024 @@ -57,6 +57,7 @@ _QWEN3_TTS_SPEECH = [ ), ] + _HIGGS_AUDIO_V3_SPEECH = [ pytest.param( OmniServerParams( @@ -84,7 +85,7 @@ _AUDEX_TTA_SPEECH = [ ] -def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None: +def _apply_batch_overrides(body: dict[str, Any], *, loc: str, overrides: dict[str, object]) -> None: if loc == "batch": body.update(overrides) return @@ -112,8 +113,8 @@ def _pcm_wav_mono_bytes(*, duration_s: float = 1.1, sample_rate: int = 44100) -> @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) -def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_audio_speech_http_request( +def test_speech_malformed_json(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_audio_speech_http_request( { "raw_body": "{", "timeout": 120, @@ -124,8 +125,8 @@ def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAICli @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) -def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_audio_speech_http_request( +def test_speech_missing_required_fields(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_audio_speech_http_request( { "json": {"model": omni_server.model}, "timeout": 120, @@ -156,8 +157,14 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: "ref_audio": None, "ref_text": None, }, - ("CustomVoice", "does not support"), - id="customvoice_task_not_supported", + # NOTE: This is delegating down the TTS adapter, so the message is both specific + # to the Qwen3TTS adapter and dependent on whether or not the model has any speakers + # (as opposed to having valida speakers, but getting an invalid speaker). it would + # be more ideal to validate that this is landing on `.validate()` on the adapter + # and ensuring that the returned message matches what we would get from calling the + # adapter validate directly. + ("no speakers configured",), + id="customvoice_invalid_voice", ), pytest.param( {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid" @@ -180,13 +187,17 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: ("mutually exclusive", "speaker_embedding", "ref_audio"), id="speaker_embedding_with_ref_audio_mutually_exclusive", ), - pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"), + pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"), pytest.param( {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1}, ("max_new_tokens", "exceed 4096"), id="max_new_tokens_above_max", ), - pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"), + pytest.param( + {"seed": -(2**63) - 1}, + ("seed", "greater_than_equal", "-9223372036854775808"), + id="seed_below_min", + ), pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"), pytest.param( {"initial_codec_chunk_frames": -1}, @@ -242,13 +253,13 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) def test_speech_invalid_field_values( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, overrides: dict[str, object], err_message: str | tuple[str, ...], ) -> None: ov = dict(overrides) embedding_only = bool(ov.pop(_SPEECH_INVALID_EMBEDDING_ONLY_BODY, False)) - body = { + body: dict[str, Any] = { "model": omni_server.model, "input": "Hello.", "voice": "clone", @@ -261,7 +272,7 @@ def test_speech_invalid_field_values( body.pop("ref_audio", None) body.pop("ref_text", None) body.update(ov) - openai_client.send_audio_speech_http_request( + online_client.send_audio_speech_http_request( {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} ) @@ -274,7 +285,9 @@ def test_speech_invalid_field_values( pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"), pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"), pytest.param( - {"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min" + {"input": "Hello world.", "max_new_tokens": 0}, + ("max_new_tokens", "greater_than_equal"), + id="max_new_tokens_below_min", ), pytest.param( { @@ -328,12 +341,12 @@ def test_speech_invalid_field_values( @pytest.mark.parametrize("omni_server", _HIGGS_AUDIO_V3_SPEECH, indirect=True) def test_speech_higgs_invalid_field_values( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, extra_json: dict[str, object], err_message: str | tuple[str, ...], ) -> None: body = {"model": omni_server.model, **extra_json} - openai_client.send_audio_speech_http_request( + online_client.send_audio_speech_http_request( {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} ) @@ -359,12 +372,12 @@ def test_speech_higgs_invalid_field_values( @pytest.mark.parametrize("omni_server", _AUDEX_TTA_SPEECH, indirect=True) def test_speech_audex_tta_invalid_field_values( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, overrides: dict[str, object], err_message: str | tuple[str, ...], ) -> None: body = {"model": omni_server.model, "input": _AUDEX_TTA_CAPTION, **overrides} - openai_client.send_audio_speech_http_request( + online_client.send_audio_speech_http_request( {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} ) @@ -375,15 +388,15 @@ def test_speech_audex_tta_invalid_field_values( @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) -def test_speech_batch_empty_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_audio_speech_batch_http_request( +def test_speech_batch_empty_items(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_audio_speech_batch_http_request( {"json": {"items": []}, "timeout": 120, "err_code": (400, 422), "err_message": ("items", "least 1 item")} ) @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) -def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_audio_speech_batch_http_request( +def test_speech_batch_missing_items(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_audio_speech_batch_http_request( { "json": {"model": omni_server.model}, "timeout": 120, @@ -553,14 +566,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: Open pytest.param( "batch", {"max_new_tokens": 0}, - ("max_new_tokens", "least 1"), + ("max_new_tokens", "greater_than_equal"), id="batch_max_new_tokens_below_min", marks=_SKIP_ISSUE_3649, ), pytest.param( "item", {"max_new_tokens": 0}, - ("max_new_tokens", "least 1"), + ("max_new_tokens", "greater_than_equal"), id="item_max_new_tokens_below_min", marks=_SKIP_ISSUE_3649, ), @@ -618,14 +631,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: Open @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) def test_speech_batch_invalid_field_values( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, loc: str, overrides: dict[str, object], err_message: str | tuple[str, ...], ) -> None: ov = dict(overrides) strip_clone = bool(ov.pop(_SPEECH_BATCH_NO_CLONE_FIELDS, False)) - body = { + body: dict[str, Any] = { "model": omni_server.model, "voice": "clone", "ref_audio": REF_AUDIO_URL, @@ -638,7 +651,7 @@ def test_speech_batch_invalid_field_values( body.pop("ref_audio", None) body.pop("ref_text", None) _apply_batch_overrides(body, loc=loc, overrides=ov) - openai_client.send_audio_speech_batch_http_request( + online_client.send_audio_speech_batch_http_request( {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} ) @@ -670,7 +683,7 @@ _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES = object() @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) def test_speech_stream_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, send_frames_spec: Any, err_message: str | tuple[str, ...], ) -> None: @@ -687,7 +700,7 @@ def test_speech_stream_invalid_requests( else: send_frames = send_frames_spec assert isinstance(send_frames, str) - openai_client.send_audio_speech_stream_ws_request( + online_client.send_audio_speech_stream_ws_request( { "send_frames": send_frames, "timeout": 120, @@ -893,7 +906,7 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None: @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) def test_voices_create_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, multipart_kind: str | None, form_template: dict[str, Any], err_message: str | tuple[str, ...], @@ -909,7 +922,7 @@ def test_voices_create_invalid_requests( files = _voices_upload_multipart_files(multipart_kind) if files is not None: cfg["files"] = files - openai_client.send_audio_voices_create_http_request(cfg) + online_client.send_audio_voices_create_http_request(cfg) @pytest.mark.parametrize( @@ -962,13 +975,13 @@ def test_voices_create_invalid_requests( @pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True) def test_voices_delete_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, voice_name: str, err_code: int | tuple[int, ...], err_message: str | tuple[str, ...], ) -> None: """DELETE ``/v1/audio/voices/{name}``: missing voice (404), odd segments (often 404), ``.``paths (405).""" - openai_client.send_audio_voices_delete_http_request( + online_client.send_audio_voices_delete_http_request( { "name": voice_name, "timeout": 120, @@ -976,3 +989,68 @@ def test_voices_delete_invalid_requests( "err_message": err_message, } ) + + +# ─── POST /v1/audio/speech · MiniMax Music 3 ─── + +# Text-to-music on the speech endpoint. Most of the speech contract does not +# apply: there is no speaker, no reference audio and no temperature, and the +# lyrics alone are not a request. Each case below pins a rejection so a caller +# who assumes the usual TTS shape finds out immediately. +_MINIMAX_MUSIC3_SPEECH = [ + pytest.param( + OmniServerParams( + model="MiniMaxAI/MiniMax-Music3", + stage_config_path=get_deploy_config_path("minimax_music3.yaml"), + server_args=_SPEECH_SERVER_ARGS, + ), + id="minimax_music3", + marks=hardware_marks(res={"cuda": "H100"}), + ), +] + +_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight" +_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor." +# 25 frames per second, capped at six minutes of song. +_MINIMAX_MUSIC3_MAX_FRAMES = 9000 + + +@pytest.mark.parametrize( + "overrides, err_message", + [ + pytest.param({"instructions": None}, ("instructions",), id="caption_missing"), + pytest.param({"instructions": " "}, ("instructions",), id="caption_blank"), + pytest.param({"input": " "}, ("input",), id="lyrics_blank"), + pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"), + pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"), + pytest.param( + {"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1}, + ("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)), + id="max_new_tokens_over_cap", + ), + pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"), + ], +) +@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True) +def test_minimax_music3_speech_invalid_field_values( + omni_server: OmniServer, + online_client: OnlineOmniClient, + overrides: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + body: dict[str, Any] = { + "model": omni_server.model, + "input": _MINIMAX_MUSIC3_LYRICS, + "instructions": _MINIMAX_MUSIC3_CAPTION, + "seed": 1, + "max_new_tokens": 250, + "response_format": "wav", + } + for key, value in overrides.items(): + if value is None: + body.pop(key, None) + else: + body[key] = value + online_client.send_audio_speech_http_request( + {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} + ) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py index 5ffd25632..9bc5e44fe 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """``POST /v1/images/edits`` multipart validation (live ``Qwen/Qwen-Image-Edit``).""" from __future__ import annotations @@ -9,7 +9,7 @@ from typing import Any import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient pytestmark = [pytest.mark.slow, pytest.mark.diffusion] @@ -72,7 +72,7 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di @pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) def test_images_edits_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, tiny_png_bytes: bytes, include_image: bool, data_template: dict[str, Any], @@ -87,4 +87,4 @@ def test_images_edits_invalid_requests( } if include_image: cfg["files"] = {"image": ("x.png", tiny_png_bytes, "image/png")} - openai_client.send_images_edits_http_request(cfg) + online_client.send_images_edits_http_request(cfg) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py index 5e9424c88..daad00957 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """``POST /v1/images/generations`` validation (live ``Qwen/Qwen-Image``).""" from __future__ import annotations @@ -9,7 +9,7 @@ from typing import Any import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient pytestmark = [pytest.mark.slow, pytest.mark.diffusion] @@ -126,7 +126,7 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]: @pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) def test_images_generations_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, body_spec: Any, err_message: str | tuple[str, ...], ) -> None: @@ -137,6 +137,6 @@ def test_images_generations_invalid_requests( assert isinstance(body_spec, dict) body = _minimal_images_gen_json(omni_server) body.update(body_spec) - openai_client.send_images_generations_http_request( + online_client.send_images_generations_http_request( {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message} ) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py index 9e39b059b..0d026a37b 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``.""" from __future__ import annotations @@ -10,7 +10,7 @@ from typing import Any import pytest from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient from tests.helpers.stage_config import get_deploy_config_path pytestmark = [pytest.mark.slow, pytest.mark.omni] @@ -72,6 +72,9 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case body["top_logprobs"] = 5 elif case_id == "speaker_unknown": body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz" + elif case_id == "audio_format_unsupported": + body["modalities"] = ["text", "audio"] + body["audio"] = {"voice": "alloy", "format": "aac"} else: raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}") return {"json": body, "timeout": 120} @@ -119,7 +122,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case pytest.param( "response_format_json_schema_incomplete", 400, - ("response_format", "value_error", "json_schema"), + ("response_format", "BadRequestError", "json_schema"), id="invalid_response_format_json_schema", ), pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649), @@ -135,12 +138,18 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case ("Invalid speaker", "Supported"), id="speaker_unknown_preset", ), + pytest.param( + "audio_format_unsupported", + 400, + ("Invalid audio format", "aac", "Supported formats"), + id="unsupported_audio_format", + ), ], ) @pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True) def test_chat_completions_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, case_id: str, err_code: int | tuple[int, ...], err_message: str | tuple[str, ...], @@ -148,7 +157,7 @@ def test_chat_completions_invalid_requests( cfg = _chat_completions_request_without_expectations(omni_server, case_id) cfg["err_code"] = err_code cfg["err_message"] = err_message - openai_client.send_chat_completions_http_request(cfg)[0] + online_client.send_chat_completions_http_request(cfg)[0] # ───────────────────────────────────────────────────────────────────────────── @@ -174,7 +183,7 @@ _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR = object() @pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True) def test_video_chat_stream_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, send_frames_spec: Any, err_message: str, ) -> None: @@ -183,7 +192,7 @@ def test_video_chat_stream_invalid_requests( else: send_frames = send_frames_spec assert isinstance(send_frames, str) - openai_client.send_video_chat_stream_ws_request( + online_client.send_video_chat_stream_ws_request( { "send_frames": send_frames, "timeout": 120, @@ -243,11 +252,12 @@ _REALTIME_WS_INVALID_AUDIO_APPEND = object() @pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True) def test_realtime_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, send_frames_spec: Any, ws_error_code: str, err_message: str | tuple[str, ...], ) -> None: + send_frames: Any if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH: send_frames = json.dumps( { @@ -262,7 +272,7 @@ def test_realtime_invalid_requests( ] else: send_frames = send_frames_spec - openai_client.send_realtime_ws_request( + online_client.send_realtime_ws_request( { "send_frames": send_frames, "timeout": 120, diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py index ad3f659f1..4be838316 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Omni sleep/wakeup JSON validation and level-2 wake contract (live ``Qwen/Qwen-Image``).""" from __future__ import annotations @@ -9,7 +9,7 @@ from typing import Any, Literal import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni] @@ -86,7 +86,7 @@ _QWEN_IMAGE = [ @pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True) def test_omni_sleep_wakeup_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, endpoint: Literal["sleep", "wakeup"], json_body: dict[str, Any], err_message: str | tuple[str, ...], @@ -98,9 +98,9 @@ def test_omni_sleep_wakeup_invalid_requests( "err_message": err_message, } if endpoint == "sleep": - openai_client.send_omni_sleep_http_request(cfg) + online_client.send_omni_sleep_http_request(cfg) else: - openai_client.send_omni_wakeup_http_request(cfg) + online_client.send_omni_wakeup_http_request(cfg) # Keep this after the invalid-JSON cases: sleep(level=2) discards weights and the @@ -108,7 +108,7 @@ def test_omni_sleep_wakeup_invalid_requests( @pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True) def test_wakeup_after_level2_sleep_fails( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ) -> None: """Regression for #4473 Repro A via HTTP: wake after sleep(level=2) → 501. @@ -116,7 +116,7 @@ def test_wakeup_after_level2_sleep_fails( Live serve surfaces ``NotImplementedError`` as OpenAI-style JSON with code 501. """ assert omni_server is not None - sleep_resps = openai_client.send_omni_sleep_http_request( + sleep_resps = online_client.send_omni_sleep_http_request( { "json": {"stage_ids": [0], "level": 2}, "timeout": 180, @@ -124,10 +124,11 @@ def test_wakeup_after_level2_sleep_fails( ) sleep_resp = sleep_resps[0] assert sleep_resp.status_code == 200, sleep_resp - assert sleep_resp.json_body is not None - assert sleep_resp.json_body.get("status") == "SUCCESS", sleep_resp.json_body + json_body = sleep_resp.json_body + assert isinstance(json_body, dict) + assert json_body.get("status") == "SUCCESS", json_body - openai_client.send_omni_wakeup_http_request( + online_client.send_omni_wakeup_http_request( { "json": {"stage_ids": [0]}, "timeout": 120, diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py index c0dd133c9..75ff814df 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """HTTP video job routes (live ``Wan-AI/Wan2.2-T2V-A14B-Diffusers``). WebSocket ``/v1/video/chat/stream`` and ``/v1/realtime`` invalid-input checks live in @@ -13,7 +13,7 @@ from typing import Any import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient pytestmark = [pytest.mark.slow, pytest.mark.diffusion] @@ -54,12 +54,12 @@ def _minimal_video_form_data() -> dict[str, str]: @pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) def test_post_videos_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, form_extra: dict[str, str] | None, err_message: str | tuple[str, ...], ) -> None: data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra} - openai_client.send_videos_create_http_request( + online_client.send_videos_create_http_request( {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message} ) @@ -85,12 +85,12 @@ def test_post_videos_invalid_requests( @pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) def test_post_videos_sync_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, form_extra: dict[str, str] | None, err_message: str | tuple[str, ...], ) -> None: data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra} - openai_client.send_videos_sync_http_request( + online_client.send_videos_sync_http_request( {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message} ) @@ -111,11 +111,11 @@ def test_post_videos_sync_invalid_requests( @pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) def test_get_videos_list_invalid_requests( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, params: dict[str, Any], err_message: str | tuple[str, ...], ) -> None: - openai_client.send_videos_list_http_request( + online_client.send_videos_list_http_request( {"params": params, "timeout": 120, "err_code": 400, "err_message": err_message} ) @@ -126,8 +126,8 @@ def test_get_videos_list_invalid_requests( @pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) -def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_video_retrieve_http_request( +def test_get_video_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_video_retrieve_http_request( { "video_id": "does-not-exist", "timeout": 120, @@ -143,8 +143,8 @@ def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: Open @pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) -def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_video_delete_http_request( +def test_delete_video_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_video_delete_http_request( { "video_id": "does-not-exist", "timeout": 120, @@ -160,8 +160,8 @@ def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: O @pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True) -def test_get_video_content_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - openai_client.send_video_content_http_request( +def test_get_video_content_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: + online_client.send_video_content_http_request( { "video_id": "does-not-exist", "timeout": 120, diff --git a/tests/dfx/reliability/test_reliability_hunyuan_image.py b/tests/dfx/reliability/test_reliability_hunyuan_image.py index 43a684d1d..ae0133503 100644 --- a/tests/dfx/reliability/test_reliability_hunyuan_image.py +++ b/tests/dfx/reliability/test_reliability_hunyuan_image.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ HunyuanImage-3.0-Instruct reliability integration tests (DiT-only deploy). @@ -35,7 +38,7 @@ from tests.dfx.reliability.helpers import ( worker_residual_timeout_after_kill_signal, ) from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OpenAIClientHandler +from tests.helpers.runtime import OnlineOmniClient RELIABILITY_SCENARIOS: list[dict[str, Any]] = [ { @@ -133,8 +136,8 @@ def _image_request_config(omni_server: Any) -> dict[str, Any]: } -def _submit_image_generation(openai_client: OpenAIClientHandler, omni_server: Any) -> None: - responses = openai_client.send_images_generations_http_request(_image_request_config(omni_server)) +def _submit_image_generation(online_client: OnlineOmniClient, omni_server: Any) -> None: + responses = online_client.send_images_generations_http_request(_image_request_config(omni_server)) resp = responses[0] if resp.status_code >= 400: raise RuntimeError( @@ -189,7 +192,7 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) - @pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_image_large_request_failure( omni_server_function, - openai_client_function, + online_client_function, ) -> None: stage_config_path = getattr(omni_server_function, "stage_config_path", None) device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path) @@ -202,7 +205,7 @@ def test_reliability_fault_gpu_oom_image_large_request_failure( ) try: try: - _submit_image_generation(openai_client_function, omni_server_function) + _submit_image_generation(online_client_function, omni_server_function) except Exception as exc: assert_fault_exception(exc, FAULT_ERROR_KEYWORDS) else: @@ -222,10 +225,10 @@ def test_reliability_fault_gpu_oom_image_large_request_failure( @pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) def test_reliability_fault_process_kill_request_failure( omni_server_after_fault_function, - openai_client_function, + online_client_function, ) -> None: try: - _submit_image_generation(openai_client_function, omni_server_after_fault_function) + _submit_image_generation(online_client_function, omni_server_after_fault_function) except Exception as exc: assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) else: @@ -244,7 +247,7 @@ def test_reliability_fault_process_kill_request_failure( @pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( omni_server_after_fault_function, - openai_client_function, + online_client_function, ) -> None: host = omni_server_after_fault_function.host port = omni_server_after_fault_function.port @@ -280,7 +283,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( start = time.monotonic() with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: futures = [ - executor.submit(_submit_image_generation, openai_client_function, omni_server_after_fault_function) + executor.submit(_submit_image_generation, online_client_function, omni_server_after_fault_function) for _ in range(3) ] done, pending = concurrent.futures.wait( @@ -335,12 +338,12 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( @pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) def test_reliability_fault_process_kill_worker_with_load_request_failure( omni_server_function, - openai_client_function, + online_client_function, fault_injector: FaultInjector, ) -> None: scenario = "kill_worker_with_load" load_result = run_fault_injection_with_rate_load( - submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function), + submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function), inject_fault=lambda: fault_injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, @@ -363,13 +366,13 @@ def test_reliability_fault_process_kill_worker_with_load_request_failure( @pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup( omni_server_function, - openai_client_function, + online_client_function, signal_name: str, ) -> None: scenario = f"kill_serve_root_with_load_{signal_name.lower()}" injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) load_result = run_fault_injection_with_rate_load( - submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function), + submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function), inject_fault=lambda: injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, @@ -448,7 +451,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup( @pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( omni_server_function, - openai_client_function, + online_client_function, signal_name: str, ) -> None: scenario = f"kill_serve_tree_with_load_{signal_name.lower()}" @@ -458,7 +461,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( inter_kill_wait_seconds=0.1, ) load_result = run_fault_injection_with_rate_load( - submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function), + submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function), inject_fault=lambda: injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, diff --git a/tests/dfx/reliability/test_reliability_qwen3_omni.py b/tests/dfx/reliability/test_reliability_qwen3_omni.py index 32c660954..10331d88d 100644 --- a/tests/dfx/reliability/test_reliability_qwen3_omni.py +++ b/tests/dfx/reliability/test_reliability_qwen3_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Qwen3-Omni reliability integration tests. """ @@ -63,7 +66,7 @@ DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vll def _default_oom_device_spec() -> str: """Use currently visible CUDA ordinals to avoid invalid device index in sidecar.""" - count = torch.accelerator.device_count() + count = torch.accelerator.device_count() or 0 if count <= 0: return "0" return ",".join(str(i) for i in range(count)) @@ -71,7 +74,7 @@ def _default_oom_device_spec() -> str: OOM_INJECTION_CONFIG = { "device": _default_oom_device_spec(), - "target_mem_ratio": 0.92, + "target_mem_ratio": 0.95, "hold_seconds": 0, "startup_timeout_sec": 20, "strict": False, @@ -414,7 +417,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech( @pytest.mark.slow @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) -def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None: +def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, online_client_function) -> None: device_spec = resolve_oom_device_spec( OOM_INJECTION_CONFIG, _deploy_config_path_from_omni_server(omni_server_function), @@ -444,7 +447,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi "key_words": {"audio": ["test"]}, } try: - openai_client_function.send_omni_request(request_config, request_num=1) + online_client_function.send_omni_request(request_config, request_num=1) except Exception as exc: assert_fault_exception(exc, FAULT_ERROR_KEYWORDS) else: @@ -456,7 +459,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi @pytest.mark.slow @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) -def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None: +def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, online_client_function) -> None: device_spec = resolve_oom_device_spec( OOM_INJECTION_CONFIG, _deploy_config_path_from_omni_server(omni_server_function), @@ -487,7 +490,7 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct "key_words": {"audio": ["test"]}, } try: - openai_client_function.send_omni_request(request_config, request_num=4) + online_client_function.send_omni_request(request_config, request_num=4) except Exception as exc: assert_fault_exception(exc, FAULT_ERROR_KEYWORDS) else: @@ -505,7 +508,7 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct ) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_process_kill_request_failure( - omni_server_after_fault_function, openai_client_function + omni_server_after_fault_function, online_client_function ) -> None: messages = dummy_messages_from_mix_data( system_prompt=_get_system_prompt(), @@ -519,7 +522,7 @@ def test_reliability_fault_process_kill_request_failure( "key_words": {"text": ["beijing"]}, } try: - openai_client_function.send_omni_request(request_config, request_num=1) + online_client_function.send_omni_request(request_config, request_num=1) except Exception as exc: assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) else: @@ -567,7 +570,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( assert elapsed <= _PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC + 1.0, ( f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s" ) - assert ff_status >= 500, ( + assert ff_status is not None and ff_status >= 500, ( f"[process_kill fast_fail] expected server-side failure after fault, " f"got status={ff_status}, body={ff_body[:200]!r}" ) @@ -652,13 +655,13 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_process_kill_worker_with_load_request_failure( omni_server_function, - openai_client_function, + online_client_function, fault_injector: FaultInjector, ) -> None: request_config = _chat_request_config(omni_server_function) scenario = "kill_worker_with_load" load_result = run_fault_injection_with_rate_load( - submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1), + submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1), inject_fault=lambda: fault_injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, @@ -704,14 +707,14 @@ def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup( omni_server_function, - openai_client_function, + online_client_function, signal_name: str, ) -> None: request_config = _chat_request_config(omni_server_function) scenario = f"kill_serve_root_with_load_{signal_name.lower()}" injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) load_result = run_fault_injection_with_rate_load( - submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1), + submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1), inject_fault=lambda: injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, @@ -764,7 +767,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup( @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( omni_server_function, - openai_client_function, + online_client_function, signal_name: str, ) -> None: request_config = _chat_request_config(omni_server_function) @@ -775,7 +778,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( inter_kill_wait_seconds=0.1, ) load_result = run_fault_injection_with_rate_load( - submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1), + submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1), inject_fault=lambda: injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, diff --git a/tests/dfx/reliability/test_reliability_voxcpm2.py b/tests/dfx/reliability/test_reliability_voxcpm2.py deleted file mode 100644 index 632c51081..000000000 --- a/tests/dfx/reliability/test_reliability_voxcpm2.py +++ /dev/null @@ -1,508 +0,0 @@ -""" -VoxCPM2 reliability integration tests. - -Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``. -""" - -from __future__ import annotations - -import concurrent.futures -import os -import time -from pathlib import Path -from typing import Any - -import pytest - -from tests.dfx.conftest import ( - assert_fault_exception, - create_reliability_omni_server_params, - resolve_oom_device_spec, -) -from tests.dfx.reliability.helpers import ( - FaultInjector, - assert_no_server_tree_process_residual_and_gpu_release, - extract_openai_error_contract_from_bytes, - get_health_raw, - inject_gpu_oom, - make_process_kill_fault_injector, - make_server_root_kill_fault_injector, - make_server_tree_kill_fault_injector, - post_json_raw, - run_fault_injection_with_rate_load, - stop_gpu_oom_hogs, - worker_residual_timeout_after_kill_signal, -) -from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url -from tests.helpers.runtime import OpenAIClientHandler - -RELIABILITY_SCENARIOS: list[dict[str, Any]] = [ - { - "test_name": "voxcpm2_reliability_default", - "server_params": { - "model": "openbmb/VoxCPM2", - "stage_config_name": "voxcpm2.yaml", - # VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False - # would otherwise override deploy YAML settings during server startup. - "server_args": ["--trust-remote-code", "--disable-log-stats"], - }, - }, -] - -DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy" -OOM_INJECTION_CONFIG = { - "target_mem_ratio": 0.96, - "hold_seconds": 0, - "startup_timeout_sec": 20, - "strict": False, -} -FAULT_ERROR_KEYWORDS = ( - "the request failed", - "oom", - "out of memory", - "cuda", - "orchestrator", - "timeout", - "connection", - "500", - "503", -) -PROCESS_KILL_ERROR_KEYWORDS = ( - "timeout", - "did not complete within", - "connection", - "engine", - "orchestrator", - "dead", - "internal", - "500", - "503", -) -RUNTIME_WORKER_PATTERN = "VLLM::" -SERVE_SIGNAL_PARAMS = [ - pytest.param("SIGTERM", id="sigterm"), - pytest.param("SIGINT", id="sigint"), - pytest.param("SIGKILL", id="sigkill"), -] -SERVE_WITH_LOAD_SIGNAL_PARAMS = [ - pytest.param("SIGTERM", id="sigterm"), - pytest.param("SIGINT", id="sigint"), - pytest.param("SIGKILL", id="sigkill"), -] -TREE_SIGNAL_PARAMS = [ - pytest.param("SIGTERM", id="sigterm"), - pytest.param("SIGKILL", id="sigkill"), -] -TREE_WITH_LOAD_SIGNAL_PARAMS = [ - pytest.param("SIGTERM", id="sigterm"), - pytest.param("SIGKILL", id="sigkill"), -] -WORKER_SIGNAL_FAULT_PARAMS = [ - pytest.param( - make_process_kill_fault_injector( - grep_patterns=RUNTIME_WORKER_PATTERN, - signal_name="SIGTERM", - limit=1, - post_kill_wait_seconds=2.0, - ), - id="runtime_process_chain_sigterm", - ), - pytest.param( - make_process_kill_fault_injector( - grep_patterns=RUNTIME_WORKER_PATTERN, - signal_name="SIGKILL", - limit=1, - post_kill_wait_seconds=2.0, - ), - id="runtime_process_chain_sigkill", - ), -] - -VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR) -INFLIGHT_INJECTION_REQUEST_RATE = 0.3 -INFLIGHT_INJECTION_REQUEST_COUNT = 10 -_SPEECH_FF_HTTP_TIMEOUT_SEC = 30 -_SPEECH_FF_MAX_ELAPSED_SEC = 30 - - -def _speech_request_config(omni_server: Any) -> dict[str, Any]: - return { - "model": omni_server.model, - "input": "The weather is nice today, perfect for a walk in the park.", - "stream": False, - "timeout": 120.0, - "response_format": "wav", - "voice": "default", - } - - -def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None: - openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1) - - -def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]: - base = "The weather is nice today, perfect for a walk in the park. " - return { - "model": omni_server.model, - "input": base * 80, - "max_new_tokens": 4096, - "ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"), - "stream": False, - "timeout": 300.0, - "response_format": "wav", - "voice": "default", - } - - -def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None: - openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1) - - -def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None: - payload = { - "model": model, - "input": "post-fault fast-fail check", - "voice": "default", - } - start = time.monotonic() - try: - status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20) - elapsed = time.monotonic() - start - assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s" - assert status >= 500, ( - f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}" - ) - except Exception: # noqa: BLE001 - elapsed = time.monotonic() - start - assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s" - - -def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None: - deadline = time.monotonic() + 20.0 - last_observation = "" - while time.monotonic() < deadline: - try: - status, body = get_health_raw(host, port, timeout_sec=5) - last_observation = f"http={status}, body={body[:200]!r}" - if status == 503: - return - except Exception as exc: # noqa: BLE001 - last_observation = f"exception={exc!r}" - return - time.sleep(0.5) - pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}") - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.skip(reason="issue#4285") -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_gpu_oom_speech_request_failure( - omni_server_function, - openai_client_function, -) -> None: - stage_config_path = getattr(omni_server_function, "stage_config_path", None) - device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path) - handle = inject_gpu_oom( - device=device_spec, - target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"], - hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"], - startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"], - strict=OOM_INJECTION_CONFIG["strict"], - ) - try: - try: - _submit_oom_speech_request(openai_client_function, omni_server_function) - except Exception as exc: - assert_fault_exception(exc, FAULT_ERROR_KEYWORDS) - else: - pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection") - finally: - stop_gpu_oom_hogs(handle) - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.skip(reason="issue#4285") -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_gpu_oom_speech_error_contract( - omni_server_function, -) -> None: - stage_config_path = getattr(omni_server_function, "stage_config_path", None) - device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path) - handle = inject_gpu_oom( - device=device_spec, - target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"], - hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"], - startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"], - strict=OOM_INJECTION_CONFIG["strict"], - ) - host = omni_server_function.host - port = omni_server_function.port - oom_config = _oom_speech_request_config(omni_server_function) - payload = { - "model": oom_config["model"], - "input": oom_config["input"], - "voice": oom_config["voice"], - "max_new_tokens": oom_config["max_new_tokens"], - "ref_audio": oom_config["ref_audio"], - } - try: - status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300) - finally: - stop_gpu_oom_hogs(handle) - - assert status >= 500, f"expected speech runtime error under OOM, got status={status}" - speech_error = extract_openai_error_contract_from_bytes(body) - assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}" - assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}" - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") -@pytest.mark.parametrize( - "fault_injector", - WORKER_SIGNAL_FAULT_PARAMS, - indirect=True, -) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_process_kill_request_failure( - omni_server_after_fault_function, - openai_client_function, -) -> None: - try: - _submit_speech_request(openai_client_function, omni_server_after_fault_function) - except Exception as exc: - assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) - else: - pytest.fail("expected /v1/audio/speech request failure after process-kill injection") - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") -@pytest.mark.parametrize( - "fault_injector", - WORKER_SIGNAL_FAULT_PARAMS, - indirect=True, -) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_process_kill_fast_fail_and_concurrent( - omni_server_after_fault_function, - openai_client_function, -) -> None: - host = omni_server_after_fault_function.host - port = omni_server_after_fault_function.port - model = omni_server_after_fault_function.model - payload = { - "model": model, - "input": "Say hello in one short sentence.", - "voice": "default", - } - - start = time.monotonic() - try: - status, body = post_json_raw( - host, - port, - "/v1/audio/speech", - payload, - timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC, - ) - elapsed = time.monotonic() - start - assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, ( - f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s" - ) - assert status >= 500, ( - "[process_kill fast_fail] expected server-side failure after fault, " - f"got status={status}, body={body[:200]!r}" - ) - except Exception: - elapsed = time.monotonic() - start - assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, ( - f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s" - ) - - start = time.monotonic() - with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: - futures = [ - executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function) - for _ in range(3) - ] - done, pending = concurrent.futures.wait( - futures, - timeout=30, - return_when=concurrent.futures.ALL_COMPLETED, - ) - - elapsed = time.monotonic() - start - assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}" - assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s" - - fault_observed = False - for future in done: - try: - future.result() - except Exception as exc: - fault_observed = True - assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) - assert fault_observed, ( - "[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault" - ) - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") -@pytest.mark.parametrize( - "fault_injector", - WORKER_SIGNAL_FAULT_PARAMS, - indirect=True, -) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_process_kill_worker_with_load_request_failure( - omni_server_function, - openai_client_function, - fault_injector: FaultInjector, -) -> None: - scenario = "kill_worker_with_load" - load_result = run_fault_injection_with_rate_load( - submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function), - inject_fault=lambda: fault_injector(omni_server_function), - num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, - request_rate=INFLIGHT_INJECTION_REQUEST_RATE, - completion_timeout_sec=120.0, - ) - assert load_result["failure_observed"], ( - f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" - ) - host = omni_server_function.host - port = omni_server_function.port - _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) - _assert_post_fault_health_terminal(host, port, scenario=scenario) - - -@pytest.mark.slow -@pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") -@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup( - omni_server_function, - signal_name: str, -) -> None: - scenario = f"kill_serve_root_no_load_{signal_name.lower()}" - injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) - injector(omni_server_function) - host = omni_server_function.host - port = omni_server_function.port - _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) - _assert_post_fault_health_terminal(host, port, scenario=scenario) - assert_no_server_tree_process_residual_and_gpu_release( - omni_server_function, - scenario=scenario, - timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), - ) - - -@pytest.mark.slow -@pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") -@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup( - omni_server_function, - openai_client_function, - signal_name: str, -) -> None: - scenario = f"kill_serve_root_with_load_{signal_name.lower()}" - injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) - load_result = run_fault_injection_with_rate_load( - submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function), - inject_fault=lambda: injector(omni_server_function), - num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, - request_rate=INFLIGHT_INJECTION_REQUEST_RATE, - completion_timeout_sec=300.0, - ) - assert load_result["failure_observed"], ( - f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" - ) - host = omni_server_function.host - port = omni_server_function.port - _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) - _assert_post_fault_health_terminal(host, port, scenario=scenario) - assert_no_server_tree_process_residual_and_gpu_release( - omni_server_function, - scenario=scenario, - timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), - ) - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") -@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup( - omni_server_function, - signal_name: str, -) -> None: - scenario = f"kill_serve_tree_no_load_{signal_name.lower()}" - injector = make_server_tree_kill_fault_injector( - signal_name=signal_name, - post_kill_wait_seconds=2.0, - inter_kill_wait_seconds=0.1, - ) - injector(omni_server_function) - host = omni_server_function.host - port = omni_server_function.port - _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) - _assert_post_fault_health_terminal(host, port, scenario=scenario) - assert_no_server_tree_process_residual_and_gpu_release( - omni_server_function, - scenario=scenario, - timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), - ) - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") -@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS) -@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) -def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( - omni_server_function, - openai_client_function, - signal_name: str, -) -> None: - scenario = f"kill_serve_tree_with_load_{signal_name.lower()}" - injector = make_server_tree_kill_fault_injector( - signal_name=signal_name, - post_kill_wait_seconds=2.0, - inter_kill_wait_seconds=0.1, - ) - load_result = run_fault_injection_with_rate_load( - submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function), - inject_fault=lambda: injector(omni_server_function), - num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, - request_rate=INFLIGHT_INJECTION_REQUEST_RATE, - completion_timeout_sec=120.0, - ) - assert load_result["failure_observed"], ( - f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" - ) - host = omni_server_function.host - port = omni_server_function.port - _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) - _assert_post_fault_health_terminal(host, port, scenario=scenario) - assert_no_server_tree_process_residual_and_gpu_release( - omni_server_function, - scenario=scenario, - timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), - ) diff --git a/tests/dfx/reliability/test_reliability_wan22.py b/tests/dfx/reliability/test_reliability_wan22.py index 84175429d..8a86c1c92 100644 --- a/tests/dfx/reliability/test_reliability_wan22.py +++ b/tests/dfx/reliability/test_reliability_wan22.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Wan2.2 reliability integration tests. """ @@ -183,7 +186,7 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) - @pytest.mark.slow @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True) -def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, openai_client_function) -> None: +def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, online_client_function) -> None: stage_config_path = getattr(omni_server_function, "stage_config_path", None) device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path) handle = inject_gpu_oom( @@ -211,7 +214,7 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct "stream": False, } try: - openai_client_function.send_video_diffusion_request(request_config, request_num=1) + online_client_function.send_video_diffusion_request(request_config, request_num=1) except Exception as exc: assert_fault_exception(exc, FAULT_ERROR_KEYWORDS) else: @@ -230,11 +233,11 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True) def test_reliability_fault_process_kill_request_failure( omni_server_after_fault_function, - openai_client_function, + online_client_function, ) -> None: request_config = _video_request_config() try: - openai_client_function.send_video_diffusion_request(request_config, request_num=1) + online_client_function.send_video_diffusion_request(request_config, request_num=1) except Exception as exc: assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) else: @@ -251,7 +254,7 @@ def test_reliability_fault_process_kill_request_failure( @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True) def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( omni_server_after_fault_function, - openai_client_function, + online_client_function, ) -> None: """Black-box: after process kill, /v1/videos fails fast and concurrent calls don't hang. @@ -308,7 +311,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( start = time.monotonic() with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: futures = [ - executor.submit(openai_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3) + executor.submit(online_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3) ] done, pending = concurrent.futures.wait( futures, @@ -367,13 +370,13 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True) def test_reliability_fault_process_kill_worker_with_load_request_failure( omni_server_function, - openai_client_function, + online_client_function, fault_injector: FaultInjector, ) -> None: request_config = _video_request_config() scenario = "kill_worker_with_load" load_result = run_fault_injection_with_rate_load( - submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1), + submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1), inject_fault=lambda: fault_injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, @@ -516,7 +519,7 @@ def test_reliability_video_oom_recovers_after_fault_removed( @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True) def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup( omni_server_function, - openai_client_function, + online_client_function, signal_name: str, ) -> None: """Black-box: during in-flight video requests, kill serve root and verify post-fault behavior/cleanup.""" @@ -524,7 +527,7 @@ def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_clean scenario = f"kill_serve_root_with_load_{signal_name.lower()}" injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) load_result = run_fault_injection_with_rate_load( - submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1), + submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1), inject_fault=lambda: injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, @@ -601,7 +604,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup( @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True) def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( omni_server_function, - openai_client_function, + online_client_function, signal_name: str, ) -> None: """Black-box: during in-flight video requests, kill server process tree and verify post-fault behavior/cleanup.""" @@ -613,7 +616,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( inter_kill_wait_seconds=0.1, ) load_result = run_fault_injection_with_rate_load( - submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1), + submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1), inject_fault=lambda: injector(omni_server_function), num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, request_rate=INFLIGHT_INJECTION_REQUEST_RATE, diff --git a/tests/dfx/stability/conftest.py b/tests/dfx/stability/conftest.py index eaa755d45..0889db885 100644 --- a/tests/dfx/stability/conftest.py +++ b/tests/dfx/stability/conftest.py @@ -37,7 +37,7 @@ def omni_server(request: pytest.FixtureRequest): with _omni_server_lock: # Same tuple and CLI composition as ``tests/dfx/perf/scripts/run_benchmark.py``; # ``serve_args`` from JSON are folded into ``extra_cli_args`` inside - # ``create_unique_server_params``. + # ``create_unique_server_pytest_params`` (JSON ``mark`` → pytest marks). test_name, model, deploy_path, stage_overrides, extra_cli_args, use_omni = request.param print(f"Starting OmniServer with test: {test_name}, model: {model}") diff --git a/tests/dfx/stability/scripts/test_stability_hunyuan_image.py b/tests/dfx/stability/scripts/test_stability_hunyuan_image.py index 38858047f..be2aec805 100644 --- a/tests/dfx/stability/scripts/test_stability_hunyuan_image.py +++ b/tests/dfx/stability/scripts/test_stability_hunyuan_image.py @@ -13,7 +13,7 @@ import pytest from tests.dfx.conftest import ( create_benchmark_indices, create_test_parameter_mapping, - create_unique_server_params, + create_unique_server_pytest_params, load_configs, ) from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop @@ -29,7 +29,7 @@ try: except FileNotFoundError: BENCHMARK_CONFIGS = [] -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) diff --git a/tests/dfx/stability/scripts/test_stability_qwen3_omni.py b/tests/dfx/stability/scripts/test_stability_qwen3_omni.py index 5f487a27c..98ab5ee89 100644 --- a/tests/dfx/stability/scripts/test_stability_qwen3_omni.py +++ b/tests/dfx/stability/scripts/test_stability_qwen3_omni.py @@ -13,7 +13,7 @@ import pytest from tests.dfx.conftest import ( create_benchmark_indices, create_test_parameter_mapping, - create_unique_server_params, + create_unique_server_pytest_params, load_configs, ) from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop @@ -29,7 +29,7 @@ try: except FileNotFoundError: BENCHMARK_CONFIGS = [] -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) diff --git a/tests/dfx/stability/scripts/test_stability_qwen3_tts.py b/tests/dfx/stability/scripts/test_stability_qwen3_tts.py index 0c8059ed9..293b723ca 100644 --- a/tests/dfx/stability/scripts/test_stability_qwen3_tts.py +++ b/tests/dfx/stability/scripts/test_stability_qwen3_tts.py @@ -13,7 +13,7 @@ import pytest from tests.dfx.conftest import ( create_benchmark_indices, create_test_parameter_mapping, - create_unique_server_params, + create_unique_server_pytest_params, load_configs, ) from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop @@ -29,7 +29,7 @@ try: except FileNotFoundError: BENCHMARK_CONFIGS = [] -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) diff --git a/tests/dfx/stability/scripts/test_stability_qwen_image.py b/tests/dfx/stability/scripts/test_stability_qwen_image.py index 7ff39b5ca..62b0cd86c 100644 --- a/tests/dfx/stability/scripts/test_stability_qwen_image.py +++ b/tests/dfx/stability/scripts/test_stability_qwen_image.py @@ -13,7 +13,7 @@ import pytest from tests.dfx.conftest import ( create_benchmark_indices, create_test_parameter_mapping, - create_unique_server_params, + create_unique_server_pytest_params, load_configs, ) from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop @@ -29,7 +29,7 @@ try: except FileNotFoundError: BENCHMARK_CONFIGS = [] -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) diff --git a/tests/dfx/stability/scripts/test_stability_voxcpm2.py b/tests/dfx/stability/scripts/test_stability_voxcpm2.py deleted file mode 100644 index e26a2050c..000000000 --- a/tests/dfx/stability/scripts/test_stability_voxcpm2.py +++ /dev/null @@ -1,69 +0,0 @@ -""" -VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration. - -Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``. -""" - -from __future__ import annotations - -from pathlib import Path - -import pytest - -from tests.dfx.conftest import ( - create_benchmark_indices, - create_test_parameter_mapping, - create_unique_server_params, - load_configs, -) -from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop - -STABILITY_DIR = Path(__file__).resolve().parent.parent -DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" -CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json") -DEFAULT_NUM_PROMPTS_PER_BATCH = 20 -STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"] - -try: - BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) -except FileNotFoundError: - BENCHMARK_CONFIGS = [] - -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] -server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} -benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) - - -@pytest.mark.slow -@pytest.mark.tts -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True) -def test_stability_voxcpm2(omni_server, stability_benchmark_params): - test_name = stability_benchmark_params["test_name"] - params = stability_benchmark_params["params"] - duration_sec = params.get("duration_sec", 300) - num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH) - request_rate = params.get("request_rate") - max_concurrency = params.get("max_concurrency") - - bench_params = { - k: v - for k, v in params.items() - if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch") - } - - result = run_stability_benchmark_loop( - host=omni_server.host, - port=omni_server.port, - model=omni_server.model, - duration_sec=duration_sec, - params=bench_params, - request_rate=request_rate, - max_concurrency=max_concurrency, - result_dir=str(STABILITY_DIR), - num_prompts_per_batch=num_prompts_per_batch, - run_one_batch=_run_one_vllm_bench_batch, - ) - - assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}" - assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed" diff --git a/tests/dfx/stability/scripts/test_stability_wan22.py b/tests/dfx/stability/scripts/test_stability_wan22.py index a363133f5..aae0ae429 100644 --- a/tests/dfx/stability/scripts/test_stability_wan22.py +++ b/tests/dfx/stability/scripts/test_stability_wan22.py @@ -13,7 +13,7 @@ import pytest from tests.dfx.conftest import ( create_benchmark_indices, create_test_parameter_mapping, - create_unique_server_params, + create_unique_server_pytest_params, load_configs, ) from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop @@ -29,7 +29,7 @@ try: except FileNotFoundError: BENCHMARK_CONFIGS = [] -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) diff --git a/tests/dfx/stability/tests/test_hunyuan_image.json b/tests/dfx/stability/tests/test_hunyuan_image.json index b78fa812c..b468ceec3 100644 --- a/tests/dfx/stability/tests/test_hunyuan_image.json +++ b/tests/dfx/stability/tests/test_hunyuan_image.json @@ -1,6 +1,18 @@ [ { "test_name": "test_hunyuan_image_stability_tp2_fp8_sp2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H800" + }, + "num_cards": 2 + } + }, + "local_model", + "diffusion" + ], "server_params": { "model": "tencent/HunyuanImage-3.0-Instruct", "serve_args": { diff --git a/tests/dfx/stability/tests/test_qwen3_omni.json b/tests/dfx/stability/tests/test_qwen3_omni.json index cbdf25e67..816e981b6 100644 --- a/tests/dfx/stability/tests/test_qwen3_omni.json +++ b/tests/dfx/stability/tests/test_qwen3_omni.json @@ -1,6 +1,18 @@ [ { "test_name": "test_qwen3_omni_stability_async_chunk", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H800" + }, + "num_cards": 2 + } + }, + "local_model", + "omni" + ], "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "extra_cli_args": ["--async-chunk"] diff --git a/tests/dfx/stability/tests/test_qwen3_tts.json b/tests/dfx/stability/tests/test_qwen3_tts.json index 0208331bd..1c9356574 100644 --- a/tests/dfx/stability/tests/test_qwen3_tts.json +++ b/tests/dfx/stability/tests/test_qwen3_tts.json @@ -1,6 +1,19 @@ [ { "test_name": "test_qwen3_tts_stability", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H800", + "npu": "A3" + }, + "num_cards": 2 + } + }, + "local_model", + "tts" + ], "server_params": { "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" }, diff --git a/tests/dfx/stability/tests/test_qwen_image.json b/tests/dfx/stability/tests/test_qwen_image.json index e28614973..0c21d0ac1 100644 --- a/tests/dfx/stability/tests/test_qwen_image.json +++ b/tests/dfx/stability/tests/test_qwen_image.json @@ -1,6 +1,18 @@ [ { "test_name": "test_qwen_image_stability", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H800" + }, + "num_cards": 1 + } + }, + "local_model", + "diffusion" + ], "server_params": { "model": "Qwen/Qwen-Image" }, diff --git a/tests/dfx/stability/tests/test_voxcpm2.json b/tests/dfx/stability/tests/test_voxcpm2.json deleted file mode 100644 index dd912da36..000000000 --- a/tests/dfx/stability/tests/test_voxcpm2.json +++ /dev/null @@ -1,35 +0,0 @@ -[ - { - "test_name": "test_voxcpm2_stability", - "server_params": { - "model": "openbmb/VoxCPM2", - "trust_remote_code": true - }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-audio-speech", - "endpoint": "/v1/audio/speech", - "trust_remote_code": true, - "duration_sec": 86400, - "max_concurrency": 2, - "num_prompts_per_batch": 50, - "random_input_len": { - "min": 0, - "max": 1000 - }, - "random_output_len": { - "min": 0, - "max": 1000 - }, - "random_range_ratio": 0.0, - "extra_body": { - "voice": "default", - "language": "English" - }, - "ignore_eos": true, - "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration" - } - ] - } -] diff --git a/tests/dfx/stability/tests/test_wan22.json b/tests/dfx/stability/tests/test_wan22.json index 7b3b54b65..e6ab0b5cf 100644 --- a/tests/dfx/stability/tests/test_wan22.json +++ b/tests/dfx/stability/tests/test_wan22.json @@ -1,6 +1,19 @@ [ { - "test_name": "test_wan22_i2v_stability_v1_videos", + "test_name": "test_wan22_i2v_stability_ulysses2_vae_patch2_hsdp_slicing_tiling", + "description": "Ulysses=2 + VAE patch parallel=2 + HSDP + VAE slicing + tiling", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H800" + }, + "num_cards": 2 + } + }, + "local_model", + "diffusion" + ], "server_params": { "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", "serve_args": { @@ -27,5 +40,45 @@ ] } ] + }, + { + "test_name": "test_wan22_i2v_stability_usp2_vae_patch2_hsdp_slicing", + "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing", + "mark": [ + { + "hardware_marks": { + "res": { + "npu": "A3" + }, + "num_cards": 2 + } + }, + "local_model", + "diffusion" + ], + "server_params": { + "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "serve_args": { + "usp": 2, + "vae-patch-parallel-size": 2, + "use-hsdp": true, + "vae-use-slicing": true + } + }, + "benchmark_params": [ + { + "dataset": "random", + "task": "i2v", + "backend": "v1/videos", + "duration_sec": 86400, + "max_concurrency": 1, + "num_prompts_per_batch": 50, + "enable_negative_prompt": true, + "random_request_config": [ + {"width": 832, "height": 480, "num_inference_steps": 4, "num_frames": 81, "fps": 16, "weight": 0.5}, + {"width": 1280, "height": 720, "num_inference_steps": 4, "num_frames": 121, "fps": 16, "weight": 0.5} + ] + } + ] } ] diff --git a/tests/diffusion/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py index 7b94fe80f..10451c134 100644 --- a/tests/diffusion/ar_diffusion/test_capability_runner.py +++ b/tests/diffusion/ar_diffusion/test_capability_runner.py @@ -18,6 +18,7 @@ from vllm_omni.experimental.ar_diffusion.capability import ( ) from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner +from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest BLOCK = 16 POS = "positive" @@ -56,6 +57,21 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec: ) +def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec: + return ARDiffusionKVCacheSpec( + num_layers=1, + num_kv_heads=1, + head_size=1, + tokens_per_frame=1, + frames_per_block=3, + window_frames=3, + sink_frames=3, + kv_branches=(ARDiffusionKVBranchSpec("main", 0),), + session_capacity=capacity, + cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),), + ) + + class CapablePipeline: def __init__(self, spec: ARDiffusionKVCacheSpec) -> None: self.spec = spec @@ -103,6 +119,7 @@ def make_runner( *, available_bytes: int = 1 << 28, step_execution: bool = False, + gpu_memory_fraction: float = 0.1, ) -> ARDiffusionModelRunner: runner = object.__new__(ARDiffusionModelRunner) runner.od_config = SimpleNamespace( @@ -113,7 +130,10 @@ def make_runner( ) runner.device = torch.device("cpu") runner.pipeline = pipeline - runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True) + runner.ar_diffusion_kv_config = ARDiffusionKVConfig( + enable=True, + gpu_memory_fraction=gpu_memory_fraction, + ) runner.kv_cache = None runner._ar_diffusion_capability = None runner._ar_diffusion_kv_cache_spec = None @@ -132,6 +152,16 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch: return state +def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str): + state = runner._get_or_create_session(session_id) + assert runner.kv_cache is not None + seq_len = BLOCK * runner.kv_cache.frames_per_block + ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx + ctx.ensure_video_slots(torch.device("cpu")) + state.commit_paged_context(kv_branch) + return state + + def test_ar_runner_rejects_pipeline_without_capability(): runner = object.__new__(ARDiffusionModelRunner) runner.od_config = SimpleNamespace(max_num_seqs=1) @@ -140,6 +170,43 @@ def test_ar_runner_rejects_pipeline_without_capability(): runner._preallocate_kv_cache(available_bytes=1 << 20) +def test_runner_uses_typed_tick_as_authoritative_session_contract(): + tick = ARDiffusionTickRequest( + session_id="world-7", + request_id="request-3", + chunk_index=3, + reset=True, + ) + req = SimpleNamespace( + request_id="request-3", + sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()), + ) + + session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req) + + assert session_id == "world-7" + assert extra_args == tick.to_extra_args() + assert parsed == tick + + +def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id(): + tick = ARDiffusionTickRequest( + session_id="world-7", + request_id="client-request-3", + chunk_index=3, + ) + req = SimpleNamespace( + request_id="engine-request-uuid", + sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()), + ) + + session_id, _, parsed = ARDiffusionModelRunner._request_session(req) + + assert session_id == "world-7" + assert req.request_id == "engine-request-uuid" + assert parsed.request_id == "client-request-3" + + def test_lingbot_like_single_branch_session_reuse_reset_and_close(): pipeline = CapablePipeline(lingbot_like_spec()) runner = make_runner(pipeline) @@ -170,6 +237,31 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close(): assert pipeline.closes == ["s1"] +def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions(): + runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2))) + kv = runner.kv_cache + assert kv is not None + + session_a = commit_one_block(runner, "world-a", "main") + a0_blocks = kv.window_block_ids(session_a.adapter("main")) + session_b = commit_one_block(runner, "world-b", "main") + b0_blocks = kv.window_block_ids(session_b.adapter("main")) + session_a_again = commit_one_block(runner, "world-a", "main") + a1_blocks = kv.window_block_ids(session_a_again.adapter("main")) + + assert session_a_again is session_a + assert session_a.adapter("main").request_id == "ar::world-a::main" + assert session_b.adapter("main").request_id == "ar::world-b::main" + assert session_a.adapter("main").completed_chunks == 6 + assert session_b.adapter("main").completed_chunks == 3 + assert len(a0_blocks) == 3 + assert len(b0_blocks) == 3 + assert len(a1_blocks) == 6 + assert set(a0_blocks) <= set(a1_blocks) + assert set(a1_blocks).isdisjoint(b0_blocks) + assert tuple(runner._sessions) == ("world-b", "world-a") + + def test_lingbot_like_sink_survives_sliding_window_eviction(): runner = make_runner(CapablePipeline(lingbot_like_spec())) kv = runner.kv_cache @@ -198,12 +290,12 @@ def test_dreamzero_like_two_branches_are_independent(): def test_lru_eviction_releases_blocks_and_notifies_pipeline(): - pipeline = CapablePipeline(lingbot_like_spec(capacity=8)) + pipeline = CapablePipeline(lingbot_like_spec(capacity=2)) runner = make_runner(pipeline) kv = runner.kv_cache assert kv is not None - assert runner._session_capacity == 1 - assert kv.session_capacity == 1 + assert runner._session_capacity == 2 + assert kv.session_capacity == 2 free_total = kv.manager.block_pool.get_num_free_blocks() old = commit_one_frame(runner, "old", "main") k = torch.randn(1, 8, 4, 64) @@ -211,13 +303,56 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline(): assert kv.manager.block_pool.get_num_free_blocks() < free_total runner._get_or_create_session("new") + assert tuple(runner._sessions) == ("old", "new") + runner._get_or_create_session("newest") - assert tuple(runner._sessions) == ("new",) + assert tuple(runner._sessions) == ("new", "newest") assert pipeline.closes == ["old"] assert "old" not in kv._cross_sessions assert kv.manager.block_pool.get_num_free_blocks() == free_total +def test_budget_reduced_capacity_drives_runner_lru(): + pipeline = CapablePipeline(tiny_spec(capacity=2)) + # One tiny LingBot-like session requires 128 bytes; two require 192. + runner = make_runner( + pipeline, + available_bytes=128, + gpu_memory_fraction=1.0, + ) + kv = runner.kv_cache + assert kv is not None + assert kv.requested_session_capacity == 2 + assert kv.session_capacity == 1 + assert runner._session_capacity == 1 + + runner._get_or_create_session("old") + runner._get_or_create_session("new") + + assert tuple(runner._sessions) == ("new",) + assert pipeline.closes == ["old"] + + +def test_dreamzero_like_requested_capacity_is_capped_by_budget(): + spec = dreamzero_like_spec(capacity=64) + # Per all-layer self-KV page: 65,536 bytes. Two resident sessions need: + # managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages, + # cross-attention=1 page/session, for 44 pages total. + page_bytes = 65_536 + pipeline = CapablePipeline(spec) + runner = make_runner( + pipeline, + available_bytes=44 * page_bytes, + gpu_memory_fraction=1.0, + ) + kv = runner.kv_cache + assert kv is not None + assert kv.requested_session_capacity == 64 + assert kv.session_capacity == 2 + assert runner._session_capacity == 2 + assert kv.cross_attention_reserved_bytes == 2 * page_bytes + + def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch): pipeline = CapablePipeline(lingbot_like_spec()) runner = make_runner(pipeline) @@ -232,7 +367,10 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat raise RuntimeError("layer exploded") monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom) - request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"})) + request = SimpleNamespace( + request_id="broken-request", + sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}), + ) with pytest.raises(RuntimeError, match="layer exploded"): runner.execute_model(request) @@ -263,7 +401,10 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch): monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation) monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom) runner.device = torch.device("cuda") - request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"})) + request = SimpleNamespace( + request_id="broken-request", + sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}), + ) with pytest.raises(RuntimeError, match="asynchronous kernel failed"): runner.execute_model(request) diff --git a/tests/diffusion/ar_diffusion/test_consumer.py b/tests/diffusion/ar_diffusion/test_consumer.py new file mode 100644 index 000000000..d1197113c --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_consumer.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import AsyncIterator, Sequence +from typing import Any + +import pytest + +from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer +from vllm_omni.experimental.ar_diffusion.session import ( + ARDiffusionSession, + ARDiffusionSessionEvent, +) +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionChunkMetadata, + ARDiffusionControlInput, + ARDiffusionTickRequest, +) +from vllm_omni.inputs.data import ( + OmniDiffusionSamplingParams, + OmniPromptType, + OmniSamplingParams, +) +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeStagePool: + def __init__(self, num_replicas: int = 1) -> None: + self.num_replicas = num_replicas + + +class FakeEngine: + def __init__(self, num_replicas: int = 1) -> None: + self.stage_pools = [FakeStagePool(num_replicas)] + + +class FakeGenerateClient: + def __init__( + self, + *, + include_metadata: bool = True, + applied_event_ids: list[int] | None = None, + output_error: str | None = None, + num_replicas: int = 1, + ) -> None: + self.engine = FakeEngine(num_replicas) + self.include_metadata = include_metadata + self.applied_event_ids = applied_event_ids + self.output_error = output_error + self.calls: list[dict[str, Any]] = [] + + async def generate( + self, + prompt: OmniPromptType, + *, + request_id: str, + sampling_params_list: Sequence[OmniSamplingParams], + output_modalities: list[str] | None = None, + ) -> AsyncIterator[OmniRequestOutput]: + self.calls.append( + { + "prompt": prompt, + "request_id": request_id, + "sampling_params_list": sampling_params_list, + "output_modalities": output_modalities, + } + ) + params = sampling_params_list[0] + assert isinstance(params, OmniDiffusionSamplingParams) + tick = ARDiffusionTickRequest.from_extra_args(params.extra_args) + assert tick is not None + chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict() + if self.applied_event_ids is not None: + chunk_metadata["applied_event_ids"] = self.applied_event_ids + metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {} + yield OmniRequestOutput( + request_id=request_id, + stage_id=0, + final_output_type="video", + images=["chunk"], # type: ignore[list-item] + _multimodal_output=metadata, + finished=True, + error=self.output_error, + ) + + +class FakeLifecycle: + async def reset_session(self, session_id: str) -> None: + pass + + async def close_session(self, session_id: str) -> None: + pass + + +def make_tick() -> ARDiffusionTickRequest: + return ARDiffusionTickRequest( + session_id="world-1", + request_id="world-1-chunk-3", + chunk_index=3, + applied_event_ids=(7, 8), + prompt="turn left", + controls=( + ARDiffusionControlInput( + track="camera", + schema="example.camera.v1", + data={"samples": [{"x": 1.0}]}, + ), + ), + ) + + +def test_consumer_rejects_replicated_ar_diffusion_stage() -> None: + client = FakeGenerateClient(num_replicas=2) + + with pytest.raises( + ValueError, + match=r"require exactly one replica.*num_replicas=2", + ): + ARDiffusionOmniTickConsumer( + client, + prompt_provider=lambda tick: tick.prompt or "", + sampling_params_list=[OmniDiffusionSamplingParams()], + diffusion_stage_id=0, + ) + + assert client.calls == [] + + +@pytest.mark.asyncio +async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None: + client = FakeGenerateClient() + template = OmniDiffusionSamplingParams( + output_type="latent", + extra_args={"deployment_option": "kept"}, + ) + consumer = ARDiffusionOmniTickConsumer( + client, + prompt_provider=lambda tick: { + "prompt": tick.prompt, + "multi_modal_data": {"image": "initial-image"}, + }, + sampling_params_list=[template], + diffusion_stage_id=0, + output_modalities=["video"], + ) + tick = make_tick() + + output = await consumer.execute_tick(tick) + + assert output.images == ["chunk"] + assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick) + assert len(client.calls) == 1 + call = client.calls[0] + assert call["request_id"] == tick.request_id + assert call["output_modalities"] == ["video"] + assert call["prompt"] == { + "prompt": "turn left", + "multi_modal_data": {"image": "initial-image"}, + } + submitted_params = call["sampling_params_list"][0] + assert submitted_params.extra_args["deployment_option"] == "kept" + assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick + assert "ar_diffusion_tick" not in template.extra_args + + +@pytest.mark.asyncio +async def test_consumer_rejects_missing_chunk_metadata() -> None: + consumer = ARDiffusionOmniTickConsumer( + FakeGenerateClient(include_metadata=False), + prompt_provider=lambda tick: tick.prompt or "", + sampling_params_list=[OmniDiffusionSamplingParams()], + diffusion_stage_id=0, + ) + + with pytest.raises(ValueError, match="multimodal_output.metadata"): + await consumer.execute_tick(make_tick()) + + +@pytest.mark.asyncio +async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None: + consumer = ARDiffusionOmniTickConsumer( + FakeGenerateClient(applied_event_ids=[8, 7]), + prompt_provider=lambda tick: tick.prompt or "", + sampling_params_list=[OmniDiffusionSamplingParams()], + diffusion_stage_id=0, + ) + + with pytest.raises(ValueError, match="unique and strictly increasing"): + await consumer.execute_tick(make_tick()) + + +@pytest.mark.asyncio +async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None: + consumer = ARDiffusionOmniTickConsumer( + FakeGenerateClient(include_metadata=False, output_error="model failed"), + prompt_provider=lambda tick: tick.prompt or "", + sampling_params_list=[OmniDiffusionSamplingParams()], + diffusion_stage_id=0, + ) + + with pytest.raises(RuntimeError, match="model failed"): + await consumer.execute_tick(make_tick()) + + +@pytest.mark.asyncio +async def test_session_commits_only_after_standard_engine_output_metadata() -> None: + client = FakeGenerateClient() + consumer = ARDiffusionOmniTickConsumer( + client, + prompt_provider=lambda tick: tick.prompt or "", + sampling_params_list=[OmniDiffusionSamplingParams()], + diffusion_stage_id=0, + ) + session = ARDiffusionSession( + "world-1", + tick_consumer=consumer, + lifecycle=FakeLifecycle(), + max_pending_events=4, + request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}", + ) + await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left")) + + output = await session.next_chunk() + + assert output.request_id == "world-1-chunk-0" + assert session.chunk_index == 1 + assert session.pending_event_count == 0 + assert client.calls[0]["request_id"] == "world-1-chunk-0" diff --git a/tests/diffusion/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py index d09c1c1b6..086666063 100644 --- a/tests/diffusion/ar_diffusion/test_kv_cache.py +++ b/tests/diffusion/ar_diffusion/test_kv_cache.py @@ -30,6 +30,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] BLOCK = 16 +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False): return ChunkWindowSpec( @@ -95,7 +97,6 @@ def test_sliding_replace_snaps_to_chunk_boundary(): def test_sink_chunks_protected(): # sink = 1 chunk (16 tokens) is never skipped. skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False) - # raw overflow = 80 - 32 - 16 = 32 -> snapped 32 assert skip == 32 @@ -201,12 +202,149 @@ def test_cross_attn_pool_deducted_from_self_attn_budget(): kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)), session_capacity=1, cross_attention_lengths={"text": L}, - device=torch.device("cpu"), ) cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers - expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2) + page_bytes = kv.spec.page_size_bytes * 2 + scratch_bytes = kv.scratch_num_blocks * page_bytes + expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes assert kv.num_blocks == expected - assert expected > 1 * (2 + 1) + 2 # above the local-slot floor, so the cross deduction is what's tested + assert kv.cross_attention_reserved_bytes == cross_bytes + assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5) + assert expected > 1 * (2 + 1) + 2 # above the local-slot minimum, so the cross deduction is what's tested + + +def _make_tiny_capacity_kv( + *, + requested_capacity: int, + available_bytes: int, + gpu_memory_fraction: float = 1.0, + model_owned_state_bytes_per_session: int = 0, +) -> ARDiffusionKVCache: + return ARDiffusionKVCache( + ARDiffusionKVConfig( + enable=True, + chunk_size=1, + window_chunks=3, + sink_chunks=3, + gpu_memory_fraction=gpu_memory_fraction, + ), + num_layers=1, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + block_size=1, + max_model_len=128, + available_bytes=available_bytes, + kv_branches=(ARDiffusionKVBranchSpec("main", 0),), + session_capacity=requested_capacity, + cross_attention_lengths={"text": 2}, + frames_per_block=3, + model_owned_state_bytes_per_session=model_owned_state_bytes_per_session, + device=torch.device("cpu"), + ) + + +def test_capacity_two_retains_both_windows_and_allocates_next_block(): + # LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires + # 17 managed + 3 scratch blocks, plus two 16-byte cross-attention + # reservations: 192 bytes exactly. + kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192) + assert kv.requested_session_capacity == 2 + assert kv.session_capacity == 2 + assert kv.managed_num_blocks == 17 + + adapters = [kv.begin_request(session_id) for session_id in ("first", "second")] + for adapter in adapters: + for _ in range(6): # three sink + three window blocks + kv.allocate_token_slots(adapter, 1) + kv.commit_chunk(adapter) + assert len(kv.window_block_ids(adapter)) == 6 + + # Only one request is in flight, so its three-frame block is counted once. + kv.allocate_token_slots(adapters[0], 3) + + +def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity(): + kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192) + + assert kv.requested_session_capacity == 64 + assert kv.session_capacity == 2 + assert kv.cross_attention_bytes_per_session == 16 + assert kv.cross_attention_reserved_bytes == 32 + assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes + + +def test_capacity_rejects_budget_that_cannot_fit_one_session(): + # One session needs 11 managed + 3 scratch blocks and 16 cross bytes: + # 14 * 8 + 16 = 128 bytes. + with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"): + _make_tiny_capacity_kv(requested_capacity=1, available_bytes=127) + + +def test_fraction_is_soft_floor_when_one_session_fits_actual_memory(): + kv = _make_tiny_capacity_kv( + requested_capacity=2, + available_bytes=1280, + gpu_memory_fraction=0.05, + ) + + assert kv.configured_memory_budget_bytes == 64 + assert kv.memory_budget_bytes == 128 + assert kv.session_capacity == 1 + + +def test_model_owned_state_reduces_effective_session_capacity(): + kv = _make_tiny_capacity_kv( + requested_capacity=2, + available_bytes=200, + model_owned_state_bytes_per_session=16, + ) + + assert kv.session_capacity == 1 + assert kv.model_owned_state_reserved_bytes == 16 + assert ( + kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes + <= kv.memory_budget_bytes + ) + + +def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache: + return ARDiffusionKVCache( + ARDiffusionKVConfig( + enable=True, + chunk_size=1560, + window_chunks=9, + sink_chunks=9, + gpu_memory_fraction=gpu_memory_fraction, + ), + num_layers=40, + num_kv_heads=40, + head_size=128, + dtype=torch.bfloat16, + block_size=1560, + max_model_len=1 << 20, + available_bytes=100 * (1 << 30), + kv_branches=(ARDiffusionKVBranchSpec("main", 0),), + session_capacity=2, + cross_attention_lengths={"text": 512}, + frames_per_block=3, + model_owned_state_bytes_per_session=7_488_000, + ) + + +def test_shipped_lingbot_geometry_default_fraction_admits_one_session(): + kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1) + + assert kv.configured_memory_budget_bytes == 10 * (1 << 30) + assert kv.memory_budget_bytes == 33_653_670_400 + assert kv.session_capacity == 1 + + +def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions(): + kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6) + + assert kv.memory_budget_bytes == 60 * (1 << 30) + assert kv.session_capacity == 2 def _make_kv( @@ -221,15 +359,24 @@ def _make_kv( if local_branches == 1 else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1)) ) + page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize + declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK + scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks) + managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2 return ARDiffusionKVCache( - ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks), + ARDiffusionKVConfig( + enable=True, + chunk_size=BLOCK, + window_chunks=window_chunks, + gpu_memory_fraction=1.0, + ), num_layers=1, num_kv_heads=4, head_size=64, dtype=torch.float32, block_size=BLOCK, max_model_len=4096, - available_bytes=1 << 16, # tiny -> the floor binds + available_bytes=(managed_blocks + scratch_blocks) * page_bytes, device=torch.device("cpu"), kv_branches=kv_branches, session_capacity=1, diff --git a/tests/diffusion/ar_diffusion/test_kv_orchestration.py b/tests/diffusion/ar_diffusion/test_kv_orchestration.py index 32603bf88..4482fa6fb 100644 --- a/tests/diffusion/ar_diffusion/test_kv_orchestration.py +++ b/tests/diffusion/ar_diffusion/test_kv_orchestration.py @@ -54,13 +54,15 @@ def test_full_request_lifecycle_and_eviction(): used = {int(s) // kv.block_size for s in slots} assert kv.null_block_id not in used assert used <= set(block_table) - # Resident window stays bounded by window + the in-flight chunk. + # Allocation may transiently add one in-flight chunk. assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1 - free_after.append(kv.manager.block_pool.get_num_free_blocks()) kv.commit_chunk(adapter) + # A successful commit immediately prunes physical ownership. + assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + free_after.append(kv.manager.block_pool.get_num_free_blocks()) # Pool memory plateaus once the window is full (eviction recycles blocks). - assert free_after[-1] == free_after[kv.spec.window_chunks] + assert free_after[-1] == free_after[kv.spec.window_chunks - 1] kv.end_request(adapter) assert kv.manager.block_pool.get_num_free_blocks() == free_total diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py index 0b66c01a0..603a2a175 100644 --- a/tests/diffusion/ar_diffusion/test_paged_attention.py +++ b/tests/diffusion/ar_diffusion/test_paged_attention.py @@ -60,9 +60,18 @@ def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor return torch.einsum("bhqk,bkhd->bqhd", probs, value) -def _cuda_flash_attn_usable() -> bool: +def _gpu_flash_attn_usable() -> bool: if not torch.cuda.is_available(): return False + if torch.version.hip is not None: + for module in ("aiter", "flash_attn"): + try: + imported = __import__(module, fromlist=["flash_attn_varlen_func"]) + if getattr(imported, "flash_attn_varlen_func", None) is not None: + return True + except ImportError: + pass + return False try: spec = find_spec("vllm.vllm_flash_attn") if spec is None or spec.origin is None: @@ -227,12 +236,11 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0) -@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required") +@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required") @pytest.mark.parametrize("history_chunks", [1, 3]) @pytest.mark.parametrize("action_len", [0, 3]) @pytest.mark.parametrize("commit_current", [False, True]) def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current): - pytest.importorskip("vllm.vllm_flash_attn") torch.manual_seed(0) device = torch.device("cuda") dtype = torch.float16 @@ -324,7 +332,7 @@ def test_block_table_padded_to_fixed_width(): def test_prepare_is_idempotent_and_layers_share_metadata(): device = torch.device("cpu") - _, st = make_state(num_layers=2) + kv, st = make_state(num_layers=2) contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False) fctx = contexts[0].forward_ctx fctx.prepare(device=device, action_len=0, query_len=BLOCK) @@ -341,6 +349,10 @@ def test_prepare_is_idempotent_and_layers_share_metadata(): assert i0.block_table is i1.block_table assert i0.seq_lens is i1.seq_lens assert i0.video_slots is i1.video_slots + assert i0.key_pool is kv._k_pools[0] + assert i0.value_pool is kv._v_pools[0] + assert i1.key_pool is kv._k_pools[1] + assert i1.value_pool is kv._v_pools[1] def test_layer_inputs_before_prepare_raises(): diff --git a/tests/diffusion/ar_diffusion/test_pipeline_bridge.py b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py index ec60290b9..ed5233800 100644 --- a/tests/diffusion/ar_diffusion/test_pipeline_bridge.py +++ b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py @@ -174,6 +174,19 @@ def test_window_reset_can_keep_named_cross_attention_cache(): assert not st.is_cross_attention_populated(POS, "text") +def test_clear_cross_attention_preserves_self_attention_history(): + kv, st = make_state(num_layers=1, cross_attn_length=8) + _prepare_and_commit(st, POS, 2) + k = torch.randn(1, 8, N_HEADS, HEAD_DIM) + st.populate_cross_attention(POS, "text", [(k, k)]) + resident_before = kv.window_block_ids(st.adapter(POS)) + + st.clear_cross_attention() + + assert kv.window_block_ids(st.adapter(POS)) == resident_before + assert not st.is_cross_attention_populated(POS, "text") + + def test_cross_attention_uses_logical_branch_when_local_slot_is_shared(): _, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True) assert st.kv_cache.num_local_kv_branches == 1 diff --git a/tests/diffusion/ar_diffusion/test_session.py b/tests/diffusion/ar_diffusion/test_session.py new file mode 100644 index 000000000..a31be388b --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_session.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from vllm_omni.experimental.ar_diffusion.session import ( + ARDiffusionChunkMetadataError, + ARDiffusionEventAcceptanceStatus, + ARDiffusionEventQueueFullError, + ARDiffusionPreparedControls, + ARDiffusionSession, + ARDiffusionSessionError, + ARDiffusionSessionEvent, + ARDiffusionSessionManager, + ARDiffusionSessionStateError, + ARDiffusionSessionStatus, + ARDiffusionStaleEventError, + ARDiffusionWorkerLifecycle, +) +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionChunkMetadata, + ARDiffusionControlInput, + ARDiffusionTickRequest, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeTickConsumer: + def __init__(self) -> None: + self.ticks: list[ARDiffusionTickRequest] = [] + self.failures_remaining = 0 + self.started: asyncio.Event | None = None + self.release: asyncio.Event | None = None + self.metadata_override: ARDiffusionChunkMetadata | None = None + + async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata: + self.ticks.append(tick) + if self.started is not None: + self.started.set() + if self.release is not None: + await self.release.wait() + if self.failures_remaining: + self.failures_remaining -= 1 + raise RuntimeError("fake tick failed") + return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick) + + def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata: + return output + + +class FakeLifecycle: + def __init__(self) -> None: + self.resets: list[str] = [] + self.closes: list[str] = [] + self.close_failures_remaining = 0 + + async def reset_session(self, session_id: str) -> None: + self.resets.append(session_id) + + async def close_session(self, session_id: str) -> None: + self.closes.append(session_id) + if self.close_failures_remaining: + self.close_failures_remaining -= 1 + raise RuntimeError("fake close failed") + + +class FailingCommitReducer: + def __init__(self) -> None: + self.commit_calls = 0 + self.reset_calls = 0 + + def prepare( + self, + *, + current_controls: dict[str, ARDiffusionControlInput], + events: tuple[ARDiffusionSessionEvent, ...], + chunk_index: int, + ) -> ARDiffusionPreparedControls: + del current_controls, events, chunk_index + return ARDiffusionPreparedControls(controls=()) + + def commit(self, prepared: ARDiffusionPreparedControls) -> None: + del prepared + self.commit_calls += 1 + raise RuntimeError("fake reducer commit failed") + + def reset(self) -> None: + self.reset_calls += 1 + + +def make_session( + *, + max_pending_events: int = 8, + consumer: FakeTickConsumer | None = None, + lifecycle: FakeLifecycle | None = None, +) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]: + resolved_consumer = consumer or FakeTickConsumer() + resolved_lifecycle = lifecycle or FakeLifecycle() + session = ARDiffusionSession( + "world-1", + tick_consumer=resolved_consumer, + lifecycle=resolved_lifecycle, + max_pending_events=max_pending_events, + request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}", + ) + return session, resolved_consumer, resolved_lifecycle + + +@pytest.mark.asyncio +async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None: + session, consumer, _ = make_session() + event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn") + event_one = ARDiffusionSessionEvent( + event_id=1, + controls=( + ARDiffusionControlInput( + track="camera", + schema="example.camera.v1", + data={"samples": [{"x": 1.0}]}, + ), + ), + ) + + accepted_two = await session.accept_event(event_two) + accepted_one = await session.accept_event(event_one) + duplicate = await session.accept_event(event_two) + + assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED + assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED + assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE + assert consumer.ticks == [] + + metadata = await session.next_chunk() + + assert metadata.applied_event_ids == (1, 2) + assert consumer.ticks[0] == ARDiffusionTickRequest( + session_id="world-1", + request_id="world-1-request-0", + chunk_index=0, + applied_event_ids=(1, 2), + prompt="turn", + controls=( + ARDiffusionControlInput( + track="camera", + schema="example.camera.v1", + data={"samples": [{"x": 1.0}]}, + ), + ), + ) + assert session.pending_event_count == 0 + + applied_duplicate = await session.accept_event(event_two) + assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE + + next_metadata = await session.next_chunk() + assert next_metadata.chunk_index == 1 + assert next_metadata.applied_event_ids == () + assert consumer.ticks[1].prompt == "turn" + assert consumer.ticks[1].controls == consumer.ticks[0].controls + + +@pytest.mark.asyncio +async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None: + session, consumer, _ = make_session(max_pending_events=1) + payload = {"samples": [{"x": 1.0}]} + composite = ARDiffusionSessionEvent( + event_id=5, + prompt="new prompt", + controls=( + ARDiffusionControlInput( + track="camera", + schema="example.camera.v1", + data=payload, + ), + ARDiffusionControlInput( + track="action", + schema="example.action.v1", + data={"keys": ["left"]}, + ), + ), + ) + await session.accept_event(composite) + payload["samples"][0]["x"] = 99.0 + + duplicate = await session.accept_event(composite) + assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE + with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"): + await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected")) + + await session.next_chunk() + assert consumer.ticks[0].applied_event_ids == (5,) + assert consumer.ticks[0].prompt == "new prompt" + assert consumer.ticks[0].controls[0].data["keys"] == ("left",) + assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},) + assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]} + assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]} + + +def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None: + with pytest.raises(ValueError, match="at most once"): + ARDiffusionSessionEvent( + event_id=1, + controls=( + ARDiffusionControlInput(track="camera", schema="one", data={}), + ARDiffusionControlInput(track="camera", schema="two", data={}), + ), + ) + + with pytest.raises(ValueError, match="transport-safe"): + ARDiffusionControlInput( + track="camera", + schema="example.camera.v1", + data={"object": object()}, + ) + + +@pytest.mark.asyncio +async def test_failed_tick_closes_worker_and_requires_reset() -> None: + consumer = FakeTickConsumer() + consumer.failures_remaining = 1 + lifecycle = FakeLifecycle() + session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle) + await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me")) + + with pytest.raises(RuntimeError, match="fake tick failed"): + await session.next_chunk() + + assert session.status is ARDiffusionSessionStatus.FAILED + assert session.chunk_index == 0 + assert session.pending_event_count == 1 + assert lifecycle.closes == ["world-1"] + with pytest.raises(ARDiffusionSessionStateError, match="failed"): + await session.next_chunk() + + await session.reset() + assert session.status is ARDiffusionSessionStatus.ACTIVE + assert session.pending_event_count == 0 + assert lifecycle.resets == ["world-1"] + await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset")) + metadata = await session.next_chunk() + assert metadata.chunk_index == 0 + assert metadata.applied_event_ids == (4,) + assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)] + + +@pytest.mark.asyncio +async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None: + consumer = FakeTickConsumer() + consumer.started = asyncio.Event() + consumer.release = asyncio.Event() + session, _, _ = make_session(consumer=consumer) + await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first")) + + first_chunk = asyncio.create_task(session.next_chunk()) + await consumer.started.wait() + accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second")) + assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED + assert consumer.ticks[0].prompt == "first" + assert consumer.ticks[0].applied_event_ids == (10,) + + with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"): + await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale")) + + consumer.release.set() + first_metadata = await first_chunk + second_metadata = await session.next_chunk() + assert first_metadata.applied_event_ids == (10,) + assert second_metadata.applied_event_ids == (11,) + assert consumer.ticks[1].prompt == "second" + + +@pytest.mark.asyncio +async def test_metadata_mismatch_does_not_commit_snapshot() -> None: + consumer = FakeTickConsumer() + consumer.metadata_override = ARDiffusionChunkMetadata( + session_id="another-world", + request_id="wrong", + chunk_index=0, + ) + lifecycle = FakeLifecycle() + session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle) + await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending")) + + with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"): + await session.next_chunk() + + assert session.chunk_index == 0 + assert session.pending_event_count == 1 + assert session.status is ARDiffusionSessionStatus.FAILED + assert lifecycle.closes == ["world-1"] + + +@pytest.mark.asyncio +async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None: + consumer = FakeTickConsumer() + lifecycle = FakeLifecycle() + reducer = FailingCommitReducer() + session = ARDiffusionSession( + "world-1", + tick_consumer=consumer, + lifecycle=lifecycle, + max_pending_events=4, + request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}", + control_reducer=reducer, + ) + await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending")) + + with pytest.raises(RuntimeError, match="fake reducer commit failed"): + await session.next_chunk() + + assert reducer.commit_calls == 1 + assert session.status is ARDiffusionSessionStatus.FAILED + assert session.chunk_index == 0 + assert session.pending_event_count == 1 + assert lifecycle.closes == ["world-1"] + with pytest.raises(ARDiffusionSessionStateError, match="failed"): + await session.next_chunk() + + await session.reset() + assert reducer.reset_calls == 1 + assert session.status is ARDiffusionSessionStatus.ACTIVE + assert session.chunk_index == 0 + assert session.pending_event_count == 0 + + +@pytest.mark.asyncio +async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None: + consumer = FakeTickConsumer() + lifecycle = FakeLifecycle() + manager = ARDiffusionSessionManager( + tick_consumer=consumer, + lifecycle=lifecycle, + max_pending_events=4, + request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}", + ) + session = await manager.create_session("world-reset") + await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset")) + await session.next_chunk() + await manager.reset_session("world-reset") + + assert lifecycle.resets == ["world-reset"] + assert session.status is ARDiffusionSessionStatus.ACTIVE + assert session.chunk_index == 0 + duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate")) + assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE + + reset_tick = await session.next_chunk() + assert reset_tick.chunk_index == 0 + assert consumer.ticks[-1].prompt is None + assert consumer.ticks[-1].controls == () + + await manager.close_session("world-reset") + assert lifecycle.closes == ["world-reset"] + assert session.status is ARDiffusionSessionStatus.CLOSED + with pytest.raises(ARDiffusionSessionStateError, match="does not exist"): + await manager.get_session("world-reset") + + disconnected = await manager.create_session("world-disconnect") + await manager.disconnect("world-disconnect") + assert disconnected.status is ARDiffusionSessionStatus.CLOSED + assert lifecycle.closes[-1] == "world-disconnect" + + +@pytest.mark.asyncio +async def test_manager_keeps_interleaved_sessions_independent() -> None: + consumer = FakeTickConsumer() + manager = ARDiffusionSessionManager( + tick_consumer=consumer, + lifecycle=FakeLifecycle(), + max_pending_events=4, + request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}", + ) + session_a = await manager.create_session("world-a") + session_b = await manager.create_session("world-b") + + await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0")) + await session_a.next_chunk() + await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0")) + await session_b.next_chunk() + await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1")) + await session_a.next_chunk() + + assert [ + ( + tick.session_id, + tick.request_id, + tick.chunk_index, + tick.applied_event_ids, + ) + for tick in consumer.ticks + ] == [ + ("world-a", "world-a-request-0", 0, (10,)), + ("world-b", "world-b-request-0", 0, (1,)), + ("world-a", "world-a-request-1", 1, (11,)), + ] + assert session_a.chunk_index == 2 + assert session_b.chunk_index == 1 + with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"): + await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a")) + assert ( + await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1")) + ).status is ARDiffusionEventAcceptanceStatus.ACCEPTED + + +@pytest.mark.asyncio +async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None: + consumer = FakeTickConsumer() + lifecycle = FakeLifecycle() + lifecycle.close_failures_remaining = 1 + manager = ARDiffusionSessionManager( + tick_consumer=consumer, + lifecycle=lifecycle, + max_pending_events=4, + ) + session = await manager.create_session("world-close-retry") + + with pytest.raises(RuntimeError, match="fake close failed"): + await manager.close_session("world-close-retry") + + assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED + assert await manager.get_session("world-close-retry") is session + with pytest.raises(ARDiffusionSessionStateError, match="already exists"): + await manager.create_session("world-close-retry") + with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"): + await manager.reset_session("world-close-retry") + + await manager.close_session("world-close-retry") + + assert lifecycle.closes == ["world-close-retry", "world-close-retry"] + assert session.status is ARDiffusionSessionStatus.CLOSED + with pytest.raises(ARDiffusionSessionStateError, match="does not exist"): + await manager.get_session("world-close-retry") + recreated = await manager.create_session("world-close-retry") + assert recreated is not session + + +@pytest.mark.asyncio +async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None: + lifecycle = FakeLifecycle() + lifecycle.close_failures_remaining = 1 + manager = ARDiffusionSessionManager( + tick_consumer=FakeTickConsumer(), + lifecycle=lifecycle, + max_pending_events=4, + ) + session = await manager.create_session("world-disconnect-retry") + + with pytest.raises(RuntimeError, match="fake close failed"): + await manager.disconnect("world-disconnect-retry") + + assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED + assert await manager.get_session("world-disconnect-retry") is session + + await manager.close_session("world-disconnect-retry") + assert session.status is ARDiffusionSessionStatus.CLOSED + + +class FakeRPCClient: + def __init__(self, result: list[Any]) -> None: + self.result = result + self.calls: list[dict[str, Any]] = [] + + async def collective_rpc(self, **kwargs: Any) -> list[Any]: + self.calls.append(kwargs) + return self.result + + +@pytest.mark.asyncio +async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None: + client = FakeRPCClient([[True]]) + lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0) + + await lifecycle.reset_session("world-1") + await lifecycle.close_session("world-1") + + assert client.calls == [ + { + "method": "reset_ar_diffusion_session", + "timeout": 3.0, + "args": ("world-1",), + "stage_ids": [2], + }, + { + "method": "close_ar_diffusion_session", + "timeout": 3.0, + "args": ("world-1",), + "stage_ids": [2], + }, + ] + + +@pytest.mark.asyncio +async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None: + lifecycle = ARDiffusionWorkerLifecycle( + FakeRPCClient([{"supported": False, "error": "not an AR stage"}]), + stage_ids=[1], + ) + + with pytest.raises(ARDiffusionSessionError, match="not an AR stage"): + await lifecycle.close_session("world-1") diff --git a/tests/diffusion/ar_diffusion/test_tick_protocol.py b/tests/diffusion/ar_diffusion/test_tick_protocol.py new file mode 100644 index 000000000..70e71ba80 --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_tick_protocol.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionChunkMetadata, + ARDiffusionControlInput, + ARDiffusionTickRequest, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None: + tick = ARDiffusionTickRequest( + session_id="world-1", + request_id="chunk-request-7", + chunk_index=7, + applied_event_ids=(10, 11), + prompt="turn left", + controls=( + ARDiffusionControlInput( + track="camera", + schema="lingbot.camera_trajectory.v1", + data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]}, + ), + ), + ) + + restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args()) + + assert restored == tick + assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == { + "session_id": "world-1", + "request_id": "chunk-request-7", + "chunk_index": 7, + "applied_event_ids": [10, 11], + } + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("session_id", "", "session_id"), + ("request_id", "", "request_id"), + ("chunk_index", -1, "chunk_index"), + ("applied_event_ids", (2, 1), "strictly increasing"), + ("applied_event_ids", (1, 1), "strictly increasing"), + ], +) +def test_tick_rejects_invalid_identity( + field: str, + value: object, + message: str, +) -> None: + values = { + "session_id": "world-1", + "request_id": "request-1", + "chunk_index": 0, + "applied_event_ids": (), + } + values[field] = value + + with pytest.raises(ValueError, match=message): + ARDiffusionTickRequest(**values) + + +def test_tick_snapshot_is_detached_and_deeply_immutable() -> None: + payload = {"samples": [{"keys": ["w"]}]} + event_ids = [1, 2] + controls = [ + ARDiffusionControlInput( + track="camera", + schema="example.camera.v1", + data=payload, + ) + ] + tick = ARDiffusionTickRequest( + session_id="world-1", + request_id="request-a", + chunk_index=0, + applied_event_ids=event_ids, # type: ignore[arg-type] + controls=controls, # type: ignore[arg-type] + ) + serialized = tick.to_extra_args() + + payload["samples"][0]["keys"].append("j") + event_ids.append(3) + controls.clear() + serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k") + + assert tick.applied_event_ids == (1, 2) + assert isinstance(tick.controls, tuple) + assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [ + { + "track": "camera", + "schema": "example.camera.v1", + "data": {"samples": [{"keys": ["w"]}]}, + } + ] + with pytest.raises(TypeError): + tick.controls[0].data["samples"] = () # type: ignore[index] + with pytest.raises(AttributeError): + tick.controls[0].data["samples"][0]["keys"].append("j") + + +def test_legacy_request_has_no_typed_tick() -> None: + assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None diff --git a/tests/diffusion/ar_diffusion/test_worker_lifecycle.py b/tests/diffusion/ar_diffusion/test_worker_lifecycle.py new file mode 100644 index 000000000..42204049c --- /dev/null +++ b/tests/diffusion/ar_diffusion/test_worker_lifecycle.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeARDiffusionRunner: + def __init__(self) -> None: + self.resets: list[str] = [] + self.closes: list[str] = [] + + def reset_session(self, session_id: str) -> None: + self.resets.append(session_id) + + def close_session(self, session_id: str) -> None: + self.closes.append(session_id) + + +class FakeRegularRunner: + pass + + +def make_worker(runner: object) -> DiffusionWorker: + worker = object.__new__(DiffusionWorker) + worker.model_runner = runner + return worker + + +def test_worker_delegates_ar_session_lifecycle_to_runner() -> None: + runner = FakeARDiffusionRunner() + worker = make_worker(runner) + + assert worker.reset_ar_diffusion_session("world-1") is True + assert worker.close_ar_diffusion_session("world-1") is True + assert runner.resets == ["world-1"] + assert runner.closes == ["world-1"] + + +def test_worker_reports_unsupported_runner_without_model_assumptions() -> None: + worker = make_worker(FakeRegularRunner()) + + assert worker.reset_ar_diffusion_session("world-1") is False + assert worker.close_ar_diffusion_session("world-1") is False + + +def test_worker_rejects_empty_session_id() -> None: + worker = make_worker(FakeARDiffusionRunner()) + + with pytest.raises(ValueError, match="session_id"): + worker.close_ar_diffusion_session("") diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 4a2c65cbb..54ad2a32b 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for per-role attention backend configuration (RFC: per-role-attention-backend). @@ -11,6 +11,7 @@ Tests cover: """ from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -90,10 +91,23 @@ class TestAttentionSpec: with pytest.raises(ValueError, match=match): AttentionSpec(**spec) + def test_fastvideo_vsa_topk_serialized(self): + spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=96) + assert spec.backend_kwargs() == {"topk": 96} + + def test_fastvideo_vsa_topk_rejected_for_other_backend(self): + with pytest.raises(ValueError, match="only supported by the FASTVIDEO_VSA"): + AttentionSpec(backend="TORCH_SDPA", fastvideo_vsa_topk=96) + def test_block_sparse_defaults_applied_when_backend_selected(self): spec = AttentionSpec(backend="RAINFUSION_ATTN") assert spec.block_sparse.sparsity == 0.8 - assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0} + assert spec.backend_kwargs() == { + "sparsity": 0.8, + "start_step": 0, + "end_step": 0, + "precision": "bf16", + } def test_block_sparse_skip_layers_selector_expanded(self): spec = AttentionSpec( @@ -104,6 +118,8 @@ class TestAttentionSpec: assert spec.backend_kwargs() == { "sparsity": 0.9, "start_step": 12, + "end_step": 0, + "precision": "bf16", "skip_layers": [0, 1, 2, 38], } @@ -113,12 +129,44 @@ class TestAttentionSpec: @pytest.mark.parametrize( "block_sparse", - [{"sparsity": 1.5}, {"start_step": -1}], + [{"sparsity": 1.5}, {"start_step": -1}, {"end_step": -1}, {"precision": "bogus"}], ) def test_block_sparse_invalid_values(self, block_sparse): with pytest.raises(ValueError): AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse) + def test_block_sparse_end_step_precision_forwarded(self): + spec = AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse={"sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8"}, + ) + assert spec.backend_kwargs() == { + "sparsity": 0.9, + "start_step": 5, + "end_step": 3, + "precision": "fp8", + } + + def test_block_sparse_precision_enum_input_normalized(self): + """RainFusionPrecision enum members must normalize to their string value. + + Regression: the str,Enum members passed the value check but str() later + produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it). + """ + from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig + from vllm_omni.diffusion.data import RainFusionPrecision + + spec = AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse={"sparsity": 0.8, "precision": RainFusionPrecision.MIX}, + ) + kwargs = spec.backend_kwargs() + assert kwargs["precision"] == "mix" + # The config path is where str(value) used to mangle enum members. + cfg = RainFusionConfig.from_backend_kwargs(kwargs) + assert cfg.precision == "mix" + assert type(cfg.precision) is str + class TestAttentionConfig: def test_empty_config(self): @@ -419,6 +467,120 @@ class TestCurrentDiffusionConfig: class TestAttentionInitUsesCurrentDiffusionConfig: + @pytest.mark.parametrize("mindiesd_available", [False, True]) + def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available): + attention = Attention.__new__(Attention) + attention._scheduler_paged_kv = True + attention.paged_kv_cache_role = "primary" + attention.backend_pref = "FLASH_ATTN" + attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN") + flash_output = torch.ones(1) + sdpa_output = torch.zeros(1) + dense_flash = Mock(return_value=flash_output) + sdpa = Mock(return_value=sdpa_output) + attention.attention = SimpleNamespace(forward=dense_flash) + attention.sdpa_fallback = SimpleNamespace(forward=sdpa) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True) + monkeypatch.setattr( + layer_mod, + "get_forward_context", + lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True), + ) + monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: True) + monkeypatch.setattr( + layer_mod.current_omni_platform, + "supports_diffusion_dense_flash_attention", + lambda: mindiesd_available, + ) + + result = attention._run_local_attention( + torch.zeros(1, dtype=torch.bfloat16), + torch.zeros(1, dtype=torch.bfloat16), + torch.zeros(1, dtype=torch.bfloat16), + AttentionMetadata(), + ) + + if mindiesd_available: + assert result is flash_output + dense_flash.assert_called_once() + sdpa.assert_not_called() + else: + assert result is sdpa_output + dense_flash.assert_not_called() + sdpa.assert_called_once() + + def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch): + attention = Attention.__new__(Attention) + attention._scheduler_paged_kv = False + attention.paged_kv_cache_role = "primary" + attention.backend_pref = "FLASH_ATTN" + attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN") + attention.attention = SimpleNamespace( + forward=Mock(side_effect=ModuleNotFoundError("No module named 'mindiesd'", name="mindiesd")) + ) + attention.sdpa_fallback = SimpleNamespace(forward=Mock(side_effect=AssertionError("must not fall back"))) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + + with pytest.raises(ModuleNotFoundError, match="mindiesd"): + attention._run_local_attention( + torch.zeros(1, dtype=torch.bfloat16), + torch.zeros(1, dtype=torch.bfloat16), + torch.zeros(1, dtype=torch.bfloat16), + AttentionMetadata(), + ) + + def test_paged_marker_does_not_change_dense_backend(self, monkeypatch): + class _FakeAttentionImpl: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def forward(self, query, key, value, attn_metadata=None): + return query + + class _DenseBackend: + supports_paged_kv = False + + @staticmethod + def get_name() -> str: + return "TORCH_SDPA" + + @staticmethod + def get_impl_cls(): + return _FakeAttentionImpl + + calls = [] + + def _fake_get_attn_backend_for_role(*args, **kwargs): + calls.append((args, kwargs)) + return _DenseBackend, None + + monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role) + monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object()) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(), + diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.DENSE_LEGACY, + parallel_config=SimpleNamespace(ring_degree=1), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config): + attention = Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + paged_kv_cache_role="primary", + ) + + assert len(calls) == 1 + assert attention.attn_backend is _DenseBackend + assert attention.backend_pref is None + def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch): class _FakeAttentionImpl: def __init__(self, **kwargs): diff --git a/tests/diffusion/attention/test_cudnn_attn.py b/tests/diffusion/attention/test_cudnn_attn.py index de1adf8c7..1ff3baafb 100644 --- a/tests/diffusion/attention/test_cudnn_attn.py +++ b/tests/diffusion/attention/test_cudnn_attn.py @@ -5,6 +5,7 @@ from contextlib import nullcontext import pytest import torch +from torch.nn.attention import SDPBackend from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl @@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] +@pytest.mark.parametrize( + ("query_sequence_length", "key_sequence_length", "expected_backend"), + [ + (1, 1, SDPBackend.MATH), + (8, 1, SDPBackend.MATH), + (1, 8, SDPBackend.MATH), + (8, 8, SDPBackend.CUDNN_ATTENTION), + ], +) +def test_cudnn_routes_single_token_attention_to_math( + monkeypatch, query_sequence_length, key_sequence_length, expected_backend +): + selected_backends = [] + + def fake_sdpa_kernel(backends): + selected_backends.append(backends) + return nullcontext() + + monkeypatch.setattr( + "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel", + fake_sdpa_kernel, + ) + monkeypatch.setattr( + torch.nn.functional, + "scaled_dot_product_attention", + lambda query, _key, _value, **_kwargs: query, + ) + impl = CuDNNAttentionImpl( + num_heads=2, + head_size=4, + softmax_scale=0.5, + ) + query = torch.randn(1, query_sequence_length, 2, 4) + key = torch.randn(1, key_sequence_length, 2, 4) + value = torch.randn_like(key) + + output = impl.forward_cuda(query, key, value) + + assert output.shape == query.shape + assert selected_backends == [[expected_backend]] + + def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch): observed = {} diff --git a/tests/diffusion/attention/test_fastvideo_vsa.py b/tests/diffusion/attention/test_fastvideo_vsa.py new file mode 100644 index 000000000..d6209bf42 --- /dev/null +++ b/tests/diffusion/attention/test_fastvideo_vsa.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import sys +import types +from typing import Any + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.fastvideo_vsa import ( + FastVideoVSAImpl, +) +from vllm_omni.diffusion.attention.backends.registry import ( + DiffusionAttentionBackendEnum, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_fastvideo_vsa_backend_is_registered(): + assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend") + + +def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch): + calls: dict[str, Any] = {} + fake_module = types.ModuleType("fastvideo_kernel") + + def fake_video_sparse_attn_bshd( + q, + k, + v, + variable_block_sizes, + q_variable_block_sizes, + compress_attn_weight, + topk, + block_size, + ): + calls["q_shape"] = tuple(q.shape) + calls["vbs"] = variable_block_sizes.detach().cpu().tolist() + calls["q_vbs"] = q_variable_block_sizes.detach().cpu().tolist() + calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum()) + calls["compress_shape"] = tuple(compress_attn_weight.shape) + calls["topk"] = topk + calls["block_size"] = block_size + return q + k + v + + setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd) + monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module) + + impl = FastVideoVSAImpl( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + backend_kwargs={ + "topk": 1, + "block_size": (4, 8, 8), + "min_seq_len": 1, + "disable_when_sp_active": False, + }, + ) + query = torch.randn(1, 300, 2, 8) + key = torch.randn_like(query) + value = torch.randn_like(query) + attn_metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10)}) + monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None) + + output = impl.forward_cuda(query, key, value, attn_metadata) + + assert output.shape == query.shape + assert calls["q_shape"] == (1, 1024, 2, 8) + assert calls["vbs"] == [192, 48, 48, 12] + assert calls["q_vbs"] == [192, 48, 48, 12] + assert calls["compress_shape"] == (1, 1024, 2, 8) + assert calls["compress_sum"] == 0.0 + assert calls["topk"] == 1 + assert calls["block_size"] == (4, 8, 8) + + +def test_fastvideo_vsa_uses_learned_gate_when_provided(monkeypatch): + calls: dict[str, Any] = {} + fake_module = types.ModuleType("fastvideo_kernel") + + def fake_video_sparse_attn_bshd( + q, + k, + v, + variable_block_sizes, + q_variable_block_sizes, + compress_attn_weight, + topk, + block_size, + ): + calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum()) + calls["compress_shape"] = tuple(compress_attn_weight.shape) + return q + k + v + + setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd) + monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module) + + impl = FastVideoVSAImpl( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + backend_kwargs={ + "topk": 1, + "block_size": (4, 8, 8), + "min_seq_len": 1, + "disable_when_sp_active": False, + }, + ) + query = torch.randn(1, 300, 2, 8) + gate = torch.ones_like(query) + metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10), "gate_compress": gate}) + monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None) + + output = impl.forward_cuda(query, query, query, metadata) + + assert output.shape == query.shape + assert calls["compress_shape"] == (1, 1024, 2, 8) + # Only the 300 real tokens carry ones; padded gate slots stay zero. + assert calls["compress_sum"] == gate.numel() + + +def test_fastvideo_vsa_falls_back_without_dit_shape(monkeypatch): + calls: dict[str, Any] = {} + + impl = FastVideoVSAImpl( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + backend_kwargs={ + "topk": 1, + "block_size": (4, 8, 8), + "min_seq_len": 1, + "disable_when_sp_active": False, + }, + ) + + def fake_fallback(query, key, value, attn_metadata, reason): + calls["reason"] = reason + return torch.zeros_like(query) + + monkeypatch.setattr(impl, "_fallback", fake_fallback) + + query = torch.randn(1, 512, 2, 8) + output = impl.forward_cuda(query, query, query, AttentionMetadata()) + + assert output.shape == query.shape + assert calls["reason"] == "vsa_dit_seq_shape metadata is required" + + +def test_fastvideo_vsa_falls_back_for_mask(monkeypatch): + calls: dict[str, Any] = {} + + impl = FastVideoVSAImpl( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + backend_kwargs={"min_seq_len": 1}, + ) + + def fake_fallback(query, key, value, attn_metadata, reason): + calls["reason"] = reason + return torch.zeros_like(query) + + monkeypatch.setattr(impl, "_fallback", fake_fallback) + + query = torch.randn(1, 512, 2, 8) + mask = torch.ones(1, 512, dtype=torch.bool) + output = impl.forward_cuda(query, query, query, AttentionMetadata(attn_mask=mask)) + + assert output.shape == query.shape + assert calls["reason"] + + +def test_fastvideo_vsa_allows_topk_equal_to_num_blocks(): + query = torch.randn(1, 512, 2, 8) + metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (4, 8, 16)}) + + all_blocks = FastVideoVSAImpl( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + backend_kwargs={"topk": 2, "block_size": (4, 8, 8), "min_seq_len": 1}, + ) + too_many = FastVideoVSAImpl( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + backend_kwargs={"topk": 3, "block_size": (4, 8, 8), "min_seq_len": 1}, + ) + + # CPU/float32 is rejected later, but k=N itself must not trigger fallback. + assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported" + assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2" diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index e3f5355bf..3b5c877cb 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -265,10 +265,15 @@ def test_fa_vs_sdpa(): print(f"Max absolute difference (valid region): {max_diff:.6f}") print(f"Mean absolute difference (valid region): {mean_diff:.6f}") - # Assert that outputs are close - # Using higher tolerance for bfloat16 and different implementations - assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01" - assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001" + # AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by + # about 1.3% relative error. An absolute-only bound misclassifies that + # expected rounding difference for larger output values (for example, + # -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere. + if current_omni_platform.is_rocm(): + torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2) + else: + assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01" + assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001" print("✓ Case 2 PASSED: FA and SDPA outputs are very close!") @@ -303,6 +308,98 @@ def test_flash_attn_func_preferred_over_varlen(): print("✓ flash_attn_func forward works correctly!") +@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU") +@pytest.mark.parametrize("k_len", [40, 256]) +def test_cross_attn_key_padding_vs_sdpa(k_len): + """ + Case 3: Cross-attention with a key-padding mask. + + The mask covers only the key sequence; every query position must attend + to the valid keys. FA and SDPA outputs should be very close. k_len=256 + covers equal Q/K lengths, which must not be mistaken for self-attention, + and valid key lengths differ per batch element. + """ + device = torch.device(current_omni_platform.device_type) + dtype = torch.bfloat16 + + batch_size = 2 + q_len = 256 + valid_lens = (17, 29) + num_heads = 8 + head_dim = 64 + + fa_impl = FlashAttentionImpl( + num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross" + ) + sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False) + + torch.manual_seed(7) + query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype) + key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype) + value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype) + + attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device) + for i, valid_len in enumerate(valid_lens): + attn_mask[i, :valid_len] = True + attn_metadata = AttentionMetadata(attn_mask=attn_mask) + + output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata) + output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata) + + max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item() + mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item() + + print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===") + print(f"Max absolute difference: {max_diff:.6f}") + print(f"Mean absolute difference: {mean_diff:.6f}") + + # Match the platform-specific comparison used by test_fa_vs_sdpa above. + if current_omni_platform.is_rocm(): + # AITER FlashAttention and PyTorch SDPA use different BF16 reduction + # implementations on ROCm, so individual values can differ by one BF16 + # step even though the outputs agree within the expected relative error. + torch.testing.assert_close(output_fa, output_sdpa, rtol=1e-2, atol=1.5e-2) + else: + assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01" + assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001" + + print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!") + + +def test_varlen_masked_routing_by_role(monkeypatch): + """The unpad route is picked by role, not Q/K length equality; runs without a GPU.""" + calls = [] + + def fake_varlen_func(q, k, v, **kwargs): + calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"])) + return q + + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func) + + batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8 + query = torch.randn(batch_size, seq_len, num_heads, head_dim) + metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]])) + + cross_impl = FlashAttentionImpl( + num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross" + ) + output = cross_impl.forward_cuda(query, query, query, metadata) + assert output.shape == query.shape + num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1] + assert num_q_rows == batch_size * seq_len + assert cu_seqlens_q.tolist() == [0, 4, 8] + assert cu_seqlens_k.tolist() == [0, 2, 5] + + self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False) + output = self_impl.forward_cuda(query, query, query, metadata) + assert output.shape == query.shape + num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1] + assert num_q_rows == 5 + assert cu_seqlens_q.tolist() == [0, 2, 5] + assert cu_seqlens_k.tolist() == [0, 2, 5] + + def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch): calls = [] diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py index 05967e28f..fc4746773 100644 --- a/tests/diffusion/attention/test_piecewise_attn.py +++ b/tests/diffusion/attention/test_piecewise_attn.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end test for ``piecewise_attn`` (CPU). Verify that running attention in segments (causal outside full-attn spans, @@ -23,8 +23,10 @@ import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_ from vllm_omni.diffusion.attention.backends.abstract import QueryRange from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( Segment, + build_paged_piecewise_plan, build_segments, piecewise_attn, + run_paged_piecewise_plan, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -84,6 +86,234 @@ def test_build_segments_full_span_retains_global_kv_end(): assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")] +def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints(): + full_attn_spans = [[(2, 5)], [(4, 7)]] + plan = build_paged_piecewise_plan( + full_attn_spans, + query_offsets=[0, 2], + query_lens=[6, 6], + seq_lens=[6, 8], + device=DEVICE, + ) + query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4) + key = query + 100 + value = query + 200 + calls: list[object] = [] + + output_buffer = torch.empty_like(query) + + def run_segment(segment_query, segment_key, segment_value, segment_batch, segment_output_buffer): + assert segment_output_buffer is None + calls.append( + ( + [segment.mode for segment in segment_batch.row_segments], + [segment.kv_end for segment in segment_batch.row_segments], + segment_batch.query_indices.tolist(), + ) + ) + torch.testing.assert_close(segment_key, segment_query + 100) + torch.testing.assert_close(segment_value, segment_query + 200) + return segment_query + + output = run_paged_piecewise_plan( + query, + key, + value, + plan, + plan.segments, + run_segment, + output_buffer=output_buffer, + ) + + assert output is output_buffer + torch.testing.assert_close(output, query) + assert all(segment.query_range is None for segment in plan.segments) + assert not plan.segments_cover_query_contiguously + assert calls == [ + (["causal", "causal"], [2, 4], [0, 1, 6, 7]), + (["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]), + (["causal", "causal"], [6, 8], [5, 11]), + ] + + +def test_paged_piecewise_runner_preserves_homogeneous_batch_without_scatter(monkeypatch: pytest.MonkeyPatch): + plan = build_paged_piecewise_plan( + [[(2, 5)], [(2, 5)]], + query_offsets=[0, 0], + query_lens=[6, 6], + seq_lens=[6, 6], + device=DEVICE, + ) + assert plan.homogeneous_batch_shape == (2, 6) + query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4) + calls: list[torch.Tensor] = [] + + def fail_index_copy(*_args, **_kwargs): + raise AssertionError("homogeneous piecewise path must not use index_copy_") + + monkeypatch.setattr(torch.Tensor, "index_copy_", fail_index_copy) + + def run_segment(segment_query, segment_key, segment_value, _metadata, segment_output_buffer): + assert segment_key is None + assert segment_value is None + assert segment_output_buffer is None + calls.append(segment_query.clone()) + return segment_query + 1 + + output_buffer = torch.empty_like(query) + output = run_paged_piecewise_plan( + query, + None, + None, + plan, + plan.segments, + run_segment, + output_buffer=output_buffer, + use_homogeneous_batch=True, + ) + + assert output is output_buffer + torch.testing.assert_close(output, query + 1) + assert [call.shape[0] for call in calls] == [4, 6, 2] + torch.testing.assert_close(calls[0], query[[0, 1, 6, 7]]) + torch.testing.assert_close(calls[1], query[[2, 3, 4, 8, 9, 10]]) + torch.testing.assert_close(calls[2], query[[5, 11]]) + + +def test_paged_piecewise_runner_keeps_indexed_fallback_for_heterogeneous_layout(): + plan = build_paged_piecewise_plan( + [[(2, 5)], [(1, 4)]], + query_offsets=[0, 0], + query_lens=[6, 6], + seq_lens=[6, 6], + device=DEVICE, + ) + assert plan.homogeneous_batch_shape is None + query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4) + output = run_paged_piecewise_plan( + query, + None, + None, + plan, + plan.segments, + lambda segment_query, *_args: segment_query, + use_homogeneous_batch=True, + ) + torch.testing.assert_close(output, query) + + +def test_paged_piecewise_runner_slices_contiguous_segments(): + plan = build_paged_piecewise_plan( + [[(2, 5)]], + query_offsets=[0], + query_lens=[6], + seq_lens=[6], + device=DEVICE, + ) + query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4) + key = query + 100 + value = query + 200 + expected_ranges = [(0, 2), (2, 5), (5, 6)] + calls: list[tuple[int, int]] = [] + + def run_segment(segment_query, segment_key, segment_value, _segment_batch, segment_output_buffer): + assert segment_output_buffer is None + expected_range = expected_ranges[len(calls)] + start, end = expected_range + calls.append(expected_range) + assert segment_query.untyped_storage().data_ptr() == query.untyped_storage().data_ptr() + assert segment_key.untyped_storage().data_ptr() == key.untyped_storage().data_ptr() + assert segment_value.untyped_storage().data_ptr() == value.untyped_storage().data_ptr() + torch.testing.assert_close(segment_query, query[start:end]) + return segment_query + + output = run_paged_piecewise_plan(query, key, value, plan, plan.segments, run_segment) + + assert calls == expected_ranges + assert [segment.query_range for segment in plan.segments] == expected_ranges + assert plan.segments_cover_query_contiguously + torch.testing.assert_close(output, query) + + +def test_paged_piecewise_runner_writes_contiguous_segments_into_output_buffer(): + plan = build_paged_piecewise_plan( + [[(2, 5)]], + query_offsets=[0], + query_lens=[6], + seq_lens=[6], + device=DEVICE, + ) + query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4) + output_buffer = torch.empty(6, 2, 3) + targets = [] + + def run_segment(segment_query, _key, _value, _metadata, segment_output_buffer): + assert segment_output_buffer is not None + assert segment_output_buffer.untyped_storage().data_ptr() == output_buffer.untyped_storage().data_ptr() + targets.append(segment_output_buffer) + return segment_output_buffer.copy_(segment_query[..., :3]) + + output = run_paged_piecewise_plan( + query, + query, + query, + plan, + plan.segments, + run_segment, + output_buffer=output_buffer, + ) + + assert output is output_buffer + assert [target.shape[0] for target in targets] == [2, 3, 1] + torch.testing.assert_close(output, query[..., :3]) + + +def test_paged_piecewise_runner_returns_single_full_segment_directly(): + plan = build_paged_piecewise_plan( + [[]], + query_offsets=[0], + query_lens=[6], + seq_lens=[6], + device=DEVICE, + ) + query = torch.randn(6, 2, 4) + runner_output = query + 1 + + output = run_paged_piecewise_plan( + query, + query, + query, + plan, + plan.segments, + lambda *_args: runner_output, + ) + + assert output is runner_output + + +def test_paged_piecewise_runner_uses_runner_output_shape(): + plan = build_paged_piecewise_plan( + [[(2, 5)]], + query_offsets=[0], + query_lens=[6], + seq_lens=[6], + device=DEVICE, + ) + query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4) + + output = run_paged_piecewise_plan( + query, + query, + query, + plan, + plan.segments, + lambda segment_query, *_args: segment_query[..., :3], + ) + + assert output.shape == (6, 2, 3) + torch.testing.assert_close(output, query[..., :3]) + + @pytest.mark.parametrize("global_spans", SPAN_CASES) @pytest.mark.parametrize("q_range", Q_RANGE_CASES) @pytest.mark.parametrize("batch_size", BATCH_CASES) diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py index f079a1758..5b5fcc405 100644 --- a/tests/diffusion/attention/test_rainfusion_plan.py +++ b/tests/diffusion/attention/test_rainfusion_plan.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2. @@ -11,18 +11,25 @@ to the always-kept prefix segment before generating the block mask. """ import dataclasses +import sys +import types +from unittest import mock import pytest import torch -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout +from vllm_omni.diffusion.attention.backends import rainfusion_attn +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan from vllm_omni.diffusion.attention.backends.rainfusion_attn import ( _BLOCK_SIZE, + RainFusionAttentionBackend, RainFusionAttentionImpl, RainFusionPlan, ) from vllm_omni.platforms import current_omni_platform +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu] + PREFIX_ROWS = 710 # 14 text rows + 696 audio rows ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows @@ -97,6 +104,62 @@ def test_video_segment_must_be_the_tail_of_packed_document_zero(): assert make_impl()._resolve_plan(metadata) is None +def test_ref2va_multi_video_spans_resolve(): + layout = VideoTokenLayout( + used_len=12000, + video_spans=( + VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"), + VideoTokenSpan(start=5000, latent_grid=(4, 16, 64), role="target"), + ), + ) + plan = make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) + + assert plan is not None + assert plan.used_len == 12000 + assert plan.video_spans == [ + {"start": 128, "latent_shape": [4, 16, 64]}, + {"start": 5000, "latent_shape": [4, 16, 64]}, + ] + + +def test_invalid_ref2va_video_spans_fall_back_to_dense(): + layout = VideoTokenLayout( + used_len=12000, + video_spans=( + VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"), + VideoTokenSpan(start=4000, latent_grid=(4, 16, 64), role="target"), + ), + ) + assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None + + +def test_validate_available_rejects_legacy_mindiesd(monkeypatch): + mindiesd = types.ModuleType("mindiesd") + + def sparse_attention(query, key, value, **kwargs): + return query + + mindiesd.sparse_attention = sparse_attention + monkeypatch.setitem(sys.modules, "mindiesd", mindiesd) + monkeypatch.setattr("importlib.util.find_spec", lambda _: object()) + + with pytest.raises(ValueError, match="video_spans"): + RainFusionAttentionBackend.validate_available() + + +def test_validate_available_accepts_new_mindiesd(monkeypatch): + mindiesd = types.ModuleType("mindiesd") + + def sparse_attention(query, key, value, *, video_spans=None, **kwargs): + return query + + mindiesd.sparse_attention = sparse_attention + monkeypatch.setitem(sys.modules, "mindiesd", mindiesd) + monkeypatch.setattr("importlib.util.find_spec", lambda _: object()) + + RainFusionAttentionBackend.validate_available() + + @pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)]) def test_short_video_stays_dense(grid): assert make_impl()._resolve_plan(make_metadata(grid)) is None @@ -169,7 +232,9 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd() heads, head_dim = 4, 128 q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3)) - impl = make_impl() + # Explicit bf16: the dense comparison exercises the plain sparse kernel, + # and older MindIE-SD releases do not accept precision= (gate would raise). + impl = make_impl(precision="bf16") impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0) plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid)) out = impl._forward_sparse_npu(q, k, v, plan) @@ -182,3 +247,90 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd() ).transpose(1, 2) error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean() assert error < 2e-3, f"mean relative error {error:.4%} against dense attention" + + +# end_step tail fallback: the last ``end_step`` denoise steps must stay dense. + + +def test_end_step_tail_window_stays_dense(): + """step_idx inside the last end_step window must resolve no sparse plan.""" + impl = make_impl(end_step=3) + fc = mock.Mock() + fc.denoise_step_idx = 47 + fc.total_denoise_steps = 50 + with ( + mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True), + mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc), + ): + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None + + +def test_end_step_before_tail_window_still_sparse(): + """step_idx before the tail window must still resolve a sparse plan.""" + impl = make_impl(end_step=3) + fc = mock.Mock() + fc.denoise_step_idx = 46 + fc.total_denoise_steps = 50 + with ( + mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True), + mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc), + ): + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None + + +def test_end_step_zero_never_triggers_tail_fallback(): + """end_step=0 must not fall back even on the final denoise step.""" + impl = make_impl(end_step=0) + fc = mock.Mock() + fc.denoise_step_idx = 49 + fc.total_denoise_steps = 50 + with ( + mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True), + mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc), + ): + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None + + +# precision capability gate: a mindiesd without sparse_attention(precision=) +# must fail fast instead of silently running the BF16 path. + + +def _fake_mindiesd_module(): + """Install a minimal stand-in so the import inside _forward_sparse_npu succeeds.""" + import types + + fake = types.ModuleType("mindiesd") + fake.sparse_attention = lambda *args, **kwargs: None + return fake + + +def test_precision_non_bf16_requires_mindiesd_support(): + """precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError.""" + import sys + + impl = make_impl(precision="mix") + sys.modules["mindiesd"] = _fake_mindiesd_module() + try: + with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False): + with pytest.raises(RuntimeError, match="requires MindIE-SD"): + impl._forward_sparse_npu(None, None, None, None) + finally: + sys.modules.pop("mindiesd", None) + + +def test_precision_non_bf16_passes_gate_when_supported(): + """precision != bf16 with a capable mindiesd must not raise from the gate.""" + import sys + + impl = make_impl(precision="mix") + sys.modules["mindiesd"] = _fake_mindiesd_module() + try: + with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True): + # q/k/v shapes: [B, S, N, D]; plan geometry must match S. + q = torch.randn(1, 8, 4, 128) + plan = RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2]) + # The gate must pass; the fake mindiesd returns None so no crash. + out = impl._forward_sparse_npu(q, q, q, plan) + assert out is None + finally: + sys.modules.pop("mindiesd", None) diff --git a/tests/diffusion/attention/test_trtllm_attn.py b/tests/diffusion/attention/test_trtllm_attn.py index 76daf9d3b..336f7d037 100644 --- a/tests/diffusion/attention/test_trtllm_attn.py +++ b/tests/diffusion/attention/test_trtllm_attn.py @@ -1,14 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math +from unittest.mock import Mock import pytest import torch import torch.nn.functional as F from vllm_omni.diffusion.attention.backends import trtllm_attn as tg -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, + PackedPaddingMetadata, +) from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] @@ -35,6 +39,15 @@ def _impl(**bk): return TrtllmAttentionImpl(8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=bk) +def _packed_padding(cu_seqlens_q, cu_seqlens_k, q_length, kv_length): + return PackedPaddingMetadata( + q_length=q_length, + kv_length=kv_length, + cu_seqlens_q=cu_seqlens_q[:2], + cu_seqlens_k=cu_seqlens_k[:2], + ) + + def test_skip_config_pure_resolution(): from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig @@ -75,7 +88,7 @@ def test_quant_quantize_requires_flashinfer_routine(monkeypatch): def test_quant_quantize_calls_routine_and_shapes_sfs(): from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig - captured = {} + captured: dict[str, object] = {} def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype): captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype) @@ -209,62 +222,80 @@ def test_masked_layer_raises(): _impl().forward_cuda(q, k, v, _Meta()) -def test_packed_metadata_trims_padding_mask(monkeypatch): - b, s, h, d = 1, 8, 8, 128 - q, k, v = (torch.zeros(b, s, h, d) for _ in range(3)) - mask = torch.tensor([[True, True, True, True, True, True, False, False]]) - cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32) +def test_generic_packed_metadata_preserves_ragged_batch(monkeypatch): + b, h, d = 1, 8, 128 + q = torch.zeros(b, 5, h, d) + k, v = (torch.zeros(b, 7, h, d) for _ in range(2)) metadata = AttentionMetadata( - attn_mask=mask, extra={ - "cu_seqlens_q": cu_seqlens, - "cu_seqlens_k": cu_seqlens, - "max_seqlen_q": 6, - "max_seqlen_k": 6, + "cu_seqlens_q": torch.tensor([0, 2, 5], dtype=torch.int32), + "cu_seqlens_k": torch.tensor([0, 3, 7], dtype=torch.int32), + "max_seqlen_q": 3, + "max_seqlen_k": 4, + } + ) + fake_attention = Mock(return_value=torch.zeros_like(q.reshape(5, h, d))) + + monkeypatch.setattr(tg, "HAS_FLASHINFER", True) + monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention) + monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0))) + + _impl().forward_cuda(q, k, v, metadata) + captured = fake_attention.call_args.kwargs + + assert captured["query"].shape[0] == 5 + assert captured["key"].shape[0] == 7 + assert captured["value"].shape[0] == 7 + assert captured["batch_size"] == 2 + assert captured["seq_lens"].tolist() == [3, 4] + assert tg.TrtllmAttentionBackend.supports_multi_doc_packed_varlen() + + +def test_packed_padding_metadata_trims_without_mask(monkeypatch): + q_len, kv_len = 8, 9 + valid_q_tokens, valid_kv_tokens = 5, 6 + b, h, d = 1, 8, 128 + q = torch.zeros(b, q_len, h, d) + k, v = (torch.zeros(b, kv_len, h, d) for _ in range(2)) + cu_seqlens_q = torch.tensor([0, valid_q_tokens, q_len], dtype=torch.int32) + cu_seqlens_k = torch.tensor([0, valid_kv_tokens, kv_len], dtype=torch.int32) + metadata = AttentionMetadata( + packed_padding=_packed_padding( + cu_seqlens_q, + cu_seqlens_k, + valid_q_tokens, + valid_kv_tokens, + ), + extra={ + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "max_seqlen_q": valid_q_tokens, + "max_seqlen_k": valid_kv_tokens, + "valid_kv_length": valid_kv_tokens, }, ) - captured = {} - - def fake_attention(**kwargs): - captured.update(kwargs) - return torch.zeros_like(kwargs["query"]) + fake_attention = Mock(return_value=torch.ones(valid_q_tokens, h, d)) monkeypatch.setattr(tg, "HAS_FLASHINFER", True) monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention) monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0))) out = _impl().forward_cuda(q, k, v, metadata) + captured = fake_attention.call_args.kwargs + assert tg.TrtllmAttentionBackend.supports_packed_mask_free() assert out.shape == q.shape - assert captured["query"].shape[0] == 6 - assert captured["key"].shape[0] == 6 - assert captured["value"].shape[0] == 6 + assert captured["query"].shape[0] == valid_q_tokens + assert captured["key"].shape[0] == valid_kv_tokens + assert captured["value"].shape[0] == valid_kv_tokens assert captured["batch_size"] == 1 - assert captured["seq_lens"].tolist() == [6] - assert captured["cum_seq_lens_q"].tolist() == [0, 6] - assert captured["cum_seq_lens_kv"].tolist() == [0, 6] - assert captured["max_q_len"] == 6 - assert captured["max_kv_len"] == 6 - assert torch.count_nonzero(out[:, 6:]) == 0 - - -def test_packed_metadata_rejects_non_prefix_mask(monkeypatch): - b, s, h, d = 1, 8, 8, 128 - q, k, v = (torch.zeros(b, s, h, d) for _ in range(3)) - cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32) - metadata = AttentionMetadata( - attn_mask=torch.tensor([[True, True, False, True, True, True, False, False]]), - extra={ - "cu_seqlens_q": cu_seqlens, - "cu_seqlens_k": cu_seqlens, - "max_seqlen_q": 6, - "max_seqlen_k": 6, - }, - ) - monkeypatch.setattr(tg, "HAS_FLASHINFER", True) - - with pytest.raises(ValueError, match="prefix-valid"): - _impl().forward_cuda(q, k, v, metadata) + assert captured["seq_lens"].tolist() == [valid_kv_tokens] + assert captured["cum_seq_lens_q"].tolist() == [0, valid_q_tokens] + assert captured["cum_seq_lens_kv"].tolist() == [0, valid_kv_tokens] + assert captured["max_q_len"] == valid_q_tokens + assert captured["max_kv_len"] == valid_kv_tokens + assert torch.count_nonzero(out[:, :valid_q_tokens] != 1) == 0 + assert torch.count_nonzero(out[:, valid_q_tokens:]) == 0 def test_packed_metadata_must_be_complete(): @@ -393,12 +424,13 @@ def test_bf16_packed_padding_matches_sdpa(): q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3)) cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda") metadata = AttentionMetadata( - attn_mask=torch.arange(s, device="cuda")[None] < used, + packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used), extra={ "cu_seqlens_q": cu_seqlens, "cu_seqlens_k": cu_seqlens, "max_seqlen_q": used, "max_seqlen_k": used, + "valid_kv_length": used, }, ) @@ -443,12 +475,13 @@ def test_sage_packed_non_aligned_length_matches_sdpa(): q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3)) cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda") metadata = AttentionMetadata( - attn_mask=torch.arange(s, device="cuda")[None] < used, + packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used), extra={ "cu_seqlens_q": cu_seqlens, "cu_seqlens_k": cu_seqlens, "max_seqlen_q": used, "max_seqlen_k": used, + "valid_kv_length": used, }, ) diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index 184807d98..77e00edda 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/tests/diffusion/attention/test_ulysses_uaa.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import os -import socket import tempfile from types import SimpleNamespace @@ -12,6 +11,7 @@ import numpy as np import pytest import torch +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention from vllm_omni.diffusion.config import set_current_diffusion_config @@ -27,20 +27,6 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - @pytest.mark.core_model @pytest.mark.cpu def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None: @@ -71,7 +57,7 @@ def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None: def _run_attention_case( local_rank: int, world_size: int, - master_port: int, + init_method: str, input_file: str, output_file: str, num_heads: int, @@ -85,10 +71,9 @@ def _run_attention_case( device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) - _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port) os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" - init_distributed_environment(world_size=world_size, rank=local_rank) + init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method) initialize_model_parallel( data_parallel_size=1, cfg_parallel_size=1, @@ -202,10 +187,8 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head batch_size = 2 head_size = 8 - base_port = _find_free_port() - sp_port = _find_free_port() - while sp_port == base_port: - sp_port = _find_free_port() + base_init_method = get_distributed_init_method() + sp_init_method = get_distributed_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name @@ -224,7 +207,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head # Baseline (no SP) torch.multiprocessing.spawn( _run_attention_case, - args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict"), + args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"), nprocs=1, ) @@ -233,7 +216,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head _run_attention_case, args=( sp_world_size, - sp_port, + sp_init_method, input_file, sp_file, num_heads, @@ -278,10 +261,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5 split_sizes = [3, 2, 3, 2] - base_port = _find_free_port() - sp_port = _find_free_port() - while sp_port == base_port: - sp_port = _find_free_port() + base_init_method = get_distributed_init_method() + sp_init_method = get_distributed_init_method() with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in: input_file = f_in.name @@ -300,7 +281,19 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: # Baseline (no SP) torch.multiprocessing.spawn( _run_attention_case, - args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict", 1, None, "mem_efficient"), + args=( + 1, + base_init_method, + input_file, + baseline_file, + num_heads, + head_size, + 1, + "strict", + 1, + None, + "mem_efficient", + ), nprocs=1, ) @@ -309,7 +302,7 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: _run_attention_case, args=( sp_world_size, - sp_port, + sp_init_method, input_file, sp_file, num_heads, diff --git a/tests/diffusion/batching/test_diffusion_batching.py b/tests/diffusion/batching/test_diffusion_batching.py index 60136ce4d..c9ef78383 100644 --- a/tests/diffusion/batching/test_diffusion_batching.py +++ b/tests/diffusion/batching/test_diffusion_batching.py @@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni. This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a diffusion model and validates that every caller receives its correct individual result. When the underlying diffusion stage is configured with -``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the +``max_num_seqs > 1``, the requests will be batched internally. Even without explicit batching config this test is useful for verifying @@ -260,7 +260,7 @@ async def compare_single_vs_parallel( ) -> None: """Run the same prompts sequentially then in parallel and print a comparison.""" - omni = AsyncOmni(model=model, diffusion_batch_size=batch_size) + omni = AsyncOmni(model=model, max_num_seqs=batch_size) try: await warmup(omni, WARMUP_PROMPTS) single_time = await run_single(omni, prompts) @@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) -> await compare_single_vs_parallel(model, prompts, batch_size=batch_size) return - omni = AsyncOmni(model=model, diffusion_batch_size=batch_size) + omni = AsyncOmni(model=model, max_num_seqs=batch_size) try: await warmup(omni, WARMUP_PROMPTS) @@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str): fired via asyncio.gather. Each request_id must appear in the results.""" async def _inner(): - omni = AsyncOmni(model=model_name, diffusion_batch_size=1) + omni = AsyncOmni(model=model_name, max_num_seqs=1) try: prompts = TEST_PROMPTS[:4] sp = _default_sampling_params() @@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str): """ async def _inner(): - omni = AsyncOmni(model=model_name, diffusion_batch_size=4) + omni = AsyncOmni(model=model_name, max_num_seqs=4) try: prompts = TEST_PROMPTS[:4] sp = _default_sampling_params() @@ -524,7 +524,12 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation") parser.add_argument("--model", type=str, required=True, help="Model name or path") parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run") - parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)") + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Scheduler in-flight width (--max-num-seqs)", + ) parser.add_argument( "--mode", choices=["batch", "single", "compare", "validate"], diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index fb912b08e..85bee131b 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -365,13 +365,10 @@ class TestCacheDiTBackend: @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") - def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter): - """Test DreamID uses pipeline.transformer for cache enable/refresh. - - NOTE: DreamID no longer has a custom enabler, so this tests against the generic path. - """ + def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter): + """Generic cache path should pick up ``transformer.fused_blocks``.""" mock_pipeline = Mock() - mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline" + mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline" mock_pipeline.transformer = Mock() mock_pipeline.transformer.fused_blocks = Mock() mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig( diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 2799f1607..165a98cc6 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -1,18 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Model specific tests for CacheDiT enablement. """ import ast -import sys from pathlib import Path from unittest.mock import Mock, patch import pytest import torch from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2 +from vllm.distributed import parallel_state import vllm_omni.diffusion.cache.cachedit as cd_backend import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific @@ -24,19 +24,9 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel from vllm_omni.platforms import current_omni_platform -# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect -# the class even though the dependency may not be set up correctly; this is ok for these -# tests because we just inspect it and never initialize the model. -for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"): - sys.modules.setdefault(mod, Mock()) -# isort: split -from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel # noqa: E402 - pytestmark = [pytest.mark.core_model, pytest.mark.cpu] SEPARATE_CFG_TRANSFORMERS = [ - DreamIdOmniModel, - LTX2VideoTransformer3DModel, HeliosTransformer3DModel, LongCatImageTransformer2DModel, Cosmos3VFMTransformer, @@ -62,6 +52,18 @@ def test_custom_cache_dit_enablers_are_registered_explicitly(): assert cd_backend.CUSTOM_DIT_ENABLERS == expected_enablers +@pytest.fixture() +def init_fake_tp_group(mocker): + """Provide a fake TP group so vLLM linear layers can be instantiated.""" + mock_tp = mocker.MagicMock() + mock_tp.world_size = 1 + mock_tp.rank_in_group = 0 + old = parallel_state._TP + parallel_state._TP = mock_tp + yield + parallel_state._TP = old + + def test_wan22_vace_uses_wan22_custom_cache_dit_enabler(): assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_model_specific.enable_cache_for_wan22 @@ -111,7 +113,9 @@ def test_cachedit_consumers_use_package_api(): for module in modules: if module == legacy_module or module.startswith(internal_prefix): - invalid_imports.append(f"{source_path.relative_to(package_root)}:{node.lineno}: {module}") + invalid_imports.append( + f"{source_path.relative_to(package_root)}:{getattr(node, 'lineno', '?')}: {module}" + ) assert not invalid_imports, "Cache-DiT consumers must use the package API:\n" + "\n".join(invalid_imports) @@ -124,6 +128,13 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model): assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True +def test_ltx2_cache_dit_uses_one_forward_per_denoise_step(): + """LTX batches all guidance passes into one Transformer invocation.""" + adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config + + assert adapter_config.has_separate_cfg is False + + @patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter") @patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit") def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter): diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index fdeca8e17..d821541d0 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/tests/diffusion/cache/test_teacache_extractors.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for TeaCache extractor functions. @@ -18,6 +18,7 @@ Currently implemented: - TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor """ +import math from abc import ABC, abstractmethod from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch @@ -650,6 +651,29 @@ class TestMiniMaxH3Extractor(BaseExtractorTest): torch.testing.assert_close(actual_video, expected_video) torch.testing.assert_close(actual_audio, expected_audio) + def test_extractor_uses_prepared_rope_table( + self, + minimax_h3_module, + sample_inputs, + monkeypatch, + ): + """TeaCache must consume the denoise-loop RoPE table without rebuilding it.""" + rope_table = minimax_h3_module.prepare_rope_table( + sample_inputs["img_position_ids"], + seq_len=sample_inputs["x"].shape[1], + ) + + def unexpected_rope_build(*args, **kwargs): + raise AssertionError("TeaCache extractor rebuilt the prepared RoPE table") + + monkeypatch.setattr(minimax_h3_module, "prepare_rope_table", unexpected_rope_build) + monkeypatch.setattr(minimax_h3_module.rope, "forward", unexpected_rope_build) + context = extract_minimax_h3_context( + minimax_h3_module, + **{**sample_inputs, "rope_table": rope_table}, + ) + context.run_transformer_blocks() + def test_run_transformer_blocks_forwards_packed_layout( self, minimax_h3_module, @@ -705,6 +729,222 @@ class TestMiniMaxH3Extractor(BaseExtractorTest): assert video_logits.shape[0] == sample_inputs["img_pos_info"]["position_ids"].shape[0] assert audio_logits.shape[0] == sample_inputs["audio_pos_info"]["position_ids"].shape[0] + @pytest.mark.sp + def test_teacache_strict_sp_keeps_cache_state_rank_local( + self, + minimax_h3_module, + sample_inputs, + monkeypatch, + ): + """Strict SP must not mix local TeaCache state with gathered block rows.""" + from vllm_omni.diffusion.attention.ops import minimax_h3_modulation + from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig + from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook + from vllm_omni.diffusion.distributed import parallel_state + + seq_len = sample_inputs["x"].shape[1] + local_len = seq_len // 2 + hidden_size = minimax_h3_module.hidden_size + video_width = minimax_h3_module.arch.latents_dim * math.prod(minimax_h3_module.arch.patch_size) + audio_width = minimax_h3_module.arch.audio_latents_dim + captured_modulation_indices: list[torch.Tensor] = [] + + class _StrictSPPrepare(nn.Module): + def __init__(self): + super().__init__() + self.global_indices: list[torch.Tensor] = [] + + def forward(self, hidden, rope_table, combined_indices): + assert hidden.shape[0] == local_len + assert rope_table.shape[0] == local_len + self.global_indices.append(combined_indices.clone()) + return hidden, rope_table, combined_indices.narrow(0, 0, local_len) + + class _StrictSPGather(nn.Module): + def __init__(self): + super().__init__() + self.inputs: list[torch.Tensor] = [] + self.outputs: list[torch.Tensor] = [] + + def forward(self, tensor): + self.inputs.append(tensor.clone()) + if tensor.shape[-1] == hidden_size: + raise AssertionError("TeaCache must not gather hidden states in strict SP") + assert tensor.shape == (local_len, video_width + audio_width) + # Model the all-gather of compact logits from every SP rank. + # postprocess() indexes the resulting global layout with + # infer_out_pos/audio_pos, including rows owned by another rank. + gathered = tensor.repeat(seq_len // local_len, 1) + assert gathered.shape == (seq_len, video_width + audio_width) + self.outputs.append(gathered) + return gathered + + class _StrictSPGroup: + def __init__(self): + super().__init__() + self.local_decisions: list[int] = [] + self.ops = [] + + def all_reduce(self, decision, op): + self.local_decisions.append(int(decision.item())) + self.ops.append(op) + # Simulate another rank requiring the second step to execute + # the collective block path even though this rank would cache. + if len(self.local_decisions) == 2: + decision.fill_(1) + return decision + + class _StrictSPBlock(nn.Module): + def __init__(self): + super().__init__() + self.block_indices: list[torch.Tensor] = [] + + def adaln_proj(self, t_emb): + del t_emb + zeros = torch.zeros(3, hidden_size, dtype=torch.float32) + return zeros, zeros + + def norm1(self, hidden): + return hidden.clone() + + def forward(self, hidden, *, combined_indices, rope_table, **kwargs): + del kwargs + assert hidden.shape[0] == local_len + assert rope_table.shape[0] == local_len + self.block_indices.append(combined_indices.clone()) + return hidden + 1 + + class _StrictSPFinalLayer(nn.Module): + def forward(self, hidden, *, t_emb, inverse_indices): + del t_emb, inverse_indices + assert hidden.shape == (local_len, hidden_size) + return hidden[:, :video_width].float(), hidden[:, :audio_width].float() + + local_prepare = _StrictSPPrepare() + gather = _StrictSPGather() + sp_group = _StrictSPGroup() + block = _StrictSPBlock() + minimax_h3_module.local_sp_prepare = local_prepare + minimax_h3_module.sp_gather = gather + minimax_h3_module.blocks = nn.ModuleList([block]) + minimax_h3_module.final_layer = _StrictSPFinalLayer() + monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, local_len)) + monkeypatch.setattr( + minimax_h3_module, + "_embed", + lambda **kwargs: ( + torch.arange(local_len * hidden_size, dtype=torch.float32).reshape(local_len, hidden_size), + torch.zeros(3, hidden_size, dtype=torch.float32), + ), + ) + monkeypatch.setattr( + minimax_h3_module, + "prepare_rope_table", + lambda *args, **kwargs: torch.zeros(local_len, 6, dtype=torch.float32), + ) + monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group) + + original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_ + + def capture_indexed_scale_shift(hidden, shift, scale, indices): + captured_modulation_indices.append(indices.clone()) + return original_indexed_scale_shift(hidden, shift, scale, indices) + + monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift) + + local_img_pos = torch.tensor([0, 1]) + local_audio_pos = torch.tensor([2, 3]) + token_tags = torch.full((seq_len,), -1, dtype=torch.long) + token_tags[local_img_pos] = 0 + token_tags[local_audio_pos] = 2 + inverse_indices = torch.zeros(seq_len, dtype=torch.long) + inputs = { + **sample_inputs, + "inverse_indices": inverse_indices, + "token_tags": token_tags, + "update_mask": torch.ones(local_img_pos.numel(), dtype=torch.bool), + "prompt_embeds": torch.empty(0, 6), + "img_pos_info": {"position_ids": local_img_pos}, + "audio_pos_info": {"position_ids": local_audio_pos}, + "text_pos_info": {"position_ids": torch.empty(0, dtype=torch.long)}, + "img_pos_for_infer_output_info": {"position_ids": local_img_pos}, + "refiner_packed_seq_params": { + "cu_seqlens_q": torch.tensor([0, 0], dtype=torch.int32), + "max_seqlen_q": 0, + }, + } + expected_local_indices = (inverse_indices * 3 + token_tags.clamp(min=0))[:local_len] + + hook = TeaCacheHook( + TeaCacheConfig( + transformer_type="MiniMaxH3DiTModel", + rel_l1_thresh=0.3, + ) + ) + hook.initialize_hook(minimax_h3_module) + + hook.new_forward(minimax_h3_module, **inputs) + hook.new_forward(minimax_h3_module, **inputs) + video_logits, audio_logits = hook.new_forward(minimax_h3_module, **inputs) + state = hook.state_manager.get_state() + + assert state.previous_residual is not None + assert state.previous_residual.shape == (local_len, hidden_size) + assert len(captured_modulation_indices) == 3 + assert all(torch.equal(indices, expected_local_indices) for indices in captured_modulation_indices) + assert sp_group.local_decisions == [1, 0, 0] + assert len(sp_group.ops) == 3 + assert all(op == torch.distributed.ReduceOp.MAX for op in sp_group.ops) + assert len(local_prepare.global_indices) == 2 + assert all(indices.shape == (seq_len,) for indices in local_prepare.global_indices) + assert len(block.block_indices) == 2 + assert all(torch.equal(indices, expected_local_indices) for indices in block.block_indices) + assert len(gather.inputs) == 3 + assert video_logits.shape == (local_img_pos.numel(), video_width) + assert audio_logits.shape == (local_audio_pos.numel(), audio_width) + + @pytest.mark.cpu + def test_teacache_non_strict_sp_gathers_hidden_once( + self, + minimax_h3_module, + sample_inputs, + monkeypatch, + ): + """The non-strict SP cache path must not gather global hidden rows twice.""" + from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig + from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook + + class _Prepare(nn.Module): + def forward(self, hidden, rope_table, combined_indices): + return hidden, rope_table, combined_indices + + class _Gather(nn.Module): + def __init__(self): + super().__init__() + self.inputs: list[torch.Tensor] = [] + + def forward(self, tensor): + self.inputs.append(tensor.clone()) + return tensor + + seq_len = sample_inputs["x"].shape[1] + gather = _Gather() + minimax_h3_module.sp_prepare = _Prepare() + minimax_h3_module.sp_gather = gather + monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, seq_len)) + + hook = TeaCacheHook( + TeaCacheConfig( + transformer_type="MiniMaxH3DiTModel", + rel_l1_thresh=0.3, + ) + ) + hook.initialize_hook(minimax_h3_module) + hook.new_forward(minimax_h3_module, **sample_inputs) + + assert len(gather.inputs) == 1 + assert gather.inputs[0].shape == (seq_len, minimax_h3_module.hidden_size) + def test_invalid_module_raises_error(self): """Test that invalid module without blocks raises ValueError.""" invalid_module = Mock() diff --git a/tests/diffusion/conftest.py b/tests/diffusion/conftest.py new file mode 100644 index 000000000..e12d15503 --- /dev/null +++ b/tests/diffusion/conftest.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def restore_torch_default_dtype(): + """Keep process-global dtype changes from leaking between diffusion tests.""" + default_dtype = torch.get_default_dtype() + try: + yield + finally: + torch.set_default_dtype(default_dtype) diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index babc23b4a..3b094c37f 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections import namedtuple from types import SimpleNamespace @@ -9,8 +9,9 @@ import torch from diffusers import DiffusionPipeline # pyright: ignore[reportPrivateImportUsage] from PIL import Image +import vllm_omni.diffusion.data as diffusion_data from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data from vllm_omni.diffusion.data import ( DiffusionOutput, DiffusionParallelConfig, @@ -629,6 +630,46 @@ class TestPipelineArgumentsHandling: with pytest.raises(ValueError): pipeline.forward(DiffusionRequestBatch(requests=[problematic_request])) + interpolation_request = _make_request( + sampling_params=OmniDiffusionSamplingParams( + enable_frame_interpolation=True, + ), + ) + with pytest.raises(ValueError, match="Frame interpolation is not supported with the Diffusers backend"): + pipeline.forward(DiffusionRequestBatch(requests=[interpolation_request])) + + def test_diffusers_config_preserves_checkpoint_class_for_metadata(self, mocker): + class MockWanImageToVideoPipeline: + pass + + od_config = _make_od_config(model_class_name=None) + mocker.patch( + "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index", + return_value={"_class_name": "WanImageToVideoPipeline"}, + ) + mocker.patch.object( + diffusion_data, + "diffusers", + SimpleNamespace(WanImageToVideoPipeline=MockWanImageToVideoPipeline), + ) + mocker.patch.object(OmniDiffusionConfig, "update_multimodal_support") + + od_config.enrich_config() + + assert od_config.model_class_name == "WanImageToVideoPipeline" + assert od_config.diffusers_pipeline_cls is MockWanImageToVideoPipeline + + def test_diffusers_config_falls_back_to_adapter_without_pipeline_index(self, mocker): + od_config = _make_od_config(model_class_name=None) + mocker.patch( + "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index", + return_value=None, + ) + + od_config.enrich_config() + + assert od_config.model_class_name == "DiffusersAdapterPipeline" + @pytest.mark.advanced_model @hardware_test(res={"cuda": "L4"}, num_cards=1) @@ -651,7 +692,7 @@ class TestDiffusersBackendEndToEndExecution: def test_t2i_random_weights( self, omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): messages = dummy_messages_from_mix_data(content_text="a photo of an astronaut riding a horse on mars") @@ -667,8 +708,10 @@ class TestDiffusersBackendEndToEndExecution: }, } - response = openai_client.send_diffusion_request(request_config) - image: Image.Image = response[0].images[0] # pyright: ignore[reportOptionalSubscript] + response = online_client.send_diffusion_request(request_config) + images = response[0].images + assert images is not None + image: Image.Image = images[0] # Request config has incomplete width/height, so internal assertion in `send_diffusion_request` is incomplete. assert image.size == (512, 512) diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py new file mode 100644 index 000000000..f2002f1c6 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_block_tables.py @@ -0,0 +1,780 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from vllm.v1.kv_cache_interface import FullAttentionSpec +from vllm.v1.worker import block_table as native_block_table + +from vllm_omni.diffusion.diffusion_kv import model_runner_backend as model_runner_backend_module +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.metadata import ( + DiffusionKVContextMetadata, + DiffusionKVMetadata, + DiffusionKVSequenceMetadata, +) +from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend +from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _registration_backend(parallel_config: object): + model_config = SimpleNamespace(set_attention_geometry=Mock()) + vllm_config = SimpleNamespace( + attention_config=SimpleNamespace(use_non_causal=False), + compilation_config=SimpleNamespace(static_forward_context={}), + model_config=model_config, + ) + backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=SimpleNamespace(parallel_config=parallel_config), + device=torch.device("cpu"), + ) + return backend, vllm_config, model_config + + +def _parallel_config(**overrides): + values = dict(ulysses_degree=1, ring_degree=1, allgather_degree=1, ulysses_mode="strict") + values.update(overrides) + return SimpleNamespace(**values) + + +def test_cache_layer_registration_installs_adapters_with_ulysses_local_geometry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend, vllm_config, model_config = _registration_backend(_parallel_config(ulysses_degree=2)) + layer = SimpleNamespace( + num_heads=8, + softmax_scale=0.125, + skip_sequence_parallel=False, + attn_backend=SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN"), + ) + unrelated = object() + vllm_config.compilation_config.static_forward_context.update({"layer-0": layer, "unrelated": unrelated}) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=8, + dtype=torch.float16, + non_causal=True, + ) + constructor_calls = [] + + class FakeLayerAdapter: + def __init__(self, *, layer_name, layer, spec, ulysses_degree, **_kwargs) -> None: + constructor_calls.append((layer_name, layer, spec, ulysses_degree)) + self.num_heads = layer.num_heads // ulysses_degree + self.num_kv_heads = spec.num_kv_heads // ulysses_degree + self.head_size = spec.head_size + self.spec = replace(spec, num_kv_heads=self.num_kv_heads) + + monkeypatch.setattr(model_runner_backend_module, "DiffusionPagedAttentionLayerAdapter", FakeLayerAdapter) + monkeypatch.setattr(model_runner_backend_module, "set_current_vllm_config", lambda _config: nullcontext()) + + specs = backend.register_kv_cache_layers({"layer-0": (layer, spec)}) + + installed = vllm_config.compilation_config.static_forward_context["layer-0"] + assert isinstance(installed, FakeLayerAdapter) + assert vllm_config.compilation_config.static_forward_context["unrelated"] is unrelated + assert specs["layer-0"].num_kv_heads == 2 + assert constructor_calls == [("layer-0", layer, spec, 2)] + model_config.set_attention_geometry.assert_called_once_with( + num_heads=4, + num_kv_heads=2, + head_size=8, + ) + assert vllm_config.attention_config.use_non_causal is True + + +def test_cache_layer_registration_rejects_backend_without_paged_support() -> None: + backend, _, _ = _registration_backend(_parallel_config()) + layer = SimpleNamespace( + attn_backend=SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA"), + ) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=8, + dtype=torch.float16, + ) + + with pytest.raises(NotImplementedError, match="layer-0.*paged support.*SDPA"): + backend.register_kv_cache_layers({"layer-0": (layer, spec)}) + + +def test_cache_layer_registration_resolves_platform_paged_hooks(monkeypatch: pytest.MonkeyPatch) -> None: + backend, _, _ = _registration_backend(_parallel_config()) + + def unsupported_block_tables(): + raise NotImplementedError("platform has no paged BlockTables") + + monkeypatch.setattr( + model_runner_backend_module.current_omni_platform, + "get_diffusion_kv_block_tables_cls", + unsupported_block_tables, + ) + + with pytest.raises(NotImplementedError, match="no paged BlockTables"): + backend.register_kv_cache_layers({"layer-0": (object(), Mock(spec=FullAttentionSpec))}) + + +@pytest.mark.parametrize( + ("parallel_config", "message"), + [ + (_parallel_config(ring_degree=2), "Ring"), + (_parallel_config(allgather_degree=2), "AllGather-KV"), + (_parallel_config(ulysses_degree=2, ulysses_mode="advanced_uaa"), "strict Ulysses"), + ], +) +def test_cache_layer_registration_rejects_unsupported_sp_modes(parallel_config, message: str) -> None: + backend, _, _ = _registration_backend(parallel_config) + + with pytest.raises(NotImplementedError, match=message): + backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=False)) + + +def test_cache_layer_registration_ignores_sp_for_opted_out_layer() -> None: + backend, _, _ = _registration_backend( + _parallel_config(ulysses_degree=2, ring_degree=2, ulysses_mode="advanced_uaa") + ) + + assert backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=True)) == 1 + + +class _FakeSpec: + def __init__(self, block_size: int) -> None: + self.block_size = block_size + self.capacity_calls: list[tuple[object, int]] = [] + + def max_num_blocks_per_req(self, vllm_config: object, max_len: int) -> int: + self.capacity_calls.append((vllm_config, max_len)) + return math.ceil(max_len / self.block_size) + + +class _FakeStagedTable: + def __init__(self) -> None: + self.staged: list[tuple[int, tuple[int, ...]]] = [] + self.clear_calls = 0 + + def clear_staged_writes(self) -> None: + self.staged.clear() + self.clear_calls += 1 + + +class _FakeNumBlocks: + def __init__(self, num_groups: int, max_num_reqs: int) -> None: + self.np = np.zeros((num_groups, max_num_reqs), dtype=np.int32) + self.copy_calls = 0 + + def copy_to_uva(self) -> None: + self.copy_calls += 1 + + +class _FakeBlockTables: + instances: list[_FakeBlockTables] = [] + + def __init__(self, **kwargs: object) -> None: + self.kwargs = kwargs + self.block_sizes = list(kwargs["block_sizes"]) + self.max_num_reqs = int(kwargs["max_num_reqs"]) + self.max_num_batched_tokens = int(kwargs["max_num_batched_tokens"]) + self.cp_size = int(kwargs["cp_size"]) + self.cp_rank = int(kwargs["cp_rank"]) + self.cp_interleave = int(kwargs["cp_interleave"]) + self.blocks_per_kv_block = [1] * len(self.block_sizes) + self.block_tables = [_FakeStagedTable() for _ in self.block_sizes] + self.num_blocks = _FakeNumBlocks(len(self.block_sizes), self.max_num_reqs) + self.append_calls: list[tuple[int, tuple[list[int], ...], bool]] = [] + self.apply_calls = 0 + self.layout_refreshes = 0 + self.fail_on_append_call: int | None = None + self.fail_apply = False + self.committed: list[dict[int, tuple[int, ...]]] = [dict() for _ in self.block_sizes] + self.__class__.instances.append(self) + + def append_block_ids( + self, + req_index: int, + new_block_ids: tuple[list[int], ...], + overwrite: bool, + ) -> None: + call_number = len(self.append_calls) + 1 + if self.fail_on_append_call == call_number: + raise RuntimeError("fake append failure") + copied_ids = tuple(list(group_ids) for group_ids in new_block_ids) + self.append_calls.append((req_index, copied_ids, overwrite)) + for group_index, group_ids in enumerate(copied_ids): + self.num_blocks.np[group_index, req_index] = len(group_ids) + if group_ids: + self.block_tables[group_index].staged.append((req_index, tuple(group_ids))) + + def apply_staged_writes(self) -> None: + self.apply_calls += 1 + if self.fail_apply: + raise RuntimeError("fake apply failure") + for group_index, group_table in enumerate(self.block_tables): + for row, group_ids in group_table.staged: + self.committed[group_index][row] = group_ids + group_table.staged.clear() + self.num_blocks.copy_to_uva() + + def init_block_table_layout_tensors(self) -> None: + self.layout_refreshes += 1 + + +def _metadata( + request_id: str = "req-0", + generation: int = 1, + *, + seq_len: int = 9, + sequence_block_ids: tuple[list[int], ...] = ([1, 2, 3], [4, 5]), + contexts: tuple[DiffusionKVContextMetadata, ...] | None = None, +) -> DiffusionKVMetadata: + if contexts is None: + contexts = ( + DiffusionKVContextMetadata( + context_id="text", + cache_role="cross_attention", + num_tokens=5, + block_ids=([6, 7], [8]), + ), + ) + return DiffusionKVMetadata( + request_id=request_id, + allocation_generation=generation, + sequences=( + DiffusionKVSequenceMetadata( + sequence_id=0, + prefix_len=4, + target_len=5, + seq_len=seq_len, + block_ids=sequence_block_ids, + context_ids=tuple(context.context_id for context in contexts), + ), + ), + contexts=contexts, + ) + + +def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]) -> None: + _FakeBlockTables.instances.clear() + + def get_block_table_width(max_num_blocks, block_size): + events.append(("width", max_num_blocks, block_size)) + return max_num_blocks + 1 + + def init_backend(kv_cache_config, vllm_config, device): + events.append(("backend", kv_cache_config, vllm_config, device)) + group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None)) + return [[group], [group]], "cg-support", [4, 8] + + def init_cache( + runner_kv_caches, + forward_context, + kv_cache_config, + attn_groups, + device, + cache_dtype, + kernel_block_sizes, + vllm_config, + ): + events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes)) + runner_kv_caches.extend(["cache-0", "cache-1"]) + return {"layer-0": "tensor-0", "layer-1": "tensor-1"} + + monkeypatch.setattr(model_runner_backend_module, "init_attn_backend", init_backend) + monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", init_cache) + monkeypatch.setattr( + model_runner_backend_module.current_omni_platform, + "get_diffusion_kv_block_tables_cls", + lambda: _FakeBlockTables, + ) + monkeypatch.setattr(native_block_table, "get_block_table_width", get_block_table_width, raising=False) + + +def _runner( + monkeypatch: pytest.MonkeyPatch, + *, + max_rows_per_request: int = 3, + max_num_seqs: int = 2, + max_len: int = 16, + max_num_batched_tokens: int = 32, + cp_size: int = 1, + cp_rank: int = 0, + cp_interleave: int = 1, +) -> tuple[DiffusionKVModelRunnerBackend, _FakeBlockTables, list[tuple]]: + events: list[tuple] = [] + _install_native_modules(monkeypatch, events) + od_config = SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + diffusion_kv_max_rows_per_request=max_rows_per_request, + max_num_seqs=max_num_seqs, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=max_len), + scheduler_config=SimpleNamespace( + max_num_batched_tokens=max_num_batched_tokens, + max_num_seqs=max_num_seqs, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=cp_size, + cp_kv_cache_interleave_size=cp_interleave, + ), + compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}), + cache_config=SimpleNamespace(cache_dtype="auto"), + ) + monkeypatch.setattr( + model_runner_backend_module, + "get_dcp_group", + lambda: SimpleNamespace(rank_in_group=cp_rank), + ) + backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=od_config, + device="cuda:1", + ) + backend._kv_cache_layer_adapters = { + "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None), + "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None), + } + + config = SimpleNamespace( + num_blocks=32, + kv_cache_groups=[ + SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)), + SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)), + ], + ) + backend.initialize_kv_cache(config) + return backend, _FakeBlockTables.instances[-1], events + + +def test_initialize_builds_native_block_tables_from_rank_local_config(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, events = _runner(monkeypatch) + + assert events[:2] == [("width", 4, 4), ("width", 2, 8)] + assert [event[0] for event in events[2:]] == ["backend", "cache"] + assert block_tables.kwargs == { + "block_sizes": [4, 8], + "max_num_reqs": 6, + "max_num_batched_tokens": 32, + "max_num_blocks_per_group": [5, 3], + "device": "cuda:1", + "kernel_block_sizes": [4, 8], + "cp_size": 1, + "cp_rank": 0, + "cp_interleave": 1, + } + assert runner.block_tables is block_tables + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + assert runner.kv_caches == ["cache-0", "cache-1"] + assert runner.vllm_config.scheduler_config.max_num_seqs == 6 + + +def test_initialize_propagates_native_context_parallel_layout(monkeypatch: pytest.MonkeyPatch) -> None: + _, block_tables, _ = _runner( + monkeypatch, + cp_size=2, + cp_rank=1, + cp_interleave=4, + ) + + assert block_tables.kwargs["cp_size"] == 2 + assert block_tables.kwargs["cp_rank"] == 1 + assert block_tables.kwargs["cp_interleave"] == 4 + + +def test_initialize_failure_does_not_publish_partial_cache_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[tuple] = [] + _install_native_modules(monkeypatch, events) + successful_init_cache = model_runner_backend_module.init_kv_cache + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=16), + scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}), + cache_config=SimpleNamespace(cache_dtype="auto"), + ) + backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + diffusion_kv_max_rows_per_request=2, + ), + device="cuda:1", + ) + backend._kv_cache_layer_adapters = { + "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-0"), + "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-1"), + } + config = SimpleNamespace( + num_blocks=32, + kv_cache_groups=[ + SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)), + SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)), + ], + ) + + def fail_init_cache(runner_kv_caches, *_args, **_kwargs): + runner_kv_caches.append("partial-cache") + backend._kv_cache_layer_adapters["layer-0"].kv_cache = "partial-binding" + raise RuntimeError("cache binding failed") + + monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", fail_init_cache) + with pytest.raises(RuntimeError, match="cache binding failed"): + backend.initialize_kv_cache(config) + + assert backend.kv_cache_config is None + assert backend.kv_caches == [] + assert backend.block_tables is None + assert backend.paged_attention_adapter is None + assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0" + assert backend.vllm_config.scheduler_config.max_num_seqs == 1 + + monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache) + backend.initialize_kv_cache(config) + assert backend.kv_cache_config is not None + assert backend.kv_caches == ["cache-0", "cache-1"] + assert backend.vllm_config.scheduler_config.max_num_seqs == 2 + + +def test_initialize_rejects_config_for_different_registered_layers() -> None: + backend = DiffusionKVModelRunnerBackend( + vllm_config=SimpleNamespace(), + od_config=SimpleNamespace(), + device="cuda:1", + ) + backend._kv_cache_layer_adapters = { + "expected-layer": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None) + } + config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(layer_names=["other-layer"])], + ) + + with pytest.raises(ValueError, match="layer mismatch"): + backend.initialize_kv_cache(config) + + +def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + + assert runner.install_diffusion_kv_metadata(_metadata()) is True + + assert block_tables.append_calls == [ + (0, ([1, 2, 3], [4, 5]), True), + (1, ([6, 7], [8]), True), + ] + assert block_tables.apply_calls == 1 + assert runner.get_diffusion_kv_row("req-0", 0) == 0 + assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1 + sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None) + context_binding = runner._resolve_paged_attention_row("req-0", None, "text") + assert (sequence_binding.row_index, sequence_binding.max_seq_len, sequence_binding.block_ids) == ( + 0, + 9, + ((1, 2, 3), (4, 5)), + ) + assert (context_binding.row_index, context_binding.max_seq_len, context_binding.block_ids) == ( + 1, + 5, + ((6, 7), (8,)), + ) + + +def test_block_table_mutations_invalidate_prepared_attention_batches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner, _, _ = _runner(monkeypatch) + assert runner.paged_attention_adapter is not None + runner.paged_attention_adapter.invalidate_prepared_batches = Mock() + + runner.install_diffusion_kv_metadata(_metadata()) + runner.paged_attention_adapter.invalidate_prepared_batches.assert_called_once_with() + + runner.remove_diffusion_kv_requests(["req-0"]) + assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 2 + + runner.refresh_block_table_layout() + assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 3 + + +def test_cache_layer_registration_is_frozen_after_physical_initialization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner, _, _ = _runner(monkeypatch) + + with pytest.raises(RuntimeError, match="cannot be changed"): + runner.register_kv_cache_layers({}) + + +def test_request_scoped_context_uses_one_row_across_sequences(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + first_sequence = _metadata().sequences[0] + metadata = replace( + _metadata(), + sequences=( + first_sequence, + replace( + first_sequence, + sequence_id=1, + block_ids=([9, 10, 11], [12, 13]), + ), + ), + ) + + assert runner.install_diffusion_kv_metadata(metadata) is True + + assert block_tables.append_calls == [ + (0, ([1, 2, 3], [4, 5]), True), + (1, ([9, 10, 11], [12, 13]), True), + (2, ([6, 7], [8]), True), + ] + assert runner.get_diffusion_kv_row("req-0", 0, "text") == 2 + assert runner.get_diffusion_kv_row("req-0", 1, "text") == 2 + + +def test_row_block_count_uses_native_spec_semantics(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + spec = runner.kv_cache_config.kv_cache_groups[0].kv_cache_spec + spec.max_num_blocks_per_req = lambda _config, token_len: math.ceil(token_len / 8) + + metadata = _metadata( + sequence_block_ids=([1, 2], [4, 5]), + contexts=(), + ) + + assert runner.install_diffusion_kv_metadata(metadata) is True + assert block_tables.append_calls == [(0, ([1, 2], [4, 5]), True)] + + +def test_repeated_native_block_ids_are_allowed(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + metadata = _metadata( + sequence_block_ids=([1, 1, 3], [4, 4]), + contexts=(), + ) + + assert runner.install_diffusion_kv_metadata(metadata) is True + assert block_tables.append_calls == [(0, ([1, 1, 3], [4, 4]), True)] + + +def test_generation_is_idempotent_stale_safe_and_conflict_safe(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + metadata = _metadata() + runner.install_diffusion_kv_metadata(metadata) + + assert runner.install_diffusion_kv_metadata(metadata) is False + assert runner.install_diffusion_kv_metadata(replace(metadata, allocation_generation=2)) is False + assert block_tables.apply_calls == 1 + + with pytest.raises(ValueError, match="Stale Diffusion KV allocation generation"): + runner.install_diffusion_kv_metadata(metadata) + with pytest.raises(ValueError, match="Conflicting Diffusion KV allocation snapshot"): + runner.install_diffusion_kv_metadata(_metadata(generation=3, sequence_block_ids=([9, 2, 3], [4, 5]))) + assert block_tables.apply_calls == 1 + + +@pytest.mark.parametrize( + ("metadata", "message"), + [ + (_metadata(sequence_block_ids=([1, 2, 3],)), "block groups"), + (_metadata(sequence_block_ids=([1, 2], [4, 5])), "expected 3"), + (_metadata(sequence_block_ids=([1, 2, 32], [4, 5])), "outside"), + ( + _metadata( + seq_len=17, + sequence_block_ids=([1, 2, 3, 4, 5], [6, 7, 8]), + contexts=(), + ), + "exceeds row capacity", + ), + ], +) +def test_invalid_group_count_block_count_range_and_row_capacity_are_atomic( + monkeypatch: pytest.MonkeyPatch, + metadata: DiffusionKVMetadata, + message: str, +) -> None: + runner, block_tables, _ = _runner(monkeypatch) + + with pytest.raises(ValueError, match=message): + runner.install_diffusion_kv_metadata(metadata) + + assert block_tables.append_calls == [] + assert block_tables.apply_calls == 0 + assert runner._diffusion_kv_identity_to_row == {} + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + + +def test_invalid_later_context_does_not_stage_earlier_valid_row(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + bad_context = DiffusionKVContextMetadata( + context_id="text", + cache_role="cross_attention", + num_tokens=5, + block_ids=([6, 32], [8]), + ) + + with pytest.raises(ValueError, match="outside"): + runner.install_diffusion_kv_metadata(_metadata(contexts=(bad_context,))) + + assert block_tables.append_calls == [] + assert all(not group_table.staged for group_table in block_tables.block_tables) + + +def test_native_append_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + block_tables.fail_on_append_call = 2 + + with pytest.raises(RuntimeError, match="fake append failure"): + runner.install_diffusion_kv_metadata(_metadata()) + + assert np.count_nonzero(block_tables.num_blocks.np) == 0 + assert all(not group_table.staged for group_table in block_tables.block_tables) + assert runner._diffusion_kv_identity_to_row == {} + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + + +def test_native_apply_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + block_tables.fail_apply = True + + with pytest.raises(RuntimeError, match="fake apply failure"): + runner.install_diffusion_kv_metadata(_metadata()) + + assert block_tables.apply_calls == 1 + assert np.count_nonzero(block_tables.num_blocks.np) == 0 + assert all(not group_table.staged for group_table in block_tables.block_tables) + assert all(group_table.clear_calls == 1 for group_table in block_tables.block_tables) + assert runner._diffusion_kv_identity_to_row == {} + assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0] + + +def test_adapter_and_global_row_capacity_are_enforced(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1) + second_context = DiffusionKVContextMetadata( + context_id="image", + cache_role="cross_attention", + num_tokens=1, + block_ids=([9], [10]), + ) + + with pytest.raises(ValueError, match="adapter limit is 2"): + runner.install_diffusion_kv_metadata(_metadata(contexts=(_metadata().contexts[0], second_context))) + assert block_tables.append_calls == [] + + runner.install_diffusion_kv_metadata(_metadata()) + append_count = len(block_tables.append_calls) + with pytest.raises(ValueError, match="only 0 rows are free"): + runner.install_diffusion_kv_metadata(_metadata("req-1")) + assert len(block_tables.append_calls) == append_count + + +def test_empty_and_duplicate_sequence_metadata_are_rejected(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + empty = DiffusionKVMetadata(request_id="req-0", allocation_generation=0, sequences=()) + duplicate_sequence = _metadata().sequences[0] + duplicate = DiffusionKVMetadata( + request_id="req-0", + allocation_generation=0, + sequences=(duplicate_sequence, duplicate_sequence), + ) + + with pytest.raises(ValueError, match="at least one sequence"): + runner.install_diffusion_kv_metadata(empty) + with pytest.raises(ValueError, match="Duplicate diffusion KV row identity"): + runner.install_diffusion_kv_metadata(duplicate) + assert block_tables.append_calls == [] + + +@pytest.mark.parametrize( + ("metadata", "message"), + [ + (_metadata(generation=-1), "non-negative integer"), + (_metadata(request_id=""), "non-empty string"), + (_metadata(sequence_block_ids=([1, 2, True], [4, 5])), "non-integer block ID"), + ], +) +def test_request_generation_and_block_id_types_are_validated( + monkeypatch: pytest.MonkeyPatch, + metadata: DiffusionKVMetadata, + message: str, +) -> None: + runner, block_tables, _ = _runner(monkeypatch) + + with pytest.raises(ValueError, match=message): + runner.install_diffusion_kv_metadata(metadata) + + assert block_tables.append_calls == [] + assert runner._diffusion_kv_request_states == {} + + +def test_cleanup_is_idempotent_and_reuses_native_rows(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1) + runner.install_diffusion_kv_metadata(_metadata()) + + assert runner.remove_diffusion_kv_requests(["missing"]) == 0 + assert runner.remove_diffusion_kv_requests(["req-0", "req-0"]) == 2 + assert block_tables.append_calls[-2:] == [(0, ([], []), True), (1, ([], []), True)] + assert block_tables.num_blocks.np[:, :2].tolist() == [[0, 0], [0, 0]] + assert runner.remove_diffusion_kv_requests(["req-0"]) == 0 + with pytest.raises(KeyError, match="No native diffusion KV row"): + runner.get_diffusion_kv_row("req-0", 0) + + runner.install_diffusion_kv_metadata(_metadata("req-1")) + assert runner.get_diffusion_kv_row("req-1", 0) == 0 + assert runner.get_diffusion_kv_row("req-1", 0, "text") == 1 + + +def test_cached_step_does_not_reinstall_metadata() -> None: + runner = object.__new__(DiffusionModelRunner) + runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER) + runner.pipeline = object() + runner.install_diffusion_kv_metadata = Mock() + runner._supports_step_mode = Mock(return_value=False) + scheduler_output = DiffusionSchedulerOutput( + step_id=1, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData(request_ids=["req-0"]), + finished_req_ids=set(), + num_running_reqs=1, + num_waiting_reqs=0, + ) + + with pytest.raises(ValueError, match="does not support step execution"): + runner.execute_stepwise(scheduler_output) + runner.install_diffusion_kv_metadata.assert_not_called() + + +def test_dense_cleanup_and_wake_are_noops_but_paged_wake_refreshes(monkeypatch: pytest.MonkeyPatch) -> None: + runner, block_tables, _ = _runner(monkeypatch) + runner.refresh_block_table_layout() + assert block_tables.layout_refreshes == 1 + + dense_backend = DiffusionKVModelRunnerBackend( + vllm_config=SimpleNamespace(), + od_config=SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY), + device="cuda:1", + ) + dense_backend.block_tables = Mock() + assert dense_backend.remove_diffusion_kv_requests(["req-0"]) == 0 + dense_backend.refresh_block_table_layout() + dense_backend.block_tables.init_block_table_layout_tensors.assert_not_called() + with pytest.raises(ValueError, match="Dense diffusion execution"): + dense_backend.install_diffusion_kv_metadata(_metadata()) diff --git a/tests/diffusion/diffusion_kv/test_config.py b/tests/diffusion/diffusion_kv/test_config.py index f03120784..61b1b3041 100644 --- a/tests/diffusion/diffusion_kv/test_config.py +++ b/tests/diffusion/diffusion_kv/test_config.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest +from vllm_omni.config.omni_config import _DiffusionConfigProjection from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -35,10 +36,12 @@ def test_non_mapping_omni_kv_config_is_rejected() -> None: OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler") -def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None: +def test_paged_scheduler_rejects_dense_connector_kv_receive() -> None: with pytest.raises(ValueError, match="does not support imported AR KV"): OmniDiffusionConfig.from_kwargs( diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=1, + max_num_batched_tokens=1, omni_kv_config={"need_recv_cache": True}, ) @@ -47,6 +50,33 @@ def test_paged_scheduler_does_not_depend_on_model_registry() -> None: config = OmniDiffusionConfig.from_kwargs( model_class_name="FutureDiffusionModel", diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=1, + max_num_batched_tokens=1, ) assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + + +@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection]) +def test_paged_scheduler_mode_is_platform_agnostic(config_cls) -> None: + config = config_cls.from_kwargs( + diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=1, + ) + + assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + + +def test_paged_scheduler_requires_a_worker_row_limit() -> None: + with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"): + OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_scheduler") + with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"): + _DiffusionConfigProjection.from_kwargs(diffusion_kv_mode="paged_scheduler") + + +@pytest.mark.parametrize("invalid_limit", [0, -1, True, 1.5]) +def test_worker_row_limit_must_be_a_positive_integer(invalid_limit: object) -> None: + with pytest.raises((TypeError, ValueError)): + OmniDiffusionConfig.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit) + with pytest.raises((TypeError, ValueError)): + _DiffusionConfigProjection.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit) diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py new file mode 100644 index 000000000..a82dc01ff --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_initialization.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from vllm.v1.kv_cache_interface import FullAttentionSpec + +import vllm_omni.diffusion.diffusion_kv.initialization as initialization +import vllm_omni.diffusion.vllm_config as diffusion_vllm_config +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler +from vllm_omni.quantization import build_quant_config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _Executor: + def __init__(self) -> None: + self.specs = [{"layer0": object()}, {"layer0": object()}] + self.memory = [1024, 1024] + self.configs = None + self.resolved_max_model_len = None + self.profile_requests = None + + def get_kv_cache_specs(self): + return self.specs + + def determine_available_kv_memory(self, profile_requests): + self.profile_requests = profile_requests + return self.memory + + def set_kv_cache_configs(self, configs, resolved_max_model_len) -> None: + self.configs = configs + self.resolved_max_model_len = resolved_max_model_len + + +def _od_config(**overrides): + values = dict( + model="test-diffusion", + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + dtype=torch.bfloat16, + max_model_len=64, + quantization_config=None, + tf_model_config=None, + enforce_eager=True, + is_moe=False, + additional_config={}, + profiler_config=None, + parallel_config=SimpleNamespace( + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + enable_expert_parallel=False, + ), + gpu_memory_utilization=0.9, + kv_cache_memory_bytes=None, + max_num_seqs=1, + max_num_batched_tokens=64, + num_gpus=1, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None: + executor = _Executor() + od_config = _od_config(num_gpus=2, max_model_len=-1) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=256), + max_in_flight_tokens=128, + ) + worker_configs = [object(), object()] + scheduler_kv_cache_config = object() + profile_requests = [object()] + monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config) + build = Mock(return_value=(worker_configs, scheduler_kv_cache_config)) + monkeypatch.setattr( + initialization, + "build_native_kv_cache_configs", + build, + ) + monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16)) + + result = initialization.initialize_diffusion_kv_control_plane( + executor, + od_config, + profile_requests=profile_requests, + device=torch.device("cpu"), + ) + + build.assert_called_once_with(vllm_config, executor.specs, executor.memory) + assert result == (scheduler_kv_cache_config, 16, 16, vllm_config) + assert executor.configs == worker_configs + assert executor.resolved_max_model_len == vllm_config.model_config.max_model_len + assert executor.profile_requests is profile_requests + + +def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None: + executor = _Executor() + od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY) + get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device")) + monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device) + + result = initialization.initialize_diffusion_kv_control_plane(executor, od_config) + + assert result is None + assert executor.configs is None + get_device.assert_not_called() + + +def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None: + executor = _Executor() + executor.memory = [1024] + od_config = _od_config(num_gpus=2) + monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object()) + + with pytest.raises(ValueError, match="rank count mismatch"): + initialization.initialize_diffusion_kv_control_plane( + executor, + od_config, + profile_requests=[object()], + device=torch.device("cpu"), + ) + + +def test_paged_control_plane_requires_prepared_profile_request() -> None: + with pytest.raises(ValueError, match="prepared memory profile requests"): + initialization.initialize_diffusion_kv_control_plane( + _Executor(), + _od_config(), + device=torch.device("cpu"), + ) + + +def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None: + od_config = _od_config() + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config) + spec = FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + non_causal=True, + ) + num_blocks = 32 + + worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs( + vllm_config, + [{"layer0": spec}], + [spec.page_size_bytes * num_blocks], + ) + + assert len(worker_configs) == 1 + assert worker_configs[0].num_blocks == num_blocks + assert scheduler_kv_cache_config.num_blocks == num_blocks + assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups + + +def test_native_control_plane_config_initializes_scheduler() -> None: + od_config = _od_config() + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config) + spec = FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + non_causal=True, + ) + _, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs( + vllm_config, + [{"layer0": spec}], + [spec.page_size_bytes * 32], + ) + + scheduler = RequestScheduler() + scheduler.initialize( + od_config, + kv_cache_config=scheduler_kv_cache_config, + scheduler_block_size=spec.block_size, + hash_block_size=spec.block_size, + kv_vllm_config=vllm_config, + ) + + assert scheduler._diffusion_kv_manager is not None + scheduler.close() + + +def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None: + od_config = _od_config( + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, + parallel_config=SimpleNamespace( + tensor_parallel_size=1, + pipeline_parallel_size=4, + data_parallel_size=1, + enable_expert_parallel=False, + ), + gpu_memory_utilization=0.5, + kv_cache_memory_bytes=4096, + max_num_batched_tokens=64, + ) + vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config) + original_pp_size = vllm_config.parallel_config.pipeline_parallel_size + original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization + original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes + original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens + + diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config) + + assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size + assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization + assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes + assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget + + +def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None: + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config( + torch.device("cpu"), + _od_config(gpu_memory_utilization=0.42), + ) + + assert vllm_config.cache_config.gpu_memory_utilization == 0.42 + + +def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None: + quantization_config = build_quant_config( + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": [], + } + ) + model_config = diffusion_vllm_config._make_diffusion_vllm_model_config( + _od_config( + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, + quantization_config=quantization_config, + ) + ) + + assert model_config.dtype is torch.bfloat16 + assert model_config.quantization == "modelopt" + assert model_config.quantization_config is quantization_config + assert model_config.is_quantized + assert model_config.is_moe is False + assert model_config.is_nvfp4_quantized() is False + assert model_config.original_max_model_len == model_config.max_model_len + + unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config( + _od_config( + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, + quantization_config=None, + ) + ) + assert not unquantized_model_config.is_quantized + + +def test_minus_one_max_model_len_uses_model_limit() -> None: + od_config = SimpleNamespace( + max_model_len=-1, + tf_model_config=SimpleNamespace(max_position_embeddings=4096), + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + ) + + assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096 + + model_config = diffusion_vllm_config._make_diffusion_vllm_model_config( + _od_config( + max_model_len=-1, + tf_model_config=SimpleNamespace(max_position_embeddings=4096), + ) + ) + assert model_config.max_model_len == 4096 + assert model_config.original_max_model_len == -1 + + +def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None: + vllm_config = diffusion_vllm_config.create_diffusion_vllm_config( + torch.device("cpu"), + _od_config( + max_model_len=-1, + tf_model_config=SimpleNamespace(max_position_embeddings=4096), + ), + ) + spec = FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + non_causal=True, + ) + model_limit = vllm_config.model_config.max_model_len + + initialization.build_native_kv_cache_configs( + vllm_config, + [{"layer0": spec}], + [spec.page_size_bytes * 16], + ) + + assert vllm_config.model_config.original_max_model_len == -1 + # Sixteen physical pages provide ``16 * block_size`` tokens. The native + # backend chooses the block geometry, so keep this assertion backend + # independent instead of assuming the generic 16-token block size. + expected_max_model_len = min(model_limit, spec.block_size * 16) + assert vllm_config.model_config.max_model_len == expected_max_model_len diff --git a/tests/diffusion/diffusion_kv/test_manager.py b/tests/diffusion/diffusion_kv/test_manager.py new file mode 100644 index 000000000..5152c2a67 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_manager.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor + +from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +BLOCK_SIZE = 4 + + +def _config(num_blocks: int) -> KVCacheConfig: + native_kv_managers.register_all_kvcache_specs(None) + spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], + ) + + +def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest: + return DiffusionKVRequest( + f"{public_id}/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=4, + target_len=4, + seq_len=seq_len, + kv_contexts=kv_contexts, + ) + + +def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager: + return DiffusionKVCacheManager( + _config(num_blocks), + max_model_len=max_model_len, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + ) + + +def test_reserve_and_free_multi_cfg_request() -> None: + manager = _manager(8) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + requests = (_request("public", 0), _request("public", 1)) + + metadata = manager.reserve_request("public", requests) + + assert metadata is not None + assert metadata.request_id == "public" + assert metadata.allocation_generation == 1 + assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1] + assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4] + assert [sequence.target_len for sequence in metadata.sequences] == [4, 4] + assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8] + assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2] + assert all(sequence.context_ids == () for sequence in metadata.sequences) + assert metadata.contexts == () + assert metadata == manager.get_metadata("public") + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4 + + manager.free_request("public") + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None: + manager = _manager(3) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + + with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"): + manager.reserve_request( + "public", + (_request("public", 0), _request("public", 1)), + ) + + assert manager.has_request("public") is False + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None: + manager = _manager(5) + assert manager.reserve_request("running", (_request("running", 0),)) is not None + free_before = manager.native_manager.block_pool.get_num_free_blocks() + + with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"): + manager.reserve_request( + "impossible", + tuple(_request("impossible", sequence_id) for sequence_id in range(3)), + ) + + assert manager.has_request("impossible") is False + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_temporary_capacity_pressure_returns_none() -> None: + manager = _manager(3) + assert manager.reserve_request("running", (_request("running", 0),)) is not None + + allocation = manager.reserve_request("waiting", (_request("waiting", 0),)) + + assert allocation is None + assert manager.has_request("waiting") is False + + +def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None: + manager = _manager(5) + assert manager.reserve_request("running", (_request("running", 0),)) is not None + waiting = (_request("waiting", 0), _request("waiting", 1)) + + assert manager.reserve_request("waiting", waiting) is None + assert manager.has_request("waiting") is False + + manager.free_request("running") + metadata = manager.reserve_request("waiting", waiting) + + assert metadata is not None + assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2] + + +def test_reallocated_request_gets_a_new_allocation_generation() -> None: + manager = _manager(3) + request = (_request("public", 0),) + + first = manager.reserve_request("public", request) + assert first is not None + manager.free_request("public") + second = manager.reserve_request("public", request) + + assert second is not None + assert second.allocation_generation == first.allocation_generation + 1 + + +def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None: + manager = _manager(8) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + native_allocate = manager.native_manager.allocate_slots + num_calls = 0 + + def fail_second_allocation(*args, **kwargs): + nonlocal num_calls + num_calls += 1 + if num_calls == 2: + raise RuntimeError("injected allocation failure") + return native_allocate(*args, **kwargs) + + monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation) + + with pytest.raises(RuntimeError, match="injected allocation failure"): + manager.reserve_request( + "public", + (_request("public", 0), _request("public", 1)), + ) + + assert manager.has_request("public") is False + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + + +def test_rejects_independent_context_until_role_routing_is_implemented() -> None: + manager = _manager(8) + context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8) + + with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"): + manager.reserve_request( + "public", + (_request("public", 0, kv_contexts=(context,)),), + ) + + +def test_rejects_sequence_longer_than_admission_bound() -> None: + manager = _manager(8, max_model_len=8) + + with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"): + manager.reserve_request("public", (_request("public", 0, seq_len=9),)) + + +@pytest.mark.parametrize( + "requests", + [ + (_request("public", 1),), + (_request("public", 1), _request("public", 0)), + (_request("public", 0), _request("public", 2)), + ], +) +def test_rejects_ambiguous_sequence_order(requests) -> None: + manager = _manager(8) + + with pytest.raises(ValueError, match="ordered by contiguous sequence_id"): + manager.reserve_request("public", requests) + + +def test_close_releases_every_public_request() -> None: + manager = _manager(8) + free_before = manager.native_manager.block_pool.get_num_free_blocks() + assert manager.reserve_request("first", (_request("first", 0),)) is not None + assert manager.reserve_request("second", (_request("second", 0),)) is not None + + manager.close() + + assert manager.native_manager.block_pool.get_num_free_blocks() == free_before + assert manager.has_request("first") is False + assert manager.has_request("second") is False diff --git a/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py b/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py new file mode 100644 index 000000000..4848fd49a --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from importlib import import_module, util + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] + + +def test_native_block_tables_slot_mapping_handles_non_aligned_tail() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the native BlockTables slot-mapping smoke test") + module_name = "vllm.v1.worker.gpu.block_table" + try: + module_spec = util.find_spec(module_name) + except ModuleNotFoundError as exc: + if exc.name is None or not module_name.startswith(f"{exc.name}."): + raise + pytest.skip("installed vLLM does not provide GPU BlockTables") + if module_spec is None: + pytest.skip("installed vLLM does not provide GPU BlockTables") + block_table_module = import_module(module_name) + block_tables_cls = getattr(block_table_module, "BlockTables", None) + if block_tables_cls is None or not hasattr(block_tables_cls, "compute_slot_mappings"): + pytest.skip("installed vLLM does not provide the required BlockTables API") + + device = torch.device("cuda", torch.accelerator.current_device_index()) + block_tables = block_tables_cls( + block_sizes=[4], + max_num_reqs=1, + max_num_batched_tokens=16, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[4], + ) + block_tables.append_block_ids(0, ([1, 2, 3],), overwrite=True) + block_tables.apply_staged_writes() + idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 9], dtype=torch.int32, device=device) + positions = torch.arange(9, dtype=torch.int64, device=device) + + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=9, + ) + torch.accelerator.synchronize(device) + + assert slot_mappings.cpu().tolist() == [[4, 5, 6, 7, 8, 9, 10, 11, 12]] diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py new file mode 100644 index 000000000..1f242420d --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py @@ -0,0 +1,1406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import math +from contextlib import nullcontext +from dataclasses import dataclass +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from torch import nn +from vllm.v1.kv_cache_interface import FullAttentionSpec + +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module +from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( + DiffusionPagedAttentionAdapter, + DiffusionPagedAttentionMetadata, + DiffusionPagedAttentionRow, + DiffusionPagedAttentionRowBinding, + DiffusionPagedAttentionRuntime, +) +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + override_paged_kv_adapter, + set_forward_context, + set_forward_context_denoise_step_idx, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _FakeBlockTables: + def __init__(self) -> None: + self.max_num_reqs = 4 + self.max_num_batched_tokens = 16 + self.cp_size = 1 + self.cp_rank = 0 + self.cp_interleave = 1 + self.block_sizes = [4] + self.blocks_per_kv_block = [1] + self.num_blocks = SimpleNamespace(np=np.full((1, self.max_num_reqs), 4, dtype=np.int32)) + self.gather_calls: list[tuple[torch.Tensor, int]] = [] + self.slot_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = [] + + def gather_block_tables( + self, + idx_mapping: torch.Tensor, + num_reqs_padded: int, + ) -> tuple[torch.Tensor, ...]: + self.gather_calls.append((idx_mapping.clone(), num_reqs_padded)) + return (torch.tensor([[3, 4], [7, 8]], dtype=torch.int32),) + + def compute_slot_mappings( + self, + idx_mapping: torch.Tensor, + query_start_loc: torch.Tensor, + positions: torch.Tensor, + num_tokens_padded: int, + ) -> torch.Tensor: + self.slot_calls.append( + ( + idx_mapping.clone(), + query_start_loc.clone(), + positions.clone(), + num_tokens_padded, + ) + ) + return positions[:num_tokens_padded].to(torch.int64).unsqueeze(0) + + +class _FakeSpec: + non_causal: bool + block_size = 4 + + def __init__(self, *, non_causal: bool) -> None: + self.non_causal = non_causal + + def max_num_blocks_per_req(self, _vllm_config, max_len: int) -> int: + return math.ceil(max_len / self.block_size) + + +class _FakeAttentionGroup: + def __init__(self, reorder_batch_threshold: int | None) -> None: + self.builder = SimpleNamespace(reorder_batch_threshold=reorder_batch_threshold) + + def get_metadata_builder(self, _index: int): + return self.builder + + +@dataclass(frozen=True, slots=True) +class _FakeNativeMetadata: + build_id: int + causal: bool + seq_lens: torch.Tensor + query_start_loc_cpu: torch.Tensor + positions: torch.Tensor + slot_mappings: torch.Tensor + + +class _FakeLayer: + def __init__(self, *, non_causal: bool, num_heads: int = 2, num_kv_heads: int = 2) -> None: + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_size = 4 + self.head_size_v = 4 + self.spec = _FakeSpec(non_causal=non_causal) + self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]] = [] + self.native_events: list[str] = [] + self.layer_name = "layer-0" + self.kv_cache = object() + self.attn_backend = SimpleNamespace(forward_includes_kv_cache_update=False) + self.impl = _FakeNativeImpl(self) + + +class _FakeNativeImpl: + def __init__(self, layer: _FakeLayer) -> None: + self.layer = layer + + def do_kv_cache_update( + self, + _layer, + key: torch.Tensor | None, + value: torch.Tensor | None, + _kv_cache, + slot_mapping: torch.Tensor, + ) -> None: + self.layer.native_events.append("update") + self.layer.updates.append((key, value, slot_mapping)) + + def forward( + self, + _layer, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + _kv_cache, + metadata: object, + output: torch.Tensor, + ) -> torch.Tensor: + self.layer.native_events.append("forward") + self.layer.calls.append((query, key, value, metadata)) + return output.copy_(query.reshape_as(output)) + + +def _make_adapter( + monkeypatch: pytest.MonkeyPatch, + *, + non_causal: bool = True, + reorder_batch_threshold: int | None = None, + num_heads: int = 2, + num_kv_heads: int = 2, + capacity: int = 16, +) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]: + events: list[tuple] = [] + monkeypatch.setattr( + adapter_module.current_omni_platform, + "requires_diffusion_paged_kv_prewrite", + lambda: False, + ) + block_tables = _FakeBlockTables() + block_tables.max_num_batched_tokens = capacity + block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0])) + layer = _FakeLayer(non_causal=non_causal, num_heads=num_heads, num_kv_heads=num_kv_heads) + kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=layer.spec)], + ) + row_map = { + ("req-0", 0, None): DiffusionPagedAttentionRowBinding(2, capacity), + ("req-1", 1, None): DiffusionPagedAttentionRowBinding(3, capacity), + ("req-0", None, "text"): DiffusionPagedAttentionRowBinding(1, capacity), + } + + monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext()) + + def build_metadata(**kwargs): + metadata = _FakeNativeMetadata( + build_id=len(events), + causal=kwargs["causal"], + seq_lens=kwargs["seq_lens"].clone(), + query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(), + positions=kwargs["positions"].clone(), + slot_mappings=kwargs["slot_mappings"].clone(), + ) + events.append(("build", kwargs, metadata)) + return {"layer-0": metadata} + + monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata) + monkeypatch.setattr( + adapter_module, + "build_slot_mappings_by_layer", + lambda slot_mappings, _config: {"layer-0": slot_mappings[0]}, + ) + + config = SimpleNamespace( + name="vllm-config", + model_config=SimpleNamespace(dtype=torch.float32), + ) + adapter = DiffusionPagedAttentionAdapter( + vllm_config=config, + device=torch.device("cpu"), + kv_cache_config=kv_cache_config, + block_tables=block_tables, + attn_groups=[[_FakeAttentionGroup(reorder_batch_threshold)]], + layers={"layer-0": layer}, + resolve_row=lambda request_id, sequence_id, context_id: row_map[(request_id, sequence_id, context_id)], + ) + return adapter, block_tables, layer, events + + +def _run_omni_paged_backend(context): + backend = FlashAttentionImpl( + num_heads=context.layer.num_heads, + head_size=context.layer.head_size, + softmax_scale=0.5, + num_kv_heads=context.layer.num_kv_heads, + ) + return backend.forward_paged(context) + + +def test_prepare_batch_reuses_native_block_table_and_metadata_builders(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, block_tables, _, events = _make_adapter(monkeypatch) + rows = ( + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=3, + seq_len=8, + kv_start_pos=4, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=2, + seq_len=2, + ), + ) + + batch = adapter.prepare_batch(rows) + + assert batch.num_tokens == 5 + assert batch.row_indices.tolist() == [2, 3] + assert batch.query_start_loc.tolist() == [0, 3, 5] + assert batch.seq_lens.tolist() == [8, 2] + assert batch.positions.tolist() == [4, 5, 6, 0, 1] + assert block_tables.gather_calls[0][1] == 2 + assert block_tables.gather_calls[0][0].tolist() == [2, 3] + assert block_tables.slot_calls[0][3] == 5 + build_kwargs = events[0][1] + assert build_kwargs["causal"] is False + assert build_kwargs["max_query_len"] == 3 + assert build_kwargs["max_seq_len"] == 8 + assert build_kwargs["positions"] is batch.positions + assert batch.attn_metadata["layer-0"] is events[0][2] + assert events[0][2].build_id == 0 + assert batch.slot_mappings_by_layer["layer-0"].tolist() == [4, 5, 6, 0, 1] + + +def test_omni_paged_backend_consumes_context_and_restores_diffusion_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=5, + seq_len=5, + ) + ] + ) + query = torch.randn(1, 5, 2, 4) + key = torch.randn_like(query) + value = torch.randn_like(query) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + output = _run_omni_paged_backend(context) + + assert output.shape == query.shape + assert torch.equal(output, query) + assert layer.calls[0][0].shape == (5, 2, 4) + assert layer.native_events == ["update", "forward"] + assert layer.calls[0][3] is events[0][2] + + +def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch, num_heads=32, num_kv_heads=8) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + query = torch.randn(1, 3, 32, 4) + key = torch.randn(1, 3, 8, 4) + value = torch.randn_like(key) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + + assert layer.num_heads == 32 + assert layer.num_kv_heads == 8 + assert context.query.shape == (3, 32, 4) + assert context.key_write.shape == (3, 8, 4) + assert context.value_write.shape == (3, 8, 4) + + +def test_paged_adapter_preserves_strided_kv_for_platform_writer(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)] + ) + packed_qkv = torch.randn(1, 5, 3 * 2 * 4) + query, key, value = packed_qkv.split(2 * 4, dim=-1) + query = query.reshape(1, 5, 2, 4) + key = key.reshape(1, 5, 2, 4) + value = value.reshape(1, 5, 2, 4) + assert not key.is_contiguous() + assert not value.is_contiguous() + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + + assert context.query.is_contiguous() + assert not context.key_write.is_contiguous() + assert not context.value_write.is_contiguous() + assert context.key_write.data_ptr() == key.data_ptr() + assert context.value_write.data_ptr() == value.data_ptr() + + +def test_paged_adapter_packs_noncontiguous_kv_head_layout(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)] + ) + query = torch.randn(1, 5, 2, 4) + key = torch.randn(1, 5, 4, 2).transpose(-1, -2) + value = torch.randn(1, 5, 4, 2).transpose(-1, -2) + assert key.shape == (1, 5, 2, 4) + assert key.stride(-1) != 1 + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + + assert context.key_write.is_contiguous() + assert context.value_write.is_contiguous() + + +def test_omni_paged_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + layer.attn_backend.forward_includes_kv_cache_update = True + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + qkv = torch.randn(1, 3, 2, 4) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, qkv) + assert layer.native_events == ["forward"] + assert layer.updates == [] + + +def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch) + layer.attn_backend.forward_includes_kv_cache_update = True + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={}) + + with adapter.activate(batch): + context = adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, qkv) + assert layer.native_events == ["forward", "forward", "forward"] + assert layer.updates == [] + assert [call[0].shape[0] for call in layer.calls] == [1, 2, 1] + assert [call[1].shape[0] for call in layer.calls] == [1, 2, 1] + assert [call[2].shape[0] for call in layer.calls] == [1, 2, 1] + assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True)) + segment_metadata = [call[3] for call in layer.calls] + assert [metadata.causal for metadata in segment_metadata] == [True, False, True] + assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[1], [3], [4]] + assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [[0, 1], [0, 2], [0, 1]] + assert [metadata.positions.tolist() for metadata in segment_metadata] == [[0], [1, 2], [3]] + assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [[[0]], [[1, 2]], [[3]]] + for event in events[1:]: + build_kwargs = event[1] + assert build_kwargs["positions"].untyped_storage().data_ptr() == batch.positions.untyped_storage().data_ptr() + assert ( + build_kwargs["slot_mappings"].untyped_storage().data_ptr() + == batch.slot_mappings.untyped_storage().data_ptr() + ) + + +def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + layer.attn_backend.forward_includes_kv_cache_update = True + monkeypatch.setattr( + adapter_module.current_omni_platform, + "requires_diffusion_paged_kv_prewrite", + lambda: True, + ) + adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding( + row_index=2, + max_seq_len=16, + block_ids=((3, 7, 9, 11),), + ) + row = DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4) + batch = adapter.prepare_batch([row]) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={}) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, qkv) + assert layer.native_events == ["update", "forward", "forward", "forward"] + assert len(layer.updates) == 1 + assert all(call[1] is None and call[2] is None for call in layer.calls) + + +def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch) + # Exercise the Ascend-only homogeneous batch path. CUDA keeps its native + # output-buffer contract and is covered by the GPU adapter tests. + monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=6, + seq_len=9, + kv_start_pos=3, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=6, + seq_len=9, + kv_start_pos=3, + ), + ] + ) + query = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4) + key = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4) + 100 + value = key + 100 + attn_mask = torch.ones(9, 9, dtype=torch.bool).tril().repeat(2, 1, 1) + attn_mask[:, 5:8, 5:8] = True + metadata = SimpleNamespace( + attn_mask=attn_mask.unsqueeze(1), + full_attn_spans=[[(5, 8)], [(5, 8)]], + query_ranges=None, + extra={}, + ) + + with adapter.activate(batch): + context = adapter.prepare_layer_context( + "layer-0", + query, + key, + value, + omni_attn_metadata=metadata, + ) + # Every model layer sees the same spans, so the native segment metadata + # is built once per active batch and then reused. + adapter.prepare_layer_context( + "layer-0", + query, + key, + value, + omni_attn_metadata=metadata, + ) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, query) + assert layer.native_events == ["update", "forward", "forward", "forward"] + assert len(layer.updates) == 1 + assert layer.updates[0][0].shape[0] == 12 + assert torch.equal(layer.updates[0][0], key.reshape(12, 2, 4)) + assert torch.equal(layer.updates[0][1], value.reshape(12, 2, 4)) + assert [call[0].shape[0] for call in layer.calls] == [4, 6, 2] + assert [call[1].shape[0] for call in layer.calls] == [4, 6, 2] + assert [call[2].shape[0] for call in layer.calls] == [4, 6, 2] + assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True)) + assert [metadata.build_id for metadata in context.piecewise_native_metadata] == [1, 2, 3] + + # The first metadata build is the normal whole-query path. Piecewise calls + # then use causal [3, 5), full [5, 8), and causal [8, 9) segments. + assert len(events) == 4 + segment_builds = [event[1] for event in events[1:]] + assert [build["causal"] for build in segment_builds] == [True, False, True] + assert [build["seq_lens"].tolist() for build in segment_builds] == [ + [5, 5], + [8, 8], + [9, 9], + ] + assert [build["query_start_loc_cpu"].tolist() for build in segment_builds] == [ + [0, 2, 4], + [0, 3, 6], + [0, 1, 2], + ] + segment_metadata = [call[3] for call in layer.calls] + assert [metadata.causal for metadata in segment_metadata] == [True, False, True] + assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[5, 5], [8, 8], [9, 9]] + assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [ + [0, 2, 4], + [0, 3, 6], + [0, 1, 2], + ] + assert [metadata.positions.tolist() for metadata in segment_metadata] == [ + [3, 4, 3, 4], + [5, 6, 7, 5, 6, 7], + [8, 8], + ] + assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [ + [[3, 4, 3, 4]], + [[5, 6, 7, 5, 6, 7]], + [[8, 8]], + ] + + +def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32) + build_id = 0 + + def build_metadata(**kwargs): + nonlocal build_id + shared_scheduler_metadata.fill_(build_id) + metadata = SimpleNamespace( + build_id=build_id, + causal=kwargs["causal"], + seq_lens=kwargs["seq_lens"].clone(), + query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(), + positions=kwargs["positions"].clone(), + slot_mappings=kwargs["slot_mappings"].clone(), + scheduler_metadata=shared_scheduler_metadata, + prefix_scheduler_metadata=None, + ) + build_id += 1 + return {"layer-0": metadata} + + monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={}) + + with adapter.activate(batch): + context = adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + segment_schedules = [item.scheduler_metadata for item in context.piecewise_native_metadata] + assert [item.item() for item in segment_schedules] == [1, 2, 3] + assert len({item.data_ptr() for item in segment_schedules}) == 3 + + +def test_omni_paged_backend_treats_empty_full_spans_as_causal(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace( + attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool).tril(), + full_attn_spans=[[]], + query_ranges=None, + extra={}, + ) + + with adapter.activate(batch): + context = adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, qkv) + assert layer.native_events == ["update", "forward"] + assert len(events) == 2 + assert events[1][1]["causal"] is True + assert events[1][1]["seq_lens"].tolist() == [4] + + +def test_omni_paged_backend_runs_unequal_hunyuan_prefixes_without_padding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, events = _make_adapter(monkeypatch, capacity=64) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=17, + seq_len=50, + kv_start_pos=33, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=17, + seq_len=54, + kv_start_pos=37, + ), + ] + ) + query = torch.arange(2 * 17 * 2 * 4, dtype=torch.float32).reshape(2, 17, 2, 4) + key = query + 100 + value = query + 200 + metadata = SimpleNamespace( + attn_mask=torch.ones(2, 1, 17, 54, dtype=torch.bool), + full_attn_spans=[[(33, 50)], [(37, 54)]], + query_ranges=None, + extra={}, + ) + + with adapter.activate(batch): + context = adapter.prepare_layer_context( + "layer-0", + query, + key, + value, + omni_attn_metadata=metadata, + ) + output = _run_omni_paged_backend(context) + + assert torch.equal(output, query) + assert batch.query_start_loc.tolist() == [0, 17, 34] + assert batch.seq_lens.tolist() == [50, 54] + assert batch.positions.tolist() == [*range(33, 50), *range(37, 54)] + assert torch.equal(context.key_write, key.reshape(34, 2, 4)) + assert torch.equal(context.value_write, value.reshape(34, 2, 4)) + assert layer.native_events == ["update", "forward"] + assert layer.updates[0][2].tolist() == batch.positions.tolist() + assert layer.calls[0][3].seq_lens.tolist() == [50, 54] + assert layer.calls[0][3].query_start_loc_cpu.tolist() == [0, 17, 34] + assert len(events) == 2 + + +def test_forward_rejects_unaligned_heterogeneous_piecewise_spans_before_cache_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4), + DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=4, seq_len=4), + ] + ) + qkv = torch.randn(2, 4, 2, 4) + metadata = SimpleNamespace( + attn_mask=torch.ones(2, 1, 4, 4, dtype=torch.bool), + full_attn_spans=[[(1, 3)], [(2, 4)]], + query_ranges=None, + extra={}, + ) + + with adapter.activate(batch), pytest.raises(ValueError, match="produce aligned segments"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +@pytest.mark.parametrize( + "spans", + [ + [(1, 3), (2, 4)], + [(1, 5)], + ], +) +def test_forward_rejects_invalid_piecewise_spans_before_cache_write( + monkeypatch: pytest.MonkeyPatch, + spans: list[tuple[int, int]], +) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace(full_attn_spans=[spans], extra={}) + + with adapter.activate(batch), pytest.raises(ValueError, match="sorted, non-overlapping, and within the sequence"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_rejects_non_4d_mask_with_piecewise_spans_before_cache_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace( + attn_mask=torch.ones(1, 4, 4, dtype=torch.bool), + full_attn_spans=[[(1, 3)]], + extra={}, + ) + + with adapter.activate(batch), pytest.raises(NotImplementedError, match="requires a 4D tensor"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_rejects_piecewise_query_ranges_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + qkv = torch.randn(1, 4, 2, 4) + metadata = SimpleNamespace( + attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool), + full_attn_spans=[[(1, 3)]], + query_ranges=(object(),), + extra={}, + ) + + with adapter.activate(batch), pytest.raises(NotImplementedError, match="query_ranges"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_rejects_piecewise_non_suffix_query_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=6, + kv_start_pos=2, + ) + ] + ) + qkv = torch.randn(1, 2, 2, 4) + metadata = SimpleNamespace(full_attn_spans=[[(2, 4)]], extra={}) + + with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"): + adapter.prepare_layer_context( + "layer-0", + qkv, + qkv, + qkv, + omni_attn_metadata=metadata, + ) + + assert layer.native_events == [] + + +def test_forward_accepts_non_aligned_current_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=5, + kv_start_pos=3, + ) + ] + ) + query = torch.randn(1, 2, 2, 4) + key = torch.arange(1 * 2 * 2 * 4, dtype=torch.float32).reshape(1, 2, 2, 4) + value = key + 100 + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, key, value) + + assert context.query.shape == (2, 2, 4) + assert torch.equal(context.key_write, key.reshape(2, 2, 4)) + assert torch.equal(context.value_write, value.reshape(2, 2, 4)) + assert context.slot_mapping.tolist() == [3, 4] + assert layer.native_events == [] + + +def test_forward_rejects_untranslated_omni_metadata_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, layer, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + qkv = torch.randn(1, 2, 4) + metadata = SimpleNamespace(attn_mask=torch.ones(1, 1), extra={}) + + with adapter.activate(batch), pytest.raises(NotImplementedError, match="cannot translate"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata) + + assert layer.updates == [] + + +def test_forward_accepts_flattened_hidden_dimension(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + query = torch.randn(3, 8) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", query, query, query) + + assert context.query.shape == (3, 2, 4) + assert context.restore_output(context.query).shape == query.shape + + +def test_forward_accepts_one_flattened_token_for_single_head(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch, num_heads=1, num_kv_heads=1) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + qkv = torch.randn(1, 4) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + output = _run_omni_paged_backend(context) + + assert context.query.shape == (1, 1, 4) + assert context.query_token_shape == (1,) + assert not context.query_has_head_dims + assert output.shape == qkv.shape + + +def test_forward_requires_active_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + qkv = torch.randn(1, 2, 4) + + with pytest.raises(RuntimeError, match="adapter.activate"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_new_preparation_invalidates_native_buffer_views_in_older_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + old_batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)]) + + with pytest.raises(ValueError, match="stale"): + with adapter.activate(old_batch): + pass + + +def test_block_table_change_invalidates_prepared_native_buffer_views( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + + adapter.invalidate_prepared_batches() + + with pytest.raises(ValueError, match="stale"): + with adapter.activate(batch): + pass + + +def test_block_table_change_is_rejected_during_active_forward(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + + with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"): + adapter.invalidate_prepared_batches() + + +def test_prepare_batch_rejects_active_native_buffers(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + + with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"): + adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)]) + + +def test_forward_rejects_tensor_count_different_from_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)] + ) + qkv = torch.randn(2, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="token count"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_forward_rejects_dtype_different_from_model_activation_dtype(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)] + ) + qkv = torch.randn(1, 2, 4, dtype=torch.float64) + + with adapter.activate(batch), pytest.raises(ValueError, match="model activation dtype"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_forward_rejects_invalid_key_token_count(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)] + ) + query = torch.randn(4, 2, 4) + key = torch.randn(3, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="key token count"): + adapter.prepare_layer_context("layer-0", query, key, query) + + +def test_forward_rejects_batched_layout_that_does_not_match_rows(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1), + DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3), + ] + ) + qkv = torch.randn(2, 2, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="must match prepared rows"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_forward_packs_and_restores_right_padded_heterogeneous_rows(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1), + DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3), + ] + ) + qkv = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4) + + with adapter.activate(batch): + context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + assert context.output_scatter_indices.tolist() == [0, 3, 4, 5] + torch.testing.assert_close(context.query, qkv.reshape(-1, 2, 4)[[0, 3, 4, 5]]) + restored = context.restore_output(context.query) + assert restored.shape == qkv.shape + torch.testing.assert_close(restored[0, 0], qkv[0, 0]) + assert torch.count_nonzero(restored[0, 1:]) == 0 + torch.testing.assert_close(restored[1], qkv[1]) + + +def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=1, + seq_len=3, + kv_start_pos=2, + ), + DiffusionPagedAttentionRow( + request_id="req-1", + sequence_id=1, + query_len=1, + seq_len=5, + kv_start_pos=4, + ), + ] + ) + query = torch.randn(2, 1, 2, 4) + full_kv = torch.randn(2, 4, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="batched key layout must match prepared rows"): + adapter.prepare_layer_context("layer-0", query, full_kv, full_kv) + + +def test_runner_runtime_switches_from_prefill_to_reusable_denoise_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch) + metadata = DiffusionPagedAttentionMetadata( + prefill_rows=(DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=6, seq_len=6),), + denoise_rows=( + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=5, + kv_start_pos=3, + ), + ), + ) + runtime = DiffusionPagedAttentionRuntime(adapter, metadata) + + with set_forward_context(paged_kv_runtime=runtime), runtime.activate(): + set_forward_context_denoise_step_idx(0) + assert Attention._active_paged_kv_adapter() is adapter + assert adapter._active_batch.rows == metadata.prefill_rows + + set_forward_context_denoise_step_idx(1) + assert Attention._active_paged_kv_adapter() is adapter + denoise_batch = adapter._active_batch + assert denoise_batch.rows == metadata.denoise_rows + + set_forward_context_denoise_step_idx(2) + assert Attention._active_paged_kv_adapter() is adapter + assert adapter._active_batch is denoise_batch + + assert adapter._active_batch is None + + +def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False) + row = DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=8, + kv_start_pos=3, + ) + + batch = adapter.prepare_batch([row]) + qkv = torch.randn(2, 2, 4) + + with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"): + adapter.prepare_layer_context("layer-0", qkv, qkv, qkv) + + +def test_prepare_batch_rejects_native_decode_after_prefill_order(monkeypatch: pytest.MonkeyPatch) -> None: + adapter, block_tables, _, _ = _make_adapter( + monkeypatch, + non_causal=False, + reorder_batch_threshold=1, + ) + rows = [ + DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=2, seq_len=2), + DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1), + ] + + with pytest.raises(ValueError, match="decode/short-query rows before"): + adapter.prepare_batch(rows) + + assert block_tables.gather_calls == [] + assert block_tables.slot_calls == [] + + +def test_prepare_batch_rejects_seq_len_beyond_installed_row_capacity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, block_tables, _, _ = _make_adapter(monkeypatch) + block_tables.num_blocks.np[0, 2] = 1 + row = DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=5, + seq_len=5, + ) + + with pytest.raises(ValueError, match="installed Worker row contains only 1 blocks"): + adapter.prepare_batch([row]) + + assert block_tables.gather_calls == [] + assert block_tables.slot_calls == [] + + +def test_prepare_batch_rejects_seq_len_beyond_installed_logical_length( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, block_tables, _, _ = _make_adapter(monkeypatch) + adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(2, 4) + row = DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=5, + seq_len=5, + ) + + with pytest.raises(ValueError, match="installed allocation has logical length 4"): + adapter.prepare_batch([row]) + + assert block_tables.gather_calls == [] + assert block_tables.slot_calls == [] + + +def test_prepare_batch_passes_dcp_local_seq_lens_to_native_builder( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter, block_tables, _, events = _make_adapter(monkeypatch) + block_tables.cp_size = 2 + block_tables.cp_rank = 1 + block_tables.cp_interleave = 2 + block_tables.num_blocks.np[0, 2] = 1 + + batch = adapter.prepare_batch( + [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)] + ) + + assert batch.row_indices.tolist() == [2] + assert events[0][1]["dcp_local_seq_lens"].tolist() == [2] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"sequence_id": 0, "context_id": "text"}, "exactly one"), + ({}, "exactly one"), + ({"sequence_id": 0, "query_len": 0}, "query_len"), + ({"sequence_id": 0, "kv_start_pos": 4, "query_len": 2, "seq_len": 5}, "exceeds seq_len"), + ], +) +def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: str) -> None: + values = { + "request_id": "req-0", + "query_len": 1, + "seq_len": 1, + **kwargs, + } + with pytest.raises(ValueError, match=message): + DiffusionPagedAttentionRow(**values) + + +def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + selected_backends = [] + specialized_backends = [] + impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock())) + native_backend = SimpleNamespace( + get_name=lambda: "ASCEND", + indexes_kv_by_block_stride=lambda: True, + get_impl_cls=lambda: impl_cls, + forward_includes_kv_cache_update=True, + ) + + original_backend_per_kind = {"full": object()} + config = SimpleNamespace( + attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind), + parallel_config=SimpleNamespace(prefill_context_parallel_size=2), + model_config=SimpleNamespace(dtype=torch.float16), + cache_config=SimpleNamespace(cache_dtype="auto"), + ) + + def select_backend(**_kwargs): + selected_backends.append( + ( + config.attention_config.backend, + config.attention_config.backend_per_kind, + config.parallel_config.prefill_context_parallel_size, + ) + ) + return native_backend + + monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend) + monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext()) + monkeypatch.setattr( + adapter_module.current_omni_platform, + "get_diffusion_paged_kv_attn_backend", + lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend, + ) + layer = SimpleNamespace(num_heads=8, softmax_scale=0.125) + spec = FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=8, + dtype=torch.float16, + non_causal=True, + ) + + native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter( + layer_name="layer-0", + layer=layer, + spec=spec, + vllm_config=config, + device=torch.device("cpu"), + ulysses_degree=2, + ) + + assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)] + assert specialized_backends == [(native_backend, 2)] + assert config.attention_config.backend is None + assert config.attention_config.backend_per_kind is original_backend_per_kind + assert config.parallel_config.prefill_context_parallel_size == 2 + assert native_layer.num_heads == 4 + assert native_layer.num_kv_heads == 2 + assert native_layer.spec.num_kv_heads == 2 + assert native_layer.spec.indexes_kv_by_block_stride is True + assert native_layer._q_scale_float == 1.0 + assert native_layer._k_scale_float == 1.0 + assert native_layer._v_scale_float == 1.0 + impl_kwargs = impl_cls.call_args.kwargs + assert impl_kwargs["num_heads"] == 4 + assert impl_kwargs["num_kv_heads"] == 2 + + +def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None: + events: list[str] = [] + + class Strategy: + name = "ulysses" + + def pre_attention(self, query, key, value, metadata): + events.append("pre") + return query + 1, key + 2, value + 3, metadata, object() + + def post_attention(self, output, _ctx): + events.append("post") + return output + 4 + + class Adapter: + def prepare_layer_context(self, layer_name, query, key, value, *, omni_attn_metadata): + events.append("prepare") + assert layer_name == "layer-0" + assert omni_attn_metadata is None + assert torch.equal(key, original + 2) + assert torch.equal(value, original + 3) + return SimpleNamespace(query=query) + + class Backend: + def forward_paged(self, context): + events.append("kernel") + return context.query * 2 + + layer = Attention.__new__(Attention) + nn.Module.__init__(layer) + layer.prefix = "layer-0" + layer.paged_kv_cache_role = "primary" + layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN") + layer.attention = Backend() + layer.use_ring = False + layer._no_parallel_strategy = object() + layer._get_active_parallel_strategy = lambda: Strategy() + layer._with_kv_cache_dtype = lambda metadata: metadata + original = torch.zeros(1, 2, 2, 4) + + with set_forward_context(), override_paged_kv_adapter(Adapter()): + assert layer.is_paged_kv_active() + output = layer._forward_impl(original, original, original) + + assert events == ["pre", "prepare", "kernel", "post"] + assert torch.equal(output, torch.full_like(original, 6)) + + +def test_omni_attention_strips_paged_ulysses_padding_around_kernel() -> None: + events: list[tuple[str, tuple[int, ...]]] = [] + + class Strategy: + name = "ulysses" + + def pre_attention(self, query, key, value, metadata): + return query, key, value, metadata, SimpleNamespace(joint_len=0, joint_strategy="front") + + def post_attention(self, output, _ctx): + events.append(("post", tuple(output.shape))) + return output + + class Adapter: + def prepare_layer_context(self, _layer_name, query, key, value, *, omni_attn_metadata): + assert omni_attn_metadata is None + events.append(("kernel_input", tuple(query.shape))) + assert query[:, -1].tolist() == [[[3.0]]] + assert torch.equal(query, key) + assert torch.equal(query, value) + return SimpleNamespace(query=query) + + class Backend: + def forward_paged(self, context): + return context.query + + layer = Attention.__new__(Attention) + nn.Module.__init__(layer) + layer.prefix = "layer-0" + layer.paged_kv_cache_role = "primary" + layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN") + layer.attention = Backend() + layer.use_ring = False + layer._no_parallel_strategy = object() + layer._get_active_parallel_strategy = lambda: Strategy() + tensor = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1, 1) + + with set_forward_context(), override_paged_kv_adapter(Adapter()): + get_forward_context().sp_padding_size = 1 + output = layer._forward_impl(tensor, tensor, tensor) + + assert events == [("kernel_input", (1, 4, 1, 1)), ("post", (1, 5, 1, 1))] + assert output.flatten().tolist() == [0.0, 1.0, 2.0, 3.0, 0.0] + + +def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None: + class Adapter: + def prepare_layer_context(self, *args, **kwargs): + raise AssertionError("an unsupported backend must fail before adapter preparation") + + layer = Attention.__new__(Attention) + nn.Module.__init__(layer) + layer.prefix = "layer-0" + layer.paged_kv_cache_role = "primary" + layer.attn_backend = SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA") + layer.use_ring = False + layer._no_parallel_strategy = object() + layer._get_active_parallel_strategy = lambda: layer._no_parallel_strategy + + with ( + set_forward_context(), + override_paged_kv_adapter(Adapter()), + pytest.raises(NotImplementedError, match="backend with paged support"), + ): + layer._forward_impl(torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4)) + + +def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None: + events: list[str] = [] + + class Strategy: + name = "ulysses" + + def pre_attention(self, query, key, value, metadata): + events.append("pre") + return query, key, value, metadata, object() + + def post_attention(self, output, _ctx): + events.append("post") + return output + + layer = Attention.__new__(Attention) + nn.Module.__init__(layer) + layer.paged_kv_cache_role = "primary" + layer.use_ring = False + layer._no_parallel_strategy = object() + layer._get_active_parallel_strategy = lambda: Strategy() + layer._with_kv_cache_dtype = lambda metadata: metadata + layer._run_local_attention = lambda query, _key, _value, _metadata: events.append("dense") or query + qkv = torch.zeros(1, 2, 2, 4) + + assert not layer.is_paged_kv_active() + output = layer._forward_impl(qkv, qkv, qkv) + + assert events == ["pre", "dense", "post"] + assert output is qkv + + +def test_omni_attention_rejects_scheduler_paged_forward_without_active_adapter() -> None: + layer = Attention.__new__(Attention) + nn.Module.__init__(layer) + layer._scheduler_paged_kv = True + layer.paged_kv_cache_role = "primary" + layer._get_active_parallel_strategy = lambda: object() + qkv = torch.zeros(1, 2, 2, 4) + + with set_forward_context(), pytest.raises(RuntimeError, match="without an active Worker adapter"): + layer._forward_impl(qkv, qkv, qkv) diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py new file mode 100644 index 000000000..ef5d75f4d --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.nn import functional as F +from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache +from vllm.v1.worker.gpu.block_table import BlockTables + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl +from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( + DiffusionPagedAttentionAdapter, + DiffusionPagedAttentionLayerAdapter, + DiffusionPagedAttentionRow, + DiffusionPagedAttentionRowBinding, +) +from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] + +_LAYER_NAME = "model.layers.0.attn" +_NUM_HEADS = 2 +_HEAD_SIZE = 64 +_BLOCK_SIZE = 16 + + +class _SmokeDiffusionAttention(nn.Module): + def __init__(self) -> None: + super().__init__() + self.num_heads = _NUM_HEADS + self.num_kv_heads = _NUM_HEADS + self.head_size = _HEAD_SIZE + self.softmax_scale = _HEAD_SIZE**-0.5 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention") +def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None: + device = torch.device("cuda", torch.accelerator.current_device_index()) + vllm_config = VllmConfig( + cache_config=CacheConfig( + block_size=_BLOCK_SIZE, + cache_dtype="float16", + enable_prefix_caching=False, + ), + attention_config=AttentionConfig(use_non_causal=True), + ) + model_config = _DiffusionVllmModelConfig( + model=None, + dtype=torch.float16, + max_model_len=32, + original_max_model_len=32, + hf_config=SimpleNamespace(model_type="diffusion_paged_attention_smoke"), + ) + vllm_config.model_config = model_config + model_config.set_attention_geometry( + num_heads=_NUM_HEADS, + num_kv_heads=_NUM_HEADS, + head_size=_HEAD_SIZE, + ) + vllm_config.compilation_config.static_forward_context.clear() + + diffusion_layer = _SmokeDiffusionAttention().to(device) + spec = FullAttentionSpec( + block_size=_BLOCK_SIZE, + num_kv_heads=_NUM_HEADS, + head_size=_HEAD_SIZE, + dtype=torch.float16, + non_causal=True, + ) + with set_current_vllm_config(vllm_config): + native_layer = DiffusionPagedAttentionLayerAdapter( + layer_name=_LAYER_NAME, + layer=diffusion_layer, + spec=spec, + vllm_config=vllm_config, + device=device, + ) + vllm_config.compilation_config.static_forward_context[_LAYER_NAME] = native_layer + + canonical_spec = native_layer.spec + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=canonical_spec.page_size_bytes * 4, + shared_by=[_LAYER_NAME], + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec( + layer_names=[_LAYER_NAME], + kv_cache_spec=canonical_spec, + ) + ], + ) + with set_current_vllm_config(vllm_config): + attn_groups, _, kernel_block_sizes = init_attn_backend( + kv_cache_config, + vllm_config, + device, + ) + runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + init_kv_cache( + runner_kv_caches, + vllm_config.compilation_config.static_forward_context, + kv_cache_config, + attn_groups, + device, + vllm_config.cache_config.cache_dtype, + kernel_block_sizes, + vllm_config, + ) + block_tables = BlockTables( + block_sizes=[_BLOCK_SIZE], + max_num_reqs=1, + max_num_batched_tokens=32, + max_num_blocks_per_group=[2], + device=device, + kernel_block_sizes=kernel_block_sizes, + ) + + block_tables.append_block_ids(0, ([3, 1],), overwrite=True) + block_tables.apply_staged_writes() + adapter = DiffusionPagedAttentionAdapter( + vllm_config=vllm_config, + device=device, + kv_cache_config=kv_cache_config, + block_tables=block_tables, + attn_groups=attn_groups, + layers={_LAYER_NAME: native_layer}, + resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding( + row_index=0, + max_seq_len=19, + ), + ) + prefix_batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=17, + seq_len=17, + ) + ] + ) + + torch.manual_seed(1) + query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device) + packed_qkv = torch.randn( + 19, + 3 * _NUM_HEADS * _HEAD_SIZE, + dtype=torch.float16, + device=device, + ) + _, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1) + key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE) + value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE) + assert not key.is_contiguous() + assert not value.is_contiguous() + omni_backend = FlashAttentionImpl( + num_heads=_NUM_HEADS, + head_size=_HEAD_SIZE, + softmax_scale=diffusion_layer.softmax_scale, + num_kv_heads=_NUM_HEADS, + ) + with adapter.activate(prefix_batch): + prefix_context = adapter.prepare_layer_context(_LAYER_NAME, query[:17], key[:17], value[:17]) + prefix_output = omni_backend.forward_paged(prefix_context) + prefix_slot_mappings = prefix_batch.slot_mappings.clone() + + suffix_batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=2, + seq_len=19, + kv_start_pos=17, + ) + ] + ) + with adapter.activate(suffix_batch): + suffix_context = adapter.prepare_layer_context(_LAYER_NAME, query[17:], key[17:], value[17:]) + suffix_output = omni_backend.forward_paged(suffix_context) + suffix_slot_mappings = suffix_batch.slot_mappings.clone() + + piecewise_batch = adapter.prepare_batch( + [ + DiffusionPagedAttentionRow( + request_id="req-0", + sequence_id=0, + query_len=19, + seq_len=19, + ) + ] + ) + mixed_mask = torch.ones(19, 19, dtype=torch.bool, device=device).tril() + mixed_mask[5:10, 5:10] = True + mixed_metadata = AttentionMetadata( + attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0), + full_attn_spans=[[(5, 10)]], + ) + with adapter.activate(piecewise_batch): + piecewise_context = adapter.prepare_layer_context( + _LAYER_NAME, + query, + key, + value, + omni_attn_metadata=mixed_metadata, + ) + piecewise_output = omni_backend.forward_paged(piecewise_context) + + prefix_reference = ( + F.scaled_dot_product_attention( + query[:17].transpose(0, 1).unsqueeze(0), + key[:17].transpose(0, 1).unsqueeze(0), + value[:17].transpose(0, 1).unsqueeze(0), + is_causal=False, + scale=diffusion_layer.softmax_scale, + ) + .squeeze(0) + .transpose(0, 1) + ) + suffix_reference = ( + F.scaled_dot_product_attention( + query[17:].transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + is_causal=False, + scale=diffusion_layer.softmax_scale, + ) + .squeeze(0) + .transpose(0, 1) + ) + piecewise_reference = ( + F.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + key.transpose(0, 1).unsqueeze(0), + value.transpose(0, 1).unsqueeze(0), + attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0), + scale=diffusion_layer.softmax_scale, + ) + .squeeze(0) + .transpose(0, 1) + ) + torch.accelerator.synchronize(device) + + assert prefix_slot_mappings.cpu().tolist() == [list(range(3 * _BLOCK_SIZE, 4 * _BLOCK_SIZE)) + [_BLOCK_SIZE]] + assert suffix_slot_mappings.cpu().tolist() == [[_BLOCK_SIZE + 1, _BLOCK_SIZE + 2]] + torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2) diff --git a/tests/diffusion/diffusion_kv/test_request.py b/tests/diffusion/diffusion_kv/test_request.py index 43734ca31..ae5609b2a 100644 --- a/tests/diffusion/diffusion_kv/test_request.py +++ b/tests/diffusion/diffusion_kv/test_request.py @@ -107,11 +107,8 @@ def test_request_rejects_non_context_values() -> None: def _manager(*, enable_caching: bool) -> KVCacheManager: # vLLM 0.26 lazily registers its built-in KV specs. Another test may have # already registered an out-of-tree spec, making the registry non-empty - # before FullAttentionSpec is registered. Initialize the native mappings - # explicitly; vLLM 0.25 has no registration hook and needs no setup. - register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None) - if register_all is not None: - register_all(None) + # before FullAttentionSpec is registered. Initialize the native mappings. + native_kv_managers.register_all_kvcache_specs(None) spec = FullAttentionSpec( block_size=BLOCK_SIZE, diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py new file mode 100644 index 000000000..31faf4d08 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_spec_discovery.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from torch import nn +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor + +import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _Backend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return True + + +class _RunnerKVBackend: + def __init__(self) -> None: + self.kv_cache_config = None + + def register_kv_cache_layers(self, layers): + return {layer_name: spec for layer_name, (_layer, spec) in layers.items()} + + def initialize_kv_cache(self, config) -> None: + configured_layers = {layer_name for group in config.kv_cache_groups for layer_name in group.layer_names} + if configured_layers != {"image_attention"}: + raise ValueError( + "Rank-local Diffusion KVCacheConfig layer mismatch: " + f"expected=['image_attention'], configured={sorted(configured_layers)}" + ) + self.kv_cache_config = config + + +def _attention(*, enabled: bool, prefix: str = "image_attention") -> Attention: + attention = Attention.__new__(Attention) + nn.Module.__init__(attention) + attention.prefix = prefix + attention.paged_kv_cache_role = "primary" if enabled else None + attention.paged_kv_cache_dtype = torch.bfloat16 + attention.num_kv_heads = 2 + attention.head_size = 8 + attention.causal = False + attention.attn_backend = _Backend + return attention + + +def _runner(attention: Attention) -> DiffusionModelRunner: + runner = object.__new__(DiffusionModelRunner) + runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER) + runner.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=4), + model_config=SimpleNamespace(dtype=torch.float16), + ) + pipeline = nn.Module() + pipeline.image_attention = attention + runner.pipeline = pipeline + runner.diffusion_kv_backend = _RunnerKVBackend() + return runner + + +def test_runner_discovers_native_spec_from_loaded_attention() -> None: + specs = _runner(_attention(enabled=True)).get_kv_cache_spec() + + assert set(specs) == {"image_attention"} + spec = specs["image_attention"] + assert isinstance(spec, FullAttentionSpec) + assert spec.block_size == 4 + assert spec.num_kv_heads == 2 + assert spec.head_size == 8 + assert spec.dtype is torch.bfloat16 + assert spec.indexes_kv_by_block_stride is True + assert spec.non_causal is True + + +def test_runner_uses_attention_prefix_as_canonical_layer_name() -> None: + runner = _runner(_attention(enabled=True, prefix="layers.0.self_attn.image_attn.attn")) + + specs = runner.get_kv_cache_spec() + + assert set(specs) == {"layers.0.self_attn.image_attn.attn"} + + +def test_runner_rejects_duplicate_attention_prefixes() -> None: + runner = _runner(_attention(enabled=True, prefix="shared")) + second_attention = _attention(enabled=True, prefix="shared") + runner.pipeline.second_image_attention = second_attention + + with pytest.raises(RuntimeError, match="Duplicate canonical paged Diffusion Attention prefix 'shared'"): + runner.get_kv_cache_spec() + + +def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None: + runner = _runner(_attention(enabled=False)) + + with pytest.raises(RuntimeError, match="no cache-enabled Attention"): + runner.get_kv_cache_spec() + + +def test_runner_retains_matching_rank_local_config() -> None: + runner = _runner(_attention(enabled=True)) + spec = runner.get_kv_cache_spec()["image_attention"] + config = KVCacheConfig( + num_blocks=8, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)], + ) + + runner.set_kv_cache_config(config) + + assert runner.kv_cache_config is config + + +def test_runner_rejects_rank_local_config_for_different_layers() -> None: + runner = _runner(_attention(enabled=True)) + spec = runner.get_kv_cache_spec()["image_attention"] + config = KVCacheConfig( + num_blocks=8, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)], + ) + + with pytest.raises(ValueError, match="layer mismatch"): + runner.set_kv_cache_config(config) + + +def test_worker_selects_its_rank_local_config() -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 1 + worker.od_config = SimpleNamespace(num_gpus=2) + worker.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=None), + cache_config=SimpleNamespace(num_gpu_blocks=None), + ) + worker._maybe_get_memory_pool_context = lambda _tag: nullcontext() + worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config)) + configs = [SimpleNamespace(num_blocks=4), SimpleNamespace(num_blocks=8)] + + worker.set_kv_cache_configs(configs, 64) + + assert worker.installed is configs[1] + assert worker.vllm_config.model_config.max_model_len == 64 + assert worker.vllm_config.cache_config.num_gpu_blocks == 8 + with pytest.raises(ValueError, match="rank count mismatch"): + worker.set_kv_cache_configs(configs[:1], 64) + + +def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096)) + worker.init_snapshot = SimpleNamespace(free_memory=16384) + worker.requested_memory = 8192 + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_request = object() + + assert worker.determine_available_kv_memory(profile_request) == [4096] + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.init_snapshot = SimpleNamespace(free_memory=900) + worker.requested_memory = 750 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_memory_bytes=0, + gpu_memory_utilization=0.75, + ) + ) + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_result = SimpleNamespace( + non_kv_cache_memory=400, + after_profile=SimpleNamespace(free_memory=500), + ) + monkeypatch.setattr( + diffusion_worker_module, + "memory_profiling", + lambda snapshot, weights_memory: nullcontext(profile_result), + ) + profile_request = object() + + assert worker.determine_available_kv_memory(profile_request) == [350] + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +@pytest.mark.parametrize("non_kv_cache_memory", [750, 800]) +def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.init_snapshot = SimpleNamespace(free_memory=900) + worker.requested_memory = 750 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_memory_bytes=None, + gpu_memory_utilization=0.75, + ) + ) + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_result = SimpleNamespace( + non_kv_cache_memory=non_kv_cache_memory, + after_profile=SimpleNamespace(free_memory=100), + ) + monkeypatch.setattr( + diffusion_worker_module, + "memory_profiling", + lambda snapshot, weights_memory: nullcontext(profile_result), + ) + profile_request = object() + + with pytest.raises( + RuntimeError, + match=( + "rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: " + f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes" + ), + ): + worker.determine_available_kv_memory(profile_request) + + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.rank = 0 + worker.init_snapshot = SimpleNamespace(free_memory=1200) + worker.requested_memory = 1000 + worker.vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + kv_cache_memory_bytes=None, + gpu_memory_utilization=0.8, + ) + ) + worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250) + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + profile_result = SimpleNamespace( + non_kv_cache_memory=650, + after_profile=SimpleNamespace(free_memory=550), + ) + profiling_calls = [] + + def profile(snapshot, *, weights_memory): + profiling_calls.append((snapshot, weights_memory)) + return nullcontext(profile_result) + + monkeypatch.setattr( + diffusion_worker_module, + "memory_profiling", + profile, + ) + profile_request = object() + + assert worker.determine_available_kv_memory(profile_request) == [350] + assert profiling_calls == [(worker.init_snapshot, 250)] + worker.model_runner.profile_run.assert_called_once_with(profile_request) + + +def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None: + worker = object.__new__(DiffusionWorker) + worker.init_snapshot = None + worker.requested_memory = None + worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None)) + worker.model_runner = SimpleNamespace() + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value]) + + with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"): + worker.determine_available_kv_memory(object()) + + +def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None: + gathered = [] + + def gather(local_result): + gathered.append(local_result) + return [local_result, (True, "peer-result")] + + def fail_local_probe(): + raise ValueError("local failure") + + monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather) + + with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"): + diffusion_worker_module._run_and_gather_rank_values( + "test probe", + fail_local_probe, + ) + + assert gathered == [(False, "ValueError: local failure")] diff --git a/tests/diffusion/diffusion_kv/test_worker_contract.py b/tests/diffusion/diffusion_kv/test_worker_contract.py index 0c78abbaf..5f2caab06 100644 --- a/tests/diffusion/diffusion_kv/test_worker_contract.py +++ b/tests/diffusion/diffusion_kv/test_worker_contract.py @@ -42,7 +42,10 @@ def make_metadata(request_id: str = "req-0") -> DiffusionKVMetadata: def make_runner(mode: DiffusionKVCacheMode) -> DiffusionModelRunner: runner = object.__new__(DiffusionModelRunner) - runner.od_config = SimpleNamespace(diffusion_kv_mode=mode) + runner.od_config = SimpleNamespace( + diffusion_kv_mode=mode, + parallel_config=SimpleNamespace(cfg_parallel_size=1), + ) return runner @@ -85,19 +88,60 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None: assert new_req.diffusion_kv_metadata is metadata -def test_paged_request_without_metadata_reaches_request_forward() -> None: +def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> None: runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) - expected_output = DiffusionOutput(output=None) - runner._execute_request_list = Mock( - return_value=SimpleNamespace( - runner_outputs=[SimpleNamespace(result=expected_output)], - ) + metadata = DiffusionKVMetadata( + request_id="req-0", + allocation_generation=1, + sequences=( + DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)), + DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)), + ), ) + + attn_metadata = runner._build_paged_attention_metadata([metadata]) + + assert [ + (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len) + for row in attn_metadata.prefill_rows + ] == [("req-0", 0, 0, 8, 8), ("req-0", 1, 0, 9, 9)] + assert [ + (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len) + for row in attn_metadata.denoise_rows + ] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)] + + +def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None: + runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) + runner.od_config.parallel_config.cfg_parallel_size = 2 + monkeypatch.setattr( + "vllm_omni.diffusion.worker.diffusion_model_runner.get_classifier_free_guidance_rank", + lambda: 1, + ) + metadata = DiffusionKVMetadata( + request_id="req-0", + allocation_generation=1, + sequences=( + DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)), + DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)), + ), + ) + + attn_metadata = runner._build_paged_attention_metadata([metadata]) + + assert [row.sequence_id for row in attn_metadata.prefill_rows] == [1] + assert [row.sequence_id for row in attn_metadata.denoise_rows] == [1] + + +def test_paged_request_without_metadata_fails_before_request_forward() -> None: + runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) + runner._execute_request_list = Mock() req = SimpleNamespace(request_id="req-0") - assert runner.execute_model(req) is expected_output + with pytest.raises(ValueError, match="requires Diffusion KV metadata"): + runner.execute_model(req) - runner._execute_request_list.assert_called_once() + runner._execute_request_list.assert_not_called() def test_dense_request_rejects_metadata_before_request_forward() -> None: @@ -141,18 +185,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None: assert calls == [] -def test_batch_path_allows_missing_metadata_before_forward() -> None: +def test_batch_path_rejects_missing_metadata_before_forward() -> None: runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) - expected_output = object() - runner._execute_request_list = Mock(return_value=expected_output) + runner._execute_request_list = Mock() new_req = NewRequestData( request_id="req-0", req=SimpleNamespace(request_id="req-0"), ) - assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output + with pytest.raises(ValueError, match="requires Diffusion KV metadata"): + runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) - runner._execute_request_list.assert_called_once() + runner._execute_request_list.assert_not_called() def test_batch_path_rejects_envelope_request_identity_mismatch() -> None: @@ -170,7 +214,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None: runner._execute_request_list.assert_not_called() -def test_step_path_allows_missing_metadata_before_step_support_check() -> None: +def test_step_path_rejects_missing_metadata_before_step_support_check() -> None: runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER) runner.pipeline = object() runner._supports_step_mode = Mock(return_value=False) @@ -179,10 +223,10 @@ def test_step_path_allows_missing_metadata_before_step_support_check() -> None: req=SimpleNamespace(request_id="req-0"), ) - with pytest.raises(ValueError, match="does not support step execution"): + with pytest.raises(ValueError, match="requires Diffusion KV metadata"): runner.execute_stepwise(make_scheduler_output(new_req)) - runner._supports_step_mode.assert_called_once() + runner._supports_step_mode.assert_not_called() def test_step_path_rejects_envelope_request_identity_mismatch() -> None: @@ -288,3 +332,42 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None: worker.execute_model(req, SimpleNamespace()) assert calls == [(req, None)] + + +def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None: + calls: list[tuple] = [] + + class ModelRunner: + def execute_model(self, req, **kwargs): + calls.append((req, kwargs)) + return DiffusionOutput(output=None) + + worker = object.__new__(DiffusionWorker) + worker.model_runner = ModelRunner() + worker.lora_manager = None + worker._get_profiler = lambda: None + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank", + lambda: 1, + ) + req_0 = SimpleNamespace(request_id="req-0") + req_1 = SimpleNamespace(request_id="req-1") + metadata_0 = make_metadata("req-0") + metadata_1 = make_metadata("req-1") + envelopes = [ + NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0), + NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1), + ] + + result = worker.execute_model(envelopes, SimpleNamespace()) + + assert calls == [ + ( + req_1, + { + "kv_prefetch_job": None, + "diffusion_kv_metadata": metadata_1, + }, + ) + ] + assert result["dp_rank"] == 1 diff --git a/tests/diffusion/distributed/test_a2a_permute.py b/tests/diffusion/distributed/test_a2a_permute.py new file mode 100644 index 000000000..a74f98378 --- /dev/null +++ b/tests/diffusion/distributed/test_a2a_permute.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.distributed import a2a_permute + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatch) -> None: + cu13_include = tmp_path / "nvidia" / "cu13" / "include" + nccl_include = tmp_path / "nvidia" / "nccl" / "include" + nccl_lib = tmp_path / "nvidia" / "nccl" / "lib" / "libnccl.so.2" + cu13_include.mkdir(parents=True) + nccl_include.mkdir(parents=True) + nccl_lib.parent.mkdir(parents=True) + (cu13_include / "cusparse.h").touch() + (nccl_include / "nccl.h").touch() + nccl_lib.touch() + + load_kwargs = {} + monkeypatch.setattr(a2a_permute.sysconfig, "get_paths", lambda: {"purelib": str(tmp_path)}) + monkeypatch.setattr( + torch.utils.cpp_extension, + "load", + lambda **kwargs: load_kwargs.update(kwargs), + ) + monkeypatch.setattr(a2a_permute.symm_mem, "set_backend", lambda _backend: None) + monkeypatch.setattr(a2a_permute, "_BUILT", False) + + a2a_permute.ensure_a2a_permute_available() + + assert set(load_kwargs["extra_include_paths"]) == {str(cu13_include), str(nccl_include)} + assert load_kwargs["extra_ldflags"] == [str(nccl_lib)] + + +@dataclass +class _FakeAllocation: + """Byte workspace stand-in that slices and reinterprets like ``symm_mem.empty()``.""" + + size: int + device: torch.device + + def __post_init__(self) -> None: + self.buffer = torch.empty(self.size, dtype=torch.uint8) + + def __getitem__(self, item: slice) -> torch.Tensor: + return self.buffer[item] + + +def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None: + allocations: list[_FakeAllocation] = [] + all_reduce_calls: list[object] = [] + synchronize_calls: list[torch.device] = [] + stream_id = [1] + a2a_permute._SYMM_WORKSPACES.clear() + + monkeypatch.setattr(a2a_permute.torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + a2a_permute.torch.cuda, + "current_stream", + lambda _device: SimpleNamespace(cuda_stream=stream_id[0]), + ) + monkeypatch.setattr(a2a_permute, "_resolve_process_group", lambda _name: object()) + monkeypatch.setattr(a2a_permute.torch, "ones", lambda *_args, **_kwargs: object()) + monkeypatch.setattr( + a2a_permute.dist, + "all_reduce", + lambda _tensor, group: all_reduce_calls.append(group), + ) + monkeypatch.setattr( + a2a_permute.torch.accelerator, + "synchronize", + lambda device: synchronize_calls.append(device), + ) + + def empty(size, *, dtype, device): + assert dtype == torch.uint8 + allocation = _FakeAllocation(size=size, device=torch.device(device)) + allocations.append(allocation) + return allocation + + monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty) + # The handle is only retained to keep the rendezvous registration alive. + monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda _allocation, _group: object()) + + device = torch.device("cuda:0") + first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group") + smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group") + larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group") + + assert (first.shape, first.dtype) == ((2, 3), torch.float16) + assert (smaller.shape, smaller.dtype) == ((1, 4), torch.float16) + assert (larger.shape, larger.dtype) == ((4, 3), torch.float16) + # A sub-capacity request reuses the workspace; only growth allocates again. + assert smaller.untyped_storage().data_ptr() == first.untyped_storage().data_ptr() + assert larger.untyped_storage().data_ptr() != first.untyped_storage().data_ptr() + assert [allocation.size for allocation in allocations] == [12, 24] + assert len(a2a_permute._SYMM_WORKSPACES) == 1 + assert len(all_reduce_calls) == 1 + assert synchronize_calls == [device, device] + + stream_id[0] = 2 + with pytest.raises(RuntimeError, match="single CUDA stream"): + a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group") + + a2a_permute.clear_a2a_permute_workspaces() + assert not a2a_permute._SYMM_WORKSPACES + assert synchronize_calls == [device, device, device] diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py index b2ee7c10d..dcaa8e01a 100644 --- a/tests/diffusion/distributed/test_distributed_vae_executor.py +++ b/tests/diffusion/distributed/test_distributed_vae_executor.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from types import SimpleNamespace import pytest @@ -58,6 +59,11 @@ class E2EOperator: return torch.cat(tiles, dim=0) +@dataclass +class FakeWorldGroup: + device_group: object + + class DummyMixin(DistributedVaeMixin): def __init__(self): self.use_tiling = True @@ -75,11 +81,13 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch): @pytest.fixture(autouse=True) -def mock_dit_group(monkeypatch: pytest.MonkeyPatch): +def mock_world_group(monkeypatch: pytest.MonkeyPatch): + group = object() monkeypatch.setattr( - "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group", - lambda: None, + "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_world_group", + lambda: FakeWorldGroup(device_group=group), ) + return group @pytest.fixture(autouse=True) @@ -93,6 +101,11 @@ def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch): # ============================ +def test_uses_dedicated_world_device_group(mock_world_group): + executor = DistributedVaeExecutor() + assert executor.group is mock_world_group + + def test_balance_tasks(): executor = DistributedVaeExecutor() workloads = [2, 5, 13, 8, 2, 4] # each is 17 diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py index 74f30f156..bfbbca588 100644 --- a/tests/diffusion/distributed/test_expert_parallel_layout.py +++ b/tests/diffusion/distributed/test_expert_parallel_layout.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace @@ -92,9 +92,8 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch): monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group) - monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) - for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"): monkeypatch.setattr(parallel_state, name, None) for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) @@ -195,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch): assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group] +@pytest.mark.cpu +@pytest.mark.core_model +@pytest.mark.parametrize( + ("fully_shard_degree", "error_message"), + [(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")], +) +def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str): + """Reject invalid HSDP shard sizes before allocating any process groups.""" + world_size = 4 + created_groups: list[object] = [] + + def fail_if_group_created(*args, **kwargs): + del args, kwargs + created_groups.append(object()) + raise AssertionError("a process group was created before HSDP validation") + + fake_world_group = SimpleNamespace( + rank_in_group=0, + local_rank=0, + device_group=object(), + ) + monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"): + monkeypatch.setattr(parallel_state, name, None) + for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): + monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) + + with pytest.raises(ValueError, match=error_message): + parallel_state.initialize_model_parallel( + fully_shard_degree=fully_shard_degree, + use_hsdp=True, + backend="gloo", + ) + + assert created_groups == [] + + @pytest.mark.cpu @pytest.mark.core_model def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): @@ -229,9 +269,8 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) - monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) - for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"): monkeypatch.setattr(parallel_state, name, None) for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) @@ -259,6 +298,33 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): ] +@pytest.mark.cpu +@pytest.mark.core_model +def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch): + """A destroyed diffusion PP group must not block the next initialization.""" + + class _DestroyableGroup: + def __init__(self) -> None: + self.destroy_calls = 0 + + def destroy(self) -> None: + self.destroy_calls += 1 + + pipeline_group = _DestroyableGroup() + for name in ("_DP", "_CFG", "_SP"): + monkeypatch.setattr(parallel_state, name, None) + monkeypatch.setattr(parallel_state, "_PP", pipeline_group) + monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None) + for name in ("_DP", "_PCP", "_TP", "_EP"): + monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) + monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False) + + parallel_state.destroy_model_parallel() + + assert pipeline_group.destroy_calls == 1 + assert parallel_state.vllm_parallel_state._PP is None + + @pytest.mark.cpu @pytest.mark.core_model def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch): @@ -295,9 +361,8 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch): monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) - monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) - for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"): monkeypatch.setattr(parallel_state, name, None) for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) diff --git a/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py b/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py index e861bd61a..8de3a42f8 100644 --- a/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py +++ b/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py @@ -53,14 +53,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: def _extract_single_image(outputs) -> Image.Image: first_output = outputs[0] assert first_output.final_output_type == "image" - if not hasattr(first_output, "request_output") or not first_output.request_output: - raise ValueError("No request_output found in OmniRequestOutput") + if not isinstance(first_output, OmniRequestOutput): + raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}") - req_out = first_output.request_output - if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): - raise ValueError("Invalid request_output structure or missing 'images' key") - - images = req_out.images + images = first_output.images if images is None or len(images) != 1: raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}") return images[0] diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 91c6784fc..25762c356 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -1,10 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.""" import gc -import os -import socket import pytest import torch @@ -12,6 +10,7 @@ import torch.distributed as dist import torch.nn as nn from torch.distributed.tensor import DeviceMesh, DTensor +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.distributed.hsdp import ( HSDPInferenceConfig, @@ -37,35 +36,17 @@ class _PackedModel(nn.Module): self.root_weight = nn.Parameter(torch.ones(2)) -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind(("127.0.0.1", 0)) - return int(sock.getsockname()[1]) - - @pytest.fixture(scope="module") def cpu_process_group(): if dist.is_initialized(): yield return - master_port = _find_free_port() - os.environ.update( - { - "RANK": "0", - "LOCAL_RANK": "0", - "WORLD_SIZE": "1", - "MASTER_ADDR": "127.0.0.1", - "MASTER_PORT": str(master_port), - } - ) - dist.init_process_group("gloo", rank=0, world_size=1) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) try: yield finally: dist.destroy_process_group() - for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"): - os.environ.pop(key, None) gc.collect() @@ -234,7 +215,7 @@ class TestDiffusionParallelConfigHSDP: def test_hsdp_cannot_use_with_tp(self): """Test that HSDP and Tensor Parallelism cannot be used together.""" - with pytest.raises(ValueError, match="cannot be used with TP or DP"): + with pytest.raises(ValueError, match="not compatible with TP"): DiffusionParallelConfig( tensor_parallel_size=2, use_hsdp=True, @@ -243,7 +224,7 @@ class TestDiffusionParallelConfigHSDP: def test_hsdp_cannot_use_with_dp(self): """Test that HSDP and Data Parallelism cannot be used together.""" - with pytest.raises(ValueError, match="cannot be used with TP or DP"): + with pytest.raises(ValueError, match="not compatible with DP"): DiffusionParallelConfig( data_parallel_size=2, use_hsdp=True, @@ -264,158 +245,6 @@ class TestDiffusionParallelConfigHSDP: assert config.hsdp_shard_size == 2 # auto: 4 // 2 -class TestStandaloneHSDPDetection: - """Tests for standalone HSDP detection and dit_parallel_size calculation. - - These tests verify the logic used in initialize_model_parallel() to detect - standalone HSDP mode and calculate effective parallel sizes. - - Standalone HSDP is when all non-HSDP parallelism dimensions are 1. - """ - - @staticmethod - def compute_standalone_hsdp_params( - data_parallel_size: int = 1, - cfg_parallel_size: int = 1, - sequence_parallel_size: int = 1, - pipeline_parallel_size: int = 1, - tensor_parallel_size: int = 1, - fully_shard_degree: int = 1, - hsdp_replicate_size: int = 1, - ) -> dict: - """Compute standalone HSDP detection parameters. - - This mirrors the logic in initialize_model_parallel(). - """ - dit_parallel_size = ( - data_parallel_size - * cfg_parallel_size - * sequence_parallel_size - * pipeline_parallel_size - * tensor_parallel_size - ) - - # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1 - is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1 - - # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) - if is_standalone_hsdp: - effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size - else: - effective_dit_parallel_size = dit_parallel_size - - effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size - - return { - "original_dit_parallel_size": dit_parallel_size, - "is_standalone_hsdp": is_standalone_hsdp, - "effective_dit_parallel_size": effective_dit_parallel_size, - "effective_dp_size": effective_dp_size, - } - - def test_standalone_hsdp_basic(self): - """Test basic standalone HSDP detection (shard_size=4, replicate=1).""" - result = self.compute_standalone_hsdp_params( - fully_shard_degree=4, - hsdp_replicate_size=1, - ) - assert result["original_dit_parallel_size"] == 1 - assert result["is_standalone_hsdp"] is True - assert result["effective_dit_parallel_size"] == 4 - assert result["effective_dp_size"] == 4 - - def test_standalone_hsdp_with_replicate(self): - """Test standalone HSDP with replication (shard_size=4, replicate=2).""" - result = self.compute_standalone_hsdp_params( - fully_shard_degree=4, - hsdp_replicate_size=2, - ) - assert result["original_dit_parallel_size"] == 1 - assert result["is_standalone_hsdp"] is True - assert result["effective_dit_parallel_size"] == 8 # 4 * 2 - assert result["effective_dp_size"] == 8 - - def test_combined_hsdp_sp_not_standalone(self): - """Test HSDP combined with SP is NOT detected as standalone. - - This is a regression test for the bug where the condition - `dit_parallel_size == fully_shard_degree` incorrectly matched - combined modes like SP=4 + HSDP=4. - """ - result = self.compute_standalone_hsdp_params( - sequence_parallel_size=4, - fully_shard_degree=4, - hsdp_replicate_size=1, - ) - assert result["original_dit_parallel_size"] == 4 - assert result["is_standalone_hsdp"] is False - # Should NOT override dp_size for combined mode - assert result["effective_dp_size"] == 1 # original data_parallel_size - - def test_combined_hsdp_cfg_not_standalone(self): - """Test HSDP combined with CFG is NOT detected as standalone.""" - result = self.compute_standalone_hsdp_params( - cfg_parallel_size=2, - fully_shard_degree=4, - hsdp_replicate_size=1, - ) - assert result["original_dit_parallel_size"] == 2 - assert result["is_standalone_hsdp"] is False - assert result["effective_dp_size"] == 1 - - def test_combined_hsdp_pp_not_standalone(self): - """Test HSDP combined with PP is NOT detected as standalone.""" - result = self.compute_standalone_hsdp_params( - pipeline_parallel_size=2, - fully_shard_degree=4, - hsdp_replicate_size=1, - ) - assert result["original_dit_parallel_size"] == 2 - assert result["is_standalone_hsdp"] is False - assert result["effective_dp_size"] == 1 - - def test_no_hsdp_not_standalone(self): - """Test that no HSDP (fully_shard_degree=1) is NOT standalone.""" - result = self.compute_standalone_hsdp_params( - fully_shard_degree=1, - ) - assert result["original_dit_parallel_size"] == 1 - assert result["is_standalone_hsdp"] is False - assert result["effective_dp_size"] == 1 - - def test_combined_multiple_parallelism_not_standalone(self): - """Test HSDP combined with multiple parallelism is NOT standalone.""" - result = self.compute_standalone_hsdp_params( - sequence_parallel_size=2, - cfg_parallel_size=2, - fully_shard_degree=4, - hsdp_replicate_size=1, - ) - assert result["original_dit_parallel_size"] == 4 # 2 * 2 - assert result["is_standalone_hsdp"] is False - assert result["effective_dp_size"] == 1 - - def test_standalone_hsdp_large_shard(self): - """Test standalone HSDP with large shard size.""" - result = self.compute_standalone_hsdp_params( - fully_shard_degree=8, - hsdp_replicate_size=1, - ) - assert result["is_standalone_hsdp"] is True - assert result["effective_dit_parallel_size"] == 8 - assert result["effective_dp_size"] == 8 - - def test_standalone_hsdp_large_replicate(self): - """Test standalone HSDP with large replicate size.""" - result = self.compute_standalone_hsdp_params( - fully_shard_degree=4, - hsdp_replicate_size=4, - ) - assert result["is_standalone_hsdp"] is True - assert result["effective_dit_parallel_size"] == 16 # 4 * 4 - assert result["effective_dp_size"] == 16 - - class TestHSDPShardConditions: """Tests for _hsdp_shard_conditions matching logic.""" diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index 0b48ed266..cccf116c3 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/tests/diffusion/distributed/test_pipeline_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for PipelineParallelMixin and related PP communication helpers. CPU unit tests exercise mixin logic without a process group. Distributed CPU @@ -9,8 +9,6 @@ real multi-GPU NCCL collectives. from __future__ import annotations -import os -import socket from typing import Literal import pytest @@ -21,6 +19,7 @@ from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, @@ -40,31 +39,18 @@ DeviceKind = Literal["cpu", "cuda"] _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] -def _find_free_port() -> str: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return str(s.getsockname()[1]) - - -def update_environment_variables(envs_dict: dict[str, str]) -> None: - for k, v in envs_dict.items(): - os.environ[k] = v - - def _cleanup_distributed() -> None: - """Destroy omni distributed state and clear test-process leftovers. + """Destroy omni distributed state left dangling after a test. ``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP`` aliased to the destroyed group. Parent-process baselines must clear that - dangling reference (and RANK/MASTER_* env) so later unit tests that call + dangling reference so later unit tests that call ``vllm.distributed.initialize_model_parallel`` do not see stale state. """ destroy_distributed_env() import vllm.distributed.parallel_state as vllm_parallel_state vllm_parallel_state._PP = None - for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"): - os.environ.pop(key, None) # --------------------------------------------------------------------------- @@ -416,24 +402,20 @@ def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device: def init_dist( local_rank: int, world_size: int, - master_port: str, + init_method: str, device_kind: DeviceKind, ) -> torch.device: """Initialise the distributed environment for a spawned worker.""" device = _worker_device(local_rank, device_kind) if device_kind == "cuda": current_omni_platform.set_device(device) - update_environment_variables( - { - "RANK": str(local_rank), - "LOCAL_RANK": str(local_rank), - "WORLD_SIZE": str(world_size), - "MASTER_ADDR": "localhost", - "MASTER_PORT": master_port, - } - ) backend = "gloo" if device_kind == "cpu" else None - init_distributed_environment(backend=backend) + init_distributed_environment( + world_size=world_size, + rank=local_rank, + distributed_init_method=init_method, + backend=backend, + ) return device @@ -493,11 +475,11 @@ def make_pipeline_and_inputs( def isend_irecv_worker( local_rank: int, world_size: int, - master_port: str, + init_method: str, device_kind: DeviceKind, result_queue, ): - device = init_dist(local_rank, world_size, master_port, device_kind) + device = init_dist(local_rank, world_size, init_method, device_kind) initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind)) pp_group = get_pp_group() @@ -517,13 +499,13 @@ def isend_irecv_worker( _cleanup_distributed() -def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None: +def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None: mp_context = torch.multiprocessing.get_context("spawn") manager = mp_context.Manager() q = manager.Queue() torch.multiprocessing.spawn( isend_irecv_worker, - args=(pp_size, master_port, device_kind, q), + args=(pp_size, init_method, device_kind, q), nprocs=pp_size, ) results = {label: tensor for label, tensor in [q.get(), q.get()]} @@ -538,7 +520,7 @@ def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> @pytest.mark.parametrize("pp_size", [2]) def test_isend_irecv_tensor_dict(pp_size: int): """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss.""" - _run_isend_irecv(pp_size, device_kind="cpu", master_port=_find_free_port()) + _run_isend_irecv(pp_size, device_kind="cpu", init_method=get_distributed_init_method()) @pytest.mark.full_model @@ -547,7 +529,7 @@ def test_isend_irecv_tensor_dict(pp_size: int): @pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)]) def test_isend_irecv_tensor_dict_parity(pp_size: int): """Nightly: isend/irecv on real multi-GPU NCCL.""" - _run_isend_irecv(pp_size, device_kind="cuda", master_port=_find_free_port()) + _run_isend_irecv(pp_size, device_kind="cuda", init_method=get_distributed_init_method()) # --------------------------------------------------------------------------- @@ -578,7 +560,7 @@ def compute_single_gpu_baseline( if key in _baseline_cache: return _baseline_cache[key] - device = init_dist(0, 1, _find_free_port(), device_kind) + device = init_dist(0, 1, get_distributed_init_method(), device_kind) try: initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) @@ -604,7 +586,7 @@ def compute_single_gpu_baseline( def predict_noise_worker( local_rank: int, world_size: int, - master_port: str, + init_method: str, pp_size: int, cfg_size: int, do_true_cfg: bool, @@ -614,7 +596,7 @@ def predict_noise_worker( result_queue, ): """Generic predict-noise worker parameterized by PP and CFG topology.""" - device = init_dist(local_rank, world_size, master_port, device_kind) + device = init_dist(local_rank, world_size, init_method, device_kind) initialize_model_parallel( pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size, @@ -661,7 +643,7 @@ def _run_predict_noise( rtol: float, atol: float, device_kind: DeviceKind, - master_port: str, + init_method: str, ) -> None: test_config = { "num_layers": num_layers, @@ -681,7 +663,7 @@ def _run_predict_noise( world_size = pp_size * cfg_size torch.multiprocessing.spawn( predict_noise_worker, - args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q), + args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q), nprocs=world_size, ) @@ -786,7 +768,7 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, r rtol=rtol, atol=atol, device_kind="cpu", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) @@ -809,7 +791,7 @@ def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers, rtol=rtol, atol=atol, device_kind="cuda", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) @@ -824,7 +806,7 @@ def compute_scheduler_step_baseline( device_kind: DeviceKind, ) -> torch.Tensor: """Single-process reference: predict_noise + scheduler_step.""" - device = init_dist(0, 1, _find_free_port(), device_kind) + device = init_dist(0, 1, get_distributed_init_method(), device_kind) try: initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind)) @@ -853,7 +835,7 @@ def compute_scheduler_step_baseline( def scheduler_step_worker( local_rank: int, world_size: int, - master_port: str, + init_method: str, pp_size: int, cfg_size: int, do_true_cfg: bool, @@ -861,7 +843,7 @@ def scheduler_step_worker( device_kind: DeviceKind, result_queue, ): - device = init_dist(local_rank, world_size, master_port, device_kind) + device = init_dist(local_rank, world_size, init_method, device_kind) initialize_model_parallel( pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size, @@ -905,7 +887,7 @@ def _run_scheduler_step( do_true_cfg: bool, input_seed: int, device_kind: DeviceKind, - master_port: str, + init_method: str, ) -> None: test_config = { "num_layers": 4, @@ -925,7 +907,7 @@ def _run_scheduler_step( world_size = pp_size * cfg_size torch.multiprocessing.spawn( scheduler_step_worker, - args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q), + args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q), nprocs=world_size, ) @@ -957,7 +939,7 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed): do_true_cfg=do_true_cfg, input_seed=input_seed, device_kind="cpu", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) @@ -979,5 +961,5 @@ def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed): do_true_cfg=do_true_cfg, input_seed=input_seed, device_kind="cuda", - master_port=_find_free_port(), + init_method=get_distributed_init_method(), ) diff --git a/tests/diffusion/distributed/test_sequence_parallel.py b/tests/diffusion/distributed/test_sequence_parallel.py index d43d18c49..f083e1d26 100644 --- a/tests/diffusion/distributed/test_sequence_parallel.py +++ b/tests/diffusion/distributed/test_sequence_parallel.py @@ -416,8 +416,8 @@ def test_sp_correctness_advanced(model_name: str): @pytest.mark.skipif( - not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), - reason="Only tested on CUDA and XPU", + not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA, ROCM and XPU", ) @pytest.mark.diffusion @pytest.mark.parallel @@ -447,8 +447,8 @@ def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.skipif( - not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), - reason="Only tested on CUDA and XPU", + not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA, ROCM and XPU", ) @pytest.mark.diffusion @pytest.mark.parallel @@ -468,8 +468,8 @@ def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> No @pytest.mark.skipif( - not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), - reason="Only tested on CUDA and XPU", + not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA, ROCM and XPU", ) @pytest.mark.diffusion @pytest.mark.parallel diff --git a/tests/diffusion/distributed/test_sp_plan_hooks.py b/tests/diffusion/distributed/test_sp_plan_hooks.py index f95dbfe2b..4ef564571 100644 --- a/tests/diffusion/distributed/test_sp_plan_hooks.py +++ b/tests/diffusion/distributed/test_sp_plan_hooks.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the Sequence Parallelism (SP) framework. These tests verify the SP plan mechanism and hooks work correctly without @@ -17,6 +17,7 @@ import pytest import torch import torch.nn as nn +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, @@ -419,7 +420,6 @@ class TestHookRegistration: # ============================================================================= -@pytest.mark.L4 class TestModelSpPlans: """Test that model _sp_plan definitions are valid. @@ -427,6 +427,8 @@ class TestModelSpPlans: May require GPU for model imports. """ + pytestmark = hardware_marks(res={"cuda": "L4"}) + def test_zimage_transformer_sp_plan(self): """Test ZImageTransformer2DModel _sp_plan structure. diff --git a/tests/diffusion/distributed/test_tensor_parallel.py b/tests/diffusion/distributed/test_tensor_parallel.py index c50c4336d..ae97c9900 100644 --- a/tests/diffusion/distributed/test_tensor_parallel.py +++ b/tests/diffusion/distributed/test_tensor_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """DiT tensor-parallel parity and perf checks for diffusion models. @@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model. from __future__ import annotations +import gc import os import time from pathlib import Path @@ -18,8 +19,8 @@ import pytest import torch from PIL import Image -from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test +from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -75,8 +76,13 @@ def _run_generate( if num_requests < 2: raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)") + gc.collect() current_omni_platform.empty_cache() device_index = current_omni_platform.current_device() + with current_omni_platform.device(device_index): + free_bytes, total_bytes = current_omni_platform.mem_get_info() + initial_used_mb = (total_bytes - free_bytes) / (1024**2) + monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() try: @@ -113,7 +119,10 @@ def _run_generate( pass median_time_s = float(np.median(per_request_times_s)) - peak_memory_mb = monitor.peak_used_mb + # DeviceMemoryMonitor samples absolute device usage. Normalize by + # this run's starting usage so compiler and allocator caches retained + # by the preceding TP=1 run are not charged to the TP=2 run. + peak_memory_mb = max(0.0, monitor.peak_used_mb - initial_used_mb) return _extract_single_image([last_output]), median_time_s, peak_memory_mb finally: monitor.stop() diff --git a/tests/diffusion/distributed/test_ulysses_uaa_perf.py b/tests/diffusion/distributed/test_ulysses_uaa_perf.py index 2a16a9ae5..707800cc0 100644 --- a/tests/diffusion/distributed/test_ulysses_uaa_perf.py +++ b/tests/diffusion/distributed/test_ulysses_uaa_perf.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Perf smoke tests for Ulysses advanced_uaa communication overhead. This test is intended for CI monitoring only: @@ -9,8 +9,6 @@ This test is intended for CI monitoring only: from __future__ import annotations -import os -import socket from dataclasses import dataclass import pytest @@ -18,6 +16,7 @@ import torch import torch.distributed as dist from tests.helpers.mark import hardware_test +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.attention.parallel.ulysses import ( _all_gather_int, _ulysses_all_to_all_any_o, @@ -33,20 +32,6 @@ from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.platforms import current_omni_platform -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - def _max_all_reduce(pg: dist.ProcessGroup, value: float, *, device: torch.device) -> float: t = torch.tensor([value], device=device, dtype=torch.float32) dist.all_reduce(t, op=dist.ReduceOp.MAX, group=pg) @@ -77,22 +62,20 @@ def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None: if available_gpus < case.world_size: pytest.skip(f"Requires {case.world_size} GPUs, got {available_gpus}") - master_port = _find_free_port() + init_method = get_distributed_init_method() torch.multiprocessing.spawn( _perf_worker, - args=(case.world_size, master_port, case.ulysses_degree, case.ring_degree), + args=(case.world_size, init_method, case.ulysses_degree, case.ring_degree), nprocs=case.world_size, ) -def _perf_worker(local_rank: int, world_size: int, master_port: int, ulysses_degree: int, ring_degree: int) -> None: +def _perf_worker(local_rank: int, world_size: int, init_method: str, ulysses_degree: int, ring_degree: int) -> None: device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) - _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port) - try: - init_distributed_environment(world_size=world_size, rank=local_rank) + init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method) initialize_model_parallel( data_parallel_size=1, cfg_parallel_size=1, diff --git a/tests/diffusion/distributed/test_vae_decode_parallelism.py b/tests/diffusion/distributed/test_vae_decode_parallelism.py index c3996fc23..28ded4639 100644 --- a/tests/diffusion/distributed/test_vae_decode_parallelism.py +++ b/tests/diffusion/distributed/test_vae_decode_parallelism.py @@ -72,8 +72,7 @@ def _to_float_array(data: Any) -> np.ndarray: def _extract_output_array(output: Any) -> np.ndarray: - req_out = output.request_output - payload = req_out.images[0] + payload = output.images[0] return _to_float_array(payload) diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py index babbc7abc..49622aecd 100644 --- a/tests/diffusion/kernels/mot/test_mot_linear.py +++ b/tests/diffusion/kernels/mot/test_mot_linear.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # ruff: noqa: N803, E741 """Layer-level correctness & performance test for MoT parallel linear layers. @@ -237,6 +240,53 @@ def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear): return output +def _assert_gen_expert_parallel_metadata(layer): + """The secondary expert must satisfy vLLM's parallel-linear contract.""" + for attr in ( + "tp_rank", + "tp_size", + "input_size", + "input_size_per_partition", + "output_size", + "output_size_per_partition", + ): + assert getattr(layer.gen_exp, attr) == getattr(layer, attr) + + +def test_mot_gen_expert_parallel_metadata(): + """Keep secondary experts compatible with vLLM online quantizers.""" + from vllm.model_executor.layers.quantization.online.fp8 import ( + _is_tp_sharded, + ) + + with set_current_vllm_config(VllmConfig()): + layers = ( + MoTQKVParallelLinear( + hidden_size=16, + head_size=4, + total_num_heads=4, + total_num_kv_heads=2, + bias=False, + vae_bias=False, + params_dtype=torch.bfloat16, + disable_tp=True, + ), + MoTRowParallelLinear( + 16, + 8, + bias=False, + vae_bias=False, + input_is_parallel=True, + params_dtype=torch.bfloat16, + disable_tp=True, + ), + ) + + for layer in layers: + _assert_gen_expert_parallel_metadata(layer) + assert _is_tp_sharded(layer.gen_exp) is False + + def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str): """Compare outputs, print metrics, and assert correctness. @@ -335,6 +385,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool disable_tp=True, ).cuda() + _assert_gen_expert_parallel_metadata(mot_linear) + assert text_linear.output_size_per_partition == N, ( f"Expected output_size_per_partition={N}, " f"got {text_linear.output_size_per_partition}. " @@ -423,6 +475,8 @@ def test_mot_o_proj( disable_tp=True, ).cuda() + _assert_gen_expert_parallel_metadata(mot_linear) + _sync_weights(text_linear, vae_linear, mot_linear) text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE) @@ -501,6 +555,13 @@ def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool): ).cuda() with torch.no_grad(): + # vLLM linear layers leave weights uninitialized (torch.empty), so + # a freshly zeroed GPU page would make both outputs zero and the + # cosine-similarity gate degenerate. Seed deterministic values so + # und-mode is compared against a meaningful reference. + ref_linear.weight.normal_(mean=0.0, std=0.02) + if ref_linear.bias is not None: + ref_linear.bias.normal_(mean=0.0, std=0.02) mot_linear.weight.copy_(ref_linear.weight) if bias and ref_linear.bias is not None: mot_linear.bias.copy_(ref_linear.bias) @@ -549,6 +610,13 @@ def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool): ).cuda() with torch.no_grad(): + # vLLM linear layers leave weights uninitialized (torch.empty), so + # a freshly zeroed GPU page would make both outputs zero and the + # cosine-similarity gate degenerate. Seed deterministic values so + # und-mode is compared against a meaningful reference. + ref_linear.weight.normal_(mean=0.0, std=0.02) + if ref_linear.bias is not None: + ref_linear.bias.normal_(mean=0.0, std=0.02) mot_linear.weight.copy_(ref_linear.weight) if bias and ref_linear.bias is not None: mot_linear.bias.copy_(ref_linear.bias) diff --git a/tests/diffusion/layers/test_activation.py b/tests/diffusion/layers/test_activation.py new file mode 100644 index 000000000..c3205502f --- /dev/null +++ b/tests/diffusion/layers/test_activation.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.layers.activation import SiluAndMul + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.cpu +def test_silu_and_mul_native_matches_packed_reference() -> None: + packed = torch.randn(257, 256, dtype=torch.float32) + gate, up = packed.chunk(2, dim=-1) + + torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up) + + +@hardware_test(res={"npu": "A3"}, num_cards=1) +def test_silu_and_mul_npu_matches_packed_reference() -> None: + packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16) + gate, up = packed.chunk(2, dim=-1) + + torch.testing.assert_close( + SiluAndMul()(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope.py b/tests/diffusion/layers/test_fused_qk_norm_rope.py new file mode 100644 index 000000000..40888556a --- /dev/null +++ b/tests/diffusion/layers/test_fused_qk_norm_rope.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] + +_HEAD_DIM = 128 +_ROTARY_DIM = 96 +_EPS = 1e-5 + + +def _reference(q, k, q_weight, k_weight, rope_table): + q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS) + k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS) + half = _ROTARY_DIM // 2 + cos = rope_table[..., :half].unsqueeze(1) + sin = rope_table[..., half:].unsqueeze(1) + + def apply(x): + first = x[..., :half] + second = x[..., half:_ROTARY_DIM] + return torch.cat( + ( + first * cos - second * sin, + second * cos + first * sin, + x[..., _ROTARY_DIM:], + ), + dim=-1, + ) + + return apply(q), apply(k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +@pytest.mark.parametrize("seq_len", [1, 257, 1024]) +def test_fused_qk_norm_rope_matches_bf16_reference(seq_len): + from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + fused_qk_norm_rope, + ) + + torch.manual_seed(17) + heads = 14 + qkv = torch.randn( + seq_len, + heads * _HEAD_DIM * 3, + device="cuda", + dtype=torch.bfloat16, + ) + q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM) + k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view( + seq_len, + heads, + _HEAD_DIM, + ) + q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda") + rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16) + + expected_q, expected_k = _reference( + q, + k, + q_weight, + k_weight, + rope_table, + ) + actual_q, actual_k = fused_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + ) + + torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py b/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py new file mode 100644 index 000000000..35d990afa --- /dev/null +++ b/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +import types + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + +_HEAD_DIM = 128 +_ROTARY_DIM = 96 +_EPS = 1e-5 + + +def _reference( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + half = _ROTARY_DIM // 2 + + def apply(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + x = F.rms_norm(x, (_HEAD_DIM,), weight, _EPS) + cos = rope_table[..., :half].unsqueeze(1) + sin = rope_table[..., half:].unsqueeze(1) + return torch.cat( + ( + x[..., :half] * cos - x[..., half:_ROTARY_DIM] * sin, + x[..., half:_ROTARY_DIM] * cos + x[..., :half] * sin, + x[..., _ROTARY_DIM:], + ), + dim=-1, + ) + + return apply(q, q_weight), apply(k, k_weight) + + +def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + torch.manual_seed(7) + tokens, heads = 11, 4 + q = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16) + k = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16) + q_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16) + k_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16) + freqs = torch.randn(tokens, _ROTARY_DIM // 2) + rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16) + return q, k, q_weight, k_weight, rope_table + + +@pytest.mark.cpu +def test_npu_qk_norm_rope_uses_torch_npu_fused_primitives_without_mindiesd( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused + + calls: dict[str, int] = {"rms_norm": 0, "rotary_mul": 0} + + def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float): + calls["rms_norm"] += 1 + return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),) + + def npu_rotary_mul(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + calls["rotary_mul"] += 1 + assert x.shape[-1] == _ROTARY_DIM + assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM) + assert sin.shape == cos.shape + half = _ROTARY_DIM // 2 + return torch.cat( + ( + x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half], + x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:], + ), + dim=-1, + ) + + monkeypatch.setitem( + sys.modules, + "torch_npu", + types.SimpleNamespace(npu_rms_norm=npu_rms_norm, npu_rotary_mul=npu_rotary_mul), + ) + monkeypatch.setattr(fused, "find_spec", lambda _: None) + q, k, q_weight, k_weight, rope_table = _inputs() + + actual_q, actual_k = fused._npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + _ROTARY_DIM, + ) + expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table) + + assert calls == {"rms_norm": 2, "rotary_mul": 2} + torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0) + torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0) + + +@pytest.mark.cpu +def test_npu_qk_norm_rope_uses_mindiesd_and_normalizes_packed_layout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused + + calls: dict[str, int] = {"rms_norm": 0, "mindiesd_rope": 0} + + def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float): + calls["rms_norm"] += 1 + return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),) + + def rotary_position_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, **kwargs): + calls["mindiesd_rope"] += 1 + assert x.dim() == 4 + assert x.shape[0] == 1 + assert x.shape[-1] == _ROTARY_DIM + assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM) + assert sin.shape == cos.shape + assert kwargs == { + "rotated_mode": "rotated_half", + "head_first": False, + "fused": True, + } + half = _ROTARY_DIM // 2 + return torch.cat( + ( + x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half], + x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:], + ), + dim=-1, + ) + + monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm)) + monkeypatch.setitem( + sys.modules, + "mindiesd", + types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding), + ) + monkeypatch.setattr(fused, "find_spec", lambda _: object()) + q, k, q_weight, k_weight, rope_table = _inputs() + + actual_q, actual_k = fused._npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + _ROTARY_DIM, + ) + expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table) + + assert calls == {"rms_norm": 2, "mindiesd_rope": 2} + torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0) + torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0) + + +@hardware_test(res={"npu": "A3"}, num_cards=1) +def test_fused_qk_norm_rope_npu_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise the public Ascend dispatch with real torch_npu/MindIE operators.""" + from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused + + q, k, q_weight, k_weight, rope_table = (tensor.to("npu") for tensor in _inputs()) + calls = 0 + original_npu_impl = fused._npu_qk_norm_rope + + def capture_npu_impl(*args, **kwargs): + nonlocal calls + calls += 1 + return original_npu_impl(*args, **kwargs) + + monkeypatch.setattr(fused, "_npu_qk_norm_rope", capture_npu_impl) + actual_q, actual_k = fused.fused_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + head_dim=_HEAD_DIM, + rotary_dim=_ROTARY_DIM, + ) + expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table) + + assert calls == 1 + torch.testing.assert_close(actual_q, expected_q, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(actual_k, expected_k, atol=2e-2, rtol=2e-2) diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index 71125684b..4794ca1c2 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/tests/diffusion/layers/test_norm.py @@ -340,6 +340,97 @@ def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.Mo assert captured["epsilon"].item() == pytest.approx(1e-5) +# ── RMSNorm residual tests ── + + +def test_rmsnorm_matches_native_residual_contract(): + """The non-NPU path preserves separate residual-add and RMSNorm math.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + eps = 1e-6 + norm = RMSNorm(4, eps=eps, dtype=torch.bfloat16) + norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16)) + x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.bfloat16) + residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16) + + expected_residual = residual + x + expected_fp32 = expected_residual.float() + expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps) + expected_output = (expected_output * norm.weight.float()).to(x.dtype) + output, updated_residual = norm.forward_native(x, residual) + + torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "forward_name", + ["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"], +) +def test_rmsnorm_non_npu_residual_delegates_to_native( + monkeypatch: pytest.MonkeyPatch, + forward_name: str, +): + """Platform wrappers keep residual arithmetic in forward_native.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16) + x = torch.randn(2, 4, dtype=torch.bfloat16) + residual = torch.randn_like(x) + expected_output = torch.randn_like(x) + expected_residual = torch.randn_like(x) + captured: dict[str, object] = {} + + def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None): + captured.update(x=native_x, residual=native_residual) + return expected_output, expected_residual + + monkeypatch.setattr(norm, "forward_native", forward_native) + + output, updated_residual = getattr(norm, forward_name)(x, residual) + + assert output is expected_output + assert updated_residual is expected_residual + assert captured["x"] is x + assert captured["residual"] is residual + + +def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch): + """NPU keeps the direct fused AddRMSNorm path instead of IR fallback.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + captured: dict[str, object] = {} + fused_output = torch.randn(2, 4, dtype=torch.bfloat16) + fused_residual = torch.randn(2, 4, dtype=torch.bfloat16) + + def npu_add_rms_norm( + x: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + epsilon: float, + ): + captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon) + return fused_output, torch.empty(0), fused_residual + + monkeypatch.setitem( + sys.modules, + "torch_npu", + types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm), + ) + norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16) + x = torch.randn(2, 4, dtype=torch.bfloat16) + residual = torch.randn(2, 4, dtype=torch.bfloat16) + + output, updated_residual = norm.forward_npu(x, residual) + + assert output is fused_output + assert updated_residual is fused_residual + assert captured["x"] is x + assert captured["residual"] is residual + assert captured["gamma"] is norm.weight + assert captured["epsilon"] == 1e-5 + + # ── RMSNorm compile-path regression tests ── diff --git a/tests/diffusion/layers/test_norm_cuda.py b/tests/diffusion/layers/test_norm_cuda.py new file mode 100644 index 000000000..b791eab71 --- /dev/null +++ b/tests/diffusion/layers/test_norm_cuda.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.platforms import current_omni_platform + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.diffusion, + pytest.mark.cuda, + pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"), +] + + +def test_rmsnorm_cuda_residual_matches_native_contract() -> None: + eps = 1e-6 + device = torch.device(current_omni_platform.device_type) + torch.manual_seed(0) + norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device) + x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16) + residual = torch.randn_like(x) + x_before = x.clone() + residual_before = residual.clone() + + expected_output, expected_residual = norm.forward_native(x_before, residual_before) + output, updated_residual = norm.forward_cuda(x, residual) + + torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0) + torch.testing.assert_close(output, expected_output, atol=0, rtol=0) + torch.testing.assert_close(x, x_before, atol=0, rtol=0) + torch.testing.assert_close(residual, residual_before, atol=0, rtol=0) diff --git a/tests/diffusion/lora/test_base_linear.py b/tests/diffusion/lora/test_base_linear.py index 9386d0909..b2f76be05 100644 --- a/tests/diffusion/lora/test_base_linear.py +++ b/tests/diffusion/lora/test_base_linear.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice(): assert torch.allclose(out, expected) +def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch): + from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA + + layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA) + torch.nn.Module.__init__(layer) + layer.base_layer = torch.nn.Module() + layer.n_slices = 1 + layer.lora_a_stacked = (torch.zeros(2),) + layer.lora_b_stacked = (torch.zeros(3),) + + layer._set_diffusion_lora_buffer_device(torch.device("meta")) + assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked)) + + def recreate_on_cpu(self, max_loras, lora_config, model_config): + self.lora_a_stacked = (torch.zeros(4),) + self.lora_b_stacked = (torch.zeros(5),) + + monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu) + layer.create_lora_weights(max_loras=1, lora_config=object()) + + assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked)) + + def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch): # Verify that after reset_lora(), apply() skips LoRA matmuls even if the # LoRA tensors are still allocated and non-empty. diff --git a/tests/diffusion/lora/test_images_generations_lora.py b/tests/diffusion/lora/test_images_generations_lora.py index dbdf59647..a01bcdae3 100644 --- a/tests/diffusion/lora/test_images_generations_lora.py +++ b/tests/diffusion/lora/test_images_generations_lora.py @@ -24,11 +24,20 @@ from safetensors.torch import save_file from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer, OmniServerParams +from vllm_omni.platforms import current_omni_platform os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" MODEL = "Tongyi-MAI/Z-Image-Turbo" +# ROCm's default FLASH_ATTN diffusion backend resolves to AITER. On a cold CI +# worker, JIT-building its MHA kernel can take longer than the server readiness +# timeout. This test validates LoRA switching rather than attention-backend or +# compiled-execution performance, so use the eager SDPA path on ROCm only. +ROCM_SERVER_ARGS = ( + ["--enforce-eager", "--diffusion-attention-backend", "TORCH_SDPA"] if current_omni_platform.is_rocm() else [] +) + PROMPT = "a photo of a cat sitting on a laptop keyboard" SIZE = "256x256" @@ -42,7 +51,14 @@ server_params = [ server_args=[ "--num-gpus", "1", + *ROCM_SERVER_ARGS, ], + # A cold-cache Z-Image load can take slightly more than the + # orchestrator's 900s default on CI (weights alone have been + # observed at ~690s and full orchestrator readiness at ~906s). + # Match the large-model offline runner's startup allowance without + # weakening timeout detection for every online-serving test. + init_timeout=1800, ), id="zimage_turbo_lora", ), diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 723e4b13a..8edc6ce4b 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +from types import SimpleNamespace + import pytest import torch from vllm.lora.lora_weights import LoRALayerWeights @@ -20,9 +22,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _DummyLoRALayer: - def __init__(self, n_slices: int, output_slices: tuple[int, ...]): + def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1): self.n_slices = n_slices self.output_slices = output_slices + self.tp_size = tp_size self.set_calls: list[ tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor] ] = [] @@ -142,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch): assert replace_calls == ["foo"] +def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch): + import vllm_omni.diffusion.lora.manager as manager_mod + + class _DLOLoRALayer(DummyBaseLayerWithLoRA): + lora_buffer_device = None + + def _set_diffusion_lora_buffer_device(self, device): + self.lora_buffer_device = device + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer) + monkeypatch.setattr( + manager_mod, + "from_layer_diffusion", + lambda *, layer, **_kwargs: _DLOLoRALayer(layer), + ) + + pipeline = torch.nn.Module() + pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True) + pipeline.transformer = torch.nn.Module() + pipeline.transformer.foo = _FakeLinearBase() + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=torch.device("cpu"), + dtype=torch.bfloat16, + ) + + manager._replace_layers_with_lora(type("_PH", (), {"r": 1})()) + + assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu") + + def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch): import vllm_omni.diffusion.lora.manager as manager_mod @@ -236,6 +270,14 @@ def test_lora_manager_activates_fused_lora_on_packed_layer(): def test_lora_manager_activates_packed_lora_from_sublayers(): pipeline = torch.nn.Module() + bound_names = None + + def validate_binding(*, lora_model, bound_lora_names): + nonlocal bound_names + assert lora_model.id == 1 + bound_names = bound_lora_names + + pipeline._validate_diffusion_lora_binding = validate_binding pipeline.stacked_params_mapping = [ (".to_qkv", ".to_q", "q"), (".to_qkv", ".to_k", "k"), @@ -279,6 +321,70 @@ def test_lora_manager_activates_packed_lora_from_sublayers(): assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0) assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0) assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0) + assert bound_names == frozenset(loras) + + +def test_lora_manager_rolls_back_all_layers_when_activation_fails(): + class _FailingLoRALayer(_DummyLoRALayer): + fail = False + + def set_lora(self, index: int, lora_a, lora_b): + if self.fail: + raise RuntimeError("bind failed") + super().set_lora(index, lora_a, lora_b) + + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=2, + ) + first = _DummyLoRALayer(n_slices=1, output_slices=(2,)) + second = _FailingLoRALayer(n_slices=1, output_slices=(2,)) + manager._lora_modules = { + "transformer.first": first, + "transformer.second": second, + } + + def adapter(adapter_id: int): + loras = { + name: LoRALayerWeights( + module_name=name, + rank=2, + lora_alpha=2, + lora_a=torch.full((2, 2), float(adapter_id)), + lora_b=torch.full((2, 2), float(adapter_id)), + ) + for name in manager._lora_modules + } + return type( + "LM", + (), + { + "id": adapter_id, + "loras": loras, + "get_lora": lambda self, key: self.loras.get(key), + }, + )() + + manager._registered_adapters = {1: adapter(1), 2: adapter(2)} + manager._activate_adapter(1, scale=1.0) + first_calls_after_success = len(first.set_calls) + + second.fail = True + with pytest.raises(RuntimeError, match="bind failed"): + manager._activate_adapter(2, scale=1.0) + + assert manager._active_adapter_id is None + assert first.reset_calls == 1 + assert second.reset_calls == 1 + + # The failed activation must not leave the old adapter eligible for the + # fast path: reactivating it binds every layer again. + second.fail = False + manager._activate_adapter(1, scale=1.0) + assert manager._active_adapter_id == 1 + assert len(first.set_calls) == first_calls_after_success + 2 def _dummy_lora_request(adapter_id: int) -> LoRARequest: diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index f370d5b49..bd59ee94b 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -1,22 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for the DiffusersPipelineLoader. """ +import json +from pathlib import Path from types import SimpleNamespace import pytest import torch import torch.nn as nn from huggingface_hub import snapshot_download +from safetensors.torch import save_file from vllm.config.load import LoadConfig from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + HostWeightPlanResult, + TensorBinding, +) +from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module from vllm_omni.diffusion.models.helios import HeliosPipeline +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.diffusion.registry import initialize_model pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -67,6 +77,73 @@ class _DummyPipelineModel(nn.Module): return loaded +class _HWRTransformer(nn.Module): + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="test-hwr-transformer", + version="1", + ) + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)) + + def validate_restored_host_weights(self): + assert self.weight.dtype is torch.bfloat16 + + +class _HWRPipeline(nn.Module): + def __init__(self, source_root: Path): + super().__init__() + self.transformer = _HWRTransformer() + self.load_count = 0 + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=str(source_root), + subfolder=None, + revision=None, + prefix="transformer.", + fall_back_to_pt=False, + ) + ] + + def load_weights(self, weights): + loaded: set[str] = set() + params = dict(self.named_parameters()) + for name, tensor in weights: + if name in params: + params[name].data.copy_(tensor.to(dtype=params[name].dtype)) + loaded.add(name) + self.load_count += 1 + return loaded + + +def _hwr_config(model: str | Path, root: Path, *, mode: str = "preferred") -> SimpleNamespace: + return SimpleNamespace( + model=str(model), + dtype=torch.bfloat16, + host_weight_runtime_mode=mode, + host_weight_runtime_root=str(root), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + lora_path=None, + quantization_config=None, + diffusion_attention_config=None, + parallel_config=SimpleNamespace( + use_hsdp=False, + data_parallel_size=1, + sequence_parallel_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + cfg_parallel_size=1, + enable_expert_parallel=False, + ulysses_degree=1, + ring_degree=1, + allgather_degree=1, + ulysses_mode="strict", + ), + ) + + def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader: od_config = SimpleNamespace( dtype=torch.float32, @@ -86,6 +163,322 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade return loader +def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading( + tmp_path: Path, + monkeypatch, +): + canonical_root = tmp_path / "canonical" + canonical_root.mkdir() + save_file( + {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)}, + str(canonical_root / "model.safetensors"), + ) + store_root = tmp_path / "hwr-store" + hash_calls = 0 + original_sha256 = source_identity_module._sha256_file + + def counted_sha256(path: Path, state: object) -> str: + nonlocal hash_calls + hash_calls += 1 + return original_sha256(path, state) # type: ignore[arg-type] + + monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256) + + def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]: + loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root)) + pipeline = _HWRPipeline(canonical_root) + monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline) + return loader, pipeline + + cold_loader, cold_model = make_loader() + cold = cold_loader.load_model(load_device="cpu", device=torch.device("cpu")) + assert cold is cold_model + assert cold_model.load_count == 1 + assert cold_loader._hwr_state is not None + + warm_loader, warm_model = make_loader() + monkeypatch.setattr( + warm_loader, + "_process_weights_after_loading", + lambda *args, **kwargs: pytest.fail("warm HWR restore re-entered byte-changing finalization"), + ) + warm = warm_loader.load_model(load_device="cpu", device=torch.device("cpu")) + + assert warm is warm_model + assert warm_model.load_count == 0 + assert torch.equal(warm_model.transformer.weight, cold_model.transformer.weight) + from vllm_omni.diffusion.offloader.startup import take_offload_startup_state + + startup_state = take_offload_startup_state(warm) + assert startup_state is not None + warm_plan = startup_state.host_weight_plan + assert warm_plan is not None + assert warm_plan.lease_carrier is not None + warm_plan.lease_carrier.close() + assert hash_calls == 1 + assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1 + + +def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch): + from vllm_omni.diffusion.model_loader import diffusers_loader as loader_module + + loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(tmp_path, tmp_path / "store")) + models: list[_DummyPipelineModel] = [] + + def init_model(*args, **kwargs): + del args, kwargs + model = _DummyPipelineModel(source_prefix="transformer.") + models.append(model) + return model + + def commit_error(*args, **kwargs): + del args, kwargs + raise loader_module._HWRCommitError("restore commit failed") + + monkeypatch.setattr(loader, "_init_from_load_format", init_model) + monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: ()) + monkeypatch.setattr(loader, "_resolve_hwr", commit_error) + monkeypatch.setattr(loader, "load_weights", lambda *args, **kwargs: None) + monkeypatch.setattr(loader, "_process_weights_after_loading", lambda *args, **kwargs: None) + monkeypatch.setattr(loader, "_apply_skip_softmax_calibration", lambda *args, **kwargs: None) + + recovered = loader.load_model(load_device="cpu", device=torch.device("cpu")) + + assert len(models) == 2 + assert recovered is models[1] + assert loader.take_host_weight_plan() is None + + +def test_required_hwr_miss_fails_before_ordinary_loading_or_publication( + tmp_path: Path, + monkeypatch, +): + canonical_root = tmp_path / "canonical" + canonical_root.mkdir() + save_file( + {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)}, + str(canonical_root / "model.safetensors"), + ) + loader = DiffusersPipelineLoader( + LoadConfig(), + _hwr_config(canonical_root, tmp_path / "empty-store", mode="required"), + ) + pipeline = _HWRPipeline(canonical_root) + monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline) + + with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"): + loader.load_model(load_device="cpu", device=torch.device("cpu")) + + assert pipeline.load_count == 0 + assert loader.take_host_weight_plan() is None + + +def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model="", + dtype=torch.float32, + quantization_config="fp8", + parallel_config=DiffusionParallelConfig( + data_parallel_size=dp_size, + sequence_parallel_size=1, + ), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + + +@pytest.mark.parametrize( + ("dist_offload", "use_allgather", "mode"), + [ + (False, False, "preferred"), + (True, True, "preferred"), + (True, False, "disabled"), + ], +) +def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction( + monkeypatch, + tmp_path, + dist_offload, + use_allgather, + mode, +): + """Disabled and AllGather paths must never construct or probe HWR.""" + from vllm_omni.host_weight_runtime import HostWeightRuntime + + root = tmp_path / "must-not-be-touched" + loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode)) + model = _DummyPipelineModel(source_prefix="transformer.") + modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,)) + + def unexpected_store_construction(*args, **kwargs): + raise AssertionError(f"HWR store interaction was not eligible: {args}, {kwargs}") + + monkeypatch.setattr(HostWeightRuntime, "from_config", unexpected_store_construction) + assert ( + loader._resolve_hwr( + model, + modules, + dist_offload=dist_offload, + use_allgather=use_allgather, + load_format="default", + sources=tuple(model.weights_sources), + ) + is None + ) + assert not root.exists() + + +def test_required_hwr_rejects_a_model_without_a_restore_contract(tmp_path): + loader = DiffusersPipelineLoader( + LoadConfig(), + _hwr_config("dummy-model", tmp_path / "store", mode="required"), + ) + model = _DummyPipelineModel(source_prefix="transformer.") + modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,)) + + with pytest.raises(ValueError, match="restore contract"): + loader._resolve_hwr( + model, + modules, + dist_offload=True, + use_allgather=False, + load_format="default", + sources=tuple(model.weights_sources), + ) + + +@pytest.mark.parametrize("offline", [False, True]) +def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + snapshot = tmp_path / "snapshot" + transformer = snapshot / "transformer" + transformer.mkdir(parents=True) + indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.safetensors" for index in (1, 2)] + stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors" + index_path = transformer / "diffusion_pytorch_model.safetensors.index.json" + index_path.write_text( + json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}}) + ) + for filename in indexed_files + [stale_file]: + (transformer / filename).touch() + + mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", offline) + + def download_index(*, filename, **_kwargs): + if filename == "transformer/diffusion_pytorch_model.safetensors.index.json": + return str(index_path) + raise loader_mod.huggingface_hub.errors.EntryNotFoundError(filename) + + hub_api = mocker.Mock() + hub_api.hf_hub_download.side_effect = download_index + mocker.patch.object(loader_mod, "hf_api", return_value=hub_api) + indexed_download = mocker.patch.object(loader_mod, "download_weights_from_hf_specific", return_value=str(snapshot)) + generic_download = mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot)) + + loader = _make_loader_with_weights([]) + cache_dir = str(tmp_path / "cache") + loader.load_config.download_dir = cache_dir + folder, files, use_safetensors = loader._prepare_weights( + "org/model", + subfolder="transformer", + revision="revision", + fall_back_to_pt=True, + allow_patterns_overrides=None, + ) + + assert folder == str(transformer) + assert [str(transformer / filename) for filename in indexed_files] == files + assert use_safetensors + assert hub_api.hf_hub_download.call_count == 2 + hub_api.hf_hub_download.assert_any_call( + repo_id="org/model", + filename="transformer/diffusion_pytorch_model.safetensors.index.json", + cache_dir=cache_dir, + revision="revision", + local_files_only=offline, + ) + indexed_download.assert_called_once_with( + model_name_or_path="org/model", + cache_dir=cache_dir, + allow_patterns=[f"transformer/{filename}" for filename in indexed_files], + revision="revision", + ignore_patterns=loader.load_config.ignore_patterns, + require_all=True, + ) + + _, override_files, _ = loader._prepare_weights( + "org/model", + subfolder="transformer", + revision="revision", + fall_back_to_pt=True, + allow_patterns_overrides=[stale_file], + ) + assert override_files == [str(transformer / stale_file)] + generic_download.assert_called_once_with( + "org/model", + cache_dir, + [stale_file], + "revision", + subfolder="transformer", + ignore_patterns=loader.load_config.ignore_patterns, + ) + + +def test_prepare_local_weights_honors_component_index(tmp_path): + transformer = tmp_path / "transformer" + transformer.mkdir() + indexed_file = "diffusion_pytorch_model-00001-of-00001.safetensors" + stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors" + (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"weight": indexed_file}}) + ) + for filename in (indexed_file, stale_file): + (transformer / filename).touch() + + folder, files, use_safetensors = _make_loader_with_weights([])._prepare_weights( + tmp_path, + subfolder="transformer", + revision=None, + fall_back_to_pt=True, + allow_patterns_overrides=None, + ) + + assert folder == str(transformer) + assert files == [str(transformer / indexed_file)] + assert use_safetensors + + +def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path, mocker): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + snapshot = tmp_path / "snapshot" + transformer = snapshot / "transformer" + transformer.mkdir(parents=True) + for shard_count in (4, 8): + for shard in range(1, shard_count + 1): + (transformer / f"diffusion_pytorch_model-{shard:05d}-of-{shard_count:05d}.safetensors").touch() + + mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", True) + hub_api = mocker.Mock() + hub_api.hf_hub_download.side_effect = loader_mod.huggingface_hub.errors.EntryNotFoundError("index is not cached") + mocker.patch.object(loader_mod, "hf_api", return_value=hub_api) + mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot)) + + loader = _make_loader_with_weights([]) + with pytest.raises(ValueError, match="conflicting shard totals"): + loader._prepare_weights( + "org/model", + subfolder="transformer", + revision="revision", + fall_back_to_pt=True, + allow_patterns_overrides=None, + ) + + assert hub_api.hf_hub_download.call_count == 2 + assert all(call.kwargs["local_files_only"] for call in hub_api.hf_hub_download.call_args_list) + + def test_strict_check_only_validates_source_prefix_parameters(): model = _DummyPipelineModel(source_prefix="transformer.") loader = _make_loader_with_weights(["transformer.weight"]) @@ -110,6 +503,85 @@ def test_empty_source_prefix_keeps_full_model_strict_check(): loader.load_weights(model) +def test_stream_online_quant_weights_offloads_layers_after_processing(): + from vllm.model_executor.model_loader.reload.layerwise import ( + get_layerwise_info, + ) + + events: list[str] = [] + + class _OnlineQuantMethod: + uses_meta_device = True + + class _TrackedLayer(nn.Linear): + def __init__(self, name: str): + super().__init__(2, 2, bias=False) + self.name = name + self.quant_method = _OnlineQuantMethod() + get_layerwise_info(self).load_numel_total = self.weight.numel() + + def to(self, *args, **kwargs): + events.append(self.name) + return super().to(*args, **kwargs) + + class _StreamingModel(nn.Module): + def __init__(self): + super().__init__() + self.first = _TrackedLayer("first") + self.second = _TrackedLayer("second") + + model = _StreamingModel() + weights = iter( + [ + ("first.weight", torch.zeros((2, 2))), + ("second.weight", torch.zeros((2, 2))), + ] + ) + streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights) + + assert next(streamed)[0] == "first.weight" + get_layerwise_info(model.first).reset() + assert next(streamed)[0] == "second.weight" + assert events == ["first"] + + get_layerwise_info(model.second).reset() + with pytest.raises(StopIteration): + next(streamed) + assert events == ["first", "second"] + + +def test_process_weights_skips_completed_online_quant_layer(monkeypatch): + from unittest.mock import Mock + + from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase + from vllm.model_executor.model_loader.reload import layerwise + + class _TrackedLayer(nn.Linear): + def __init__(self): + super().__init__(2, 2, bias=False) + self.to_calls: list[object] = [] + + def to(self, *args, **kwargs): + self.to_calls.append(args[0] if args else kwargs.get("device")) + return super().to(*args, **kwargs) + + model = nn.Module() + model.layer = _TrackedLayer() + quant_method = Mock(spec=QuantizeMethodBase) + quant_method.uses_meta_device = True + model.layer.quant_method = quant_method + model.layer._already_called_process_weights_after_loading = True + finalize = Mock() + monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize) + + loader = _make_loader_with_weights([]) + loader._process_weights_after_loading(model, torch.device("cuda")) + + finalize.assert_called_once_with(model, model_config=None) + quant_method.process_weights_after_loading.assert_not_called() + assert model.layer.to_calls == [] + + class _ConfigAwareModel(nn.Module): def __init__(self, *, od_config): super().__init__() @@ -179,6 +651,340 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch): assert get_current_diffusion_config_or_none() is None +def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={}, + ) + loaded_ordinary_weights = False + + def load_weights(_model): + nonlocal loaded_ordinary_weights + loaded_ordinary_weights = True + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = load_weights # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(plan), + ) + + assert loader.load_model(load_device="cpu") is model + assert not loaded_ordinary_weights + assert loader.take_host_weight_plan() is plan + assert loader.take_host_weight_plan() is None + + +def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + class MixedSourceModel(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Linear(2, 2, bias=False) + self.text_encoder = nn.Linear(2, 2, bias=False) + self.weights_sources = ( + DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."), + DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."), + ) + self.loaded_weight_names: list[str] = [] + + def load_weights(self, weights): + self.loaded_weight_names = [name for name, _ in weights] + return set(self.loaded_weight_names) + + model = MixedSourceModel() + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={ + "transformer.weight": TensorBinding( + checkpoint_key="weight", + file_path="unused", + ) + }, + planned_source_prefixes=frozenset({"transformer."}), + ) + requested_prefixes: list[str] = [] + + def get_weights(source, model=None): + del model + requested_prefixes.append(source.prefix) + yield source.prefix + "weight", torch.ones(2, 2) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader._get_weights_iterator = get_weights # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(plan), + ) + + assert loader.load_model(load_device="cpu") is model + assert requested_prefixes == ["text_encoder."] + assert model.loaded_weight_names == ["text_encoder.weight"] + assert loader.take_host_weight_plan() is plan + + +def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + calls: list[str] = [] + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"), + ) + + assert loader.load_model(load_device="cpu") is model + assert calls == ["load", "process"] + assert loader.take_host_weight_plan() is None + + +def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch): + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = _make_dlo_online_quant_config() + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod) + model.transformer.quant_method.uses_meta_device = True + calls: list[object] = [] + allowlist_models: list[nn.Module] = [] + + original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods + + def check_allowlist(candidate: nn.Module) -> tuple[str, ...]: + allowlist_models.append(candidate) + return original_allowlist_check(candidate) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist) + loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign] + loader.load_weights = ( # type: ignore[method-assign] + lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu)) + ) + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + + assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model + assert allowlist_models == [model] + assert calls == [("load", True), "process"] + assert loader.take_host_weight_plan() is None + + +def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod + + od_config = _make_dlo_online_quant_config() + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod) + model.transformer.quant_method.uses_meta_device = True + calls: list[object] = [] + allowlist_models: list[nn.Module] = [] + + original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods + + def check_allowlist(candidate: nn.Module) -> tuple[str, ...]: + allowlist_models.append(candidate) + return original_allowlist_check(candidate) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist) + loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign] + loader.load_weights = ( # type: ignore[method-assign] + lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu)) + ) + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + + assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model + assert allowlist_models == [model] + assert calls == [("load", True), "process"] + assert loader.take_host_weight_plan() is None + + +def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch): + """A DLO group of one runs no weight collective, so an otherwise + unvalidated online method must not be rejected by the AllGather gate.""" + + class UnsupportedOnlineMethod: + uses_meta_device = True + + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = _make_dlo_online_quant_config(dp_size=1) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = UnsupportedOnlineMethod() + calls: list[object] = [] + allowlist_models: list[nn.Module] = [] + + original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods + + def check_allowlist(candidate: nn.Module) -> tuple[str, ...]: + allowlist_models.append(candidate) + return original_allowlist_check(candidate) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist) + loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign] + loader.load_weights = ( # type: ignore[method-assign] + lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu)) + ) + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + + assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model + assert allowlist_models == [] + assert calls == [("load", True), "process"] + + +def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + class UnsupportedOnlineMethod: + uses_meta_device = True + + od_config = _make_dlo_online_quant_config() + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = UnsupportedOnlineMethod() + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + + with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"): + loader.load_model(load_device="cpu") + + +def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch): + mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config") + NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod + + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = _make_dlo_online_quant_config() + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod) + model.transformer.quant_method.uses_meta_device = True + calls: list[object] = [] + allowlist_models: list[nn.Module] = [] + + original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods + + def check_allowlist(candidate: nn.Module) -> tuple[str, ...]: + allowlist_models.append(candidate) + return original_allowlist_check(candidate) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist) + loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign] + loader.load_weights = ( # type: ignore[method-assign] + lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu)) + ) + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult( + None, + "online quantization requires the ordinary loader", + ), + ) + + assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model + assert allowlist_models == [model] + assert calls == [("load", True), "process"] + assert loader.take_host_weight_plan() is None + + def test_hsdp_processes_quantized_weights_before_sharding(mocker): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod from vllm_omni.diffusion.offloader.module_collector import PipelineModules diff --git a/tests/diffusion/model_loader/test_final_layout_host_weights.py b/tests/diffusion/model_loader/test_final_layout_host_weights.py new file mode 100644 index 000000000..0b64286cc --- /dev/null +++ b/tests/diffusion/model_loader/test_final_layout_host_weights.py @@ -0,0 +1,1428 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Contract tests for final-layout BF16 Host Weight Runtime artifacts.""" + +from __future__ import annotations + +import dataclasses +import gc +import json +import threading +from collections.abc import Sequence +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.model_loader.host_weights import ( + FINAL_LAYOUT_BF16_POLICY, + FINAL_LAYOUT_BF16_SPEC, + FinalLayoutBF16Producer, + FinalLayoutIdentityContext, + FinalLayoutLoaderIdentity, + FinalLayoutParallelIdentity, + FinalLayoutRequest, + FinalLayoutTensorRestorer, + ImplementationIdentity, + NodeSourceDigestCache, + PreparedWeightSource, + WeightSourceKind, + build_final_layout_identity, +) +from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module +from vllm_omni.diffusion.model_loader.host_weights.contracts import ( + FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + FinalLayoutArtifactSpec, + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutTensorPolicy, + implementation_abi_fingerprint, +) +from vllm_omni.diffusion.model_loader.host_weights.tensor_layout import ( + RuntimeTensorTarget, + collect_final_layout_targets, +) +from vllm_omni.diffusion.models.host_weight_contract import ( + FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + FinalLayoutModelContract, +) +from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3DiTModel +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + CoordinationScope, + FailureCode, + HostWeightRuntime, + HostWeightRuntimeConfig, + LookupPhase, + PostLoadPublicationOutcome, + ProducerIdentity, + ProductionPolicy, + ProductionSourceMode, + ResolutionOutcome, + ResolutionStage, + RuntimeMode, + StorageDomainPolicy, + StorageScope, + TensorKind, + WaitPolicy, + WeightRepresentation, +) +from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _TinyDiT(nn.Module): + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="test-tiny-dit", + version="1", + ) + + def __init__(self, *, device: torch.device | str = "cpu") -> None: + super().__init__() + target = torch.device(device) + self.restore_validations = 0 + self.proj = nn.Linear(3, 2, dtype=torch.bfloat16, device=target) + self.fp32_gain = nn.Parameter(torch.empty(1, dtype=torch.float32, device=target)) + self.register_buffer("scale", torch.empty(1, dtype=torch.float32, device=target)) + self.register_buffer( + "derived", + torch.tensor([7.0], dtype=torch.float32), + persistent=False, + ) + + def validate_restored_host_weights(self) -> None: + self.restore_validations += 1 + assert not self.proj.weight.is_meta + assert self.proj.weight.dtype is torch.bfloat16 + assert self.fp32_gain.dtype is torch.float32 + assert self.scale.dtype is torch.float32 + + +class _TinyPipeline(nn.Module): + def __init__(self, *, dit_device: torch.device | str = "cpu") -> None: + super().__init__() + self.transformer = _TinyDiT(device=dit_device) + self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16) + self.vae = nn.Module() + self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32)) + + +class _AlternateTinyPipeline(_TinyPipeline): + pass + + +class _MultiTinyPipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer_a = _TinyDiT() + self.transformer_b = _TinyDiT() + + +class _Float32DiT(nn.Module): + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="test-float32-dit", + version="1", + ) + + def __init__(self) -> None: + super().__init__() + self.proj = nn.Linear(2, 2, dtype=torch.float32) + + def validate_restored_host_weights(self) -> None: + assert self.proj.weight.dtype is torch.float32 + + +class _Float32Pipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer = _Float32DiT() + + +_SYNTHETIC_FP32_ABI = CanonicalJson.from_value( + { + "artifact_identity": "diffusion-final-layout-identity-v1", + "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + "producer": "test-fp32-producer-v1", + "representation_policy": "test-all-fp32-v1", + "restorer": "exact-final-layout-tensor-rebind-v1", + "source_identity": "prepared-diffusion-weight-source-v2", + "tensor_contract": "complete-strided-tensor-ownership-v1", + } +) +_SYNTHETIC_FP32_SPEC = FinalLayoutArtifactSpec( + representation=WeightRepresentation( + name="test-final-layout-fp32", + dtype=str(torch.float32), + metadata=CanonicalJson.from_value({"policy": "all-fp32"}), + ), + producer=ProducerIdentity( + producer_id="test.final-layout-fp32", + version="1", + implementation_fingerprint=implementation_abi_fingerprint(_SYNTHETIC_FP32_ABI), + manifest_schema="test-final-layout-fp32-manifest-v1", + restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + ), + implementation_abi=_SYNTHETIC_FP32_ABI, +) + + +class _SyntheticFP32Policy: + @property + def spec(self) -> FinalLayoutArtifactSpec: + return _SYNTHETIC_FP32_SPEC + + def validate_request(self, request: FinalLayoutRequest) -> None: + if request.load_format != "default": + raise ValueError("synthetic policy requires default loading") + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: + del name, tensor + return "weight" if kind is TensorKind.PARAMETER else "persistent_buffer" + + def validate_target(self, target: RuntimeTensorTarget) -> None: + if target.tensor.dtype is not torch.float32: + raise ValueError("synthetic policy requires FP32 tensors") + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: + if not targets: + raise ValueError("synthetic policy requires tensors") + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: + return CanonicalJson.from_value( + { + "component_names": list(component_names), + "format": "test-final-layout-fp32", + "tensor_contract_sha256": tensor_contract_digest, + "tensor_count": tensor_count, + } + ) + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: + assert metadata == self.build_format_metadata( + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + + +_SYNTHETIC_FP32_POLICY = _SyntheticFP32Policy() + + +class _PolicyWithSpec: + def __init__(self, base: FinalLayoutTensorPolicy, spec: FinalLayoutArtifactSpec) -> None: + self._base = base + self._spec = spec + + @property + def spec(self) -> FinalLayoutArtifactSpec: + return self._spec + + def validate_request(self, request: FinalLayoutRequest) -> None: + self._base.validate_request(request) + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: + return self._base.tensor_role(name, tensor, kind) + + def validate_target(self, target: RuntimeTensorTarget) -> None: + self._base.validate_target(target) + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: + self._base.validate_collection(targets) + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: + return self._base.build_format_metadata( + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: + self._base.validate_format_metadata( + metadata, + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + + +class _FakeH3Linear(nn.Module): + def __init__( + self, + *args: object, + bias: bool = False, + params_dtype: torch.dtype | None = None, + total_num_heads: int | None = None, + total_num_kv_heads: int | None = None, + **kwargs: object, + ) -> None: + del args, kwargs + super().__init__() + dtype = params_dtype or torch.get_default_dtype() + self.weight = nn.Parameter(torch.empty(1, dtype=dtype)) + if bias: + self.bias = nn.Parameter(torch.empty(1, dtype=dtype)) + else: + self.register_parameter("bias", None) + self.num_heads = total_num_heads + self.num_kv_heads = total_num_kv_heads + + +class _FakeH3Attention(nn.Module): + def __init__(self, **kwargs: object) -> None: + del kwargs + super().__init__() + + +class _FakeFluxLinear(_FakeH3Linear): + def __init__( + self, + *args: object, + bias: bool = False, + params_dtype: torch.dtype | None = None, + total_num_heads: int | None = None, + total_num_kv_heads: int | None = None, + **kwargs: object, + ) -> None: + super().__init__( + *args, + bias=bias, + params_dtype=params_dtype, + total_num_heads=total_num_heads, + total_num_kv_heads=total_num_kv_heads, + **kwargs, + ) + self.weight.weight_loader = self._load_weight + + @staticmethod + def _load_weight( + parameter: nn.Parameter, + loaded_weight: torch.Tensor, + shard_id: str | None = None, + ) -> None: + del shard_id + parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter)) + + +class _FakeFluxAttention(nn.Module): + def __init__(self, **kwargs: object) -> None: + del kwargs + super().__init__() + self.register_buffer("beta", torch.tensor([0.5])) + self.register_buffer("eps", torch.tensor([-1e-6])) + + +class _Flux2KleinPipeline(nn.Module): + def __init__(self, transformer: nn.Module) -> None: + super().__init__() + self.transformer = transformer + self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16) + self.vae = nn.Module() + self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32)) + + +def _reduced_flux2_klein( + monkeypatch: pytest.MonkeyPatch, + *, + device: torch.device | str = "cpu", +) -> nn.Module: + from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2 + + monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear) + monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear) + monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear) + monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear) + monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention) + + transformer = flux2.Flux2Transformer2DModel( + patch_size=1, + in_channels=8, + out_channels=8, + num_layers=1, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=1, + joint_attention_dim=8, + timestep_guidance_channels=8, + mlp_ratio=1.0, + axes_dims_rope=(2, 2, 2, 2), + guidance_embeds=False, + ) + return transformer.to(device=device, dtype=torch.bfloat16) + + +class _H3Pipeline(nn.Module): + def __init__(self, transformer: nn.Module) -> None: + super().__init__() + self.transformer = transformer + + +def _small_h3_config() -> SimpleNamespace: + return SimpleNamespace( + tf_model_config={ + "num_layers": 1, + "token_refiner_num_layers": 1, + "hidden_size": 8, + "num_attention_heads": 2, + "attention_head_dim": 4, + "ffn_hidden_size": 16, + "latents_dim": 2, + "audio_latents_dim": 2, + "patch_size": (1, 2, 2), + "text_dim": 6, + "timestep_input_dim": 4, + "time_embed_hidden_size": 8, + "time_embed_dim": 4, + "adaln_out_features": 18 * 8, + "final_adaln_out_features": 2 * 8, + "rope_inv_freq_len": 2, + }, + parallel_config=SimpleNamespace(ulysses_degree=1), + ) + + +def _prepared_source( + tmp_path: Path, + *, + content: bytes = b"canonical-source-for-identity", + requested_revision: str | None = None, + resolved_revision: str | None = None, + directory: str = "canonical", + prefix: str = "transformer.", +) -> PreparedWeightSource: + source_root = tmp_path / directory + if resolved_revision is not None: + source_root = source_root / "snapshots" / resolved_revision + source_root.mkdir(parents=True, exist_ok=True) + weight_file = source_root / "model.safetensors" + if not weight_file.exists(): + weight_file.write_bytes(content) + elif weight_file.read_bytes() != content: + weight_file.write_bytes(content) + return PreparedWeightSource( + model_or_path="test-org/tiny-diffusion", + subfolder=None, + requested_revision=requested_revision, + prefix=prefix, + resolved_root=source_root, + weight_files=(weight_file,), + use_safetensors=True, + ) + + +def _request(**changes: object) -> FinalLayoutRequest: + request = FinalLayoutRequest( + model_id="test-org/tiny-diffusion", + loader=FinalLayoutLoaderIdentity( + implementation=ImplementationIdentity( + implementation_id="test-diffusion-loader", + version="1", + fingerprint="test-loader-implementation-v1", + ), + model_config_fingerprint="test-model-config-v1", + weight_transform_fingerprint="test-weight-transform-v1", + ), + ) + return dataclasses.replace(request, **changes) + + +def _identity( + model: nn.Module, + source: PreparedWeightSource, + *, + request: FinalLayoutRequest | None = None, + policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY, + source_digest_cache: NodeSourceDigestCache | None = None, +) -> FinalLayoutIdentityContext: + transformer = model.get_submodule("transformer") + return build_final_layout_identity( + model, + dit_modules=(("transformer", transformer),), + prepared_sources=(source,), + request=request or _request(), + policy=policy, + source_digest_cache=source_digest_cache, + ) + + +def _multi_identity( + model: _MultiTinyPipeline, + sources: tuple[PreparedWeightSource, ...], +) -> FinalLayoutIdentityContext: + return build_final_layout_identity( + model, + dit_modules=( + ("transformer_a", model.transformer_a), + ("transformer_b", model.transformer_b), + ), + prepared_sources=sources, + request=_request(), + policy=FINAL_LAYOUT_BF16_POLICY, + ) + + +def _runtime(root: Path) -> HostWeightRuntime: + return HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=StorageDomainPolicy( + root=root, + scope=StorageScope.NODE, + domain_id="node", + storage_class=detect_storage_class(root.parent), + ), + production=ProductionPolicy( + allow_local_build=False, + allow_post_load_publish=True, + ), + wait=WaitPolicy(coordination_timeout_seconds=5.0), + ) + ) + + +def _dit_modules(model: _TinyPipeline) -> tuple[tuple[str, nn.Module], ...]: + return (("transformer", model.transformer),) + + +def _fill_final_weights(model: _TinyPipeline) -> None: + with torch.no_grad(): + model.transformer.proj.weight.copy_(torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3)) + model.transformer.proj.bias.copy_(torch.tensor([3.0, 4.0], dtype=torch.bfloat16)) + model.transformer.fp32_gain.copy_(torch.tensor([6.5])) + model.transformer.scale.copy_(torch.tensor([2.5])) + model.text_encoder.weight.fill_(11) + model.text_encoder.bias.fill_(12) + + +def _pointer_snapshot(model: nn.Module) -> dict[str, tuple[int, int, str]]: + snapshot: dict[str, tuple[int, int, str]] = {} + for name, tensor in list(model.named_parameters()) + list(model.named_buffers()): + pointer = 0 if tensor.is_meta else tensor.untyped_storage().data_ptr() + snapshot[name] = (id(tensor), pointer, tensor.device.type) + return snapshot + + +def test_preload_identity_is_stable_for_cpu_and_meta_skeletons(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + cpu_model = _TinyPipeline() + meta_model = _TinyPipeline(dit_device="meta") + + cpu = _identity(cpu_model, source) + meta = _identity(meta_model, source) + + assert cpu.identity == meta.identity + assert cpu.tensor_names == { + "transformer.fp32_gain", + "transformer.proj.bias", + "transformer.proj.weight", + "transformer.scale", + } + assert "transformer.derived" not in cpu.tensor_names + assert all("text_encoder" not in name and "vae" not in name for name in cpu.tensor_names) + + request_fields = {field.name for field in dataclasses.fields(FinalLayoutRequest)} + assert not any("data_parallel" in name for name in request_fields) + assert not any("metadata" in name for name in request_fields) + identity_bytes = cpu.identity.canonical_bytes + assert b"data_parallel" not in identity_bytes + assert b"dlo_use_allgather" not in identity_bytes + assert b"registration" not in identity_bytes + with pytest.raises(TypeError, match="unexpected keyword argument"): + dataclasses.replace(_request(), loader_metadata=CanonicalJson.empty()) + + +def test_generic_identity_and_tensor_discovery_accept_a_second_policy(tmp_path: Path) -> None: + model = _Float32Pipeline() + source = _prepared_source(tmp_path, directory="fp32-source") + + context = _identity(model, source, policy=_SYNTHETIC_FP32_POLICY) + targets = collect_final_layout_targets( + model, + (("transformer", model.transformer),), + policy=_SYNTHETIC_FP32_POLICY, + require_materialized=True, + ) + + assert context.spec is _SYNTHETIC_FP32_SPEC + assert context.identity.representation.name == "test-final-layout-fp32" + assert context.identity.representation.dtype == str(torch.float32) + assert all(target.tensor.dtype is torch.float32 for target in targets) + assert context.identity.producer.implementation_fingerprint != ( + FINAL_LAYOUT_BF16_SPEC.producer.implementation_fingerprint + ) + + +def test_source_identity_requires_complete_multi_component_coverage(tmp_path: Path) -> None: + model = _MultiTinyPipeline() + source_a = _prepared_source( + tmp_path, + directory="source-a", + prefix="transformer_a.", + ) + + with pytest.raises(FinalLayoutContractError, match="transformer_b") as exc_info: + _multi_identity(model, (source_a,)) + assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID + + source_b = _prepared_source( + tmp_path, + directory="source-b", + prefix="transformer_b.", + ) + context = _multi_identity(model, (source_a, source_b)) + metadata = context.identity.source.metadata.to_value() + assert isinstance(metadata, dict) + bindings = metadata["target_bindings"] + assert isinstance(bindings, list) + assert {binding["source_prefix"] for binding in bindings} == { + "transformer_a.", + "transformer_b.", + } + assert {binding["target_name"].split(".", 1)[0] for binding in bindings} == { + "transformer_a", + "transformer_b", + } + + +def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path) -> None: + model = _MultiTinyPipeline() + root_source = _prepared_source( + tmp_path, + directory="root-source", + prefix="", + ) + source_a = _prepared_source( + tmp_path, + directory="specific-a", + prefix="transformer_a.", + ) + + context = _multi_identity(model, (root_source, source_a)) + reordered = _multi_identity(model, (source_a, root_source)) + assert context.identity == reordered.identity + metadata = context.identity.source.metadata.to_value() + assert isinstance(metadata, dict) + bindings = metadata["target_bindings"] + assert isinstance(bindings, list) + by_target = {binding["target_name"]: binding["source_prefix"] for binding in bindings} + assert all(prefix == "transformer_a." for name, prefix in by_target.items() if name.startswith("transformer_a.")) + assert all(prefix == "" for name, prefix in by_target.items() if name.startswith("transformer_b.")) + + conflicting_a = _prepared_source( + tmp_path, + directory="conflicting-a", + prefix="transformer_a.", + ) + with pytest.raises(FinalLayoutContractError, match="equally specific") as exc_info: + _multi_identity(model, (root_source, source_a, conflicting_a)) + assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID + + +def test_node_source_digest_cache_reuses_recovers_and_invalidates( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + cache_root = tmp_path / "store" + calls = 0 + original_sha256 = source_identity_module._sha256_file + + def counted_sha256(path: Path, state: object) -> str: + nonlocal calls + calls += 1 + return original_sha256(path, state) # type: ignore[arg-type] + + monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256) + first = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + second = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + + assert second.identity == first.identity + assert calls == 1 + + cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json")) + cache_entry.write_text("not-json", encoding="utf-8") + recovered = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + assert recovered.identity == first.identity + assert calls == 2 + + cache_document = json.loads(cache_entry.read_text(encoding="utf-8")) + cache_document["content_id"] = f"sha256:{'0' * 64}" + cache_entry.write_text(json.dumps(cache_document), encoding="utf-8") + checksum_recovered = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + assert checksum_recovered.identity == first.identity + assert calls == 3 + + source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size") + changed = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + assert changed.identity != first.identity + assert calls == 4 + + +def test_source_identity_hashes_checkpoint_shards_in_parallel( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + second_file = source.resolved_root / "model-00002.safetensors" + second_file.write_bytes(b"second-checkpoint-shard") + source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file)) + barrier = threading.Barrier(2) + thread_ids: set[int] = set() + thread_ids_lock = threading.Lock() + original_sha256 = source_identity_module._sha256_file + + def synchronized_sha256(path: Path, state: object) -> str: + with thread_ids_lock: + thread_ids.add(threading.get_ident()) + barrier.wait(timeout=5) + return original_sha256(path, state) # type: ignore[arg-type] + + monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256) + _identity(model, source) + + assert len(thread_ids) == 2 + + +def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01) + + def unavailable_lock(*_args: object, **_kwargs: object) -> object: + raise source_identity_module.FileLockTimeoutError("busy") + + monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock) + context = _identity(model, source, source_digest_cache=cache) + + assert context.identity.source.revision.startswith("content-") + assert not tuple(cache.entries.glob("*.json")) + + +def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + original_mkdir = Path.mkdir + + def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None: + if "source-digests-v1" in path.parts: + raise OSError("injected cache directory failure") + original_mkdir(path, *args, **kwargs) # type: ignore[arg-type] + + monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir) + cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01) + context = _identity(model, source, source_digest_cache=cache) + + assert context.identity.source.revision.startswith("content-") + + +def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None: + model = _TinyPipeline() + source_root = tmp_path / "local-symlink-source" + blobs_root = tmp_path / "local-blobs" + source_root.mkdir() + blobs_root.mkdir() + blob_path = blobs_root / ("a" * 64) + blob_path.write_bytes(b"a" * 32) + weight_path = source_root / "model.safetensors" + weight_path.symlink_to(blob_path) + source = PreparedWeightSource( + model_or_path=str(source_root), + subfolder=None, + requested_revision=None, + prefix="transformer.", + resolved_root=source_root, + weight_files=(weight_path,), + use_safetensors=True, + ) + + first = _identity(model, source) + first_metadata = first.identity.source.metadata.to_value() + assert isinstance(first_metadata, dict) + first_sources = first_metadata["sources"] + assert isinstance(first_sources, list) + assert first_sources[0]["files"][0]["content_id"].startswith("sha256:") + + blob_path.write_bytes(b"b" * 32) + second = _identity(model, source) + + assert second.identity != first.identity + + +def test_hf_blob_shortcut_requires_validated_snapshot_topology(tmp_path: Path) -> None: + model = _TinyPipeline() + model_id = "test-org/tiny-diffusion" + repo_root = tmp_path / "models--test-org--tiny-diffusion" + blobs_root = repo_root / "blobs" + revision = "1" * 40 + snapshot_root = repo_root / "snapshots" / revision + blobs_root.mkdir(parents=True) + snapshot_root.mkdir(parents=True) + blob_name = "2" * 64 + blob_path = blobs_root / blob_name + blob_path.write_bytes(b"trusted-hf-blob") + weight_path = snapshot_root / "model.safetensors" + weight_path.symlink_to(Path("..") / ".." / "blobs" / blob_name) + source = PreparedWeightSource( + model_or_path=model_id, + subfolder=None, + requested_revision="main", + prefix="transformer.", + resolved_root=snapshot_root, + weight_files=(weight_path,), + use_safetensors=True, + source_kind=WeightSourceKind.HUGGING_FACE_HUB, + ) + + context = _identity(model, source) + metadata = context.identity.source.metadata.to_value() + assert isinstance(metadata, dict) + sources = metadata["sources"] + assert isinstance(sources, list) + assert sources[0]["files"][0]["content_id"] == f"immutable-blob:{blob_name}" + + +def test_identity_uses_resolved_revision_and_exact_semantics(tmp_path: Path) -> None: + model = _TinyPipeline() + commit = "0123456789abcdef0123456789abcdef01234567" + source = _prepared_source(tmp_path, requested_revision="main", resolved_revision=commit) + by_alias = _identity(model, source) + by_commit = _identity(model, dataclasses.replace(source, requested_revision=commit)) + + assert by_alias.identity == by_commit.identity + + base = _request() + variants = ( + dataclasses.replace( + base, + loader=dataclasses.replace(base.loader, model_config_fingerprint="test-model-config-v2"), + ), + dataclasses.replace( + base, + loader=dataclasses.replace( + base.loader, + implementation=dataclasses.replace(base.loader.implementation, version="2"), + ), + ), + dataclasses.replace( + base, + parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0), + ), + dataclasses.replace( + base, + parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1), + ), + dataclasses.replace( + base, + parallel=FinalLayoutParallelIdentity( + sequence_parallel_size=2, + ulysses_degree=2, + ulysses_mode="strict", + ), + ), + ) + base_identity = _identity(model, source, request=base).identity + assert all(_identity(model, source, request=variant).identity != base_identity for variant in variants) + + model_v2 = _TinyPipeline() + model_v2.transformer.host_weight_restore_contract = dataclasses.replace( + model_v2.transformer.host_weight_restore_contract, + version="2", + ) + assert _identity(model_v2, source).identity != base_identity + + changed_source = _prepared_source( + tmp_path, + directory="changed", + content=b"different-canonical-source", + ) + assert _identity(model, changed_source).identity != base_identity + source_with_adapter = dataclasses.replace( + source, + checkpoint_adapter=ImplementationIdentity( + implementation_id="test-checkpoint-adapter", + version="1", + fingerprint="test-checkpoint-adapter-v1", + ), + ) + assert _identity(model, source_with_adapter).identity != base_identity + + abi_v2 = CanonicalJson.from_value( + { + "artifact_identity": "diffusion-final-layout-identity-v1", + "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + "producer": "finalized-bf16-tensor-writer-v1", + "representation_policy": "bf16-with-preserved-fp32-v1", + "restorer": "exact-final-layout-tensor-rebind-v2", + "source_identity": "prepared-diffusion-weight-source-v2", + "tensor_contract": "complete-strided-tensor-ownership-v1", + } + ) + with pytest.raises(ValueError, match="fingerprint differs"): + dataclasses.replace(FINAL_LAYOUT_BF16_SPEC, implementation_abi=abi_v2) + spec_v2 = dataclasses.replace( + FINAL_LAYOUT_BF16_SPEC, + producer=dataclasses.replace( + FINAL_LAYOUT_BF16_SPEC.producer, + version="2", + implementation_fingerprint=implementation_abi_fingerprint(abi_v2), + ), + implementation_abi=abi_v2, + ) + identity_v2 = _identity( + model, + source, + policy=_PolicyWithSpec(FINAL_LAYOUT_BF16_POLICY, spec_v2), + ).identity + assert identity_v2.representation == base_identity.representation + assert identity_v2.key != base_identity.key + + +@pytest.mark.parametrize( + "parallel", + [ + FinalLayoutParallelIdentity(), + FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0), + FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1), + FinalLayoutParallelIdentity(sequence_parallel_size=2, ulysses_degree=2), + FinalLayoutParallelIdentity(sequence_parallel_size=2, ring_degree=2), + FinalLayoutParallelIdentity( + tensor_parallel_size=2, + tensor_parallel_rank=1, + sequence_parallel_size=2, + ulysses_degree=2, + ), + ], +) +def test_parallel_identity_matrix_is_exact_for_tp_and_sp(tmp_path: Path, parallel: FinalLayoutParallelIdentity) -> None: + """TP rank/size and SP backend are identity coordinates for HWR artifacts.""" + model = _TinyPipeline() + source = _prepared_source(tmp_path) + context = _identity(model, source, request=_request(parallel=parallel)) + layout = context.identity.layout + + assert layout.tensor_parallel_size == parallel.tensor_parallel_size + assert layout.tensor_parallel_rank == parallel.tensor_parallel_rank + assert layout.sequence_parallel_size == parallel.sequence_parallel_size + assert layout.sequence_parallel_backend == parallel.sequence_parallel_backend + assert "data_parallel" not in context.identity.layout.metadata.to_value().get("parallel", {}) + + +@pytest.mark.parametrize( + ("semantic_request", "message"), + [ + (_request(load_format="diffusers"), "load_format='default'"), + ( + _request(parallel=FinalLayoutParallelIdentity(pipeline_parallel_size=2)), + "pipeline-parallel", + ), + ( + _request(parallel=FinalLayoutParallelIdentity(cfg_parallel_size=2)), + "CFG-parallel", + ), + (_request(parallel=FinalLayoutParallelIdentity(use_hsdp=True)), "HSDP"), + ( + _request(parallel=FinalLayoutParallelIdentity(enable_expert_parallel=True)), + "expert-parallel", + ), + ( + _request(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256")), + "LoRA", + ), + ], +) +def test_bf16_policy_fails_closed_for_unsupported_semantics( + tmp_path: Path, + semantic_request: FinalLayoutRequest, + message: str, +) -> None: + model = _TinyPipeline() + with pytest.raises(ValueError, match=message): + _identity(model, _prepared_source(tmp_path), request=semantic_request) + + +def test_tensor_ownership_is_complete_mixed_precision_and_alias_free(tmp_path: Path) -> None: + model = _TinyPipeline() + records = collect_final_layout_targets( + model, + _dit_modules(model), + policy=FINAL_LAYOUT_BF16_POLICY, + require_materialized=True, + ) + by_name = {record.name: record for record in records} + + assert set(by_name) == { + "transformer.fp32_gain", + "transformer.proj.bias", + "transformer.proj.weight", + "transformer.scale", + } + assert by_name["transformer.proj.weight"].tensor.dtype is torch.bfloat16 + assert by_name["transformer.fp32_gain"].tensor.dtype is torch.float32 + assert by_name["transformer.scale"].role == "persistent_buffer" + + model.transformer.register_parameter("alias", model.transformer.proj.weight) + with pytest.raises(FinalLayoutContractError, match="aliases tensor object"): + _identity(model, _prepared_source(tmp_path)) + + +def test_identity_requires_explicit_model_restore_contract(tmp_path: Path) -> None: + model = _TinyPipeline() + model.transformer.host_weight_restore_contract = "unsupported-contract" + + with pytest.raises(FinalLayoutContractError, match="does not declare"): + _identity(model, _prepared_source(tmp_path)) + + +def test_publication_is_warm_only_and_restore_is_transactional( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = _prepared_source(tmp_path) + cold_model = _TinyPipeline() + _fill_final_weights(cold_model) + context = _identity(cold_model, source) + runtime = _runtime(tmp_path / "store") + producer = FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model), max_shard_bytes=16) + assert producer.spec.lookup_phase is LookupPhase.POST_LOAD_ONLY + assert producer.spec.source_mode is ProductionSourceMode.FINALIZED_MODEL + assert producer.spec.coordination_scope is CoordinationScope.SINGLE_PROCESS + + miss = runtime.resolve(context.identity) + assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + cold_pointers = _pointer_snapshot(cold_model) + cold_values = {name: tensor.detach().clone() for name, tensor in cold_model.named_parameters()} + + publication = runtime.publish_after_load(context.identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED + assert publication.failure is None + assert _pointer_snapshot(cold_model) == cold_pointers + assert all(torch.equal(dict(cold_model.named_parameters())[name], value) for name, value in cold_values.items()) + assert cold_model.transformer.restore_validations == 1 + + def unexpected_produce(_writer: object) -> object: + raise AssertionError("already-present publication must not invoke the producer") + + monkeypatch.setattr(producer, "produce", unexpected_produce) + already_present = runtime.publish_after_load(context.identity, producer=producer) + assert already_present.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT + + warm_model = _TinyPipeline(dit_device="meta") + warm_context = _identity(warm_model, source) + assert warm_context.identity == context.identity + hit = runtime.resolve(warm_context.identity) + assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT + assert hit.lease is not None + + before_plan = _pointer_snapshot(warm_model) + text_weight = warm_model.text_encoder.weight + vae_gain = warm_model.vae.gain + plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease) + assert _pointer_snapshot(warm_model) == before_plan + assert warm_model.transformer.restore_validations == 0 + + plan.commit() + with pytest.raises(RuntimeError, match="already committed"): + plan.commit() + + expected = { + "transformer.proj.weight": torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3), + "transformer.proj.bias": torch.tensor([3.0, 4.0], dtype=torch.bfloat16), + "transformer.fp32_gain": torch.tensor([6.5]), + "transformer.scale": torch.tensor([2.5]), + } + restored = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers()) + assert all(torch.equal(restored[name], value) for name, value in expected.items()) + assert all(not restored[name].is_meta for name in warm_context.tensor_names) + assert warm_model.transformer.restore_validations == 1 + assert warm_model.text_encoder.weight is text_weight + assert warm_model.vae.gain is vae_gain + assert warm_model.transformer.derived.item() == 7.0 + for name in warm_context.tensor_names: + assert restored[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr() + + del restored, warm_model, plan + gc.collect() + hit.lease.close() + assert isinstance(runtime.store, FilesystemHostWeightStore) + assert runtime.store.cleanup(context.identity) is None + + +def test_independent_publications_are_content_deterministic(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + model = _TinyPipeline() + _fill_final_weights(model) + context = _identity(model, source) + first_runtime = _runtime(tmp_path / "first-store") + second_runtime = _runtime(tmp_path / "second-store") + + first = first_runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=16), + ) + second = second_runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=1024), + ) + assert first.outcome is PostLoadPublicationOutcome.PUBLISHED + assert second.outcome is PostLoadPublicationOutcome.PUBLISHED + + first_hit = first_runtime.resolve(context.identity) + second_hit = second_runtime.resolve(context.identity) + assert first_hit.lease is not None and second_hit.lease is not None + assert first_hit.lease.manifest.artifact_content_sha256 == second_hit.lease.manifest.artifact_content_sha256 + assert first_hit.lease.manifest.format_metadata == second_hit.lease.manifest.format_metadata + assert {entry.name: entry.sha256 for entry in first_hit.lease.manifest.tensors} == { + entry.name: entry.sha256 for entry in second_hit.lease.manifest.tensors + } + first_hit.lease.close() + second_hit.lease.close() + + +def test_source_replacement_prevents_publication(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + model = _TinyPipeline() + _fill_final_weights(model) + context = _identity(model, source) + runtime = _runtime(tmp_path / "store") + source.weight_files[0].write_bytes(b"replacement-source") + + publication = runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, model, _dit_modules(model)), + ) + + assert publication.outcome is PostLoadPublicationOutcome.FAILED + assert publication.failure is not None + assert publication.failure.stage is ResolutionStage.CANONICAL_LOADING + assert publication.failure.code is FailureCode.CANONICAL_SOURCE_FAILED + assert publication.failure.details.to_value() == {"final_layout_code": "source_changed"} + assert not context.sources_unchanged() + assert runtime.resolve(context.identity).report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + + +def test_source_replacement_between_restore_plan_and_commit_causes_no_mutation(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + cold_model = _TinyPipeline() + _fill_final_weights(cold_model) + context = _identity(cold_model, source) + runtime = _runtime(tmp_path / "store") + assert ( + runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)), + ).outcome + is PostLoadPublicationOutcome.PUBLISHED + ) + hit = runtime.resolve(context.identity) + assert hit.lease is not None + + warm_model = _TinyPipeline() + warm_context = _identity(warm_model, source) + plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease) + before = _pointer_snapshot(warm_model) + source.weight_files[0].write_bytes(b"source-changed-before-commit") + + with pytest.raises(FinalLayoutContractError, match="canonical source changed"): + plan.commit() + assert _pointer_snapshot(warm_model) == before + del plan + gc.collect() + hit.lease.close() + + +def test_restore_rejects_wrong_identity_and_coverage_without_mutation(tmp_path: Path) -> None: + source = _prepared_source(tmp_path) + cold_model = _TinyPipeline() + _fill_final_weights(cold_model) + context = _identity(cold_model, source) + runtime = _runtime(tmp_path / "store") + assert ( + runtime.publish_after_load( + context.identity, + producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)), + ).outcome + is PostLoadPublicationOutcome.PUBLISHED + ) + hit = runtime.resolve(context.identity) + assert hit.lease is not None + + alternate = _AlternateTinyPipeline() + _fill_final_weights(alternate) + with pytest.raises(ValueError, match="producer model implementation differs"): + FinalLayoutBF16Producer(context, alternate, _dit_modules(alternate)) + with pytest.raises(ValueError, match="restore model implementation differs"): + FinalLayoutTensorRestorer(context).plan_restore(alternate, hit.lease) + + model = _TinyPipeline() + wrong_context = _identity( + model, + source, + request=_request( + parallel=FinalLayoutParallelIdentity( + tensor_parallel_size=2, + tensor_parallel_rank=1, + ) + ), + ) + with pytest.raises(ValueError, match="semantic identity differs"): + FinalLayoutTensorRestorer(wrong_context).plan_restore(model, hit.lease) + + exact_context = _identity(model, source) + before = _pointer_snapshot(model) + model.transformer.register_buffer("new_persistent_state", torch.tensor([1.0])) + with pytest.raises(FinalLayoutContractError, match="tensor contract differs"): + FinalLayoutTensorRestorer(exact_context).plan_restore(model, hit.lease) + after = _pointer_snapshot(model) + assert all(after[name] == value for name, value in before.items()) + hit.lease.close() + + +def test_reduced_minimax_h3_satisfies_real_tensor_ownership_contract( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "QKVParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "RowParallelLinear", _FakeH3Linear) + monkeypatch.setattr(h3, "Attention", _FakeH3Attention) + monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1) + + transformer = h3.MiniMaxH3DiTModel(_small_h3_config(), quant_config=None) + pipeline = _H3Pipeline(transformer) + source = _prepared_source(tmp_path, directory="minimax-h3-source") + context = _identity(pipeline, source) + targets = collect_final_layout_targets( + pipeline, + (("transformer", transformer),), + policy=FINAL_LAYOUT_BF16_POLICY, + require_materialized=True, + ) + by_name = {target.name: target for target in targets} + + assert context.model_contracts == (MiniMaxH3DiTModel.host_weight_restore_contract,) + assert "transformer.video_patch_proj.weight" in by_name + assert "transformer.blocks.0.attn.qkv_proj.weight" in by_name + assert "transformer.rope.inv_freq" in by_name + assert by_name["transformer.video_patch_proj.weight"].tensor.dtype is torch.float32 + assert by_name["transformer.blocks.0.attn.qkv_proj.weight"].tensor.dtype is torch.bfloat16 + assert by_name["transformer.rope.inv_freq"].role == "persistent_buffer" + transformer.validate_restored_host_weights() + + +def test_minimax_h3_declares_the_final_layout_restore_contract() -> None: + contract = MiniMaxH3DiTModel.host_weight_restore_contract + assert isinstance(contract, FinalLayoutModelContract) + assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA + assert contract.implementation_id == "minimax-h3-dit" + assert callable(MiniMaxH3DiTModel.validate_restored_host_weights) + assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__ + + +def test_reduced_flux2_klein_declares_complete_final_layout_contract( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import ( + Flux2Transformer2DModel, + ) + + transformer = _reduced_flux2_klein(monkeypatch) + pipeline = _Flux2KleinPipeline(transformer) + source = _prepared_source(tmp_path, directory="flux2-klein-source") + context = _identity(pipeline, source) + targets = collect_final_layout_targets( + pipeline, + (("transformer", transformer),), + policy=FINAL_LAYOUT_BF16_POLICY, + require_materialized=True, + ) + target_names = {target.name for target in targets} + + contract = Flux2Transformer2DModel.host_weight_restore_contract + assert isinstance(contract, FinalLayoutModelContract) + assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA + assert contract.implementation_id == "flux2-klein-dit" + assert contract.version == "1" + assert context.model_contracts == (contract,) + assert transformer.stacked_params_mapping == [ + (".to_qkv.", ".to_q.", "q"), + (".to_qkv.", ".to_k.", "k"), + (".to_qkv.", ".to_v.", "v"), + (".add_kv_proj", ".add_q_proj", "q"), + (".add_kv_proj", ".add_k_proj", "k"), + (".add_kv_proj", ".add_v_proj", "v"), + ] + loaded = transformer.load_weights( + [ + ("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)), + ("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)), + ( + "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight", + torch.tensor([3.0], dtype=torch.bfloat16), + ), + ] + ) + assert loaded == { + "transformer_blocks.0.attn.to_qkv.weight", + "transformer_blocks.0.attn.add_kv_proj.weight", + "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight", + } + assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names) + assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names) + assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names + assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names + assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names + assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names + assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names + assert all("text_encoder" not in name and "vae" not in name for name in target_names) + transformer.validate_restored_host_weights() + + +def test_flux2_klein_restore_validator_rejects_non_ready_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + transformer = _reduced_flux2_klein(monkeypatch) + + mapping = transformer.stacked_params_mapping + transformer.stacked_params_mapping = None + with pytest.raises(ValueError, match="packed QKV mapping"): + transformer.validate_restored_host_weights() + transformer.stacked_params_mapping = mapping + + weight = transformer.x_embedder.weight + weight.data = weight.data.float() + with pytest.raises(ValueError, match="must stay bf16"): + transformer.validate_restored_host_weights() + weight.data = weight.data.to(torch.bfloat16) + + weight.data = weight.data.T + with pytest.raises(ValueError, match="contiguous strided layout"): + transformer.validate_restored_host_weights() + weight.data = weight.data.contiguous() + + attention = transformer.transformer_blocks[0].attn.attn + attention._non_persistent_buffers_set.add("beta") + with pytest.raises(ValueError, match="must be persistent"): + transformer.validate_restored_host_weights() + attention._non_persistent_buffers_set.remove("beta") + + attention.beta.data = attention.beta.data.float() + with pytest.raises(ValueError, match="scalar bf16"): + transformer.validate_restored_host_weights() + + +def test_flux2_klein_publication_restore_fallback_and_teardown( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = _prepared_source(tmp_path, directory="flux2-klein-source") + cold_transformer = _reduced_flux2_klein(monkeypatch) + cold_model = _Flux2KleinPipeline(cold_transformer) + with torch.no_grad(): + for index, parameter in enumerate(cold_transformer.parameters(), start=1): + parameter.fill_(index) + for name, buffer in cold_transformer.named_buffers(): + if name.endswith((".beta", ".eps")): + buffer.fill_(0.5 if name.endswith(".beta") else -1e-6) + + context = _identity(cold_model, source) + runtime = _runtime(tmp_path / "flux2-klein-store") + miss = runtime.resolve(context.identity) + assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + + producer = FinalLayoutBF16Producer( + context, + cold_model, + (("transformer", cold_transformer),), + max_shard_bytes=64, + ) + publication = runtime.publish_after_load(context.identity, producer=producer) + assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED + + warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta") + warm_model = _Flux2KleinPipeline(warm_transformer) + warm_context = _identity(warm_model, source) + assert warm_context.identity == context.identity + hit = runtime.resolve(warm_context.identity) + assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT + assert hit.lease is not None + + plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease) + assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer.")) + plan.commit() + warm_transformer.validate_restored_host_weights() + + cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers()) + warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers()) + for name in warm_context.tensor_names: + assert torch.equal(warm_tensors[name], cold_tensors[name]) + assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr() + assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping + + del warm_tensors, warm_model, warm_transformer, plan + gc.collect() + hit.lease.close() + assert hit.lease.closed + assert isinstance(runtime.store, FilesystemHostWeightStore) + assert runtime.store.cleanup(context.identity) is None diff --git a/tests/diffusion/model_loader/test_hub_prefetch.py b/tests/diffusion/model_loader/test_hub_prefetch.py new file mode 100644 index 000000000..39eb2b277 --- /dev/null +++ b/tests/diffusion/model_loader/test_hub_prefetch.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import contextlib + +import huggingface_hub +import pytest + +from vllm_omni.diffusion.model_loader import hub_prefetch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_prefetch_subfolders_propagates_revision(monkeypatch): + calls = [] + + def fake_snapshot_download(**kwargs): + calls.append(kwargs) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext()) + + hub_prefetch.prefetch_subfolders( + "org/model", + ["transformer"], + revision="pinned-revision", + ) + + assert calls == [ + { + "repo_id": "org/model", + "revision": "pinned-revision", + "allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"], + } + ] + + +def test_from_pretrained_retry_preserves_revision(monkeypatch): + factory_calls = [] + prefetch_calls = [] + + def fake_factory(model, **kwargs): + factory_calls.append((model, kwargs)) + if len(factory_calls) == 1: + raise OSError("partial cache") + return "loaded" + + def fake_prefetch(model, subfolders, **kwargs): + prefetch_calls.append((model, tuple(subfolders), kwargs)) + + monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch) + monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None) + + result = hub_prefetch.from_pretrained_with_prefetch( + fake_factory, + "org/model", + subfolder="transformer", + prefetch_list=["transformer", "vae"], + revision="pinned-revision", + max_attempts=2, + ) + + assert result == "loaded" + assert len(factory_calls) == 2 + assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls) + assert prefetch_calls == [ + ( + "org/model", + ("transformer", "vae"), + {"local_files_only": False, "revision": "pinned-revision"}, + ) + ] diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py index d38be97fa..999fdbb41 100644 --- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py +++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """L1 unit tests for the native Boogu-Image pipeline. @@ -133,6 +133,51 @@ def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies): assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"] +def test_constructor_forwards_revision_to_all_component_loaders(mock_dependencies, mocker): + from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import ( + BooguImagePipeline, + ) + + revision = "hotfix-1k-20260708" + prefetch = mocker.patch(f"{_MODULE}.prefetch_subfolders") + scheduler_loader = mocker.patch( + f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained", + return_value=mock_dependencies["scheduler"], + ) + mllm_loader = mocker.patch( + f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained", + return_value=mock_dependencies["mllm_wrapper"], + ) + processor_loader = mocker.patch( + f"{_MODULE}.Qwen3VLProcessor.from_pretrained", + return_value=mock_dependencies["processor"], + ) + vae_loader = mocker.patch( + f"{_MODULE}.AutoencoderKL.from_pretrained", + return_value=mock_dependencies["vae"], + ) + od_config = OmniDiffusionConfig( + model="dummy-boogu", + revision=revision, + tf_model_config=TransformerConfig(params={}), + dtype=torch.float32, + num_gpus=1, + ) + + pipeline = BooguImagePipeline(od_config=od_config) + + (source,) = pipeline.weights_sources + assert source.revision == revision + prefetch.assert_called_once_with( + "dummy-boogu", + ["scheduler", "vae", "mllm", "processor"], + local_files_only=True, + revision=revision, + ) + for loader in (scheduler_loader, mllm_loader, processor_loader, vae_loader): + assert loader.call_args.kwargs["revision"] == revision + + def test_constructor_weights_sources(boogu_pipeline): (source,) = boogu_pipeline.weights_sources assert source.model_or_path == "dummy-boogu" diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 796de06f1..f2527bd99 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -1019,11 +1019,19 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": True})) assert cfg is not None assert list(cfg.hints) == ["edge"] + assert cfg.hints["edge"].control_weight == 1.0 + assert cfg.normalized_control_weights == [1.0] assert cfg.guidance_scale == 3.0 assert cfg.control_guidance == 1.5 assert cfg.flow_shift == 10.0 assert cfg.num_video_frames_per_chunk == 93 assert cfg.share_vision_temporal_positions is True + assert cfg.emphasize_control_in_prompt is True + prompt_ablation_cfg = transfer.resolve_transfer_config( + make_sampling_params(extra_args={"edge": True, "emphasize_control_in_prompt": False}) + ) + assert prompt_ablation_cfg is not None + assert prompt_ablation_cfg.emphasize_control_in_prompt is False # fps omitted (no fps/frame_rate on the sampling params) -> wsm preset default (10) applies. defaulted_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"wsm": True})) assert defaulted_fps_cfg is not None @@ -1087,11 +1095,46 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest ) additional = preprocess(request).prompt["additional_information"] assert (request.sampling_params.height, request.sampling_params.width) == (192, 320) + assert request.sampling_params.extra_args["_cosmos3_transfer_requested_size"] == {"width": 32, "height": 16} assert tuple(additional["preprocessed_transfer_video"].shape) == (1, 3, 4, 192, 320) assert additional["transfer_input_fps"] == 12.5 assert "preprocessed_video" not in additional +def test_transfer_control_weight_validation_and_normalization() -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + + single = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": 0.4}})) + assert single is not None + assert single.hints["edge"].control_weight == 0.4 + assert single.normalized_control_weights == [1.0] + + multiple = transfer.resolve_transfer_config( + make_sampling_params( + extra_args={ + "edge": {"control_weight": 1.0}, + "depth": {"control_weight": 3.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)}, + } + ) + ) + assert multiple is not None + assert multiple.normalized_control_weights == [0.25, 0.75] + + with pytest.raises(ValueError, match="Unsupported.*weight"): + transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"weight": 0.5}})) + with pytest.raises(ValueError, match="finite and non-negative"): + transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": -0.5}})) + with pytest.raises(ValueError, match="positive sum"): + transfer.resolve_transfer_config( + make_sampling_params( + extra_args={ + "edge": {"control_weight": 0.0}, + "depth": {"control_weight": 0.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)}, + } + ) + ) + + def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None: """When fps and frame_rate differ, transfer must pick frame_rate -- the same precedence as OmniDiffusionSamplingParams.resolved_frame_rate. Uses the real @@ -1352,6 +1395,63 @@ def test_format_and_tokenize_prompts_applies_video_templates_and_system_override ] +def test_format_and_tokenize_prompts_applies_transfer_prompt_contract(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE, + COSMOS3_TRANSFER_SYSTEM_PROMPT, + ) + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + directive = COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names="edge") + + pipeline._format_and_tokenize_prompts( + "A robot", + "bad", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace( + extra_args={ + "use_resolution_template": True, + "system_prompt": "request-level system prompt", + } + ), + use_system_prompt=True, + is_t2i=False, + system_prompt=COSMOS3_TRANSFER_SYSTEM_PROMPT, + prompt_suffix=directive, + use_duration_template=True, + use_resolution_template=True, + negative_metadata_mode="same", + ) + + metadata = "The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution." + assert calls[0]["text"] == f"A robot. {metadata} {directive}" + assert calls[1]["text"] == f"bad. {metadata}" + assert all(call["use_system_prompt"] is True for call in calls) + assert all(call["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT for call in calls) + + +def test_format_and_tokenize_prompts_rejects_unknown_negative_metadata_mode(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + + with pytest.raises(ValueError, match="negative_metadata_mode"): + pipeline._format_and_tokenize_prompts( + "A robot", + "bad", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + negative_metadata_mode="unexpected", + ) + + def test_format_and_tokenize_prompts_uses_image_templates_for_t2i(make_cosmos3_pipeline) -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_T2I_SYSTEM_PROMPT @@ -1421,6 +1521,140 @@ def test_format_and_tokenize_prompts_rewrites_json_object_metadata(make_cosmos3_ ) +def test_format_and_tokenize_prompts_transfer_aspect_ratio_override(make_cosmos3_pipeline) -> None: + import json + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + '{"caption": "A robot", "aspect_ratio": "4,3"}', + "", + num_frames=48, + frame_rate=24, + height=480, + width=832, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={"aspect_ratio": "4:3"}), + is_t2i=False, + aspect_ratio_override="16,9", + ) + + assert json.loads(calls[0]["text"])["aspect_ratio"] == "16,9" + + +@pytest.mark.parametrize("rank_zero", [True, False]) +def test_format_and_tokenize_prompts_corrects_conflicting_aspect_ratio_metadata_on_every_rank( + make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, + mocker, + rank_zero: bool, +) -> None: + import json + + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + monkeypatch.setattr(pipeline_cosmos3, "_is_rank_zero", lambda: rank_zero) + warning = mocker.patch.object(pipeline_cosmos3.logger, "warning") + + pipeline._format_and_tokenize_prompts( + '{"caption": "A robot", "aspect_ratio": "4,3"}', + "", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + is_t2i=False, + ) + + assert json.loads(calls[0]["text"])["aspect_ratio"] == "16,9" + if rank_zero: + rendered_warning = warning.call_args.args[0] % warning.call_args.args[1:] + assert "JSON prompt aspect_ratio='4,3' conflicts with the generated 1280x720 canvas" in rendered_warning + else: + warning.assert_not_called() + + +@pytest.mark.parametrize( + ("width", "height", "aspect_ratio"), + [ + (1104, 816, "4,3"), + (832, 468, "16,9"), + ], +) +def test_format_and_tokenize_prompts_keeps_nearest_canonical_aspect_ratio( + make_cosmos3_pipeline, + mocker, + width: int, + height: int, + aspect_ratio: str, +) -> None: + import json + + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + warning = mocker.patch.object(pipeline_cosmos3.logger, "warning") + + pipeline._format_and_tokenize_prompts( + json.dumps({"caption": "A robot", "aspect_ratio": aspect_ratio}), + "", + num_frames=48, + frame_rate=24, + height=height, + width=width, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + is_t2i=False, + ) + + assert json.loads(calls[0]["text"])["aspect_ratio"] == aspect_ratio + warning.assert_not_called() + + +def test_transfer_bucket_selection_warns_on_size_and_aspect_ratio_conflicts( + make_cosmos3_pipeline, + mocker, +) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + warning = mocker.patch.object(pipeline_cosmos3.logger, "warning") + sp = make_sampling_params( + width=832, + height=480, + extra_args={ + "resolution": "480", + "aspect_ratio": "4:3", + "_cosmos3_transfer_requested_size": {"width": 1280, "height": 720}, + }, + ) + + height, width, aspect_ratio = pipeline._transfer_bucket_size(sp, (480, 832)) + assert (height, width, aspect_ratio) == (480, 832, "16,9") + + pipeline._warn_transfer_bucket_conflicts( + sp, + "A transfer prompt", + source_hw=(480, 832), + height=height, + width=width, + aspect_ratio=aspect_ratio, + ) + + rendered_warnings = [call.args[0] % call.args[1:] for call in warning.call_args_list] + assert any("ignores requested size=1280x720 (WxH)" in message for message in rendered_warnings) + assert any( + "requested aspect_ratio='4:3' conflicts with the control-selected 16,9 bucket" in message + for message in rendered_warnings + ) + + def test_format_and_tokenize_prompts_removes_video_metadata_from_t2i_json(make_cosmos3_pipeline) -> None: import json @@ -1740,6 +1974,7 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_ control_guidance=1.5, control_guidance_interval=None, control_latents=[torch.zeros_like(latents)], + control_weights=[1.0], shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, velocity_mask=velocity_mask, condition_latents=torch.zeros_like(latents), @@ -1750,6 +1985,9 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_ (2, False), (1, True), ] + assert pipeline.transformer.calls[0]["kwargs"]["control_weights"] == [1.0] + assert "control_weights" not in pipeline.transformer.calls[1]["kwargs"] + assert pipeline.transformer.calls[2]["kwargs"]["control_weights"] == [1.0] torch.testing.assert_close(result, torch.full_like(latents, 254.0)) @@ -1857,14 +2095,36 @@ def test_diffuse_transfer_interval_switches_branch_counts(make_cosmos3_pipeline, torch.testing.assert_close(result, torch.full_like(latents, 508.0)) -@pytest.mark.parametrize(("hint_key", "expected_fps"), [("edge", 8.0), ("wsm", 10.0)]) -def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint_key: str, expected_fps: float) -> None: +@pytest.mark.parametrize( + ("hint_key", "emphasize_control", "expected_fps", "negative_prompt"), + [ + ("edge", None, 8.0, None), + ("wsm", None, 10.0, ""), + ("edge", False, 8.0, "custom negative"), + ], +) +def test_forward_transfer_uses_transfer_prompt_contract_and_source_fps_except_wsm( + make_cosmos3_pipeline, + hint_key: str, + emphasize_control: bool | None, + expected_fps: float, + negative_prompt: str | None, +) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE, + COSMOS3_TRANSFER_SYSTEM_PROMPT, + ) + pipeline = make_cosmos3_pipeline() captured: dict[str, Any] = {} def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): - del prompt, negative_prompt, num_frames, height, width, args, kwargs + del num_frames, height, width, args + captured["prompt"] = prompt + captured["negative_prompt"] = negative_prompt captured["format_frame_rate"] = frame_rate + captured["format_kwargs"] = kwargs + captured["format_aspect_ratio"] = kwargs["aspect_ratio_override"] return _ids(2), _mask(), _ids(1), _mask() def fake_encode(video: torch.Tensor) -> torch.Tensor: @@ -1887,7 +2147,7 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint captured["shared_kwargs"] = kwargs["shared_kwargs"] return kwargs["latents"] - pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1") pipeline._format_and_tokenize_prompts = fake_format pipeline._encode_video_tensor = fake_encode pipeline._prepare_transfer_latents = fake_prepare @@ -1901,33 +2161,49 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta") control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8) + extra_args = { + hint_key: {"control": control}, + "max_frames": 5, + "num_video_frames_per_chunk": 5, + "show_control_condition": True, + } + if emphasize_control is not None: + extra_args["emphasize_control_in_prompt"] = emphasize_control + prompt_data = { + "prompt": "transfer", + "modalities": ["video"], + "additional_information": { + "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16), + "transfer_input_fps": 8.0, + }, + } + if negative_prompt is not None: + prompt_data["negative_prompt"] = negative_prompt request = SimpleNamespace( - prompts=[ - { - "prompt": "transfer", - "modalities": ["video"], - "additional_information": { - "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16), - "transfer_input_fps": 8.0, - }, - } - ], + prompts=[prompt_data], sampling_params=make_sampling_params( height=16, width=16, # fps omitted (no frame_rate) -> non-wsm uses the source video fps (8), wsm uses its preset (10). - extra_args={ - hint_key: {"control": control}, - "max_frames": 5, - "num_video_frames_per_chunk": 5, - "show_control_condition": True, - }, + extra_args=extra_args, ), ) output = pipeline.forward(request) assert captured["format_frame_rate"] == expected_fps + assert captured["format_aspect_ratio"] == "1,1" + expected_negative_prompt = "" if negative_prompt is None else negative_prompt + assert captured["negative_prompt"] == expected_negative_prompt + assert captured["format_kwargs"]["use_system_prompt"] is True + assert captured["format_kwargs"]["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT + assert captured["format_kwargs"]["use_duration_template"] is True + assert captured["format_kwargs"]["use_resolution_template"] is True + assert captured["format_kwargs"]["negative_metadata_mode"] == "same" + expected_suffix = ( + None if emphasize_control is False else COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names=hint_key) + ) + assert captured["format_kwargs"]["prompt_suffix"] == expected_suffix assert captured["shared_kwargs"]["fps"] == expected_fps assert captured["flow_shifts"] == [10.0] # Transfer applies the V2V flow shift when building its timestep schedule. @@ -1943,7 +2219,7 @@ def test_forward_transfer_runs_multichunk_overlap_path( pipeline = make_cosmos3_pipeline() captured: dict[str, Any] = {"targets": [], "conditional_frames": []} - pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1") pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask()) pipeline._set_flow_shift = lambda target, **_kwargs: captured.setdefault("flow_shifts", []).append(target) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 5dd8a875a..2bfa7b558 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -504,6 +504,35 @@ def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device( assert generator_param.device == accelerator_device +def test_multi_control_attention_weights_target_outputs() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention + + class ControlValueAttention: + def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + del k + # Each per-control K/V layout is [text, control_i, target]. + return v[:, 1:2].expand(q.shape[0], q.shape[1], -1, -1) + + attention = SimpleNamespace(multi_control_attn=ControlValueAttention()) + q = torch.zeros(1, 3, 1, 1) + k = torch.zeros_like(q) + v = torch.tensor([[[[10.0]], [[20.0]], [[30.0]]]]) + text_kv = torch.zeros(1, 1, 1, 1) + + output = Cosmos3CrossAttention._forward_multi_control( + attention, + q, + k, + v, + text_kv, + text_kv, + control_token_sizes=(1, 1), + control_weights=(0.25, 0.75), + ) + + torch.testing.assert_close(output, torch.tensor([[[10.0], [20.0], [17.5]]])) + + def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 @@ -521,9 +550,21 @@ def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatc video_shape=(1, 2, 2), fps=24.0, control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], + control_weights=[1.0, 3.0], ) assert tuple(output.shape) == tuple(hidden_states.shape) + with pytest.raises(ValueError, match="control_weights length"): + model( + hidden_states=hidden_states, + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 2, 2), + fps=24.0, + control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], + control_weights=[1.0], + ) with pytest.raises(ValueError, match="control latent shape"): model( hidden_states=hidden_states, diff --git a/tests/cosmos3/test_session_memory_equivalence.py b/tests/diffusion/models/cosmos3/test_session_memory_equivalence.py similarity index 100% rename from tests/cosmos3/test_session_memory_equivalence.py rename to tests/diffusion/models/cosmos3/test_session_memory_equivalence.py diff --git a/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py b/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py deleted file mode 100644 index c7e977ecf..000000000 --- a/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py +++ /dev/null @@ -1,27 +0,0 @@ -import pytest -import torch.nn as nn - -pytest.importorskip("dreamid_omni.modules.model") - -from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks(): - model = object.__new__(FusionModel) - nn.Module.__init__(model) - model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) - model.attn = nn.Identity() - - conditions = getattr(model, "_hsdp_shard_conditions", None) - - assert conditions is not None - assert len(conditions) == 1 - - matched = [] - for name, module in model.named_modules(): - if any(cond(name, module) for cond in conditions): - matched.append(name) - - assert matched == ["fused_blocks.0", "fused_blocks.1"] diff --git a/tests/diffusion/models/dreamzero/test_pipeline_state.py b/tests/diffusion/models/dreamzero/test_pipeline_state.py index 366876ef7..ff349d6af 100644 --- a/tests/diffusion/models/dreamzero/test_pipeline_state.py +++ b/tests/diffusion/models/dreamzero/test_pipeline_state.py @@ -6,7 +6,10 @@ from types import SimpleNamespace import pytest -from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline +from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import ( + DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION, + DreamZeroPipeline, +) from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -32,6 +35,37 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None: assert session_b.call_count == 3 +def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None: + pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline) + pipeline.transformer = SimpleNamespace( + frame_seqlen=16, + blocks=[ + SimpleNamespace( + self_attn=SimpleNamespace( + max_attention_size=96, + tp_num_heads=4, + ) + ) + ], + text_len=8, + model_type="t2v", + num_layers=2, + num_heads=4, + dim=256, + num_frame_per_block=4, + num_action_per_block=2, + num_state_per_block=1, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size", + lambda: 1, + ) + + spec = pipeline.ar_diffusion_kv_cache_spec() + + assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION + + def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None: pipeline = _empty_pipeline() diff --git a/tests/diffusion/models/flux2/test_flux2_quantization.py b/tests/diffusion/models/flux2/test_flux2_quantization.py new file mode 100644 index 000000000..3c461acf5 --- /dev/null +++ b/tests/diffusion/models/flux2/test_flux2_quantization.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Tests for FLUX.2-dev quantization routing.""" + +import pytest +from torch import nn + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _RecordingBlock(nn.Module): + def __init__(self, *, quant_config=None, prefix: str = "", **kwargs): + super().__init__() + self.quant_config = quant_config + self.prefix = prefix + + +def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch): + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + from vllm_omni.diffusion.models.flux2 import flux2_transformer + + monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock) + monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock) + + quant_config = Fp8Config( + is_checkpoint_fp8_serialized=False, + activation_scheme="dynamic", + ) + model = flux2_transformer.Flux2Transformer2DModel( + in_channels=4, + out_channels=4, + num_layers=2, + num_single_layers=3, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=16, + timestep_guidance_channels=8, + axes_dims_rope=(2, 2, 2, 2), + quant_config=quant_config, + ) + + assert [block.quant_config for block in model.transformer_blocks] == [quant_config, quant_config] + assert [block.prefix for block in model.transformer_blocks] == [ + "transformer_blocks.0", + "transformer_blocks.1", + ] + assert [block.quant_config for block in model.single_transformer_blocks] == [ + quant_config, + quant_config, + quant_config, + ] + assert [block.prefix for block in model.single_transformer_blocks] == [ + "single_transformer_blocks.0", + "single_transformer_blocks.1", + "single_transformer_blocks.2", + ] + + assert isinstance(model.x_embedder, nn.Linear) + assert isinstance(model.context_embedder, nn.Linear) + assert isinstance(model.proj_out, nn.Linear) diff --git a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py index 3900c0c94..ff9e0e110 100644 --- a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py +++ b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py @@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import ( Flux2Pipeline, _resolve_component_quant_config, ) +from vllm_omni.quantization import ComponentQuantizationConfig pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_resolve_component_quant_config_routes_components(): text_encoder_config = object() - component_config = Mock() - component_config.resolve.side_effect = { - "text_encoder": text_encoder_config, - "transformer": None, - }.get + component_config = ComponentQuantizationConfig( + component_configs={"text_encoder": text_encoder_config}, + default_config=None, + ) assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config assert _resolve_component_quant_config(component_config, "transformer") is None diff --git a/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py b/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py new file mode 100644 index 000000000..e5581cd77 --- /dev/null +++ b/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from cache_dit import ForwardPattern +from torch import nn + +from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend +from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import ( + HiDreamO1TextModel, +) +from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import ( + HiDreamO1ImagePipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _empty_text_model() -> HiDreamO1TextModel: + model = HiDreamO1TextModel.__new__(HiDreamO1TextModel) + nn.Module.__init__(model) + model.layers = nn.ModuleList([nn.Identity()]) + return model + + +def test_cache_dit_targets_the_decoder_module_that_owns_layers(): + text_model = _empty_text_model() + pipeline = SimpleNamespace( + _dit_modules=HiDreamO1ImagePipeline._dit_modules, + model=SimpleNamespace( + model=SimpleNamespace(language_model=text_model), + ), + ) + + with ( + patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") as block_adapter_cls, + patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") as cache_dit, + ): + backend = CacheDiTBackend() + backend.enable(pipeline) + backend.refresh(pipeline, num_inference_steps=50) + + adapter_config = HiDreamO1TextModel._cache_dit_adapter_config + assert isinstance(adapter_config, CacheDiTAdapterConfig) + assert adapter_config.block_forward_patterns == {"layers": ForwardPattern.Pattern_3} + assert adapter_config.has_separate_cfg is True + + block_adapter_cls.assert_called_once_with( + transformer=text_model, + blocks=[text_model.layers], + forward_pattern=[ForwardPattern.Pattern_3], + has_separate_cfg=True, + check_forward_pattern=True, + ) + cache_dit.enable_cache.assert_called_once() + enable_call = cache_dit.enable_cache.call_args + assert enable_call.args == (block_adapter_cls.return_value,) + assert enable_call.kwargs["cache_config"].Fn_compute_blocks == 1 + assert enable_call.kwargs["calibrator_config"] is None + cache_dit.refresh_context.assert_called_once_with( + text_model, + num_inference_steps=50, + verbose=True, + ) diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py index dfb383e87..b62fe950e 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace @@ -270,6 +270,17 @@ def test_grouped_denoise_allows_sdpa_attention_backend(): pipeline._ensure_grouped_attention_backend_supported(2) +def test_scheduler_paged_step_execution_is_rejected(): + pipeline = _pipeline() + pipeline.od_config.diffusion_kv_mode = hy3_module.DiffusionKVCacheMode.PAGED_SCHEDULER + state = _state("paged-step", 0) + + with pytest.raises(ValueError, match="request-level execution only"): + pipeline.prepare_encode(state) + with pytest.raises(ValueError, match="request-level execution only"): + pipeline.denoise_step(InputBatch.make_batch([state])) + + def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches(): pipeline = _pipeline() pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {}) diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py index 543b93a25..65fc0e67b 100644 --- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py +++ b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for ImageKVCacheManager. Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP, cross-request isolation. @@ -19,8 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] _TRANSFORMER_MODULE = "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer" -NUM_HEADS = 4 -NUM_KV_HEADS = 2 +NUM_HEADS = 32 +NUM_KV_HEADS = 8 HEAD_DIM = 16 IMAGE_TOKEN_LEN = 8 SUFFIX_TOKEN_LEN = 3 @@ -35,8 +35,23 @@ SCALING = 1.0 / math.sqrt(HEAD_DIM) class MockAttention(nn.Module): def __init__(self, num_heads, head_size, causal=False, softmax_scale=None, num_kv_heads=None, **kwargs): super().__init__() + self.paged_kv_active = False + self.calls = [] + self.paged_calls = [] + self.paged_metadata = [] + + def is_paged_kv_active(self): + return self.paged_kv_active def forward(self, query, key, value, attn_metadata=None, **kwargs): + self.calls.append((query, key, value)) + from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + + if is_forward_context_available() and get_forward_context().paged_kv_adapter is not None: + self.paged_calls.append((query, key, value)) + self.paged_metadata.append(attn_metadata) + if attn_metadata is not None and attn_metadata.joint_query is not None: + return torch.cat((attn_metadata.joint_query, query), dim=1) return query @@ -99,6 +114,7 @@ def _call_mgr( shard_image_size=None, gen_timestep_scatter_index=None, position_ids=None, + full_attn_spans=None, ): query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM) attn_mask = torch.zeros(bs, 1, seq_len, seq_len) @@ -115,9 +131,127 @@ def _call_mgr( shard_image_size=shard_image_size, gen_timestep_scatter_index=gen_timestep_scatter_index, position_ids=position_ids, + full_attn_spans=full_attn_spans, ) +def test_cache_manager_registers_attention_without_adding_dense_state() -> None: + mgr = _make_cache_mgr() + + assert isinstance(mgr, nn.Module) + assert dict(mgr.named_modules())["attn"] is mgr.attn + assert mgr.state_dict() == {} + + +def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None: + from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context + + mgr = _make_cache_mgr() + mgr.attn.paged_kv_active = True + bs = 2 + prefix_lens = torch.tensor([[2], [4]], dtype=torch.long) + first_q_len = 12 + first_k, first_v = _make_known_kv(bs * first_q_len, base=1.0) + full_attn_spans = [[(2, 12)], [(2, 12)]] + + with set_forward_context(), override_paged_kv_adapter(object()): + _call_mgr( + mgr, + bs, + q_len=first_q_len, + seq_len=first_q_len, + key_flat=first_k, + value_flat=first_v, + first_step=True, + gen_timestep_scatter_index=prefix_lens, + full_attn_spans=full_attn_spans, + ) + + assert mgr.image_kv_cache_map is None + assert mgr.image_kv_cache_lens is None + assert len(mgr.attn.paged_calls) == 1 + assert mgr.attn.paged_calls[0][0].shape == (bs, first_q_len, NUM_HEADS, HEAD_DIM) + assert mgr.attn.paged_calls[0][1].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM) + assert mgr.attn.paged_calls[0][2].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM) + assert mgr.attn.paged_metadata[0].full_attn_spans == full_attn_spans + + current_q_len = IMAGE_TOKEN_LEN + current_k, current_v = _make_known_kv(bs * current_q_len, base=50.0) + _call_mgr( + mgr, + bs, + q_len=current_q_len, + seq_len=prefix_lens[-1].item() + current_q_len, + key_flat=current_k, + value_flat=current_v, + position_ids=torch.arange(4, 4 + current_q_len).repeat(bs, 1), + full_attn_spans=full_attn_spans, + ) + + assert len(mgr.attn.paged_calls) == 2 + assert mgr.attn.paged_calls[1][0].shape == (bs, current_q_len, NUM_HEADS, HEAD_DIM) + assert mgr.attn.paged_calls[1][1].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM) + assert mgr.attn.paged_calls[1][2].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM) + assert mgr.attn.paged_metadata[1].full_attn_spans == full_attn_spans + + +def test_cache_manager_rejects_dense_imported_prefix_for_paged_attention() -> None: + from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context + + mgr = _make_cache_mgr() + mgr.attn.paged_kv_active = True + prefix_len = 4 + prefix_key, prefix_value = _make_known_kv(prefix_len, base=100.0) + mgr._injected_ar_kv = [(prefix_key, prefix_value)] + current_key, current_value = _make_known_kv(3, base=1.0) + + with set_forward_context(), override_paged_kv_adapter(object()): + with pytest.raises(NotImplementedError, match="does not support imported AR KV"): + _call_mgr( + mgr, + bs=1, + q_len=3, + seq_len=7, + key_flat=current_key, + value_flat=current_value, + first_step=True, + gen_timestep_scatter_index=torch.tensor([[0]]), + full_attn_spans=[[(4, 7)]], + ) + + +@pytest.mark.parametrize( + ("paged_kv_active", "expected_kv_heads"), + [(False, NUM_HEADS), (True, NUM_KV_HEADS)], +) +def test_gqa_kv_stays_compressed_for_paged_attention( + paged_kv_active: bool, + expected_kv_heads: int, +) -> None: + mgr = _make_cache_mgr() + mgr.attn.paged_kv_active = paged_kv_active + q_len = 3 + key, value = _make_known_kv(q_len) + + _call_mgr( + mgr, + bs=1, + q_len=q_len, + seq_len=q_len, + key_flat=key, + value_flat=value, + first_step=True, + num_image_tokens=0, + gen_timestep_scatter_index=_gen_timestep_index(1, 0), + full_attn_spans=[[(0, q_len)]] if paged_kv_active else None, + ) + + _, key_input, value_input = mgr.attn.calls[-1] + expected_shape = (1, q_len, expected_kv_heads, HEAD_DIM) + assert key_input.shape == expected_shape + assert value_input.shape == expected_shape + + # ============================================================ # Test 1: No AR KV — basic cache → reuse # ============================================================ diff --git a/vllm_omni/diffusion/models/magi_human/__init__.py b/tests/diffusion/models/lingbot_world/__init__.py similarity index 100% rename from vllm_omni/diffusion/models/magi_human/__init__.py rename to tests/diffusion/models/lingbot_world/__init__.py diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture new file mode 100644 index 000000000..015a5aa91 --- /dev/null +++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture @@ -0,0 +1,13 @@ +{ + "_class_name": "LingBotWorldCausalDMDPipeline", + "_diffusers_version": "0.35.0.dev0", + "boundary_ratio": 0.9, + "image_encoder": [null, null], + "image_processor": [null, null], + "scheduler": ["diffusers", "UniPCMultistepScheduler"], + "text_encoder": ["transformers", "UMT5EncoderModel"], + "tokenizer": ["transformers", "T5TokenizerFast"], + "transformer": ["diffusers", "CausalLingBotWorldTransformer3DModel"], + "transformer_2": [null, null], + "vae": ["diffusers", "AutoencoderKLWan"] +} diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture new file mode 100644 index 000000000..d07408a04 --- /dev/null +++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture @@ -0,0 +1,42 @@ +{ + "sources": { + "config": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/config.json", + "shard_1_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00001-of-00008.safetensors", + "shard_8_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00008-of-00008.safetensors" + }, + "dtype": "F32", + "shard_representatives": { + "diffusion_pytorch_model-00001-of-00008.safetensors": "blocks.0.cam_injector_layer1.bias", + "diffusion_pytorch_model-00002-of-00008.safetensors": "blocks.4.cam_injector_layer1.bias", + "diffusion_pytorch_model-00003-of-00008.safetensors": "blocks.10.cam_injector_layer1.bias", + "diffusion_pytorch_model-00004-of-00008.safetensors": "blocks.14.cam_scale_layer.bias", + "diffusion_pytorch_model-00005-of-00008.safetensors": "blocks.20.cam_injector_layer1.bias", + "diffusion_pytorch_model-00006-of-00008.safetensors": "blocks.25.cam_injector_layer1.bias", + "diffusion_pytorch_model-00007-of-00008.safetensors": "blocks.30.cam_injector_layer1.bias", + "diffusion_pytorch_model-00008-of-00008.safetensors": "blocks.35.cam_injector_layer1.bias" + }, + "parameter_shapes": { + "blocks.0.cam_injector_layer1.bias": [5120], + "blocks.4.cam_injector_layer1.bias": [5120], + "blocks.10.cam_injector_layer1.bias": [5120], + "blocks.14.cam_scale_layer.bias": [5120], + "blocks.20.cam_injector_layer1.bias": [5120], + "blocks.25.cam_injector_layer1.bias": [5120], + "blocks.30.cam_injector_layer1.bias": [5120], + "blocks.35.cam_injector_layer1.bias": [5120], + "patch_embedding.weight": [5120, 36, 1, 2, 2], + "patch_embedding_wancamctrl.weight": [5120, 1536], + "c2ws_hidden_states_layer1.weight": [5120, 5120], + "c2ws_hidden_states_layer2.weight": [5120, 5120], + "blocks.0.modulation": [1, 6, 5120], + "blocks.0.self_attn.q.weight": [5120, 5120], + "blocks.0.ffn.0.weight": [13824, 5120], + "blocks.0.ffn.2.weight": [5120, 13824], + "blocks.0.cam_injector_layer1.weight": [5120, 5120], + "blocks.0.cam_injector_layer2.weight": [5120, 5120], + "blocks.0.cam_scale_layer.weight": [5120, 5120], + "blocks.0.cam_shift_layer.weight": [5120, 5120], + "head.modulation": [1, 2, 5120], + "head.head.weight": [64, 5120] + } +} diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture new file mode 100644 index 000000000..950d26fae --- /dev/null +++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture @@ -0,0 +1,30 @@ +{ + "_class_name": "UniPCMultistepScheduler", + "_diffusers_version": "0.35.0.dev0", + "beta_end": 0.02, + "beta_schedule": "linear", + "beta_start": 0.0001, + "disable_corrector": [], + "dynamic_thresholding_ratio": 0.995, + "final_sigmas_type": "zero", + "flow_shift": 5.0, + "lower_order_final": true, + "num_train_timesteps": 1000, + "predict_x0": true, + "prediction_type": "flow_prediction", + "rescale_betas_zero_snr": false, + "sample_max_value": 1.0, + "solver_order": 2, + "solver_p": null, + "solver_type": "bh2", + "steps_offset": 0, + "thresholding": false, + "time_shift_type": "exponential", + "timestep_spacing": "linspace", + "trained_betas": null, + "use_beta_sigmas": false, + "use_dynamic_shifting": false, + "use_exponential_sigmas": false, + "use_flow_sigmas": true, + "use_karras_sigmas": false +} diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture new file mode 100644 index 000000000..2a2c93efd --- /dev/null +++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture @@ -0,0 +1,112 @@ +{ + "source": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model.safetensors.index.json", + "index_parameter_count": 1421, + "canonical_name_set_sha256": "5d31b7f64fd5e5e271b789b2dbe50fb7bbf146bd8b750646bdc9abb4e5b8139d", + "canonical_name_to_shard_sha256": "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992", + "canonical_serialization": "sorted lines formatted as parameter_name=shard_filename with LF separators and no trailing LF", + "shards": { + "diffusion_pytorch_model-00001-of-00008.safetensors": { + "parameter_count": 180, + "parameter_names_sha256": "81a3d4778de604e75609aae7dac3d8dcbc081811e7748267fe953bbeb375e79d" + }, + "diffusion_pytorch_model-00002-of-00008.safetensors": { + "parameter_count": 180, + "parameter_names_sha256": "4e36e2596a6468f29179fa5795396f2e9a7f7d8debc912b34a6f6c2c7c7c610d" + }, + "diffusion_pytorch_model-00003-of-00008.safetensors": { + "parameter_count": 180, + "parameter_names_sha256": "6f74a7b2cf7444c2daa2b3fc32bd3787589b35543c230c5c13b89f0f0f2a0998" + }, + "diffusion_pytorch_model-00004-of-00008.safetensors": { + "parameter_count": 180, + "parameter_names_sha256": "ec5c967ed5271ba7834337c4ee93a73f8ed0eead4492d4920624251c65c058ac" + }, + "diffusion_pytorch_model-00005-of-00008.safetensors": { + "parameter_count": 180, + "parameter_names_sha256": "ad1fdab5830dd9ce43ea007f112b4be223327b7c8197c766d9cf9219e74d619b" + }, + "diffusion_pytorch_model-00006-of-00008.safetensors": { + "parameter_count": 180, + "parameter_names_sha256": "fd58f0d86be42d8afcf55d3e07d8a22f7f6ef5064da931a600c890b71b6dfbae" + }, + "diffusion_pytorch_model-00007-of-00008.safetensors": { + "parameter_count": 180, + "parameter_names_sha256": "305a1714b5f4fb54d5877c593b451273ab1490d90fd0896a299d96e140295047" + }, + "diffusion_pytorch_model-00008-of-00008.safetensors": { + "parameter_count": 161, + "parameter_names_sha256": "4c18d94d96168682eb212f93e3cdf7313928776d8d9e6e4de0255fe899decb25" + } + }, + "top_level_modules": [ + "blocks", + "c2ws_hidden_states_layer1", + "c2ws_hidden_states_layer2", + "head", + "patch_embedding", + "patch_embedding_wancamctrl", + "text_embedding", + "time_embedding", + "time_projection" + ], + "non_block_parameters": [ + "c2ws_hidden_states_layer1.bias", + "c2ws_hidden_states_layer1.weight", + "c2ws_hidden_states_layer2.bias", + "c2ws_hidden_states_layer2.weight", + "head.head.bias", + "head.head.weight", + "head.modulation", + "patch_embedding.bias", + "patch_embedding.weight", + "patch_embedding_wancamctrl.bias", + "patch_embedding_wancamctrl.weight", + "text_embedding.0.bias", + "text_embedding.0.weight", + "text_embedding.2.bias", + "text_embedding.2.weight", + "time_embedding.0.bias", + "time_embedding.0.weight", + "time_embedding.2.bias", + "time_embedding.2.weight", + "time_projection.1.bias", + "time_projection.1.weight" + ], + "block_parameter_suffixes": [ + "cam_injector_layer1.bias", + "cam_injector_layer1.weight", + "cam_injector_layer2.bias", + "cam_injector_layer2.weight", + "cam_scale_layer.bias", + "cam_scale_layer.weight", + "cam_shift_layer.bias", + "cam_shift_layer.weight", + "cross_attn.k.bias", + "cross_attn.k.weight", + "cross_attn.norm_k.weight", + "cross_attn.norm_q.weight", + "cross_attn.o.bias", + "cross_attn.o.weight", + "cross_attn.q.bias", + "cross_attn.q.weight", + "cross_attn.v.bias", + "cross_attn.v.weight", + "ffn.0.bias", + "ffn.0.weight", + "ffn.2.bias", + "ffn.2.weight", + "modulation", + "norm3.bias", + "norm3.weight", + "self_attn.k.bias", + "self_attn.k.weight", + "self_attn.norm_k.weight", + "self_attn.norm_q.weight", + "self_attn.o.bias", + "self_attn.o.weight", + "self_attn.q.bias", + "self_attn.q.weight", + "self_attn.v.bias", + "self_attn.v.weight" + ] +} diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py new file mode 100644 index 000000000..df012777f --- /dev/null +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio + +import pytest +import torch + +from vllm_omni.diffusion.models.lingbot_world.actions import ( + LINGBOT_CAMERA_ACTION_SCHEMA, + LINGBOT_CAMERA_TRAJECTORY_SCHEMA, + LingBotCameraControlReducer, + integrate_lingbot_camera_actions, + parse_lingbot_camera_action_frames, +) +from vllm_omni.experimental.ar_diffusion.session import ( + ARDiffusionSession, + ARDiffusionSessionEvent, +) +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionChunkMetadata, + ARDiffusionControlInput, + ARDiffusionTickRequest, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _action_control(mode: str, **data) -> ARDiffusionControlInput: + return ARDiffusionControlInput( + track="camera", + schema=LINGBOT_CAMERA_ACTION_SCHEMA, + data={"mode": mode, **data}, + ) + + +def _prepared_frames(prepared) -> tuple[tuple[str, ...], ...]: + assert len(prepared.controls) == 1 + return prepared.controls[0].data["frames"] + + +def test_state_controls_hold_release_and_preserve_short_pulse() -> None: + reducer = LingBotCameraControlReducer() + pressed = reducer.prepare( + current_controls={}, + events=( + ARDiffusionSessionEvent( + event_id=1, + controls=( + _action_control( + "state", + transitions=[{"actions": ["W"], "client_ts_ms": 10}], + ), + ), + ), + ), + chunk_index=0, + ) + assert _prepared_frames(pressed) == (("w",), ("w",), ("w",)) + assert pressed.controls[0].to_dict()["data"]["frames"] == [["w"], ["w"], ["w"]] + reducer.commit(pressed) + + held = reducer.prepare( + current_controls={"camera": pressed.controls[0]}, + events=(), + chunk_index=1, + ) + assert _prepared_frames(held) == (("w",), ("w",), ("w",)) + reducer.commit(held) + + pulse = reducer.prepare( + current_controls={"camera": held.controls[0]}, + events=( + ARDiffusionSessionEvent( + event_id=2, + controls=( + _action_control( + "state", + transitions=[ + {"actions": ["j"], "client_ts_ms": 20}, + {"actions": [], "client_ts_ms": 21}, + ], + ), + ), + ), + ), + chunk_index=2, + ) + assert _prepared_frames(pulse) == (("j",), (), ()) + reducer.commit(pulse) + + released = reducer.prepare( + current_controls={"camera": pulse.controls[0]}, + events=(), + chunk_index=3, + ) + assert _prepared_frames(released) == ((), (), ()) + + +def test_script_is_transactional_consumed_once_and_neutral_padded() -> None: + reducer = LingBotCameraControlReducer() + event = ARDiffusionSessionEvent( + event_id=1, + controls=( + _action_control( + "script", + frames=[["w"], ["d"], ["l"], ["i"]], + ), + ), + ) + first_attempt = reducer.prepare( + current_controls={}, + events=(event,), + chunk_index=0, + ) + assert _prepared_frames(first_attempt) == (("w",), ("d",), ("l",)) + + # An uncommitted failed attempt must not consume the reducer's script. + second_attempt = reducer.prepare( + current_controls={}, + events=(event,), + chunk_index=0, + ) + assert _prepared_frames(second_attempt) == (("w",), ("d",), ("l",)) + reducer.commit(second_attempt) + + tail = reducer.prepare( + current_controls={"camera": second_attempt.controls[0]}, + events=(), + chunk_index=1, + ) + assert _prepared_frames(tail) == (("i",), (), ()) + reducer.commit(tail) + exhausted = reducer.prepare( + current_controls={"camera": tail.controls[0]}, + events=(), + chunk_index=2, + ) + assert _prepared_frames(exhausted) == ((), (), ()) + + +def test_trajectory_control_remains_an_expert_override() -> None: + reducer = LingBotCameraControlReducer() + trajectory = ARDiffusionControlInput( + track="camera", + schema=LINGBOT_CAMERA_TRAJECTORY_SCHEMA, + data={"poses": [[[1.0]]], "intrinsics": [[1.0]]}, + ) + prepared = reducer.prepare( + current_controls={}, + events=( + ARDiffusionSessionEvent( + event_id=1, + controls=(trajectory,), + ), + ), + chunk_index=0, + ) + assert prepared.controls == (trajectory,) + + +def test_action_integrator_matches_lingbot_motion_and_calibration() -> None: + trajectory, pitch = integrate_lingbot_camera_actions( + [["w"], ["d"], ["l"]], + width=832, + height=480, + ) + + assert pitch == 0.0 + torch.testing.assert_close( + trajectory.poses[0, :3, 3], + torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64), + atol=1e-6, + rtol=0, + ) + torch.testing.assert_close( + trajectory.poses[1, :3, 3], + torch.tensor([0.05, 0.0, 0.05], dtype=torch.float64), + atol=1e-6, + rtol=0, + ) + expected_yaw = torch.tensor( + [ + [0.9945219, 0.0, 0.10452846], + [0.0, 1.0, 0.0], + [-0.10452846, 0.0, 0.9945219], + ], + dtype=torch.float64, + ) + torch.testing.assert_close( + trajectory.poses[2, :3, :3], + expected_yaw, + atol=1e-6, + rtol=0, + ) + torch.testing.assert_close( + trajectory.intrinsics[0], + torch.tensor([500.0, 500.0, 416.0, 240.0]), + ) + + resized, _ = integrate_lingbot_camera_actions( + [[]], + width=416, + height=240, + ) + scaled = resized.intrinsics[0] * torch.tensor([416 / 832, 240 / 480, 416 / 832, 240 / 480]) + torch.testing.assert_close( + scaled, + torch.tensor([500.0, 500.0, 208.0, 120.0]), + ) + + +@pytest.mark.parametrize( + ("data", "message"), + [ + ({"mode": "frames", "frames": [["x"], [], []]}, "W/A/S/D/I/J/K/L"), + ({"mode": "frames", "frames": [["w"]]}, "exactly one action list"), + ({"mode": "state", "frames": [[], [], []]}, "mode='frames'"), + ], +) +def test_model_tick_action_payload_validation(data, message: str) -> None: + with pytest.raises(ValueError, match=message): + parse_lingbot_camera_action_frames(data, expected_frames=3) + + +class _TickConsumer: + def __init__(self) -> None: + self.ticks: list[ARDiffusionTickRequest] = [] + + async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata: + self.ticks.append(tick) + return ARDiffusionChunkMetadata.from_tick(tick) + + def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata: + return output + + +class _Lifecycle: + async def reset_session(self, session_id: str) -> None: + del session_id + + async def close_session(self, session_id: str) -> None: + del session_id + + +def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None: + async def exercise() -> None: + consumer = _TickConsumer() + reducer = LingBotCameraControlReducer() + session = ARDiffusionSession( + "world-actions", + tick_consumer=consumer, + lifecycle=_Lifecycle(), + max_pending_events=8, + request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}", + control_reducer=reducer, + ) + await session.accept_event( + ARDiffusionSessionEvent( + event_id=1, + prompt="turn toward the window", + controls=( + _action_control( + "state", + transitions=[{"actions": ["d"]}], + ), + ), + ) + ) + await session.next_chunk() + + tick = consumer.ticks[0] + assert tick.prompt == "turn toward the window" + assert tick.applied_event_ids == (1,) + assert tick.controls[0].data["frames"] == (("d",), ("d",), ("d",)) + + await session.reset() + await session.accept_event( + ARDiffusionSessionEvent( + event_id=2, + controls=( + _action_control( + "state", + transitions=[{"actions": []}], + ), + ), + ) + ) + await session.next_chunk() + assert consumer.ticks[1].chunk_index == 0 + assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ()) + + asyncio.run(exercise()) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py new file mode 100644 index 000000000..f5567a706 --- /dev/null +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py @@ -0,0 +1,543 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib.util +import math +import sys +from pathlib import Path +from types import ModuleType + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/transformer.py" +_MODULE_NAME = "_lingbot_world_attention_under_test" +_MISSING_MODULE = object() +_STUBBED_MODULE_NAMES = ( + "vllm", + "vllm.distributed", + "vllm.model_executor", + "vllm.model_executor.layers", + "vllm.model_executor.layers.conv", + "vllm.model_executor.layers.linear", + "vllm.model_executor.model_loader", + "vllm.model_executor.model_loader.weight_utils", + "vllm.model_executor.utils", + "vllm_omni", + "vllm_omni.diffusion", + "vllm_omni.diffusion.attention", + "vllm_omni.diffusion.attention.layer", + "vllm_omni.diffusion.layers", + "vllm_omni.diffusion.layers.norm", + "vllm_omni.diffusion.layers.rope", + "vllm_omni.experimental", + "vllm_omni.experimental.ar_diffusion", + "vllm_omni.experimental.ar_diffusion.kv_cache", + "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention", +) + + +@pytest.fixture(autouse=True) +def _inference_context(): + with torch.inference_mode(): + yield + + +def _install_vllm_stubs() -> None: + for name in _STUBBED_MODULE_NAMES: + sys.modules[name] = ModuleType(name) + + distributed = sys.modules["vllm.distributed"] + distributed.get_tensor_model_parallel_rank = lambda: 0 + distributed.get_tensor_model_parallel_world_size = lambda: 1 + distributed.tensor_model_parallel_all_reduce = lambda value: value + + def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None: + for name, value in attrs.items(): + setattr(weight, name, value) + + sys.modules["vllm.model_executor.utils"].set_weight_attrs = set_weight_attrs + sys.modules["vllm.model_executor.model_loader.weight_utils"].default_weight_loader = lambda param, loaded_weight: ( + param.data.copy_(loaded_weight) + ) + + class _Linear(nn.Module): + def __init__( + self, + input_size: int, + output_size: int, + *, + bias: bool = True, + return_bias: bool = False, + **kwargs, + ) -> None: + super().__init__() + del kwargs + self.return_bias = return_bias + self.weight = nn.Parameter(torch.empty(output_size, input_size)) + self.bias = nn.Parameter(torch.empty(output_size)) if bias else None + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + nn.init.zeros_(self.bias) + self.calls = 0 + + def forward(self, value: torch.Tensor): + self.calls += 1 + output = F.linear(value, self.weight, self.bias) + return (output, self.bias) if self.return_bias else output + + linear = sys.modules["vllm.model_executor.layers.linear"] + linear.ColumnParallelLinear = _Linear + linear.RowParallelLinear = _Linear + + class _QKVLinear(nn.Module): + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int | None = None, + *, + bias: bool = True, + **kwargs, + ) -> None: + super().__init__() + del kwargs + self.num_heads = total_num_heads + self.num_kv_heads = total_num_kv_heads or total_num_heads + output_size = (self.num_heads + 2 * self.num_kv_heads) * head_size + self.weight = nn.Parameter(torch.empty(output_size, hidden_size)) + self.bias = nn.Parameter(torch.empty(output_size)) if bias else None + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + nn.init.zeros_(self.bias) + self.calls = 0 + + def load_shard(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None: + shard_index = {"q": 0, "k": 1, "v": 2}[shard_id] + shard_size = param.shape[0] // 3 + param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight) + + self.weight.weight_loader = load_shard + if self.bias is not None: + self.bias.weight_loader = load_shard + + def forward(self, value: torch.Tensor): + self.calls += 1 + return F.linear(value, self.weight, self.bias), None + + linear.QKVParallelLinear = _QKVLinear + + class _Conv3dLayer(nn.Conv3d): + def __init__(self, in_channels: int, out_channels: int, kernel_size, stride, **kwargs) -> None: + super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs) + + sys.modules["vllm.model_executor.layers.conv"].Conv3dLayer = _Conv3dLayer + + class _LayerNorm(nn.LayerNorm): + def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True) -> None: + super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine) + + sys.modules["vllm_omni.diffusion.layers.norm"].LayerNorm = _LayerNorm + + class _Attention(nn.Module): + def __init__(self, *args, softmax_scale: float, **kwargs) -> None: + super().__init__() + del args + self.softmax_scale = softmax_scale + self.causal = kwargs["causal"] + self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata=None, + ) -> torch.Tensor: + del attn_metadata + self.calls.append((query.detach().clone(), key.detach().clone(), value.detach().clone())) + scores = torch.einsum("bqhd,bkhd->bhqk", query, key) * self.softmax_scale + weights = scores.softmax(dim=-1) + return torch.einsum("bhqk,bkhd->bqhd", weights, value) + + sys.modules["vllm_omni.diffusion.attention.layer"].Attention = _Attention + + class _RotaryEmbeddingWan(nn.Module): + def __init__(self, *args, **kwargs) -> None: + super().__init__() + del args, kwargs + self.calls: list[tuple[torch.Tensor, torch.Tensor]] = [] + + def forward(self, value: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + self.calls.append((cos.detach().clone(), sin.detach().clone())) + return value + + sys.modules["vllm_omni.diffusion.layers.rope"].RotaryEmbeddingWan = _RotaryEmbeddingWan + + paged_attention = sys.modules["vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"] + + class _PagedLayerContext: + pass + + class _PagedLayerInputs: + pass + + paged_attention.ARDiffusionPagedLayerContext = _PagedLayerContext + paged_attention.ARDiffusionPagedLayerInputs = _PagedLayerInputs + paged_attention.ar_diffusion_paged_attention = lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("CUDA paged attention was not expected in this CPU test") + ) + paged_attention.paged_write_attn = lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("paged path was not expected in this direct-cache test") + ) + + +def _load_module(): + assert _MODULE_PATH.exists(), "LingBot attention module has not been implemented" + module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME) + previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names} + try: + _install_vllm_stubs() + spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + finally: + for name, original in previous.items(): + if original is _MISSING_MODULE: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + + +def test_load_module_restores_process_module_state() -> None: + module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME) + previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names} + + _load_module() + + for name, original in previous.items(): + assert sys.modules.get(name, _MISSING_MODULE) is original + + +def test_allocate_lingbot_cache_creates_request_local_layer_storage() -> None: + module = _load_module() + + cache = module.allocate_lingbot_cache( + batch_size=2, + num_layers=3, + max_tokens=7, + num_local_heads=2, + head_dim=4, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert len(cache.self_attention) == 3 + assert len(cache.cross_attention) == 3 + assert cache.cross_attention == [None, None, None] + assert cache.self_attention[0].key.shape == (2, 7, 2, 4) + assert cache.self_attention[0].value.shape == (2, 7, 2, 4) + assert cache.self_attention[0].key.dtype == torch.float32 + assert cache.self_attention[0].end == 0 + assert cache.self_attention[0].key.data_ptr() != cache.self_attention[1].key.data_ptr() + + +def _allocate_single_layer(module, *, max_tokens: int): + return module.allocate_lingbot_cache( + batch_size=1, + num_layers=1, + max_tokens=max_tokens, + num_local_heads=1, + head_dim=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + +def _set_identity_attention(attention: nn.Module) -> None: + qkv = getattr(attention, "qkv", None) + if qkv is not None: + projections = (qkv, attention.o) + else: + projections = (attention.q, attention.k, attention.v, attention.o) + for projection in projections: + with torch.no_grad(): + identity = torch.eye(2) + projection.weight.copy_(torch.cat((identity, identity, identity)) if projection is qkv else identity) + if projection.bias is not None: + projection.bias.zero_() + attention.norm_q = nn.Identity() + attention.norm_k = nn.Identity() + + +def _tokens(*values: float) -> torch.Tensor: + return torch.tensor([[[value, 0.0] for value in values]]) + + +def _record_inputs(module: nn.Module) -> list[tuple]: + records: list[tuple] = [] + module.register_forward_pre_hook(lambda _module, args: records.append(args)) + return records + + +def _record_outputs(module: nn.Module) -> list[torch.Tensor]: + records: list[torch.Tensor] = [] + + def record(_module, _args, output) -> None: + records.append(output[0] if isinstance(output, tuple) else output) + + module.register_forward_hook(record) + return records + + +def _cache_snapshot(cache) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]: + return ( + cache.key.detach().view(torch.uint8).clone(), + cache.value.detach().view(torch.uint8).clone(), + (cache.end, cache.absolute_end, cache.last_start, cache.sink_end), + ) + + +def _assert_cache_unchanged(cache, snapshot) -> None: + key, value, metadata = snapshot + assert torch.equal(cache.key.detach().view(torch.uint8), key) + assert torch.equal(cache.value.detach().view(torch.uint8), value) + assert (cache.end, cache.absolute_end, cache.last_start, cache.sink_end) == metadata + + +def test_self_attention_repeated_offset_overwrites_then_later_offset_appends() -> None: + module = _load_module() + attention = module.LingBotSelfAttention(dim=2, num_heads=1) + _set_identity_attention(attention) + cache = _allocate_single_layer(module, max_tokens=6).self_attention[0] + + attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0) + attention(_tokens(10, 20), cache=cache, current_start=0, sink_tokens=0) + + assert cache.end == 2 + assert cache.absolute_end == 2 + torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0])) + + attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0) + + assert cache.end == 4 + assert cache.absolute_end == 4 + torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0, 3.0, 4.0])) + + +def test_self_attention_uses_one_fused_qkv_projection() -> None: + module = _load_module() + attention = module.LingBotSelfAttention(dim=2, num_heads=1) + cache = _allocate_single_layer(module, max_tokens=2).self_attention[0] + + attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0) + + assert attention.qkv.calls == 1 + + +def test_self_attention_routes_paged_inputs_through_engine_owned_cache() -> None: + module = _load_module() + attention = module.LingBotSelfAttention(dim=2, num_heads=1) + _set_identity_attention(attention) + paged_inputs = module.ARDiffusionPagedLayerInputs() + calls = [] + + def paged_write(inputs, query, key, value, k_act, v_act, scale): + calls.append((inputs, query.clone(), key.clone(), value.clone(), k_act, v_act, scale)) + return query + + module.paged_write_attn = paged_write + output = attention( + _tokens(1, 2), + cache=paged_inputs, + current_start=0, + sink_tokens=0, + update_cache=False, + ) + + assert output.shape == (1, 2, 2) + assert len(calls) == 1 + assert calls[0][0] is paged_inputs + assert calls[0][4:6] == (None, None) + assert attention.attn.calls == [] + assert not any(hasattr(attention, name) for name in ("q", "k", "v")) + + +def test_self_attention_is_chunk_causal_without_masking_inside_current_chunk() -> None: + module = _load_module() + attention = module.LingBotSelfAttention(dim=2, num_heads=1) + _set_identity_attention(attention) + cache = _allocate_single_layer(module, max_tokens=8).self_attention[0] + attention_inputs = _record_inputs(attention.attn) + + attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0) + first_keys = attention_inputs[-1][1].detach().clone() + attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0) + second_keys = attention_inputs[-1][1].detach().clone() + + assert attention.attn.causal is False + torch.testing.assert_close(first_keys[0, :, 0, 0], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(second_keys[0, :, 0, 0], torch.tensor([1.0, 2.0, 3.0, 4.0])) + + +def test_self_attention_retains_sink_and_latest_local_history_after_eviction() -> None: + module = _load_module() + attention = module.LingBotSelfAttention(dim=2, num_heads=1) + _set_identity_attention(attention) + cache = _allocate_single_layer(module, max_tokens=4).self_attention[0] + attention_inputs = _record_inputs(attention.attn) + + attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=1) + attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=1) + attention(_tokens(5, 6), cache=cache, current_start=4, sink_tokens=1) + + visible_keys = attention_inputs[-1][1].detach().clone() + assert cache.end == 4 + assert cache.absolute_end == 6 + torch.testing.assert_close(visible_keys[0, :, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0])) + torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0])) + + +def test_cross_attention_projects_encoder_kv_once_per_request() -> None: + module = _load_module() + attention = module.LingBotCrossAttention(dim=2, num_heads=1) + _set_identity_attention(attention) + key_outputs = _record_outputs(attention.k) + value_outputs = _record_outputs(attention.v) + attention_inputs = _record_inputs(attention.attn) + + output, cache = attention(_tokens(1), _tokens(2, 3), cache=None) + second_output, reused_cache = attention(_tokens(4), _tokens(20, 30), cache=cache) + + assert output.shape == second_output.shape == (1, 1, 2) + assert reused_cache is cache + assert len(key_outputs) == 1 + assert len(value_outputs) == 1 + assert cache.end == 2 + torch.testing.assert_close(cache.key[0, :, 0, 0], torch.tensor([2.0, 3.0])) + torch.testing.assert_close(attention_inputs[-1][1], cache.key) + + +def test_self_attention_cache_is_isolated_between_requests() -> None: + module = _load_module() + attention = module.LingBotSelfAttention(dim=2, num_heads=1) + _set_identity_attention(attention) + first = _allocate_single_layer(module, max_tokens=4).self_attention[0] + second = _allocate_single_layer(module, max_tokens=4).self_attention[0] + + attention(_tokens(1, 2), cache=first, current_start=0, sink_tokens=0) + + assert first.end == 2 + assert second.end == 0 + assert torch.count_nonzero(second.key) == 0 + attention(_tokens(8), cache=second, current_start=0, sink_tokens=0) + torch.testing.assert_close(first.key[0, : first.end, 0, 0], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(second.key[0, : second.end, 0, 0], torch.tensor([8.0])) + + +def test_cross_attention_keeps_checkpoint_qkvo_and_norm_parameter_names() -> None: + module = _load_module() + attention = module.LingBotCrossAttention(dim=4, num_heads=2) + + names = set(attention.state_dict()) + + assert { + "q.weight", + "k.weight", + "v.weight", + "o.weight", + "norm_q.weight", + "norm_k.weight", + } <= names + assert callable(attention.norm_q.weight.weight_loader) + + +def test_tp_world_size_one_attention_output_shapes() -> None: + module = _load_module() + self_attention = module.LingBotSelfAttention(dim=4, num_heads=2) + cross_attention = module.LingBotCrossAttention(dim=4, num_heads=2) + cache = module.allocate_lingbot_cache( + batch_size=2, + num_layers=1, + max_tokens=5, + num_local_heads=2, + head_dim=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + hidden_states = torch.randn(2, 3, 4) + encoder_hidden_states = torch.randn(2, 5, 4) + + self_output = self_attention(hidden_states, cache=cache.self_attention[0], current_start=0, sink_tokens=0) + cross_output, cache.cross_attention[0] = cross_attention( + hidden_states, + encoder_hidden_states, + cache=cache.cross_attention[0], + ) + + assert self_output.shape == hidden_states.shape + assert cross_output.shape == hidden_states.shape + assert cache.cross_attention[0] is not None + + +def test_self_attention_applies_rotary_embedding_to_current_query_and_key() -> None: + module = _load_module() + attention = module.LingBotSelfAttention(dim=2, num_heads=1) + cache = _allocate_single_layer(module, max_tokens=2).self_attention[0] + cos = torch.ones(2, 1) + sin = torch.zeros(2, 1) + rotary_inputs = _record_inputs(attention.rotary_embedding) + + attention(_tokens(1, 2), cache=cache, current_start=0, rotary_emb=(cos, sin), sink_tokens=0) + + assert len(rotary_inputs) == 2 + torch.testing.assert_close(rotary_inputs[0][1], cos) + torch.testing.assert_close(rotary_inputs[1][2], sin) + + +def test_attention_block_owns_checkpoint_parent_namespaces() -> None: + module = _load_module() + block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0") + + names = set(block.state_dict()) + expected = { + *(f"self_attn.{child}.weight" for child in ("qkv", "o", "norm_q", "norm_k")), + *(f"cross_attn.{child}.weight" for child in ("q", "k", "v", "o", "norm_q", "norm_k")), + } + + assert expected <= names + + +def test_cross_attention_caches_are_isolated_between_requests() -> None: + module = _load_module() + attention = module.LingBotCrossAttention(dim=2, num_heads=1) + + _, first = attention(_tokens(1), _tokens(2, 3), cache=None) + _, second = attention(_tokens(1), _tokens(8, 9), cache=None) + + assert first.key.untyped_storage().data_ptr() != second.key.untyped_storage().data_ptr() + assert first.value.untyped_storage().data_ptr() != second.value.untyped_storage().data_ptr() + second_key = second.key.clone() + first.key.zero_() + torch.testing.assert_close(second.key, second_key) + + +def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_module() + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1) + norm = module._LingBotRMSNorm(2, eps=1e-6) + + norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0])) + + torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0])) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py new file mode 100644 index 000000000..121364c07 --- /dev/null +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py @@ -0,0 +1,525 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import importlib.util +import math +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/camera.py" +_SPEC = importlib.util.spec_from_file_location("_lingbot_world_camera_under_test", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_CAMERA_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _CAMERA_MODULE +_SPEC.loader.exec_module(_CAMERA_MODULE) + +CameraTrajectory = _CAMERA_MODULE.CameraTrajectory +build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding +interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory +load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory +resolve_trusted_action_directory = _CAMERA_MODULE.resolve_trusted_action_directory + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _identity_poses(num_frames: int) -> np.ndarray: + return np.repeat(np.eye(4, dtype=np.float64)[None], num_frames, axis=0) + + +def _write_trajectory( + directory: Path, + poses: np.ndarray, + intrinsics: np.ndarray, +) -> None: + np.save(directory / "poses.npy", poses) + np.save(directory / "intrinsics.npy", intrinsics) + + +def _trusted(directory: Path): + return resolve_trusted_action_directory(".", directory) + + +def _trajectory( + poses: torch.Tensor, + intrinsics: torch.Tensor | None = None, +) -> CameraTrajectory: + if intrinsics is None: + intrinsics = torch.tensor([[832.0 / 3.0, 160.0, 416.0, 240.0]]).repeat(poses.shape[0], 1) + return CameraTrajectory(poses=poses, intrinsics=intrinsics) + + +@pytest.mark.parametrize("present_file", [None, "poses.npy"]) +def test_load_camera_trajectory_rejects_missing_files(tmp_path: Path, present_file: str | None) -> None: + if present_file == "poses.npy": + np.save(tmp_path / present_file, _identity_poses(1)) + + missing_file = "poses.npy" if present_file is None else "intrinsics.npy" + with pytest.raises(FileNotFoundError, match=missing_file): + load_camera_trajectory(_trusted(tmp_path)) + + +@pytest.mark.parametrize( + ("poses", "intrinsics", "message"), + [ + (np.zeros((2, 3, 4)), np.zeros((2, 4)), "poses.npy"), + (_identity_poses(2), np.zeros((2, 3)), "intrinsics.npy"), + (_identity_poses(0), np.zeros((0, 4)), "1 and 4096 source frames"), + (_identity_poses(2), np.zeros((1, 4)), "same number of frames"), + ], +) +def test_load_camera_trajectory_rejects_invalid_shapes( + tmp_path: Path, + poses: np.ndarray, + intrinsics: np.ndarray, + message: str, +) -> None: + _write_trajectory(tmp_path, poses, intrinsics) + + with pytest.raises(ValueError, match=message): + load_camera_trajectory(_trusted(tmp_path)) + + +@pytest.mark.parametrize("array_name", ["poses", "intrinsics"]) +def test_load_camera_trajectory_rejects_non_finite_values(tmp_path: Path, array_name: str) -> None: + poses = _identity_poses(1) + intrinsics = np.ones((1, 4), dtype=np.float64) + if array_name == "poses": + poses[0, 0, 0] = np.nan + else: + intrinsics[0, 0] = np.inf + _write_trajectory(tmp_path, poses, intrinsics) + + with pytest.raises(ValueError, match="finite"): + load_camera_trajectory(_trusted(tmp_path)) + + +def test_load_camera_trajectory_preserves_raw_c2w_poses(tmp_path: Path) -> None: + first_c2w = np.array( + [ + [0.0, -1.0, 0.0, 2.0], + [1.0, 0.0, 0.0, 3.0], + [0.0, 0.0, 1.0, 4.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + second_c2w = np.array( + [ + [0.0, 0.0, 1.0, 5.0], + [0.0, 1.0, 0.0, 4.0], + [-1.0, 0.0, 0.0, 6.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + poses = np.stack((first_c2w, second_c2w)) + intrinsics = np.array([[100.0, 110.0, 120.0, 130.0]] * 2) + _write_trajectory(tmp_path, poses, intrinsics) + + trajectory = load_camera_trajectory(_trusted(tmp_path)) + + assert trajectory.poses.dtype == torch.float32 + assert trajectory.intrinsics.dtype == torch.float32 + torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float()) + torch.testing.assert_close(trajectory.intrinsics, torch.from_numpy(intrinsics).float()) + + +def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_path: Path) -> None: + poses = _identity_poses(1).astype(object) + intrinsics = np.ones((1, 4), dtype=np.float32) + _write_trajectory(tmp_path, poses, intrinsics) + + with pytest.raises(ValueError, match="poses.npy.*real numeric"): + load_camera_trajectory(_trusted(tmp_path)) + + +def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None: + _write_trajectory( + tmp_path, + _identity_poses(269), + np.arange(269 * 4, dtype=np.float32).reshape(269, 4), + ) + + trajectory = load_camera_trajectory(_trusted(tmp_path)) + + assert trajectory.poses.shape == (117, 4, 4) + assert trajectory.intrinsics.shape == (117, 4) + np.testing.assert_array_equal( + trajectory.intrinsics.numpy(), + np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117], + ) + + +def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None: + _write_trajectory( + tmp_path, + _identity_poses(4097), + np.ones((4097, 4), dtype=np.float32), + ) + + with pytest.raises(ValueError, match="poses.npy.*bounded camera-array limit|4096 source frames"): + load_camera_trajectory(_trusted(tmp_path)) + + +def test_load_camera_trajectory_rejects_truncated_payload_before_loading(tmp_path: Path) -> None: + _write_trajectory(tmp_path, _identity_poses(1), np.ones((1, 4), dtype=np.float32)) + poses_path = tmp_path / "poses.npy" + poses_path.write_bytes(poses_path.read_bytes()[:-1]) + + with pytest.raises(ValueError, match="poses.npy byte size"): + load_camera_trajectory(_trusted(tmp_path)) + + +def test_load_camera_trajectory_rejects_oversized_npy_header(tmp_path: Path) -> None: + header_size = 20_000 + (tmp_path / "poses.npy").write_bytes( + np.lib.format.magic(2, 0) + header_size.to_bytes(4, "little") + b" " * (header_size - 1) + b"\n" + ) + np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32)) + + with pytest.raises(ValueError, match="poses.npy.*oversized NPY header"): + load_camera_trajectory(_trusted(tmp_path)) + + +def test_load_camera_trajectory_rejects_symlink_swap_after_resolution(tmp_path: Path) -> None: + trusted_root = tmp_path / "trusted" + action_directory = trusted_root / "action" + outside = tmp_path / "outside" + action_directory.mkdir(parents=True) + outside.mkdir() + _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32)) + _write_trajectory(outside, _identity_poses(1), np.ones((1, 4), dtype=np.float32)) + resolved = resolve_trusted_action_directory("action", trusted_root) + action_directory.rename(trusted_root / "original-action") + action_directory.symlink_to(outside, target_is_directory=True) + + with pytest.raises(ValueError, match="securely open.*action directory"): + load_camera_trajectory(resolved) + + +def test_load_camera_trajectory_rejects_symlinked_camera_file(tmp_path: Path) -> None: + trusted_root = tmp_path / "trusted" + action_directory = trusted_root / "action" + outside = tmp_path / "outside" + action_directory.mkdir(parents=True) + outside.mkdir() + _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32)) + np.save(outside / "poses.npy", _identity_poses(1)) + resolved = resolve_trusted_action_directory("action", trusted_root) + (action_directory / "poses.npy").unlink() + (action_directory / "poses.npy").symlink_to(outside / "poses.npy") + + with pytest.raises(ValueError, match="poses.npy.*opened securely"): + load_camera_trajectory(resolved) + + +@pytest.mark.skipif(not hasattr(__import__("os"), "mkfifo"), reason="requires POSIX FIFO support") +def test_load_camera_trajectory_rejects_fifo_without_blocking(tmp_path: Path) -> None: + import os + + os.mkfifo(tmp_path / "poses.npy") + np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32)) + script = """ +import importlib.util +import sys + +spec = importlib.util.spec_from_file_location("_lingbot_camera_fifo_test", sys.argv[1]) +module = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = module +spec.loader.exec_module(module) +action = module.resolve_trusted_action_directory(".", sys.argv[2]) +try: + module.load_camera_trajectory(action) +except ValueError as error: + print(error) +else: + raise AssertionError("FIFO camera input was accepted") +""" + + completed = subprocess.run( + [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + assert "poses.npy must be a regular NPY file" in completed.stdout + + +def test_load_camera_trajectory_decodes_the_validated_file_snapshot(tmp_path: Path, monkeypatch) -> None: + original_poses = _identity_poses(1) + replacement_poses = original_poses.copy() + replacement_poses[0, 0, 3] = 123.0 + _write_trajectory(tmp_path, original_poses, np.ones((1, 4), dtype=np.float32)) + poses_path = tmp_path / "poses.npy" + original_load = _CAMERA_MODULE.np.load + mutation_count = 0 + + def mutate_source_then_decode(source, *args, **kwargs): + nonlocal mutation_count + if mutation_count == 0: + mutation_count += 1 + np.save(poses_path, replacement_poses) + return original_load(source, *args, **kwargs) + + monkeypatch.setattr(_CAMERA_MODULE.np, "load", mutate_source_then_decode) + + trajectory = load_camera_trajectory(_trusted(tmp_path)) + + assert mutation_count == 1 + torch.testing.assert_close(trajectory.poses, torch.from_numpy(original_poses).float()) + + +def test_load_camera_trajectory_ignores_non_camera_event_files(tmp_path: Path) -> None: + poses = _identity_poses(1) + intrinsics = np.ones((1, 4), dtype=np.float32) + _write_trajectory(tmp_path, poses, intrinsics) + np.save(tmp_path / "wasd_events.npy", np.array([{"key": "w"}], dtype=object)) + + trajectory = load_camera_trajectory(_trusted(tmp_path)) + + torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float()) + + +def test_interpolate_camera_trajectory_uses_linear_and_spherical_interpolation() -> None: + poses = torch.eye(4).repeat(2, 1, 1) + poses[1, :3, :3] = torch.tensor( + [ + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + poses[1, :3, 3] = torch.tensor([2.0, 4.0, 6.0]) + intrinsics = torch.tensor([[10.0, 20.0, 30.0, 40.0], [20.0, 40.0, 60.0, 80.0]]) + trajectory = CameraTrajectory(poses=poses, intrinsics=intrinsics) + + result = interpolate_camera_trajectory(trajectory, 3) + + torch.testing.assert_close(result.poses[0], poses[0]) + torch.testing.assert_close(result.poses[-1], poses[-1]) + torch.testing.assert_close(result.poses[1, :3, 3], torch.tensor([1.0, 2.0, 3.0])) + expected_mid_rotation = torch.tensor( + [ + [math.sqrt(0.5), -math.sqrt(0.5), 0.0], + [math.sqrt(0.5), math.sqrt(0.5), 0.0], + [0.0, 0.0, 1.0], + ] + ) + torch.testing.assert_close(result.poses[1, :3, :3], expected_mid_rotation) + torch.testing.assert_close(result.intrinsics[1], torch.tensor([15.0, 30.0, 45.0, 60.0])) + + +def test_build_plucker_embedding_has_forward_center_ray() -> None: + trajectory = _trajectory(torch.eye(4).unsqueeze(0)) + + result = build_plucker_embedding( + trajectory, + height=3, + width=3, + target_height=3, + target_width=3, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + torch.testing.assert_close(result[0, :3, 1, 1], torch.zeros(3)) + torch.testing.assert_close(result[0, 3:, 1, 1], torch.tensor([0.0, 0.0, 1.0])) + + +def test_build_plucker_embedding_scales_reference_intrinsics() -> None: + trajectory = _trajectory( + torch.eye(4).unsqueeze(0), + torch.tensor([[416.0, 240.0, 416.0, 240.0]]), + ) + + result = build_plucker_embedding( + trajectory, + height=4, + width=8, + target_height=4, + target_width=8, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + expected_corner_direction = torch.tensor([-0.875, -0.75, 1.0]) + expected_corner_direction /= torch.linalg.vector_norm(expected_corner_direction) + torch.testing.assert_close(result[0, 3:, 0, 0], expected_corner_direction) + + +def test_build_plucker_embedding_samples_target_pixel_centers() -> None: + trajectory = _trajectory( + torch.eye(4).unsqueeze(0), + torch.tensor([[416.0, 240.0, 416.0, 240.0]]), + ) + + result = build_plucker_embedding( + trajectory, + height=6, + width=10, + target_height=2, + target_width=5, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + # Target centers map through the source scale: (0.5 * 2, 0.5 * 3). + expected_direction = torch.tensor([-0.8, -0.5, 1.0]) + expected_direction /= torch.linalg.vector_norm(expected_direction) + torch.testing.assert_close(result[0, 3:, 0, 0], expected_direction) + + +def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_translations() -> None: + poses = torch.eye(4).repeat(3, 1, 1) + poses[0, :3, 3] = torch.tensor([5.0, 7.0, 11.0]) + poses[1, :3, :3] = torch.tensor( + [ + [0.0, -1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + poses[1, :3, 3] = torch.tensor([5.0, 9.0, 11.0]) + poses[2, :3, :3] = torch.tensor( + [ + [-1.0, 0.0, 0.0], + [0.0, -1.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + poses[2, :3, 3] = torch.tensor([3.0, 9.0, 11.0]) + trajectory = _trajectory( + poses, + torch.tensor([[416.0, 480.0, 208.0, 240.0]]).repeat(3, 1), + ) + + result = build_plucker_embedding( + trajectory, + height=1, + width=2, + target_height=1, + target_width=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + expected_direction = torch.tensor([0.0, 1.0, 1.0]) / math.sqrt(2.0) + torch.testing.assert_close(result[0, :3, 0, 1], torch.zeros(3)) + torch.testing.assert_close(result[1, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0])) + torch.testing.assert_close(result[2, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0])) + torch.testing.assert_close(result[1, 3:, 0, 1], expected_direction) + torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction) + + +def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None: + poses = torch.eye(4).repeat(3, 1, 1) + poses[:, :3, 3] = torch.tensor( + [ + [10.0, -2.0, 5.0], + [13.0, 2.0, 5.0], + [19.0, 10.0, 5.0], + ] + ) + + result = build_plucker_embedding( + _trajectory(poses), + height=1, + width=1, + target_height=1, + target_width=1, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]]) + torch.testing.assert_close(result[:, :3, 0, 0], expected_origins) + + +def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None: + poses = torch.eye(4).repeat(2, 1, 1) + poses[0, :3, 3] = torch.tensor([20.0, 4.0, -2.0]) + poses[1, :3, 3] = torch.tensor([26.0, 12.0, -2.0]) + raw = _trajectory(poses) + + interpolated = interpolate_camera_trajectory(raw, 3) + result = build_plucker_embedding( + interpolated, + height=1, + width=1, + target_height=1, + target_width=1, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + expected_delta = torch.tensor([0.6, 0.8, 0.0]) + torch.testing.assert_close(result[0, :3, 0, 0], torch.zeros(3)) + torch.testing.assert_close(result[1, :3, 0, 0], expected_delta) + torch.testing.assert_close(result[2, :3, 0, 0], expected_delta) + + +def test_build_plucker_embedding_uses_half_pixel_centers_and_origin_direction_channel_order() -> None: + poses = torch.eye(4).repeat(2, 1, 1) + poses[1, :3, 3] = torch.tensor([2.0, 0.0, 0.0]) + trajectory = _trajectory( + poses, + torch.tensor([[832.0, 480.0, 0.0, 0.0]]).repeat(2, 1), + ) + + result = build_plucker_embedding( + trajectory, + height=1, + width=2, + target_height=1, + target_width=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + expected_direction = torch.tensor([0.25, 0.5, 1.0]) + expected_direction /= torch.linalg.vector_norm(expected_direction) + torch.testing.assert_close(result[1, :3, 0, 0], torch.tensor([1.0, 0.0, 0.0])) + torch.testing.assert_close(result[1, 3:, 0, 0], expected_direction) + + +def test_build_plucker_embedding_honors_shape_dtype_and_device() -> None: + trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1)) + + result = build_plucker_embedding( + trajectory, + height=12, + width=20, + target_height=3, + target_width=5, + device=torch.device("cpu"), + dtype=torch.float64, + ) + + assert result.shape == (2, 6, 3, 5) + assert result.dtype == torch.float64 + assert result.device == torch.device("cpu") + + +def test_build_plucker_embedding_is_deterministic() -> None: + trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1)) + kwargs = { + "height": 12, + "width": 20, + "target_height": 3, + "target_width": 5, + "device": torch.device("cpu"), + "dtype": torch.float32, + } + + first = build_plucker_embedding(trajectory, **kwargs) + second = build_plucker_embedding(trajectory, **kwargs) + + torch.testing.assert_close(first, second, rtol=0.0, atol=0.0) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py new file mode 100644 index 000000000..e2e2a08b6 --- /dev/null +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py @@ -0,0 +1,685 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import hashlib +import inspect +import json +import math +from pathlib import Path + +import pytest +import torch + +from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture" +_SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture" + + +@pytest.fixture(autouse=True) +def _inference_context(): + with torch.inference_mode(): + yield + + +def _tiny_model( + module, + *, + num_layers: int = 2, + num_frames_per_block: int = 1, + sliding_window_num_frames: int = 3, +): + return module.CausalLingBotWorldTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=2, + attention_head_dim=2, + in_channels=36, + out_channels=2, + text_dim=6, + freq_dim=4, + ffn_dim=8, + num_layers=num_layers, + cross_attn_norm=True, + eps=1e-6, + rope_max_seq_len=16, + sink_size=1, + num_frames_per_block=num_frames_per_block, + sliding_window_num_frames=sliding_window_num_frames, + local_attn_size=-1, + ) + + +def _cache(module, model, *, max_tokens: int | None = None): + if max_tokens is None: + window_frames = ( + model.config.local_attn_size + if model.config.local_attn_size != -1 + else model.config.sliding_window_num_frames + ) + max_tokens = window_frames * 2 * 2 + return module.allocate_lingbot_cache( + batch_size=1, + num_layers=len(model.blocks), + max_tokens=max_tokens, + num_local_heads=2, + head_dim=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + +def _checkpoint_parameter_specs(model) -> dict[str, tuple[torch.Size, torch.dtype]]: + specs: dict[str, tuple[torch.Size, torch.dtype]] = {} + for name, parameter in model.named_parameters(): + marker = ".self_attn.qkv." + if marker not in name: + specs[name] = (parameter.shape, parameter.dtype) + continue + shard_shape = torch.Size((parameter.shape[0] // 3, *parameter.shape[1:])) + for projection in ("q", "k", "v"): + specs[name.replace(marker, f".self_attn.{projection}.")] = (shard_shape, parameter.dtype) + return specs + + +def _checkpoint_weights_for_model(model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + checkpoint_weights: dict[str, torch.Tensor] = {} + expected_parameters: dict[str, torch.Tensor] = {} + next_value = 1 + for name, parameter in model.named_parameters(): + marker = ".self_attn.qkv." + if marker not in name: + value = torch.full_like(parameter, next_value) + checkpoint_weights[name] = value + expected_parameters[name] = value + next_value += 1 + continue + shards = [] + shard_shape = (parameter.shape[0] // 3, *parameter.shape[1:]) + for projection in ("q", "k", "v"): + value = torch.full(shard_shape, next_value, dtype=parameter.dtype, device=parameter.device) + checkpoint_weights[name.replace(marker, f".self_attn.{projection}.")] = value + shards.append(value) + next_value += 1 + expected_parameters[name] = torch.cat(shards) + return checkpoint_weights, expected_parameters + + +def _self_cache_snapshot(cache) -> list[tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]]: + return [attention_tests._cache_snapshot(layer_cache) for layer_cache in cache.self_attention] + + +def _assert_self_cache_unchanged(cache, snapshot) -> None: + for layer_cache, layer_snapshot in zip(cache.self_attention, snapshot, strict=True): + attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot) + + +def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None: + torch.manual_seed(7) + module = attention_tests._load_module() + model = _tiny_model(module).eval() + cache = _cache(module, model) + encoder_hidden_states = torch.randn(1, 3, 6) + camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4) + cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks] + text_projection_outputs = attention_tests._record_outputs(model.text_embedding) + outputs = [] + + assert model.patch_embedding.in_channels == 36 + assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2 + + for start_frame in range(4): + hidden_states = torch.randn(1, 36, 1, 4, 4) + timestep = torch.tensor([float(start_frame + 1)]) + snapshot = _self_cache_snapshot(cache) + + transient_output = model( + hidden_states, + timestep, + encoder_hidden_states, + camera_hidden_states, + cache=cache, + start_frame=start_frame, + update_cache=False, + ) + + assert transient_output.shape == (1, 2, 1, 4, 4) + _assert_self_cache_unchanged(cache, snapshot) + assert all(layer_cache is not None for layer_cache in cache.cross_attention) + + committed_output = model( + hidden_states, + timestep, + encoder_hidden_states, + camera_hidden_states, + cache=cache, + start_frame=start_frame, + update_cache=True, + ) + torch.testing.assert_close(committed_output, transient_output) + outputs.append(committed_output) + + expected_token_end = (start_frame + 1) * 2 * 2 + assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention) + + assert len(outputs) == 4 + assert all(len(outputs) == 1 for outputs in cross_key_outputs) + assert len(text_projection_outputs) == 1 + + alternate_cache = _cache(module, model) + alternate_output = model( + torch.zeros(1, 36, 1, 4, 4), + torch.tensor([1.0]), + encoder_hidden_states, + torch.zeros_like(camera_hidden_states), + cache=alternate_cache, + start_frame=0, + update_cache=False, + ) + camera_output = model( + torch.zeros(1, 36, 1, 4, 4), + torch.tensor([1.0]), + encoder_hidden_states, + torch.ones_like(camera_hidden_states), + cache=_cache(module, model), + start_frame=0, + update_cache=False, + ) + assert not torch.equal(camera_output, alternate_output) + + +def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None: + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=2).eval() + cache = _cache(module, model) + hidden_states = torch.randn(1, 36, 1, 4, 4) + timestep = torch.tensor([1.0]) + encoder_hidden_states = torch.randn(1, 3, 6) + camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4) + + model( + hidden_states, + timestep, + encoder_hidden_states, + camera_hidden_states, + cache=cache, + start_frame=0, + update_cache=False, + ) + cache.cross_attention[1] = None + text_projection_outputs = attention_tests._record_outputs(model.text_embedding) + cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks] + + model( + hidden_states, + timestep, + encoder_hidden_states, + camera_hidden_states, + cache=cache, + start_frame=0, + update_cache=False, + ) + + assert len(text_projection_outputs) == 1 + assert [len(outputs) for outputs in cross_key_outputs] == [0, 1] + assert all(layer_cache is not None for layer_cache in cache.cross_attention) + + +@pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks")) +def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None: + module = attention_tests._load_module() + model = _tiny_model( + module, + num_layers=1, + num_frames_per_block=3, + sliding_window_num_frames=6, + ) + cache = _cache(module, model) + snapshot = _self_cache_snapshot(cache) + patch_inputs = attention_tests._record_inputs(model.patch_embedding) + + with pytest.raises(ValueError, match="exactly 3 post-patch frames"): + model( + torch.randn(1, 36, frames, 4, 4), + torch.tensor([1.0]), + torch.randn(1, 3, 6), + torch.randn(1, 6 * 8 * 8, frames, 4, 4), + cache=cache, + start_frame=0, + update_cache=True, + ) + + assert patch_inputs == [] + _assert_self_cache_unchanged(cache, snapshot) + assert cache.cross_attention == [None] + + +def test_forward_accepts_exactly_one_configured_frame_block() -> None: + module = attention_tests._load_module() + model = _tiny_model( + module, + num_layers=1, + num_frames_per_block=3, + sliding_window_num_frames=6, + ) + + output = model( + torch.randn(1, 36, 3, 4, 4), + torch.tensor([1.0]), + torch.randn(1, 3, 6), + torch.randn(1, 6 * 8 * 8, 3, 4, 4), + cache=_cache(module, model), + start_frame=0, + update_cache=False, + ) + + assert output.shape == (1, 2, 3, 4, 4) + + +def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None: + module = attention_tests._load_module() + model = _tiny_model( + module, + num_layers=2, + num_frames_per_block=3, + sliding_window_num_frames=6, + ) + + cache = model.allocate_cache( + batch_size=2, + latent_height=8, + latent_width=12, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert len(cache.self_attention) == 2 + assert len(cache.cross_attention) == 2 + assert cache.cross_attention == [None, None] + assert all(layer.key.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention) + assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention) + + +def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None: + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1, num_frames_per_block=2) + hidden_states = torch.zeros(1, 36, 2, 2, 4) + hidden_states[0, 0] = torch.tensor( + [ + [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]], + ] + ) + with torch.no_grad(): + model.patch_embedding.weight.zero_() + model.patch_embedding.bias.zero_() + model.patch_embedding.weight[0, 0, 0, 0, 0] = 1 + + tokens = model.patch_embedding(hidden_states).flatten(2).transpose(1, 2) + + torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0])) + + +def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None: + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1, num_frames_per_block=2) + hidden_states = ( + torch.stack( + ( + torch.arange(0, 8), + torch.arange(10, 18), + torch.arange(20, 28), + torch.arange(30, 38), + ) + ) + .unsqueeze(0) + .float() + ) + expected = torch.tensor( + [ + [ + [[0, 2, 10, 12], [4, 6, 14, 16]], + [[20, 22, 30, 32], [24, 26, 34, 36]], + ], + [ + [[1, 3, 11, 13], [5, 7, 15, 17]], + [[21, 23, 31, 33], [25, 27, 35, 37]], + ], + ], + dtype=torch.float32, + ).unsqueeze(0) + + output = model._unpatchify( + hidden_states, + batch_size=1, + frames=2, + height=1, + width=2, + ) + + torch.testing.assert_close(output, expected) + + +def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None: + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1, num_frames_per_block=2) + model.head.norm = torch.nn.Identity() + with torch.no_grad(): + model.head.modulation.zero_() + model.head.head.weight.zero_() + model.head.head.bias.zero_() + model.head.head.weight[:, 0] = 1 + hidden_states = torch.ones(1, 4, model.dim) + timestep_embedding = torch.tensor([[[10.0, 0.0, 0.0, 0.0], [20.0, 0.0, 0.0, 0.0]]]) + expected = torch.tensor([21.0, 21.0, 41.0, 41.0]).view(1, 4, 1).expand(1, 4, 8) + + output = model.head(hidden_states, timestep_embedding) + + torch.testing.assert_close(output, expected) + + +def test_constructor_defaults_match_official_checkpoint_config() -> None: + module = attention_tests._load_module() + parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters + expected = { + "in_channels": 36, + "out_channels": 16, + "num_layers": 40, + "num_attention_heads": 40, + "attention_head_dim": 128, + "patch_size": (1, 2, 2), + "text_dim": 4096, + "ffn_dim": 13824, + "freq_dim": 256, + "rope_max_seq_len": 1024, + "eps": 1e-6, + "cross_attn_norm": True, + "sink_size": 9, + "num_frames_per_block": 3, + "sliding_window_num_frames": 18, + "local_attn_size": -1, + } + + assert {name: parameters[name].default for name in expected} == expected + + +def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None: + module = attention_tests._load_module() + model = _tiny_model(module) + + assert model.dtype == next(model.parameters()).dtype + + +def test_transformer_declares_regional_compile_block() -> None: + module = attention_tests._load_module() + + assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"] + + +def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None: + module = attention_tests._load_module() + reduced_values: list[torch.Tensor] = [] + + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2) + + def all_reduce(value: torch.Tensor) -> torch.Tensor: + reduced_values.append(value.detach().clone()) + return value * 2 + + monkeypatch.setattr(module, "tensor_model_parallel_all_reduce", all_reduce) + norm = module._LingBotRMSNorm(hidden_size=2, eps=0.0) + value = torch.tensor([[3.0, 4.0]]) + + output = norm(value) + + torch.testing.assert_close(output, value / math.sqrt(12.5)) + assert len(reduced_values) == 1 + torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]])) + + +def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None: + module = attention_tests._load_module() + monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3) + + with pytest.raises(ValueError, match="num_heads=2.*tp_size=3"): + module.LingBotSelfAttention(dim=4, num_heads=2) + + +def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None: + module = attention_tests._load_module() + + with pytest.raises(ValueError, match="qk_norm.*rms_norm_across_heads"): + module.CausalLingBotWorldTransformer3DModel(qk_norm="layer_norm") + + +@pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"]) +def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None: + module = attention_tests._load_module() + + with pytest.raises(ValueError, match=field): + module.CausalLingBotWorldTransformer3DModel(**{field: 4}) + + +def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None: + module = attention_tests._load_module() + + with pytest.raises(RuntimeError, match="quant_config.*not supported"): + module.CausalLingBotWorldTransformer3DModel(quant_config=object()) + + +def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None: + module = attention_tests._load_module() + + class ConfigProbe(module.CausalLingBotWorldTransformer3DModel): + def __init__(self, **kwargs) -> None: + self.received_kwargs = kwargs + + config = { + "_class_name": "CausalLingBotWorldTransformer3DModel", + "_diffusers_version": "0.35.0.dev0", + "patch_size": [1, 2, 2], + "num_attention_heads": 40, + "attention_head_dim": 128, + "in_channels": 36, + "out_channels": 16, + "text_dim": 4096, + "freq_dim": 256, + "ffn_dim": 13824, + "num_layers": 40, + "cross_attn_norm": True, + "eps": 1e-6, + "image_dim": None, + "added_kv_proj_dim": None, + "rope_max_seq_len": 1024, + "pos_embed_seq_len": None, + "qk_norm": "rms_norm_across_heads", + "sink_size": 9, + "num_frames_per_block": 3, + "sliding_window_num_frames": 18, + "local_attn_size": -1, + } + + model = ConfigProbe.from_config(config, prefix="transformer") + + assert model.received_kwargs["patch_size"] == (1, 2, 2) + assert model.received_kwargs["num_layers"] == 40 + assert model.received_kwargs["prefix"] == "transformer" + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("in_channels", 35), + ("num_layers", 39), + ("num_attention_heads", 32), + ("attention_head_dim", 64), + ("ffn_dim", 10240), + ("num_frames_per_block", 4), + ("sliding_window_num_frames", 24), + ("sink_size", 0), + ], +) +def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None: + module = attention_tests._load_module() + config = { + "_class_name": "CausalLingBotWorldTransformer3DModel", + "_diffusers_version": "0.35.0.dev0", + "patch_size": [1, 2, 2], + "num_attention_heads": 40, + "attention_head_dim": 128, + "in_channels": 36, + "out_channels": 16, + "text_dim": 4096, + "freq_dim": 256, + "ffn_dim": 13824, + "num_layers": 40, + "cross_attn_norm": True, + "eps": 1e-6, + "image_dim": None, + "added_kv_proj_dim": None, + "rope_max_seq_len": 1024, + "pos_embed_seq_len": None, + "qk_norm": "rms_norm_across_heads", + "sink_size": 9, + "num_frames_per_block": 3, + "sliding_window_num_frames": 18, + "local_attn_size": -1, + } + config[field] = value + + with pytest.raises(ValueError, match=field): + module.CausalLingBotWorldTransformer3DModel.from_config(config) + + +def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None: + module = attention_tests._load_module() + config = { + "_class_name": "CausalLingBotWorldTransformer3DModel", + "architectures": ["CausalLingBotWorldTransformer3DModel"], + "future_diffusers_metadata": {"version": 1}, + } + + with torch.device("meta"): + model = module.CausalLingBotWorldTransformer3DModel.from_config(config) + + assert model.config.in_channels == 36 + assert model.config.num_layers == 40 + + +def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None: + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1) + checkpoint_weights, expected_parameters = _checkpoint_weights_for_model(model) + loader_calls: list[str] = [] + qkv_weight = model.blocks[0].self_attn.qkv.weight + + def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None: + loader_calls.append(shard_id) + shard_index = {"q": 0, "k": 1, "v": 2}[shard_id] + shard_size = param.shape[0] // 3 + param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight) + + qkv_weight.weight_loader = record_loader + loaded = model.load_weights(iter(checkpoint_weights.items())) + + assert set(checkpoint_weights) <= loaded + assert set(dict(model.named_parameters())) <= loaded + assert loader_calls == ["q", "k", "v"] + for name, param in model.named_parameters(): + torch.testing.assert_close(param, expected_parameters[name]) + + first_name = next(iter(checkpoint_weights)) + partial = _tiny_model(module, num_layers=1) + partial_loaded = partial.load_weights([(first_name, checkpoint_weights[first_name])]) + expected_partial = {first_name} + for projection in ("q", "k", "v"): + marker = f".self_attn.{projection}." + if marker in first_name: + expected_partial.add(first_name.replace(marker, ".self_attn.qkv.")) + break + assert partial_loaded == expected_partial + + with pytest.raises(KeyError, match="unexpected_model.weight"): + model.load_weights([("unexpected_model.weight", torch.ones(1))]) + + +def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None: + module = attention_tests._load_module() + model = _tiny_model(module, num_layers=1) + first_name, first_param = next(iter(model.named_parameters())) + loaded_weight = torch.full_like(first_param, 17) + + def weights(): + yield first_name, loaded_weight + torch.testing.assert_close(first_param, loaded_weight) + + assert model.load_weights(weights()) == {first_name} + + +def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None: + utils = pytest.importorskip("vllm.model_executor.models.utils") + module = attention_tests._load_module() + transformer = _tiny_model(module, num_layers=1) + checkpoint_weights, _ = _checkpoint_weights_for_model(transformer) + + class Pipeline(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer = transformer + + pipeline = Pipeline() + loaded = utils.AutoWeightsLoader(pipeline).load_weights( + (f"transformer.{name}", value) for name, value in checkpoint_weights.items() + ) + + assert set(dict(pipeline.named_parameters())) <= loaded + + +def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None: + module = attention_tests._load_module() + fixture = json.loads(_FIXTURE_PATH.read_text()) + model = _tiny_model(module, num_layers=1) + parameter_names = set(_checkpoint_parameter_specs(model)) + + assert {name.split(".", 1)[0] for name in parameter_names} == set(fixture["top_level_modules"]) + assert {name for name in parameter_names if not name.startswith("blocks.")} == set(fixture["non_block_parameters"]) + assert {name.removeprefix("blocks.0.") for name in parameter_names if name.startswith("blocks.0.")} == set( + fixture["block_parameter_suffixes"] + ) + expanded_names = set(fixture["non_block_parameters"]) + expanded_names.update( + f"blocks.{layer}.{suffix}" for layer in range(40) for suffix in fixture["block_parameter_suffixes"] + ) + canonical_names = "\n".join(sorted(expanded_names)) + + assert len(expanded_names) == fixture["index_parameter_count"] == 1421 + assert hashlib.sha256(canonical_names.encode()).hexdigest() == fixture["canonical_name_set_sha256"] + assert fixture["canonical_name_to_shard_sha256"] == ( + "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992" + ) + assert list(fixture["shards"]) == [ + f"diffusion_pytorch_model-{index:05}-of-00008.safetensors" for index in range(1, 9) + ] + assert sum(shard["parameter_count"] for shard in fixture["shards"].values()) == 1421 + assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values()) + + +def test_official_default_shapes_match_public_safetensors_header_fixture() -> None: + module = attention_tests._load_module() + fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text()) + + with torch.device("meta"): + model = module.CausalLingBotWorldTransformer3DModel() + parameter_specs = _checkpoint_parameter_specs(model) + expected_dtype = {"F32": torch.float32}[fixture["dtype"]] + + assert len(fixture["shard_representatives"]) == 8 + assert set(fixture["shard_representatives"].values()) <= set(fixture["parameter_shapes"]) + + for name, expected_shape in fixture["parameter_shapes"].items(): + shape, dtype = parameter_specs[name] + assert tuple(shape) == tuple(expected_shape), name + assert dtype == expected_dtype diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py new file mode 100644 index 000000000..56c25fee3 --- /dev/null +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_ROOT = Path(__file__).parents[4] +_PRODUCTION_FILES = ( + _ROOT / "vllm_omni/diffusion/models/lingbot_world/camera.py", + _ROOT / "vllm_omni/diffusion/models/lingbot_world/transformer.py", + _ROOT / "vllm_omni/diffusion/models/lingbot_world/pipeline.py", +) + + +def test_all_lingbot_production_functions_have_complete_annotations() -> None: + missing: list[str] = [] + for path in _PRODUCTION_FILES: + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + continue + arguments = (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs) + for argument in arguments: + if argument.arg not in {"self", "cls"} and argument.annotation is None: + missing.append(f"{path.name}:{node.lineno}:{node.name}:{argument.arg}") + if node.args.vararg is not None and node.args.vararg.annotation is None: + missing.append(f"{path.name}:{node.lineno}:{node.name}:*{node.args.vararg.arg}") + if node.args.kwarg is not None and node.args.kwarg.annotation is None: + missing.append(f"{path.name}:{node.lineno}:{node.name}:**{node.args.kwarg.arg}") + if node.returns is None: + missing.append(f"{path.name}:{node.lineno}:{node.name}:return") + + assert missing == [] diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py new file mode 100644 index 000000000..1a08e8270 --- /dev/null +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -0,0 +1,1765 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import gc +import weakref +from pathlib import Path +from threading import Lock +from types import SimpleNamespace + +import pytest +import torch +from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor +from PIL import Image +from torch import nn + +import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline +from tests.diffusion.models.wan2_2.conftest import noop_progress_bar +from vllm_omni.diffusion.models.lingbot_world.actions import ( + integrate_lingbot_camera_actions, +) +from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as _CameraTrajectory +from vllm_omni.diffusion.models.lingbot_world.camera import ( + build_plucker_embedding as _real_build_plucker_embedding, +) +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionControlInput, + ARDiffusionTickRequest, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_ROOT = Path(__file__).parents[4] +_LINGBOT_INIT_PATH = _ROOT / "vllm_omni/diffusion/models/lingbot_world/__init__.py" +_MODEL_INDEX_FIXTURE = Path(__file__).with_name("fixtures") / "lingbot_world_model_index.json.fixture" +_SCHEDULER_FIXTURE = Path(__file__).with_name("fixtures") / "lingbot_world_scheduler_config.json.fixture" + + +def _scheduler_config(**overrides): + values = { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 5.0, + "prediction_type": "flow_prediction", + "predict_x0": True, + "use_flow_sigmas": True, + "use_dynamic_shifting": False, + "use_beta_sigmas": False, + "use_exponential_sigmas": False, + "use_karras_sigmas": False, + "final_sigmas_type": "zero", + "timestep_spacing": "linspace", + "solver_order": 2, + "solver_type": "bh2", + "lower_order_final": True, + "disable_corrector": [], + "time_shift_type": "exponential", + } + values.update(overrides) + return values + + +class _AutoWeightsLoader: + def __init__(self, module): + self.module = module + + def load_weights(self, weights): + loaded = self.module.transformer.load_weights( + (name.removeprefix("transformer."), value) for name, value in weights if name.startswith("transformer.") + ) + return {f"transformer.{name}" for name in loaded} + + +class _FakePretrained: + @classmethod + def from_pretrained(cls, *args, **kwargs): + del args, kwargs + return cls() + + def to(self, *args, **kwargs): + self.to_calls = getattr(self, "to_calls", []) + self.to_calls.append((args, kwargs)) + return self + + +class _FakeScheduler: + def __init__(self, **kwargs): + self.config = SimpleNamespace(**kwargs) + + +class _FakeTransformerFactory: + @classmethod + def from_config(cls, config, *, quant_config=None, prefix=""): + cls.last_call = (config, quant_config, prefix) + return _RecordingTransformer() + + +class _Cache: + def __init__(self, *, num_layers: int, max_tokens: int, num_local_heads: int, head_dim: int): + shape = (1, max_tokens, num_local_heads, head_dim) + self.self_attention = [ + SimpleNamespace(key=torch.zeros(shape), value=torch.zeros(shape)) for _ in range(num_layers) + ] + self.cross_attention = [None] * num_layers + + +class _RecordingTransformer(nn.Module): + def __init__(self, *, raise_on_call: int | None = None, dtype: torch.dtype = torch.float32): + super().__init__() + self.config = SimpleNamespace( + patch_size=(1, 2, 2), + in_channels=36, + out_channels=16, + text_dim=8, + num_layers=2, + num_attention_heads=2, + attention_head_dim=4, + num_frames_per_block=3, + sliding_window_num_frames=6, + local_attn_size=-1, + sink_size=3, + ) + self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()]) + for block in self.blocks: + block.self_attn = SimpleNamespace(num_local_heads=2, head_dim=4) + self.calls: list[dict] = [] + self.cache_allocations: list[dict] = [] + self.raise_on_call = raise_on_call + self._dtype = dtype + self.loaded_weights: list[tuple[str, torch.Tensor]] = [] + + @property + def dtype(self) -> torch.dtype: + return self._dtype + + def forward(self, **kwargs): + call = { + "hidden_states": kwargs["hidden_states"].detach().clone(), + "timestep": kwargs["timestep"].detach().clone(), + "encoder_hidden_states": kwargs["encoder_hidden_states"].detach().clone(), + "camera_hidden_states": kwargs["camera_hidden_states"].detach().clone(), + "cache_id": id(kwargs["cache"]), + "start_frame": kwargs["start_frame"], + "update_cache": kwargs["update_cache"], + } + self.calls.append(call) + if self.raise_on_call == len(self.calls): + raise RuntimeError("forced transformer failure") + return torch.ones_like(kwargs["hidden_states"][:, :16]) + + def allocate_cache(self, **kwargs): + self.cache_allocations.append(dict(kwargs)) + patch_height, patch_width = self.config.patch_size[1:] + window_frames = ( + self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames + ) + max_tokens = window_frames * (kwargs["latent_height"] // patch_height) * (kwargs["latent_width"] // patch_width) + return _Cache( + num_layers=self.config.num_layers, + max_tokens=max_tokens, + num_local_heads=2, + head_dim=4, + ) + + def load_weights(self, weights): + self.loaded_weights = list(weights) + return {name for name, _ in self.loaded_weights} + + +class _StubVAE(_FakePretrained): + dtype = torch.float32 + + def __init__(self): + self.config = SimpleNamespace( + z_dim=16, + scale_factor_temporal=4, + scale_factor_spatial=8, + latents_mean=[float(index) - 3.0 for index in range(16)], + latents_std=[1.0 + index / 4.0 for index in range(16)], + ) + self.encode_inputs: list[torch.Tensor] = [] + self.decode_inputs: list[torch.Tensor] = [] + self.on_decode = None + + def encode(self, video: torch.Tensor): + self.encode_inputs.append(video.detach().clone()) + latent_frames = (video.shape[2] - 1) // 4 + 1 + latents = torch.zeros( + video.shape[0], + 16, + latent_frames, + video.shape[-2] // 8, + video.shape[-1] // 8, + dtype=video.dtype, + device=video.device, + ) + latents[:, :, 0] = 2.0 + return SimpleNamespace(latents=latents) + + def decode(self, latents: torch.Tensor, return_dict: bool = False): + del return_dict + if self.on_decode is not None: + self.on_decode() + self.decode_inputs.append(latents.detach().clone()) + pixel_frames = (latents.shape[2] - 1) * 4 + 1 + decoded = torch.zeros( + latents.shape[0], + 3, + pixel_frames, + latents.shape[-2] * 8, + latents.shape[-1] * 8, + dtype=latents.dtype, + device=latents.device, + ) + return (decoded,) + + +class _SamplingParams: + def __init__( + self, + *, + height: int | None = 16, + width: int | None = 16, + num_frames: int = 9, + num_inference_steps: int | None = 4, + num_outputs_per_prompt: int = 1, + seed: int | None = 17, + generator: torch.Generator | None = None, + output_type: str | None = "latent", + max_sequence_length: int | None = 512, + extra_args: dict | None = None, + include_action: bool = True, + ) -> None: + self.height = height + self.width = width + self.num_frames = num_frames + self.num_inference_steps = num_inference_steps + self.num_outputs_per_prompt = num_outputs_per_prompt + self.seed = seed + self.generator = ( + generator + if generator is not None + else (torch.Generator(device="cpu").manual_seed(seed) if seed is not None else None) + ) + self.output_type = output_type + self.max_sequence_length = max_sequence_length + self.extra_args = {"action_path": "."} if include_action else {} + self.extra_args["_lingbot_camera_trajectory"] = _CameraTrajectory( + poses=torch.eye(4).repeat(32, 1, 1), + intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(32, 1), + ) + if extra_args is not None: + self.extra_args.update(extra_args) + self.latents = None + + +class _RequestBatch: + def __init__(self, prompt, sampling_params: _SamplingParams, *, num_reqs: int = 1): + self._prompts = [prompt] * num_reqs + self._sampling = sampling_params + self.num_reqs = num_reqs + + @property + def prompts(self): + return self._prompts + + @property + def sampling_params(self): + return self._sampling + + @property + def request_id(self): + return "legacy-lingbot-request" + + +def _load_pipeline_module(): + return lingbot_pipeline + + +@pytest.fixture(autouse=True) +def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: + """Replace only heavyweight component construction, not the import graph.""" + + loader_state = SimpleNamespace(prefetch_calls=[]) + + def prefetch_subfolders(model, subfolders, *, local_files_only): + loader_state.prefetch_calls.append((model, tuple(subfolders), local_files_only)) + + def from_pretrained_with_prefetch(callable_, model, **kwargs): + kwargs.pop("prefetch_list", None) + return callable_(model, **kwargs) + + def load_transformer_config(model, subfolder, local_files_only): + del model, subfolder, local_files_only + return { + "_class_name": "CausalLingBotWorldTransformer3DModel", + "patch_size": [1, 2, 2], + "in_channels": 36, + "out_channels": 16, + "text_dim": 8, + "num_layers": 2, + "num_attention_heads": 2, + "attention_head_dim": 4, + "num_frames_per_block": 3, + "sliding_window_num_frames": 6, + "local_attn_size": -1, + } + + def retrieve_latents(value, sample_mode="argmax"): + assert sample_mode == "argmax" + return value.latents + + def load_json(model, filename, local_files_only): + del model, local_files_only + assert filename == "scheduler/scheduler_config.json" + return _scheduler_config() + + trajectory = _CameraTrajectory( + poses=torch.eye(4).repeat(32, 1, 1), + intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(32, 1), + ) + + def load_camera_trajectory(action_path): + assert action_path + return trajectory + + def interpolate_camera_trajectory(value, num_frames): + return SimpleNamespace( + poses=value.poses[:num_frames], + intrinsics=value.intrinsics[:num_frames], + ) + + def build_plucker_embedding(value, *, height, width, target_height, target_width, device, dtype): + del target_height, target_width + frames = value.poses.shape[0] + data = torch.arange(frames * 6 * height * width, device=device, dtype=torch.float32) + return data.reshape(frames, 6, height, width).to(dtype=dtype) + + replacements = { + "AutoTokenizer": _FakePretrained, + "UMT5EncoderModel": _FakePretrained, + "AutoWeightsLoader": _AutoWeightsLoader, + "DistributedAutoencoderKLWan": _StubVAE, + "FlowUniPCMultistepScheduler": _FakeScheduler, + "CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory, + "get_local_device": lambda: torch.device("cpu"), + "set_forward_context_denoise_step_idx": lambda index: None, + "prefetch_subfolders": prefetch_subfolders, + "from_pretrained_with_prefetch": from_pretrained_with_prefetch, + "load_transformer_config": load_transformer_config, + "retrieve_latents": retrieve_latents, + "_load_json": load_json, + "load_camera_trajectory": load_camera_trajectory, + "interpolate_camera_trajectory": interpolate_camera_trajectory, + "build_plucker_embedding": build_plucker_embedding, + "randn_tensor": _diffusers_randn_tensor, + } + for name, value in replacements.items(): + monkeypatch.setattr(lingbot_pipeline, name, value) + monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False) + + +def _od_config(**overrides): + values = { + "model": "checkpoint", + "dtype": torch.float32, + "flow_shift": None, + "quantization_config": None, + "enable_diffusion_pipeline_profiler": False, + "model_config": { + "lingbot_action_root": str(_ROOT), + "ar_diffusion_height": 16, + "ar_diffusion_width": 16, + }, + "enable_cpu_offload": False, + "enable_layerwise_offload": False, + "enforce_eager": True, + "parallel_config": SimpleNamespace( + pipeline_parallel_size=1, + sequence_parallel_size=1, + cfg_parallel_size=1, + vae_patch_parallel_size=1, + use_hsdp=False, + enable_expert_parallel=False, + ), + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _prompt(*, action_path: str | None = None, images=None): + if images is None: + images = Image.new("RGB", (16, 16), color=(255, 128, 0)) + prompt = { + "prompt": "move through the room", + "multi_modal_data": {"image": images}, + "additional_information": {}, + } + if action_path is not None: + prompt["additional_information"]["action_path"] = action_path + return prompt + + +def _pipeline(module, *, transformer=None, od_config=None): + pipeline = module.LingBotWorldCausalDMDPipeline(od_config=od_config or _od_config()) + if transformer is not None: + pipeline.transformer = transformer + pipeline.encode_prompt = lambda *args, **kwargs: torch.ones(1, 512, 8) + pipeline.progress_bar = noop_progress_bar + return pipeline + + +def _preprocess_request(module, *, prompt=None, sampling=None, od_config=None): + request = SimpleNamespace( + prompt=_prompt() if prompt is None else prompt, + sampling_params=_SamplingParams() if sampling is None else sampling, + ) + preprocess = module.get_lingbot_world_pre_process_func(od_config or _od_config()) + return preprocess(request) + + +@pytest.mark.parametrize("offload_field", ["enable_cpu_offload", "enable_layerwise_offload"]) +def test_pipeline_respects_loader_managed_component_placement(offload_field: str) -> None: + module = _load_pipeline_module() + + pipeline = _pipeline(module, od_config=_od_config(**{offload_field: True})) + + assert getattr(pipeline.text_encoder, "to_calls", []) == [] + assert getattr(pipeline.vae, "to_calls", []) == [] + + +def test_ar_diffusion_capability_uses_fixed_tp_local_lingbot_geometry() -> None: + module = _load_pipeline_module() + module.get_tensor_model_parallel_world_size = lambda: 1 + pipeline = _pipeline(module) + + spec = pipeline.ar_diffusion_kv_cache_spec() + + assert spec.num_layers == 2 + assert spec.num_kv_heads == 2 + assert spec.head_size == 4 + assert spec.tokens_per_frame == 1 + assert spec.frames_per_block == 3 + assert spec.window_frames == 3 + assert spec.sink_frames == 3 + assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)] + assert spec.cross_attention_lengths == {"text": 512} + assert spec.model_owned_state_bytes_per_session == 9_600 + + +def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None: + module = _load_pipeline_module() + action_root = tmp_path / "trusted-actions" + action_dir = action_root / "forward" + action_dir.mkdir(parents=True) + image_path = tmp_path / "first-frame.png" + Image.new("RGBA", (16, 16), color=(1, 2, 3, 4)).save(image_path) + sampling = _SamplingParams(extra_args={"action_path": "forward"}) + sampling.extra_args.pop("_lingbot_camera_trajectory") + request = SimpleNamespace( + prompt={"prompt": "move", "multi_modal_data": {"image": str(image_path)}}, + sampling_params=sampling, + ) + trajectory = _CameraTrajectory( + poses=torch.eye(4).repeat(9, 1, 1), + intrinsics=torch.ones(9, 4), + ) + load_calls = [] + module.load_camera_trajectory = lambda action: load_calls.append(action) or trajectory + + preprocess = module.get_lingbot_world_pre_process_func( + _od_config(model_config={"lingbot_action_root": str(action_root)}) + ) + result = preprocess(request) + + assert result is request + assert isinstance(request.prompt["multi_modal_data"]["image"], Image.Image) + assert request.prompt["multi_modal_data"]["image"].mode == "RGB" + assert sampling.extra_args["action_path"] == "forward" + assert sampling.extra_args["_lingbot_camera_trajectory"] is trajectory + assert len(load_calls) == 1 + + +def test_preprocess_materializes_camera_from_typed_tick_without_action_path() -> None: + module = _load_pipeline_module() + poses = torch.eye(4).repeat(9, 1, 1) + intrinsics = torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(9, 1) + tick = ARDiffusionTickRequest( + session_id="world-1", + request_id="tick-request-0", + chunk_index=0, + controls=( + ARDiffusionControlInput( + track="camera", + schema="lingbot.camera_trajectory.v1", + data={ + "poses": poses.tolist(), + "intrinsics": intrinsics.tolist(), + }, + ), + ), + ) + sampling = _SamplingParams(include_action=False) + sampling.extra_args.update(tick.to_extra_args()) + request = SimpleNamespace( + request_id="tick-request-0", + prompt=_prompt(), + sampling_params=sampling, + ) + + result = module.get_lingbot_world_pre_process_func(_od_config())(request) + + trajectory = result.sampling_params.extra_args["_lingbot_camera_trajectory"] + torch.testing.assert_close(trajectory.poses, poses) + torch.testing.assert_close(trajectory.intrinsics, intrinsics) + + +def test_preprocess_materializes_chunk_sized_actions_from_typed_tick() -> None: + module = _load_pipeline_module() + tick = ARDiffusionTickRequest( + session_id="world-actions", + request_id="tick-request-0", + chunk_index=0, + controls=( + ARDiffusionControlInput( + track="camera", + schema="lingbot.camera_actions.v1", + data={ + "mode": "frames", + "frames": [["w"], ["w", "j"], []], + }, + ), + ), + ) + sampling = _SamplingParams(include_action=False) + sampling.extra_args.update(tick.to_extra_args()) + request = SimpleNamespace( + request_id="tick-request-0", + prompt=_prompt(), + sampling_params=sampling, + ) + + result = module.get_lingbot_world_pre_process_func(_od_config())(request) + + assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None + assert result.sampling_params.extra_args["_lingbot_camera_actions"] == ( + ("w",), + ("w", "j"), + (), + ) + + +def _request(*, sampling=None, prompt=None, num_reqs: int = 1): + return _RequestBatch( + _prompt() if prompt is None else prompt, + _SamplingParams() if sampling is None else sampling, + num_reqs=num_reqs, + ) + + +def _resolve_pipeline_through_real_registry(pipeline_module): + from vllm_omni.diffusion import registry as registry_module + + resolved = registry_module.DiffusionModelRegistry._try_load_model_cls("LingBotWorldCausalDMDPipeline") + entry = registry_module._DIFFUSION_MODELS["LingBotWorldCausalDMDPipeline"] + cache_acceleration_disabled = "LingBotWorldCausalDMDPipeline" in registry_module._NO_CACHE_ACCELERATION + preprocess_name = registry_module._DIFFUSION_PRE_PROCESS_FUNCS["LingBotWorldCausalDMDPipeline"] + preprocess = registry_module.get_diffusion_pre_process_func( + SimpleNamespace( + model_class_name="LingBotWorldCausalDMDPipeline", + model_config={"lingbot_action_root": str(_ROOT)}, + ) + ) + return resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess + + +def test_component_discovery_uses_official_checkpoint_contract() -> None: + module = _load_pipeline_module() + pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config()) + + assert pipeline._dit_modules == ["transformer"] + assert pipeline._encoder_modules == ["text_encoder"] + assert pipeline._vae_modules == ["vae"] + assert pipeline.dummy_run_num_frames == 0 + assert pipeline.weights_sources == [ + module.DiffusersPipelineLoader.ComponentSource("checkpoint", "transformer", None, "transformer.", True) + ] + assert _FakeTransformerFactory.last_call[1:] == (None, "transformer") + assert pipeline.scheduler.config.shift == 5.0 + assert pipeline.scheduler.config.num_train_timesteps == 1000 + assert module._loader_state.prefetch_calls == [("checkpoint", ("tokenizer", "text_encoder", "vae"), False)] + + +@pytest.mark.parametrize( + ("field", "value", "feature"), + [ + ("pipeline_parallel_size", 2, "pipeline parallelism"), + ("sequence_parallel_size", 2, "sequence parallelism"), + ("cfg_parallel_size", 2, "CFG parallelism"), + ("vae_patch_parallel_size", 2, "VAE parallelism"), + ("use_hsdp", True, "HSDP"), + ("enable_expert_parallel", True, "expert parallelism"), + ], +) +def test_unsupported_parallel_modes_fail_before_component_loading(field: str, value: object, feature: str) -> None: + module = _load_pipeline_module() + parallel_config = _od_config().parallel_config + setattr(parallel_config, field, value) + + with pytest.raises(NotImplementedError, match=feature): + module.LingBotWorldCausalDMDPipeline(od_config=_od_config(parallel_config=parallel_config)) + + assert module._loader_state.prefetch_calls == [] + + +def test_unsupported_quantization_fails_before_component_loading() -> None: + module = _load_pipeline_module() + + with pytest.raises(NotImplementedError, match="quantization"): + module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=object())) + + assert module._loader_state.prefetch_calls == [] + + +def test_official_scheduler_config_matches_fixed_dmd_contract() -> None: + import json + + module = _load_pipeline_module() + scheduler_config = json.loads(_SCHEDULER_FIXTURE.read_text()) + module._load_json = lambda *args, **kwargs: scheduler_config.copy() + + pipeline = _pipeline(module) + + assert pipeline.scheduler.config.num_train_timesteps == 1000 + assert pipeline.scheduler.config.shift == 5.0 + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("_class_name", "FlowMatchEulerDiscreteScheduler"), + ("num_train_timesteps", 999), + ("prediction_type", "epsilon"), + ("predict_x0", False), + ("use_flow_sigmas", False), + ("use_dynamic_shifting", True), + ("final_sigmas_type", "sigma_min"), + ], +) +def test_scheduler_config_rejects_semantic_drift(field: str, value: object) -> None: + module = _load_pipeline_module() + module._load_json = lambda *args, **kwargs: _scheduler_config(**{field: value}) + + with pytest.raises(ValueError, match=field): + _pipeline(module) + + +def test_official_model_index_discovers_only_declared_components() -> None: + import json + + module = _load_pipeline_module() + model_index = json.loads(_MODEL_INDEX_FIXTURE.read_text()) + resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry( + module + ) + + assert model_index["_class_name"] == "LingBotWorldCausalDMDPipeline" + assert resolved is module.LingBotWorldCausalDMDPipeline + assert entry == ("lingbot_world", "pipeline", "LingBotWorldCausalDMDPipeline") + assert cache_acceleration_disabled + assert preprocess_name == "get_lingbot_world_pre_process_func" + assert callable(preprocess) + assert model_index["tokenizer"] == ["transformers", "T5TokenizerFast"] + assert model_index["text_encoder"] == ["transformers", "UMT5EncoderModel"] + assert model_index["vae"] == ["diffusers", "AutoencoderKLWan"] + assert model_index["scheduler"] == ["diffusers", "UniPCMultistepScheduler"] + assert model_index["transformer"] == ["diffusers", "CausalLingBotWorldTransformer3DModel"] + assert model_index["image_encoder"] == [None, None] + assert model_index["image_processor"] == [None, None] + assert model_index["transformer_2"] == [None, None] + + +def test_postprocess_returns_latents_without_video_conversion() -> None: + module = _load_pipeline_module() + postprocess = module.get_lingbot_world_post_process_func(_od_config()) + latents = torch.randn(1, 16, 3, 2, 2) + + output = postprocess(latents, sampling_params=SimpleNamespace(output_type="latent")) + + assert output is latents + + +def test_postprocess_uses_the_standard_diffusion_output_envelope(monkeypatch) -> None: + import diffusers.video_processor as video_processor_module + + module = _load_pipeline_module() + processed_video = object() + calls = [] + + class VideoProcessor: + def __init__(self, *, vae_scale_factor): + assert vae_scale_factor == 8 + + def postprocess_video(self, video, *, output_type): + calls.append((video, output_type)) + return processed_video + + monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor) + postprocess = module.get_lingbot_world_post_process_func(_od_config()) + video = torch.randn(1, 3, 9, 8, 8) + + output = postprocess(video, output_type="np") + + assert output == {"payload": {"video": processed_video}, "metadata": {}} + assert calls == [(video, "np")] + + +def test_forward_propagates_pipeline_profiler_stage_durations() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + pipeline._profiler_lock = Lock() + pipeline._stage_durations = {"_generate_block": 0.25, "vae.decode": 0.5} + + output = pipeline(_request()) + + assert output.stage_durations == pipeline.stage_durations + + +def test_pipeline_weight_loader_preserves_component_parameter_prefixes() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + weight = torch.randn(2, 2) + + loaded = pipeline.load_weights([("transformer.blocks.0.weight", weight)]) + + assert loaded == {"transformer.blocks.0.weight"} + assert pipeline.transformer.loaded_weights == [("blocks.0.weight", weight)] + + +def test_request_requires_runner_provided_generator() -> None: + module = _load_pipeline_module() + sampling = _SamplingParams(seed=None, generator=None) + + with pytest.raises(ValueError, match="runner-provided torch.Generator"): + _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling)) + + +@pytest.mark.parametrize("unsupported", ["generator-list", "caller-latents"]) +def test_request_rejects_unsupported_sampling_state(unsupported: str) -> None: + module = _load_pipeline_module() + sampling = _SamplingParams() + if unsupported == "generator-list": + sampling.generator = [sampling.generator] + message = "generator list" + else: + sampling.latents = torch.zeros(1, 16, 3, 2, 2) + message = "caller-provided latents" + + with pytest.raises(ValueError, match=message): + _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling)) + + +def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer(dtype=torch.bfloat16) + pipeline = _pipeline(module, transformer=transformer) + requested_noise_dtypes: list[torch.dtype] = [] + + def randn(shape, *, generator, device, dtype): + del generator + assert device == pipeline.device + requested_noise_dtypes.append(dtype) + return torch.zeros(shape, device=device, dtype=dtype) + + module.randn_tensor = randn + result = pipeline(_request()) + + assert requested_noise_dtypes == [torch.float32] * 4 + assert result.output.dtype == torch.float32 + assert all(call["hidden_states"].dtype == torch.bfloat16 for call in transformer.calls) + + +def test_path_image_rejects_oversized_source_before_decode_or_convert(monkeypatch, tmp_path: Path) -> None: + module = _load_pipeline_module() + source_path = tmp_path / "oversized-compressed.png" + Image.new("1", (4097, 4097)).save(source_path) + decode_calls: list[str] = [] + + def forbidden_convert(*args, **kwargs): + del args, kwargs + decode_calls.append("convert") + raise AssertionError("oversized source reached convert") + + def forbidden_load(*args, **kwargs): + del args, kwargs + decode_calls.append("load") + raise AssertionError("oversized source reached load") + + monkeypatch.setattr(module.PIL.Image.Image, "convert", forbidden_convert) + monkeypatch.setattr(module.PIL.Image.Image, "load", forbidden_load) + + with pytest.raises(ValueError, match="source image.*4096.*4096"): + _preprocess_request(module, prompt=_prompt(images=str(source_path))) + + assert decode_calls == [] + + +def test_normal_path_image_is_decoded_after_source_size_validation(tmp_path: Path) -> None: + module = _load_pipeline_module() + source_path = tmp_path / "normal.png" + Image.new("RGBA", (32, 24), color=(1, 2, 3, 128)).save(source_path) + + request = _preprocess_request(module, prompt=_prompt(images=str(source_path))) + parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params)) + + assert module._MAX_SOURCE_IMAGE_PIXELS == 4096 * 4096 + assert isinstance(parsed.image, Image.Image) + assert parsed.image.mode == "RGB" + assert parsed.image.size == (32, 24) + + +@pytest.mark.parametrize( + ("source_size", "expected_size"), + [ + ((832, 480), (464, 832)), + ((480, 832), (832, 480)), + ((512, 512), (624, 624)), + ], +) +def test_default_resolution_matches_official_480p_aspect_derivation(source_size, expected_size) -> None: + module = _load_pipeline_module() + sampling = _SamplingParams(height=None, width=None) + + parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(images=Image.new("RGB", source_size)), sampling)) + + assert (parsed.height, parsed.width) == expected_size + assert parsed.height % 16 == 0 + assert parsed.width % 16 == 0 + assert parsed.height * parsed.width <= 480 * 832 + + +def test_pil_and_tensor_images_share_official_bicubic_preprocess() -> None: + module = _load_pipeline_module() + height, width = 7, 11 + source = torch.arange(3 * height * width, dtype=torch.uint8).reshape(3, height, width) + pil_image = Image.fromarray(source.permute(1, 2, 0).numpy(), mode="RGB") + sampling = _SamplingParams(height=16, width=32) + + pil_pipeline = _pipeline(module) + tensor_pipeline = _pipeline(module) + pil_inputs = pil_pipeline._parse_request(_RequestBatch(_prompt(images=pil_image), sampling)) + tensor_inputs = tensor_pipeline._parse_request( + _RequestBatch(_prompt(images=source.clone()), _SamplingParams(height=16, width=32)) + ) + pil_condition = pil_pipeline._prepare_image_tensor(pil_inputs.image, height=16, width=32) + tensor_condition = tensor_pipeline._prepare_image_tensor(tensor_inputs.image, height=16, width=32) + expected = torch.nn.functional.interpolate( + source.unsqueeze(0).float() / 255.0, + size=(16, 32), + mode="bicubic", + align_corners=False, + ) + expected = expected * 2.0 - 1.0 + + torch.testing.assert_close(pil_condition, expected) + torch.testing.assert_close(tensor_condition, expected) + + +def test_path_image_decode_error_is_sanitized(monkeypatch) -> None: + module = _load_pipeline_module() + unsafe_path = "/private/source/customer-secret.png" + monkeypatch.setattr( + module.PIL.Image, + "open", + lambda path: (_ for _ in ()).throw(module.PIL.Image.DecompressionBombError(f"unsafe {path}")), + ) + + with pytest.raises(ValueError, match="Unable to load multi_modal_data.image") as exc_info: + _preprocess_request(module, prompt=_prompt(images=unsafe_path)) + + assert unsafe_path not in str(exc_info.value) + assert exc_info.value.__cause__ is None + + +@pytest.mark.parametrize("source_size", [(4096, 4096), (4097, 4097)]) +def test_supplied_pil_image_obeys_documented_source_pixel_ceiling(source_size) -> None: + module = _load_pipeline_module() + source_image = Image.new("1", source_size) + close_calls: list[bool] = [] + source_image.close = lambda: close_calls.append(True) + + if source_size[0] * source_size[1] <= 4096 * 4096: + request = _preprocess_request(module, prompt=_prompt(images=source_image)) + parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params)) + assert parsed.image is not source_image + assert parsed.image.mode == "RGB" + assert close_calls == [] + else: + with pytest.raises(ValueError, match="source image.*4096.*4096"): + _preprocess_request(module, prompt=_prompt(images=source_image)) + assert close_calls == [] + + +def test_supplied_pil_decode_error_is_sanitized_without_closing_caller(monkeypatch) -> None: + module = _load_pipeline_module() + source_image = Image.new("RGB", (16, 16)) + unsafe_detail = "/private/source/customer-secret.png" + close_calls: list[bool] = [] + monkeypatch.setattr(source_image, "convert", lambda mode: (_ for _ in ()).throw(OSError(unsafe_detail))) + monkeypatch.setattr(source_image, "close", lambda: close_calls.append(True)) + + with pytest.raises(ValueError, match="Unable to load multi_modal_data.image") as exc_info: + _preprocess_request(module, prompt=_prompt(images=source_image)) + + assert unsafe_detail not in str(exc_info.value) + assert exc_info.value.__cause__ is None + assert close_calls == [] + + +def test_action_path_resolves_from_sampling_extra_args(tmp_path: Path) -> None: + module = _load_pipeline_module() + action_root = tmp_path / "trusted-actions" + contained_action = action_root / "actions" + contained_action.mkdir(parents=True) + sampling = _SamplingParams(extra_args={"action_path": "actions"}) + prompt = _prompt() + original_prompt = prompt.copy() + trajectory = _CameraTrajectory(torch.eye(4).repeat(9, 1, 1), torch.ones(9, 4)) + resolved_actions = [] + module.load_camera_trajectory = lambda action: resolved_actions.append(action) or trajectory + + request = _preprocess_request( + module, + prompt=prompt, + sampling=sampling, + od_config=_od_config(model_config={"lingbot_action_root": str(action_root)}), + ) + parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params)) + + assert resolved_actions[0].root == action_root.resolve() + assert resolved_actions[0].relative == Path("actions") + assert parsed.camera_trajectory is trajectory + assert sampling.extra_args["action_path"] == "actions" + assert sampling.extra_args["_lingbot_camera_trajectory"] is trajectory + assert prompt == original_prompt + + +@pytest.mark.parametrize("with_extra", [False, True], ids=["additional-only", "extra-and-additional"]) +def test_additional_information_action_path_is_not_an_input_source(with_extra: bool) -> None: + module = _load_pipeline_module() + sampling = _SamplingParams( + extra_args={"action_path": "."} if with_extra else None, + include_action=with_extra, + ) + + with pytest.raises(ValueError, match="additional_information.*action_path.*not supported"): + _preprocess_request(module, prompt=_prompt(action_path="legacy-actions"), sampling=sampling) + + +def test_action_path_is_required_in_sampling_extra_args() -> None: + module = _load_pipeline_module() + + with pytest.raises(ValueError, match="sampling_params.extra_args.action_path"): + _preprocess_request(module, sampling=_SamplingParams(include_action=False)) + + +@pytest.mark.parametrize("escape_kind", ["traversal", "absolute", "symlink"]) +def test_online_action_path_rejects_escape_from_trusted_root(escape_kind: str, tmp_path: Path) -> None: + module = _load_pipeline_module() + root = tmp_path / "trusted" + root.mkdir() + outside = tmp_path / "outside" + outside.mkdir() + if escape_kind == "traversal": + action_path = "../outside" + elif escape_kind == "absolute": + action_path = str(outside) + else: + (root / "escape-link").symlink_to(outside, target_is_directory=True) + action_path = "escape-link" + with pytest.raises(ValueError, match="trusted.*root|contained"): + _preprocess_request( + module, + sampling=_SamplingParams(extra_args={"action_path": action_path}), + od_config=_od_config(model_config={"lingbot_action_root": str(root)}), + ) + + +def test_online_action_path_uses_environment_root_fallback(monkeypatch, tmp_path: Path) -> None: + module = _load_pipeline_module() + root = tmp_path / "trusted" + action_dir = root / "forward" + action_dir.mkdir(parents=True) + monkeypatch.setenv("VLLM_OMNI_LINGBOT_ACTION_ROOT", str(root)) + resolved_actions = [] + module.load_camera_trajectory = lambda action: ( + resolved_actions.append(action) or _CameraTrajectory(torch.eye(4).repeat(9, 1, 1), torch.ones(9, 4)) + ) + + request = _preprocess_request( + module, + sampling=_SamplingParams(extra_args={"action_path": "forward"}), + od_config=_od_config(model_config={}), + ) + + assert request.sampling_params.extra_args["_lingbot_camera_trajectory"] is not None + assert resolved_actions[0].root == root.resolve() + assert resolved_actions[0].relative == Path("forward") + + +def test_online_action_path_error_suppresses_path_bearing_filesystem_cause(tmp_path: Path) -> None: + module = _load_pipeline_module() + root = tmp_path / "trusted" + root.mkdir() + + with pytest.raises(ValueError, match="trusted action root") as exc_info: + _preprocess_request( + module, + sampling=_SamplingParams(extra_args={"action_path": "does-not-exist"}), + od_config=_od_config(model_config={"lingbot_action_root": str(root)}), + ) + + assert str(root) not in str(exc_info.value) + assert exc_info.value.__cause__ is None + + +@pytest.mark.parametrize("source", ["root", "candidate"]) +def test_online_action_path_unknown_user_is_sanitized(source: str, tmp_path: Path) -> None: + module = _load_pipeline_module() + unknown_user_path = "~__vllm_omni_user_that_does_not_exist__/actions" + root = tmp_path / "trusted" + root.mkdir() + configured_root = unknown_user_path if source == "root" else str(root) + action_path = "actions" if source == "root" else unknown_user_path + + with pytest.raises(ValueError, match="trusted action root") as exc_info: + _preprocess_request( + module, + sampling=_SamplingParams(extra_args={"action_path": action_path}), + od_config=_od_config(model_config={"lingbot_action_root": configured_root}), + ) + + assert "__vllm_omni_user_that_does_not_exist__" not in str(exc_info.value) + assert exc_info.value.__cause__ is None + + +def test_action_path_requires_a_trusted_root_for_every_request() -> None: + module = _load_pipeline_module() + + with pytest.raises(ValueError, match="lingbot_action_root|VLLM_OMNI_LINGBOT_ACTION_ROOT"): + _preprocess_request(module, od_config=_od_config(model_config={})) + + +@pytest.mark.parametrize( + ("request_batch", "message"), + [ + (_request(num_reqs=2), "single prompt"), + (_request(prompt=_prompt(images=[Image.new("RGB", (16, 16))])), "image.*list"), + (_request(prompt=_prompt(images=[])), "image.*list"), + ( + _request(prompt=_prompt(images=[Image.new("RGB", (16, 16)), Image.new("RGB", (16, 16))])), + "image.*list", + ), + (_request(sampling=_SamplingParams(num_outputs_per_prompt=2)), "num_outputs_per_prompt"), + (_request(sampling=_SamplingParams(num_inference_steps=5)), "num_inference_steps"), + (_request(sampling=_SamplingParams(height=15)), "height.*divisible"), + (_request(sampling=_SamplingParams(width=15)), "width.*divisible"), + (_request(sampling=_SamplingParams(height=None, width=16)), "height and width.*both"), + (_request(sampling=_SamplingParams(num_frames=13)), "num_frames.*three-frame"), + ], +) +def test_request_validation_rejects_unsupported_contracts(request_batch, message: str) -> None: + module = _load_pipeline_module() + + with pytest.raises(ValueError, match=message): + _pipeline(module)._parse_request(request_batch) + + +def test_resource_limits_accept_exact_documented_boundaries() -> None: + module = _load_pipeline_module() + sampling = _SamplingParams(height=480, width=832, num_frames=117, max_sequence_length=512) + + parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling)) + + assert (parsed.height, parsed.width, parsed.num_frames, parsed.max_sequence_length) == (480, 832, 117, 512) + + +@pytest.mark.parametrize( + ("sampling", "message"), + [ + (_SamplingParams(height=480, width=848), "pixel area|480.*832"), + (_SamplingParams(num_frames=129), "num_frames.*117"), + (_SamplingParams(max_sequence_length=511), "max_sequence_length.*512"), + (_SamplingParams(max_sequence_length=513), "max_sequence_length.*512"), + (_SamplingParams(max_sequence_length=512.0), "max_sequence_length.*512"), + ], +) +def test_resource_limits_reject_oversize_before_any_component_call(sampling, message: str) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + calls: list[str] = [] + pipeline.encode_prompt = lambda *args, **kwargs: calls.append("text") + pipeline._prepare_condition = lambda *args, **kwargs: calls.append("vae") + pipeline._prepare_camera = lambda *args, **kwargs: calls.append("camera") + pipeline.transformer.allocate_cache = lambda *args, **kwargs: calls.append("cache") + + with pytest.raises(ValueError, match=message): + pipeline(_RequestBatch(_prompt(), sampling)) + + assert calls == [] + + +def test_request_validation_rejects_insufficient_camera_frames() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + short_trajectory = _CameraTrajectory( + poses=torch.eye(4).repeat(8, 1, 1), + intrinsics=torch.ones(8, 4), + ) + sampling = _SamplingParams(extra_args={"_lingbot_camera_trajectory": short_trajectory}) + + with pytest.raises(ValueError, match="camera.*frames.*num_frames"): + pipeline(_request(sampling=sampling)) + + +def test_camera_load_error_is_actionable_without_echoing_path_contents() -> None: + module = _load_pipeline_module() + module.load_camera_trajectory = lambda path: (_ for _ in ()).throw(FileNotFoundError(f"missing {path}/poses.npy")) + + with pytest.raises(ValueError, match="camera trajectory.*action_path") as exc_info: + _preprocess_request(module) + + assert str(_ROOT) not in str(exc_info.value) + assert exc_info.value.__cause__ is None + + +def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + module.randn_tensor = lambda shape, **kwargs: torch.full( + shape, + -99.0, + device=kwargs["device"], + dtype=kwargs["dtype"], + ) + + result = pipeline(_request()) + + assert result.output.shape == (1, 16, 3, 2, 2) + first_input = transformer.calls[0]["hidden_states"] + assert first_input.shape == (1, 36, 3, 2, 2) + mean = torch.tensor(pipeline.vae.config.latents_mean).view(1, 16, 1, 1) + std = torch.tensor(pipeline.vae.config.latents_std).view(1, 16, 1, 1) + expected_first = (torch.full((1, 16, 2, 2), 2.0) - mean) / std + expected_future = (torch.zeros(1, 16, 2, 2, 2) - mean.unsqueeze(2)) / std.unsqueeze(2) + torch.testing.assert_close(first_input[:, :16], torch.full((1, 16, 3, 2, 2), -99.0)) + torch.testing.assert_close(first_input[:, 16:20, 0], torch.ones(1, 4, 2, 2)) + torch.testing.assert_close(first_input[:, 16:20, 1:], torch.zeros(1, 4, 2, 2, 2)) + torch.testing.assert_close(first_input[:, 20:36, 0], expected_first) + torch.testing.assert_close(first_input[:, 20:36, 1:], expected_future) + + raw_camera = module.build_plucker_embedding( + SimpleNamespace(poses=torch.empty(3, 4, 4)), + height=16, + width=16, + target_height=16, + target_width=16, + device=torch.device("cpu"), + dtype=torch.float32, + ) + expected_camera = module._fold_camera_embedding(raw_camera) + reference_camera = torch.nn.functional.pixel_unshuffle(raw_camera, 8).permute(1, 0, 2, 3).unsqueeze(0) + torch.testing.assert_close(expected_camera, reference_camera) + torch.testing.assert_close(transformer.calls[0]["camera_hidden_states"], expected_camera) + assert expected_camera.shape == (1, 384, 3, 2, 2) + + +def test_fixed_dmd_transition_and_cache_commit_trace() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + + result = pipeline(_request()) + + generator = torch.Generator(device="cpu").manual_seed(17) + current = torch.randn((1, 16, 3, 2, 2), generator=generator) + warped_schedule = ( + (1000.0, 1.0), + (937.5, 0.9375), + (2500.0 / 3.0, 5.0 / 6.0), + (625.0, 0.625), + ) + for index, (_, sigma) in enumerate(warped_schedule): + x0 = current - sigma + if index + 1 < len(module.LINGBOT_DMD_TIMESTEPS): + next_sigma = warped_schedule[index + 1][1] + noise = torch.randn(current.shape, generator=generator) + current = (1.0 - next_sigma) * x0 + next_sigma * noise + else: + current = x0 + torch.testing.assert_close(result.output, current) + + torch.testing.assert_close( + torch.cat([call["timestep"] for call in transformer.calls]), + torch.tensor([*(timestep for timestep, _ in warped_schedule), 0.0]), + ) + assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] + assert [call["start_frame"] for call in transformer.calls] == [0, 0, 0, 0, 0] + assert len({call["cache_id"] for call in transformer.calls}) == 1 + torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], result.output) + + +@pytest.mark.parametrize( + ("flow_shift", "expected_timesteps"), + [ + (5.0, (1000.0, 937.5, 2500.0 / 3.0, 625.0)), + (2.5, (1000.0, 15000.0 / 17.0, 5000.0 / 7.0, 5000.0 / 11.0)), + ], +) +def test_request_flow_shift_warps_transformer_timesteps(flow_shift: float, expected_timesteps) -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + + pipeline( + _request( + sampling=_SamplingParams(extra_args={"action_path": ".", "flow_shift": flow_shift}), + ) + ) + + torch.testing.assert_close( + torch.cat([call["timestep"] for call in transformer.calls[:4]]), + torch.tensor(expected_timesteps), + ) + + +def test_request_tiny_positive_flow_shift_keeps_schedule_finite() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + + pipeline( + _request( + sampling=_SamplingParams(extra_args={"action_path": ".", "flow_shift": 1e-20}), + ) + ) + + timesteps = torch.cat([call["timestep"] for call in transformer.calls[:4]]) + assert torch.isfinite(timesteps).all() + assert timesteps[0].item() == 1000.0 + + +def test_request_flow_shift_override_has_precedence_without_mutating_scheduler() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module, od_config=_od_config(flow_shift=3.0)) + + default_inputs = pipeline._parse_request(_request()) + sampling = _SamplingParams(extra_args={"flow_shift": 2.0}) + override_inputs = pipeline._parse_request(_RequestBatch(_prompt(), sampling)) + + assert default_inputs.flow_shift == 3.0 + assert override_inputs.flow_shift == 2.0 + assert sampling.extra_args["action_path"] == "." + assert sampling.extra_args["flow_shift"] == 2.0 + assert pipeline.scheduler.config.shift == 3.0 + + +def test_checkpoint_flow_shift_default_and_engine_request_override_precedence() -> None: + module = _load_pipeline_module() + checkpoint_scheduler_config = _scheduler_config(flow_shift=6.25, shift=7.5) + module._load_json = lambda *args, **kwargs: checkpoint_scheduler_config.copy() + + checkpoint_default = _pipeline(module) + engine_override = _pipeline(module, od_config=_od_config(flow_shift=3.5)) + request_override = checkpoint_default._parse_request( + _RequestBatch( + _prompt(), + _SamplingParams(extra_args={"flow_shift": 2.25}), + ) + ) + + assert checkpoint_default.scheduler.config.shift == 6.25 + assert checkpoint_default._parse_request(_request()).flow_shift == 6.25 + assert engine_override.scheduler.config.shift == 3.5 + assert engine_override._parse_request(_request()).flow_shift == 3.5 + assert request_override.flow_shift == 2.25 + + +@pytest.mark.parametrize( + ("scheduler_config", "expected"), + [ + ({"flow_shift": 6.0, "shift": 7.0}, 6.0), + ({"shift": 7.0}, 7.0), + ({}, 5.0), + ], +) +def test_checkpoint_scheduler_shift_fallback_order(scheduler_config, expected: float) -> None: + module = _load_pipeline_module() + checkpoint_config = _scheduler_config() + checkpoint_config.pop("flow_shift") + checkpoint_config.update(scheduler_config) + module._load_json = lambda *args, **kwargs: checkpoint_config.copy() + + pipeline = _pipeline(module) + + assert pipeline.scheduler.config.shift == expected + + +@pytest.mark.parametrize("flow_shift", [0, -1, float("nan"), float("inf"), "invalid"]) +def test_request_flow_shift_must_be_positive_and_finite(flow_shift) -> None: + module = _load_pipeline_module() + sampling = _SamplingParams(extra_args={"flow_shift": flow_shift}) + + with pytest.raises(ValueError, match="flow_shift.*positive.*finite"): + _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling)) + + +@pytest.mark.parametrize( + "flow_shift", + [True, False, 10**10000], + ids=["true", "false", "overflowing-integer"], +) +def test_request_flow_shift_rejects_booleans_and_normalizes_integer_overflow(flow_shift) -> None: + module = _load_pipeline_module() + sampling = _SamplingParams(extra_args={"flow_shift": flow_shift}) + + with pytest.raises(ValueError, match="flow_shift.*positive.*finite"): + _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling)) + + +def test_flow_shift_is_request_local_without_mutating_scheduler() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + + def generate(flow_shift=None): + extra_args = {} if flow_shift is None else {"flow_shift": flow_shift} + return pipeline( + _RequestBatch( + _prompt(), + _SamplingParams(num_frames=21, extra_args=extra_args), + ) + ).output + + default_before = generate() + shifted_two = generate(2.0) + shifted_seven = generate(7.0) + default_after = generate() + + torch.testing.assert_close(default_before, default_after) + assert not torch.equal(default_before, shifted_two) + assert not torch.equal(shifted_two, shifted_seven) + assert pipeline.scheduler.config.shift == 5.0 + + +def test_encode_prompt_zeroes_padded_umt5_states_to_exactly_512_tokens() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + attention_mask = torch.zeros(1, 512, dtype=torch.long) + attention_mask[:, :3] = 1 + pipeline.tokenizer = lambda *args, **kwargs: SimpleNamespace( + input_ids=torch.arange(512).view(1, 512), + attention_mask=attention_mask, + ) + raw_states = torch.arange(512 * 8, dtype=torch.float32).view(1, 512, 8) + 1.0 + pipeline.text_encoder = lambda input_ids, mask: SimpleNamespace(last_hidden_state=raw_states.clone()) + + encoded = module.LingBotWorldCausalDMDPipeline.encode_prompt( + pipeline, + "move", + max_sequence_length=512, + dtype=torch.float32, + ) + + assert encoded.shape == (1, 512, 8) + torch.testing.assert_close(encoded[:, :3], raw_states[:, :3]) + torch.testing.assert_close(encoded[:, 3:], torch.zeros_like(encoded[:, 3:])) + + +def test_multi_chunk_generation_uses_one_request_local_cache_and_decodes_accumulated_latents() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + allocations = [] + original_allocate = transformer.allocate_cache + + def allocate(**kwargs): + cache = original_allocate(**kwargs) + allocations.append((kwargs, weakref.ref(cache))) + return cache + + transformer.allocate_cache = allocate + sampling = _SamplingParams(num_frames=21, output_type="np") + + result = pipeline(_request(sampling=sampling)) + + assert result.output.shape == (1, 3, 21, 16, 16) + assert len(transformer.calls) == 10 + assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5 + expected_timesteps = torch.tensor([1000.0, 937.5, 2500.0 / 3.0, 625.0, 0.0] * 2) + torch.testing.assert_close( + torch.cat([call["timestep"] for call in transformer.calls]), + expected_timesteps, + ) + assert len(allocations) == 1 + kwargs, cache_ref = allocations[0] + assert kwargs == { + "batch_size": 1, + "latent_height": 2, + "latent_width": 2, + "device": torch.device("cpu"), + "dtype": torch.float32, + } + assert not hasattr(pipeline, "cache") + assert not hasattr(pipeline, "transformer_cache") + assert pipeline.vae.decode_inputs[0].shape == (1, 16, 6, 2, 2) + normalized_latents = torch.cat( + (transformer.calls[4]["hidden_states"][:, :16], transformer.calls[9]["hidden_states"][:, :16]), + dim=2, + ) + mean = torch.tensor(pipeline.vae.config.latents_mean).view(1, 16, 1, 1, 1) + std = torch.tensor(pipeline.vae.config.latents_std).view(1, 16, 1, 1, 1) + torch.testing.assert_close(pipeline.vae.decode_inputs[0], normalized_latents * std + mean) + transformer.calls.clear() + gc.collect() + assert cache_ref() is None + + +def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room"): + return ARDiffusionTickRequest( + session_id="world-1", + request_id="legacy-lingbot-request", + chunk_index=chunk_index, + applied_event_ids=(chunk_index,), + prompt=prompt, + controls=( + ARDiffusionControlInput( + track="camera", + schema="lingbot.camera_trajectory.v1", + data={"poses": [], "intrinsics": []}, + ), + ), + ).to_extra_args() + + +def test_typed_ticks_generate_one_global_block_and_return_standard_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + pipeline._ar_diffusion_kv_state = object() + cross_calls = [] + generated = [] + + def ar_text_caches(prompt_embeds, *, invalidate): + del prompt_embeds + cross_calls.append(invalidate) + return [SimpleNamespace()] + + def generate_block(**kwargs): + block = torch.randn( + (1, 16, 3, 2, 2), + generator=kwargs["generator"], + ) + generated.append( + { + "start_frame": kwargs["start_frame"], + "condition": kwargs["condition"].clone(), + "block": block.clone(), + } + ) + return block + + monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches) + monkeypatch.setattr(pipeline, "_generate_block", generate_block) + + first_sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=0)) + first = pipeline(_request(sampling=first_sampling)) + switched_prompt = _prompt() + switched_prompt["prompt"] = "enter the snowy valley" + second_sampling = _SamplingParams( + extra_args=_tick_extra_args( + chunk_index=1, + prompt=switched_prompt["prompt"], + ) + ) + second = pipeline(_request(sampling=second_sampling, prompt=switched_prompt)) + + assert generated[0]["start_frame"] == 0 + assert generated[1]["start_frame"] == 3 + assert torch.count_nonzero(generated[0]["condition"]) > 0 + torch.testing.assert_close( + generated[1]["condition"][:, :4], + torch.zeros_like(generated[1]["condition"][:, :4]), + ) + assert torch.count_nonzero(generated[1]["condition"][:, 4:]) > 0 + assert cross_calls == [False, True] + assert first.output["payload"]["latents"].shape == (1, 16, 3, 2, 2) + assert second.output["metadata"]["ar_diffusion"] == { + "session_id": "world-1", + "request_id": "legacy-lingbot-request", + "chunk_index": 1, + "applied_event_ids": [1], + } + assert pipeline._ar_sessions["world-1"].next_chunk_index == 2 + assert pipeline._ar_sessions["world-1"].prompt == "enter the snowy valley" + assert pipeline._ar_sessions["world-1"].camera_tail is not None + assert pipeline._ar_sessions["world-1"].camera_tail.poses.shape == (1, 4, 4) + expected_generator = torch.Generator(device="cpu").manual_seed(17) + expected_first = torch.randn((1, 16, 3, 2, 2), generator=expected_generator) + expected_second = torch.randn((1, 16, 3, 2, 2), generator=expected_generator) + torch.testing.assert_close(generated[0]["block"], expected_first) + torch.testing.assert_close(generated[1]["block"], expected_second) + + +def test_typed_action_ticks_integrate_camera_across_chunks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + pipeline._ar_diffusion_kv_state = object() + monkeypatch.setattr( + pipeline, + "_ar_text_caches", + lambda *args, **kwargs: [SimpleNamespace()], + ) + monkeypatch.setattr( + pipeline, + "_generate_block", + lambda **kwargs: torch.zeros_like(kwargs["condition"][:, :16]), + ) + + def action_sampling(chunk_index: int, frames: list[list[str]]) -> _SamplingParams: + tick = ARDiffusionTickRequest( + session_id="world-actions", + request_id="legacy-lingbot-request", + chunk_index=chunk_index, + controls=( + ARDiffusionControlInput( + track="camera", + schema="lingbot.camera_actions.v1", + data={"mode": "frames", "frames": frames}, + ), + ), + ) + sampling = _SamplingParams(extra_args=tick.to_extra_args()) + sampling.extra_args["_lingbot_camera_trajectory"] = None + sampling.extra_args["_lingbot_camera_actions"] = tuple(tuple(frame) for frame in frames) + return sampling + + pipeline(_request(sampling=action_sampling(0, [["w"], ["w"], ["w"]]))) + first_tail = pipeline._ar_sessions["world-actions"].camera_tail.poses.clone() + pipeline(_request(sampling=action_sampling(1, [["d"], ["d"], ["d"]]))) + state = pipeline._ar_sessions["world-actions"] + + assert state.next_chunk_index == 2 + assert state.camera_tail is not None + torch.testing.assert_close(state.camera_tail.poses[0, 2, 3], first_tail[0, 2, 3]) + assert state.camera_tail.poses[0, 0, 3] > first_tail[0, 0, 3] + + +def test_first_typed_yaw_action_uses_pre_action_identity_anchor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + monkeypatch.setattr( + module, + "build_plucker_embedding", + _real_build_plucker_embedding, + ) + action_trajectory, _ = integrate_lingbot_camera_actions( + [["j"], [], []], + width=16, + height=16, + ) + neutral_trajectory, _ = integrate_lingbot_camera_actions( + [[], [], []], + width=16, + height=16, + ) + + def inputs(trajectory, actions): + return SimpleNamespace( + camera_trajectory=trajectory, + camera_actions=actions, + num_latent_frames=3, + num_frames=9, + height=16, + width=16, + ) + + action_embedding, tail = pipeline._prepare_camera( + inputs(action_trajectory, (("j",), (), ())), + dtype=torch.float32, + ) + neutral_embedding, _ = pipeline._prepare_camera( + inputs(neutral_trajectory, ((), (), ())), + dtype=torch.float32, + ) + + identity_anchor = torch.eye(4, dtype=action_trajectory.poses.dtype).unsqueeze(0) + explicit_trajectory = _CameraTrajectory( + poses=torch.cat((identity_anchor, action_trajectory.poses), dim=0), + intrinsics=torch.cat( + (action_trajectory.intrinsics[:1], action_trajectory.intrinsics), + dim=0, + ), + ) + explicit = _real_build_plucker_embedding( + explicit_trajectory, + height=16, + width=16, + target_height=16, + target_width=16, + device=torch.device("cpu"), + dtype=torch.float32, + )[1:] + + assert not torch.equal(action_embedding, neutral_embedding) + torch.testing.assert_close(action_embedding, module._fold_camera_embedding(explicit)) + torch.testing.assert_close(tail.poses, action_trajectory.poses[-1:]) + + +def test_typed_tick_rejects_non_contiguous_chunk_index( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + pipeline._ar_diffusion_kv_state = object() + monkeypatch.setattr( + pipeline, + "_ar_text_caches", + lambda *args, **kwargs: [SimpleNamespace()], + ) + + sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=2)) + with pytest.raises(ValueError, match="must be contiguous"): + pipeline(_request(sampling=sampling)) + + +def test_typed_tick_rejects_chunk_beyond_realtime_condition_horizon( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + pipeline._ar_diffusion_kv_state = object() + pipeline._ar_sessions["world-1"] = module._LingBotARSessionState( + next_chunk_index=10, + ) + monkeypatch.setattr( + pipeline, + "encode_prompt", + lambda *args, **kwargs: pytest.fail("the realtime horizon must be validated before prompt encoding"), + ) + + sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=10)) + with pytest.raises( + ValueError, + match=r"at most 10 ticks.*chunk_index 0 through 9.*117 pixel frames", + ): + pipeline(_request(sampling=sampling)) + + +def test_request_cache_is_released_before_vae_decode() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + cache_refs: list[weakref.ReferenceType] = [] + original_allocate = transformer.allocate_cache + + def allocate(**kwargs): + cache = original_allocate(**kwargs) + cache_refs.append(weakref.ref(cache)) + return cache + + transformer.allocate_cache = allocate + + def assert_cache_released() -> None: + gc.collect() + assert len(cache_refs) == 1 + assert cache_refs[0]() is None + + pipeline.vae.on_decode = assert_cache_released + + result = pipeline(_request(sampling=_SamplingParams(num_frames=21, output_type="np"))) + + assert result.output.shape[2] == 21 + + +def test_117_frame_request_generates_ten_complete_latent_blocks() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + trajectory = _CameraTrajectory( + poses=torch.eye(4).repeat(117, 1, 1), + intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(117, 1), + ) + sampling = _SamplingParams( + num_frames=117, + output_type="latent", + extra_args={"_lingbot_camera_trajectory": trajectory}, + ) + + result = pipeline(_request(sampling=sampling)) + + assert result.output.shape == (1, 16, 30, 2, 2) + assert len(transformer.calls) == 50 + assert [call["start_frame"] for call in transformer.calls[::5]] == list(range(0, 30, 3)) + + +def test_request_cache_becomes_unreachable_after_transformer_error() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer(raise_on_call=2) + pipeline = _pipeline(module, transformer=transformer) + cache_refs = [] + original_allocate = transformer.allocate_cache + + def allocate(**kwargs): + cache = original_allocate(**kwargs) + cache_refs.append(weakref.ref(cache)) + return cache + + transformer.allocate_cache = allocate + + with pytest.raises(RuntimeError, match="forced transformer failure") as exc_info: + pipeline(_request()) + + exc_info.value.__traceback__ = None + del exc_info + transformer.calls.clear() + gc.collect() + assert len(cache_refs) == 1 + assert cache_refs[0]() is None + assert not hasattr(pipeline, "cache") + assert not hasattr(pipeline, "transformer_cache") + + +def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None: + module = _load_pipeline_module() + resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry( + module + ) + + assert entry == ("lingbot_world", "pipeline", "LingBotWorldCausalDMDPipeline") + assert resolved is module.LingBotWorldCausalDMDPipeline + assert cache_acceleration_disabled + assert preprocess_name == "get_lingbot_world_pre_process_func" + request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams()) + assert preprocess(request) is request + assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory) + assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline" + lingbot_init = _LINGBOT_INIT_PATH.read_text() + assert "from .pipeline import" in lingbot_init + assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init + assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py new file mode 100644 index 000000000..4b1fef7d5 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py @@ -0,0 +1,1041 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Focused LTX-2.5 pipeline correctness tests.""" + +import json +from dataclasses import replace +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime +from vllm_omni.diffusion.models.ltx2.ltx2_components import ( + LTX25_DISTILLED_COMPONENT_PROFILE, + LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, + LTX25_FULL_COMPONENT_PROFILE, + LTX25_TWO_STAGE_COMPONENT_PROFILE, + detect_ltx_model_version, + resolve_ltx_checkpoint_kind, + resolve_ltx_component_profile, +) +from vllm_omni.diffusion.models.ltx2.ltx2_denoise import ( + LTXPhaseResult, + _first_frame_keyframes_mask, + _official_ltx_sigmas, +) +from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( + LTX_GUIDANCE_EXECUTOR, + LTXGuidanceSpec, + LTXModalityGuidance, + velocity_from_x0, + x0_from_velocity, +) +from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState +from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX2_ONE_STAGE_RECIPE, + LTX23_ONE_STAGE_RECIPE, + LTX25_DISTILLED_ONE_STAGE_RECIPE, + LTX25_DISTILLED_TWO_STAGE_RECIPE, + LTX25_FULL_RECIPE, + LTX25_TWO_STAGE_RECIPE, + LTX_DISTILLED_SIGMAS, + LTX_POSITIVE_ONLY_RECIPE, + resolve_ltx_pipeline_recipe, +) +from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs +from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2DistilledOneStagePipeline, + LTX2Pipeline, +) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( + LTX2DistilledTwoStagePipeline, + LTX2TwoStagePipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"]) +def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version): + pipe = object.__new__(LTX2Pipeline) + pipe.model_version = model_version + pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)) + ts = torch.tensor([750.0, 500.0]) + + kwargs = LTXRuntime._denoise_timestep_kwargs( + pipe, + ts, + SimpleNamespace(), + SimpleNamespace(), + video_token_count=4, + audio_token_count=3, + ) + + torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4)) + torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3)) + assert kwargs["sigma"] is ts + assert kwargs["audio_sigma"] is ts + + +def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch): + state = LTXAVState( + video=torch.zeros(1, 2, 2), + audio=torch.zeros(1, 4, 2), + ) + forward_ctx = SimpleNamespace( + original_audio_num_frames=3, + sampler="euler_ancestral", + ) + denoise_ctx = SimpleNamespace(latents=None, audio_latents=None) + pipe = SimpleNamespace( + _predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio)) + ) + + def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args): + assert actual_forward_ctx.sampler == "euler_ancestral" + updated_audio = torch.ones_like(state.audio) + updated_audio[:, 3:] = 99.0 + return torch.ones_like(state.video), updated_audio + + monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step) + + actual = LTXRuntime._denoise_step( + pipe, + 0, + torch.tensor(1.0), + state, + forward_ctx, + denoise_ctx, + ) + + torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3])) + torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:])) + + +def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding(): + reference = torch.zeros(2, 12, 8) + + mask = _first_frame_keyframes_mask(reference, latent_num_frames=3) + + assert mask.shape == (2, 12, 1) + torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1)) + torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1)) + + +def _make_ltx_request_pipe(cls): + pipe = object.__new__(cls) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 99 + pipe.vae_spatial_compression_ratio = 32 + pipe.vae_temporal_compression_ratio = 8 + return pipe + + +def _resolve_request_inputs_for_test( + pipe, + req, + *, + guidance_scale=4.0, + negative_prompt=None, +): + return pipe._resolve_request_inputs( + req, + prompt=None, + negative_prompt=negative_prompt, + height=None, + width=None, + num_frames=None, + frame_rate=None, + num_inference_steps=None, + guidance_scale=guidance_scale, + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=None, + ) + + +def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch): + pipe = SimpleNamespace( + component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None), + use_diffusion_decoder=True, + ) + od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16) + + monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None) + monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object()) + + with pytest.raises(ImportError) as exc_info: + ltx2_components.initialize_pipeline_components(pipe, od_config) + + assert str(exc_info.value) == ( + "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15." + ) + + +def test_ltx_converted_component_loading_propagates_revision(monkeypatch): + revision = "pinned-revision" + calls: dict[str, Any] = {"components": []} + natten_processor = object() + + class FakeDiffusionDecoder: + def set_attn_processor(self, processor): + calls["diffusion_decoder_processor"] = processor + + def enable_tiling(self): + calls["diffusion_decoder_tiling"] = True + + def set_parallel_size(self, parallel_size, mode): + calls["diffusion_decoder_parallel"] = (parallel_size, mode) + + profile = replace( + LTX25_DISTILLED_COMPONENT_PROFILE, + text_encoder_cls=object, + video_vae_cls=object, + vocoder_cls=object, + vocoder_fallback_cls=None, + scheduler_use_dynamic_shifting=False, + ) + pipeline = SimpleNamespace( + component_profile=profile, + pipeline_kind="distilled_two_stage", + use_diffusion_decoder=True, + ) + od_config = SimpleNamespace( + model="org/converted-ltx25", + revision=revision, + dtype=torch.bfloat16, + quantization_config=None, + extras={}, + vae_use_tiling=False, + parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"), + ) + + monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None) + monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None) + + def fake_prefetch(model, subfolders, **kwargs): + calls["prefetch"] = (model, tuple(subfolders), kwargs) + + def fake_tokenizer(*args, **kwargs): + calls["tokenizer"] = (args, kwargs) + return SimpleNamespace(model_max_length=1024) + + def fake_component(component_cls, model, subfolder, **kwargs): + calls["components"].append((component_cls, model, subfolder, kwargs)) + if subfolder == "vae": + return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8) + if subfolder == "audio_vae": + return SimpleNamespace( + mel_compression_ratio=4, + temporal_compression_ratio=4, + config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160), + ) + return object() + + def fake_native_diffusion_decoder(model, **kwargs): + calls["native_diffusion_decoder"] = (model, kwargs) + return FakeDiffusionDecoder() + + def fake_transformer_config(model, subfolder, local_files_only, *, revision): + calls["transformer_config"] = (model, subfolder, local_files_only, revision) + return {"component": "transformer"} + + monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch) + monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer) + monkeypatch.setattr(ltx2_components, "_load_component", fake_component) + monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder) + monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config) + monkeypatch.setattr( + ltx2_components, + "LTX2VideoVaeNeighborhoodNattenProcessor", + lambda: natten_processor, + ) + monkeypatch.setattr( + ltx2_components, + "create_transformer_from_config", + lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)), + ) + + def fake_scheduler(*args, **kwargs): + calls["scheduler"] = (args, kwargs) + return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None}) + + monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler) + + ltx2_components.initialize_pipeline_components(pipeline, od_config) + + assert pipeline.weights_sources[0].revision == revision + assert calls["prefetch"][2]["revision"] == revision + assert calls["tokenizer"][1]["revision"] == revision + assert {call[2] for call in calls["components"]} == { + "text_encoder", + "connectors", + "vae", + "audio_vae", + "vocoder", + "latent_upsampler", + } + assert all(call[3]["revision"] == revision for call in calls["components"]) + assert "diffusion_decoder" not in calls["prefetch"][1] + assert calls["native_diffusion_decoder"] == ( + od_config.model, + {"local_files_only": False, "dtype": torch.bfloat16, "revision": revision}, + ) + assert calls["diffusion_decoder_processor"] is natten_processor + assert calls["diffusion_decoder_tiling"] is True + assert calls["diffusion_decoder_parallel"] == (2, "tile") + assert calls["transformer_config"] == ( + od_config.model, + profile.transformer_subfolder, + False, + revision, + ) + assert calls["scheduler"][1]["revision"] == revision + + +def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch): + revision = "diffusers-revision" + config = {"decoder_stage_channels": [8, 8, 8, 8, 8]} + calls: dict[str, Any] = {} + + class FakeDecoder: + def init_distributed(self): + calls["init_distributed"] = True + + expected = FakeDecoder() + + def fake_load_config(model, **kwargs): + calls["config"] = (model, kwargs) + return config + + def fake_resolve(model, repo_id, filename, **kwargs): + calls["artifact"] = (model, repo_id, filename, kwargs) + return "/models/native-diffvae.safetensors" + + def fake_load(cls, path, passed_config, dtype): + calls["load"] = (path, passed_config, dtype) + return expected + + monkeypatch.setattr( + ltx2_components.DistributedLTX2VideoDiffusionDecoderModel, + "load_config", + staticmethod(fake_load_config), + ) + monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve) + monkeypatch.setattr( + ltx2_components.DistributedLTX2VideoDiffusionDecoderModel, + "from_ltx25_native_checkpoint", + classmethod(fake_load), + ) + + actual = ltx2_components._load_ltx25_native_diffusion_decoder( + "Lightricks/LTX-2.5-Diffusers", + local_files_only=False, + dtype=torch.bfloat16, + revision=revision, + ) + + assert actual is expected + assert calls["config"] == ( + "Lightricks/LTX-2.5-Diffusers", + { + "subfolder": "diffusion_decoder", + "local_files_only": False, + "revision": revision, + }, + ) + assert calls["artifact"] == ( + "Lightricks/LTX-2.5-Diffusers", + "Lightricks/LTX-2.5", + "vae/ltx-2.5-video-vae-bf16.safetensors", + { + "model_revision": revision, + "artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION, + }, + ) + assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16) + assert calls["init_distributed"] is True + + +def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch): + class MissingNattenProcessor: + def __init__(self): + raise ImportError("kernels is unavailable") + + monkeypatch.setattr( + ltx2_components, + "LTX2VideoVaeNeighborhoodNattenProcessor", + MissingNattenProcessor, + ) + + with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"): + ltx2_components._create_ltx25_natten_processor() + + +def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path): + (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"})) + (tmp_path / "transformer").mkdir() + (tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False})) + + assert detect_ltx_model_version(str(tmp_path)) == "2.3" + + +def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text( + json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]}) + ) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False)) + + assert pipe.model_version == "2.5" + assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX25_FULL_RECIPE + assert pipe.preserve_sp_padded_audio_duration + assert pipe.reports_stage_durations + + +def test_ltx25_full_recipe_matches_official_sft_defaults(): + (phase,) = LTX25_FULL_RECIPE.phases + + assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960) + assert LTX25_FULL_RECIPE.num_frames == 121 + assert LTX25_FULL_RECIPE.frame_rate == 24.0 + assert LTX25_FULL_RECIPE.num_inference_steps == 30 + assert phase.guidance.video == LTXModalityGuidance( + cfg_scale=3.0, + stg_scale=1.0, + modality_scale=3.0, + rescale_scale=0.7, + stg_blocks=(28,), + ) + assert phase.guidance.audio == LTXModalityGuidance( + cfg_scale=7.0, + stg_scale=1.0, + modality_scale=3.0, + rescale_scale=0.7, + stg_blocks=(28,), + ) + assert phase.sigmas is None + assert phase.noise_scale == 1.0 + assert phase.use_official_sigma_schedule + assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full" + assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting + assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1 + + +def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor(): + scheduler = SimpleNamespace( + config={ + "base_image_seq_len": 1024, + "max_image_seq_len": 4096, + "base_shift": 0.95, + "max_shift": 2.05, + "shift_terminal": 0.1, + } + ) + + sigmas = _official_ltx_sigmas( + scheduler, + steps=30, + device=torch.device("cpu"), + ) + + torch.testing.assert_close( + sigmas[[0, 1, -2, -1]], + torch.tensor([1.0, 0.9949570, 0.1, 0.0]), + rtol=1e-5, + atol=1e-6, + ) + + +def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text( + json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]}) + ) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = LTX2DistilledTwoStagePipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + enable_diffusion_pipeline_profiler=False, + ) + ) + + assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE + + +def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"})) + for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")): + path = tmp_path / subfolder + path.mkdir() + (path / "config.json").write_text(json.dumps({"marker": marker})) + + config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline") + config.enrich_config() + + assert config.tf_model_config.get("marker") == "full" + + +def test_ltx25_four_public_pipeline_semantics_are_disjoint(): + assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE + assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE + assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE + assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE + + assert LTX2Pipeline.pipeline_kind == "one_stage" + assert LTX2TwoStagePipeline.pipeline_kind == "two_stage" + assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage" + assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage" + + +@pytest.mark.parametrize( + "pipeline_cls", + [ + LTX2Pipeline, + LTX2TwoStagePipeline, + LTX2DistilledOneStagePipeline, + LTX2DistilledTwoStagePipeline, + ], +) +@pytest.mark.parametrize( + ("extras", "use_diffusion_decoder"), + [({}, True), ({"ltx2_use_conv_vae": True}, False)], +) +def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in( + tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder +): + (tmp_path / "model_index.json").write_text( + json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]}) + ) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = pipeline_cls( + od_config=SimpleNamespace( + model=str(tmp_path), + extras=extras, + enable_diffusion_pipeline_profiler=False, + ) + ) + + assert pipe.use_diffusion_decoder is use_diffusion_decoder + assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder) + + +@pytest.mark.parametrize( + ("pipeline_kind", "expected"), + [ + ("one_stage", "regular"), + ("two_stage", "regular"), + ("distilled_one_stage", "distilled"), + ("distilled_two_stage", "distilled"), + ("dmd2", None), + ], +) +def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected): + assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected + + +def test_ltx_checkpoint_kind_rejects_unknown_pipeline(): + with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"): + resolve_ltx_checkpoint_kind("unknown") + + +def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution(): + stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases + + assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920) + assert stage1.spatial_downscale == 2 + assert stage1.sigmas == LTX_DISTILLED_SIGMAS + assert stage1.sampler == "euler_ancestral" + assert stage2.sampler == "euler" + assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0) + assert stage2.input_transform == "spatial_upsample" + + assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == "" + assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt + + +def test_ltx25_distilled_one_stage_uses_native_low_resolution(): + (phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases + + assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960) + assert phase.spatial_downscale == 1 + assert phase.sigmas == LTX_DISTILLED_SIGMAS + assert phase.sampler == "euler_ancestral" + assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit + + +def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora(): + stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases + + assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920) + assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30 + assert stage1.spatial_downscale == 2 + assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance + assert stage2.adapter_slot == "ltx_distilled" + assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas + assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full" + assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None + + +@pytest.mark.parametrize( + "recipe", + [ + LTX2_ONE_STAGE_RECIPE, + LTX23_ONE_STAGE_RECIPE, + LTX_POSITIVE_ONLY_RECIPE, + ], +) +def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe): + assert recipe.supports_cache_dit + + +@pytest.mark.parametrize( + "recipe", + [ + LTX25_FULL_RECIPE, + LTX25_DISTILLED_ONE_STAGE_RECIPE, + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX25_DISTILLED_TWO_STAGE_RECIPE, + ], +) +def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe): + assert not recipe.supports_cache_dit + + +def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip(): + sample = torch.full((1, 1), 1e8, dtype=torch.float32) + velocity = torch.ones((1, 1), dtype=torch.bfloat16) + sigma = torch.tensor(0.5) + + actual = LTX_GUIDANCE_EXECUTOR._guide_modality( + sample, + {"cond": velocity}, + sigma, + LTXModalityGuidance(), + ) + expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma) + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +def test_ltx_velocity_from_x0_materializes_official_scalar_sigma(): + item_calls = [] + + class TrackingTensor(torch.Tensor): + @staticmethod + def __new__(cls): + return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False) + + def item(self): + item_calls.append(True) + return super().item() + + sample = torch.tensor([[1.0]], dtype=torch.bfloat16) + x0 = torch.tensor([[0.5]], dtype=torch.bfloat16) + + velocity_from_x0(sample, x0, TrackingTensor()) + + assert item_calls == [True] + + +def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity(): + sample = torch.full((1, 1), 1e8, dtype=torch.float32) + velocity = torch.ones((1, 1), dtype=torch.bfloat16) + + actual = LTX_GUIDANCE_EXECUTOR._guide_modality( + sample, + {"cond": velocity}, + torch.tensor(0.5), + LTXModalityGuidance(), + preserve_positive_velocity=True, + ) + + assert actual is velocity + + +@pytest.mark.parametrize("model_version", ["2.5", "2.3"]) +def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version): + expected_dtype = torch.bfloat16 + captured = {} + + def prepare_latents(**kwargs): + captured["video_dtype"] = kwargs["dtype"] + return torch.empty(1, 1, 1, dtype=kwargs["dtype"]) + + def prepare_audio_latents(*_args, **kwargs): + captured["audio_dtype"] = kwargs["dtype"] + return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1 + + pipeline = SimpleNamespace( + model_version=model_version, + transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)), + prepare_latents=prepare_latents, + prepare_audio_latents=prepare_audio_latents, + audio_sampling_rate=48_000, + audio_hop_length=160, + audio_vae_temporal_compression_ratio=4, + audio_vae_mel_compression_ratio=4, + audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)), + _resolve_audio_latent_length=lambda length, _latents: length, + ) + request = SimpleNamespace( + num_videos_per_prompt=1, + height=64, + width=64, + num_frames=1, + frame_rate=24.0, + generator=None, + latents=None, + audio_latents=None, + audio_latents_normalized=False, + ) + prompt = SimpleNamespace( + batch_size=1, + positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16), + positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16), + ) + + LTXRuntime._prepare_video_latents_stage( + pipeline, + request, + prompt, + device=torch.device("cpu"), + noise_scale=1.0, + ) + LTXRuntime._prepare_audio_latents_stage( + pipeline, + request, + prompt, + device=torch.device("cpu"), + noise_scale=1.0, + ) + + assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype} + + +def test_ltx_seeded_latents_match_official_packed_rng_layout(): + pipeline = SimpleNamespace( + model_version="2.5", + vae_spatial_compression_ratio=2, + vae_temporal_compression_ratio=2, + transformer_spatial_patch_size=1, + transformer_temporal_patch_size=1, + audio_vae_mel_compression_ratio=2, + od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)), + ) + device = torch.device("cpu") + actual_generator = torch.Generator(device=device).manual_seed(123) + expected_generator = torch.Generator(device=device).manual_seed(123) + + actual_video = ltx2_latents.prepare_video_latents( + pipeline, + batch_size=1, + num_channels_latents=2, + height=4, + width=6, + num_frames=3, + dtype=torch.float32, + device=device, + generator=actual_generator, + ) + actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents( + pipeline, + batch_size=1, + num_channels_latents=2, + audio_latent_length=3, + num_mel_bins=4, + dtype=torch.float32, + device=device, + generator=actual_generator, + ) + + expected_video = torch.randn((1, 12, 2), generator=expected_generator) + expected_audio = torch.randn((1, 3, 4), generator=expected_generator) + + torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0) + torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0) + assert (original_length, padded_length) == (3, 3) + + +def test_ltx25_distilled_two_stage_executes_custom_phase_schedules(): + request_inputs = LTXRequestInputs( + prompt="prompt", + negative_prompt="", + height=64, + width=64, + num_frames=1, + frame_rate=24.0, + num_inference_steps=8, + guidance=LTXGuidanceSpec.positive_only(), + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.25, + decode_noise_scale=0.5, + output_type="np", + max_sequence_length=16, + ) + prompt_context = object() + phase_calls: list[Any] = [] + + def resolve_request_inputs(req, **kwargs): + return request_inputs + + def run_phase(req, inputs, *, prompt_context=None, **kwargs): + phase_recipe = kwargs["phase_recipe"] + assert kwargs["image"] is source_image + phase_calls.append((phase_recipe, inputs, prompt_context)) + if len(phase_calls) == 1: + assert prompt_context is None + assert inputs.decode_timestep == 0.25 + assert inputs.decode_noise_scale == 0.5 + context = prompt_context_sentinel + assert phase_recipe.name == "generate_lowres" + assert (inputs.height, inputs.width) == (32, 32) + assert inputs.num_inference_steps == 2 + assert kwargs["sigmas"] == stage_1_sigmas + assert kwargs["noise_scale"] == 1.0 + video = torch.ones(1, 128, 1, 1, 1) + audio = torch.full((1, 8, 1, 2), 2.0) + audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0) + else: + assert prompt_context is prompt_context_sentinel + assert phase_recipe.name == "refine" + assert (inputs.height, inputs.width) == (64, 64) + assert inputs.latents.shape == (1, 128, 1, 2, 2) + torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0)) + assert inputs.audio_latents_normalized + assert inputs.guidance_scale == 1.0 + assert inputs.num_inference_steps == 2 + assert kwargs["sigmas"] == stage_2_sigmas + assert kwargs["noise_scale"] == stage_2_sigmas[0] + assert inputs.decode_timestep == 0.0 + assert inputs.decode_noise_scale is None + context = prompt_context + video = torch.full((1, 128, 1, 2, 2), 3.0) + audio = torch.full((1, 8, 1, 2), 4.0) + audio_for_next_phase = None + return LTXPhaseResult( + forward_context=SimpleNamespace(prompt_context=context), + video=video, + audio=audio, + audio_for_next_phase=audio_for_next_phase, + ) + + def decode_phase(phase): + return DiffusionOutput(output=(phase.video, phase.audio)) + + class FakeUpsampler(torch.nn.Module): + dtype = torch.float32 + + def forward(self, latents): + return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) + + prompt_context_sentinel = prompt_context + stage_1_sigmas = [1.0, 0.5, 0.0] + stage_2_sigmas = [0.8, 0.2, 0.0] + pipeline = object.__new__(LTX2TwoStagePipeline) + torch.nn.Module.__init__(pipeline) + pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE + pipeline.device = torch.device("cpu") + pipeline.vae_spatial_compression_ratio = 32 + pipeline.vae_temporal_compression_ratio = 8 + pipeline.latent_upsampler = FakeUpsampler() + object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) + object.__setattr__(pipeline, "run_phase", run_phase) + object.__setattr__(pipeline, "decode_phase", decode_phase) + + source_image = object() + output = pipeline.forward( + SimpleNamespace(sampling_params_list=[], prompts=[]), + image=source_image, + stage_1_sigmas=stage_1_sigmas, + stage_2_sigmas=stage_2_sigmas, + ) + + assert len(phase_calls) == 2 + assert phase_calls[1][2] is prompt_context_sentinel + torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0)) + torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0)) + + +class TestLTXRequestParsing: + def test_request_resolves_custom_sigmas_from_extra_args(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sigmas = [1.0, 0.5, 0.0] + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}), + request_id="ltx-custom-sigmas", + ) + ] + ) + + assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas + + def test_request_resolves_per_stage_sigmas_from_extra_args(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + stage_1_sigmas = [1.0, 0.4, 0.0] + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}), + request_id="ltx-custom-phase-sigmas", + ) + ] + ) + + resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas( + req, [1.0, 0.0], [0.8, 0.0] + ) + + assert resolved == (stage_1_sigmas, [0.8, 0.0]) + + @pytest.mark.parametrize("value", [-1, 52, 1.5, True]) + def test_request_rejects_invalid_image_crf(self, value): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}), + request_id="ltx-invalid-image-crf", + ) + ] + ) + + with pytest.raises(ValueError, match="image_crf"): + _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req) + + +class TestLTXForwardStages: + def test_gemma_tokenization_matches_official_bos_and_left_padding(self): + pipe = _make_ltx_request_pipe(LTX2Pipeline) + + class FakeTokenizer: + bos_token_id = 2 + eos_token_id = 1 + eos_token = "" + pad_token_id = 0 + pad_token = "" + padding_side = "right" + + def __init__(self): + self.calls = [] + + def __call__(self, text, **kwargs): + self.calls.append((text, kwargs)) + ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id] + return SimpleNamespace(input_ids=ids) + + def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask): + assert padding == "max_length" + assert return_tensors == "pt" + assert return_attention_mask is True + rows = [] + masks = [] + for ids in encoded["input_ids"]: + pad = max_length - len(ids) + rows.append([self.pad_token_id] * pad + ids) + masks.append([0] * pad + [1] * len(ids)) + return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks)) + + class FakeTextEncoder(torch.nn.Module): + dtype = torch.float32 + + def __init__(self): + super().__init__() + self.seen_input_ids = None + self.seen_attention_mask = None + + def forward(self, *, input_ids, attention_mask, output_hidden_states): + assert output_hidden_states is True + self.seen_input_ids = input_ids.clone() + self.seen_attention_mask = attention_mask.clone() + return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),)) + + tokenizer = FakeTokenizer() + text_encoder = FakeTextEncoder() + object.__setattr__(pipe, "tokenizer", tokenizer) + object.__setattr__(pipe, "text_encoder", text_encoder) + + pipe._get_gemma_prompt_embeds( + [" long ", "short"], + max_sequence_length=4, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]])) + torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]])) + assert tokenizer.padding_side == "left" + assert [call[0] for call in tokenizer.calls] == ["long", "short"] diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py index e309dbf47..aa5654c74 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for LTX image-to-video input and conditioning behavior.""" +import sys from types import SimpleNamespace import pytest @@ -23,7 +24,7 @@ def _make_ltx23_request_pipe(cls): class TestLTXImageToVideoForwardStages: - def test_i2v_pil_preprocessing_matches_official_center_crop(self): + def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self): from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images pixels = torch.zeros(4, 8, 3, dtype=torch.uint8) @@ -35,6 +36,136 @@ class TestLTXImageToVideoForwardStages: torch.testing.assert_close(actual, expected) + def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning + + pixels = torch.zeros(4, 8, 3, dtype=torch.uint8) + pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8) + image = Image.fromarray(pixels.numpy()) + monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array) + + actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18) + expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0 + + torch.testing.assert_close(actual, expected) + + def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + captured_crfs = [] + captured_prepare_kwargs = [] + + def fake_preprocess(image, *, height, width, crf, device, dtype): + del image, device + captured_crfs.append(crf) + return torch.zeros(1, 3, height, width, dtype=dtype) + + def fake_prepare_latents(**kwargs): + captured_prepare_kwargs.append(kwargs) + return kwargs["image"], None + + monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess) + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4))) + object.__setattr__(pipe, "prepare_latents", fake_prepare_latents) + request_inputs = SimpleNamespace( + height=4, + width=4, + num_frames=3, + num_videos_per_prompt=1, + generator=None, + latents=None, + ) + prompt_context = SimpleNamespace( + batch_size=1, + positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16), + ) + + pipe._prepare_video_latents_stage( + request_inputs, + prompt_context, + device=torch.device("cpu"), + noise_scale=0.0, + image=Image.new("RGB", (8, 4)), + ) + + assert captured_crfs == [18] + assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16 + assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16 + + def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + + with pytest.raises(ValueError, match="requires PIL images"): + pipe._prepare_video_latents_stage( + SimpleNamespace(), + SimpleNamespace(), + device=torch.device("cpu"), + noise_scale=0.0, + image=torch.zeros(3, 4, 4), + ) + + def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + captured_prepare_kwargs = [] + + def fake_prepare_latents(**kwargs): + captured_prepare_kwargs.append(kwargs) + return kwargs["image"], None + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4))) + object.__setattr__(pipe, "prepare_latents", fake_prepare_latents) + request_inputs = SimpleNamespace( + height=4, + width=4, + num_frames=3, + num_videos_per_prompt=1, + generator=None, + latents=None, + image_crf=0, + ) + prompt_context = SimpleNamespace( + batch_size=1, + positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16), + ) + + pipe._prepare_video_latents_stage( + request_inputs, + prompt_context, + device=torch.device("cpu"), + noise_scale=0.0, + image=torch.zeros(3, 4, 4), + ) + + assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4) + + def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf + + image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy() + monkeypatch.setitem(sys.modules, "av", None) + + assert _apply_image_conditioning_crf(image, 18) is image + + def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf + + monkeypatch.setitem(sys.modules, "av", None) + + with pytest.raises(ImportError, match="PyAV with a libx264 encoder"): + _apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18) + def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -98,7 +229,7 @@ class TestLTXImageToVideoForwardStages: with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"): pipe._resolve_request_image(mixed_req, None, request_inputs) - def test_denoise_timestep_kwargs_masks_video_only(self): + def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pipe = object.__new__(LTX2Pipeline) @@ -127,6 +258,39 @@ class TestLTXImageToVideoConditioning: assert LTX2Pipeline.support_image_input is True + def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.use_diffusion_decoder = True + pipe.vae = SimpleNamespace( + encode=lambda image: image, + latents_mean=torch.full((2,), 100.0), + latents_std=torch.full((2,), 10.0), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipe.diffusion_decoder = SimpleNamespace( + latents_mean=torch.tensor([1.0, 3.0]), + latents_std=torch.tensor([2.0, 4.0]), + config=SimpleNamespace(scaling_factor=2.0), + ) + monkeypatch.setattr( + ltx2_conditioning, + "retrieve_latents", + lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]), + ) + + actual = pipe._encode_i2v_image_latents( + torch.zeros(1, 3, 1, 1), + batch_size=1, + generator=None, + dtype=torch.float32, + ) + + torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0)) + def test_ltx23_i2v_rejects_multi_image_prompt_list(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline @@ -189,13 +353,18 @@ class TestLTXImageToVideoConditioning: torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]])) - def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch): + @pytest.mark.parametrize( + ("model_version", "sampled_shape"), + [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))], + ) + def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape): import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) + pipe.model_version = model_version pipe.vae_spatial_compression_ratio = 1 pipe.vae_temporal_compression_ratio = 1 pipe.transformer_spatial_patch_size = 1 @@ -237,14 +406,20 @@ class TestLTXImageToVideoConditioning: torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]])) - assert sampled_shapes == [(1, 3, 2)] + assert sampled_shapes == [sampled_shape] - def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch): + @pytest.mark.parametrize( + ("model_version", "sampled_shape"), + [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))], + ) + def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape): import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning + import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) + pipe.model_version = model_version pipe.vae_spatial_compression_ratio = 1 pipe.vae_temporal_compression_ratio = 1 pipe.transformer_spatial_patch_size = 1 @@ -266,6 +441,7 @@ class TestLTXImageToVideoConditioning: return torch.ones(shape, device=device, dtype=dtype) monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor) + monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor) out, conditioning_mask = pipe.prepare_latents( image=torch.zeros(1, 3, 1, 1), @@ -279,10 +455,54 @@ class TestLTXImageToVideoConditioning: device=torch.device("cpu"), ) - assert sampled_shapes == [(1, 3, 2)] + assert sampled_shapes == [sampled_shape] torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]])) + def test_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.model_version = "2.5" + pipe.vae_spatial_compression_ratio = 1 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + object.__setattr__( + pipe, + "_encode_i2v_image_latents", + lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16), + ) + captured = {} + + def fake_create_noised_state(latents, noise_scale, generator=None): + del generator + captured["noise_scale"] = noise_scale + return latents + + monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state) + pipe.prepare_latents( + image=torch.zeros(1, 3, 1, 1), + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=3, + noise_scale=0.909375, + dtype=torch.bfloat16, + device=torch.device("cpu"), + latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16), + ) + + assert captured["noise_scale"] == 0.909375 + def test_ltx23_i2v_video_step_preserves_conditioning_frame(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline @@ -385,6 +605,99 @@ class TestLTXImageToVideoConditioning: torch.testing.assert_close(audio, torch.tensor([[[-2.0]]])) assert pipeline.step_indices == [0, 1] + def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self): + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter + + device = torch.device("cpu") + sigmas = torch.tensor([1.0, 0.5, 0.0]) + pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas)) + mask = torch.tensor([[1.0, 0.0]]) + adapter = LTXVideoAudioStepAdapter( + pipeline, + SimpleNamespace(sigmas=sigmas.clone()), + latent_num_frames=1, + latent_height=1, + latent_width=1, + image_conditioned=True, + sampler="euler_ancestral", + generator=torch.Generator(device=device).manual_seed(42), + conditioning_mask=mask, + ) + video = torch.tensor([[[3.0], [2.0]]]) + audio = torch.tensor([[[4.0]]]) + video_velocity = torch.tensor([[[7.0], [0.25]]]) + audio_velocity = torch.tensor([[[0.5]]]) + + reference_generator = torch.Generator(device=device).manual_seed(10042) + video_noise = torch.randn(video.shape, generator=reference_generator) + audio_noise = torch.randn(audio.shape, generator=reference_generator) + + def official_step(sample, velocity, noise): + sigma = sigmas[0] + sigma_next = sigmas[1] + denoised = sample.float() - velocity.float() * sigma + sigma_down = sigma_next * (sigma_next / sigma) + ratio = sigma_down / sigma + deterministic = ratio * sample.float() + (1.0 - ratio) * denoised + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt() + return (alpha_next / alpha_down) * deterministic + noise * coeff + + expected_video = official_step(video, video_velocity, video_noise) + expected_video[:, :1] = video[:, :1] + expected_audio = official_step(audio, audio_velocity, audio_noise) + + ((actual_video, actual_audio),) = adapter.step( + (video_velocity, audio_velocity), + (sigmas[0], sigmas[0]), + (video, audio), + ) + torch.testing.assert_close(actual_video, expected_video) + torch.testing.assert_close(actual_audio, expected_audio) + + terminal_video_velocity = torch.full_like(actual_video, 0.125) + terminal_audio_velocity = torch.full_like(actual_audio, 0.25) + ((terminal_video, terminal_audio),) = adapter.step( + (terminal_video_velocity, terminal_audio_velocity), + (sigmas[1], sigmas[1]), + (actual_video, actual_audio), + ) + expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1] + expected_terminal_video[:, :1] = actual_video[:, :1] + expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1] + torch.testing.assert_close(terminal_video, expected_terminal_video) + torch.testing.assert_close(terminal_audio, expected_terminal_audio) + + def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self): + from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity + + sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16) + velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16) + sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32) + generator = torch.Generator(device="cpu").manual_seed(10042) + reference_generator = torch.Generator(device="cpu").manual_seed(10042) + + sigma, sigma_next = sigmas[:2] + denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float() + sigma_down = sigma_next * (sigma_next / sigma) + sigma_down_ratio = sigma_down / sigma + expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt() + noise = torch.randn( + sample.shape, + generator=reference_generator, + dtype=sample.dtype, + device=sample.device, + ) + expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype) + + actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self): from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index 36f7e9687..fbac0bfe5 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace import pytest @@ -23,7 +26,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1): return pipeline -def test_prepare_video_latents_samples_directly_in_packed_token_space(): +def test_prepare_video_latents_matches_official_values_and_token_major_layout(): pipeline = _make_pipeline(LTX2Pipeline) pipeline.vae_spatial_compression_ratio = 8 pipeline.vae_temporal_compression_ratio = 8 @@ -44,6 +47,40 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space(): ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert actual.stride()[1:] == (1, actual.shape[1]) + + +def test_prepare_video_latents_uses_diffusion_decoder_statistics(): + pipeline = _make_pipeline(LTX2Pipeline) + pipeline.use_diffusion_decoder = True + pipeline.vae_spatial_compression_ratio = 1 + pipeline.vae_temporal_compression_ratio = 1 + pipeline.transformer_spatial_patch_size = 1 + pipeline.transformer_temporal_patch_size = 1 + pipeline.vae = SimpleNamespace( + latents_mean=torch.full((2,), 100.0), + latents_std=torch.full((2,), 10.0), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipeline.diffusion_decoder = SimpleNamespace( + latents_mean=torch.tensor([1.0, 3.0]), + latents_std=torch.tensor([2.0, 4.0]), + config=SimpleNamespace(scaling_factor=2.0), + ) + + actual = pipeline.prepare_latents( + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=1, + noise_scale=0.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]), + ) + + torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0)) def test_prepare_audio_latents_samples_directly_in_packed_token_space(): @@ -81,6 +118,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(): assert original_num_frames == 1 assert padded_num_frames == 2 assert latents.shape == (1, 2, 128) + torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:])) + + +def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding(): + pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1) + pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4) + generator_sp1 = torch.Generator().manual_seed(42) + generator_sp4 = torch.Generator().manual_seed(42) + + latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=3, + dtype=torch.float32, + device=torch.device("cpu"), + generator=generator_sp1, + ) + latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=3, + dtype=torch.float32, + device=torch.device("cpu"), + generator=generator_sp4, + ) + + torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0) + torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:])) + torch.testing.assert_close( + torch.randn(8, generator=generator_sp4), + torch.randn(8, generator=generator_sp1), + rtol=0, + atol=0, + ) def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents(): @@ -137,7 +210,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(): assert original_num_frames == 10 assert padded_num_frames == 12 assert padded.shape == (1, 12, 8) - torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents)) + packed = latent_ops.pack_audio_latents(latents) + torch.testing.assert_close(padded[:, :10], packed[:, :10]) + torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:])) def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(): @@ -163,3 +238,68 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length(): device=torch.device("cpu"), latents=latents, ) + + +def test_create_noised_state_matches_official_fp32_lerp(): + latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128) + expected_generator = torch.Generator().manual_seed(42) + noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype) + expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to( + latents.dtype + ) + + actual = latent_ops.create_noised_state( + latents, + torch.tensor(0.15, dtype=torch.bfloat16), + torch.Generator().manual_seed(42), + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_create_conditioned_noised_state_matches_official_two_lerps(): + latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8) + clean_latents = latents.clone() + clean_latents[:, 0] = 3.0 + denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]]) + expected_generator = torch.Generator().manual_seed(42) + noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype) + noised = torch.lerp(latents.float(), noise.float(), 0.909375) + expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype) + + actual = latent_ops.create_conditioned_noised_state( + latents, + clean_latents, + denoise_mask, + 0.909375, + torch.Generator().manual_seed(42), + ) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_prepare_supplied_video_latents_uses_official_token_major_layout(): + pipeline = _make_pipeline(LTX2Pipeline) + supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16) + actual_video = pipeline.prepare_latents( + batch_size=1, + num_channels_latents=4, + height=64, + width=64, + num_frames=9, + dtype=torch.float32, + device=torch.device("cpu"), + latents=supplied, + ) + + torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0) + assert actual_video.stride()[1:] == (1, actual_video.shape[1]) + + +def test_clear_audio_padding_keeps_padding_outside_sampler_state(): + updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]]) + + actual = latent_ops.clear_audio_padding(updated, 2) + + torch.testing.assert_close(actual[:, :2], updated[:, :2]) + torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:])) diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 8d8c22e79..0573dfb37 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -12,9 +12,31 @@ import torch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker): + from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0 + + sigma = mocker.MagicMock() + sigma.to.return_value.item.return_value = 0.5 + + actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma) + + sigma.to.assert_called_once_with(torch.float32) + sigma.to.return_value.item.assert_called_once_with() + torch.testing.assert_close(actual, torch.tensor([2.0])) + + class TestCFGParallelHelpers: """Test LTX-2.3 CFG helper math without loading model weights.""" + def test_single_guidance_pass_preserves_official_token_layout(self): + from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch + + tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2) + repeated = _repeat_batch(tokens, 1) + + assert repeated is tokens + assert repeated.stride() == tokens.stride() + def test_combine_cfg_noise_matches_x0_space_formula(self): from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline @@ -199,6 +221,7 @@ class TestCFGParallelHelpers: ), attention_kwargs=None, audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])), + original_audio_num_frames=state.audio.shape[1], ) actual_video, actual_audio = executor.predict_parallel_guidance( diff --git a/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py new file mode 100644 index 000000000..f9b3bedce --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for the fixed LTX refinement-phase adapter.""" + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from safetensors.torch import save_file + +from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter +from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import ( + AdapterTarget, + parse_ltx_adapter, +) +from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE +from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime +from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_TWO_STAGE_RECIPE, + LTX_DISTILLED_ADAPTER_SLOT, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _target( + lora_a, + lora_b, + *, + source_module="proj", + module="proj", + target_slice=None, +): + return AdapterTarget( + source_module=source_module, + module=module, + slice=target_slice, + a_key="a", + b_key="b", + a_shape=tuple(lora_a.shape), + b_shape=tuple(lora_b.shape), + ) + + +def _wrap( + layer, + lora_a, + lora_b, + *, + source_module="proj", + module="proj", + target_slice=None, + layer_fused=False, +): + target = _target( + lora_a, + lora_b, + source_module=source_module, + module=module, + target_slice=target_slice, + ) + wrapper = ltx2_phase_adapter._PhaseAdapterLinear( + layer, + [target], + adapter_dtype=lora_a.dtype, + layer_fused=layer_fused, + ) + wrapper.load_target(target, lora_a, lora_b) + wrapper.set_enabled(True) + return wrapper + + +class _LoRAAttention(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.to_qkv = torch.nn.Identity() + + +class _LoRATransformer(torch.nn.Module): + stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),) + + def __init__(self) -> None: + super().__init__() + self.proj_in = torch.nn.Linear(4, 6, bias=False) + block = torch.nn.Module() + block.attn1 = _LoRAAttention() + self.transformer_blocks = torch.nn.ModuleList([block]) + + +def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path): + adapter_path = tmp_path / "distilled.safetensors" + save_file( + { + "diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4), + "diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4), + "diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2), + "diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4), + "diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2), + }, + adapter_path, + ) + + manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path)) + assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == { + ("proj_in", "proj_in", None), + ("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"), + } + + +def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path): + adapter_path = tmp_path / "distilled.safetensors" + save_file( + {"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)}, + adapter_path, + ) + + with pytest.raises(ValueError, match="Missing paired LoRA A tensors"): + parse_ltx_adapter(_LoRATransformer(), str(adapter_path)) + + +def _pipeline(*, dtype, quantization_config): + return SimpleNamespace( + pipeline_recipe=LTX2_TWO_STAGE_RECIPE, + component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE, + transformer=object(), + od_config=SimpleNamespace( + model="unused", + lora_path=None, + dtype=dtype, + quantization_config=quantization_config, + ), + ) + + +@pytest.mark.parametrize( + ("dtype", "quantization_config", "layer_fused"), + [ + (torch.float32, object(), False), + (torch.bfloat16, None, True), + ], +) +def test_ltx_phase_adapter_selects_execution_from_quantization( + monkeypatch, + dtype, + quantization_config, + layer_fused, +): + pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config) + manifest = object() + captured: dict[str, Any] = {} + + def resolve_artifact(*_args, **kwargs): + assert kwargs == { + "model_revision": None, + "artifact_revision": None, + } + return "adapter" + + def parse(transformer, path): + assert transformer is pipeline.transformer + assert path == "adapter" + return manifest + + class Runtime: + def __init__(self, transformer, parsed_manifest, **kwargs): + captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs) + + monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact) + monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse) + monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime) + + runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline) + + assert isinstance(runtime, Runtime) + assert captured == { + "transformer": pipeline.transformer, + "manifest": manifest, + "dtype": dtype, + "layer_fused": layer_fused, + } + + +def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch): + pipeline = _pipeline(dtype=torch.float16, quantization_config=None) + monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter") + + with pytest.raises(ValueError, match="bfloat16"): + ltx2_phase_adapter.build_ltx_phase_adapter(pipeline) + + +def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path): + class TinyTransformer(torch.nn.Module): + stacked_params_mapping = () + + def __init__(self) -> None: + super().__init__() + self.proj = torch.nn.Linear(2, 2, bias=False) + + adapter_path = tmp_path / "distilled.safetensors" + save_file( + { + "diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]), + "diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]), + }, + adapter_path, + ) + transformer = TinyTransformer() + runtime = LTXPhaseAdapterRuntime( + transformer, + parse_ltx_adapter(transformer, str(adapter_path)), + dtype=torch.float32, + ) + + assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight" + assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"} + with torch.no_grad(): + transformer.proj.base_layer.weight.copy_(torch.eye(2)) + runtime.finalize() + + x = torch.tensor([[1.0, 1.0]]) + runtime.activate(None) + torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]])) + runtime.activate(LTX_DISTILLED_ADAPTER_SLOT) + torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]])) + + +def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion(): + layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16) + base_weight = torch.tensor( + [ + [0.5, -0.25, 0.125, 0.75], + [-0.375, 0.625, 0.25, -0.5], + [0.875, -0.125, -0.75, 0.375], + ], + dtype=torch.bfloat16, + ) + lora_a = torch.tensor( + [ + [0.1015625, -0.203125, 0.3046875, -0.40625], + [0.5078125, -0.609375, 0.7109375, -0.8125], + ], + dtype=torch.bfloat16, + ) + lora_b = torch.tensor( + [[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]], + dtype=torch.bfloat16, + ) + with torch.no_grad(): + layer.weight.copy_(base_weight) + layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16)) + wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True) + + expected_weight = torch.matmul(lora_b, lora_a) + expected_weight.add_(base_weight) + x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16) + expected = torch.nn.functional.linear(x, expected_weight, layer.bias) + + assert torch.equal(wrapper(x), expected) + assert torch.equal(layer.weight, base_weight) + + +def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + class FakeRow(torch.nn.Module): + input_size = 4 + output_size = 2 + input_size_per_partition = 2 + input_is_parallel = True + tp_rank = 0 + tp_size = 2 + reduce_results = True + skip_bias_add = False + return_bias = False + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16)) + self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16)) + self.quant_method = UnquantizedLinearMethod() + + monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow) + reductions = [] + + def all_reduce(value): + reductions.append(value.clone()) + return value + 1 + + monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce) + layer = FakeRow() + lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16) + lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16) + wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True) + + expected_weight = torch.matmul(lora_b, lora_a[:, :2]) + expected_weight.add_(layer.weight) + x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16) + expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias) + + actual = wrapper(x) + assert len(reductions) == 1 + assert torch.equal(reductions[0], expected_local) + assert torch.equal(actual, expected_local + 1) + + +def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + class FakeColumn(torch.nn.Module): + input_size = 3 + output_size = 4 + output_size_per_partition = 2 + tp_rank = 1 + tp_size = 2 + gather_output = False + skip_bias_add = False + return_bias = False + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16) + ) + self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16)) + self.quant_method = UnquantizedLinearMethod() + + monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn) + layer = FakeColumn() + base_weight = layer.weight.detach().clone() + lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16) + lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16) + wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True) + + expected_weight = torch.matmul(lora_b[2:], lora_a) + expected_weight.add_(base_weight) + x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16) + expected = torch.nn.functional.linear(x, expected_weight, layer.bias) + + assert torch.equal(wrapper(x), expected) + assert torch.equal(layer.weight, base_weight) + + +def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + class FakeColumn(torch.nn.Module): + pass + + class FakeQKV(FakeColumn): + input_size = 4 + head_size = 2 + v_head_size = 2 + total_num_heads = 2 + total_num_kv_heads = 2 + num_heads = 1 + num_kv_heads = 1 + num_kv_head_replicas = 1 + tp_rank = 1 + tp_size = 2 + gather_output = False + skip_bias_add = True + return_bias = True + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter( + torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16 + ) + self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8) + self.quant_method = UnquantizedLinearMethod() + + monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn) + monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV) + layer = FakeQKV() + base_weight = layer.weight.detach().clone() + lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16) + lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16) + wrapper = _wrap( + layer, + lora_a, + lora_b, + source_module="to_k", + module="to_qkv", + target_slice="k", + layer_fused=True, + ) + + expected_weight = base_weight.clone() + expected_k = torch.matmul(lora_b[2:], lora_a) + expected_k.add_(base_weight[2:4]) + expected_weight[2:4].copy_(expected_k) + x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16) + expected_output = torch.nn.functional.linear(x, expected_weight) + + actual_output, returned_bias = wrapper(x) + assert torch.equal(actual_output, expected_output) + assert returned_bias is layer.bias + assert torch.equal(layer.weight, base_weight) + + +@pytest.mark.parametrize( + ("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"), + [ + ( + lambda: type( + "FakeColumn", + (torch.nn.Module,), + { + "input_size": 4, + "output_size": 4, + "output_size_per_partition": 2, + "tp_rank": 1, + "gather_output": False, + "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)), + }, + )(), + None, + torch.tensor([[1.0, 10.0, 100.0, 1000.0]]), + torch.tensor([[1.0], [2.0], [3.0], [4.0]]), + torch.tensor([[12963.0, 17284.0]]), + ), + ( + lambda: type( + "FakeRow", + (torch.nn.Module,), + { + "input_size": 4, + "output_size": 2, + "input_size_per_partition": 2, + "tp_rank": 1, + "reduce_results": False, + "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)), + }, + )(), + None, + torch.tensor([[1.0, 10.0, 100.0, 1000.0]]), + torch.tensor([[2.0], [3.0]]), + torch.tensor([[8600.0, 12900.0]]), + ), + ( + lambda: type( + "FakeQKV", + (torch.nn.Module,), + { + "input_size": 4, + "head_size": 2, + "v_head_size": 2, + "total_num_heads": 2, + "total_num_kv_heads": 2, + "num_heads": 1, + "num_kv_heads": 1, + "num_kv_head_replicas": 1, + "tp_rank": 1, + "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)), + }, + )(), + "k", + torch.tensor([[1.0, 10.0, 100.0, 1000.0]]), + torch.tensor([[1.0], [2.0], [3.0], [4.0]]), + torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]), + ), + ], +) +def test_ltx_phase_adapter_uses_local_tp_lora_slices( + monkeypatch, + layer_factory, + target_slice, + expected_a, + expected_b, + expected_output, +): + layer = layer_factory() + layer.register_buffer("device_anchor", torch.empty(0)) + layer_type = type(layer) + if target_slice == "k": + monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type) + elif "Column" in layer_type.__name__: + monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type) + else: + monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type) + + wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice) + torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output) + expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size) + assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index d70820f41..9cafa22b6 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the shared LTX pipeline runtime and public contracts.""" @@ -8,18 +8,28 @@ import os import tempfile from dataclasses import replace from types import SimpleNamespace +from typing import Any import pytest import torch from vllm_omni.diffusion.data import DiffusionOutput -from vllm_omni.diffusion.models.ltx2 import ltx2_components +from vllm_omni.diffusion.models.ltx2 import ( + ltx2_components, + ltx2_transformer, +) from vllm_omni.diffusion.models.ltx2.ltx2_components import ( LTX2_COMPONENT_PROFILE, LTX2_DISTILLED_COMPONENT_PROFILE, + LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, + LTX2_TWO_STAGE_COMPONENT_PROFILE, LTX23_COMPONENT_PROFILE, + LTX23_DISTILLED_COMPONENT_PROFILE, + LTX23_TWO_STAGE_COMPONENT_PROFILE, _install_connector_attention, detect_ltx_model_version, + resolve_ltx_artifact, + resolve_ltx_component_profile, ) from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin from vllm_omni.diffusion.models.ltx2.ltx2_denoise import ( @@ -33,25 +43,40 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( LTXGuidancePlan, LTXGuidanceSpec, LTXModalityGuidance, + _repeat_batch, build_perturbation_kwargs, combine_guided_x0, ) from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_DISTILLED_ONE_STAGE_RECIPE, LTX2_DISTILLED_TWO_STAGE_RECIPE, LTX2_ONE_STAGE_RECIPE, + LTX2_TWO_STAGE_RECIPE, + LTX23_DISTILLED_TWO_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE, + LTX23_TWO_STAGE_RECIPE, LTX_DEFAULT_NEGATIVE_PROMPT, LTX_POSITIVE_ONLY_RECIPE, + resolve_ltx_pipeline_recipe, +) +from vllm_omni.diffusion.models.ltx2.ltx2_request import ( + LTXRequestInputs, + validate_ltx_checkpoint, + validate_pipeline_request, ) -from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2DistilledOneStagePipeline, LTX2I2VDMD2Pipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( + LTX2DistilledPipeline, + LTX2DistilledTwoStagePipeline, + LTX2TwoStagePipeline, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -104,10 +129,28 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries(): ) assert issubclass(LTX2Pipeline, LTXRuntime) + assert issubclass(LTX2TwoStagePipeline, LTXRuntime) + assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime) + assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime) assert issubclass(LTX2DistilledPipeline, LTXRuntime) + assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE + assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE + assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE + assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE + assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE + for pipeline_cls, profile in ( + (LTX2Pipeline, LTX2_COMPONENT_PROFILE), + (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE), + (LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE), + (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE), + ): + assert pipeline_cls._dit_modules == list(profile.dit_modules) + assert pipeline_cls._encoder_modules == list(profile.encoder_modules) + assert pipeline_cls._vae_modules == list(profile.vae_modules) + assert pipeline_cls._resident_modules == list(profile.resident_modules) assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video @@ -123,6 +166,131 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path): assert detect_ltx_model_version(str(tmp_path)) == "2" +@pytest.mark.parametrize( + ("model", "repo_id", "expected_revision"), + [ + ("Lightricks/test", "Lightricks/test", "model-revision"), + ("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"), + ], +) +def test_ltx_artifact_uses_source_revision_and_hub_fallback( + monkeypatch, + model, + repo_id, + expected_revision, +): + filename = "ltx-sidecar.safetensors" + calls = [] + + def fake_download(**kwargs): + calls.append(kwargs) + return "/cache/ltx-sidecar.safetensors" + + monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) + + assert ( + resolve_ltx_artifact( + model, + repo_id, + filename, + model_revision="model-revision", + artifact_revision="artifact-revision", + ) + == "/cache/ltx-sidecar.safetensors" + ) + assert calls == [ + { + "repo_id": repo_id, + "filename": filename, + "revision": expected_revision, + } + ] + + +def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): + filename = "ltx-sidecar.safetensors" + expected = tmp_path / filename + expected.write_bytes(b"sidecar") + monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup")) + + assert resolve_ltx_artifact( + str(tmp_path), + "Lightricks/test", + filename, + model_revision="unused-model-revision", + artifact_revision="unused-artifact-revision", + ) == str(expected) + + +def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch): + calls = [] + + def fake_download(**kwargs): + calls.append(kwargs) + return "/cache/ltx-sidecar.safetensors" + + monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) + + assert ( + resolve_ltx_artifact( + str(tmp_path), + "Lightricks/test", + "ltx-sidecar.safetensors", + model_revision="local-model-revision", + artifact_revision="artifact-revision", + ) + == "/cache/ltx-sidecar.safetensors" + ) + assert calls == [ + { + "repo_id": "Lightricks/test", + "filename": "ltx-sidecar.safetensors", + "revision": "artifact-revision", + } + ] + + +@pytest.mark.parametrize( + ("expected_kind", "scheduler_config"), + [ + ("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}), + ("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}), + ], +) +def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config): + validate_ltx_checkpoint( + scheduler_config, + expected_kind=expected_kind, + pipeline_name="TestLTXPipeline", + ) + + +@pytest.mark.parametrize( + ("expected_kind", "scheduler_config", "error"), + [ + ("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"), + ("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"), + ], +) +def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error): + with pytest.raises(ValueError, match=error): + validate_ltx_checkpoint( + scheduler_config, + expected_kind=expected_kind, + pipeline_name="TestLTXPipeline", + ) + + +def test_ltx_rejects_advanced_uaa_before_component_initialization(): + od_config = SimpleNamespace( + model="unused", + parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"), + ) + + with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"): + LTX2Pipeline(od_config=od_config) + + def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch): from vllm_omni.diffusion.models.ltx2 import ltx2_runtime @@ -144,16 +312,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m assert pipe.reports_stage_durations -def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts(): +def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]})) + initialized = False + + def stub_components(*_args, **_kwargs): + nonlocal initialized + initialized = True + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + + with pytest.raises(ValueError, match="not qualified for this LTX recipe"): + LTX2TwoStagePipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + cache_backend="cache_dit", + ) + ) + + assert not initialized + + +def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]})) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace()) + pipe = LTX2TwoStagePipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + enable_diffusion_pipeline_profiler=False, + lora_path=None, + model_config={}, + ) + ) + + assert pipe.model_version == "2.3" + assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE + + +def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]})) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) + + pipe = LTX2DistilledPipeline( + od_config=SimpleNamespace( + model=str(tmp_path), + enable_diffusion_pipeline_profiler=False, + ) + ) + + assert pipe.model_version == "2.3" + assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE + assert pipe._phase_adapter is None + assert pipe.preserve_sp_padded_audio_duration + assert pipe.reports_stage_durations + + +def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts(): assert LTX2Pipeline.supports_request_batch assert LTX2Pipeline.support_image_input assert LTX2Pipeline.unified_text_image_entry assert LTX2Pipeline.distributed_video_decode + assert not LTX2TwoStagePipeline.supports_request_batch + assert LTX2TwoStagePipeline.support_image_input + assert LTX2TwoStagePipeline.unified_text_image_entry assert not LTX2DistilledPipeline.supports_request_batch assert LTX2DistilledPipeline.support_image_input assert LTX2DistilledPipeline.unified_text_image_entry +def test_ltx_two_stage_resolves_request_embedded_images(): + pipe = object.__new__(LTX2TwoStagePipeline) + request_inputs = SimpleNamespace(latents=None) + image = torch.zeros(3, 8, 8) + req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}]) + + assert pipe._resolve_request_image(req, None, request_inputs) is image + + def test_ltx_dmd2_entries_retain_their_task_boundaries(): assert not LTX2T2VDMD2Pipeline.support_image_input assert LTX2I2VDMD2Pipeline.support_image_input @@ -330,6 +586,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics(): assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",) +def test_ltx_positive_only_guidance_preserves_token_major_layout(): + tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2) + + repeated = _repeat_batch(tensor, 1) + + assert repeated is tensor + assert repeated.stride() == tensor.stride() + + def test_ltx_guidance_rescale_is_invariant_to_request_batching(): cond = torch.tensor( [ @@ -380,6 +645,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction(): torch.testing.assert_close(actual, torch.zeros_like(cond)) +def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens(): + logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]]) + logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]]) + guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0) + expected = combine_guided_x0( + cond=logical_cond, + uncond_text=logical_uncond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1) + padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1) + + actual = combine_guided_x0( + cond=padded_cond, + uncond_text=padded_uncond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + rescale_token_count=2, + ) + + torch.testing.assert_close(actual[:, :2], expected) + + def test_ltx_guidance_builds_per_sample_stg_and_modality_masks(): plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1)) @@ -488,15 +779,48 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch): assert len(installed) == 2 assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed) - assert all(processor.has_learned_registers for processor in installed) assert all(kernel["role"] == "ltx2.connector" for kernel in kernels) assert all(kernel["role_category"] == "self" for kernel in kernels) assert all(kernel["skip_sequence_parallel"] for kernel in kernels) assert all(kernel["disable_kv_quant"] for kernel in kernels) + assert all(processor.has_learned_registers for processor in installed) + assert not any(processor.preserve_learned_register_mask for processor in installed) -@pytest.mark.parametrize("has_learned_registers", [False, True]) -def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers): +def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch): + installed = [] + + class OmniAttention: + def __init__(self, **_kwargs): + pass + + class Attention: + heads = 32 + head_dim = 128 + inner_kv_dim = 4096 + + def set_processor(self, processor): + installed.append(processor) + + monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention) + connectors = SimpleNamespace( + video_connector=SimpleNamespace( + learnable_registers=object(), + transformer_blocks=[SimpleNamespace(attn1=Attention())], + ), + audio_connector=SimpleNamespace( + learnable_registers=object(), + transformer_blocks=[SimpleNamespace(attn1=Attention())], + ), + ) + + _install_connector_attention(connectors, preserve_learned_register_mask=True) + + assert all(processor.has_learned_registers for processor in installed) + assert all(processor.preserve_learned_register_mask for processor in installed) + + +def test_ltx_connector_attention_dispatches_through_omni_kernel(): calls = [] class Backend: @@ -511,12 +835,16 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi calls.append((query, key, value, metadata)) return query + class UpcastingNorm(torch.nn.Module): + def forward(self, tensor): + return tensor.float() + attention = SimpleNamespace( heads=2, head_dim=4, inner_kv_dim=8, - norm_q=torch.nn.Identity(), - norm_k=torch.nn.Identity(), + norm_q=UpcastingNorm(), + norm_k=UpcastingNorm(), to_q=torch.nn.Identity(), to_k=torch.nn.Identity(), to_v=torch.nn.Identity(), @@ -525,33 +853,46 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi rope_type="split", omni_attention=OmniAttention(), ) - hidden_states = torch.randn(2, 3, 8) + hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16) additive_mask = torch.zeros(2, 1, 3, 3) + rotary = ( + torch.randn(2, 2, 3, 2, dtype=torch.float32), + torch.randn(2, 2, 3, 2, dtype=torch.float32), + ) - output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)( + output = ltx2_components._LTXConnectorAttnProcessor()( attention, hidden_states, attention_mask=additive_mask, + query_rotary_emb=rotary, ) query, key, value, metadata = calls[0] assert query.shape == key.shape == value.shape == (2, 3, 2, 4) - if has_learned_registers: - assert metadata is None - else: - assert metadata.attn_mask.shape == (2, 3) - assert metadata.attn_mask.dtype == torch.bool - assert metadata.attn_mask.all() - torch.testing.assert_close(output, hidden_states) + assert query.dtype == key.dtype == value.dtype == torch.bfloat16 + assert metadata.attn_mask.shape == (2, 3) + assert metadata.attn_mask.dtype == torch.bool + assert metadata.attn_mask.all() + expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary) + expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4) + torch.testing.assert_close(output, expected, rtol=0, atol=0) def test_ltx2_distilled_pipeline_shares_recipe_runtime(): assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin) - assert issubclass(LTX2DistilledPipeline, LTX2Pipeline) + assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline) assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe assert "_forward_request" not in LTX2DistilledPipeline.__dict__ +def test_ltx_two_stage_entry_is_a_thin_pipeline_variant(): + assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__ + assert "_forward_request" not in LTX2TwoStagePipeline.__dict__ + assert "load_weights" not in LTX2TwoStagePipeline.__dict__ + assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin) + assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline) + + def test_ltx_variants_share_denoise_loop_and_i2v_conditioning(): assert "_denoise_loop" not in LTX2Pipeline.__dict__ assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin) @@ -594,7 +935,8 @@ def test_denoise_executor_owns_progress_and_interrupt(): torch.testing.assert_close(state.audio, torch.tensor(11.0)) -def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): +@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline]) +def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls): request_inputs = LTXRequestInputs( prompt="prompt", negative_prompt="", @@ -618,7 +960,7 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): max_sequence_length=16, ) prompt_context = object() - phase_calls = [] + phase_calls: list[Any] = [] def resolve_request_inputs(req, **kwargs): return request_inputs @@ -671,12 +1013,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) prompt_context_sentinel = prompt_context - pipeline = object.__new__(LTX2DistilledPipeline) + pipeline = object.__new__(pipeline_cls) torch.nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") pipeline.vae_spatial_compression_ratio = 32 pipeline.vae_temporal_compression_ratio = 8 pipeline.latent_upsampler = FakeUpsampler() + activated_slots: list[Any] = [] + if pipeline_cls is LTX2TwoStagePipeline: + pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append) object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) object.__setattr__(pipeline, "run_phase", run_phase) object.__setattr__(pipeline, "decode_phase", decode_phase) @@ -686,14 +1031,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): assert len(phase_calls) == 2 assert phase_calls[1][2] is prompt_context_sentinel + if pipeline_cls is LTX2TwoStagePipeline: + assert activated_slots == [None, "ltx_distilled"] torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0)) - torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0)) + expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0 + torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio)) -def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning(): +@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline]) +def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls): from vllm_omni.diffusion.models.ltx2 import ltx2_latents - pipeline = object.__new__(LTX2DistilledPipeline) + pipeline = object.__new__(pipeline_cls) torch.nn.Module.__init__(pipeline) pipeline.vae_spatial_compression_ratio = 32 pipeline.vae_temporal_compression_ratio = 8 @@ -750,6 +1099,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only(): assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536) +def test_ltx23_full_distilled_route_uses_merged_weights_without_lora(): + assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE + assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE + assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE + assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases) + assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None + assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls + assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == ( + "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" + ) + + @pytest.mark.parametrize( ("request_inputs", "error"), [ @@ -832,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas(): ) +def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested(): + request = LTXRequestInputs( + prompt="prompt", + negative_prompt="", + height=128, + width=128, + num_frames=9, + frame_rate=24.0, + num_inference_steps=30, + guidance=LTX2_ONE_STAGE_RECIPE.request_guidance, + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=16, + ) + + with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"): + validate_pipeline_request( + request, + pipeline_recipe=LTX2_ONE_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2Pipeline", + min_video_latent_spatial_size=(7, 7), + ) + + # ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid. + validate_pipeline_request( + request, + pipeline_recipe=LTX2_ONE_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2Pipeline", + ) + + validate_pipeline_request( + replace(request, height=256, width=256), + pipeline_recipe=LTX2_ONE_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2Pipeline", + min_video_latent_spatial_size=(7, 7), + ) + + @pytest.mark.parametrize( ("direct_kwargs", "sampling_kwargs", "error"), [ @@ -958,6 +1372,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps) +@pytest.mark.parametrize( + ("recipe", "steps", "stg_block"), + [ + (LTX2_TWO_STAGE_RECIPE, 40, 29), + (LTX23_TWO_STAGE_RECIPE, 30, 28), + ], +) +def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block): + stage1, stage2 = recipe.phases + assert (recipe.width, recipe.height) == (1536, 1024) + assert recipe.num_inference_steps == steps + assert stage1.spatial_downscale == 2 + assert stage1.adapter_slot is None + assert stage1.guidance.video.stg_blocks == (stg_block,) + assert stage1.guidance.audio.stg_blocks == (stg_block,) + assert stage2.input_transform == "spatial_upsample" + assert stage2.adapter_slot == "ltx_distilled" + assert stage2.guidance == LTXGuidanceSpec.positive_only() + assert stage2.num_inference_steps == 3 + assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0) + # Official Stage 2 refines video only; its audio result is discarded. + assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0) + + +def test_ltx_two_stage_entries_select_the_official_adapter_slots(): + phase_switches: list[Any] = [] + ordinary_pipeline = object.__new__(LTX2TwoStagePipeline) + ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append) + for phase in LTX2_TWO_STAGE_RECIPE.phases: + ordinary_pipeline._enter_phase(phase) + assert phase_switches == [None, "ltx_distilled"] + + distilled_pipeline = object.__new__(LTX2DistilledPipeline) + for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases: + distilled_pipeline._enter_phase(phase) + + +def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch): + class FakeColumn(torch.nn.Module): + def __init__( + self, + input_size, + output_size, + *, + bias, + gather_output, + return_bias, + quant_config, + prefix, + ): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.bias_enabled = bias + self.gather_output = gather_output + self.return_bias = return_bias + self.quant_config = quant_config + self.prefix = prefix + + class FakeRow(torch.nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + class FakeAttention(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + + monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn) + monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow) + monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention) + monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2) + + attention = ltx2_transformer.LTX2Attention( + query_dim=16, + heads=4, + dim_head=4, + qk_norm="rms_norm", + apply_gated_attention=True, + pack_qkv=False, + prefix="transformer_blocks.0.attn1", + ) + + gate = attention.to_gate_logits + assert isinstance(gate, FakeColumn) + assert gate.input_size == 16 + assert gate.output_size == 4 + assert gate.gather_output is False + assert gate.return_bias is False + assert gate.quant_config is None + assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits" + + class TestLTXRequestParsing: def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self): from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -1366,6 +1872,7 @@ class TestLTXForwardStages: assert seen["request_inputs"].output_type == "latent" assert seen["kwargs"] == { "request_sigmas": request_sigmas, + "request_phase_sigmas": None, "image": None, } @@ -1432,6 +1939,7 @@ class TestPipelineComponents: def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self): assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler") + assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler") class TestLTX23DecodeConditioning: @@ -1477,6 +1985,21 @@ class TestRegistryIntegration: from vllm_omni.diffusion.registry import _DIFFUSION_MODELS assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline") + assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2TwoStagePipeline", + ) + assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == ( + "ltx2", + "pipeline_ltx2", + "LTX2DistilledOneStagePipeline", + ) + assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2DistilledTwoStagePipeline", + ) assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == ( "ltx2", "pipeline_ltx2_two_stage", @@ -1488,6 +2011,9 @@ class TestRegistryIntegration: "LTX23ImageToVideoPipeline", "LTX2TwoStagesPipeline", "LTX2ImageToVideoTwoStagesPipeline", + "LTX2ImageToVideoTwoStagePipeline", + "LTX23TwoStagePipeline", + "LTX23ImageToVideoTwoStagePipeline", } assert removed_entries.isdisjoint(_DIFFUSION_MODELS) @@ -1497,6 +2023,9 @@ class TestRegistryIntegration: expected = [ "LTX2Pipeline", + "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", @@ -1509,6 +2038,9 @@ class TestRegistryIntegration: "model_class_name", [ "LTX2Pipeline", + "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index 73e10463e..e9bba1f1f 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -1,17 +1,60 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest +import torch from torch import nn from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +from vllm_omni.diffusion.models.ltx2 import ltx2_transformer +from vllm_omni.diffusion.models.ltx2.ltx2_transformer import ( + LTX2FeedForward, + LTX2VideoTransformer3DModel, + LTX2VideoTransformerBlock, + _make_rms_norm, + apply_interleaved_rotary_emb, + apply_keyframes_absolute_embedding, + apply_split_rotary_emb, +) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic(): + torch.manual_seed(0) + x = torch.randn(2, 3, 8, dtype=torch.bfloat16) + cos = torch.randn(2, 3, 8, dtype=x.dtype) + sin = torch.randn(2, 3, 8, dtype=x.dtype) + x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1) + rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) + expected = x * cos + rotated * sin + + actual = apply_interleaved_rotary_emb(x, (cos, sin)) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_ltx_split_rope_matches_official_model_dtype_arithmetic(): + torch.manual_seed(1) + x = torch.randn(2, 3, 8, dtype=torch.bfloat16) + cos = torch.randn(2, 3, 4, dtype=x.dtype) + sin = torch.randn(2, 3, 4, dtype=x.dtype) + split_x = x.reshape(2, 3, 2, 4) + first_x, second_x = split_x.chunk(2, dim=-2) + cos_u = cos.unsqueeze(-2) + sin_u = sin.unsqueeze(-2) + expected = split_x * cos_u + expected[..., :1, :].addcmul_(-sin_u, second_x) + expected[..., 1:, :].addcmul_(sin_u, first_x) + expected = expected.reshape_as(x) + + actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + def test_ltx_rms_norm_no_affine_matches_official_torch_module(): norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False) @@ -29,11 +72,53 @@ def test_ltx_rms_norm_affine_weight_remains_parameter(): assert "weight" in norm.state_dict() -def test_ltx_transformer_has_separate_cfg_cache_dit_config(): +def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms(): + calls = [] + + class UpcastingNorm(nn.Module): + def forward(self, tensor): + return tensor.float() + + class RecordingAttention: + def __call__(self, query, key, value, metadata): + calls.append((query, key, value, metadata)) + return query + + attention = SimpleNamespace( + heads=2, + head_dim=4, + to_qkv=None, + to_q=nn.Identity(), + to_k=nn.Identity(), + to_v=nn.Identity(), + norm_q=UpcastingNorm(), + norm_k=UpcastingNorm(), + to_gate_logits=None, + attn=RecordingAttention(), + to_out=(nn.Identity(), nn.Identity()), + rope_type="split", + ) + hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16) + rotary = ( + torch.ones(2, 2, 3, 2, dtype=torch.float32), + torch.zeros(2, 2, 3, 2, dtype=torch.float32), + ) + + processor = ltx2_transformer.LTX2AudioVideoAttnProcessor() + processor._slice_rope_for_tp = lambda rope, _: rope + output = processor(attention, hidden_states, query_rotary_emb=rotary) + + query, key, value, metadata = calls[0] + assert query.dtype == key.dtype == value.dtype == torch.bfloat16 + assert metadata is None + torch.testing.assert_close(output, hidden_states) + + +def test_ltx_transformer_has_fused_cfg_cache_dit_config(): adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config") assert isinstance(adapter_config, CacheDiTAdapterConfig) - assert adapter_config.has_separate_cfg + assert not adapter_config.has_separate_cfg def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks(): @@ -66,3 +151,332 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type): assert video_timestep.expected_dims == audio_timestep.expected_dims == 2 assert not video_timestep.split_output assert not audio_timestep.split_output + + +def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch): + processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor + monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True)) + processor = processor_cls() + hidden_states = torch.zeros(1, 3, 4) + encoder_hidden_states = torch.zeros(1, 4, 4) + key_padding_mask = torch.tensor([[True, True, True, False]]) + + actual = processor._prepare_attention_mask( + SimpleNamespace(), + hidden_states, + encoder_hidden_states, + key_padding_mask, + batch_size=1, + sequence_length=4, + ) + + torch.testing.assert_close(actual, key_padding_mask[:, None, None, :]) + + +def _fake_ltx_parallel_attention(): + calls = [] + + class Backend: + @staticmethod + def get_name(): + return "FLASH_ATTN" + + @staticmethod + def supports_attention_mask(): + return True + + def run(name): + def forward(query, _key, _value, _metadata): + calls.append(name) + return torch.zeros_like(query) + + return forward + + layer = object.__new__(ltx2_transformer._LTX2ParallelAttention) + nn.Module.__init__(layer) + layer.attn_backend = Backend + layer.attention = SimpleNamespace(forward=run("native")) + layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa")) + layer.backend_pref = "FLASH_ATTN" + return layer, calls + + +@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)]) +def test_ltx_masked_attention_uses_sdpa(query_length, key_length): + layer, calls = _fake_ltx_parallel_attention() + query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16) + key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16) + key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]]) + metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :]) + + layer._run_local_attention(query, key, value, metadata) + + assert calls == ["sdpa"] + + +def test_ltx_unmasked_attention_keeps_native_backend(): + layer, calls = _fake_ltx_parallel_attention() + query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16) + + layer._run_local_attention(query, key, value, None) + + assert calls == ["native"] + + +def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch): + class FakeParallelLinear(nn.Linear): + def __init__(self, dim_in, dim_out, *, bias=True, **kwargs): + super().__init__(dim_in, dim_out, bias=bias) + + monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear) + monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear) + + without_bias = LTX2FeedForward(8, inner_dim=16, bias=False) + with_bias = LTX2FeedForward(8, inner_dim=16, bias=True) + + assert without_bias.net[0].proj.bias is None + assert without_bias.net[2].bias is None + assert with_bias.net[0].proj.bias is not None + assert with_bias.net[2].bias is not None + + +def _build_tiny_ltx_transformer(monkeypatch, **overrides): + class FakeBlock(nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.init_kwargs = kwargs + + def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs): + return hidden_states, audio_hidden_states + + monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock) + kwargs = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 4, + "cross_attention_dim": 8, + "vae_scale_factors": (1, 1, 1), + "pos_embed_max_pos": 4, + "base_height": 8, + "base_width": 8, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "audio_scale_factor": 1, + "audio_pos_embed_max_pos": 4, + "audio_sampling_rate": 16, + "audio_hop_length": 4, + "num_layers": 1, + "caption_channels": 8, + "use_prompt_embeddings": False, + "cross_attn_mod": True, + "audio_cross_attn_mod": True, + } + kwargs.update(overrides) + return LTX2VideoTransformer3DModel(**kwargs) + + +def test_ltx25_keyframe_absolute_embedding_matches_official(): + hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8) + keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]]) + embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8) + + actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding) + expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states + + +def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch): + legacy_model = _build_tiny_ltx_transformer(monkeypatch) + + assert legacy_model.config.ff_bias is True + assert legacy_model.config.audio_ff_bias is True + assert legacy_model.config.use_prompt_adaln_single is True + assert legacy_model.config.use_keyframes_abs_pos_embedding is False + assert hasattr(legacy_model, "prompt_adaln") + assert hasattr(legacy_model, "audio_prompt_adaln") + assert not hasattr(legacy_model, "keyframes_abs_pos_embedding") + assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True + assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True + + ltx25_model = _build_tiny_ltx_transformer( + monkeypatch, + ff_bias=False, + audio_ff_bias=True, + use_prompt_adaln_single=False, + use_keyframes_abs_pos_embedding=True, + ) + + assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8) + torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8)) + assert not hasattr(ltx25_model, "prompt_adaln") + assert not hasattr(ltx25_model, "audio_prompt_adaln") + assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False + assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True + + monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1) + expected = torch.arange(8, dtype=torch.float32).view(1, 8) + loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)]) + + assert loaded == {"keyframes_abs_pos_embedding"} + torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected) + + +def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch): + class CaptureAda(nn.Module): + def __init__(self): + super().__init__() + self.calls = [] + + def forward(self, timestep, *, batch_size, hidden_dtype): + self.calls.append(timestep.detach().clone()) + embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype) + return embedding, embedding + + class CaptureRope(nn.Module): + def __init__(self): + super().__init__() + self.output_dtypes = [] + + def forward(self, _coords, *, device, out_dtype=None): + self.output_dtypes.append(out_dtype) + empty = torch.empty(0, device=device, dtype=out_dtype) + return empty, empty + + model = _build_tiny_ltx_transformer(monkeypatch) + rope = CaptureRope() + model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope + model.prompt_modulation = False + video_scale_shift = CaptureAda() + video_gate = CaptureAda() + audio_scale_shift = CaptureAda() + audio_gate = CaptureAda() + model.av_cross_attn_video_scale_shift = video_scale_shift + model.av_cross_attn_video_a2v_gate = video_gate + model.av_cross_attn_audio_scale_shift = audio_scale_shift + model.av_cross_attn_audio_v2a_gate = audio_gate + + model( + hidden_states=torch.zeros(1, 2, 4), + audio_hidden_states=torch.zeros(1, 1, 4), + encoder_hidden_states=torch.zeros(1, 1, 8), + audio_encoder_hidden_states=torch.zeros(1, 1, 8), + timestep=torch.tensor([[0.0, 2.0]]), + audio_timestep=torch.tensor([[5.0]]), + sigma=torch.tensor([2.0]), + audio_sigma=torch.tensor([5.0]), + video_coords=torch.zeros(1, 3, 2, 2), + audio_coords=torch.zeros(1, 1, 1, 2), + ) + + torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0])) + torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0])) + torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0])) + torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0])) + assert len(rope.output_dtypes) == 4 + assert all(dtype is torch.float32 for dtype in rope.output_dtypes) + + +def test_ltx25_static_prompt_table_is_used_without_prompt_adaln(): + class CaptureAttention(nn.Module): + def __init__(self): + super().__init__() + self.encoder_hidden_states = None + + def forward(self, hidden_states, encoder_hidden_states=None, **kwargs): + if encoder_hidden_states is not None: + self.encoder_hidden_states = encoder_hidden_states.detach().clone() + return torch.zeros_like(hidden_states) + + class ZeroOutput(nn.Module): + def forward(self, hidden_states): + return torch.zeros_like(hidden_states) + + block = object.__new__(LTX2VideoTransformerBlock) + nn.Module.__init__(block) + block.video_cross_attn_adaln = True + block.audio_cross_attn_adaln = True + block.cross_attn_adaln = True + block.scale_shift_table = nn.Parameter(torch.zeros(9, 2)) + block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2)) + block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]])) + block.norm1 = nn.Identity() + block.audio_norm1 = nn.Identity() + block.norm2 = nn.Identity() + block.audio_norm2 = nn.Identity() + block.norm3 = nn.Identity() + block.audio_norm3 = nn.Identity() + block.attn1 = CaptureAttention() + block.audio_attn1 = CaptureAttention() + block.attn2 = CaptureAttention() + block.audio_attn2 = CaptureAttention() + block.ff = ZeroOutput() + block.audio_ff = ZeroOutput() + + hidden_states = torch.ones(1, 1, 2) + audio_hidden_states = torch.ones(1, 1, 2) + encoder_hidden_states = torch.ones(1, 1, 2) + audio_encoder_hidden_states = torch.ones(1, 1, 2) + block( + hidden_states, + audio_hidden_states, + encoder_hidden_states, + audio_encoder_hidden_states, + temb=torch.zeros(1, 1, 18), + temb_audio=torch.zeros(1, 1, 18), + temb_ca_scale_shift=torch.empty(0), + temb_ca_audio_scale_shift=torch.empty(0), + temb_ca_gate=torch.empty(0), + temb_ca_audio_gate=torch.empty(0), + temb_prompt=None, + temb_prompt_audio=None, + use_a2v_cross_attention=False, + use_v2a_cross_attention=False, + ) + + torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]])) + torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]])) + + +def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch): + column_calls = [] + + class FakeColumnParallelLinear(nn.Linear): + def __init__(self, input_size, output_size, *, bias=True, **kwargs): + column_calls.append((output_size, kwargs)) + super().__init__(input_size, output_size, bias=bias) + + class FakeRowParallelLinear(nn.Linear): + def __init__(self, input_size, output_size, *, bias=True, **_kwargs): + super().__init__(input_size, output_size, bias=bias) + + class FakeParallelAttention(nn.Module): + def __init__(self, **_kwargs): + super().__init__() + + monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear) + monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear) + monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention) + monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2) + + attention = ltx2_transformer.LTX2Attention( + query_dim=8, + heads=4, + kv_heads=4, + dim_head=2, + qk_norm="rms_norm", + apply_gated_attention=True, + pack_qkv=False, + ) + + gate_output_size, gate_kwargs = column_calls[-1] + assert gate_output_size == 4 + assert gate_kwargs["gather_output"] is False + assert gate_kwargs["return_bias"] is False + assert attention.query_num_heads == 2 diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 108bb960d..964bb7d06 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior.""" +"""Unit tests for LTX video VAE tiling and distributed decode behavior.""" from types import SimpleNamespace +from typing import Any import pytest import torch @@ -11,6 +12,466 @@ import torch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime + + class FakeBaseVocoder(torch.nn.Module): + def __init__(self): + super().__init__() + self.input_dtype = None + + def forward(self, value): + self.input_dtype = value.dtype + return value + + monkeypatch.setattr( + ltx_runtime.torch, + "autocast", + lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")), + ) + vocoder = FakeBaseVocoder() + + output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16)) + + assert vocoder.input_dtype == torch.bfloat16 + assert output.dtype == torch.bfloat16 + + +class TestLTXDiffusionDecoder: + @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"]) + def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None) + + with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"): + model.set_parallel_size(2, mode=mode) + + def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + calls = [] + model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.distributed_executor = SimpleNamespace( + set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode)) + ) + + model.set_parallel_size(4, mode="tile") + + assert calls == [(4, "tile")] + + def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + LTX2VideoDiffusionDecoder3d, + LTX2VideoDiffusionDecoderModel, + ) + + class IdentityAttention: + def build_block_mask(self, _hidden_states): + return None + + class IdentityBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.attn = IdentityAttention() + + def forward(self, hidden_states, _block_mask): + return hidden_states + + class IdentityUpsample(torch.nn.Module): + def forward(self, hidden_states, *, drop_leading_frame): + assert drop_leading_frame + return hidden_states + + decoder = object.__new__(LTX2VideoDiffusionDecoder3d) + torch.nn.Module.__init__(decoder) + decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])]) + decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()]) + decoder.trailing_pad_latent_frames = 2 + decoder.temporal_compression_ratio = 8 + decoder.stage5_kernel = (11, 3, 3) + + # 17 frames contain one real frame followed by 16 propagated ghost + # frames. The reference keeps 11 for stage-5 NATTEN instead of + # cropping the tensor immediately back to one. + context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4)) + assert context.shape == (1, 11, 3, 3, 4) + + class FakeDecoder(torch.nn.Module): + def forward(self, z, *, generator, num_inference_steps): + assert generator is None + assert num_inference_steps is None + return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2) + + model = object.__new__(LTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.decoder = FakeDecoder() + model.use_tiling = False + model.tile_sample_min_num_frames = 80 + model.tile_sample_min_height = 768 + model.tile_sample_min_width = 768 + model.temporal_compression_ratio = 8 + model.spatial_compression_ratio = 2 + + output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0] + assert output.shape == (1, 3, 1, 4, 6) + + @pytest.mark.parametrize( + ("model_version", "extras", "expected"), + [ + ("2.5", {}, True), + ("2.5", {"ltx2_use_conv_vae": True}, False), + ("2.3", {}, False), + ], + ) + def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected): + from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder + + assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected + + def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self): + from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder + + with pytest.raises(TypeError, match="ltx2_use_conv_vae"): + _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5") + + def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + convert_ltx25_native_diffusion_decoder_state_dict, + ) + + qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4) + qkv_bias = torch.arange(6, dtype=torch.float32) + projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2) + projection_bias = torch.ones(2) + gate = torch.tensor([2.0, 3.0]) + mean = torch.tensor([0.1, 0.2]) + std = torch.tensor([0.3, 0.4]) + native = { + "encoder.ignored": torch.ones(1), + "per_channel_statistics.mean-of-means": mean, + "vae.per_channel_statistics.std-of-means": std, + "decoder.type_emb": torch.ones(2), + "decoder.coarse_head.weight": torch.ones(1), + "decoder.det_stages.0.0.attn.qkv.weight": qkv_weight, + "decoder.det_stages.0.0.attn.qkv.bias": qkv_bias, + "decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]), + "decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]), + "decoder.diff_blocks.0.attn.proj.weight": projection_weight, + "decoder.diff_blocks.0.attn.proj.bias": projection_bias, + "decoder.diff_blocks.0.gate_msa": gate, + "vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2), + } + + converted = convert_ltx25_native_diffusion_decoder_state_dict(native) + + torch.testing.assert_close(converted["latents_mean"], mean) + torch.testing.assert_close(converted["latents_std"], std) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:]) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0])) + torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0])) + torch.testing.assert_close( + converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight + ) + torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias) + assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted + assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted) + + def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + convert_ltx25_native_diffusion_decoder_state_dict, + ) + + with pytest.raises(ValueError, match="not divisible by 3"): + convert_ltx25_native_diffusion_decoder_state_dict( + {"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)} + ) + + def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + seen = {} + + class ConvVae: + dtype = torch.float32 + config = SimpleNamespace(timestep_conditioning=True) + + def decode(self, *_args, **_kwargs): + raise AssertionError("the convolutional VAE decoder must not run") + + class DiffusionDecoder: + dtype = torch.float32 + + def decode(self, latents, *, generator, return_dict): + seen["latents"] = latents + seen["generator"] = generator + seen["return_dict"] = return_dict + return (latents + 1,) + + class AudioVae: + dtype = torch.float32 + + def decode(self, audio_latents, *, return_dict): + assert return_dict is False + return (audio_latents + 2,) + + class VideoProcessor: + def postprocess_video(self, video, *, output_type): + seen["output_type"] = output_type + return video + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.use_diffusion_decoder = True + pipe.vae = ConvVae() + pipe.diffusion_decoder = DiffusionDecoder() + pipe.audio_vae = AudioVae() + pipe.vocoder = torch.nn.Identity() + pipe.video_processor = VideoProcessor() + + generator = torch.Generator().manual_seed(123) + latents = torch.ones(1, 1) + output = pipe._decode_output( + latents=latents, + audio_latents=torch.ones(1, 1), + output_type="pt", + connector_prompt_embeds=torch.ones(1, 1), + generator=generator, + device=torch.device("cpu"), + decode_timestep=0.7, + decode_noise_scale=0.5, + prompt_batch_size=1, + ) + + torch.testing.assert_close(seen["latents"], latents) + assert seen["generator"] is generator + assert seen["return_dict"] is False + assert seen["output_type"] == "pt" + torch.testing.assert_close(output.output[0], latents + 1) + torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0)) + + def test_diffusion_decode_denormalizes_with_decoder_statistics(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.use_diffusion_decoder = True + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.vae = SimpleNamespace( + latents_mean=torch.full((2,), 100.0), + latents_std=torch.full((2,), 10.0), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipe.diffusion_decoder = SimpleNamespace( + latents_mean=torch.tensor([10.0, 20.0]), + latents_std=torch.tensor([2.0, 4.0]), + config=SimpleNamespace(scaling_factor=2.0), + ) + pipe.audio_vae = SimpleNamespace( + latents_mean=torch.tensor(0.0), + latents_std=torch.tensor(1.0), + ) + forward_ctx = SimpleNamespace( + latent_num_frames=1, + latent_height=1, + latent_width=1, + original_audio_num_frames=1, + latent_mel_bins=2, + ) + + video, audio = pipe._unpack_and_denormalize_stage( + forward_ctx, + torch.tensor([[[1.0, 2.0]]]), + torch.zeros(1, 1, 2), + ) + + torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0])) + torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2)) + + def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + class ConvVae: + config = SimpleNamespace(timestep_conditioning=False) + + def is_distributed_enabled(self): + raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives") + + class DiffusionDecoder: + dtype = torch.float32 + + def decode(self, *_args, **_kwargs): + raise AssertionError("only the output rank may run DiffVAE decode") + + monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1) + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.use_diffusion_decoder = True + pipe.vae = ConvVae() + pipe.diffusion_decoder = DiffusionDecoder() + + output = pipe._decode_output( + latents=torch.ones(1, 1), + audio_latents=torch.ones(1, 1), + output_type="pt", + connector_prompt_embeds=torch.ones(1, 1), + generator=None, + device=torch.device("cpu"), + decode_timestep=0.0, + decode_noise_scale=None, + prompt_batch_size=1, + ) + + assert output.output[0].numel() == 0 + assert output.output[1].numel() == 0 + + def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + + class ConvVae: + config = SimpleNamespace(timestep_conditioning=False) + + def is_distributed_enabled(self): + raise AssertionError("DiffVAE decode must query the active decoder") + + class DiffusionDecoder: + dtype = torch.float32 + + def __init__(self): + self.decode_calls = 0 + + def is_distributed_enabled(self): + return True + + def decode(self, *_args, **_kwargs): + self.decode_calls += 1 + return (torch.empty(0, 3, 0, 0, 0),) + + monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1) + pipe = object.__new__(LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.use_diffusion_decoder = True + pipe.vae = ConvVae() + pipe.diffusion_decoder = DiffusionDecoder() + + output = pipe._decode_output( + latents=torch.ones(1, 1), + audio_latents=torch.ones(1, 1), + output_type="pt", + connector_prompt_embeds=torch.ones(1, 1), + generator=None, + device=torch.device("cpu"), + decode_timestep=0.0, + decode_noise_scale=None, + prompt_batch_size=1, + ) + + assert pipe.diffusion_decoder.decode_calls == 1 + assert output.output[0].numel() == 0 + assert output.output[1].numel() == 0 + + def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + expected = torch.ones(1, 3, 1, 2, 2) + seen = {} + + def native_tiled_decode(self, z, generator=None, num_inference_steps=None): + seen["z"] = z + seen["generator"] = generator + seen["num_inference_steps"] = num_inference_steps + return expected + + monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode) + model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel) + torch.nn.Module.__init__(model) + model.is_distributed_enabled = lambda: False + generator = torch.Generator().manual_seed(17) + z = torch.zeros(1, 1, 1, 1, 1) + + result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode( + model, + z, + generator=generator, + num_inference_steps=1, + ) + + assert result is expected + assert seen["z"] is z + assert seen["generator"] is generator + assert seen["num_inference_steps"] == 1 + + def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch): + from diffusers.utils.torch_utils import randn_tensor + + from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + DistributedLTX2VideoDiffusionDecoderModel, + ) + + model = DistributedLTX2VideoDiffusionDecoderModel( + out_channels=1, + latent_channels=1, + patch_size=1, + decoder_head_dim=8, + decoder_stage_channels=(8, 8, 8, 8, 8), + decoder_stage_depths=(1, 1, 1, 1, 1), + decoder_stage_kernels=((1, 1, 1),) * 4, + decoder_upsample_strides=((1, 1, 1),) * 4, + decoder_upsample_channel_reductions=(1, 1, 1, 1), + decoder_stage5_kernel=(1, 1, 1), + spatial_compression_ratio=1, + temporal_compression_ratio=1, + ) + model.tile_sample_min_num_frames = 2 + model.tile_sample_stride_num_frames = 1 + model.tile_sample_min_height = 2 + model.tile_sample_stride_height = 1 + model.tile_sample_min_width = 2 + model.tile_sample_stride_width = 1 + model.distributed_executor = SimpleNamespace(group=None) + model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1) + monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None) + + z = torch.zeros(1, 1, 2, 3, 3) + distributed_generator = torch.Generator().manual_seed(123) + reference_generator = torch.Generator().manual_seed(123) + tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1) + + assert grid_spec.grid_shape == (1, 3, 3) + assert grid_spec.split_dims == (2, 3, 4) + for task in tasks: + tile_shape = model._tiled_pixel_shape_from_features( + task.tensor, + drop_leading_frame=task.drop_leading_frame, + crop_trailing_ghost=task.crop_trailing_ghost, + ) + expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype) + actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype) + torch.testing.assert_close(actual, expected) + + assert torch.equal(distributed_generator.get_state(), reference_generator.get_state()) + + class TestLTXOutputRank: @pytest.mark.parametrize( ("distributed_vae_state", "expected_decode_calls"), @@ -113,7 +574,7 @@ class TestLTX23VaeDistributedDecode: }, output_dtype=torch.float32, ) - seen = {} + seen: dict[str, Any] = {} def exec_tile(task): return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1)) diff --git a/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py new file mode 100644 index 000000000..d3d885d57 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CUDA regression tests for LTX-2 BWE vocoder precision.""" + +import pytest +import torch + +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +class _BWEConvVocoder(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv1d( + 1, + 1, + kernel_size=1, + bias=False, + device="cuda", + dtype=torch.bfloat16, + ) + self.bwe_generator = torch.nn.Identity() + self.input_dtype = None + self.conv_output_dtype = None + + def forward(self, value): + self.input_dtype = value.dtype + output = self.conv(value) + self.conv_output_dtype = output.dtype + return output + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32(): + from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder + + vocoder = _BWEConvVocoder() + generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16) + + output = _run_ltx_vocoder(vocoder, generated_mel) + + assert next(vocoder.parameters()).dtype == torch.bfloat16 + assert vocoder.input_dtype == torch.float32 + assert vocoder.conv_output_dtype == torch.float32 + assert output.dtype == torch.bfloat16 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 60ed7090f..382ecf9aa 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import json +import sys from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace -from unittest.mock import Mock +from typing import Any, TypeVar +from unittest.mock import Mock, patch import numpy as np import pytest @@ -13,6 +15,14 @@ import torch.nn as nn pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +_ItemT = TypeVar("_ItemT") +_ResultT = TypeVar("_ResultT") + + +def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ResultT: + items.append(item) + return result + def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -28,6 +38,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): pipeline.default_audio_shift = 3.0 pipeline.device = torch.device("cpu") pipeline.od_config = SimpleNamespace() + pipeline.text_encoder = object() cache_spec = object() quality_plan = SimpleNamespace(cache_dit=cache_spec) pipeline._quality_policy = Mock() @@ -46,7 +57,8 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): return torch.zeros(1), torch.zeros(1) pipeline.diffuse = diffuse - pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + decoded_audio = torch.zeros(1) + pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), decoded_audio)) sampling = OmniDiffusionSamplingParams( quality="high", width=1344, @@ -75,7 +87,9 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): num_inference_steps=50, extra_args={"task": "t2va", "aspect_ratio": "16:9"}, ) - assert output.output == pipeline.decode.return_value + assert output.output[0].shape == (1, 1, 1, 1, 3) + assert output.output[0].dtype == torch.uint8 + assert output.output[1] is decoded_audio def test_pipeline_import_registry_and_component_discovery(): @@ -100,6 +114,40 @@ def test_pipeline_import_registry_and_component_discovery(): assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"] +def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = torch.nn.Linear(1, 1) + pipeline.video_vae = torch.nn.Linear(1, 1) + pipeline.audio_vae = torch.nn.Linear(1, 1) + pipeline.text_encoder = None + pipeline._dit_modules = ["transformer"] + pipeline._encoder_modules = [] + pipeline._vae_modules = ["video_vae", "audio_vae"] + + discovered = ModuleDiscovery.discover(pipeline) + + assert discovered.dit_names == ["transformer"] + assert discovered.encoder_names == [] + assert discovered.vae_names == ["video_vae", "audio_vae"] + + +def test_encoder_free_stage_skips_missing_component_during_dlo_registration(): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + cache = Mock() + video_vae = Mock() + audio_vae = Mock() + + pipeline_module._register_dlo_component_cache(cache, None, video_vae, audio_vae) + + video_vae.set_omni_component_cache.assert_called_once_with(cache) + audio_vae.set_omni_component_cache.assert_called_once_with(cache) + + def _write_partition_index(path, *, partition, tasks): path.mkdir(parents=True) (path / "model_index.json").write_text( @@ -184,6 +232,7 @@ def test_startup_task_rejects_unsupported_value(): def test_combined_task_inference_and_transformer_routing(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.errors import OmniClientError pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) @@ -195,6 +244,9 @@ def test_combined_task_inference_and_transformer_routing(): assert pipeline._resolve_task(None, {"image": object()}) == "fl2va" assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va" assert pipeline._resolve_task(None, {"video": object()}) == "ref2va" + assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va" + with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"): + pipeline._resolve_task("ref2va", {}, has_turbo_lora=True) pipeline.partition = "ref2va" pipeline.supported_tasks = frozenset({"ref2va"}) assert pipeline._resolve_task(None, {"image": object()}) == "ref2va" @@ -259,7 +311,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once( ): (tmp_path / partition_name / component).mkdir() - created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} + created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []} class FakeModule(torch.nn.Module): def __init__(self, *args, **kwargs): @@ -277,15 +329,15 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once( return create - tokenizer_calls = [] - processor_calls = [] + tokenizer_calls: list[Any] = [] + processor_calls: list[Any] = [] tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast) - tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: ( - tokenizer_calls.append((args, kwargs)) or object() + tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return( + tokenizer_calls, (args, kwargs), object() ) processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor) - processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: ( - processor_calls.append((args, kwargs)) or object() + processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return( + processor_calls, (args, kwargs), object() ) monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT) monkeypatch.setattr( @@ -362,8 +414,18 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once( if expected_partition == "ref2va": assert pipeline._transformer_for_task("ref2va") is pipeline.transformer assert [source.model_or_path for source in pipeline.weights_sources] == [ - str(tmp_path / partition_name) for partition_name in source_partitions + *(str(tmp_path / partition_name) for partition_name in source_partitions), + str(component_path), ] + assert [source.subfolder for source in pipeline.weights_sources] == [ + *("transformer" for _ in source_partitions), + "text_encoder", + ] + expected_prefixes = ["transformer."] + if expected_dits == 2: + expected_prefixes.append("transformers_ref.") + expected_prefixes.append("text_encoder.") + assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes expected_patterns = ( pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS if expected_partition == "combined" @@ -434,18 +496,41 @@ def test_joint_postprocess_is_multiprocessing_picklable(): postprocess = get_minimax_h3_post_process_func(SimpleNamespace()) postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess)) - video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5) + # decode() hands over quantised frames already laid out as (B, T, H, W, C). + video = torch.arange(2 * 2 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 2, 4, 5, 3) audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6) result = postprocess((video, audio), output_type="np") assert isinstance(result["video"], list) + assert result["video"][0].dtype == np.uint8 assert result["video"][0].shape == (2, 4, 5, 3) + np.testing.assert_array_equal(result["video"][0], video[0].numpy()) np.testing.assert_array_equal(result["audio"], audio.numpy()) assert result["audio_sample_rate"] == 32000 assert result["fps"] == 24 +def test_joint_video_output_is_quantized_before_transfer(): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _minimax_h3_post_process, + _prepare_minimax_h3_video_output, + ) + + video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5) + expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1) + prepared = _prepare_minimax_h3_video_output(video) + + assert prepared.dtype == torch.uint8 + assert prepared.is_contiguous() + assert prepared.shape == (2, 2, 4, 5, 3) + torch.testing.assert_close(prepared, expected, rtol=0, atol=0) + + result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6))) + assert result["video"][0].dtype == np.uint8 + np.testing.assert_array_equal(result["video"][0], expected[0].numpy()) + + def test_cfg_parallel_is_rejected_for_distilled_checkpoint(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -513,7 +598,7 @@ def test_base_schedule_overrides_the_uniform_sigma_positions(): @pytest.mark.parametrize( "base_schedule", [ - [0.9, 0.5, 0.0], + [1.5, 0.5, 0.0], [1.0, 0.5, 0.1], [1.0, 0.5, 0.5, 0.0], [1.0], @@ -551,6 +636,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition): pipeline.default_audio_shift = 3.0 pipeline.device = torch.device("cpu") pipeline.od_config = SimpleNamespace() + pipeline.text_encoder = Mock() pipeline._base_schedule_by_partition = schedules pipeline._quality_policy = Mock() pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None) @@ -567,7 +653,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition): return torch.zeros(1), torch.zeros(1) pipeline.diffuse = diffuse - pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), torch.zeros(1))) return pipeline @@ -615,7 +701,7 @@ def test_partially_constructed_pipeline_falls_back_to_the_uniform_schedule(): def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries(): - diffuse_calls = [] + diffuse_calls: list[dict[str, Any]] = [] base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0] pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": base_schedule, "ref2va": None}) @@ -632,7 +718,7 @@ def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries(): def test_distilled_forward_accepts_the_matching_explicit_step_count(): - diffuse_calls = [] + diffuse_calls: list[dict[str, Any]] = [] pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None}) pipeline.forward(_t2va_batch(num_inference_steps=4)) @@ -702,8 +788,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): ) assert output.shape == q.shape - assert attention.attention.metadata.attn_mask is None - assert attention.attention.metadata.extra["valid_kv_length"] == 5 + metadata = attention.attention.metadata + assert metadata is not None + assert metadata.attn_mask is None + assert metadata.extra["valid_kv_length"] == 5 def test_packed_attention_skips_mask_for_packed_mask_free_backend(): @@ -743,9 +831,17 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend(): packed_total=8, ) - assert attention.attention.metadata.attn_mask is None - assert attention.attention.metadata.extra["valid_kv_length"] == 5 - assert attention.attention.metadata.extra["npu_attn_varlen"] is True + metadata = attention.attention.metadata + assert metadata is not None + assert metadata.attn_mask is None + assert metadata.extra["valid_kv_length"] == 5 + assert metadata.extra["npu_attn_varlen"] is True + packed_padding = metadata.packed_padding + assert packed_padding is not None + assert packed_padding.q_length == 5 + assert packed_padding.kv_length == 5 + assert packed_padding.cu_seqlens_q.tolist() == [0, 5] + assert packed_padding.cu_seqlens_k.tolist() == [0, 5] def test_packed_attention_keeps_padding_mask_for_other_backends(): @@ -785,12 +881,90 @@ def test_packed_attention_keeps_padding_mask_for_other_backends(): packed_total=8, ) + metadata = attention.attention.metadata + assert metadata is not None assert torch.equal( - attention.attention.metadata.attn_mask, + metadata.attn_mask, torch.tensor([[True, True, True, True, True, False, False, False]]), ) +def _fake_packed_attention( + backend_name: str, + *, + prefix_kv_slicing: bool = False, + supports_multi_doc: bool = False, +): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + class FakeBackend: + supports_prefix_kv_slicing = prefix_kv_slicing + + @staticmethod + def get_name() -> str: + return backend_name + + @classmethod + def supports_multi_doc_packed_varlen(cls) -> bool: + return supports_multi_doc + + class FakeAttention(torch.nn.Module): + attn_backend = FakeBackend + + def __init__(self): + super().__init__() + self.metadata = None + + def forward(self, query, key, value, metadata): + self.metadata = metadata + return query + + attention = object.__new__(MiniMaxH3Attention) + torch.nn.Module.__init__(attention) + attention.attention = FakeAttention() + return attention + + +def test_packed_attention_drops_prefix_semantics_for_a_co_batched_layout(): + """A batch's valid rows are block-diagonal, so no prefix can describe them.""" + attention = _fake_packed_attention("FLASH_ATTN", supports_multi_doc=True) + q = torch.randn(12, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + # Two requests: rows + padding tail each. + cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32), + max_seqlen=5, + packed_total=12, + num_requests=2, + ) + + assert attention.attention.metadata.attn_mask is None + assert attention.attention.metadata.extra["valid_kv_length"] == 12 + + +@pytest.mark.parametrize("backend_name", ["TORCH_SDPA", "FLASH_ATTN"]) +def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_name): + """Backend names alone must not allow cross-request attention.""" + attention = _fake_packed_attention(backend_name, supports_multi_doc=False) + q = torch.randn(12, 2, 4) + + with pytest.raises(ValueError, match=backend_name): + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32), + max_seqlen=5, + packed_total=12, + num_requests=2, + ) + + def test_reference_image_resize_contract(): from PIL import Image @@ -852,6 +1026,80 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit(): assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9 +def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + + class FakeEncoderGroup: + rank_in_group = 0 + world_size = 2 + + def all_reduce(self, tensor): + del tensor + + config = SimpleNamespace( + hidden_size=8, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=2, + rms_norm_eps=1e-6, + ) + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + attention = encoder_module.MiniMaxH3Qwen3VLTextAttention(FakeEncoderGroup(), config, torch.float32) + # This CPU contract test only needs projection shapes and head reshaping. + # Avoid dispatching platform-specific norm and linear kernels (for example, + # rocm_unquantized_gemm when the test suite is collected on ROCm). + attention.q_norm = nn.Identity() + attention.k_norm = nn.Identity() + with torch.no_grad(): + attention.qkv_proj.weight.zero_() + attention.o_proj = nn.Linear( + attention.qkv_proj.local_num_heads * attention.head_dim, + config.hidden_size, + bias=False, + ) + attn_call: dict[str, object] = {} + + def fake_attention(query, key, value): + attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape) + return query + + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", fake_attention) + hidden_states = torch.randn(1, 3, config.hidden_size) + cos = torch.ones(1, 3, config.head_dim) + sin = torch.zeros_like(cos) + + output = attention(hidden_states, (cos, sin)) + + assert attn_call["query_shape"] == (1, config.num_attention_heads // 2, 3, config.head_dim) + assert attn_call["key_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim) + assert attn_call["value_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim) + assert output.shape == hidden_states.shape + + +def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + + sdpa_call: dict[str, object] = {} + + def fake_sdpa(query, key, value, **kwargs): + sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs) + return query + + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + output = encoder_module._scaled_dot_product_attention( + torch.randn(1, 4, 3, 8), + torch.randn(1, 2, 3, 8), + torch.randn(1, 2, 3, 8), + ) + + assert sdpa_call["query_shape"] == sdpa_call["key_shape"] == sdpa_call["value_shape"] == (1, 4, 3, 8) + assert sdpa_call["kwargs"] == {"dropout_p": 0.0, "is_causal": True} + assert output.shape == (1, 4, 3, 8) + + def test_encoder_forward_uses_hook_compatible_encode_entrypoint(): from vllm_omni.diffusion.models.minimax_h3.encoder import ( MiniMaxH3Qwen3VLEncoder, @@ -909,7 +1157,7 @@ def test_encoder_forward_forwards_video_inputs(): def test_reference_video_shape_uses_h3_adapt_shape_policy(): - from vllm_omni.diffusion.models.minimax_h3.reference_video import ( + from vllm_omni.model_executor.models.minimax_h3.reference_video import ( _reference_video_shape, ) @@ -935,6 +1183,66 @@ def test_text_encoder_stub_constructs_without_group_or_weights(): assert list(encoder.named_parameters()) == [] +def test_global_quant_config_is_shared_by_dit_and_encoder(): + from vllm_omni.quantization import build_quant_config, resolve_component_quant_config + + global_config = build_quant_config("fp8") + + assert resolve_component_quant_config(global_config, "transformer") is global_config + assert resolve_component_quant_config(global_config, "text_encoder") is global_config + + +def test_minimax_h3_fp8_component_selection(): + from vllm_omni.quantization import build_quant_config, resolve_component_quant_config + + dit_only = build_quant_config({"transformer": {"method": "fp8"}}) + encoder_only = build_quant_config({"text_encoder": {"method": "fp8"}}) + + assert resolve_component_quant_config(dit_only, "transformer").get_name() == "fp8" + assert resolve_component_quant_config(dit_only, "text_encoder") is None + assert resolve_component_quant_config(encoder_only, "transformer") is None + assert resolve_component_quant_config(encoder_only, "text_encoder").get_name() == "fp8" + + +def test_text_encoder_rejects_serialized_fp8(): + from vllm_omni.diffusion.models.minimax_h3.encoder import _validate_encoder_quant_config + + config = SimpleNamespace( + get_name=lambda: "fp8", + is_checkpoint_fp8_serialized=True, + ) + with pytest.raises(ValueError, match="online FP8 only"): + _validate_encoder_quant_config(config) + + +def test_text_encoder_linear_delegates_quantization_to_vllm_factory(): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module + + group = SimpleNamespace(rank_in_group=0, world_size=1) + method = UnquantizedLinearMethod() + quant_config = Mock() + quant_config.get_quant_method.return_value = method + prefix = "text_encoder.text_model.layers.0.mlp.gate_up_proj" + + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + linear = encoder_module.MiniMaxH3Qwen3VLMergedColumnParallelLinear( + group, + input_size=4, + intermediate_size=8, + dtype=torch.bfloat16, + quant_config=quant_config, + prefix=prefix, + ) + + assert linear.quant_method is method + quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix) + + def test_no_offload_keeps_text_encoder_resident(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline @@ -967,6 +1275,7 @@ def test_model_offload_uses_hooked_text_encoder_call(): ) expected = torch.ones(2, 3) pipeline.text_encoder = Mock(return_value=expected) + pipeline._model_cpu_offload_modules = [pipeline.text_encoder] input_ids = torch.tensor([1, 2]) vision_kwargs = {"pixel_values": torch.ones(1, 4)} @@ -977,15 +1286,21 @@ def test_model_offload_uses_hooked_text_encoder_call(): pipeline.text_encoder.load_to_device.assert_not_called() -def test_layerwise_offload_releases_text_encoder(): +@pytest.mark.parametrize( + "offload_flag", + ["enable_layerwise_offload", "enable_distributed_layerwise_offload"], +) +def test_layerwise_offload_releases_text_encoder(offload_flag): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline pipeline = object.__new__(MiniMaxH3Pipeline) torch.nn.Module.__init__(pipeline) pipeline.od_config = SimpleNamespace( enable_cpu_offload=False, - enable_layerwise_offload=True, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, ) + setattr(pipeline.od_config, offload_flag, True) pipeline.text_encoder = Mock() expected = torch.ones(2, 3) pipeline.text_encoder.encode_ids.return_value = expected @@ -1324,7 +1639,7 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch): lambda path: f"frames:{path}", ) - rows, shapes = pipeline._encode_video_conditions(None, count=2) + rows, shapes = pipeline._encode_video_conditions_resident(None, count=2) assert pipeline.video_vae.calls == [ "frames:video-1.mp4", @@ -1406,12 +1721,12 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti "audio_codecs": (), "file_size": 1024, } - transcode_calls = [] + transcode_calls: list[Any] = [] monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata) monkeypatch.setattr( reference_video_module, "_transcode_reference_video", - lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4", + lambda source, **kwargs: _append_and_return(transcode_calls, (source, kwargs), "prepared.mp4"), ) prepared = reference_video_module.prepare_reference_videos( @@ -1428,34 +1743,305 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti assert transcode_calls[0][1]["target_frame_count"] == 209 -def test_ref2va_qwen_sampling_decodes_prepared_video_once(monkeypatch): +def test_ref2va_transcode_zero_frame_count_keeps_video_stream(monkeypatch, tmp_path): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module - decoded = np.arange(25 * 2 * 2 * 3, dtype=np.uint8).reshape(25, 2, 2, 3) - load_calls = [] + commands: list[list[str]] = [] + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda command, **_kwargs: commands.append(command), + ) + + reference_video_module._transcode_reference_video( + "input.mp4", + target_width=768, + target_height=448, + target_frame_count=0, + workdir=str(tmp_path), + duration_seconds=4.0, + ) + + assert "-frames:v" not in commands[0] + + +def test_prepared_reference_video_descriptor_round_trip(tmp_path): + from vllm_omni.model_executor.models.minimax_h3.reference_video import ( + deserialize_prepared_reference_videos, + serialize_prepared_reference_videos, + ) + + item = { + "original_path": str(tmp_path / "source.mp4"), + "prepared_path": str(tmp_path / "prepared.mp4"), + "input_has_audio": True, + "width": 1344, + "height": 768, + "start_time_seconds": 0.0, + "duration_seconds": 5.2, + "audio_duration_seconds": 5.2, + } + + descriptor = serialize_prepared_reference_videos([item], str(tmp_path)) + + assert deserialize_prepared_reference_videos(descriptor) == (str(tmp_path), [item]) + + +def test_prepared_reference_video_descriptor_rejects_unknown_fields(): + from vllm_omni.model_executor.models.minimax_h3.reference_video import ( + deserialize_prepared_reference_videos, + ) + + with pytest.raises(ValueError, match="invalid MiniMax H3 prepared-reference-video descriptor"): + deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}') + + +def test_stage_one_reuses_existing_prepared_reference_video(tmp_path): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _reuse_prepared_reference_videos, + ) + + prepared_path = tmp_path / "prepared.mp4" + prepared_path.touch() + prepared = [{"prepared_path": str(prepared_path)}] + + assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared + + +def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _reuse_prepared_reference_videos, + ) + + with pytest.raises(ValueError, match="count does not match"): + _reuse_prepared_reference_videos([], expected_count=1) + with pytest.raises(ValueError, match="prepared reference video is unavailable"): + _reuse_prepared_reference_videos( + [{"prepared_path": str(tmp_path / "missing.mp4")}], + expected_count=1, + ) + + +def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3) + decode_calls: list[Any] = [] monkeypatch.setattr( reference_video_module, "_probe_video", - lambda _path: {"frame_count": len(decoded)}, + lambda _path: { + "frame_count": 25, + "width": 2, + "height": 2, + }, ) monkeypatch.setattr( reference_video_module, - "load_video_frames", - lambda path: load_calls.append(path) or decoded, + "_decode_video_frames_ffmpeg", + lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), decoded), ) sampled = reference_video_module.sample_reference_video_frames("prepared.mp4") - assert load_calls == ["prepared.mp4"] - for frame, expected in zip( - sampled["frames"], - (decoded[0], decoded[12], decoded[24]), - strict=True, - ): + assert decode_calls == [ + ( + "prepared.mp4", + { + "frame_count": 25, + "indices": [0, 12, 24], + "width": 2, + "height": 2, + }, + ) + ] + for frame, expected in zip(sampled["frames"], decoded, strict=True): np.testing.assert_array_equal(frame, expected) assert sampled["block_timestamps"] == [0.25, 1.0] +@pytest.mark.parametrize( + ("indices", "frame_count"), + [([0, 12], 25), (None, 2)], +) +def test_ref2va_ffmpeg_decode_reads_exact_rgb_frames(monkeypatch, indices, frame_count): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + expected = np.arange(2 * 2 * 3 * 3, dtype=np.uint8).reshape(2, 2, 3, 3) + observed = [] + + def fake_run(command, **kwargs): + observed.append((command, kwargs)) + return SimpleNamespace(stdout=expected.tobytes()) + + monkeypatch.setattr(reference_video_module.subprocess, "run", fake_run) + + actual = reference_video_module._decode_video_frames_ffmpeg( + "prepared.mp4", + frame_count=frame_count, + indices=indices, + width=3, + height=2, + ) + + np.testing.assert_array_equal(actual, expected) + command, kwargs = observed[0] + assert command[0] == "ffmpeg" + assert command[command.index("-frames:v") + 1] == "2" + if indices is None: + assert "-vf" not in command + else: + assert command[command.index("-vf") + 1] == "select=eq(n\\,0)+eq(n\\,12)" + assert command[-1] == "pipe:1" + assert kwargs == {"check": True, "capture_output": True} + + +def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + expected = np.zeros((2, 2, 3, 3), dtype=np.uint8) + decode_calls: list[Any] = [] + monkeypatch.setitem(sys.modules, "decord", None) + monkeypatch.setattr( + reference_video_module, + "_probe_video", + lambda _path: { + "frame_count": 2, + "width": 3, + "height": 2, + }, + ) + monkeypatch.setattr( + reference_video_module, + "_decode_video_frames_ffmpeg", + lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), expected), + ) + + actual = reference_video_module.load_video_frames("prepared.mp4") + + assert actual is expected + assert decode_calls == [ + ( + "prepared.mp4", + { + "frame_count": 2, + "width": 3, + "height": 2, + }, + ) + ] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"frame_count": 0, "width": 3, "height": 2}, "video has no frames"), + ({"frame_count": 1, "width": 0, "height": 2}, "invalid dimensions"), + ( + {"frame_count": 1, "indices": [], "width": 3, "height": 2}, + "invalid frame indices", + ), + ], +) +def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError, match=message): + reference_video_module._decode_video_frames_ffmpeg("prepared.mp4", **kwargs) + + +def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + calls: list[Any] = [] + probe = { + "streams": [ + { + "codec_type": "video", + "width": 4, + "height": 2, + "r_frame_rate": "24/1", + "nb_frames": "3", + "duration": "0.125", + "sample_aspect_ratio": "1:1", + } + ], + "format": {"duration": "0.125", "size": "123"}, + } + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda command, **kwargs: _append_and_return( + calls, + (command, kwargs), + SimpleNamespace(stdout=json.dumps(probe)), + ), + ) + + metadata = reference_video_module._probe_video("prepared.mp4") + + assert metadata["frame_count"] == 3 + assert len(calls) == 1 + assert "-count_frames" not in calls[0][0] + + +def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + calls: list[Any] = [] + stream = { + "codec_type": "video", + "width": 4, + "height": 2, + "r_frame_rate": "24/1", + "duration": "0.125", + "sample_aspect_ratio": "1:1", + } + first_probe = { + "streams": [stream], + "format": {"duration": "0.125", "size": "123"}, + } + second_probe = { + "streams": [{**stream, "nb_read_frames": "3"}], + "format": {"duration": "0.125", "size": "123"}, + } + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda command, **kwargs: _append_and_return( + calls, + (command, kwargs), + SimpleNamespace(stdout=json.dumps(first_probe if not calls else second_probe)), + ), + ) + + metadata = reference_video_module._probe_video("prepared.mp4") + + assert metadata["frame_count"] == 3 + assert len(calls) == 2 + assert "-count_frames" in calls[1][0] + + +def test_ref2va_ffmpeg_decode_rejects_incomplete_output(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + from vllm_omni.errors import OmniClientError + + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace(stdout=b"short"), + ) + + with pytest.raises(OmniClientError, match="decoded 5 bytes.*expected 18"): + reference_video_module._decode_video_frames_ffmpeg( + "prepared.mp4", + frame_count=1, + width=3, + height=2, + ) + + def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module @@ -1487,12 +2073,16 @@ def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_p "file_size": 1024, }, } - transcode_calls = [] + transcode_calls: list[Any] = [] monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)]) monkeypatch.setattr( reference_video_module, "_transcode_reference_video", - lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4", + lambda source, **kwargs: _append_and_return( + transcode_calls, + (source, kwargs), + f"{source}.prepared.mp4", + ), ) prepared = reference_video_module.prepare_reference_videos( @@ -1671,7 +2261,7 @@ def test_g4_reference_image_file_format_and_size_contract(tmp_path): def test_g4_standalone_audio_duration_and_total_duration_contract(): - from vllm_omni.diffusion.models.minimax_h3.reference_video import ( + from vllm_omni.model_executor.models.minimax_h3.reference_video import ( validate_reference_audio_waveforms, ) @@ -1702,7 +2292,7 @@ def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch): waveform = torch.zeros(1, 10) with pytest.raises(ValueError, match="max_duration_seconds must be positive"): - pipeline._encode_audio_conditions( + pipeline._encode_audio_conditions_resident( [(waveform, 10)], max_duration_seconds=0, ) @@ -1723,7 +2313,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration(): pipeline.audio_vae = FakeAudioVAE() waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40) - encoded, lengths = pipeline._encode_audio_conditions( + encoded, lengths = pipeline._encode_audio_conditions_resident( [(waveform, 10)], max_duration_seconds=2.5, ) @@ -1746,7 +2336,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration(): ], ) def test_g4_reference_video_metadata_validation(field, value, message, tmp_path): - from vllm_omni.diffusion.models.minimax_h3.reference_video import ( + from vllm_omni.model_executor.models.minimax_h3.reference_video import ( _validate_reference_video_metadata, ) @@ -1802,21 +2392,25 @@ def _one_layer_text_encoder(): group = _SingleRankEncoderGroup(0) layer = nn.Module() layer.self_attn = nn.Module() - layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear( - group, - hidden_size=_ENCODER_HIDDEN, - num_heads=_ENCODER_NUM_HEADS, - num_kv_heads=_ENCODER_NUM_KV_HEADS, - head_dim=_ENCODER_HEAD_DIM, - dtype=torch.float32, - ) layer.mlp = nn.Module() - layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear( - group, - input_size=_ENCODER_HIDDEN, - intermediate_size=_ENCODER_INTERMEDIATE, - dtype=torch.float32, - ) + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear( + group, + hidden_size=_ENCODER_HIDDEN, + num_heads=_ENCODER_NUM_HEADS, + num_kv_heads=_ENCODER_NUM_KV_HEADS, + head_dim=_ENCODER_HEAD_DIM, + dtype=torch.float32, + ) + layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear( + group, + input_size=_ENCODER_HIDDEN, + intermediate_size=_ENCODER_INTERMEDIATE, + dtype=torch.float32, + ) layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN) encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) @@ -1849,6 +2443,7 @@ def _write_text_encoder_checkpoint(path, *, omit=()): safetensors.torch.save_file(tensors, str(path / "model.safetensors")) index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")} (path / "model.safetensors.index.json").write_text(json.dumps(index)) + return tensors @pytest.mark.parametrize( @@ -1870,10 +2465,10 @@ def test_encoder_load_reports_missing_fused_source_shards( tmp_path, omitted_sources, expected_detail, unreported_target ): encoder = _one_layer_text_encoder() - _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources) + weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources) with pytest.raises(RuntimeError) as excinfo: - encoder._load_weights(str(tmp_path)) + encoder.load_weights(weights.items()) message = str(excinfo.value) assert expected_detail in message @@ -1882,17 +2477,17 @@ def test_encoder_load_reports_missing_fused_source_shards( def test_encoder_load_rejects_unloaded_plain_param(tmp_path): encoder = _one_layer_text_encoder() - _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",)) + weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",)) with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"): - encoder._load_weights(str(tmp_path)) + encoder.load_weights(weights.items()) def test_encoder_load_places_every_source_in_its_own_rows(tmp_path): encoder = _one_layer_text_encoder() - _write_text_encoder_checkpoint(tmp_path) + weights = _write_text_encoder_checkpoint(tmp_path) - encoder._load_weights(str(tmp_path)) + encoder.load_weights(weights.items()) params = dict(encoder.named_parameters()) q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py new file mode 100644 index 000000000..b0ebf1038 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _OffloadAbort(BaseException): + pass + + +def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker): + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + + class VisionStack(torch.nn.Module): + def __init__(self): + super().__init__() + self.patch_embed = torch.nn.Linear(2, 2) + self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)]) + + class TextStack(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed_tokens = torch.nn.Embedding(4, 2) + self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)]) + + encoder = object.__new__(encoder_module.MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + encoder.device_target = torch.device("cpu") + encoder.vision = VisionStack() + encoder.text_model = TextStack() + hook = mocker.Mock() + hook.pin_memory = False + encoder._omni_layerwise_hooks = [hook] + encoder._omni_layerwise_enabled = True + cache = mocker.Mock() + encoder.set_omni_component_cache(cache) + stager = mocker.Mock() + stager_cls = mocker.Mock(return_value=stager) + monkeypatch.setattr(encoder_module, "PinnedModuleStager", stager_cls) + + encoder.load_to_device() + encoder.offload_to_cpu() + + stager_cls.assert_called_once_with( + [encoder.vision.patch_embed, encoder.text_model.embed_tokens], + torch.device("cpu"), + pin_memory=False, + cache_retention=cache, + ) + stager.load.assert_called_once_with() + stager.offload.assert_called_once_with() + hook.offload_layer.assert_called_once_with() + + +def test_manual_component_failure_forces_retained_cache_release(mocker): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = mocker.Mock() + pipeline.od_config.enable_layerwise_offload = False + pipeline.od_config.enable_distributed_layerwise_offload = True + pipeline._model_cpu_offload_modules = [] + pipeline._dlo_component_cache = mocker.Mock() + component = mocker.Mock() + component.offload_to_cpu.side_effect = _OffloadAbort("offload failed") + + with pytest.raises(RuntimeError, match="component failed"): + with pipeline._component_on_device(component): + raise RuntimeError("component failed") + + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_called_once_with() + pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True) + + +def test_manual_component_offload_failure_forces_retained_cache_release(mocker): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = mocker.Mock() + pipeline.od_config.enable_layerwise_offload = False + pipeline.od_config.enable_distributed_layerwise_offload = True + pipeline._model_cpu_offload_modules = [] + pipeline._dlo_component_cache = mocker.Mock() + component = mocker.Mock() + component.offload_to_cpu.side_effect = [None, _OffloadAbort("offload failed")] + + with pipeline._component_on_device(component): + pass + with pytest.raises(_OffloadAbort, match="offload failed"): + with pipeline._component_on_device(component): + pass + + assert component.load_to_device.call_count == 2 + assert component.offload_to_cpu.call_count == 2 + pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py new file mode 100644 index 000000000..b169d918e --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections.abc import Iterator + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLEncoder + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture(autouse=True) +def restore_cudnn_sdp_state() -> Iterator[None]: + original = torch.backends.cuda.cudnn_sdp_enabled() + try: + yield + finally: + torch.backends.cuda.enable_cudnn_sdp(original) + + +def _encoder_without_weights() -> MiniMaxH3Qwen3VLEncoder: + encoder = MiniMaxH3Qwen3VLEncoder("unused", device=torch.device("cpu"), load_model=False) + encoder.text_model = nn.Linear(1, 1) + return encoder + + +@pytest.mark.parametrize("initial_state", [False, True]) +def test_encode_ids_restores_cudnn_sdp_after_success(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None: + encoder = _encoder_without_weights() + expected = torch.tensor([1.0]) + + def encode(*args, **kwargs): + assert torch.backends.cuda.cudnn_sdp_enabled() + return expected + + monkeypatch.setattr(encoder, "_encode", encode) + torch.backends.cuda.enable_cudnn_sdp(initial_state) + + actual = encoder.encode_ids(torch.tensor([1])) + + torch.testing.assert_close(actual, expected) + assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state + + +@pytest.mark.parametrize("initial_state", [False, True]) +def test_encode_ids_restores_cudnn_sdp_after_failure(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None: + encoder = _encoder_without_weights() + + def fail(*args, **kwargs): + assert torch.backends.cuda.cudnn_sdp_enabled() + raise RuntimeError("encode failed") + + monkeypatch.setattr(encoder, "_encode", fail) + torch.backends.cuda.enable_cudnn_sdp(initial_state) + + with pytest.raises(RuntimeError, match="encode failed"): + encoder.encode_ids(torch.tensor([1])) + + assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py new file mode 100644 index 000000000..8945b9b47 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py @@ -0,0 +1,498 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from vllm_omni.diffusion.models.minimax_h3.fasth3 import ( + FASTH3_BASE_MODEL, + FASTH3_BASE_SCHEDULE, + FASTH3_DENOISE_STEPS, + FASTH3_FORMAT, + FastH3AdapterError, + FastH3WeightFusion, + resolve_fasth3_fusion, +) +from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv +from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas +from vllm_omni.errors import OmniClientError + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_RANK = 2 +_HIDDEN = 4 +_HEAD_DIM = 2 +_HEADS = 3 +_INNER = _HEAD_DIM * _HEADS # attention inner size +_FFN = 5 + + +def _factors(out_dim: int, in_dim: int) -> tuple[torch.Tensor, torch.Tensor]: + """A rank-2 pair whose product is reproducible and non-symmetric.""" + a = torch.arange(_RANK * in_dim, dtype=torch.float32).reshape(_RANK, in_dim) / (in_dim * _RANK) + b = torch.arange(out_dim * _RANK, dtype=torch.float32).reshape(out_dim, _RANK) / (out_dim * _RANK) + return a, b + + +def _write_adapter( + path, + *, + tensors=None, + drop: str | None = None, + blocks: int = 1, + metadata: dict[str, str | None] | None = None, +) -> None: + """Write an artifact in the published ``fastvideo-lora-v2`` shape.""" + payload: dict[str, torch.Tensor] = {} + for block in range(blocks): + for suffix, (out_dim, in_dim) in { + "attn.to_q": (_INNER, _HIDDEN), + "attn.to_k": (_INNER, _HIDDEN), + "attn.to_v": (_INNER, _HIDDEN), + "attn.to_out.0": (_HIDDEN, _INNER), + "ff.net.0.proj": (2 * _FFN, _HIDDEN), + "ff.net.2": (_HIDDEN, _FFN), + }.items(): + a, b = _factors(out_dim, in_dim) + payload[f"transformer_blocks.{block}.{suffix}.lora_A.weight"] = a + payload[f"transformer_blocks.{block}.{suffix}.lora_B.weight"] = b + payload.update(tensors or {}) + if drop is not None: + payload.pop(drop, None) + path.parent.mkdir(parents=True, exist_ok=True) + # The published writer records what it emitted; mirror it so the fixtures + # are as self-describing as the real artifact. + declared = { + "format": FASTH3_FORMAT, + "finetuned_model": "FastVideo/FastVideo-FastH3-Dense-4-step-v1", + "base_model": FASTH3_BASE_MODEL, + "rank": str(_RANK), + "low_rank_tensors": str(sum(1 for key in payload if key.endswith((".lora_A.weight", ".lora_B.weight")))), + "diff_tensors": str(sum(1 for key in payload if key.endswith((".diff", ".diff_b")))), + "set_weight_tensors": str(sum(1 for key in payload if key.endswith(".set_weight"))), + } + # safetensors metadata is string-only, so a ``None`` override drops the key. + declared.update(metadata or {}) + save_file(payload, str(path), metadata={key: value for key, value in declared.items() if value is not None}) + + +def _claim(path, **kwargs) -> FastH3WeightFusion | None: + """Whatever from_path decides, including declining the artifact.""" + kwargs.setdefault("head_dim", _HEAD_DIM) + kwargs.setdefault("num_blocks", 1) + kwargs.setdefault("num_refiner_blocks", 0) + return FastH3WeightFusion.from_path(path, **kwargs) + + +def _load(path, **kwargs) -> FastH3WeightFusion: + """The same, for the tests that only make sense on a claimed adapter.""" + fusion = _claim(path, **kwargs) + assert fusion is not None, f"{path} was not claimed as a FastH3 adapter" + return fusion + + +def test_only_an_artifact_carrying_the_release_identity_is_claimed(tmp_path): + plain = tmp_path / "peft" / "adapter_model.safetensors" + plain.parent.mkdir(parents=True) + save_file({"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(plain)) + # No fastvideo-lora-v2 metadata: this is somebody else's LoRA and has to + # stay on the dynamic route rather than being fused. + assert _claim(plain.parent) is None + assert _claim(tmp_path / "missing") is None + + claimed = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(claimed) + assert _claim(claimed.parent) is not None + + +@pytest.mark.parametrize( + "identity", + [ + # An ordinary H3 adapter out of FastVideo's own extraction tools. + {"finetuned_model": "someone/minimax-h3-style-lora"}, + # Somebody else's adapter that merely names the student it imitates. + {"finetuned_model": "someone/fasth3-style-lora-for-h3"}, + {"finetuned_model": ""}, + {"finetuned_model": None}, + # A FastH3 name over a different base model. + {"base_model": "Wan-AI/Wan2.2-TI2V-5B"}, + ], +) +def test_a_generic_fastvideo_h3_adapter_stays_on_the_dynamic_route(identity, tmp_path): + """The fixture edits every block, so only the identity separates it from FastH3.""" + path = tmp_path / "generic" / "adapter_model.safetensors" + _write_adapter(path, metadata=identity) + + assert _claim(path.parent) is None + + +def test_a_claimed_artifact_must_declare_every_tensor_count(tmp_path): + """Declining is for other people's adapters; a FastH3 file is held to its word.""" + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, metadata={"diff_tensors": None}) + + with pytest.raises(FastH3AdapterError, match="omits diff_tensors"): + _claim(path.parent) + + +def test_the_published_bundle_root_is_refused_rather_than_guessed(tmp_path): + root = tmp_path / "FastVideo-FastH3-4-step-Preview-v1-LoRA" + for slug in ("dense-datafree", "vsa-datafree"): + _write_adapter(root / slug / "adapter_model.safetensors") + (root / "adapter_manifest.json").write_text(json.dumps({"schema_version": "fasth3-lora-bundle-v1"})) + + with pytest.raises(FastH3AdapterError, match="point --lora-path at one variant"): + _load(root) + # One variant inside it is unambiguous. + assert _claim(root / "dense-datafree") is not None + + +def test_a_multi_shard_non_fasth3_adapter_stays_on_the_dynamic_route(tmp_path): + # A plain PEFT LoRA saved as several shards must not hard-fail startup with a + # FastH3-specific message; it has to fall through to the dynamic LoRA route. + directory = tmp_path / "peft-sharded" + directory.mkdir() + for shard in ("adapter_model-00001-of-00002.safetensors", "adapter_model-00002-of-00002.safetensors"): + save_file( + {"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(directory / shard) + ) + + assert _claim(directory) is None + + +def test_low_rank_factors_reach_the_fused_projections(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + fusion = _load(path.parent) + + qkv = torch.zeros((3 * _INNER, _HIDDEN), dtype=torch.float32) + fused = fusion.fuse("blocks.0.attn.qkv_proj.weight", qkv).cpu() + # The checkpoint stores one head group at a time as [q, k, v], so the delta + # has to survive the loader's own unpacking as three separate projections. + q, k, v = torch.split( + _reorder_grouped_qkv_to_qkv(fused, num_query_groups=_HEADS, heads_per_group=1, head_dim=_HEAD_DIM), + [_INNER, _INNER, _INNER], + ) + a, b = _factors(_INNER, _HIDDEN) + expected = b @ a + for got in (q, k, v): + assert torch.allclose(got, expected, atol=1e-5) + + +def test_the_fused_mlp_delta_is_swapped_into_gate_first_order(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + fusion = _load(path.parent) + + fused = fusion.fuse("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))).cpu() + a, b = _factors(2 * _FFN, _HIDDEN) + value, gate = (b @ a).chunk(2, dim=0) + # The diffusers export is value-first; H3's fc1 is gate-first. + assert torch.allclose(fused, torch.cat((gate, value), dim=0), atol=1e-5) + + +def test_diff_and_diff_b_edit_weights_and_biases(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter( + path, + tensors={ + "transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.25), + "transformer_blocks.0.adaln_proj.linear.diff_b": torch.full((_HIDDEN,), -0.5), + "context_embedder.diff": torch.full((_HIDDEN, _HIDDEN), 0.125), + "context_embedder.diff_b": torch.full((_HIDDEN,), 2.0), + }, + ) + fusion = _load(path.parent) + + # A full-rank delta on an RMSNorm vector is exactly what a LoRA layer + # cannot express, and why this release is fused instead of switched. + assert torch.allclose(fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 1.25)) + assert torch.allclose( + fusion.fuse("blocks.0.adaln_proj.linear.bias", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 0.5) + ) + assert torch.allclose( + fusion.fuse("condition_proj.weight", torch.zeros((_HIDDEN, _HIDDEN))).cpu(), + torch.full((_HIDDEN, _HIDDEN), 0.125), + ) + assert torch.allclose(fusion.fuse("condition_proj.bias", torch.zeros(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 2.0)) + + +def test_a_streamed_checkpoint_keeps_the_parameters_the_adapter_never_edits(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 3.0)}) + fusion = _load(path.parent) + + untouched = torch.arange(6, dtype=torch.bfloat16).reshape(2, 3) + streamed = dict( + fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN)), ("blocks.0.attn.q_norm.weight", untouched)]) + ) + assert torch.allclose(streamed["blocks.0.norm1.weight"].cpu(), torch.full((_HIDDEN,), 3.0)) + assert streamed["blocks.0.attn.q_norm.weight"] is untouched + + # validate_fully_applied() releases the deltas, so a second stream would + # fuse nothing at all; it has to fail rather than serve base H3 weights. + with pytest.raises(FastH3AdapterError, match="already been fused"): + next(fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN))])) + + +def test_the_fused_result_keeps_the_parameter_dtype(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.5)}) + fusion = _load(path.parent) + + fused = fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN, dtype=torch.bfloat16)) + assert fused.dtype == torch.bfloat16 + + +def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path): + dense = tmp_path / "dense" / "adapter_model.safetensors" + _write_adapter(dense) + assert _load(dense.parent).requires_vsa is False + + sparse = tmp_path / "vsa" / "adapter_model.safetensors" + _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))}) + assert _load(sparse.parent).requires_vsa is True + + +def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}) + # Dropping it silently would load a model that is not the distilled student. + with pytest.raises(FastH3AdapterError, match="name no known"): + _load(path.parent) + + +def test_an_unpaired_low_rank_factor_is_an_error(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, drop="transformer_blocks.0.attn.to_q.lora_B.weight") + + with pytest.raises(FastH3AdapterError, match="unpaired factor"): + _load(path.parent) + + +def test_a_delta_the_checkpoint_never_offered_is_reported(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + fusion = _load(path.parent) + + fusion.fuse("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN))) + # An unapplied delta is the failure that matters: the model would load and + # generate, just not as the distilled student. + with pytest.raises(FastH3AdapterError, match="never provided"): + fusion.validate_fully_applied() + + for name, shape in ( + ("blocks.0.attn.out_proj.weight", (_HIDDEN, _INNER)), + ("blocks.0.mlp.fc1.weight", (2 * _FFN, _HIDDEN)), + ("blocks.0.mlp.fc2.weight", (_HIDDEN, _FFN)), + ): + fusion.fuse(name, torch.zeros(shape)) + fusion.validate_fully_applied() + + +def _pipeline_stub(fusion=None): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "combined" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + pipeline._turbo_lora_adapter_ids = set() + pipeline._fasth3 = fusion + pipeline.od_config = SimpleNamespace() + pipeline.default_video_shift, pipeline.default_audio_shift = 12.0, 3.0 + return pipeline + + +def _sampling(steps=FASTH3_DENOISE_STEPS, **kwargs) -> SimpleNamespace: + fields = {"num_inference_steps": steps, "extra_args": {}, "lora_request": None} + return SimpleNamespace(**{**fields, **kwargs}) + + +def _check_request(fusion, sampling) -> None: + fusion.check_request(sampling, video_shift=12.0, audio_shift=3.0) + + +def _fusion(tmp_path, **kwargs) -> FastH3WeightFusion: + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, **kwargs) + return _load(path.parent) + + +@pytest.mark.parametrize("task", ["fl2va", "ref2va"]) +def test_an_active_fasth3_fusion_restricts_requests_to_t2va(task, tmp_path): + pipeline = _pipeline_stub(_fusion(tmp_path)) + + with pytest.raises(OmniClientError, match="distills \\['t2va'\\] only"): + pipeline._resolve_task(task, {}) + assert pipeline._resolve_task("t2va", {}) == "t2va" + assert _pipeline_stub()._resolve_task(task, {}) == task + + +@pytest.mark.parametrize("num_inference_steps", [None, 5, 50]) +def test_fasth3_requires_one_forward_per_sigma_interval(num_inference_steps, tmp_path): + fusion = _fusion(tmp_path) + + # The five sigma points bound four forwards, and a request states forwards: + # the interval contract the pinned schedules use, and the count the step + # scheduler admits a request on before any pipeline hook runs. + with pytest.raises(OmniClientError, match=f"num_inference_steps={FASTH3_DENOISE_STEPS}"): + _check_request(fusion, _sampling(num_inference_steps)) + _check_request(fusion, _sampling()) + # The pinned branch coerces the string form; this one cannot disagree. + _check_request(fusion, _sampling(str(FASTH3_DENOISE_STEPS))) + + +def test_the_request_denoises_on_the_release_ladder(tmp_path): + pipeline = _pipeline_stub(_fusion(tmp_path)) + + positions, num_steps = pipeline._resolve_sigma_positions("t2va", _sampling()) + + # dmd_denoising_steps [999, 749, 500, 250] out of 1000, closed with the + # terminal 0.0: the pre-shift positions the student was distilled at, not + # the uniform ladder num_inference_steps would otherwise derive. + assert positions == (0.999, 0.749, 0.5, 0.25, 0.0) + # Five sigma points bound four transformer forwards, and forwards is the + # unit num_steps carries downstream (Cache-DiT, quality hints). + assert num_steps == FASTH3_DENOISE_STEPS == len(positions) - 1 + # H3's own per-modality shifts still apply on top, and are what turn those + # positions into the noise levels the four forwards run at. + assert minimax_h3_time_shift_sigmas(shift_scale=12.0, base_schedule=positions) == pytest.approx( + [0.999917, 0.972833, 0.923077, 0.8, 0.0], abs=1e-6 + ) + assert minimax_h3_time_shift_sigmas(shift_scale=3.0, base_schedule=positions) == pytest.approx( + [0.999666, 0.899520, 0.75, 0.5, 0.0], abs=1e-6 + ) + # Without a fused adapter the same request keeps the undistilled ladder. + assert _pipeline_stub()._resolve_sigma_positions("t2va", _sampling()) == (None, FASTH3_DENOISE_STEPS) + + +def test_the_contract_leaves_the_checkpoint_shifts_alone(tmp_path): + pipeline = _pipeline_stub(_fusion(tmp_path)) + pipeline.partition = "fl2va" + + pipeline._fasth3.check_serving_contract( + partition=pipeline.partition, + od_config=pipeline.od_config, + video_shift=pipeline.default_video_shift, + audio_shift=pipeline.default_audio_shift, + ) + + # The ladder is stated pre-shift, so H3's 12/3 shifts are what place the + # student on the levels it was distilled at; the contract must not touch them. + assert pipeline.default_video_shift == 12.0 + assert pipeline.default_audio_shift == 3.0 + + +@pytest.mark.parametrize( + "extra", + [{"flow_shift": 1.0}, {"audio_flow_shift": 1.0}, {"flow_shift": "bad"}], +) +def test_a_request_may_not_move_the_student_off_its_rungs(extra, tmp_path): + fusion = _fusion(tmp_path) + + with pytest.raises(OmniClientError, match="FastH3 requires"): + _check_request(fusion, _sampling(extra_args=extra)) + + _check_request(fusion, _sampling(extra_args={"flow_shift": 12.0, "audio_flow_shift": 3.0})) + + +def test_a_request_may_not_carry_a_lora_on_a_fused_server(tmp_path): + fusion = _fusion(tmp_path) + + # The dynamic LoRA manager is skipped for a fused adapter, so an adapter + # asked for per request would be neither applied nor reported. + with pytest.raises(OmniClientError, match="per-request lora is unavailable"): + _check_request(fusion, _sampling(lora_request=SimpleNamespace(lora_int_id=1))) + + +def test_offload_is_refused_because_it_bypasses_the_fusion(tmp_path): + fusion = _fusion(tmp_path) + for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload"): + # A host-weight plan installs the transformer without load_weights(), + # so the fusion and its completeness check would both be skipped. + with pytest.raises(ValueError, match="cannot be combined with"): + fusion.check_serving_contract( + partition="fl2va", + od_config=SimpleNamespace(**{flag: True}), + video_shift=12.0, + audio_shift=3.0, + ) + + +def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path): + sparse = tmp_path / "vsa" / "adapter_model.safetensors" + _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))}) + contract = {"od_config": SimpleNamespace(), "video_shift": 12.0, "audio_shift": 3.0} + with pytest.raises(ValueError, match="Video Sparse Attention variant"): + _load(sparse.parent).check_serving_contract(partition="fl2va", **contract) + + with pytest.raises(ValueError, match="cannot serve a Ref2VA partition"): + _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract) + + +def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path): + # Only low-rank factors are placed into H3's fused QKV and gate/up layouts, + # so a .diff aimed at one would otherwise be added transposed. + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, tensors={"transformer_blocks.0.attn.to_q.diff": torch.ones((_INNER, _HIDDEN))}) + with pytest.raises(FastH3AdapterError, match="fused layout"): + _claim(path.parent) + + +def test_an_adapter_that_leaves_blocks_untouched_is_refused(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, blocks=2) + + # Every tensor in this file is well formed and pairs up, so the fusion's + # own completeness check passes; only the model's depth reveals that a + # third of it would have been served as base H3 weights. + with pytest.raises(FastH3AdapterError, match="missing=\\[2\\]"): + _claim(path.parent, num_blocks=3) + with pytest.raises(FastH3AdapterError, match="unknown=\\[1\\]"): + _claim(path.parent, num_blocks=1) + assert _claim(path.parent, num_blocks=2) is not None + + +def test_an_adapter_shorter_than_it_declares_is_refused(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path, metadata={"low_rank_tensors": "724", "diff_tensors": "85"}) + + # The writer records what it emitted, which is the file's one statement + # about its own completeness. + with pytest.raises(FastH3AdapterError, match="declares low_rank_tensors=724 but carries 12"): + _claim(path.parent) + + +def test_only_a_configured_adapter_path_is_read(tmp_path): + path = tmp_path / "fasth3" / "adapter_model.safetensors" + _write_adapter(path) + transformer = SimpleNamespace( + arch=SimpleNamespace(attention_head_dim=_HEAD_DIM, num_layers=1, token_refiner_num_layers=0) + ) + unbuilt = object() # a transformer whose arch must not be touched + + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=str(path)), transformer) is not None + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path)]), transformer) is not None + # Nothing configured, or more than one artifact, stays on the dynamic route + # without reading anything off the model. + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=None), unbuilt) is None + assert resolve_fasth3_fusion(SimpleNamespace(), unbuilt) is None + assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path), str(path)]), unbuilt) is None + + +def test_a_pipeline_that_fused_its_adapter_needs_no_lora_manager(tmp_path): + assert _pipeline_stub(_fusion(tmp_path)).lora_is_fused is True + assert _pipeline_stub().lora_is_fused is False + + +def test_the_ladder_is_the_one_the_release_publishes(): + # FASTH3_DENOISE_STEPS is derived from it, so the two cannot drift apart. + assert FASTH3_BASE_SCHEDULE.base_schedule == (0.999, 0.749, 0.5, 0.25, 0.0) + assert FASTH3_BASE_SCHEDULE.num_inference_steps == FASTH3_DENOISE_STEPS == 4 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py new file mode 100644 index 000000000..5f98a88d5 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py @@ -0,0 +1,416 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file +from vllm.lora.lora_weights import PackedLoRALayerWeights + +from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager +from vllm_omni.diffusion.models.minimax_h3 import lora as lora_module +from vllm_omni.diffusion.models.minimax_h3.lora import ( + load_minimax_h3_turbo_lora, +) +from vllm_omni.errors import OmniClientError +from vllm_omni.lora.request import LoRARequest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_TINY_TARGET_DIMS = { + "attn.to_q": (2, 3), + "attn.to_k": (2, 3), + "attn.to_v": (2, 3), + "attn.to_out.0": (3, 2), + "ff.net.0.proj": (2, 4), + "ff.net.2": (2, 2), +} + + +@pytest.fixture(autouse=True) +def _use_tiny_h3_dimensions(monkeypatch): + monkeypatch.setattr(lora_module, "_TURBO_TARGET_DIMS", _TINY_TARGET_DIMS) + + +def _request(path) -> LoRARequest: + return LoRARequest( + lora_name="turbo", + lora_int_id=1, + lora_path=str(path), + ) + + +def _write_tiny_turbo( + path, + *, + alpha: str = "128", + key_format: str = "minimax-h3-diffusers", + omit_target: str | None = None, + shape_overrides: dict[str, tuple[int, int]] | None = None, +) -> None: + rank = 128 + tensors = {} + overrides = shape_overrides or {} + target_suffixes = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "ff.net.0.proj", + "ff.net.2", + ) + for prefix, block_count in ( + ("transformer_blocks", 50), + ("token_refiner.refiner_blocks", 2), + ): + for block_index in range(block_count): + for suffix in target_suffixes: + target = f"{prefix}.{block_index}.{suffix}" + if target == omit_target: + continue + input_dim, output_dim = _TINY_TARGET_DIMS[suffix] + a_name = f"{target}.lora_A.default.weight" + b_name = f"{target}.lora_B.default.weight" + tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim))) + if suffix == "ff.net.0.proj": + if b_name in overrides: + tensors[b_name] = torch.ones(overrides[b_name]) + else: + tensors[b_name] = torch.cat( + ( + torch.ones(output_dim // 2, rank), + torch.full((output_dim // 2, rank), 2.0), + ), + dim=0, + ) + else: + tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank))) + save_file( + tensors, + str(path), + metadata={"alpha": alpha, "key_format": key_format}, + ) + + +def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path) + + loaded = load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + assert loaded is not None + lora_model, peft_helper = loaded + assert peft_helper.r == 128 + assert peft_helper.lora_alpha == 128 + assert len(lora_model.loras) == 312 + assert "blocks.0.attn.to_q" in lora_model.loras + assert "blocks.0.mlp.fc1" in lora_model.loras + fc1 = lora_model.get_lora("blocks.0.mlp.fc1") + assert isinstance(fc1, PackedLoRALayerWeights) + assert len(fc1.lora_a) == 2 + assert len(fc1.lora_b) == 2 + torch.testing.assert_close(fc1.lora_a[0], torch.ones(128, 2)) + torch.testing.assert_close(fc1.lora_a[1], torch.ones(128, 2)) + torch.testing.assert_close(fc1.lora_b[0], torch.full((2, 128), 2.0)) + torch.testing.assert_close(fc1.lora_b[1], torch.ones(2, 128)) + + +@pytest.mark.parametrize( + ("tensor_suffix", "bad_shape"), + [ + ("attn.to_q.lora_A.default.weight", (128, 1)), + ("attn.to_q.lora_A.default.weight", (128, 3)), + ("attn.to_q.lora_B.default.weight", (2, 128)), + ("attn.to_q.lora_B.default.weight", (4, 128)), + ], +) +def test_h3_turbo_rejects_under_and_oversized_global_shapes(tmp_path, tensor_suffix, bad_shape): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + tensor_name = f"transformer_blocks.0.{tensor_suffix}" + _write_tiny_turbo(path, shape_overrides={tensor_name: bad_shape}) + + with pytest.raises(ValueError, match="invalid global shape"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(tmp_path): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path) + + class _Pipeline: + def _load_diffusion_lora_adapter(self, **kwargs): + return load_minimax_h3_turbo_lora(partition="fl2va", **kwargs) + + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = _Pipeline() + manager.dtype = torch.float32 + manager._expected_lora_modules = {"to_q", "fc1"} + + lora_model, peft_helper = manager._load_adapter(_request(path)) + + assert lora_model.id == 1 + assert peft_helper.lora_alpha == 128 + assert len(lora_model.loras) == 312 + assert "blocks.0.attn.to_q" in lora_model.loras + assert "blocks.0.mlp.fc1" in lora_model.loras + + +def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path): + wrong_alpha = tmp_path / "wrong_alpha" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + wrong_alpha.parent.mkdir() + _write_tiny_turbo(wrong_alpha, alpha="8") + with pytest.raises(ValueError, match="requires alpha=128"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(wrong_alpha), + lora_path=wrong_alpha, + dtype=torch.float32, + ) + + valid = tmp_path / "valid" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + valid.parent.mkdir() + _write_tiny_turbo(valid) + with pytest.raises(ValueError, match="supports FL2VA/T2VA only"): + load_minimax_h3_turbo_lora( + partition="ref2va", + lora_request=_request(valid), + lora_path=valid, + dtype=torch.float32, + ) + + +@pytest.mark.parametrize( + "offload_mode", + [ + "model-level CPU offload (--enable-cpu-offload)", + "layerwise offload (--enable-layerwise-offload)", + ], +) +def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path) + + with pytest.raises(ValueError, match="does not support"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + unsupported_offload_mode=offload_mode, + ) + + +def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import ( + pipeline_minimax_h3 as pipeline_module, + ) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + captured = {} + + def load_turbo(**kwargs): + captured.update(kwargs) + return object(), object() + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo) + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=_request("turbo"), + lora_path="turbo", + dtype=torch.bfloat16, + ) + + assert loaded is not None + assert captured["unsupported_offload_mode"] is None + + +def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path): + unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors" + _write_tiny_turbo(unsupported) + + with pytest.raises(ValueError, match="supports only"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(unsupported), + lora_path=unsupported, + dtype=torch.float32, + ) + + supported = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(supported) + # Directory resolution selects the declared artifact even when another + # v1.0 checkpoint is present beside it. + assert ( + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(tmp_path), + lora_path=tmp_path, + dtype=torch.float32, + ) + is not None + ) + + +def test_h3_turbo_rejects_a_truncated_declared_artifact(tmp_path): + path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(path, omit_target="transformer_blocks.49.ff.net.2") + + with pytest.raises(ValueError, match="target set does not match"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_non_h3_checkpoint_falls_back_to_the_generic_peft_loader(tmp_path): + path = tmp_path / "other.safetensors" + _write_tiny_turbo(path, key_format="other") + + assert ( + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + is None + ) + + declared = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" + _write_tiny_turbo(declared, key_format="other") + with pytest.raises(ValueError, match="requires safetensors metadata"): + load_minimax_h3_turbo_lora( + partition="fl2va", + lora_request=_request(declared), + lora_path=declared, + dtype=torch.float32, + ) + + +def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import ( + pipeline_minimax_h3 as pipeline_module, + ) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "combined" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + request = _request("generic-peft") + + def load(): + return pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path=request.lora_path, + dtype=torch.float32, + ) + + def resolve(scale): + sampling = SimpleNamespace(lora_request=request, lora_scale=scale) + return pipeline._resolve_task( + "ref2va", + {}, + has_turbo_lora=pipeline._has_active_turbo_lora(sampling), + ) + + recognized = (object(), object()) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: recognized) + assert load() is recognized + assert resolve(0.0) == "ref2va" + with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"): + resolve(1.0) + + # Simulate manager eviction followed by a generic PEFT adapter reusing the + # same client-supplied ID. A real reload must replace the stale kind. + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + assert load() is None + assert resolve(1.0) == "ref2va" + + +def test_h3_turbo_requires_all_loaded_targets_to_bind(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline._turbo_lora_adapter_ids = {1} + lora_model = SimpleNamespace(id=1, loras={"to_q": object(), "to_k": object()}) + + pipeline._validate_diffusion_lora_binding( + lora_model=lora_model, + bound_lora_names=frozenset(lora_model.loras), + ) + with pytest.raises(ValueError, match="bound=1/2"): + pipeline._validate_diffusion_lora_binding( + lora_model=lora_model, + bound_lora_names=frozenset({"to_q"}), + ) + + +@pytest.mark.parametrize( + ("num_inference_steps", "extra_args", "error"), + [ + (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), + (4, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), + ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"), + (5, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"), + (5, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"), + (5, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"), + (5, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"), + ], +) +def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, error): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + sampling = SimpleNamespace( + num_inference_steps=num_inference_steps, + extra_args=extra_args, + ) + + with pytest.raises(OmniClientError, match=error): + pipeline._validate_turbo_sampling(sampling) + + +def test_h3_turbo_accepts_five_sigma_points_for_four_nfe(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline._validate_turbo_sampling( + SimpleNamespace( + num_inference_steps=5, + extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0}, + ) + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py new file mode 100644 index 000000000..c4e581381 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py @@ -0,0 +1,864 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file +from vllm.lora.lora_weights import PackedLoRALayerWeights + +from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager +from vllm_omni.diffusion.models.minimax_h3.npu import lora as lora_module +from vllm_omni.diffusion.models.minimax_h3.npu.lora import load_minimax_h3_native_lora +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.lora.request import LoRARequest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_TINY_HIDDEN = 2 +_TINY_FFN = 2 +_TINY_TIME = 2 +_TINY_BLOCK_ADALN = 4 +_TINY_FINAL_ADALN = 3 +_TINY_RANK = 4 +_TINY_NUM_QUERY_GROUPS = 2 +_TINY_HEADS_PER_GROUP = 1 +_TINY_HEAD_DIM = 1 +_TINY_QKV_SLICE = _TINY_NUM_QUERY_GROUPS * _TINY_HEADS_PER_GROUP * _TINY_HEAD_DIM +# Grouped qkv stores (heads_per_group + 2) * head_dim rows per query group, so +# the attention inner size follows from the head layout and cannot be picked +# independently: 3 * inner must equal the grouped row count. +_TINY_INNER = _TINY_QKV_SLICE + + +@pytest.fixture(autouse=True) +def _use_tiny_native_dimensions(monkeypatch): + monkeypatch.setattr(lora_module, "_NATIVE_RANK", _TINY_RANK) + monkeypatch.setattr(lora_module, "_NATIVE_ALPHA", _TINY_RANK) + monkeypatch.setattr(lora_module, "_NATIVE_HIDDEN_SIZE", _TINY_HIDDEN) + monkeypatch.setattr(lora_module, "_NATIVE_ATTENTION_INNER_SIZE", _TINY_INNER) + monkeypatch.setattr(lora_module, "_NATIVE_FFN_HIDDEN_SIZE", _TINY_FFN) + monkeypatch.setattr(lora_module, "_NATIVE_TIME_EMBED_DIM", _TINY_TIME) + monkeypatch.setattr(lora_module, "_NATIVE_BLOCK_ADALN_OUT", _TINY_BLOCK_ADALN) + monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_OUT", _TINY_FINAL_ADALN) + monkeypatch.setattr(lora_module, "_NATIVE_NUM_QUERY_GROUPS", _TINY_NUM_QUERY_GROUPS) + monkeypatch.setattr(lora_module, "_NATIVE_HEADS_PER_GROUP", _TINY_HEADS_PER_GROUP) + monkeypatch.setattr(lora_module, "_NATIVE_HEAD_DIM", _TINY_HEAD_DIM) + monkeypatch.setattr(lora_module, "_NATIVE_QKV_SLICE", _TINY_QKV_SLICE) + monkeypatch.setattr( + lora_module, + "_NATIVE_TARGET_DIMS", + { + "attn.qkv_proj": (_TINY_HIDDEN, 3 * _TINY_INNER), + "attn.out_proj": (_TINY_INNER, _TINY_HIDDEN), + "mlp.fc1": (_TINY_HIDDEN, 2 * _TINY_FFN), + "mlp.fc2": (_TINY_FFN, _TINY_HIDDEN), + "adaln_proj.linear": (_TINY_TIME, _TINY_BLOCK_ADALN), + }, + ) + monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_DIMS", (_TINY_TIME, _TINY_FINAL_ADALN)) + monkeypatch.setattr( + lora_module, + "_NATIVE_EXPECTED_TARGETS", + frozenset( + [ + *( + f"blocks.{block_index}.{suffix}" + for block_index in range(50) + for suffix in lora_module._NATIVE_TARGET_SUFFIXES + ) + ] + + [ + *( + f"token_refiner.blocks.{block_index}.{suffix}" + for block_index in range(2) + for suffix in lora_module._NATIVE_TOKEN_REFINER_SUFFIXES + ) + ] + + ["final_layer.adaln_proj.linear"] + ), + ) + + +def _request(path) -> LoRARequest: + return LoRARequest( + lora_name="flashgen", + lora_int_id=7, + lora_path=str(path), + ) + + +def _write_tiny_native( + path, + *, + omit_target: str | None = None, + shape_overrides: dict[str, tuple[int, int]] | None = None, + metadata: dict[str, str] | None = None, +) -> None: + rank = _TINY_RANK + tensors = {} + overrides = shape_overrides or {} + for target in sorted(lora_module._NATIVE_EXPECTED_TARGETS): + if target == omit_target: + continue + input_dim, output_dim = lora_module._native_target_dims(target) + raw_target = f"transformer.{target}" + a_name = f"{raw_target}.lora_A.default.weight" + b_name = f"{raw_target}.lora_B.default.weight" + tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim))) + if target.endswith(".attn.qkv_proj"): + grouped_rows = [] + for group in range(_TINY_NUM_QUERY_GROUPS): + grouped_rows.extend( + [ + torch.full((1, rank), float(group * 10 + 1)), + torch.full((1, rank), float(group * 10 + 2)), + torch.full((1, rank), float(group * 10 + 3)), + ] + ) + tensors[b_name] = torch.cat(grouped_rows, dim=0) + elif target.endswith(".mlp.fc1"): + tensors[b_name] = torch.cat( + ( + torch.full((output_dim // 2, rank), 2.0), + torch.full((output_dim // 2, rank), 1.0), + ), + dim=0, + ) + else: + tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank))) + save_file( + tensors, + str(path), + metadata=metadata + or { + "format": "pt", + "key_format": "minimax-h3-native", + "qkv_layout": "grouped", + "lora_rank": str(rank), + "lora_alpha": str(rank), + "base_schedule": "1.0,0.7,0.4,0.15,0.0", + "tasks": "t2va", + }, + ) + + +def test_h3_native_loads_and_packs_qkv_and_fc1(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + assert loaded is not None + lora_model, peft_helper, schedule = loaded + assert peft_helper.r == _TINY_RANK + assert schedule.num_inference_steps == 4 + assert len(lora_model.loras) == 259 + assert "blocks.0.adaln_proj.linear" in lora_model.loras + assert "final_layer.adaln_proj.linear" in lora_model.loras + + # The writer tags group g with q=g*10+1, k=g*10+2, v=g*10+3, so a correct + # grouped -> q/k/v permutation yields slices that are uniform per role. + qkv = lora_model.get_lora("blocks.0.attn.qkv_proj") + assert isinstance(qkv, PackedLoRALayerWeights) + torch.testing.assert_close(qkv.lora_b[0], torch.tensor([[1.0] * _TINY_RANK, [11.0] * _TINY_RANK])) + torch.testing.assert_close(qkv.lora_b[1], torch.tensor([[2.0] * _TINY_RANK, [12.0] * _TINY_RANK])) + torch.testing.assert_close(qkv.lora_b[2], torch.tensor([[3.0] * _TINY_RANK, [13.0] * _TINY_RANK])) + + fc1 = lora_model.get_lora("blocks.0.mlp.fc1") + assert isinstance(fc1, PackedLoRALayerWeights) + torch.testing.assert_close(fc1.lora_b[0], torch.full((_TINY_FFN, _TINY_RANK), 2.0)) + torch.testing.assert_close(fc1.lora_b[1], torch.full((_TINY_FFN, _TINY_RANK), 1.0)) + + +def test_h3_native_rejects_bad_metadata_and_ref2va(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, metadata={"key_format": "minimax-h3-native", "qkv_layout": "runtime"}) + with pytest.raises(ValueError, match="qkv_layout='grouped'"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + # The exact-filename guard runs before the partition check, so the artifact + # has to keep the published name for this to reach the Ref2VA rejection. + valid_dir = tmp_path / "valid" + valid_dir.mkdir() + valid = valid_dir / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(valid) + with pytest.raises(ValueError, match="supports T2VA only"): + load_minimax_h3_native_lora( + partition="ref2va", + lora_request=_request(valid), + lora_path=valid, + dtype=torch.float32, + ) + + renamed = tmp_path / "renamed.safetensors" + _write_tiny_native(renamed) + with pytest.raises(ValueError, match="supports only"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(renamed), + lora_path=renamed, + dtype=torch.float32, + ) + + +def test_h3_native_rejects_truncated_artifact(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, omit_target="blocks.49.mlp.fc2") + with pytest.raises(ValueError, match="target set does not match"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path): + """A mislabeled schedule must fail closed instead of silently changing steps.""" + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native( + path, + metadata={ + "format": "pt", + "key_format": "minimax-h3-native", + "qkv_layout": "grouped", + "lora_rank": str(_TINY_RANK), + "lora_alpha": str(_TINY_RANK), + "base_schedule": "1.0,0.5,0.0", + "tasks": "t2va", + }, + ) + with pytest.raises(ValueError, match="4-interval base_schedule"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +@pytest.mark.parametrize( + ("raw_schedule", "message"), + [ + (None, "requires safetensors metadata base_schedule"), + ("", "must not be empty"), + (",,", "must not be empty"), + ("1.0,oops,0.15,0.0", "is malformed"), + # A capped opening rung below 1.0 is a legal DMD2 schedule, so the + # out-of-range cases are a first position above 1.0 and a tail that + # never reaches clean latents. + ("1.5,0.7,0.4,0.15,0.0", "is malformed"), + ("1.0,0.7,0.4,0.15", "is malformed"), + ], +) +def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message): + """safetensors metadata is string-only, so the loader owns the parse contract.""" + metadata = { + "format": "pt", + "key_format": "minimax-h3-native", + "qkv_layout": "grouped", + "lora_rank": str(_TINY_RANK), + "lora_alpha": str(_TINY_RANK), + "tasks": "t2va", + } + if raw_schedule is not None: + metadata["base_schedule"] = raw_schedule + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, metadata=metadata) + + with pytest.raises(ValueError, match=message): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_h3_native_declared_file_fails_closed_on_invalid_metadata(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path, metadata={"key_format": "other"}) + with pytest.raises(ValueError, match="requires safetensors metadata"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + +def test_pipeline_native_schedule_and_task_validation(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"}) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])} + pipeline._base_schedule_by_partition = {"fl2va": None} + + sampling = SimpleNamespace( + lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"), + lora_scale=1.0, + num_inference_steps=4, + extra_args={}, + ) + assert pipeline._sigma_schedule_for_request(sampling, "t2va").num_inference_steps == 4 + with pytest.raises(OmniClientError, match="num_inference_steps must be 4"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5), + task="t2va", + ) + with pytest.raises(OmniClientError, match="supports T2VA requests only"): + pipeline._resolve_task( + "fl2va", + {}, + has_turbo_lora=False, + has_native_lora=True, + ) + + pipeline._base_schedule_by_partition = {"fl2va": DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])} + with pytest.raises(OmniClientError, match="already pins base_schedule"): + pipeline._sigma_schedule_for_request(sampling, "t2va") + + +def test_pipeline_native_step_count_follows_adapter_schedule(): + """The accepted step count is the adapter's interval count, never a literal.""" + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])} + request = LoRARequest("flashgen", 7, "/tmp/native.safetensors") + + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=request, num_inference_steps=2), + task="t2va", + ) + with pytest.raises(OmniClientError, match="num_inference_steps must be 2 or omitted, not 3"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=request, num_inference_steps=3), + task="t2va", + ) + with pytest.raises(OmniClientError, match="requires num_inference_steps=2"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=request, num_inference_steps=4), + task="t2va", + ) + + +@pytest.mark.parametrize("step_execution", [False, True]) +def test_pipeline_native_omitted_step_count_follows_execution_mode(step_execution): + """Only request mode can resolve an omitted count from the adapter schedule. + + ``StepScheduler`` reads ``num_inference_steps`` off the request at admission, + before this pipeline sees it, so under step execution an omitted value can + only disagree with the denoise loop. The advertised contract has to match. + """ + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])} + pipeline.od_config = SimpleNamespace(step_execution=step_execution) + sampling = SimpleNamespace( + lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"), + num_inference_steps=None, + ) + + if not step_execution: + # Request mode defaults to the adapter's own interval count. + pipeline._validate_native_sampling(sampling, task="t2va") + return + + with pytest.raises(OmniClientError, match="explicit num_inference_steps=4 under step execution"): + pipeline._validate_native_sampling(sampling, task="t2va") + # The step-mode message must not offer omission it cannot accept. + with pytest.raises(OmniClientError, match=r"must be 4, not 5$"): + pipeline._validate_native_sampling( + SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5), + task="t2va", + ) + + +def test_pipeline_replaces_native_classification_after_reload(monkeypatch, tmp_path): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + request = _request(path) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path=path, + dtype=torch.float32, + ) + assert loaded is not None + assert request.lora_int_id in pipeline._native_lora_adapter_ids + assert request.lora_int_id in pipeline._lora_sigma_schedules + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None) + assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None + assert request.lora_int_id not in pipeline._native_lora_adapter_ids + assert request.lora_int_id not in pipeline._lora_sigma_schedules + + +@pytest.mark.parametrize( + "offload_mode", + [ + "model-level CPU offload (--enable-cpu-offload)", + "layerwise offload (--enable-layerwise-offload)", + ], +) +def test_h3_native_rejects_offload_modes(tmp_path, offload_mode): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + with pytest.raises(ValueError, match="does not support"): + load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + unsupported_offload_mode=offload_mode, + ) + + +def test_h3_native_allows_distributed_layerwise_offload(monkeypatch): + """DLO keeps LoRA A/B buffers resident, so native must not fail closed.""" + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + captured: dict[str, object] = {} + schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0]) + + def load_native(**kwargs): + captured.update(kwargs) + return object(), object(), schedule + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", load_native) + + request = _request("flashgen") + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path="flashgen", + dtype=torch.bfloat16, + ) + + assert loaded is not None + assert captured["unsupported_offload_mode"] is None + assert request.lora_int_id in pipeline._native_lora_adapter_ids + assert pipeline._lora_sigma_schedules[request.lora_int_id] is schedule + + +def test_h3_native_qkv_reorder_matches_base_loader_contract(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv + + num_query_groups = 56 + heads_per_group = 1 + head_dim = 128 + grouped = torch.arange(num_query_groups * (heads_per_group + 2) * head_dim, dtype=torch.float32) + grouped = grouped.reshape(num_query_groups, (heads_per_group + 2) * head_dim) + reordered = _reorder_grouped_qkv_to_qkv( + grouped.reshape(-1, 1), + num_query_groups=num_query_groups, + heads_per_group=heads_per_group, + head_dim=head_dim, + ).reshape(-1) + q_size = num_query_groups * heads_per_group * head_dim + k_size = num_query_groups * head_dim + q, k, v = torch.split(reordered, [q_size, k_size, k_size]) + assert q.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 0] + assert k.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, head_dim] + assert v.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 2 * head_dim] + + +def test_native_target_pattern_matches_component_qualified_names(): + """The manager matches ".", so the pattern needs the prefix.""" + import regex as re + + pattern = lora_module._NATIVE_TARGET_PATTERN + matched = [ + target + for target in lora_module._NATIVE_EXPECTED_TARGETS + if re.search(pattern, f"transformer.{target}") is not None + ] + assert len(matched) == 259 + assert re.search(pattern, "transformer.blocks.0.attn.q_proj") is None + assert re.search(pattern, "transformer.token_refiner.blocks.0.adaln_proj.linear") is None + + +def test_replace_layers_binds_every_native_target(monkeypatch): + """Regression for bound=0/259: layer replacement must reach all 259 modules.""" + from vllm.lora.layers import BaseLayerWithLoRA + from vllm.lora.peft_helper import PEFTHelper + + from vllm_omni.diffusion.lora import manager as manager_module + + class _StubLoRALayer(BaseLayerWithLoRA): + def __init__(self, base_layer): + super().__init__() + self.base_layer = base_layer + + def _linear() -> torch.nn.Module: + return torch.nn.Linear(_TINY_HIDDEN, _TINY_HIDDEN) + + def _block(*, with_adaln: bool) -> torch.nn.Module: + block = torch.nn.Module() + block.attn = torch.nn.Module() + block.attn.qkv_proj = _linear() + block.attn.out_proj = _linear() + block.mlp = torch.nn.Module() + block.mlp.fc1 = _linear() + block.mlp.fc2 = _linear() + if with_adaln: + block.adaln_proj = torch.nn.Module() + block.adaln_proj.linear = _linear() + return block + + transformer = torch.nn.Module() + transformer.blocks = torch.nn.ModuleList([_block(with_adaln=True) for _ in range(50)]) + transformer.token_refiner = torch.nn.Module() + transformer.token_refiner.blocks = torch.nn.ModuleList([_block(with_adaln=False) for _ in range(2)]) + transformer.final_layer = torch.nn.Module() + transformer.final_layer.adaln_proj = torch.nn.Module() + transformer.final_layer.adaln_proj.linear = _linear() + pipeline = torch.nn.Module() + pipeline.transformer = transformer + + monkeypatch.setattr(manager_module, "from_layer_diffusion", lambda *, layer, **_: _StubLoRALayer(layer)) + + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = pipeline + manager.dtype = torch.float32 + manager.max_cached_adapters = 1 + manager._lora_modules = {} + manager._max_lora_rank = 64 + manager._resident_lora_device = None + + peft_helper = PEFTHelper.from_dict( + { + "r": _TINY_RANK, + "lora_alpha": _TINY_RANK, + "target_modules": lora_module._NATIVE_TARGET_PATTERN, + } + ) + manager._replace_layers_with_lora(peft_helper) + + assert set(manager._lora_modules) == {f"transformer.{target}" for target in lora_module._NATIVE_EXPECTED_TARGETS} + + +def test_legacy_manager_uses_native_loader(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + class _Pipeline: + def _load_diffusion_lora_adapter(self, **kwargs): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + MiniMaxH3Pipeline, + ) + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + return pipeline._load_diffusion_lora_adapter(**kwargs) + + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = _Pipeline() + manager.dtype = torch.float32 + manager._expected_lora_modules = {"qkv_proj", "fc1", "adaln_proj.linear"} + + lora_model, peft_helper = manager._load_adapter(_request(path)) + assert lora_model.id == 7 + assert peft_helper.lora_alpha == _TINY_RANK + assert len(lora_model.loras) == 259 + + +def test_pipeline_schedule_inactive_when_scale_zero(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = {7} + pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])} + pipeline._base_schedule_by_partition = {"fl2va": None} + + sampling = SimpleNamespace( + lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"), + lora_scale=0.0, + num_inference_steps=4, + ) + assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None + assert not pipeline._has_active_native_lora(sampling) + + +def test_pipeline_schedule_falls_back_after_eviction(monkeypatch, tmp_path): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module + + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + request = _request(path) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._turbo_lora_adapter_ids = set() + pipeline._native_lora_adapter_ids = set() + pipeline._lora_sigma_schedules = {} + pipeline._base_schedule_by_partition = {"fl2va": None} + + loaded = pipeline._load_diffusion_lora_adapter( + lora_request=request, + lora_path=path, + dtype=torch.float32, + ) + assert loaded is not None + assert request.lora_int_id in pipeline._lora_sigma_schedules + + monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None) + monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None) + assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None + assert request.lora_int_id not in pipeline._lora_sigma_schedules + + sampling = SimpleNamespace( + lora_request=request, + lora_scale=1.0, + num_inference_steps=4, + ) + assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None + + +def test_native_packed_qkv_slices_are_tp2_divisible(): + assert lora_module._NATIVE_QKV_SLICE % 2 == 0 + assert lora_module._NATIVE_FFN_HIDDEN_SIZE % 2 == 0 + + +def test_lora_manager_activates_native_packed_qkv(tmp_path): + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + assert loaded is not None + lora_model, _, _ = loaded + packed = lora_model.get_lora("blocks.0.attn.qkv_proj") + assert isinstance(packed, PackedLoRALayerWeights) + assert len(packed.lora_b) == 3 + assert all(b.shape[0] == _TINY_QKV_SLICE for b in packed.lora_b) + + class _DummyLoRALayer: + def __init__(self): + self.n_slices = 3 + self.output_slices = (_TINY_QKV_SLICE, _TINY_QKV_SLICE, _TINY_QKV_SLICE) + self.set_calls: list[tuple[list, list]] = [] + self.reset_calls = 0 + + def reset_lora(self, index: int): + self.reset_calls += 1 + + def set_lora(self, index: int, lora_a, lora_b): + self.set_calls.append((lora_a, lora_b)) + + layer = _DummyLoRALayer() + manager = object.__new__(DiffusionLoRAManager) + manager.pipeline = torch.nn.Module() + manager._lora_modules = {"blocks.0.attn.qkv_proj": layer} + manager._registered_adapters = { + lora_model.id: lora_model, + } + manager._active_adapter_id = None + manager._adapter_scales = {} + manager._activate_adapter(lora_model.id, scale=0.5) + + assert manager._active_adapter_id == lora_model.id + + assert len(layer.set_calls) == 1 + lora_a_list, lora_b_list = layer.set_calls[0] + assert len(lora_a_list) == 3 + assert len(lora_b_list) == 3 + torch.testing.assert_close(lora_b_list[0], packed.lora_b[0] * 0.5) + torch.testing.assert_close(lora_b_list[1], packed.lora_b[1] * 0.5) + torch.testing.assert_close(lora_b_list[2], packed.lora_b[2] * 0.5) + + +def test_lora_manager_tp2_splits_native_packed_qkv_per_rank(): + """Each Q/K/V slice must be divisible by TP2 so rank-local layers can shard rows.""" + + tp_size = 2 + slice_rows = lora_module._NATIVE_QKV_SLICE + assert slice_rows % tp_size == 0 + local_rows = slice_rows // tp_size + + class _TpQkvLoRALayer: + def __init__(self, tp_rank: int): + self.tp_rank = tp_rank + self.n_slices = 3 + self.output_slices = (local_rows, local_rows, local_rows) + self.set_calls: list[tuple[list, list]] = [] + + def reset_lora(self, index: int): + return + + def set_lora(self, index: int, lora_a, lora_b): + self.set_calls.append((lora_a, lora_b)) + + for tp_rank in range(tp_size): + layer = _TpQkvLoRALayer(tp_rank) + start = tp_rank * local_rows + end = start + local_rows + full_packed = PackedLoRALayerWeights( + module_name="blocks.0.attn.qkv_proj", + rank=lora_module._NATIVE_RANK, + lora_alphas=[64, 64, 64], + lora_a=[torch.ones(lora_module._NATIVE_RANK, lora_module._NATIVE_HIDDEN_SIZE)] * 3, + lora_b=[torch.full((slice_rows, lora_module._NATIVE_RANK), float(i + 1)) for i in range(3)], + scaling=[1.0, 1.0, 1.0], + ) + local_b = [b[start:end].contiguous() for b in full_packed.lora_b] + layer.set_lora(0, full_packed.lora_a, local_b) + assert len(layer.set_calls) == 1 + _, lora_b_list = layer.set_calls[0] + assert all(b.shape[0] == local_rows for b in lora_b_list) + assert lora_b_list[0][0, 0].item() == 1.0 + assert lora_b_list[1][0, 0].item() == 2.0 + assert lora_b_list[2][0, 0].item() == 3.0 + + +# 259 modules, of which 52 qkv expand to 3 slices and 52 fc1 expand to 2. +_EXPECTED_NATIVE_SLICE_PAIRS = 415 + + +def _iter_native_slice_pairs(lora_model): + """Yield ``(module_name, lora_a, lora_b)`` per slice, flattening packed weights.""" + for module_name, weights in lora_model.loras.items(): + lora_a = weights.lora_a if isinstance(weights.lora_a, list) else [weights.lora_a] + lora_b = weights.lora_b if isinstance(weights.lora_b, list) else [weights.lora_b] + for a, b in zip(lora_a, lora_b, strict=True): + yield module_name, a, b + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_native_lora_load_is_platform_agnostic(tmp_path, dtype): + """The artifact is trained on Ascend, but loading it is not Ascend-only. + + This whole module runs in the CPU job with no ``torch_npu`` present, so a + passing run is itself the evidence. The assertions pin the property that + makes it portable: tensors land on CPU in exactly the requested dtype and + the loader never picks a device, leaving placement to the pipeline. + """ + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=dtype, + ) + + assert loaded is not None + lora_model, _, _ = loaded + tensors = { + f"{module_name}.{side}": tensor + for module_name, a, b in _iter_native_slice_pairs(lora_model) + for side, tensor in (("lora_a", a), ("lora_b", b)) + } + assert sum(1 for _ in _iter_native_slice_pairs(lora_model)) == _EXPECTED_NATIVE_SLICE_PAIRS + assert {name: tensor.device.type for name, tensor in tensors.items() if tensor.device.type != "cpu"} == {} + assert {name: tensor.dtype for name, tensor in tensors.items() if tensor.dtype != dtype} == {} + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +def test_native_lora_delta_matches_between_cpu_and_cuda(tmp_path): + """The packed adapter produces the same delta weight on CUDA as on CPU. + + Covers the GPU side of the same artifact: every ``lora_b @ lora_a`` product, + including the reordered q/k/v slices, is recomputed on CUDA and compared + against the CPU result, so a GPU deployment cannot silently diverge. + """ + path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" + _write_tiny_native(path) + + loaded = load_minimax_h3_native_lora( + partition="fl2va", + lora_request=_request(path), + lora_path=path, + dtype=torch.float32, + ) + + assert loaded is not None + lora_model, _, _ = loaded + checked = 0 + for module_name, a, b in _iter_native_slice_pairs(lora_model): + cpu_delta = b @ a + cuda_delta = b.cuda() @ a.cuda() + assert cuda_delta.device.type == "cuda" + torch.testing.assert_close(cuda_delta.cpu(), cpu_delta, msg=f"delta mismatch for {module_name}") + checked += 1 + assert checked == _EXPECTED_NATIVE_SLICE_PAIRS diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py new file mode 100644 index 000000000..c993de492 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from PIL import Image + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_h3_profiler_targets_cover_aggregate_reference_encoders(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + targets = MiniMaxH3Pipeline._PROFILER_TARGETS + assert "_encode_visual_conditions" in targets + assert "_encode_reference_audio_conditions" in targets + assert "_encode_video_conditions" not in targets + assert "_encode_video_audio_conditions" not in targets + assert "_encode_audio_conditions" not in targets + + +@pytest.mark.parametrize("distributed", [False, True]) +def test_manual_component_offload_wins_over_requested_model_offload(distributed): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=not distributed, + enable_distributed_layerwise_offload=distributed, + ) + pipeline._model_cpu_offload_modules = [] + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + component = Mock() + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + with pipeline._component_on_device(component): + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_not_called() + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + component.offload_to_cpu.assert_called_once_with() + + +def test_requested_model_offload_without_backend_uses_resident_components(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=False, + ) + pipeline._model_cpu_offload_modules = [] + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.encode_ids.assert_called_once() + + +def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = torch.nn.Linear(2, 2) + pipeline.transformers_ref = torch.nn.Linear(2, 2) + pipeline.text_encoder = torch.nn.Linear(2, 2) + pipeline.video_vae = torch.nn.Linear(2, 2) + pipeline.audio_vae = torch.nn.Linear(2, 2) + apply_offload = Mock() + remove_offload = Mock() + monkeypatch.setattr(module, "apply_sequential_offload", apply_offload) + monkeypatch.setattr(module, "remove_sequential_offload", remove_offload) + + pipeline.enable_omni_model_cpu_offload( + device=torch.device("cpu"), + pin_memory=False, + use_hsdp=False, + ) + + dits = [pipeline.transformer, pipeline.transformers_ref] + stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae] + apply_offload.assert_called_once_with( + dit_modules=dits, + encoder_modules=stages, + device=torch.device("cpu"), + pin_memory=False, + use_hsdp=False, + offload_initial_dits=True, + ) + + pipeline.disable_omni_model_cpu_offload() + + remove_offload.assert_called_once_with([*dits, *stages]) + + +@pytest.mark.parametrize("decode_fails", [False, True]) +def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + component = torch.nn.Linear(2, 2) + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + monkeypatch.setattr(module, "sequential_offload_component", record_component) + pipeline._model_cpu_offload_modules = [component] + + def decode(): + with pipeline._component_on_device(component): + events.append(("decode", component)) + if decode_fails: + raise RuntimeError("decode failed") + + if decode_fails: + with pytest.raises(RuntimeError, match="decode failed"): + decode() + else: + decode() + + assert events == [ + ("activate", component), + ("decode", component), + ("offload", component), + ] + + +@pytest.mark.parametrize("encode_fails", [False, True]) +def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode_fails): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.video_vae = Mock() + pipeline._model_cpu_offload_modules = [pipeline.video_vae] + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + def encode_image(image): + events.append(("encode", image.size)) + if encode_fails and image.width == 32: + raise RuntimeError("encode failed") + return torch.ones(image.width // 16, 4) + + pipeline.video_vae.encode_image.side_effect = encode_image + monkeypatch.setattr(module, "sequential_offload_component", record_component) + monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1)) + monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value) + images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))] + + if encode_fails: + with pytest.raises(RuntimeError, match="encode failed"): + pipeline._encode_visual_conditions(images, None, video_count=0) + else: + rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0) + assert rows.shape == (3, 4) + assert shapes == [(1, 1, 1), (1, 1, 2)] + + assert events == [ + ("activate", pipeline.video_vae), + ("encode", (16, 16)), + ("encode", (32, 16)), + ("offload", pipeline.video_vae), + ] + + +def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.video_vae = Mock() + pipeline.video_vae.is_distributed_enabled.return_value = False + pipeline._model_cpu_offload_modules = [pipeline.video_vae] + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + def encode_image(image): + events.append(("image", image.size)) + return torch.ones(1, 4) + + def encode_video(frames): + events.append(("video", frames)) + return torch.ones(3, 4), (3, 2, 2) + + pipeline.video_vae.encode_image.side_effect = encode_image + pipeline.video_vae.encode_video.side_effect = encode_video + monkeypatch.setattr(module, "sequential_offload_component", record_component) + monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1)) + monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value) + monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}") + images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))] + prepared_videos = [{"prepared_path": "reference.mp4"}] + + rows, shapes = pipeline._encode_visual_conditions( + images, + prepared_videos, + video_count=1, + ) + + assert rows.shape == (5, 4) + assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)] + assert events == [ + ("activate", pipeline.video_vae), + ("image", (16, 16)), + ("image", (32, 16)), + ("video", "frames:reference.mp4"), + ("offload", pipeline.video_vae), + ] + + +@pytest.mark.parametrize("encode_fails", [False, True]) +def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeypatch, encode_fails): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = Mock() + pipeline._model_cpu_offload_modules = [pipeline.audio_vae] + events = [] + + @contextmanager + def record_component(value): + events.append(("activate", value)) + try: + yield + finally: + events.append(("offload", value)) + + def encode_embedded(*args, **kwargs): + events.append(("embedded",)) + if encode_fails: + raise RuntimeError("audio encode failed") + return torch.ones(1, 2), [1] + + def encode_standalone(*args, **kwargs): + events.append(("standalone",)) + return torch.ones(1, 2), [1] + + pipeline._encode_video_audio_conditions_resident = encode_embedded + pipeline._encode_audio_conditions_resident = encode_standalone + monkeypatch.setattr(module, "sequential_offload_component", record_component) + + def call(): + return pipeline._encode_reference_audio_conditions( + [{"input_has_audio": True}], + has_audio=[True], + standalone_audios=[(torch.ones(1), 16_000)], + max_duration_seconds=1.0, + ) + + if encode_fails: + with pytest.raises(RuntimeError, match="audio encode failed"): + call() + else: + result = call() + assert result[0].shape == (1, 2) + assert result[2].shape == (1, 2) + + assert events[0] == ("activate", pipeline.audio_vae) + assert events[-1] == ("offload", pipeline.audio_vae) + assert events.count(("activate", pipeline.audio_vae)) == 1 + assert events.count(("offload", pipeline.audio_vae)) == 1 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py index efe53f826..53348f579 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py @@ -74,6 +74,8 @@ def test_t2va_and_fl2va_packing_keep_update_rows_separate(): assert fl2va["update_mask"].sum().item() == 12 assert (~fl2va["update_mask"]).sum().item() == 6 assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6 + assert t2va["video_spans"] == ({"start": 10, "latent_grid": (2, 2, 3), "role": "target"},) + assert fl2va["video_spans"] == ({"start": 16, "latent_grid": (2, 2, 3), "role": "target"},) def test_ref2va_packing_tracks_video_and_audio_update_masks(): @@ -101,6 +103,33 @@ def test_ref2va_packing_tracks_video_and_audio_update_masks(): assert packed["cu_seqlens"].tolist() == [0, 30, 64] +def test_ref2va_packing_publishes_physical_video_spans_only(): + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence_ref2va_blocks, + ) + + packed = minimax_h3_packed_sequence_ref2va_blocks( + text_len=4, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ref_blocks=[ + {"kind": "image", "latent_h": 4, "latent_w": 4}, + {"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 6}, + {"kind": "audio", "ref_audio_t": 1}, + {"kind": "video", "ref_audio_t": 0, "latent_t": 1, "latent_h": 4, "latent_w": 4}, + ], + ) + + # Image and all audio rows intentionally stay out of the sparse spans. + assert packed["video_spans"] == ( + {"start": 12, "latent_grid": (2, 2, 3), "role": "reference"}, + {"start": 26, "latent_grid": (1, 2, 2), "role": "reference"}, + {"start": 36, "latent_grid": (2, 2, 3), "role": "target"}, + ) + + def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one(): from vllm_omni.diffusion.models.minimax_h3.condition_noise import ( minimax_h3_audio_cond_noise_aug_rows, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py index 2b168499e..45f9229d6 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch @@ -36,7 +36,13 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp(): assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"] assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3 assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg - assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"} + assert set(MiniMaxH3DiTModel._sp_plan) == { + "sp_prepare", + "local_sp_prepare", + "sp_gather", + } + assert set(MiniMaxH3DiTModel._sp_plan["sp_prepare"]) == {0, 1, 2} + assert set(MiniMaxH3DiTModel._sp_plan["local_sp_prepare"]) == {2} model = object.__new__(MiniMaxH3DiTModel) nn.Module.__init__(model) @@ -89,6 +95,419 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims(): torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0) +def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + _build_rope_table, + ) + + local_freqs_half = torch.randn(2, 48) + local_freqs = torch.cat((local_freqs_half, local_freqs_half), dim=-1) + + actual = _build_rope_table(local_freqs) + expected = torch.cat( + (torch.cos(local_freqs_half), torch.sin(local_freqs_half)), + dim=-1, + ).to(torch.bfloat16) + + assert actual.shape == (2, 96) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def test_h3_prepared_rope_table_is_rank_local_and_validated(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + class Rope(nn.Module): + def forward(self, position_ids): + freqs_half = position_ids[0].float() + return torch.cat((freqs_half, freqs_half), dim=-1) + + model = object.__new__(h3.MiniMaxH3DiTModel) + nn.Module.__init__(model) + model.arch = h3.MiniMaxH3DiTArchConfig(rope_inv_freq_len=1) + model.rope = Rope() + model.local_sp_prepare = nn.Identity() + monkeypatch.setattr(h3, "_sequence_parallel_local_span", lambda *args, **kwargs: (1, 2)) + + position_ids = torch.arange(12, dtype=torch.long).reshape(1, 4, 3) + actual = model.prepare_rope_table(position_ids, seq_len=4) + expected = h3._build_rope_table(model.rope(position_ids[:, 1:3])) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + model._validate_prepared_rope_table(actual, local_len=2, device=torch.device("cpu")) + with pytest.raises(ValueError, match="local_seq_len"): + model._validate_prepared_rope_table(actual[:1], local_len=2, device=torch.device("cpu")) + + +def test_denoise_branch_prepares_rope_table_once_per_npu_run(monkeypatch: pytest.MonkeyPatch): + from vllm_omni.diffusion.models.minimax_h3 import denoise_loop + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import ( + MiniMaxH3DenoiseBranch, + minimax_h3_denoise_loop, + ) + + monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: True) + + packed = { + "seq_len": torch.tensor(3), + "img_pos": torch.tensor([1]), + "audio_pos": torch.tensor([2]), + "text_pos": torch.tensor([0]), + "update_mask": torch.tensor([True]), + "cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32), + "img_position_ids": torch.zeros(3, 3, dtype=torch.long), + "latent_grid": torch.tensor([1, 1, 1]), + "video_row_start": torch.tensor(1), + } + branch = MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=torch.zeros(1, 2), + token_tags=torch.zeros(3, dtype=torch.long), + device=torch.device("cpu"), + ) + + class Model: + def __init__(self): + self.rope_table = torch.zeros(3, 6, dtype=torch.bfloat16) + self.prepare_calls = 0 + self.forward_calls = 0 + + def prepare_rope_table(self, img_position_ids, *, seq_len): + assert img_position_ids is branch.static_kwargs["img_position_ids"] + assert seq_len == 3 + self.prepare_calls += 1 + return self.rope_table + + def __call__(self, **kwargs): + assert kwargs["rope_table"] is self.rope_table + self.forward_calls += 1 + return torch.zeros(1, 96), torch.zeros(1, 32) + + model = Model() + minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=torch.zeros(1, 96), + initial_audio_rows=torch.zeros(1, 32), + keyframe_cond_rows=None, + sigmas_video=[1.0, 0.5, 0.0], + sigmas_audio=[1.0, 0.5, 0.0], + device=torch.device("cpu"), + ) + + assert model.prepare_calls == 1 + assert model.forward_calls == 2 + + +def test_denoise_branch_keeps_rope_construction_on_the_reference_path_off_npu( + monkeypatch: pytest.MonkeyPatch, +): + from vllm_omni.diffusion.models.minimax_h3 import denoise_loop + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import ( + MiniMaxH3DenoiseBranch, + minimax_h3_denoise_loop, + ) + + monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: False) + packed = { + "seq_len": torch.tensor(3), + "img_pos": torch.tensor([1]), + "audio_pos": torch.tensor([2]), + "text_pos": torch.tensor([0]), + "update_mask": torch.tensor([True]), + "cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32), + "img_position_ids": torch.zeros(3, 3, dtype=torch.long), + "latent_grid": torch.tensor([1, 1, 1]), + "video_row_start": torch.tensor(1), + } + branch = MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=torch.zeros(1, 2), + token_tags=torch.zeros(3, dtype=torch.long), + device=torch.device("cpu"), + ) + + class Model: + def __init__(self): + self.forward_calls = 0 + + def prepare_rope_table(self, *_args, **_kwargs): + raise AssertionError("the reference path must not pre-build a rope_table") + + def __call__(self, **kwargs): + assert "rope_table" not in kwargs + self.forward_calls += 1 + return torch.zeros(1, 96), torch.zeros(1, 32) + + model = Model() + minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=torch.zeros(1, 96), + initial_audio_rows=torch.zeros(1, 32), + keyframe_cond_rows=None, + sigmas_video=[1.0, 0.5, 0.0], + sigmas_audio=[1.0, 0.5, 0.0], + device=torch.device("cpu"), + ) + + assert model.forward_calls == 2 + + +def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch): + from vllm_omni.diffusion import forward_context + from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True) + monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: "strict") + monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8) + monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3) + monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 8) + monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1) + monkeypatch.setattr(parallel_state, "get_allgather_parallel_world_size", lambda: 1) + + assert h3._sequence_parallel_local_span(16, hooks_applied=True) == (6, 2) + + +@pytest.mark.parametrize( + ("hooks_applied", "mode", "seq_len", "ulysses", "ring", "allgather"), + [ + (False, "strict", 16, 8, 1, 1), + (True, "advanced_uaa", 16, 8, 1, 1), + (True, "strict", 17, 8, 1, 1), + (True, "strict", 16, 4, 2, 1), + (True, "strict", 16, 1, 1, 8), + ], +) +def test_local_span_falls_back_when_local_embedding_is_unsafe( + monkeypatch, + hooks_applied, + mode, + seq_len, + ulysses, + ring, + allgather, +): + from vllm_omni.diffusion import forward_context + from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True) + monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: mode) + monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8) + monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3) + monkeypatch.setattr( + parallel_state, + "get_ulysses_parallel_world_size", + lambda: ulysses, + ) + monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: ring) + monkeypatch.setattr( + parallel_state, + "get_allgather_parallel_world_size", + lambda: allgather, + ) + + assert h3._sequence_parallel_local_span(seq_len, hooks_applied=hooks_applied) == ( + 0, + seq_len, + ) + + +def test_local_embedding_spans_reassemble_multimodal_rows(): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + class _IdentityProjection(nn.Module): + def forward(self, x): + return x, None + + class _IdentityRefiner(nn.Module): + def forward(self, x, **kwargs): + return x + + model = object.__new__(h3.MiniMaxH3DiTModel) + nn.Module.__init__(model) + model.hidden_size = 2 + model.video_patch_proj = _IdentityProjection() + model.audio_patch_proj = _IdentityProjection() + model.condition_proj = _IdentityProjection() + model.token_refiner = _IdentityRefiner() + model.time_embedder = nn.Identity() + + seq_len = 16 + x = torch.arange(seq_len * 2, dtype=torch.float32).reshape(1, seq_len, 2) + audio_x = x + 100 + text = torch.tensor( + [[200.0, 201.0], [202.0, 203.0], [204.0, 205.0], [206.0, 207.0]], + dtype=torch.bfloat16, + ) + common = { + "x": x, + "audio_x": audio_x, + "text_embeddings_selected": text, + "unique_timesteps": torch.tensor([0.2, 0.3]), + "img_pos": torch.tensor([1, 2, 6, 7, 11, 12]), + "audio_pos": torch.tensor([3, 4, 8, 9, 13, 14]), + "text_pos": torch.tensor([0, 5, 10, 15]), + "refiner_cu_seqlens": torch.tensor([0, 4, 4], dtype=torch.int32), + "refiner_max_seqlen": 4, + "seq_len": seq_len, + "device": torch.device("cpu"), + "local_span": (0, seq_len), + } + + full, full_t_emb = model._embed(**common) + local_spans = [] + for rank in range(8): + local_span = (rank * (seq_len // 8), seq_len // 8) + local, local_t_emb = model._embed( + **{ + **common, + "local_span": local_span, + } + ) + assert local.shape == (2, 2) + torch.testing.assert_close(local_t_emb, full_t_emb) + local_spans.append(local) + + torch.testing.assert_close(torch.cat(local_spans), full, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "local_projection", + [False, True], + ids=["fallback", "strict-local"], +) +def test_sp_output_boundary_preserves_logits_and_minimizes_gather( + monkeypatch, + local_projection, +): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + class _StaticRope(nn.Module): + input_rows = None + + def forward(self, position_ids): + self.input_rows = position_ids.shape[1] + return torch.zeros(position_ids.shape[1], 6) + + class _Prepare(nn.Module): + def __init__(self, shard): + super().__init__() + self.shard = shard + + def forward(self, hidden, rope, combined): + local_len = hidden.shape[0] // 2 if self.shard else hidden.shape[0] + return hidden[:local_len], rope[:local_len], combined[:local_len] + + class _Gather(nn.Module): + def __init__(self, remote_rows, replace): + super().__init__() + self.remote_rows = remote_rows + self.replace = replace + self.input_width = None + + def forward(self, tensor): + self.input_width = tensor.shape[-1] + if self.replace: + return self.remote_rows + return torch.cat((tensor, self.remote_rows), dim=0) + + class _FinalLayer(nn.Module): + def forward(self, hidden, *, t_emb, inverse_indices): + del t_emb + inverse = inverse_indices.to(torch.float32).unsqueeze(-1) + return ( + hidden[:, :2].to(torch.float32) + inverse, + hidden[:, 2:3].to(torch.float32) - inverse, + ) + + seq_len = 8 + full_hidden = torch.arange(seq_len * 4, dtype=torch.bfloat16).reshape(seq_len, 4) + inverse_indices = torch.tensor([0, 1, 0, 1, 1, 0, 1, 0]) + final_layer = _FinalLayer() + remote_video, remote_audio = final_layer( + full_hidden[seq_len // 2 :], + t_emb=torch.empty(0), + inverse_indices=inverse_indices[seq_len // 2 :], + ) + if local_projection: + remote_rows = torch.cat((remote_video, remote_audio), dim=-1) + else: + remote_rows = full_hidden + gather = _Gather(remote_rows, replace=not local_projection) + + model = object.__new__(h3.MiniMaxH3DiTModel) + nn.Module.__init__(model) + model.arch = h3.MiniMaxH3DiTArchConfig( + latents_dim=2, + audio_latents_dim=1, + patch_size=(1, 1, 1), + ) + rope = _StaticRope() + model.rope = rope + model.sp_prepare = _Prepare(shard=True) + model.local_sp_prepare = _Prepare(shard=False) + model.sp_gather = gather + model.blocks = nn.ModuleList() + model.final_layer = final_layer + monkeypatch.setattr( + model, + "_embed", + lambda **kwargs: ( + full_hidden[: seq_len // 2] if local_projection else full_hidden, + torch.tensor([[0.5]], dtype=torch.float32), + ), + ) + + def _local_span(*args, **kwargs): + del args, kwargs + return (0, seq_len // 2) if local_projection else (0, seq_len) + + monkeypatch.setattr( + h3, + "_sequence_parallel_local_span", + _local_span, + ) + + img_pos = torch.tensor([0, 2, 4, 6]) + audio_pos = torch.tensor([1, 3, 5, 7]) + video, audio = model( + x=torch.zeros(1, seq_len, 2), + audio_x=torch.zeros(1, seq_len, 1), + img_position_ids=torch.zeros(1, seq_len, 3, dtype=torch.long), + unique_timesteps=torch.tensor([0.5]), + inverse_indices=inverse_indices, + update_mask=torch.ones(img_pos.numel()), + token_tags=torch.zeros(seq_len, dtype=torch.long), + prompt_embeds=torch.empty(0, 2), + img_pos_info={"position_ids": img_pos}, + audio_pos_info={"position_ids": audio_pos}, + text_pos_info={"position_ids": torch.empty(0, dtype=torch.long)}, + img_pos_for_infer_output_info={"position_ids": img_pos}, + packed_seq_params={ + "cu_seqlens_q": torch.tensor([0, seq_len, seq_len], dtype=torch.int32), + "max_seqlen_q": seq_len, + }, + refiner_packed_seq_params={ + "cu_seqlens_q": torch.tensor([0, 0, 0], dtype=torch.int32), + "max_seqlen_q": 0, + }, + ) + + expected_video, expected_audio = final_layer( + full_hidden, + t_emb=torch.empty(0), + inverse_indices=inverse_indices, + ) + torch.testing.assert_close(video, expected_video.index_select(0, img_pos)) + torch.testing.assert_close(audio, expected_audio.index_select(0, audio_pos)) + expected_gather_width = 3 if local_projection else full_hidden.shape[-1] + expected_rope_rows = seq_len // 2 if local_projection else seq_len + assert gather.input_width == expected_gather_width + assert rope.input_rows == expected_rope_rows + + @pytest.mark.parametrize( ("tp_size", "message"), [ diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index b9ae3335c..92bd09ac2 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -191,6 +191,49 @@ def test_model_load_weights_transforms_before_calling_vllm_loader(): ] +def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch): + from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( + get_direct_mmap_adapter, + ) + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "Attention", _FakeAttention) + + arch = h3.MiniMaxH3DiTArchConfig( + hidden_size=1, + num_attention_heads=2, + attention_head_dim=1, + rope_inv_freq_len=1, + ) + attention = h3.MiniMaxH3Attention( + arch, + quant_config=None, + prefix="blocks.0.attn", + ) + transformer = object.__new__(h3.MiniMaxH3DiTModel) + nn.Module.__init__(transformer) + transformer.blocks = nn.ModuleList([nn.Module()]) + transformer.blocks[0].attn = attention + pipeline = nn.Module() + pipeline.transformer = transformer + + adapter = get_direct_mmap_adapter(pipeline) + assert adapter is not None + policy = adapter.policy_for( + "transformer.blocks.0.attn.qkv_proj.weight", + attention.qkv_proj.weight, + ) + assert policy is not None + assert policy.allow_custom_loader + assert policy.transform is not None + checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1) + + assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5] + assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform") + + def test_pipeline_resolves_transformer_component_quant_config(): from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( _resolve_component_quant_config, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index 80d99895f..0fb043a1e 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from pathlib import Path @@ -142,12 +142,6 @@ def _generate_joint_output(omni, config: QualityTestConfig): first = outputs[0] worker_peak_memory_mb = float(getattr(first, "peak_memory_mb", 0.0) or 0.0) - if hasattr(first, "request_output") and isinstance(first.request_output, list): - first = first.request_output[0] - worker_peak_memory_mb = max( - worker_peak_memory_mb, - float(getattr(first, "peak_memory_mb", 0.0) or 0.0), - ) if not hasattr(first, "images") or not first.images: raise ValueError("Could not extract video frames from H3 output") frames = first.images[0] @@ -193,6 +187,7 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig): "model": model_ref, "enforce_eager": True, "tensor_parallel_size": 2, + # The fused BF16 baseline only fits on H100-80GB with encoder TP. "text_encoder_tp_size": 2, "vae_use_tiling": True, } diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py new file mode 100644 index 000000000..45a109b28 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_minimax_h3_npu_sdpa_patch_is_independent_and_idempotent(monkeypatch) -> None: + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + def sdpa_sentinel(query, key, value): + del key, value + return query + + def rope_sentinel(q, k, cos, sin): + del cos, sin + return q, k + + monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False) + monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False) + monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel) + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel) + + npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch() + + assert encoder_module._scaled_dot_product_attention is npu_minimax_h3._scaled_dot_product_attention_npu + assert encoder_module._apply_rotary_pos_emb is rope_sentinel + assert not npu_minimax_h3._ROPE_PATCHED + assert npu_minimax_h3._SDPA_PATCHED + + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel) + npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch() + + assert encoder_module._scaled_dot_product_attention is sdpa_sentinel + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize( + ("num_heads", "num_kv_heads", "expected_enable_gqa"), + [(2, 2, False), (4, 2, True)], +) +def test_minimax_h3_npu_sdpa_preserves_compressed_heads_and_selects_gqa( + monkeypatch, + num_heads: int, + num_kv_heads: int, + expected_enable_gqa: bool, +) -> None: + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + captured: dict[str, object] = {} + + def fake_sdpa(query, key, value, **kwargs): + captured.update(query=query, key=key, value=value, kwargs=kwargs) + return query + + monkeypatch.setattr(F, "scaled_dot_product_attention", fake_sdpa) + query = torch.randn(1, num_heads, 3, 8) + key = torch.randn(1, num_kv_heads, 3, 8) + value = torch.randn_like(key) + + output = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value) + + assert captured["query"] is query + assert captured["key"] is key + assert captured["value"] is value + assert captured["kwargs"] == { + "dropout_p": 0.0, + "is_causal": True, + "enable_gqa": expected_enable_gqa, + } + assert output is query + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize( + ("num_heads", "num_key_heads", "num_value_heads", "message"), + [ + (4, 2, 1, "key and value to have the same number of heads"), + (4, 0, 0, "at least one KV head"), + (3, 2, 2, "query heads to be a multiple of KV heads"), + ], +) +def test_minimax_h3_npu_sdpa_rejects_invalid_head_configurations( + num_heads: int, + num_key_heads: int, + num_value_heads: int, + message: str, +) -> None: + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + query = torch.randn(1, num_heads, 3, 8) + key = torch.randn(1, num_key_heads, 3, 8) + value = torch.randn(1, num_value_heads, 3, 8) + + with pytest.raises(ValueError, match=message): + npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_minimax_h3_npu_causal_gqa_matches_expanded_kv() -> None: + """Exercise the real torch-npu causal GQA operator.""" + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + torch.manual_seed(0) + batch_size, seq_len = 1, 16 + num_heads, num_kv_heads, head_size = 8, 2, 128 + device = torch.device("npu") + dtype = torch.bfloat16 + + query = torch.randn(batch_size, num_heads, seq_len, head_size, device=device, dtype=dtype) + key = torch.randn(batch_size, num_kv_heads, seq_len, head_size, device=device, dtype=dtype) + value = torch.randn_like(key) + + actual = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value) + + repeat_num = num_heads // num_kv_heads + expected = F.scaled_dot_product_attention( + query, + key.repeat_interleave(repeat_num, dim=1), + value.repeat_interleave(repeat_num, dim=1), + dropout_p=0.0, + is_causal=True, + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py new file mode 100644 index 000000000..7422c0b5e --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_rope_patch_is_independent_and_idempotent(monkeypatch) -> None: + from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module + from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3 + + def rope_sentinel(q, k, cos, sin): + del cos, sin + return q, k + + def sdpa_sentinel(query, key, value): + del key, value + return query + + monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False) + monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False) + monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel) + monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel) + + npu_minimax_h3.apply_minimax_h3_qwen3vl_patch() + + assert encoder_module._apply_rotary_pos_emb is npu_minimax_h3._apply_rotary_pos_emb_npu + assert encoder_module._scaled_dot_product_attention is sdpa_sentinel + assert npu_minimax_h3._ROPE_PATCHED + assert not npu_minimax_h3._SDPA_PATCHED + + monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel) + npu_minimax_h3.apply_minimax_h3_qwen3vl_patch() + + assert encoder_module._apply_rotary_pos_emb is rope_sentinel + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize("num_q_heads,num_kv_heads,seq_len", [(8, 2, 16), (64, 8, 1)]) +def test_qwen3_vl_npu_rope_matches_bnsd_reference( + num_q_heads: int, + num_kv_heads: int, + seq_len: int, +) -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu + + device = torch.device("npu") + dtype = torch.bfloat16 + q = torch.randn(1, num_q_heads, seq_len, 128, device=device, dtype=dtype) + k = torch.randn(1, num_kv_heads, seq_len, 128, device=device, dtype=dtype) + cos = torch.randn(1, seq_len, 128, device=device, dtype=dtype) + sin = torch.randn(1, seq_len, 128, device=device, dtype=dtype) + + q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin) + cos_bnsd = cos.unsqueeze(1) + sin_bnsd = sin.unsqueeze(1) + + def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + torch.testing.assert_close( + q_out, + q * cos_bnsd + rotate_half(q) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + torch.testing.assert_close( + k_out, + k * cos_bnsd + rotate_half(k) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_rope_preserves_batched_mrope_positions() -> None: + """Keep per-sample Qwen3-VL M-RoPE frequencies on the BNSD NPU path.""" + from types import SimpleNamespace + + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLTextRotaryEmbedding, + ) + from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu + + device = torch.device("npu") + dtype = torch.bfloat16 + batch_size, seq_len, head_dim = 2, 16, 128 + config = SimpleNamespace( + hidden_size=head_dim, + num_attention_heads=1, + rope_parameters={ + "rope_theta": 1_000_000.0, + "mrope_section": [24, 20, 20], + }, + ) + rope = MiniMaxH3Qwen3VLTextRotaryEmbedding(config).to(device) + positions = torch.arange(seq_len, device=device) + position_ids = torch.stack( + ( + torch.stack((positions, positions + 7)), + torch.stack((positions * 2, positions * 2 + 3)), + torch.stack((positions * 3, positions * 3 + 5)), + ) + ) + q = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn(batch_size, 2, seq_len, head_dim, device=device, dtype=dtype) + cos, sin = rope(q, position_ids) + + q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin) + cos_bnsd = cos.unsqueeze(1) + sin_bnsd = sin.unsqueeze(1) + + def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + assert not torch.equal(cos[0], cos[1]) + torch.testing.assert_close( + q_out, + q * cos_bnsd + rotate_half(q) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) + torch.testing.assert_close( + k_out, + k * cos_bnsd + rotate_half(k) * sin_bnsd, + atol=2e-2, + rtol=2e-2, + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py new file mode 100644 index 000000000..3ed7649c6 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_marks +from vllm_omni.platforms import current_omni_platform + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.diffusion, + *hardware_marks(res={"npu": "A3"}, num_cards=1), +] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize("shape", [(1, 16, 256), (2, 1, 256)]) +def test_qwen3_vl_npu_swiglu_matches_split_reference( + shape: tuple[int, int, int], +) -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed + + packed = torch.randn(*shape, device="npu", dtype=torch.bfloat16) + gate, up = packed.chunk(2, dim=-1) + + torch.testing.assert_close( + npu_swiglu_from_packed(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None: + from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed + + x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16) + weight = torch.randn(256, 128, device="npu", dtype=torch.bfloat16) + packed = F.linear(x, weight) + gate = F.linear(x, weight[:128]) + up = F.linear(x, weight[128:]) + + torch.testing.assert_close( + npu_swiglu_from_packed(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py index 551fd68a5..65a0221ca 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py @@ -1,10 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import patch + +import pytest import torch +import torch.nn as nn from vllm_omni.diffusion.layers.norm import RMSNorm -from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm +from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLRMSNorm, + MiniMaxH3Qwen3VLTextDecoderLayer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None: @@ -24,3 +34,43 @@ def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None: assert norm.weight.dtype == torch.bfloat16 assert set(norm.state_dict()) == {"weight"} torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0) + + +def test_qwen3_vl_decoder_preserves_post_attention_residual_contract() -> None: + class ConstantAttention(nn.Module): + def forward(self, hidden_states, position_embeddings): + del position_embeddings + return torch.full_like(hidden_states, 2.0) + + config = SimpleNamespace( + hidden_size=8, + intermediate_size=16, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=2, + rms_norm_eps=1e-6, + ) + group = SimpleNamespace(rank_in_group=0, world_size=1) + with ( + patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0), + patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1), + ): + layer = MiniMaxH3Qwen3VLTextDecoderLayer(group, config, torch.float32) + layer.input_layernorm = nn.Identity() + layer.self_attn = ConstantAttention() + layer.mlp = nn.Identity() + + hidden_states = torch.arange(24, dtype=torch.float32).reshape(1, 3, 8) + attention_output = torch.full_like(hidden_states, 2.0) + expected_residual = hidden_states + attention_output + expected_normalized = RMSNorm.forward_native(layer.post_attention_layernorm, expected_residual) + + with patch.object( + layer.post_attention_layernorm, + "forward", + wraps=layer.post_attention_layernorm.forward_native, + ): + output = layer(hidden_states, (torch.empty(0), torch.empty(0))) + + assert isinstance(layer.post_attention_layernorm, MiniMaxH3Qwen3VLRMSNorm) + torch.testing.assert_close(output, expected_residual + expected_normalized) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py new file mode 100644 index 000000000..63c3a42f1 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -0,0 +1,473 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 step-wise execution (continuous batching) contract tests. + +These run on CPU against a stand-in DiT, so they cover the packing, the +per-request scheduler math, and the runner-facing state wiring without needing +checkpoint weights. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_HIDDEN = 8 + + +class _SegmentMeanModel: + """Stand-in DiT whose output depends on packed-document boundaries. + + Every row is mixed with the mean of its own ``cu_seqlens`` document and with + its own timestep, so a wrong document boundary, row offset, or timestep + assignment shows up as a numeric difference instead of passing silently. + """ + + def __call__(self, **kwargs): + x = kwargs["x"][0] + audio_x = kwargs["audio_x"][0] + bounds = kwargs["packed_seq_params"]["cu_seqlens_q"].tolist() + if kwargs["packed_seq_params"].get("num_requests", 1) > 1: + assert all(start < stop for start, stop in zip(bounds[:-1], bounds[1:])) + img_pos = kwargs["img_pos_info"]["position_ids"] + audio_pos = kwargs["audio_pos_info"]["position_ids"] + + pooled_video = torch.zeros_like(x) + pooled_audio = torch.zeros_like(audio_x) + for start, stop in zip(bounds[:-1], bounds[1:]): + if stop <= start: + continue + pooled_video[start:stop] = x[start:stop].mean(dim=0, keepdim=True) + pooled_audio[start:stop] = audio_x[start:stop].mean(dim=0, keepdim=True) + + row_timesteps = kwargs["unique_timesteps"][kwargs["inverse_indices"]].unsqueeze(-1) + video = (pooled_video + 0.5 * x + row_timesteps)[img_pos] + audio = (pooled_audio + 0.5 * audio_x + row_timesteps)[audio_pos] + if not kwargs.get("skip_mask_out_condition", False): + video = video * kwargs["update_mask"].view(-1).unsqueeze(-1) + return video, audio + + +def _make_branch(*, text_len: int, latent_t: int, latent_h: int, latent_w: int, audio_t: int, seed: int): + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import MiniMaxH3DenoiseBranch + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import minimax_h3_packed_sequence + + packed = minimax_h3_packed_sequence( + text_len=text_len, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + include_keyframe_cond=False, + ) + generator = torch.Generator().manual_seed(seed) + text_embeddings = torch.randn(text_len, _HIDDEN, generator=generator, dtype=torch.float32) + branch = MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=text_embeddings, + token_tags=packed["token_tags"], + device=torch.device("cpu"), + ) + video_rows = torch.randn(int(branch.img_pos.shape[0]), 96, generator=generator, dtype=torch.float32) + audio_rows = torch.randn(int(branch.audio_pos.shape[0]), 32, generator=generator, dtype=torch.float32) + return branch, video_rows, audio_rows + + +def _sigmas(num_points: int, shift: float) -> list[float]: + from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas + + return minimax_h3_time_shift_sigmas(num_steps=num_points, shift_scale=shift) + + +def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.diffusion.worker.utils import StepRequestState + + state = StepRequestState(request_id=request_id, sampling=SimpleNamespace()) + state.latents = video_rows.clone() + state.timesteps = torch.tensor([1.0 - sigma for sigma in sigmas_video[:-1]], dtype=torch.float32) + state.step_index = 0 + state.extra = { + mod._STEP_BRANCH: branch, + # Co-batched requests must share one DiT instance, or denoise_step() + # treats the batch as mixed-task and falls back to one forward each. + mod._STEP_TRANSFORMER: model, + mod._STEP_AUDIO_ROWS: audio_rows.clone(), + mod._STEP_COND_ANCHOR: None, + mod._STEP_AUDIO_ANCHOR: None, + mod._STEP_SIGMAS_VIDEO: sigmas_video, + mod._STEP_SIGMAS_AUDIO: sigmas_audio, + } + return state + + +def _step_pipeline(model, *, packed_batch_supported: bool = True): + """A pipeline instance carrying only what the step methods touch.""" + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + pipeline.transformer = model + pipeline.device = torch.device("cpu") + pipeline._transformer_for_task = lambda task: model + pipeline._packed_batch_supported = lambda transformer: packed_batch_supported + return pipeline + + +def test_step_execution_matches_request_mode_denoise_loop(): + """Stepping through the contract must reproduce the request-mode loop.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop + + model = _SegmentMeanModel() + branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5) + sigmas_video = _sigmas(6, 12.0) + sigmas_audio = _sigmas(6, 3.0) + + reference_video, reference_audio = minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=video_rows, + initial_audio_rows=audio_rows, + keyframe_cond_rows=None, + sigmas_video=sigmas_video, + sigmas_audio=sigmas_audio, + device=torch.device("cpu"), + ) + + pipeline = _step_pipeline(model) + state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + input_batch = SimpleNamespace(states=(state,)) + + steps = 0 + while not state.denoise_completed: + noise_pred = pipeline.denoise_step(input_batch, states=[state]) + pipeline.step_scheduler(state, noise_pred) + steps += 1 + + assert steps == len(sigmas_video) - 1 + assert state.total_steps == steps + torch.testing.assert_close(state.latents, reference_video) + torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio) + + +def test_batched_step_execution_matches_independent_requests(): + """Two co-batched requests must land where they would have landed alone.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + model = _SegmentMeanModel() + specs = [ + # Exactly 64 packed rows exercises the no-padding-tail boundary case. + dict(text_len=46, latent_t=7, latent_h=2, latent_w=4, audio_t=2, seed=7), + dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6), + ] + # Different step counts, so the batch composition changes mid-flight. + schedules = [(_sigmas(6, 12.0), _sigmas(6, 3.0)), (_sigmas(4, 12.0), _sigmas(4, 3.0))] + + pipeline = _step_pipeline(model) + + alone: list[tuple[torch.Tensor, torch.Tensor]] = [] + for spec, (sigmas_video, sigmas_audio) in zip(specs, schedules): + branch, video_rows, audio_rows = _make_branch(**spec) + state = _make_state("solo", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + while not state.denoise_completed: + pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state])) + alone.append((state.latents, state.extra[mod._STEP_AUDIO_ROWS])) + + states = [] + for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)): + branch, video_rows, audio_rows = _make_branch(**spec) + if index == 0: + assert branch.used_len == branch.seq_len + states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)) + + active = list(states) + while active: + noise_pred = pipeline.denoise_step(SimpleNamespace(states=tuple(active)), states=active) + offset = 0 + for state in active: + rows = state.latents.shape[0] + pipeline.step_scheduler(state, noise_pred[offset : offset + rows]) + offset += rows + assert offset == noise_pred.shape[0] + # Finished requests leave the batch, exactly like the runner drops them. + active = [state for state in active if not state.denoise_completed] + + for state, (expected_video, expected_audio) in zip(states, alone): + torch.testing.assert_close(state.latents, expected_video) + torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], expected_audio) + + +def test_both_modes_publish_denoise_progress_for_gated_attention(): + """TRTLLM's skip gate and RAINFUSION's warmup stay dense without these.""" + from vllm_omni.diffusion import forward_context as fc + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop + + model = _SegmentMeanModel() + branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=14) + sigmas_video = _sigmas(4, 12.0) + sigmas_audio = _sigmas(4, 3.0) + + published: list[tuple[int | None, float | None]] = [] + + class _Recorder: + def __init__(self): + self.denoise_step_idx = None + self.denoise_timestep = None + + def __setattr__(self, name, value): + object.__setattr__(self, name, value) + if name == "denoise_timestep": + published.append((self.denoise_step_idx, value)) + + recorder = _Recorder() + published.clear() # drop the pair __init__ recorded + original = fc._forward_context + fc._forward_context = recorder + try: + minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=video_rows, + initial_audio_rows=audio_rows, + keyframe_cond_rows=None, + sigmas_video=sigmas_video, + sigmas_audio=sigmas_audio, + device=torch.device("cpu"), + ) + request_mode = list(published) + + published.clear() + pipeline = _step_pipeline(model) + state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio) + while not state.denoise_completed: + pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state])) + step_mode = list(published) + finally: + fc._forward_context = original + + num_steps = len(sigmas_video) - 1 + expected = [(step, sigmas_video[step]) for step in range(num_steps)] + assert request_mode == expected + [(None, None)] + # Step mode has no loop to close, so it publishes only the per-step pairs. + assert step_mode == expected + + +def test_mixed_step_batch_leaves_gated_attention_dense(): + """A batch spanning different steps has no single progress point to publish.""" + from vllm_omni.diffusion import forward_context as fc + + model = _SegmentMeanModel() + first, first_video, first_audio = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=15) + second, second_video, second_audio = _make_branch( + text_len=5, latent_t=3, latent_h=6, latent_w=4, audio_t=2, seed=16 + ) + sigmas_video, sigmas_audio = _sigmas(6, 12.0), _sigmas(6, 3.0) + states = [ + _make_state("req-0", model, first, first_video, first_audio, sigmas_video, sigmas_audio), + _make_state("req-1", model, second, second_video, second_audio, sigmas_video, sigmas_audio), + ] + states[1].step_index = 2 # admitted later, so it trails the batch + + recorder = SimpleNamespace(denoise_step_idx="unset", denoise_timestep="unset") + original = fc._forward_context + fc._forward_context = recorder + try: + _step_pipeline(model).denoise_step(SimpleNamespace(states=tuple(states)), states=states) + finally: + fc._forward_context = original + + assert recorder.denoise_step_idx is None + assert recorder.denoise_timestep is None + + +def test_prepare_encode_seeds_runner_visible_state(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8) + sigmas_video = _sigmas(6, 12.0) + sigmas_audio = _sigmas(6, 3.0) + context = { + "height": 96, + "width": 64, + "latent_t": 2, + "latent_h": 4, + "latent_w": 6, + "audio_t": 3, + **{key: None for key in mod._MINIMAX_H3_DENOISE_INPUT_KEYS}, + } + + pipeline = _step_pipeline(_SegmentMeanModel()) + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {}))) + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context) + monkeypatch.setattr( + mod.MiniMaxH3Pipeline, + "_build_denoise_inputs", + lambda self, **_: { + "branch": branch, + "video_rows": video_rows, + "audio_rows": audio_rows, + "cond_anchor": None, + "audio_anchor": None, + "sigmas_video": sigmas_video, + "sigmas_audio": sigmas_audio, + }, + ) + + from vllm_omni.diffusion.worker.utils import StepRequestState + + state = StepRequestState( + request_id="req-0", + sampling=SimpleNamespace(num_outputs_per_prompt=1), + prompt="a prompt", + ) + pipeline.prepare_encode(state) + + # The runner slices the batched velocity by this row count. + assert state.latents.shape == (int(branch.img_pos.shape[0]), 96) + assert state.total_steps == len(sigmas_video) - 1 + assert state.step_index == 0 + assert state.do_true_cfg is False + torch.testing.assert_close(state.current_timestep, torch.tensor(1.0 - sigmas_video[0])) + assert state.extra[mod._STEP_BRANCH] is branch + assert state.extra[mod._STEP_SHAPE]["height"] == 96 + + +def test_prepare_encode_rejects_request_mode_only_features(): + """Multi-output and DLO have no representation in the step contract.""" + from vllm_omni.diffusion.worker.utils import StepRequestState + from vllm_omni.errors import OmniClientError + + # A request state carries exactly one latent tensor. + multi_output = StepRequestState( + request_id="req-0", + sampling=SimpleNamespace(num_outputs_per_prompt=2), + prompt="a prompt", + ) + with pytest.raises(OmniClientError, match="one output per request"): + _step_pipeline(_SegmentMeanModel()).prepare_encode(multi_output) + + # Distributed layerwise offload streams the DiT around a whole denoise loop. + single_output = StepRequestState( + request_id="req-1", + sampling=SimpleNamespace(num_outputs_per_prompt=1), + prompt="a prompt", + ) + dlo_pipeline = _step_pipeline(_SegmentMeanModel()) + dlo_pipeline._dlo_residency_controller = object() + with pytest.raises(ValueError, match="distributed layerwise offload"): + dlo_pipeline.prepare_encode(single_output) + + +def test_prepare_encode_rejects_high_quality_cache_dit(): + """quality=high installs a transformer-scoped Cache-DiT profile that + would leak across interleaved or co-batched step-mode requests.""" + from vllm_omni.diffusion.worker.utils import StepRequestState + from vllm_omni.errors import OmniClientError + + state = StepRequestState( + request_id="req-hi", + sampling=SimpleNamespace(num_outputs_per_prompt=1, quality="high"), + prompt="a prompt", + ) + with pytest.raises(OmniClientError, match="quality=high"): + _step_pipeline(_SegmentMeanModel()).prepare_encode(state) + + +def _fake_attention_module( + *, + use_ring: bool, + backend: str = "FLASH_ATTN", + supports_multi_doc: bool = True, +): + """Build a bare MiniMaxH3Attention with the attributes ``_packed_batch_supported`` reads. + + ``supports_multi_doc`` models the platform-dependent capability probe on + ``AttentionBackend`` (e.g. FLASH_ATTN returns True on CUDA/ROCm/MUSA and + False on NPU/XPU); ``backend`` remains only for logging/back-compat. + """ + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3Attention + + attn = object.__new__(MiniMaxH3Attention) + attn.attention = SimpleNamespace( + attn_backend=SimpleNamespace( + get_name=lambda: backend, + supports_multi_doc_packed_varlen=lambda: supports_multi_doc, + ), + use_ring=use_ring, + ) + return attn + + +class _FakeTransformer: + """The minimal ``modules()`` protocol ``_packed_batch_supported`` walks.""" + + def __init__(self, attentions): + self._attentions = list(attentions) + + def modules(self): + return iter([self, *self._attentions]) + + +@pytest.mark.parametrize( + "attention", + [ + _fake_attention_module(use_ring=True), + _fake_attention_module(use_ring=False, backend="XFORMERS", supports_multi_doc=False), + _fake_attention_module(use_ring=False, backend="FLASH_ATTN", supports_multi_doc=False), + ], +) +def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention): + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline + + assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False + + +def test_broadcast_rank0_exception_single_rank_reraises(): + """Single-rank execution has no group; the helper just reraises.""" + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _broadcast_rank0_exception, + ) + from vllm_omni.errors import OmniClientError + + _broadcast_rank0_exception(None) + with pytest.raises(OmniClientError, match="bad ref"): + _broadcast_rank0_exception(OmniClientError("bad ref")) + + +def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch): + """A rank-0 error becomes a matching client error on every other DiT rank.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + from vllm_omni.errors import OmniClientError + + def fake_rank_world(rank): + return lambda: (object(), rank, 4) + + def make_broadcast(rank0_payload): + def fake_broadcast(payload_list, *, src, group): + payload_list[0] = rank0_payload + + return fake_broadcast + + monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0)) + err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError") + rank0_payload = { + "type": type(err).__name__, + "message": str(err), + "status_code": err.status_code, + "error_type": err.error_type, + } + monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload)) + with pytest.raises(OmniClientError) as rank0_info: + mod._broadcast_rank0_exception(err) + assert rank0_info.value is err + + monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2)) + with pytest.raises(OmniClientError) as rank2_info: + mod._broadcast_rank0_exception(None) + assert rank2_info.value.status_code == 422 + assert rank2_info.value.error_type == "UnprocessableEntityError" + assert "invalid reference-video file" in str(rank2_info.value) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py new file mode 100644 index 000000000..7a1c54f7b --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py @@ -0,0 +1,385 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.gpu, pytest.mark.diffusion] + + +def _selected_operators(): + from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import ( + resolve_h3_vae_operators, + ) + + if not current_omni_platform.is_available(): + pytest.skip("No accelerator is available") + device = current_omni_platform.get_torch_device() + operators = resolve_h3_vae_operators(device) + if operators is None: + pytest.skip("No H3 VAE operator implementation is registered for this target") + return device, operators + + +def _operator_set(*, supports=lambda _device: True): + from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import H3VAEOperatorSet + + return H3VAEOperatorSet( + supports=supports, + qk_norm_rope=lambda *_args: None, + scaled_residual=lambda *_args: None, + ) + + +def _qk_reference(x, cos, sin): + normalized = F.rms_norm(x.float(), (64,), None, 1e-5).to(x.dtype) + rotary = normalized[..., :48] + first, second = rotary.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return torch.cat( + (rotary * cos + rotated * sin, normalized[..., 48:]), + dim=-1, + ) + + +def _vit_norm_input(_module, hidden_states): + return hidden_states.float() + + +def _same_dtype_norm_input(_module, hidden_states): + return hidden_states + + +def _failing_norm_input(_module, _hidden_states): + raise RuntimeError("unsupported remote normalization semantics") + + +@pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)]) +def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence): + device, operators = _selected_operators() + + torch.manual_seed(17) + qkv = torch.randn( + batch, + sequence, + 32, + 192, + device=device, + dtype=torch.float16, + ) + q, k, _ = qkv.chunk(3, dim=-1) + cos = torch.randn(batch, sequence, 1, 48, device=device, dtype=torch.float16) + sin = torch.randn_like(cos) + + expected_q = _qk_reference(q, cos, sin) + expected_k = _qk_reference(k, cos, sin) + with torch.inference_mode(): + actual = operators.qk_norm_rope(q, k, (cos, sin), 1e-5) + + assert actual is not None + actual_q, actual_k = actual + assert torch.equal(actual_q, expected_q) + assert torch.equal(actual_k, expected_k) + + +def test_h3_vae_scaled_residual_is_bit_exact(): + device, operators = _selected_operators() + + torch.manual_seed(29) + residual = torch.randn(195, 2048, device=device, dtype=torch.float32) + branch = torch.randn(195, 2048, device=device, dtype=torch.float16) + scale = torch.randn(2048, device=device, dtype=torch.float32) + expected = residual + branch * scale + + with torch.inference_mode(): + actual = operators.scaled_residual(residual, branch, scale) + + assert actual is not None + assert torch.equal(actual, expected) + + +def test_h3_vae_exact_ops_reject_unsupported_inputs(): + from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import ( + try_qk_norm_rope_exact, + ) + from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import try_scaled_residual_exact + + q = torch.randn(1, 2, 32, 64, dtype=torch.float32) + cos = torch.randn(1, 2, 1, 48, dtype=torch.float32) + residual = torch.randn(2, 2048) + branch = torch.randn(2, 2048, dtype=torch.float16) + scale = torch.randn(2048) + with torch.inference_mode(): + assert try_qk_norm_rope_exact(q, q, (cos, cos), 1e-5) is None + assert try_scaled_residual_exact(residual, branch, scale) is None + + +def _make_decoder(): + class Attention(nn.Module): + def __init__(self): + super().__init__() + self.to_qkv = nn.Linear(8, 24) + self.to_out = nn.Linear(8, 8) + self.norm_q = nn.RMSNorm(8, elementwise_affine=False) + self.norm_k = nn.RMSNorm(8, elementwise_affine=False) + self.spatial_parallel = False + self.dim_head = 8 + + def perform_attention(self, query, _key, _value, _pack_info): + return query + + def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None): + return hidden_states + + class FeedForward(nn.Module): + def __init__(self): + super().__init__() + self.w1 = nn.Linear(8, 32) + self.w2 = nn.Linear(16, 8) + self.use_gated = True + self.act_fn = nn.SiLU() + self._compile_forward_enabled = False + self._compile_forward_fatal = False + self._compiled_forward = None + + def forward(self, hidden_states): + hidden_states = self.w1(hidden_states) + gate, hidden_states = hidden_states.chunk(2, dim=-1) + return self.w2(self.act_fn(gate) * hidden_states) + + class Block(nn.Module): + def __init__(self): + super().__init__() + self.attn = Attention() + self.ff = FeedForward() + self.norm1 = nn.RMSNorm(8) + self.norm2 = nn.RMSNorm(8) + self.scale1 = nn.Parameter(torch.zeros(8)) + self.scale2 = nn.Parameter(torch.zeros(8)) + self.use_scale = True + + def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None): + return hidden_states + + decoder = nn.Module() + decoder.transformer_blocks = nn.ModuleList([Block(), Block()]) + decoder.proj_out = nn.Linear(8, 8) + return decoder + + +def test_video_vae_installs_exact_optimizations(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.decoder = torch.nn.Module() + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = FakeModel() + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + }, + ) + monkeypatch.setattr( + vae_module, + "_load_remote_component", + lambda _path, _config: FakeRemote(), + ) + install = Mock(return_value=True) + monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install) + monkeypatch.setattr(vae_module, "PinnedModuleStager", Mock()) + + video_vae = vae_module.MiniMaxH3VideoVAE( + "unused", + device=torch.device("cuda"), + load_device=torch.device("cpu"), + ) + + install.assert_called_once_with( + video_vae.model.decoder, + device=torch.device("cuda"), + ) + + +def test_h3_vae_install_precasts_only_block_linears(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + operators = _operator_set() + monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators) + decoder = _make_decoder() + + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + + for block in decoder.transformer_blocks: + assert block.attn.to_qkv.weight.dtype == torch.float16 + assert block.attn.to_out.weight.dtype == torch.float16 + assert block.ff.w1.weight.dtype == torch.float16 + assert block.ff.w2.weight.dtype == torch.float16 + assert block.forward.__func__.__name__ == "_optimized_transformer_block" + assert block.attn.forward.__func__.__name__ == "_optimized_attention" + assert block.ff.forward.__func__.__name__ == "_optimized_feed_forward" + assert block.attn._omni_qk_norm_rope is operators.qk_norm_rope + assert block._omni_scaled_residual is operators.scaled_residual + assert decoder.proj_out.weight.dtype == torch.float32 + + # Repeated installation is idempotent. + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + + +def test_h3_vae_install_accepts_remote_integer_parallel_flag(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr( + vae_ops, + "resolve_h3_vae_operators", + lambda _device: _operator_set(), + ) + decoder = _make_decoder() + for block in decoder.transformer_blocks: + # The official checkpoint loads this JSON boolean-like field as int 0. + block.attn.spatial_parallel = 0 + + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert getattr(decoder, "_omni_h3_vae_optimizations_installed", False) + + +def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + device, operators = _selected_operators() + monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators) + decoder = _make_decoder().to(device) + feed_forward = decoder.transformer_blocks[0].ff + reference_forward = type(feed_forward).forward + assert vae_ops.install_h3_vae_optimizations( + decoder, + device=device, + ) + + hidden_states = torch.randn(4, 8, device=device, dtype=torch.float32) + with torch.inference_mode(), torch.autocast(device.type, dtype=torch.float16): + expected = reference_forward(feed_forward, hidden_states) + actual = feed_forward(hidden_states) + + assert expected.dtype == torch.float16 + assert torch.equal(actual, expected) + + +def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: None) + decoder = _make_decoder() + + assert not vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed") + for block in decoder.transformer_blocks: + assert block.attn.to_qkv.weight.dtype == torch.float32 + assert block.ff.w1.weight.dtype == torch.float32 + assert block.forward.__func__.__name__ == "forward" + assert block.attn.forward.__func__.__name__ == "forward" + + +@pytest.mark.parametrize("norm_input", [None, _same_dtype_norm_input, _failing_norm_input]) +def test_h3_vae_install_requires_fp32_attention_norm_semantics(monkeypatch, norm_input): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr( + vae_ops, + "resolve_h3_vae_operators", + lambda _device: _operator_set(), + ) + decoder = _make_decoder() + forward_globals = type(decoder.transformer_blocks[0].attn).forward.__globals__ + monkeypatch.setitem(forward_globals, "_vit_norm_input", norm_input) + + assert not vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed") + assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32 + + +@pytest.mark.parametrize( + "break_contract", + [ + lambda block: delattr(block.attn, "spatial_parallel"), + lambda block: setattr(block.attn, "spatial_parallel", 2), + lambda block: setattr(block.attn, "perform_attention", None), + lambda block: setattr(block, "scale1", nn.Parameter(torch.zeros(7))), + lambda block: setattr(block.ff, "_compile_forward_enabled", True), + ], +) +def test_h3_vae_install_rejects_incompatible_remote_contract(monkeypatch, break_contract): + from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops + + monkeypatch.setattr( + vae_ops, + "resolve_h3_vae_operators", + lambda _device: _operator_set(), + ) + decoder = _make_decoder() + break_contract(decoder.transformer_blocks[0]) + + assert not vae_ops.install_h3_vae_optimizations( + decoder, + device=torch.device("meta"), + ) + assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32 + + +def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch + + first = _operator_set(supports=lambda _device: False) + added = _operator_set(supports=lambda device: device.type == "meta") + monkeypatch.setattr(dispatch, "H3_VAE_OPERATOR_TABLE", (first, added)) + + assert dispatch.resolve_h3_vae_operators(torch.device("meta")) is added + assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None + + +def test_h3_vae_dispatch_selects_supported_cuda_capabilities(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch + + platform = Mock() + platform.is_cuda.return_value = True + platform.is_available.return_value = True + monkeypatch.setattr(dispatch, "HAS_TRITON", True) + monkeypatch.setattr(dispatch, "current_omni_platform", platform) + + for capability in (90, 100, 103): + platform.get_device_capability.return_value.to_int.return_value = capability + assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is not None + + for capability in (89, 101, 110): + platform.get_device_capability.return_value.to_int.return_value = capability + assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is None diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py new file mode 100644 index 000000000..ef04019f6 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CPU tests for the decoder tile-shortage guard in the MiniMax-H3 video VAE.""" + +import sys +import types + +import pytest +import torch + +from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +RATIO = 16 +# The checkpoint keeps two independent grids -- ``split_tiles`` selects +# ``decoder_tile_size``/``decoder_tile_overlap_min`` when ``is_decoder`` is set +# and ``tile_size``/``tile_overlap_min`` otherwise. The shipped H3 config leaves +# both at 256/64, so the two happen to coincide there, but that is a config +# coincidence. The double deliberately gives them different values so that +# passing the wrong flag changes the tile count and fails a test. +TILE_SIZE, OVERLAP = 256, 64 # decoder, matching the shipped config +ENCODER_TILE_SIZE, ENCODER_OVERLAP = 128, 32 + + +class _FakeCheckpointModel: + """The parts of the checkpoint's AutoencoderKL the guard touches.""" + + def __init__(self): + self.vae_ratio = RATIO + self.parallel_tiling = True + self.decoded_with = [] + + def split_tiles(self, input_len, is_decoder=False): + """Same grid arithmetic as the checkpoint, on pixel dimensions.""" + tile = TILE_SIZE if is_decoder else ENCODER_TILE_SIZE + overlap = OVERLAP if is_decoder else ENCODER_OVERLAP + if tile >= input_len: + return [0], [input_len], [] + n = -(-input_len // tile) + while tile * n - overlap * (n - 1) - input_len < 0: + n += 1 + return list(range(n)), [tile] * n, [overlap] * (n - 1) + + +def _vae(parallel_size, state): + """A stand-in instance: the guard only reads .model/.remote/.parallel_size.""" + module = types.ModuleType("fake_ckpt.parallel") + module.get_parallel_state = lambda: state + sys.modules.setdefault("fake_ckpt", types.ModuleType("fake_ckpt")) + sys.modules["fake_ckpt.parallel"] = module + + remote = type("Remote", (), {"__module__": "fake_ckpt.klvae"})() + + vae = object.__new__(MiniMaxH3VideoVAE) + vae.model = _FakeCheckpointModel() + vae.remote = remote + vae.parallel_size = parallel_size + return vae + + +def _latent(h, w): + return torch.zeros(1, 24, 9, h, w) + + +@pytest.mark.parametrize( + ("h", "w", "expected"), + [ + (16, 16, 1), # 256x256 px -> single tile + (24, 16, 2), # 384x256 px + (24, 24, 4), # 384x384 px + (48, 84, 28), # 768x1344 px, the shipped H3 request + (96, 168, 112), # 1536x2688 px, a larger canvas the request path also allows + ], +) +def test_tile_count_matches_the_checkpoint_grid(h, w, expected): + assert MiniMaxH3VideoVAE._decoder_tile_count(_vae(1, {}), _latent(h, w)) == expected + + +@pytest.mark.parametrize(("h", "w", "decoder_tiles", "encoder_tiles"), [(16, 16, 1, 9), (24, 24, 4, 16)]) +def test_tile_count_uses_the_decoder_grid_not_the_encoder_one(h, w, decoder_tiles, encoder_tiles): + """Passing ``is_decoder=False`` would silently change the count, so pin it.""" + vae = _vae(1, {}) + assert MiniMaxH3VideoVAE._decoder_tile_count(vae, _latent(h, w)) == decoder_tiles + + px_h, px_w = h * RATIO, w * RATIO + as_encoder = len(vae.model.split_tiles(px_h, False)[0]) * len(vae.model.split_tiles(px_w, False)[0]) + assert as_encoder == encoder_tiles != decoder_tiles + + +def test_rank_local_tiling_restores_the_group_state(): + state = {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"} + vae = _vae(4, state) + + with vae._rank_local_tiling(): + assert state["sp_size"] == 1 + assert state["sp_rank"] == 0 + assert state["sp_enabled"] is False + assert state["sp_process_group"] is None + assert vae.model.parallel_tiling is False + + assert state == {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"} + assert vae.model.parallel_tiling is True + + +def test_rank_local_tiling_restores_after_an_exception(): + state = {"sp_size": 4, "sp_rank": 1, "sp_enabled": True, "sp_process_group": "pg"} + vae = _vae(4, state) + + with pytest.raises(RuntimeError, match="decode failed"): + with vae._rank_local_tiling(): + raise RuntimeError("decode failed") + + assert state["sp_size"] == 4 + assert vae.model.parallel_tiling is True + + +# --------------------------------------------------------------------------- +# Dispatch: every one of these calls ``decode_latent``, so deleting the guard +# turns the whole matrix red rather than leaving it green. +# --------------------------------------------------------------------------- + + +class _RecordingProcessor: + @staticmethod + def revert_tensor(decoded): + return decoded + + +def _dispatch_vae(parallel_size): + """A VAE whose ``decode_base`` records the parallel state it ran under.""" + state = { + "group_size": parallel_size, + "group_rank": 0, + "local_process_group": object(), + "sp_size": parallel_size, + "sp_rank": 0, + "sp_enabled": True, + "sp_process_group": object(), + "tp_size": 1, + "tp_rank": 0, + } + vae = _vae(parallel_size, state) + vae.config_dict = { + "latent_channels": 24, + "latents_mean": [0.0] * 24, + "latents_std": [1.0] * 24, + } + vae.model.processor = _RecordingProcessor() + + seen = {} + + def decode_base(sample): + seen["sp_size"] = state["sp_size"] + seen["sp_enabled"] = state["sp_enabled"] + seen["sp_process_group"] = state["sp_process_group"] + seen["parallel_tiling"] = vae.model.parallel_tiling + # the guard never inspects the decoded shape, so keep it small + return torch.zeros(1, 3, 2, 4, 4) + + vae.model.decode_base = decode_base + return vae, state, seen + + +@pytest.mark.parametrize( + ("parallel_size", "h", "w", "falls_back"), + [ + (4, 24, 16, True), # 2 tiles, ranks 2 and 3 would get none + (4, 16, 16, True), # 1 tile + (4, 24, 24, False), # 4 tiles, exactly enough + (4, 48, 84, False), # 28 tiles, the shipped 768x1344 request + (4, 96, 168, False), # 112 tiles, 1536x2688 + (2, 24, 16, False), # 2 tiles across 2 ranks + (1, 16, 16, False), # never parallel, nothing to guard + ], +) +def test_decode_latent_falls_back_exactly_when_tiles_are_short(parallel_size, h, w, falls_back): + """Fewer tiles than ranks is the hang condition; equal or more is fine.""" + vae, _, seen = _dispatch_vae(parallel_size) + + MiniMaxH3VideoVAE.decode_latent(vae, _latent(h, w)) + + # only the guard clears ``parallel_tiling``, so it is the signal that it fired + assert (seen["parallel_tiling"] is False) is falls_back + assert seen["sp_size"] == (1 if falls_back else parallel_size) + + +def test_decode_latent_isolates_and_restores_the_whole_group_state(): + """1 tile, 4 ranks: ``decode_base`` must observe rank-local tiling.""" + vae, state, seen = _dispatch_vae(4) + + frames = MiniMaxH3VideoVAE.decode_latent(vae, _latent(16, 16)) + + assert seen["sp_size"] == 1 + assert seen["sp_enabled"] is False + assert seen["sp_process_group"] is None + assert seen["parallel_tiling"] is False + assert frames.ndim == 5 and frames.dtype is torch.float32 + # the group state is restored once decode returns + assert state["sp_size"] == 4 and state["sp_enabled"] is True + assert vae.model.parallel_tiling is True + + +def test_decode_latent_leaves_the_group_state_alone_when_the_grid_is_large_enough(): + """The shipped 768x1344 request is 28 tiles: the guard must not fire.""" + vae, state, seen = _dispatch_vae(4) + + MiniMaxH3VideoVAE.decode_latent(vae, _latent(48, 84)) + + assert seen["sp_size"] == 4 + assert seen["sp_enabled"] is True + assert seen["sp_process_group"] is state["sp_process_group"] + assert seen["parallel_tiling"] is True diff --git a/vllm_omni/model_executor/models/ming_tts/fm/__init__.py b/tests/diffusion/models/pi0/__init__.py similarity index 100% rename from vllm_omni/model_executor/models/ming_tts/fm/__init__.py rename to tests/diffusion/models/pi0/__init__.py diff --git a/tests/diffusion/models/pi0/test_pi0_parity.py b/tests/diffusion/models/pi0/test_pi0_parity.py new file mode 100644 index 000000000..3b6358b50 --- /dev/null +++ b/tests/diffusion/models/pi0/test_pi0_parity.py @@ -0,0 +1,366 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""π0 LeRobot parity (in-process): the bit-for-bit correctness oracle. + +Verifies that vllm-omni's ``Pi0ForActionPrediction`` produces bit-for-bit +matching action chunks with LeRobot's ``PI0Policy`` when fed: + * the same weights (``lerobot/pi0_base``) + * the same pre-processed inputs (images, masks, tokens, state) + * the same initial noise tensor + +π0 is a flow-matching model (Euler-integrated ODE from t=1 → t=0 with a fixed +``num_steps``), so the output is deterministic once the noise is fixed and +``torch.allclose`` on the final action chunk is a valid oracle. This is the +authoritative correctness oracle for the vllm-omni π0 port (max|Δ| < 1e-4). + +The OpenPI websocket online-serving e2e lives separately in +``tests/e2e/online_serving/test_pi0_expansion.py``. + +Run in a SEPARATE ``lerobot[pi]`` venv (avoids dep conflict with the vllm-omni +env), with the vllm-omni ``pi0`` package importable:: + + python -m pytest tests/diffusion/models/pi0/test_pi0_parity.py -v -s + +Skipped automatically when LeRobot is not installed (e.g. the vllm-omni env). +The run is CPU/float32 with fixed defaults; the only override is +``PI0_PARITY_MODEL_PATH`` (a local pi0_base dir in LeRobot format, to skip the +HF download; defaults to ``lerobot/pi0_base``). +""" + +from __future__ import annotations + +import copy +import importlib.util +import os + +import pytest +import torch + +# local_model: needs real weights + a lerobot venv (transformers 5.3.0), so it +# runs locally rather than in the ready-CI (see docs/contributing/ci/CI_5levels.md). +# Additionally gated on lerobot being importable. +_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None + +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion] + + +# ─── Config (fixed; matches LeRobot defaults for ``lerobot/pi0_base``) ── +DEVICE = "cpu" +DTYPE_STR = "float32" +ATOL = 1e-4 +NUM_STEPS = 10 +BATCH_SIZE = 2 +ACTION_DIM = 32 +STATE_DIM = 32 +ACTION_HORIZON = 50 +MAX_TOKEN_LEN = 48 + +# The only knob: point at a local pi0_base dir (LeRobot format) to skip the HF +# download; defaults to the HF repo id. +MODEL_PATH = os.environ.get("PI0_PARITY_MODEL_PATH", "lerobot/pi0_base") + + +def _resolve_checkpoint_dir() -> str: + """Return a local dir containing the pi0_base checkpoint (download if needed).""" + if os.path.isdir(MODEL_PATH): + return MODEL_PATH + from huggingface_hub import snapshot_download + + return snapshot_download(repo_id=MODEL_PATH, repo_type="model") + + +# ─── Dummy dataset stats (identity) ────────────────────────────────── +def _dummy_dataset_stats() -> dict: + return { + "observation.state": { + "mean": torch.zeros(STATE_DIM), + "std": torch.ones(STATE_DIM), + "q01": torch.zeros(STATE_DIM), + "q99": torch.ones(STATE_DIM), + }, + "action": { + "mean": torch.zeros(ACTION_DIM), + "std": torch.ones(ACTION_DIM), + "q01": torch.zeros(ACTION_DIM), + "q99": torch.ones(ACTION_DIM), + }, + "images": { + cam: { + "mean": torch.zeros(3, 224, 224), + "std": torch.ones(3, 224, 224), + "q01": torch.zeros(3, 224, 224), + "q99": torch.ones(3, 224, 224), + } + for cam in ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb") + }, + } + + +def _create_dummy_batch(batch_size: int = BATCH_SIZE, device: str = DEVICE) -> dict: + """Reproducible dummy inputs — identical across both implementations.""" + g = torch.Generator(device="cpu").manual_seed(0) + prompt = "Pick up the red block and place it in the bin" + return { + "observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device), + "action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device), + "observation.images.base_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to( + device + ), + "observation.images.left_wrist_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to( + device + ), + "observation.images.right_wrist_0_rgb": torch.rand( + batch_size, 3, 224, 224, generator=g, dtype=torch.float32 + ).to(device), + "task": [prompt for _ in range(batch_size)], + } + + +# ─── LeRobot instantiation ──────────────────────────────────────────── +def _instantiate_lerobot(): + from lerobot.policies.pi0 import PI0Policy + from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors + + policy = PI0Policy.from_pretrained(MODEL_PATH, strict=True) + policy.to(DEVICE) + policy.config.device = DEVICE + policy.eval() + + pre, post = make_pi0_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats()) + return policy, pre, post + + +# ─── vllm-omni instantiation ────────────────────────────────────────── +def _instantiate_vllm_omni(): + """Build the vllm-omni π0 model in isolation (no pipeline, no engine).""" + from vllm_omni.diffusion.models.pi0 import Pi0Config, Pi0ForActionPrediction + + cfg = Pi0Config( + max_action_dim=ACTION_DIM, + max_state_dim=STATE_DIM, + chunk_size=ACTION_HORIZON, + num_inference_steps=NUM_STEPS, + dtype=DTYPE_STR, + ) + model = Pi0ForActionPrediction(cfg) + model.to(DEVICE).eval() + _load_lerobot_weights(model) + return model + + +def _load_lerobot_weights(model): + """Feed the ``lerobot/pi0_base`` safetensors into the vllm-omni model. + + ``Pi0ForActionPrediction.load_weights`` handles the leading ``model.`` + prefix and the flat→nested / lm_head→embed_tokens remaps, so we can pass + the raw checkpoint dict. + """ + import safetensors.torch + + cache_dir = _resolve_checkpoint_dir() + path = os.path.join(cache_dir, "model.safetensors") + state = safetensors.torch.load_file(path) + model.load_weights(list(state.items())) + + +# ─── Helpers to extract LeRobot's pre-processed inputs ──────────────── +def _extract_lerobot_model_inputs(lerobot_policy, processed_batch): + """Mimic what ``PI0Policy.predict_action_chunk`` feeds into + ``self.model.sample_actions``. We use these *exact* tensors for vllm-omni + so any divergence must come from the core model, not preprocessing. + """ + images, img_masks = lerobot_policy._preprocess_images(processed_batch) + from lerobot.utils.constants import ( + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + ) + + lang_tokens = processed_batch[OBS_LANGUAGE_TOKENS] + lang_masks = processed_batch[OBS_LANGUAGE_ATTENTION_MASK] + state = lerobot_policy.prepare_state(processed_batch) + return images, img_masks, lang_tokens, lang_masks, state + + +# ─── Shared fixed-noise sampler ─────────────────────────────────────── +def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor: + g = torch.Generator(device="cpu").manual_seed(42) + return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device) + + +# ─── Main test ──────────────────────────────────────────────────────── +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +def test_pi0_vllm_omni_vs_lerobot(): + print("\n[parity] Instantiating LeRobot…") + lerobot_policy, lerobot_pre, _ = _instantiate_lerobot() + + print("[parity] Instantiating vllm-omni…") + omni_model = _instantiate_vllm_omni() + + print("[parity] Preparing shared inputs…") + raw_batch = _create_dummy_batch() + processed_batch = lerobot_pre(copy.deepcopy(raw_batch)) + images, img_masks, lang_tokens, lang_masks, state = _extract_lerobot_model_inputs(lerobot_policy, processed_batch) + noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE) + + print(f"[parity] state.shape={state.shape} lang_tokens.shape={lang_tokens.shape}") + print(f"[parity] images[0].shape={images[0].shape} (num_cams={len(images)})") + print(f"[parity] noise.shape={noise.shape} dtype={noise.dtype}") + + # ── LeRobot forward ── + print("[parity] Running LeRobot sample_actions…") + with torch.no_grad(): + lerobot_actions = lerobot_policy.model.sample_actions( + images, + img_masks, + lang_tokens, + lang_masks, + state, + noise=noise, + num_steps=NUM_STEPS, + ) + print( + f"[parity] LeRobot actions: shape={lerobot_actions.shape} " + f"mean={lerobot_actions.mean().item():.6f} std={lerobot_actions.std().item():.6f}" + ) + + # ── vllm-omni forward ── + print("[parity] Running vllm-omni sample_actions…") + with torch.no_grad(): + omni_actions = omni_model.sample_actions( + images=images, + image_masks=img_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + state=state, + noise=noise, + num_steps=NUM_STEPS, + ) + print( + f"[parity] vllm-omni actions: shape={omni_actions.shape} " + f"mean={omni_actions.mean().item():.6f} std={omni_actions.std().item():.6f}" + ) + + # ── Compare ── + diff = (lerobot_actions.float() - omni_actions.float()).abs() + max_diff = diff.max().item() + mean_diff = diff.mean().item() + print(f"[parity] |Δ| max={max_diff:.2e} mean={mean_diff:.2e} atol={ATOL:.1e}") + close = torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL) + print(f"[parity] torch.allclose(atol={ATOL}): {close}") + + if not close: + print("\n[parity] ⚠️ Outputs diverge — running per-stage diagnostics…") + _diagnose_divergence( + lerobot_policy.model, + omni_model, + images, + img_masks, + lang_tokens, + lang_masks, + state, + noise, + ) + + assert close, ( + f"vllm-omni vs LeRobot actions differ beyond atol={ATOL}. max_diff={max_diff:.2e} mean_diff={mean_diff:.2e}" + ) + + +# ─── Per-stage divergence diagnostics ───────────────────────────────── +@torch.no_grad() +def _diagnose_divergence( + lerobot_flow_model, + omni_model, + images, + img_masks, + lang_tokens, + lang_masks, + state, + noise, +): + """Localize a numerical mismatch to a specific pipeline stage: + 1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang. + 2. Prefix KV cache layer 0 — PaliGemma LM attention. + 3. A single denoise_step velocity at t=1.0 — action expert forward. + """ + from vllm_omni.diffusion.models.pi0.modeling_pi0 import ( + make_att_2d_masks, + prepare_attention_masks_4d, + ) + + # ── Stage 1: prefix embeddings ── + lr_prefix_embs, lr_prefix_pad, lr_prefix_att = lerobot_flow_model.embed_prefix( + images, img_masks, lang_tokens, lang_masks + ) + sg_prefix_embs, sg_prefix_pad, sg_prefix_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks) + total_diff = (lr_prefix_embs.float() - sg_prefix_embs.float()).abs().max().item() + print(f"[diag] prefix_embs max |Δ| = {total_diff:.2e} (shape={tuple(sg_prefix_embs.shape)})") + print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_prefix_pad, sg_prefix_pad)}") + print(f"[diag] prefix_att_masks equal: {torch.equal(lr_prefix_att.bool(), sg_prefix_att.bool())}") + + num_cams = len(images) + img_len = 256 * num_cams + lang_len = lr_prefix_embs.shape[1] - img_len + img_diff = (lr_prefix_embs[:, :img_len].float() - sg_prefix_embs[:, :img_len].float()).abs().max().item() + lang_diff = (lr_prefix_embs[:, img_len:].float() - sg_prefix_embs[:, img_len:].float()).abs().max().item() + print(f"[diag] image slice [:{img_len}] max|Δ| = {img_diff:.2e} (num_cams={num_cams})") + print(f"[diag] lang slice [{img_len}:] max|Δ| = {lang_diff:.2e} (len={lang_len})") + + def _stats(name, t): + t = t.float() + print( + f"[diag] {name} prefix_embs: mean={t.mean().item():+.4f} " + f"std={t.std().item():.4f} min={t.min().item():+.2f} " + f"max={t.max().item():+.2f}" + ) + + _stats("LeRobot ", lr_prefix_embs) + _stats("vllm-omni", sg_prefix_embs) + + # ── Stage 2: prefix KV cache ── + prefix_att_2d = make_att_2d_masks(sg_prefix_pad, sg_prefix_att) + prefix_pos = torch.cumsum(sg_prefix_pad, dim=1) - 1 + prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d) + + lerobot_flow_model.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager" + _, lr_kv = lerobot_flow_model.paligemma_with_expert.forward( + attention_mask=prefix_att_4d, + position_ids=prefix_pos, + past_key_values=None, + inputs_embeds=[lr_prefix_embs, None], + use_cache=True, + ) + _, sg_kv = omni_model.paligemma_with_expert.forward( + attention_mask=prefix_att_4d, + position_ids=prefix_pos, + past_key_values=None, + inputs_embeds=[sg_prefix_embs, None], + use_cache=True, + ) + + def _layer0_kv(cache): + if isinstance(cache, list): + return cache[0] + layer0 = cache.layers[0] + return layer0.keys, layer0.values + + try: + lr_k, lr_v = _layer0_kv(lr_kv) + sg_k, sg_v = _layer0_kv(sg_kv) + dk = (lr_k.float() - sg_k.float()).abs().max().item() + dv = (lr_v.float() - sg_v.float()).abs().max().item() + print(f"[diag] prefix KV layer0 K max|Δ|={dk:.2e} V max|Δ|={dv:.2e}") + except Exception as e: # noqa: BLE001 + print(f"[diag] could not extract prefix KV for comparison: {e}") + + # ── Stage 3: a single denoise_step at t=1.0 ── + bsize = state.shape[0] + t = torch.ones(bsize, dtype=torch.float32, device=state.device) + lr_vt = lerobot_flow_model.denoise_step(state, sg_prefix_pad, lr_kv, noise, t) + sg_vt = omni_model.denoise_step(state, sg_prefix_pad, sg_kv, noise, t) + print(f"[diag] denoise_step(t=1) v_t max|Δ| = {(lr_vt.float() - sg_vt.float()).abs().max().item():.2e}") + + +if __name__ == "__main__": + test_pi0_vllm_omni_vs_lerobot() diff --git a/tests/diffusion/models/pi0/test_pi0_units.py b/tests/diffusion/models/pi0/test_pi0_units.py new file mode 100644 index 000000000..a1fa18f1a --- /dev/null +++ b/tests/diffusion/models/pi0/test_pi0_units.py @@ -0,0 +1,585 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU unit tests for the π0 (Pi-Zero) VLA model. + +Deliberately lightweight: pure utility code (config plumbing, mask building, +sinusoidal embeddings, normalization, image preprocessing, weight-load remap) +that runs on CPU with no weights and no lerobot. The full LeRobot parity check +lives in ``tests/diffusion/models/pi0/test_pi0_parity.py`` and requires a model download. + + pytest tests/diffusion/models/pi0/test_pi0_units.py -v +""" + +from __future__ import annotations + +import logging +import math + +import pytest +import torch + +from vllm_omni.diffusion.models.pi0.config import Pi0Config +from vllm_omni.diffusion.models.pi0.modeling_pi0 import ( + OPENPI_ATTENTION_MASK_VALUE, + Pi0ForActionPrediction, + _apply_norm, + _build_norm_buffers, + create_sinusoidal_pos_embedding, + get_gemma_config, + make_att_2d_masks, + prepare_attention_masks_4d, +) +from vllm_omni.diffusion.models.pi0.processor_pi0 import ( + Pi0ImageProcessor, + build_model_inputs, + pil_image_to_tensor, + resize_with_pad, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# ---------------------------------------------------------------------------- +# Pi0Config dataclass + checkpoint resolver +# ---------------------------------------------------------------------------- +_LEROBOT_CFG = { + "type": "pi0", + "paligemma_variant": "gemma_2b", + "action_expert_variant": "gemma_300m", + "chunk_size": 50, + "max_action_dim": 32, + "max_state_dim": 32, + "num_inference_steps": 10, + "image_resolution": [224, 224], + "tokenizer_max_length": 48, + "dtype": "float32", + # Training-only keys that must be filtered out of the dataclass. + "optimizer_lr": 2.5e-5, + "scheduler_warmup_steps": 1000, + "input_features": { + "observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]}, + "observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]}, + "observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]}, + "observation.state": {"type": "STATE", "shape": [32]}, + }, +} + +_EXPECTED_CAMERA_ORDER = [ + "observation.images.base_0_rgb", + "observation.images.left_wrist_0_rgb", + "observation.images.right_wrist_0_rgb", +] + + +def test_config_camera_order_from_input_features(): + c = Pi0Config.from_model_config(_LEROBOT_CFG) + assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER + + +def test_config_filters_training_only_keys(): + c = Pi0Config.from_model_config(_LEROBOT_CFG) + assert not hasattr(c, "optimizer_lr") + assert not hasattr(c, "scheduler_warmup_steps") + + +def test_config_coerces_image_resolution_to_tuple(): + c = Pi0Config.from_model_config(_LEROBOT_CFG) + assert c.image_resolution == (224, 224) + assert isinstance(c.image_resolution, tuple) + + +def test_config_non_square_resolution_raises(): + with pytest.raises(ValueError): + Pi0Config(image_resolution=(224, 256)) + + +@pytest.mark.parametrize("bad", [(224,), (224, 224, 3), 224]) +def test_config_malformed_resolution_raises(bad): + with pytest.raises(ValueError): + Pi0Config(image_resolution=bad) + + +def test_config_defaults_and_empty(): + c = Pi0Config.from_model_config(None) + assert c.chunk_size == 50 + assert c.max_action_dim == 32 + assert c.max_state_dim == 32 + assert c.num_inference_steps == 10 + assert c.image_resolution == (224, 224) + assert c.tokenizer_max_length == 48 + assert c.max_cameras == 3 + assert c.norm_stats is None + assert c.image_key_map == {} + + +def test_config_deploy_yaml_override(): + """Deploy-yaml model_config is authoritative (e.g. force dtype/variant).""" + c = Pi0Config.from_model_config({**_LEROBOT_CFG, "dtype": "bfloat16", "num_inference_steps": 4}) + assert c.dtype == "bfloat16" + assert c.num_inference_steps == 4 + + +# ---------------------------------------------------------------------------- +# Kernel utilities: sinusoidal timestep embedding +# ---------------------------------------------------------------------------- +def test_sinusoidal_embedding_shape(): + emb = create_sinusoidal_pos_embedding(torch.tensor([0.0, 0.5, 1.0]), dimension=64) + assert emb.shape == (3, 64) + + +def test_sinusoidal_embedding_values_at_zero(): + emb = create_sinusoidal_pos_embedding(torch.tensor([0.0]), dimension=4) + assert torch.allclose(emb[0, :2].float(), torch.zeros(2), atol=1e-6) + assert torch.allclose(emb[0, 2:].float(), torch.ones(2), atol=1e-6) + + +def test_sinusoidal_embedding_odd_dim_raises(): + with pytest.raises(ValueError): + create_sinusoidal_pos_embedding(torch.tensor([1.0]), dimension=3) + + +def test_sinusoidal_embedding_non_1d_raises(): + with pytest.raises(ValueError): + create_sinusoidal_pos_embedding(torch.tensor([[0.0, 0.5]]), dimension=4) + + +def test_sinusoidal_embedding_distinct_timesteps_differ(): + emb = create_sinusoidal_pos_embedding(torch.tensor([0.1, 0.2, 0.5, 0.9]), dimension=16) + for i in range(emb.shape[0]): + for j in range(i + 1, emb.shape[0]): + assert not torch.allclose(emb[i], emb[j]) + + +# ---------------------------------------------------------------------------- +# Kernel utilities: attention mask construction +# ---------------------------------------------------------------------------- +def test_make_att_2d_masks_bidirectional(): + result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.zeros(1, 4, dtype=torch.long)) + assert result.all() + + +def test_make_att_2d_masks_causal(): + result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.ones(1, 4, dtype=torch.long)) + expected = torch.tril(torch.ones(4, 4, dtype=torch.bool)) + assert torch.equal(result[0], expected) + + +def test_make_att_2d_masks_prefix_lm(): + result = make_att_2d_masks(torch.ones(1, 5, dtype=torch.bool), torch.tensor([[0, 0, 0, 1, 1]], dtype=torch.long)) + assert result[0, 0, 0] + assert not result[0, 0, 3] + assert result[0, 3, 0] + assert not result[0, 3, 4] + + +def test_make_att_2d_masks_respects_padding(): + pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.bool) # last token is padding + att = torch.zeros(1, 4, dtype=torch.long) # full bidirectional + result = make_att_2d_masks(pad, att) + assert not result[0, :, 3].any(), "Nothing should attend to padded key" + assert not result[0, 3, :].any(), "Padded query should attend to nothing" + + +def test_prepare_attention_masks_4d(): + mask_2d = torch.tensor([[[True, False], [True, True]]]) + mask_4d = prepare_attention_masks_4d(mask_2d) + assert mask_4d.shape == (1, 1, 2, 2) + assert mask_4d[0, 0, 0, 0].item() == 0.0 + assert mask_4d[0, 0, 1, 1].item() == 0.0 + expected = torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=mask_4d.dtype) + assert torch.equal(mask_4d[0, 0, 0, 1], expected) + + +def test_attention_mask_value_finite_and_negative(): + assert math.isfinite(OPENPI_ATTENTION_MASK_VALUE) + assert OPENPI_ATTENTION_MASK_VALUE < -1e30 + + +# ---------------------------------------------------------------------------- +# Normalization round-trips +# ---------------------------------------------------------------------------- +def test_build_norm_buffers_missing(): + assert _build_norm_buffers(None, "state") is None + assert _build_norm_buffers({}, "state") is None + assert _build_norm_buffers({"state": None}, "state") is None + + +def test_build_norm_buffers_mean_std(): + stats = {"state": {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 0.5]}} + out = _build_norm_buffers(stats, "state") + assert out["mode"] == "mean_std" + assert torch.equal(out["mean"], torch.tensor([1.0, 2.0])) + assert torch.equal(out["std"], torch.tensor([0.5, 0.5])) + + +def test_build_norm_buffers_min_max(): + stats = {"action": {"mode": "min_max", "min": [-1.0, 0.0], "max": [1.0, 2.0]}} + out = _build_norm_buffers(stats, "action") + assert out["mode"] == "min_max" + assert torch.equal(out["min"], torch.tensor([-1.0, 0.0])) + assert torch.equal(out["max"], torch.tensor([1.0, 2.0])) + + +def test_apply_norm_mean_std_forward_inverse(): + stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.5, -0.2], "std": [2.0, 0.1]}}, "s") + x = torch.tensor([[1.0, 0.3]]) + z = _apply_norm(x, stats, inverse=False) + assert torch.allclose(z, torch.tensor([[0.25, 5.0]]), atol=1e-5) + x_back = _apply_norm(z, stats, inverse=True) + assert torch.allclose(x_back, x, atol=1e-5) + + +def test_apply_norm_min_max_forward_inverse(): + stats = _build_norm_buffers({"s": {"mode": "min_max", "min": [0.0, -1.0], "max": [1.0, 1.0]}}, "s") + x = torch.tensor([[0.25, 0.5]]) + z = _apply_norm(x, stats, inverse=False) + assert torch.allclose(z, torch.tensor([[-0.5, 0.5]]), atol=1e-5) + x_back = _apply_norm(z, stats, inverse=True) + assert torch.allclose(x_back, x, atol=1e-5) + + +def test_apply_norm_preserves_padded_tail(): + stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.0, 0.0], "std": [1.0, 1.0]}}, "s") + x = torch.tensor([[1.0, 2.0, 99.0, -99.0]]) + z = _apply_norm(x, stats, inverse=False) + assert torch.allclose(z[:, :2], x[:, :2]) + assert torch.allclose(z[:, 2:], x[:, 2:]) + + +def test_apply_norm_noop_when_stats_none(): + x = torch.randn(2, 5) + assert torch.equal(_apply_norm(x, None, inverse=False), x) + assert torch.equal(_apply_norm(x, None, inverse=True), x) + + +# ---------------------------------------------------------------------------- +# Gemma variant configs +# ---------------------------------------------------------------------------- +def test_gemma_2b_dims(): + c = get_gemma_config("gemma_2b") + assert c.width == 2048 + assert c.depth == 18 + + +def test_gemma_300m_dims(): + c = get_gemma_config("gemma_300m") + assert c.width == 1024 + assert c.depth == 18 + + +def test_gemma_unknown_variant_raises(): + with pytest.raises(ValueError): + get_gemma_config("gemma_7b") + + +# ---------------------------------------------------------------------------- +# load_weights remap rules (tiny model; needs torch + transformers) +# ---------------------------------------------------------------------------- +def _tiny_pi0_model(): + """Construct a tiny π0 model for remap tests (full Gemma depth but small + action/state dims — the PaliGemma backbone is always full-size).""" + cfg = Pi0Config(max_action_dim=8, max_state_dim=8, chunk_size=2) + return Pi0ForActionPrediction(cfg) + + +@pytest.fixture(autouse=True) +def _silence_pi0_loader(caplog): + # The remap tests feed a single synthetic weight; the loader's + # "missing params" warning is expected noise here. + logging.getLogger("vllm_omni.diffusion.models.pi0.modeling_pi0").setLevel(logging.ERROR) + yield + + +@pytest.mark.slow +def test_remap_paligemma_submodules_to_nested_layout(): + model = _tiny_pi0_model() + params = dict(model.named_parameters()) + target = next(k for k in params if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.")) + flat_key = "model." + target.replace( + "paligemma_with_expert.paligemma.model.vision_tower.", + "paligemma_with_expert.paligemma.vision_tower.", + ) + payload = torch.full(params[target].shape, 0.12345) + before = params[target].detach().clone() + model.load_weights([(flat_key, payload)]) + after = dict(model.named_parameters())[target] + assert not torch.equal(before, after), "flat→nested remap broken" + assert torch.allclose(after, payload) + + +@pytest.mark.slow +def test_lm_head_remap_to_embed_tokens(): + model = _tiny_pi0_model() + params = dict(model.named_parameters()) + target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + assert target in params + payload = torch.full(params[target].shape, 0.5) + before = params[target].detach().clone() + model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)]) + after = dict(model.named_parameters())[target] + assert not torch.equal(before, after), "lm_head→embed_tokens remap broken" + assert torch.allclose(after, payload) + + +@pytest.mark.slow +def test_vision_tower_version_robust_remap(): + """transformers >=5.4 flattens SigLIP from vision_tower.vision_model.* to + vision_tower.*; the lerobot/pi0_base checkpoint uses the vision_model infix. + load_weights must adapt so vision weights aren't silently dropped (which + would leave the vision encoder at random init).""" + model = _tiny_pi0_model() + params = dict(model.named_parameters()) + # A representative vision-tower param as it exists on the model. + target = next( + k + for k in params + if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.") and "patch_embedding.weight" in k + ) + vt = "paligemma_with_expert.paligemma.model.vision_tower." + rest = target[len(vt) :] + # Build the checkpoint key in whichever layout differs from the model's. + if rest.startswith("vision_model."): + ckpt_key = vt + rest[len("vision_model.") :] + else: + ckpt_key = vt + "vision_model." + rest + payload = torch.full(params[target].shape, 0.073) + before = params[target].detach().clone() + model.load_weights([("model." + ckpt_key, payload)]) + after = dict(model.named_parameters())[target] + assert not torch.equal(before, after), "vision_tower version-robust remap broken" + assert torch.allclose(after, payload) + + +# ---------------------------------------------------------------------------- +# Image processor +# ---------------------------------------------------------------------------- +def test_resize_with_pad_square(): + out = resize_with_pad(torch.ones(1, 3, 224, 224), 224, 224) + assert out.shape == (1, 3, 224, 224) + + +def test_resize_with_pad_wide(): + """Wide source (H < W): padding applied to top/bottom (rows).""" + out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224) + assert out.shape == (1, 3, 224, 224) + assert out[0, 0, 0, 112].item() == -1.0 # top row is padded + + +def test_resize_with_pad_tall(): + """Tall source (H > W): padding applied to left/right (cols).""" + out = resize_with_pad(torch.zeros(1, 3, 400, 100), 224, 224) + assert out.shape == (1, 3, 224, 224) + assert out[0, 0, 112, 0].item() == -1.0 # left-most column padded + + +def test_resize_with_pad_rejects_non_4d(): + with pytest.raises(ValueError): + resize_with_pad(torch.zeros(3, 100, 200), 224, 224) + + +def test_pil_to_tensor(): + import numpy as np + from PIL import Image + + img = Image.fromarray(np.full((10, 10, 3), 128, dtype=np.uint8)) + t = pil_image_to_tensor(img) + assert t.shape == (1, 3, 10, 10) + expected = 128.0 / 255.0 * 2.0 - 1.0 + assert abs(t[0, 0, 5, 5].item() - expected) < 1e-4 + + +def test_pil_to_tensor_rgba_to_rgb(): + import numpy as np + from PIL import Image + + rgba = np.full((8, 8, 4), 200, dtype=np.uint8) + img = Image.fromarray(rgba, mode="RGBA") + t = pil_image_to_tensor(img) + assert t.shape == (1, 3, 8, 8) + + +def test_make_empty_image(): + proc = Pi0ImageProcessor(image_size=224) + empty = proc.make_empty_image() + assert empty.shape == (1, 3, 224, 224) + assert (empty == -1.0).all() + + +def test_preprocess_single_resizes_ndarray(): + import numpy as np + + proc = Pi0ImageProcessor(image_size=224) + arr = np.full((120, 120, 3), 100, dtype=np.uint8) + t = proc.preprocess_single(arr) + assert t.shape == (1, 3, 224, 224) + + +# ---------------------------------------------------------------------------- +# build_model_inputs: camera ordering, padding, masks, state +# ---------------------------------------------------------------------------- +class _FakeTokenizer: + """Minimal stand-in for the PaliGemma tokenizer (no download).""" + + def __call__( + self, text, padding=None, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None + ): + ids = [2] + [ord(c) % 100 for c in text][: max_length - 1] + ids = ids[:max_length] + attn = [1] * len(ids) + while len(ids) < max_length: + ids.append(0) + attn.append(0) + return {"input_ids": ids, "attention_mask": attn} + + +def test_build_model_inputs_camera_order_and_padding(): + import numpy as np + + config = Pi0Config.from_model_config(_LEROBOT_CFG) # 3-camera config + device = torch.device("cpu") + # Provide only 2 of the 3 declared cameras. + robot_obs = { + "images": { + "observation.images.base_0_rgb": np.full((224, 224, 3), 50, dtype=np.uint8), + "observation.images.left_wrist_0_rgb": np.full((224, 224, 3), 80, dtype=np.uint8), + }, + "state": [0.1] * 14, # raw state shorter than max_state_dim=32 + "prompt": "pick up the block", + } + images, image_masks, lang_tokens, lang_masks, state = build_model_inputs( + robot_obs, config, _FakeTokenizer(), device + ) + assert len(images) == 3 + assert all(img.shape == (1, 3, 224, 224) for img in images) + masks = [bool(m.item()) for m in image_masks] + assert masks == [True, True, False] + # The padded (3rd) camera must be all -1. + assert (images[2] == -1.0).all() + assert lang_tokens.shape == (1, 48) + assert lang_masks.shape == (1, 48) + assert state.shape == (1, 32) + # State padded with zeros beyond the 14 real entries. + assert torch.allclose(state[0, 14:], torch.zeros(18)) + + +def test_build_model_inputs_state_truncated(): + config = Pi0Config.from_model_config(_LEROBOT_CFG) + robot_obs = {"images": {}, "state": [0.5] * 40, "prompt": "x"} + _, _, _, _, state = build_model_inputs(robot_obs, config, _FakeTokenizer(), torch.device("cpu")) + assert state.shape == (1, 32) + + +def test_build_model_inputs_image_key_map(): + import numpy as np + + cfg = Pi0Config.from_model_config({**_LEROBOT_CFG, "image_key_map": {"cam_high": "observation.images.base_0_rgb"}}) + robot_obs = { + "images": {"cam_high": np.full((224, 224, 3), 60, dtype=np.uint8)}, + "prompt": "go", + } + _, image_masks, _, _, _ = build_model_inputs(robot_obs, cfg, _FakeTokenizer(), torch.device("cpu")) + masks = [bool(m.item()) for m in image_masks] + assert masks == [True, False, False] # only base_0_rgb (mapped) is present + + +# ---------------------------------------------------------------------------- +# Version-stability / self-consistency (gated on a real pi0_base checkpoint) +# +# The LeRobot parity gate (test_pi0_parity.py) only validates correctness at +# transformers 5.3.0 (where LeRobot loads pi0_base cleanly). Production vllm-omni +# ships transformers 5.6.2, and LeRobot is not a usable reference there (it +# incorrectly loads the SigLIP vision tower at >=5.4). So we pin the kernel's deterministic +# output to golden values established at 5.3.0 and re-check them here: running on +# the production transformers proves the port is version-stable and guards the +# vision_tower / embed_scale version-robust paths against regressions. +# +# PI0_SELFCONSIST_CKPT=/path/to/pi0_base \ +# pytest tests/diffusion/models/pi0/test_pi0_units.py -k version_stable -v -s +# ---------------------------------------------------------------------------- +import os # noqa: E402 + +_SELFCONSIST_CKPT = os.environ.get("PI0_SELFCONSIST_CKPT") + +# Golden values established at transformers 5.3.0 (where parity = 7.15e-07). +_GOLDEN_SUM = -67.808113 +_GOLDEN_STD = 0.252789 +_GOLDEN_MEAN = -0.021190 +_GOLDEN_EMBED_IMAGE_STD = 2.17603 + + +def _selfconsist_inputs(device): + # NB: draw order is load-bearing — the golden values were established with + # state drawn first, then the 3 images, then lang (shared RNG, seed 0). + g = torch.Generator(device="cpu").manual_seed(0) + b = 2 + state = torch.randn(b, 32, generator=g) + images = [torch.rand(b, 3, 224, 224, generator=g) for _ in range(3)] + masks = [torch.ones(b, dtype=torch.bool) for _ in range(3)] + lang = torch.randint(0, 1000, (b, 48), generator=g) + lmask = torch.ones(b, 48, dtype=torch.bool) + ng = torch.Generator(device="cpu").manual_seed(42) + noise = torch.randn(b, 50, 32, generator=ng) + return ( + [im.to(device) for im in images], + [m.to(device) for m in masks], + lang.to(device), + lmask.to(device), + state.to(device), + noise.to(device), + ) + + +def _selfconsist_model(device): + import safetensors.torch + + model = Pi0ForActionPrediction( + Pi0Config(max_action_dim=32, max_state_dim=32, chunk_size=50, num_inference_steps=10, dtype="float32") + ) + state = safetensors.torch.load_file(os.path.join(_SELFCONSIST_CKPT, "model.safetensors")) + filled = model.load_weights(list(state.items())) + expected = {n for n, _ in model.named_parameters() if "rotary_emb" not in n and not n.endswith(".inv_freq")} + missing = expected - (filled or set()) + assert not missing, f"{len(missing)} params received no weight (e.g. {sorted(missing)[:3]})" + return model.to(device).eval() + + +@pytest.mark.slow +@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.") +@torch.no_grad() +def test_version_stable_actions(): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = _selfconsist_model(device) + images, masks, lang, lmask, state, noise = _selfconsist_inputs(device) + actions = ( + model.sample_actions( + images=images, + image_masks=masks, + lang_tokens=lang, + lang_masks=lmask, + state=state, + noise=noise, + num_steps=10, + ) + .float() + .cpu() + ) + s, m, std = actions.sum().item(), actions.mean().item(), actions.std().item() + assert abs(s - _GOLDEN_SUM) < 0.5, f"action sum drifted: {s} vs {_GOLDEN_SUM}" + assert abs(std - _GOLDEN_STD) < 0.01, f"action std drifted: {std} vs {_GOLDEN_STD}" + assert abs(m - _GOLDEN_MEAN) < 0.01, f"action mean drifted: {m} vs {_GOLDEN_MEAN}" + + +@pytest.mark.slow +@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.") +@torch.no_grad() +def test_version_stable_embed_image(): + """Directly guards the SigLIP vision-tower load across the >=5.4 flattening.""" + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = _selfconsist_model(device) + g = torch.Generator(device="cpu").manual_seed(7) + img = torch.rand(1, 3, 224, 224, generator=g, dtype=torch.float32).to(device) + std = model.paligemma_with_expert.embed_image(img).float().std().item() + assert abs(std - _GOLDEN_EMBED_IMAGE_STD) < 0.05, ( + f"embed_image std drifted: {std} vs {_GOLDEN_EMBED_IMAGE_STD} — " + "likely a vision_tower weight-load regression under this transformers version." + ) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py index 62d17f705..23f0bd467 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import json from types import SimpleNamespace diff --git a/tests/diffusion/models/sana_video/test_pipeline_sana_video.py b/tests/diffusion/models/sana_video/test_pipeline_sana_video.py new file mode 100644 index 000000000..f6f469aff --- /dev/null +++ b/tests/diffusion/models/sana_video/test_pipeline_sana_video.py @@ -0,0 +1,992 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import nullcontext +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image +from torch import nn + +from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _DummyComponent: + pass + + +def _make_request_batch(prompt, request_id="sana-video-test", **sampling_overrides): + sampling = OmniDiffusionSamplingParams(**sampling_overrides) + request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=request_id) + return DiffusionRequestBatch([request]) + + +def test_sana_video_pipeline_import_and_registry(): + from vllm_omni.diffusion.models.sana_video import ( + SanaImageToVideoPipeline, + SanaVideoPipeline, + SanaVideoTransformer3DModel, + get_sana_video_i2v_post_process_func, + get_sana_video_i2v_pre_process_func, + get_sana_video_post_process_func, + ) + from vllm_omni.diffusion.registry import ( + _DIFFUSION_MODELS, + _DIFFUSION_POST_PROCESS_FUNCS, + _DIFFUSION_PRE_PROCESS_FUNCS, + ) + + assert SanaImageToVideoPipeline is not None + assert SanaVideoPipeline is not None + assert SanaVideoTransformer3DModel is not None + assert get_sana_video_i2v_post_process_func is not None + assert get_sana_video_i2v_pre_process_func is not None + assert get_sana_video_post_process_func is not None + assert _DIFFUSION_MODELS["SanaVideoPipeline"] == ( + "sana_video", + "pipeline_sana_video", + "SanaVideoPipeline", + ) + assert _DIFFUSION_POST_PROCESS_FUNCS["SanaVideoPipeline"] == "get_sana_video_post_process_func" + assert _DIFFUSION_MODELS["SanaImageToVideoPipeline"] == ( + "sana_video", + "pipeline_sana_video_i2v", + "SanaImageToVideoPipeline", + ) + assert _DIFFUSION_POST_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_post_process_func" + assert _DIFFUSION_PRE_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_pre_process_func" + + +def test_component_discovery_declarations(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + assert SanaVideoPipeline._dit_modules == ["transformer"] + assert SanaVideoPipeline._encoder_modules == ["text_encoder"] + assert SanaVideoPipeline._vae_modules == ["vae"] + assert SanaVideoPipeline.supports_step_execution is False + + +def test_sana_video_loads_concrete_gemma_tokenizer(monkeypatch): + from vllm_omni.diffusion.models.sana_video import pipeline_sana_video + + captured: dict[str, object] = {} + expected = object() + + def fake_from_pretrained_with_prefetch(loader, model, **kwargs): + captured.update(loader=loader, model=model, kwargs=kwargs) + return expected + + monkeypatch.setattr( + pipeline_sana_video, + "from_pretrained_with_prefetch", + fake_from_pretrained_with_prefetch, + ) + + result = pipeline_sana_video._load_sana_tokenizer( + "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + ["tokenizer", "text_encoder"], + local_files_only=False, + ) + + assert result is expected + assert captured["loader"] == pipeline_sana_video.GemmaTokenizer.from_pretrained + assert captured["model"] == "Efficient-Large-Model/SANA-Video_2B_480p_diffusers" + assert captured["kwargs"] == { + "subfolder": "tokenizer", + "prefetch_list": ["tokenizer", "text_encoder"], + "local_files_only": False, + } + + +def test_sana_video_declares_extra_body_params(): + from vllm_omni.model_extras import get_extra_body_params + + assert get_extra_body_params("SanaVideoPipeline") == { + "clean_caption", + "motion_score", + "use_resolution_binning", + } + assert get_extra_body_params("SanaImageToVideoPipeline") == { + "clean_caption", + "motion_score", + "use_resolution_binning", + } + + +def test_sana_video_i2v_preprocesses_image_and_preserves_aspect_ratio(): + from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func + + request = OmniDiffusionRequest( + prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (640, 360))}}, + sampling_params=OmniDiffusionSamplingParams(), + request_id="sana-video-i2v-preprocess", + ) + processed = get_sana_video_i2v_pre_process_func(SimpleNamespace())(request) + + assert processed.sampling_params.height == 448 + assert processed.sampling_params.width == 832 + assert processed.prompt["multi_modal_data"]["image"].size == (832, 448) + + +def test_sana_video_i2v_remote_720p_uses_loaded_transformer_config(): + from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func + + od_config = SimpleNamespace( + model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + tf_model_config={"sample_size": 22}, + ) + request = OmniDiffusionRequest( + prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (1280, 704))}}, + sampling_params=OmniDiffusionSamplingParams(), + request_id="sana-video-i2v-remote-720p", + ) + + processed = get_sana_video_i2v_pre_process_func(od_config)(request) + + assert processed.sampling_params.height == 704 + assert processed.sampling_params.width == 1280 + assert processed.prompt["multi_modal_data"]["image"].size == (1280, 704) + + +@pytest.mark.parametrize( + "frames", + [ + [[Image.new("RGB", (16, 16)) for _ in range(2)]], + np.zeros((1, 2, 16, 16, 3), dtype=np.float32), + [[np.zeros((16, 16, 3), dtype=np.float32) for _ in range(2)]], + ], +) +@pytest.mark.parametrize( + "factory_name", + ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"], +) +def test_sana_video_diffusers_preserves_postprocessed_frames(frames, factory_name): + from vllm_omni.diffusion.models import sana_video + + post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers")) + + result = post_process(frames) + + assert result["payload"]["video"] is frames + + +@pytest.mark.parametrize( + "factory_name", + ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"], +) +def test_sana_video_diffusers_still_postprocesses_native_tensor(factory_name): + from vllm_omni.diffusion.models import sana_video + + decoded_video = torch.zeros(1, 3, 2, 16, 16) + post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers")) + + result = post_process(decoded_video) + + assert isinstance(result["payload"]["video"], np.ndarray) + assert result["payload"]["video"].shape == (1, 2, 16, 16, 3) + + +@pytest.mark.parametrize( + "factory_name", + ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"], +) +def test_sana_video_postprocess_preserves_requested_latents(factory_name): + from vllm_omni.diffusion.models import sana_video + + latents = torch.randn(1, 4, 2, 2, 2) + post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="native")) + + result = post_process(latents, sampling_params=SimpleNamespace(output_type="latent")) + + assert result is latents + + +def test_sana_video_720p_model_id_fallback(monkeypatch): + from vllm_omni.diffusion.models.sana_video import pipeline_sana_video + + monkeypatch.setattr(pipeline_sana_video, "get_hf_file_to_dict", lambda *_args, **_kwargs: None) + od_config = SimpleNamespace( + model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + tf_model_config={}, + ) + + assert pipeline_sana_video.resolve_sana_video_sample_size(od_config) == 22 + + +@pytest.mark.parametrize( + ("requested_output_type", "expected_internal_output_type"), + [ + (None, "raw"), + ("np", "raw"), + ("latent", "latent"), + ], +) +def test_sana_video_i2v_forward_maps_image_request(requested_output_type, expected_internal_output_type): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate_i2v(**kwargs): + calls.append(kwargs) + return torch.zeros(kwargs["num_videos_per_prompt"], 3, 9, 192, 320) + + pipeline._generate_i2v = fake_generate_i2v + image = Image.new("RGB", (320, 192)) + custom_timesteps = torch.tensor([900, 100]) + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}}, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + timesteps=custom_timesteps, + guidance_scale=4.5, + num_outputs_per_prompt=2, + eta=0.25, + seed=42, + output_type=requested_output_type, + extra_args={"clean_caption": True, "motion_score": 30, "use_resolution_binning": False}, + ) + + output = pipeline.forward(req) + + assert output.output.shape == (2, 3, 9, 192, 320) + assert calls[0]["image"] is image + assert calls[0]["prompt"] == "a robot walks motion score: 30." + assert calls[0]["negative_prompt"] == "blurry" + assert calls[0]["frames"] == 9 + assert calls[0]["timesteps"] is custom_timesteps + assert calls[0]["sigmas"] is None + assert calls[0]["num_videos_per_prompt"] == 2 + assert calls[0]["eta"] == 0.25 + assert calls[0]["clean_caption"] is True + assert calls[0]["generator"].initial_seed() == 42 + assert calls[0]["output_type"] == expected_internal_output_type + + +def test_sana_video_i2v_forward_maps_custom_sigmas(): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate_i2v(**kwargs): + calls.append(kwargs) + return torch.zeros(1, 3, 9, 192, 320) + + pipeline._generate_i2v = fake_generate_i2v + sigmas = [1.0, 0.5] + output = pipeline.forward( + _make_request_batch( + { + "prompt": "a robot walks", + "multi_modal_data": {"image": Image.new("RGB", (320, 192))}, + }, + height=192, + width=320, + num_frames=9, + sigmas=sigmas, + generator_device="cpu", + extra_args={"use_resolution_binning": False}, + ) + ) + + assert output.output.shape == (1, 3, 9, 192, 320) + assert calls[0]["timesteps"] is None + assert calls[0]["sigmas"] is sigmas + + +def test_sana_video_i2v_omitted_num_frames_uses_model_default(): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate_i2v(**kwargs): + calls.append(kwargs) + return torch.zeros(1, 3, 81, 2, 2) + + pipeline._generate_i2v = fake_generate_i2v + req = _make_request_batch( + { + "prompt": "a robot walks", + "multi_modal_data": {"image": Image.new("RGB", (320, 192))}, + }, + height=192, + width=320, + num_inference_steps=1, + generator_device="cpu", + extra_args={"use_resolution_binning": False}, + ) + + assert req.sampling_params.num_frames == 1 + output = pipeline.forward(req) + + assert output.output.shape == (1, 3, 81, 2, 2) + assert calls[0]["frames"] == 81 + + +def test_sana_video_i2v_clean_caption_request_is_explicitly_rejected(): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + pipeline.text_encoder = SimpleNamespace(dtype=torch.float32) + pipeline.tokenizer = None + pipeline.check_inputs = lambda **_kwargs: None + + req = _make_request_batch( + { + "prompt": "a robot walks", + "multi_modal_data": {"image": Image.new("RGB", (320, 192))}, + }, + height=192, + width=320, + num_frames=9, + generator_device="cpu", + extra_args={"clean_caption": True, "use_resolution_binning": False}, + ) + + with pytest.raises(ValueError, match="does not support `clean_caption=True`"): + pipeline.forward(req) + + +@pytest.mark.parametrize( + ("custom_timesteps", "custom_sigmas"), + [ + (torch.tensor([4.0, 2.0]), None), + (None, [1.0, 0.5]), + ], +) +def test_sana_video_i2v_consumes_batch_schedule_and_step_kwargs(custom_timesteps, custom_sigmas): + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + calls = SimpleNamespace( + encode_prompt=[], + prepare_latents=[], + set_timesteps=[], + step=[], + transformer=[], + ) + + class StubTransformer: + dtype = torch.float32 + config = SimpleNamespace( + sample_size=30, + in_channels=4, + out_channels=4, + patch_size=(1, 1, 1), + ) + + def __call__( + self, + hidden_states, + *, + encoder_hidden_states, + encoder_attention_mask, + timestep, + return_dict, + ): + calls.transformer.append( + ( + hidden_states.shape, + encoder_hidden_states.shape, + encoder_attention_mask.shape, + timestep.shape, + return_dict, + ) + ) + return (torch.zeros_like(hidden_states),) + + class StubScheduler: + order = 1 + + def set_timesteps( + self, + num_inference_steps=None, + *, + device=None, + timesteps=None, + sigmas=None, + ): + calls.set_timesteps.append( + { + "num_inference_steps": num_inference_steps, + "timesteps": timesteps, + "sigmas": sigmas, + } + ) + schedule_length = len(timesteps if timesteps is not None else sigmas) + self.timesteps = torch.arange(schedule_length, 0, -1, dtype=torch.float32, device=device) + + def step( + self, + _noise_pred, + _timestep, + current_latents, + *, + eta, + generator, + return_dict, + ): + calls.step.append({"eta": eta, "generator": generator, "return_dict": return_dict}) + return (current_latents + 1,) + + class VaeMustNotBeUsed: + @property + def dtype(self): + pytest.fail("latent output must return before accessing the VAE") + + def decode(self, *_args, **_kwargs): + pytest.fail("latent output must not call VAE decode") + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = StubTransformer() + pipeline.scheduler = StubScheduler() + pipeline.vae = VaeMustNotBeUsed() + pipeline.check_inputs = lambda **_kwargs: None + pipeline.video_processor = SimpleNamespace( + preprocess=lambda _image, height, width: torch.zeros(1, 3, height, width), + ) + pipeline.progress_bar = lambda **_kwargs: nullcontext(SimpleNamespace(update=lambda: None)) + + def capture_encode_prompt(*_args, **kwargs): + calls.encode_prompt.append(kwargs) + batch_size = kwargs["num_videos_per_prompt"] + return ( + torch.zeros(batch_size, 1, 1), + torch.ones(batch_size, 1, dtype=torch.bool), + torch.zeros(batch_size, 1, 1), + torch.ones(batch_size, 1, dtype=torch.bool), + ) + + pipeline.encode_prompt = capture_encode_prompt + initial_latents = torch.zeros(2, 4, 2, 2, 2) + + def capture_prepare_latents(*_args, **kwargs): + calls.prepare_latents.append(kwargs) + return initial_latents + + pipeline._prepare_i2v_latents = capture_prepare_latents + generator = torch.Generator(device="cpu").manual_seed(42) + + output = pipeline._generate_i2v( + image=Image.new("RGB", (2, 2)), + prompt="test", + negative_prompt="", + height=2, + width=2, + frames=2, + num_inference_steps=50, + timesteps=custom_timesteps, + sigmas=custom_sigmas, + guidance_scale=4.5, + num_videos_per_prompt=2, + eta=0.25, + generator=generator, + latents=None, + clean_caption=True, + use_resolution_binning=False, + max_sequence_length=1, + output_type="latent", + ) + + assert calls.encode_prompt[0]["num_videos_per_prompt"] == 2 + assert calls.encode_prompt[0]["clean_caption"] is True + assert calls.prepare_latents[0]["batch_size"] == 2 + assert len(calls.set_timesteps) == 1 + assert calls.set_timesteps[0]["num_inference_steps"] is None + assert calls.set_timesteps[0]["timesteps"] is custom_timesteps + assert calls.set_timesteps[0]["sigmas"] is custom_sigmas + assert len(calls.step) == 2 + assert all(call == {"eta": 0.25, "generator": generator, "return_dict": False} for call in calls.step) + assert all( + call + == ( + torch.Size([4, 4, 2, 2, 2]), + torch.Size([4, 1, 1]), + torch.Size([4, 1]), + torch.Size([4, 1, 2, 2, 2]), + False, + ) + for call in calls.transformer + ) + expected = initial_latents.clone() + expected[:, :, 1:] += 2 + torch.testing.assert_close(output, expected) + + +def test_sana_video_i2v_uses_diffusers_complex_instruction_default(): + import inspect + + from diffusers import SanaImageToVideoPipeline as DiffusersSanaImageToVideoPipeline + + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + reference_default = ( + inspect.signature(DiffusersSanaImageToVideoPipeline.__call__).parameters["complex_human_instruction"].default + ) + native_default = ( + inspect.signature(SanaImageToVideoPipeline._generate_i2v).parameters["complex_human_instruction"].default + ) + assert native_default == reference_default + + class StopAfterPromptEncodingError(Exception): + pass + + pipeline = object.__new__(SanaImageToVideoPipeline) + pipeline.device = torch.device("cpu") + pipeline.check_inputs = lambda **_kwargs: None + captured = {} + + def capture_encode_prompt(*_args, **kwargs): + captured.update(kwargs) + raise StopAfterPromptEncodingError + + pipeline.encode_prompt = capture_encode_prompt + with pytest.raises(StopAfterPromptEncodingError): + pipeline._generate_i2v( + image=Image.new("RGB", (320, 192)), + prompt="a robot walks", + negative_prompt="blurry", + height=192, + width=320, + frames=9, + num_inference_steps=2, + guidance_scale=4.5, + generator=torch.Generator(device="cpu").manual_seed(42), + latents=None, + use_resolution_binning=False, + max_sequence_length=300, + ) + + assert captured["complex_human_instruction"] == reference_default + + +def test_sana_video_t2v_complex_instruction_has_no_mutable_default(): + import inspect + + from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION, + SanaVideoPipeline, + ) + + parameter = inspect.signature(SanaVideoPipeline._generate).parameters["complex_human_instruction"] + assert parameter.default is None + assert isinstance(SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION, tuple) + + +def test_diffusers_adapter_t2v_omitted_num_frames_uses_sana_default(): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "negative_prompt", "frames", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None) + + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry"}, + seed=42, + generator_device="cpu", + ) + assert req.sampling_params.num_frames == 1 + + kwargs = adapter._build_call_kwargs(req) + + assert kwargs["frames"] == 81 + assert "num_frames" not in kwargs + assert kwargs["generator"].initial_seed() == 42 + + +def test_diffusers_adapter_i2v_omitted_num_frames_uses_sana_default(): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "negative_prompt", "image", "frames", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None) + image = Image.new("RGB", (320, 192)) + + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}}, + seed=42, + generator_device="cpu", + ) + assert req.sampling_params.num_frames == 1 + + kwargs = adapter._build_call_kwargs(req) + + assert kwargs["image"] is image + assert kwargs["prompt"] == "a robot walks" + assert kwargs["negative_prompt"] == "blurry" + assert kwargs["frames"] == 81 + assert kwargs["generator"].initial_seed() == 42 + + +@pytest.mark.parametrize( + ("sampling_overrides", "expected_frames"), + [ + ({}, 49), + ({"num_frames": 9}, 9), + ], +) +def test_diffusers_adapter_sana_resolves_frame_precedence(sampling_overrides, expected_frames): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "frames", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={"frames": 49}, output_type=None) + + kwargs = adapter._build_call_kwargs( + _make_request_batch( + "a robot walks", + seed=42, + generator_device="cpu", + **sampling_overrides, + ) + ) + + assert kwargs["frames"] == expected_frames + + +@pytest.mark.parametrize("pipeline_class_name", ["SanaVideoPipeline", "SanaImageToVideoPipeline"]) +def test_diffusers_adapter_selects_sana_pipeline_utils(pipeline_class_name): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import ( + SanaVideoPipelineUtils, + get_pipeline_utils, + ) + + assert isinstance(get_pipeline_utils(pipeline_class_name), SanaVideoPipelineUtils) + + +def test_diffusers_adapter_resolves_requested_sana_i2v_pipeline(): + from diffusers import SanaImageToVideoPipeline, SanaVideoPipeline + + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import ( + SanaVideoPipelineUtils, + get_pipeline_utils_for_config, + resolve_diffusers_pipeline_class, + ) + + od_config = SimpleNamespace( + model_class_name="SanaImageToVideoPipeline", + diffusers_pipeline_cls=SanaVideoPipeline, + ) + + assert isinstance(get_pipeline_utils_for_config(od_config), SanaVideoPipelineUtils) + assert resolve_diffusers_pipeline_class(od_config) is SanaImageToVideoPipeline + + +def test_diffusers_adapter_pipeline_override_is_sana_specific(): + from diffusers import DiffusionPipeline + + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import resolve_diffusers_pipeline_class + + class OtherPipeline(DiffusionPipeline): + pass + + od_config = SimpleNamespace( + model_class_name="DiffusersAdapterPipeline", + diffusers_pipeline_cls=OtherPipeline, + ) + + assert resolve_diffusers_pipeline_class(od_config) is OtherPipeline + + +def test_diffusers_adapter_loads_resolved_pipeline_class(monkeypatch): + from vllm_omni.diffusion.models.diffusers_adapter import pipeline_diffusers_adapter + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + + loaded = [] + expected_pipeline = object() + + class ResolvedPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + loaded.append((model_id, kwargs)) + return expected_pipeline + + monkeypatch.setattr( + pipeline_diffusers_adapter, + "resolve_diffusers_pipeline_class", + lambda _od_config: ResolvedPipeline, + ) + adapter = object.__new__(DiffusersAdapterPipeline) + adapter.od_config = SimpleNamespace(diffusers_pipeline_cls=object) + + result = adapter._load_pipeline_from_pretrained("sana-checkpoint", {"torch_dtype": torch.bfloat16}) + + assert result is expected_pipeline + assert loaded == [("sana-checkpoint", {"torch_dtype": torch.bfloat16})] + + +def test_diffusers_adapter_sana_i2v_override_reports_image_capability(): + from diffusers import SanaVideoPipeline + + from vllm_omni.diffusion.io_support import supports_multimodal_input + + od_config = SimpleNamespace( + diffusion_load_format="diffusers", + model_class_name="SanaImageToVideoPipeline", + diffusers_pipeline_cls=SanaVideoPipeline, + ) + + assert supports_multimodal_input(od_config) == (True, False) + + +def test_diffusers_adapter_sana_i2v_warmup_supplies_image(): + from diffusers import SanaVideoPipeline + + from vllm_omni.diffusion.diffusion_engine import DiffusionEngine + + engine = object.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace( + diffusion_load_format="diffusers", + model_class_name="SanaImageToVideoPipeline", + diffusers_pipeline_cls=SanaVideoPipeline, + ) + captured_requests = [] + engine.pre_process_func = lambda request: request + + def capture_request(request): + captured_requests.append(request) + return SimpleNamespace(error=None) + + engine.add_req_and_wait_for_response = capture_request + + engine._dummy_run() + + assert len(captured_requests) == 1 + prompt = captured_requests[0].prompt + assert isinstance(prompt, dict) + assert isinstance(prompt["multi_modal_data"]["image"], Image.Image) + + +def test_diffusers_adapter_disables_resolution_binning_for_warmup(): + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline + from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils + + adapter = object.__new__(DiffusersAdapterPipeline) + adapter._accept_call_kwargs = {"prompt", "frames", "use_resolution_binning", "generator"} + adapter._pipeline_utils = SanaVideoPipelineUtils() + adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None) + + kwargs = adapter._build_call_kwargs( + _make_request_batch( + "dummy run", + request_id=DUMMY_DIFFUSION_REQUEST_ID, + generator_device="cpu", + ) + ) + + assert kwargs["use_resolution_binning"] is False + assert kwargs["frames"] == 1 + + +def test_pipeline_is_torch_module_and_supports_eval(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + nn.Module.__init__(pipeline) + + assert isinstance(pipeline, nn.Module) + assert pipeline.eval() is pipeline + assert pipeline.training is False + + +@pytest.mark.parametrize( + ("vae_scale", "patch_size", "valid_size", "invalid_size", "alignment"), + [ + (8, (1, 2, 2), 624, 632, 16), + (32, (1, 1, 1), 704, 712, 32), + ], +) +def test_check_inputs_uses_variant_spatial_alignment(vae_scale, patch_size, valid_size, invalid_size, alignment): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.vae_scale_factor_spatial = vae_scale + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(patch_size=patch_size)) + + pipeline.check_inputs(prompt="test", height=valid_size, width=valid_size) + with pytest.raises(ValueError, match=f"divisible by {alignment}"): + pipeline.check_inputs(prompt="test", height=invalid_size, width=valid_size) + + +@pytest.mark.parametrize( + ("requested_output_type", "expected_internal_output_type"), + [ + (None, "raw"), + ("np", "raw"), + ("latent", "latent"), + ], +) +def test_forward_maps_omni_request_to_sana_generation_args(requested_output_type, expected_internal_output_type): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 9, 192, 320)) + + pipeline._generate = fake_generate + req = _make_request_batch( + {"prompt": "a robot walks", "negative_prompt": "blurry"}, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + guidance_scale=4.5, + seed=42, + output_type=requested_output_type, + extra_args={"motion_score": 30, "use_resolution_binning": False}, + ) + + output = pipeline.forward(req) + + assert output.output.shape == (1, 3, 9, 192, 320) + assert calls[0]["prompt"] == "a robot walks motion score: 30." + assert calls[0]["negative_prompt"] == "blurry" + assert calls[0]["height"] == 192 + assert calls[0]["width"] == 320 + assert calls[0]["frames"] == 9 + assert calls[0]["num_inference_steps"] == 2 + assert calls[0]["guidance_scale"] == 4.5 + assert calls[0]["use_resolution_binning"] is False + assert calls[0]["generator"].initial_seed() == 42 + assert calls[0]["output_type"] == expected_internal_output_type + + +def test_sana_video_t2v_omitted_num_frames_uses_model_default(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 81, 2, 2)) + + pipeline._generate = fake_generate + req = _make_request_batch( + "a robot walks", + height=192, + width=320, + num_inference_steps=1, + generator_device="cpu", + ) + + assert req.sampling_params.num_frames == 1 + output = pipeline.forward(req) + + assert output.output.shape == (1, 3, 81, 2, 2) + assert calls[0]["frames"] == 81 + + +def test_sana_video_t2v_dummy_run_keeps_single_frame(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 1, 192, 320)) + + pipeline._generate = fake_generate + + output = pipeline.forward( + _make_request_batch( + "dummy run", + request_id=DUMMY_DIFFUSION_REQUEST_ID, + height=192, + width=320, + num_frames=1, + num_inference_steps=1, + generator_device="cpu", + ) + ) + + assert output.output.shape == (1, 3, 1, 192, 320) + assert calls[0]["frames"] == 1 + + +@pytest.mark.parametrize( + ("sample_size", "expected_height", "expected_width"), + [ + (30, 480, 832), + (22, 704, 1280), + ], +) +def test_t2v_uses_variant_default_resolution(sample_size, expected_height, expected_width): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=sample_size)) + calls = [] + + def fake_generate(**kwargs): + calls.append(kwargs) + return SimpleNamespace(frames=torch.zeros(1, 3, 9, expected_height, expected_width)) + + pipeline._generate = fake_generate + output = pipeline.forward( + _make_request_batch( + "a robot walks", + num_frames=9, + num_inference_steps=1, + guidance_scale=1.0, + seed=42, + generator_device="cpu", + ) + ) + + assert output.output.shape == (1, 3, 9, expected_height, expected_width) + assert calls[0]["height"] == expected_height + assert calls[0]["width"] == expected_width + + +def test_forward_requires_exactly_one_nonempty_prompt(): + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + pipeline = object.__new__(SanaVideoPipeline) + + with pytest.raises(ValueError, match="Prompt is required"): + pipeline.forward(_make_request_batch("")) + + first = OmniDiffusionRequest( + prompt="first", + sampling_params=OmniDiffusionSamplingParams(), + request_id="first", + ) + second = OmniDiffusionRequest( + prompt="second", + sampling_params=OmniDiffusionSamplingParams(), + request_id="second", + ) + with pytest.raises(ValueError, match="exactly one prompt"): + pipeline.forward(DiffusionRequestBatch([first, second])) diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video.py b/tests/diffusion/models/sana_video/test_transformer_sana_video.py new file mode 100644 index 000000000..804692d7a --- /dev/null +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video.py @@ -0,0 +1,602 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_TINY_CONFIG = { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 12, + "num_layers": 1, + "num_cross_attention_heads": 2, + "cross_attention_head_dim": 12, + "cross_attention_dim": 24, + "caption_channels": 8, + "mlp_ratio": 2.0, + "patch_size": (1, 2, 2), + "sample_size": 4, + "rope_max_seq_len": 32, +} + +_GOLDEN_PREFIX = torch.tensor( + [ + -0.1159307063, + 0.2003862262, + -0.3754127920, + -0.2403523028, + -1.2376221418, + -0.6513092518, + -0.3358457983, + -0.9894289970, + 0.7858567238, + 0.5089095831, + -0.5362522602, + -0.3732060790, + ] +) + + +@pytest.mark.parametrize( + ("elementwise_affine", "bias", "expected_state_keys"), + [ + (False, False, set()), + (False, True, set()), + (True, False, {"weight"}), + (True, True, {"weight", "bias"}), + ], +) +@pytest.mark.parametrize( + ("input_dtype", "parameter_dtype"), + [ + (torch.float16, torch.float16), + (torch.bfloat16, torch.bfloat16), + (torch.float32, torch.float32), + (torch.float64, torch.float64), + (torch.float16, torch.float32), + (torch.bfloat16, torch.float32), + (torch.float32, torch.float16), + (torch.float32, torch.bfloat16), + ], +) +def test_sana_rms_norm_matches_diffusers( + elementwise_affine, + bias, + expected_state_keys, + input_dtype, + parameter_dtype, +): + from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaRMSNorm + + reference = DiffusersRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to( + dtype=parameter_dtype + ) + actual = SanaRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to(dtype=parameter_dtype) + + assert actual.eps == reference.eps + assert actual.elementwise_affine == reference.elementwise_affine + assert actual.dim == reference.dim + assert set(actual.state_dict()) == expected_state_keys + assert set(actual.state_dict()) == set(reference.state_dict()) + + if elementwise_affine: + weight = torch.linspace(0.5, 1.5, 8, dtype=parameter_dtype) + actual.weight.data.copy_(weight) + reference.weight.data.copy_(weight) + if bias: + norm_bias = torch.linspace(-0.25, 0.25, 8, dtype=parameter_dtype) + actual.bias.data.copy_(norm_bias) + reference.bias.data.copy_(norm_bias) + else: + assert actual.weight is None + assert actual.bias is None + + hidden_states = torch.linspace(-2.0, 2.0, 48, dtype=input_dtype).reshape(2, 3, 8) + expected = reference(hidden_states) + result = actual(hidden_states) + + assert result.dtype == expected.dtype + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("num_channels", "flip_sin_to_cos", "downscale_freq_shift", "scale"), + [ + (256, True, 0, 1), + (7, False, 1, 0.5), + (12, True, 0.5, 2), + ], +) +@pytest.mark.parametrize("timestep_dtype", [torch.int64, torch.float32, torch.float64]) +def test_sana_timesteps_matches_diffusers( + num_channels, + flip_sin_to_cos, + downscale_freq_shift, + scale, + timestep_dtype, +): + from diffusers.models.embeddings import Timesteps + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaTimesteps + + reference = Timesteps( + num_channels=num_channels, + flip_sin_to_cos=flip_sin_to_cos, + downscale_freq_shift=downscale_freq_shift, + scale=scale, + ) + actual = SanaTimesteps( + num_channels=num_channels, + flip_sin_to_cos=flip_sin_to_cos, + downscale_freq_shift=downscale_freq_shift, + scale=scale, + ) + timesteps = torch.tensor([0, 1, 500, 999], dtype=timestep_dtype) + + assert actual.state_dict() == reference.state_dict() == {} + assert actual.num_channels == reference.num_channels + assert actual.flip_sin_to_cos == reference.flip_sin_to_cos + assert actual.downscale_freq_shift == reference.downscale_freq_shift + assert actual.scale == reference.scale + torch.testing.assert_close(actual(timesteps), reference(timesteps), rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("kwargs", "sample_shape", "condition_shape"), + [ + ({}, (2, 256), None), + ( + { + "out_dim": 12, + "post_act_fn": "silu", + "cond_proj_dim": 6, + "sample_proj_bias": False, + }, + (2, 256), + (2, 6), + ), + ], +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_sana_timestep_embedding_matches_diffusers( + kwargs, + sample_shape, + condition_shape, + dtype, +): + from diffusers.models.embeddings import TimestepEmbedding + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaTimestepEmbedding, + ) + + torch.manual_seed(17) + reference = TimestepEmbedding(256, 24, **kwargs).to(dtype=dtype) + actual = SanaTimestepEmbedding(256, 24, **kwargs).to(dtype=dtype) + assert set(actual.state_dict()) == set(reference.state_dict()) + actual.load_state_dict(reference.state_dict()) + + sample = torch.randn(sample_shape, dtype=dtype) + condition = torch.randn(condition_shape, dtype=dtype) if condition_shape else None + expected = reference(sample, condition) + result = actual(sample, condition) + + assert result.shape == expected.shape + assert result.dtype == expected.dtype + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +def test_sana_timestep_embedding_rejects_unsupported_activation(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaTimestepEmbedding, + ) + + with pytest.raises(ValueError, match="act_fn='silu'"): + SanaTimestepEmbedding(8, 16, act_fn="gelu") + with pytest.raises(ValueError, match="post_act_fn"): + SanaTimestepEmbedding(8, 16, post_act_fn="gelu") + + module = SanaTimestepEmbedding(8, 16) + with pytest.raises(ValueError, match="cond_proj_dim"): + module(torch.randn(1, 8), torch.randn(1, 4)) + + +@pytest.mark.parametrize("use_additional_conditions", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_sana_ada_layer_norm_single_matches_diffusers( + use_additional_conditions, + dtype, +): + from diffusers.models.normalization import AdaLayerNormSingle + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaAdaLayerNormSingle, + ) + + torch.manual_seed(19) + reference = AdaLayerNormSingle( + 24, + use_additional_conditions=use_additional_conditions, + ).to(dtype=dtype) + actual = SanaAdaLayerNormSingle( + 24, + use_additional_conditions=use_additional_conditions, + ).to(dtype=dtype) + assert set(actual.state_dict()) == set(reference.state_dict()) + actual.load_state_dict(reference.state_dict()) + + timestep = torch.tensor([10, 500]) + added_cond_kwargs = None + if use_additional_conditions: + added_cond_kwargs = { + "resolution": torch.tensor([[480, 832], [704, 1280]]), + "aspect_ratio": torch.tensor([[832 / 480], [1280 / 704]]), + } + + expected = reference( + timestep, + added_cond_kwargs=added_cond_kwargs, + batch_size=2, + hidden_dtype=dtype, + ) + result = actual( + timestep, + added_cond_kwargs=added_cond_kwargs, + batch_size=2, + hidden_dtype=dtype, + ) + + assert len(result) == len(expected) == 2 + for actual_tensor, expected_tensor in zip(result, expected): + assert actual_tensor.shape == expected_tensor.shape + assert actual_tensor.dtype == expected_tensor.dtype + torch.testing.assert_close(actual_tensor, expected_tensor, rtol=0, atol=0) + + +@pytest.mark.parametrize("act_fn", ["gelu_tanh", "silu", "silu_fp32"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_sana_pixart_text_projection_matches_diffusers(act_fn, dtype): + from diffusers.models.embeddings import PixArtAlphaTextProjection + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaPixArtAlphaTextProjection, + ) + + torch.manual_seed(23) + reference = PixArtAlphaTextProjection( + in_features=8, + hidden_size=24, + out_features=12, + act_fn=act_fn, + ).to(dtype=dtype) + actual = SanaPixArtAlphaTextProjection( + in_features=8, + hidden_size=24, + out_features=12, + act_fn=act_fn, + ).to(dtype=dtype) + assert set(actual.state_dict()) == set(reference.state_dict()) + actual.load_state_dict(reference.state_dict()) + + caption = torch.randn(2, 5, 8, dtype=dtype) + expected = reference(caption) + result = actual(caption) + + assert result.shape == expected.shape + assert result.dtype == expected.dtype + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("freqs_dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize(("dim", "max_seq_len", "theta"), [(4, 8, 10000.0), (12, 32, 256.0)]) +def test_native_rope_matches_diffusers(dim, max_seq_len, theta, freqs_dtype): + from diffusers.models.embeddings import get_1d_rotary_pos_embed + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + _get_1d_rotary_pos_embed, + ) + + expected_cos, expected_sin = get_1d_rotary_pos_embed( + dim, + max_seq_len, + theta, + use_real=True, + repeat_interleave_real=True, + freqs_dtype=freqs_dtype, + ) + actual_cos, actual_sin = _get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype) + + torch.testing.assert_close(actual_cos, expected_cos, rtol=0, atol=0) + torch.testing.assert_close(actual_sin, expected_sin, rtol=0, atol=0) + + +def test_tiny_transformer_matches_diffusers_and_frozen_output(): + from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer + from diffusers.configuration_utils import ConfigMixin + from diffusers.models.modeling_utils import ModelMixin + + from vllm_omni.diffusion.attention.layer import Attention + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaLinearAttention, + SanaRMSNorm, + SanaVideoTransformerConfig, + SanaVideoTransformerOutput, + ) + + torch.manual_seed(7) + reference = DiffusersTransformer(**_TINY_CONFIG).eval() + model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval() + assert not isinstance(model, ModelMixin) + assert not isinstance(model, ConfigMixin) + assert isinstance(model.config, SanaVideoTransformerConfig) + assert set(model.state_dict()) == set(reference.state_dict()) + model.load_state_dict(reference.state_dict()) + + block = model.transformer_blocks[0] + assert isinstance(block.attn1, SanaLinearAttention) + assert isinstance(block.attn2.attn, Attention) + assert block.attn2.attn.role == "cross" + assert isinstance(block.attn1.norm_q, SanaRMSNorm) + assert isinstance(block.attn1.norm_k, SanaRMSNorm) + assert isinstance(block.attn2.norm_q, SanaRMSNorm) + assert isinstance(block.attn2.norm_k, SanaRMSNorm) + assert isinstance(model.caption_norm, SanaRMSNorm) + + torch.manual_seed(11) + hidden_states = torch.randn(1, 4, 3, 4, 4) + encoder_hidden_states = torch.randn(1, 5, 8) + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]]) + timestep = torch.tensor([500.0]) + + with torch.no_grad(): + expected = reference( + hidden_states, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + ).sample + actual_output = model( + hidden_states, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + ) + actual = actual_output.sample + + assert isinstance(actual_output, SanaVideoTransformerOutput) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(actual.flatten()[: len(_GOLDEN_PREFIX)], _GOLDEN_PREFIX, rtol=1e-5, atol=1e-5) + + +def test_encoder_attention_mask_uses_bool_sdpa_contract(monkeypatch): + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + torch.manual_seed(37) + model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval() + captured = {} + omni_attention = model.transformer_blocks[0].attn2.attn + + def fake_sdpa_forward(query, key, value, attn_metadata): + captured["query_shape"] = query.shape + captured["key_shape"] = key.shape + captured["mask"] = attn_metadata.attn_mask.detach().clone() + return torch.zeros_like(query) + + def unexpected_backend_forward(*args, **kwargs): + raise AssertionError("masked SANA cross-attention must use SDPA") + + monkeypatch.setattr(omni_attention.sdpa_fallback, "forward", fake_sdpa_forward) + monkeypatch.setattr(omni_attention, "forward", unexpected_backend_forward) + + hidden_states = torch.randn(2, 4, 3, 4, 4) + encoder_hidden_states = torch.randn(2, 5, 8) + input_mask = torch.tensor( + [ + [1, 1, 1, 0, 0], + [1, 1, 1, 1, 0], + ], + dtype=torch.int64, + ) + + with torch.no_grad(): + output = model( + hidden_states, + encoder_hidden_states, + torch.tensor([100.0, 500.0]), + encoder_attention_mask=input_mask, + ).sample + + assert output.shape == hidden_states.shape + assert captured["query_shape"][:2] == (2, 12) + assert captured["key_shape"][:2] == (2, 5) + assert captured["query_shape"][1] != captured["key_shape"][1] + assert captured["mask"].shape == (2, 5) + assert captured["mask"].dtype == torch.bool + assert torch.equal(captured["mask"], input_mask.bool()) + + +def test_encoder_attention_mask_rejects_additive_bias(): + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval() + additive_mask = torch.tensor([[0.0, 0.0, 0.0, -10000.0, -10000.0]]) + + with pytest.raises(TypeError, match="not a floating-point additive attention bias"): + model( + torch.randn(1, 4, 3, 4, 4), + torch.randn(1, 5, 8), + torch.tensor([500.0]), + encoder_attention_mask=additive_mask, + ) + + +def test_linear_attention_requires_rotary_embeddings(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention + + attention = SanaLinearAttention( + dim=24, + num_heads=2, + head_dim=12, + dropout=0.0, + bias=True, + qk_norm="rms_norm_across_heads", + ) + + with pytest.raises(ValueError, match="requires rotary_emb"): + attention(torch.randn(1, 4, 24)) + + +def test_linear_attention_bfloat16_matches_diffusers_mixed_precision(): + from diffusers.models.attention_processor import Attention as DiffusersAttention + from diffusers.models.transformers.transformer_sana_video import SanaLinearAttnProcessor3_0 + + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention + + torch.manual_seed(17) + reference = DiffusersAttention( + query_dim=24, + heads=2, + dim_head=12, + kv_heads=2, + qk_norm="rms_norm_across_heads", + dropout=0.0, + bias=True, + cross_attention_dim=None, + processor=SanaLinearAttnProcessor3_0(), + ).to(dtype=torch.bfloat16) + actual = SanaLinearAttention( + dim=24, + num_heads=2, + head_dim=12, + dropout=0.0, + bias=True, + qk_norm="rms_norm_across_heads", + ).to(dtype=torch.bfloat16) + actual.load_state_dict(reference.state_dict()) + + hidden_states = torch.randn(1, 5, 24, dtype=torch.bfloat16) + rotary_emb = ( + torch.randn(1, 5, 1, 12, dtype=torch.bfloat16), + torch.randn(1, 5, 1, 12, dtype=torch.bfloat16), + ) + + with torch.no_grad(): + expected = reference(hidden_states, rotary_emb=rotary_emb) + result = actual(hidden_states, rotary_emb=rotary_emb) + + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +def test_transformer_block_without_cross_attention_keeps_feed_forward_norm(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerBlock + + block = SanaVideoTransformerBlock( + dim=24, + num_attention_heads=2, + attention_head_dim=12, + num_cross_attention_heads=2, + cross_attention_head_dim=12, + cross_attention_dim=None, + mlp_ratio=2.0, + ) + + assert block.attn2 is None + assert isinstance(block.norm2, torch.nn.LayerNorm) + + +def test_guidance_transformer_matches_diffusers_and_tuple_output(): + from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer + + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import ( + SanaTimestepEmbedding, + SanaTimesteps, + ) + + config = _TINY_CONFIG | {"guidance_embeds": True} + torch.manual_seed(29) + reference = DiffusersTransformer(**config).eval() + model = SanaVideoTransformer3DModel(**config).eval() + assert set(model.state_dict()) == set(reference.state_dict()) + model.load_state_dict(reference.state_dict()) + assert isinstance(model.time_embed.time_proj, SanaTimesteps) + assert isinstance(model.time_embed.timestep_embedder, SanaTimestepEmbedding) + assert isinstance(model.time_embed.guidance_condition_proj, SanaTimesteps) + assert isinstance(model.time_embed.guidance_embedder, SanaTimestepEmbedding) + + torch.manual_seed(31) + hidden_states = torch.randn(1, 4, 3, 4, 4) + encoder_hidden_states = torch.randn(1, 5, 8) + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]]) + timestep = torch.tensor([500.0]) + guidance = torch.tensor([6.0]) + + with torch.no_grad(): + expected = reference( + hidden_states, + encoder_hidden_states, + timestep, + guidance=guidance, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + actual = model( + hidden_states, + encoder_hidden_states, + timestep, + guidance=guidance, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + + assert isinstance(actual, tuple) + assert len(actual) == len(expected) == 1 + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + + +def test_native_transformer_config_filters_diffusers_metadata(): + from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerConfig + + config = SanaVideoTransformerConfig.from_dict( + _TINY_CONFIG | {"_class_name": "SanaVideoTransformer3DModel", "_diffusers_version": "0.38.0"} + ) + assert config.patch_size == (1, 2, 2) + + with pytest.raises(ValueError, match="unsupported_field"): + SanaVideoTransformerConfig.from_dict(_TINY_CONFIG | {"unsupported_field": True}) + + +def test_dual_vae_variant_resolution(): + from diffusers.pipelines.pipeline_utils import DiffusionPipeline + + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, + ) + from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin + from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + SanaVideoPipeline, + _resolve_vae_class_and_dtype, + ) + + assert not issubclass(SanaVideoPipeline, DiffusionPipeline) + assert issubclass(SanaVideoPipeline, ProgressBarMixin) + + vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLWan", torch.bfloat16) + assert vae_class is DistributedAutoencoderKLWan + assert dtype is torch.float32 + + vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLLTX2Video", torch.bfloat16) + assert vae_class is DistributedAutoencoderKLLTX2Video + assert dtype is torch.bfloat16 + + with pytest.raises(ValueError, match="Unsupported SANA-Video VAE"): + _resolve_vae_class_and_dtype("UnknownVAE", torch.bfloat16) diff --git a/tests/diffusion/models/sana_wm/__init__.py b/tests/diffusion/models/sana_wm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py b/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py new file mode 100644 index 000000000..cea8aeb76 --- /dev/null +++ b/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Equivalence tests for SANA-WM's camera-branch reduce-scatter. + +``out_proj_cam`` is row-parallel, so every rank produces a partial sum over the +full hidden width, while the GDN stream it is added to is TP-local. Instead of +all-reducing and then discarding all but one slice, the branch reduce-scatters +straight to the local shard. These tests pin the shard each rank ends up with, +including the bias slice, on both code paths: + +* the process-group path, which drives ``torch.distributed.reduce_scatter_tensor`` + directly because vLLM-Omni's diffusion ``GroupCoordinator`` does not implement + ``reduce_scatter``; +* the all-reduce-and-slice fallback taken when the coordinator exposes no + ``device_group``. + +CPU-only: the collectives and the rank/world-size lookups are mocked, so a +single process can simulate an arbitrary TP degree. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmSelfAttention + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer" + +HIDDEN = 24 + + +def _host(bias: torch.Tensor | None) -> SimpleNamespace: + """A stand-in ``self`` exposing only what the method reads.""" + return SimpleNamespace(out_proj_cam=SimpleNamespace(bias=bias)) + + +def _fake_reduce_scatter_tensor(rank: int, tp_size: int): + """Model ``reduce_scatter_tensor``: split dim 0 and keep chunk ``rank``. + + The test drives a single process, so the other ranks' partials are already + folded into the input; scattering is all this has to reproduce. + """ + + def _impl(output: torch.Tensor, input_: torch.Tensor, group=None) -> None: + output.copy_(input_.chunk(tp_size, dim=0)[rank]) + + return _impl + + +def _tp_group(with_device_group: bool): + return SimpleNamespace(device_group=object() if with_device_group else None) + + +@pytest.mark.parametrize("tp_size", [2, 4]) +@pytest.mark.parametrize("rank", [0, 1]) +def test_process_group_path_matches_all_reduce_then_slice(tp_size: int, rank: int) -> None: + """The scattered shard equals the slice the all-reduce path would take.""" + local_width = HIDDEN // tp_size + torch.manual_seed(0) + contrib = torch.randn(2, 3, HIDDEN) + bias = torch.randn(HIDDEN) + reference = torch.zeros(2, 3, local_width) + + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size), + patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank), + patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(True)), + patch("torch.distributed.reduce_scatter_tensor", side_effect=_fake_reduce_scatter_tensor(rank, tp_size)), + ): + out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference) + + start = rank * local_width + expected = (contrib + bias)[..., start : start + local_width] + assert out.shape == (2, 3, local_width) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_fallback_path_matches_the_process_group_path(rank: int) -> None: + """Without a device_group the method all-reduces and slices, same result.""" + tp_size = 2 + local_width = HIDDEN // tp_size + torch.manual_seed(1) + contrib = torch.randn(2, 3, HIDDEN) + bias = torch.randn(HIDDEN) + reference = torch.zeros(2, 3, local_width) + + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size), + patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank), + patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(False)), + patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t), + ): + out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference) + + start = rank * local_width + torch.testing.assert_close(out, (contrib + bias)[..., start : start + local_width]) + + +def test_tp1_adds_bias_without_a_collective() -> None: + contrib = torch.randn(2, 3, HIDDEN) + bias = torch.randn(HIDDEN) + + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1), + patch("torch.distributed.reduce_scatter_tensor") as collective, + ): + out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, contrib) + + assert collective.call_count == 0 + torch.testing.assert_close(out, contrib + bias) + + +def test_rejects_a_width_that_does_not_shard() -> None: + """A local width that is not 1/tp of the contribution is a wiring bug.""" + contrib = torch.randn(1, 2, HIDDEN) + reference = torch.zeros(1, 2, HIDDEN // 3) # tp=2 would need HIDDEN // 2 + + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2), + patch("torch.distributed.reduce_scatter_tensor"), + pytest.raises(ValueError, match="camera contribution width mismatch"), + ): + SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(None), contrib, reference) diff --git a/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py b/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py new file mode 100644 index 000000000..b01ef02ce --- /dev/null +++ b/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Equivalence tests for SANA-WM's TensorParallelRMSNorm q/k fusion. + +`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor +and issues a single tensor-parallel all-reduce instead of two. These tests lock +in that it is numerically identical to applying `norm_q`/`norm_k` separately, +and that it really issues exactly one (fused) all-reduce. + +CPU-only: the collective is mocked, so a single process simulates one rank +(all-reduce is identity). This exercises the cat/split logic a real multi-rank +run depends on. +""" + +from unittest.mock import patch + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import ( + TensorParallelRMSNorm, + fused_qk_rms_norm, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer" + + +def _make_norms(dim: int, tp_size: int) -> tuple[TensorParallelRMSNorm, TensorParallelRMSNorm]: + """Two norms with distinct random weights so a swapped q/k slice is caught.""" + torch.manual_seed(0) + norm_q = TensorParallelRMSNorm(dim, tp_size=tp_size) + norm_k = TensorParallelRMSNorm(dim, tp_size=tp_size) + norm_q.weight.data.normal_() + norm_k.weight.data.normal_() + return norm_q, norm_k + + +def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor: + """Single-rank all-reduce: the reduction over one rank is the identity.""" + return tensor + + +def test_fused_matches_separate() -> None: + dim = 64 + norm_q, norm_k = _make_norms(dim, tp_size=2) + q = torch.randn(2, 5, dim) + k = torch.randn(2, 5, dim) + + with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce): + ref_q, ref_k = norm_q(q), norm_k(k) + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + + torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0) + torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0) + + +def test_fused_issues_one_all_reduce() -> None: + dim = 32 + norm_q, norm_k = _make_norms(dim, tp_size=4) + q = torch.randn(1, 3, dim) + k = torch.randn(1, 3, dim) + + with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective: + fused_qk_rms_norm(norm_q, norm_k, q, k) + assert collective.call_count == 1 + # Both sums travel in one tensor: last dim is 2, not 1. + assert collective.call_args.args[0].shape[-1] == 2 + + collective.reset_mock() + norm_q(q), norm_k(k) + assert collective.call_count == 2 + + +@pytest.mark.parametrize( + "norms", + [ + (RMSNorm(16), RMSNorm(16)), + (nn.Identity(), nn.Identity()), + ], + ids=["plain_rms_norm", "identity"], +) +def test_falls_back_for_non_tp_norms(norms: tuple[nn.Module, nn.Module]) -> None: + """qk_norm=False and TP=1 both hand in norms the fused path cannot use.""" + norm_q, norm_k = norms + q = torch.randn(2, 4, 16) + k = torch.randn(2, 4, 16) + + with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective: + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + assert collective.call_count == 0 + + torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0) + torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0) + + +def test_tp1_skips_the_collective() -> None: + norm_q, norm_k = _make_norms(16, tp_size=1) + q = torch.randn(2, 4, 16) + k = torch.randn(2, 4, 16) + + with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective: + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + assert collective.call_count == 0 + + torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0) + torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0) diff --git a/tests/diffusion/models/sana_wm/test_request.py b/tests/diffusion/models/sana_wm/test_request.py new file mode 100644 index 000000000..7e40eedaf --- /dev/null +++ b/tests/diffusion/models/sana_wm/test_request.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for SANA-WM request normalization.""" + +import pytest + +from vllm_omni.diffusion.models.sana_wm.request import ( + SANA_WM_DEFAULT_HEIGHT, + SANA_WM_DEFAULT_NUM_FRAMES, + SANA_WM_DEFAULT_WIDTH, + normalize_sana_wm_payload, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _prompt(**sana_wm): + """A valid request, with the action rollout kept in step with num_frames.""" + block = {"num_frames": SANA_WM_DEFAULT_NUM_FRAMES, **sana_wm} + block.setdefault("action", f"w-{block['num_frames'] - 1}") + return { + "prompt": "a quiet city street", + "multi_modal_data": {"image": object()}, + "sana_wm": block, + } + + +def _identity_poses(count): + return [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]] for _ in range(count)] + + +def _camera_prompt(**camera): + prompt = _prompt(num_frames=9) + prompt["sana_wm"].pop("action") + prompt["sana_wm"]["camera"] = {"poses": _identity_poses(9), **camera} + return prompt + + +def test_defaults_are_latent_aligned(): + payload = normalize_sana_wm_payload(_prompt())["additional_information"]["sana_wm"] + assert payload["num_frames"] == SANA_WM_DEFAULT_NUM_FRAMES + assert payload["height"] == SANA_WM_DEFAULT_HEIGHT + assert payload["width"] == SANA_WM_DEFAULT_WIDTH + + +@pytest.mark.parametrize("num_frames", [9, 33, 161]) +def test_accepts_aligned_geometry(num_frames): + payload = normalize_sana_wm_payload(_prompt(num_frames=num_frames))["additional_information"]["sana_wm"] + assert payload["num_frames"] == num_frames + + +@pytest.mark.parametrize("num_frames", [8, 10, 160]) +def test_rejects_unaligned_num_frames(num_frames): + # The VAE compresses 8x temporally; an unaligned count would be floored. + with pytest.raises(ValueError, match="num_frames"): + normalize_sana_wm_payload(_prompt(num_frames=num_frames)) + + +@pytest.mark.parametrize(("height", "width"), [(700, 1280), (704, 1270), (703, 1279)]) +def test_rejects_unaligned_resolution(height, width): + with pytest.raises(ValueError, match="divisible by 32"): + normalize_sana_wm_payload(_prompt(height=height, width=width)) + + +def test_normalization_is_idempotent(): + once = normalize_sana_wm_payload(_prompt(num_frames=33)) + twice = normalize_sana_wm_payload(once) + assert once["additional_information"]["sana_wm"] == twice["additional_information"]["sana_wm"] + + +def test_rejects_action_shorter_than_num_frames(): + # Used to shrink the camera tensor silently, then fail deep in the + # transformer on a camera/latent frame mismatch. + with pytest.raises(ValueError, match="rolls out 33 frames but num_frames is 161"): + normalize_sana_wm_payload(_prompt(action="w-32")) + + +def test_rejects_action_longer_than_num_frames(): + with pytest.raises(ValueError, match="rolls out 161 frames but num_frames is 9"): + normalize_sana_wm_payload(_prompt(num_frames=9, action="w-160")) + + +def test_rejects_oversized_action_without_expanding_it(): + # A 12-byte request body: expanding this rollout before checking the frame + # count allocated ~16 GB. Rejection has to be arithmetic, so this test also + # fails (by hanging/OOM) if the duration is ever expanded again. + with pytest.raises(ValueError, match="rolls out 1000000001 frames"): + normalize_sana_wm_payload(_prompt(action="w-1000000000")) + + +@pytest.mark.parametrize("action", ["", " ", "w-0", "w-abc", "w", "x-4"]) +def test_rejects_malformed_action(action): + with pytest.raises(ValueError): + normalize_sana_wm_payload(_prompt(action=action)) + + +@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")]) +def test_rejects_unusable_focal_length(value): + # compute_raymap divides by fx/fy; a zero or non-finite focal length turns + # the whole Plucker map into NaN and the response into a NaN video. + intrinsics = {"fx": value, "fy": 640, "cx": 640, "cy": 352} + with pytest.raises(ValueError, match="fx"): + normalize_sana_wm_payload(_prompt(intrinsics=intrinsics)) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_rejects_non_finite_principal_point(value): + intrinsics = {"fx": 640, "fy": 640, "cx": value, "cy": 352} + with pytest.raises(ValueError, match="must be finite"): + normalize_sana_wm_payload(_prompt(intrinsics=intrinsics)) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_rejects_non_finite_speeds(value): + # `parsed <= 0` is False for NaN, so the positivity check alone let it pass. + with pytest.raises(ValueError, match="must be finite"): + normalize_sana_wm_payload(_prompt(translation_speed=value)) + + +def test_rejects_non_finite_pose_entry(): + prompt = _camera_prompt() + prompt["sana_wm"]["camera"]["poses"][3][2][1] = float("nan") + with pytest.raises(ValueError, match="must be finite"): + normalize_sana_wm_payload(prompt) + + +def test_accepts_default_camera_metadata(): + payload = normalize_sana_wm_payload(_camera_prompt())["additional_information"]["sana_wm"] + assert payload["camera"]["poses"] + # The unread fields are not carried into the canonical payload. + assert set(payload["camera"]) == {"poses"} + + +@pytest.mark.parametrize( + ("field", "value"), + [("format", "w2c_4x4"), ("coordinate_system", "opengl")], +) +def test_rejects_unsupported_camera_convention(field, value): + # Nothing downstream reads these, so accepting one would mean interpreting + # the poses under the wrong convention. + with pytest.raises(ValueError, match=field.split("_")[0]): + normalize_sana_wm_payload(_camera_prompt(**{field: value})) diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py new file mode 100644 index 000000000..26af9a259 --- /dev/null +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression test: SenseNovaU1ForCausalLM instantiation under the diffusion +config shim must not crash on missing ``head_dtype`` and the resulting +LogitsProcessor must have ``head_dtype is None`` (= use model dtype).""" + +import os + +import pytest +import torch +from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config +from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, +) + +from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import ( + SenseNovaU1ForCausalLM, +) +from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig +from vllm_omni.transformers_utils.configs.sensenova_u1 import ( + SenseNovaU1Config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _build_tiny_sensenova_u1(quant_config=None, prefix="model"): + """Build a minimal SenseNovaU1ForCausalLM for unit tests.""" + config = SenseNovaU1Config( + llm_config={ + "hidden_size": 64, + "num_attention_heads": 2, + "num_hidden_layers": 1, + "intermediate_size": 128, + "vocab_size": 32, + "max_position_embeddings": 128, + "max_position_embeddings_hw": 128, + }, + ) + return SenseNovaU1ForCausalLM(config.llm_config, quant_config=quant_config, prefix=prefix) + + +def test_logits_processor_head_dtype_under_diffusion_shim(): + # Initialize the distributed environment + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29543") + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") + initialize_model_parallel() + + # Create a vLLM config with a wrapped diffusion config + fake_diff_config = _DiffusionVllmModelConfig( + model="sensenova-u1-test", + dtype=torch.bfloat16, + max_model_len=8192, + original_max_model_len=8192, + ) + vllm_config = VllmConfig( + device_config=DeviceConfig(device="cpu"), + ) + vllm_config.modelconfig = fake_diff_config # type: ignore[assignment] + + # Build a tiny version of the Causal LM component + with set_current_vllm_config(vllm_config): + model = _build_tiny_sensenova_u1() + + # Ensure that we have a set head_dtype attribute. Currently, the head_dtype + # is set to None + head_dtype = model.logits_processor.head_dtype + assert head_dtype is None + cleanup_dist_env_and_memory() diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py index 9d5b11d3b..12c0ebf45 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py @@ -8,15 +8,28 @@ import torch from PIL import Image from torch import nn -from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar +from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import ( Wan22I2VPipeline, + get_wan22_i2v_post_process_func, get_wan22_i2v_pre_process_func, ) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +def test_wan22_i2v_postprocess_honors_request_output_type() -> None: + video = torch.zeros(1, 4, 1, 2, 2) + + output = get_wan22_i2v_post_process_func(SimpleNamespace())( + video, + sampling_params=SimpleNamespace(output_type="latent"), + ) + + assert output is video + + def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: pipeline = object.__new__(Wan22I2VPipeline) nn.Module.__init__(pipeline) @@ -31,6 +44,29 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: return pipeline +def _make_i2v_sampling(**overrides): + values = { + "height": 16, + "width": 16, + "num_frames": 5, + "num_inference_steps": 1, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "guidance_scale_2": None, + "guidance_scale_2_provided": False, + "boundary_ratio": None, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 8, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None: preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) request = SimpleNamespace( @@ -53,6 +89,88 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None: preprocess(missing_image) +def _make_i2v_preprocess_request(last_image): + return SimpleNamespace( + prompt={ + "prompt": "p", + "multi_modal_data": { + "image": Image.new("RGB", (320, 160), "red"), + "last_image": last_image, + }, + }, + sampling_params=SimpleNamespace(height=16, width=16), + ) + + +def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + result = preprocess(_make_i2v_preprocess_request([])) + + assert result.prompt["multi_modal_data"]["last_image"] is None + assert result.batch_compatibility_key == ("wan22_i2v_last_image", False) + + +def test_i2v_preprocess_unwraps_single_last_image_list() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + last_image = Image.new("RGB", (16, 16), "blue") + + result = preprocess(_make_i2v_preprocess_request([last_image])) + + assert result.prompt["multi_modal_data"]["last_image"] is last_image + assert result.batch_compatibility_key == ("wan22_i2v_last_image", True) + + +@pytest.mark.parametrize( + "last_image", + [ + Image.new("RGB", (16, 16), "blue"), + torch.zeros(3, 16, 16), + ], +) +def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + result = preprocess(_make_i2v_preprocess_request(last_image)) + + assert result.prompt["multi_modal_data"]["last_image"] is last_image + assert result.batch_compatibility_key == ("wan22_i2v_last_image", True) + + +def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + path = tmp_path / "last.png" + Image.new("RGB", (16, 16), "blue").save(path) + + result = preprocess(_make_i2v_preprocess_request(str(path))) + + last_image = result.prompt["multi_modal_data"]["last_image"] + assert isinstance(last_image, Image.Image) + assert last_image.mode == "RGB" + assert result.batch_compatibility_key == ("wan22_i2v_last_image", True) + + +def test_i2v_preprocess_rejects_multiple_last_images() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + with pytest.raises(ValueError, match="at most one last_image"): + preprocess( + _make_i2v_preprocess_request( + [ + Image.new("RGB", (16, 16), "blue"), + Image.new("RGB", (16, 16), "green"), + ] + ) + ) + + +def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + with pytest.raises(TypeError, match="Unsupported last_image format"): + preprocess(_make_i2v_preprocess_request({"not": "an image"})) + + def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None: pipeline = _make_i2v_pipeline(expand_timesteps=True) latents = torch.zeros(1, 4, 2, 4, 4) @@ -126,3 +244,169 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N assert first_frame_mask.shape == (1, 1, 2, 2, 2) assert first_frame_mask[:, :, 0].sum() == 0 assert first_frame_mask[:, :, 1].sum() == 4 + + +def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None: + pipeline = _make_i2v_pipeline(expand_timesteps=True) + generators = [ + torch.Generator(device="cpu").manual_seed(1), + torch.Generator(device="cpu").manual_seed(2), + ] + + latents, condition, first_frame_mask = pipeline.prepare_latents( + image=torch.zeros(2, 3, 16, 16), + batch_size=2, + num_channels_latents=4, + height=16, + width=16, + num_frames=5, + dtype=torch.float32, + device=torch.device("cpu"), + generator=generators, + ) + + assert latents.shape == (2, 4, 2, 2, 2) + assert condition.shape == (2, 4, 1, 2, 2) + assert first_frame_mask.shape == (2, 1, 2, 2, 2) + + +def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_i2v_pipeline(expand_timesteps=True) + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=5.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 5.0 + pipeline.boundary_ratio = 0.875 + pipeline.has_image_encoder = False + pipeline._guidance_scale = None + pipeline._guidance_scale_2 = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + prepare_call = {} + + def _fake_encode_prompt(**kwargs): + batch_size = len(kwargs["prompt"]) + n = batch_size * kwargs["num_videos_per_prompt"] + return torch.zeros(n, 2, 3), None + + def _fake_prepare_latents(**kwargs): + prepare_call.update(kwargs) + batch_size = kwargs["batch_size"] + return ( + kwargs["latents"], + torch.zeros(batch_size, 4, 1, 2, 2), + torch.ones(batch_size, 1, 2, 2, 2), + ) + + pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign] + pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign] + pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign] + + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + latents_a = torch.zeros(2, 4, 2, 2, 2) + latents_b = torch.ones(2, 4, 2, 2, 2) + image_a = torch.zeros(1, 3, 16, 16) + image_b = torch.ones(1, 3, 16, 16) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first", "multi_modal_data": {"image": image_a}}, + sampling_params=_make_i2v_sampling( + generator=gen_a, + latents=latents_a, + num_outputs_per_prompt=2, + ), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "multi_modal_data": {"image": image_b}}, + sampling_params=_make_i2v_sampling( + generator=gen_b, + latents=latents_b, + num_outputs_per_prompt=2, + ), + ), + ] + ) + + outputs = pipeline.forward(batch) + + assert prepare_call["batch_size"] == 4 + assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] + torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b])) + torch.testing.assert_close( + prepare_call["image"], + torch.cat([image_a, image_a, image_b, image_b]), + ) + assert len(outputs) == 2 + torch.testing.assert_close(outputs[0].output, latents_a) + torch.testing.assert_close(outputs[1].output, latents_b) + + +def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_i2v_pipeline(expand_timesteps=True) + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=5.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 5.0 + pipeline.boundary_ratio = 0.875 + pipeline.has_image_encoder = False + pipeline._guidance_scale = None + pipeline._guidance_scale_2 = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}}, + sampling_params=_make_i2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}}, + sampling_params=_make_i2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="image condition"): + pipeline.forward(batch) + + +def test_i2v_forward_rejects_mixed_last_image_presence() -> None: + pipeline = _make_i2v_pipeline(expand_timesteps=True) + image = torch.zeros(1, 3, 16, 16) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "last_image": image}, + }, + sampling_params=_make_i2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "multi_modal_data": {"image": image}}, + sampling_params=_make_i2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="mix of provided and missing last_image"): + pipeline.forward(batch) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 8c06df034..de5abeab0 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import importlib from contextlib import contextmanager from types import SimpleNamespace @@ -9,6 +10,7 @@ import torch from torch import nn from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline +from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -82,6 +84,7 @@ def _make_pipeline() -> Wan22Pipeline: pipeline.vae_scale_factor_spatial = 8 pipeline.boundary_ratio = 0.875 pipeline.expand_timesteps = False + pipeline.is_dmd = False pipeline._guidance_scale = None pipeline._guidance_scale_2 = None pipeline._num_timesteps = None @@ -93,6 +96,29 @@ def _make_pipeline() -> Wan22Pipeline: return pipeline +def _make_sampling(**overrides): + values: dict[str, object] = { + "height": None, + "width": None, + "num_frames": 1, + "num_inference_steps": 2, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "guidance_scale_2": None, + "guidance_scale_2_provided": False, + "boundary_ratio": None, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 32, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + @pytest.mark.parametrize( ("sampling_params_kwargs", "expected_low", "expected_high"), [ @@ -129,9 +155,10 @@ def test_forward_delegates_denoising_to_diffuse( ) batch = DiffusionRequestBatch(requests=[mock_req]) - output = pipeline.forward(batch) + outputs = pipeline.forward(batch) - assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8))) + assert len(outputs) == 1 + assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8))) assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8)) assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps) assert captured["guidance_low"] == expected_low @@ -142,6 +169,134 @@ def test_forward_delegates_denoising_to_diffuse( assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))] +def test_forward_batches_text_generators_latents_and_splits_outputs() -> None: + pipeline = _make_pipeline() + encode_call = {} + prepare_call = {} + + def _fake_encode_prompt(**kwargs): + encode_call.update(kwargs) + batch_size = len(kwargs["prompt"]) + n = batch_size * kwargs["num_videos_per_prompt"] + return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1) + + def _fake_prepare_latents(**kwargs): + prepare_call.update(kwargs) + return kwargs["latents"] + + pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign] + pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign] + pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign] + + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + latents_a = torch.zeros(2, 4, 1, 2, 2) + latents_b = torch.ones(2, 4, 1, 2, 2) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first", "negative_prompt": "bad first"}, + sampling_params=_make_sampling( + generator=gen_a, + latents=latents_a, + num_outputs_per_prompt=2, + ), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "negative_prompt": "bad second"}, + sampling_params=_make_sampling( + generator=gen_b, + latents=latents_b, + num_outputs_per_prompt=2, + ), + ), + ] + ) + + outputs = pipeline.forward(batch) + + assert encode_call["prompt"] == ["first", "second"] + assert encode_call["negative_prompt"] == ["bad first", "bad second"] + assert prepare_call["batch_size"] == 4 + assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] + torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b])) + assert len(outputs) == 2 + torch.testing.assert_close(outputs[0].output, latents_a) + torch.testing.assert_close(outputs[1].output, latents_b) + + +def test_forward_batches_precomputed_prompt_embeddings() -> None: + pipeline = _make_pipeline() + diffuse_call = {} + pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run") # type: ignore[method-assign] + pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2) # type: ignore[method-assign] + + def _fake_diffuse(**kwargs): + diffuse_call.update(kwargs) + return kwargs["latents"] + + pipeline.diffuse = _fake_diffuse # type: ignore[method-assign] + embeds_a = torch.zeros(3, 4) + embeds_b = torch.ones(3, 4) + negative_a = torch.full((3, 4), 2.0) + negative_b = torch.full((3, 4), 3.0) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a}, + sampling_params=_make_sampling(guidance_scale=4.0), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b}, + sampling_params=_make_sampling(guidance_scale=4.0), + ), + ] + ) + + outputs = pipeline.forward(batch) + + torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b])) + torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b])) + assert len(outputs) == 2 + + +def test_prepare_latents_with_request_generators_matches_single_generation() -> None: + pipeline = _make_pipeline() + kwargs = { + "num_channels_latents": 4, + "height": 16, + "width": 16, + "num_frames": 5, + "dtype": torch.float32, + "device": torch.device("cpu"), + } + + batched = Wan22Pipeline.prepare_latents( + pipeline, + batch_size=2, + generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)], + **kwargs, + ) + singles = torch.cat( + [ + Wan22Pipeline.prepare_latents( + pipeline, + batch_size=1, + generator=torch.Generator().manual_seed(seed), + **kwargs, + ) + for seed in (1, 2) + ] + ) + + torch.testing.assert_close(batched, singles) + assert not torch.equal(batched[0], batched[1]) + + def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None: pipeline = _make_pipeline() latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32) @@ -186,3 +341,125 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None: (3.0, 3, 28.0, False), ] assert torch.equal(result, torch.full_like(latents, 10.0)) + + +class _StubDMDScheduler: + def __init__(self) -> None: + self.predict_clean_calls: list[tuple[float, float, float]] = [] + self.add_noise_calls: list[tuple[float, float, float]] = [] + + def predict_clean(self, model_output, sample, timestep): + self.predict_clean_calls.append((float(model_output.mean()), float(sample.mean()), float(timestep))) + return sample - model_output + + def add_noise(self, clean_sample, noise, timestep): + self.add_noise_calls.append((float(clean_sample.mean()), float(noise.mean()), float(timestep))) + return clean_sample + 10.0 + + +def test_diffuse_dmd_predicts_clean_and_renoises_between_steps(monkeypatch) -> None: + pipeline = _make_pipeline() + pipeline.is_dmd = True + pipeline.scheduler = _StubDMDScheduler() + latents = torch.zeros((1, 1, 1, 1, 1), dtype=torch.float32) + timesteps = torch.tensor([1000.0, 757.0, 522.0]) + + pipeline.predict_noise_maybe_with_cfg = lambda **kwargs: torch.ones_like(latents) # type: ignore[method-assign] + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.randn_tensor", + lambda *args, **kwargs: torch.full(args[0], 2.0, dtype=kwargs["dtype"]), + ) + + result = pipeline.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=torch.zeros(1, 8), + negative_prompt_embeds=None, + guidance_low=1.0, + guidance_high=1.0, + boundary_timestep=None, + dtype=torch.float32, + attention_kwargs={}, + generator=torch.Generator(device="cpu").manual_seed(1), + ) + + assert pipeline.scheduler.predict_clean_calls == [ + (1.0, 0.0, 1000.0), + (1.0, 9.0, 757.0), + (1.0, 18.0, 522.0), + ] + assert pipeline.scheduler.add_noise_calls == [ + (-1.0, 2.0, 757.0), + (8.0, 2.0, 522.0), + ] + torch.testing.assert_close(result, torch.tensor([[[[[17.0]]]]])) + + +def _make_gate_loading_pipeline(): + pipeline = Wan22Pipeline.__new__(Wan22Pipeline) + nn.Module.__init__(pipeline) + gate = WanSelfAttention.__new__(WanSelfAttention) + nn.Module.__init__(gate) + gate.to_gate_compress = nn.Linear(1, 1) + pipeline.gate_holder = gate + return pipeline, gate + + +@pytest.mark.parametrize( + ("module_name", "class_name"), + [ + ("pipeline_wan2_2", "Wan22Pipeline"), + ("pipeline_wan2_2_i2v", "Wan22I2VPipeline"), + ("pipeline_wan2_2_s2v", "Wan22S2VPipeline"), + ("pipeline_wan2_2_vace", "Wan22VACEPipeline"), + ], +) +def test_wan_pipeline_loaders_share_optional_gate_cleanup(monkeypatch, module_name, class_name) -> None: + module = importlib.import_module(f"vllm_omni.diffusion.models.wan2_2.{module_name}") + pipeline_cls = getattr(module, class_name) + pipeline = pipeline_cls.__new__(pipeline_cls) + expected = {"loaded"} + + def fake_loader(model, weights): + assert model is pipeline + assert list(weights) == [("weight", torch.ones(1))] + return expected + + monkeypatch.setattr(module, "load_wan_weights_with_optional_gate", fake_loader) + + assert pipeline_cls.load_weights(pipeline, iter((("weight", torch.ones(1)),))) is expected + + +def test_load_weights_removes_unloaded_vsa_gate(monkeypatch) -> None: + pipeline, gate = _make_gate_loading_pipeline() + + class _Loader: + def __init__(self, model): + del model + + def load_weights(self, weights): + return {name for name, _ in weights} + + monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader) + pipeline.load_weights(iter((("other.weight", torch.ones(1)),))) + + assert pipeline.has_gate_compress_weights is False + assert gate.to_gate_compress is None + + +def test_load_weights_keeps_trained_vsa_gate(monkeypatch) -> None: + pipeline, gate = _make_gate_loading_pipeline() + original_gate = gate.to_gate_compress + + class _Loader: + def __init__(self, model): + del model + + def load_weights(self, weights): + return {name for name, _ in weights} + + monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader) + pipeline.load_weights(iter((("gate_holder.to_gate_compress.weight", torch.ones(1)),))) + + assert pipeline.has_gate_compress_weights is True + assert gate.to_gate_compress is original_gate diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py index 314717869..5a37e14f2 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py @@ -10,6 +10,7 @@ import torch import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( create_transformer_from_config, + get_wan22_post_process_func, load_transformer_config, retrieve_latents, ) @@ -26,6 +27,17 @@ class _LatentDist: return torch.tensor([2.0]) +def test_wan22_postprocess_honors_request_output_type() -> None: + video = torch.zeros(1, 4, 1, 2, 2) + + output = get_wan22_post_process_func(SimpleNamespace())( + video, + sampling_params=SimpleNamespace(output_type="latent"), + ) + + assert output is video + + def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None: generator = torch.Generator(device="cpu") diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py index 4440a8f35..e5bae65df 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -1,14 +1,307 @@ +from types import SimpleNamespace from unittest.mock import MagicMock, patch +import numpy as np +import PIL.Image import pytest import torch import torch.nn as nn +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import ( + Wan22S2VPipeline, + _make_clip_generators, +) from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _make_s2v_sampling(**overrides): + values = { + "height": 16, + "width": 16, + "num_frames": 8, + "num_inference_steps": 1, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 8, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _make_s2v_validation_pipeline() -> Wan22S2VPipeline: + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(dtype=torch.float32) + pipeline.vae_scale_factor_spatial = 8 + pipeline.resolution_divisor = 16 + pipeline.motion_frames = 7 + pipeline.drop_first_motion = True + pipeline._DEFAULT_INFER_FRAMES = 8 + pipeline._guidance_scale = None + pipeline.check_inputs = lambda *args, **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + return pipeline + + +def test_s2v_predict_noise_keeps_batch_dimension() -> None: + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + transformer = MagicMock() + transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),) + pipeline.transformer = transformer + + result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2)) + + assert result.shape == (1, 16, 2, 2, 2) + + +def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None: + explicit_generator = torch.Generator(device="cpu").manual_seed(999) + + generators = _make_clip_generators( + seeds=[1234, 5678, None, None], + request_generators=[ + torch.Generator(device="cpu").manual_seed(1234), + torch.Generator(device="cpu").manual_seed(5678), + explicit_generator, + None, + ], + clip_index=2, + device=torch.device("cpu"), + ) + + assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2] + assert generators[2] is explicit_generator + + +def test_s2v_forward_rejects_different_num_repeat() -> None: + pipeline = _make_s2v_validation_pipeline() + pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16) # type: ignore[method-assign] + image = PIL.Image.new("RGB", (16, 16)) + audio = np.zeros(16000, dtype=np.float32) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"num_repeat": 1}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"num_repeat": 2}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="num_repeat"): + pipeline.forward(batch) + + +def test_s2v_forward_rejects_different_audio_time_lengths() -> None: + pipeline = _make_s2v_validation_pipeline() + pipeline.encode_audio = lambda audio, **kwargs: ( # type: ignore[method-assign] + torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24), + 1, + 16 if audio[0] == 0 else 24, + ) + image = PIL.Image.new("RGB", (16, 16)) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="audio embedding time lengths"): + pipeline.forward(batch) + + +def test_s2v_forward_rejects_different_init_first_frame() -> None: + pipeline = _make_s2v_validation_pipeline() + image = PIL.Image.new("RGB", (16, 16)) + audio = np.zeros(16000, dtype=np.float32) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"init_first_frame": True}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"init_first_frame": False}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="init_first_frame"): + pipeline.forward(batch) + + +def test_s2v_forward_rejects_different_raw_audio_shapes() -> None: + pipeline = _make_s2v_validation_pipeline() + pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16) # type: ignore[method-assign] + image = PIL.Image.new("RGB", (16, 16)) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="raw audio shapes and sample rates"): + pipeline.forward(batch) + + +def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = MagicMock() + pipeline.transformer.dtype = torch.float32 + pipeline.transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.float32)]) + pipeline.transformer.casual_audio_encoder = None + pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio} + pipeline.vae = MagicMock() + pipeline.vae.dtype = torch.float32 + pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0])) + pipeline.vae_scale_factor_spatial = 8 + pipeline.resolution_divisor = 16 + pipeline.motion_frames = 7 + pipeline.drop_first_motion = True + pipeline._DEFAULT_INFER_FRAMES = 8 + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline.check_inputs = lambda *args, **kwargs: None + pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3))) + pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8)) + pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2)) + pipeline.prepare_motion_latents = MagicMock( + side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2) + ) + pipeline._denormalize_latents = lambda latents: latents + pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"]) + + image = PIL.Image.new("RGB", (16, 16)) + audio_a = np.zeros(16000, dtype=np.float32) + audio_b = np.ones(16000, dtype=np.float32) + latents_a = torch.zeros(2, 16, 2, 2, 2) + latents_b = torch.ones(2, 16, 2, 2, 2) + generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)] + generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)] + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "negative_prompt": "negative first", + "multi_modal_data": {"image": image, "audio": audio_a}, + }, + sampling_params=_make_s2v_sampling( + num_outputs_per_prompt=2, + generator=generators_a, + latents=latents_a, + ), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "negative_prompt": "negative second", + "multi_modal_data": {"image": image, "audio": audio_b}, + }, + sampling_params=_make_s2v_sampling( + num_outputs_per_prompt=2, + generator=generators_b, + latents=latents_b, + ), + ), + ] + ) + + with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform: + platform.is_available.return_value = False + outputs = pipeline.forward(batch) + + assert len(outputs) == 2 + assert outputs[0].output[0].shape[0] == 2 + assert outputs[1].output[0].shape[0] == 2 + assert outputs[0].output[1].shape == (16000,) + assert outputs[1].output[1].shape == (16000,) + np.testing.assert_array_equal(outputs[0].output[1], audio_a) + np.testing.assert_array_equal(outputs[1].output[1], audio_b) + assert np.shares_memory(outputs[0].output[1], audio_a) + assert np.shares_memory(outputs[1].output[1], audio_b) + torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b])) + assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b + assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"] + assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"] + + def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks(): model = object.__new__(WanS2VTransformer3DModel) nn.Module.__init__(model) @@ -260,31 +553,37 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean # -- Build request -- - from vllm_omni.diffusion.request import OmniDiffusionRequest - - sampling_params = MagicMock() - sampling_params.height = 704 - sampling_params.width = 1024 - sampling_params.num_frames = 80 - sampling_params.num_inference_steps = 5 - sampling_params.guidance_scale = 4.5 - sampling_params.guidance_scale_provided = True - sampling_params.generator = None - sampling_params.seed = 42 + sampling_params = _make_s2v_sampling( + height=704, + width=1024, + num_frames=80, + num_inference_steps=5, + guidance_scale=4.5, + generator=torch.Generator(device="cpu").manual_seed(42), + max_sequence_length=512, + ) ref_image = PIL.Image.new("RGB", (1024, 704)) audio_data = np.zeros(16000, dtype=np.float32) - req = MagicMock(spec=OmniDiffusionRequest) - req.prompts = [ - { - "prompt": "test prompt", - "negative_prompt": "bad quality", - "multi_modal_data": {"image": ref_image, "audio": audio_data}, - "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False}, - } - ] - req.sampling_params = sampling_params + req = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="test", + prompt={ + "prompt": "test prompt", + "negative_prompt": "bad quality", + "multi_modal_data": {"image": ref_image, "audio": audio_data}, + "additional_information": { + "audio_path": audio_data, + "pose_video": None, + "init_first_frame": False, + }, + }, + sampling_params=sampling_params, + ) + ] + ) # -- Mock methods that use complex internal state -- pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80)) @@ -295,7 +594,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False) # Mock predict_noise_maybe_with_cfg from CFGParallelMixin - pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128)) + pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128)) pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline) # Mock platform methods @@ -306,7 +605,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): with patch( "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000) ): - result = Wan22S2VPipeline.forward(pipeline, req=req) + result = Wan22S2VPipeline.forward(pipeline, req=req)[0] # -- Assertions -- # Text encoder was called diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py index 714b79d50..6534b7cff 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py @@ -8,13 +8,14 @@ import torch from PIL import Image from torch import nn -from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar +from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import ( Wan22VACEPipeline, create_vace_transformer_from_config, get_wan22_vace_pre_process_func, ) from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -32,6 +33,27 @@ def _make_vace_pipeline() -> Wan22VACEPipeline: return pipeline +def _make_vace_sampling(**overrides): + values = { + "height": 16, + "width": 16, + "num_frames": 5, + "num_inference_steps": 1, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "boundary_ratio": None, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 8, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None: preprocess = get_wan22_vace_pre_process_func(SimpleNamespace()) ref = Image.new("RGB", (320, 160), "green") @@ -138,3 +160,189 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None: assert calls[0]["positive_kwargs"]["vace_context"] is vace_context assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75 torch.testing.assert_close(result, torch.ones_like(latents)) + + +def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_vace_pipeline() + pipeline.transformer.vace_patch_embedding = None + pipeline.transformer_2 = None + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=3.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 3.0 + pipeline.boundary_ratio = None + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + prepare_call = {} + + def _fake_encode_prompt(**kwargs): + batch_size = len(kwargs["prompt"]) + n = batch_size * kwargs["num_videos_per_prompt"] + return torch.zeros(n, 2, 3), None + + def _fake_prepare_latents(**kwargs): + prepare_call.update(kwargs) + return kwargs["latents"] + + pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign] + pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign] + pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign] + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + latents_a = torch.zeros(2, 4, 2, 2, 2) + latents_b = torch.ones(2, 4, 2, 2, 2) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first"}, + sampling_params=_make_vace_sampling( + generator=gen_a, + latents=latents_a, + num_outputs_per_prompt=2, + ), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second"}, + sampling_params=_make_vace_sampling( + generator=gen_b, + latents=latents_b, + num_outputs_per_prompt=2, + ), + ), + ] + ) + + outputs = pipeline.forward(batch) + + assert prepare_call["batch_size"] == 4 + assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] + torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b])) + assert len(outputs) == 2 + torch.testing.assert_close(outputs[0].output, latents_a) + torch.testing.assert_close(outputs[1].output, latents_b) + + +def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_vace_pipeline() + pipeline.transformer.vace_patch_embedding = object() + pipeline.transformer_2 = None + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=3.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 3.0 + pipeline.boundary_ratio = None + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + + def _fake_preprocess(video, **kwargs): + del kwargs + mask = torch.ones_like(video) + return video, mask, [[]] + + pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign] + pipeline.prepare_video_latents = ( # type: ignore[method-assign] + lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2) + ) + pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2) # type: ignore[method-assign] + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)}, + }, + sampling_params=_make_vace_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)}, + }, + sampling_params=_make_vace_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="VACE condition"): + pipeline.forward(batch) + + +def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_vace_pipeline() + pipeline.transformer.vace_patch_embedding = object() + pipeline.transformer_2 = None + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=3.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 3.0 + pipeline.boundary_ratio = None + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + + def _fake_preprocess(video, mask, reference_images, **kwargs): + del kwargs + return video, mask, [reference_images] + + pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign] + pipeline.prepare_video_latents = ( # type: ignore[method-assign] + lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2) + ) + pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2) # type: ignore[method-assign] + video = torch.zeros(1, 3, 5, 16, 16) + mask = torch.ones_like(video) + reference = torch.zeros(3, 16, 16) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "additional_information": { + "source_video": video, + "mask": mask, + "reference_images": [reference], + }, + }, + sampling_params=_make_vace_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "additional_information": { + "source_video": video, + "mask": mask, + "reference_images": [reference, reference], + }, + }, + sampling_params=_make_vace_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="same number of reference images"): + pipeline.forward(batch) diff --git a/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py b/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py new file mode 100644 index 000000000..8c9208014 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_dmd_predict_clean_and_add_noise_match_flow_formulas() -> None: + scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0) + timestep = torch.tensor(757.0) + sigma = scheduler.sigma_for_timestep(timestep) + sample = torch.tensor([2.0, -1.0]) + model_output = torch.tensor([0.5, 0.25]) + noise = torch.tensor([-0.5, 1.5]) + + clean = scheduler.predict_clean(model_output, sample, timestep) + expected_clean = sample - sigma * model_output + torch.testing.assert_close(clean, expected_clean) + + renoised = scheduler.add_noise(clean, noise, timestep) + expected_renoised = (1.0 - sigma) * expected_clean + sigma * noise + torch.testing.assert_close(renoised, expected_renoised) + + +def test_dmd_sigma_requires_scalar_timestep() -> None: + scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0) + with pytest.raises(ValueError, match="scalar timestep"): + scheduler.sigma_for_timestep(torch.tensor([757.0, 522.0])) diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py new file mode 100644 index 000000000..a3efbb40a --- /dev/null +++ b/tests/diffusion/offloader/test_cuda_host_registration.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import ctypes +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module +from vllm_omni.diffusion.offloader.cuda_host_registration import ( + CudaHostRegistration, + _coalesce_ranges, + _mapped_regions, +) +from vllm_omni.diffusion.offloader.host_registration import ( + HostRegistrationBudgetError, + HostRegistrationCleanupError, + HostRegistrationError, + register_host_mappings, +) +from vllm_omni.host_weight_runtime import MappedHostRegion + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +class _FakeRuntime: + def __init__( + self, + register_results: list[int | Exception], + unregister_results: list[int | Exception] | None = None, + *, + read_only_supported: bool = True, + attribute_result: int = 0, + ) -> None: + self._register_results = iter(register_results) + self._unregister_results = iter(unregister_results or []) + self._read_only_supported = read_only_supported + self._attribute_result = attribute_result + self.attribute_queries: list[tuple[int, int]] = [] + self.registered: list[tuple[int, int, int]] = [] + self.unregistered: list[int] = [] + + def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int: + self.attribute_queries.append((attribute, device)) + ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_supported) + return self._attribute_result + + def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int: + self.registered.append((pointer, size, flags)) + result = next(self._register_results) + if isinstance(result, Exception): + raise result + return result + + def cudaHostUnregister(self, pointer: int) -> int: + self.unregistered.append(pointer) + result = next(self._unregister_results, 0) + if isinstance(result, Exception): + raise result + return result + + @staticmethod + def cudaGetErrorString(error: int) -> str: + return f"error-{error}" + + +def _region( + file_name: str, + address: int, + length: int, + *, + read_only: bool = True, +) -> MappedHostRegion: + return MappedHostRegion(file_name=file_name, address=address, length=length, read_only=read_only) + + +@pytest.fixture(autouse=True) +def _cuda_runtime_prerequisites(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0) + + +def test_platform_dispatch_and_unsupported_device(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + regions = (_region("weights", 0x1000, 4096),) + + def create(actual_regions, *, max_bytes): + assert actual_regions is regions + assert max_bytes == 4096 + return sentinel + + monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create)) + assert register_host_mappings(regions, device=torch.device("cuda"), max_bytes=4096) is sentinel + with pytest.raises(HostRegistrationError, match="not supported on cpu"): + register_host_mappings(regions, device=torch.device("cpu"), max_bytes=None) + + +def test_ranges_align_per_file_without_merging_adjacent_files() -> None: + assert _coalesce_ranges([(0x1003, 4096), (0x2800, 1024)], page_size=4096) == ( + registration_module._AddressRange(0x1000, 0x3000), + ) + assert _mapped_regions( + ( + _region("first.safetensors", 0x1000, 4096), + _region("second.safetensors", 0x2000, 4096), + ) + ) == ( + registration_module._AddressRange(0x1000, 0x2000), + registration_module._AddressRange(0x2000, 0x3000), + ) + + +def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + + with pytest.raises(HostRegistrationError, match="is writable"): + CudaHostRegistration.create((_region("weights", 0x1000, 1, read_only=False),), max_bytes=None) + with pytest.raises(HostRegistrationBudgetError, match="exceeding"): + CudaHostRegistration.create((_region("weights", 0x1003, 4096),), max_bytes=4096) + assert runtime.registered == [] + + +def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0], read_only_supported=False) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + + with pytest.raises(HostRegistrationError, match="does not support read-only"): + CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None) + assert runtime.attribute_queries == [ + (registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0) + ] + assert runtime.registered == [] + + +def test_capability_error_is_consumed_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0], attribute_result=7) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + consumed: list[tuple[object, int]] = [] + monkeypatch.setattr( + registration_module, + "_consume_last_cuda_error", + lambda actual_runtime, error: consumed.append((actual_runtime, error)), + ) + + with pytest.raises(HostRegistrationError, match="error-7"): + CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None) + assert consumed == [(runtime, 7)] + + +def test_partial_registration_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0, 7]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None) + + with pytest.raises(HostRegistrationError, match="error-7"): + CudaHostRegistration.create( + ( + _region("first", 0x1000, 4096), + _region("second", 0x9000, 4096), + ), + max_bytes=None, + ) + assert runtime.unregistered == [0x1000] + + +def test_failed_rollback_exposes_active_registration_for_retry(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0, 7], unregister_results=[9, 0]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None) + + with pytest.raises(HostRegistrationCleanupError, match="rollback errors") as error: + CudaHostRegistration.create( + ( + _region("first", 0x1000, 4096), + _region("second", 0x9000, 4096), + ), + max_bytes=None, + ) + registration = error.value.active_registration + assert registration is not None + assert registration.close() == () + assert runtime.unregistered == [0x1000, 0x1000] + + +def test_successful_registration_retries_failed_close(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0], unregister_results=[9, 0]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None) + registration = CudaHostRegistration.create((_region("weights", 0x1003, 1),), max_bytes=None) + + assert registration.total_bytes == 4096 + assert registration.region_count == 1 + assert registration.close() == ("cudaHostUnregister(0x1000) failed: error-9",) + assert registration.close() == () + assert runtime.unregistered == [0x1000, 0x1000] + + +@pytest.mark.parametrize("pending_error", [0, 801]) +def test_consume_last_cuda_error_accepts_cleared_or_matching_state( + monkeypatch: pytest.MonkeyPatch, + pending_error: int, +) -> None: + runtime = SimpleNamespace(cudaGetLastError=lambda: pending_error) + monkeypatch.setattr( + registration_module.ctypes, + "CDLL", + lambda _name: pytest.fail("the cudart handle should provide cudaGetLastError"), + ) + + registration_module._consume_last_cuda_error(runtime, 801) + + +def test_consume_last_cuda_error_falls_back_to_global_symbol(monkeypatch: pytest.MonkeyPatch) -> None: + global_runtime = SimpleNamespace(cudaGetLastError=lambda: 801) + monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: global_runtime) + + registration_module._consume_last_cuda_error(SimpleNamespace(), 801) + + +def test_clean_rollback_remains_recoverable_when_error_state_cannot_be_cleared( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runtime = _FakeRuntime([0, 7]) + monkeypatch.setattr(runtime, "cudaGetLastError", None, raising=False) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + + def missing_global_symbol(_name: object) -> None: + raise OSError("cudart symbols are local") + + monkeypatch.setattr( + registration_module.ctypes, + "CDLL", + missing_global_symbol, + ) + + with pytest.raises(HostRegistrationError, match="cannot clear") as error: + CudaHostRegistration.create( + ( + _region("first", 0x1000, 4096), + _region("second", 0x9000, 4096), + ), + max_bytes=None, + ) + + assert not isinstance(error.value, HostRegistrationCleanupError) + assert runtime.unregistered == [0x1000] diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 810f0dd52..80a03180f 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -1,25 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import gc +from typing import Any import numpy as np import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test +from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform -AUDIO_MODEL = { +AUDIO_MODEL: dict[str, dict[str, int | None]] = { "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None}, } -IMAGE_MODELS = { +IMAGE_MODELS: dict[str, dict[str, int | None]] = { "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100}, } -MODELS = {**AUDIO_MODEL, **IMAGE_MODELS} +MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_MODELS} MODEL_MARKS = { "riverclouds/qwen_image_random": pytest.mark.core_model, @@ -29,12 +33,12 @@ MODEL_MARKS = { _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"} -AUDIO_MODEL_PARAMS = { +AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = { "runner_params": {}, "sampler_params": {}, } -IMAGE_MODEL_PARAMS = { +IMAGE_MODEL_PARAMS: dict[str, dict[str, Any]] = { "runner_params": {}, "sampler_params": { "height": 256, @@ -43,7 +47,7 @@ IMAGE_MODEL_PARAMS = { } -def inference(model_name: str, offload: bool = True): +def inference(model_name: str, offload: bool = True) -> tuple[float, Any]: gc.collect() current_omni_platform.empty_cache() device_index = current_omni_platform.current_device() @@ -133,6 +137,7 @@ def test_cpu_offload_diffusion_model(model_name: str): threshold = MODELS[model_name][platform] if threshold is None: pytest.skip(f"Threshold not defined for {platform} on {model_name}") + assert threshold is not None assert offload_peak_memory + threshold < no_offload_peak_memory, ( f"Offload peak memory {offload_peak_memory} MB should be less than " diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index 26944b202..8e083605f 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -1,36 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import gc +from typing import Any + import numpy as np import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test +from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform -AUDIO_MODEL = { - "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 1500}, +AUDIO_MODEL: dict[str, dict[str, int | None]] = { + # The inference peak includes backend workspaces as well as model weights. + # On ROCm, large MIOpen workspaces and allocator fragmentation can mask the + # resident-weight reduction, so use a conservative floor that still catches + # a disabled/no-op layerwise offloader. + "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 512}, } -IMAGE_VIDEO_MODELS = { +IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = { "riverclouds/qwen_image_random": {"cuda": 4500, "rocm": None}, # "Wan-AI/Wan2.2-T2V-A14B-Diffusers": {"cuda": 45000, "rocm": None}, } -MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} +MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} MODEL_MARKS = { "riverclouds/qwen_image_random": pytest.mark.core_model, "stabilityai/stable-audio-open-1.0": pytest.mark.full_model, } -AUDIO_MODEL_PARAMS = { +AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = { "runner_params": {}, "sampler_params": {}, } -IMAGE_VIDEO_MODELS_PARAMS = { +IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = { "runner_params": {"boundary_ratio": 0.875, "flow_shift": 5.0}, "sampler_params": {"height": 480, "width": 640, "num_frames": 5}, } @@ -55,11 +65,12 @@ def run_inference( model_name: str, layerwise_offload: bool = False, num_inference_steps: int = 3, -) -> float: +) -> tuple[float, Any]: current_omni_platform.empty_cache() device_index = current_omni_platform.current_device() - monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) - monitor.start() + with current_omni_platform.device(device_index): + free_bytes, total_bytes = current_omni_platform.mem_get_info() + initial_used_mb = (total_bytes - free_bytes) / (1024**2) if model_name in AUDIO_MODEL: params = AUDIO_MODEL_PARAMS @@ -73,22 +84,37 @@ def run_inference( # cache_backend="cache_dit", **params["runner_params"], ) as runner: + # Measure steady-state inference memory, not model construction. Enabling + # layerwise offload first loads the model and then replaces each block's + # device storage with CPU-backed weights. Monitoring that transition + # captures both the original model and temporary staging allocations, + # which is not representative of layerwise-offloaded inference. + monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) current_omni_platform.reset_peak_memory_stats() + monitor.start() - # Refer to tests/e2e/offline_inference/test_wan22.py - # Use minimal settings for testing - output = runner.omni.generate( - "A cat sitting on a table", - OmniDiffusionSamplingParams( - generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), - guidance_scale=1.0, - num_inference_steps=num_inference_steps, - **params["sampler_params"], - ), - ) + try: + # Refer to tests/e2e/offline_inference/test_wan22.py + # Use minimal settings for testing + output = runner.omni.generate( + "A cat sitting on a table", + OmniDiffusionSamplingParams( + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + guidance_scale=1.0, + num_inference_steps=num_inference_steps, + **params["sampler_params"], + ), + ) + finally: + monitor.stop() - peak = monitor.peak_used_mb - monitor.stop() + # DeviceMemoryMonitor reports absolute device usage. Subtract this run's + # starting usage so process-wide compiler/workspace caches retained from a + # previous run do not make the second measurement order-dependent. + peak = max(0.0, monitor.peak_used_mb - initial_used_mb) + + gc.collect() + current_omni_platform.empty_cache() return peak, output @@ -140,10 +166,13 @@ def test_layerwise_offload_diffusion_model(model_name: str): if expected_saved_memory is None: pytest.skip(f"Threshold not defined for {platform} on {model_name}") + assert expected_saved_memory is not None # Verify that layerwise offloading significantly reduces memory usage - # Passes only if the actual savings exceeds the expected savings - assert layerwise_offload_peak_memory + expected_saved_memory < no_offload_peak_memory, ( + # Passes only if the actual savings meets the expected savings + actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory + assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, ( f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB " - f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB" + f"should be at least {expected_saved_memory} MB less than no offload peak memory " + f"{no_offload_peak_memory} MB (actual savings: {actual_saved_memory} MB)" ) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index fd004d056..934fe3299 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -1,12 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for DistributedLayerwiseOffloadHook and backend utilities.""" import gc import json -import os -import socket +import weakref from contextlib import contextmanager from types import SimpleNamespace @@ -18,6 +17,13 @@ from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module +from tests.helpers.runtime import get_distributed_init_method +from vllm_omni.diffusion.data import validate_dlo_host_registration_options +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + TensorBinding, + build_checkpoint_mmap_plan, +) from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy from vllm_omni.diffusion.offloader.block_discovery import ( get_blocks_attr_names, @@ -29,8 +35,17 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( DistributedLayerwiseOffloadHook, PinnedResidentLayerGroup, ) +from vllm_omni.diffusion.offloader.host_registration import ( + HostRegistrationCleanupError, + HostRegistrationError, +) from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager -from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan +from vllm_omni.diffusion.offloader.offload_plan import ( + OffloadPlan, + get_offload_plan, +) +from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state +from vllm_omni.host_weight_runtime import MappedHostRegion from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -48,33 +63,19 @@ class DummyEvent: def record(self, _stream) -> None: return None + def synchronize(self) -> None: + return None + @contextmanager def dummy_stream(_stream): yield None -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - def _cleanup_distributed() -> None: if dist.is_initialized(): dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - gc.collect() if current_omni_platform.is_available(): current_omni_platform.empty_cache() @@ -83,10 +84,7 @@ def _cleanup_distributed() -> None: @pytest.fixture(scope="module") def dist_group(): - master_port = _find_free_port() - _set_dist_env(rank=0, world_size=1, master_port=master_port) - - dist.init_process_group("gloo", rank=0, world_size=1) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) try: yield finally: @@ -139,6 +137,268 @@ class TestDistributedLayerwiseOffloadHook: assert shard.numel() == 4 assert torch.equal(shard, _make_values(10.0)) + def test_prefetch_preserves_transposed_weight_stride( + self, + dist_group, + patched_offload_runtime, + ): + """Online-FP8 Cutlass weights must retain their transposed layout.""" + + class StridedBlock(nn.Module): + def __init__(self, start: float): + super().__init__() + base = torch.arange(start, start + 12).reshape(3, 4) + self.weight = nn.Parameter(base.t(), requires_grad=False) + + current_block = StridedBlock(1.0) + next_block = StridedBlock(20.0) + expected = next_block.weight.detach().clone() + expected_stride = next_block.weight.stride() + assert expected_stride == (1, 4) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + hook.initialize_hook(current_block) + hook.prefetch_layer(slot=0, non_blocking=False) + + assert next_block.weight.stride() == expected_stride + assert torch.equal(next_block.weight, expected) + + def test_allgather_reconstructs_online_fp8_weight_and_scale( + self, + patched_offload_runtime, + monkeypatch, + ): + """DLO must reconstruct finalized online-FP8 tensors byte-for-byte.""" + + class OnlineFp8Block(nn.Module): + def __init__(self): + super().__init__() + logical_weight = torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + dtype=torch.float32, + ).to(torch.float8_e4m3fn) + self.weight = nn.Parameter(logical_weight.t(), requires_grad=False) + self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False) + + current_block = OnlineFp8Block() + rank0_block = OnlineFp8Block() + rank1_block = OnlineFp8Block() + expected_weight = rank0_block.weight.detach().clone() + expected_scale = rank0_block.weight_scale.detach().clone() + expected_stride = rank0_block.weight.stride() + + rank0_hook = DistributedLayerwiseOffloadHook( + next_block=rank0_block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + rank1_hook = DistributedLayerwiseOffloadHook( + next_block=rank1_block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=1, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + rank0_hook.initialize_hook(current_block) + rank1_hook.initialize_hook(OnlineFp8Block()) + rank0_hook.gpu_shard_buffers = [ + {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2) + ] + + def fake_allgather(output, local_shard, *, group): + del group + remote_shard = rank1_hook.cpu_shards[local_shard.dtype] + shard_size = local_shard.numel() + output[:shard_size].copy_(local_shard) + output[shard_size : 2 * shard_size].copy_(remote_shard) + + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather) + rank0_hook.prefetch_layer(slot=0, non_blocking=False) + + assert rank0_block.weight.stride() == expected_stride + torch.testing.assert_close( + rank0_block.weight.float(), + expected_weight.float(), + rtol=0, + atol=0, + ) + torch.testing.assert_close(rank0_block.weight_scale, expected_scale) + + def test_allgather_reconstructs_online_int8_weight_and_scale( + self, + patched_offload_runtime, + monkeypatch, + ): + """DLO must reconstruct finalized online-INT8 tensors byte-for-byte. + + One block carries both finalized online-INT8 layouts: the NPU method + leaves a contiguous pre-transposed (K, N) int8 weight with a 1-D fp32 + per-channel scale, while the CUDA method leaves a transposed-view + weight (stride (1, K)) with a scalar fp32 scale. The int8 and fp32 + dtype groups must survive the flat-shard AllGather independently. + """ + + class OnlineInt8Block(nn.Module): + def __init__(self): + super().__init__() + logical = torch.arange(1.0, 13.0).reshape(3, 4) + # NPU online int8: qweight.t().contiguous() -> (K, N), 1-D scale + self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.int8), requires_grad=False) + self.npu_weight_scale = nn.Parameter(torch.tensor([2.0, 3.0, 4.0]), requires_grad=False) + # CUDA online int8: Parameter(qweight.t().data) -> transposed + # view (stride (1, K)) over the (N, K) storage, scalar scale + self.cuda_weight = nn.Parameter((logical + 100).to(torch.int8).t(), requires_grad=False) + self.cuda_weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False) + + current_block = OnlineInt8Block() + rank0_block = OnlineInt8Block() + rank1_block = OnlineInt8Block() + expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()} + expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()} + + def make_hook(block: OnlineInt8Block, rank: int) -> DistributedLayerwiseOffloadHook: + return DistributedLayerwiseOffloadHook( + next_block=block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=rank, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + + rank0_hook = make_hook(rank0_block, 0) + rank1_hook = make_hook(rank1_block, 1) + rank0_hook.initialize_hook(current_block) + rank1_hook.initialize_hook(OnlineInt8Block()) + rank0_hook.gpu_shard_buffers = [ + {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2) + ] + + def fake_allgather(output, local_shard, *, group): + del group + remote_shard = rank1_hook.cpu_shards[local_shard.dtype] + shard_size = local_shard.numel() + output[:shard_size].copy_(local_shard) + output[shard_size : 2 * shard_size].copy_(remote_shard) + + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather) + rank0_hook.prefetch_layer(slot=0, non_blocking=False) + + for name, tensor in rank0_block.named_parameters(): + assert tensor.stride() == expected_strides[name] + torch.testing.assert_close( + tensor.float(), + expected[name].float(), + rtol=0, + atol=0, + ) + + def test_allgather_reconstructs_online_mxfp8_weight_and_scale( + self, + patched_offload_runtime, + monkeypatch, + ): + """DLO must reconstruct finalized online-MXFP8 tensors byte-for-byte. + + One block carries both finalized online-MXFP8 layouts: the NPU method + leaves a contiguous pre-transposed (K, N) fp8 weight with a contiguous + (K_groups/2, N, 2) e8m0 block scale, while the vLLM-kernel method + (XPU) leaves an (N, K) fp8 weight whose e8m0 scale is a transposed + view over a contiguous (K/32, N) buffer. The fp8 and e8m0 dtype + groups must survive the flat-shard AllGather independently. + """ + + class OnlineMxfp8Block(nn.Module): + def __init__(self): + super().__init__() + logical = torch.arange(1.0, 13.0).reshape(3, 4) + # Build scales from raw e8m0 bit patterns (1 byte/elem). + npu_scale_bits = torch.arange(1, 13, dtype=torch.uint8).reshape(2, 3, 2) + kernel_scale_bits = torch.arange(1, 7, dtype=torch.uint8).reshape(2, 3) + # NPU online mxfp8: (N, K).t().contiguous() -> (K, N) fp8 + # weight, contiguous (K_groups/2, N, 2) e8m0 scale + self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.float8_e4m3fn), requires_grad=False) + self.npu_weight_scale = nn.Parameter(npu_scale_bits.view(torch.float8_e8m0fnu), requires_grad=False) + # vLLM-kernel online mxfp8 (XPU): (N, K) fp8 weight, e8m0 + # scale kept as a transposed view (stride (1, K/32)) over the + # contiguous (K/32, N) storage + self.kernel_weight = nn.Parameter((logical + 100).to(torch.float8_e4m3fn), requires_grad=False) + self.kernel_weight_scale = nn.Parameter( + kernel_scale_bits.view(torch.float8_e8m0fnu).t(), requires_grad=False + ) + + current_block = OnlineMxfp8Block() + rank0_block = OnlineMxfp8Block() + rank1_block = OnlineMxfp8Block() + expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()} + expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()} + + def make_hook(block: OnlineMxfp8Block, rank: int) -> DistributedLayerwiseOffloadHook: + return DistributedLayerwiseOffloadHook( + next_block=block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=rank, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + + rank0_hook = make_hook(rank0_block, 0) + rank1_hook = make_hook(rank1_block, 1) + rank0_hook.initialize_hook(current_block) + rank1_hook.initialize_hook(OnlineMxfp8Block()) + rank0_hook.gpu_shard_buffers = [ + {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2) + ] + + def fake_allgather(output, local_shard, *, group): + del group + remote_shard = rank1_hook.cpu_shards[local_shard.dtype] + shard_size = local_shard.numel() + output[:shard_size].copy_(local_shard) + output[shard_size : 2 * shard_size].copy_(remote_shard) + + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather) + rank0_hook.prefetch_layer(slot=0, non_blocking=False) + + for name, tensor in rank0_block.named_parameters(): + assert tensor.stride() == expected_strides[name] + # Compare raw bytes in logical order: float8/e8m0 CPU upcasts are + # version-sensitive, while contiguous().view(uint8) is an exact + # byte-level comparison for both contiguous and transposed-view + # parameters. + torch.testing.assert_close( + tensor.detach().contiguous().view(torch.uint8), + expected[name].contiguous().view(torch.uint8), + rtol=0, + atol=0, + ) + def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime): """DTensor wrapper should be preserved through prefetch/offload cycle.""" current_block = TinyBlock(_make_values(1.0)) @@ -316,8 +576,9 @@ class TestDistributedLayerwiseOffloadHook: block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) next_block = nn.Module() next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32)) - next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0) - next_block.weight.mmap_weight_transform_pending = True + + def transform(tensor): + return tensor.flip(0) hook = DistributedLayerwiseOffloadHook( next_block=next_block, @@ -326,17 +587,17 @@ class TestDistributedLayerwiseOffloadHook: dp_size=2, rank=0, pin_memory=False, + tensor_transforms={id(next_block.weight): transform}, ) hook.initialize_hook(block) assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0])) - def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime): + def test_transform_is_not_applied_without_a_loader_plan(self, dist_group, patched_offload_runtime): block = nn.Module() block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) next_block = nn.Module() next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32)) - next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0) hook = DistributedLayerwiseOffloadHook( next_block=next_block, @@ -350,6 +611,82 @@ class TestDistributedLayerwiseOffloadHook: assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0])) + def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime): + current_block = nn.Linear(2, 2, bias=False) + next_block = nn.Linear(2, 2, bias=False) + next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2)) + source_storage = next_block.weight.untyped_storage().data_ptr() + + def transform(tensor): + return tensor.t() + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + rank_local_mmap=True, + tensor_transforms={id(next_block.weight): transform}, + ) + hook.initialize_hook(current_block) + + source = hook.cpu_sources[torch.float32][0]["tensor"] + assert source.untyped_storage().data_ptr() == source_storage + assert hook.cpu_shards == {} + assert next_block.weight.numel() == 0 + + hook.cpu_staging_buffers = [ + {torch.float32: torch.empty(4)}, + {torch.float32: torch.empty(4)}, + ] + hook.prefetch_layer(slot=0, non_blocking=False) + + assert torch.equal( + next_block.weight, + torch.tensor([[0.0, 2.0], [1.0, 3.0]]), + ) + assert next_block.weight.stride() == (1, 2) + assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2)) + + def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime): + hooks = [] + for size in (4, 9): + current_block = nn.Linear(1, 1, bias=False) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32)) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + shared_buffers=[None, None], + rank_local_mmap=True, + ) + hook.initialize_hook(current_block) + hooks.append(hook) + + resident_block = nn.Module() + resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32)) + resident_group = PinnedResidentLayerGroup( + [resident_block], + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) + staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers( + hooks, + resident_group, + ) + + assert len(staging) == 2 + assert staging[0][torch.float32].numel() == 12 + assert staging[1][torch.float32].numel() == 12 + class TestPinnedResidentLayerGroup: def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime): @@ -386,23 +723,59 @@ class TestPinnedResidentLayerGroup: group.offload() - def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self): - module = nn.Linear(2, 2, bias=False) + def test_load_preserves_noncontiguous_weight_stride(self, patched_offload_runtime): + block = nn.Module() + weight = torch.arange(6, dtype=torch.float32).reshape(2, 3).t() + expected = weight.detach().clone() + expected_stride = weight.stride() + block.weight = nn.Parameter(weight) + group = PinnedResidentLayerGroup( + [block], + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + ) - def transform(tensor): - return tensor.flip(0) + group.load() + assert torch.equal(block.weight, expected) + assert block.weight.stride() == expected_stride + group.offload() - module.weight.mmap_weight_transform = transform + def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime): + blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)] + expected = [] + source_storage = [] + for index, block in enumerate(blocks): + block.weight.data.fill_(index + 1) + expected.append(block.weight.detach().clone()) + source_storage.append(block.weight.untyped_storage().data_ptr()) - backend = object.__new__(DistributedLayerwiseOffloadBackend) - backend._remember_mmap_param_attrs(module) - module.to_empty(device="meta") + group = PinnedResidentLayerGroup( + blocks, + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) - replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float()) - backend._attach_mmap_param_attrs("weight", replacement, module.weight) + for index, state in enumerate(group._states): + source = state["cpu_sources"][torch.float32][0]["tensor"] + assert source.untyped_storage().data_ptr() == source_storage[index] + assert state["cpu_shards"] == {} + assert len(group._cpu_staging_buffers) == 2 - assert replacement.mmap_weight_transform is transform - assert replacement.mmap_weight_transform_pending is True + group.load() + for block, weight in zip(blocks, expected): + assert torch.equal(block.weight, weight) + + blocks[0].weight.data.zero_() + group.offload() + group.load() + for block, weight in zip(blocks, expected): + assert torch.equal(block.weight, weight) + + group.offload() + assert all(block.weight.numel() == 0 for block in blocks) def test_shared_allgather_output_is_narrowed_to_current_block( self, dist_group, patched_offload_runtime, monkeypatch @@ -484,6 +857,271 @@ class _SingleBlockModel(nn.Module): self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) +class _HWRPipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=3) + + +class _FakeHostWeightLease: + def __init__(self, resolution_id: str, events: list[str] | None = None): + self.closed = False + self.provenance = SimpleNamespace(resolution_id=resolution_id) + self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),) + self._events = events + + def close(self): + if self._events is not None: + self._events.append("lease") + self.closed = True + + +class _FakeLeaseCarrier: + def __init__(self, lease: _FakeHostWeightLease): + self.lease = lease + self.taken = False + + @property + def closed(self): + return self.taken or self.lease.closed + + def take(self): + if self.taken: + raise RuntimeError("already taken") + self.taken = True + return self.lease + + def close(self): + if not self.taken: + self.lease.close() + + +def _fake_hwr_plan(resolution_id: str = "hwr-test"): + lease = _FakeHostWeightLease(resolution_id) + carrier = _FakeLeaseCarrier(lease) + plan = HostWeightPlan( + backing_kind="host_weight_runtime", + bindings={}, + lease_carrier=carrier, + ) + return plan, carrier, lease + + +def _hwr_backend(): + plan, carrier, lease = _fake_hwr_plan() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + host_weight_plan=plan, + ) + return _HWRPipeline(), backend, carrier, lease + + +def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime): + current_block = nn.Linear(2, 2, bias=False) + next_block = nn.Linear(2, 2, bias=False) + expected = torch.arange(4, dtype=torch.float32).view(2, 2) + next_block.weight.data.copy_(expected) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + rank_local_mmap=True, + ) + hook.initialize_hook(current_block) + hook.registered_mmap = True + hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging") # type: ignore[method-assign] + + hook.prefetch_layer(slot=0, non_blocking=True) + + assert torch.equal(next_block.weight, expected) + + +def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime): + block = nn.Linear(2, 2, bias=False) + expected = torch.arange(4, dtype=torch.float32).view(2, 2) + block.weight.data.copy_(expected) + group = PinnedResidentLayerGroup( + [block], + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) + group.registered_mmap = True + group._cpu_staging_buffers.clear() + + group.load() + + assert torch.equal(block.weight, expected) + + +def test_hwr_registration_uses_transport_budget( + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, +): + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + plan, _, lease = _fake_hwr_plan() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dp_size=1, + dlo_use_allgather=False, + dlo_host_registration_limit_gib=1.5, + ), + torch.device("cuda"), + host_weight_plan=plan, + ) + backend._host_weight_lease = lease # type: ignore[assignment] + calls: list[tuple[object, int | None]] = [] + + class Registration: + total_bytes = 4096 + region_count = 1 + + @staticmethod + def close() -> tuple[str, ...]: + return () + + def register(regions, *, device, max_bytes): + assert device == torch.device("cuda") + calls.append((regions, max_bytes)) + return Registration() + + monkeypatch.setattr(dist_backend_module, "register_host_mappings", register) + pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True) + + assert backend._try_register_hwr_mmap((pinned_source,)) # type: ignore[arg-type] + assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))] + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)] + + backend._release_registered_mmap() + + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [] + + +def test_hwr_registration_budget_validation_is_transport_scoped(): + assert ( + validate_dlo_host_registration_options( + limit_gib=1.5, + enable_dlo=True, + use_allgather=False, + hwr_mode="preferred", + ) + == 1.5 + ) + invalid = ( + dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"), + dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"), + dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"), + dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"), + dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"), + ) + for options in invalid: + with pytest.raises(ValueError): + validate_dlo_host_registration_options(**options) + + +def test_unregistration_precedes_lease_close_and_retries_failures( + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, +): + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + events: list[str] = [] + responses = iter([("busy",), ()]) + + class Registration: + @staticmethod + def close() -> tuple[str, ...]: + events.append("unregister") + return next(responses) + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + ) + lease = _FakeHostWeightLease("retry", events) + backend._host_weight_lease = lease # type: ignore[assignment] + backend._host_registration = Registration() + backend._using_rank_local_mmap = True + backend._using_registered_mmap = True + backend.enabled = True + monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None) + + with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"): + backend.disable() + assert not lease.closed + assert backend._host_registration is not None + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)] + + backend.disable() + + assert lease.closed + assert backend._host_registration is None + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [] + assert events == ["unregister", "unregister", "lease"] + + +def test_failed_unregistration_retains_lease_after_backend_is_dropped( + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, +): + events: list[str] = [] + + class Registration: + @staticmethod + def close() -> tuple[str, ...]: + events.append("unregister") + return ("busy",) + + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + ) + lease = _FakeHostWeightLease("retained", events) + registration = Registration() + backend._host_weight_lease = lease # type: ignore[assignment] + backend._host_registration = registration + backend._using_rank_local_mmap = True + backend.enabled = True + monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None) + lease_ref = weakref.ref(lease) + + with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"): + backend.disable() + del backend + del lease + gc.collect() + + retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1] + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)] + assert lease_ref() is retained_lease + assert not retained_lease.closed + + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + retained_lease.close() + + class _MultiBlockModel(nn.Module): _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] @@ -545,7 +1183,6 @@ class TestMmapWeightLoading: OffloadConfig( strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, - model_path=str(tmp_path), ), torch.device("cpu"), ) @@ -553,13 +1190,167 @@ class TestMmapWeightLoading: dits=[pipeline.transformer], dit_names=["transformer"], ) + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={ + name: TensorBinding( + checkpoint_key=name, + file_path=str(weight_file), + ) + for name in weights + }, + ) - backend._load_weights_via_mmap(pipeline, modules) + backend._load_weights_via_mmap(pipeline, modules, plan) assert pipeline.transformer.post_load_calls == 1 assert pipeline.transformer.time_embedder.weight.dtype == torch.float32 assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16 + def test_allgather_flag_with_group_size_one_uses_rank_local_mmap( + self, + tmp_path, + patched_offload_runtime, + ): + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList( + [ + nn.Linear(2, 2, bias=False), + nn.Linear(2, 2, bias=False), + ] + ) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = Pipeline() + weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()} + checkpoint_file = tmp_path / "model.safetensors" + save_file(weights, str(checkpoint_file)) + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + assert result.plan is not None + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=1, + dlo_use_allgather=True, + ), + torch.device("cpu"), + host_weight_plan=result.plan, + ) + + backend.enable(pipeline) + + assert backend._using_rank_local_mmap + assert backend.dp_group is None + assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group) + backend.disable() + + def test_hwr_carrier_is_taken_and_released_after_bounded_staging(self, patched_offload_runtime): + """A warm HWR plan uses the existing two-slot rank-local transport.""" + pipeline, backend, carrier, lease = _hwr_backend() + model_bytes = sum(param.numel() * param.element_size() for param in pipeline.transformer.parameters()) + + backend.enable(pipeline) + + assert carrier.taken + assert backend._using_rank_local_mmap + assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group) + staging_bytes = sum( + buffer.numel() * buffer.element_size() + for buffer in backend._all_hook_groups[0][0].cpu_staging_buffers[0].values() + ) + assert staging_bytes * 2 < model_bytes + assert not lease.closed + + backend.disable() + assert lease.closed + + def test_hwr_registration_failure_falls_back_to_bounded_staging( + self, + monkeypatch, + patched_offload_runtime, + ): + plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback") + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + host_weight_plan=plan, + ) + pipeline = _HWRPipeline() + original_staging_allocator = backend._allocate_shared_cpu_staging_buffers + + def fail_registration(*args, **kwargs): + del args, kwargs + raise HostRegistrationError("registration unavailable in CPU test") + + def allocate_unpinned_staging(hooks, resident_group=None): + # Registration selection still observes pin_cpu_memory=True. The + # CPU-only test avoids asking the host for CUDA-pinned allocations. + for hook in hooks: + hook.pin_memory = False + return original_staging_allocator(hooks, resident_group) + + monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration) + monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging) + + backend.enable(pipeline) + + assert carrier.taken + assert backend._using_rank_local_mmap + assert not backend._using_registered_mmap + assert backend._host_registration is None + assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group) + assert not lease.closed + + backend.disable() + assert lease.closed + + def test_hwr_backend_failure_drains_partial_setup_before_lease_close( + self, + monkeypatch, + patched_offload_runtime, + ): + pipeline, backend, carrier, lease = _hwr_backend() + monkeypatch.setattr( + backend, + "_allocate_shared_buffers", + lambda hooks: (_ for _ in ()).throw(RuntimeError("device buffer allocation failed")), + ) + + with pytest.raises(RuntimeError, match="device buffer allocation failed"): + backend.enable(pipeline) + + assert carrier.taken + assert lease.closed + assert not backend.enabled + assert not backend._blocks + assert not backend._all_hook_groups + class TestGetBlocksFromDit: def test_get_blocks_from_dit_single_block_attr(self): @@ -825,23 +1616,123 @@ class TestOffloadPlan: class TestMmapValidation: - """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights.""" + """Tests for loader preflight and backend plan realization.""" - def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch): - """DLO+AllGather should reject when TP world size > 1.""" - import json - from types import SimpleNamespace - - from safetensors.torch import save_file - - from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy - - # Mock TP world size = 2 - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size", - lambda: 2, + def test_component_source_prefix_builds_model_namespace(self, tmp_path): + source_root = tmp_path / "partition" + checkpoint_dir = source_root / "transformer" + checkpoint_dir.mkdir(parents=True) + checkpoint_file = checkpoint_dir / "model.safetensors" + save_file( + {"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)}, + str(checkpoint_file), ) + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Module() + self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)]) + self.weights_sources = [ + SimpleNamespace( + model_or_path=str(source_root), + subfolder="transformer", + revision=None, + prefix="transformer.", + ) + ] + + pipeline = Pipeline() + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=tuple(pipeline.weights_sources), + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.fallback_reason is None + assert result.plan is not None + assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding( + checkpoint_key="blocks.0.weight", + file_path=str(checkpoint_file), + ) + assert result.plan.planned_source_prefixes == frozenset({"transformer."}) + + def test_non_dedicated_component_source_falls_back_before_discovery(self): + pipeline = nn.Module() + pipeline.transformer = nn.Linear(2, 2, bias=False) + sources = ( + SimpleNamespace( + model_or_path="unused", + subfolder=None, + revision=None, + prefix="", + ), + ) + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=sources, + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is None + assert "dedicated component weight source" in result.fallback_reason + + def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path): + source_root = tmp_path / "partition" + checkpoint_dir = source_root / "transformer" + checkpoint_dir.mkdir(parents=True) + save_file( + {"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)}, + str(checkpoint_dir / "model.safetensors"), + ) + + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)]) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + self.weights_sources = [ + SimpleNamespace( + model_or_path=str(source_root), + subfolder="transformer", + revision=None, + prefix="transformer.", + ) + ] + + pipeline = Pipeline() + original_weight = pipeline.transformer.blocks[0].weight + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=tuple(pipeline.weights_sources), + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is None + assert "1 required DiT tensors" in result.fallback_reason + assert pipeline.transformer.blocks[0].weight is original_weight + assert not original_weight.is_meta + + def test_tp_falls_back_before_checkpoint_discovery(self): class SimpleModel(nn.Module): _layerwise_offload_blocks_attrs = ["blocks"] @@ -859,41 +1750,118 @@ class TestMmapValidation: return key pipeline = SimplePipeline() - weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta} - save_file(weights, str(tmp_path / "model.safetensors")) - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {k: "model.safetensors" for k in weights}}) + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path="unused", + tensor_parallel_size=2, + use_hsdp=False, + online_quantization=False, ) - backend = DistributedLayerwiseOffloadBackend( - OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, - pin_cpu_memory=False, - model_path=str(tmp_path), - ), - torch.device("cpu"), - ) - modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + assert result.plan is None + assert result.fallback_reason == "TP=2 requires the ordinary loader" - with pytest.raises(ValueError, match="Tensor Parallel"): - backend._load_weights_via_mmap(pipeline, modules) + @pytest.mark.parametrize( + ("use_hsdp", "online_quantization", "expected_reason"), + [ + (True, False, "HSDP requires the ordinary loader"), + (False, True, "online quantization requires the ordinary loader"), + ], + ) + def test_runtime_layouts_that_need_loader_callbacks_fall_back( + self, + use_hsdp, + online_quantization, + expected_reason, + ): + pipeline = nn.Module() + pipeline.transformer = nn.Linear(2, 2, bias=False) - def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch): - """DLO+AllGather should work when TP world size = 1, even if params - have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1).""" - import json - from types import SimpleNamespace - - from safetensors.torch import save_file - - from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy - - # Mock TP world size = 1 (default) - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size", - lambda: 1, + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path="unused", + tensor_parallel_size=1, + use_hsdp=use_hsdp, + online_quantization=online_quantization, ) + assert result.plan is None + assert result.fallback_reason == expected_reason + + @pytest.mark.parametrize( + ("checkpoint_tensor", "expected_reason"), + [ + (torch.ones((3, 2), dtype=torch.float32), "shape mismatch"), + (torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"), + ], + ) + def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason): + checkpoint_file = tmp_path / "model.safetensors" + save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file)) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Linear(2, 2, bias=False) + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = Pipeline() + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is None + assert expected_reason in result.fallback_reason + + def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path): + class Transformer(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(2, 2)) + self.register_buffer("derived", torch.ones(2), persistent=False) + + checkpoint_file = tmp_path / "model.safetensors" + save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file)) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = Pipeline() + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is not None + assert set(result.plan.bindings) == {"transformer.weight"} + + def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path): def custom_weight_loader(param, weight): param.data.copy_(weight) @@ -923,27 +1891,67 @@ class TestMmapValidation: json.dumps({"weight_map": {k: "model.safetensors" for k in weights}}) ) - backend = DistributedLayerwiseOffloadBackend( - OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, - pin_cpu_memory=False, - model_path=str(tmp_path), - ), - torch.device("cpu"), + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, ) - modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) - # Should NOT raise — TP=1 is allowed even with custom weight_loader - backend._load_weights_via_mmap(pipeline, modules) + assert result.plan is None + assert "custom weight_loader" in result.fallback_reason + + def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path): + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + Cosmos3VFMTransformer, + ) + + transformer = object.__new__(Cosmos3VFMTransformer) + nn.Module.__init__(transformer) + transformer.proj_in = nn.Linear(2, 2, bias=False) + transformer.proj_in.weight.weight_loader = lambda *_args: None + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = transformer + + source_root = tmp_path / "partition" + checkpoint_dir = source_root / "transformer" + checkpoint_dir.mkdir(parents=True) + checkpoint_file = checkpoint_dir / "model.safetensors" + save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file)) + sources = ( + SimpleNamespace( + model_or_path=str(source_root), + subfolder="transformer", + revision=None, + prefix="transformer.", + ), + ) + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", transformer),), + sources=sources, + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is not None + assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding( + checkpoint_key="proj_in.weight", + file_path=str(checkpoint_file), + ) def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime): """validate_loaded_weights should be called after mmap loading.""" - import json - from types import SimpleNamespace - - from safetensors.torch import save_file - - from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy class ModelWithValidate(nn.Module): _layerwise_offload_blocks_attrs = ["blocks"] @@ -985,19 +1993,153 @@ class TestMmapValidation: backend = DistributedLayerwiseOffloadBackend( OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path) + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, ), torch.device("cpu"), ) modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) - backend._load_weights_via_mmap(pipeline, modules) + assert result.plan is not None + backend._load_weights_via_mmap(pipeline, modules, result.plan) assert pipeline.transformer.validate_called, "validate_loaded_weights should be called" class TestConfigValidation: """Tests for configuration validation in OffloadConfig / execute_request.""" + @pytest.mark.parametrize( + ("dp", "sp", "tp", "use_allgather", "expected_dlo_group"), + [ + (1, 1, 1, True, 1), + (2, 1, 1, True, 2), + (1, 4, 1, True, 4), + (2, 2, 2, True, 2), + (2, 2, 2, False, 1), + (1, 4, 2, False, 1), + ], + ) + def test_dlo_group_selection_covers_dp_sp_tp_matrix( + self, + dp, + sp, + tp, + use_allgather, + expected_dlo_group, + ): + """DLO sharding follows DP, falls back to SP, and never uses TP.""" + + config = OffloadConfig.from_od_config( + SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=use_allgather, + dlo_resident_layers=0, + pin_cpu_memory=False, + parallel_config=SimpleNamespace( + data_parallel_size=dp, + use_hsdp=False, + hsdp_shard_size=-1, + hsdp_replicate_size=1, + sequence_parallel_size=sp, + tensor_parallel_size=tp, + ), + ) + ) + assert config.dp_size == expected_dlo_group + assert config.dlo_use_allgather is use_allgather + + def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch): + import vllm_omni.diffusion.offloader as offloader_module + + od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + dlo_resident_layers=0, + pin_cpu_memory=False, + model="unused", + parallel_config=SimpleNamespace( + use_hsdp=False, + data_parallel_size=1, + sequence_parallel_size=1, + ), + ) + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={}, + ) + monkeypatch.setattr( + offloader_module.current_omni_platform, + "supports_cpu_offload", + lambda: False, + ) + + with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"): + offloader_module.get_offload_backend( + od_config, + device=torch.device("cpu"), + host_weight_plan=plan, + ) + + def test_startup_recovery_stays_inside_offloader_boundary(self, monkeypatch): + import vllm_omni.diffusion.offloader as offloader_module + + class FailingBackend: + def __init__(self): + self.disabled = False + + def enable(self, pipeline): + del pipeline + raise RuntimeError("initial prefetch failed") + + def disable(self): + self.disabled = True + + warm_pipeline = nn.Module() + canonical_pipeline = nn.Module() + plan, carrier, _lease = _fake_hwr_plan("hwr-recovery") + failing_backend = FailingBackend() + calls = [] + + def fake_backend(od_config, device, host_weight_plan): + del od_config, device + calls.append(host_weight_plan) + return failing_backend if host_weight_plan is plan else None + + attach_offload_startup_state( + warm_pipeline, + OffloadStartupState( + host_weight_plan=plan, + fresh_model_loader=lambda: canonical_pipeline, + allow_fresh_retry=True, + ), + ) + monkeypatch.setattr(offloader_module, "get_offload_backend", fake_backend) + + recovered, backend = offloader_module.enable_offload_backend( + SimpleNamespace(), + warm_pipeline, + device=torch.device("cpu"), + ) + + assert recovered is canonical_pipeline + assert backend is None + assert failing_backend.disabled + assert carrier.closed + assert calls == [plan, None] + def test_hsdp_with_allgather_rejected(self): """HSDP + DLO + AllGather should raise ValueError (double sharding).""" from vllm_omni.diffusion.offloader.base import OffloadConfig @@ -1195,8 +2337,9 @@ class TestDynamicSlotTracking: prev_idx = (i - 1) % num_blocks # Dynamic slot tracking: read from prev's prefetched slot - if prefetched_slots[prev_idx] is not None: - current_slots[i] = prefetched_slots[prev_idx] + prefetched_slot = prefetched_slots[prev_idx] + if prefetched_slot is not None: + current_slots[i] = prefetched_slot read_slot = current_slots[i] diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 5fd80e75c..9253f6d03 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities.""" import gc -import os -import socket from contextlib import contextmanager import pytest @@ -15,6 +13,7 @@ from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module +from tests.helpers.runtime import get_distributed_init_method from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook from vllm_omni.platforms import current_omni_platform @@ -39,27 +38,10 @@ def dummy_stream(_stream): yield None -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: - os.environ["RANK"] = str(rank) - os.environ["LOCAL_RANK"] = str(rank) - os.environ["WORLD_SIZE"] = str(world_size) - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(master_port) - - def _cleanup_distributed() -> None: if dist.is_initialized(): dist.destroy_process_group() - for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: - os.environ.pop(key, None) - gc.collect() if current_omni_platform.is_available(): current_omni_platform.empty_cache() @@ -68,10 +50,7 @@ def _cleanup_distributed() -> None: @pytest.fixture(scope="module") def dist_group(): - master_port = _find_free_port() - _set_dist_env(rank=0, world_size=1, master_port=master_port) - - dist.init_process_group("gloo", rank=0, world_size=1) + dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method()) try: yield finally: @@ -79,11 +58,11 @@ def dist_group(): @pytest.fixture -def patched_offload_runtime(mocker): - mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream) - mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent) - mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) - mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream) +def patched_offload_runtime(monkeypatch): + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) + monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream) class TinyBlock(nn.Module): @@ -128,6 +107,33 @@ class TestLayerwiseOffloadHook: assert current_block.weight.to_local().shape == torch.Size([4]) assert not hook.is_materialized + def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime): + """Online-FP8 Cutlass weights must retain their transposed layout.""" + + class StridedBlock(nn.Module): + def __init__(self, start: float): + super().__init__() + base = torch.arange(start, start + 12).reshape(3, 4) + self.weight = nn.Parameter(base.t(), requires_grad=False) + + current_block = StridedBlock(1.0) + next_block = StridedBlock(20.0) + expected = next_block.weight.detach().clone() + expected_stride = next_block.weight.stride() + assert expected_stride == (1, 4) + + hook = LayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + stream=DummyStream(), + pin_memory=False, + ) + hook.initialize_hook(current_block) + hook.prefetch_layer(non_blocking=False) + + assert next_block.weight.stride() == expected_stride + assert torch.equal(next_block.weight, expected) + class _DummyBlock(nn.Module): def __init__(self): diff --git a/tests/diffusion/offloader/test_module_residency.py b/tests/diffusion/offloader/test_module_residency.py new file mode 100644 index 000000000..adf3c3afc --- /dev/null +++ b/tests/diffusion/offloader/test_module_residency.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import contextmanager + +import pytest +import torch +from torch import nn + +import vllm_omni.diffusion.offloader.module_residency as residency_module +from vllm_omni.diffusion.offloader.module_residency import ( + BoundedAllocatorCache, + PinnedModuleStager, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +class _DummyStream: + def wait_stream(self, _stream) -> None: + return None + + def wait_event(self, _event) -> None: + return None + + +class _DummyEvent: + def __init__(self) -> None: + self.record_count = 0 + + def record(self, _stream) -> None: + self.record_count += 1 + + +@contextmanager +def _dummy_stream(_stream): + yield + + +@pytest.fixture +def patched_runtime(monkeypatch, mocker): + empty_cache = mocker.Mock() + synchronize = mocker.Mock() + monkeypatch.setattr(residency_module.current_omni_platform, "Stream", _DummyStream) + monkeypatch.setattr(residency_module.current_omni_platform, "Event", _DummyEvent) + monkeypatch.setattr(residency_module.current_omni_platform, "current_stream", _DummyStream) + monkeypatch.setattr(residency_module.current_omni_platform, "stream", _dummy_stream) + monkeypatch.setattr(residency_module.current_omni_platform, "empty_cache", empty_cache) + monkeypatch.setattr(residency_module.current_omni_platform, "synchronize", synchronize) + return empty_cache, synchronize + + +def _aliased_modules() -> tuple[nn.Module, nn.Module]: + parameter_storage = torch.arange(12, dtype=torch.float32) + buffer_storage = torch.arange(8, dtype=torch.int64) + left = nn.Module() + right = nn.Module() + left.weight = nn.Parameter(parameter_storage[:8].view(2, 4)) + right.weight = nn.Parameter(parameter_storage[4:].view(2, 4)) + left.register_buffer("state", buffer_storage[:6]) + right.register_buffer("state", buffer_storage[2:]) + return left, right + + +def test_module_group_preserves_parameter_and_buffer_storage_aliases(patched_runtime, mocker): + left, right = _aliased_modules() + expected_left_weight = left.weight.detach().clone() + expected_right_weight = right.weight.detach().clone() + expected_left_state = left.state.clone() + expected_right_state = right.state.clone() + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + [left, right], + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + + master_weight_ptr = left.weight.untyped_storage().data_ptr() + master_state_ptr = left.state.untyped_storage().data_ptr() + assert right.weight.untyped_storage().data_ptr() == master_weight_ptr + assert right.state.untyped_storage().data_ptr() == master_state_ptr + assert (left.weight.storage_offset(), right.weight.storage_offset()) == (0, 4) + assert (left.state.storage_offset(), right.state.storage_offset()) == (0, 2) + + stager.load() + + assert left.weight.untyped_storage().data_ptr() != master_weight_ptr + assert right.weight.untyped_storage().data_ptr() == left.weight.untyped_storage().data_ptr() + assert right.state.untyped_storage().data_ptr() == left.state.untyped_storage().data_ptr() + left.weight.data[1, 0] = -1 + left.state[2] = -1 + assert right.weight[0, 0].item() == -1 + assert right.state[0].item() == -1 + + stager.offload() + + assert left.weight.untyped_storage().data_ptr() == master_weight_ptr + assert right.weight.untyped_storage().data_ptr() == master_weight_ptr + assert left.state.untyped_storage().data_ptr() == master_state_ptr + assert right.state.untyped_storage().data_ptr() == master_state_ptr + assert torch.equal(left.weight, expected_left_weight) + assert torch.equal(right.weight, expected_right_weight) + assert torch.equal(left.state, expected_left_state) + assert torch.equal(right.state, expected_right_state) + retention.release_if_needed.assert_called_once_with(force=False) + + +def test_staging_transitions_are_idempotent_and_reuse_one_event(patched_runtime, mocker): + module = nn.Linear(2, 2) + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + ready_event = stager._ready_event + + stager.load() + first_storage = module.weight.untyped_storage().data_ptr() + stager.load() + assert module.weight.untyped_storage().data_ptr() == first_storage + assert stager._ready_event is ready_event + assert ready_event.record_count == 1 + + stager.offload() + stager.offload() + assert not stager.loaded + retention.release_if_needed.assert_called_once_with(force=False) + + +def test_failed_load_restores_master_and_forces_cache_release(monkeypatch, patched_runtime, mocker): + module = nn.Linear(2, 2) + expected = module.weight.detach().clone() + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + master_ptr = module.weight.untyped_storage().data_ptr() + + def fail_after_rebind() -> None: + device_storages = [torch.empty_like(group.master) for group in stager._groups] + stager._bind(device_storages) + module.weight.data.zero_() + raise RuntimeError("copy failed") + + monkeypatch.setattr(stager, "_load_once", fail_after_rebind) + with pytest.raises(RuntimeError, match="copy failed"): + stager.load() + + assert not stager.loaded + assert module.weight.untyped_storage().data_ptr() == master_ptr + assert torch.equal(module.weight, expected) + retention.release_if_needed.assert_called_once_with(force=True) + + +def test_out_of_memory_releases_cache_and_retries_once(monkeypatch, patched_runtime, mocker): + module = nn.Linear(2, 2) + retention = mocker.Mock(spec=BoundedAllocatorCache) + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + cache_retention=retention, + ) + load_once = stager._load_once + attempts = 0 + + def fail_once() -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise torch.OutOfMemoryError("memory pressure") + load_once() + + monkeypatch.setattr(stager, "_load_once", fail_once) + stager.load() + + assert stager.loaded + assert attempts == 2 + retention.release_if_needed.assert_called_once_with(force=True) + + +@pytest.mark.parametrize( + ("reserved", "allocated", "free", "expected_release"), + [ + (20, 10, 80, False), + (40, 10, 80, True), + (20, 10, 4, True), + ], +) +def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds( + monkeypatch, + patched_runtime, + reserved, + allocated, + free, + expected_release, +): + empty_cache, _ = patched_runtime + monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved) + monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated) + monkeypatch.setattr( + residency_module.current_omni_platform, + "get_device_memory", + lambda _device: (free, 100), + ) + retention = BoundedAllocatorCache(torch.device("cpu")) + + released = retention.release_if_needed() + + assert released is expected_release + assert empty_cache.call_count == int(expected_release) + + +def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime): + empty_cache, _ = patched_runtime + + def unavailable(_device): + raise NotImplementedError + + monkeypatch.setattr(torch.accelerator, "memory_reserved", unavailable) + retention = BoundedAllocatorCache(torch.device("cpu")) + + assert retention.release_if_needed() + empty_cache.assert_called_once_with() diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index dc74248e5..ac3d3181d 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -8,7 +8,13 @@ import torch from torch import nn from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy -from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook +from vllm_omni.diffusion.offloader.sequential_backend import ( + ModelLevelOffloadBackend, + SequentialOffloadHook, + apply_sequential_offload, + remove_sequential_offload, + sequential_offload_component, +) from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -76,6 +82,52 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None: assert pipeline.disable_called is True +def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None: + dit = _create_simple_module() + encoder = _create_simple_module() + offloaded: list[nn.Module] = [] + monkeypatch.setattr( + SequentialOffloadHook, + "_to_cpu", + lambda self, module: offloaded.append(module), + ) + + apply_sequential_offload( + dit_modules=[dit], + encoder_modules=[encoder], + device=torch.device("cpu"), + offload_initial_dits=True, + ) + + assert offloaded == [dit] + remove_sequential_offload([dit, encoder]) + + +def test_direct_component_activation_failure_still_offloads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + component = _create_simple_module() + apply_sequential_offload( + dit_modules=[_create_simple_module()], + encoder_modules=[component], + device=torch.device("cpu"), + ) + hook = component._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME) + offloaded: list[nn.Module] = [] + + def fail_activation(module: nn.Module) -> None: + raise RuntimeError("activation failed") + + monkeypatch.setattr(hook, "pre_forward", fail_activation) + monkeypatch.setattr(hook, "_to_cpu", lambda module: offloaded.append(module)) + + with pytest.raises(RuntimeError, match="activation failed"): + with sequential_offload_component(component): + pass + + assert offloaded == [component] + + class TestMoveParamsPinMemory: def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """DTensor should skip pin_memory to avoid RuntimeError.""" @@ -85,7 +137,12 @@ class TestMoveParamsPinMemory: original_isinstance = isinstance def fake_isinstance(obj, cls): - if cls.__name__ == "DTensor": + # torchada's Tensor.to() checks tuples such as + # ``(str, torch.device)``. Keep the DTensor probe compatible with + # the normal isinstance contract instead of breaking those calls. + if type(cls) is tuple: + return any(fake_isinstance(obj, candidate) for candidate in cls) + if getattr(cls, "__name__", None) == "DTensor": return True return original_isinstance(obj, cls) diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py index c8b3837e2..943411779 100644 --- a/tests/diffusion/quantization/test_component_routing.py +++ b/tests/diffusion/quantization/test_component_routing.py @@ -9,6 +9,9 @@ import torch from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) from vllm.model_executor.models.utils import WeightsMapper from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( @@ -16,10 +19,11 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( ) from vllm_omni.quantization.component_config import ( ComponentQuantizationConfig, + resolve_component_quant_config, ) from vllm_omni.quantization.inc_config import OmniINCConfig -pytestmark = [pytest.mark.core_model] +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # --------------------------------------------------------------------------- @@ -343,10 +347,62 @@ class TestTalkerVisualRouting: # =================================================================== -# 4. ComponentQuantizationConfig.resolve +# 4. ComponentQuantizationConfig.apply_vllm_mapper # =================================================================== +class TestComponentApplyVllmMapper: + def test_maps_compressed_tensors_component_and_default(self): + thinker = CompressedTensorsConfig( + target_scheme_map={}, + ignore=["thinker.model.layers.0.mlp.gate"], + quant_format="pack-quantized", + ) + default = CompressedTensorsConfig( + target_scheme_map={}, + ignore=["talker.model.layers.0.mlp.gate"], + quant_format="pack-quantized", + ) + config = ComponentQuantizationConfig( + component_configs={"thinker": thinker, "code2wav": None}, + default_config=default, + ) + mapper = WeightsMapper( + orig_to_new_prefix={ + "thinker.model": "thinker.language_model.model", + "talker.model": "talker.language_model.model", + } + ).get_unstacked_mapper() + + config.apply_vllm_mapper(mapper) + + assert thinker.ignore == ["thinker.language_model.model.layers.0.mlp.gate"] + assert default.ignore == ["talker.language_model.model.layers.0.mlp.gate"] + + +# =================================================================== +# 5. ComponentQuantizationConfig.resolve +# =================================================================== + + +class TestResolveComponentQuantConfig: + def test_global_config_applies_to_each_quantization_aware_component(self): + config = _MockQuantConfig("fp8") + + assert resolve_component_quant_config(config, "transformer") is config + assert resolve_component_quant_config(config, "text_encoder") is config + + def test_component_map_is_the_only_scope_override(self): + transformer = _MockQuantConfig("fp8") + config = ComponentQuantizationConfig( + component_configs={"transformer": transformer}, + default_config=None, + ) + + assert resolve_component_quant_config(config, "transformer") is transformer + assert resolve_component_quant_config(config, "text_encoder") is None + + class TestComponentResolve: def test_longest_prefix_match(self): a = _MockQuantConfig("a") diff --git a/tests/diffusion/quantization/test_gguf_diffusion.py b/tests/diffusion/quantization/test_gguf_diffusion.py index 243282b0c..ad2e72edc 100644 --- a/tests/diffusion/quantization/test_gguf_diffusion.py +++ b/tests/diffusion/quantization/test_gguf_diffusion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end tests for GGUF quantization on diffusion models. Validates that GGUF-quantized diffusion models generate valid images and @@ -19,8 +19,8 @@ import numpy as np import pytest import torch -from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test +from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py index 482235ff2..ce939dbf3 100644 --- a/tests/diffusion/quantization/test_hsdp_fp8.py +++ b/tests/diffusion/quantization/test_hsdp_fp8.py @@ -36,15 +36,23 @@ class _ToyKernel: # --- helper to build a toy module simulating online-FP8 post-load state --- -def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32): +def _make_fp8_toy_module( + out_features: int = 16, + in_features: int = 32, + *, + quant_method_cls=None, +): """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose ``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``. """ from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod + if quant_method_cls is None: + quant_method_cls = Fp8LinearMethod + # Create the quant-method instance without calling __init__ (avoids # heavyweight upstream config dependencies). isinstance() still works. - qm = object.__new__(Fp8LinearMethod) + qm = object.__new__(quant_method_cls) qm.fp8_linear = _ToyKernel() # randn does not support float8 on CPU; zeros does and the test only @@ -104,6 +112,28 @@ def test_rewrites_non_contiguous_weight(): assert module.weight.shape == (16, 32) +def test_rewrites_per_tensor_online_fp8_weight(): + """Regression test for vLLM's new online FP8 frontend (#45463).""" + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + module = _make_fp8_toy_module( + 16, + 32, + quant_method_cls=Fp8PerTensorOnlineLinearMethod, + ) + + n_rewritten_layers = prepare_fp8_layers_for_fsdp(module) + + assert n_rewritten_layers == 1 + assert module.weight.is_contiguous() + assert module.weight.shape == (16, 32) + w, *_ = module.quant_method.fp8_linear._get_layer_params(module) + assert w.shape == (32, 16) + assert w.stride() == (1, 32) + + def test_patched_get_layer_params_returns_transposed_view(): module = _make_fp8_toy_module(16, 32) prepare_fp8_layers_for_fsdp(module) diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 617083398..0da5b6dd3 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -199,7 +199,7 @@ def _generate_bagel_image( config_path = get_deploy_config_path("ci/bagel.yaml") omni_kwargs: dict[str, Any] = { "model": "ByteDance-Seed/BAGEL-7B-MoT", - "stage_configs_path": config_path, + "deploy_config": config_path, "stage_init_timeout": 300, } if diffusion_quantization_config: diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index fdcef37b5..47908daa0 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Quantization quality gate for diffusion models. @@ -53,6 +53,7 @@ _BENCH_MODULE_NAME = "benchmarks.diffusion.quantization_quality" if _BENCH_MODULE_NAME not in sys.modules: _spec = importlib.util.spec_from_file_location(_BENCH_MODULE_NAME, _BENCH_MODULE_PATH) + assert _spec is not None and _spec.loader is not None _mod = importlib.util.module_from_spec(_spec) sys.modules[_BENCH_MODULE_NAME] = _mod _spec.loader.exec_module(_mod) @@ -80,7 +81,7 @@ class QualityTestConfig: num_frames: int = 5 # only for t2v seed: int = 42 gpu: str = "H100" # minimum GPU requirement - negative_prompt: str = "" + negative_prompt: str | None = "" guidance_scale: float | None = None sigmas: list[float] | None = None enable_cpu_offload: bool = False @@ -93,7 +94,7 @@ class QualityTestConfig: return self.quantized_model return self.model or "" - def quantization_ref(self) -> str | None: + def quantization_ref(self) -> str | dict[str, object] | None: if self.quantized_model is not None: return None return self.quantization @@ -211,24 +212,14 @@ QUALITY_CONFIGS = [ quantization="fp8", task="t2v", prompt="A serene lakeside sunrise with mist over the water", - max_lpips=0.10, - height=256, - width=256, - num_frames=25, - num_inference_steps=8, - # Preserve the final scheduler trajectory used when this FP8 quality - # gate was established. Explicit LTX sigmas bypass dynamic shifting. - sigmas=[ - 1.0, - 0.92185378074646, - 0.8327768445014954, - 0.7303033471107483, - 0.6111654043197632, - 0.4709382653236389, - 0.30347853899002075, - 0.10000002384185791, - 0.0, - ], + max_lpips=0.20, + height=384, + width=512, + num_frames=73, + num_inference_steps=10, + # Do not override the recipe's negative conditioning, guidance, or + # scheduler trajectory: this gate follows the supported LTX default. + negative_prompt=None, ), ] @@ -295,8 +286,11 @@ def _generate_image(omni, config: QualityTestConfig): ).manual_seed(config.seed) torch.accelerator.reset_peak_memory_stats() + prompt = {"prompt": config.prompt} + if config.negative_prompt is not None: + prompt["negative_prompt"] = config.negative_prompt outputs = omni.generate( - {"prompt": config.prompt, "negative_prompt": config.negative_prompt}, + prompt, OmniDiffusionSamplingParams( height=config.height, width=config.width, @@ -327,8 +321,11 @@ def _generate_video(omni, config: QualityTestConfig): ).manual_seed(config.seed) torch.accelerator.reset_peak_memory_stats() + prompt = {"prompt": config.prompt} + if config.negative_prompt is not None: + prompt["negative_prompt"] = config.negative_prompt outputs = omni.generate( - {"prompt": config.prompt, "negative_prompt": config.negative_prompt}, + prompt, OmniDiffusionSamplingParams( height=config.height, width=config.width, @@ -464,6 +461,7 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch def test_generate_video_forwards_sigmas(monkeypatch): + from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False) @@ -474,7 +472,12 @@ def test_generate_video_forwards_sigmas(monkeypatch): class DummyOmni: def generate(self, _prompt, sampling): captured.sampling = sampling - return [SimpleNamespace(images=[np.zeros((1, 2, 2, 3), dtype=np.float32)])] + return [ + OmniRequestOutput.from_diffusion( + request_id="req", + images=[np.zeros((1, 2, 2, 3), dtype=np.float32)], + ) + ] config = QualityTestConfig( id="ltx-sigmas", @@ -490,6 +493,41 @@ def test_generate_video_forwards_sigmas(monkeypatch): assert captured.sampling.sigmas == [1.0, 0.5] +@pytest.mark.core_model +@pytest.mark.cpu +def test_ltx_quality_gate_uses_official_eager_defaults(monkeypatch): + from vllm_omni.outputs import OmniRequestOutput + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False) + monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False) + monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False) + captured = SimpleNamespace(prompt=None, sampling=None) + + class DummyOmni: + def generate(self, prompt, sampling): + captured.prompt = prompt + captured.sampling = sampling + return [ + OmniRequestOutput.from_diffusion( + request_id="req", + images=[np.zeros((1, 2, 2, 3), dtype=np.float32)], + ) + ] + + config = next(config for config in QUALITY_CONFIGS if config.id == "fp8_ltx2") + _generate_video(DummyOmni(), config) + + assert (config.width, config.height, config.num_frames, config.num_inference_steps) == (512, 384, 73, 10) + assert config.max_lpips == 0.20 + assert config.sigmas is None + assert config.guidance_scale is None + assert config.negative_prompt is None + assert "negative_prompt" not in captured.prompt + assert captured.sampling.sigmas is None + assert captured.sampling.guidance_scale is None + + _marks = hardware_marks(res={"cuda": "H100"}) _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None diff --git a/tests/diffusion/quantization/test_svdquant_config.py b/tests/diffusion/quantization/test_svdquant_config.py new file mode 100644 index 000000000..adbb88ed0 --- /dev/null +++ b/tests/diffusion/quantization/test_svdquant_config.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from unittest.mock import Mock + +import pytest +import torch +from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) + +from vllm_omni.quantization import svdquant_config +from vllm_omni.quantization.factory import ( + SUPPORTED_QUANTIZATION_METHODS, + build_quant_config, +) +from vllm_omni.quantization.svdquant_config import ( + DiffusionSVDQuantConfig, + DiffusionSVDQuantLinearMethod, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_config_from_checkpoint_dict() -> None: + config = DiffusionSVDQuantConfig.from_config( + { + "rank": 32, + "precision": "nvfp4", + "act_unsigned": False, + "modules_to_not_convert": ["blocks.0.adaln_proj.linear"], + } + ) + + assert config.rank == 32 + assert config.precision == "nvfp4" + assert config.modules_to_not_convert == ["blocks.0.adaln_proj.linear"] + assert config.get_supported_act_dtypes() == [torch.bfloat16] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"rank": 0}, "rank must be positive"), + ({"precision": "int4"}, "NVFP4 checkpoints only"), + ({"act_unsigned": True}, "unsigned activations"), + ], +) +def test_config_rejects_unsupported_phase1_options( + kwargs: dict[str, object], + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + DiffusionSVDQuantConfig(**kwargs) + + +def test_factory_builds_svdquant_with_extra_checkpoint_fields() -> None: + config = build_quant_config( + { + "quant_method": "svdquant", + "rank": 16, + "precision": "nvfp4", + "w4a16_modules": ["adaln_proj.linear"], + "w4a16_group_size": 64, + } + ) + + assert isinstance(config, DiffusionSVDQuantConfig) + assert config.rank == 16 + assert "svdquant" in SUPPORTED_QUANTIZATION_METHODS + + +def test_skipped_and_non_linear_modules_remain_unquantized() -> None: + config = DiffusionSVDQuantConfig(modules_to_not_convert=["blocks.0.adaln_proj.linear"]) + linear = Mock(spec=LinearBase) + + assert isinstance( + config.get_quant_method(linear, "blocks.0.adaln_proj.linear"), + UnquantizedLinearMethod, + ) + assert config.get_quant_method(torch.nn.ReLU(), "blocks.0.act") is None + + +def test_linear_method_creates_canonical_partitioned_weights( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None) + method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig(rank=32)) + layer = torch.nn.Module() + + method.create_weights( + layer, + input_size_per_partition=128, + output_partition_sizes=[64, 64], + input_size=128, + output_size=128, + params_dtype=torch.bfloat16, + ) + + assert layer.qweight.shape == (128, 64) + assert layer.qweight.dtype == torch.int8 + assert layer.wscales.shape == (8, 128) + assert layer.wscales.dtype == torch.float8_e4m3fn + assert layer.proj_down.shape == (128, 32) + assert layer.proj_up.shape == (128, 32) + assert layer.smooth_factor.shape == (128,) + assert layer.wcscales.shape == (128,) + assert layer.wtscale.shape == (1,) + + assert layer.qweight.input_dim == 1 + assert layer.qweight.output_dim == 0 + assert layer.wscales.input_dim == 0 + assert layer.wscales.output_dim == 1 + assert layer.proj_down.input_dim == 0 + assert layer.proj_up.output_dim == 0 + assert layer.wcscales.output_dim == 0 + + +def test_linear_method_rejects_misaligned_input_partition( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None) + method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig()) + + with pytest.raises(ValueError, match="divisible by the block size 16"): + method.create_weights( + torch.nn.Module(), + input_size_per_partition=24, + output_partition_sizes=[32], + input_size=24, + output_size=32, + params_dtype=torch.bfloat16, + ) + + +def test_active_linear_uses_svdquant_method( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None) + config = DiffusionSVDQuantConfig() + linear = Mock(spec=LinearBase) + + method = config.get_quant_method(linear, "blocks.0.attn.out_proj") + + assert isinstance(method, DiffusionSVDQuantLinearMethod) diff --git a/tests/diffusion/quantization/test_svdquant_linear.py b/tests/diffusion/quantization/test_svdquant_linear.py new file mode 100644 index 000000000..bdca88d9d --- /dev/null +++ b/tests/diffusion/quantization/test_svdquant_linear.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.quantization import svdquant_config as svdquant + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +class _FakeNvfp4Kernel: + def __init__(self, base_weight: torch.Tensor | None = None) -> None: + self.base_weight = base_weight + self.processed = False + self.last_input: torch.Tensor | None = None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + del layer + self.processed = True + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + del layer + assert bias is None + assert self.base_weight is not None + self.last_input = x.clone() + return torch.mm(x, self.base_weight) + + +def _register_parameter( + layer: torch.nn.Module, + name: str, + value: torch.Tensor, +) -> None: + layer.register_parameter( + name, + torch.nn.Parameter(value, requires_grad=False), + ) + + +def test_supports_only_validated_datacenter_blackwell() -> None: + assert not svdquant._supports_capability(SimpleNamespace(major=10, minor=0)) + assert svdquant._supports_capability(SimpleNamespace(major=10, minor=3)) + assert not svdquant._supports_capability(SimpleNamespace(major=11, minor=0)) + assert not svdquant._supports_capability(SimpleNamespace(major=12, minor=0)) + assert not svdquant._supports_capability(None) + + +def test_rejects_nvfp4_backend_with_incompatible_layer_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + svdquant._nvfp4_kernel.cache_clear() + monkeypatch.setattr( + svdquant, + "init_nvfp4_linear_kernel", + lambda: _FakeNvfp4Kernel(), + ) + + with pytest.raises(RuntimeError, match="incompatible with the SVDQuant checkpoint"): + svdquant._nvfp4_kernel() + + svdquant._nvfp4_kernel.cache_clear() + + +def test_prepare_weights_uses_existing_nvfp4_layout_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = _FakeNvfp4Kernel() + monkeypatch.setattr(svdquant, "_assert_supported", lambda: None) + monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel) + method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig()) + layer = torch.nn.Module() + qweight = torch.arange(16, dtype=torch.int8).reshape(4, 4) + wscales = torch.arange(8, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(2, 4) + _register_parameter(layer, "qweight", qweight) + _register_parameter(layer, "wscales", wscales) + _register_parameter(layer, "wtscale", torch.tensor([2.0], dtype=torch.bfloat16)) + _register_parameter( + layer, + "wcscales", + torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16), + ) + + method.process_weights_after_loading(layer) + + assert kernel.processed + assert not hasattr(layer, "qweight") + assert not hasattr(layer, "wscales") + assert layer.weight.dtype == torch.uint8 + assert torch.equal(layer.weight.view(torch.int8), qweight) + assert torch.equal(layer.weight_scale, wscales.transpose(0, 1)) + torch.testing.assert_close(layer.alpha, torch.tensor([2.0])) + assert layer.input_global_scale_inv.dtype == torch.float32 + assert layer.output_channel_scale.shape == (4,) + + +def test_prepare_weights_drops_identity_output_scale( + monkeypatch: pytest.MonkeyPatch, +) -> None: + kernel = _FakeNvfp4Kernel() + monkeypatch.setattr(svdquant, "_assert_supported", lambda: None) + monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel) + method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig()) + layer = torch.nn.Module() + _register_parameter(layer, "qweight", torch.zeros(4, 4, dtype=torch.int8)) + _register_parameter( + layer, + "wscales", + torch.ones(2, 4, dtype=torch.float8_e4m3fn), + ) + _register_parameter(layer, "wtscale", torch.ones(1, dtype=torch.bfloat16)) + _register_parameter(layer, "wcscales", torch.ones(4, dtype=torch.bfloat16)) + + method.process_weights_after_loading(layer) + + assert layer.output_channel_scale is None + + +def test_compatibility_path_applies_qkv_scale_and_rank_correction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + base_weight = torch.tensor( + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + dtype=torch.bfloat16, + ) + kernel = _FakeNvfp4Kernel(base_weight) + monkeypatch.setattr(svdquant, "_assert_supported", lambda: None) + monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel) + method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig()) + + layer = torch.nn.Module() + layer.smooth_factor = torch.tensor([2.0, 4.0], dtype=torch.bfloat16) + layer.proj_down = torch.tensor([[1.0], [2.0]], dtype=torch.bfloat16) + layer.proj_up = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16) + layer.output_channel_scale = torch.tensor([1.0, 2.0, 4.0], dtype=torch.bfloat16) + layer.output_size_per_partition = 3 + x = torch.tensor([[2.0, 4.0], [4.0, 8.0]], dtype=torch.bfloat16) + bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16) + + actual = method.apply(layer, x, bias) + + smoothed = x / layer.smooth_factor + expected = torch.mm(smoothed, base_weight) + expected = expected * layer.output_channel_scale + expected = expected + torch.mm(torch.mm(x, layer.proj_down), layer.proj_up.T) + expected = expected + bias + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(kernel.last_input, smoothed) + + +def test_compatibility_path_rejects_fp16_activations( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(svdquant, "_assert_supported", lambda: None) + method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig()) + layer = torch.nn.Module() + with pytest.raises(ValueError, match="requires BF16 activations"): + method.apply( + layer, + torch.zeros(1, 16, dtype=torch.float16), + ) diff --git a/tests/diffusion/quantization/test_svdquant_tp_loading.py b/tests/diffusion/quantization/test_svdquant_tp_loading.py new file mode 100644 index 000000000..14138dc1c --- /dev/null +++ b/tests/diffusion/quantization/test_svdquant_tp_loading.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) + +from vllm_omni.quantization import svdquant_config as svdquant + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _linear_without_init(linear_type: type[torch.nn.Module], tp_rank: int): + layer = object.__new__(linear_type) + torch.nn.Module.__init__(layer) + layer.tp_rank = tp_rank + return layer + + +def _values(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + values = torch.arange(torch.Size(shape).numel(), dtype=torch.int64).remainder(127) + return values.reshape(shape).to(dtype) + + +def _method(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(svdquant, "_assert_supported", lambda: None) + return svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig(rank=8)) + + +def test_column_parallel_loader_shards_svdquant_outputs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + layer = _linear_without_init(ColumnParallelLinear, tp_rank=1) + method = _method(monkeypatch) + method.create_weights( + layer, + input_size_per_partition=128, + output_partition_sizes=[32], + input_size=128, + output_size=64, + params_dtype=torch.bfloat16, + weight_loader=layer.weight_loader, + ) + + loaded_tensors = { + "qweight": _values((64, 64), torch.int8), + "wscales": _values((8, 64), torch.float8_e4m3fn), + "proj_up": _values((64, 8), torch.bfloat16), + "wcscales": _values((64,), torch.bfloat16), + } + for name, loaded in loaded_tensors.items(): + parameter = getattr(layer, name) + parameter.weight_loader(parameter, loaded) + expected = loaded.narrow(parameter.output_dim, 32, 32) + torch.testing.assert_close(parameter, expected) + + +def test_row_parallel_loader_shards_svdquant_inputs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + layer = _linear_without_init(RowParallelLinear, tp_rank=1) + method = _method(monkeypatch) + method.create_weights( + layer, + input_size_per_partition=64, + output_partition_sizes=[64], + input_size=128, + output_size=64, + params_dtype=torch.bfloat16, + weight_loader=layer.weight_loader, + ) + + loaded_tensors = { + "qweight": _values((64, 64), torch.int8), + "wscales": _values((8, 64), torch.float8_e4m3fn), + "proj_down": _values((128, 8), torch.bfloat16), + "smooth_factor": _values((128,), torch.bfloat16), + } + for name, loaded in loaded_tensors.items(): + parameter = getattr(layer, name) + parameter.weight_loader(parameter, loaded) + shard_size = parameter.shape[parameter.input_dim] + expected = loaded.narrow(parameter.input_dim, shard_size, shard_size) + torch.testing.assert_close(parameter, expected) + + +def test_qkv_parallel_loader_shards_fused_svdquant_outputs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + layer = _linear_without_init(QKVParallelLinear, tp_rank=1) + layer.total_num_heads = 4 + layer.total_num_kv_heads = 2 + layer.num_heads = 2 + layer.num_kv_heads = 1 + layer.num_kv_head_replicas = 1 + layer.head_size = 8 + layer.v_head_size = 8 + method = _method(monkeypatch) + method.create_weights( + layer, + input_size_per_partition=128, + output_partition_sizes=[16, 8, 8], + input_size=128, + output_size=64, + params_dtype=torch.bfloat16, + weight_loader=layer.weight_loader, + ) + + loaded_tensors = { + "qweight": _values((64, 64), torch.int8), + "wscales": _values((8, 64), torch.float8_e4m3fn), + "proj_up": _values((64, 8), torch.bfloat16), + "wcscales": _values((64,), torch.bfloat16), + } + expected_global_outputs = [slice(16, 32), slice(40, 48), slice(56, 64)] + for name, loaded in loaded_tensors.items(): + parameter = getattr(layer, name) + parameter.weight_loader(parameter, loaded) + expected = torch.cat( + [ + loaded.narrow(parameter.output_dim, part.start, part.stop - part.start) + for part in expected_global_outputs + ], + dim=parameter.output_dim, + ) + torch.testing.assert_close(parameter, expected) diff --git a/tests/diffusion/quantization/test_wan_autoround_mxfp4.py b/tests/diffusion/quantization/test_wan_autoround_mxfp4.py new file mode 100644 index 000000000..89c317970 --- /dev/null +++ b/tests/diffusion/quantization/test_wan_autoround_mxfp4.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Wan AutoRound-format MXFP4 configuration and layer-mapping tests.""" + +from unittest.mock import patch + +import pytest +from torch import nn +from vllm.model_executor.layers.linear import LinearBase +from vllm.model_executor.model_loader.utils import configure_quant_config + +from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import ( + WanTransformer3DModel, +) +from vllm_omni.quantization import build_quant_config +from vllm_omni.quantization.inc_config import OmniINCConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _autoround_mxfp4_config() -> dict: + return { + "bits": 4, + "act_bits": 4, + "data_type": "mx_fp", + "act_data_type": "mx_fp", + "group_size": 32, + "act_group_size": 32, + "sym": True, + "act_sym": True, + "act_dynamic": True, + "enable_quanted_input": False, + "block_name_to_quantize": "blocks", + "quant_method": "auto-round", + "packing_format": "auto_round:llm_compressor", + "autoround_version": "0.14.2", + } + + +def test_build_wan_autoround_mxfp4_config(): + config = build_quant_config(_autoround_mxfp4_config()) + + assert isinstance(config, OmniINCConfig) + assert config.get_name() == "inc" + assert config.weight_bits == 4 + assert config.data_type == "mx_fp" + assert config.group_size == 32 + assert config.packing_format == "auto_round:llm_compressor" + assert config.block_name_to_quantize == ["blocks"] + + +def test_wan_autoround_config_uses_runtime_layer_names(): + config = build_quant_config(_autoround_mxfp4_config()) + configure_quant_config(config, WanTransformer3DModel) + + assert config.packed_modules_mapping is WanTransformer3DModel.packed_modules_mapping + assert config.block_name_to_quantize == ["blocks"] + + layer = object.__new__(LinearBase) + nn.Module.__init__(layer) + for prefix in ( + "blocks.0.ffn.net_0.proj", + "blocks.0.ffn.net_2", + "blocks.0.attn1.to_out", + ): + layer_config = config.config_parser.resolve(layer, prefix) + assert layer_config.quantized + assert layer_config.is_mxfp4 + + +def test_wan_autoround_scope_dispatches_only_blocks_to_mxfp4(): + config = build_quant_config(_autoround_mxfp4_config()) + configure_quant_config(config, WanTransformer3DModel) + + layer = object.__new__(LinearBase) + nn.Module.__init__(layer) + layer_config = config.config_parser.resolve(layer, "blocks.0.ffn.net_0.proj") + outside_config = config.config_parser.resolve(layer, "condition_embedder.time_proj") + + assert layer_config.quantized + assert layer_config.is_mxfp4 + assert not outside_config.quantized + + fake_method = object() + with patch( + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_scheme.INCMxfp4Scheme.get_linear_method", + return_value=fake_method, + ): + method = config.get_quant_method(layer, "blocks.0.ffn.net_0.proj") + assert method is fake_method diff --git a/tests/diffusion/sched/test_dmd2_sigma_schedule.py b/tests/diffusion/sched/test_dmd2_sigma_schedule.py index 108cbe68c..8fbad959b 100644 --- a/tests/diffusion/sched/test_dmd2_sigma_schedule.py +++ b/tests/diffusion/sched/test_dmd2_sigma_schedule.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest from vllm_omni.diffusion.sched import DMD2SigmaSchedule @@ -18,7 +18,8 @@ def test_step_count_is_intervals_not_boundaries(): [ [], [1.0], - [0.9, 0.5, 0.0], + [1.5, 0.5, 0.0], + [0.0, 0.5, 0.0], [1.0, 0.5, 0.1], [1.0, 0.4, 0.6, 0.0], [1.0, 0.5, 0.5, 0.0], @@ -31,6 +32,15 @@ def test_malformed_positions_are_rejected(positions): DMD2SigmaSchedule.from_positions(positions) +def test_a_capped_first_position_is_a_schedule_not_a_malformed_one(): + # A DMD2 student trained with max_timestep_ratio < 1 never sees pure noise, + # so its opening rung sits just below 1.0 (FastH3: 999/1000). + schedule = DMD2SigmaSchedule.from_positions([0.999, 0.749, 0.5, 0.25, 0.0]) + + assert schedule.base_schedule[0] == 0.999 + assert schedule.num_inference_steps == 4 + + def test_one_base_schedule_drives_several_shift_scales(): schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0]) diff --git a/tests/diffusion/test_async_output_timeout.py b/tests/diffusion/test_async_output_timeout.py new file mode 100644 index 000000000..2875ed02f --- /dev/null +++ b/tests/diffusion/test_async_output_timeout.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the async-output wait bound (_async_output_timeout).""" + +import logging + +import pytest + +from vllm_omni.diffusion.diffusion_engine import ( + _ASYNC_OUTPUT_TIMEOUT_DEFAULT, + _ASYNC_OUTPUT_TIMEOUT_ENV, + _async_output_timeout, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestAsyncOutputTimeoutIsConfigurable: + def test_default_is_generous(self, monkeypatch): + """The bound covers a copy queued behind a full denoise step, so the + default must not be tight enough to abort a slow but healthy render. + """ + monkeypatch.delenv(_ASYNC_OUTPUT_TIMEOUT_ENV, raising=False) + assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT == 600.0 + + def test_env_var_overrides_the_default(self, monkeypatch): + monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "1800") + assert _async_output_timeout() == 1800.0 + + def test_env_var_accepts_a_float(self, monkeypatch): + monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "45.5") + assert _async_output_timeout() == 45.5 + + def test_value_is_not_frozen_at_import(self, monkeypatch): + """Resolved per call rather than captured in a module constant, so the + value never goes stale relative to os.environ and no module reload is + needed to exercise it. + """ + monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "60") + assert _async_output_timeout() == 60.0 + monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "120") + assert _async_output_timeout() == 120.0 + + +class TestAsyncOutputTimeoutRejectsBadValues: + """This runs on the request path, so a typo in the environment must degrade + to the default rather than fail the generation. + """ + + @pytest.mark.parametrize("value", ["", "abc", "30s", "None"]) + def test_non_numeric_falls_back_to_the_default(self, monkeypatch, caplog, value): + monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value) + with caplog.at_level(logging.WARNING): + assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT + + @pytest.mark.parametrize("value", ["0", "-1", "-0.5"]) + def test_non_positive_falls_back_to_the_default(self, monkeypatch, caplog, value): + monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value) + with caplog.at_level(logging.WARNING): + assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py index da684a4f5..a4bff4625 100644 --- a/tests/diffusion/test_async_output_worker.py +++ b/tests/diffusion/test_async_output_worker.py @@ -3,6 +3,8 @@ """Unit tests for WorkerProc async output logic.""" import queue +import threading +import time from unittest.mock import MagicMock import pytest @@ -29,9 +31,59 @@ def _make_worker_proc(step_execution=False): proc._running = True proc._async_output_queue = queue.Queue() if not step_execution else None proc._async_output_thread = None + proc._async_output_done = threading.Condition() + proc._async_output_pending = 0 + proc._result_mq_lock = threading.Lock() return proc +class _OverlapDetectingQueue: + """Fails if two threads are inside enqueue() at the same time. + + Mirrors vLLM's MessageQueue, whose acquire_write() reads current_idx, + writes that block, then advances it. Overlapping writers can target the + same ring block and silently drop a message. + """ + + def __init__(self): + self._active = 0 + self._guard = threading.Lock() + self.overlaps = 0 + self.messages = [] + + def enqueue(self, msg, timeout=None): + with self._guard: + self._active += 1 + if self._active > 1: + self.overlaps += 1 + time.sleep(0.001) # widen the window a real ring-buffer write would have + with self._guard: + self.messages.append(msg) + self._active -= 1 + + +class TestResultQueueSingleWriter: + """Worker must serialize writes to the single-writer result queue.""" + + def test_concurrent_enqueues_do_not_overlap(self): + proc = _make_worker_proc(step_execution=False) + detector = _OverlapDetectingQueue() + proc.result_mq = detector + + def writer(tag): + for i in range(20): + proc._enqueue_result(f"{tag}-{i}") + + threads = [threading.Thread(target=writer, args=(tag,)) for tag in ("compute", "output")] + for t in threads: + t.start() + for t in threads: + t.join() + + assert detector.overlaps == 0, f"{detector.overlaps} overlapping writes to result_mq" + assert len(detector.messages) == 40 + + class TestGenerateAsyncOutputId: """Test _generate_async_output_id uniqueness.""" @@ -192,6 +244,87 @@ class TestAsyncOutputLoopLogic: assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls" +class TestDrainAsyncOutputs: + """Async outputs must finish packing before device memory is released.""" + + def test_drain_returns_immediately_when_nothing_in_flight(self): + proc = _make_worker_proc(step_execution=False) + + assert proc.drain_async_outputs(timeout=0.01) is True + + def test_return_result_keeps_output_pending_until_loop_completes(self, mocker): + """_return_result marks work pending; the bg loop clears it.""" + proc = _make_worker_proc(step_execution=False) + + mock_platform = mocker.MagicMock() + mock_platform.record_device_event.return_value = MagicMock() + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform", + mock_platform, + ) + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm") + mock_torch = mocker.MagicMock() + mock_torch.accelerator.current_accelerator.return_value.type = "cpu" + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch) + + proc._return_result(DiffusionOutput(output="data"), rpc_id="1") + + # Packing has not run yet, so a sleep here would free tensors in flight. + assert proc._async_output_pending == 1 + assert proc.drain_async_outputs(timeout=0.05) is False + + thread = threading.Thread(target=proc._async_output_loop, daemon=True) + thread.start() + try: + assert proc.drain_async_outputs(timeout=5.0) is True + finally: + proc._async_output_queue.put(None) + thread.join(timeout=2.0) + + def test_sleep_rpc_drains_before_releasing_memory(self, mocker): + proc = _make_worker_proc(step_execution=False) + proc.worker = MagicMock() + + call_order: list[str] = [] + mocker.patch.object( + proc, + "drain_async_outputs", + side_effect=lambda *a, **kw: call_order.append("drain") or True, + ) + proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append("sleep") + + proc._execute_rpc( + { + "method": "handle_sleep_task", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": False, + "collect_rank_status": False, + } + ) + + assert call_order == ["drain", "sleep"] + + def test_non_sleep_rpc_does_not_drain(self, mocker): + proc = _make_worker_proc(step_execution=False) + proc.worker = MagicMock() + drain = mocker.patch.object(proc, "drain_async_outputs", return_value=True) + + proc._execute_rpc( + { + "method": "execute_model", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": False, + "collect_rank_status": False, + } + ) + + drain.assert_not_called() + + class TestWorkerProcShutdown: def test_shutdown_stops_async_thread_and_releases_queues(self): import threading diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py index 665aa6825..3cf8f1428 100644 --- a/tests/diffusion/test_compile.py +++ b/tests/diffusion/test_compile.py @@ -6,6 +6,7 @@ import torch.nn as nn import vllm_omni.diffusion.compile as compile_module from vllm_omni.diffusion.compile import regionally_compile +from vllm_omni.diffusion.hooks import HookRegistry, ModelHook pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -76,6 +77,37 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards +def test_regionally_compile_keeps_hook_dispatch_outside_compiled_graph(monkeypatch): + model = _ModelWithWrappedRepeatedBlocks() + block = model.transformer_blocks[0] + registry = HookRegistry.get_or_create(block) + hook = ModelHook() + registry.register_hook("test", hook) + + wrapped_forward = block.forward + original_forward = block._omni_original_forward + compile_calls = [] + + def _compile(fn, *args, **kwargs): + compile_calls.append(fn) + + def _compiled(*fn_args, **fn_kwargs): + return f"compiled:{fn(*fn_args, **fn_kwargs)}" + + return _compiled + + monkeypatch.setattr(compile_module.torch, "compile", _compile) + + regionally_compile(model) + + assert compile_calls[0] is original_forward + assert compile_calls[0] is not wrapped_forward + assert block.forward is wrapped_forward + assert block._omni_original_forward is not original_forward + assert hook.fn_ref.original_forward is block._omni_original_forward + assert block("ok") == "compiled:ok" + + def test_compiled_block_preserves_forward_signature_for_inspection(monkeypatch): """cache-dit matches blocks via inspect.signature(block.forward). diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index e25d6497c..60d0c63b8 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -92,6 +92,30 @@ class TestParallelConfigPropagation: od = _roundtrip_diffusion_config(model="x", parallel_config=pc) assert od.num_gpus == 2 + def test_dp_is_inferred_from_num_gpus(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2) + od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8) + assert od.parallel_config.data_parallel_size == 2 + assert od.parallel_config.world_size == 8 + + def test_explicit_dp_is_validated_against_num_gpus(self): + pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2) + od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4) + assert od.parallel_config.data_parallel_size == 2 + + with pytest.raises(ValueError, match="does not match WORLD-derived value"): + OmniDiffusionConfig.from_kwargs( + model="x", + parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2), + num_gpus=8, + ) + + def test_hsdp_does_not_infer_ordinary_dp(self): + pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4) + od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4) + assert od.parallel_config.data_parallel_size == 1 + assert od.parallel_config.world_size == 4 + class TestCreateDefaultDiffusion: """Verify engine_args structure from create_default_diffusion.""" @@ -125,10 +149,33 @@ class TestCreateDefaultDiffusion: assert ea["enforce_eager"] is True assert ea["lora_path"] == "/tmp/lora" - def test_diffusion_kv_mode_roundtrip(self): - od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler") + def test_diffusion_kv_mode_roundtrip(self, monkeypatch): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + od = _roundtrip_diffusion_config( + model="x", + diffusion_kv_mode="paged_scheduler", + diffusion_kv_max_rows_per_request=2, + ) assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + assert od.diffusion_kv_max_rows_per_request == 2 + + def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch): + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500) + od = _roundtrip_diffusion_config( + model="x", + kv_cache_memory_bytes=4096, + gpu_memory_utilization=0.75, + max_num_batched_tokens=2048, + max_model_len=4096, + ) + + assert od.kv_cache_memory_bytes == 4096 + assert od.gpu_memory_utilization == 0.75 + assert od.max_num_batched_tokens == 2048 + assert od.max_model_len == 4096 def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index a0c6acb02..04d528af4 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import queue @@ -197,6 +197,59 @@ class TestRequestBatchCapability: assert diffusion_engine_module.supports_request_batch(od_config) is True + def test_engine_rejects_multi_seq_diffusers_using_adapter_capability( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace( + model_class_name="QwenImagePipeline", + custom_pipeline_args=None, + diffusion_load_format="diffusers", + streaming_output=False, + max_num_seqs=2, + ) + registry_load = mocker.Mock( + side_effect=lambda model_class_name: ( + _SingleRequestPipeline if model_class_name == "DiffusersAdapterPipeline" else _BatchCapablePipeline + ) + ) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + registry_load, + ) + + engine = DiffusionEngine.__new__(DiffusionEngine) + with pytest.raises(ValueError, match="max_num_seqs=1"): + engine._resolve_execution_mode(od_config) + + registry_load.assert_called_once_with("DiffusersAdapterPipeline") + + def test_engine_prefers_batch_capable_custom_pipeline_over_diffusers_adapter( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace( + model_class_name="QwenImagePipeline", + custom_pipeline_args={"pipeline_class": _BatchCapablePipeline}, + diffusion_load_format="diffusers", + streaming_output=False, + max_num_seqs=2, + ) + registry_load = mocker.Mock(return_value=_SingleRequestPipeline) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + registry_load, + ) + + engine = DiffusionEngine.__new__(DiffusionEngine) + + assert engine._resolve_execution_mode(od_config) is DiffusionExecutionMode.REQUEST_BATCH + registry_load.assert_not_called() + def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None: od_config = SimpleNamespace( model_class_name="SinglePipeline", diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index fb8858a48..ba7d7d5d2 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -12,6 +12,7 @@ import pytest from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -88,6 +89,50 @@ def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None: assert engine.scheduler.get_request_state(request_id) is None +def _make_kv_cleanup_engine(mode: DiffusionKVCacheMode) -> tuple[DiffusionEngine, Mock]: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(diffusion_kv_mode=mode) + cleanup = Mock() + engine.executor = SimpleNamespace(remove_diffusion_kv_requests=cleanup) + return engine, cleanup + + +def test_paged_terminal_output_clears_worker_rows() -> None: + engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER) + engine._finalize_finished_request = lambda request_id, *_args: request_id + engine._put_output = lambda *_args: None + + engine._emit_finished_outputs({"req-0", "req-1"}) + + cleanup.assert_called_once() + assert set(cleanup.call_args.args[0]) == {"req-0", "req-1"} + + +def test_paged_abort_clears_worker_rows_before_scheduler_free() -> None: + engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER) + engine.scheduler = SimpleNamespace( + get_request_state=lambda _request_id: None, + finish_requests=Mock(side_effect=AssertionError("unknown request must not be logically freed")), + ) + + engine._abort_requests(["req-idle", "req-idle"]) + + cleanup.assert_called_once_with(["req-idle"]) + engine.scheduler.finish_requests.assert_not_called() + + +def test_dense_terminal_and_abort_paths_skip_worker_row_cleanup() -> None: + engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.DENSE_LEGACY) + engine._finalize_finished_request = lambda request_id, *_args: request_id + engine._put_output = lambda *_args: None + engine.scheduler = SimpleNamespace(get_request_state=lambda _request_id: None) + + engine._emit_finished_outputs({"req-0"}) + engine._abort_requests(["req-idle"]) + + cleanup.assert_not_called() + + def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None: od_config = SimpleNamespace( custom_pipeline_args=None, @@ -161,9 +206,83 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa assert exc_info.value is initialization_error custom_scheduler.initialize.assert_called_once_with(od_config) custom_scheduler.close.assert_called_once_with() + + +def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_executor = SimpleNamespace(shutdown=Mock()) + monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None) + monkeypatch.setattr( + DiffusionEngine, + "_resolve_execution_mode", + lambda self, config: DiffusionExecutionMode.REQUEST_BATCH, + ) + monkeypatch.setattr( + DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor) + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane", + Mock(side_effect=RuntimeError("KV initialization failed")), + ) + + with pytest.raises(RuntimeError, match="KV initialization failed"): + DiffusionEngine(SimpleNamespace()) + fake_executor.shutdown.assert_called_once_with() +def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None: + events = [] + expected_profile_requests = [object()] + fake_executor = SimpleNamespace(shutdown=Mock()) + kv_cache_config = object() + kv_vllm_config = object() + od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler") + + monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None) + monkeypatch.setattr( + DiffusionEngine, + "_resolve_execution_mode", + lambda self, config: DiffusionExecutionMode.REQUEST_BATCH, + ) + monkeypatch.setattr( + DiffusionEngine, + "_init_executor", + lambda self, config: setattr(self, "executor", fake_executor), + ) + monkeypatch.setattr( + DiffusionEngine, + "_prepare_diffusion_kv_profile_requests", + lambda self: events.append("prepare-profile") or expected_profile_requests, + ) + + def initialize(executor, config, *, profile_requests): + events.append("profile-workers") + assert executor is fake_executor + assert config is od_config + assert profile_requests is expected_profile_requests + return kv_cache_config, 16, 16, kv_vllm_config + + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane", + initialize, + ) + + def init_scheduler(self, config, scheduler, *args, **kwargs): + events.append("initialize-scheduler") + self.scheduler = SimpleNamespace(close=Mock()) + + monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler) + monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None) + monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None) + monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None) + + DiffusionEngine(od_config) + + assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"] + + @pytest.mark.asyncio async def test_step_compatibility_wrapper_returns_final_batch() -> None: engine = _make_engine() diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 50dde3083..247fee029 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -1,9 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import Mock + import pytest from vllm_omni.diffusion import io_support +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -55,3 +63,151 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk ) assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2 + + +def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None: + assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3 + assert io_support.get_dummy_run_num_image_inputs("unknown") == 1 + + +def test_dense_mode_does_not_build_kv_profile_request() -> None: + engine = object.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY) + engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile")) + + assert engine._prepare_diffusion_kv_profile_requests() is None + engine._make_dummy_request.assert_not_called() + + +@pytest.mark.parametrize( + ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"), + [ + (DiffusionExecutionMode.STEP_BATCH, False, 2, 2), + (DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1), + (DiffusionExecutionMode.STEP_BATCH, True, 3, 3), + ], +) +def test_paged_kv_profile_requests_match_per_rank_batch( + execution_mode: DiffusionExecutionMode, + uses_dlo_dp: bool, + max_num_seqs: int, + expected_profile_requests: int, +) -> None: + engine = object.__new__(DiffusionEngine) + engine.execution_mode = execution_mode + engine.od_config = SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + model_class_name="HunyuanImage3ForCausalMM", + max_num_seqs=max_num_seqs, + parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1), + enable_distributed_layerwise_offload=uses_dlo_dp, + dlo_use_allgather=True, + ) + request = OmniDiffusionRequest( + prompt="profile", + request_id="profile-request", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + ) + engine._make_dummy_request = Mock(return_value=request) + prepared_layout = object() + scheduler_kv_state = ( + DiffusionKVRequest( + "profile-request/diffusion-kv/0", + sequence_id=0, + prefix_len=3, + target_len=8, + seq_len=12, + ), + DiffusionKVRequest( + "profile-request/diffusion-kv/1", + sequence_id=1, + prefix_len=4, + target_len=8, + seq_len=13, + ), + ) + + def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest: + req.prepared_layout = prepared_layout + req.diffusion_kv_requests = scheduler_kv_state + return req + + engine._prepare_request_for_admission = Mock(side_effect=preprocess) + + result = engine._prepare_diffusion_kv_profile_requests() + + assert result is not None + assert len(result) == expected_profile_requests + assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests + assert all(profile_request.is_dummy_run() for profile_request in result) + assert all(profile_request.prepared_layout is prepared_layout for profile_request in result) + assert all(profile_request.diffusion_kv_requests is None for profile_request in result) + assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests + assert engine._diffusion_kv_profile_limits == (2, 13, 8) + assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2 + engine._make_dummy_request.assert_called_once_with( + height=1024, + width=1024, + guidance_scale=5.0, + num_image_inputs=3, + ) + engine._prepare_request_for_admission.assert_called_once_with(request) + + +@pytest.mark.parametrize( + ("num_sequences", "seq_len", "target_len"), + [ + (3, 16, 8), + (1, 17, 8), + (1, 16, 9), + ], +) +def test_paged_kv_admission_rejects_shape_beyond_profile_envelope( + num_sequences: int, + seq_len: int, + target_len: int, +) -> None: + engine = object.__new__(DiffusionEngine) + engine._diffusion_kv_profile_limits = (2, 16, 8) + engine.pre_process_func = lambda request: request + request = OmniDiffusionRequest( + prompt="too large", + request_id="too-large", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + diffusion_kv_requests=tuple( + DiffusionKVRequest( + f"too-large/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=0, + target_len=target_len, + seq_len=seq_len, + ) + for sequence_id in range(num_sequences) + ), + ) + + with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"): + engine._prepare_request_for_admission(request) + + +def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None: + engine = object.__new__(DiffusionEngine) + engine._diffusion_kv_profile_limits = (2, 16, 8) + engine.pre_process_func = lambda request: request + request = OmniDiffusionRequest( + prompt="fits", + request_id="fits", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + diffusion_kv_requests=tuple( + DiffusionKVRequest( + f"fits/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=8, + target_len=8, + seq_len=16, + ) + for sequence_id in range(2) + ), + ) + + assert engine._prepare_request_for_admission(request) is request diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index 8f434c0e2..dc909a2b9 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Source-level regression tests for diffusion output/engine helpers. These tests verify naming conventions and patterns by inspecting source code @@ -37,10 +37,50 @@ _FORMATTER_PATH = os.path.normpath( "output_formatter.py", ) ) +_METRICS_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "metrics", + "utils.py", + ) +) +_INLINE_CLIENT_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "diffusion", + "inline_stage_diffusion_client.py", + ) +) +_STAGE_PROC_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "diffusion", + "stage_diffusion_proc.py", + ) +) +_ORCHESTRATOR_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "engine", + "orchestrator.py", + ) +) def _read_source(path: str) -> str: - with open(path) as f: + with open(path, encoding="utf-8") as f: return f.read() @@ -100,9 +140,116 @@ class TestMetricKeys: raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()") assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source - assert '"diffusion_engine_total_time_ms": step_total_ms' in step_streaming_source + # step_total_ms is no longer emitted as a metric (no family consumes + # it); it lives only in the debug log breakdown below. + assert '"diffusion_engine_total_time_ms"' not in step_streaming_source assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source + def test_step_streaming_piggybacks_scheduler_waiting_snapshot(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "diffusion_scheduler_waiting_metrics(" in step_streaming_source + assert "_scheduler_num_waiting_reqs" in step_streaming_source + + def test_step_streaming_preserves_snapshot_on_postprocess_error(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "scheduler_metrics = diffusion_scheduler_waiting_metrics(" in step_streaming_source + assert 'setattr(exc, "diffusion_metrics", scheduler_metrics)' in step_streaming_source + assert step_streaming_source.index("scheduler_metrics =") < step_streaming_source.index( + "self.postprocess_output(" + ) + + def test_abort_keeps_output_consumers_alive_for_terminal_snapshot(self) -> None: + inline_source = _read_source(_INLINE_CLIENT_PATH) + abort_source = _get_function_source(inline_source, "InlineStageDiffusionClient", "abort_requests_async") + assert ".cancel(" not in abort_source + + proc_source = _read_source(_STAGE_PROC_PATH) + run_loop_source = _get_function_source(proc_source, "StageDiffusionProc", "run_loop") + abort_branch = run_loop_source[run_loop_source.index('elif msg_type == "abort":') :] + abort_branch = abort_branch[: abort_branch.index('elif msg_type == "collective_rpc":')] + assert ".cancel(" not in abort_branch + + def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None: + """A metrics-only sentinel contributes its queue depth and is not routed. + + Absorption lives in ``_absorb_diffusion_metrics`` rather than inline in a + loop, so both orchestration loops share one implementation. The ordering + is asserted where it now lives: the snapshot inside the helper, and the + absorb-before-route guard inside every loop that polls diffusion output. + """ + source = _read_source(_ORCHESTRATOR_PATH) + + # Snapshot before the verdict, so a sentinel still reports its waiting + # depth on the way out instead of being dropped unaccounted. + absorb_source = _get_function_source(source, "Orchestrator", "_absorb_diffusion_metrics") + snapshot_pos = absorb_source.index("_update_stage_replica_waiting(") + sentinel_pos = absorb_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID") + assert snapshot_pos < sentinel_pos + + # Absorb before routing, or a sentinel reaches the downstream consumer + # as if it were a real request output. + for loop_name in ("_orchestration_loop", "_orchestration_loop_event_driven"): + loop_source = _get_function_source(source, "Orchestrator", loop_name) + absorb_pos = loop_source.index("self._absorb_diffusion_metrics(") + route_pos = loop_source.index("record_output_timestamps(") + assert absorb_pos < route_pos, loop_name + + +class TestVaeDecodeEmit: + """VAE decode timing is sourced from ``DiffusionOutput.stage_durations`` + (populated by the diffusion pipeline profiler) and emitted as + ``vae_decode_time_ms`` so the ``_MS_TO_S`` accumulator picks it up. + """ + + def test_extract_vae_decode_ms_helper_exists(self) -> None: + source = _read_source(_METRICS_PATH) + helper_src = _get_function_source(source, None, "extract_diffusion_vae_decode_ms") + assert ".vae.decode" in helper_src, "helper must key on the '.vae.decode' suffix" + + def test_step_streaming_emits_vae_decode_time_ms(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "extract_diffusion_vae_decode_ms(" in step_streaming_source, ( + "step_streaming must call extract_diffusion_vae_decode_ms to source VAE decode timing" + ) + assert '"vae_decode_time_ms"' in step_streaming_source, ( + "step_streaming must emit vae_decode_time_ms key for the _MS_TO_S accumulator" + ) + + +class TestDiffuseForwardEmit: + """Forward-only timing → forward_time_ms, mirroring TestVaeDecodeEmit.""" + + def test_extract_diffuse_ms_helper_exists(self) -> None: + source = _read_source(_METRICS_PATH) + helper_src = _get_function_source(source, None, "extract_diffusion_denoise_ms") + assert ".diffuse" in helper_src, "helper must key on the '.diffuse' suffix" + + def test_step_streaming_emits_forward_time_ms(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "extract_diffusion_denoise_ms(" in step_streaming_source, ( + "step_streaming must call extract_diffusion_denoise_ms to source forward-only timing" + ) + assert '"forward_time_ms"' in step_streaming_source, ( + "step_streaming must emit forward_time_ms key for the _MS_TO_S accumulator" + ) + + +class TestKvRecvEmit: + """KV-recv timing → kv_recv_time_ms, sourced from DiffusionOutput.kv_recv_ms + (set by the runner's _prepare_request_for_forward, not the profiler).""" + + def test_step_streaming_emits_kv_recv_time_ms(self) -> None: + source = _read_source(_ENGINE_PATH) + step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming") + assert "kv_recv_ms" in step_streaming_source, "step_streaming must read output.kv_recv_ms" + assert '"kv_recv_time_ms"' in step_streaming_source, ( + "step_streaming must emit kv_recv_time_ms key for the _MS_TO_S accumulator" + ) + class TestDummyRunAllocation: """Verify _dummy_run generates exact-sized audio arrays.""" diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index 1d4c2d17e..119939ae2 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -123,6 +123,7 @@ class _ConcurrencyTrackingExecutor: def _make_request(tag: str): sampling_params = dict(_SAMPLING_KEY_DEFAULTS) sampling_params["num_inference_steps"] = 1 + sampling_params["extra_args"] = {} return SimpleNamespace( request_id=tag, prompt=f"prompt_{tag}", @@ -472,6 +473,38 @@ async def test_cancelled_future_does_not_kill_busy_loop(): _stop_engine(engine) +@pytest.mark.asyncio +async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + + def fail_reservation(*args, **kwargs): + raise RuntimeError("native allocation bug") + + engine.scheduler._diffusion_kv_manager = SimpleNamespace( + has_request=lambda request_id: False, + reserve_request=fail_reservation, + free_request=lambda request_id: None, + ) + request = _make_request("kv-error") + request.diffusion_kv_requests = (object(),) + try: + response_stream = engine.async_add_req_and_stream_response(request) + output_queue = engine._out_streams[request.request_id] + for _ in range(300): + if not output_queue.empty(): + break + await asyncio.sleep(0.01) + assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream" + + output = await anext(response_stream) + assert output.error == "native allocation bug" + assert output.finished + assert engine.worker_thread.is_alive() + finally: + _stop_engine(engine) + + # ─────────────────────────── _RpcTask basics ─────────────────────────────── diff --git a/tests/diffusion/test_diffusion_ipc.py b/tests/diffusion/test_diffusion_ipc.py index 3e4c587f1..aec062057 100644 --- a/tests/diffusion/test_diffusion_ipc.py +++ b/tests/diffusion/test_diffusion_ipc.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import contextlib import multiprocessing as mp @@ -7,6 +7,7 @@ import queue import threading import time from multiprocessing import shared_memory +from multiprocessing.connection import Connection import numpy as np import pytest @@ -29,7 +30,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def _large_numel(dtype: torch.dtype) -> int: - return (_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1 + return int(_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1 def _cleanup_shm_handle(value: object) -> None: @@ -38,7 +39,7 @@ def _cleanup_shm_handle(value: object) -> None: _unpack_if_shm_handle(value) -def _result_queue_worker(connection, rank: int) -> None: +def _result_queue_worker(connection: Connection, rank: int) -> None: """Send one nested NumPy result through a worker-owned MessageQueue.""" result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0]) try: @@ -180,7 +181,7 @@ def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None: def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None: - # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output. + # LTX2 returns (video, audio) tuples as DiffusionOutput.output. video = torch.arange(300_000, dtype=torch.float32) audio = torch.arange(300_000, dtype=torch.float32) * 3 output = DiffusionOutput(output=(video, audio)) @@ -305,6 +306,48 @@ def test_batch_runner_output_round_trips_nested_results_through_shm() -> None: assert output["req-error"].result.error == "boom" +def test_per_request_views_of_a_batch_tensor_do_not_ship_whole_storage() -> None: + """Regression: per-request slices used to pickle the whole batch storage. + + Each view is under the packing threshold, but pickle writes the entire + shared storage per view. Packing must key off storage nbytes so those + views become SHM handles instead of going on the wire. + """ + num_requests = 16 + # Each row stays under the threshold; the shared storage is well over it. + row_numel = _SHM_TENSOR_THRESHOLD // (2 * torch.empty((), dtype=torch.float32).element_size()) + batch = torch.arange(num_requests * row_numel, dtype=torch.float32).reshape(num_requests, row_numel) + per_request = batch[0] + assert per_request.nelement() * per_request.element_size() <= _SHM_TENSOR_THRESHOLD + assert per_request.untyped_storage().nbytes() > _SHM_TENSOR_THRESHOLD + + packed_view = _pack_value_if_large(per_request) + try: + assert isinstance(packed_view, dict) + assert packed_view["__tensor_shm__"] is True + finally: + _cleanup_shm_handle(packed_view) + + output = BatchRunnerOutput.from_list( + [ + RunnerOutput( + request_id=f"req-{i}", + finished=True, + result=DiffusionOutput(trajectory_latents=batch[i]), + ) + for i in range(num_requests) + ] + ) + + pack_diffusion_output_shm(output) + + assert output.runner_outputs[0].result.trajectory_latents["__tensor_shm__"] is True + + unpack_diffusion_output_shm(output) + for i in range(num_requests): + torch.testing.assert_close(output[f"req-{i}"].result.trajectory_latents, batch[i]) + + def test_pack_value_keeps_tensor_at_threshold_inline() -> None: tensor = torch.arange( _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(), @@ -337,31 +380,37 @@ def test_pack_value_recurses_nested_dicts_and_lists_without_mutating_inline_valu large = torch.arange(_large_numel(torch.float32), dtype=torch.float32) small = torch.arange(8, dtype=torch.float32) list_tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + original_list = [list_tensor] payload = { "media": { "large": large, "small": small, }, - "list_value": [list_tensor], + "list_value": original_list, "metadata": {"prompt": "keep inline"}, } packed = _pack_value_if_large(payload) try: + assert isinstance(packed, dict) assert packed is not payload - assert packed["media"] is not payload["media"] - assert packed["media"]["large"]["__tensor_shm__"] is True - assert packed["media"]["small"] is small + media = packed["media"] + list_value = packed["list_value"] + assert isinstance(media, dict) + assert isinstance(list_value, list) + assert media is not payload["media"] + assert media["large"]["__tensor_shm__"] is True + assert media["small"] is small # Lists are recursed too: the large tensor inside is packed and a new # list is returned, while the input payload is left untouched. - assert packed["list_value"] is not payload["list_value"] - assert packed["list_value"][0]["__tensor_shm__"] is True - assert payload["list_value"][0] is list_tensor + assert list_value is not original_list + assert list_value[0]["__tensor_shm__"] is True + assert original_list[0] is list_tensor assert packed["metadata"] == {"prompt": "keep inline"} - torch.testing.assert_close(_unpack_if_shm_handle(packed["media"]["large"]), large) - torch.testing.assert_close(_unpack_if_shm_handle(packed["list_value"][0]), list_tensor) + torch.testing.assert_close(_unpack_if_shm_handle(media["large"]), large) + torch.testing.assert_close(_unpack_if_shm_handle(list_value[0]), list_tensor) finally: if isinstance(packed, dict): _cleanup_shm_handle(packed.get("media", {}).get("large")) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 665196ad9..814450d92 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import contextmanager from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -10,6 +11,7 @@ import torch import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request @@ -192,6 +194,11 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) + runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend( + vllm_config=runner.vllm_config, + od_config=runner.od_config, + device=runner.device, + ) runner.kv_transfer_manager = SimpleNamespace( receive_kv_cache=lambda req, target_device=None: None, receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None, @@ -512,6 +519,102 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin assert runner.pipeline.last_req.num_reqs == 1 +@pytest.mark.core_model +@pytest.mark.cpu +def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name="none") + request = _make_request() + runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission")) + record_names = [] + original_execute = runner._execute_request_list + + def execute(*args, **kwargs): + record_names.append(kwargs["record_name"]) + return original_execute(*args, **kwargs) + + runner._execute_request_list = execute + forward_context_calls = [] + + @contextmanager + def record_forward_context(*args, **kwargs): + forward_context_calls.append((args, kwargs)) + yield + + monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context) + reset_peak_memory_stats = Mock() + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "reset_peak_memory_stats", + reset_peak_memory_stats, + ) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + synchronize = Mock() + monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize) + + runner.profile_run([request]) + + assert runner.pipeline.forward_calls == 1 + assert record_names == ["pipeline_memory_profile"] + assert len(forward_context_calls) == 1 + assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True + runner._validate_diffusion_kv_metadata.assert_not_called() + reset_peak_memory_stats.assert_not_called() + synchronize.assert_called_once_with() + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name=None) + runner.pipeline = _FinalOnlyStepPipeline() + runner.od_config.step_execution = True + requests = [_make_request(), _make_request()] + requests[0].request_id = "profile-0" + requests[1].request_id = "profile-1" + for request in requests: + request.sampling_params.num_inference_steps = 1 + + observed_batch_rows = [] + original_denoise_step = runner.pipeline.denoise_step + + def denoise_step(input_batch, states): + observed_batch_rows.append((len(states), int(input_batch.latents.shape[0]))) + return original_denoise_step(input_batch, states) + + runner.pipeline.denoise_step = denoise_step + runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission")) + forward_context_calls = [] + + @contextmanager + def record_forward_context(*args, **kwargs): + forward_context_calls.append((args, kwargs)) + yield + + monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context) + reset_peak_memory_stats = Mock() + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "reset_peak_memory_stats", + reset_peak_memory_stats, + ) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + synchronize = Mock() + monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize) + + runner.profile_run(requests) + + assert observed_batch_rows == [(2, 2)] + assert runner.state_cache == {} + assert runner.input_batch is None + assert len(forward_context_calls) == 1 + assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True + runner._validate_diffusion_kv_metadata.assert_not_called() + reset_peak_memory_stats.assert_not_called() + synchronize.assert_called_once_with() + + class _BatchPipeline: """Pipeline returning a configurable list of outputs from forward().""" @@ -779,7 +882,11 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) - monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None) + monkeypatch.setattr( + model_runner_module, + "enable_offload_backend", + lambda od_config, pipeline, device: (pipeline, None), + ) monkeypatch.setattr( model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend ) @@ -791,8 +898,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): assert dummy_cache_backend.enabled is False -@pytest.mark.core_model -@pytest.mark.cpu def test_set_forward_context_enters_vllm_config_contexts(monkeypatch): """Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions.""" import vllm.config.vllm as vllm_config_module diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py index dc21351ff..4578d8368 100644 --- a/tests/diffusion/test_diffusion_output_formatter.py +++ b/tests/diffusion/test_diffusion_output_formatter.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -76,7 +76,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata( result = results[0] assert result.images == ["frame-0"] assert result.prompt == "prompt-0" - assert result.final_output_type == "image" + assert result.final_output_type == "video" assert result.multimodal_output == { "metadata": { "audio": {"sample_rate": 48000}, @@ -315,7 +315,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions( assert result.images == ["frame-0"] assert result.prompt == "watch and listen" - assert result.final_output_type == "image" + assert result.final_output_type == "video" assert result.multimodal_output == { "metadata": { "audio": {"sample_rate": 16000}, @@ -328,6 +328,25 @@ def test_formatter_preserves_audio_model_video_audio_and_actions( } +def test_formatter_uses_native_model_identity_for_diffusers_video_frames( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + frames = [["frame-0", "frame-1"]] + postprocess_output = normalize_diffusion_postprocess_output(frames) + + [result] = format_diffusion_outputs( + request=_request("make a video"), + od_config=_config("WanImageToVideoPipeline"), + diffusion_output=DiffusionOutput(output=frames), + output_data=frames, + postprocess_output=postprocess_output, + ) + + assert result.images == frames + assert result.final_output_type == "video" + + def test_formatter_preserves_audio_only_postprocess_dict( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/diffusion/test_diffusion_output_metadata.py b/tests/diffusion/test_diffusion_output_metadata.py index ee3559f47..b1228cebd 100644 --- a/tests/diffusion/test_diffusion_output_metadata.py +++ b/tests/diffusion/test_diffusion_output_metadata.py @@ -25,7 +25,6 @@ def test_validate_accepts_canonical_metadata_fields() -> None: }, "audio": { "sample_rate": 48000, - "audiox_task": "tts", }, "actions": { "raw_action_dim": 7, @@ -55,7 +54,6 @@ def test_validate_accepts_canonical_metadata_fields() -> None: ({"video": {"video_fps_multiplier": 0}}, ValueError), ({"video": {"shape": [1, 0, 64]}}, ValueError), ({"audio": {"sample_rate": 0}}, ValueError), - ({"audio": {"audiox_task": ""}}, ValueError), ({"actions": {"raw_action_dim": 0}}, ValueError), ({"actions": {"action_mode": ""}}, ValueError), ({"actions": {"domain_id": "7"}}, TypeError), diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py index a437a572a..65ad9ccc7 100644 --- a/tests/diffusion/test_diffusion_plugin_hooks.py +++ b/tests/diffusion/test_diffusion_plugin_hooks.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for diffusion engine plugin extensibility hooks. @@ -20,6 +20,8 @@ from vllm_omni.diffusion.registry import ( _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS, + get_diffusion_post_process_func, + get_diffusion_pre_process_func, register_diffusion_model, ) from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum @@ -40,6 +42,12 @@ class TestPlatformDiffusionHooks: result = OmniPlatform.get_diffusion_model_runner_cls() assert result == "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner" + def test_get_diffusion_kv_block_tables_cls_default(self): + """Test the default paged-KV BlockTables implementation.""" + from vllm.v1.worker.gpu.block_table import BlockTables + + assert OmniPlatform.get_diffusion_kv_block_tables_cls() is BlockTables + def test_oot_enum_exists(self): """Test OOT is a valid platform enum value.""" assert OmniPlatformEnum.OOT.value == "oot" @@ -106,6 +114,49 @@ class TestRegisterDiffusionModel: assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process" +@pytest.mark.parametrize( + ("model_class_name", "get_process_func"), + [ + ("QwenImagePipeline", get_diffusion_post_process_func), + ("WanImageToVideoPipeline", get_diffusion_post_process_func), + ("WanImageToVideoPipeline", get_diffusion_pre_process_func), + ], +) +def test_diffusers_backend_skips_native_process_hooks(model_class_name, get_process_func): + od_config = SimpleNamespace( + model_class_name=model_class_name, + custom_pipeline_args=None, + diffusion_load_format="diffusers", + ) + + with patch("vllm_omni.diffusion.registry._load_process_func") as load_process_func: + assert get_process_func(od_config) is None + + load_process_func.assert_not_called() + + +@pytest.mark.parametrize( + ("model_class_name", "get_process_func"), + [ + ("QwenImagePipeline", get_diffusion_post_process_func), + ("WanImageToVideoPipeline", get_diffusion_post_process_func), + ("WanImageToVideoPipeline", get_diffusion_pre_process_func), + ], +) +def test_custom_pipeline_preserves_native_process_hooks(model_class_name, get_process_func): + process_func = Mock() + od_config = SimpleNamespace( + model_class_name=model_class_name, + custom_pipeline_args={"pipeline_class": "custom.Pipeline"}, + diffusion_load_format="diffusers", + ) + + with patch("vllm_omni.diffusion.registry._load_process_func", return_value=process_func) as load_process_func: + assert get_process_func(od_config) is process_func + + load_process_func.assert_called_once() + + class TestWorkerUsesHook: """Test that DiffusionWorker resolves model runner via platform hook.""" diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py index 259492c1e..7476ec9c1 100644 --- a/tests/diffusion/test_diffusion_request.py +++ b/tests/diffusion/test_diffusion_request.py @@ -52,6 +52,7 @@ def test_dummy_run_request_is_identified_by_reserved_request_id(): assert req.is_dummy_run() assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID) + assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1") def test_non_dummy_request_is_not_identified_as_dummy_run(): diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index ce9f0ebe1..ae9b86f80 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -8,12 +8,14 @@ from types import SimpleNamespace import pytest import torch +import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers from pytest_mock import MockerFixture +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode -from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import ( BaseScheduler, @@ -23,7 +25,7 @@ from vllm_omni.diffusion.sched import ( StepScheduler, ) from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData -from vllm_omni.diffusion.worker.utils import RunnerOutput +from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -87,6 +89,52 @@ def _cached_ids(sched_output) -> list[str]: return list(sched_output.scheduled_cached_reqs.request_ids) +def _initialize_paged_scheduler( + scheduler: BaseScheduler, + *, + num_blocks: int = 64, + max_num_seqs: int = 1, +) -> None: + native_kv_managers.register_all_kvcache_specs(None) + spec = FullAttentionSpec( + block_size=4, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + ) + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], + ) + scheduler.initialize( + SimpleNamespace( + diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, + max_model_len=64, + max_num_seqs=max_num_seqs, + ), + kv_cache_config=config, + scheduler_block_size=4, + hash_block_size=4, + kv_vllm_config=SimpleNamespace( + model_config=SimpleNamespace(max_model_len=64), + max_in_flight_tokens=64, + ), + ) + + +def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None: + request.diffusion_kv_requests = ( + DiffusionKVRequest( + f"{request.request_id}/diffusion-kv/0", + sequence_id=0, + prefix_len=4, + target_len=4, + seq_len=seq_len, + ), + ) + + class _StubScheduler: def __init__(self, request: OmniDiffusionRequest, output) -> None: self._request = request @@ -217,13 +265,21 @@ class TestGetRequestBatchSamplingParamsKey: num_inference_steps: int = 2, seed: int | None = 123, generator: torch.Generator | None = None, + extra_args: dict | None = None, + condition_key: tuple | None = None, ) -> OmniDiffusionRequest: sp = OmniDiffusionSamplingParams( num_inference_steps=num_inference_steps, seed=seed, generator=generator, + extra_args=extra_args or {}, + ) + return OmniDiffusionRequest( + prompt="prompt", + sampling_params=sp, + request_id=f"req-{num_inference_steps}", + batch_compatibility_key=condition_key, ) - return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}") def test_distinguishes_num_inference_steps(self) -> None: scheduler = RequestScheduler() @@ -240,6 +296,57 @@ class TestGetRequestBatchSamplingParamsKey: self._make(seed=1, generator=gen_a) ) == scheduler._build_sampling_params_key(self._make(seed=2, generator=gen_b)) + @pytest.mark.parametrize( + ("first_extra_args", "second_extra_args"), + [ + ({"sample_solver": "unipc"}, {"sample_solver": "euler"}), + ({"flow_shift": 3.0}, {"flow_shift": 5.0}), + ], + ) + def test_distinguishes_wan_scheduler_structure( + self, + first_extra_args: dict, + second_extra_args: dict, + ) -> None: + scheduler = RequestScheduler() + + assert scheduler._build_sampling_params_key( + self._make(extra_args=first_extra_args) + ) != scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args)) + + @pytest.mark.parametrize( + ("first_extra_args", "second_extra_args"), + [ + ({"sample_solver": " Euler "}, {"sample_solver": "euler"}), + ({"flow_shift": "5.0"}, {"flow_shift": 5.0}), + ], + ) + def test_normalizes_equivalent_wan_scheduler_structure( + self, + first_extra_args: dict, + second_extra_args: dict, + ) -> None: + scheduler = RequestScheduler() + + assert scheduler._build_sampling_params_key( + self._make(extra_args=first_extra_args) + ) == scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args)) + + def test_uses_none_for_unspecified_wan_scheduler_structure(self) -> None: + scheduler = RequestScheduler() + + key = scheduler._build_sampling_params_key(self._make()) + + assert key.sample_solver is None + assert key.flow_shift is None + + def test_distinguishes_pipeline_condition_structure(self) -> None: + scheduler = RequestScheduler() + + assert scheduler._build_sampling_params_key( + self._make(condition_key=("wan22_s2v_condition", True)) + ) != scheduler._build_sampling_params_key(self._make(condition_key=("wan22_s2v_condition", False))) + class TestRequestScheduler: def setup_method(self) -> None: @@ -341,18 +448,16 @@ class TestRequestScheduler: assert self.scheduler.has_requests() is False def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None: - self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + _initialize_paged_scheduler(self.scheduler) request = _make_request("diffusion-kv") prepared_layout = object() request.prepared_layout = prepared_layout - text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8) kv_request = DiffusionKVRequest( "diffusion-kv/diffusion-kv/0", sequence_id=0, prefix_len=4, target_len=8, seq_len=16, - kv_contexts=(text_context,), ) request.diffusion_kv_requests = (kv_request,) @@ -362,12 +467,179 @@ class TestRequestScheduler: assert state is not None assert state.diffusion_kv_requests == (kv_request,) - assert state.diffusion_kv_requests[0].kv_contexts == (text_context,) assert request.diffusion_kv_requests is None assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout + metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata + assert metadata is not None + assert metadata.request_id == request.request_id + assert len(metadata.sequences[0].block_ids[0]) == 4 + + def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2) + first_request = _make_request("first") + second_request = _make_request("second") + _attach_diffusion_kv(first_request) + _attach_diffusion_kv(second_request) + self.scheduler.add_request(first_request) + self.scheduler.add_request(second_request) + + first_output = self.scheduler.schedule() + + assert _new_ids(first_output) == ["first"] + assert first_output.num_waiting_reqs == 1 + + self.scheduler.update_from_output(first_output, _make_request_output("first")) + second_output = self.scheduler.schedule() + + assert _new_ids(second_output) == ["second"] + assert second_output.num_waiting_reqs == 0 + # The same output carries both the release and the newly admitted block + # table; the future Worker data plane must process finished ids first. + assert second_output.finished_req_ids == {"first"} + assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None + + def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2) + impossible = _make_request("impossible") + schedulable = _make_request("schedulable") + _attach_diffusion_kv(impossible, seq_len=8) + schedulable.diffusion_kv_requests = ( + DiffusionKVRequest( + "schedulable/diffusion-kv/0", + sequence_id=0, + prefix_len=0, + target_len=4, + seq_len=4, + ), + ) + self.scheduler.add_request(impossible) + self.scheduler.add_request(schedulable) + + sched_output = self.scheduler.schedule() + + failed_state = self.scheduler.get_request_state("impossible") + assert failed_state is not None + assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR + assert failed_state.error is not None + assert "cannot fit even when the block pool is empty" in failed_state.error + assert sched_output.finished_req_ids == {"impossible"} + assert _new_ids(sched_output) == ["schedulable"] + + finished = self.scheduler.update_from_output( + sched_output, + _make_request_output("schedulable"), + ) + + assert finished == {"impossible", "schedulable"} + + def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2) + running = _make_request("running") + _attach_diffusion_kv(running) + self.scheduler.add_request(running) + assert _new_ids(self.scheduler.schedule()) == ["running"] + + impossible = _make_request("impossible") + impossible.diffusion_kv_requests = tuple( + DiffusionKVRequest( + f"impossible/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + prefix_len=0, + target_len=4, + seq_len=4, + ) + for sequence_id in range(4) + ) + schedulable = _make_request("schedulable") + schedulable.diffusion_kv_requests = ( + DiffusionKVRequest( + "schedulable/diffusion-kv/0", + sequence_id=0, + prefix_len=0, + target_len=4, + seq_len=4, + ), + ) + self.scheduler.add_request(impossible) + self.scheduler.add_request(schedulable) + + sched_output = self.scheduler.schedule() + + failed_state = self.scheduler.get_request_state("impossible") + assert failed_state is not None + assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR + assert failed_state.error is not None + assert "required_blocks=4, available_blocks=3" in failed_state.error + assert sched_output.finished_req_ids == {"impossible"} + assert _new_ids(sched_output) == ["schedulable"] + + def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None: + _initialize_paged_scheduler(self.scheduler) + request = _make_request("native-error") + _attach_diffusion_kv(request) + self.scheduler.add_request(request) + manager = self.scheduler._diffusion_kv_manager + assert manager is not None + + def raise_native_error(*args, **kwargs): + raise ValueError("native allocation bug") + + monkeypatch.setattr( + manager, + "reserve_request", + raise_native_error, + ) + + sched_output = self.scheduler.schedule() + + state = self.scheduler.get_request_state("native-error") + assert state is not None + assert state.status == DiffusionRequestStatus.FINISHED_ERROR + assert state.error == "native allocation bug" + assert sched_output.finished_req_ids == {"native-error"} + assert sched_output.scheduled_request_ids == [] + + def test_diffusion_kv_preemption_retains_allocation(self) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=3) + request = _make_request("preempted") + _attach_diffusion_kv(request) + self.scheduler.add_request(request) + self.scheduler.schedule() + manager = self.scheduler._diffusion_kv_manager + assert manager is not None + metadata = manager.get_metadata("preempted") + + assert self.scheduler.preempt_request("preempted") is True + resumed_output = self.scheduler.schedule() + + assert resumed_output.scheduled_new_reqs == [] + assert _cached_ids(resumed_output) == ["preempted"] + assert manager.get_metadata("preempted") == metadata + + @pytest.mark.parametrize( + "status", + [ + DiffusionRequestStatus.FINISHED_COMPLETED, + DiffusionRequestStatus.FINISHED_ABORTED, + DiffusionRequestStatus.FINISHED_ERROR, + ], + ) + def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None: + _initialize_paged_scheduler(self.scheduler, num_blocks=3) + request = _make_request("terminal") + _attach_diffusion_kv(request) + self.scheduler.add_request(request) + self.scheduler.schedule() + manager = self.scheduler._diffusion_kv_manager + assert manager is not None + assert manager.has_request("terminal") is True + + self.scheduler.finish_requests("terminal", status) + + assert manager.has_request("terminal") is False def test_paged_scheduler_rejects_missing_kv_request(self) -> None: - self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + _initialize_paged_scheduler(self.scheduler) request = _make_request("diffusion-kv") with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"): @@ -384,7 +656,7 @@ class TestRequestScheduler: ], ) def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None: - self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + _initialize_paged_scheduler(self.scheduler) request = _make_request("diffusion-kv") owner = request if owner_name == "request" else request.sampling_params setattr(owner, field_name, object()) @@ -501,6 +773,19 @@ class TestRequestScheduler: assert third.num_running_reqs == 1 assert third.num_waiting_reqs == 0 + def test_records_initial_scheduler_queue_wait(self, mocker: MockerFixture) -> None: + perf_counter = mocker.patch( + "vllm_omni.diffusion.sched.base_scheduler.time.perf_counter", + side_effect=[10.0, 10.125], + ) + request = _make_step_request("queue-wait", num_inference_steps=2) + + self.scheduler.add_request(request) + self.scheduler.schedule() + + assert request.scheduler_queue_wait_ms == pytest.approx(125.0) + assert perf_counter.call_count == 2 + def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=2)) @@ -602,6 +887,38 @@ class TestRequestScheduler: assert first.num_running_reqs == 1 assert first.num_waiting_reqs == 1 + def test_batches_incompatible_pipeline_conditions_separately(self) -> None: + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + + request_a = OmniDiffusionRequest( + prompt="a", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=123), + request_id="a", + batch_compatibility_key=("wan22_s2v_condition", True), + ) + request_b = OmniDiffusionRequest( + prompt="b", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=456), + request_id="b", + batch_compatibility_key=("wan22_s2v_condition", False), + ) + scheduler.add_request(request_a) + scheduler.add_request(request_b) + + first = scheduler.schedule() + + assert _new_ids(first) == ["a"] + assert first.num_running_reqs == 1 + assert first.num_waiting_reqs == 1 + + scheduler.update_from_output(first, _make_request_output("a")) + second = scheduler.schedule() + + assert _new_ids(second) == ["b"] + assert second.num_running_reqs == 1 + assert second.num_waiting_reqs == 0 + def test_batches_different_request_local_seed_together(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=2)) @@ -625,6 +942,47 @@ class TestRequestScheduler: assert first.num_running_reqs == 2 assert first.num_waiting_reqs == 0 + @pytest.mark.parametrize( + ("first_extra_args", "second_extra_args"), + [ + ({"sample_solver": "unipc"}, {"sample_solver": "euler"}), + ({"flow_shift": 3.0}, {"flow_shift": 5.0}), + ({"sample_solver": "unipc"}, {}), + ({"flow_shift": 3.0}, {}), + ], + ) + def test_batches_incompatible_wan_scheduler_structure_separately( + self, + first_extra_args: dict, + second_extra_args: dict, + ) -> None: + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + first_id = scheduler.add_request( + _make_step_request( + "a", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args=first_extra_args, + ), + ) + ) + scheduler.add_request( + _make_step_request( + "b", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args=second_extra_args, + ), + ) + ) + + first = scheduler.schedule() + + assert _new_ids(first) == [first_id] + assert first.num_running_reqs == 1 + assert first.num_waiting_reqs == 1 + def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=3)) @@ -789,6 +1147,45 @@ class TestDiffusionEngine: assert engine.scheduler.max_num_running_reqs == 1 fake_executor_cls.assert_called_once_with(od_config) + def test_initializes_paged_scheduler_from_native_control_plane( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False) + fake_executor_cls = mocker.Mock(return_value=mocker.Mock()) + scheduler = mocker.Mock() + kv_cache_config = object() + kv_vllm_config = object() + + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", + lambda *args, **kwargs: fake_executor_cls, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane", + lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config), + ) + + engine = DiffusionEngine(od_config, scheduler=scheduler) + + assert engine.scheduler is scheduler + scheduler.initialize.assert_called_once_with( + od_config, + kv_cache_config=kv_cache_config, + scheduler_block_size=16, + hash_block_size=16, + kv_vllm_config=kv_vllm_config, + ) + def test_scheduler_alias_keeps_default_request_scheduler(self) -> None: scheduler = Scheduler() scheduler.initialize(SimpleNamespace()) @@ -855,6 +1252,21 @@ class TestDiffusionEngine: assert output.aborted is True assert output.abort_message == "Request req-finalize aborted." + def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.scheduler = StepScheduler() + engine.scheduler.initialize(SimpleNamespace()) + + req_id = engine.scheduler.add_request(_make_request("req-error")) + engine.scheduler._finish_requests( + {req_id: DiffusionRequestStatus.FINISHED_ERROR}, + {req_id: "KV request cannot fit"}, + ) + + output = engine._finalize_finished_request(req_id) + + assert output.error == "KV request cannot fit" + @pytest.mark.asyncio async def test_streaming_runner_output_notifies_each_chunk(self) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) @@ -1160,6 +1572,39 @@ class TestStepScheduler: assert sched_output.num_running_reqs == 2 assert sched_output.num_waiting_reqs == 0 + def test_mixed_batch_error_finishes_only_failed_request(self) -> None: + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + + req_bad = scheduler.add_request(_make_step_request("bad", num_inference_steps=2)) + req_good = scheduler.add_request(_make_step_request("good", num_inference_steps=2)) + first = scheduler.schedule() + + finished = scheduler.update_from_output( + first, + BatchRunnerOutput.from_list( + [ + _make_step_output( + req_bad, + step_index=0, + finished=True, + error="invalid request input", + ), + _make_step_output(req_good, step_index=1), + ] + ), + ) + + assert finished == {req_bad} + assert scheduler.get_request_state(req_bad).status == DiffusionRequestStatus.FINISHED_ERROR + assert scheduler.get_request_state(req_good).status == DiffusionRequestStatus.RUNNING + assert scheduler.get_request_state(req_good).req.sampling_params.step_index == 1 + + second = scheduler.schedule() + assert second.finished_req_ids == {req_bad} + assert _new_ids(second) == [] + assert _cached_ids(second) == [req_good] + def test_step_batch_rejects_different_sampling_key(self) -> None: scheduler = StepScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=3)) diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 0964833d7..94c3eb48f 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for step-level diffusion execution across runner / worker / executor / engine.""" import contextlib @@ -17,6 +17,8 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D from vllm_omni.diffusion.distributed.parallel_state import ( @@ -45,6 +47,7 @@ from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -100,6 +103,22 @@ class _StepPipeline: return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32)) +class _PerRequestErrorStepPipeline(_StepPipeline): + def prepare_encode(self, state, **kwargs): + if state.prompt == "fail-prepare": + raise OmniClientError( + "invalid request input", + status_code=422, + error_type="UnprocessableEntityError", + ) + return super().prepare_encode(state, **kwargs) + + def denoise_step(self, input_batch, **kwargs): + del kwargs + self.denoise_calls += 1 + return torch.ones_like(input_batch.latents) + + class _ProfilingStepPipeline(_StepPipeline): enable_diffusion_pipeline_profiler = True @@ -306,11 +325,14 @@ def _make_vllm_config(): ) -def _make_runner(): +def _make_runner( + diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY, +): runner = object.__new__(DiffusionModelRunner) runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( cache_backend=None, + diffusion_kv_mode=diffusion_kv_mode, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) @@ -330,6 +352,7 @@ def _make_distributed_runner(mode: str, device: torch.device): runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( cache_backend=None, + diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) @@ -502,6 +525,81 @@ def test_step_profiler_reports_denoise_step_as_diffuse(): class TestRunner: """DiffusionModelRunner.execute_stepwise""" + @pytest.mark.parametrize( + ("diffusion_kv_mode", "should_cleanup"), + [ + (DiffusionKVCacheMode.DENSE_LEGACY, False), + (DiffusionKVCacheMode.PAGED_SCHEDULER, True), + ], + ) + def test_scheduler_finished_cleanup_only_runs_for_paged_kv( + self, + mocker: MockerFixture, + diffusion_kv_mode: DiffusionKVCacheMode, + should_cleanup: bool, + ): + runner = _make_runner(diffusion_kv_mode) + cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests") + expected_output = object() + mocker.patch.object(runner, "_execute_stepwise_core", return_value=expected_output) + scheduler_output = DiffusionSchedulerOutput( + step_id=1, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids={"finished-request"}, + num_running_reqs=0, + num_waiting_reqs=0, + ) + + output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + + assert output is expected_output + if should_cleanup: + cleanup.assert_called_once_with(["finished-request"]) + else: + cleanup.assert_not_called() + + @pytest.mark.parametrize( + ("diffusion_kv_mode", "should_cleanup"), + [ + (DiffusionKVCacheMode.DENSE_LEGACY, False), + (DiffusionKVCacheMode.PAGED_SCHEDULER, True), + ], + ) + def test_terminal_cleanup_only_runs_for_paged_kv( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + diffusion_kv_mode: DiffusionKVCacheMode, + should_cleanup: bool, + ): + runner = _make_runner(diffusion_kv_mode) + cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests") + install = mocker.patch.object(runner, "install_diffusion_kv_metadata") + req = _make_step_request() + scheduler_output = _make_scheduler_output(req, step_id=0) + if should_cleanup: + scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata = DiffusionKVMetadata( + request_id=req.request_id, + allocation_generation=1, + sequences=(), + ) + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + first_output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert first_output.get_request_output(req.request_id).finished is False + cleanup.assert_not_called() + + output = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1)) + + assert output.get_request_output(req.request_id).finished is True + if should_cleanup: + install.assert_called_once_with(scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata) + cleanup.assert_called_once_with([req.request_id]) + else: + install.assert_not_called() + cleanup.assert_not_called() + def test_completes_request_and_clears_state(self, monkeypatch): runner = _make_runner() req = _make_step_request() @@ -616,6 +714,72 @@ class TestRunner: result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) assert len(result) == 2 + def test_prepare_error_isolated_from_healthy_request(self, monkeypatch): + runner = _make_runner() + runner.pipeline = _PerRequestErrorStepPipeline() + bad_req = _make_step_request() + bad_req.request_id = "req-bad" + bad_req.prompt = "fail-prepare" + good_req = _make_step_request() + good_req.request_id = "req-good" + good_req.prompt = "valid" + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + result = DiffusionModelRunner.execute_stepwise( + runner, + _make_batch_scheduler_output([bad_req, good_req]), + ) + + bad_output = result.get_request_output("req-bad") + assert bad_output.finished is True + assert bad_output.result is not None + assert bad_output.result.error == "invalid request input" + assert bad_output.result.error_status_code == 422 + assert bad_output.result.error_type == "UnprocessableEntityError" + good_output = result.get_request_output("req-good") + assert good_output.finished is False + assert good_output.step_index == 1 + assert good_output.result is None + assert set(runner.state_cache) == {"req-good"} + assert runner.pipeline.denoise_calls == 1 + assert runner.pipeline.scheduler_calls == 1 + + def test_prepare_error_on_other_rank_is_isolated_via_all_reduce(self, monkeypatch): + """P1-B: rank-0-only pipeline work (e.g. H3 reference-video prep) + can raise on rank 0 while non-zero ranks return cleanly and then hang + on a downstream ``dist.broadcast``. Beyond the pipeline-side guard, + the runner all-reduces a per-request failure flag across the DiT + group so a rank that ``prepare_encode`` returned cleanly on still + skips the request and does not enter the denoise step batch. + """ + runner = _make_runner() + runner.pipeline = _StepPipeline() # locally always succeeds + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + # Simulate a peer DiT rank having reported failure. We patch the + # helper directly (rather than ``torch.distributed`` broadly) because + # ``execute_stepwise`` also queries ``torch.distributed.get_rank``. + monkeypatch.setattr( + model_runner_module, + "_dit_any_rank_failed", + lambda local_failed: True, + ) + + req = _make_step_request(num_inference_steps=1) + result = DiffusionModelRunner.execute_stepwise( + runner, + _make_scheduler_output(req), + ) + + output = result.get_request_output("req-1") + assert output.finished is True + assert output.result is not None + assert output.result.error is not None + assert "another DiT rank" in output.result.error + # A rank that skipped the request must not have driven a denoise step. + assert runner.pipeline.denoise_calls == 0 + assert runner.state_cache == {} + def test_receives_kv_payload_before_prepare_encode(self, monkeypatch): runner = _make_runner() captured: dict[str, object] = {} @@ -715,7 +879,11 @@ class TestRunner: monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler) - monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None) + monkeypatch.setattr( + model_runner_module, + "enable_offload_backend", + lambda _config, pipeline, **_kwargs: (pipeline, None), + ) monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) with pytest.raises(ValueError, match="RequestOnlyPipeline"): diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py index d84f49b31..c833ae7a1 100644 --- a/tests/diffusion/test_diffusion_streaming_output.py +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Integration tests for diffusion pipeline streaming output mode. Covers: @@ -311,8 +311,8 @@ class TestPipelineStreamingOutputToEntrypoint: lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(), ) mocker.patch( - "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", - return_value="diffusion", + "vllm_omni.entrypoints.openai.serving_video_output_stream.is_video_generation_pipeline", + return_value=True, ) mocker.patch( "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", @@ -441,7 +441,6 @@ class TestPipelineStreamingOutputToEntrypoint: model="test_model", od_config=od_config, metadata=metadata, - batch_size=1, ) @staticmethod @@ -526,6 +525,7 @@ class TestSupportedPipelines: runner = object.__new__(DiffusionModelRunner) runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( + max_num_seqs=1, enable_cpu_offload=False, enable_layerwise_offload=False, enforce_eager=True, @@ -544,7 +544,11 @@ class TestSupportedPipelines: monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler) - monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None) + monkeypatch.setattr( + model_runner_module, + "enable_offload_backend", + lambda _config, pipeline, **_kwargs: (pipeline, None), + ) monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) with pytest.raises(ValueError, match="NoStepPipeline"): diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 6a49d0909..0d8d5e92a 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -13,11 +13,7 @@ import pytest import torch from pytest_mock import MockerFixture -from vllm_omni.diffusion.worker.diffusion_worker import ( - DiffusionWorker, - _create_diffusion_worker_vllm_config, - _make_diffusion_vllm_model_config, -) +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] @@ -58,34 +54,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config): return worker -def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods(): - from types import SimpleNamespace - - from vllm_omni.quantization import build_quant_config - - od_config = SimpleNamespace( - model="dummy", - dtype=torch.bfloat16, - quantization_config=build_quant_config( - { - "quant_method": "modelopt", - "quant_algo": "FP8", - "ignore": [], - } - ), - tf_model_config=SimpleNamespace(), - enforce_eager=True, - is_moe=False, - ) - - model_config = _make_diffusion_vllm_model_config(od_config) - - assert model_config.dtype is torch.bfloat16 - assert model_config.quantization == "modelopt" - assert model_config.quantization_config is od_config.quantization_config - assert model_config.is_quantized() - - class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" @@ -329,49 +297,3 @@ class TestDiffusionWorkerWakeUp: mock_buffer2.data.copy_.assert_not_called() assert result is True - - -class TestWorkerVllmConfigAdditionalConfig: - """Test worker-side VllmConfig construction with additional_config.""" - - def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture): - mock_vllm_config = mocker.Mock() - mock_vllm_cls = mocker.patch( - "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", - return_value=mock_vllm_config, - ) - od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}}) - - result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) - - assert result is mock_vllm_config - assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config - - def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config( - self, mocker: MockerFixture - ): - mock_vllm_config = mocker.Mock() - mock_vllm_cls = mocker.patch( - "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", - side_effect=[ - TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"), - mock_vllm_config, - ], - ) - od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}}) - - result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) - - assert result is mock_vllm_config - assert mock_vllm_cls.call_count == 2 - assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config - - def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture): - mocker.patch( - "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig", - side_effect=TypeError("additional_config values must be hashable"), - ) - od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}}) - - with pytest.raises(TypeError, match="hashable"): - _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config) diff --git a/tests/diffusion/test_diffusion_worker_broadcast.py b/tests/diffusion/test_diffusion_worker_broadcast.py new file mode 100644 index 000000000..750d4c30d --- /dev/null +++ b/tests/diffusion/test_diffusion_worker_broadcast.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +import time +from unittest.mock import Mock + +import pytest +from vllm.distributed.device_communicators.shm_broadcast import MessageQueue + +import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module +from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _DelayedMultipartSocket: + def __init__( + self, + socket, + send_started: threading.Event, + release_send: threading.Event, + ): + self._socket = socket + self._send_started = send_started + self._release_send = release_send + self._delay_next_send = True + + def __getattr__(self, name): + return getattr(self._socket, name) + + def send_multipart(self, *args, **kwargs): + if self._delay_next_send: + self._delay_next_send = False + self._send_started.set() + if not self._release_send.wait(timeout=2): + raise TimeoutError("Timed out waiting to release delayed overflow payload") + return self._socket.send_multipart(*args, **kwargs) + + +class _ObservedPollSocket: + def __init__(self, socket, poll_started: threading.Event): + self._socket = socket + self._poll_started = poll_started + + def __getattr__(self, name): + return getattr(self._socket, name) + + def poll(self, *args, **kwargs): + self._poll_started.set() + return self._socket.poll(*args, **kwargs) + + +class _RecordingWorker: + def __init__(self): + self.calls: list[str] = [] + + def execute_method(self, method, *args, **kwargs): + assert method == "record" + self.calls.append(args[0]) + + +def test_recv_message_waits_indefinitely(): + worker_proc = WorkerProc.__new__(WorkerProc) + worker_proc.mq = Mock() + worker_proc.mq.dequeue.return_value = {"type": "rpc"} + + result = worker_proc.recv_message() + + assert result == {"type": "rpc"} + worker_proc.mq.dequeue.assert_called_once_with(indefinite=True) + + +def test_busy_loop_releases_failed_rpc_traceback_and_device_cache(monkeypatch): + original = RuntimeError("model execution failed") + worker_proc = WorkerProc.__new__(WorkerProc) + worker_proc.gpu_id = 0 + worker_proc.mq = Mock() + worker_proc.mq.dequeue.side_effect = [ + { + "type": "rpc", + "method": "execute_model", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": False, + "collect_rank_status": False, + }, + {"type": "shutdown"}, + ] + worker_proc.worker = Mock() + worker_proc.worker.execute_method.side_effect = original + worker_proc.result_mq = object() + worker_proc._return_result = Mock() + worker_proc.wake_event = None + worker_proc._running = True + + gc_collect = Mock() + mock_platform = Mock() + monkeypatch.setattr(diffusion_worker_module.gc, "collect", gc_collect) + monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", mock_platform) + + worker_proc._worker_busy_loop() + + worker_proc._return_result.assert_called_once_with( + {"status": "error", "error": "model execution failed", "wave_id": None} + ) + assert original.__traceback__ is None + gc_collect.assert_called_once_with() + mock_platform.empty_cache.assert_called_once_with() + + +def test_busy_loop_preserves_order_when_overflow_payload_is_delayed(monkeypatch): + """A consumed overflow marker must not be abandoned on a socket timeout.""" + writer = MessageQueue( + n_reader=1, + n_local_reader=1, + local_reader_ranks=[0], + max_chunk_bytes=1024, + max_chunks=4, + ) + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0) + writer.wait_until_ready() + reader.wait_until_ready() + + overflow_send_started = threading.Event() + release_overflow_send = threading.Event() + overflow_poll_started = threading.Event() + writer.local_socket = _DelayedMultipartSocket( + writer.local_socket, + overflow_send_started, + release_overflow_send, + ) + reader.local_socket = _ObservedPollSocket(reader.local_socket, overflow_poll_started) + + recording_worker = _RecordingWorker() + worker_proc = WorkerProc.__new__(WorkerProc) + worker_proc.gpu_id = 0 + worker_proc.mq = reader + worker_proc.worker = recording_worker + worker_proc.result_mq = None + worker_proc.wake_event = None + worker_proc._running = True + + # Make the pre-fix implementation fail quickly: it read this module-level + # value directly from _worker_busy_loop. The fixed implementation ignores it + # and delegates to recv_message(), which waits indefinitely. + monkeypatch.setattr( + diffusion_worker_module, + "_BROADCAST_DEQUEUE_TIMEOUT_S", + 0.02, + raising=False, + ) + + overflow_message = { + "type": "rpc", + "method": "record", + "args": ("overflow-1",), + "padding": b"x" * 2048, + } + inline_message = { + "type": "rpc", + "method": "record", + "args": ("inline-2",), + } + shutdown_message = {"type": "shutdown"} + thread_errors: list[BaseException] = [] + + def send_messages(): + try: + writer.enqueue(overflow_message) + writer.enqueue(inline_message) + writer.enqueue(shutdown_message) + except BaseException as exc: + thread_errors.append(exc) + + def run_busy_loop(): + try: + worker_proc._worker_busy_loop() + except BaseException as exc: + thread_errors.append(exc) + + writer_thread = threading.Thread(target=send_messages, daemon=True) + worker_thread = threading.Thread(target=run_busy_loop, daemon=True) + + try: + writer_thread.start() + worker_thread.start() + + assert overflow_send_started.wait(timeout=2) + assert overflow_poll_started.wait(timeout=2) + + # Keep the multipart unavailable beyond the pre-fix 20 ms timeout. + time.sleep(0.1) + release_overflow_send.set() + + writer_thread.join(timeout=2) + worker_thread.join(timeout=2) + + assert not writer_thread.is_alive() + assert not worker_thread.is_alive() + assert thread_errors == [] + assert recording_worker.calls == ["overflow-1", "inline-2"] + finally: + release_overflow_send.set() + worker_proc._running = False + reader.shutdown() + writer.shutdown() + writer_thread.join(timeout=1) + worker_thread.join(timeout=1) diff --git a/tests/diffusion/test_diffusion_worker_process_title.py b/tests/diffusion/test_diffusion_worker_process_title.py new file mode 100644 index 000000000..83863bb86 --- /dev/null +++ b/tests/diffusion/test_diffusion_worker_process_title.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib.util +import shutil +import subprocess +import sys +import textwrap +import time +from dataclasses import dataclass +from unittest.mock import Mock + +import pytest +from pytest_mock import MockerFixture + +import vllm_omni.diffusion.worker.diffusion_worker as worker_module + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@dataclass(frozen=True) +class _FakeGroup: + world_size: int = 1 + rank_in_group: int = 0 + + +_GROUP_GETTERS = ( + "get_dp_group", + "get_pp_group", + "get_sp_group", + "get_cfg_group", + "get_tp_group", + "get_fs_group", + "get_hsdp_replicate_group", + "get_ep_group", +) + + +def _patch_groups( + mocker: MockerFixture, + groups: dict[str, _FakeGroup], +) -> dict[str, Mock]: + patched_getters = {} + for getter_name in _GROUP_GETTERS: + patched_getters[getter_name] = mocker.patch.object( + worker_module, + getter_name, + return_value=groups.get(getter_name, _FakeGroup()), + ) + return patched_getters + + +def test_setup_uses_default_name_before_model_parallel_init( + mocker: MockerFixture, +) -> None: + mocker.patch.object( + worker_module, + "model_parallel_is_initialized", + return_value=False, + ) + group_getters = _patch_groups(mocker, {}) + set_process_title = mocker.patch.object(worker_module, "set_process_title") + decorate_logs = mocker.patch.object(worker_module, "decorate_logs") + + worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=True, use_hsdp=True) + + set_process_title.assert_called_once_with( + name="DiffusionWorker", + prefix="vLLM-Omni", + ) + decorate_logs.assert_called_once_with("DiffusionWorker") + for group_getter in group_getters.values(): + group_getter.assert_not_called() + + +@pytest.mark.parametrize( + ("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"), + [ + ({}, False, False, 1, "DiffusionWorker"), + ({"get_tp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_TP1"), + ({"get_dp_group": _FakeGroup(4, 2)}, False, False, 1, "DiffusionWorker_DP2"), + ({"get_pp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_PP1"), + ( + { + "get_cfg_group": _FakeGroup(2, 1), + "get_tp_group": _FakeGroup(2, 1), + }, + False, + False, + 1, + "DiffusionWorker_CFG1_TP1", + ), + ( + { + "get_dp_group": _FakeGroup(2, 1), + "get_sp_group": _FakeGroup(4, 2), + "get_tp_group": _FakeGroup(2, 1), + }, + False, + False, + 1, + "DiffusionWorker_DP1_SP2_TP1", + ), + ( + {"get_fs_group": _FakeGroup(4, 3)}, + False, + True, + 1, + "DiffusionWorker_FS3", + ), + ( + { + "get_sp_group": _FakeGroup(2, 1), + "get_fs_group": _FakeGroup(4, 3), + }, + False, + True, + 1, + "DiffusionWorker_SP1_FS3", + ), + ( + { + "get_fs_group": _FakeGroup(1, 0), + "get_hsdp_replicate_group": _FakeGroup(2, 1), + }, + False, + True, + 2, + "DiffusionWorker_RP1", + ), + ({"get_ep_group": _FakeGroup(4, 2)}, True, False, 1, "DiffusionWorker_EP2"), + ({}, True, False, 1, "DiffusionWorker"), + ], + ids=[ + "all-singleton", + "tp-only", + "dp-only", + "pp-only", + "tp-cfg", + "dp-sp-tp", + "hsdp-only", + "hsdp-sp", + "hsdp-replicated", + "expert-parallel", + "singleton-ep", + ], +) +def test_setup_uses_initialized_parallel_groups( + mocker: MockerFixture, + groups: dict[str, _FakeGroup], + enable_ep: bool, + use_hsdp: bool, + hsdp_replicate_size: int, + expected_name: str, +) -> None: + mocker.patch.object( + worker_module, + "model_parallel_is_initialized", + return_value=True, + ) + group_getters = _patch_groups(mocker, groups) + set_process_title = mocker.patch.object(worker_module, "set_process_title") + decorate_logs = mocker.patch.object(worker_module, "decorate_logs") + + worker_module._setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep=enable_ep, + use_hsdp=use_hsdp, + hsdp_replicate_size=hsdp_replicate_size, + ) + + set_process_title.assert_called_once_with( + name=expected_name, + prefix="vLLM-Omni", + ) + decorate_logs.assert_called_once_with(expected_name) + if use_hsdp: + group_getters["get_fs_group"].assert_called_once_with() + if hsdp_replicate_size > 1: + group_getters["get_hsdp_replicate_group"].assert_called_once_with() + else: + group_getters["get_hsdp_replicate_group"].assert_not_called() + else: + group_getters["get_fs_group"].assert_not_called() + group_getters["get_hsdp_replicate_group"].assert_not_called() + if enable_ep: + group_getters["get_ep_group"].assert_called_once_with() + else: + group_getters["get_ep_group"].assert_not_called() + + +def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None: + mocker.patch.object( + worker_module, + "model_parallel_is_initialized", + return_value=False, + ) + mocker.patch.dict(sys.modules, {"setproctitle": None}) + decorate_logs = mocker.patch.object(worker_module, "decorate_logs") + + worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False) + + decorate_logs.assert_called_once_with("DiffusionWorker") + + +@pytest.mark.skipif( + sys.platform != "linux" or shutil.which("ps") is None or importlib.util.find_spec("setproctitle") is None, + reason="requires Linux ps and setproctitle", +) +def test_process_title_is_visible_through_ps() -> None: + expected_name = "vLLM-Omni::DiffusionWorker_TP1" + child_script = textwrap.dedent( + """ + import time + + from vllm.utils.system_utils import set_process_title + + + set_process_title( + name="DiffusionWorker_TP1", + prefix="vLLM-Omni", + ) + time.sleep(30) + """ + ) + process = subprocess.Popen( + [sys.executable, "-c", child_script], + stderr=subprocess.PIPE, + text=True, + ) + + observed_title = "" + try: + for _ in range(100): + if process.poll() is not None: + stderr = process.stderr.read() if process.stderr is not None else "" + pytest.fail(f"child process exited before title check: {stderr}") + result = subprocess.run( + ["ps", "-o", "args=", "-p", str(process.pid)], + check=False, + capture_output=True, + text=True, + ) + observed_title = result.stdout.strip() + if expected_name in observed_title: + break + time.sleep(0.1) + else: + pytest.fail(f"expected {expected_name!r} in process title, got {observed_title!r}") + finally: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index 824c1483a..e39c18ff4 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -47,11 +47,16 @@ def client(mock_engine): with patch.object(InlineStageDiffusionClient, "_enrich_config"): od_config = MagicMock(spec=OmniDiffusionConfig) od_config.streaming_output = False - c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1) + c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata) yield c c.shutdown() +def test_inline_client_implements_shared_prompt_hook_protocol(client): + assert client.prompt_transform_func is None + assert client.prompt_expand_func is None + + @pytest.mark.asyncio async def test_inline_dispatch_request_success(client, mock_engine): mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()]) @@ -229,5 +234,4 @@ def test_inline_client_requires_replica_id(mock_engine): model="test_model", od_config=od_config, metadata=metadata, - batch_size=1, ) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index aa40d0770..855f38c4a 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import multiprocessing as mp @@ -14,6 +14,7 @@ import torch import zmq from vllm.v1.engine.exceptions import EngineDeadError +import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor @@ -413,6 +414,26 @@ class TestRequestModeDispatch: assert [output.result.error for output in result.runner_outputs] == ["A", "B"] executor.collective_rpc.assert_called_once() + def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self): + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")]) + scheduler_output = _make_sched_output("A", "B") + for new_req in scheduler_output.scheduled_new_reqs: + new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id) + + result = executor.execute_request(scheduler_output) + + assert [output.result.error for output in result.runner_outputs] == ["A", "B"] + forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0] + assert forwarded_envelopes is scheduler_output.scheduled_new_reqs + assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"] + @pytest.mark.parametrize( ("field", "value"), [ @@ -802,6 +823,31 @@ class TestWorkerProcRpcRankStatus: assert status["bool_result"] is None assert "local boom" in status["traceback"] + def test_execute_rpc_collect_exception_releases_traceback_and_device_cache(self, monkeypatch): + proc = self._make_worker_proc() + original = RuntimeError("local boom") + proc.worker.execute_method = Mock(side_effect=original) + gc_collect = Mock() + mock_platform = Mock() + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + monkeypatch.setattr(diffusion_worker_module.gc, "collect", gc_collect) + monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", mock_platform) + + proc._execute_rpc( + { + "method": "add_lora", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": True, + "collect_rank_status": True, + } + ) + + assert original.__traceback__ is None + gc_collect.assert_called_once_with() + mock_platform.empty_cache.assert_called_once_with() + def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self): proc = self._make_worker_proc() @@ -1033,6 +1079,42 @@ class TestStageDiffusionClientErrorPropagation: assert client.get_diffusion_output_nowait() is None + def test_error_response_preserves_scheduler_metrics(self): + from vllm_omni.metrics import definitions as metric_defs + + client = self._make_client() + client._response_socket.recv.side_effect = [b"message", zmq.Again()] + client._decoder.decode.return_value = { + "type": "error", + "request_id": "req-error", + "error": "gpu fault", + "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0}, + } + + output = client.get_diffusion_output_nowait() + + assert output is not None + assert output.error == "gpu fault" + assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0 + + def test_metrics_response_creates_metrics_only_output(self): + from vllm_omni.metrics import definitions as metric_defs + from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID + + client = self._make_client() + client._response_socket.recv.side_effect = [b"message", zmq.Again()] + client._decoder.decode.return_value = { + "type": "metrics", + "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0}, + } + + output = client.get_diffusion_output_nowait() + + assert output is not None + assert output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID + assert output.error is None + assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0 + def test_check_health_raises_when_proc_dead(self): """``check_health`` detects a dead subprocess via the manager's proc and raises ``EngineDeadError``, setting ``_engine_dead`` as a @@ -1089,7 +1171,6 @@ class TestStageDiffusionClientErrorPropagation: metadata, "tcp://req", "tcp://resp", - batch_size=1, ) @pytest.mark.asyncio diff --git a/tests/diffusion/test_omni_diffusion_config_master_port.py b/tests/diffusion/test_omni_diffusion_config_master_port.py index 8addd1c8c..0fe2dd08d 100644 --- a/tests/diffusion/test_omni_diffusion_config_master_port.py +++ b/tests/diffusion/test_omni_diffusion_config_master_port.py @@ -1,22 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for OmniDiffusionConfig master_port resolution (issue #3794).""" import socket import pytest +from tests.helpers.runtime import get_open_port from vllm_omni.diffusion.data import OmniDiffusionConfig pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _get_available_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind(("", 0)) - return sock.getsockname()[1] - - def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket: for _ in range(100): sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -36,27 +31,27 @@ def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket: class TestOmniDiffusionConfigMasterPort: def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None: - port = _get_available_port() + port = get_open_port() monkeypatch.setenv("MASTER_PORT", str(port)) config = OmniDiffusionConfig(model="test") assert config.master_port == port def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("MASTER_PORT", raising=False) - port = _get_available_port() + port = get_open_port() config = OmniDiffusionConfig(model="test", master_port=port) assert config.master_port == port def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None: - env_port = _get_available_port() - kwarg_port = _get_available_port() + env_port = get_open_port() + kwarg_port = get_open_port() monkeypatch.setenv("MASTER_PORT", str(env_port)) config = OmniDiffusionConfig(model="test", master_port=kwarg_port) assert config.master_port == env_port def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("MASTER_PORT", raising=False) - port = _get_available_port() + port = get_open_port() ports = {OmniDiffusionConfig(model="test", master_port=port).master_port for _ in range(5)} assert ports == {port} diff --git a/tests/diffusion/test_prompt_update.py b/tests/diffusion/test_prompt_update.py index cdee35979..f08825d1a 100644 --- a/tests/diffusion/test_prompt_update.py +++ b/tests/diffusion/test_prompt_update.py @@ -387,7 +387,6 @@ class TestPromptUpdateIntegration: model="test_model", od_config=od_config, metadata=metadata, - batch_size=1, ) @staticmethod diff --git a/tests/diffusion/test_result_pump.py b/tests/diffusion/test_result_pump.py index cf00cb6cc..fa172fd0c 100644 --- a/tests/diffusion/test_result_pump.py +++ b/tests/diffusion/test_result_pump.py @@ -6,6 +6,7 @@ import concurrent.futures import queue import threading import time +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -264,6 +265,119 @@ class TestResultPumpDispatch: fut = executor.wait_output_ready(async_output_id) assert fut.done() assert fut.result(timeout=1.0) is output + assert async_output_id not in executor._output_futures + + def test_output_ready_atomic_resolution_when_future_already_waiting(self): + """When OUTPUT_READY arrives and a future is already waiting, resolve it directly.""" + executor = _make_executor() + async_output_id = "abc123" + output = DiffusionOutput(output="waiting") + fut = executor.wait_output_ready(async_output_id) + assert not fut.done() + + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id=async_output_id, + output=output, + ) + _feed_one_msg_to_pump(executor, msg) + + assert fut.done() + assert fut.result(timeout=1.0) is output + assert async_output_id not in executor._output_futures + assert async_output_id not in executor._completed_outputs + + +class _FakeBatchOutput: + """Batch-level output exposing per-request results.""" + + def __init__(self, results): + self._results = results + + def get_request_output(self, req_id): + result = self._results.get(req_id) + if result is None: + return None + return SimpleNamespace(result=result) + + +def _make_scheduler_output(req_ids): + return SimpleNamespace( + scheduled_new_reqs=[SimpleNamespace(request_id=rid, req=SimpleNamespace()) for rid in req_ids] + ) + + +class TestBatchSplitDelivery: + """execute_batch must resolve per-request futures in either arrival order.""" + + @staticmethod + def _run(executor, req_ids, batch_id, deliver_early): + # Keep execute_batch on the fused request-batch path (not DLO DP). + executor.od_config.parallel_config.data_parallel_size = 1 + executor.od_config.enable_distributed_layerwise_offload = False + executor._ensure_open = lambda: None + + outputs = {rid: DiffusionOutput(output=f"img-{rid}") for rid in req_ids} + ready = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id=batch_id, + output=_FakeBatchOutput(outputs), + ) + + def fake_collective_rpc(*args, **kwargs): + if deliver_early: + # Worker's background D2H/SHM thread wins the race: OUTPUT_READY + # is pumped before execute_batch registers the split map. + _feed_one_msg_to_pump(executor, ready) + return AsyncDiffusionOutput( + kind=AsyncOutputKind.COMPUTE_DONE, + rpc_id="1", + async_output_id=batch_id, + ) + + executor.collective_rpc = fake_collective_rpc + batch = executor.execute_batch(_make_scheduler_output(req_ids)) + if not deliver_early: + _feed_one_msg_to_pump(executor, ready) + return batch, outputs + + def test_output_ready_after_split_map(self): + executor = _make_executor() + req_ids = ["r0", "r1", "r2"] + _, outputs = self._run(executor, req_ids, "batch-1", deliver_early=False) + + for rid in req_ids: + fut = executor.wait_output_ready(f"batch-1/{rid}") + assert fut.done() + assert fut.result(timeout=1.0) is outputs[rid] + + def test_output_ready_before_split_map(self): + """Regression: the whole batch used to be lost, hanging every request.""" + executor = _make_executor() + req_ids = ["r0", "r1", "r2"] + _, outputs = self._run(executor, req_ids, "batch-1", deliver_early=True) + + for rid in req_ids: + fut = executor.wait_output_ready(f"batch-1/{rid}") + assert fut.done(), f"request {rid} never resolved" + assert fut.result(timeout=1.0) is outputs[rid] + + # No stale batch-level state left behind. + assert executor._batch_split_map == {} + assert executor._completed_outputs == {} + + def test_engine_waiter_before_output_is_resolved(self): + """Consumers already blocked in wait_output_ready must be woken.""" + executor = _make_executor() + req_ids = ["r0", "r1"] + batch_id = "batch-2" + futures = {rid: executor.wait_output_ready(f"{batch_id}/{rid}") for rid in req_ids} + + _, outputs = self._run(executor, req_ids, batch_id, deliver_early=True) + + for rid in req_ids: + assert futures[rid].done(), f"request {rid} never resolved" + assert futures[rid].result(timeout=1.0) is outputs[rid] class TestShutdownCleansUpFutures: @@ -304,3 +418,77 @@ class TestShutdownCleansUpFutures: assert len(executor._rpc_futures) == 0 assert len(executor._output_futures) == 0 + + +class _RacyFuture(concurrent.futures.Future): + # done() always lies and reports False, to deterministically force the + # pump's check-then-act race window without needing real thread timing. + def done(self) -> bool: + return False + + +def _racy_cancelled_future() -> concurrent.futures.Future: + fut = _RacyFuture() + fut.cancel() + assert fut.cancelled() + return fut + + +class TestResultPumpCancelledFutureRace: + """Regression for #5793: a future cancelled concurrently with the pump's + resolve call (e.g. an asyncio.wait_for timeout, or a request abort) must + be dropped, not raise InvalidStateError and kill the pump thread. + """ + + def test_compute_done_racing_cancel_does_not_crash_pump(self): + executor = _make_executor() + fut = _racy_cancelled_future() + with executor._futures_lock: + executor._rpc_futures["1"] = fut + + msg = AsyncDiffusionOutput(kind=AsyncOutputKind.COMPUTE_DONE, rpc_id="1", async_output_id="abc") + _feed_one_msg_to_pump(executor, msg) + + assert fut.cancelled() + + def test_output_ready_racing_cancel_does_not_crash_pump(self): + executor = _make_executor() + fut = _racy_cancelled_future() + with executor._futures_lock: + executor._output_futures["abc123"] = fut + + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id="abc123", + output=DiffusionOutput(output="late"), + ) + _feed_one_msg_to_pump(executor, msg) + + assert fut.cancelled() + + def test_batch_split_racing_cancel_does_not_crash_pump(self): + executor = _make_executor() + cancelled_fut = _racy_cancelled_future() + healthy_fut = concurrent.futures.Future() + with executor._futures_lock: + executor._output_futures["batch-1/r-aborted"] = cancelled_fut + executor._output_futures["batch-1/r-healthy"] = healthy_fut + executor._batch_split_map["batch-1"] = { + "batch-1/r-aborted": "r-aborted", + "batch-1/r-healthy": "r-healthy", + } + + outputs = { + "r-aborted": DiffusionOutput(output="late"), + "r-healthy": DiffusionOutput(output="ok"), + } + msg = AsyncDiffusionOutput( + kind=AsyncOutputKind.OUTPUT_READY, + async_output_id="batch-1", + output=_FakeBatchOutput(outputs), + ) + _feed_one_msg_to_pump(executor, msg) + + assert cancelled_fut.cancelled() + assert healthy_fut.done() + assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"] diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index 697465097..90e753c9b 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import time @@ -7,6 +7,8 @@ from dataclasses import asdict, dataclass import pytest +import vllm_omni.diffusion.stage_diffusion_proc as stage_diffusion_proc +import vllm_omni.plugins as omni_plugins from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -26,6 +28,33 @@ BASE_INFER_STEPS = 10 DELAY_BASE = 0.01 +def test_run_diffusion_proc_sets_lifecycle_before_loading_plugins(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + + class StopProcessError(Exception): + pass + + class TestStageDiffusionProc(StageDiffusionProc): + def __init__(self, model, od_config): + events.append("proc") + raise StopProcessError + + monkeypatch.setattr(omni_plugins, "load_omni_general_plugins", lambda: events.append("plugins")) + monkeypatch.setattr(stage_diffusion_proc, "set_death_signal", lambda _: events.append("death_signal")) + monkeypatch.setattr(stage_diffusion_proc.signal, "signal", lambda *_: events.append("signal_handler")) + + with pytest.raises(StopProcessError): + TestStageDiffusionProc.run_diffusion_proc( + model="test-model", + od_config=None, + handshake_address="test-address", + local_client=True, + headless=False, + ) + + assert events == ["death_signal", "signal_handler", "signal_handler", "plugins", "proc"] + + class MockDiffusionEngine: async def step_streaming(self, request): def simulate_step_delay(height, width, num_inference_steps) -> float: diff --git a/tests/diffusion/test_uniproc_executor.py b/tests/diffusion/test_uniproc_executor.py new file mode 100644 index 000000000..dfc1d2eb1 --- /dev/null +++ b/tests/diffusion/test_uniproc_executor.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Single-GPU (uniproc) diffusion executor: selection and RPC behaviour.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from vllm.v1.engine.exceptions import EngineDeadError + +from vllm_omni.diffusion.executor.abstract import DiffusionExecutor +from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor +from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeODConfig: + def __init__(self, num_gpus: int = 1, backend: str | None = None) -> None: + self.num_gpus = num_gpus + self.distributed_executor_backend = backend + self.worker_extension_cls = None + self.custom_pipeline_args = None + + +def _od_config(num_gpus: int = 1, backend: str | None = None) -> _FakeODConfig: + return _FakeODConfig(num_gpus=num_gpus, backend=backend) + + +def test_single_gpu_defaults_to_uniproc(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=1)) is UniProcDiffusionExecutor + + +def test_explicit_mp_is_honored_on_single_gpu(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="mp")) is MultiprocDiffusionExecutor + + +def test_multi_gpu_defaults_to_multiproc(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=2)) is MultiprocDiffusionExecutor + + +def test_explicit_uni_backend_is_honored(): + assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="uni")) is UniProcDiffusionExecutor + + +@pytest.fixture +def executor(monkeypatch: pytest.MonkeyPatch): + """A ``UniProcDiffusionExecutor`` with a mocked worker (no model load).""" + worker = MagicMock() + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls", + MagicMock(return_value="vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker"), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname", + MagicMock(return_value=object), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase", + MagicMock(return_value=worker), + ) + return UniProcDiffusionExecutor(_od_config(num_gpus=1)), worker + + +def test_rejects_multi_gpu_config(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls", + MagicMock(return_value="unused"), + ) + with pytest.raises(ValueError, match="single GPU only"): + UniProcDiffusionExecutor(_od_config(num_gpus=4)) + + +def test_shutdown_is_safe_on_partially_constructed_executor(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls", + MagicMock(return_value="unused"), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname", + MagicMock(return_value=object), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase", + MagicMock(side_effect=RuntimeError("model load blew up")), + ) + constructed = UniProcDiffusionExecutor.__new__(UniProcDiffusionExecutor) + with pytest.raises(RuntimeError, match="model load blew up"): + constructed.__init__(_od_config(num_gpus=1)) + + constructed.shutdown() + + +def test_collective_rpc_calls_worker_directly(executor): + exec_, worker = executor + worker.execute_method.return_value = "result" + + out = exec_.collective_rpc("some_method", args=(1, 2), kwargs={"k": "v"}, unique_reply_rank=0) + + assert out == "result" + worker.execute_method.assert_called_once_with("some_method", 1, 2, k="v") + + +def test_collective_rpc_returns_list_when_no_reply_rank(executor): + exec_, worker = executor + worker.execute_method.return_value = "result" + + assert exec_.collective_rpc("some_method") == ["result"] + + +def test_collective_rpc_propagates_worker_exceptions(executor): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + +def test_recoverable_worker_failure_keeps_the_executor_alive(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("CUDA out of memory") + monkeypatch.setattr(exec_, "_device_is_usable", lambda: True) + died = MagicMock() + exec_.register_failure_callback(died) + + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + assert exec_.is_dead is False + died.assert_not_called() + exec_.check_health() + + +def test_poisoned_cuda_context_marks_the_executor_dead(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("an illegal memory access was encountered") + monkeypatch.setattr(exec_, "_device_is_usable", lambda: False) + died = MagicMock() + exec_.register_failure_callback(died) + + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + assert exec_.is_dead is True + died.assert_called_once_with() + with pytest.raises(EngineDeadError): + exec_.check_health() + + +def test_failure_is_latched_and_callbacks_fire_once(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(exec_, "_device_is_usable", lambda: False) + died = MagicMock() + exec_.register_failure_callback(died) + + for _ in range(3): + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + died.assert_called_once_with() + + +def test_device_probe_is_skipped_when_cuda_was_never_initialized(executor): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + + with pytest.raises(RuntimeError): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + if not torch.cuda.is_initialized(): + assert exec_.is_dead is False + + +def test_device_probe_skipped_when_accelerator_unavailable(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(torch.accelerator, "is_available", lambda: False) + sync = MagicMock() + monkeypatch.setattr(torch.accelerator, "synchronize", sync) + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + sync.assert_not_called() + assert exec_.is_dead is False + + +def test_device_probe_skipped_when_device_not_initialized(executor, monkeypatch): + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu")) + monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: False), raising=False) + sync = MagicMock() + monkeypatch.setattr(torch.accelerator, "synchronize", sync) + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + sync.assert_not_called() + assert exec_.is_dead is False + + +def test_device_probe_latches_sticky_fault_on_npu(executor, monkeypatch): + """NPU must not short-circuit the probe the way a CUDA-only guard did.""" + exec_, worker = executor + worker.execute_method.side_effect = RuntimeError("boom") + monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu")) + monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: True), raising=False) + monkeypatch.setattr( + torch.accelerator, + "synchronize", + MagicMock(side_effect=RuntimeError("NPU context poisoned")), + ) + died = MagicMock() + exec_.register_failure_callback(died) + + with pytest.raises(RuntimeError, match="boom"): + exec_.collective_rpc("some_method", unique_reply_rank=0) + + assert exec_.is_dead is True + died.assert_called_once_with() + + +def test_check_health_ok_then_dead(executor): + exec_, _ = executor + exec_.check_health() + + exec_._is_failed = True + with pytest.raises(EngineDeadError): + exec_.check_health() + + +def test_shutdown_is_idempotent_and_closes_executor(executor, monkeypatch): + exec_, worker = executor + collect = MagicMock() + empty_cache = MagicMock() + monkeypatch.setattr("vllm_omni.diffusion.executor.uniproc_executor.gc.collect", collect) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.is_available", + lambda: True, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.empty_cache", + empty_cache, + ) + exec_.register_failure_callback(MagicMock()) + + exec_.shutdown() + exec_.shutdown() + + worker.shutdown.assert_called_once() + assert exec_.driver_worker is None + assert exec_._failure_callbacks == [] + collect.assert_called_once_with() + empty_cache.assert_called_once_with() + with pytest.raises(RuntimeError, match="closed"): + exec_.collective_rpc("some_method", unique_reply_rank=0) diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index 170faec80..709fc7283 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for WorkerWrapperBase class. @@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation: - Dynamic worker class extension """ +from types import SimpleNamespace from typing import Any import pytest @@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation: wrapper.worker.shutdown.assert_called_once() assert result is None + def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture): + events: list[str] = [] + offload_backend = mocker.Mock() + offload_backend.disable.side_effect = lambda: events.append("offload") + kv_manager = mocker.Mock() + kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv") + destroy = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env", + side_effect=lambda: events.append("distributed"), + ) + worker = DiffusionWorker.__new__(DiffusionWorker) + worker.model_runner = SimpleNamespace( + offload_backend=offload_backend, + kv_transfer_manager=kv_manager, + ) + + worker.shutdown() + + assert events == ["offload", "kv", "distributed"] + destroy.assert_called_once_with() + # ------------------------------------------------------------------------- # Tests: execute_method diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py index 57997905c..c31b80a6f 100644 --- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end test for Bagel text2img with the Mooncake inter-stage connector. @@ -22,7 +22,7 @@ import pytest from PIL import Image from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner +from tests.helpers.runtime import OmniRunner, get_open_port from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.omni import Omni from vllm_omni.outputs import OmniRequestOutput @@ -55,15 +55,6 @@ PIXEL_TOLERANCE = 10 DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" -def _find_free_port() -> int: - """Find and return a free ephemeral port by binding to port 0.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - s.listen(1) - port = s.getsockname()[1] - return port - - def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list: """Configure sampling parameters for Bagel text2img generation. @@ -277,9 +268,9 @@ def test_bagel_text2img_mooncake_connector(run_level): "mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)" ) MOONCAKE_HOST = "127.0.0.1" - MOONCAKE_RPC_PORT = _find_free_port() - MOONCAKE_HTTP_PORT = _find_free_port() - MOONCAKE_METRICS_PORT = _find_free_port() + MOONCAKE_RPC_PORT = get_open_port() + MOONCAKE_HTTP_PORT = get_open_port() + MOONCAKE_METRICS_PORT = get_open_port() mooncake_master_proc = None temp_config_file = None @@ -314,7 +305,7 @@ def test_bagel_text2img_mooncake_connector(run_level): temp_config_file = _resolve_deploy_config(temp_config_file, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", - stage_configs_path=temp_config_file, + deploy_config=temp_config_file, stage_init_timeout=300, ) as runner: generated_image = _generate_bagel_image(runner.omni) diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py index b97161335..6b23ba321 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end tests for Bagel with shared memory connector: img2img and text2img. @@ -13,7 +13,6 @@ End-to-end tests for Bagel with shared memory connector: img2img and text2img. import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -import socket from typing import Any import pytest @@ -82,15 +81,6 @@ def _load_input_image() -> Image.Image: return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") -def _find_free_port() -> int: - """Find and return a free ephemeral port by binding to port 0.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - s.listen(1) - port = s.getsockname()[1] - return port - - def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list: """Configure sampling parameters for Bagel img2img generation. @@ -251,7 +241,7 @@ def test_bagel_img2img_shared_memory_connector(run_level): config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", - stage_configs_path=config_path, + deploy_config=config_path, ) as runner: generated_image = _generate_bagel_img2img(runner.omni, input_image) if run_level == "advanced_model": @@ -266,7 +256,7 @@ def test_bagel_text2img_shared_memory_connector(run_level): config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", - stage_configs_path=config_path, + deploy_config=config_path, ) as runner: generated_image = _generate_bagel_text2img(runner.omni) if run_level == "advanced_model": diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index da4e81995..0e3194190 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -1,18 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import threading from collections import deque -from types import SimpleNamespace -from unittest.mock import patch +from types import MethodType, SimpleNamespace +from unittest.mock import Mock, patch import pytest import torch from pytest_mock import MockerFixture +from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats -from vllm.v1.request import RequestStatus +from vllm.v1.request import Request, RequestStatus +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt @@ -32,9 +34,14 @@ class DummyWaitingQueue(list): def add_request(self, request): self.append(request) + def remove_requests(self, requests): + remove = set(requests) + self[:] = [request for request in self if request not in remove] + def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None): - return SimpleNamespace( + request = Mock( + client_index=0, request_id=req_id, external_req_id=external_req_id or req_id, status=status, @@ -44,7 +51,22 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None) num_output_placeholders=0, prefill_stats=None, additional_information=None, - is_finished=lambda: status == RequestStatus.FINISHED_STOPPED, + resumable=False, + ) + request.is_finished = lambda: RequestStatus.is_finished(request.status) + return request + + +def _streaming_request(mocker: MockerFixture, num_computed_tokens: int) -> SimpleNamespace: + prompt_token_ids = [0] * num_computed_tokens + return SimpleNamespace( + _all_token_ids=prompt_token_ids.copy(), + _output_token_ids=[], + prompt_token_ids=prompt_token_ids, + num_computed_tokens=num_computed_tokens, + num_prompt_tokens=num_computed_tokens, + num_output_placeholders=0, + update_block_hashes=mocker.Mock(), ) @@ -75,30 +97,219 @@ def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) request.update_block_hashes.assert_called_once_with() -def test_streaming_payload_can_append_exact_prompt_length(mocker: MockerFixture) -> None: +def test_turn_start_replacement_ignores_accumulated_prompt_capacity(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=4064) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "replace_streaming_prompt": True, + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + } + + replaced = construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=10, + previous_condition_seq=7, + condition_seq=8, + recompute_previous_chunks=1, + ) + + assert replaced is True + assert request.prompt_token_ids == [0] * 10 + assert payload["meta"]["streaming_prompt_recompute"] is False + assert "streaming_prompt_previous_codes" not in payload["ids"] + + +@pytest.mark.parametrize("reserve", [True, 0, -1, 1.5, "26", None]) +def test_streaming_prompt_rejects_invalid_generation_reserve( + mocker: MockerFixture, + reserve: object, +) -> None: + request = _streaming_request(mocker, num_computed_tokens=10) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": reserve, + }, + } + + with pytest.raises(ValueError, match="next_stage_generation_tokens must be a positive integer"): + construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + +def test_capacity_managed_streaming_prompt_at_limit_still_appends(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=4060) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + } + + replaced = construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + assert replaced is False + assert request.num_computed_tokens == 4060 + assert request.num_prompt_tokens == 4070 + + +def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixture) -> None: + previous_condition_len = 10 + previous_codes = list(range(25)) request = SimpleNamespace( - _all_token_ids=[0, 0, 7, 8], - _output_token_ids=[7, 8], - prompt_token_ids=[0, 0], - num_computed_tokens=4, - num_prompt_tokens=2, + _all_token_ids=[0] * 4039 + previous_codes + [999], + _output_token_ids=previous_codes + [999], + prompt_token_ids=[0] * 4039, + num_computed_tokens=4065, + num_prompt_tokens=4039, + num_output_placeholders=1, update_block_hashes=mocker.Mock(), ) payload = { - "ids": {"prompt": [1, 2, 3]}, - "meta": {"next_stage_prompt_len": 3}, + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, } - construct_next_stage_streaming_input_prompt(payload, request) + replaced = construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=previous_condition_len, + previous_condition_seq=7, + condition_seq=8, + recompute_previous_chunks=1, + ) - assert request.prompt_token_ids == [0, 0, 7, 8, 0, 0, 0] - assert request._all_token_ids == [0, 0, 7, 8, 0, 0, 0] - assert request._output_token_ids == [] - assert request.num_computed_tokens == 4 - assert request.num_prompt_tokens == 7 + assert replaced is True + assert request.prompt_token_ids == [0] * 45 + assert request.num_prompt_tokens == 45 + assert request.num_computed_tokens == 0 + assert payload["ids"]["streaming_prompt_previous_codes"] == previous_codes + assert payload["meta"] == { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + "streaming_prompt_recompute": True, + "streaming_condition_seq": 8, + } request.update_block_hashes.assert_called_once_with() +def test_streaming_window_recomputes_each_condition_without_accumulating(mocker: MockerFixture) -> None: + request = SimpleNamespace( + _all_token_ids=[0] * 10 + [101, 102], + _output_token_ids=[101, 102], + prompt_token_ids=[0] * 10, + num_computed_tokens=12, + num_prompt_tokens=10, + num_output_placeholders=0, + update_block_hashes=mocker.Mock(), + ) + second = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 12, + "next_stage_generation_tokens": 26, + }, + } + + assert construct_next_stage_streaming_input_prompt( + second, + request, + max_model_len=4096, + previous_condition_len=10, + previous_condition_seq=0, + condition_seq=1, + recompute_previous_chunks=1, + ) + assert request.num_prompt_tokens == 24 + assert second["ids"]["streaming_prompt_previous_codes"] == [101, 102] + + request._all_token_ids.extend([201, 202, 203]) + request._output_token_ids.extend([201, 202, 203]) + request.num_computed_tokens = 27 + third = { + "ids": {"prompt": [2]}, + "meta": { + "next_stage_prompt_len": 8, + "next_stage_generation_tokens": 26, + }, + } + + assert construct_next_stage_streaming_input_prompt( + third, + request, + max_model_len=4096, + previous_condition_len=12, + previous_condition_seq=1, + condition_seq=2, + recompute_previous_chunks=1, + ) + assert request.num_prompt_tokens == 23 + assert third["ids"]["streaming_prompt_previous_codes"] == [201, 202, 203] + assert third["meta"]["streaming_condition_seq"] == 2 + + +def test_capacity_rollover_requires_explicit_window_contract(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=4064) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + } + + with pytest.raises(ValueError, match="requires window_size=1"): + construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=10, + previous_condition_seq=0, + condition_seq=1, + ) + + +@pytest.mark.parametrize("next_stage_prompt_len", [0, -1]) +def test_capacity_managed_streaming_prompt_rejects_nonpositive_length( + mocker: MockerFixture, next_stage_prompt_len: int +) -> None: + request = _streaming_request(mocker, num_computed_tokens=0) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": next_stage_prompt_len, + "next_stage_generation_tokens": 26, + }, + } + + with pytest.raises(ValueError, match="positive next_stage_prompt_len"): + construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + +def test_capacity_managed_streaming_prompt_rejects_fresh_overflow(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=0) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 4071, + "next_stage_generation_tokens": 26, + }, + } + + with pytest.raises(ValueError, match="exceeds max_model_len"): + construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + @pytest.fixture def build_adapter(monkeypatch, mocker: MockerFixture): def _build( @@ -106,6 +317,10 @@ def build_adapter(monkeypatch, mocker: MockerFixture): stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2, + max_model_len: int = 0, + tts_max_model_len: int = 0, + tts_attention_type: str = "full_attention", + flat_tts_config: bool = False, active_stream_window: int = 0, connector_extra: dict | None = None, ): @@ -122,6 +337,10 @@ def build_adapter(monkeypatch, mocker: MockerFixture): self._cancelled_load_reqs = set() self._pending_save_reqs = deque() self._finished_save_reqs = set() + self._send_failures = {} + self._send_failure_lock = threading.Lock() + self._receive_failures = {} + self._receive_failure_lock = threading.Lock() self.stop_event = threading.Event() self._recv_cond = threading.Condition() self._save_cond = threading.Condition() @@ -133,9 +352,23 @@ def build_adapter(monkeypatch, mocker: MockerFixture): classmethod(lambda cls, _model_config: connector), ) + hf_config = SimpleNamespace( + model_type="minicpmtts", + max_position_embeddings=tts_max_model_len, + attention_type=tts_attention_type, + ) + if not flat_tts_config: + hf_config = SimpleNamespace( + tts_config=SimpleNamespace( + max_position_embeddings=tts_max_model_len, + attention_type=tts_attention_type, + ) + ) model_config = SimpleNamespace( worker_type=model_mode, max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + hf_config=hf_config, active_stream_window=active_stream_window, stage_connector_config={ "name": "SharedMemoryConnector", @@ -151,6 +384,27 @@ def build_adapter(monkeypatch, mocker: MockerFixture): return _build +def _dequeue_load_entry(adapter, request): + """Register and take the entry as the background recv loop would.""" + adapter.load_async(request) + entry = adapter._registered_load_entries[request.request_id] + try: + adapter._pending_load_reqs.remove(entry) + except ValueError: + pass + return entry + + +@pytest.mark.parametrize( + ("attention_type", "expected_previous_chunks"), + [("full_attention", 0), ("sliding_recompute", 1)], +) +def test_talker_attention_policy_controls_streaming_recompute(build_adapter, attention_type, expected_previous_chunks): + adapter, _ = build_adapter(tts_attention_type=attention_type) + + assert adapter._streaming_prompt_previous_chunks == expected_previous_chunks + + @pytest.mark.parametrize( ("raw_cfg", "expected_name", "expected_extra"), [ @@ -191,15 +445,258 @@ def test_load_poll(build_adapter): "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, } connector.get.return_value = (payload, 16) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) + assert request.additional_information is None + assert "req-1" in adapter._finished_load_reqs + request.status = RequestStatus.WAITING_FOR_CHUNK + adapter.process_pending_chunks( + DummyWaitingQueue(), + [request], + scheduler_requests={request.request_id: request}, + ) assert request.additional_information == payload assert adapter.get_req_chunk["req-1"] == 1 - assert "req-1" in adapter._finished_load_reqs + assert "req-1" not in adapter._finished_load_reqs assert "req-1" in adapter.upstream_exhausted_requests assert "req-1" not in adapter._pending_load_reqs +def test_load_async_does_not_requeue_registered_inflight_request(build_adapter): + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-inflight", RequestStatus.RUNNING, external_req_id="ext-inflight") + + entry = _dequeue_load_entry(adapter, request) + + # The recv loop owns the popped request until polling completes. A + # scheduler-side duplicate registration must not append a second copy. + adapter.load_async(request) + assert adapter._pending_load_reqs == deque() + + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": {"finished": True}, + }, + 1, + ) + assert adapter._poll_single_request(entry) is True + + # Successful polling releases the registration for the next chunk. + adapter.load_async(request) + assert [queued.request for queued in adapter._pending_load_reqs] == [request] + + +def test_cleanup_then_reregister_drops_stale_queued_entry(build_adapter): + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-resume", RequestStatus.WAITING, external_req_id="ext-resume") + + adapter.load_async(request) + stale_entry = adapter._pending_load_reqs[0] + adapter.cleanup_receiver(request.request_id) + adapter.load_async(request) + fresh_entry = adapter._pending_load_reqs[-1] + assert len(adapter._pending_load_reqs) == 2 + assert adapter._pending_load_reqs.popleft() is stale_entry + assert adapter._pending_load_reqs.popleft() is fresh_entry + + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": {"finished": True}, + }, + 1, + ) + assert adapter._poll_single_request(stale_entry) is True + connector.get.assert_not_called() + assert adapter._registered_load_entries[request.request_id] is fresh_entry + + assert adapter._poll_single_request(fresh_entry) is True + connector.get.assert_called_once() + assert request.request_id not in adapter._registered_load_entries + assert adapter.get_req_chunk[request.request_id] == 1 + + +def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter): + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-replace", RequestStatus.RUNNING, external_req_id="external-replace") + request.resumable = True + request.prompt_token_ids = [0] * 4 + request._all_token_ids = [0] * 4 + request._output_token_ids = [] + request.num_prompt_tokens = 4 + request.num_computed_tokens = 4 + request.update_block_hashes = Mock() + adapter.get_req_chunk[request.request_id] = 1 + adapter.requests_num_chunks_sent[request.external_req_id] = 4 + adapter.request_ids_mapping[request.request_id] = request.external_req_id + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "replace_streaming_prompt": True, + "finished": False, + }, + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + request.status = RequestStatus.WAITING_FOR_CHUNK + running_queue = [request] + waiting_queue = DummyWaitingQueue() + adapter.process_pending_chunks( + waiting_queue, + running_queue, + scheduler_requests={request.request_id: request}, + ) + assert request.num_computed_tokens == 0 + assert request.prompt_token_ids == [0] * 10 + assert request.request_id in adapter.replaced_streaming_prompt_ids + assert running_queue == [] + assert waiting_queue == [request] + assert request.status == RequestStatus.WAITING + + adapter.requests_num_chunks_sent[request.external_req_id] = 9 + adapter.postprocess_scheduler_output( + SimpleNamespace( + scheduled_new_reqs=[SimpleNamespace(req_id=request.request_id)], + scheduled_cached_reqs=SimpleNamespace(req_ids=[]), + ) + ) + assert request.request_id not in adapter.replaced_streaming_prompt_ids + assert request.external_req_id not in adapter.requests_num_chunks_sent + + +@pytest.mark.parametrize("flat_tts_config", [False, True]) +def test_load_poll_ar_recomputes_native_duplex_prompt_each_condition(build_adapter, flat_tts_config): + adapter, connector = build_adapter( + stage_id=1, + model_mode="ar", + max_model_len=8192, + tts_max_model_len=4096, + tts_attention_type="sliding_recompute", + flat_tts_config=flat_tts_config, + ) + request = _req("req-rollover", RequestStatus.RUNNING, external_req_id="external-rollover") + request.resumable = True + previous_codes = list(range(25)) + request.prompt_token_ids = [0] * 10 + request._all_token_ids = [0] * 10 + previous_codes + request._output_token_ids = previous_codes.copy() + request.num_prompt_tokens = 10 + request.num_computed_tokens = 35 + request.num_output_placeholders = 0 + request.update_block_hashes = Mock() + adapter.get_req_chunk[request.request_id] = 1 + adapter.request_ids_mapping[request.request_id] = request.external_req_id + adapter._streaming_condition_lengths[request.request_id] = 10 + adapter._streaming_condition_seqs[request.request_id] = 0 + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + "finished": True, + }, + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + request.status = RequestStatus.WAITING_FOR_CHUNK + running_queue = [request] + waiting_queue = DummyWaitingQueue() + adapter.process_pending_chunks( + waiting_queue, + running_queue, + scheduler_requests={request.request_id: request}, + ) + + assert request.num_computed_tokens == 0 + assert request.prompt_token_ids == [0] * 45 + assert request.additional_information["ids"]["streaming_prompt_previous_codes"] == previous_codes + assert request.additional_information["meta"]["streaming_prompt_recompute"] is True + assert request.resumable is False + assert request.request_id in adapter.replaced_streaming_prompt_ids + assert adapter._max_model_len == 4096 + assert running_queue == [] + assert waiting_queue == [request] + + +def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> None: + adapter, connector = build_adapter( + stage_id=1, + model_mode="ar", + tts_attention_type="sliding_recompute", + ) + request = _req("req-condition-seq", RequestStatus.RUNNING) + + def receive(payload): + connector.get.return_value = (payload, 1) + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + adapter.requests_with_ready_chunks.add(request.request_id) + adapter._apply_pending_ar_prompt_updates({request.request_id: request}) + adapter.requests_with_ready_chunks.discard(request.request_id) + return request.additional_information + + first = receive( + { + "ids": {"prompt": [1]}, + "meta": {"next_stage_prompt_len": 10}, + } + ) + assert first["meta"]["streaming_condition_seq"] == 0 + assert adapter._streaming_condition_seqs[request.request_id] == 0 + + control = receive({"meta": {"is_segment_finished": True}}) + assert "streaming_condition_seq" not in control["meta"] + assert adapter._streaming_condition_seqs[request.request_id] == 0 + + second = receive( + { + "ids": {"prompt": [2]}, + "meta": {"next_stage_prompt_len": 12}, + } + ) + assert second["meta"]["streaming_condition_seq"] == 1 + assert adapter._streaming_condition_lengths[request.request_id] == 12 + + +def test_load_poll_ar_reports_invalid_capacity_metadata(build_adapter): + adapter, connector = build_adapter(stage_id=1, model_mode="ar", max_model_len=4096) + request = _req("req-invalid-capacity", RequestStatus.RUNNING) + request.resumable = True + adapter.get_req_chunk[request.request_id] = 1 + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 4096, + "next_stage_generation_tokens": 26, + "finished": False, + }, + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + request.status = RequestStatus.WAITING_FOR_CHUNK + running_queue = [request] + adapter.process_pending_chunks( + DummyWaitingQueue(), + running_queue, + scheduler_requests={request.request_id: request}, + ) + assert adapter.collect_failed_receive_request_ids() == { + request.request_id: "fresh streaming prompt plus generation reserve exceeds max_model_len: " + "prompt=4096, reserve=26, limit=4096" + } + assert request.request_id not in adapter._finished_load_reqs + + def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter): adapter, connector = build_adapter(stage_id=1, model_mode="generation") request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor") @@ -214,7 +711,7 @@ def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter) } connector.get.return_value = (payload, 16) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) assert request.prompt_token_ids == [0] assert request.num_computed_tokens == 0 @@ -244,12 +741,13 @@ def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapte } connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)] - assert adapter._poll_single_request(request) is False + entry = _dequeue_load_entry(adapter, request) + assert adapter._poll_single_request(entry) is False assert request.request_id not in adapter._finished_load_reqs assert request.request_id not in adapter.requests_with_ready_chunks assert adapter.get_req_chunk[request.request_id] == 1 - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(entry) is True assert request.request_id in adapter._finished_load_reqs assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"]) assert adapter.get_req_chunk[request.request_id] == 2 @@ -280,6 +778,126 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad assert len(adapter._pending_save_reqs) == 1 +def test_segment_boundary_starts_new_send_watermark_before_background_flush(build_adapter): + """A queued boundary owns the end of its deduplication generation. + + The next segment can start before the save thread sends the old boundary. + Sending that old task later must not erase the new segment's watermark. + """ + adapter, _ = build_adapter(stage_id=1) + request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") + request.resumable = True + request._omni_segment_generation = 1 + request.num_computed_tokens = 0 + adapter.requests_num_chunks_sent["ext-stream"] = 26 + + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=True, + confirmed_num_computed_tokens=26, + segment_generation=0, + ) + + assert len(adapter._pending_save_reqs) == 1 + boundary_task = adapter._pending_save_reqs.popleft() + assert boundary_task["segment_generation"] == 0 + assert "ext-stream" not in adapter.requests_num_chunks_sent + + request.num_computed_tokens = 3 + request._omni_segment_generation = 1 + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=False, + ) + + assert len(adapter._pending_save_reqs) == 1 + assert adapter.requests_num_chunks_sent["ext-stream"] == 3 + + adapter._send_single_request(boundary_task) + + assert adapter.requests_num_chunks_sent["ext-stream"] == 3 + + +def test_background_send_uses_enqueued_request_snapshot(build_adapter): + """A queued segment must not observe later in-place request mutations.""" + adapter, _ = build_adapter(stage_id=1) + request = Request( + request_id="req-stream", + prompt_token_ids=[1, 2], + sampling_params=SamplingParams(max_tokens=8), + pooling_params=None, + resumable=True, + ) + request.external_req_id = "ext-stream" + ref_audio = torch.tensor([0.1, -0.1]) + request.additional_information = {"codes": {"ref": ref_audio}, "meta": {"segment": "old"}} + request.append_output_token_ids([7]) + seen_requests = [] + + def recording_processor(**kwargs): + queued = kwargs["request"] + seen_requests.append( + ( + queued.additional_information["meta"]["segment"], + list(queued.prompt_token_ids), + list(queued.output_token_ids), + list(queued.all_token_ids), + queued.additional_information["codes"]["ref"] is ref_audio, + ) + ) + return OmniPayloadStruct() + + adapter.custom_process_next_stage_input_func = recording_processor + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=True, + ) + request.additional_information["meta"]["segment"] = "next" + request.prompt_token_ids.append(3) + request.append_output_token_ids([8]) + request._omni_segment_generation = 1 + adapter.save_async( + multimodal_output=None, + request=request, + is_segment_finished=True, + ) + request.additional_information["meta"]["segment"] = "later" + request.prompt_token_ids.append(4) + request.append_output_token_ids([9]) + + assert len(adapter._pending_save_reqs) == 2 + first_task = adapter._pending_save_reqs.popleft() + second_task = adapter._pending_save_reqs.popleft() + assert first_task["request"].additional_information is not request.additional_information + assert first_task["request"].additional_information["meta"] is not request.additional_information["meta"] + + adapter._send_single_request(first_task) + adapter._send_single_request(second_task) + + assert seen_requests == [ + ("old", [1, 2], [7], [1, 2, 7], True), + ("next", [1, 2, 3], [7, 8], [1, 2, 7, 8], True), + ] + + +def test_save_without_custom_processor_does_not_snapshot_request(build_adapter, mocker): + adapter, _ = build_adapter(stage_id=1) + request = _req("req-direct", RequestStatus.RUNNING, external_req_id="ext-direct") + snapshot = mocker.patch.object( + adapter, + "_snapshot_processor_request", + side_effect=AssertionError("unexpected snapshot"), + ) + + adapter.save_async(multimodal_output=None, request=request) + + snapshot.assert_not_called() + assert adapter._pending_save_reqs.popleft()["request"] is request + + def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch): """A terminal stop is not a segment boundary (#5383), but the producer-side processor must still receive the flush signal on the terminal chunk. @@ -489,7 +1107,7 @@ def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter): assert second.codes is None -def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter): +def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter): adapter, _ = build_adapter(stage_id=1) request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") request.resumable = True @@ -508,6 +1126,51 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt assert adapter.requests_num_chunks_sent["ext-stream"] == 0 +def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter): + adapter, _ = build_adapter(stage_id=1) + request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") + request.resumable = True + request._omni_segment_generation = 0 + adapter.requests_num_chunks_sent["ext-stream"] = 111 + + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True) + + late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream") + late.resumable = True + late.num_computed_tokens = 111 + late._omni_segment_generation = 0 + adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False) + + next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream") + next_segment.resumable = True + next_segment.num_computed_tokens = 3 + next_segment._omni_segment_generation = 1 + adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False) + + assert len(adapter._pending_save_reqs) == 2 + assert adapter.requests_num_chunks_sent["ext-stream"] == 3 + assert adapter._pending_save_reqs[-1]["request"] is next_segment + + +def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter): + adapter, connector = build_adapter(stage_id=1) + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct()) + request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old") + request.resumable = True + request._omni_segment_generation = 0 + + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False) + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True) + + old_task = adapter._pending_save_reqs.popleft() + adapter._send_single_request(old_task) + connector.put.assert_called_once() + + boundary_task = adapter._pending_save_reqs.popleft() + adapter._send_single_request(boundary_task) + assert connector.put.call_count == 2 + + def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch): adapter, _ = build_adapter(stage_id=1) request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished") @@ -528,6 +1191,77 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon assert args[1] == "ext-finished" +def test_old_turn_terminal_send_preserves_new_turn_state_and_sender(build_adapter): + adapter, connector = build_adapter( + stage_id=1, + model_mode="ar", + max_model_len=4096, + tts_attention_type="sliding_recompute", + ) + request = _req("req-live", RequestStatus.RUNNING, external_req_id="ext-live") + request.resumable = True + request._all_token_ids = [] + request._output_token_ids = [] + request.update_block_hashes = Mock() + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct( + meta=MetaStruct( + finished=torch.tensor(True, dtype=torch.bool), + codec_streaming=True, + ) + ) + + put_started = threading.Event() + release_put = threading.Event() + put_keys = [] + + def blocking_first_put(*_args, **kwargs): + put_keys.append(kwargs["put_key"]) + if len(put_keys) == 1: + put_started.set() + assert release_put.wait(timeout=5) + return True, 1, {} + + connector.put.side_effect = blocking_first_put + adapter.save_async(multimodal_output=None, request=request) + first_task = adapter._pending_save_reqs.popleft() + first_sender_token = first_task["sender_token"] + send_thread = threading.Thread(target=lambda: adapter._send_single_request(first_task)) + send_thread.start() + assert put_started.wait(timeout=5) + + entry = _dequeue_load_entry(adapter, request) + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": { + "replace_streaming_prompt": True, + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + }, + 1, + ) + assert adapter._poll_single_request(entry) is True + adapter.requests_with_ready_chunks.add(request.request_id) + adapter._apply_pending_ar_prompt_updates({request.request_id: request}) + assert adapter._streaming_condition_seqs[request.request_id] == 0 + + request.num_computed_tokens = 1 + adapter.save_async(multimodal_output=None, request=request) + second_task = adapter._pending_save_reqs.popleft() + assert second_task["sender_token"] is first_sender_token + + release_put.set() + send_thread.join(timeout=5) + assert not send_thread.is_alive() + assert adapter._streaming_condition_seqs[request.request_id] == 0 + assert adapter._sender_tokens[request.external_req_id] is first_sender_token + assert first_sender_token.cancelled is False + + adapter._send_single_request(second_task) + assert put_keys == ["ext-live_1_0", "ext-live_1_1"] + + def test_load_poll_non_ar_merges_into_existing_additional_information(build_adapter): adapter, connector = build_adapter(stage_id=2, model_mode="diffusion") request = _req("req-non-ar", RequestStatus.WAITING, external_req_id="ext-non-ar") @@ -547,7 +1281,7 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap } connector.get.return_value = (payload, 8) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True assert request.prompt_token_ids == [7, 8] assert request.num_computed_tokens == 0 @@ -565,6 +1299,92 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap assert "req-non-ar" in adapter.upstream_exhausted_requests +def test_load_poll_generation_segment_marker_replaces_previous_chunk(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="generation") + request = _req("req-marker", RequestStatus.WAITING, external_req_id="external-marker") + request.additional_information = { + "codes": {"audio": torch.tensor([1, 2])}, + "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True}, + } + connector.get.return_value = ( + { + "codes": { + "audio": torch.tensor([7, 8], dtype=torch.long), + "ref": torch.tensor([0.1, -0.1]), + }, + "meta": { + "finished": torch.tensor(False, dtype=torch.bool), + "is_segment_finished": torch.tensor(True, dtype=torch.bool), + "request_finished": torch.tensor(False, dtype=torch.bool), + "replace_runtime_additional_information": True, + }, + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + + assert request.prompt_token_ids == [7, 8] + assert "audio" not in request.additional_information["codes"] + torch.testing.assert_close( + request.additional_information["codes"]["ref"], + torch.tensor([0.1, -0.1]), + ) + assert not {"cache_epoch", "chunk_seq", "last_chunk"}.intersection(request.additional_information["meta"]) + assert request.request_id in adapter.segment_finished_requests + + +def test_load_poll_generation_empty_replacement_snapshot_is_ready(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="generation") + request = _req("req-empty-marker", RequestStatus.WAITING, external_req_id="external-empty-marker") + request.additional_information = { + "codes": {"audio": torch.tensor([1, 2])}, + "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True}, + } + connector.get.return_value = ( + { + "meta": { + "is_segment_finished": torch.tensor(True, dtype=torch.bool), + "replace_runtime_additional_information": True, + } + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + + assert request.prompt_token_ids == [0] + assert "codes" not in request.additional_information + assert request.additional_information["meta"]["replace_runtime_additional_information"] is True + assert request.request_id in adapter.segment_finished_requests + assert request.request_id in adapter._finished_load_reqs + + +def test_load_poll_generation_without_snapshot_marker_keeps_incremental_state(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="generation") + request = _req("req-incremental", RequestStatus.WAITING, external_req_id="external-incremental") + request.additional_information = { + "codes": {"audio": torch.tensor([1, 2])}, + "meta": {"cache_epoch": 3, "chunk_seq": 2}, + } + connector.get.return_value = ( + { + "meta": { + "finished": torch.tensor(False, dtype=torch.bool), + "phase": "decode", + } + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is False + + assert torch.equal(request.additional_information["codes"]["audio"], torch.tensor([1, 2])) + assert request.additional_information["meta"]["cache_epoch"] == 3 + assert request.additional_information["meta"]["chunk_seq"] == 2 + assert request.additional_information["meta"]["phase"] == "decode" + + def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter): adapter, connector = build_adapter(stage_id=2, model_mode="ar") request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged") @@ -581,9 +1401,17 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt } connector.get.return_value = (payload, 8) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) - # AR mode now forwards the latest payload directly. + # The recv thread records the latest payload; scheduler-side queue + # restoration publishes it onto the Request. + assert torch.equal(request.additional_information["hidden_states"]["output"], torch.tensor([[1.0]])) + request.status = RequestStatus.WAITING_FOR_CHUNK + adapter.process_pending_chunks( + DummyWaitingQueue(), + [request], + scheduler_requests={request.request_id: request}, + ) assert request.additional_information == payload assert request.additional_information["meta"]["finished"].item() is True @@ -601,7 +1429,7 @@ def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter) }, 8, ) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True OmniGenerationScheduler._record_prefill_stats(request) first_chunk_stats = request.prefill_stats request.prefill_stats = None @@ -613,7 +1441,7 @@ def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter) }, 8, ) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True assert isinstance(request.prefill_stats, PrefillStats) OmniGenerationScheduler._record_prefill_stats(request) second_chunk_stats = request.prefill_stats @@ -712,7 +1540,7 @@ def test_non_active_waiting_request_is_held_off_scheduler(build_adapter): assert waiting_queue == [] assert active.status == RequestStatus.WAITING_FOR_CHUNK assert non_active.status == RequestStatus.WAITING - assert list(adapter._pending_load_reqs) == [active] + assert [entry.request for entry in adapter._pending_load_reqs] == [active] assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active] @@ -863,8 +1691,19 @@ def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"): adapter.get_req_chunk[req_id] = 3 adapter.requests_with_ready_chunks.add(req_id) adapter.request_ids_mapping[req_id] = ext_id - adapter._pending_load_reqs.append(SimpleNamespace(request_id=req_id)) + entry = SimpleNamespace(request_id=req_id) + adapter._pending_load_reqs.append(entry) + adapter._registered_load_entries[req_id] = entry adapter._finished_load_reqs.add(req_id) + adapter._pending_ar_prompt_updates[req_id] = ( + SimpleNamespace(request_id=req_id), + {}, + False, + False, + False, + ) + adapter._streaming_condition_lengths[req_id] = 10 + adapter._streaming_condition_seqs[req_id] = 2 adapter.put_req_chunk[ext_id] = 5 adapter.request_payload[ext_id] = {"hidden": [1, 2]} @@ -886,12 +1725,71 @@ def test_cleanup_clears_all_state(build_adapter): assert req_id not in adapter.request_ids_mapping assert req_id in adapter._cancelled_load_reqs assert req_id not in adapter._finished_load_reqs + assert req_id not in adapter._pending_ar_prompt_updates + assert req_id not in adapter._streaming_condition_lengths + assert req_id not in adapter._streaming_condition_seqs + assert req_id not in adapter._registered_load_entries assert ext_id not in adapter.put_req_chunk assert ext_id not in adapter.request_payload assert ext_id not in adapter.code_prompt_token_ids +def test_cleanup_and_reregister_discards_inflight_old_segment_chunk(build_adapter): + """A late old-segment chunk cannot overwrite a resumed registration.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-late", RequestStatus.RUNNING, external_req_id="ext-late") + request.resumable = True + entry = _dequeue_load_entry(adapter, request) + + get_started = threading.Event() + release_get = threading.Event() + + def blocking_get(*_args): + get_started.set() + assert release_get.wait(timeout=5) + return ( + { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + }, + }, + 1, + ) + + connector.get.side_effect = blocking_get + poll_results = [] + poll_thread = threading.Thread(target=lambda: poll_results.append(adapter._poll_single_request(entry))) + poll_thread.start() + assert get_started.wait(timeout=5) + + adapter.cleanup_receiver(request.request_id) + fresh_entry = _dequeue_load_entry(adapter, request) + release_get.set() + poll_thread.join(timeout=5) + + assert not poll_thread.is_alive() + assert poll_results == [True] + assert request.additional_information is None + assert request.request_id not in adapter._pending_ar_prompt_updates + assert request.request_id not in adapter._streaming_condition_seqs + assert adapter.get_req_chunk[request.request_id] == 0 + assert adapter._registered_load_entries[request.request_id] is fresh_entry + + connector.get.side_effect = None + connector.get.return_value = ( + { + "ids": {"prompt": [2]}, + "meta": {"finished": True}, + }, + 1, + ) + assert adapter._poll_single_request(fresh_entry) is True + assert adapter.get_req_chunk[request.request_id] == 1 + assert request.request_id not in adapter._registered_load_entries + + def test_cleanup_infers_external_id(build_adapter): """When external_req_id is None, cleanup should look it up from the mapping.""" adapter, _ = build_adapter(stage_id=1) @@ -950,6 +1848,67 @@ def test_cleanup_preserves_pending_save(build_adapter): assert len(adapter._pending_save_reqs) == 1 +def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter): + adapter, connector = build_adapter(stage_id=1) + request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused") + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long)) + ) + + adapter.save_async(multimodal_output=None, request=request) + stale_task = adapter._pending_save_reqs.popleft() + adapter.cleanup(request.request_id, request.external_req_id) + adapter._send_single_request(stale_task) + + connector.put.assert_not_called() + assert request.external_req_id not in adapter.put_req_chunk + + replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id) + adapter.save_async(multimodal_output=None, request=replacement) + current_task = adapter._pending_save_reqs.popleft() + adapter._send_single_request(current_task) + + connector.put.assert_called_once() + assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0" + + +def test_finish_requests_does_not_wait_for_inflight_send(build_adapter): + adapter, connector = build_adapter(stage_id=1) + first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first") + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct() + put_started = threading.Event() + release_put = threading.Event() + + def blocking_put(**kwargs): + put_started.set() + release_put.wait(timeout=2) + return True, 1, {} + + connector.put.side_effect = blocking_put + adapter.save_async(multimodal_output=None, request=first) + task = adapter._pending_save_reqs.popleft() + sender = threading.Thread(target=adapter._send_single_request, args=(task,)) + sender.start() + assert put_started.wait(timeout=1) + cleaner = threading.Thread( + target=adapter.finish_requests, + args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}), + ) + cleaner.start() + try: + cleaner.join(timeout=0.5) + assert not cleaner.is_alive() + assert task["sender_token"].cancelled + assert first.external_req_id in adapter._sender_tokens + finally: + release_put.set() + sender.join(timeout=1) + cleaner.join(timeout=1) + assert not sender.is_alive() + assert first.external_req_id not in adapter._sender_tokens + assert first.external_req_id not in adapter.put_req_chunk + + def test_cleanup_only_affects_target_request(build_adapter): """Cleanup for one request must not affect another request's state.""" adapter, _ = build_adapter(stage_id=1) @@ -979,7 +1938,7 @@ def test_cleanup_after_poll_flow(build_adapter): "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, } connector.get.return_value = (payload, 8) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) assert "req-flow" in adapter.upstream_exhausted_requests assert adapter.get_req_chunk["req-flow"] == 1 @@ -1006,12 +1965,30 @@ def test_finish_requests_restores_status(build_adapter): prior = RequestStatus.RUNNING request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK) adapter.requests_origin_status[req_id] = prior + adapter.waiting_for_chunk_running_requests.append(request) requests_map = {req_id: request} adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map) assert request.status == prior assert req_id not in adapter.requests_origin_status + assert not adapter.waiting_for_chunk_running_requests + + +def test_finish_requests_does_not_restore_stale_status_without_connector_ownership(build_adapter): + adapter, _ = build_adapter(stage_id=1) + request = _req("req-restored", RequestStatus.WAITING_FOR_STREAMING_REQ) + request.resumable = True + adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING + + adapter.finish_requests( + [request.request_id], + RequestStatus.FINISHED_ABORTED, + {request.request_id: request}, + ) + + assert request.status == RequestStatus.WAITING_FOR_STREAMING_REQ + assert request.request_id not in adapter.requests_origin_status def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter): @@ -1061,6 +2038,108 @@ def test_finish_requests_releases_active_stream_slot(build_adapter): assert waiting.status == RequestStatus.WAITING_FOR_CHUNK +@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler]) +@pytest.mark.parametrize( + ("placement", "origin_status", "initial_status", "streaming_counter_owned"), + [ + ("hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, False), + ("running", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False), + ("waiting", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False), + ("skipped", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, True), + ("skipped_streaming", RequestStatus.RUNNING, RequestStatus.WAITING_FOR_STREAMING_REQ, True), + ("skipped_hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, True), + ("skipped_non_streaming", RequestStatus.WAITING, RequestStatus.WAITING, False), + ], +) +def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity( + build_adapter, + scheduler_cls, + placement, + origin_status, + initial_status, + streaming_counter_owned, +): + adapter, _ = build_adapter(stage_id=1, active_stream_window=2) + request = _req( + "req-segment-stop", + initial_status, + external_req_id="ext-segment-stop", + ) + request.resumable = True + adapter.requests_origin_status[request.request_id] = origin_status + if placement in {"hidden", "skipped_hidden"}: + adapter.waiting_for_chunk_running_requests.append(request) + adapter._active_streams[request.request_id] = request + + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.chunk_transfer_adapter = adapter + scheduler.input_coordinator = None + scheduler.requests = {request.request_id: request} + scheduler.running = [request] if placement == "running" else [] + scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else []) + scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else []) + scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned) + + freed = [] + + def fake_free_request(self, live, delay_free_blocks=False): + del delay_free_blocks + freed.append(live.request_id) + self.requests.pop(live.request_id) + return None, None + + scheduler._free_request = MethodType(fake_free_request, scheduler) + + first = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED) + second = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED) + + assert len(first) == 1 + assert second == [] + assert freed == [request.request_id] + assert request.request_id not in scheduler.requests + assert request.request_id not in adapter._active_streams + assert not adapter.waiting_for_chunk_running_requests + assert scheduler.running == [] + assert list(scheduler.waiting) == [] + assert list(scheduler.skipped_waiting) == [] + assert scheduler.num_waiting_for_streaming_input == 0 + + fresh_a = _req("req-fresh-a", RequestStatus.WAITING) + fresh_b = _req("req-fresh-b", RequestStatus.WAITING) + assert adapter._ensure_active_stream(fresh_a) + assert adapter._ensure_active_stream(fresh_b) + assert set(adapter._active_streams) == {fresh_a.request_id, fresh_b.request_id} + + +@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler]) +def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapter, scheduler_cls): + adapter, _ = build_adapter(stage_id=1, active_stream_window=2) + request = _req("req-off-queue", RequestStatus.FINISHED_STOPPED) + request.resumable = True + adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING + + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.chunk_transfer_adapter = adapter + scheduler.input_coordinator = None + scheduler.requests = {request.request_id: request} + scheduler.running = [] + scheduler.waiting = DummyWaitingQueue() + scheduler.skipped_waiting = DummyWaitingQueue() + scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed") + + assert ( + scheduler_cls.finish_requests( + scheduler, + [request.request_id], + RequestStatus.FINISHED_ABORTED, + ) + == [] + ) + assert scheduler.requests == {request.request_id: request} + assert request.status == RequestStatus.FINISHED_STOPPED + assert request.request_id not in adapter.requests_origin_status + + def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter): adapter, _ = build_adapter(stage_id=1) zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK) @@ -1128,6 +2207,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc scheduler.structured_output_manager.should_advance.return_value = False scheduler.finished_req_ids_dict = {} scheduler.kv_cache_manager.take_events.return_value = None + scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0 scheduler.kv_event_publisher = mocker.MagicMock() request = _HashableRequest( @@ -1208,6 +2288,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF scheduler.finished_req_ids_dict = {} scheduler.kv_cache_manager = mocker.MagicMock() scheduler.kv_cache_manager.take_events.return_value = None + scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0 scheduler.kv_event_publisher = mocker.MagicMock() scheduler.waiting_for_transfer_free = set() scheduler.transfer_triggered_requests = set() @@ -1363,6 +2444,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs): scheduler.structured_output_manager.should_advance.return_value = False scheduler.finished_req_ids_dict = {} scheduler.kv_cache_manager.take_events.return_value = None + scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0 scheduler.kv_event_publisher = mocker.MagicMock() scheduler._pending_finish_reqs = list(pending_finish_reqs) @@ -1728,3 +2810,183 @@ def test_purge_is_noop_on_empty_deques(build_adapter): adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={}) assert running_queue == [] assert waiting_queue == [] + + +# --------------------------------------------------------------------------- # +# Chunk-wait deadline (RFC #4855 R1.1, issue #3833) +# +# Before this, the async-chunk path had no deadline of any kind, so a dropped +# terminal chunk or an upstream stage that died mid-stream parked the request +# in WAITING_FOR_CHUNK forever. None of these scenarios were covered. +# --------------------------------------------------------------------------- # + + +def _park_in_chunk_wait(adapter, request, *, waiting=True): + """Drive a request through one process_pending_chunks round into the wait.""" + queue = DummyWaitingQueue([request]) if waiting else [request] + if waiting: + adapter.process_pending_chunks(queue, []) + else: + adapter.process_pending_chunks(DummyWaitingQueue(), queue) + return queue + + +def test_chunk_wait_clock_starts_when_a_request_parks(build_adapter): + adapter, _ = build_adapter(stage_id=1, model_mode="ar") + request = _req("r1", RequestStatus.WAITING) + + _park_in_chunk_wait(adapter, request) + + assert request.status == RequestStatus.WAITING_FOR_CHUNK + assert "r1" in adapter._waiting_since + + +def test_chunk_wait_does_not_expire_before_the_deadline(build_adapter): + adapter, _ = build_adapter(stage_id=1, model_mode="ar") + _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING)) + + assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set() + assert "r1" in adapter._waiting_since + + +def test_dropped_terminal_chunk_expires_the_request(build_adapter): + """The #3833 scenario: upstream stops sending and never marks the stream done.""" + adapter, _ = build_adapter(stage_id=1, model_mode="ar") + _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING)) + + adapter._waiting_since["r1"] -= 601.0 + + assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"r1"} + # Cleared, so a second sweep does not re-report the same request. + assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set() + + +def test_arriving_chunk_resets_the_clock(build_adapter): + """A slow but healthy stream must never expire: the deadline measures stall + time between chunks, not the lifetime of the stream.""" + adapter, _ = build_adapter(stage_id=1, model_mode="ar") + request = _req("r1", RequestStatus.WAITING) + queue = _park_in_chunk_wait(adapter, request) + adapter.restore_queues(queue, [], scheduler_requests={"r1": request}) + + # Age the wait to just under the deadline, then deliver a chunk. + adapter._waiting_since["r1"] -= 599.0 + adapter._finished_load_reqs.add("r1") + adapter.process_pending_chunks(queue, []) + + assert "r1" not in adapter._waiting_since + assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set() + + +def test_a_disabled_deadline_never_expires(build_adapter): + adapter, _ = build_adapter(stage_id=1, model_mode="ar") + _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING)) + adapter._waiting_since["r1"] -= 10_000.0 + + assert adapter.collect_timed_out_request_ids(timeout_s=0.0) == set() + assert adapter.collect_timed_out_request_ids(timeout_s=-1.0) == set() + + +def test_finished_request_leaves_no_stale_timestamp(build_adapter): + """Otherwise a completed request would be reported as timed out later.""" + adapter, _ = build_adapter(stage_id=1, model_mode="ar") + request = _req("r1", RequestStatus.WAITING) + _park_in_chunk_wait(adapter, request) + + adapter.finish_requests(["r1"], RequestStatus.FINISHED_STOPPED, {"r1": request}) + + assert "r1" not in adapter._waiting_since + assert adapter.collect_timed_out_request_ids(timeout_s=0.001) == set() + + +def test_aborted_request_leaves_no_stale_timestamp(build_adapter): + adapter, _ = build_adapter(stage_id=1, model_mode="ar") + _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING)) + + adapter.cleanup_receiver("r1") + + assert "r1" not in adapter._waiting_since + + +def test_expiry_is_per_request(build_adapter): + """One stalled stream must not take down its healthy neighbours.""" + adapter, _ = build_adapter(stage_id=1, model_mode="ar", max_num_seqs=4) + stalled = _req("stalled", RequestStatus.WAITING) + healthy = _req("healthy", RequestStatus.WAITING) + adapter.process_pending_chunks(DummyWaitingQueue([stalled, healthy]), []) + + adapter._waiting_since["stalled"] -= 601.0 + + assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"} + assert "healthy" in adapter._waiting_since + + +def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter): + from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import ( + talker2code2wav_async_chunk, + ) + + adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1}) + external_req_id = "ext-native-abort" + first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id) + first.resumable = True + first.model_intermediate_buffer = None + first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1} + first_frame = torch.full((1, 31), 101, dtype=torch.long) + + talker2code2wav_async_chunk( + adapter, + {"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}}, + first, + is_finished=False, + ) + + adapter.requests_num_chunks_sent[external_req_id] = 1 + adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}) + + assert external_req_id not in adapter.request_payload + assert external_req_id not in adapter.requests_num_chunks_sent + + replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id) + replacement.resumable = True + replacement.model_intermediate_buffer = None + replacement.additional_information = first.additional_information + replacement_payload = talker2code2wav_async_chunk( + adapter, + {"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}}, + replacement, + is_finished=False, + ) + + assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long)) + + +def test_save_async_boundary_holds_generation_without_request_counter(build_adapter): + """A boundary must not advance past a generation nothing can catch up to. + + When the producing stage maintains ``_omni_segment_generation`` on the + request, a boundary may start the next segment's watermark at + ``generation + 1`` because the following frame will carry it. On paths + where the request has no such counter, that speculative advance makes every + subsequent frame look stale and starves the stream permanently + (vllm-project/vllm-omni#6816). + """ + adapter, _ = build_adapter(stage_id=1) + + request = _req("req-nogen", RequestStatus.WAITING, external_req_id="ext-nogen") + request.resumable = True + request._omni_segment_generation = None + + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True) + assert adapter._segment_generation["ext-nogen"] == 0 + + follow_up = _req("req-nogen-next", RequestStatus.WAITING, external_req_id="ext-nogen") + follow_up.resumable = True + follow_up.num_computed_tokens = 3 + follow_up._omni_segment_generation = None + queued_before = len(adapter._pending_save_reqs) + + adapter.save_async(multimodal_output=None, request=follow_up, is_segment_finished=False) + + assert len(adapter._pending_save_reqs) == queued_before + 1 + assert adapter._pending_save_reqs[-1]["request"] is follow_up diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py index 60653b3cd..66f0327f7 100755 --- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py +++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Integration tests for MooncakeTransferEngineConnector. @@ -18,6 +18,7 @@ import pytest import torch from tests.helpers.mark import hardware_test +from tests.helpers.runtime import get_open_port from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( ManagedBuffer, MooncakeTransferEngineConnector, @@ -164,14 +165,6 @@ RDMA_HOST = get_rdma_host() RDMA_DEVICE = _detect_rdma_device() -def _free_port() -> int: - import socket - - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind((RDMA_HOST, 0)) - return s.getsockname()[1] - - def _connector_config( zmq_port: int, pool_size: int = 16 * 1024 * 1024, @@ -204,7 +197,7 @@ class TestBasicConnector: """Verify connector initialization, put, cleanup, and health check.""" def test_initialization(self): - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c: assert c.rpc_port != 0 assert c.pool_size == 1024 * 1024 @@ -214,7 +207,7 @@ class TestBasicConnector: def test_put_tensor_bytes_object(self): """Put tensor / bytes / dict and verify metadata.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port)) as c: ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100)) assert ok @@ -229,7 +222,7 @@ class TestBasicConnector: assert not meta["is_fast_path"] def test_cleanup_releases_buffer(self): - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port)) as c: c.put("s0", "s1", "r1", torch.randn(100)) key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1") @@ -239,7 +232,7 @@ class TestBasicConnector: def test_pool_exhaustion_and_recovery(self): """Fill pool, verify failure, free, verify recovery.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024)) as c: ids = [] for i in range(10): @@ -264,8 +257,8 @@ class TestEndToEnd: """E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types.""" def _pair(self, pool_size=16 * 1024 * 1024): - p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) - c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) + p = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) return p, c def test_tensor_e2e(self): @@ -401,7 +394,7 @@ class TestEndToEnd: def test_concurrent_put(self): """10 concurrent puts should all succeed.""" - port = _free_port() + port = get_open_port(RDMA_HOST) conn = MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024 * 1024)) errors: list[str] = [] lock = threading.Lock() @@ -457,20 +450,20 @@ class TestLifecycle: """Close, context manager, double-close safety.""" def test_close_releases_resources(self): - c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024)) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size=1024 * 1024)) c.put("s0", "s1", "x", torch.randn(100)) c.close() assert c._stop_event.is_set() assert len(c._local_buffers) == 0 def test_context_manager(self): - with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c: + with MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST))) as c: ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50)) assert ok assert c._stop_event.is_set() def test_double_close_safe(self): - c = MooncakeTransferEngineConnector(_connector_config(_free_port())) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST))) c.close() c.close() @@ -489,12 +482,12 @@ class TestGPUPool: return _connector_config(port, pool_size, pool_device="cuda:0") def test_gpu_pool_init(self): - with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c: + with MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) as c: assert c.pool_device == "cuda:0" assert c.pool.is_cuda def test_gpu_pool_put_cpu_and_gpu_tensor(self): - with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c: + with MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) as c: ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256)) assert ok assert meta["is_fast_path"] @@ -504,8 +497,8 @@ class TestGPUPool: assert meta["is_fast_path"] def test_gpu_e2e_transfer(self): - p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) - c = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) + p = MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) + c = MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) try: orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0") ok, _, meta = p.put("s0", "s1", "ge", orig) @@ -536,8 +529,8 @@ class TestStressCorrectness: """ def _pair(self, pool_size=64 * 1024 * 1024): - p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) - c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size)) + p = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) + c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size)) return p, c # -- Concurrent put + get with data integrity -- @@ -712,7 +705,7 @@ class TestStressCorrectness: def test_empty_bytes_rejected(self): """Connector should gracefully reject empty bytes payload.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c: ok, sz, meta = c.put("s0", "s1", "empty_b", b"") assert not ok, "Empty bytes should be rejected by connector" @@ -744,7 +737,7 @@ class TestStressCorrectness: def test_rapid_alloc_free_cycle(self): """Put + cleanup in tight loop to stress allocator under real connector.""" - port = _free_port() + port = get_open_port(RDMA_HOST) with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c: for i in range(50): rid = f"cycle_{i}" diff --git a/tests/docs/test_example_output_accessors.py b/tests/docs/test_example_output_accessors.py new file mode 100644 index 000000000..e7a618890 --- /dev/null +++ b/tests/docs/test_example_output_accessors.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Guard the public accessor that examples and docs use to read omni outputs. + +``OmniRequestOutput`` used to hold a nested ``RequestOutput`` in a +``request_output`` field. Since #5146 it *inherits* ``RequestOutput``, so the +generation content lives on the object itself and the nested field is gone. + +Call sites that kept the old accessor only fail once a real model runs, which +means a GPU nightly job is the first thing to notice. These checks pin both +halves of the contract on CPU instead: the field is really gone, and no doc, +example, test or tool still reaches for it. +""" + +import re +from pathlib import Path + +import pytest + +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +ROOT_DIR = Path(__file__).parents[2] +SCANNED_DIRS = ("docs", "examples", "tests", "tools") +SCANNED_SUFFIXES = (".py", ".md") + +# ``.request_outputs`` (plural, an OutputProcessor batch) is a different and +# still-valid attribute, so require a non-identifier character after the name. +NESTED_ACCESSOR = re.compile(r"\.request_output(?![\w])") + + +def _iter_scanned_files(): + for directory in SCANNED_DIRS: + for path in sorted((ROOT_DIR / directory).rglob("*")): + if path.suffix in SCANNED_SUFFIXES and path.resolve() != Path(__file__).resolve(): + yield path + + +def test_omni_request_output_has_no_nested_request_output(): + assert not hasattr(OmniRequestOutput(request_id="req-accessor"), "request_output") + + +def test_no_source_reads_the_nested_request_output(): + offenders = [] + for path in _iter_scanned_files(): + for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): + if NESTED_ACCESSOR.search(line): + offenders.append(f"{path.relative_to(ROOT_DIR)}:{lineno}: {line.strip()}") + + assert not offenders, ( + "OmniRequestOutput no longer nests a RequestOutput; read the content off the " + "output itself (e.g. `outputs[0].images`) instead of `.request_output`:\n" + "\n".join(offenders) + ) diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index 99880fc73..832d34557 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/tests/e2e/accuracy/helpers.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import json import os import re @@ -14,7 +17,7 @@ from hashlib import sha1 from io import BytesIO from mimetypes import guess_type from pathlib import Path -from typing import Any +from typing import IO, Any from urllib.parse import urlparse import numpy as np @@ -35,11 +38,14 @@ class IdentityFtfy: def apply_ftfy_mock(*, wan_i2v_module: Any | None = None) -> None: """Mock ftfy library for test cases in the main process""" - if wan_i2v_module is None: - from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module + module: Any = wan_i2v_module + if module is None: + from diffusers.pipelines.wan import pipeline_wan_i2v as wan_pipeline - if not hasattr(wan_i2v_module, "ftfy"): - wan_i2v_module.ftfy = IdentityFtfy() + module = wan_pipeline + + if not hasattr(module, "ftfy"): + module.ftfy = IdentityFtfy() print("ftfy (text encoding sanitizer) is not installed. Using mock ftfy implementation (identity function)") else: print("ftfy (text encoding sanitizer) is installed. Using actual ftfy implementation.") @@ -118,6 +124,7 @@ def probe_binary(binary: str, *, test_name: str = "video similarity e2e test") - resolved = shutil.which(binary) if resolved is None: pytest.skip(f"{binary} is required for {test_name}.") + assert resolved is not None return resolved @@ -235,7 +242,7 @@ def video_artifact_dir(result_root: Path, source: str) -> Path: def send_video_request_with_timeout( - openai_client, + online_client, request_config: dict[str, Any], *, timeout_seconds: int, @@ -266,7 +273,7 @@ def send_video_request_with_timeout( normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference}) start_time = time.perf_counter() - create_url = openai_client._build_url("/v1/videos") + create_url = online_client._build_url("/v1/videos") response = requests.post( create_url, data=normalized_form_data, @@ -280,8 +287,8 @@ def send_video_request_with_timeout( raise requests.HTTPError(f"{exc}; response body: {response.text}", response=response) from exc job_data = response.json() video_id = job_data["id"] - openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds) - video_content = openai_client._download_video_content(video_id) + online_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds) + video_content = online_client._download_video_content(video_id) print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}") return video_content @@ -571,11 +578,11 @@ class SemanticSimilarityScorer: raise ImportError("FlagEmbedding not installed. Install with: pip install FlagEmbedding") return self._text_model - def text_similarity(self, text1: str, text2: str) -> float: + def text_similarity(self, text1: str, text2: str) -> dict[str, float | int]: """ Compute semantic similarity between two texts using BGE-M3 """ - results = {} + results: dict[str, float | int] = {} # 1. Text prefix matching (maximum matching character length) match_count = 0 for char_vllm, char_ref in zip(text1, text2): @@ -653,7 +660,7 @@ def dreamzero_cfg_parallel_size() -> int: def _dreamzero_parity_vllm_env() -> dict[str, str]: - from tests.helpers.env import pick_least_used_device_indices + from tests.helpers.clean import pick_least_used_device_indices from tests.helpers.runtime import get_open_port parallel_size = dreamzero_cfg_parallel_size() @@ -683,8 +690,9 @@ def dreamzero_parity_vllm_params(model: str | None = None): ] if parallel_size > 1: server_args.extend(["--cfg-parallel-size", str(parallel_size)]) + resolved_model = model or os.environ.get("VLLM_DREAMZERO_MODEL") or _DREAMZERO_DEFAULT_MODEL return OmniServerParams( - model=model or os.environ.get("VLLM_DREAMZERO_MODEL", _DREAMZERO_DEFAULT_MODEL), + model=resolved_model, stage_config_path=get_deploy_config_path("dreamzero_tp1_cfg2.yaml"), server_args=server_args, env_dict=_dreamzero_parity_vllm_env(), @@ -725,9 +733,9 @@ def dreamzero_parity_vllm_client( def _dreamzero_create_openpi_vllm_client(host: str, port: int) -> Any: from openpi_client import msgpack_numpy - _, WebsocketClientPolicy = load_dreamzero_upstream_modules() + _, websocket_client_policy_cls = load_dreamzero_upstream_modules() - class _OpenPIWebsocketClientPolicy(WebsocketClientPolicy): + class _OpenPIWebsocketClientPolicy(websocket_client_policy_cls): # type: ignore[misc, valid-type] def __init__( self, host: str = "127.0.0.1", @@ -869,7 +877,7 @@ class _DreamZeroParityServer(ABC): self.port = get_open_port() if port is None else port self.ready_timeout_s = ready_timeout_s self.proc: subprocess.Popen[str] | None = None - self._log_file = None + self._log_file: IO[str] | None = None @abstractmethod def _build_argv(self) -> list[str]: @@ -932,7 +940,7 @@ class DreamZeroUpstreamServer(_DreamZeroParityServer): ] def _build_env(self) -> dict[str, str]: - from tests.helpers.env import pick_least_used_device_indices + from tests.helpers.clean import pick_least_used_device_indices dreamzero_repo, _ = require_dreamzero_parity_env() env = os.environ.copy() diff --git a/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py index b830cdc6c..eb8f2d769 100644 --- a/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py +++ b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os @@ -149,7 +152,7 @@ def _generate_offline_video(*, image_source: str) -> Path: def _generate_online_video( *, omni_server, - openai_client, + online_client, image_source: str, timeout_seconds: int, ) -> Path: @@ -171,7 +174,7 @@ def _generate_online_video( "image_reference": build_online_image_reference(online_image_source), } online_video_bytes = send_video_request_with_timeout( - openai_client, + online_client, request_config, timeout_seconds=timeout_seconds, ) @@ -204,7 +207,7 @@ def test_hunyuanvideo15_i2v_diffusers_offline_generates_video( @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_hunyuanvideo15_i2v_online_serving_generates_video( omni_server, - openai_client, + online_client, hunyuanvideo15_i2v_image_source: str | None, hunyuanvideo15_online_timeout_seconds: int, ) -> None: @@ -216,7 +219,7 @@ def test_hunyuanvideo15_i2v_online_serving_generates_video( validate_image_source(image_source) online_path = _generate_online_video( omni_server=omni_server, - openai_client=openai_client, + online_client=online_client, image_source=image_source, timeout_seconds=hunyuanvideo15_online_timeout_seconds, ) diff --git a/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py index f2c14e2b0..f5b90822d 100644 --- a/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py +++ b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os @@ -115,7 +118,7 @@ def _generate_offline_video() -> Path: return offline_path -def _generate_online_video(*, omni_server, openai_client, timeout_seconds: int) -> Path: +def _generate_online_video(*, omni_server, online_client, timeout_seconds: int) -> Path: online_path, _ = _artifact_paths() request_config = { "model": omni_server.model, @@ -132,7 +135,7 @@ def _generate_online_video(*, omni_server, openai_client, timeout_seconds: int) }, } online_video_bytes = send_video_request_with_timeout( - openai_client, + online_client, request_config, timeout_seconds=timeout_seconds, ) @@ -161,7 +164,7 @@ def test_hunyuanvideo15_t2v_diffusers_offline_generates_video() -> None: @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_hunyuanvideo15_t2v_online_serving_generates_video( omni_server, - openai_client, + online_client, hunyuanvideo15_online_timeout_seconds: int, ) -> None: if not torch.cuda.is_available(): @@ -170,7 +173,7 @@ def test_hunyuanvideo15_t2v_online_serving_generates_video( probe_binary("ffprobe") online_path = _generate_online_video( omni_server=omni_server, - openai_client=openai_client, + online_client=online_client, timeout_seconds=hunyuanvideo15_online_timeout_seconds, ) assert online_path.exists(), f"Expected online video artifact at {online_path}" diff --git a/tests/e2e/accuracy/ltx/run_ltx25_reference.py b/tests/e2e/accuracy/ltx/run_ltx25_reference.py new file mode 100644 index 000000000..cbc8917a8 --- /dev/null +++ b/tests/e2e/accuracy/ltx/run_ltx25_reference.py @@ -0,0 +1,591 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate raw LTX video and audio outputs with the official or Omni runtime.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--backend", choices=("official", "omni"), required=True) + parser.add_argument("--request", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model") + parser.add_argument("--model-class-name") + parser.add_argument("--official-root", type=Path) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--gemma-root", type=Path) + parser.add_argument("--transformer-path", type=Path) + parser.add_argument("--text-encoder-path", type=Path) + parser.add_argument("--video-vae-path", type=Path) + parser.add_argument("--audio-vae-path", type=Path) + parser.add_argument("--spatial-upsampler-path", type=Path) + parser.add_argument("--distilled-lora-path", type=Path) + parser.add_argument("--connector-model", type=Path) + parser.add_argument( + "--official-pipeline", + choices=("distilled_two_stage", "full_one_stage", "full_two_stage"), + default="distilled_two_stage", + ) + parser.add_argument("--enable-layerwise-offload", action="store_true") + return parser.parse_args() + + +def _save_outputs( + output_dir: Path, + *, + video: torch.Tensor, + audio: torch.Tensor, + audio_sample_rate: int, + fps: float, + metadata: dict[str, Any], +) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + video = video.detach().float().cpu().clamp(0.0, 1.0) + audio = audio.detach().float().cpu() + np.save(output_dir / "video.npy", video.numpy()) + np.save(output_dir / "audio.npy", audio.numpy()) + + from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes + + frames_u8 = video.mul(255.0).round().to(torch.uint8).numpy() + audio_array = audio.numpy() + while audio_array.ndim > 2 and audio_array.shape[0] == 1: + audio_array = audio_array[0] + (output_dir / "output.mp4").write_bytes( + mux_video_audio_bytes( + frames_u8, + audio_array, + fps=float(fps), + audio_sample_rate=int(audio_sample_rate), + video_codec_options={"preset": "ultrafast"}, + ) + ) + + frame_indices = sorted({0, video.shape[0] // 2, video.shape[0] - 1}) + for index in frame_indices: + frame = video[index].mul(255.0).round().to(torch.uint8).numpy() + Image.fromarray(frame).save(output_dir / f"frame_{index:04d}.png") + + metadata.update( + { + "video_shape": list(video.shape), + "audio_shape": list(audio.shape), + "audio_sample_rate": int(audio_sample_rate), + "frame_indices": frame_indices, + "mp4": "output.mp4", + } + ) + (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n") + + +def _reset_accelerator_peak_memory_stats() -> None: + if torch.accelerator.is_available(): + torch.accelerator.reset_peak_memory_stats() + + +def _accelerator_peak_memory_stats() -> dict[str, float]: + if not torch.accelerator.is_available(): + return {} + torch.accelerator.synchronize() + mib = 1024**2 + return { + "peak_memory_allocated_mb": torch.accelerator.max_memory_allocated() / mib, + "peak_memory_reserved_mb": torch.accelerator.max_memory_reserved() / mib, + } + + +def _insert_official_paths(official_root: Path) -> None: + for relative_path in ("packages/ltx-core/src", "packages/ltx-pipelines/src"): + path = str((official_root / relative_path).resolve()) + if path not in sys.path: + sys.path.insert(0, path) + + +def _configure_official_sdpa(pipeline: Any) -> None: + """Pin official connector and denoiser attention to cuDNN, with MATH fallback.""" + from ltx_core.loader.attention_ops import set_attention_module_op + from ltx_core.model.transformer.attention import PytorchAttention + from torch.nn.attention import SDPBackend + + attention = PytorchAttention(priority=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]) + module_op = set_attention_module_op( + attention=attention, + masked_attention=attention, + ) + transformer_owners = [ + owner + for attribute in ("stage", "stage_1", "stage_2") + if (owner := getattr(pipeline, attribute, None)) is not None + ] + owners_and_attributes = [ + *((owner, "_transformer_builder") for owner in transformer_owners), + (pipeline.prompt_encoder, "_embeddings_processor_builder"), + ] + for owner, attribute in owners_and_attributes: + builder = getattr(owner, attribute) + setattr( + owner, + attribute, + builder.with_module_ops((*builder.module_ops, module_op)), + ) + + +def _use_diffusers_connector_weights(prompt_encoder: Any, model: Path) -> None: + """Run official connector code with the checkpoint under test.""" + from safetensors import safe_open + + connector_root = model / "connectors" + shards = sorted(connector_root.glob("*.safetensors")) + if not shards: + raise ValueError(f"No connector safetensors found under {connector_root}") + + original_build_embeddings_processor = prompt_encoder._build_embeddings_processor + + def build_embeddings_processor(): + processor = original_build_embeddings_processor() + parameters = dict(processor.named_parameters()) + expected = {name for name in parameters if name.startswith(("video_connector.", "audio_connector."))} + loaded: set[str] = set() + with torch.no_grad(): + for shard in shards: + with safe_open(str(shard), framework="pt", device="cpu") as weights: + for source_name in weights.keys(): + if not source_name.startswith(("video_connector.", "audio_connector.")): + continue + target_name = ( + source_name.replace(".transformer_blocks.", ".transformer_1d_blocks.") + .replace(".attn1.norm_q.", ".attn1.q_norm.") + .replace(".attn1.norm_k.", ".attn1.k_norm.") + ) + parameter = parameters[target_name] + parameter.copy_(weights.get_tensor(source_name).to(parameter.device, parameter.dtype)) + loaded.add(target_name) + if loaded != expected: + missing = sorted(expected - loaded) + unexpected = sorted(loaded - expected) + raise ValueError(f"Connector weight mismatch: missing={missing}, unexpected={unexpected}") + return processor + + prompt_encoder._build_embeddings_processor = build_embeddings_processor + + +@torch.inference_mode() +def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: + if args.official_root is None: + raise ValueError("Official backend requires --official-root") + _insert_official_paths(args.official_root) + # Keep Gemma on the same SDPA path in both subprocesses. DiT attention is + # still pinned independently to the explicit cuDNN backend below. + torch.backends.cuda.enable_cudnn_sdp(False) + _reset_accelerator_peak_memory_stats() + + from ltx_pipelines.utils.args import ImageConditioningInput + from ltx_pipelines.utils.model_paths import ModelPaths + from ltx_pipelines.utils.types import OffloadMode + + offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE + image_path = request.get("image") + images = ( + [] + if image_path is None + else [ + ImageConditioningInput( + path=str(image_path), + frame_idx=0, + strength=1.0, + crf=request.get("image_crf", 18), + ) + ] + ) + + if args.transformer_path is not None: + required = { + "--text-encoder-path": args.text_encoder_path, + "--video-vae-path": args.video_vae_path, + "--audio-vae-path": args.audio_vae_path, + } + if args.official_pipeline in {"distilled_two_stage", "full_two_stage"}: + required["--spatial-upsampler-path"] = args.spatial_upsampler_path + if args.official_pipeline == "full_two_stage": + required["--distilled-lora-path"] = args.distilled_lora_path + missing = [name for name, value in required.items() if value is None] + if missing: + raise ValueError(f"Official LTX-2.5 split backend is missing: {', '.join(missing)}") + + model_paths = ModelPaths.from_split( + transformer_path=str(args.transformer_path), + text_encoder_path=str(args.text_encoder_path), + video_vae_path=str(args.video_vae_path), + audio_vae_path=str(args.audio_vae_path), + ) + if args.official_pipeline == "distilled_two_stage": + from ltx_pipelines.distilled import DistilledPipeline + + pipeline = DistilledPipeline( + model_paths=model_paths, + spatial_upsampler_path=str(args.spatial_upsampler_path), + loras=(), + offload_mode=offload_mode, + ) + if args.connector_model is not None: + _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model) + _configure_official_sdpa(pipeline) + video, audio, _, _ = pipeline( + prompt=request["prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + images=images, + ) + pipeline_name = "DistilledPipeline" + elif args.official_pipeline == "full_two_stage": + from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps + from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline + + pipeline = TI2VidTwoStagesPipeline( + model_paths=model_paths, + distilled_lora=[ + LoraPathStrengthAndSDOps( + path=str(args.distilled_lora_path), + strength=1.0, + sd_ops=LTXV_LORA_COMFY_RENAMING_MAP, + ) + ], + spatial_upsampler_path=str(args.spatial_upsampler_path), + loras=[], + offload_mode=offload_mode, + ) + if args.connector_model is not None: + _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model) + _configure_official_sdpa(pipeline) + video, audio, _, _ = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + stage_1_sigmas=torch.tensor(request["stage_1_sigmas"], dtype=torch.float32), + stage_2_sigmas=torch.tensor(request["stage_2_sigmas"], dtype=torch.float32), + ) + pipeline_name = "TI2VidTwoStagesPipeline" + else: + from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + + pipeline = TI2VidOneStagePipeline( + model_paths=model_paths, + loras=(), + offload_mode=offload_mode, + ) + if args.connector_model is not None: + _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model) + _configure_official_sdpa(pipeline) + request_sigmas = request.get("sigmas") + video, audio, _ = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + sigmas=None if request_sigmas is None else torch.tensor(request_sigmas, dtype=torch.float32), + ) + pipeline_name = "TI2VidOneStagePipeline" + checkpoint_metadata = { + "transformer_path": str(args.transformer_path), + "video_vae_path": str(args.video_vae_path), + "connector_model": None if args.connector_model is None else str(args.connector_model), + "pipeline": pipeline_name, + } + else: + if args.checkpoint is None or args.gemma_root is None: + raise ValueError("Official monolith backend requires --checkpoint and --gemma-root") + + from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + + pipeline = TI2VidOneStagePipeline( + model_paths=ModelPaths.from_monolith(str(args.checkpoint), str(args.gemma_root)), + loras=(), + offload_mode=offload_mode, + ) + _configure_official_sdpa(pipeline) + video, audio, _ = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + ) + checkpoint_metadata = { + "checkpoint": str(args.checkpoint), + "pipeline": "TI2VidOneStagePipeline", + } + + video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0) + _save_outputs( + args.output_dir, + video=_canonical_video(video_tensor), + audio=audio.waveform, + audio_sample_rate=audio.sampling_rate, + fps=float(request["fps"]), + metadata={ + "backend": "official", + "attention_backend": "torch_sdpa_cudnn_bf16", + "official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"), + "seed": request["seed"], + "sigmas": request.get("sigmas"), + "stage_1_sigmas": request.get("stage_1_sigmas"), + "stage_2_sigmas": request.get("stage_2_sigmas"), + **_accelerator_peak_memory_stats(), + **checkpoint_metadata, + }, + ) + + +def _omni_worker_peak_memory_mb(output: Any) -> float: + result = output[0] if isinstance(output, list) and output else output + return float(getattr(result, "peak_memory_mb", 0.0) or 0.0) + + +def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]: + from vllm_omni.outputs import OmniRequestOutput + + audio = None + audio_sample_rate = None + frames = output[0] if isinstance(output, list) and output else output + if isinstance(frames, OmniRequestOutput): + multimodal_output = frames.multimodal_output or {} + audio = multimodal_output.get("audio") + audio_sample_rate = multimodal_output.get("audio_sample_rate") + nested_output = getattr(frames, "request_output", None) + if frames.is_pipeline_output and isinstance(nested_output, OmniRequestOutput): + frames = nested_output + multimodal_output = frames.multimodal_output or {} + audio = multimodal_output.get("audio", audio) + audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate) + if isinstance(frames, OmniRequestOutput): + if not frames.images: + raise ValueError("No video frames found in OmniRequestOutput") + frames = frames.images + + if isinstance(frames, list) and len(frames) == 1: + frames = frames[0] + if isinstance(frames, tuple) and len(frames) == 2: + frames, audio = frames + elif isinstance(frames, dict): + audio = frames.get("audio", audio) + audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate) + frames = frames.get("frames", frames.get("video")) + + if frames is None or audio is None or audio_sample_rate is None: + raise ValueError("Omni output did not contain video, audio, and audio_sample_rate") + return frames, audio, int(audio_sample_rate) + + +def _canonical_video(video: Any) -> torch.Tensor: + if isinstance(video, list): + if len(video) == 1: + return _canonical_video(video[0]) + tensors = [torch.as_tensor(np.asarray(item) if not isinstance(item, torch.Tensor) else item) for item in video] + if all(item.ndim == 3 for item in tensors): + video = torch.stack(tensors) + elif all(item.ndim == 4 for item in tensors): + video = torch.cat(tensors) + else: + raise ValueError(f"Cannot combine video items with dimensions {[item.ndim for item in tensors]}") + tensor = torch.as_tensor(np.asarray(video) if not isinstance(video, torch.Tensor) else video).detach().cpu() + if tensor.ndim == 5: + tensor = tensor[0] + if tensor.ndim != 4: + raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}") + if tensor.shape[-1] in (3, 4): + tensor = tensor[..., :3] + elif tensor.shape[1] in (3, 4): + tensor = tensor[:, :3].permute(0, 2, 3, 1) + elif tensor.shape[0] in (3, 4): + tensor = tensor[:3].permute(1, 2, 3, 0) + else: + raise ValueError(f"Cannot infer video channel dimension from {tuple(tensor.shape)}") + tensor = tensor.float() + if tensor.numel() and tensor.max() > 1: + tensor = tensor / 255.0 + elif tensor.numel() and tensor.min() < 0: + tensor = tensor.clamp(-1.0, 1.0).add(1.0).mul(0.5) + return tensor.clamp(0.0, 1.0) + + +@torch.inference_mode() +def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: + if args.model is None or args.model_class_name is None: + raise ValueError("Omni backend requires --model and --model-class-name") + + # vLLM imports can otherwise change the process-wide Gemma dispatch. + torch.backends.cuda.enable_cudnn_sdp(False) + _reset_accelerator_peak_memory_stats() + + attention_config = {"default": {"backend": "CUDNN_ATTN"}} + + from vllm_omni.diffusion.data import DiffusionParallelConfig + from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.model_extras import get_extra_body_params + from vllm_omni.platforms import current_omni_platform + + generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request["seed"]) + omni = Omni( + model=args.model, + model_class_name=args.model_class_name, + enforce_eager=True, + enable_layerwise_offload=args.enable_layerwise_offload, + diffusion_attention_config=attention_config, + parallel_config=DiffusionParallelConfig(), + ) + try: + # The explicit public class argument selects the worker pipeline. The + # parent Omni proxy can only expose the checkpoint's base class, which + # is insufficient to distinguish LTX execution modes. + model_class_name = args.model_class_name + sampling_params = OmniDiffusionSamplingParams( + height=request["height"], + width=request["width"], + generator=generator, + guidance_scale=None, + num_inference_steps=request["num_inference_steps"], + num_frames=request["num_frames"], + fps=request["fps"], + frame_rate=float(request["fps"]), + output_type="np", + ) + extra_args = { + key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_") + } + for key in ("sigmas", "stage_1_sigmas", "stage_2_sigmas", "image_crf"): + if key in request: + extra_args[key] = request[key] + apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), extra_args) + prompt: dict[str, Any] = {"prompt": request["prompt"]} + if "negative_prompt" in request: + prompt["negative_prompt"] = request["negative_prompt"] + image_path = request.get("image") + if image_path is not None: + with Image.open(str(image_path)) as source_image: + image = source_image.convert("RGB") + image.load() + prompt["multi_modal_data"] = {"image": image} + output = omni.generate(prompt, sampling_params) + video, audio, audio_sample_rate = _unwrap_omni_output(output) + audio_tensor = torch.as_tensor(np.asarray(audio) if not isinstance(audio, torch.Tensor) else audio) + _save_outputs( + args.output_dir, + video=_canonical_video(video), + audio=audio_tensor, + audio_sample_rate=audio_sample_rate, + fps=float(request["fps"]), + metadata={ + "backend": "omni", + "attention_backend": "torch_sdpa_cudnn_bf16", + "seed": request["seed"], + "sigmas": request.get("sigmas"), + "model": args.model, + "stage_1_sigmas": request.get("stage_1_sigmas"), + "stage_2_sigmas": request.get("stage_2_sigmas"), + "model_class_name": model_class_name, + "worker_peak_memory_mb": _omni_worker_peak_memory_mb(output), + **_accelerator_peak_memory_stats(), + }, + ) + finally: + omni.shutdown() + + +def main() -> None: + args = _parse_args() + request = json.loads(args.request.read_text()) + if args.backend == "official": + _run_official(args, request) + else: + _run_omni(args, request) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py index a4a180047..ccabd4596 100644 --- a/tests/e2e/accuracy/ltx/run_ltx_reference.py +++ b/tests/e2e/accuracy/ltx/run_ltx_reference.py @@ -20,6 +20,11 @@ from PIL import Image def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--backend", choices=("official", "omni"), required=True) + parser.add_argument( + "--pipeline-kind", + choices=("one_stage", "distilled", "two_stage"), + default="one_stage", + ) parser.add_argument("--request", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--model") @@ -27,6 +32,8 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--official-root", type=Path) parser.add_argument("--checkpoint", type=Path) parser.add_argument("--gemma-root", type=Path) + parser.add_argument("--spatial-upsampler", type=Path) + parser.add_argument("--distilled-lora", type=Path) parser.add_argument("--enable-layerwise-offload", action="store_true") return parser.parse_args() @@ -91,9 +98,13 @@ def _configure_official_sdpa(pipeline: Any) -> None: attention=attention, masked_attention=attention, ) - owners_and_attributes = ( - (pipeline.stage, "_transformer_builder"), + owners_and_attributes = [ (pipeline.prompt_encoder, "_embeddings_processor_builder"), + ] + owners_and_attributes.extend( + (stage, "_transformer_builder") + for name in ("stage", "stage_1", "stage_2") + if (stage := getattr(pipeline, name, None)) is not None ) for owner, attribute in owners_and_attributes: builder = getattr(owner, attribute) @@ -104,6 +115,12 @@ def _configure_official_sdpa(pipeline: Any) -> None: ) +def _require_path(path: Path | None, description: str) -> str: + if path is None or not path.is_file(): + raise ValueError(f"Official {description} is required and must exist: {path}") + return str(path) + + @torch.inference_mode() def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: if args.official_root is None or args.checkpoint is None or args.gemma_root is None: @@ -113,16 +130,54 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: torch.backends.cuda.enable_cudnn_sdp(False) from ltx_core.components.guiders import MultiModalGuiderParams + from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps + from ltx_pipelines.distilled import DistilledPipeline from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline + from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline from ltx_pipelines.utils.args import ImageConditioningInput from ltx_pipelines.utils.types import OffloadMode - pipeline = TI2VidOneStagePipeline( - checkpoint_path=str(args.checkpoint), - gemma_root=str(args.gemma_root), - loras=(), - offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE, - ) + if request.get("pipeline_kind", "one_stage") != args.pipeline_kind: + raise ValueError( + f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}" + ) + checkpoint = _require_path(args.checkpoint, "checkpoint") + gemma_root = str(args.gemma_root) + if not Path(gemma_root).is_dir(): + raise ValueError(f"Official Gemma root is required and must be a directory: {gemma_root}") + offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE + pipeline: Any + if args.pipeline_kind == "one_stage": + pipeline = TI2VidOneStagePipeline( + checkpoint_path=checkpoint, + gemma_root=gemma_root, + loras=(), + offload_mode=offload_mode, + ) + elif args.pipeline_kind == "distilled": + pipeline = DistilledPipeline( + distilled_checkpoint_path=checkpoint, + spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"), + gemma_root=gemma_root, + loras=(), + offload_mode=offload_mode, + ) + else: + distilled_lora = [ + LoraPathStrengthAndSDOps( + path=_require_path(args.distilled_lora, "distilled LoRA"), + strength=1.0, + sd_ops=LTXV_LORA_COMFY_RENAMING_MAP, + ) + ] + pipeline = TI2VidTwoStagesPipeline( + checkpoint_path=checkpoint, + distilled_lora=distilled_lora, + spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"), + gemma_root=gemma_root, + loras=(), + offload_mode=offload_mode, + ) _configure_official_sdpa(pipeline) image_path = request.get("image") images = ( @@ -140,34 +195,45 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: ) ] ) - video, audio = pipeline( - prompt=request["prompt"], - negative_prompt=request["negative_prompt"], - seed=request["seed"], - height=request["height"], - width=request["width"], - num_frames=request["num_frames"], - frame_rate=request["fps"], - num_inference_steps=request["num_inference_steps"], - video_guider_params=MultiModalGuiderParams( - cfg_scale=request["video_cfg_scale"], - stg_scale=request["video_stg_scale"], - rescale_scale=request["video_rescale_scale"], - modality_scale=request["video_modality_scale"], - skip_step=0, - stg_blocks=request["video_stg_blocks"], - ), - audio_guider_params=MultiModalGuiderParams( - cfg_scale=request["audio_cfg_scale"], - stg_scale=request["audio_stg_scale"], - rescale_scale=request["audio_rescale_scale"], - modality_scale=request["audio_modality_scale"], - skip_step=0, - stg_blocks=request["audio_stg_blocks"], - ), - images=images, - max_batch_size=4, - ) + if args.pipeline_kind == "distilled": + video, audio = pipeline( + prompt=request["prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + images=images, + ) + else: + video, audio = pipeline( + prompt=request["prompt"], + negative_prompt=request["negative_prompt"], + seed=request["seed"], + height=request["height"], + width=request["width"], + num_frames=request["num_frames"], + frame_rate=request["fps"], + num_inference_steps=request["num_inference_steps"], + video_guider_params=MultiModalGuiderParams( + cfg_scale=request["video_cfg_scale"], + stg_scale=request["video_stg_scale"], + rescale_scale=request["video_rescale_scale"], + modality_scale=request["video_modality_scale"], + skip_step=0, + stg_blocks=request["video_stg_blocks"], + ), + audio_guider_params=MultiModalGuiderParams( + cfg_scale=request["audio_cfg_scale"], + stg_scale=request["audio_stg_scale"], + rescale_scale=request["audio_rescale_scale"], + modality_scale=request["audio_modality_scale"], + skip_step=0, + stg_blocks=request["audio_stg_blocks"], + ), + images=images, + max_batch_size=4, + ) video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0) _save_outputs( args.output_dir, @@ -179,6 +245,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: "attention_backend": "torch_sdpa", "official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"), "checkpoint": str(args.checkpoint), + "pipeline_kind": args.pipeline_kind, }, ) @@ -193,15 +260,9 @@ def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]: multimodal_output = frames.multimodal_output or {} audio = multimodal_output.get("audio") audio_sample_rate = multimodal_output.get("audio_sample_rate") - if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput): - frames = frames.request_output - multimodal_output = frames.multimodal_output or {} - audio = multimodal_output.get("audio", audio) - audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate) - if isinstance(frames, OmniRequestOutput): - if not frames.images: - raise ValueError("No video frames found in OmniRequestOutput") - frames = frames.images + if not frames.images: + raise ValueError("No video frames found in OmniRequestOutput") + frames = frames.images if isinstance(frames, list) and len(frames) == 1: frames = frames[0] @@ -253,6 +314,10 @@ def _canonical_video(video: Any) -> torch.Tensor: def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: if args.model is None or args.model_class_name is None: raise ValueError("Omni backend requires --model and --model-class-name") + if request.get("pipeline_kind", "one_stage") != args.pipeline_kind: + raise ValueError( + f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}" + ) attention_config = {"default": {"backend": "TORCH_SDPA"}} @@ -273,7 +338,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: parallel_config=DiffusionParallelConfig(), ) try: - model_class_name = get_model_class_name(omni) + detected_model_class_name = get_model_class_name(omni) + model_class_name = args.model_class_name sampling_params = OmniDiffusionSamplingParams( height=request["height"], width=request["width"], @@ -289,10 +355,9 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_") } apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance) - prompt: dict[str, Any] = { - "prompt": request["prompt"], - "negative_prompt": request["negative_prompt"], - } + prompt: dict[str, Any] = {"prompt": request["prompt"]} + if request.get("negative_prompt"): + prompt["negative_prompt"] = request["negative_prompt"] image_path = request.get("image") if image_path is not None: with Image.open(str(image_path)) as source_image: @@ -312,6 +377,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None: "attention_backend": "torch_sdpa", "model": args.model, "model_class_name": model_class_name, + "model_config_class_name": detected_model_class_name, + "pipeline_kind": args.pipeline_kind, }, ) finally: diff --git a/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py new file mode 100644 index 000000000..c8ce29c37 --- /dev/null +++ b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py @@ -0,0 +1,748 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""E2E accuracy guards against pinned official Lightricks LTX references. + +The LTX-2/2.3 comparison runs both runtimes through PyTorch SDPA and uses +``max_batch_size=4`` in the official reference to match Omni's fused guidance +batch. Video and audio guidance use the official non-HQ one-stage defaults; +only the generation shape and step count are reduced for CI runtime. +LTX-2.5 runs the official split-artifact pipelines with connector weights from +the same Diffusers checkpoint under test. The distilled case covers the fixed +two-stage schedule; the Full/SFT case covers raw dev weights and an explicit +shared one-stage schedule. +""" + +from __future__ import annotations + +import json +import math +import os +import shutil +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from pathlib import Path + +import numpy as np +import pytest +import torch +from huggingface_hub import hf_hub_download, snapshot_download +from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure + +from tests.e2e.accuracy.helpers import reset_artifact_dir +from tests.helpers import skip_if_gated_repo_inaccessible +from tests.helpers.mark import hardware_test + +OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git" +LTX25_OFFICIAL_REVISION = "7954dcb0d986bdc36ef272564a9789ade07fcc65" +LTX25_OMNI_MODEL_ID = "Lightricks/LTX-2.5-Diffusers" +LTX25_OMNI_MODEL_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d" +LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5" +LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1" +LTX25_OFFICIAL_COMMON_FILES = { + "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors", + "video_vae": "vae/ltx-2.5-video-vae-bf16.safetensors", + "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors", +} +LTX25_OFFICIAL_VARIANT_FILES = { + "distilled": { + "transformer": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors", + "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", + }, + "full": { + "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", + }, + "full_two_stage": { + "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors", + "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors", + "distilled_lora": "loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors", + }, +} +# Version selected by the pinned official source's uv.lock. Keep it isolated +# from Omni's runtime and development dependencies. +OFFICIAL_OPENIMAGEIO_VERSION = "3.1.11.0" +LTX25_PROMPT = "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside." +LTX25_I2V_IMAGE_REPO = "huggingface/documentation-images" +LTX25_I2V_IMAGE_FILENAME = "diffusers/svd/rocket.png" +LTX25_I2V_IMAGE_REVISION = "645d8364f0c7a101180b364811b5a11a362e4010" +NEGATIVE_PROMPT = ( + "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, " + "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, " + "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, " + "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of field, " + "background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting " + "direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, " + "incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, mismatched lip sync, " + "silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, incorrect " + "dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural " + "transitions, inconsistent framing, tilted camera, flat lighting, inconsistent tone, cinematic oversaturation, " + "stylized filters, or AI artifacts." +) + +# Release parity exercises the accelerated BF16 cuDNN path in both runtimes. +# Degenerate one-token attention falls back to PyTorch SDPA MATH. +ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16" +LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0] + + +@dataclass(frozen=True) +class LTXAccuracyThresholds: + video_ssim_mean: float + video_ssim_min: float + video_psnr_mean_db: float + audio_relative_l2: float + audio_cosine_similarity: float + + +STRICT_THRESHOLDS = LTXAccuracyThresholds( + video_ssim_mean=0.99, + video_ssim_min=0.99, + video_psnr_mean_db=40.0, + audio_relative_l2=0.10, + audio_cosine_similarity=0.99, +) + + +def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None: + from vllm_omni.outputs import OmniRequestOutput + + from .run_ltx25_reference import _omni_worker_peak_memory_mb, _unwrap_omni_output + + frame = object() + audio = object() + output = OmniRequestOutput( + stage_id=0, + images=[frame], + _multimodal_output={"audio": audio, "audio_sample_rate": 48_000}, + ) + + frames, actual_audio, sample_rate = _unwrap_omni_output(output) + + assert frames is frame + assert actual_audio is audio + assert sample_rate == 48_000 + output.peak_memory_mb = 12_345.0 + assert _omni_worker_peak_memory_mb([output]) == 12_345.0 + + +def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None: + start = time.perf_counter() + subprocess.run(command, env=env, timeout=timeout, check=True) + print(f"{' '.join(command[:3])} finished in {time.perf_counter() - start:.1f}s") + + +def _clone_pinned_source(root: Path, repository: str, revision: str) -> None: + root.mkdir(parents=True) + subprocess.run(["git", "init", "-q", str(root)], check=True) + subprocess.run(["git", "-C", str(root), "remote", "add", "origin", repository], check=True) + last_error: subprocess.CalledProcessError | None = None + for attempt in range(3): + try: + subprocess.run( + ["git", "-C", str(root), "fetch", "--depth", "1", "origin", revision], + check=True, + ) + last_error = None + break + except subprocess.CalledProcessError as error: + last_error = error + if attempt < 2: + time.sleep(5 * (attempt + 1)) + if last_error is not None: + raise last_error + subprocess.run(["git", "-C", str(root), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True) + + +def _git_revision(root: Path) -> str | None: + try: + return subprocess.check_output( + ["git", "-C", str(root), "rev-parse", "HEAD"], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except subprocess.CalledProcessError: + return None + + +def _ltx25_official_source(artifact_root: Path) -> tuple[Path, str]: + repository = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REPOSITORY", OFFICIAL_REPOSITORY) + revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REVISION", LTX25_OFFICIAL_REVISION) + configured_root = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_ROOT") + root = Path(configured_root) if configured_root else artifact_root / f"official-source-{revision[:12]}" + actual_revision = _git_revision(root) if root.exists() else None + if actual_revision != revision and configured_root: + raise AssertionError(f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}") + if actual_revision != revision: + if root.exists(): + shutil.rmtree(root) + _clone_pinned_source(root, repository, revision) + actual_revision = _git_revision(root) + assert actual_revision == revision, f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}" + return root, revision + + +def _official_runner_prefix() -> list[str]: + """Run the reference with its missing binary dependency isolated from CI.""" + uv = shutil.which("uv") + assert uv is not None, "uv is required to run the pinned official LTX reference" + return [ + uv, + "run", + "--no-project", + "--with", + f"openimageio=={OFFICIAL_OPENIMAGEIO_VERSION}", + "--python", + sys.executable, + "python", + ] + + +def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") -> tuple[Path, str, str | None]: + if transformer_subfolder not in {"transformer", "transformer_full"}: + raise ValueError(f"Unsupported LTX-2.5 transformer subfolder: {transformer_subfolder!r}") + configured_model = os.environ.get("VLLM_TEST_LTX25_MODEL") + configured_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION") + if configured_model and Path(configured_model).exists(): + model = Path(configured_model).resolve() + snapshot_revision = model.name if model.parent.name == "snapshots" else None + assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}" + assert (model / "diffusion_decoder" / "config.json").is_file(), ( + f"LTX-2.5 DiffVAE config not found: {model / 'diffusion_decoder' / 'config.json'}" + ) + return model, configured_model, configured_revision or snapshot_revision + model_id = configured_model or LTX25_OMNI_MODEL_ID + revision = configured_revision + if revision is None and model_id == LTX25_OMNI_MODEL_ID: + revision = LTX25_OMNI_MODEL_REVISION + model = Path( + snapshot_download( + repo_id=model_id, + revision=revision, + allow_patterns=[ + "model_index.json", + "audio_vae/*", + "connectors/config.json", + "connectors/diffusion_pytorch_model.safetensors.index.json", + "connectors/diffusion_pytorch_model-*.safetensors", + "diffusion_decoder/config.json", + "scheduler/*", + "text_encoder/config.json", + "text_encoder/generation_config.json", + "text_encoder/model*", + "tokenizer/*", + f"{transformer_subfolder}/*", + "vae/*", + "vocoder/*", + ], + ) + ).resolve() + snapshot_revision = model.name if model.parent.name == "snapshots" else None + return model, model_id, snapshot_revision or revision + + +def _resolve_ltx25_official_artifacts( + *, checkpoint_variant: str = "distilled" +) -> tuple[dict[str, Path], str, str | None]: + try: + variant_files = LTX25_OFFICIAL_VARIANT_FILES[checkpoint_variant] + except KeyError as exc: + raise ValueError(f"Unsupported official LTX-2.5 checkpoint variant: {checkpoint_variant!r}") from exc + files = {**LTX25_OFFICIAL_COMMON_FILES, **variant_files} + configured_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL") + configured_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION") + if configured_model and Path(configured_model).exists(): + root = Path(configured_model).resolve() + model_source = configured_model + snapshot_revision = root.name if root.parent.name == "snapshots" else None + resolved_revision = configured_revision or snapshot_revision + else: + model_source = configured_model or LTX25_OFFICIAL_MODEL_ID + revision = configured_revision + if revision is None and model_source == LTX25_OFFICIAL_MODEL_ID: + revision = LTX25_OFFICIAL_MODEL_REVISION + root = Path( + snapshot_download( + repo_id=model_source, + revision=revision, + allow_patterns=list(files.values()), + ) + ).resolve() + snapshot_revision = root.name if root.parent.name == "snapshots" else None + resolved_revision = snapshot_revision or revision + + artifacts = {name: root / relative_path for name, relative_path in files.items()} + missing = [f"{name}: {path}" for name, path in artifacts.items() if not path.is_file()] + assert not missing, f"Official LTX-2.5 artifacts are missing: {', '.join(missing)}" + return artifacts, model_source, resolved_revision + + +def _resolve_gemma_root(model: Path) -> Path: + configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT") + if configured_root: + root = Path(configured_root) + assert root.is_dir(), f"Gemma root not found: {root}" + return root + return model + + +def _resolve_ltx25_image() -> Path: + configured_image = os.environ.get("VLLM_TEST_LTX25_I2V_IMAGE") + if configured_image: + image = Path(configured_image) + assert image.is_file(), f"LTX-2.5 I2V conditioning image not found: {image}" + return image + return Path( + hf_hub_download( + repo_id=LTX25_I2V_IMAGE_REPO, + repo_type="dataset", + filename=LTX25_I2V_IMAGE_FILENAME, + revision=LTX25_I2V_IMAGE_REVISION, + ) + ) + + +def _require_ltx25_model_access() -> None: + configured_omni_model = os.environ.get("VLLM_TEST_LTX25_MODEL") + if not (configured_omni_model and Path(configured_omni_model).exists()): + omni_model = configured_omni_model or LTX25_OMNI_MODEL_ID + omni_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION") + if omni_revision is None and omni_model == LTX25_OMNI_MODEL_ID: + omni_revision = LTX25_OMNI_MODEL_REVISION + skip_if_gated_repo_inaccessible( + omni_model, + revision=omni_revision, + filename="model_index.json", + ) + + configured_official_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL") + if not (configured_official_model and Path(configured_official_model).exists()): + official_model = configured_official_model or LTX25_OFFICIAL_MODEL_ID + official_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION") + if official_revision is None and official_model == LTX25_OFFICIAL_MODEL_ID: + official_revision = LTX25_OFFICIAL_MODEL_REVISION + skip_if_gated_repo_inaccessible( + official_model, + revision=official_revision, + filename=LTX25_OFFICIAL_COMMON_FILES["text_encoder"], + ) + + +def _ltx25_request(image: Path | None = None) -> dict[str, object]: + request: dict[str, object] = { + "prompt": LTX25_PROMPT, + # Keep the nightly parity test small enough to run the pinned official + # pipeline and Omni sequentially on one H100. The public recipe covers + # the release-resolution 1920x1088x121 configuration. + "width": 1024, + "height": 576, + "num_frames": 25, + "fps": 24, + "num_inference_steps": 8, + "seed": 42, + } + if image is not None: + request.update(image=str(image.resolve()), image_crf=18) + return request + + +def _ltx25_full_sigmas(num_inference_steps: int) -> list[float]: + """Materialize the official LTX-2.5 one-stage schedule in FP32.""" + sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1) + base_anchor = 1024 + max_anchor = 4096 + base_shift = 0.95 + max_shift = 2.05 + slope = (max_shift - base_shift) / (max_anchor - base_anchor) + sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor) + exp_shift = math.exp(sigma_shift) + sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0) + + non_zero = sigmas != 0 + one_minus_sigmas = 1.0 - sigmas[non_zero] + scale = one_minus_sigmas[-1] / (1.0 - 0.1) + sigmas[non_zero] = 1.0 - one_minus_sigmas / scale + return [float(sigma) for sigma in sigmas] + + +def _ltx25_full_request(image: Path | None = None) -> dict[str, object]: + num_inference_steps = 30 + request: dict[str, object] = { + "prompt": LTX25_PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + # Full guidance expands to a four-way batch. Keep the decoded output + # small enough for the pinned official and Omni runs to share one H100. + "width": 512, + "height": 384, + "num_frames": 25, + "fps": 24, + "num_inference_steps": num_inference_steps, + "seed": 42, + "video_cfg_scale": 3.0, + "audio_cfg_scale": 7.0, + "video_stg_scale": 1.0, + "audio_stg_scale": 1.0, + "video_modality_scale": 3.0, + "audio_modality_scale": 3.0, + "video_rescale_scale": 0.7, + "audio_rescale_scale": 0.7, + "video_stg_blocks": [28], + "audio_stg_blocks": [28], + "sigmas": _ltx25_full_sigmas(num_inference_steps), + } + if image is not None: + request.update(image=str(image.resolve()), image_crf=18) + return request + + +def test_ltx25_full_request_pins_official_schedule() -> None: + request = _ltx25_full_request() + sigmas = request["sigmas"] + + assert isinstance(sigmas, list) + assert isinstance(request["num_inference_steps"], int) + assert len(sigmas) == request["num_inference_steps"] + 1 + assert math.isclose(sigmas[0], 1.0, rel_tol=0.0, abs_tol=1e-6) + assert sigmas[-1] == 0.0 + assert all(sigmas[index] > sigmas[index + 1] for index in range(len(sigmas) - 1)) + + +def test_ltx25_i2v_requests_pin_official_crf() -> None: + image = Path("conditioning.png") + for request in (_ltx25_request(image), _ltx25_full_request(image)): + assert request["image"] == str(image.resolve()) + assert request["image_crf"] == 18 + + +def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]: + assert reference.shape == prediction.shape + assert reference.ndim == 4 and reference.shape[-1] == 3 + ssim_scores: list[float] = [] + psnr_scores: list[float] = [] + for reference_frame, prediction_frame in zip(reference, prediction, strict=True): + reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0) + prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0) + ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor))) + psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor))) + difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64)) + return { + "ssim_mean": float(np.mean(ssim_scores)), + "ssim_min": float(np.min(ssim_scores)), + "psnr_mean_db": float(np.mean(psnr_scores)), + "max_abs": float(difference.max()), + "mean_abs": float(difference.mean()), + } + + +def _canonical_audio(audio: np.ndarray) -> np.ndarray: + while audio.ndim > 2 and audio.shape[0] == 1: + audio = audio[0] + return audio.astype(np.float64) + + +def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float | bool]: + reference = _canonical_audio(reference) + prediction = _canonical_audio(prediction) + assert reference.shape == prediction.shape + difference = reference - prediction + reference_norm = max(float(np.linalg.norm(reference)), 1e-12) + prediction_norm = max(float(np.linalg.norm(prediction)), 1e-12) + return { + "bitwise_equal": bool(np.array_equal(reference, prediction)), + "max_abs": float(np.abs(difference).max()), + "mean_abs": float(np.abs(difference).mean()), + "relative_l2": float(np.linalg.norm(difference) / reference_norm), + "cosine_similarity": float(np.vdot(reference.ravel(), prediction.ravel()) / (reference_norm * prediction_norm)), + } + + +def _assert_strict_similarity( + video_metrics: dict[str, float], + audio_metrics: dict[str, float | bool], +) -> None: + assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean + assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min + assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db + assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2 + assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity + + +@pytest.mark.slow +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("task", ["t2v", "i2v"]) +def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None: + """Compare Omni with official canonical DiffVAE and matching connector weights.""" + _require_ltx25_model_access() + artifact_parent = accuracy_artifact_root / "ltx_official" + output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}") + official_root, official_revision = _ltx25_official_source(artifact_parent) + official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts() + omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model() + image = _resolve_ltx25_image() if task == "i2v" else None + request_path = output_root / "request.json" + request_path.write_text(json.dumps(_ltx25_request(image), indent=2) + "\n") + + runner = Path(__file__).with_name("run_ltx25_reference.py") + runner_args = [ + str(runner), + "--request", + str(request_path), + "--enable-layerwise-offload", + ] + env = os.environ.copy() + env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision + env["PYTHONUNBUFFERED"] = "1" + repository_root = Path(__file__).resolve().parents[4] + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}" + ) + + official_output = output_root / "official" + _run( + _official_runner_prefix() + + runner_args + + [ + "--backend", + "official", + "--output-dir", + str(official_output), + "--official-root", + str(official_root), + "--transformer-path", + str(official_artifacts["transformer"]), + "--text-encoder-path", + str(official_artifacts["text_encoder"]), + "--video-vae-path", + str(official_artifacts["video_vae"]), + "--audio-vae-path", + str(official_artifacts["audio_vae"]), + "--spatial-upsampler-path", + str(official_artifacts["spatial_upsampler"]), + "--connector-model", + str(omni_model), + ], + env=env, + ) + + omni_output = output_root / "omni" + _run( + [sys.executable] + + runner_args + + [ + "--backend", + "omni", + "--output-dir", + str(omni_output), + "--model", + str(omni_model), + "--model-class-name", + "LTX2DistilledTwoStagePipeline", + ], + env=env, + ) + + official_metadata = json.loads((official_output / "metadata.json").read_text()) + omni_metadata = json.loads((omni_output / "metadata.json").read_text()) + assert official_metadata["official_revision"] == official_revision + assert official_metadata["pipeline"] == "DistilledPipeline" + assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"]) + assert official_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["connector_model"] == str(omni_model) + assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline" + assert omni_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"] + video_metrics = _video_metrics( + np.load(official_output / "video.npy"), + np.load(omni_output / "video.npy"), + ) + audio_metrics = _audio_metrics( + np.load(official_output / "audio.npy"), + np.load(omni_output / "audio.npy"), + ) + result = { + "case": "ltx2_5_distilled", + "task": task, + "attention_backend": ATTENTION_BACKEND, + "official_revision": official_revision, + "official_model": official_model_source, + "official_model_revision": official_model_revision, + "official_connector_model": omni_model_source, + "official_connector_revision": omni_model_revision, + "omni_model": omni_model_source, + "omni_model_revision": omni_model_revision, + "resolved_omni_model_path": str(omni_model), + "thresholds": asdict(STRICT_THRESHOLDS), + "peak_memory_mb": { + "official_allocated": official_metadata["peak_memory_allocated_mb"], + "official_reserved": official_metadata["peak_memory_reserved_mb"], + "omni_worker_reported": omni_metadata["worker_peak_memory_mb"], + "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"], + "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"], + }, + "video": video_metrics, + "audio": audio_metrics, + } + (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + _assert_strict_similarity(video_metrics, audio_metrics) + + +@pytest.mark.slow +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("task", ["t2v", "i2v"]) +@pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"]) +def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None: + """Compare raw official Full/SFT and DiffVAE weights with Omni.""" + _require_ltx25_model_access() + artifact_parent = accuracy_artifact_root / "ltx_official" + output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}") + model_class_name = { + "full_one_stage": "LTX2Pipeline", + "full_two_stage": "LTX2TwoStagePipeline", + }[pipeline_mode] + official_root, official_revision = _ltx25_official_source(artifact_parent) + official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts( + checkpoint_variant="full" if pipeline_mode == "full_one_stage" else "full_two_stage" + ) + omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model( + transformer_subfolder="transformer_full" + ) + image = _resolve_ltx25_image() if task == "i2v" else None + request = _ltx25_full_request(image) + if pipeline_mode == "full_two_stage": + request["stage_1_sigmas"] = request.pop("sigmas") + request["stage_2_sigmas"] = LTX25_STAGE_2_SIGMAS + request_path = output_root / "request.json" + request_path.write_text(json.dumps(request, indent=2) + "\n") + + runner = Path(__file__).with_name("run_ltx25_reference.py") + runner_args = [ + str(runner), + "--request", + str(request_path), + "--enable-layerwise-offload", + ] + env = os.environ.copy() + env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision + env["PYTHONUNBUFFERED"] = "1" + repository_root = Path(__file__).resolve().parents[4] + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}" + ) + + official_output = output_root / "official" + official_sidecar_args: list[str] = [] + if pipeline_mode == "full_two_stage": + official_sidecar_args = [ + "--spatial-upsampler-path", + str(official_artifacts["spatial_upsampler"]), + "--distilled-lora-path", + str(official_artifacts["distilled_lora"]), + ] + + _run( + _official_runner_prefix() + + runner_args + + [ + "--backend", + "official", + "--output-dir", + str(official_output), + "--official-root", + str(official_root), + "--official-pipeline", + pipeline_mode, + "--transformer-path", + str(official_artifacts["transformer"]), + "--text-encoder-path", + str(official_artifacts["text_encoder"]), + "--video-vae-path", + str(official_artifacts["video_vae"]), + "--audio-vae-path", + str(official_artifacts["audio_vae"]), + "--connector-model", + str(omni_model), + ] + + official_sidecar_args, + env=env, + ) + + omni_output = output_root / "omni" + _run( + [sys.executable] + + runner_args + + [ + "--backend", + "omni", + "--output-dir", + str(omni_output), + "--model", + str(omni_model), + "--model-class-name", + model_class_name, + ], + env=env, + ) + + official_metadata = json.loads((official_output / "metadata.json").read_text()) + omni_metadata = json.loads((omni_output / "metadata.json").read_text()) + assert official_metadata["official_revision"] == official_revision + assert ( + official_metadata["pipeline"] + == { + "full_one_stage": "TI2VidOneStagePipeline", + "full_two_stage": "TI2VidTwoStagesPipeline", + }[pipeline_mode] + ) + assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"]) + assert official_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["connector_model"] == str(omni_model) + assert omni_metadata["model_class_name"] == model_class_name + assert omni_metadata["attention_backend"] == ATTENTION_BACKEND + assert official_metadata["seed"] == omni_metadata["seed"] + for schedule_name in ("sigmas", "stage_1_sigmas", "stage_2_sigmas"): + assert official_metadata[schedule_name] == omni_metadata[schedule_name] + assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"] + video_metrics = _video_metrics( + np.load(official_output / "video.npy"), + np.load(omni_output / "video.npy"), + ) + audio_metrics = _audio_metrics( + np.load(official_output / "audio.npy"), + np.load(omni_output / "audio.npy"), + ) + result = { + "case": f"ltx2_5_{pipeline_mode}", + "task": task, + "attention_backend": ATTENTION_BACKEND, + "official_revision": official_revision, + "official_model": official_model_source, + "official_model_revision": official_model_revision, + "official_connector_model": omni_model_source, + "official_connector_revision": omni_model_revision, + "omni_model": omni_model_source, + "omni_model_revision": omni_model_revision, + "resolved_omni_model_path": str(omni_model), + "thresholds": asdict(STRICT_THRESHOLDS), + "peak_memory_mb": { + "official_allocated": official_metadata["peak_memory_allocated_mb"], + "official_reserved": official_metadata["peak_memory_reserved_mb"], + "omni_worker_reported": omni_metadata["worker_peak_memory_mb"], + "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"], + "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"], + }, + "video": video_metrics, + "audio": audio_metrics, + } + (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + _assert_strict_similarity(video_metrics, audio_metrics) diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py index 647ece7a7..63ce73923 100644 --- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py +++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py @@ -3,10 +3,8 @@ """E2E accuracy guard against a pinned Lightricks LTX pipeline revision. -The comparison runs both runtimes through PyTorch SDPA and uses -``max_batch_size=4`` in the official reference to match Omni's fused guidance -batch. Video and audio guidance use the official non-HQ one-stage defaults; -only the generation shape and step count are reduced for CI runtime. +Four complementary default-shape cases cover both model versions, both +two-stage weight families, and T2V/I2V without duplicating one-stage coverage. """ from __future__ import annotations @@ -17,8 +15,9 @@ import shutil import subprocess import sys import time -from dataclasses import dataclass +from dataclasses import asdict, dataclass, replace from pathlib import Path +from typing import Literal import numpy as np import pytest @@ -53,73 +52,198 @@ NEGATIVE_PROMPT = ( # Both runtimes use PyTorch SDPA with the current Torch dispatch defaults. ATTENTION_BACKEND = "torch_sdpa" -VIDEO_SSIM_MEAN_THRESHOLD = 0.95 -VIDEO_SSIM_MIN_THRESHOLD = 0.90 -VIDEO_PSNR_MEAN_THRESHOLD = 30.0 -AUDIO_RELATIVE_L2_THRESHOLD = 0.2 -AUDIO_COSINE_THRESHOLD = 0.95 + + +@dataclass(frozen=True) +class LTXArtifact: + repo_id: str + filename: str + revision: str + env: str + repo_type: Literal["model", "dataset"] = "model" + + +@dataclass(frozen=True) +class LTXAccuracyThresholds: + video_ssim_mean: float + video_ssim_min: float + video_psnr_mean_db: float + audio_relative_l2: float + audio_cosine_similarity: float + + +STRICT_THRESHOLDS = LTXAccuracyThresholds( + video_ssim_mean=0.99, + video_ssim_min=0.99, + video_psnr_mean_db=40.0, + audio_relative_l2=0.10, + audio_cosine_similarity=0.99, +) @dataclass(frozen=True) class LTXAccuracyCase: name: str + pipeline_kind: Literal["distilled", "two_stage"] model_id: str model_revision: str model_env: str model_class_name: str - checkpoint_repo: str - checkpoint_filename: str - checkpoint_revision: str - checkpoint_env: str - stg_block: int - prompt: str = PROMPT - image_repo: str | None = None - image_filename: str | None = None - image_revision: str | None = None + checkpoint: LTXArtifact + width: int + height: int + num_frames: int + num_inference_steps: int + seed: int + stg_block: int | None + gemma_model_id: str | None = None + gemma_model_revision: str | None = None + gemma_model_env: str | None = None + image: LTXArtifact | None = None + spatial_upsampler: LTXArtifact | None = None + distilled_lora: LTXArtifact | None = None + # Official uses block streaming; Omni uses layerwise offload. Both preserve + # bf16 arithmetic while keeping full-resolution two-stage cases on one H100. + enable_layerwise_offload: bool = False -CASES = ( +LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a" +LTX23_REVISION = "4229404625088d21c4f112eb640fb04a0900ee25" +LTX2_CHECKPOINT = LTXArtifact( + repo_id="Lightricks/LTX-2", + filename="ltx-2-19b-dev.safetensors", + revision=LTX2_REVISION, + env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT", +) +LTX2_DISTILLED_CHECKPOINT = replace( + LTX2_CHECKPOINT, + filename="ltx-2-19b-distilled.safetensors", + env="VLLM_TEST_LTX2_DISTILLED_OFFICIAL_CHECKPOINT", +) +LTX2_UPSAMPLER = replace( + LTX2_CHECKPOINT, + filename="ltx-2-spatial-upscaler-x2-1.0.safetensors", + env="VLLM_TEST_LTX2_UPSAMPLER", +) +LTX2_DISTILLED_LORA = replace( + LTX2_CHECKPOINT, + filename="ltx-2-19b-distilled-lora-384.safetensors", + env="VLLM_TEST_LTX2_DISTILLED_LORA", +) +LTX23_CHECKPOINT = LTXArtifact( + repo_id="Lightricks/LTX-2.3", + filename="ltx-2.3-22b-dev.safetensors", + revision=LTX23_REVISION, + env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT", +) +LTX23_DISTILLED_CHECKPOINT = replace( + LTX23_CHECKPOINT, + # The pinned Diffusers checkpoint contains the original merged distilled + # Transformer. Keep the reference Transformer on the same version; the + # spatial upsampler is versioned independently and remains on 1.1 below. + filename="ltx-2.3-22b-distilled.safetensors", + env="VLLM_TEST_LTX23_DISTILLED_OFFICIAL_CHECKPOINT", +) +LTX23_UPSAMPLER = replace( + LTX23_CHECKPOINT, + filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", + env="VLLM_TEST_LTX23_UPSAMPLER", +) +LTX23_DISTILLED_LORA = replace( + LTX23_CHECKPOINT, + filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors", + env="VLLM_TEST_LTX23_DISTILLED_LORA", +) +I2V_IMAGE = LTXArtifact( + repo_id="huggingface/documentation-images", + filename="diffusers/svd/rocket.png", + revision="645d8364f0c7a101180b364811b5a11a362e4010", + env="VLLM_TEST_LTX_I2V_IMAGE", + repo_type="dataset", +) + + +TWO_STAGE_CASES = ( LTXAccuracyCase( - name="ltx2", + name="ltx2_distilled_t2v", + pipeline_kind="distilled", + model_id="rootonchair/LTX-2-19b-distilled", + model_revision="388e2846f54aae51687498ffb6b27c7c2c9ce9e5", + model_env="VLLM_TEST_LTX2_DISTILLED_MODEL", + model_class_name="LTX2DistilledPipeline", + checkpoint=LTX2_DISTILLED_CHECKPOINT, + spatial_upsampler=LTX2_UPSAMPLER, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=8, + seed=10, + stg_block=None, + gemma_model_id="Lightricks/LTX-2", + gemma_model_revision=LTX2_REVISION, + gemma_model_env="VLLM_TEST_LTX2_MODEL", + ), + LTXAccuracyCase( + name="ltx23_distilled_i2v", + pipeline_kind="distilled", + model_id="diffusers/LTX-2.3-Distilled-Diffusers", + model_revision="432e0d3c2d1769aaa4d295f9243f7062bf6b47ee", + model_env="VLLM_TEST_LTX23_DISTILLED_MODEL", + model_class_name="LTX2DistilledPipeline", + checkpoint=LTX23_DISTILLED_CHECKPOINT, + spatial_upsampler=LTX23_UPSAMPLER, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=8, + seed=10, + stg_block=None, + gemma_model_id="diffusers/LTX-2.3-Diffusers", + gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf", + gemma_model_env="VLLM_TEST_LTX23_MODEL", + image=I2V_IMAGE, + ), + LTXAccuracyCase( + name="ltx2_two_stage_layer_fused_t2v", + pipeline_kind="two_stage", model_id="Lightricks/LTX-2", - model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a", + model_revision=LTX2_REVISION, model_env="VLLM_TEST_LTX2_MODEL", - model_class_name="LTX2Pipeline", - checkpoint_repo="Lightricks/LTX-2", - checkpoint_filename="ltx-2-19b-dev.safetensors", - checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a", - checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT", + model_class_name="LTX2TwoStagePipeline", + checkpoint=LTX2_CHECKPOINT, + spatial_upsampler=LTX2_UPSAMPLER, + distilled_lora=LTX2_DISTILLED_LORA, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=40, + seed=10, stg_block=29, + enable_layerwise_offload=True, ), LTXAccuracyCase( - name="ltx2_3", + name="ltx23_two_stage_layer_fused_i2v", + pipeline_kind="two_stage", model_id="diffusers/LTX-2.3-Diffusers", model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf", model_env="VLLM_TEST_LTX23_MODEL", - model_class_name="LTX2Pipeline", - checkpoint_repo="Lightricks/LTX-2.3", - checkpoint_filename="ltx-2.3-22b-dev.safetensors", - checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25", - checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT", + model_class_name="LTX2TwoStagePipeline", + checkpoint=LTX23_CHECKPOINT, + spatial_upsampler=LTX23_UPSAMPLER, + distilled_lora=LTX23_DISTILLED_LORA, + width=1536, + height=1024, + num_frames=121, + num_inference_steps=30, + seed=10, stg_block=28, - ), - LTXAccuracyCase( - name="ltx2_3_i2v", - model_id="diffusers/LTX-2.3-Diffusers", - model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf", - model_env="VLLM_TEST_LTX23_MODEL", - model_class_name="LTX2Pipeline", - checkpoint_repo="Lightricks/LTX-2.3", - checkpoint_filename="ltx-2.3-22b-dev.safetensors", - checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25", - checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT", - stg_block=28, - image_repo="huggingface/documentation-images", - image_filename="diffusers/svd/rocket.png", - image_revision="645d8364f0c7a101180b364811b5a11a362e4010", + enable_layerwise_offload=True, + image=I2V_IMAGE, ), ) +WEEKLY_CASES = TWO_STAGE_CASES + def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None: start = time.perf_counter() @@ -209,6 +333,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path: "model_index.json", "audio_vae/*", "connectors/*", + "latent_upsampler/*", "processor/*", "scheduler/*", "text_encoder/config.json", @@ -223,74 +348,124 @@ def _resolve_model(case: LTXAccuracyCase) -> Path: ) -def _resolve_gemma_root(model: Path) -> Path: +def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path: configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT") if configured_root: root = Path(configured_root) assert root.is_dir(), f"Gemma root not found: {root}" return root + if case.gemma_model_id is not None: + configured_model = os.environ.get(case.gemma_model_env or "") + if configured_model and Path(configured_model).is_dir(): + return Path(configured_model) + return Path( + snapshot_download( + repo_id=case.gemma_model_id, + revision=case.gemma_model_revision, + allow_patterns=[ + "processor/*", + "text_encoder/*", + "tokenizer/*", + ], + ) + ) return model -def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path: - configured_checkpoint = os.environ.get(case.checkpoint_env) - if configured_checkpoint: - checkpoint = Path(configured_checkpoint) - assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}" - return checkpoint - model_checkpoint = model / case.checkpoint_filename - if model_checkpoint.is_file(): - return model_checkpoint +def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path: + configured_path = os.environ.get(artifact.env) + if configured_path: + path = Path(configured_path) + assert path.is_file(), f"Configured LTX artifact not found: {path}" + return path + if model is not None: + model_path = model / artifact.filename + if model_path.is_file(): + return model_path return Path( hf_hub_download( - repo_id=case.checkpoint_repo, - filename=case.checkpoint_filename, - revision=case.checkpoint_revision, + repo_id=artifact.repo_id, + repo_type=None if artifact.repo_type == "model" else artifact.repo_type, + filename=artifact.filename, + revision=artifact.revision, ) ) def _resolve_image(case: LTXAccuracyCase) -> Path | None: - if case.image_filename is None: + if case.image is None: return None - if case.image_repo is None or case.image_revision is None: - raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.") - configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE") - if configured_image: - image = Path(configured_image) - assert image.is_file(), f"LTX I2V conditioning image not found: {image}" - return image - return Path( - hf_hub_download( - repo_id=case.image_repo, - repo_type="dataset", - filename=case.image_filename, - revision=case.image_revision, - ) - ) + return _resolve_artifact(case.image) + + +def _omni_model_with_pinned_artifacts( + model: Path, + output_root: Path, + artifacts: tuple[Path | None, ...], +) -> Path: + """Expose the already-resolved sidecars to Omni without copying the model.""" + resolved_artifacts = tuple(artifact for artifact in artifacts if artifact is not None) + if not resolved_artifacts: + return model + + overlay = output_root / "omni-model" + overlay.mkdir() + for source in model.iterdir(): + (overlay / source.name).symlink_to(source.resolve(), target_is_directory=source.is_dir()) + for artifact in resolved_artifacts: + target = overlay / artifact.name + if target.exists(): + assert target.samefile(artifact), f"Model sidecar does not match pinned artifact: {target} != {artifact}" + else: + target.symlink_to(artifact.resolve()) + return overlay + + +def test_omni_model_overlay_pins_artifacts(tmp_path: Path) -> None: + model = tmp_path / "model" + model.mkdir() + component = model / "model_index.json" + component.write_text("{}\n") + artifact = tmp_path / "sidecar.safetensors" + artifact.write_bytes(b"pinned") + output_root = tmp_path / "output" + output_root.mkdir() + + overlay = _omni_model_with_pinned_artifacts(model, output_root, (artifact,)) + + assert (overlay / component.name).samefile(component) + assert (overlay / artifact.name).samefile(artifact) def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]: request: dict[str, object] = { - "prompt": case.prompt, - "negative_prompt": NEGATIVE_PROMPT, - "width": 512, - "height": 384, - "num_frames": 25, + "pipeline_kind": case.pipeline_kind, + "prompt": PROMPT, + "negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT, + "width": case.width, + "height": case.height, + "num_frames": case.num_frames, "fps": 24, - "num_inference_steps": 20, - "seed": 42, - "video_cfg_scale": 3.0, - "audio_cfg_scale": 7.0, - "video_stg_scale": 1.0, - "audio_stg_scale": 1.0, - "video_modality_scale": 3.0, - "audio_modality_scale": 3.0, - "video_rescale_scale": 0.7, - "audio_rescale_scale": 0.7, - "video_stg_blocks": [case.stg_block], - "audio_stg_blocks": [case.stg_block], + "num_inference_steps": case.num_inference_steps, + "seed": case.seed, } + if case.pipeline_kind != "distilled": + request.update( + video_cfg_scale=3.0, + audio_cfg_scale=7.0, + video_modality_scale=3.0, + audio_modality_scale=3.0, + ) + if case.pipeline_kind == "two_stage": + assert case.stg_block is not None + request.update( + video_stg_scale=1.0, + audio_stg_scale=1.0, + video_rescale_scale=0.7, + audio_rescale_scale=0.7, + video_stg_blocks=[case.stg_block], + audio_stg_blocks=[case.stg_block], + ) if image is not None: request["image"] = str(image.resolve()) return request @@ -299,20 +474,30 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]: def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]: assert reference.shape == prediction.shape assert reference.ndim == 4 and reference.shape[-1] == 3 + device = torch.device("cpu") + ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device) + psnr = PeakSignalNoiseRatio(data_range=1.0).to(device) ssim_scores: list[float] = [] psnr_scores: list[float] = [] - for reference_frame, prediction_frame in zip(reference, prediction, strict=True): - reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0) - prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0) - ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor))) - psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor))) - difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64)) + max_abs = 0.0 + absolute_error_sum = 0.0 + with torch.inference_mode(): + for reference_frame, prediction_frame in zip(reference, prediction, strict=True): + reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device) + prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device) + ssim_scores.append(float(ssim(prediction_tensor, reference_tensor))) + psnr_scores.append(float(psnr(prediction_tensor, reference_tensor))) + ssim.reset() + psnr.reset() + difference = np.abs(reference_frame.astype(np.float64) - prediction_frame.astype(np.float64)) + max_abs = max(max_abs, float(difference.max())) + absolute_error_sum += float(difference.sum()) return { "ssim_mean": float(np.mean(ssim_scores)), "ssim_min": float(np.min(ssim_scores)), "psnr_mean_db": float(np.mean(psnr_scores)), - "max_abs": float(difference.max()), - "mean_abs": float(difference.mean()), + "max_abs": max_abs, + "mean_abs": absolute_error_sum / reference.size, } @@ -338,29 +523,42 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f } -@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name) -@pytest.mark.full_model +@pytest.mark.parametrize("case", WEEKLY_CASES, ids=lambda case: case.name) +@pytest.mark.slow @pytest.mark.benchmark @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None: - """Compare official and Omni raw AV outputs from the same E2E request.""" +def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None: + """Compare official and Omni raw AV outputs from the same single-H100 request.""" + configured_artifact_root = os.environ.get("VLLM_TEST_LTX_ARTIFACT_ROOT") + if configured_artifact_root: + accuracy_artifact_root = Path(configured_artifact_root) + accuracy_artifact_root.mkdir(parents=True, exist_ok=True) output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name) official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official") model = _resolve_model(case) - gemma_root = _resolve_gemma_root(model) - checkpoint = _resolve_checkpoint(case, model) + gemma_root = _resolve_gemma_root(case, model) + checkpoint = _resolve_artifact(case.checkpoint, model) + spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None + distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None + omni_model = _omni_model_with_pinned_artifacts(model, output_root, (spatial_upsampler, distilled_lora)) image = _resolve_image(case) + request = _request(case, image) request_path = output_root / "request.json" - request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n") + request_path.write_text(json.dumps(request, indent=2) + "\n") runner = Path(__file__).with_name("run_ltx_reference.py") runner_args = [ str(runner), "--request", str(request_path), + "--pipeline-kind", + case.pipeline_kind, ] - if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}: + enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get( + "VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "" + ).lower() in {"1", "true", "yes", "on"} + if enable_layerwise_offload: runner_args.append("--enable-layerwise-offload") env = os.environ.copy() env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision @@ -372,10 +570,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact ) official_output = output_root / "official" - _run( - _official_runner_prefix() - + runner_args - + [ + official_args = _official_runner_prefix() + runner_args + official_args.extend( + [ "--backend", "official", "--output-dir", @@ -386,12 +583,16 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact str(checkpoint), "--gemma-root", str(gemma_root), - ], - env=env, + ] ) + if spatial_upsampler is not None: + official_args.extend(["--spatial-upsampler", str(spatial_upsampler)]) + if distilled_lora is not None: + official_args.extend(["--distilled-lora", str(distilled_lora)]) + _run(official_args, env=env) omni_output = output_root / "omni" - _run( + omni_args = ( [sys.executable] + runner_args + [ @@ -400,12 +601,12 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact "--output-dir", str(omni_output), "--model", - str(model), + str(omni_model), "--model-class-name", case.model_class_name, - ], - env=env, + ] ) + _run(omni_args, env=env) official_metadata = json.loads((official_output / "metadata.json").read_text()) omni_metadata = json.loads((omni_output / "metadata.json").read_text()) @@ -423,18 +624,25 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact result = { "case": case.name, "task": "i2v" if image is not None else "t2v", + "pipeline_kind": case.pipeline_kind, + "model_class_name": case.model_class_name, "attention_backend": ATTENTION_BACKEND, "official_revision": official_revision, "model_revision": case.model_revision, - "checkpoint_revision": case.checkpoint_revision, + "checkpoint_revision": case.checkpoint.revision, + "spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None), + "distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None, + "enable_layerwise_offload": enable_layerwise_offload, + "thresholds": asdict(STRICT_THRESHOLDS), + "request": request, "video": video_metrics, "audio": audio_metrics, } (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n") print(json.dumps(result, indent=2)) - assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD - assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD - assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD - assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD - assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD + assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean + assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min + assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db + assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2 + assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity diff --git a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py index 50474a644..e688b9360 100644 --- a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py +++ b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniCPM-o 4.5 Daily-Omni + Seed-TTS accuracy regression coverage. Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches @@ -21,13 +22,25 @@ from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import ( build_acc_benchmark_cli_argv, find_vllm_cli, ) -from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import SERVER_PARAMS as DUPLEX_TEST_PARAMS +from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( + DEPLOY_CONFIG as _DUPLEX_DEPLOY_CONFIG, +) +from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( + SERVER_PARAMS as DUPLEX_TEST_PARAMS, +) from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config _MODEL = os.environ.get("VLLM_TEST_MINICPMO_4_5_MODEL", "openbmb/MiniCPM-o-4_5") _DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml") +_DUPLEX_BF16_DEPLOY_CONFIG = modify_stage_config( + _DUPLEX_DEPLOY_CONFIG, + updates={ + "base_config": _DEPLOY_CONFIG, + "connectors.connector_of_shared_memory.extra.code2wav_bfloat16_attention_cache": True, + }, +) _RESULT_DIR = Path( os.environ.get( "ACC_BENCH_RESULT_DIR", @@ -91,6 +104,20 @@ seed_test_params = [ ] +duplex_accuracy_test_params = [ + *DUPLEX_TEST_PARAMS, + pytest.param( + OmniServerParams( + model=_MODEL, + stage_config_path=_DUPLEX_BF16_DEPLOY_CONFIG, + use_stage_cli=False, + server_args=list(_SEED_TTS_SERVER_ARGS), + ), + id="three-stage-single-gpu-bf16-attention-cache", + ), +] + + def _require_vllm_cli() -> None: try: find_vllm_cli() @@ -174,7 +201,7 @@ def test_minicpmo_4_5_seed_tts_wer_bench(omni_server) -> None: @hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.parametrize("omni_server", DUPLEX_TEST_PARAMS, indirect=True) +@pytest.mark.parametrize("omni_server", duplex_accuracy_test_params, indirect=True) def test_minicpmo_4_5_duplex_seed_tts_wer_bench(omni_server) -> None: """Gate Seed-TTS WER through the explicit Realtime TTS contract.""" _require_vllm_cli() diff --git a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py b/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py index b00e90d5e..60f829bee 100644 --- a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py +++ b/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -119,6 +119,9 @@ def _model_name(model_env_var: str = MODEL_ENV_VAR, component: str = "FL2VA") -> def _server_args() -> list[str]: return [ "--trust-remote-code", + # Both official-reference requests use fixed shapes. Keep regional + # compilation static so cold Inductor codegen is stable across CI nodes. + "--no-diffusion-compile-dynamic", "--num-gpus", "4", "--usp", diff --git a/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py index dfc53a06e..d612b96a5 100644 --- a/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py +++ b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared helpers for Qwen3-Omni Daily-Omni / Seed-TTS ``vllm bench serve --omni`` accuracy runs. Local dataset paths are **optional**. When ``VLLM_DAILY_OMNI_QA_JSON`` + ``VLLM_DAILY_OMNI_VIDEO_DIR`` @@ -28,7 +29,7 @@ DEFAULT_SEED_TTS_HF_REPO = "zhaochenyang20/seed-tts-eval" class OmniBenchServerEndpoint(Protocol): - """Anything with ``host`` / ``port`` / ``model`` (e.g. :class:`tests.conftest.OmniServer`).""" + """Anything with ``host`` / ``port`` / ``model`` (e.g. :class:`tests.helpers.runtime.OmniServer`).""" host: str port: int diff --git a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py index 121551664..7c48053e1 100644 --- a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py +++ b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen3-Omni accuracy benchmarks (Daily-Omni MCQ + Seed-TTS WER) via ``vllm bench serve --omni``. Starts a **module-scoped** Omni OpenAI-compatible server (same pattern as ``tests/dfx/perf`` and @@ -67,12 +68,8 @@ def get_chunk_config(config_path: str | None = None): """Load the qwen3_omni CI deploy yaml with async_chunk modifications for streaming mode.""" if config_path is None: config_path = _CI_DEPLOY - # TODO: remove this workaround once legacy `stage_args` path is deleted. - # The pipeline (qwen3_omni/pipeline.py) already wires - # thinker2talker_async_chunk / talker2code2wav_async_chunk on stage 0/1, - # so only async_chunk needs flipping. Writing nested `engine_args:` into - # the new-schema overlay trips _parse_stage_deploy's legacy branch and - # drops flat fields (load_format, max_num_seqs, ...). + # Pipeline already wires thinker2talker_async_chunk / talker2code2wav_async_chunk + # on stage 0/1; only the top-level async_chunk flag needs flipping here. return modify_stage_config(config_path, updates={"async_chunk": True}) diff --git a/tests/e2e/accuracy/sana_video/README.md b/tests/e2e/accuracy/sana_video/README.md new file mode 100644 index 000000000..cab38b7a2 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/README.md @@ -0,0 +1,63 @@ +# SANA-Video golden outputs + +The v2 golden set covers both text-to-video and image-to-video for each +checkpoint variant. Generate references with the Diffusers version pinned by +this repository and the immutable Hub revisions embedded in +`generate_goldens.py`: + +```bash +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 480p --task t2v +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 480p --task i2v --input-image /path/to/fixed-input.png +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 720p --task t2v +python tests/e2e/accuracy/sana_video/generate_goldens.py \ + --variant 720p --task i2v --input-image /path/to/fixed-input.png +``` + +Use the same reviewed input image for both I2V variants. The generator +normalizes it to an RGB PNG and includes that exact file and its SHA256 in the +case manifest, so an I2V reference is self-contained. + +Each case is written under: + +```text +/// +``` + +It contains: + +- `transformer_case.safetensors`: deterministic transformer inputs and output. +- `pipeline.mp4`: the encoded reference used by online similarity tests. +- `pipeline_reference.safetensors`: final denoised latents and decoded, + pre-MP4 frames quantized to the exact uint8 video input. +- `input.png`: the normalized conditioning image for I2V only. +- `metadata.json`: model, generator, scheduler, input and sampling provenance. +- `manifest.json`: the size and SHA256 of every artifact. + +## Publication requirements + +A canonical golden must be generated from the pinned Hugging Face model ID and +immutable commit. The generator also requires a clean vLLM-Omni worktree and +records the repository revision plus the generator script hash. + +`--model` is available for local development. It hashes every file in the local +checkpoint directory and records optional `--local-source-model` and +`--local-source-revision` claims, but marks the result `publishable: false`. +The golden test rejects such a manifest. A local directory hash proves which +bytes were used; it does not prove that those bytes match a claimed upstream +commit. + +After reviewing the metadata, input image and generated videos, upload each +publishable case directory without renaming files to: + +```text +s3://vllm-public-assets/omni-assets/sana-video/v2/// +``` + +Set `SANA_VIDEO_GOLDEN_BASE_URL` to the corresponding `v2` HTTP base URL when +running the tests. Uploading is intentionally separate from generation so test +runs cannot overwrite a frozen reference. The storage location must prevent +in-place mutation; publish a new version prefix instead of replacing any v2 +file when a model, dependency, input or sampling parameter changes. diff --git a/tests/e2e/accuracy/sana_video/generate_goldens.py b/tests/e2e/accuracy/sana_video/generate_goldens.py new file mode 100644 index 000000000..8ab25dd18 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/generate_goldens.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Generate frozen SANA-Video transformer and pipeline reference artifacts.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import subprocess +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from diffusers import SanaImageToVideoPipeline, SanaVideoPipeline +from diffusers.utils import export_to_video +from PIL import Image +from safetensors.torch import save_file + +PROMPT = "A cat walking on the grass, facing the camera. motion score: 30." +NEGATIVE_PROMPT = ( + "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, " + "jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, " + "jitter, and ghosting effects, creating a disorienting visual experience." +) +SCHEMA_VERSION = 2 +SEED = 42 +NUM_FRAMES = 81 +FPS = 16 +NUM_INFERENCE_STEPS = 50 +GUIDANCE_SCALE = 6.0 +VARIANTS = { + "480p": { + "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + "revision": "fed3bce411c58a0f688a31afe8f52e61acc2b15f", + "height": 480, + "width": 832, + "vae_dtype": torch.float32, + "transformer_case_shape": (1, 16, 3, 8, 8), + }, + "720p": { + "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + "revision": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e", + "height": 704, + "width": 1280, + "vae_dtype": torch.bfloat16, + "transformer_case_shape": (1, 128, 2, 4, 4), + }, +} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _json_sha256(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _directory_fingerprint(root: Path) -> dict[str, Any]: + """Fingerprint every local checkpoint file without recording host paths.""" + if not root.is_dir(): + raise ValueError(f"Local checkpoint must be a directory: {root}") + files = [] + digest = hashlib.sha256() + for path in sorted(candidate for candidate in root.rglob("*") if candidate.is_file()): + relative = path.relative_to(root).as_posix() + file_sha256 = _sha256(path) + size = path.stat().st_size + entry = f"{relative}\0{size}\0{file_sha256}\n".encode() + digest.update(entry) + files.append({"path": relative, "size": size, "sha256": file_sha256}) + if not files: + raise ValueError(f"Local checkpoint contains no files: {root}") + return { + "algorithm": "sha256-tree-v1", + "sha256": digest.hexdigest(), + "file_count": len(files), + "total_size": sum(entry["size"] for entry in files), + } + + +def _repository_provenance() -> dict[str, Any]: + repo_root = Path(__file__).resolve().parents[4] + revision = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + dirty = bool( + subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=repo_root, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + ) + return { + "repository": "vllm-project/vllm-omni", + "revision": revision, + "dirty": dirty, + "generator_sha256": _sha256(Path(__file__)), + } + + +def _cpu(tensor: torch.Tensor) -> torch.Tensor: + return tensor.detach().cpu().contiguous() + + +def _normalize_input_image(source: Path, destination: Path) -> Image.Image: + if not source.is_file(): + raise ValueError(f"I2V input image does not exist: {source}") + with Image.open(source) as image: + normalized = image.convert("RGB") + normalized.load() + normalized.save(destination, format="PNG", optimize=False) + return normalized + + +def _frames_to_uint8(frames: np.ndarray) -> torch.Tensor: + tensor = torch.from_numpy(np.asarray(frames)) + if tensor.is_floating_point(): + # Match diffusers.utils.export_to_video, which truncates rather than + # rounds when converting float frames to the encoder's uint8 input. + tensor = tensor.clamp(0, 1) * 255 + return tensor.to(torch.uint8).contiguous() + + +def _model_provenance( + config: dict[str, Any], + model: str | None, + local_source_model: str | None, + local_source_revision: str | None, +) -> tuple[dict[str, Any], bool]: + if model is None: + return ( + { + "kind": "huggingface_hub", + "model_id": config["model"], + "revision": config["revision"], + }, + True, + ) + + local_path = Path(model).expanduser().resolve() + if bool(local_source_model) != bool(local_source_revision): + raise ValueError("--local-source-model and --local-source-revision must be specified together.") + provenance = { + "kind": "local_checkpoint", + "checkpoint_fingerprint": _directory_fingerprint(local_path), + "claimed_source_model": local_source_model, + "claimed_source_revision": local_source_revision, + "verification": "unverified-local-copy", + } + # A directory hash proves which local bytes were used, but not that those + # bytes match a claimed Hub commit. Such output is useful for development + # comparisons but must never be uploaded as the canonical frozen reference. + return provenance, False + + +def generate( + variant: str, + task: str, + model: str | None, + input_image: Path | None, + output_root: Path, + local_source_model: str | None = None, + local_source_revision: str | None = None, +) -> Path: + config = dict(VARIANTS[variant]) + model_provenance, publishable = _model_provenance( + config, + model, + local_source_model, + local_source_revision, + ) + repository_provenance = _repository_provenance() + if publishable and repository_provenance["dirty"]: + raise RuntimeError("Refusing to generate publishable golden assets from a dirty vllm-omni worktree.") + if model is not None: + config["model"] = str(Path(model).expanduser().resolve()) + config["revision"] = None + + output_dir = output_root / variant / task + output_dir.mkdir(parents=True, exist_ok=True) + normalized_image: Image.Image | None = None + input_path: Path | None = None + if task == "i2v": + if input_image is None: + raise ValueError("--input-image is required for task=i2v.") + input_path = output_dir / "input.png" + normalized_image = _normalize_input_image(input_image, input_path) + elif input_image is not None: + raise ValueError("--input-image is only valid for task=i2v.") + + pipeline_class = SanaImageToVideoPipeline if task == "i2v" else SanaVideoPipeline + pipe = pipeline_class.from_pretrained( + config["model"], + revision=config["revision"], + torch_dtype=torch.bfloat16, + ) + pipe.text_encoder.to(torch.bfloat16) + pipe.vae.to(config["vae_dtype"]) + pipe.to("cuda") + + torch.manual_seed(1234) + hidden_states = torch.randn(config["transformer_case_shape"], dtype=torch.bfloat16, device="cuda") + encoder_hidden_states = torch.randn(1, 8, 2304, dtype=torch.bfloat16, device="cuda") + encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 0]], device="cuda") + timestep = torch.tensor([500.0], device="cuda") + with torch.inference_mode(): + transformer_output = pipe.transformer( + hidden_states, + encoder_hidden_states, + timestep, + encoder_attention_mask=encoder_attention_mask, + ).sample + + transformer_path = output_dir / "transformer_case.safetensors" + save_file( + { + "hidden_states": _cpu(hidden_states), + "encoder_hidden_states": _cpu(encoder_hidden_states), + "encoder_attention_mask": _cpu(encoder_attention_mask), + "timestep": _cpu(timestep), + "output": _cpu(transformer_output), + }, + transformer_path, + ) + + final_latents: dict[str, torch.Tensor] = {} + + def capture_final_latents(_pipe, step, _timestep, callback_kwargs): + if step == NUM_INFERENCE_STEPS - 1: + final_latents["latents"] = _cpu(callback_kwargs["latents"]) + return callback_kwargs + + torch.accelerator.reset_peak_memory_stats() + generator = torch.Generator(device="cuda").manual_seed(SEED) + generation_kwargs = { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": config["height"], + "width": config["width"], + "frames": NUM_FRAMES, + "guidance_scale": GUIDANCE_SCALE, + "num_inference_steps": NUM_INFERENCE_STEPS, + "generator": generator, + "output_type": "np", + "callback_on_step_end": capture_final_latents, + "callback_on_step_end_tensor_inputs": ["latents"], + } + if normalized_image is not None: + generation_kwargs["image"] = normalized_image + + torch.accelerator.synchronize() + started = time.perf_counter() + with torch.inference_mode(): + frames = pipe(**generation_kwargs).frames[0] + torch.accelerator.synchronize() + generation_seconds = time.perf_counter() - started + if "latents" not in final_latents: + raise RuntimeError("Diffusers did not expose the final denoised latents to the generation callback.") + + video_path = output_dir / "pipeline.mp4" + export_to_video(frames, video_path, fps=FPS) + reference_path = output_dir / "pipeline_reference.safetensors" + save_file( + { + "final_latents": final_latents["latents"], + "decoded_frames_uint8": _frames_to_uint8(frames), + }, + reference_path, + ) + + scheduler_config = dict(pipe.scheduler.config) + input_metadata = None + if input_path is not None: + input_metadata = { + "filename": input_path.name, + "sha256": _sha256(input_path), + "size": input_path.stat().st_size, + "format": "png-rgb", + } + metadata = { + "schema_version": SCHEMA_VERSION, + "variant": variant, + "task": task, + "publishable": publishable, + "reference_implementation": { + "library": "diffusers", + "version": __import__("diffusers").__version__, + "pipeline_class": pipeline_class.__name__, + "torch_version": torch.__version__, + }, + "model_provenance": model_provenance, + "generator_provenance": repository_provenance, + "scheduler": { + "class": type(pipe.scheduler).__name__, + "config_sha256": _json_sha256(scheduler_config), + "config": scheduler_config, + }, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "seed": SEED, + "generator_device": "cuda", + "dtype": "bfloat16", + "vae_dtype": str(config["vae_dtype"]).removeprefix("torch."), + "height": config["height"], + "width": config["width"], + "num_frames": NUM_FRAMES, + "fps": FPS, + "num_inference_steps": NUM_INFERENCE_STEPS, + "guidance_scale": GUIDANCE_SCALE, + "input_image": input_metadata, + "generation_seconds": generation_seconds, + "peak_reserved_gib": torch.accelerator.max_memory_reserved() / 1024**3, + } + metadata_path = output_dir / "metadata.json" + metadata_path.write_text(json.dumps(metadata, indent=2, default=str) + "\n") + + artifact_paths = [transformer_path, video_path, reference_path, metadata_path] + if input_path is not None: + artifact_paths.append(input_path) + manifest = { + "schema_version": SCHEMA_VERSION, + "variant": variant, + "task": task, + "publishable": publishable, + "model_provenance": model_provenance, + "metadata_sha256": _sha256(metadata_path), + "files": {path.name: {"sha256": _sha256(path), "size": path.stat().st_size} for path in artifact_paths}, + } + manifest_path = output_dir / "manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + return manifest_path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--variant", choices=VARIANTS, required=True) + parser.add_argument("--task", choices=("t2v", "i2v"), required=True) + parser.add_argument( + "--model", + help=( + "Optional local checkpoint. Local output is fingerprinted and marked development-only; " + "it cannot be published as the canonical golden." + ), + ) + parser.add_argument("--input-image", type=Path, help="Fixed local RGB input for task=i2v.") + parser.add_argument("--local-source-model", help="Claimed upstream model ID for local checkpoint audit records.") + parser.add_argument("--local-source-revision", help="Claimed upstream commit for local checkpoint audit records.") + parser.add_argument("--output-root", type=Path, default=Path("/tmp/sana-video-goldens-v2")) + args = parser.parse_args() + print( + generate( + args.variant, + args.task, + args.model, + args.input_image, + args.output_root, + args.local_source_model, + args.local_source_revision, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_golden.py b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py new file mode 100644 index 000000000..c0c468315 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import base64 +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import pytest +import requests +import torch +from safetensors.torch import load_file +from torch import nn + +from tests.e2e.accuracy.helpers import assert_video_metadata, assert_video_similarity_metrics, probe_video +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +GOLDEN_BASE_URL = os.getenv("SANA_VIDEO_GOLDEN_BASE_URL") +SCHEMA_VERSION = 2 +PROMPT = "A cat walking on the grass, facing the camera. motion score: 30." +NEGATIVE_PROMPT = ( + "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, " + "jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, " + "jitter, and ghosting effects, creating a disorienting visual experience." +) +VARIANTS = { + "480p": ("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", 480, 832), + "720p": ("Efficient-Large-Model/SANA-Video_2B_720p_diffusers", 704, 1280), +} +REVISIONS = { + "480p": "fed3bce411c58a0f688a31afe8f52e61acc2b15f", + "720p": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e", +} +SSIM_THRESHOLDS = { + ("480p", "t2v"): 0.91, + ("480p", "i2v"): 0.93, + ("720p", "t2v"): 0.93, + ("720p", "i2v"): 0.93, +} + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.skipif( + not GOLDEN_BASE_URL, + reason="Set SANA_VIDEO_GOLDEN_BASE_URL after publishing the frozen v2 S3 assets.", + ), +] + + +class _TransformerCheckpoint(nn.Module): + def __init__(self, transformer: nn.Module, model: str, revision: str) -> None: + super().__init__() + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + + self.transformer = transformer + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="transformer", + revision=revision, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _json_sha256(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _assert_hex_digest(value: Any, *, length: int = 64) -> None: + assert isinstance(value, str) and len(value) == length + assert all(character in "0123456789abcdef" for character in value) + + +def _validate_manifest(payload: Any, variant: str, task: str) -> dict[str, Any]: + assert isinstance(payload, dict) + assert payload["schema_version"] == SCHEMA_VERSION + assert payload["variant"] == variant + assert payload["task"] == task + # Local checkpoint generation is deliberately marked non-publishable. Do + # not permit such output to become the canonical reference by accident. + assert payload["publishable"] is True + assert payload["model_provenance"] == { + "kind": "huggingface_hub", + "model_id": VARIANTS[variant][0], + "revision": REVISIONS[variant], + } + _assert_hex_digest(payload["metadata_sha256"]) + + expected_files = { + "metadata.json", + "pipeline.mp4", + "pipeline_reference.safetensors", + "transformer_case.safetensors", + } + if task == "i2v": + expected_files.add("input.png") + files = payload["files"] + assert isinstance(files, dict) + assert set(files) == expected_files + for filename, expected in files.items(): + # Prevent a compromised manifest from writing outside output_dir. + assert filename == Path(filename).name + assert isinstance(expected, dict) and set(expected) == {"sha256", "size"} + _assert_hex_digest(expected["sha256"]) + assert isinstance(expected["size"], int) and expected["size"] > 0 + return payload + + +def _validate_metadata(payload: Any, manifest: dict[str, Any], variant: str, task: str) -> None: + assert isinstance(payload, dict) + _, height, width = VARIANTS[variant] + assert payload["schema_version"] == SCHEMA_VERSION + assert payload["variant"] == variant + assert payload["task"] == task + assert payload["publishable"] is True + assert payload["model_provenance"] == manifest["model_provenance"] + assert payload["prompt"] == PROMPT + assert payload["negative_prompt"] == NEGATIVE_PROMPT + assert payload["seed"] == 42 + assert payload["generator_device"] == "cuda" + assert payload["dtype"] == "bfloat16" + assert payload["height"] == height + assert payload["width"] == width + assert payload["num_frames"] == 81 + assert payload["fps"] == 16 + assert payload["num_inference_steps"] == 50 + assert payload["guidance_scale"] == 6.0 + + implementation = payload["reference_implementation"] + assert implementation["library"] == "diffusers" + assert implementation["pipeline_class"] == ("SanaImageToVideoPipeline" if task == "i2v" else "SanaVideoPipeline") + assert implementation["version"] + assert implementation["torch_version"] + + generator = payload["generator_provenance"] + assert generator["repository"] == "vllm-project/vllm-omni" + _assert_hex_digest(generator["revision"], length=40) + _assert_hex_digest(generator["generator_sha256"]) + assert generator["dirty"] is False + + scheduler = payload["scheduler"] + assert scheduler["class"] == "DPMSolverMultistepScheduler" + _assert_hex_digest(scheduler["config_sha256"]) + assert scheduler["config_sha256"] == _json_sha256(scheduler["config"]) + + input_image = payload["input_image"] + if task == "i2v": + assert input_image["filename"] == "input.png" + _assert_hex_digest(input_image["sha256"]) + assert input_image["size"] == manifest["files"]["input.png"]["size"] + assert input_image["sha256"] == manifest["files"]["input.png"]["sha256"] + assert input_image["format"] == "png-rgb" + else: + assert input_image is None + + +def _download_case(variant: str, task: str, output_dir: Path, filenames: set[str]) -> dict[str, Any]: + output_dir.mkdir(parents=True, exist_ok=True) + base_url = f"{GOLDEN_BASE_URL.rstrip('/')}/{variant}/{task}" + response = requests.get(f"{base_url}/manifest.json", timeout=60) + response.raise_for_status() + payload = _validate_manifest(response.json(), variant, task) + filenames = filenames | {"metadata.json"} + if task == "i2v": + filenames.add("input.png") + assert filenames <= payload["files"].keys() + for filename in filenames: + expected = payload["files"][filename] + path = output_dir / filename + response = requests.get(f"{base_url}/{filename}", timeout=300) + response.raise_for_status() + path.write_bytes(response.content) + assert path.stat().st_size == expected["size"] + assert _sha256(path) == expected["sha256"] + assert _sha256(output_dir / "metadata.json") == payload["metadata_sha256"] + metadata = json.loads((output_dir / "metadata.json").read_text()) + _validate_metadata(metadata, payload, variant, task) + return payload + + +@pytest.mark.benchmark +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize( + "_hardware", + [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))], +) +def test_sana_video_transformer_matches_frozen_golden( + variant: str, + _hardware, + accuracy_artifact_root: Path, +) -> None: + del _hardware + from vllm.config import LoadConfig + from vllm.transformers_utils.config import get_hf_file_to_dict + + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + + model, _, _ = VARIANTS[variant] + revision = REVISIONS[variant] + output_dir = accuracy_artifact_root / "sana_video" / variant / "t2v" + _download_case(variant, "t2v", output_dir, {"transformer_case.safetensors"}) + case = load_file(output_dir / "transformer_case.safetensors") + transformer_config = get_hf_file_to_dict("transformer/config.json", model, revision=revision) + assert transformer_config is not None + transformer = SanaVideoTransformer3DModel.from_config(transformer_config).to( + device="cuda", + dtype=torch.bfloat16, + ) + checkpoint = _TransformerCheckpoint(transformer, model, revision) + od_config = OmniDiffusionConfig( + model=model, + model_class_name="SanaVideoPipeline", + dtype=torch.bfloat16, + revision=revision, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config) + transformer.load_weights( + (name.removeprefix("transformer."), tensor) + for name, tensor in loader.get_all_weights(checkpoint) + if name.startswith("transformer.") + ) + transformer.eval() + with torch.inference_mode(): + actual = transformer( + case["hidden_states"].to("cuda"), + case["encoder_hidden_states"].to("cuda"), + case["timestep"].to("cuda"), + encoder_attention_mask=case["encoder_attention_mask"].to("cuda"), + ).sample.float() + expected = case["output"].to("cuda").float() + error = actual - expected + max_abs = error.abs().max() + relative_l2 = torch.linalg.vector_norm(error) / torch.linalg.vector_norm(expected) + cosine = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0) + # CUDNN_ATTN BF16 reduction order is not bitwise deterministic on + # Blackwell. Bound rare outliers separately while keeping the aggregate + # relative-L2 and cosine checks strict. + max_abs_threshold = 0.2 if variant == "480p" else 0.1 + relative_l2_threshold = 0.017 if variant == "480p" else 0.015 + assert max_abs.item() <= max_abs_threshold + assert relative_l2.item() <= relative_l2_threshold + assert cosine.item() >= 0.999 + + +SERVER_CASES = [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--model-class-name", + pipeline_class, + ], + ), + id=f"{variant}-{task}", + marks=hardware_marks(res={"cuda": "H100"}), + ) + for variant, (model, _, _) in VARIANTS.items() + for task, pipeline_class in ( + ("t2v", "SanaVideoPipeline"), + ("i2v", "SanaImageToVideoPipeline"), + ) +] + + +@pytest.mark.benchmark +@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) +def test_sana_video_pipeline_matches_frozen_golden( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + accuracy_artifact_root: Path, +) -> None: + variant = "720p" if "720p" in omni_server.model else "480p" + task = "i2v" if "SanaImageToVideoPipeline" in omni_server.serve_args else "t2v" + _, height, width = VARIANTS[variant] + output_dir = accuracy_artifact_root / "sana_video" / variant / task + _download_case(variant, task, output_dir, {"pipeline.mp4"}) + + request_config = { + "model": omni_server.model, + "form_data": { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": height, + "width": width, + "num_frames": 81, + "fps": 16, + "num_inference_steps": 50, + "guidance_scale": 6.0, + "seed": 42, + }, + } + if task == "i2v": + encoded = base64.b64encode((output_dir / "input.png").read_bytes()).decode() + request_config["image_reference"] = f"data:image/png;base64,{encoded}" + result = openai_client.send_video_diffusion_request(request_config)[0] + actual_path = output_dir / "actual.mp4" + actual_path.write_bytes(result.videos[0]) + golden_path = output_dir / "pipeline.mp4" + + metadata = probe_video(actual_path) + assert_video_metadata(metadata, width=width, height=height, fps=16, frame_count=81) + assert_video_similarity_metrics( + label=f"sana_video_{variant}_{task}", + online_path=actual_path, + offline_path=golden_path, + # The native prompt cross-attention uses vLLM-Omni Attention rather + # than Diffusers SDPA. Frozen stage checks prove exact inputs and + # scheduler state, while the BF16 backend reduction difference grows + # across CFG denoising. Calibrate only the affected 480p T2V case; + # its frozen result is 0.919861 SSIM / 28.7442 dB PSNR. + ssim_threshold=SSIM_THRESHOLDS[(variant, task)], + psnr_threshold=28.0, + ) diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py new file mode 100644 index 000000000..8dd35eca2 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from tests.helpers.mark import hardware_marks + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.benchmark, +] + +VARIANTS = { + "480p": { + "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + "vae_class": "AutoencoderKLWan", + "native_vae_class": "DistributedAutoencoderKLWan", + "dtype": torch.float32, + "latent_channels": 16, + "spatial_scale": 8, + "temporal_scale": 4, + "frames": 5, + "max_abs": 1e-6, + "relative_l2": 1e-6, + }, + "720p": { + "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + "vae_class": "AutoencoderKLLTX2Video", + "native_vae_class": "DistributedAutoencoderKLLTX2Video", + "dtype": torch.bfloat16, + "latent_channels": 128, + "spatial_scale": 32, + "temporal_scale": 8, + "frames": 9, + "max_abs": 0.04, + "relative_l2": 0.004, + }, +} + + +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize( + "_hardware", + [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))], +) +def test_sana_video_i2v_first_frame_latents_match_diffusers(variant: str, _hardware) -> None: + del _hardware + from diffusers import ( + AutoencoderKLLTX2Video, + AutoencoderKLWan, + ) + from diffusers import ( + SanaImageToVideoPipeline as DiffusersSanaImageToVideoPipeline, + ) + + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, + ) + from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline + + config = VARIANTS[variant] + vae_classes = { + "AutoencoderKLLTX2Video": AutoencoderKLLTX2Video, + "AutoencoderKLWan": AutoencoderKLWan, + } + native_vae_classes = { + "DistributedAutoencoderKLLTX2Video": DistributedAutoencoderKLLTX2Video, + "DistributedAutoencoderKLWan": DistributedAutoencoderKLWan, + } + reference_vae = ( + vae_classes[config["vae_class"]] + .from_pretrained( + config["model"], + subfolder="vae", + torch_dtype=config["dtype"], + ) + .to("cuda") + ) + native_vae = ( + native_vae_classes[config["native_vae_class"]] + .from_config(dict(reference_vae.config)) + .to(device="cuda", dtype=config["dtype"]) + ) + native_vae.load_state_dict(reference_vae.state_dict()) + reference_vae.eval() + native_vae.eval() + + height, width = 192, 320 + num_latent_frames = (config["frames"] - 1) // config["temporal_scale"] + 1 + image = torch.linspace( + -1, + 1, + 3 * height * width, + dtype=torch.float32, + device="cuda", + ).reshape(1, 3, height, width) + noise = torch.randn( + ( + 1, + config["latent_channels"], + num_latent_frames, + height // config["spatial_scale"], + width // config["spatial_scale"], + ), + generator=torch.Generator(device="cuda").manual_seed(7), + device="cuda", + dtype=torch.float32, + ) + reference_pipeline = SimpleNamespace( + vae=reference_vae, + vae_scale_factor_temporal=config["temporal_scale"], + vae_scale_factor_spatial=config["spatial_scale"], + ) + native_pipeline = SimpleNamespace( + vae=native_vae, + vae_scale_factor_temporal=config["temporal_scale"], + vae_scale_factor_spatial=config["spatial_scale"], + device=torch.device("cuda"), + ) + + with torch.inference_mode(): + expected = DiffusersSanaImageToVideoPipeline.prepare_latents( + reference_pipeline, + image, + 1, + config["latent_channels"], + height, + width, + config["frames"], + torch.float32, + torch.device("cuda"), + None, + noise.clone(), + ) + actual = SanaImageToVideoPipeline._prepare_i2v_latents( + native_pipeline, + image, + 1, + config["latent_channels"], + height, + width, + config["frames"], + torch.float32, + None, + noise.clone(), + ) + + # Conditioning replaces only the first latent frame. + torch.testing.assert_close(actual[:, :, 1:], noise[:, :, 1:], rtol=0, atol=0) + error = (actual[:, :, :1] - expected[:, :, :1]).float() + expected_first_frame = expected[:, :, :1].float() + max_abs = error.abs().max() + relative_l2 = torch.linalg.vector_norm(error) / torch.linalg.vector_norm(expected_first_frame) + cosine = torch.nn.functional.cosine_similarity( + actual[:, :, :1].flatten(), + expected[:, :, :1].flatten(), + dim=0, + ) + assert max_abs.item() <= config["max_abs"] + assert relative_l2.item() <= config["relative_l2"] + assert cosine.item() >= 0.9999 diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py new file mode 100644 index 000000000..a88aec7e2 --- /dev/null +++ b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py @@ -0,0 +1,432 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import gc +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_marks + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.benchmark, +] + +PROMPT = "A small red panda walks through a bamboo forest." +NEGATIVE_PROMPT = "blur, distortion, abrupt cuts" +SEED = 42 +GUIDANCE_SCALE = 6.0 +NUM_INFERENCE_STEPS = 2 +HEIGHT = 192 +WIDTH = 320 +NUM_FRAMES = 9 + +VARIANTS = { + "480p": { + "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + "revision": "fed3bce411c58a0f688a31afe8f52e61acc2b15f", + "model_env": "SANA_VIDEO_480P_MODEL", + "max_abs": 0.2, + "relative_l2": 0.02, + "cosine": 0.999, + "cfg_max_abs": 0.5, + "cfg_relative_l2": 0.05, + "cfg_cosine": 0.999, + "step_max_abs": 0.075, + "step_relative_l2": 0.012, + "step_cosine": 0.9999, + }, + "720p": { + "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", + "revision": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e", + "model_env": "SANA_VIDEO_720P_MODEL", + "max_abs": 0.1, + "relative_l2": 0.015, + "cosine": 0.999, + "cfg_max_abs": 0.6, + "cfg_relative_l2": 0.065, + "cfg_cosine": 0.998, + "step_max_abs": 0.08, + "step_relative_l2": 0.014, + "step_cosine": 0.9999, + }, +} + + +@dataclass(frozen=True) +class ErrorMetrics: + shape: tuple[int, ...] + actual_dtype: torch.dtype + expected_dtype: torch.dtype + max_abs: float + mean_abs: float + relative_l2: float + cosine: float + + def __str__(self) -> str: + return ( + f"shape={self.shape}, actual_dtype={self.actual_dtype}, " + f"expected_dtype={self.expected_dtype}, max_abs={self.max_abs:.8g}, " + f"mean_abs={self.mean_abs:.8g}, relative_l2={self.relative_l2:.8g}, " + f"cosine={self.cosine:.8g}" + ) + + +def _metrics(actual: torch.Tensor, expected: torch.Tensor) -> ErrorMetrics: + assert actual.shape == expected.shape + actual_float = actual.detach().float() + expected_float = expected.detach().float() + error = actual_float - expected_float + expected_norm = torch.linalg.vector_norm(expected_float) + relative_l2 = torch.linalg.vector_norm(error) / expected_norm.clamp_min(torch.finfo(torch.float32).tiny) + if actual_float.numel() == 1: + cosine = torch.tensor( + 1.0 if bool(torch.equal(actual_float, expected_float)) else 0.0, + device=actual.device, + ) + else: + cosine = F.cosine_similarity(actual_float.flatten(), expected_float.flatten(), dim=0) + return ErrorMetrics( + shape=tuple(actual.shape), + actual_dtype=actual.dtype, + expected_dtype=expected.dtype, + max_abs=error.abs().max().item(), + mean_abs=error.abs().mean().item(), + relative_l2=relative_l2.item(), + cosine=cosine.item(), + ) + + +def _assert_stage( + stage: str, + actual: torch.Tensor, + expected: torch.Tensor, + *, + max_abs: float, + relative_l2: float, + cosine: float, +) -> ErrorMetrics: + result = _metrics(actual, expected) + print(f"\nSANA-Video alignment [{stage}]: {result}") + assert result.actual_dtype == result.expected_dtype, f"{stage}: {result}" + assert result.max_abs <= max_abs, f"{stage}: {result}" + assert result.relative_l2 <= relative_l2, f"{stage}: {result}" + assert result.cosine >= cosine, f"{stage}: {result}" + return result + + +def _model_source(variant: str) -> tuple[str, str | None]: + config = VARIANTS[variant] + model = os.getenv(config["model_env"], config["model"]) + revision = None if Path(model).exists() else config["revision"] + return model, revision + + +def _native_method_proxy(reference_pipeline: Any): + """Build the minimum native-pipeline surface needed by stage helpers. + + Sharing the tokenizer and text encoder makes this a comparison of pipeline + orchestration rather than a second copy of the same multi-billion-parameter + encoder. Transformer execution still uses separate Diffusers and native + modules below. + """ + from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline + + proxy = object.__new__(SanaVideoPipeline) + torch.nn.Module.__init__(proxy) + proxy.tokenizer = reference_pipeline.tokenizer + proxy.text_encoder = reference_pipeline.text_encoder + proxy.device = torch.device("cuda") + proxy.vae_scale_factor_temporal = reference_pipeline.vae_scale_factor_temporal + proxy.vae_scale_factor_spatial = reference_pipeline.vae_scale_factor_spatial + return proxy + + +def _cfg_prediction( + prediction: torch.Tensor, + *, + guidance_scale: float, + latent_channels: int, + out_channels: int, +) -> torch.Tensor: + unconditional, conditional = prediction.float().chunk(2) + prediction = unconditional + guidance_scale * (conditional - unconditional) + if out_channels // 2 == latent_channels: + prediction = prediction.chunk(2, dim=1)[0] + return prediction + + +@pytest.mark.parametrize("variant", VARIANTS) +@pytest.mark.parametrize( + "_hardware", + [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))], +) +def test_sana_video_quick_pipeline_stages_match_diffusers( + variant: str, + _hardware, +) -> None: + """Compare the first denoising step without replacing native attention. + + This deliberately stops after the first scheduler step. It keeps the test + useful as a stage-localizing PR/nightly check while the frozen 81-frame, + 50-step golden test covers the final latent and decoded-video boundary. + """ + del _hardware + from diffusers import SanaVideoPipeline as DiffusersSanaVideoPipeline + from diffusers.pipelines.sana_video.pipeline_sana_video import ( + retrieve_timesteps as diffusers_retrieve_timesteps, + ) + from diffusers.schedulers import DPMSolverMultistepScheduler + + from vllm_omni.diffusion.models.sana_video import ( + SanaVideoPipeline, + SanaVideoTransformer3DModel, + ) + from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + retrieve_timesteps as native_retrieve_timesteps, + ) + + model, revision = _model_source(variant) + config = VARIANTS[variant] + load_kwargs: dict[str, Any] = { + "torch_dtype": torch.bfloat16, + # No stage in this test decodes latents. Avoid loading a multi-GB VAE + # that would add memory and runtime without increasing coverage. + "vae": None, + } + if revision is not None: + load_kwargs["revision"] = revision + + reference = DiffusersSanaVideoPipeline.from_pretrained(model, **load_kwargs).to("cuda") + assert type(reference) is DiffusersSanaVideoPipeline + assert type(reference.transformer).__module__.startswith("diffusers.") + native_proxy = _native_method_proxy(reference) + assert type(native_proxy) is SanaVideoPipeline + assert type(native_proxy).__module__.startswith("vllm_omni.") + try: + with torch.inference_mode(): + reference_prompt = reference.encode_prompt( + PROMPT, + True, + negative_prompt=NEGATIVE_PROMPT, + device=torch.device("cuda"), + clean_caption=False, + max_sequence_length=300, + complex_human_instruction=None, + ) + native_prompt = SanaVideoPipeline.encode_prompt( + native_proxy, + PROMPT, + True, + negative_prompt=NEGATIVE_PROMPT, + device=torch.device("cuda"), + clean_caption=False, + max_sequence_length=300, + complex_human_instruction=None, + ) + prompt_stage_names = ( + "prompt_embeddings", + "prompt_attention_mask", + "negative_prompt_embeddings", + "negative_prompt_attention_mask", + ) + for stage, actual, expected in zip(prompt_stage_names, native_prompt, reference_prompt, strict=True): + _assert_stage( + stage, + actual, + expected, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + + reference_prompt_embeds, reference_prompt_mask, reference_negative_embeds, reference_negative_mask = ( + reference_prompt + ) + encoder_hidden_states = torch.cat([reference_negative_embeds, reference_prompt_embeds]) + encoder_attention_mask = torch.cat([reference_negative_mask, reference_prompt_mask]) + + latent_channels = reference.transformer.config.in_channels + reference_latents = reference.prepare_latents( + 1, + latent_channels, + HEIGHT, + WIDTH, + NUM_FRAMES, + torch.float32, + torch.device("cuda"), + torch.Generator(device="cuda").manual_seed(SEED), + ) + native_latents = SanaVideoPipeline.prepare_latents( + native_proxy, + 1, + latent_channels, + HEIGHT, + WIDTH, + NUM_FRAMES, + torch.float32, + torch.device("cuda"), + torch.Generator(device="cuda").manual_seed(SEED), + ) + _assert_stage( + "initial_noise_latents", + native_latents, + reference_latents, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + + reference_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config) + native_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config) + isolated_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config) + reference_timesteps, reference_num_steps = diffusers_retrieve_timesteps( + reference_scheduler, + NUM_INFERENCE_STEPS, + torch.device("cuda"), + ) + native_timesteps, native_num_steps = native_retrieve_timesteps( + native_scheduler, + NUM_INFERENCE_STEPS, + torch.device("cuda"), + ) + isolated_retrieve_timesteps, isolated_num_steps = native_retrieve_timesteps( + isolated_scheduler, + NUM_INFERENCE_STEPS, + torch.device("cuda"), + ) + assert reference_num_steps == native_num_steps == isolated_num_steps + _assert_stage( + "timesteps", + native_timesteps, + reference_timesteps, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + _assert_stage( + "sigmas", + native_scheduler.sigmas, + reference_scheduler.sigmas, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + torch.testing.assert_close(isolated_retrieve_timesteps, reference_timesteps, rtol=0, atol=0) + + native_transformer = SanaVideoTransformer3DModel.from_config(dict(reference.transformer.config)).to( + device="cuda", + dtype=torch.bfloat16, + ) + incompatible = native_transformer.load_state_dict(reference.transformer.state_dict(), strict=True) + assert not incompatible.missing_keys + assert not incompatible.unexpected_keys + assert type(native_transformer) is SanaVideoTransformer3DModel + assert type(native_transformer).__module__.startswith("vllm_omni.") + assert type(native_transformer) is not type(reference.transformer) + reference_parameter = next(reference.transformer.parameters()) + native_parameter = next(native_transformer.parameters()) + assert native_parameter.data_ptr() != reference_parameter.data_ptr() + torch.testing.assert_close(native_parameter, reference_parameter, rtol=0, atol=0) + native_transformer.eval() + reference.transformer.eval() + + latent_model_input = torch.cat([reference_latents, reference_latents]) + timestep = reference_timesteps[0].expand(latent_model_input.shape[0]) + with torch.inference_mode(): + reference_raw_prediction = reference.transformer( + latent_model_input.to(dtype=torch.bfloat16), + encoder_hidden_states=encoder_hidden_states.to(dtype=torch.bfloat16), + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + return_dict=False, + )[0] + native_raw_prediction = native_transformer( + latent_model_input.to(dtype=torch.bfloat16), + encoder_hidden_states=encoder_hidden_states.to(dtype=torch.bfloat16), + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + return_dict=False, + )[0] + + _assert_stage( + "first_transformer_raw_prediction", + native_raw_prediction, + reference_raw_prediction, + max_abs=config["max_abs"], + relative_l2=config["relative_l2"], + cosine=config["cosine"], + ) + reference_cfg_prediction = _cfg_prediction( + reference_raw_prediction, + guidance_scale=GUIDANCE_SCALE, + latent_channels=latent_channels, + out_channels=reference.transformer.config.out_channels, + ) + native_cfg_prediction = _cfg_prediction( + native_raw_prediction, + guidance_scale=GUIDANCE_SCALE, + latent_channels=latent_channels, + out_channels=native_transformer.config.out_channels, + ) + _assert_stage( + "first_cfg_prediction", + native_cfg_prediction, + reference_cfg_prediction, + max_abs=config["cfg_max_abs"], + relative_l2=config["cfg_relative_l2"], + cosine=config["cfg_cosine"], + ) + + reference_first_step = reference_scheduler.step( + reference_cfg_prediction, + reference_timesteps[0], + reference_latents, + return_dict=False, + )[0] + # Isolate scheduler orchestration from native-attention error by + # feeding the Diffusers prediction through an independently configured + # scheduler initialized by the native helper. + isolated_first_step = isolated_scheduler.step( + reference_cfg_prediction, + isolated_retrieve_timesteps[0], + reference_latents, + return_dict=False, + )[0] + _assert_stage( + "first_scheduler_step_isolated", + isolated_first_step, + reference_first_step, + max_abs=0.0, + relative_l2=0.0, + cosine=0.999999, + ) + + native_first_step = native_scheduler.step( + native_cfg_prediction, + native_timesteps[0], + native_latents, + return_dict=False, + )[0] + _assert_stage( + "first_scheduler_step_end_to_end", + native_first_step, + reference_first_step, + max_abs=config["step_max_abs"], + relative_l2=config["step_relative_l2"], + cosine=config["step_cosine"], + ) + finally: + del native_proxy + del reference + gc.collect() + torch.accelerator.empty_cache() diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index 68f890c64..be3d7f84c 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Run vLLM-Omni with diffusers backend, then run diffusers directly. Compare the output similarity. @@ -39,9 +39,9 @@ from tests.e2e.accuracy.helpers import ( from tests.e2e.accuracy.helpers import ( run_ffmpeg_similarity as _run_ffmpeg_similarity, ) -from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.clean import cleanup_test_environment from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniServer, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OnlineOmniClient from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import ( CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS, ) @@ -128,7 +128,7 @@ def _run_vllm_omni_wan22_i2v( "seed": SEED, } with OmniServer(model, server_args, env_dict=env_to_apply_ftfy_mock_in_subproc(), use_omni=True) as omni_server: - client = OpenAIClientHandler( + client = OnlineOmniClient( host=omni_server.host, port=omni_server.port, run_level="full_model", @@ -139,16 +139,20 @@ def _run_vllm_omni_wan22_i2v( "image_reference": f"data:image/png;base64,{pil_to_base64(conditioning_image, 'png')}", } result = client.send_video_diffusion_request(request_config)[0] - video_bytes = result.videos[0] # pyright: ignore[reportOptionalSubscript] # Guaranteed not None + videos = result.videos + assert videos is not None + video_bytes = videos[0] output_path.write_bytes(video_bytes) - return result.e2e_latency # pyright: ignore[reportReturnType] # Guaranteed not None + latency = result.e2e_latency + assert latency is not None + return latency def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_image: Image.Image) -> float: from diffusers import WanImageToVideoPipeline # pyright: ignore[reportPrivateImportUsage] from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler - run_pre_test_cleanup() + cleanup_test_environment() apply_ftfy_mock() pipe: WanImageToVideoPipeline | None = None try: @@ -190,7 +194,7 @@ def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_imag gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() - run_post_test_cleanup() + cleanup_test_environment() def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> tuple[Image.Image, float]: @@ -248,7 +252,7 @@ def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.Image, float]: - run_pre_test_cleanup() + cleanup_test_environment() pipe: DiffusionPipeline | None = None try: pipe = DiffusionPipeline.from_pretrained( @@ -288,7 +292,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() - run_post_test_cleanup() + cleanup_test_environment() @pytest.mark.benchmark diff --git a/tests/e2e/accuracy/test_hidream_o1_image.py b/tests/e2e/accuracy/test_hidream_o1_image.py new file mode 100644 index 000000000..923e53a47 --- /dev/null +++ b/tests/e2e/accuracy/test_hidream_o1_image.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +from PIL import Image + +from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + +MODEL_ID = "HiDream-ai/HiDream-O1-Image" +PROMPT = "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" +BASELINE_PATH = Path(__file__).resolve().parents[2] / "assets" / "hidream_o1" / "hidream_o1_baseline.png" +SSIM_THRESHOLD = 0.8 +PSNR_THRESHOLD = 19.0 + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +def test_hidream_o1_image_tp2_matches_baseline( + accuracy_artifact_root: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA") + model = os.environ.get("HIDREAM_O1_IMAGE_MODEL", MODEL_ID) + with OmniRunner( + model, + parallel_config=DiffusionParallelConfig(tensor_parallel_size=2), + enforce_eager=True, + ) as runner: + outputs = runner.omni.generate( + {"prompt": PROMPT}, + OmniDiffusionSamplingParams( + height=2048, + width=2048, + num_inference_steps=50, + guidance_scale=5.0, + seed=42, + ), + ) + images = extract_images_from_outputs(outputs) + assert images + image = images[0].convert("RGB") + + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + image.save(output_dir / "vllm_omni_tp2.png") + assert BASELINE_PATH.exists() + baseline = Image.open(BASELINE_PATH).convert("RGB") + resized = image.resize(baseline.size, Image.Resampling.LANCZOS) + assert_similarity( + model_name=MODEL_ID, + vllm_image=resized, + diffusers_image=baseline, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index 6f3ed7e71..a09014c8d 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import base64 @@ -13,7 +16,7 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline from PIL import Image from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir -from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.clean import cleanup_test_environment from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer @@ -29,14 +32,14 @@ HEIGHT = 512 NUM_INFERENCE_STEPS = 20 TRUE_CFG_SCALE = 4.0 SEED = 42 -SSIM_THRESHOLD = 0.97 -# The FA3-hub-matched comparison clears 0.97 with margin. The SDPA-matched -# fallback (hub kernel unavailable for the image's torch) landed at 0.9687 in -# build 2954 — the attention math matches but the surrounding stacks (vLLM -# compiled path vs eager diffusers) still differ slightly. Keep the strict -# gate for the FA3 path and a slightly looser one for the fallback only. +SSIM_THRESHOLD = 0.94 +# With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8 +# vs Diffusers (see #5734 / #5963), so the primary gates are 0.94 / 27.0. +# The SDPA-matched fallback (hub kernel unavailable for the image's torch) +# landed at SSIM 0.9687 in build 2954 — keep a slightly tighter SSIM gate for +# that path only. SSIM_THRESHOLD_SDPA_FALLBACK = 0.96 -PSNR_THRESHOLD = 30.0 +PSNR_THRESHOLD = 27.0 MODEL_2512_ID = "Qwen/Qwen-Image-2512" MODEL_2512_ENV_VAR = "QWEN_IMAGE_2512_MODEL" @@ -149,7 +152,7 @@ def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: - run_pre_test_cleanup() + cleanup_test_environment() pipe: DiffusionPipeline | None = None try: pipe = DiffusionPipeline.from_pretrained( @@ -179,7 +182,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() - run_post_test_cleanup() + cleanup_test_environment() def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image: @@ -211,7 +214,7 @@ def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image: - run_pre_test_cleanup() + cleanup_test_environment() pipe: DiffusionPipeline | None = None try: pipe = DiffusionPipeline.from_pretrained( @@ -241,7 +244,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() - run_post_test_cleanup() + cleanup_test_environment() @pytest.mark.benchmark diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py index 9ee9c2c87..9c3b9dd9c 100644 --- a/tests/e2e/accuracy/test_qwen_image_edit.py +++ b/tests/e2e/accuracy/test_qwen_image_edit.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import gc @@ -12,12 +15,12 @@ from PIL import Image from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir -from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.clean import cleanup_test_environment from tests.helpers.mark import hardware_test from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServer -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.diffusion] SINGLE_MODEL = "Qwen/Qwen-Image-Edit" @@ -109,7 +112,7 @@ def _run_diffusers_image_edit( input_images: list[Image.Image], output_path: Path, ) -> Image.Image: - run_pre_test_cleanup() + cleanup_test_environment() pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None try: images = input_images[0] if len(input_images) == 1 else input_images @@ -140,7 +143,7 @@ def _run_diffusers_image_edit( gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() - run_post_test_cleanup() + cleanup_test_environment() def _load_edit_2511_input_images() -> list[Image.Image]: @@ -189,7 +192,7 @@ def _run_diffusers_image_edit_2511( input_images: list[Image.Image], output_path: Path, ) -> Image.Image: - run_pre_test_cleanup() + cleanup_test_environment() pipe: QwenImageEditPlusPipeline | None = None try: pipe = QwenImageEditPlusPipeline.from_pretrained( @@ -222,7 +225,7 @@ def _run_diffusers_image_edit_2511( gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() - run_post_test_cleanup() + cleanup_test_environment() def _vllm_omni_output_single_image( @@ -253,6 +256,7 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image ) +@pytest.mark.slow @pytest.mark.benchmark @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_edit_single_matches_diffusers( @@ -278,6 +282,7 @@ def test_qwen_image_edit_single_matches_diffusers( ) +@pytest.mark.slow @pytest.mark.benchmark @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None: diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py index f2bc1b6aa..53354a739 100644 --- a/tests/e2e/accuracy/test_qwen_image_layered.py +++ b/tests/e2e/accuracy/test_qwen_image_layered.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import base64 @@ -13,11 +16,11 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline from PIL import Image from tests.e2e.accuracy.helpers import assert_image_sequence_similarity, model_output_dir -from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.clean import cleanup_test_environment from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] MODEL_ID = "Qwen/Qwen-Image-Layered" @@ -93,7 +96,7 @@ def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, o def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]: - run_pre_test_cleanup() + cleanup_test_environment() pipe: DiffusionPipeline | None = None try: pipe = DiffusionPipeline.from_pretrained( @@ -126,7 +129,7 @@ def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, o gc.collect() if torch.cuda.is_available(): torch.accelerator.empty_cache() - run_post_test_cleanup() + cleanup_test_environment() @pytest.mark.benchmark diff --git a/tests/e2e/accuracy/test_video_streaming_output_similarity.py b/tests/e2e/accuracy/test_video_streaming_output_similarity.py index 799e6e125..cdb2c4a28 100644 --- a/tests/e2e/accuracy/test_video_streaming_output_similarity.py +++ b/tests/e2e/accuracy/test_video_streaming_output_similarity.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Online serving smoke test for streaming video output similarity. """ @@ -9,9 +12,9 @@ import pytest from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_video from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler +from tests.helpers.runtime import DiffusionResponse, OmniServer, OnlineOmniClient, get_model_prefix -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "BestWishYsh/Helios-Distilled" PROMPT = "A serene lakeside sunrise with mist over the water." @@ -67,8 +70,8 @@ def _request_config(model: str, *, serialize_extra_params: bool) -> dict: } -def _client_for(server: OmniServer, *, run_level: str) -> OpenAIClientHandler: - return OpenAIClientHandler( +def _client_for(server: OmniServer, *, run_level: str) -> OnlineOmniClient: + return OnlineOmniClient( host=server.host, port=server.port, api_key="EMPTY", @@ -95,12 +98,11 @@ def _write_video_artifact(response: DiffusionResponse, output_path: Path) -> Pat ) def test_helios_streaming_output_matches_non_streaming( model: str, - model_prefix: str, run_level: str, tmp_path: Path, record_property, ): - model = model_prefix + model + model = get_model_prefix() + model with OmniServer(model, _server_args(streaming_output=True), use_omni=True) as server: client = _client_for(server, run_level=run_level) diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py index a09d2330a..a995a9cf9 100644 --- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py +++ b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os @@ -338,7 +341,7 @@ def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]: def _generate_online_video( *, omni_server, - openai_client, + online_client, image_source: str, online_timeout_seconds: int, ) -> Path: @@ -360,7 +363,7 @@ def _generate_online_video( "image_reference": _build_online_image_reference(image_source), } online_video_bytes = _send_video_request_with_timeout( - openai_client, + online_client, request_config, timeout_seconds=_online_timeout_seconds(online_timeout_seconds), ) @@ -412,7 +415,7 @@ def test_wan22_i2v_diffusers_offline_generates_video( @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_wan22_i2v_online_serving_generates_video( omni_server, - openai_client, + online_client, wan22_i2v_image_source: str | None, wan22_i2v_online_timeout_seconds: int, ) -> None: @@ -424,7 +427,7 @@ def test_wan22_i2v_online_serving_generates_video( _validate_image_source(image_source) online_path = _generate_online_video( omni_server=omni_server, - openai_client=openai_client, + online_client=online_client, image_source=image_source, online_timeout_seconds=wan22_i2v_online_timeout_seconds, ) diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py index 5dcb9c449..61cf76dcd 100644 --- a/tests/e2e/features/comfyui/test_comfyui_integration.py +++ b/tests/e2e/features/comfyui/test_comfyui_integration.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Integration tests for ComfyUI nodes that use the Omni API client, with a mocked AsyncOmni and a real API server running in a background process. These tests cover the integration between ComfyUI node and the API server, without actual model inference logic. @@ -39,6 +42,7 @@ from pytest_mock import MockerFixture from vllm import SamplingParams from vllm.outputs import CompletionOutput, RequestOutput +from tests.helpers.runtime import get_open_port from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni from vllm_omni.entrypoints.cli.serve import OmniServeCommand from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams @@ -515,7 +519,7 @@ def mock_async_omni( @pytest.fixture -def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni, tmp_path): +def api_server(server_case: ServerCase, mock_async_omni, tmp_path): """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni.""" # Override the STORAGE_MANAGER path to a writable temp directory before # forking the server subprocess. The default /tmp/storage may be owned @@ -533,7 +537,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni cmd = OmniServeCommand() cmd.subparser_init(subparsers) - port = unused_tcp_port_factory() + port = get_open_port(host="0.0.0.0") args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)]) def run_server(): @@ -550,11 +554,12 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni wait_poll_interval = 1 for _ in range(wait_time // wait_poll_interval): try: - response = requests.get(f"http://127.0.0.1:{port}/health") + response = requests.get(f"http://127.0.0.1:{port}/health", timeout=wait_poll_interval) if response.status_code == 200: break - except requests.ConnectionError: - time.sleep(wait_poll_interval) + except requests.RequestException: + pass + time.sleep(wait_poll_interval) else: if server_process.is_alive(): server_process.terminate() @@ -809,7 +814,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_ ServerCase( served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", stage_list=["diffusion"], - stage_configs=[{"stage_type": "diffusion"}], + stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}], outputs=[_build_diffusion_video_output()], ), "Wan-AI/Wan2.2-T2V-A14B-Diffusers", @@ -820,7 +825,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_ ServerCase( served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", stage_list=["diffusion"], - stage_configs=[{"stage_type": "diffusion"}], + stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}], outputs=[_build_diffusion_video_output()], ), "Wan-AI/Wan2.2-I2V-A14B-Diffusers", @@ -881,7 +886,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b ServerCase( served_model="MiniMaxAI/MiniMax-H3", stage_list=["diffusion"], - stage_configs=[{"stage_type": "diffusion"}], + stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}], outputs=[_build_diffusion_video_output()], ), SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None), @@ -892,7 +897,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b ServerCase( served_model="MiniMaxAI/MiniMax-H3", stage_list=["diffusion"], - stage_configs=[{"stage_type": "diffusion"}], + stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}], outputs=[_build_diffusion_video_output()], ), SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None), diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py index 04e83caeb..a11c9b7e4 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py @@ -166,8 +166,8 @@ async def test_sleep_memory_reclaimed_custom_pipeline(): The fix moves custom_pipeline init outside the CUDA context so all weights go through the caching allocator and are therefore fully reclaimed by - sleep(level=1). A non-zero ``CuMemAllocator.get_current_usage()`` after - sleep is the direct signal that the bypass is still occurring. + sleep(level=1). The tracked allocations must transition to the allocator's + asleep state and report physical reclamation. """ with ExitStack() as after: engine = AsyncOmni( @@ -187,62 +187,72 @@ async def test_sleep_memory_reclaimed_custom_pipeline(): free_before, total = torch.cuda.mem_get_info() used_before_gib = (total - free_before) / 1024**3 - # Measure CuMemAllocator-tracked usage before sleep. In inline mode - # the worker runs in a thread pool inside this process, so the allocator - # singleton is shared and can be read directly. + # Measure CuMemAllocator-tracked usage before sleep. In inline / uni + # mode the worker runs in this process, so the allocator singleton is + # shared and can be read directly. allocator = None tracked_before = 0 + tracked_ptrs: set[int] = set() try: from vllm.device_allocator.cumem import CuMemAllocator allocator = CuMemAllocator.get_instance() - tracked_before = allocator.get_current_usage() + tracked_ptrs = {ptr for ptr, data in allocator.pointer_to_data.items() if not data.is_asleep} + tracked_before = sum(allocator.pointer_to_data[ptr].handle[1] for ptr in tracked_ptrs) except Exception: pass - # Put the engine to sleep; all weights should be offloaded via the pool. - acks = await engine.sleep(level=1) - await asyncio.sleep(0.5) # allow the CUDA driver to settle - torch.accelerator.synchronize() + slept = False + try: + # Put the engine to sleep; this engine's weights should be + # offloaded via the pool. + acks = await engine.sleep(level=1) + slept = True + await asyncio.sleep(0.5) # allow the CUDA driver to settle + torch.accelerator.synchronize() - # Measure after sleep. - free_after, _ = torch.cuda.mem_get_info() - used_after_gib = (total - free_after) / 1024**3 - drop_gib = used_before_gib - used_after_gib + # Measure after sleep. + free_after, _ = torch.cuda.mem_get_info() + used_after_gib = (total - free_after) / 1024**3 + drop_gib = used_before_gib - used_after_gib - # --- Primary assertion: allocator reports zero tracked memory. --- - # If this fails it means weights were allocated outside the CuMem pool - # (safetensors direct-to-GPU bypass) — the exact regression this test - # is designed to catch. - if allocator is not None: - tracked_after = allocator.get_current_usage() - assert tracked_after == 0, ( - f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB " - f"after sleep(level=1) on custom_pipeline path " - f"(was {tracked_before / 1024**3:.3f} GiB before sleep). " - "Weights were allocated outside the CuMem pool via the " - "safetensors direct-to-GPU fast path — loader-context fix " - "may not be applied." + # get_current_usage() includes unmapped sleeping handles, so it + # does not decrease on sleep. Verify that allocations mapped + # before the call instead transitioned to the asleep state. + if allocator is not None: + assert tracked_before > 0, ( + "Expected custom_pipeline weights to use the CuMem pool, " + "but no mapped allocations were tracked before sleep." + ) + still_awake = [ + ptr + for ptr in tracked_ptrs + if ptr in allocator.pointer_to_data and not allocator.pointer_to_data[ptr].is_asleep + ] + assert not still_awake, ( + f"{len(still_awake)} CuMem allocation(s) remained mapped " + "after sleep(level=1) on the custom_pipeline path." + ) + + # Physical VRAM or ACK freed_bytes confirms reclamation at the + # driver level. + total_freed_bytes = sum( + (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0)) + for ack in acks + if ack is not None + ) + freed_gib = total_freed_bytes / 1024**3 + assert freed_gib > 0 or drop_gib > 0, ( + f"Expected GPU memory to be reclaimed after sleep(level=1) on " + f"custom_pipeline + enable_sleep_mode=True. " + f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, " + f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB." ) - # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms - # reclamation at the driver level. - total_freed_bytes = sum( - (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0)) - for ack in acks - if ack is not None - ) - freed_gib = total_freed_bytes / 1024**3 - assert freed_gib > 0 or drop_gib > 0, ( - f"Expected GPU memory to be reclaimed after sleep(level=1) on " - f"custom_pipeline + enable_sleep_mode=True. " - f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, " - f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB." - ) - - # Engine must report it is sleeping. - assert await engine.is_sleeping() - - # Wake up and confirm the engine is functional again. - await engine.wake_up() - assert not await engine.is_sleeping() + assert await engine.is_sleeping() + await engine.wake_up() + slept = False + assert not await engine.is_sleeping() + finally: + if slept: + await engine.wake_up() diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py index b24f364e4..12ad4f1b3 100644 --- a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py +++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py @@ -17,9 +17,12 @@ from functools import lru_cache import numpy as np import pytest +import torch.distributed as dist from transformers import AutoTokenizer from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor +from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -131,6 +134,19 @@ async def _generate_once( return last_output +def _assert_live_uniproc_worker(engine: AsyncOmni) -> None: + """Default single-GPU path must use the in-process executor, not mp IPC.""" + clients = engine.engine.stage_clients + assert len(clients) == 1 + client = clients[0] + assert isinstance(client, InlineStageDiffusionClient) + executor = client._engine.executor + assert isinstance(executor, UniProcDiffusionExecutor) + assert executor.driver_worker is not None + assert dist.is_initialized() + assert dist.get_world_size() == 1 + + def _assert_valid_image_output(output: OmniRequestOutput) -> None: assert output.final_output_type == "image" assert output.images, "Expected at least one generated image" @@ -158,6 +174,7 @@ async def test_async_omni_generate(): ) after.callback(engine.shutdown) + _assert_live_uniproc_worker(engine) output = await _generate_once( engine, "a beautiful sunset over the ocean with vibrant orange and purple clouds " diff --git a/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py b/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py index abb387573..65d4f1ca9 100644 --- a/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py +++ b/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py @@ -8,6 +8,7 @@ from dataclasses import FrozenInstanceError import pytest +from tests.helpers.stage_config import get_deploy_duplex_max_sessions from vllm_omni.config.stage_config import ( DuplexSessionRuntimeConfig, load_deploy_config, @@ -66,3 +67,18 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"): load_deploy_config(deploy_path) + + +@pytest.mark.parametrize( + ("deploy_yaml", "expected"), + [ + ("minicpmo_4_5.yaml", 4), + ("minicpmo_4_5_8x4090.yaml", 1), + ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4), + ], +) +def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None: + # Guards the admission probe. A capacity edit, a config that declares none, + # or an overlay inheriting one from its base must surface here rather than + # as another nightly duplex admission timeout. + assert get_deploy_duplex_max_sessions(deploy_yaml) == expected diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py new file mode 100644 index 000000000..5790b11f1 --- /dev/null +++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest + +from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import ( + NemotronVoiceChatDataPlaneContext, + NemotronVoiceChatDataPlaneSession, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +_RUNTIME = { + "nvc_text_bos_id": 0, + "nvc_text_eos_id": 1, + "nvc_text_pad_id": 12, + "nvc_function_sotc_id": 20, + "nvc_function_eotc_id": 21, + "nvc_tokenizer_ref": "test-tokenizer", +} + + +def _projector(encode_audio=lambda *_: None) -> NemotronVoiceChatDataPlaneSession: + projector = NemotronVoiceChatDataPlaneSession(encode_audio) + projector.configure_runtime(_RUNTIME, tokenizer=SimpleNamespace(decode=lambda *_args, **_kwargs: "text")) + return projector + + +def _stage0_output(request_id: str = "req-0", *, text_token: int = 12, function_token: int | None = None) -> object: + metadata: dict[str, object] = {"nvc_text_token_ids": [text_token]} + if function_token is not None: + metadata["nvc_function_token"] = [function_token] + completion = SimpleNamespace(multimodal_output=metadata) + return SimpleNamespace( + stage_id=0, + request_output=SimpleNamespace(request_id=request_id, outputs=[completion]), + ) + + +def test_function_channel_projects_completed_call_without_ending_speech() -> None: + projector = _projector() + projector._decode = lambda token_ids: ( + '[{"name":"weather","arguments":{"city":"Shanghai"}}]' if token_ids == [99] else "" + ) + context = NemotronVoiceChatDataPlaneContext(epoch=0) + + events = [] + for function_token in (20, 99, 21): + events.extend(projector.project_output(_stage0_output(function_token=function_token), context=context)) + + listen = [event for event in events if event.get("is_listen") is True] + function = [event for event in events if event.get("function_call") is True] + assert len(listen) == 3 + assert len(function) == 1 + assert function[0]["name"] == "weather" + assert function[0]["arguments"] == '{"city":"Shanghai"}' + + projector._decode = lambda _token_ids: "not-json" + bad_events = [] + for function_token in (20, 98, 21): + bad_events.extend(projector.project_output(_stage0_output(function_token=function_token))) + assert any(event.get("error_code") == "nemotron_function_call_parse_error" for event in bad_events) + + +def test_new_epoch_request_does_not_inherit_partial_function_or_frame_state() -> None: + projector = _projector(lambda *_: "audio") + old_request = duplex_resource_request_id(DuplexFence("sid", epoch=1), "stage0") + new_request = duplex_resource_request_id(DuplexFence("sid", epoch=2), "stage0") + + projector.begin_request(old_request) + list(projector.project_output(_stage0_output(old_request, text_token=42, function_token=20))) + projector.begin_request(new_request) + + # EOTC in the new epoch must not close the old epoch's partial function. + new_events = list(projector.project_output(_stage0_output(new_request, text_token=1, function_token=21))) + assert not [event for event in new_events if event.get("function_call") is True] + assert not [event for event in new_events if event.get("end_of_turn") is True] + + audio = SimpleNamespace( + stage_id=2, + request_output=SimpleNamespace( + request_id=new_request, + outputs=[ + SimpleNamespace(multimodal_output={"model_outputs": [np.ones(1764, dtype=np.float32)], "sr": [22050]}) + ], + ), + ) + audio_events = list(projector.project_output(audio)) + assert sum(event.get("end_of_turn") is True for event in audio_events) == 1 diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py new file mode 100644 index 000000000..84c8b123a --- /dev/null +++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import base64 + +import numpy as np +import pytest + +from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import ( + NemotronVoiceChatPcmAppendBuffer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _payload(samples: np.ndarray, *, sample_rate_hz: int = 16000, fmt: str = "pcm_f32le") -> dict[str, object]: + return { + "type": "audio", + "audio": base64.b64encode(np.asarray(samples, dtype=np.float32).tobytes()).decode("ascii"), + "format": fmt, + "sample_rate_hz": sample_rate_hz, + } + + +def _append(buffer, samples, operation_id): + return buffer.prepare_append( + _payload(np.asarray(samples, dtype=np.float32)), + operation_id=operation_id, + chunk_period_ms=80, + allow_emit=True, + ) + + +def test_irregular_browser_packets_emit_exact_ordered_80ms_frame() -> None: + buffer = NemotronVoiceChatPcmAppendBuffer() + assert _append(buffer, np.arange(1000), "packet-1") is None + + reservation = _append(buffer, np.arange(1000, 1280), "packet-2") + + assert reservation is not None + assert reservation.operation_id == "packet-2" + decoded = np.frombuffer(base64.b64decode(reservation.payload["audio"]), dtype=np.float32) + np.testing.assert_array_equal(decoded, np.arange(1280, dtype=np.float32)) + assert buffer.pending_byte_count == 0 + + +def test_append_rejects_invalid_frame_contract() -> None: + cases = ( + (1281, 16000, "pcm_f32le", 0.0, "at most 1280"), + (1, 8000, "pcm_f32le", 0.0, "sample_rate_hz"), + (1, 16000, "audio/pcm", 0.0, "format"), + (1, 16000, "pcm_f32le", np.nan, "finite"), + ) + for samples, sample_rate_hz, fmt, value, message in cases: + with pytest.raises(ValueError, match=message): + NemotronVoiceChatPcmAppendBuffer().prepare_append( + _payload(np.full(samples, value, dtype=np.float32), sample_rate_hz=sample_rate_hz, fmt=fmt), + operation_id="invalid", + chunk_period_ms=80, + allow_emit=True, + ) + + +def test_terminal_commit_rejects_more_than_one_frame_tail() -> None: + buffer = NemotronVoiceChatPcmAppendBuffer() + for index, samples in enumerate((1000, 281)): + buffer.prepare_append( + _payload(np.zeros(samples, dtype=np.float32)), + operation_id=f"tail-{index}", + chunk_period_ms=80, + allow_emit=False, + ) + with pytest.raises(ValueError, match="more than one 80 ms frame"): + buffer.prepare_commit(operation_id="commit", chunk_period_ms=80) diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py new file mode 100644 index 000000000..0b32322c0 --- /dev/null +++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import base64 +from types import SimpleNamespace + +import numpy as np +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.experimental.fullduplex.engine.contracts import DuplexInputMode +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.nemotron_voicechat.runtime import ( + NemotronVoiceChatDuplexRuntimeExtension, +) +from vllm_omni.experimental.fullduplex.nemotron_voicechat.serving_adapter import ( + NemotronVoiceChatServingRuntimeAdapter, + _render_tool_response, +) +from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import ( + NemotronVoiceChatThinkerForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _frame() -> dict[str, object]: + raw = np.zeros(1280, dtype=np.float32).tobytes() + return { + "type": "audio", + "audio": base64.b64encode(raw).decode("ascii"), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + } + + +def _plan(extension, *, input_seq: int): + return extension.plan_append( + request_id="req", + fence=DuplexFence("sid", incarnation=2, epoch=3), + session_config={}, + runtime_config={ + "nvc_prompt_token_ids": [0, 42, 1], + "nvc_text_pad_id": 12, + "nvc_max_model_len": 8192, + }, + seq=input_seq, + turn_seq=input_seq, + mode=DuplexInputMode.APPEND_AUDIO_CHUNK, + payload=_frame(), + final=False, + sampling_params=SamplingParams(), + ) + + +def test_first_append_prefills_prompt_then_each_append_consumes_one_frame() -> None: + extension = NemotronVoiceChatDuplexRuntimeExtension() + + params = extension.configure_sampling_params( + runtime_config={}, defaults=(SamplingParams(temperature=0.7, max_tokens=4),) + )[0] + assert params.temperature == 0.0 and params.max_tokens == 1 and params.top_p == 1.0 and params.top_k == 0 + + first = _plan(extension, input_seq=1) + later = _plan(extension, input_seq=2) + + assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12] + assert later.prompt["prompt_token_ids"] == [12] + assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1 + assert later.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2 + + +def test_append_rejects_stage0_context_overflow() -> None: + with pytest.raises(ValueError, match="max_model_len"): + _plan(NemotronVoiceChatDuplexRuntimeExtension(), input_seq=8190) + + +def test_function_output_becomes_versioned_nvidia_channel_tokens() -> None: + encoded: list[str] = [] + adapter = NemotronVoiceChatServingRuntimeAdapter(lambda *_: None) + adapter._tokenizer = SimpleNamespace( + encode=lambda text, **_kwargs: encoded.append(text) or [31, 32, 33], + ) + + first = adapter.runtime_config_for_function_output( + {"type": "function_call_output", "call_id": "call-1", "output": '{"result":20}'}, + {}, + ) + second = adapter.runtime_config_for_function_output( + {"type": "function_call_output", "call_id": "call-2", "output": "plain text"}, + first, + ) + + assert _render_tool_response('{"result":20}') == '[{"result":20}]' + assert encoded == [ + '[{"result":20}]', + '["plain text"]', + ] + assert first["nvc_function_response_generation"] == 1 + assert second["nvc_function_response_generation"] == 2 + assert second["nvc_function_response_token_ids"] == [31, 32, 33] + assert second["nvc_function_response_call_id"] == "call-2" + assert second["nvc_function_response_batches"] == [ + {"generation": 1, "call_id": "call-1", "token_ids": [31, 32, 33]}, + {"generation": 2, "call_id": "call-2", "token_ids": [31, 32, 33]}, + ] + + session = {"function_response_generation": 0, "forced_function_tokens": []} + NemotronVoiceChatThinkerForConditionalGeneration._sync_forced_function_response(session, second) + + assert session["function_response_generation"] == 2 + assert session["forced_function_token"] == 31 + assert session["forced_function_tokens"] == [31, 32, 33, 31, 32, 33] diff --git a/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py b/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py index f4a3c994e..f2970a240 100644 --- a/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py +++ b/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py @@ -18,6 +18,10 @@ from vllm_omni.experimental.fullduplex.engine.duplex_runtime import ( DuplexInputMode, ) from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ( + ServingRuntimeConfigError, +) +from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA from vllm_omni.experimental.fullduplex.personaplex.data_plane import ( PersonaPlexDataPlaneContext, PersonaPlexDataPlaneSession, @@ -242,6 +246,53 @@ def test_personaplex_capabilities_are_honest() -> None: assert multi.supports_barge_in is False +def test_personaplex_runtime_config_update_rejects_changed_persona() -> None: + adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None) + short_persona = "You are a helpful assistant." + longer_persona = ( + "You are now a swashbuckling pirate captain who speaks entirely in " + "nautical slang, references buried treasure constantly, and never " + "breaks character no matter what the user asks." + ) + current = {"personaplex_persona": short_persona, "personaplex_voice_prompt": "NATF2.pt"} + + with pytest.raises(ServingRuntimeConfigError, match="persona"): + adapter.runtime_config_for_update( + SimpleNamespace(instructions=longer_persona, extra_body={}), + current, + ) + + unchanged = adapter.runtime_config_for_update( + SimpleNamespace(instructions=short_persona, extra_body={}), + current, + ) + assert unchanged["personaplex_persona"] == short_persona + + +def test_personaplex_runtime_config_update_rejects_changed_voice() -> None: + adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None) + current = {"personaplex_persona": DEFAULT_PERSONA, "personaplex_voice_prompt": "NATF2.pt"} + + with pytest.raises(ServingRuntimeConfigError, match="voice"): + adapter.runtime_config_for_update( + SimpleNamespace(instructions=None, voice="OtherVoice.pt", extra_body={}), + current, + ) + + unchanged = adapter.runtime_config_for_update( + SimpleNamespace(instructions=None, voice="NATF2.pt", extra_body={}), + current, + ) + assert unchanged["personaplex_voice_prompt"] == "NATF2.pt" + + +def test_personaplex_rejects_minicpmo_only_opt_out_flag() -> None: + adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None) + + with pytest.raises(ServingRuntimeConfigError, match="minicpmo45_native_duplex"): + adapter.validate_client_extra_body({"minicpmo45_native_duplex": False}) + + def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None: buffer = PersonaPlexPcmAppendBuffer() reservation = buffer.prepare_append( diff --git a/tests/e2e/features/fullduplex/test_client.py b/tests/e2e/features/fullduplex/test_client.py index d90c9e5aa..53965b89e 100644 --- a/tests/e2e/features/fullduplex/test_client.py +++ b/tests/e2e/features/fullduplex/test_client.py @@ -1,3 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio import base64 import wave from urllib.parse import parse_qs, urlsplit @@ -14,13 +18,17 @@ from vllm_omni.experimental.fullduplex.client import ( from vllm_omni.experimental.fullduplex.minicpmo45.policy import ( MiniCPMO45DuplexPolicy, ) +from vllm_omni.experimental.fullduplex.video_stacking import ( + concat_frames, + unit_subframe_offsets, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_realtime_client_builds_explicit_native_duplex_url(): url = build_realtime_url( - "ws://localhost:8099/v1/realtime?custom=1", + "ws://localhost:8099/v1/realtime?custom=1&duplex=0&model=stale&minicpmo45_native_duplex=0&session_id=stale", "openbmb/MiniCPM-o-4_5", session_id="session-a", ) @@ -50,7 +58,7 @@ def test_seed_tts_initial_text_is_part_of_native_duplex_context(): def test_realtime_client_builds_resume_only_url_when_autostart_disabled(): url = build_realtime_url( - "ws://localhost:8099/v1/realtime?duplex=1", + "ws://localhost:8099/v1/realtime?duplex=1&autostart=1", "openbmb/MiniCPM-o-4_5", autostart=False, ) @@ -73,10 +81,15 @@ async def test_realtime_client_configure_omits_ref_audio_by_default(): client = Client() - await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1) + await client.configure( + "openbmb/MiniCPM-o-4_5", + idle_timeout_s=900, + timeout_s=1, + ) session = client.sent[0]["session"] assert "ref_audio" not in session + assert session["idle_timeout_s"] == 900 @pytest.mark.asyncio @@ -100,6 +113,74 @@ async def test_realtime_client_configure_sends_explicit_ref_audio(): session = client.sent[0]["session"] assert session["ref_audio"] == "data:audio/wav;base64,AAAA" + assert "idle_timeout_s" not in session + + +@pytest.mark.asyncio +async def test_realtime_client_reports_reader_exit_to_event_waiters(): + client = RealtimeDuplexClient("ws://unused") + + async def stopped(): + return None + + client._reader_task = asyncio.create_task(stopped()) + await client._reader_task + + with pytest.raises(ConnectionError, match="closed before"): + client.raise_if_reader_stopped() + + +@pytest.mark.asyncio +async def test_realtime_client_configure_reports_reader_failure_immediately(): + class Client(RealtimeDuplexClient): + async def send(self, event): + async def fail(): + raise RuntimeError("reader failed") + + self._reader_task = asyncio.create_task(fail()) + await asyncio.sleep(0) + + client = Client("ws://unused") + with pytest.raises(ConnectionError, match="reader failed"): + await client.configure("model", timeout_s=10.0) + + +@pytest.mark.asyncio +async def test_realtime_client_close_reports_reader_cancellation_immediately(): + class Client(RealtimeDuplexClient): + async def send(self, event): + async def wait_forever(): + await asyncio.Future() + + self._reader_task = asyncio.create_task(wait_forever()) + self._reader_task.cancel() + with pytest.raises(asyncio.CancelledError): + await self._reader_task + + client = Client("ws://unused") + with pytest.raises(ConnectionError, match="reader was cancelled"): + await client.close_session(timeout_s=10.0) + + +@pytest.mark.asyncio +async def test_realtime_client_close_waits_for_a_new_session_closed_event(): + class Client(RealtimeDuplexClient): + async def send(self, event): + assert event["type"] == "session.close" + + async def acknowledge_close(): + await asyncio.sleep(0.02) + self.events.add({"type": "session.closed"}) + + asyncio.create_task(acknowledge_close()) + + client = Client("ws://unused") + client.events.add({"type": "session.closed", "reason": "old"}) + started_at = asyncio.get_running_loop().time() + + await client.close_session(timeout_s=0.2) + + assert asyncio.get_running_loop().time() - started_at >= 0.01 @pytest.mark.asyncio @@ -235,6 +316,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): "output_token_count": 4, "ttft_ms": 120.0, "tpot_ms": 15.0, + "itls_ms": [10.0, 14.0, 18.0], "inter_token_interval_ms": { "count": 3, "mean": 14.0, @@ -340,3 +422,62 @@ def test_realtime_client_pcm16_wav_round_trip(tmp_path): assert wav_file.getnchannels() == 1 assert wav_file.getframerate() == 16_000 assert read_pcm16_wav(path) == pcm16 + + +def _stack_test_image(color, size=(40, 30)): + from PIL import Image + + return Image.new("RGB", size, color) + + +def test_video_stacking_tiles_a_units_subframes_into_one_image(): + frames = [_stack_test_image((255, 0, 0)), _stack_test_image((0, 255, 0))] + + composite = concat_frames(frames) + + # Landscape cells stack into a column because that canvas is squarer, and + # the interior seam carries the separator band that lets the model tell the + # sub-frames apart. + assert composite.size == (40, 30 * 2 + 6) + assert composite.getpixel((20, 30 + 3)) == (0, 0, 0) + + +def test_video_stacking_picks_the_squarest_grid(): + portrait = [_stack_test_image((255, 0, 0), size=(20, 60)) for _ in range(2)] + four = [_stack_test_image((0, 0, 255)) for _ in range(4)] + + # A portrait clip tiles side by side, four frames always tile 2x2. + assert concat_frames(portrait).size == (20 * 2 + 6, 60) + assert concat_frames(four).size == (40 * 2 + 6, 30 * 2 + 6) + + +def test_video_stacking_skips_the_subframe_that_duplicates_the_base_frame(): + # Official samples stack_frames=5 as 0.2/0.4/0.6/0.8 s: offset 0 is the base + # frame, already sent as frame_list[0]. + assert unit_subframe_offsets(5) == pytest.approx([0.2, 0.4, 0.6, 0.8]) + assert unit_subframe_offsets(1) == [] + + +def test_realtime_client_sends_each_units_composite_beside_its_base_frame(): + sent: list[dict] = [] + + class _Client(RealtimeDuplexClient): + def __init__(self): + pass + + async def send(self, event): + sent.append(event) + + asyncio.run( + _Client().stream_pcm16( + b"\x01\x00" * (16_000 * 3), + chunk_ms=200, + realtime=False, + video_frames=["f0", "f1"], + stacked_video_frames=["s0", None], + ) + ) + + # A composite belongs to the unit it was captured in, so it rides the same + # append as that unit's base frame; a unit without one sends the base alone. + assert [event["video_frames"] for event in sent if "video_frames" in event] == [["f0", "s0"], ["f1"]] diff --git a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py new file mode 100644 index 000000000..192859de5 --- /dev/null +++ b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import argparse +import base64 +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest + +from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner +from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample +from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge +from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli +from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark] + + +def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys): + manifest = tmp_path / "manifest.json" + manifest.write_text( + json.dumps( + [ + { + "id": "sample-1", + "split": "PR_correction", + "question_text": "What changed?", + "answer1": "The object moved.", + } + ] + ), + encoding="utf-8", + ) + response_root = tmp_path / "responses" + score_root = tmp_path / "scores" + + async def fake_generate(sample, *, output_root, **kwargs): + output = Path(output_root) / sample.split / f"{sample.id}.json" + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8") + output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8") + return output + + class FakeJudge: + def __init__(self, *args, **kwargs): + pass + + def chat(self, *args, **kwargs): + return '{"success_score": 1, "is_relevant": 1}' + + monkeypatch.setattr(cli, "generate_sample", fake_generate) + monkeypatch.setattr(cli, "DuplexJudge", FakeJudge) + + parser = argparse.ArgumentParser() + OmniDuplexEvalSubcommand.add_cli_args(parser) + + common = ["--dataset", str(manifest), "--family", "pr"] + generate = parser.parse_args( + [ + "generate", + *common, + "--model", + "mock", + "--ref-audio", + str(manifest), + "--response-root", + str(response_root), + "--concurrency", + "2", + ] + ) + OmniDuplexEvalSubcommand.cmd(generate) + evaluate = parser.parse_args( + [ + "evaluate", + *common, + "--response-root", + str(response_root), + "--score-root", + str(score_root), + "--judge-model", + "mock-judge", + "--eval-workers", + "2", + ] + ) + OmniDuplexEvalSubcommand.cmd(evaluate) + OmniDuplexEvalSubcommand.cmd(parser.parse_args(["summarize", "--score-root", str(score_root)])) + + summary = json.loads(capsys.readouterr().out) + assert summary["samples"] == 1 + assert summary["pr"]["mean_all_success"] == 1.0 + + +@pytest.mark.asyncio +async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + import websockets + + received = [] + + async def handler(websocket): + response_started = False + async for raw in websocket: + event = json.loads(raw) + received.append(event) + if event["type"] == "session.update": + await websocket.send(json.dumps({"type": "session.created"})) + elif event["type"] == "input_audio_buffer.append" and not response_started: + response_started = True + await websocket.send(json.dumps({"type": "response.created", "response": {"id": "r1"}})) + await websocket.send( + json.dumps({"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."}) + ) + await websocket.send( + json.dumps( + { + "type": "response.audio.delta", + "response_id": "r1", + "delta": base64.b64encode(b"\0\0" * 240).decode(), + "sample_rate_hz": 24_000, + } + ) + ) + elif event["type"] == "input_audio_buffer.commit": + await websocket.send(json.dumps({"type": "response.done", "response": {"id": "r1"}})) + elif event["type"] == "session.close": + await websocket.send(json.dumps({"type": "session.closed"})) + return + + monkeypatch.setattr(runner, "read_audio_pcm16", lambda path: b"\0\0" * (16_000 * 8 // 10)) + monkeypatch.setattr(runner, "video_duration", lambda path: 0.8) + monkeypatch.setattr(runner, "iter_jpegs", lambda *args, **kwargs: iter([(0.0, b"jpeg")])) + audio = tmp_path / "question.wav" + video = tmp_path / "video.mp4" + ref = tmp_path / "ref.wav" + for path in (audio, video, ref): + path.write_bytes(b"data") + sample = DuplexSample("sample", "PR_correction", "pr", "correction", video, audio) + + async with websockets.serve(handler, "127.0.0.1", 0) as server: + port = server.sockets[0].getsockname()[1] + output = await runner.generate_sample( + sample, + url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1", + model="mock", + ref_audio=ref, + output_root=tmp_path / "responses", + ) + + assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}] + meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8")) + assert meta["response_done"] is True + assert meta["drain_timeout"] is None + assert any(event["type"] == "playback.ack" for event in received) + + +def test_judge_exercises_openai_http_schema(): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + length = int(self.headers["Content-Length"]) + requests.append((self.path, self.headers["Authorization"], json.loads(self.rfile.read(length)))) + payload = json.dumps({"choices": [{"message": {"content": '{"success_score": 1}'}}]}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + judge = DuplexJudge(f"http://127.0.0.1:{server.server_port}", "judge", api_key="token") + assert judge.chat("prompt") == '{"success_score": 1}' + finally: + server.shutdown() + thread.join() + server.server_close() + + path, authorization, payload = requests[0] + assert path == "/v1/chat/completions" + assert authorization == "Bearer token" + assert payload["model"] == "judge" + assert payload["messages"] == [{"role": "user", "content": "prompt"}] diff --git a/tests/e2e/features/fullduplex/test_personaplex_elastic.py b/tests/e2e/features/fullduplex/test_personaplex_elastic.py index 453d90153..1b2f47984 100644 --- a/tests/e2e/features/fullduplex/test_personaplex_elastic.py +++ b/tests/e2e/features/fullduplex/test_personaplex_elastic.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """GPU e2e for PersonaPlex elastic batching (bit-exact isolation + recycle). Greedy decoding makes every cross-slot leak observable as a bit difference: @@ -24,12 +24,17 @@ from __future__ import annotations import numpy as np import pytest +from tests.helpers.mark import hardware_marks + torch = pytest.importorskip("torch") -# slow + H100 + omni routes this module into the weekly Buildkite E2E group -# ("Omni · H100" in .buildkite/cuda/test-weekly.yml: -m "slow and H100 and -# omni" --run-level full_model); the ready/merge CPU sweeps exclude it by mark. -pytestmark = [pytest.mark.full_model, pytest.mark.omni, pytest.mark.slow, pytest.mark.H100] +# slow + H100 + omni + cards_1 routes this module into weekly +# "Omni · H100 · Single-GPU" (test-weekly.yml). +pytestmark = [ + pytest.mark.omni, + pytest.mark.slow, + *hardware_marks(res={"cuda": "H100"}), +] if not torch.cuda.is_available(): # noqa: E402 pytest.skip("needs CUDA", allow_module_level=True) diff --git a/tests/e2e/offline_inference/test_audiox_model_expansion.py b/tests/e2e/offline_inference/test_audiox_model_expansion.py deleted file mode 100644 index af944b144..000000000 --- a/tests/e2e/offline_inference/test_audiox_model_expansion.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import os - -import numpy as np -import pytest -import torch - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner -from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform - -# Prefer a tiny/random checkpoint for CI. -# Override in CI if needed: AUDIOX_TEST_MODEL= -models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")] - -# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize -_OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models] - -pytestmark = [ - pytest.mark.slow, - pytest.mark.diffusion, - pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True, ids=models), -] - - -@hardware_test(res={"cuda": "L4", "xpu": "B60"}) -def test_audiox_model(omni_runner: OmniRunner) -> None: - # Keep runtime short for CI. - seconds_total = 2.0 - # AudioXPipeline always emits 44.1 kHz stereo (advertised via class-level - # ``audio_sample_rate``); the trimmed output should match this rate. - sample_rate = 44100 - - outputs = omni_runner.omni.generate( - prompts={"prompt": "A dog barking in a quiet park."}, - sampling_params_list=OmniDiffusionSamplingParams( - num_inference_steps=4, - guidance_scale=6.0, - generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), - num_outputs_per_prompt=1, - extra_args={ - "audiox_task": "t2a", - "seconds_start": 0.0, - "seconds_total": seconds_total, - }, - ), - ) - - assert outputs is not None - first_output = outputs[0] - assert first_output.final_output_type == "audio" - assert isinstance(first_output, OmniRequestOutput) and first_output - - req_out = first_output - assert isinstance(req_out, OmniRequestOutput) - assert req_out.final_output_type == "audio" - assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output - - audio = req_out.multimodal_output.get("audio") - assert isinstance(audio, np.ndarray) - # audio shape: (batch, channels, samples) - assert audio.ndim == 3 - assert audio.shape[0] == 1 - assert audio.shape[1] == 2 - assert audio.shape[2] > 0 - expected_samples = int(seconds_total * sample_rate) - assert abs(audio.shape[2] - expected_samples) <= 2 * 1024 diff --git a/tests/e2e/offline_inference/test_bagel.py b/tests/e2e/offline_inference/test_bagel.py index 9ba7bdac4..580a46742 100644 --- a/tests/e2e/offline_inference/test_bagel.py +++ b/tests/e2e/offline_inference/test_bagel.py @@ -46,7 +46,7 @@ REFERENCE_TEXT_IMG2TEXT = ( _OMNI_RUNNER_PARAM = (MODEL_NAME, STAGE_CONFIG) pytestmark = [ - pytest.mark.advanced_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] diff --git a/tests/e2e/offline_inference/test_bagel_expansion.py b/tests/e2e/offline_inference/test_bagel_expansion.py index 980fc633f..4c31be722 100644 --- a/tests/e2e/offline_inference/test_bagel_expansion.py +++ b/tests/e2e/offline_inference/test_bagel_expansion.py @@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT" BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml") DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" -# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` +# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` _OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG) @@ -138,7 +138,7 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest: pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] diff --git a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py index fa1bdfead..7e57433ef 100644 --- a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py +++ b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py @@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s ] -# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` +# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)`` _cosy_deployment = get_deploy_config_path("cosyvoice3.yaml") _cosy_model_path = str(_resolve_model_dir()) _OMNI_RUNNER_PARAMS = [ diff --git a/tests/e2e/offline_inference/test_covo_audio_expansion.py b/tests/e2e/offline_inference/test_covo_audio_expansion.py index 11672686e..e743814d3 100644 --- a/tests/e2e/offline_inference/test_covo_audio_expansion.py +++ b/tests/e2e/offline_inference/test_covo_audio_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ E2E tests for Covo-Audio-Chat model with audio input and audio/text output. """ @@ -25,7 +28,7 @@ def get_question(): @pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_audio_to_audio(omni_runner, omni_runner_handler) -> None: +def test_audio_to_audio(omni_runner, offline_client) -> None: """ Test audio input → text + audio output. Deploy Setting: default yaml @@ -44,4 +47,4 @@ def test_audio_to_audio(omni_runner, omni_runner_handler) -> None: "audios": (audio, 16000), "modalities": ["audio"], } - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) diff --git a/tests/e2e/offline_inference/test_dots_tts_expansion.py b/tests/e2e/offline_inference/test_dots_tts_expansion.py new file mode 100644 index 000000000..1c9362a23 --- /dev/null +++ b/tests/e2e/offline_inference/test_dots_tts_expansion.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for dots.tts native AR offline inference.""" + +from collections.abc import Mapping + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path + +DOTS_TTS_MODEL = "rednote-hilab/dots.tts-soar" +DEPLOY_CONFIG = get_deploy_config_path("dots_tts.yaml") +SAMPLE_RATE = 48000 + +# dots.tts ships a custom tokenizer carrying the added ``<|audio_gen_start|>`` +# token, so remote code must be explicitly enabled. +_OMNI_RUNNER_PARAM = (DOTS_TTS_MODEL, DEPLOY_CONFIG, {"trust_remote_code": True}) + +pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) + + +def _build_prompt(omni_runner: OmniRunner, text: str) -> dict: + """Wrap ``text`` in the prefix scaffold dots.tts prefill expects. + + Unlike most TTS models here, a bare ``{"prompt": ...}`` string does not + work: prefill needs ``[文本][文本对应语音]<|audio_gen_start|>`` as + real Qwen2 token ids. + """ + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import ( + build_dots_tts_prompt, + ) + + tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) + return build_dots_tts_prompt(tokenizer=tokenizer, text=text) + + +def _extract_audio(multimodal_output: dict) -> torch.Tensor: + """Extract the final complete audio tensor from multimodal output.""" + assert isinstance(multimodal_output, (dict, Mapping)), f"Expected dict/Mapping, got {type(multimodal_output)}" + + # The talker emits delta chunks; the output processor accumulates them + # under "model_outputs" ("audio" kept for backwards compat). + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + assert audio is not None, f"No audio key, got {list(multimodal_output.keys())}" + + if isinstance(audio, list): + valid = [torch.as_tensor(x).float().cpu().reshape(-1) for x in audio if x is not None] + assert valid, "No valid audio tensors in output list" + audio = torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] + + assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}" + return audio.float().cpu().reshape(-1) + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_dots_tts_zero_shot_001(omni_runner: OmniRunner) -> None: + """Zero-shot synthesis produces audio of a plausible duration.""" + prompt = _build_prompt(omni_runner, "Hello, this is a test of dots TTS running on vLLM Omni.") + + outputs = omni_runner.omni.generate([prompt]) + assert len(outputs) == 1 + + audio = _extract_audio(outputs[0].outputs[0].multimodal_output) + duration_s = audio.shape[0] / SAMPLE_RATE + assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s" + assert torch.isfinite(audio).all(), "Audio contains NaN/Inf" + assert audio.abs().max() > 1e-3, "Audio is silent" + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_dots_tts_prefill_decode_mixed_batch_002(omni_runner: OmniRunner) -> None: + """Mixed prefill+decode batches must keep per-request rows aligned. + + ``compute_logits`` fills ``logits[i]`` positionally from + ``_results_queue``, so every request in the batch has to push exactly one + entry per step, prefills included. When the prefill branch pushed nothing, + a shorter request's stop signal could be read into another request's row. + Regression guard for that pairing (PR #4765 review). + """ + long_text = ( + "This is a deliberately long sentence that stays in the decode phase " + "for many steps, so the shorter requests behind it keep entering " + "prefill alongside it and reproduce the mixed batch pattern." + ) + short_texts = ["Hello one.", "Hello two.", "Hello three.", "Hello four."] + prompts = [_build_prompt(omni_runner, long_text)] + [_build_prompt(omni_runner, t) for t in short_texts] + + outputs = omni_runner.omni.generate(prompts) + assert len(outputs) == len(prompts) + + for i, out in enumerate(outputs): + audio = _extract_audio(out.outputs[0].multimodal_output) + duration_s = audio.shape[0] / SAMPLE_RATE + assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s" + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_dots_tts_deterministic_flow_matching_noise_003(omni_runner: OmniRunner) -> None: + """The same text twice must yield the same audio. + + The DiT flow-matching noise is drawn from a per-request generator seeded + from ``blake2b(seed:request_key:noise_step)`` rather than the global CUDA + RNG. Drawing from the global RNG made output non-reproducible and let + concurrent requests perturb each other's noise streams (PR #4765 review). + """ + text = "Reading aloud every evening helps children build a lasting love of language." + + first = _extract_audio( + omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output + ) + second = _extract_audio( + omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output + ) + + assert first.shape == second.shape, f"Length differs across runs: {first.shape} vs {second.shape}" + assert torch.allclose(first, second, atol=1e-4), "Audio differs across runs; flow-matching noise is not seeded" diff --git a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py deleted file mode 100644 index e22aa1174..000000000 --- a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py +++ /dev/null @@ -1,290 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""E2E test for DreamID-Omni offline inference with Cache-Dit. - -Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of -synthetic inputs so the coverage matches the upstream demo structure. -""" - -from __future__ import annotations - -import importlib.util -import json -import os -import re -import time -from functools import lru_cache -from pathlib import Path -from typing import NamedTuple - -import numpy as np -import pytest -import soundfile as sf -from PIL import Image - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" - -pytestmark = [pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow] - -REPO_ROOT = Path(__file__).resolve().parents[3] -DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni" -MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH" -EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT" -NUM_INFERENCE_STEPS = 45 - -CACHE_CONFIG = { - "Fn_compute_blocks": 1, - "Bn_compute_blocks": 0, - "max_warmup_steps": 4, - "max_cached_steps": 20, - "residual_diff_threshold": 0.24, - "max_continuous_cached_steps": 3, - "enable_taylorseer": False, - "taylorseer_order": 1, - "scm_steps_mask_policy": None, - "scm_steps_policy": "dynamic", -} - - -class DreamIDGenerationResult(NamedTuple): - video: np.ndarray - audio: np.ndarray - fps: int - duration_s: float - elapsed_s: float - - -class DreamIDExampleCase(NamedTuple): - prompt: str - image_paths: tuple[Path, ...] - audio_paths: tuple[Path, ...] - - -def _require_model_path() -> str: - model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH)) - if not Path(model_path).exists(): - pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.") - - if importlib.util.find_spec("dreamid_omni") is None: - pytest.skip( - "dreamid_omni dependency package is not importable. " - "Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first." - ) - - return model_path - - -def _dreamid_repo_root_candidates() -> list[Path]: - candidates: list[Path] = [] - example_root = os.environ.get(EXAMPLE_ROOT_ENV) - if example_root: - candidates.append(Path(example_root).expanduser().resolve()) - - spec = importlib.util.find_spec("dreamid_omni") - if spec is not None and spec.origin: - candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip") - - return candidates - - -@lru_cache(maxsize=1) -def _load_example_case() -> DreamIDExampleCase: - for candidate_root in _dreamid_repo_root_candidates(): - caption_path = candidate_root / "captions" / "20.json" - image_paths = ( - candidate_root / "imgs" / "20" / "0.png", - candidate_root / "imgs" / "20" / "1.png", - ) - audio_paths = ( - candidate_root / "audios" / "20" / "0.wav", - candidate_root / "audios" / "20" / "1.wav", - ) - if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths): - prompt = _load_prompt_from_json(caption_path) - return DreamIDExampleCase( - prompt=prompt, - image_paths=image_paths, - audio_paths=audio_paths, - ) - - searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or "" - pytest.skip( - "DreamID-Omni official example assets not found. " - f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. " - f"Searched: {searched}" - ) - - -def _load_prompt_from_json(path: Path) -> str: - raw_text = path.read_text(encoding="utf-8").strip() - parsed = json.loads(raw_text) - if isinstance(parsed, str): - prompt = parsed - elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str): - prompt = parsed["content"] - else: - raise AssertionError(f"Unexpected DreamID caption format in {path}") - - prompt = re.sub( - r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", - "", - prompt, - flags=re.DOTALL, - ).strip() - prompt = re.sub( - r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", - "", - prompt, - flags=re.DOTALL, - ).strip() - prompt = re.sub(r"\[[A-Z_]+\]", "", prompt) - prompt = re.sub(r"\n\s*\n", "\n", prompt).strip() - return prompt - - -def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]: - images: list[Image.Image] = [] - for path in example_case.image_paths: - with Image.open(path) as image: - images.append(image.convert("RGB")) - return images - - -def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]: - audios: list[np.ndarray] = [] - for path in example_case.audio_paths: - audio_array, sample_rate = sf.read(str(path), dtype="float32") - audio_np = np.asarray(audio_array, dtype=np.float32) - if audio_np.ndim == 2: - audio_np = np.mean(audio_np, axis=1) - audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)] - audios.append(audio_np.reshape(-1)) - return audios - - -def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams: - return OmniDiffusionSamplingParams( - height=704, - width=1280, - num_inference_steps=num_inference_steps, - seed=42, - extra_args={ - "solver_name": "unipc", - "shift": 5.0, - }, - ) - - -def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]: - return { - "prompt": example_case.prompt, - "video_negative_prompt": "jitter, bad hands, blur, distortion", - "audio_negative_prompt": "robotic, muffled, echo, distorted", - "multi_modal_data": { - "image": _load_example_images(example_case), - "audio": _load_example_audios(example_case), - }, - } - - -def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult: - """Extract video and audio from raw Omni outputs""" - assert outputs, "DreamID-Omni returned no outputs" - - result = outputs[0] - assert hasattr(result, "images") and result.images, "not return video frames" - - video = result.images[0] - assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}" - assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}" - - multimodal_output = getattr(result, "multimodal_output", None) or {} - assert "audio" in multimodal_output, "not return audio" - audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1) - assert audio.size > 0, "returned empty audio" - - audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000)) - fps = multimodal_output.get("fps", 24) - duration_s = float(audio.size) / float(audio_sample_rate) - - generation = DreamIDGenerationResult( - video=video.astype(np.float32, copy=False), - audio=audio, - fps=fps, - duration_s=duration_s, - elapsed_s=elapsed_s, - ) - return generation - - -def _run_generation( - *, - model_path: str, - example_case: DreamIDExampleCase, - use_cache_dit: bool, -) -> DreamIDGenerationResult: - runner_kwargs = { - "model_type": "dreamid-omni", - "parallel_config": DiffusionParallelConfig(cfg_parallel_size=1), - } - if use_cache_dit: - runner_kwargs["cache_backend"] = "cache_dit" - runner_kwargs["cache_config"] = CACHE_CONFIG - - with OmniRunner(model_path, **runner_kwargs) as runner: - warmup_payload = _request_payload(example_case) - _ = runner.omni.generate(warmup_payload, _sampling_params(2)) - - elapsed_samples: list[float] = [] - extracted_result: DreamIDGenerationResult | None = None - - request_payload = _request_payload(example_case) - start = time.perf_counter() - outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS)) - elapsed_s = time.perf_counter() - start - elapsed_samples.append(elapsed_s) - if extracted_result is None: - extracted_result = _extract_generation(outputs, elapsed_s) - - assert extracted_result is not None - median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64))) - return extracted_result._replace(elapsed_s=median_elapsed_s) - - -def _video_mae(a: np.ndarray, b: np.ndarray) -> float: - assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}" - return float(np.mean(np.abs(a - b))) - - -def _audio_rms(audio: np.ndarray) -> float: - audio = np.asarray(audio, dtype=np.float32).reshape(-1) - return float(np.sqrt(np.mean(audio**2))) - - -def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float: - rms_a = _audio_rms(a.audio) - rms_b = _audio_rms(b.audio) - return rms_b / max(rms_a, 1e-6) - - -@hardware_test(res={"cuda": "H100"}) -def test_dreamid_omni_cache_dit_e2e() -> None: - model_path = _require_model_path() - example_case = _load_example_case() - - baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False) - cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True) - - video_mae = _video_mae(baseline.video, cached.video) - audio_rms_ratio = _audio_rms_ratio(baseline, cached) - speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf") - - assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}" - assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}" - assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x" diff --git a/tests/e2e/offline_inference/test_ernie_image_expansion.py b/tests/e2e/offline_inference/test_ernie_image_expansion.py index 19793ff00..c986a6366 100644 --- a/tests/e2e/offline_inference/test_ernie_image_expansion.py +++ b/tests/e2e/offline_inference/test_ernie_image_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end test for ErnieImage text-to-image generation. @@ -11,13 +11,13 @@ Equivalent to running: import pytest from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunnerHandler +from tests.helpers.runtime import OfflineOmniClient from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "baidu/ERNIE-Image" -# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize +# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize _OMNI_RUNNER_PARAM = ( MODEL, None, @@ -37,9 +37,9 @@ pytestmark = [ @hardware_test(res={"cuda": "L4"}, num_cards=2) -def test_ernie_image_text_to_image(omni_runner_handler: OmniRunnerHandler) -> None: +def test_ernie_image_text_to_image(offline_client: OfflineOmniClient) -> None: request_config = { - "model": omni_runner_handler.runner.model_name, + "model": offline_client.runner.model_name, "prompt": "A photo of a cat sitting on a laptop", "sampling_params": OmniDiffusionSamplingParams( height=512, @@ -49,4 +49,4 @@ def test_ernie_image_text_to_image(omni_runner_handler: OmniRunnerHandler) -> No seed=42, ), } - omni_runner_handler.send_diffusion_request(request_config) + offline_client.send_diffusion_request(request_config) diff --git a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py index ec48a49f7..92a710f6e 100644 --- a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for FLUX AutoRound W4A16 quantized inference. These tests require: @@ -19,8 +19,8 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test +from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput diff --git a/tests/e2e/offline_inference/test_flux1_kontext_expansion.py b/tests/e2e/offline_inference/test_flux1_kontext_expansion.py index abf7bbc2c..61a47a3b0 100644 --- a/tests/e2e/offline_inference/test_flux1_kontext_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_kontext_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for FluxKontext model pipeline. @@ -15,7 +15,7 @@ import pytest from vllm.assets.image import ImageAsset from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunnerHandler +from tests.helpers.runtime import OfflineOmniClient from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -47,18 +47,18 @@ def _sampling_512() -> OmniDiffusionSamplingParams: @hardware_test(res={"cuda": "H100"}, num_cards=2) -def test_flux_kontext_text_to_image(omni_runner_handler: OmniRunnerHandler): +def test_flux_kontext_text_to_image(offline_client: OfflineOmniClient): """Test FluxKontext text-to-image generation with real model.""" request_config = { "model": MODEL, "prompt": "A photo of a cat sitting on a laptop", "sampling_params": _sampling_512(), } - omni_runner_handler.send_diffusion_request(request_config) + offline_client.send_diffusion_request(request_config) @hardware_test(res={"cuda": "H100"}, num_cards=2) -def test_flux_kontext_image_edit(omni_runner_handler: OmniRunnerHandler): +def test_flux_kontext_image_edit(offline_client: OfflineOmniClient): """Test FluxKontext image-to-image editing with real model.""" input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") request_config = { @@ -67,4 +67,4 @@ def test_flux_kontext_image_edit(omni_runner_handler: OmniRunnerHandler): "modalities": ["img2img"], "sampling_params": _sampling_512(), } - omni_runner_handler.send_diffusion_request(request_config) + offline_client.send_diffusion_request(request_config) diff --git a/tests/e2e/offline_inference/test_flux1_schnell_expansion.py b/tests/e2e/offline_inference/test_flux1_schnell_expansion.py index de6a758f6..9b61f876f 100644 --- a/tests/e2e/offline_inference/test_flux1_schnell_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_schnell_expansion.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for Flux1 Schnell.""" import pytest from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunnerHandler +from tests.helpers.runtime import OfflineOmniClient from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "black-forest-labs/FLUX.1-schnell" @@ -21,7 +21,7 @@ pytestmark = [ @hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_flux_schnell_text_to_image(omni_runner_handler: OmniRunnerHandler): +def test_flux_schnell_text_to_image(offline_client: OfflineOmniClient): """Test FLUX.1-schnell text-to-image generation.""" request_config = { "model": MODEL, @@ -33,4 +33,4 @@ def test_flux_schnell_text_to_image(omni_runner_handler: OmniRunnerHandler): seed=42, ), } - omni_runner_handler.send_diffusion_request(request_config) + offline_client.send_diffusion_request(request_config) diff --git a/tests/e2e/offline_inference/test_flux2_klein.py b/tests/e2e/offline_inference/test_flux2_klein.py index 9656772c9..ed0cbaf91 100644 --- a/tests/e2e/offline_inference/test_flux2_klein.py +++ b/tests/e2e/offline_inference/test_flux2_klein.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end test for Flux2 Klein inpainting. -Inpainting uses ``omni_runner_handler.send_diffusion_request`` with +Inpainting uses ``offline_client.send_diffusion_request`` with ``multi_modal_data`` containing ``image`` and ``mask_image``; see -:meth:`OmniRunnerHandler.send_diffusion_request` in ``tests.helpers.runtime``. +:meth:`OfflineOmniClient.send_diffusion_request` in ``tests.helpers.runtime``. """ import pytest @@ -14,7 +14,7 @@ import torch from PIL import Image, ImageDraw from tests.helpers.mark import hardware_test -from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler +from tests.helpers.runtime import DiffusionResponse, OfflineOmniClient from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -47,15 +47,17 @@ def _create_test_inputs(color: tuple = (100, 150, 200)): def _images_from_response(response: DiffusionResponse) -> list[Image.Image]: - if isinstance(response.images[0], list): - return [f for fr in response.images for f in fr] - return list(response.images) + images = response.images + assert images is not None + if images and isinstance(images[0], list): + return [frame for frames in images for frame in frames] + return list(images) def _send_inpaint_with_generator( - omni_runner_handler: OmniRunnerHandler, prompt: str, input_image, mask_image, generator: torch.Generator + offline_client: OfflineOmniClient, prompt: str, input_image, mask_image, generator: torch.Generator ) -> DiffusionResponse: - return omni_runner_handler.send_diffusion_request( + return offline_client.send_diffusion_request( { "model": MODEL, "prompt": prompt, @@ -72,18 +74,19 @@ def _send_inpaint_with_generator( ) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion +@pytest.mark.advanced_model @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_flux2_klein_inpaint_deterministic(omni_runner_handler: OmniRunnerHandler): +def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient): input_image, mask_image = _create_test_inputs() seed = 12345 gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(seed) gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(seed) - r1 = _send_inpaint_with_generator(omni_runner_handler, "A red flower in a field", input_image, mask_image, gen1) - r2 = _send_inpaint_with_generator(omni_runner_handler, "A red flower in a field", input_image, mask_image, gen2) + r1 = _send_inpaint_with_generator(offline_client, "A red flower in a field", input_image, mask_image, gen1) + r2 = _send_inpaint_with_generator(offline_client, "A red flower in a field", input_image, mask_image, gen2) images1 = _images_from_response(r1) images2 = _images_from_response(r2) @@ -93,17 +96,18 @@ def test_flux2_klein_inpaint_deterministic(omni_runner_handler: OmniRunnerHandle ) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion +@pytest.mark.advanced_model @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_flux2_klein_inpaint_different_seeds_different_output(omni_runner_handler: OmniRunnerHandler): +def test_flux2_klein_inpaint_different_seeds_different_output(offline_client: OfflineOmniClient): input_image, mask_image = _create_test_inputs() gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(42) gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(99999) - r1 = _send_inpaint_with_generator(omni_runner_handler, "A beautiful landscape", input_image, mask_image, gen1) - r2 = _send_inpaint_with_generator(omni_runner_handler, "A beautiful landscape", input_image, mask_image, gen2) + r1 = _send_inpaint_with_generator(offline_client, "A beautiful landscape", input_image, mask_image, gen1) + r2 = _send_inpaint_with_generator(offline_client, "A beautiful landscape", input_image, mask_image, gen2) images1 = _images_from_response(r1) images2 = _images_from_response(r2) diff --git a/tests/e2e/offline_inference/test_gepard_tts.py b/tests/e2e/offline_inference/test_gepard_tts.py new file mode 100644 index 000000000..b201842b4 --- /dev/null +++ b/tests/e2e/offline_inference/test_gepard_tts.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for the Gepard-1.0 single-stage native-AR TTS. + +Zero-shot only (the model's default learned voice). These tests need a GPU and +the NeMo NanoCodec, so the whole file runs in the weekly tier -- see the comment +above ``pytestmark``. The window arithmetic behind the streaming decode is +covered on CPU in ``tests/model_executor/models/test_gepard_window.py``. +""" + +from __future__ import annotations + +import functools +import tempfile +from pathlib import Path +from typing import NamedTuple + +import pytest +import soundfile as sf +import torch +from transformers import AutoTokenizer + +from tests.helpers.mark import hardware_test +from tests.helpers.media import convert_audio_file_to_text +from tests.helpers.runtime import OmniRunner # noqa: F401 (indirect fixture) +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni import Omni +from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig +from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids + +# The codec decoder is an optional dependency; skip rather than fail the engine +# inside the worker process when it is absent. The CI image does not carry it, +# so the weekly step installs it first -- see .buildkite/cuda/test-weekly.yml. +# Install recipe: examples/offline_inference/text_to_speech/README.md. +pytest.importorskip("nemo.collections.tts.models") + +MODEL_NAME = "nineninesix/gepard-1.0" +STAGE_CONFIG = get_deploy_config_path("gepard.yaml") +SAMPLE_RATE = 22050 +# NanoCodec upsamples one FSQ frame to exactly this many samples. +SAMPLES_PER_FRAME = 1024 + + +@functools.cache +def _config() -> GepardConfig: + """The checkpoint's own layout, not the class defaults. + + The worker builds its config from the ``gepard_config.json`` sidecar, so a + prompt assembled from the defaults would silently use a different layout + the day the sidecar moves — and a mismatched layout degrades the speech + rather than failing, which is the hardest way for this to be found. + """ + return GepardConfig.from_checkpoint(MODEL_NAME) + + +# Prompt -> a word that appears in no other prompt, so a clip can be traced +# back to the request that asked for it. +# +# These are not interchangeable. Batching perturbs every request's sampling +# trajectory, and on some trajectories this checkpoint's stop head never fires; +# this set was screened over 10 batched runs with no runaway. Re-screen the +# whole set when changing one, not just the row changed. +_DISTINGUISHABLE_PROMPTS = { + "He drinks coffee every morning.": "coffee", + "Machine learning is interesting.": "learning", + "Please close the window before leaving.": "window", + "My favorite color is purple.": "purple", +} + +# Above text_repetition.apply_below, so the layout is voiced once instead of +# repeated. Every other prompt here takes the repeated branch. +_LONG_TEXT = ( + "The morning train was late again, so I walked the last two miles along the " + "river and reached the office just before nine." +) + +_OMNI_RUNNER_PARAM = ( + MODEL_NAME, + STAGE_CONFIG, + {"trust_remote_code": True}, +) + +# Only a common word is checked, which needs less ASR than the shared helper's +# ``small`` default; a miss escalates to it. +_ASR_MODEL = "base" +_ASR_ESCALATION_MODEL = "small" + +# Transcription dominates this file's runtime -- Whisper runs on CPU on a +# single-GPU runner -- so only two of the four concurrent clips are read as +# speech. The pairwise waveform comparison covers all four. +_TRANSCRIBED_CLIPS = 2 + +# No level marker at module scope: each test carries the one matching how often +# its claim needs re-checking, though the weekly step collects the file whole. +# There is deliberately no core_model row -- under dummy weights the stop head is +# random, so every request would run to gepard.yaml's max_tokens. +pytestmark = [ + pytest.mark.tts, + pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), +] + + +def _build_request(text: str) -> dict: + """Assemble the [speaker slots, SOT, text, EOT, SOS] prompt.""" + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + prompt_token_ids = build_gepard_prompt_ids( + tokenizer(text, add_special_tokens=False)["input_ids"], + config=_config(), + ) + return { + "prompt_token_ids": prompt_token_ids, + "additional_information": {"text": [text]}, + } + + +def _extract_audio(mm) -> torch.Tensor | None: + # The consolidation path renames "model_outputs" to "audio". + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if isinstance(audio, list): + audio = audio[0] if len(audio) else None + return audio + + +class _Clip(NamedTuple): + wav: torch.Tensor + # Counts the frames the codec was handed using only what the engine + # returned, never the model's internals. + committed_frames: int + + +def _synthesize_all(omni: Omni, texts: list[str]) -> list[_Clip]: + """Submit every text in one call; return one clip per request, in + submission order. + + ``Omni.generate`` yields in completion order and never sorts, so the pairing + is restored from the request id prefix it assigns (``f"{i}_{uuid4()}"``). + Without that, the only testable claim is that the batch as a whole looks + plausible, which a stable swap between two requests satisfies. + + Nothing here authors a SamplingParams: a fresh one would replace the stage + defaults rather than merge over them, dropping the pipeline's stop token. + """ + outputs = omni.generate([_build_request(t) for t in texts]) + assert len(outputs) == len(texts), f"expected {len(texts)} outputs, got {len(outputs)}" + + by_index: dict[int, _Clip] = {} + for stage_outputs in outputs: + index = int(str(stage_outputs.request_id).split("_", 1)[0]) + waveform = None + frames = 0 + for out in stage_outputs.outputs: + # Explicit None check: `x or []` would ask an array for a truth value. + token_ids = out.token_ids if out.token_ids is not None else [] + # STOP sits just past head0's range, so a token equal to it is + # never a committed frame. + frames += sum(1 for t in token_ids if int(t) != _config().stop_token) + mm = out.multimodal_output + if mm is None: + continue + audio = _extract_audio(mm) + if audio is None: + continue + sr_t = mm.get("sr") + if hasattr(sr_t, "item"): + assert int(sr_t.item()) == SAMPLE_RATE + waveform = audio.cpu().float() + assert waveform is not None, "no audio output produced" + by_index[index] = _Clip(wav=waveform, committed_frames=frames) + + assert sorted(by_index) == list(range(len(texts))), ( + f"request ids did not cover every submitted prompt exactly once: got {sorted(by_index)}" + ) + return [by_index[i] for i in range(len(texts))] + + +def _synthesize(omni: Omni, text: str) -> _Clip: + return _synthesize_all(omni, [text])[0] + + +def _assert_clip_is_sane(clip: _Clip) -> None: + """The checks every clip has to pass, whatever produced it.""" + wav = clip.wav + assert wav.dim() == 1, f"expected mono 1-D, got shape {tuple(wav.shape)}" + assert wav.numel() > 0, "empty waveform" + assert torch.isfinite(wav).all(), "non-finite samples" + + # Frame conservation against the engine's own token count, not just + # divisibility. This is what catches a payload carrying non-audio data or a + # dropped tail; both pass every other check here. + assert wav.numel() == clip.committed_frames * SAMPLES_PER_FRAME, ( + f"{wav.numel()} samples for {clip.committed_frames} committed frames — " + f"expected exactly {clip.committed_frames * SAMPLES_PER_FRAME}" + ) + + # Anything far outside the waveform range is not codec output. + assert float(wav.abs().max()) <= 1.5, f"peak {float(wav.abs().max()):.2f} is out of range" + + +def _assert_stopped_on_its_own(clip: _Clip) -> None: + """The request must stop on its own rather than run to ``max_tokens``. + + Separate from the structural checks because it reads the stop head's + judgement, not the payload's shape: under ``load_format: dummy`` that head + is random and never fires, so every clip runs to the token cap. + """ + seconds = clip.wav.numel() / SAMPLE_RATE + assert 0.5 < seconds < 30.0, f"implausible duration {seconds:.2f}s — did STOP fire?" + + +def _transcribe(wav: torch.Tensor, tmp_dir: str, name: str, model_size: str = _ASR_MODEL) -> str: + path = str(Path(tmp_dir) / f"{name}.wav") + sf.write(path, wav.numpy(), SAMPLE_RATE) + return convert_audio_file_to_text(path, model_size=model_size, language="en").lower() + + +def _transcribe_for(wav: torch.Tensor, tmp_dir: str, name: str, accepts) -> str: + """Transcribe with the fast model, retrying once with the stronger one. + + ``accepts`` decides whether the fast transcript is usable. Escalating only + on a miss keeps the average clip at ``_ASR_MODEL`` while a genuine model + artifact still fails the caller's assertion, since the stronger ASR + mistranscribes it too. + """ + transcript = _transcribe(wav, tmp_dir, name) + if accepts(transcript): + return transcript + return _transcribe(wav, tmp_dir, name, model_size=_ASR_ESCALATION_MODEL) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize( + ("text", "keyword"), + [ + # The canonical path: one request, one transcript, the cheapest check + # that reads the audio as speech. + pytest.param( + "Hello, this is Gepard speaking.", + "hello", + marks=pytest.mark.advanced_model, + id="canonical", + ), + # No keyword: this row carries the structural and stopping checks over a + # second text and leaves the transcript to the row above. + pytest.param( + "The quick brown fox jumps over the lazy dog.", + None, + marks=pytest.mark.full_model, + id="second_text", + ), + ], +) +def test_gepard_offline_zero_shot(omni_runner, run_level: str, text: str, keyword: str | None) -> None: + """Zero-shot synthesis produces a finite mono waveform at 22.05 kHz. + + The transcript check is the one assertion here that reads the audio as + speech rather than as a tensor. It stands in for a parity comparison, which + needs a reference implementation this PR cannot run — noise, a repeated + loop or the wrong clip all pass every numeric check above it. + """ + clip = _synthesize(omni_runner, text) + _assert_clip_is_sane(clip) + + if run_level not in {"advanced_model", "full_model"}: + return # dummy weights: the audio is structural only, so do not read it + + _assert_stopped_on_its_own(clip) + if keyword is None: + return + + with tempfile.TemporaryDirectory() as tmp_dir: + transcript = _transcribe_for(clip.wav, tmp_dir, "zero_shot", lambda t: keyword in t) + assert keyword in transcript, f"expected {keyword!r} in transcript, got {transcript!r}" + + +@pytest.mark.full_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_gepard_offline_long_text_skips_repetition(omni_runner, run_level: str) -> None: + """The other branch of the prompt layout, end to end. + + A text at or above ``apply_below`` tokens is voiced once; shorter ones are + repeated to carry enough mass against the speaker prefix. The layout + assertion keeps this row honest — if a tokenizer change pushed this text + back under the threshold it would silently retest the repeated path. + """ + prompt_ids = _build_request(_LONG_TEXT)["prompt_token_ids"] + assert prompt_ids.count(_config().start_of_text) == 1, "_LONG_TEXT is no longer above apply_below" + + clip = _synthesize(omni_runner, _LONG_TEXT) + _assert_clip_is_sane(clip) + + if run_level not in {"advanced_model", "full_model"}: + return + + _assert_stopped_on_its_own(clip) + with tempfile.TemporaryDirectory() as tmp_dir: + transcript = _transcribe_for(clip.wav, tmp_dir, "long_text", lambda t: "river" in t) + assert "river" in transcript, f"expected 'river' in transcript, got {transcript!r}" + + +@pytest.mark.slow +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_gepard_offline_concurrent_requests_stay_isolated(omni_runner, run_level: str) -> None: + """Four requests in one batch, each holding its own generation state. + + The prompts are deliberately unlike each other: cross-talk between the + per-request frame history, the codec window or the output routing shows up + as a transcript matching the wrong prompt, which similar prompts would + hide. ``max_num_seqs`` is 4, so these really do run together. + """ + texts = list(_DISTINGUISHABLE_PROMPTS) + clips = _synthesize_all(omni_runner, texts) + + for clip in clips: + _assert_clip_is_sane(clip) + + if run_level not in {"advanced_model", "full_model"}: + return # dummy weights: the audio is structural only, so do not read it + + for clip in clips: + _assert_stopped_on_its_own(clip) + + # All four clips, not just the transcribed ones: delivering one request's + # audio for another leaves two byte-identical waveforms. + for i in range(len(clips)): + for j in range(i + 1, len(clips)): + assert not torch.equal(clips[i].wav, clips[j].wav), ( + f"clips {i} and {j} are the same waveform — one request's audio was delivered for another" + ) + + # Each clip must say what its own prompt asked for: accepting any of the + # four would pass a stable permutation, which is what a mix-up produces. + keywords = set(_DISTINGUISHABLE_PROMPTS.values()) + with tempfile.TemporaryDirectory() as tmp_dir: + for i, clip in enumerate(clips[:_TRANSCRIBED_CLIPS]): + expected = _DISTINGUISHABLE_PROMPTS[texts[i]] + # Escalate on the expected word being absent, not on "no keyword + # at all" -- a neighbour's word is the failure under test. + transcript = _transcribe_for(clip.wav, tmp_dir, f"concurrent_{i}", lambda t: expected in t) + present = sorted(k for k in keywords if k in transcript) + assert present == [expected], ( + f"clip for {texts[i]!r} transcribed as {transcript!r}: expected {expected!r} and " + f"no other prompt's keyword, found {present} — requests bled into each other" + ) diff --git a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py index 28eaab9b3..b38e26993 100644 --- a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for GLM-Image AutoRound W4A16 quantized inference. These tests cover text-to-image and image-to-image generation with @@ -19,10 +19,10 @@ import pytest from PIL import Image from vllm import SamplingParams -from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import OmniRunnerHandler +from tests.helpers.monitor import DeviceMemoryMonitor +from tests.helpers.runtime import OfflineOmniClient from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -109,7 +109,7 @@ def _diffusion_sampling_params( pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, ] @@ -121,7 +121,7 @@ pytestmark = [ @pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) @hardware_test(res={"cuda": "H100"}, num_cards=2) -def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunnerHandler): +def test_glm_image_autoround_w4a16_generates_image(offline_client: OfflineOmniClient): """Load the W4A16 quantized GLM-Image model and verify it produces a valid image.""" gc.collect() current_omni_platform.empty_cache() @@ -152,7 +152,7 @@ def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunn seed=42, ) - outputs = omni_runner_handler.runner.generate( + outputs = offline_client.runner.generate( [prompt_dict], [ar_params, diffusion_params], ) @@ -186,7 +186,7 @@ def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunn @pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) @hardware_test(res={"cuda": "H100"}, num_cards=2) -def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunnerHandler): +def test_glm_image_autoround_w4a16_image_to_image(offline_client: OfflineOmniClient): """Load the W4A16 quantized GLM-Image and verify image-to-image generation works.""" ref_image_arr = generate_synthetic_image(WIDTH, HEIGHT)["np_array"] @@ -217,7 +217,7 @@ def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunne seed=42, ) - outputs = omni_runner_handler.runner.generate( + outputs = offline_client.runner.generate( [prompt_dict], [ar_params, diffusion_params], ) diff --git a/tests/e2e/offline_inference/test_glm_tts_expansion.py b/tests/e2e/offline_inference/test_glm_tts_expansion.py index e54dd5480..2bde38e6b 100644 --- a/tests/e2e/offline_inference/test_glm_tts_expansion.py +++ b/tests/e2e/offline_inference/test_glm_tts_expansion.py @@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None: with OmniRunner( MODEL, - stage_configs_path=_get_deploy_config(async_chunk=async_chunk), + deploy_config=_get_deploy_config(async_chunk=async_chunk), stage_init_timeout=600, trust_remote_code=True, ) as omni_runner: diff --git a/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py b/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py index 00482a132..53134446d 100644 --- a/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py +++ b/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Offline expansion tests for higgs-audio v2. Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py: build Stage-0 prompt_token_ids with the model's own processor + plain-text chat template, drive the engine via ``Omni.generate``, and inspect the -Stage-1 PCM payload directly. The OmniRunnerHandler.send_audio_speech_request +Stage-1 PCM payload directly. The OfflineOmniClient.send_audio_speech_request helper applies a Qwen-style ChatML wrap that is not valid for higgs, so this test bypasses it (same approach as test_voxtral_tts.py). diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py index 4bdb42c34..7e7c63900 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py @@ -1,21 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Smoke test for HunyuanImage-3.0 Image-to-Text (I2T) pipeline.""" from collections.abc import Generator -from pathlib import Path import pytest -import torch from PIL import Image +from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path from vllm_omni import Omni from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct" -REPO_ROOT = Path(__file__).resolve().parents[3] -STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_i2t.yaml" +DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3_ar.yaml") # First 20 generated token IDs from the HF greedy reference on this input. # vllm-omni AR output matches this prefix bitwise; the two implementations @@ -52,13 +51,13 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] def omni() -> Generator[Omni, None, None]: with OmniRunner( MODEL_NAME, - stage_configs_path=str(STAGE_CONFIG_PATH), + deploy_config=DEPLOY_CONFIG_PATH, trust_remote_code=True, ) as runner: yield runner.omni -@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs.") +@hardware_test(res={"cuda": "H100"}, num_cards=4) def test_i2t_generates_text(omni: Omni) -> None: """Verify I2T output's first 20 token IDs match the HF greedy baseline.""" # Solid-color image keeps the input self-contained and reproducible. diff --git a/tests/e2e/offline_inference/test_indextts2_expansion.py b/tests/e2e/offline_inference/test_indextts2_expansion.py index ba5b984f9..c95de98ea 100644 --- a/tests/e2e/offline_inference/test_indextts2_expansion.py +++ b/tests/e2e/offline_inference/test_indextts2_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference tests for IndexTTS2 two-stage pipeline. Stage 0 (GPT AR) → Stage 1 (S2Mel + BigVGAN). Output is 22050 Hz mono WAV. @@ -8,12 +8,11 @@ Covers: basic TTS. from __future__ import annotations -from pathlib import Path - import pytest import torch from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path from vllm_omni import Omni @@ -29,8 +28,7 @@ _OMNI_RUNNER_PARAM = ( STAGE_CONFIG, ) -ASSETS_DIR = Path(__file__).resolve().parents[2] / "assets" / "indextts2" -REF_AUDIO = str(ASSETS_DIR / "ref_audio.wav") +REF_AUDIO = str(get_asset_path("indextts2/ref_audio.wav")) pytestmark = [ pytest.mark.slow, @@ -55,9 +53,12 @@ def _sample_rate_from_mm(mm: dict) -> int: sr = mm.get("sr") if isinstance(sr, list): sr = sr[-1] if sr else None - if hasattr(sr, "item"): - return int(sr.item()) - return int(sr) if sr is not None else SAMPLE_RATE + if sr is None: + return SAMPLE_RATE + item = getattr(sr, "item", None) + if callable(item): + return int(item()) + return int(sr) def _build_request(text: str, ref_audio: str, **extra) -> dict: diff --git a/tests/e2e/offline_inference/test_krea2_expansion.py b/tests/e2e/offline_inference/test_krea2_expansion.py index a8f63032b..372a36afa 100644 --- a/tests/e2e/offline_inference/test_krea2_expansion.py +++ b/tests/e2e/offline_inference/test_krea2_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """L4 functionality tests for the Krea 2 text-to-image diffusion pipeline (``Krea2Pipeline``). @@ -15,7 +15,7 @@ import numpy as np import pytest from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunnerHandler +from tests.helpers.runtime import OfflineOmniClient from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id @@ -49,9 +49,9 @@ def _sampling() -> OmniDiffusionSamplingParams: @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) -def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None: +def test_krea2_text_to_image_001(offline_client: OfflineOmniClient) -> None: """Basic functional smoke: a single prompt produces a decoded image.""" - omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) @hardware_test(res={"cuda": "H100"}) @@ -60,22 +60,22 @@ def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None [(MODEL, None, {"enable_layerwise_offload": True})], indirect=True, ) -def test_krea2_layerwise_offload(omni_runner_handler: OmniRunnerHandler) -> None: +def test_krea2_layerwise_offload(offline_client: OfflineOmniClient) -> None: """Exercise layerwise CPU offload on the DiT (SupportsComponentDiscovery + _layerwise_offload_blocks_attrs).""" - omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) @hardware_test(res={"cuda": "H100"}) @pytest.mark.cache @pytest.mark.parametrize("omni_runner", [(MODEL, None, {"cache_backend": "cache_dit"})], indirect=True) -def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None: +def test_krea2_cache_dit(offline_client: OfflineOmniClient) -> None: """Exercise Cache-DiT on Krea 2 via the custom Krea2Pipeline enabler (``enable_cache_for_krea2``). Validates the docs' Cache-DiT support claim for Krea 2. ``has_separate_cfg`` is checkpoint-aware (False for the distilled Turbo no-CFG path, True for the Raw CFG path); the default Turbo checkpoint exercises the no-CFG branch. """ - omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) @hardware_test(res={"cuda": "H100"}, num_cards=2) @@ -88,13 +88,13 @@ def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None: [(MODEL, None, {"use_hsdp": True, "hsdp_shard_size": 2})], indirect=True, ) -def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None: +def test_krea2_hsdp(offline_client: OfflineOmniClient) -> None: """Exercise Hybrid Sharded Data Parallel (HSDP) weight sharding on the DiT across 2 GPUs. - Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU (distributed_cuda) - shard by ``num_cards=2``; skipped automatically on boxes with fewer than 2 CUDA devices. + Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU + shard by ``num_cards=2`` (``not cards_1``); skipped automatically on boxes with fewer than 2 CUDA devices. """ - omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) @hardware_test(res={"cuda": "H100"}, num_cards=2) @@ -107,27 +107,29 @@ def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None: [(MODEL, None, {"tensor_parallel_size": 2, "vae_patch_parallel_size": 2})], indirect=True, ) -def test_krea2_vae_patch_parallel(omni_runner_handler: OmniRunnerHandler) -> None: +def test_krea2_vae_patch_parallel(offline_client: OfflineOmniClient) -> None: """Exercise VAE patch-parallel (decode) across 2 GPUs. Validates the docs' VAE-patch-parallel (decode) support claim for Krea 2. Scheduled on the - multi-GPU (distributed_cuda) shard by ``num_cards=2``; skipped automatically on boxes with + multi-GPU shard by ``num_cards=2`` (``not cards_1``); skipped automatically on boxes with fewer than 2 CUDA devices. """ - omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) -def _generate(handler: OmniRunnerHandler, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray: +def _generate(handler: OfflineOmniClient, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray: sp = _sampling() sp.lora_request = lora_request sp.lora_scale = lora_scale resp = handler.send_diffusion_request({"model": MODEL, "prompt": LORA_PROMPT, "sampling_params": sp}) - return np.asarray(resp.images[0].convert("RGB"), dtype=np.int16) + images = resp.images + assert images is not None + return np.asarray(images[0].convert("RGB"), dtype=np.int16) @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) -def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None: +def test_krea2_lora(offline_client: OfflineOmniClient) -> None: """Validate diffusion LoRA on Krea 2: visible effect, scale sensitivity, clean deactivation. Uses ``LORA`` (a vLLM-Omni-compatible PEFT repackaging of ``krea/Krea-2-LoRA-darkbrush``, @@ -137,10 +139,10 @@ def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None: """ lora_request = LoRARequest(lora_name="darkbrush", lora_int_id=stable_lora_int_id(LORA), lora_path=LORA) - baseline = _generate(omni_runner_handler, None) - img_1x = _generate(omni_runner_handler, lora_request, lora_scale=1.0) - img_2x = _generate(omni_runner_handler, lora_request, lora_scale=2.0) - restored = _generate(omni_runner_handler, None) + baseline = _generate(offline_client, None) + img_1x = _generate(offline_client, lora_request, lora_scale=1.0) + img_2x = _generate(offline_client, lora_request, lora_scale=2.0) + restored = _generate(offline_client, None) diff_1x = np.abs(baseline - img_1x).mean() diff_2x = np.abs(baseline - img_2x).mean() diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2.py b/tests/e2e/offline_inference/test_lingbot_world_v2.py new file mode 100644 index 000000000..39d9b9d95 --- /dev/null +++ b/tests/e2e/offline_inference/test_lingbot_world_v2.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end smoke test for LingBot-World v2 generation.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import numpy as np +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunnerHandler +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = os.environ.get( + "VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH", + "robbyant/lingbot-world-v2-14b-causal-fast-diffusers", +) +_IMAGE_PATH = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH") +_ACTION_DIR_VALUE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR") +_ACTION_DIR = Path(_ACTION_DIR_VALUE).expanduser().resolve() if _ACTION_DIR_VALUE else None +_HAS_INPUT_ASSETS = _IMAGE_PATH is not None and _ACTION_DIR is not None + +_OMNI_RUNNER_PARAM = ( + MODEL, + None, + { + "flow_shift": 5.0, + "parallel_config": DiffusionParallelConfig(tensor_parallel_size=1), + "model_config": { + "lingbot_action_root": str(_ACTION_DIR.parent) if _ACTION_DIR is not None else "", + }, + }, +) + +pytestmark = [ + pytest.mark.slow, + pytest.mark.diffusion, + pytest.mark.skipif( + not _HAS_INPUT_ASSETS, + reason=("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH and VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR are required"), + ), + pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), +] + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None: + """Load the real checkpoint and generate one causal block on one GPU.""" + + assert _IMAGE_PATH is not None + assert _ACTION_DIR is not None + response = omni_runner_handler.send_diffusion_request( + { + "model": omni_runner_handler.runner.model_name, + "prompt": "The camera moves slowly forward through the scene.", + "multi_modal_data": {"image": str(Path(_IMAGE_PATH).expanduser().resolve())}, + "sampling_params": OmniDiffusionSamplingParams( + height=480, + width=832, + num_frames=9, + num_inference_steps=4, + max_sequence_length=512, + seed=42, + fps=16, + extra_args={ + "action_path": _ACTION_DIR.name, + "flow_shift": 5.0, + }, + ), + } + ) + + assert response.success + assert response.images is not None + assert len(response.images) == 1 + video = np.asarray(response.images[0]) + assert video.shape == (1, 9, 480, 832, 3) + assert np.isfinite(video).all() + assert float(np.ptp(video)) > 1e-6 + assert float(np.mean(np.abs(np.diff(video, axis=1)))) > 1e-6 diff --git a/tests/e2e/offline_inference/test_longcat_video_avatar.py b/tests/e2e/offline_inference/test_longcat_video_avatar.py index f0fa81eb3..404ad9485 100644 --- a/tests/e2e/offline_inference/test_longcat_video_avatar.py +++ b/tests/e2e/offline_inference/test_longcat_video_avatar.py @@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]: with OmniRunner( model, seed=42, - stage_configs_path=None, model_class_name="LongCatVideoAvatarPipeline", additional_config={ "model_type": "avatar-v1.5", diff --git a/tests/e2e/offline_inference/test_magi_human_expansion.py b/tests/e2e/offline_inference/test_magi_human_expansion.py deleted file mode 100644 index e6b3c6a42..000000000 --- a/tests/e2e/offline_inference/test_magi_human_expansion.py +++ /dev/null @@ -1,162 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""End-to-end tests for MagiHuman pipeline via vLLM-Omni.""" - -import io - -import av -import numpy as np -import pytest - -from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunner -from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -MODEL = "SII-GAIR/daVinci-MagiHuman-Base-1080p" - -# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime -_OMNI_RUNNER_PARAM = ( - MODEL, - None, - { - "init_timeout": 1200, - "tensor_parallel_size": 2, - "vae_use_tiling": True, - }, -) - -_OMNI_RUNNER_PARAM_OFFLOAD = ( - MODEL, - None, - { - "init_timeout": 1200, - "tensor_parallel_size": 2, - "enable_cpu_offload": True, - "vae_use_tiling": True, - }, -) - -pytestmark = [ - pytest.mark.slow, - pytest.mark.diffusion, - pytest.mark.parametrize( - "omni_runner", - [_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD], - indirect=True, - ), -] - - -def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None: - """Validate that the MP4 contains meaningful video and audio tracks.""" - container = av.open(io.BytesIO(video_bytes)) - - v_streams = [s for s in container.streams if s.type == "video"] - assert len(v_streams) >= 1, "No video stream found in MP4" - - a_streams = [s for s in container.streams if s.type == "audio"] - assert len(a_streams) >= 1, "No audio stream found in MP4" - - v_stream = v_streams[0] - assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}" - assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}" - - frame_count = 0 - for frame in container.decode(video=0): - frame_count += 1 - if frame_count >= min_frames: - break - assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})" - - container.close() - - -@hardware_test(res={"cuda": "H100"}, num_cards=2) -def test_magi_human_e2e(omni_runner: OmniRunner): - """End-to-end test for MagiHuman generating video and audio.""" - prompt = ( - "A young woman with long, wavy golden blonde hair and bright blue eyes, " - "wearing a fitted ivory silk blouse with a delicate lace collar, sits " - "stationary in front of a softly lit, blurred warm-toned interior. Her " - "overall disposition is warm, composed, and gently confident. The camera " - "holds a static medium close-up, framing her from the shoulders up, " - "with shallow depth of field keeping her face in sharp focus. Soft " - "directional key light falls from the upper left, casting a gentle " - "highlight along her cheekbone and nose bridge. She draws a quiet breath, " - "the levator labii superiors relaxing as her lips part. She speaks in " - "clear, warm, unhurried American English: " - "\"The most beautiful things in life aren't things at all — " - "they're moments, feelings, and the people who make you feel truly alive.\" " - "Her jaw descends smoothly on each stressed syllable; the orbicularis oris " - "shapes each vowel with precision. A faint, genuine smile engages the " - "zygomaticus major, lifting her lip corners fractionally. Her brows rest " - "in a soft, neutral arch throughout. She maintains steady, forward-facing " - "eye contact. Head position remains level; no torso displacement occurs.\n\n" - "Dialogue:\n" - ": " - "\"The most beautiful things in life aren't things at all — " - "they're moments, feelings, and the people who make you feel truly alive.\"\n\n" - "Background Sound:\n" - "" - ) - - sampling_params = OmniDiffusionSamplingParams( - height=256, - width=448, - num_inference_steps=8, - seed=52, - extra_args={ - "seconds": 5, - "sr_height": 1080, - "sr_width": 1920, - "sr_num_inference_steps": 5, - }, - ) - - omni = omni_runner.omni - outputs = list( - omni.generate( - prompts=[prompt], - sampling_params_list=[sampling_params], - ) - ) - - assert len(outputs) > 0, "No outputs returned" - first = outputs[0] - - assert hasattr(first, "images") and first.images, "No video frames in output" - video_frames = first.images[0] - assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}" - assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}" - - mm = first.multimodal_output - assert mm, "multimodal_output is empty or missing" - - audio_waveform = mm.get("audio") - assert audio_waveform is not None, "No audio waveform in multimodal_output" - - audio_sample_rate = mm.get("audio_sample_rate") - assert audio_sample_rate is not None, ( - "audio_sample_rate not found in multimodal_output; model post-process must propagate it" - ) - assert isinstance(audio_sample_rate, (int, float)), ( - f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}" - ) - assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}" - - fps = mm.get("fps") - assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it" - assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}" - assert int(fps) > 0, f"fps must be positive, got {fps}" - - video_bytes = mux_video_audio_bytes( - video_frames, - audio_waveform, - fps=float(fps), - audio_sample_rate=int(audio_sample_rate), - ) - assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}" - assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)" - - _validate_mp4(video_bytes) diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py index ec1545f6a..f570a6ea7 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py @@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16 MODEL_PATH = "bytedance-research/MammothModa2-Preview" T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml") -_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG}) +_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG) # Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate. _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json" diff --git a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py index 6d4bd197a..f8d43eaba 100644 --- a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py +++ b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import os @@ -53,7 +53,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]] @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) -def test_text_to_text(omni_runner, omni_runner_handler) -> None: +def test_text_to_text(omni_runner, offline_client) -> None: """ Test text-only input processing and text output generation. Input Modal: text @@ -62,14 +62,14 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None: prompt = build_prompt("请详细介绍鹦鹉的生活习性。") request_config = {"prompts": prompt, "modalities": ["text"]} - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) -def test_image_to_text(omni_runner, omni_runner_handler) -> None: +def test_image_to_text(omni_runner, offline_client) -> None: """ Test image understanding with text output. Input Modal: image + text @@ -79,14 +79,14 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None: prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.") request_config = {"prompts": prompt, "images": image, "modalities": ["text"]} - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) -def test_audio_to_text(omni_runner, omni_runner_handler) -> None: +def test_audio_to_text(omni_runner, offline_client) -> None: """ Test audio understanding with text output. Input Modal: audio + text @@ -98,14 +98,14 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None: prompt = build_prompt(f"{AUDIO_TOKEN}Please recognize the language of this speech and transcribe it. Format: oral.") request_config = {"prompts": prompt, "audios": audio, "modalities": ["text"]} - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) -def test_video_to_text(omni_runner, omni_runner_handler) -> None: +def test_video_to_text(omni_runner, offline_client) -> None: """ Test video understanding with text output. Input Modal: video + text @@ -115,14 +115,14 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None: prompt = build_prompt(f"{VIDEO_TOKEN}Describe what is happening in this video.") request_config = {"prompts": prompt, "videos": video, "modalities": ["text"]} - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) -def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: +def test_mixed_to_text(omni_runner, offline_client) -> None: """ Test mixed modality input (image + audio) with text output. Input Modal: image + audio + text @@ -135,14 +135,14 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: prompt = build_prompt(f"{IMAGE_TOKEN}{AUDIO_TOKEN}Describe the image and transcribe the audio.") request_config = {"prompts": prompt, "images": image, "audios": audio, "modalities": ["text"]} - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_text_to_audio(omni_runner, omni_runner_handler) -> None: +def test_text_to_audio(omni_runner, offline_client) -> None: """ Test text input with audio output via the thinker+talker pipeline. Input Modal: text @@ -151,14 +151,14 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None: prompt = build_prompt("请简单介绍一下北京。") request_config = {"prompts": prompt, "modalities": ["audio"]} - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_image_to_audio(omni_runner, omni_runner_handler) -> None: +def test_image_to_audio(omni_runner, offline_client) -> None: """ Test image + text input with audio output via the thinker+talker pipeline. Input Modal: image + text @@ -168,4 +168,4 @@ def test_image_to_audio(omni_runner, omni_runner_handler) -> None: prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.") request_config = {"prompts": prompt, "images": image, "modalities": ["audio"]} - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) diff --git a/tests/e2e/offline_inference/test_ming_tts_expansion.py b/tests/e2e/offline_inference/test_ming_tts_expansion.py index 8dcefd182..e2e848ee0 100644 --- a/tests/e2e/offline_inference/test_ming_tts_expansion.py +++ b/tests/e2e/offline_inference/test_ming_tts_expansion.py @@ -35,7 +35,6 @@ MIN_AUDIO_SAMPLES = 1000 pytestmark = [ pytest.mark.slow, pytest.mark.tts, - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4704"), ] diff --git a/tests/e2e/offline_inference/test_minicpmo_4_5.py b/tests/e2e/offline_inference/test_minicpmo_4_5.py index 6725ce2aa..5a518a7f7 100644 --- a/tests/e2e/offline_inference/test_minicpmo_4_5.py +++ b/tests/e2e/offline_inference/test_minicpmo_4_5.py @@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"] _CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml") -test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models] +test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models] def get_question(prompt_type: str = "text") -> str: diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py index b894fe4ea..41e2229d3 100644 --- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py @@ -25,7 +25,7 @@ from vllm_omni import Omni MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano" STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml") -# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize +# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize _OMNI_RUNNER_PARAM = ( MODEL_NAME, STAGE_CONFIG, diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py index 892a395ba..34d36118d 100644 --- a/tests/e2e/offline_inference/test_moss_tts_realtime.py +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -96,7 +96,7 @@ def _get_test_config() -> str: pytestmark = [ pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.tts, pytest.mark.parametrize( "omni_runner", @@ -271,7 +271,6 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor # --------------------------------------------------------------------------- -@pytest.mark.advanced_model @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None: """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" @@ -283,7 +282,6 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) assert not torch.all(audio == 0), "Audio is silence" -@pytest.mark.advanced_model @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None: """MossTTSRealtime: Chinese input produces non-empty audio.""" diff --git a/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py index 25c671e05..487a6bfa6 100644 --- a/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py @@ -28,10 +28,14 @@ _OMNI_RUNNER_PARAM = ( {"stage_init_timeout": 600, "trust_remote_code": True}, ) +_SKIP_ISSUE_6417 = pytest.mark.skip( + reason="https://github.com/vllm-project/vllm-omni/issues/6417", +) + pytestmark = [ - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4643"), pytest.mark.slow, pytest.mark.tts, + _SKIP_ISSUE_6417, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py index 5b4c7a935..dd65cbdcc 100644 --- a/tests/e2e/offline_inference/test_omnivoice_expansion.py +++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py @@ -20,13 +20,11 @@ from tests.helpers.stage_config import get_deploy_config_path MODEL = "k2-fsa/OmniVoice" DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml") -# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs. -# The migrated test passes deploy_config through extra Omni kwargs. +# OmniRunner tuple: model, deploy config path, extra Omni kwargs. _OMNI_RUNNER_PARAM = ( MODEL, - None, + DEPLOY_CONFIG, { - "deploy_config": DEPLOY_CONFIG, "trust_remote_code": True, "log_stats": True, }, diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py index 6be54969c..41a84d677 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for Qwen2.5-Omni AutoRound W4A16 quantized inference. These tests cover text, audio, image, video, and mixed-modality inputs @@ -36,8 +36,8 @@ BASELINE_MODEL = os.environ.get("QWEN2_5_OMNI_BASELINE_MODEL", BASELINE_MODEL) _CI_DEPLOY = get_deploy_config_path("ci/qwen2_5_omni.yaml") -_SKIP_ISSUE_4663 = pytest.mark.skip( - reason="https://github.com/vllm-project/vllm-omni/issues/4663", +_SKIP_ISSUE_5652 = pytest.mark.skip( + reason="https://github.com/vllm-project/vllm-omni/issues/5652", ) @@ -65,16 +65,16 @@ quant_params = [(QUANTIZED_MODEL, stage_config)] # ------------------------------------------------------------------ -@_SKIP_ISSUE_4663 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_text_to_text(omni_runner, omni_runner_handler): +def test_text_to_text(omni_runner, offline_client): """Text input → text output with W4A16 quantized Qwen2.5-Omni.""" request_config = { "prompts": "What is the capital of China?", "modalities": ["text"], } - response = omni_runner_handler.send_omni_request(request_config) + response = offline_client.send_omni_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -84,10 +84,10 @@ def test_text_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_4663 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_audio_to_text(omni_runner, omni_runner_handler): +def test_audio_to_text(omni_runner, offline_client): """Audio input → text output with W4A16 quantized Qwen2.5-Omni.""" audio = generate_synthetic_audio(1, 1, 16000)["np_array"] if len(audio.shape) == 2: @@ -98,7 +98,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler): "audios": (audio, 16000), "modalities": ["text"], } - response = omni_runner_handler.send_omni_request(request_config) + response = offline_client.send_omni_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -108,10 +108,10 @@ def test_audio_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_4663 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_image_to_text(omni_runner, omni_runner_handler): +def test_image_to_text(omni_runner, offline_client): """Image input → text output with W4A16 quantized Qwen2.5-Omni.""" image = generate_synthetic_image(16, 16)["np_array"] @@ -120,7 +120,7 @@ def test_image_to_text(omni_runner, omni_runner_handler): "images": image, "modalities": ["text"], } - response = omni_runner_handler.send_omni_request(request_config) + response = offline_client.send_omni_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -130,10 +130,10 @@ def test_image_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_4663 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_video_to_text(omni_runner, omni_runner_handler): +def test_video_to_text(omni_runner, offline_client): """Video input → text output with W4A16 quantized Qwen2.5-Omni.""" video = generate_synthetic_video(16, 16, 30)["np_array"] @@ -142,7 +142,7 @@ def test_video_to_text(omni_runner, omni_runner_handler): "videos": video, "modalities": ["text"], } - response = omni_runner_handler.send_omni_request(request_config) + response = offline_client.send_omni_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -152,10 +152,10 @@ def test_video_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_4663 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_mix_to_audio(omni_runner, omni_runner_handler): +def test_mix_to_audio(omni_runner, offline_client): """Mixed-modality input → audio output with W4A16 quantized Qwen2.5-Omni.""" video = generate_synthetic_video(16, 16, 30)["np_array"] image = generate_synthetic_image(16, 16)["np_array"] @@ -170,5 +170,5 @@ def test_mix_to_audio(omni_runner, omni_runner_handler): "audios": (audio, 16000), "modalities": ["audio"], } - response = omni_runner_handler.send_omni_request(request_config) + response = offline_client.send_omni_request(request_config) assert response.success, "Request failed" diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py b/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py index 9d368a688..e2465234a 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ E2E tests for Qwen2.5-Omni model with mixed modality inputs, audio and text output. """ @@ -55,7 +58,7 @@ def get_question(prompt_type="mix"): @pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3}) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_mix_to_audio(omni_runner, omni_runner_handler) -> None: +def test_mix_to_audio(omni_runner, offline_client) -> None: """ Test multi-modal input processing and text/audio output generation via OpenAI API. Deploy Setting: default yaml @@ -79,14 +82,14 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None: } # Test single completion - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3}) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_text_to_text(omni_runner, omni_runner_handler) -> None: +def test_text_to_text(omni_runner, offline_client) -> None: """ Test text input processing and text output generation via OpenAI API. Deploy Setting: default yaml @@ -99,4 +102,4 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None: request_config = {"prompts": get_question("text_only"), "modalities": ["text"]} # Test single completion - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index fd3a27d2e..32c345194 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -1,5 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ E2E offline tests for Omni model with video input and audio output. + +Abort / sleep-admission lives in ``test_qwen3_omni_colocate_async.py`` so it +does not overlap this module's OmniRunners. """ import os @@ -57,23 +63,23 @@ def get_question(prompt_type="video"): @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) -def test_video_to_audio(omni_runner, omni_runner_handler) -> None: +def test_video_to_audio(omni_runner, offline_client) -> None: """Test processing video, generating audio output.""" video = generate_synthetic_video(224, 224, 300)["np_array"] request_config = {"prompts": get_question(), "videos": video, "modalities": ["audio"]} # Test single completion - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) @pytest.mark.advanced_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_runner", thinker_test_params, indirect=True) -def test_thinker_only_model_request(omni_runner, omni_runner_handler) -> None: +def test_thinker_only_model_request(omni_runner, offline_client) -> None: """Test that we can load and run a request through a model that only has the thinker stage.""" request_config = {"prompts": "what color is the sky?", "modalities": ["text"]} # Test single completion - omni_runner_handler.send_omni_request(request_config) + offline_client.send_omni_request(request_config) diff --git a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py index 9111740ea..b1f17356b 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for Qwen3-Omni AutoRound W4A16 quantized inference. These tests cover text, audio, image, video, and mixed-modality inputs @@ -27,8 +27,8 @@ pytestmark = [ pytest.mark.omni, ] -_SKIP_ISSUE_3195 = pytest.mark.skip( - reason="https://github.com/vllm-project/vllm-omni/issues/3195", +_SKIP_ISSUE_5652 = pytest.mark.skip( + reason="https://github.com/vllm-project/vllm-omni/issues/5652", ) QUANTIZED_MODEL = "Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound" @@ -77,16 +77,16 @@ quant_params = [(QUANTIZED_MODEL, stage_config)] # ------------------------------------------------------------------ -@_SKIP_ISSUE_3195 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_text_to_text(omni_runner, omni_runner_handler): +def test_text_to_text(omni_runner, offline_client): """Text input → text output with W4A16 quantized Qwen3-Omni.""" request_config = { "prompts": "What is the capital of France?", "modalities": ["text"], } - response = omni_runner_handler.send_request(request_config) + response = offline_client.send_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -96,10 +96,10 @@ def test_text_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_3195 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_audio_to_text(omni_runner, omni_runner_handler): +def test_audio_to_text(omni_runner, offline_client): """Audio input → text output with W4A16 quantized Qwen3-Omni.""" audio = generate_synthetic_audio(1, 1, 16000)["np_array"] if len(audio.shape) == 2: @@ -110,7 +110,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler): "audios": (audio, 16000), "modalities": ["text"], } - response = omni_runner_handler.send_request(request_config) + response = offline_client.send_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -120,10 +120,10 @@ def test_audio_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_3195 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_image_to_text(omni_runner, omni_runner_handler): +def test_image_to_text(omni_runner, offline_client): """Image input → text output with W4A16 quantized Qwen3-Omni.""" image = generate_synthetic_image(16, 16)["np_array"] @@ -132,7 +132,7 @@ def test_image_to_text(omni_runner, omni_runner_handler): "images": image, "modalities": ["text"], } - response = omni_runner_handler.send_request(request_config) + response = offline_client.send_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -142,10 +142,10 @@ def test_image_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_3195 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_video_to_text(omni_runner, omni_runner_handler): +def test_video_to_text(omni_runner, offline_client): """Video input → text output with W4A16 quantized Qwen3-Omni.""" video = generate_synthetic_video(224, 224, 300)["np_array"] @@ -154,7 +154,7 @@ def test_video_to_text(omni_runner, omni_runner_handler): "videos": video, "modalities": ["text"], } - response = omni_runner_handler.send_request(request_config) + response = offline_client.send_request(request_config) assert response.success, "Request failed" assert response.text_content and len(response.text_content.strip()) > 0 @@ -164,10 +164,10 @@ def test_video_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_3195 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_video_to_audio(omni_runner, omni_runner_handler): +def test_video_to_audio(omni_runner, offline_client): """Video input → audio output with W4A16 quantized Qwen3-Omni.""" video = generate_synthetic_video(224, 224, 300)["np_array"] @@ -176,7 +176,7 @@ def test_video_to_audio(omni_runner, omni_runner_handler): "videos": video, "modalities": ["audio"], } - response = omni_runner_handler.send_request(request_config) + response = offline_client.send_request(request_config) assert response.success, "Request failed" @@ -185,10 +185,10 @@ def test_video_to_audio(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ -@_SKIP_ISSUE_3195 +@_SKIP_ISSUE_5652 @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_mix_to_audio(omni_runner, omni_runner_handler): +def test_mix_to_audio(omni_runner, offline_client): """Mixed-modality input → audio output with W4A16 quantized Qwen3-Omni.""" video = generate_synthetic_video(224, 224, 300)["np_array"] image = generate_synthetic_image(16, 16)["np_array"] @@ -203,5 +203,5 @@ def test_mix_to_audio(omni_runner, omni_runner_handler): "audios": (audio, 16000), "modalities": ["audio"], } - response = omni_runner_handler.send_request(request_config) + response = offline_client.send_request(request_config) assert response.success, "Request failed" diff --git a/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py b/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py new file mode 100644 index 000000000..54e9b8f62 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Isolated AsyncOmni abort / sleep-admission regression for Qwen3-Omni. + +This lives in its own module so the L3 merge job can run it *after* +``test_qwen3_omni.py`` (offline) and the online serving file. Those modules keep +module-scoped OmniRunner / OmniServer instances until teardown; overlapping a +CuMem sleep-mode engine with them kills StageEngineCoreProc at handshake +(``Failed core proc(s): {}``). +""" + +import asyncio +import os +from contextlib import ExitStack + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest +from vllm import SamplingParams +from vllm.inputs import TokensPrompt + +from tests.helpers.clean import wait_for_gpu_memory_to_clear +from tests.helpers.mark import hardware_test +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.platforms import current_omni_platform + +MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + + +def _output_token_ids(output) -> list[int]: + if not getattr(output, "outputs", None): + return [] + completion = output.outputs[0] + token_ids = list(completion.token_ids or []) + if token_ids: + return token_ids + return list(getattr(completion, "cumulative_token_ids", None) or []) + + +def _finish_reason(output) -> str | None: + if not getattr(output, "outputs", None): + return None + return getattr(output.outputs[0], "finish_reason", None) + + +def _colocate_async_deploy() -> str: + """Thinker-only Instruct deploy with sleep mode on a single GPU. + + Abort / sleep admission is a thinker control-plane contract. The 3-stage + CI yaml plus CuMem on every stage fails engine-core startup in the L3 + merge job: Stage 1 resolves ``devices: "1"`` while only GPU 0 is visible. + """ + return modify_stage_config( + get_deploy_config_path("qwen3_omni_moe_thinking.yaml"), + updates={ + "stages": { + 0: { + "enable_sleep_mode": True, + "enforce_eager": True, + "trust_remote_code": True, + "tensor_parallel_size": 1, + "devices": "0", + "gpu_memory_utilization": 0.9, + "max_num_seqs": 1, + "max_model_len": 8192, + "max_num_batched_tokens": 8192, + "enable_prefix_caching": False, + "skip_mm_profiling": True, + }, + }, + }, + ) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) +@pytest.mark.usefixtures("clean_gpu_memory_between_tests") +@pytest.mark.asyncio +async def test_colocate_async_abort_tokens_and_sleep_admission() -> None: + """**required** regression. Control-plane APIs (``abort`` / ``sleep`` / + ``wake_up`` / ``resume_generation``) are not in ``send_omni_request``. + + On ``origin/main`` this test fails: + - abort: ``generate()`` never yields ``finish_reason="abort"`` with the + tokens produced so far, so resume has an empty prefix (or the stream + hangs after frontend state is dropped). + - sleep: ``generate()`` is admitted into EngineCore while weights/KV are + offloaded, so the task errors (asleep / corrupted ADD frame) instead of + waiting for ``resume_generation()``. + + On this PR both legs pass: abort returns a resumeable prefix, and sleep + holds ``generate()`` until wake + resume. + """ + prompt = "What color is the sky? Write a long, detailed explanation." + device_count = current_omni_platform.device_count() + if device_count > 0: + # Fail closed if a prior module-scoped runner is still on the GPU. + # ``clean_gpu_memory_between_tests`` only logs a note on timeout. + wait_for_gpu_memory_to_clear( + devices=list(range(device_count)), + threshold_ratio=0.15, + timeout_s=120, + ) + + with ExitStack() as after: + engine = AsyncOmni( + model=MODEL, + deploy_config=_colocate_async_deploy(), + enable_sleep_mode=True, + ) + after.callback(engine.shutdown) + + request_id = "qwen3-abort-partial" + outputs: list = [] + + async def _generate(max_tokens: int, req_id: str, gen_prompt) -> None: + async for output in engine.generate( + prompt=gen_prompt, + request_id=req_id, + sampling_params=SamplingParams(temperature=0.0, max_tokens=max_tokens), + output_modalities=["text"], + ): + outputs.append(output) + + abort_task = asyncio.create_task(_generate(256, request_id, prompt)) + prefix: list[int] = [] + prompt_token_ids: list[int] = [] + for _ in range(600): + if abort_task.done(): + break + if outputs: + latest = outputs[-1] + prefix = _output_token_ids(latest) + prompt_token_ids = list(getattr(latest, "prompt_token_ids", None) or []) + if prefix: + break + await asyncio.sleep(0.1) + + assert prefix, "generate produced no tokens before abort" + assert not abort_task.done(), "generate finished before abort; raise max_tokens" + await engine.abort(request_id) + await asyncio.wait_for(abort_task, timeout=90) + + final = outputs[-1] + assert final.finished + assert _finish_reason(final) == "abort" + abort_prefix = _output_token_ids(final) + assert abort_prefix, "abort dropped the generated prefix (main-branch behavior)" + + resume_prompt: str | TokensPrompt = ( + TokensPrompt(prompt_token_ids=prompt_token_ids + abort_prefix) if prompt_token_ids else prompt + ) + outputs.clear() + async for output in engine.generate( + prompt=resume_prompt, + request_id="qwen3-abort-resume", + sampling_params=SamplingParams(temperature=0.0, max_tokens=8), + output_modalities=["text"], + ): + outputs.append(output) + assert outputs + assert any(_output_token_ids(out) for out in outputs) + + # Trainer order is pause → abort → sleep. Pause here after the + # resume generate so EngineCore is idle before CuMem offload. + await engine.pause_generation(mode="abort", clear_cache=True) + await engine.sleep(level=1) + outputs.clear() + sleep_task = asyncio.create_task(_generate(8, "qwen3-sleep-admission", prompt)) + await asyncio.sleep(1.0) + assert not sleep_task.done(), "generate() ran while EngineCore was sleeping (main-branch admission race)" + await engine.wake_up() + await asyncio.sleep(0.5) + assert not sleep_task.done(), "generate() resumed before resume_generation()" + await engine.resume_generation() + await asyncio.wait_for(sleep_task, timeout=180) + assert outputs + assert _finish_reason(outputs[-1]) != "abort" + await asyncio.sleep(5) diff --git a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py index 5d1b3dda7..ad6935e71 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py +++ b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for Qwen3-Omni ModelOpt NVFP4 W4A4 — Marlin fallback path. Loads the published full-thinker NVFP4 W4A4 checkpoint on H100 (sm_90), @@ -74,7 +74,7 @@ quant_params = [(QUANTIZED_MODEL, _stage_config())] @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) -def test_marlin_fallback_no_nan_collapse(omni_runner, omni_runner_handler): +def test_marlin_fallback_no_nan_collapse(omni_runner, offline_client): """H100 Marlin FP4 fallback: text in -> coherent text out. On H100 the upstream NVFP4 PWAL routes through @@ -89,7 +89,7 @@ def test_marlin_fallback_no_nan_collapse(omni_runner, omni_runner_handler): degenerate-output signature that indicates NaN propagation through the FP4 GEMM. """ - response = omni_runner_handler.send_omni_request({"prompts": "Why is the sky blue?", "modalities": ["text"]}) + response = offline_client.send_omni_request({"prompts": "Why is the sky blue?", "modalities": ["text"]}) assert response.success, "request failed" text = (response.text_content or "").strip() assert text, "empty response — likely server returned only EOS" diff --git a/tests/e2e/offline_inference/test_qwen3_tts_base.py b/tests/e2e/offline_inference/test_qwen3_tts_base.py index a076526a4..f800a1163 100644 --- a/tests/e2e/offline_inference/test_qwen3_tts_base.py +++ b/tests/e2e/offline_inference/test_qwen3_tts_base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E offline tests for Qwen3-TTS Base model with text input and audio output. @@ -15,12 +15,12 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" # See tests/e2e/online_serving/test_qwen3_tts_base.py for the vendored-asset rationale. -REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." @@ -70,7 +70,7 @@ def get_prompt(): @pytest.mark.tts @hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1) @pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: +def test_text_to_audio_001(omni_runner, offline_client) -> None: """ Test text input processing and audio output via offline Omni runner. Deploy Setting: qwen3_tts_no_async_chunk.yaml + enforce_eager=true @@ -87,4 +87,4 @@ def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - omni_runner_handler.send_audio_speech_request(request_config) + offline_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py index 4fa2d4b6f..b303b1cda 100644 --- a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py +++ b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E offline tests for Qwen3-TTS CustomVoice model with text input and audio output. @@ -64,7 +64,7 @@ def get_prompt(): @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: +def test_text_to_audio_001(omni_runner, offline_client) -> None: """ Test text input processing and audio output via offline Omni runner. Deploy Setting: qwen3_tts_no_async_chunk.yaml + enforce_eager=true @@ -74,4 +74,4 @@ def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: Datasets: few requests """ request_config = {"input": get_prompt(), "voice": "vivian"} - omni_runner_handler.send_audio_speech_request(request_config) + offline_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/offline_inference/test_sana_wm.py b/tests/e2e/offline_inference/test_sana_wm.py new file mode 100644 index 000000000..8b1c460c2 --- /dev/null +++ b/tests/e2e/offline_inference/test_sana_wm.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import os +from typing import Any + +import numpy as np +import pytest + +pytestmark = [ + pytest.mark.advanced_model, + pytest.mark.diffusion, + pytest.mark.gpu, +] + + +def _make_sana_wm_e2e_image() -> Any: + from PIL import Image + + from vllm_omni.diffusion.models.sana_wm import SANA_WM_OUTPUT_HEIGHT, SANA_WM_OUTPUT_WIDTH + + return Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160)) + + +def _sana_wm_e2e_tensor_parallel_size() -> int: + return int( + os.environ.get( + "SANA_WM_E2E_TENSOR_PARALLEL_SIZE", + os.environ.get("SANA_WM_E2E_TP", "1"), + ) + ) + + +def _coerce_video_array(video: Any) -> np.ndarray: + import torch + + if isinstance(video, list): + assert video, "SANA-WM e2e produced an empty video list." + video = video[0] + if isinstance(video, torch.Tensor): + video = video.detach().cpu().float().numpy() + video_array = np.asarray(video) + if video_array.ndim == 5: + assert video_array.shape[0] == 1 + video_array = video_array[0] + assert video_array.ndim == 4 + return video_array + + +def _assert_sana_wm_e2e_shape( + video: np.ndarray, + *, + output_type: str, + num_frames: int, +) -> None: + if output_type == "latent": + # Latent output is channel-first: (C, T, H, W) after removing the + # batch dimension. Stage-1 emits LTX-2 latents with 128 channels. + assert video.shape[0] == 128 + assert 0 < video.shape[1] <= num_frames + return + assert 0 < video.shape[0] <= num_frames + + +def _run_sana_wm_e2e(*, output_type: str) -> np.ndarray: + import torch + + from vllm_omni.diffusion.models.sana_wm import ( + SANA_WM_MODEL_ID, + SANA_WM_OUTPUT_HEIGHT, + SANA_WM_OUTPUT_WIDTH, + ) + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.outputs import OmniRequestOutput + + if not torch.cuda.is_available(): + pytest.skip("Sana-WM e2e requires CUDA.") + + model = os.environ.get("SANA_WM_E2E_MODEL", SANA_WM_MODEL_ID) + model_class_name = os.environ.get("SANA_WM_E2E_MODEL_CLASS", "SanaWmPipeline") + num_frames = int(os.environ.get("SANA_WM_E2E_NUM_FRAMES", "9")) + stage1_steps = int( + os.environ.get( + "SANA_WM_E2E_STAGE1_STEPS", + os.environ.get("SANA_WM_E2E_NUM_INFERENCE_STEPS", "1"), + ) + ) + action = os.environ.get("SANA_WM_E2E_ACTION", f"w-{max(num_frames - 1, 1)}") + tensor_parallel_size = _sana_wm_e2e_tensor_parallel_size() + image = _make_sana_wm_e2e_image() + omni = Omni( + model=model, + model_class_name=model_class_name, + enforce_eager=True, + tensor_parallel_size=tensor_parallel_size, + ) + extra_args: dict[str, Any] = {} + # The native Stage-1 path is latent-token capped; raise it for full-size + # e2e runs (704x1280 exceeds the default 4096-token cap). + native_max_tokens = os.environ.get("SANA_WM_E2E_NATIVE_MAX_TOKENS", "").strip() + if native_max_tokens: + extra_args["sana_wm_native_max_tokens"] = int(native_max_tokens) + # Return the requested output type so the shape assertion matches the + # pipeline output instead of the raw Stage-1 latent default. + + output = omni.generate( + { + "prompt": "A slow forward camera move through a quiet city street.", + "multi_modal_data": {"image": image}, + "sana_wm": { + "action": action, + "num_frames": num_frames, + "translation_speed": 0.055, + "rotation_speed_deg": 1.2, + # Avoid the optional Pi3X dependency by passing deterministic + # camera intrinsics directly. + "intrinsics": { + "fx": SANA_WM_OUTPUT_WIDTH / 2, + "fy": SANA_WM_OUTPUT_WIDTH / 2, + "cx": SANA_WM_OUTPUT_WIDTH / 2, + "cy": SANA_WM_OUTPUT_HEIGHT / 2, + }, + }, + }, + OmniDiffusionSamplingParams( + height=SANA_WM_OUTPUT_HEIGHT, + width=SANA_WM_OUTPUT_WIDTH, + num_frames=num_frames, + seed=0, + fps=16, + num_inference_steps=stage1_steps, + guidance_scale=1.0, + guidance_scale_provided=True, + output_type=output_type, + extra_args=extra_args, + ), + ) + + request_output = output[0] if isinstance(output, list) else output + assert isinstance(request_output, OmniRequestOutput) + assert request_output.error is None + assert request_output.images + + frames = request_output.images[0] + video = _coerce_video_array(frames) + _assert_sana_wm_e2e_shape(video, output_type=output_type, num_frames=num_frames) + return video + + +def test_sana_wm_native_generates_video() -> None: + output_type = os.environ.get("SANA_WM_E2E_OUTPUT_TYPE", "np") + video = _run_sana_wm_e2e(output_type=output_type) + if output_type == "latent": + assert video.ndim == 4 + else: + assert video.shape[-1] in (3, 4) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py index 4f5c97d11..42510d7c6 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py @@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None: """Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation.""" lora_request = _make_file_lora_request(tmp_path / "sensenova_lora") - with OmniRunner(MODEL, stage_configs_path=None) as runner: + with OmniRunner(MODEL) as runner: omni = runner.omni baseline = _generate(omni) img_1x = _generate(omni, lora_request, lora_scale=1.0) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py index 1ac76da28..c2739844c 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py @@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img( def test_sensenova_u1_img2img(run_level): """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML).""" input_image = _load_input_image() - with OmniRunner( - "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, - ) as runner: + with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner: generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image) if run_level == "full_model" or run_level == "advanced_model": _validate_pixels(generated_image) @@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit(): input_image = _load_input_image() with OmniRunner( "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, cache_backend="cache_dit", ) as runner: _generate_sensenova_u1_img2img(runner.omni, input_image) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py index 092da16b9..2e07e2076 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py @@ -151,10 +151,7 @@ def _generate_sensenova_u1_image( @hardware_test(res={"cuda": "H100"}) def test_sensenova_u1_text2img(run_level): """Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML).""" - with OmniRunner( - "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, - ) as runner: + with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner: generated_image = _generate_sensenova_u1_image(runner.omni) if run_level == "full_model" or run_level == "advanced_model": _validate_pixels(generated_image) @@ -168,7 +165,6 @@ def test_sensenova_u1_text2img_cache_dit(): """Test SenseNova-U1 text2img with Cache-DiT enabled.""" with OmniRunner( "SenseNova/SenseNova-U1-8B-MoT", - stage_configs_path=None, cache_backend="cache_dit", ) as runner: _generate_sensenova_u1_image(runner.omni) diff --git a/tests/e2e/offline_inference/test_soulxsinger_expansion.py b/tests/e2e/offline_inference/test_soulxsinger_expansion.py deleted file mode 100644 index c3a6091d6..000000000 --- a/tests/e2e/offline_inference/test_soulxsinger_expansion.py +++ /dev/null @@ -1,241 +0,0 @@ -"""E2E offline inference tests for SoulX-Singer (single-stage, preprocess inline).""" - -import functools -import importlib -import json -import os -import shutil -from pathlib import Path - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - -import numpy as np -import pytest - -from tests.helpers.mark import hardware_test -from tests.helpers.media import get_asset_path -from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import get_deploy_config_path -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -PROMPT_AUDIO = get_asset_path("soulxsinger/zh_prompt.mp3") -TARGET_AUDIO = get_asset_path("soulxsinger/music.mp3") -PROMPT_METADATA = get_asset_path("soulxsinger/zh_prompt.json") -TARGET_METADATA = get_asset_path("soulxsinger/music.json") -SAMPLE_RATE = 24_000 - -# phone_set.json (phoneme vocab) is not shipped on HuggingFace; SVS loads it -# from the model dir. Stage it from this pinned upstream commit so the SVS DiT -# can run in CI; the SVS tests skip gracefully if it cannot be fetched. -_PHONE_SET_URL = ( - "https://raw.githubusercontent.com/Soul-AILab/SoulX-Singer/" - "81aeb3ae772c70093c3de74dc23c92d983801ae4/soulxsinger/utils/phoneme/phone_set.json" -) - -if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file(): - pytest.skip( - f"Missing SoulX-Singer audio assets: {PROMPT_AUDIO.name}, {TARGET_AUDIO.name}", - allow_module_level=True, - ) - -pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.tts] - -_CASES = ( - pytest.param( - "SoulXSingerPipeline", - "soulxsinger_svs.yaml", - { - "language": "Mandarin", - "vocal_sep": False, - "control": "score", - "auto_shift": False, - "pitch_shift": 0, - }, - ("g2pM", "g2p_en"), - id="svs", - ), - pytest.param( - "SoulXSingerSVCPipeline", - "soulxsinger_svc.yaml", - {"vocal_sep": False, "auto_shift": False, "pitch_shift": 0}, - (), - id="svc", - ), -) - - -@functools.lru_cache(maxsize=1) -def _resolve_weights() -> tuple[Path, Path, Path]: - from huggingface_hub import snapshot_download - - base = Path(snapshot_download("Soul-AILab/SoulX-Singer", allow_patterns=["*"])) - - # phone_set.json is not on HF; best-effort stage it from pinned upstream so - # SVS can load. SVS tests skip (not fail) if both this and a manual copy are absent. - phone_set = base / "phoneme" / "phone_set.json" - if not phone_set.is_file() and not (base / "phone_set.json").is_file(): - try: - from urllib.request import urlretrieve - - phone_set.parent.mkdir(exist_ok=True) - urlretrieve(_PHONE_SET_URL, phone_set) - except Exception: - pass - - # make a temporary svc directory - svc_dir = base.parent / "SoulX-Singer-SVC" - svc_dir.mkdir(parents=True, exist_ok=True) - shutil.copyfile(base / "config.yaml", svc_dir / "config.yaml") - shutil.move(base / "model-svc.pt", svc_dir / "model-svc.pt") - - if raw := os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR"): - pre = Path(raw).expanduser().resolve() - if (pre / "rmvpe" / "rmvpe.pt").is_file(): - return base, svc_dir, pre - - from huggingface_hub import snapshot_download - - pre = Path(snapshot_download("Soul-AILab/SoulX-Singer-Preprocess", allow_patterns=["*"])) - return base, svc_dir, pre - - -@pytest.fixture(scope="session") -def soulx_weights() -> tuple[Path, Path, Path]: - try: - return _resolve_weights() - except Exception as exc: - pytest.skip(f"Set SOULXSINGER_MODEL_DIR / SOULX_PREPROCESS_WEIGHTS_DIR. ({exc})") - - -def _flatten_audio(audio_val) -> np.ndarray: - import torch - - if isinstance(audio_val, list): - chunks = [c.detach().cpu().float().numpy().reshape(-1) for c in audio_val if c is not None] - return np.concatenate(chunks) if chunks else np.array([], dtype=np.float32) - if isinstance(audio_val, torch.Tensor): - return audio_val.detach().cpu().float().numpy().reshape(-1) - return np.asarray(audio_val, dtype=np.float32).reshape(-1) - - -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("architecture,deploy_yaml,extra_args,py_deps", _CASES) -def test_soulxsinger_multistage_from_audio( - soulx_weights: tuple[Path, Path, Path], - architecture: str, - deploy_yaml: str, - extra_args: dict, - py_deps: tuple[str, ...], -) -> None: - for mod in py_deps: - try: - importlib.import_module(mod) - except ImportError as exc: - pytest.fail(f"SoulX SVS requires {mod}: {exc}") - - base_dir, svc_dir, preprocess_dir = soulx_weights - - # SVS mode requires phone_set.json in the model directory - if architecture == "SoulXSingerPipeline": - json.dump( - {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1}, - open(base_dir / "config.json", "w"), - indent=4, - ) - model = str(base_dir) - if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file(): - pytest.skip( - "SoulX-Singer SVS test requires phoneme/phone_set.json. " - "Copy it from github.com/Soul-AILab/SoulX-Singer into the model dir. " - "See `examples/offline_inference/text_to_speech/README.md` for details." - ) - if architecture == "SoulXSingerSVCPipeline": - json.dump( - {"model_type": "soulxsinger", "architectures": ["SoulXSingerSVCPipeline"], "max_num_seqs": 1}, - open(svc_dir / "config.json", "w"), - indent=4, - ) - model = str(svc_dir) - - with OmniRunner( - model, - stage_configs_path=get_deploy_config_path(deploy_yaml), - async_chunk=False, - ) as runner: - sampling = OmniDiffusionSamplingParams( - num_inference_steps=4, - guidance_scale=3.0, - seed=42, - extra_args={ - "prompt_audio": str(PROMPT_AUDIO), - "target_audio": str(TARGET_AUDIO), - "preprocess_weights_dir": str(preprocess_dir), - **extra_args, - }, - ) - prompt = {"prompt_token_ids": [0]} - outputs = runner.generate([prompt], sampling) - - assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output" - mm = outputs[0].multimodal_output - assert isinstance(mm, dict) and "audio" in mm - audio = _flatten_audio(mm["audio"]) - assert 12_000 <= audio.size - assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4 - duration_s = audio.size / SAMPLE_RATE - assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s" - - -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> None: - """SVS from precomputed note/lyric metadata. - - Exercises the SVS DiT + vocoder serving path while skipping the heavy - upstream preprocess (funasr ASR + ROSVOT note transcription), so it runs in - the merge gate without those non-pip dependencies. - """ - base_dir, _svc_dir, preprocess_dir = soulx_weights - - if not PROMPT_METADATA.is_file() or not TARGET_METADATA.is_file(): - pytest.skip(f"Missing SoulX-Singer metadata fixtures: {PROMPT_METADATA.name}, {TARGET_METADATA.name}") - if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file(): - pytest.skip( - "SoulX-Singer SVS requires phoneme/phone_set.json (upstream fetch unavailable). " - "See `examples/offline_inference/text_to_speech/README.md`." - ) - - json.dump( - {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1}, - open(base_dir / "config.json", "w"), - indent=4, - ) - - with OmniRunner( - str(base_dir), - stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"), - async_chunk=False, - ) as runner: - sampling = OmniDiffusionSamplingParams( - num_inference_steps=4, - guidance_scale=3.0, - seed=42, - extra_args={ - "language": "Mandarin", - "control": "score", - "prompt_metadata_path": str(PROMPT_METADATA), - "target_metadata_path": str(TARGET_METADATA), - "audio_path": str(PROMPT_AUDIO), - "preprocess_weights_dir": str(preprocess_dir), - }, - ) - prompt = {"prompt_token_ids": [0]} - outputs = runner.generate([prompt], sampling) - - assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output" - mm = outputs[0].multimodal_output - assert isinstance(mm, dict) and "audio" in mm - audio = _flatten_audio(mm["audio"]) - assert 12_000 <= audio.size - assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4 - duration_s = audio.size / SAMPLE_RATE - assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s" diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index 00ecd09d9..759c6c500 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stable Audio offline e2e: real weights, FP8 + TeaCache (single job to save GPU). @@ -15,7 +15,6 @@ import numpy as np import pytest import torch -from tests.helpers import skip_if_gated_repo_inaccessible from tests.helpers.assertions import assert_audio_valid from tests.helpers.mark import hardware_test from vllm_omni import Omni @@ -25,8 +24,6 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.slow, pytest.mark.diffusion] -_MODEL_REPO = "stabilityai/stable-audio-open-1.0" - _SAMPLE_RATE = 44100 _CLIP_DURATION_S = 2.0 @@ -80,12 +77,10 @@ def test_stable_audio_quantization_and_teacache() -> None: CI should provide ``HF_TOKEN`` if the checkpoint is gated. """ - skip_if_gated_repo_inaccessible(_MODEL_REPO) # ``model_class_name`` must be passed explicitly: the default-stage-cfg # factory in ``async_omni_engine.py`` reads it out of ``kwargs`` when # deciding ``final_output_type`` (#2077), and at construction time the - # auto-resolution from ``model_index.json`` has not run yet. AudioX's - # offline test follows the same pattern. + # auto-resolution from ``model_index.json`` has not run yet. m = Omni( model="stabilityai/stable-audio-open-1.0", model_class_name="StableAudioPipeline", diff --git a/tests/e2e/offline_inference/test_step_audio2_expansion.py b/tests/e2e/offline_inference/test_step_audio2_expansion.py index 483396b23..3a1b60c94 100644 --- a/tests/e2e/offline_inference/test_step_audio2_expansion.py +++ b/tests/e2e/offline_inference/test_step_audio2_expansion.py @@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput MODEL = "stepfun-ai/Step-Audio-2-mini" DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml") -TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})] +TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})] pytestmark = [pytest.mark.slow, pytest.mark.tts] diff --git a/tests/e2e/offline_inference/test_voxcpm2_tts.py b/tests/e2e/offline_inference/test_voxcpm2_tts.py index c10694a40..af55c51aa 100644 --- a/tests/e2e/offline_inference/test_voxcpm2_tts.py +++ b/tests/e2e/offline_inference/test_voxcpm2_tts.py @@ -39,8 +39,7 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor: return audio -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None: @@ -53,8 +52,7 @@ def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None: assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s" -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None: @@ -95,8 +93,7 @@ def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None: assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s" -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voxcpm2_prefill_decode_mixed_batch_003(omni_runner: OmniRunner) -> None: diff --git a/tests/e2e/offline_inference/test_voxtral_tts_expansion.py b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py index 2441c46dc..dcff41cda 100644 --- a/tests/e2e/offline_inference/test_voxtral_tts_expansion.py +++ b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py @@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming(): async def _run(): async_omni = AsyncOmni( model=MODEL, - stage_configs_path=STAGE_CONFIG, + deploy_config=STAGE_CONFIG, stage_init_timeout=300, enforce_eager=True, ) diff --git a/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py b/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py new file mode 100644 index 000000000..184f0f7d2 --- /dev/null +++ b/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E smoke test for Wan2.1 AutoRound MXFP4 inference.""" + +import os + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +QUANTIZED_MODEL = os.environ.get( + "WAN21_AUTOROUND_MXFP4_MODEL", + "INCModel/Wan2.1-T2V-1.3B-Diffusers-MXFP4-AutoRound", +) + +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + +@hardware_test(res={"xpu": "B60"}) +@pytest.mark.parametrize( + "omni_runner", + [(QUANTIZED_MODEL, None)], + indirect=True, +) +def test_wan21_autoround_mxfp4_generates_video(omni_runner, omni_runner_handler): + sampling_params = OmniDiffusionSamplingParams( + height=256, + width=256, + num_frames=5, + num_inference_steps=1, + guidance_scale=1.0, + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + ) + response = omni_runner_handler.send_diffusion_request( + { + "prompt": "A red toy car driving on a table", + "sampling_params": sampling_params, + } + ) + + assert response.success, response.error_message + assert response.images is not None and len(response.images) == 1 + frames = np.asarray(response.images[0]) + assert frames.shape[1:4] == (5, 256, 256) + assert frames.std() > 0.01 diff --git a/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py index 044435cb5..4ac51a344 100644 --- a/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E tests for Wan2.2 AutoRound W4A16 quantized inference. These tests cover I2V (image-to-video) and T2V (text-to-video) generation @@ -19,8 +19,8 @@ import pytest import torch from PIL import Image -from tests.helpers.env import DeviceMemoryMonitor from tests.helpers.mark import hardware_test +from tests.helpers.monitor import DeviceMemoryMonitor from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -91,7 +91,7 @@ def _create_test_image(width: int = WIDTH, height: int = HEIGHT) -> Image.Image: return Image.fromarray(arr) -def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table, smooth motion") -> tuple: +def _generate_i2v_video(offline_client, prompt: str = "A cat sitting on a table, smooth motion") -> tuple: """Generate one I2V video, return (frames, peak_memory_mb).""" gc.collect() current_omni_platform.empty_cache() @@ -101,7 +101,7 @@ def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t monitor.start() image = _create_test_image() - response = omni_runner_handler.send_diffusion_request( + response = offline_client.send_diffusion_request( { "prompt": prompt, "images": image, @@ -112,7 +112,7 @@ def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t peak = monitor.peak_used_mb monitor.stop() - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.images is not None and len(response.images) > 0, "Expected image output" frames = response.images[0] @@ -122,7 +122,7 @@ def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t return frames, peak -def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table") -> tuple: +def _generate_t2v_video(offline_client, prompt: str = "A cat sitting on a table") -> tuple: """Generate one T2V video, return (frames, peak_memory_mb).""" gc.collect() current_omni_platform.empty_cache() @@ -131,7 +131,7 @@ def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() - response = omni_runner_handler.send_diffusion_request( + response = offline_client.send_diffusion_request( { "prompt": prompt, "sampling_params": _sampling_params_t2v(), @@ -141,7 +141,7 @@ def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t peak = monitor.peak_used_mb monitor.stop() - assert response.success, f"Request failed: {response.error_message}" + assert response.success, "Request failed" assert response.images is not None and len(response.images) > 0, "Expected image output" frames = response.images[0] @@ -158,9 +158,9 @@ def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True) -def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler): +def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, offline_client): """Load the W4A16 quantized Wan2.2 I2V model and verify it produces a valid video.""" - frames, _ = _generate_i2v_video(omni_runner_handler) + frames, _ = _generate_i2v_video(offline_client) assert frames is not None, "Expected video frames output" assert hasattr(frames, "shape"), "Expected frames to have a shape attribute" @@ -182,9 +182,9 @@ def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, omni_runner_hand @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True) -def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler): +def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, offline_client): """Load the W4A16 quantized Wan2.2 T2V model and verify it produces a valid video.""" - frames, _ = _generate_t2v_video(omni_runner_handler) + frames, _ = _generate_t2v_video(offline_client) assert frames is not None, "Expected video frames output" assert hasattr(frames, "shape"), "Expected frames to have a shape attribute" @@ -204,9 +204,9 @@ def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, omni_runner_hand @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True) -def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler): +def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, offline_client): """Measure peak GPU memory of W4A16 quantized I2V model.""" - frames, peak = _generate_i2v_video(omni_runner_handler) + frames, peak = _generate_i2v_video(offline_client) assert frames is not None, "Expected video frames output" _memory_results["quant_i2v"] = peak @@ -220,9 +220,9 @@ def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler): @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", baseline_i2v_params, indirect=True) -def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler): +def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, offline_client): """Measure peak GPU memory of BF16 baseline I2V model.""" - frames, peak = _generate_i2v_video(omni_runner_handler) + frames, peak = _generate_i2v_video(offline_client) assert frames is not None, "Expected video frames output" _memory_results["baseline_i2v"] = peak @@ -261,9 +261,9 @@ def test_wan22_i2v_autoround_w4a16_memory_savings(): @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True) -def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler): +def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, offline_client): """Measure peak GPU memory of W4A16 quantized T2V model.""" - frames, peak = _generate_t2v_video(omni_runner_handler) + frames, peak = _generate_t2v_video(offline_client) assert frames is not None, "Expected video frames output" _memory_results["quant_t2v"] = peak @@ -277,9 +277,9 @@ def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler): @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_runner", baseline_t2v_params, indirect=True) -def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler): +def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, offline_client): """Measure peak GPU memory of BF16 baseline T2V model.""" - frames, peak = _generate_t2v_video(omni_runner_handler) + frames, peak = _generate_t2v_video(offline_client) assert frames is not None, "Expected video frames output" _memory_results["baseline_t2v"] = peak diff --git a/tests/e2e/offline_inference/test_wan22_t2v.py b/tests/e2e/offline_inference/test_wan22_t2v.py index f3d03f399..805fbd1fe 100644 --- a/tests/e2e/offline_inference/test_wan22_t2v.py +++ b/tests/e2e/offline_inference/test_wan22_t2v.py @@ -1,7 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunnerHandler +from tests.helpers.runtime import OfflineOmniClient from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" @@ -13,7 +16,7 @@ NEGATIVE_PROMPT = "low quality, blurry, watermark, text" @pytest.mark.diffusion @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1}) @pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) -def test_text_to_video_001(omni_runner_handler: OmniRunnerHandler): +def test_text_to_video_001(offline_client: OfflineOmniClient): sampling = OmniDiffusionSamplingParams( height=512, width=512, @@ -29,4 +32,4 @@ def test_text_to_video_001(omni_runner_handler: OmniRunnerHandler): "negative_prompt": NEGATIVE_PROMPT, "sampling_params": sampling, } - omni_runner_handler.send_diffusion_request(request_config) + offline_client.send_diffusion_request(request_config) diff --git a/tests/e2e/offline_inference/test_z_image.py b/tests/e2e/offline_inference/test_z_image.py index 1770ccd85..ede9ea1ba 100644 --- a/tests/e2e/offline_inference/test_z_image.py +++ b/tests/e2e/offline_inference/test_z_image.py @@ -1,21 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest from tests.helpers.mark import hardware_test -from tests.helpers.runtime import OmniRunnerHandler +from tests.helpers.runtime import OfflineOmniClient from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "Tongyi-MAI/Z-Image-Turbo" PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard." -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 1, "rocm": 1, "xpu": 2}) @pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) -def test_zimage(omni_runner_handler: OmniRunnerHandler): +def test_zimage(offline_client: OfflineOmniClient): # high resolution may cause OOM on L4 sampling = OmniDiffusionSamplingParams( height=256, @@ -30,4 +30,4 @@ def test_zimage(omni_runner_handler: OmniRunnerHandler): "prompt": PROMPT, "sampling_params": sampling, } - omni_runner_handler.send_diffusion_request(request_config) + offline_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py index cb77c40e7..01afe689a 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py +++ b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Shared fixtures and argument builders for MiniCPM-o 4.5 duplex E2E tests.""" from __future__ import annotations @@ -11,21 +14,12 @@ from types import SimpleNamespace import pytest from huggingface_hub import snapshot_download -from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from tests.helpers.runtime import OmniServerParams, get_model_prefix +from tests.helpers.stage_config import get_deploy_config_path, get_deploy_duplex_max_sessions MODEL = "openbmb/MiniCPM-o-4_5" -DEPLOY_CONFIG = modify_stage_config( - get_deploy_config_path("minicpmo_4_5_duplex.yaml"), - updates={ - "base_config": get_deploy_config_path("minicpmo_4_5.yaml"), - # Talker context is 4096 (tts_config.max_position_embeddings); KV sizing - # is left automatic so duplex matches the simplex deploy profiles. - "stages": { - 1: {"max_model_len": 4096}, - }, - }, -) +DEPLOY_CONFIG_REL = "minicpmo_4_5.yaml" +DEPLOY_CONFIG = get_deploy_config_path(DEPLOY_CONFIG_REL) ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5" RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav" RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941" @@ -46,6 +40,16 @@ SERVER_PARAMS = [ ] +def deploy_max_sessions() -> int: + """Concurrent duplex sessions the deploy config under test admits. + + The admission probe has to expect the capacity the server is actually + started with. Hardcoding it silently drifts the moment the deploy config + changes its capacity, turning that change into an unrelated probe timeout. + """ + return get_deploy_duplex_max_sessions(DEPLOY_CONFIG_REL) + + def validated_wav(path: Path, expected_sha256: str) -> Path: actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest() if actual_sha256 != expected_sha256: @@ -73,10 +77,11 @@ def validated_soft_interrupt_wav() -> Path: return validated_wav(SOFT_INTERRUPT_WAV, SOFT_INTERRUPT_SHA256) -def resolve_ref_audio(model_prefix: str) -> Path: +def resolve_ref_audio() -> Path: # ``MODEL`` may be a local checkpoint directory instead of a Hub repo id # (``Path("/prefix") / "/abs/path"`` collapses to the absolute path), so only # fall back to the Hub cache when no such directory exists on disk. + model_prefix = get_model_prefix() model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL) if not model_root.is_dir(): model_root = Path(snapshot_download(MODEL, local_files_only=True)) @@ -86,6 +91,20 @@ def resolve_ref_audio(model_prefix: str) -> Path: return ref_audio +def duplex_camera_frames(*, seconds: int, cache_dir: Path) -> list[str]: + """Base64 JPEG frames of a moving synthetic scene, one per second. + + Mirrors what the realtime web client captures (1 fps, no client-side + resize: the server normalizes at ``scale_resolution=448``). + """ + from tests.e2e.online_serving.helpers.minicpmo_realtime_duplex_scenarios import _video_frames_from_file + from tests.helpers.media import generate_synthetic_video + + video = generate_synthetic_video(448, 448, 30 * seconds, cache_dir=cache_dir) + frames, _ = _video_frames_from_file(Path(video["file_path"])) + return frames + + def realtime_url(omni_server) -> str: return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1" @@ -108,6 +127,8 @@ def demo_args( output_audio_format="pcm16", chunk_ms=200, realtime_input=True, + input_video=None, + stack_frames=1, first_turn_ms=1400, turn_duration_ms=[], first_turn_transcript="duplex CI speech", diff --git a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py index 17e0cb2eb..7776f2a0f 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py +++ b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Strict MiniCPM-o 4.5 Realtime duplex E2E scenario harness. This script is intentionally scenario-based instead of a generic chat client. @@ -34,8 +37,10 @@ import json import sys import time import wave +from collections.abc import Sequence from dataclasses import dataclass, field from pathlib import Path +from types import SimpleNamespace _REPO_ROOT = Path(__file__).resolve().parents[4] if str(_REPO_ROOT) not in sys.path: @@ -52,19 +57,100 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 PCM16_SAMPLE_RATE, RealtimeEventCollector, build_realtime_url, + duplex_unit_boundary_ms, read_pcm16_wav, summarize_session_request_metrics, ) +from vllm_omni.experimental.fullduplex.client import ( # noqa: E402 + reference_audio_data_url as _ref_audio_data_url, +) +from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402 + concat_frames_b64, + unit_subframe_offsets, +) _url_with_model = build_realtime_url _read_wav_pcm16 = read_pcm16_wav +DemoArgs = argparse.Namespace | SimpleNamespace -def _ref_audio_data_url(path: str | None) -> str | None: - if path is None: - return None - ref_path = Path(path).expanduser() - return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii") +def _video_frames_from_file( + path: Path, + *, + fps: float = 1.0, + jpeg_quality: int = 70, + stack_frames: int = 1, +) -> tuple[list[str], list[str | None]]: + """Decode ``path`` into base64 JPEG frames at the omni duplex cadence. + + One frame per second matches what the realtime web client captures, so the + stream stays representative of a live camera. ``stack_frames > 1`` also + samples each unit's interior sub-frames and tiles them into one composite, + the official way to raise visual refresh rate without changing the audio + cadence. Returns ``(base_frames, stacked_frames)`` with the tracks parallel. + """ + import cv2 + + offsets = unit_subframe_offsets(stack_frames) + capture = cv2.VideoCapture(str(path)) + if not capture.isOpened(): + raise ValueError(f"cannot open video: {path}") + + def encode(frame_bgr, index: int) -> str: + ok, jpeg = cv2.imencode(".jpg", frame_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality]) + if not ok: + raise ValueError(f"cannot JPEG-encode frame {index} of {path}") + return base64.b64encode(jpeg.tobytes()).decode("ascii") + + try: + source_fps = capture.get(cv2.CAP_PROP_FPS) or 0.0 + stride = max(1, round(source_fps / fps)) if source_fps > 0 else 1 + # Sub-frame k of a unit sits ``offset * stride`` frames past its base. + sub_strides = {max(1, round(offset * stride)): position for position, offset in enumerate(offsets)} + frames: list[str] = [] + interiors: list[list[str]] = [] + index = 0 + while True: + ok, frame_bgr = capture.read() + if not ok: + break + position, within_unit = divmod(index, stride) + if within_unit == 0: + frames.append(encode(frame_bgr, index)) + interiors.append([]) + elif within_unit in sub_strides and position < len(interiors): + interiors[position].append(encode(frame_bgr, index)) + index += 1 + finally: + capture.release() + if not frames: + raise ValueError(f"video has no decodable frames: {path}") + stacked = [concat_frames_b64(interior) if interior else None for interior in interiors] + return frames, stacked + + +def _resolve_video_frames(args: DemoArgs) -> tuple[list[str], list[str | None]]: + """Base64 JPEG frames for the duplex camera track, newest last. + + Returns the base track plus the parallel stacked composites. Only a video + source has sub-frames to stack; an explicit frame list or a still image + stacks nothing. + """ + frames = list(getattr(args, "video_frames_b64", None) or []) + if frames: + stacked = list(getattr(args, "video_stacked_frames_b64", None) or []) + stacked += [None] * (len(frames) - len(stacked)) + return frames, stacked[: len(frames)] + input_video = getattr(args, "input_video", None) + if input_video: + return _video_frames_from_file( + Path(input_video).expanduser(), + stack_frames=max(1, int(getattr(args, "stack_frames", 1) or 1)), + ) + frame_image = getattr(args, "frame_image", None) + if frame_image: + return [base64.b64encode(Path(frame_image).expanduser().read_bytes()).decode("ascii")], [None] + return [], [] @dataclass @@ -374,17 +460,20 @@ class DemoState: commit_indices = [ index for index, event in enumerate(self.events) if event.get("type") == "input_audio_buffer.committed" ] - decision_indices = [ - index - for index, event in enumerate(self.events) - if event.get("type") == "response.created" - or ( - event.get("type") == "response.listen" - and isinstance(event.get("response"), dict) - and isinstance(event["response"].get("metadata"), dict) - and event["response"]["metadata"].get("model_listen") is True - ) - ] + decision_indices: list[int] = [] + for index, event in enumerate(self.events): + event_type = event.get("type") + if event_type == "response.created": + decision_indices.append(index) + continue + if event_type != "response.listen": + continue + response = event.get("response") + if not isinstance(response, dict): + continue + metadata = response.get("metadata") + if isinstance(metadata, dict) and metadata.get("model_listen") is True: + decision_indices.append(index) first_input_index = self.first_index("input_audio_buffer.speech_started") if first_input_index is None: first_input_index = commit_indices[0] if commit_indices else None @@ -501,7 +590,7 @@ class DemoState: return stale -def _session_update_event(args: argparse.Namespace) -> dict[str, object]: +def _session_update_event(args: DemoArgs) -> dict[str, object]: session_payload: dict[str, object] = { "model": args.model, "modalities": ["audio", "text"], @@ -837,7 +926,8 @@ async def _send_pcm16( hints: dict[str, object] | None = None, first_chunk_hints: dict[str, object] | None = None, on_model_unit_ready=None, - frame_b64: str | None = None, + frames_b64: Sequence[str] | None = None, + stacked_frames_b64: Sequence[str | None] | None = None, ) -> None: hints = hints or {} first_chunk_hints = first_chunk_hints or {} @@ -845,6 +935,8 @@ async def _send_pcm16( model_unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE audio_ms = 0 next_model_unit_bytes = model_unit_bytes + frames = list(frames_b64 or []) + stacked = list(stacked_frames_b64 or []) frames_sent = 0 for offset in range(0, len(pcm16), chunk_bytes): chunk = pcm16[offset : offset + chunk_bytes] @@ -853,9 +945,18 @@ async def _send_pcm16( chunk_hints = dict(hints) if offset == 0: chunk_hints.update(first_chunk_hints) - if frame_b64 is not None and audio_ms > frames_sent * 1000: - # Omni duplex cadence: one camera frame per 1 s of audio. - chunk_hints["video_frames"] = [frame_b64] + if frames and audio_ms >= duplex_unit_boundary_ms(frames_sent): + # Omni duplex cadence: one camera frame per model unit, riding the + # append that closes it, so the frame and the second of audio it was + # captured during enter the same unit. A video advances one frame per + # unit and holds its last frame when the audio outlives the clip; a + # still image is a one-element list and therefore repeats. + # A stacked composite of that unit's interior sub-frames rides the + # same append right after the base frame (wire maximum of 2). + index = min(frames_sent, len(frames) - 1) + composite = stacked[index] if index < len(stacked) else None + base = frames[index] + chunk_hints["video_frames"] = [base] if composite is None else [base, composite] frames_sent += 1 await ws.send( json.dumps( @@ -936,7 +1037,8 @@ async def _send_clean_turn( realtime_input: bool = False, model_policy_settle_s: float = 2.0, commit_input: bool = True, - frame_b64: str | None = None, + frames_b64: Sequence[str] | None = None, + stacked_frames_b64: Sequence[str | None] | None = None, ) -> tuple[str | None, str]: before_created = state.count("response.created") before_model_listen = state.model_listen_count @@ -958,7 +1060,8 @@ async def _send_clean_turn( chunk_ms=chunk_ms, realtime_delay=realtime_input, hints={"transcript": transcript} if send_transcript_hint else {}, - frame_b64=frame_b64, + frames_b64=frames_b64, + stacked_frames_b64=stacked_frames_b64, ) if commit_input: state.input_commit_sent_at_s.append(time.monotonic()) @@ -1164,7 +1267,8 @@ async def _send_listen_only_overlap_pair( model_policy_settle_s: float, validation_mode: str, silence_ms: int, - frame_b64: str | None = None, + frames_b64: Sequence[str] | None = None, + stacked_frames_b64: Sequence[str | None] | None = None, ) -> tuple[list[str | None], list[str], bool]: if not realtime_input: raise ValueError("listen-only-overlap requires --realtime-input") @@ -1195,7 +1299,8 @@ async def _send_listen_only_overlap_pair( realtime_delay=True, hints={"transcript": transcripts[0]} if send_transcript_hint else {}, on_model_unit_ready=continue_until_first_speak, - frame_b64=frame_b64, + frames_b64=frames_b64, + stacked_frames_b64=stacked_frames_b64, ) state.input_commit_sent_at_s.append(time.monotonic()) await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) @@ -1234,7 +1339,8 @@ async def _send_listen_only_overlap_pair( realtime_delay=True, hints={"transcript": transcripts[1]} if send_transcript_hint else {}, on_model_unit_ready=record_model_unit_ready, - frame_b64=frame_b64, + frames_b64=frames_b64, + stacked_frames_b64=stacked_frames_b64, ) state.input_commit_sent_at_s.append(time.monotonic()) await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) @@ -1298,10 +1404,8 @@ async def _send_listen_only_overlap_pair( return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok -async def run_demo(args: argparse.Namespace) -> dict[str, object]: - demo_frame_b64: str | None = None - if getattr(args, "frame_image", None): - demo_frame_b64 = base64.b64encode(Path(args.frame_image).read_bytes()).decode("ascii") +async def run_demo(args: DemoArgs) -> dict[str, object]: + demo_frames_b64, demo_stacked_b64 = _resolve_video_frames(args) turn_input_paths = _turn_input_paths( Path(args.input_wav), list(getattr(args, "turn_input_wav", []) or []), @@ -1351,6 +1455,9 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: try: await ws.send(json.dumps(_session_update_event(args))) await _wait_for(state, lambda: state.count("session.created") > 0, timeout_s=20, label="session.created") + start_barrier = getattr(args, "start_barrier", None) + if start_barrier is not None: + await start_barrier.wait() turn_transcripts = _turn_transcripts(args.first_turn_transcript, turns=args.turns) turn_specs = list(zip(turn_transcripts, turn_durations, strict=True)) @@ -1369,7 +1476,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000), validation_mode=validation_mode, silence_ms=max(0, args.silence_ms), - frame_b64=demo_frame_b64, + frames_b64=demo_frames_b64, + stacked_frames_b64=demo_stacked_b64, ) for turn_index in range(2, len(turn_specs)): transcript, duration_ms = turn_specs[turn_index] @@ -1387,7 +1495,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: realtime_input=realtime_input, model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000), commit_input=not continuous_input, - frame_b64=demo_frame_b64, + frames_b64=demo_frames_b64, + stacked_frames_b64=demo_stacked_b64, ) turn_response_ids.append(response_id) turn_outcomes.append(outcome) @@ -1407,7 +1516,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: realtime_input=realtime_input, model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000), commit_input=not continuous_input, - frame_b64=demo_frame_b64, + frames_b64=demo_frames_b64, + stacked_frames_b64=demo_stacked_b64, ) turn_response_ids.append(response_id) turn_outcomes.append(outcome) @@ -1582,6 +1692,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: "continuous_input": continuous_input, "continuous_input_ok": continuous_input_ok, "input_chunk_ms": args.chunk_ms, + "video_frame_count": len(demo_frames_b64), + "video_stacked_frame_count": sum(frame is not None for frame in demo_stacked_b64), "model_policy_settle_ms": args.model_policy_settle_ms, "turn_outcomes": turn_outcomes, "turn_outcomes_ok": turn_outcomes_ok, @@ -1639,6 +1751,19 @@ def parse_args() -> argparse.Namespace: default=None, help="Optional image file sent as an omni-duplex camera frame (one per 1 s of audio)", ) + parser.add_argument( + "--input-video", + default=None, + help="Optional video file streamed as omni-duplex camera frames (decoded at 1 fps, " + "overrides --frame-image; the last frame is held if the audio outlives the clip)", + ) + parser.add_argument( + "--stack-frames", + type=int, + default=1, + choices=[1, 2], + help="Frames per append: 2 also resends the previous frame for motion context.", + ) parser.add_argument( "--turn-input-wav", action="append", diff --git a/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py new file mode 100644 index 000000000..b623446ad --- /dev/null +++ b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py @@ -0,0 +1,364 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import argparse +import asyncio +import base64 +import json +import math +import uuid +import wave +from collections.abc import Sequence +from contextlib import asynccontextmanager, suppress +from pathlib import Path +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import numpy as np +from scipy.signal import resample_poly + +from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, wait_for, write_pcm16_wav + +INPUT_SAMPLE_RATE_HZ = 16_000 +OUTPUT_SAMPLE_RATE_HZ = 22_050 +FRAME_SAMPLES = 1_280 +FRAME_PERIOD_S = FRAME_SAMPLES / INPUT_SAMPLE_RATE_HZ +DEFAULT_FUNCTION_TOOLS = [ + { + "type": "function", + "name": "generate_random_number", + "description": "Generate a random integer between min and max (inclusive).", + "parameters": { + "type": "object", + "properties": { + "min": {"type": "integer", "description": "Minimum value (inclusive)"}, + "max": {"type": "integer", "description": "Maximum value (inclusive)"}, + }, + "required": ["min", "max"], + }, + } +] +DEFAULT_INSTRUCTIONS = "You are NVIDIA Voice Chat. Answer briefly. Start by greeting the user." +DEFAULT_FUNCTION_INSTRUCTIONS = ( + "You are NVIDIA Voice Chat. If the user's request matches an available tool, " + "you MUST call that tool instead of answering from your own knowledge. " + "Use only argument values spoken by the user and never invent missing values." +) + + +def _url(base_url: str, model: str, session_id: str) -> str: + parts = urlsplit(base_url) + query = dict(parse_qsl(parts.query, keep_blank_values=True)) + query.update(duplex="1", model=model, autostart="0", session_id=session_id) + return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)) + + +def _read_wav(path: Path, *, input_channel: int = 0) -> np.ndarray: + with wave.open(str(path), "rb") as wav_file: + channels = wav_file.getnchannels() + width = wav_file.getsampwidth() + source_rate = wav_file.getframerate() + raw = wav_file.readframes(wav_file.getnframes()) + dtypes = {1: np.uint8, 2: np.dtype(" 1: + pcm = pcm.reshape(-1, channels)[:, input_channel] + if source_rate != INPUT_SAMPLE_RATE_HZ: + divisor = math.gcd(source_rate, INPUT_SAMPLE_RATE_HZ) + pcm = resample_poly(pcm, up=INPUT_SAMPLE_RATE_HZ // divisor, down=source_rate // divisor) + return np.ascontiguousarray(pcm, dtype=" int: + count = math.ceil(pcm.size / FRAME_SAMPLES) + if max_frames is not None: + count = min(count, max_frames) + for seq in range(count): + frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES] + frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size)).astype(" list[dict[str, object]]: + return [event for event in client.events.events if event.get("type") == event_type] + + +async def _return_function_output_when_ready( + client: RealtimeDuplexClient, + *, + output: str, + timeout_s: float, +) -> tuple[str, int]: + await wait_for( + lambda: bool(client.events.errors()) or client.events.count("response.function_call_arguments.done") > 0, + timeout_s=timeout_s, + label="function call to execute", + ) + if client.events.errors(): + raise AssertionError(f"function call failed before tool execution: {client.events.errors()}") + function_done = _events(client, "response.function_call_arguments.done")[-1] + call_id = function_done.get("call_id") + if not isinstance(call_id, str) or not call_id: + raise AssertionError(f"completed function call has no call_id: {function_done}") + event_count_before_output = len(client.events.events) + await client.send( + { + "type": "conversation.item.create", + "item": { + "type": "function_call_output", + "call_id": call_id, + "output": output, + }, + } + ) + return call_id, event_count_before_output + + +def _write_events(path: Path, client: RealtimeDuplexClient) -> None: + path.write_text( + "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events), + encoding="utf-8", + ) + + +@asynccontextmanager +async def _managed_client(client: RealtimeDuplexClient, *, timeout_s: float): + async with client: + try: + yield + finally: + if client.events.count("session.created") and not client.events.count("session.closed"): + with suppress(Exception): + await client.close_session(timeout_s=min(timeout_s, 30.0)) + + +async def run(args: argparse.Namespace) -> dict[str, object]: + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + if args.expect_function_call and args.instructions == DEFAULT_INSTRUCTIONS: + instructions = DEFAULT_FUNCTION_INSTRUCTIONS + else: + instructions = args.instructions + tools = DEFAULT_FUNCTION_TOOLS if args.expect_function_call else None + session_id = f"nemotron-voicechat-{uuid.uuid4().hex}" + client = RealtimeDuplexClient(_url(args.url, args.model, session_id)) + async with _managed_client(client, timeout_s=args.timeout_s): + session_payload: dict[str, object] = { + "session_id": session_id, + "model": args.model, + "modalities": ["audio", "text"], + "input_audio_format": "pcm_f32le", + "output_audio_format": "pcm16", + "instructions": instructions, + "idle_timeout_s": args.timeout_s, + "turn_detection": None, + "extra_body": {"auto_response": True}, + } + if tools is not None: + session_payload["tools"] = tools + await client.send({"type": "session.update", "session": session_payload}) + await wait_for( + lambda: client.events.count("session.created") > 0 or bool(client.events.errors()), + timeout_s=args.timeout_s, + label="session.created", + ) + if client.events.errors(): + raise AssertionError(f"session setup failed: {client.events.errors()}") + created = _events(client, "session.created")[-1] + session = created.get("session") + capabilities = session.get("capabilities") if isinstance(session, dict) else None + expected = dict( + implementation_level="model_native_duplex", + chunk_period_ms=80, + supports_core_resumable_request=True, + supports_core_kv_lease=False, + supports_multi_session=False, + ) + if not isinstance(capabilities, dict) or any(capabilities.get(key) != value for key, value in expected.items()): + raise AssertionError(f"unexpected capabilities: {capabilities}") + + function_output_task = ( + asyncio.create_task( + _return_function_output_when_ready( + client, + output=args.function_output, + timeout_s=args.timeout_s, + ) + ) + if args.function_output is not None + else None + ) + pcm = _read_wav(Path(args.input_wav), input_channel=args.input_channel) + frame_count = await _stream(client, pcm, max_frames=args.max_frames, realtime=not args.no_realtime) + completed_responses_at_commit = client.events.count("response.done") + await client.send({"type": "input_audio_buffer.commit", "final": True}) + await wait_for( + lambda: ( + bool(client.events.errors()) + or ( + client.events.count("response.function_call_arguments.done") > 0 + if args.expect_function_call + else ( + client.events.count("response.audio.delta") >= args.minimum_audio_chunks + and (client.events.count("response.done") > completed_responses_at_commit) + ) + ) + ), + timeout_s=args.timeout_s, + label="model output", + ) + await asyncio.sleep(args.drain_s) + if client.events.errors(): + raise AssertionError(f"Realtime session emitted errors: {client.events.errors()}") + done_events = _events(client, "response.done") + if not args.expect_function_call: + response = done_events[-1].get("response") if done_events else None + status = response.get("status") if isinstance(response, dict) else None + if status != "completed": + raise AssertionError(f"response did not complete successfully: {done_events[-1:]}") + + function_events = [ + event for event in client.events.events if str(event.get("type", "")).startswith("response.function_call") + ] + function_items = [ + event + for event in _events(client, "response.output_item.done") + if isinstance(event.get("item"), dict) and event["item"].get("type") == "function_call" + ] + if args.expect_function_call and not any( + event.get("type") == "response.function_call_arguments.done" for event in function_events + ): + raise AssertionError(f"no completed function call: {function_events}") + if args.expect_function_call: + matching_items = [ + event for event in function_items if event["item"].get("name") == args.expected_function_name + ] + if not matching_items: + raise AssertionError(f"expected {args.expected_function_name!r}, got {function_items}") + function_item = matching_items[-1]["item"] + try: + function_arguments = json.loads(str(function_item.get("arguments", ""))) + except json.JSONDecodeError as exc: + raise AssertionError(f"function arguments are not JSON: {function_item}") from exc + if args.expected_function_arguments is not None: + expected_arguments = json.loads(args.expected_function_arguments) + if function_arguments != expected_arguments: + raise AssertionError( + f"function arguments differ: expected={expected_arguments!r}, actual={function_arguments!r}" + ) + + if args.function_output is not None: + call_id = function_item.get("call_id") + if not isinstance(call_id, str) or not call_id: + raise AssertionError(f"function item has no call_id: {function_item}") + assert function_output_task is not None + returned_call_id, event_count_before_output = await function_output_task + assert returned_call_id == call_id + + def tool_result_completed() -> bool: + later = client.events.events[event_count_before_output:] + transcript = "".join( + str(event.get("delta", "")) + for event in later + if event.get("type") == "response.audio_transcript.delta" + ).lower() + return ( + any(event.get("type") == "response.audio.delta" for event in later) + and any(event.get("type") == "response.done" for event in later) + and (args.expected_post_tool_text is None or args.expected_post_tool_text.lower() in transcript) + ) + + await wait_for( + lambda: bool(client.events.errors()) or tool_result_completed(), + timeout_s=args.timeout_s, + label="completed response after function output", + ) + if client.events.errors(): + raise AssertionError(f"function output failed: {client.events.errors()}") + await asyncio.sleep(args.drain_s) + + audio = client.events.audio_bytes() + audio_events = _events(client, "response.audio.delta") + rates = {event.get("sample_rate_hz") for event in audio_events} + if not args.expect_function_call and audio and rates != {OUTPUT_SAMPLE_RATE_HZ}: + raise AssertionError(f"unexpected output sample rates: {rates}") + if not args.expect_function_call and args.minimum_audio_chunks and not audio: + raise AssertionError("model produced no audio") + expected_bytes = 2 * OUTPUT_SAMPLE_RATE_HZ * expected["chunk_period_ms"] // 1000 + packet_sizes = [len(base64.b64decode(str(event.get("delta", "")), validate=True)) for event in audio_events] + if not args.expect_function_call and any(size != expected_bytes for size in packet_sizes): + raise AssertionError(f"audio deltas are not fixed 80 ms PCM16 packets: {packet_sizes}") + audio_pcm = np.frombuffer(audio, dtype=" argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--url", default="ws://127.0.0.1:8125/v1/realtime") + parser.add_argument("--model", required=True) + parser.add_argument("--input-wav", required=True) + parser.add_argument("--input-channel", type=int, default=0) + parser.add_argument("--output-dir", default="/tmp/nemotron-voicechat-duplex") + parser.add_argument("--instructions", default=DEFAULT_INSTRUCTIONS) + parser.add_argument("--max-frames", type=int) + parser.add_argument("--minimum-audio-chunks", type=int, default=1) + parser.add_argument("--minimum-audio-rms", type=float, default=1e-4) + parser.add_argument("--expect-function-call", action="store_true") + parser.add_argument("--expected-function-name", default="generate_random_number") + parser.add_argument("--expected-function-arguments") + parser.add_argument("--function-output") + parser.add_argument("--expected-post-tool-text") + parser.add_argument("--no-realtime", action="store_true") + parser.add_argument("--drain-s", type=float, default=2.0) + parser.add_argument("--timeout-s", type=float, default=600.0) + return parser.parse_args(argv) + + +def main() -> None: + print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py index 331f0d20f..910a5a180 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Concurrent MiniCPM-o Realtime duplex and resumable-session E2E driver.""" from __future__ import annotations @@ -14,6 +17,7 @@ from types import SimpleNamespace from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import websockets +from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed SCRIPT_DIR = Path(__file__).resolve().parent @@ -42,6 +46,67 @@ async def run_demo(args): return await _scenario_module().run_demo(args) +class _SynchronizedStartGate: + """One-shot, failure-propagating start gate for Python 3.10+.""" + + def __init__(self, parties: int, *, timeout_s: float): + if parties < 1: + raise ValueError("synchronized start parties must be positive") + if timeout_s <= 0: + raise ValueError("synchronized start timeout must be positive") + self._parties = parties + self._timeout_s = timeout_s + self._arrived = 0 + self._released = False + self._failure: BaseException | None = None + self._event = asyncio.Event() + self._lock = asyncio.Lock() + + async def wait(self) -> None: + async with self._lock: + if self._failure is not None: + raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure + if self._released: + return + self._arrived += 1 + if self._arrived == self._parties: + self._released = True + self._event.set() + return + + try: + await asyncio.wait_for(self._event.wait(), timeout=self._timeout_s) + except asyncio.TimeoutError as exc: + failure = TimeoutError( + f"synchronized start timed out after {self._timeout_s:.1f}s " + f"({self._arrived}/{self._parties} sessions ready)" + ) + await self.abort(failure) + if self._released and self._failure is None: + return + raise failure from exc + + if self._failure is not None: + raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure + + async def abort(self, failure: BaseException) -> None: + async with self._lock: + if self._released or self._failure is not None: + return + self._failure = failure + self._event.set() + + +async def _run_demo_with_start_gate(args: SimpleNamespace): + start_gate = getattr(args, "start_barrier", None) + try: + return await run_demo(args) + except BaseException as exc: + if start_gate is not None: + await start_gate.abort(exc) + raise + + def _with_resume_mode(url: str) -> str: parts = urlsplit(url) query = dict(parse_qsl(parts.query, keep_blank_values=True)) @@ -101,14 +166,24 @@ def _validate_semantic_isolation( return True if len(expected_tokens) != len(results): return False - for result, expected_token in zip(results, expected_tokens, strict=True): + normalized_expected_tokens = [token.strip().casefold() for token in expected_tokens] + if any(not token for token in normalized_expected_tokens): + return False + for result_index, (result, expected_token) in enumerate(zip(results, normalized_expected_tokens, strict=True)): details = result.get("transcript_integrity") transcripts = ( [str(item.get("transcript", "")) for item in details if isinstance(item, dict)] if isinstance(details, list) else [] ) - if expected_token not in "".join(transcripts): + joined_transcripts = "".join(transcripts).casefold() + if expected_token not in joined_transcripts: + return False + if any( + other_token in joined_transcripts + for other_index, other_token in enumerate(normalized_expected_tokens) + if other_index != result_index + ): return False return True @@ -130,10 +205,14 @@ async def _receive_until(ws, event_type: str, *, timeout_s: float) -> tuple[dict return await asyncio.wait_for(receive(), timeout=timeout_s) +def _server_event_sequences(events: list[dict[str, object]]) -> list[int]: + return [sequence for event in events if isinstance(sequence := event.get("server_event_seq"), int)] + + async def _open_admission_session( args: argparse.Namespace, session_id: str, -) -> tuple[object, dict[str, object]]: +) -> tuple[ClientConnection, dict[str, object]]: url = _url_with_model( args.url, args.model, @@ -159,7 +238,7 @@ async def _open_admission_session( return ws, created -async def _close_admission_session(ws, *, timeout_s: float) -> None: +async def _close_admission_session(ws: ClientConnection, *, timeout_s: float) -> None: await ws.send(json.dumps({"type": "session.close"})) await _receive_until(ws, "session.closed", timeout_s=timeout_s) await ws.close() @@ -169,8 +248,8 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str, if limit < 1: raise ValueError("admission limit must be positive") prefix = f"admission-{uuid.uuid4().hex}" - accepted: list[tuple[object, dict[str, object]]] = [] - replacement: tuple[object, dict[str, object]] | None = None + accepted: list[tuple[ClientConnection, dict[str, object]]] = [] + replacement: tuple[ClientConnection, dict[str, object]] | None = None overflow_code = None try: for index in range(limit): @@ -274,14 +353,7 @@ async def _resume_probe( generation = created.get("attachment_generation") if not isinstance(token, str) or not isinstance(incarnation, int): raise RuntimeError("session.created omitted resumable credentials") - last_seq = max( - ( - event.get("server_event_seq", 0) - for event in first_events - if isinstance(event.get("server_event_seq"), int) - ), - default=0, - ) + last_seq = max(_server_event_sequences(first_events), default=0) delay_s = args.expire_after_s if expect_expired else args.resume_after_ms / 1000 if delay_s > 0: @@ -324,7 +396,7 @@ async def _resume_probe( await second.send(json.dumps({"type": "session.close"})) closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s) - replay_sequences = [event["server_event_seq"] for event in replay if isinstance(event.get("server_event_seq"), int)] + replay_sequences = _server_event_sequences(replay) return { "ok": ( resumed.get("session_id") == session_id @@ -386,14 +458,7 @@ async def _takeover_probe( generation = created.get("attachment_generation") if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int): raise RuntimeError("session.created omitted takeover credentials") - last_seq = max( - ( - event.get("server_event_seq", 0) - for event in first_events - if isinstance(event.get("server_event_seq"), int) - ), - default=0, - ) + last_seq = max(_server_event_sequences(first_events), default=0) second = await websockets.connect(_with_resume_mode(url), max_size=64 * 1024 * 1024) await second.send( @@ -457,7 +522,9 @@ async def _takeover_probe( await first.close() -def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace: +def _demo_args( + args: argparse.Namespace, index: int, start_barrier: _SynchronizedStartGate | None = None +) -> SimpleNamespace: validation_mode = "response-required" if args.response_required else "model-policy" input_wav = args.session_input_wav[index] if args.session_input_wav else args.input_wav return SimpleNamespace( @@ -477,14 +544,16 @@ def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace: omit_transcript_hints=True, validation_mode=validation_mode, temperature=args.temperature, - scenario="sequential", + scenario=getattr(args, "scenario", "sequential"), + silence_ms=getattr(args, "silence_ms", 500), require_audio=args.response_required, require_distinct_inputs=False, - expect_empty_turn=[], + expect_empty_turn=list(getattr(args, "expect_empty_turn", []) or []), short_ack_ms=350, turns=args.turns, timeout_s=args.timeout_s, model_policy_settle_ms=args.model_policy_settle_ms, + start_barrier=start_barrier, ) @@ -511,8 +580,16 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]: ) lifecycle_result = await run_lifecycle_probes(args) + start_barrier = ( + _SynchronizedStartGate( + args.sessions, + timeout_s=min(float(args.timeout_s), 30.0), + ) + if getattr(args, "synchronized_start", False) + else None + ) session_results = await asyncio.gather( - *(run_demo(_demo_args(args, index)) for index in range(args.sessions)), + *(_run_demo_with_start_gate(_demo_args(args, index, start_barrier)) for index in range(args.sessions)), return_exceptions=True, ) failures = [repr(result) for result in session_results if isinstance(result, BaseException)] @@ -596,7 +673,14 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--turns", type=int, default=1) parser.add_argument("--first-turn-ms", type=int, default=1400) parser.add_argument("--turn-duration-ms", type=int, action="append", default=[]) + parser.add_argument("--expect-empty-turn", type=int, action="append", default=[]) parser.add_argument("--response-required", action="store_true") + parser.add_argument( + "--scenario", + choices=["sequential", "listen-only-overlap"], + default="sequential", + ) + parser.add_argument("--silence-ms", type=int, default=500) parser.add_argument("--temperature", type=float, default=None) parser.add_argument("--disconnect-session-index", type=int) parser.add_argument("--takeover-session-index", type=int) @@ -604,6 +688,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--expire-session-index", type=int) parser.add_argument("--expire-after-s", type=float, default=6.0) parser.add_argument("--verify-admission-limit", type=int) + parser.add_argument( + "--synchronized-start", + action="store_true", + help="Wait until every session is created before starting input.", + ) parser.add_argument("--model-policy-settle-ms", type=int, default=600) parser.add_argument("--timeout-s", type=float, default=120.0) args = parser.parse_args() @@ -615,6 +704,13 @@ def parse_args() -> argparse.Namespace: parser.error("provide exactly one --session-expected-token per session") if args.session_expected_token and not args.session_input_wav: parser.error("--session-expected-token requires --session-input-wav") + normalized_expected_tokens = [token.strip().casefold() for token in args.session_expected_token] + for token_index, token in enumerate(normalized_expected_tokens): + if any( + token in other_token or other_token in token + for other_token in normalized_expected_tokens[token_index + 1 :] + ): + parser.error("--session-expected-token values must not overlap after normalization") return args diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py new file mode 100644 index 000000000..4047a0901 --- /dev/null +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py @@ -0,0 +1,85 @@ +import base64 +from pathlib import Path + +from vllm_omni.experimental.fullduplex.client import ( + PCM16_BYTES_PER_SAMPLE, + PCM16_SAMPLE_RATE, + RealtimeDuplexClient, + RealtimeEventCollector, + build_realtime_url, + read_pcm16_wav, + wait_for, +) + + +def _events(events: list[dict[str, object]], response_id: str, kind: str) -> list[dict[str, object]]: + return [e for e in events if e.get("type") == kind and RealtimeEventCollector.response_id(e) == response_id] + + +async def run_server_vad_interrupt(args) -> dict[str, object]: + initial = read_pcm16_wav(Path(args.input_wav)) + interrupt = read_pcm16_wav(Path(args.interrupt_wav))[: PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * 47 // 10] + session_id = f"server-vad-hard-interrupt-{id(args)}" + client = RealtimeDuplexClient(build_realtime_url(args.url, args.model, autostart=False, session_id=session_id)) + + async def until(predicate, label: str) -> None: + await wait_for(predicate, timeout_s=args.timeout_s, label=label) + + try: + await client.__aenter__() + await client.configure( + args.model, + ref_audio="data:audio/wav;base64," + base64.b64encode(Path(args.ref_audio).read_bytes()).decode(), + session_id=session_id, + turn_detection={"type": "server_vad", "interrupt_response": True}, + timeout_s=args.timeout_s, + ) + await client.stream_pcm16(initial + bytes(16_000 * 2 * 1600 // 1000), chunk_ms=200, realtime=True) + await client.commit() + await until(lambda: any(e.get("type") == "response.created" for e in client.events.events), "response.created") + target_id = RealtimeEventCollector.response_id( + next(e for e in client.events.events if e.get("type") == "response.created") + ) + assert target_id + await until(lambda: _events(client.events.events, target_id, "response.audio.delta"), "response.audio.delta") + cursor = len(client.events.events) + await client.stream_pcm16(interrupt + bytes(16_000 * 2 * 800 // 1000), chunk_ms=args.chunk_ms, realtime=True) + await until(lambda: _events(client.events.events, target_id, "response.done"), "cancelled response.done") + await client.commit() + await until( + lambda: any( + e.get("type") == "response.done" and RealtimeEventCollector.response_id(e) != target_id + for e in client.events.events[cursor:] + ), + "follow-up response.done", + ) + events = client.events.events + finally: + try: + await client.close_session(timeout_s=args.timeout_s) + finally: + await client.__aexit__(None, None, None) + + done = _events(events, target_id, "response.done") + terminal = done[0] if len(done) == 1 else None + trailing = events[events.index(terminal) + 1 :] if terminal is not None else events + stale = any( + e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) == target_id for e in trailing + ) + followup_audio = any( + e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) != target_id + for e in events[cursor:] + ) + ok = ( + terminal is not None + and isinstance(terminal.get("response"), dict) + and terminal["response"].get("status") == "cancelled" + and terminal["response"].get("status_details", {}).get("reason") == "turn_detected" + and not stale + and sum(e.get("type") == "input_audio_buffer.speech_started" for e in events[cursor:]) == 1 + and sum(e.get("type") == "input_audio_buffer.speech_stopped" for e in events[cursor:]) == 1 + and any(e.get("type") == "input_audio_buffer.committed" for e in events[cursor:]) + and followup_audio + and not any(e.get("type") == "error" for e in events) + ) + return {"ok": ok} diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py index d1b42cba4..3aff0cf8d 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py @@ -252,6 +252,19 @@ def summarize_artifacts( (summary["done_indices"][0] for summary in reversed(response_summaries) if summary.get("done_indices")), None, ) + # A soft interrupt's follow-up response can drain past the commit (residual + # model unit), so its response.done lands after input_audio_buffer.committed. + # Anchor the "listen after last done, before commit" sandwich on the last + # turn that closed its floor strictly before the commit, not the + # globally-last done (which may sit after the commit and leave the interval empty). + last_done_before_commit_index = next( + ( + summary["done_indices"][0] + for summary in reversed(response_summaries) + if summary.get("done_indices") and commit_index is not None and summary["done_indices"][0] < commit_index + ), + None, + ) listen_indices = [index for index, event in enumerate(events) if event.get("type") == "response.listen"] effective_min_responses = min_responses if validation_mode == "response-required" else 1 enough_responses = len(response_summaries) >= effective_min_responses @@ -279,9 +292,9 @@ def summarize_artifacts( ) listen_after_last_done = last_done_index is not None and any(index > last_done_index for index in listen_indices) listen_after_response_before_commit = ( - last_done_index is not None + last_done_before_commit_index is not None and commit_index is not None - and any(last_done_index < index < commit_index for index in listen_indices) + and any(last_done_before_commit_index < index < commit_index for index in listen_indices) ) final_listen_after_commit = commit_index is not None and any(index > commit_index for index in listen_indices) transcript = "".join(str(summary.get("transcript") or "") for summary in response_summaries) diff --git a/tests/e2e/online_serving/test_audiox_expansion.py b/tests/e2e/online_serving/test_audiox_expansion.py deleted file mode 100644 index 612f397e0..000000000 --- a/tests/e2e/online_serving/test_audiox_expansion.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E online serving test for AudioX text-to-audio diffusion. - -Mirrors `tests/e2e/online_serving/test_sd3_expansion.py` for image diffusion: spin up -`vllm-omni` with `--model-class-name AudioXPipeline` and validate via the standard -`send_diffusion_request` helper (now audio-aware). -""" - -import os - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data - -# Tiny / random checkpoint usable in CI; override to a real bundle locally. -AUDIOX_TEST_MODEL = os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random") -T2A_PROMPT = "A quiet living room with soft fabric rustle and gentle cat breathing." - -SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) - - -def _audiox_server_cases(model: str): - return [ - pytest.param( - OmniServerParams( - model=model, - server_args=["--model-class-name", "AudioXPipeline"], - ), - id="t2a", - marks=SINGLE_CARD_FEATURE_MARKS, - ), - ] - - -@pytest.mark.slow -@pytest.mark.diffusion -@pytest.mark.parametrize("omni_server", _audiox_server_cases(AUDIOX_TEST_MODEL), indirect=True) -def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - """AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`. - - AudioX is registered in ``vllm_omni/model_extras`` (``AudioXPipeline``), so the - ``audiox_task`` / ``seconds_*`` / ``sigma_*`` keys below flow through the standard - ``apply_declared_extra_args`` server path into ``sampling_params.extra_args`` rather - than the old top-level ``extra_args`` escape hatch. - """ - request_config = { - "model": omni_server.model, - "messages": dummy_messages_from_mix_data(content_text=T2A_PROMPT), - "extra_body": { - "num_inference_steps": 4, - "guidance_scale": 6.0, - "seed": 42, - "audiox_task": "t2a", - "seconds_start": 0.0, - "seconds_total": 2.0, - "sigma_min": 0.03, - "sigma_max": 1000.0, - }, - } - openai_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_bagel.py b/tests/e2e/online_serving/test_bagel.py index b22762acd..84c2302cc 100644 --- a/tests/e2e/online_serving/test_bagel.py +++ b/tests/e2e/online_serving/test_bagel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end online serving test for Bagel text2img and img2img generation. @@ -75,12 +75,11 @@ def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]: ] -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_bagel_text2img_online(omni_server, openai_client) -> None: +def test_bagel_text2img_online(omni_server, online_client) -> None: """Test Bagel text2img via OpenAI-compatible chat completions API.""" request_config = { "model": omni_server.model, @@ -95,15 +94,14 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None: }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_bagel_img2img_online(omni_server, openai_client) -> None: +def test_bagel_img2img_online(omni_server, online_client) -> None: """Test Bagel img2img via OpenAI-compatible chat completions API.""" input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") buffer = BytesIO() @@ -121,4 +119,4 @@ def test_bagel_img2img_online(omni_server, openai_client) -> None: }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index fa17791e5..e22614a6b 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/tests/e2e/online_serving/test_bagel_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Comprehensive tests of diffusion features that are available in online serving mode and are supported by the following models: @@ -19,10 +22,10 @@ assert_diffusion_response validates successful generation and the expected import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution." NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark" @@ -183,7 +186,7 @@ def _get_diffusion_feature_cases(model: str): ) def test_bagel( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """L4 diffusion feature coverage for Bagel on H100. @@ -218,4 +221,4 @@ def test_bagel( }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_bagel_multi_replicas.py b/tests/e2e/online_serving/test_bagel_multi_replicas.py index cb7aeeb94..c360b63e8 100644 --- a/tests/e2e/online_serving/test_bagel_multi_replicas.py +++ b/tests/e2e/online_serving/test_bagel_multi_replicas.py @@ -1,6 +1,11 @@ -"""Nightly-only e2e coverage for BAGEL diffusion multi-replica serving. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Weekly e2e coverage for BAGEL diffusion multi-replica serving. This test needs 4 H100 GPUs; keep it out of generic test_bagel_* jobs. +It is covered by the weekly Diffusion · H100 · Multi-GPU lane +(``.buildkite/cuda/test-weekly.yml``, ``mirror_hardwares: h100_4``). """ import os @@ -17,7 +22,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT" MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/bagel_multi_replicas_4gpu.yaml") ROUTE_STRESS_REQUESTS = 6 -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] test_params = [ OmniServerParams( @@ -31,7 +36,7 @@ test_params = [ @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, openai_client) -> None: +def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, online_client) -> None: request_config = { "model": omni_server.model, "messages": dummy_messages_from_mix_data( @@ -47,7 +52,7 @@ def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, openai_ }, } - responses = openai_client.send_diffusion_request(request_config, request_num=ROUTE_STRESS_REQUESTS) + responses = online_client.send_diffusion_request(request_config, request_num=ROUTE_STRESS_REQUESTS) assert len(responses) == ROUTE_STRESS_REQUESTS assert all(response.success for response in responses) diff --git a/tests/e2e/online_serving/test_boogu_image.py b/tests/e2e/online_serving/test_boogu_image.py index 5a5494d7d..440e031f1 100644 --- a/tests/e2e/online_serving/test_boogu_image.py +++ b/tests/e2e/online_serving/test_boogu_image.py @@ -16,7 +16,7 @@ parity vs. the Diffusers baseline is covered separately (see the support plan, s From ``tests/``:: pytest -s -v e2e/online_serving/test_boogu_image.py \ - -m "full_model and diffusion and H100" --run-level=full_model + -m "slow and diffusion and H100" """ import os @@ -28,7 +28,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "Boogu/Boogu-Image-0.1-Base" T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting" diff --git a/tests/e2e/online_serving/test_boogu_image_edit.py b/tests/e2e/online_serving/test_boogu_image_edit.py index 45f999d26..118b46335 100644 --- a/tests/e2e/online_serving/test_boogu_image_edit.py +++ b/tests/e2e/online_serving/test_boogu_image_edit.py @@ -1,15 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ -Online serving tests for Boogu-Image-0.1-Edit (image-to-image via chat completions). +Online serving tests for Boogu-Image-0.1-Edit and Edit-Turbo. -The Edit checkpoint shares the ``BooguImagePipeline`` class with the Base +Both Edit checkpoints share the ``BooguImagePipeline`` class with the Base text-to-image checkpoint; the editing (TI2I) path activates when the request -carries a reference image. Text guidance rides on ``guidance_scale`` (upstream -default 4.0), and image guidance rides on ``guidance_scale_2`` (upstream default -1.0 = off; a value > 1 enables the double-guidance path). Only a single -reference image is supported. +carries a reference image. The regular Edit checkpoint uses text guidance via +``guidance_scale`` and optional image guidance via ``guidance_scale_2``. +Edit-Turbo is guidance-distilled and uses four steps with guidance scale 1.0. +Only a single reference image is supported. - ``test_single_image_to_image_001``: one reference image, text guidance only. - ``test_double_guidance_001``: one reference image, text + image guidance. @@ -19,7 +19,7 @@ reference image is supported. From ``tests/``:: pytest -s -v e2e/online_serving/test_boogu_image_edit.py \ - -m "full_model and diffusion and H100" --run-level=full_model + -m "slow and diffusion and H100" """ import base64 @@ -35,11 +35,15 @@ from tests.helpers.runtime import ( dummy_messages_from_mix_data, ) -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] EDIT_PROMPT = "Change the style to a colored pencil drawing." SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) +EDIT_MODEL = "Boogu/Boogu-Image-0.1-Edit" +EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo" +EDIT_TURBO_REVISION = "hotfix-1k-20260708" + def _get_diffusion_feature_cases(model: str): """Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``).""" @@ -54,13 +58,44 @@ def _get_diffusion_feature_cases(model: str): ] +def _get_single_edit_cases(): + """Return the regular and distilled Edit server/request configurations.""" + return [ + pytest.param( + OmniServerParams(model=EDIT_MODEL), + { + "num_inference_steps": 2, + "guidance_scale": 5.0, + }, + id="edit", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=EDIT_TURBO_MODEL, + server_args=["--revision", EDIT_TURBO_REVISION], + ), + { + "num_inference_steps": 4, + "guidance_scale": 1.0, + }, + id="edit-turbo-hotfix-1k", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + @pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"), - indirect=True, + ("omni_server", "sampling_params"), + _get_single_edit_cases(), + indirect=["omni_server"], ) -def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Single-reference text-guided edit smoke for ``Boogu/Boogu-Image-0.1-Edit``.""" +def test_single_image_to_image_001( + omni_server: OmniServer, + sampling_params: dict, + openai_client: OpenAIClientHandler, +): + """Single-reference text-guided edit smoke for regular and Turbo Edit.""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) @@ -69,9 +104,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA "model": omni_server.model, "messages": messages, "extra_body": { - "num_inference_steps": 2, - # Boogu text guidance (upstream default 4.0); > 1 enables CFG. - "guidance_scale": 5.0, + **sampling_params, "seed": 42, }, } @@ -81,7 +114,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"), + _get_diffusion_feature_cases(EDIT_MODEL), indirect=True, ) def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): @@ -105,15 +138,20 @@ def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClien openai_client.send_diffusion_request(request_config) -def _edits_request_config(model: str, *, guidance_scale_2: float | None = None) -> dict: +def _edits_request_config( + model: str, + *, + num_inference_steps: int = 2, + guidance_scale: float = 5.0, + guidance_scale_2: float | None = None, +) -> dict: """Multipart ``/v1/images/edits`` config: one synthetic reference + smoke-depth params.""" image_bytes = base64.b64decode(generate_synthetic_image(512, 512)["base64"]) data = { "model": model, "prompt": EDIT_PROMPT, - "num_inference_steps": 2, - # Boogu text guidance (upstream default 4.0); > 1 enables CFG. - "guidance_scale": 5.0, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, "seed": 42, } if guidance_scale_2 is not None: @@ -133,19 +171,25 @@ def _assert_edits_response_has_image(responses) -> None: @pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"), - indirect=True, + ("omni_server", "sampling_params"), + _get_single_edit_cases(), + indirect=["omni_server"], ) -def test_images_edits_endpoint_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Single-reference text-guided edit via the OpenAI-compatible ``/v1/images/edits`` endpoint.""" - responses = openai_client.send_images_edits_http_request(_edits_request_config(omni_server.model)) +def test_images_edits_endpoint_001( + omni_server: OmniServer, + sampling_params: dict, + openai_client: OpenAIClientHandler, +): + """Single-reference regular and Turbo Edit via ``/v1/images/edits``.""" + responses = openai_client.send_images_edits_http_request( + _edits_request_config(omni_server.model, **sampling_params) + ) _assert_edits_response_has_image(responses) @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"), + _get_diffusion_feature_cases(EDIT_MODEL), indirect=True, ) def test_images_edits_endpoint_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): diff --git a/tests/e2e/online_serving/test_boogu_image_edit_expansion.py b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py index efc968eb6..df7650b65 100644 --- a/tests/e2e/online_serving/test_boogu_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ -L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Edit`` (image editing / TI2I). +L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (image editing / TI2I). The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image text-guided editing and the text+image double-guidance path. This expansion @@ -16,16 +16,17 @@ module adds the two edit paths smoke does not exercise: ``height`` / ``width``: asserts the output resolution follows the reference (upstream ``align_res``, on by default for a single reference). -The Edit checkpoint shares ``BooguImagePipeline`` with the Base checkpoint; the -TI2I path activates when the request carries a reference image. Only a single -reference image is supported. Cases stay smoke-depth (``num_inference_steps=2``). -CPU offload, Cache-DiT, and multi-GPU parallelism are not yet supported for this -model and are intentionally omitted. +Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint; +the TI2I path activates when the request carries a reference image. Only a +single reference image is supported. The regular Edit cases stay smoke-depth +(``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step +configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet +supported for these models and are intentionally omitted. From ``tests/``:: pytest -s -v e2e/online_serving/test_boogu_image_edit_expansion.py \ - -m "full_model and diffusion and H100" --run-level=full_model + -m "slow and diffusion and H100" """ import os @@ -38,9 +39,11 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "Boogu/Boogu-Image-0.1-Edit" +EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo" +EDIT_TURBO_REVISION = "hotfix-1k-20260708" EDIT_PROMPT = "Change the style to a colored pencil drawing." SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) @@ -57,6 +60,33 @@ def _get_diffusion_feature_cases(model: str): ] +def _get_align_res_cases(): + """Return regular and Turbo Edit configurations for align-res coverage.""" + return [ + pytest.param( + OmniServerParams(model=MODEL), + { + "num_inference_steps": 2, + "guidance_scale": 5.0, + }, + id="edit", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=EDIT_TURBO_MODEL, + server_args=["--revision", EDIT_TURBO_REVISION], + ), + { + "num_inference_steps": 4, + "guidance_scale": 1.0, + }, + id="edit-turbo-hotfix-1k", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: """Image-only guidance: text CFG off (1.0), image CFG on (2.0) -> image-only CFG branch.""" @@ -76,8 +106,16 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien openai_client.send_diffusion_request(request_config) -@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: +@pytest.mark.parametrize( + ("omni_server", "sampling_params"), + _get_align_res_cases(), + indirect=["omni_server"], +) +def test_align_res_output_size( + omni_server: OmniServer, + sampling_params: dict, + openai_client: OpenAIClientHandler, +) -> None: """Output resolution follows the reference image (upstream ``align_res``). No ``height`` / ``width`` is requested. The reference is 512x384 — both @@ -92,8 +130,7 @@ def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAICli "model": omni_server.model, "messages": messages, "extra_body": { - "num_inference_steps": 2, - "guidance_scale": 5.0, + **sampling_params, "seed": 42, }, } diff --git a/tests/e2e/online_serving/test_boogu_image_expansion.py b/tests/e2e/online_serving/test_boogu_image_expansion.py index ef6d3a62e..a56e5a338 100644 --- a/tests/e2e/online_serving/test_boogu_image_expansion.py +++ b/tests/e2e/online_serving/test_boogu_image_expansion.py @@ -30,7 +30,7 @@ omitted (add rows here when those features land). From ``tests/``:: pytest -s -v e2e/online_serving/test_boogu_image_expansion.py \ - -m "full_model and diffusion and H100" --run-level=full_model + -m "slow and diffusion and H100" """ import os @@ -42,7 +42,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "Boogu/Boogu-Image-0.1-Base" T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting" diff --git a/tests/e2e/online_serving/test_cosmos3.py b/tests/e2e/online_serving/test_cosmos3.py index 09cb2365b..f28e2ae26 100644 --- a/tests/e2e/online_serving/test_cosmos3.py +++ b/tests/e2e/online_serving/test_cosmos3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Online serving L2 smoke tests for ``nvidia/Cosmos3-Nano``. @@ -8,6 +8,7 @@ broader image/video similarity checks stay in ``tests/e2e/accuracy``. """ import os +from typing import Any import pytest @@ -16,7 +17,7 @@ from tests.helpers.media import generate_synthetic_image, generate_synthetic_vid from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, ) os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -56,9 +57,9 @@ def _get_diffusion_feature_cases(model: str): @pytest.mark.core_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: +def test_text_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: """Default Cosmos3 T2I smoke through ``/v1/images/generations``.""" - responses = openai_client.send_images_generations_http_request( + responses = online_client.send_images_generations_http_request( { "json": { "model": omni_server.model, @@ -85,12 +86,12 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH @pytest.mark.core_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: +def test_video_generation_modes_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: """Default Cosmos3 T2V, I2V, and V2V smoke through the async ``/v1/videos`` endpoint.""" image_reference = f"data:image/jpeg;base64,{generate_synthetic_image(256, 256, seed=42)['base64']}" video_reference = f"data:video/mp4;base64,{generate_synthetic_video(256, 256, 5)['base64']}" - video_cases = [ + video_cases: list[dict[str, Any]] = [ { "form_data": { **VIDEO_SHAPE, @@ -118,8 +119,8 @@ def test_video_generation_modes_001(omni_server: OmniServer, openai_client: Open ] for case in video_cases: - request_config = { + request_config: dict[str, Any] = { "model": omni_server.model, **case, } - openai_client.send_video_diffusion_request(request_config) + online_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py index 4b21f956e..544afc7bd 100644 --- a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py +++ b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for CosyVoice3 TTS model with voice cloning. @@ -14,18 +14,26 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.slow, pytest.mark.tts] +if current_omni_platform.is_rocm(): + os.environ.setdefault("COSYVOICE3_TRT", "0") + +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.core_model, +] MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" # Official CosyVoice zero-shot prompt audio and its transcript. Vendored under # tests/assets/ so the server does not depend on raw.githubusercontent.com being # reachable at request time (same rationale as issue #3263 for Qwen3-TTS). -REF_AUDIO_URL = load_test_audio_data_url("cosyvoice3/zero_shot_prompt.wav") +REF_AUDIO_URL = get_asset_path("cosyvoice3/zero_shot_prompt.wav", as_data_url=True) REF_TEXT = "希望你以后能够做的比我还好呦。" @@ -42,6 +50,12 @@ def get_prompt(prompt_type="zh"): return prompts.get(prompt_type, prompts["zh"]) +def _set_rocm_request_timeout(request_config: dict) -> None: + if current_omni_platform.is_rocm(): + # Cold AITER MHA builds on loaded MI300 agents can exceed 20 minutes. + request_config["timeout"] = 1800.0 + + tts_server_params = [ pytest.param( OmniServerParams( @@ -65,10 +79,9 @@ tts_async_chunk_server_params = [ ] -@hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644") +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_clone_zh_001(omni_server, openai_client) -> None: +def test_voice_clone_zh_001(omni_server, online_client) -> None: """ Test voice cloning TTS with Chinese text via OpenAI API. Deploy Setting: default yaml @@ -85,13 +98,13 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config) + _set_rocm_request_timeout(request_config) + online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644") +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True) -def test_voice_clone_zh_002(omni_server, openai_client) -> None: +def test_voice_clone_zh_002(omni_server, online_client) -> None: """ Test voice cloning TTS with Chinese text via async_chunk streaming. Deploy Setting: cosyvoice3.yaml with default ``async_chunk: true`` @@ -109,13 +122,13 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config) + _set_rocm_request_timeout(request_config) + online_client.send_audio_speech_request(request_config) -@hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644") +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_clone_en_001(omni_server, openai_client) -> None: +def test_voice_clone_en_001(omni_server, online_client) -> None: """ Test voice cloning TTS with English text via OpenAI API. Deploy Setting: default yaml @@ -132,4 +145,5 @@ def test_voice_clone_en_001(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config) + _set_rocm_request_timeout(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_covo_audio_expansion.py b/tests/e2e/online_serving/test_covo_audio_expansion.py index 33771d647..6a387269b 100644 --- a/tests/e2e/online_serving/test_covo_audio_expansion.py +++ b/tests/e2e/online_serving/test_covo_audio_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for Covo-Audio-Chat model. """ @@ -37,7 +37,7 @@ test_params = [ @pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_audio_to_audio_001(omni_server, openai_client) -> None: +def test_audio_to_audio_001(omni_server, online_client) -> None: """ Input Modal: text + audio Output Modal: audio @@ -58,4 +58,4 @@ def test_audio_to_audio_001(omni_server, openai_client) -> None: "stream": False, } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py index 1f16a584b..c4cd40ff8 100644 --- a/tests/e2e/online_serving/test_dynin_omni_expansion.py +++ b/tests/e2e/online_serving/test_dynin_omni_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Example online tests for Dynin-Omni model. """ @@ -64,7 +64,7 @@ def _build_i2i_messages(prompt: str) -> list[dict]: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_i2i_request_001(omni_server, openai_client) -> None: +def test_send_i2i_request_001(omni_server, online_client) -> None: request_config = { "model": omni_server.model, "messages": _build_i2i_messages(I2I_PROMPT), @@ -73,27 +73,27 @@ def test_send_i2i_request_001(omni_server, openai_client) -> None: "sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)], }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_t2i_request_001(omni_server, openai_client) -> None: +def test_send_t2i_request_001(omni_server, online_client) -> None: request_config = { "model": omni_server.model, "messages": _build_t2i_messages(T2I_PROMPT), "modalities": ["image"], } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_t2s_request_001(omni_server, openai_client) -> None: +def test_send_t2s_request_001(omni_server, online_client) -> None: request_config = { "model": omni_server.model, "messages": _build_t2s_messages(T2S_PROMPT), "modalities": ["audio"], "audio_ref_text": T2S_PROMPT, } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) diff --git a/tests/e2e/online_serving/test_ernie_image_expansion.py b/tests/e2e/online_serving/test_ernie_image_expansion.py index 848fec4e6..b06267b39 100644 --- a/tests/e2e/online_serving/test_ernie_image_expansion.py +++ b/tests/e2e/online_serving/test_ernie_image_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ End-to-end test for ErnieImage text-to-image generation via online serving. @@ -11,7 +11,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) @@ -45,7 +45,7 @@ def _get_diffusion_feature_cases(model: str): _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_ernie_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_ernie_image(omni_server: OmniServer, online_client: OnlineOmniClient): """Test each diffusion feature with ERNIE-Image (text-to-image), one feature per case.""" messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT) request_config = { @@ -60,4 +60,4 @@ def test_ernie_image(omni_server: OmniServer, openai_client: OpenAIClientHandler "seed": 42, }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py index 2659cb927..d7c20e832 100644 --- a/tests/e2e/online_serving/test_flux2_expansion.py +++ b/tests/e2e/online_serving/test_flux2_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Tests for Flux2 Klein; currently Dev is implemented separately, but ideally these models will fold together in the future. @@ -15,7 +18,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) @@ -113,7 +116,7 @@ def _get_diffusion_feature_cases(model: str): ), indirect=True, ) -def test_flux2_klein(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_flux2_klein(omni_server: OmniServer, online_client: OnlineOmniClient): messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT) request_config = { "model": omni_server.model, @@ -128,4 +131,4 @@ def test_flux2_klein(omni_server: OmniServer, openai_client: OpenAIClientHandler }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index 19a14ce39..932b3cfa5 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ End-to-end diffusion coverage for FLUX.2-dev in online serving mode. @@ -19,7 +22,7 @@ image resolution. import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data pytestmark = [pytest.mark.diffusion, pytest.mark.slow] @@ -116,7 +119,7 @@ def _get_flux_2_dev_feature_cases(model: str): ) def test_flux_2_dev( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """Validate FLUX.2-dev online serving with CPU offload.""" @@ -135,7 +138,7 @@ def test_flux_2_dev( }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.parametrize( @@ -145,14 +148,14 @@ def test_flux_2_dev( ) def test_flux_2_dev_batched_chat_completions( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """Validate batched chat completions for diffusion models.""" messages = [ [{"role": "user", "content": PROMPT}], [{"role": "user", "content": "A sunset over the ocean."}], ] - responses = openai_client.send_batched_chat_completions_http_request( + responses = online_client.send_batched_chat_completions_http_request( { "json": { "model": omni_server.model, @@ -167,5 +170,7 @@ def test_flux_2_dev_batched_chat_completions( assert responses and len(responses) == 1 resp = responses[0] assert resp.success - choices = resp.json_body["choices"] + body = resp.json_body + assert isinstance(body, dict) + choices = body["choices"] assert len(choices) == len(messages) diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py index ede5ff56b..79460c631 100644 --- a/tests/e2e/online_serving/test_flux_kontext_expansion.py +++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Comprehensive tests of diffusion features that are available in online serving mode and are supported by the FluxKontext model. @@ -7,7 +10,7 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data pytestmark = [pytest.mark.diffusion, pytest.mark.slow] @@ -32,7 +35,7 @@ def _get_diffusion_feature_cases(model: str): _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_flux_kontext_text_to_image(omni_server: OmniServer, online_client: OnlineOmniClient): """Test text-to-image generation with FluxKontext in regular end-user scenarios.""" messages = dummy_messages_from_mix_data(content_text="A photo of a cat sitting on a laptop keyboard") @@ -47,7 +50,7 @@ def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: Open }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.parametrize( @@ -55,7 +58,7 @@ def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: Open _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_flux_kontext_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient): """Test image editing with FluxKontext in regular end-user scenarios.""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" @@ -74,7 +77,7 @@ def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIC }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.parametrize( @@ -82,7 +85,7 @@ def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIC _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, online_client: OnlineOmniClient): """Test image editing with FluxKontext without negative prompt.""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" @@ -99,7 +102,7 @@ def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_cli }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.parametrize( @@ -107,7 +110,7 @@ def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_cli _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_flux_kontext_high_resolution(omni_server: OmniServer, online_client: OnlineOmniClient): """Test high-resolution generation with FluxKontext.""" messages = dummy_messages_from_mix_data(content_text="A beautiful landscape with mountains and a lake") @@ -122,7 +125,7 @@ def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: Op }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.parametrize( @@ -130,7 +133,7 @@ def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: Op _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_flux_kontext_multiple_outputs(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_flux_kontext_multiple_outputs(omni_server: OmniServer, online_client: OnlineOmniClient): """Test generating multiple outputs with FluxKontext.""" messages = dummy_messages_from_mix_data(content_text="A photo of a cat sitting on a laptop keyboard") @@ -146,4 +149,4 @@ def test_flux_kontext_multiple_outputs(omni_server: OmniServer, openai_client: O }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_glm_image_expansion.py b/tests/e2e/online_serving/test_glm_image_expansion.py index d2a08f633..290272364 100644 --- a/tests/e2e/online_serving/test_glm_image_expansion.py +++ b/tests/e2e/online_serving/test_glm_image_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Comprehensive tests of diffusion features that are available in online serving mode and are supported by the following model: @@ -25,12 +28,12 @@ from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = os.environ.get("GLM_IMAGE_MODEL_PATH", "zai-org/GLM-Image") @@ -152,7 +155,7 @@ MODES = ["t2i", "i2i"] def test_glm_image( omni_server: OmniServer, mode: str, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """GLM-Image diffusion feature coverage for t2i and i2i on two H100s.""" @@ -180,4 +183,4 @@ def test_glm_image( }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_glm_tts_expansion.py b/tests/e2e/online_serving/test_glm_tts_expansion.py index c558e0c40..89a98f6fb 100644 --- a/tests/e2e/online_serving/test_glm_tts_expansion.py +++ b/tests/e2e/online_serving/test_glm_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E expansion tests for GLM-TTS online serving (nightly CI). @@ -14,7 +14,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path @@ -24,7 +24,7 @@ MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS") REF_TEXT = "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml") -REF_AUDIO_URL = load_test_audio_data_url("glm_tts/jiayan_zh.wav") +REF_AUDIO_URL = get_asset_path("glm_tts/jiayan_zh.wav", as_data_url=True) SYNC_EXTRA_ARGS = [ "--trust-remote-code", @@ -64,7 +64,7 @@ tts_async_chunk_server_params = [ @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True) -def test_voice_clone_zh_sync(omni_server, openai_client) -> None: +def test_voice_clone_zh_sync(omni_server, online_client) -> None: """ Test voice cloning TTS with Chinese text. Deploy Setting: glm_tts.yaml with ``--no-async-chunk`` (sync two-stage) @@ -80,12 +80,12 @@ def test_voice_clone_zh_sync(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True) -def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None: +def test_voice_clone_zh_async_chunk(omni_server, online_client) -> None: """ Test voice cloning TTS with Chinese text via async_chunk streaming. Deploy Setting: glm_tts.yaml default ``async_chunk: true`` @@ -102,12 +102,12 @@ def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True) -def test_models_endpoint(omni_server, openai_client) -> None: +def test_models_endpoint(omni_server, online_client) -> None: """Test the /v1/models endpoint returns loaded model.""" - models = openai_client.client.models.list() + models = online_client.client.models.list() assert len(models.data) > 0 diff --git a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py index a362d4902..dea3b3c9a 100644 --- a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py +++ b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ End-to-end diffusion coverage for HiDream-I1-Full in online serving mode. @@ -12,7 +15,7 @@ image resolution. import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data pytestmark = [pytest.mark.diffusion, pytest.mark.slow] @@ -49,7 +52,7 @@ def _get_hidream_i1_image_feature_cases(model: str): @pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4662") def test_hidream_i1_image( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """Validate HiDream-I1-Full online serving with CPU offload.""" @@ -67,4 +70,4 @@ def test_hidream_i1_image( }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py b/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py index 8927eb039..e549c6bfc 100644 --- a/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py +++ b/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py @@ -1,17 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online expansion tests for higgs-audio v2 against /v1/audio/speech. v1 scope is plain text -> 24 kHz speech plus shallow voice clone via ref_audio + ref_text (inline) or voice= (after POST /v1/audio/voices). -The model-aware request validator -(vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2::validate) -rejects multi-speaker tags, language overrides, task_type, and bare -voice= for names that do not match an uploaded speaker — this suite -exercises both the happy path (plain text in, audio bytes out) and the -validator rejections, so a regression that loosens the schema will fail -loudly. +Model-aware validator rejections live in L1 +``tests/entrypoints/openai_api/test_tts_adapter.py`` +(``test_higgs_audio_v2_validate_*``). """ from __future__ import annotations @@ -27,7 +23,7 @@ os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0") import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path @@ -63,7 +59,7 @@ _MIN_AUDIO_BYTES = 20_000 # speech) + its transcript. See tests/e2e/online_serving/test_qwen3_tts_base.py # for the asset rationale — keeping a single shared reference clip across TTS # tests avoids duplicating WAVs in the repo. -_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +_REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) _REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." @@ -72,9 +68,9 @@ class TestHiggsAudioV2OnlineHappyPath: """Plain-text -> audio happy paths against the live HTTP server.""" @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_plain_text_wav(self, omni_server, openai_client) -> None: + def test_plain_text_wav(self, omni_server, online_client) -> None: """Single non-streaming WAV request — covers the canonical TTS happy path.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Hello world.", @@ -86,7 +82,7 @@ class TestHiggsAudioV2OnlineHappyPath: ) @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: + def test_plain_text_pcm_streaming(self, omni_server, online_client) -> None: """Streaming PCM via the shared-memory connector's codec_streaming path. higgs_audio_v2.yaml pins ``codec_streaming: true`` + ``async_chunk: false``, @@ -102,7 +98,7 @@ class TestHiggsAudioV2OnlineHappyPath: content correctness is separately guarded by the full_model transcript check. """ - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Streaming the quick brown fox over the lazy dog.", @@ -116,9 +112,9 @@ class TestHiggsAudioV2OnlineHappyPath: ) @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None: + def test_plain_text_with_max_new_tokens(self, omni_server, online_client) -> None: """max_new_tokens is one of the few extra fields the higgs validator accepts.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Innovation distinguishes between a leader and a follower.", @@ -131,9 +127,9 @@ class TestHiggsAudioV2OnlineHappyPath: ) @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_concurrent_plain_text(self, omni_server, openai_client) -> None: + def test_concurrent_plain_text(self, omni_server, online_client) -> None: """Three concurrent non-streaming requests — guards the per-slot audio state.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "It was the night before my birthday.", @@ -146,104 +142,6 @@ class TestHiggsAudioV2OnlineHappyPath: ) -# --------------------------------------------------------------------------- -# Validator rejections served over HTTP. Each case targets one out-of-scope -# field; the validator returns 4xx via the OpenAI error path. We keep the -# error-message substring loose so a phrasing tweak in the validator does not -# break CI, but tight enough to catch a regression that silently accepts the -# field. -# --------------------------------------------------------------------------- - - -@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -class TestHiggsAudioV2OnlineValidatorRejections: - """Out-of-scope fields must come back as 4xx with a higgs-named message.""" - - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None: - """Voice clone needs the transcript too — half-supplied is a 4xx.""" - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "Hello world.", - "ref_audio": _REF_AUDIO_URL, - "response_format": "wav", - "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "status_code": (400, 422), - "err_message": "ref_text", - } - ) - - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None: - """Symmetric guard: ref_text alone is not enough — must come with ref_audio.""" - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "Hello world.", - "ref_text": "some transcript", - "response_format": "wav", - "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "status_code": (400, 422), - "err_message": "ref_audio", - } - ) - - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_rejects_task_type(self, omni_server, openai_client) -> None: - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "Hello world.", - "task_type": "Base", - "response_format": "wav", - "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "status_code": (400, 422), - "err_message": "task_type", - } - ) - - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_rejects_language_override(self, omni_server, openai_client) -> None: - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "Hello world.", - "language": "Chinese", - "response_format": "wav", - "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "status_code": (400, 422), - "err_message": "language", - } - ) - - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None: - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "[SPEAKER0] hi", - "response_format": "wav", - "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "status_code": (400, 422), - "err_message": "multi-speaker", - } - ) - - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_rejects_empty_input(self, omni_server, openai_client) -> None: - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": " ", - "response_format": "wav", - "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "status_code": (400, 422), - "err_message": "empty", - } - ) - - @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) class TestHiggsAudioV2OnlineVoiceClone: """Shallow voice clone: ref_audio + ref_text -> speech in the cloned voice. @@ -256,9 +154,9 @@ class TestHiggsAudioV2OnlineVoiceClone: """ @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_voice_clone_basic(self, omni_server, openai_client) -> None: + def test_voice_clone_basic(self, omni_server, online_client) -> None: """Single non-streaming WAV request driven by the qwen3_tts ref clip.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Hello world.", @@ -272,9 +170,9 @@ class TestHiggsAudioV2OnlineVoiceClone: ) @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None: + def test_voice_clone_pcm_streaming(self, omni_server, online_client) -> None: """Voice clone over the streaming PCM path.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Innovation distinguishes a leader from a follower.", diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py index d26e6bff4..e9ba68909 100644 --- a/tests/e2e/online_serving/test_higgs_audio_v3.py +++ b/tests/e2e/online_serving/test_higgs_audio_v3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end online tests for higgs-audio v3 against /v1/audio/speech. Mirrors the higgs_audio_v2 test layout. Covers the plain-text-in / audio-out @@ -21,7 +21,7 @@ os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0") import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path @@ -51,7 +51,7 @@ _MIN_AUDIO_BYTES = 20_000 # Reuse the shared TTS reference clip (clean ~5 s 24 kHz mono human speech) # vendored under tests/assets/qwen3_tts/. Keeps a single WAV across TTS # tests rather than duplicating asset bytes. -_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +_REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) _REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." @@ -59,13 +59,12 @@ _REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know w class TestHiggsAudioV3OnlineHappyPath: """Plain-text -> audio happy paths against the live HTTP server.""" - @pytest.mark.core_model - @pytest.mark.advanced_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_plain_text_wav(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_plain_text_wav(self, omni_server, online_client) -> None: """Single non-streaming WAV request - canonical TTS happy path.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Hello world.", @@ -79,31 +78,29 @@ class TestHiggsAudioV3OnlineHappyPath: } ) - @pytest.mark.core_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_plain_text_with_max_new_tokens(self, omni_server, online_client) -> None: """``max_new_tokens`` is one of the few extra fields the v3 validator accepts.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Innovation distinguishes between a leader and a follower.", "stream": False, "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "max_new_tokens": 32, + "max_new_tokens": 1024, "min_audio_bytes": _MIN_AUDIO_BYTES, } ) - @pytest.mark.skip(reason="issue#4411") - @pytest.mark.core_model - @pytest.mark.advanced_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_concurrent_plain_text(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_concurrent_plain_text(self, omni_server, online_client) -> None: """Three concurrent non-streaming requests - guards per-slot audio state and Stage-0 PC under batching.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "it was the night before my birthday", @@ -114,11 +111,10 @@ class TestHiggsAudioV3OnlineHappyPath: request_num=3, ) - @pytest.mark.skip(reason="issue#4411") - @pytest.mark.advanced_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_plain_text_pcm_streaming(self, omni_server, online_client) -> None: """Streaming PCM via talker2code2wav_async_chunk + the connector's codec_streaming path. The deploy YAML pins ``async_chunk: true`` and ``codec_streaming: true`` so chunks @@ -134,7 +130,7 @@ class TestHiggsAudioV3OnlineHappyPath: ~0.16 dB margin over measured single-request output (1.16 dB on L4), so flake risk is too high there. """ - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "The quick brown fox jumps over the lazy dog.", @@ -146,11 +142,10 @@ class TestHiggsAudioV3OnlineHappyPath: } ) - @pytest.mark.skip(reason="issue#4411") - @pytest.mark.advanced_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_concurrent_pcm_streaming(self, omni_server, online_client) -> None: """Three concurrent streaming requests - guards per-request frame cursors in ``talker2code2wav_async_chunk`` and per-slot delay-pattern state under batched AR. @@ -160,7 +155,7 @@ class TestHiggsAudioV3OnlineHappyPath: over the measured worst-of-3 on L4 (-0.42 dB) while still well above the catastrophic-failure region. """ - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "She sells seashells by the seashore.", @@ -184,17 +179,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens: tokens (pause, sfx) go inline. SFX tokens must pair with their written onomatopoeia. - These tests check the *serving surface* - the validator accepts the - payload, the engine produces audio, and the WAV has non-trivial size. - They do not assert audio quality (out of scope for CI). + These tests check the serving surface and spoken text. They do not assert + whether the requested control changes audio quality. """ - @pytest.mark.core_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_inline_emotion_and_expressive(self, omni_server, online_client) -> None: """Delivery tokens (emotion + expressive_high) at the start of input.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": ("<|emotion:amusement|><|prosody:expressive_high|>Wait, that was actually hilarious."), @@ -202,15 +196,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens: "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, "min_audio_bytes": _MIN_AUDIO_BYTES, + "transcript_expected_text": "Wait, that was actually hilarious.", } ) - @pytest.mark.core_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_inline_style_whispering(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_inline_style_whispering(self, omni_server, online_client) -> None: """Style token at the start - ``whispering``.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "<|style:whispering|>It is just between you and me, alright?", @@ -218,15 +213,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens: "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, "min_audio_bytes": _MIN_AUDIO_BYTES, + "transcript_expected_text": "It is just between you and me, alright?", } ) - @pytest.mark.core_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_inline_prosody_speed_and_pitch(self, omni_server, online_client) -> None: """Two prosody tokens at the start - slow speed plus low pitch.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": ( @@ -236,15 +232,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens: "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, "min_audio_bytes": _MIN_AUDIO_BYTES, + "transcript_expected_text": "The radar shows a storm approaching from the east.", } ) - @pytest.mark.core_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_inline_pause_mid_text(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_inline_pause_mid_text(self, omni_server, online_client) -> None: """Positional pause token placed inline between two clauses.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Hold on a moment <|prosody:pause|> let me think about it.", @@ -252,15 +249,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens: "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, "min_audio_bytes": _MIN_AUDIO_BYTES, + "transcript_expected_text": "Hold on a moment. Let me think about it.", } ) - @pytest.mark.core_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_inline_sfx_with_onomatopoeia(self, omni_server, online_client) -> None: """SFX token paired with its written onomatopoeia (``<|sfx:laughter|>Hehe``).""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": ( @@ -271,6 +269,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens: "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, "min_audio_bytes": _MIN_AUDIO_BYTES, + "transcript_expected_text": "I cannot believe that just happened. I am still recovering from it.", } ) @@ -279,13 +278,12 @@ class TestHiggsAudioV3OnlineInlineControlTokens: class TestHiggsAudioV3OnlineVoiceClone: """Voice clone via the two payload shapes the model serves.""" - @pytest.mark.core_model - @pytest.mark.advanced_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_voice_clone_ref_audio_ref_text(self, omni_server, online_client) -> None: """Canonical vllm-omni voice clone via ``ref_audio`` + ``ref_text``.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Hello world.", @@ -297,17 +295,17 @@ class TestHiggsAudioV3OnlineVoiceClone: } ) - @pytest.mark.core_model + @pytest.mark.slow @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_voice_clone_references_alias(self, omni_server, openai_client) -> None: + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) + def test_voice_clone_references_alias(self, omni_server, online_client) -> None: """BosonAI cookbook payload: ``references=[{audio_path, text}]``. ``normalize_references_alias`` (``protocol/audio.py``) is supposed to translate the cookbook field into ``ref_audio`` / ``ref_text``; without it the request would silently fall through to zero-shot synthesis. """ - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Hello world.", diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py index a448335c6..0bc3f449f 100644 --- a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py +++ b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py @@ -1,20 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests of common diffusion feature combinations in online serving mode for HunyuanVideo-1.5-T2V (480p). Coverage (H100, since model cannot fit L4): -- CacheDiT + Layerwise CPU offloading (1 GPU) -- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs) +- CPU offloading (1 GPU) — ``core_model`` + ``advanced_model`` +- CacheDiT + Layerwise CPU offloading (1 GPU) — ``full_model`` +- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs) — ``full_model`` + +HunyuanVideo-1.5 is a high-priority model, so only the most basic single-card deployment +row runs on every PR (L2) and on merge (L3). The heavyweight feature combinations stay +nightly-only (L4), together with the video similarity suites in +``tests/e2e/accuracy/hunyuanvideo15_{t2v,i2v}/``. + +From ``tests/``:: + + pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "core_model and diffusion" --run-level=core_model + pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model and diffusion" --run-level=advanced_model + pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and diffusion" --run-level=full_model """ import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient -pytestmark = [pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.diffusion] PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion." NEGATIVE_PROMPT = "low quality, blurry, distorted" @@ -24,9 +36,14 @@ MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +# Only the most basic single-card deployment row is cheap enough for PR (L2) / merge (L3) +# CI; the CacheDiT / parallel feature combinations run nightly (L4). +_CORE_MARKS = [pytest.mark.core_model, pytest.mark.advanced_model] +_FULL_MARKS = [pytest.mark.full_model] + def _get_diffusion_feature_cases(model: str): - """Return L4 diffusion feature cases for HunyuanVideo-1.5. + """Return diffusion feature cases for HunyuanVideo-1.5. Designed for 2x H100 environment per issue #1832. """ @@ -40,7 +57,7 @@ def _get_diffusion_feature_cases(model: str): ], ), id="single_card_cpu_offload", - marks=SINGLE_CARD_MARKS, + marks=SINGLE_CARD_MARKS + _CORE_MARKS, ), # (1 GPU) CacheDiT + Layerwise CPU offloading pytest.param( @@ -53,7 +70,7 @@ def _get_diffusion_feature_cases(model: str): ], ), id="single_card_cachedit_layerwise", - marks=SINGLE_CARD_MARKS, + marks=SINGLE_CARD_MARKS + _FULL_MARKS, ), # (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2 pytest.param( @@ -70,7 +87,7 @@ def _get_diffusion_feature_cases(model: str): ], ), id="parallel_cachedit_tp2_vae2", - marks=PARALLEL_MARKS, + marks=PARALLEL_MARKS + _FULL_MARKS, ), ] @@ -82,9 +99,9 @@ def _get_diffusion_feature_cases(model: str): ) def test_hunyuan_video_15_t2v( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): - """L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100.""" + """Diffusion feature coverage for HunyuanVideo-1.5-T2V on H100.""" form_data = { "prompt": PROMPT, "negative_prompt": NEGATIVE_PROMPT, @@ -101,4 +118,4 @@ def test_hunyuan_video_15_t2v( "form_data": form_data, } - openai_client.send_video_diffusion_request(request_config) + online_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_indextts2_expansion.py b/tests/e2e/online_serving/test_indextts2_expansion.py index d971ef166..da9bb6dd8 100644 --- a/tests/e2e/online_serving/test_indextts2_expansion.py +++ b/tests/e2e/online_serving/test_indextts2_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E online serving tests for IndexTTS2 via /v1/audio/speech endpoint. Two-stage pipeline: GPT AR → S2Mel + BigVGAN. Output is 22050 Hz mono WAV. @@ -15,7 +15,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path @@ -23,7 +23,7 @@ pytestmark = [pytest.mark.slow, pytest.mark.tts] MODEL = "IndexTeam/IndexTTS-2" -REF_AUDIO_URL = load_test_audio_data_url("indextts2/ref_audio.wav") +REF_AUDIO_URL = get_asset_path("indextts2/ref_audio.wav", as_data_url=True) tts_server_params = [ pytest.param( @@ -39,7 +39,7 @@ tts_server_params = [ @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_basic_english(omni_server, openai_client) -> None: +def test_basic_english(omni_server, online_client) -> None: """ Basic English TTS: text + ref_audio, no emotion. Deploy Setting: default yaml @@ -55,12 +55,12 @@ def test_basic_english(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "min_audio_bytes": 1024, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_streaming(omni_server, openai_client) -> None: +def test_streaming(omni_server, online_client) -> None: """ stream=True compatibility: IndexTTS2 itself is non-async-chunking, but the OpenAI speech endpoint should still accept the streaming request path and @@ -80,4 +80,4 @@ def test_streaming(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "min_audio_bytes": 1024, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_krea2_expansion.py b/tests/e2e/online_serving/test_krea2_expansion.py index f76a567f4..eb1d6d078 100644 --- a/tests/e2e/online_serving/test_krea2_expansion.py +++ b/tests/e2e/online_serving/test_krea2_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """L4 online-serving tests for the Krea 2 text-to-image diffusion pipeline. @@ -29,7 +29,7 @@ import requests from PIL import Image from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data pytestmark = [pytest.mark.diffusion, pytest.mark.slow] @@ -76,7 +76,7 @@ def _get_diffusion_feature_cases(model: str): _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_krea2(omni_server: OmniServer, online_client: OnlineOmniClient): """Serve Krea 2 and validate the OpenAI-compatible image-generation route. guidance_scale resolution is checkpoint-aware inside the pipeline (distilled -> no-CFG, @@ -96,7 +96,7 @@ def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler): "seed": 42, }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) def _basic_lora_payload(lora: dict | None = None) -> dict: diff --git a/tests/e2e/online_serving/test_lance.py b/tests/e2e/online_serving/test_lance.py index 9c1a40be9..54e782e10 100644 --- a/tests/e2e/online_serving/test_lance.py +++ b/tests/e2e/online_serving/test_lance.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end online serving test for Lance (single-stage). Verifies that the Lance pipeline can serve text-to-image and image-edit @@ -100,7 +100,7 @@ def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]: @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_lance_text2img_online(omni_server, openai_client) -> None: +def test_lance_text2img_online(omni_server, online_client) -> None: """Lance text2img via the OpenAI-compatible chat completions API.""" request_config = { "model": omni_server.model, @@ -115,7 +115,7 @@ def test_lance_text2img_online(omni_server, openai_client) -> None: }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.core_model @@ -123,7 +123,7 @@ def test_lance_text2img_online(omni_server, openai_client) -> None: @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_lance_img2img_online(omni_server, openai_client) -> None: +def test_lance_img2img_online(omni_server, online_client) -> None: """Lance image_edit via the OpenAI-compatible chat completions API.""" input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") buffer = BytesIO() @@ -141,4 +141,4 @@ def test_lance_img2img_online(omni_server, openai_client) -> None: }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_lingbot_video_moe.py b/tests/e2e/online_serving/test_lingbot_video_moe.py index 777880732..7e4d16568 100644 --- a/tests/e2e/online_serving/test_lingbot_video_moe.py +++ b/tests/e2e/online_serving/test_lingbot_video_moe.py @@ -63,7 +63,7 @@ def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientH ) -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True) def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: diff --git a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py index 7c7e0788c..af6bbac50 100644 --- a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Recommended tests of diffusion features that are available in online serving mode and are supported by the following model: @@ -16,7 +19,7 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data pytestmark = [pytest.mark.diffusion, pytest.mark.slow] @@ -66,7 +69,7 @@ def _get_diffusion_feature_cases(model: str): _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_longcat_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_longcat_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient): """Test the recommended feature combinations for LongCat-Image-Edit.""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) @@ -84,4 +87,4 @@ def test_longcat_image_edit(omni_server: OmniServer, openai_client: OpenAIClient }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_longcat_image_expansion.py b/tests/e2e/online_serving/test_longcat_image_expansion.py index 5ed7bb156..5616708c1 100644 --- a/tests/e2e/online_serving/test_longcat_image_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Recommended tests of diffusion features that are available in online serving mode and are supported by the following model: @@ -15,7 +18,7 @@ This validates: import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data pytestmark = [pytest.mark.diffusion, pytest.mark.slow] @@ -66,7 +69,7 @@ def _get_diffusion_feature_cases(model: str): _get_diffusion_feature_cases(MODEL), indirect=True, ) -def test_longcat_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_longcat_image(omni_server: OmniServer, online_client: OnlineOmniClient): """Test the recommended feature combinations for LongCat-Image.""" messages = dummy_messages_from_mix_data(content_text=TEXT_TO_IMAGE_PROMPT) @@ -83,4 +86,4 @@ def test_longcat_image(omni_server: OmniServer, openai_client: OpenAIClientHandl }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_ltx.py b/tests/e2e/online_serving/test_ltx.py new file mode 100644 index 000000000..e4d8f51a8 --- /dev/null +++ b/tests/e2e/online_serving/test_ltx.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Online-serving smoke tests for the public LTX-2 pipeline entries.""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] + +DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers") +BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers") +PROMPT = "A serene lake at sunset with mountains in the background." +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) +HSDP_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) + + +def _server( + model: str, + model_class_name: str, + *extra_args: str, +) -> OmniServerParams: + return OmniServerParams( + model=model, + server_args=[ + "--enforce-eager", + *extra_args, + "--model-class-name", + model_class_name, + ], + ) + + +def _cases(): + return [ + pytest.param( + _server(BASE_MODEL, "LTX2Pipeline"), + 3, + id="one_stage", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + _server(DISTILLED_MODEL, "LTX2DistilledPipeline"), + 8, + id="distilled", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + _server( + DISTILLED_MODEL, + "LTX2DistilledPipeline", + "--use-hsdp", + "--hsdp-shard-size", + "2", + ), + 8, + id="distilled_hsdp2", + marks=HSDP_MARKS, + ), + pytest.param( + _server( + BASE_MODEL, + "LTX2TwoStagePipeline", + "--enable-layerwise-offload", + ), + 3, + id="two_stage", + marks=SINGLE_CARD_MARKS, + ), + ] + + +@pytest.mark.parametrize( + ("omni_server", "num_inference_steps"), + _cases(), + indirect=["omni_server"], +) +def test_ltx_pipeline_entries( + omni_server: OmniServer, + num_inference_steps: int, + openai_client: OpenAIClientHandler, + subtests, +): + for task in ("t2v", "i2v"): + with subtests.test(task=task): + form_data = { + "prompt": PROMPT, + "model": omni_server.model, + "height": 128, + "width": 128, + "num_frames": 9, + "fps": 8, + "num_inference_steps": num_inference_steps, + "seed": 42, + } + request_config = { + "model": omni_server.model, + "form_data": form_data, + } + if task == "i2v": + request_config["image_reference"] = ( + f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + ) + + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_ltx25.py b/tests/e2e/online_serving/test_ltx25.py new file mode 100644 index 000000000..4975abcc0 --- /dev/null +++ b/tests/e2e/online_serving/test_ltx25.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Online serving smokes for all four canonical LTX-2.5 pipelines.""" + +import os + +import pytest + +from tests.helpers import skip_if_gated_repo_inaccessible +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +DEFAULT_MODEL = "Lightricks/LTX-2.5-Diffusers" +DEFAULT_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d" +MODEL = os.getenv("VLLM_TEST_LTX25_MODEL", DEFAULT_MODEL) +MODEL_REVISION = os.getenv("VLLM_TEST_LTX25_MODEL_REVISION", DEFAULT_REVISION if MODEL == DEFAULT_MODEL else "") +PROMPT = "A red fox walks through a snowy forest while the camera remains fixed." + +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + + +@pytest.fixture(scope="module", autouse=True) +def require_ltx25_model_access() -> None: + if not os.path.isdir(MODEL): + skip_if_gated_repo_inaccessible( + MODEL, + revision=MODEL_REVISION or None, + filename="model_index.json", + ) + + +def _server(model_class_name: str) -> OmniServerParams: + return OmniServerParams( + model=MODEL, + server_args=[ + *(["--revision", MODEL_REVISION] if MODEL_REVISION else []), + "--model-class-name", + model_class_name, + "--enforce-eager", + "--enable-layerwise-offload", + "--diffusion-attention-backend", + "CUDNN_ATTN", + ], + ) + + +def _cases(): + return [ + pytest.param(_server("LTX2Pipeline"), 30, id="full_one_stage", marks=SINGLE_CARD_MARKS), + pytest.param(_server("LTX2TwoStagePipeline"), 30, id="full_two_stage", marks=SINGLE_CARD_MARKS), + pytest.param(_server("LTX2DistilledOneStagePipeline"), 8, id="distilled_one_stage", marks=SINGLE_CARD_MARKS), + pytest.param(_server("LTX2DistilledTwoStagePipeline"), 8, id="distilled_two_stage", marks=SINGLE_CARD_MARKS), + ] + + +@pytest.mark.parametrize(("omni_server", "num_inference_steps"), _cases(), indirect=["omni_server"]) +def test_ltx25_pipeline_entries( + omni_server: OmniServer, + num_inference_steps: int, + openai_client: OpenAIClientHandler, + subtests, +) -> None: + """Generate T2V and first-frame I2V through the synchronous video API.""" + for task in ("t2v", "i2v"): + with subtests.test(task=task): + request_config = { + "model": omni_server.model, + "form_data": { + "model": omni_server.model, + "prompt": PROMPT, + "height": 256, + "width": 256, + "num_frames": 9, + "fps": 24, + "num_inference_steps": num_inference_steps, + "seed": 42, + }, + } + if task == "i2v": + request_config["image_reference"] = ( + f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + ) + + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py deleted file mode 100644 index 4d0eb6c90..000000000 --- a/tests/e2e/online_serving/test_ltx2_expansion.py +++ /dev/null @@ -1,91 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Expansion tests for LTX-2 two-stage pipelines. - -Coverage: -- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video) -""" - -import os - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler - -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] - -MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled") -PROMPT = "A serene lake at sunset with mountains in the background." -PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) -HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"] - - -def _cases(): - cases = [] - - cases.append( - pytest.param( - OmniServerParams( - model=MODEL, - server_args=[ - *HSDP_ARGS, - "--model-class-name", - "LTX2DistilledPipeline", - ], - ), - False, - id="t2v_hsdp", - marks=PARALLEL_MARKS, - ) - ) - - cases.append( - pytest.param( - OmniServerParams( - model=MODEL, - server_args=[ - *HSDP_ARGS, - "--model-class-name", - "LTX2DistilledPipeline", - ], - ), - True, - id="i2v_hsdp", - marks=PARALLEL_MARKS, - ) - ) - - return cases - - -@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"]) -def test_ltx2_two_stage_hsdp( - omni_server: OmniServer, - is_i2v: bool, - openai_client: OpenAIClientHandler, -): - # Keep CI small while exercising the fixed 8+3-step distilled recipe. - # Height and width describe the final output; Stage 1 runs at half size. - form_data = { - "prompt": PROMPT, - "model": omni_server.model, - "height": 128, - "width": 128, - "num_frames": 9, - "fps": 8, - "num_inference_steps": 8, - "seed": 42, - } - - request_config = { - "model": omni_server.model, - "form_data": form_data, - } - - if is_i2v: - request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" - - openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_mimo_audio_expansion.py b/tests/e2e/online_serving/test_mimo_audio_expansion.py index e5722ddb8..51f4d7ed9 100644 --- a/tests/e2e/online_serving/test_mimo_audio_expansion.py +++ b/tests/e2e/online_serving/test_mimo_audio_expansion.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ -E2E expansion tests for MiMo-Audio online serving (nightly Omni · Function Test with H100). +E2E expansion tests for MiMo-Audio online serving (nightly Omni · Function Test with H100 · Single-GPU). """ import os @@ -77,7 +77,7 @@ def get_max_batch_size(size_type="few"): @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_001(omni_server, openai_client) -> None: +def test_text_to_text_001(omni_server, online_client) -> None: """ Test text input processing and text-only output generation via OpenAI API. Deploy Setting: default yaml @@ -95,13 +95,13 @@ def test_text_to_text_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.skip(reason="CI failed 8571") -def test_audio_to_text_audio_001(omni_server, openai_client) -> None: +def test_audio_to_text_audio_001(omni_server, online_client) -> None: """ Test audio and text input processing and text/audio output generation via OpenAI API. Deploy Setting: default yaml @@ -126,4 +126,4 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None: }, } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) diff --git a/tests/e2e/online_serving/test_ming_flash_omni_expansion.py b/tests/e2e/online_serving/test_ming_flash_omni_expansion.py index 8ab327022..97a634906 100644 --- a/tests/e2e/online_serving/test_ming_flash_omni_expansion.py +++ b/tests/e2e/online_serving/test_ming_flash_omni_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E online serving tests for Ming-flash-omni-2.0 model (Thinker stage). Tests multimodal understanding via OpenAI-compatible API. @@ -20,7 +20,7 @@ from tests.helpers.stage_config import get_deploy_config_path os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -pytestmark = [pytest.mark.omni, pytest.mark.full_model] +pytestmark = [pytest.mark.omni, pytest.mark.slow] _SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip( reason="Requires 4x H100 GPUs; skipped in CI for now.", @@ -70,7 +70,7 @@ def get_max_batch_size(size_type="few"): @_SKIP_NEED_4_H100_NOT_CI @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_001(omni_server, openai_client) -> None: +def test_text_to_text_001(omni_server, online_client) -> None: """ Input Modal: text Output Modal: text @@ -90,13 +90,13 @@ def test_text_to_text_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @_SKIP_NEED_4_H100_NOT_CI @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_stream_001(omni_server, openai_client) -> None: +def test_text_to_text_stream_001(omni_server, online_client) -> None: """ Input Modal: text Output Modal: text @@ -116,13 +116,13 @@ def test_text_to_text_stream_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @_SKIP_NEED_4_H100_NOT_CI @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_image_to_text_001(omni_server, openai_client) -> None: +def test_image_to_text_001(omni_server, online_client) -> None: """ Input Modal: image + text Output Modal: text @@ -143,13 +143,13 @@ def test_image_to_text_001(omni_server, openai_client) -> None: "modalities": ["text"], } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @_SKIP_NEED_4_H100_NOT_CI @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_audio_to_text_001(omni_server, openai_client) -> None: +def test_audio_to_text_001(omni_server, online_client) -> None: """ Input Modal: audio + text Output Modal: text @@ -170,13 +170,13 @@ def test_audio_to_text_001(omni_server, openai_client) -> None: "modalities": ["text"], } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @_SKIP_NEED_4_H100_NOT_CI @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_video_to_text_001(omni_server, openai_client) -> None: +def test_video_to_text_001(omni_server, online_client) -> None: """ Input Modal: video + text Output Modal: text @@ -197,13 +197,13 @@ def test_video_to_text_001(omni_server, openai_client) -> None: "modalities": ["text"], } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @_SKIP_NEED_4_H100_NOT_CI @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_mix_to_text_001(omni_server, openai_client) -> None: +def test_mix_to_text_001(omni_server, online_client) -> None: """ Input Modal: text + audio + image + video Output Modal: text @@ -228,4 +228,4 @@ def test_mix_to_text_001(omni_server, openai_client) -> None: "modalities": ["text"], } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) diff --git a/tests/e2e/online_serving/test_ming_tts_expansion.py b/tests/e2e/online_serving/test_ming_tts_expansion.py index d4fbb7026..4352a8966 100644 --- a/tests/e2e/online_serving/test_ming_tts_expansion.py +++ b/tests/e2e/online_serving/test_ming_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E expansion tests for Ming-omni-tts online serving (nightly CI). @@ -18,7 +18,6 @@ from tests.helpers.stage_config import get_deploy_config_path pytestmark = [ pytest.mark.slow, pytest.mark.tts, - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4704"), ] MODEL = "inclusionAI/Ming-omni-tts-0.5B" @@ -57,7 +56,7 @@ def get_prompt(prompt_type="zh"): @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", no_async_chunk_params, indirect=True) -def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None: +def test_text_to_audio_non_streaming_001(omni_server, online_client) -> None: """ Deploy Setting: ming_tts.yaml with --no-async-chunk Input Modal: text @@ -72,12 +71,12 @@ def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None: "response_format": "wav", "timeout": 300.0, } - openai_client.send_audio_speech_request(request_config, request_num=2) + online_client.send_audio_speech_request(request_config, request_num=2) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", async_chunk_params, indirect=True) -def test_text_to_audio_streaming_001(omni_server, openai_client) -> None: +def test_text_to_audio_streaming_001(omni_server, online_client) -> None: """ Deploy Setting: ming_tts.yaml (async_chunk=true) Input Modal: text + voice @@ -94,4 +93,4 @@ def test_text_to_audio_streaming_001(omni_server, openai_client) -> None: "response_format": "wav", "timeout": 300.0, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py index 6953ef291..cb8995206 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5.py @@ -1,12 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output. Exercises async chunk streaming (``--async-chunk``) across separate Thinker, Talker, and Code2Wav stages. """ +import json import os +import tempfile +from copy import deepcopy +from pathlib import Path import pytest +from vllm.logger import DEFAULT_LOGGING_CONFIG from tests.helpers.mark import hardware_test from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video @@ -19,10 +27,36 @@ _MODEL = "openbmb/MiniCPM-o-4_5" _CI_DEPLOY = modify_stage_config( get_deploy_config_path("minicpmo_4_5.yaml"), updates={ - "stages": {0: {"default_sampling_params.max_tokens": 64}, 1: {"default_sampling_params.max_tokens": 1024}} + "stages": { + 0: {"default_sampling_params.max_tokens": 64}, + 1: { + "default_sampling_params.max_tokens": 1024, + # Content-consistency assertions must not depend on a random + # codec trajectory: different valid TTS samples can append an + # audible tail even when Thinker text is identical. + "default_sampling_params.temperature": 0.0, + }, + }, }, ) +_PROMPT_LOG_TEMP_DIR = tempfile.TemporaryDirectory(prefix="minicpmo45-prompt-selection-") +_PROMPT_SELECTION_LOG = Path(_PROMPT_LOG_TEMP_DIR.name) / "prompt-selection.log" +_PROMPT_LOGGING_CONFIG_PATH = Path(_PROMPT_LOG_TEMP_DIR.name) / "logging.json" +_PROMPT_LOGGING_CONFIG = deepcopy(DEFAULT_LOGGING_CONFIG) +_PROMPT_LOGGING_CONFIG["handlers"]["minicpmo_prompt_selection"] = { + "class": "logging.FileHandler", + "filename": str(_PROMPT_SELECTION_LOG), + "formatter": "vllm", + "level": "DEBUG", +} +_PROMPT_LOGGING_CONFIG["loggers"]["vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav"] = { + "handlers": ["minicpmo_prompt_selection"], + "level": "DEBUG", + "propagate": True, +} +_PROMPT_LOGGING_CONFIG_PATH.write_text(json.dumps(_PROMPT_LOGGING_CONFIG), encoding="utf-8") + test_params = [ pytest.param( OmniServerParams( @@ -30,6 +64,10 @@ test_params = [ stage_config_path=_CI_DEPLOY, use_stage_cli=False, server_args=["--trust-remote-code", "--async-chunk"], + env_dict={ + "VLLM_CONFIGURE_LOGGING": "1", + "VLLM_LOGGING_CONFIG_PATH": str(_PROMPT_LOGGING_CONFIG_PATH), + }, ), id="async_chunk", ), @@ -65,6 +103,27 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) +# Close and emit <|tts_bos|> so Talker speaks the answer, not reasoning. +_TTS_EXTRA_BODY = { + "chat_template_kwargs": { + "use_tts_template": True, + "enable_thinking": False, + } +} + + +def _prompt_selection_log_offset() -> int: + return _PROMPT_SELECTION_LOG.stat().st_size if _PROMPT_SELECTION_LOG.exists() else 0 + + +def _read_prompt_selection_log(offset: int) -> str: + if not _PROMPT_SELECTION_LOG.exists(): + return "" + with open(_PROMPT_SELECTION_LOG, "rb") as log_file: + log_file.seek(offset) + return log_file.read().decode("utf-8") + + @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.omni @@ -111,11 +170,57 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: "messages": messages, "stream": True, "key_words": {"audio": ["Beijing"]}, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config) +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_audio_with_reference_audio(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "messages": dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()), + "stream": True, + "modalities": ["text", "audio"], + "extra_body": { + **_TTS_EXTRA_BODY, + "ref_audio": f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}", + }, + } + + log_offset = _prompt_selection_log_offset() + responses = openai_client.send_omni_request(request_config) + prompt_log = _read_prompt_selection_log(log_offset) + + assert responses[0].success + assert responses[0].audio_bytes + assert "prompt_cache_id=runtime-ref-" in prompt_log + assert "minicpmo45_ref_" in prompt_log + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_audio_with_default_reference(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "messages": dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()), + "stream": True, + "modalities": ["text", "audio"], + "extra_body": _TTS_EXTRA_BODY, + } + + responses = openai_client.send_omni_request(request_config) + + assert responses[0].success + assert responses[0].audio_bytes + + @pytest.mark.full_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @@ -141,12 +246,7 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None: "stream": True, "key_words": {"text": ["Beijing"]}, "modalities": ["text", "audio"], - "extra_body": { - "chat_template_kwargs": { - "use_tts_template": True, - "enable_thinking": False, - } - }, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @@ -175,6 +275,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages, "stream": True, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @@ -203,6 +304,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages, "stream": True, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @@ -236,6 +338,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages, "stream": True, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py index d63811022..bbdf85c09 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """CI coverage for the MiniCPM-o 4.5 native-duplex Realtime API.""" from __future__ import annotations @@ -13,6 +16,7 @@ import websockets from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( SERVER_PARAMS, demo_args, + duplex_camera_frames, multi_session_args, realtime_url, resolve_ref_audio, @@ -27,10 +31,16 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import ) from tests.helpers.mark import hardware_test from vllm_omni.experimental.fullduplex.client import build_realtime_url +from vllm_omni.experimental.fullduplex.video_stacking import concat_frames_b64 pytestmark = pytest.mark.omni +def _assert_positive_int(value: object) -> None: + assert isinstance(value, int) + assert value > 0 + + def _assert_request_metrics(metrics: object, *, expected_count: int) -> None: assert isinstance(metrics, list) assert len(metrics) == expected_count @@ -107,8 +117,8 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[ @pytest.mark.core_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) -def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None: - ref_audio = resolve_ref_audio(model_prefix) +def test_duplex_websocket_protocol_smoke(omni_server) -> None: + ref_audio = resolve_ref_audio() events = asyncio.run( _run_protocol_smoke( url=realtime_url(omni_server), @@ -125,11 +135,11 @@ def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None @pytest.mark.advanced_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) -def test_duplex_single_session_response_required(omni_server, model_prefix: str, tmp_path: Path) -> None: +def test_duplex_single_session_response_required(omni_server, tmp_path: Path) -> None: args = demo_args( omni_server=omni_server, input_wav=validated_input_wav(), - ref_audio=resolve_ref_audio(model_prefix), + ref_audio=resolve_ref_audio(), output_dir=tmp_path / "single_session", ) args.turns = 2 @@ -139,7 +149,7 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str, args.turn_duration_ms = [args.first_turn_ms] * args.turns result = asyncio.run(run_demo(args)) assert result["ok"] is True - assert result["audio_delta_count"] > 0 + _assert_positive_int(result["audio_delta_count"]) assert result["done_count"] == 2 assert result["error_count"] == 0 assert result["all_audio_responses_have_transcript"] is True @@ -151,13 +161,51 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str, @pytest.mark.advanced_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) -def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str, tmp_path: Path) -> None: +def test_duplex_single_session_video_input(omni_server, tmp_path: Path) -> None: + """Audio plus a 1 fps camera track, the omni-duplex video contract. + + Frames ride the audio appends, so the response contract is unchanged: what + this covers is that interleaved vision input keeps the turn intact instead + of stalling or erroring out mid-segment. + """ + args = demo_args( + omni_server=omni_server, + input_wav=validated_input_wav(), + ref_audio=resolve_ref_audio(), + output_dir=tmp_path / "video_input", + ) + args.turns = 2 + args.turn_duration_ms = [args.first_turn_ms] * args.turns + frames = duplex_camera_frames(seconds=4, cache_dir=tmp_path / "camera") + args.video_frames_b64 = frames + # Each unit also carries a composite of its interior sub-frames, so the + # append exercises the official two-image frame_list that carries motion. + args.video_stacked_frames_b64 = [concat_frames_b64([frames[index]] * 2) for index in range(len(frames))] + + result = asyncio.run(run_demo(args)) + + assert result["ok"] is True + assert result["video_frame_count"] == 4 + assert result["video_stacked_frame_count"] == 4 + _assert_positive_int(result["audio_delta_count"]) + assert result["done_count"] == 2 + assert result["error_count"] == 0 + assert result["all_audio_responses_have_transcript"] is True + assert result["transcript_delta_done_ok"] is True + _assert_request_metrics(result["request_metrics"], expected_count=2) + _assert_session_metrics(result["session_metrics"], expected_count=2) + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_two_sessions_resume_and_takeover(omni_server, tmp_path: Path) -> None: result = asyncio.run( run_multi_session( multi_session_args( omni_server=omni_server, input_wav=validated_input_wav(), - ref_audio=resolve_ref_audio(model_prefix), + ref_audio=resolve_ref_audio(), output_dir=tmp_path / "multi_session", response_required=True, ) diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py index 68204fe28..5b35c4b53 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Nightly lifecycle coverage for the MiniCPM-o 4.5 native-duplex API.""" from __future__ import annotations @@ -12,6 +15,7 @@ import pytest from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( SERVER_PARAMS, SOFT_INTERRUPT_SHA256, + deploy_max_sessions, multi_session_args, realtime_url, resolve_ref_audio, @@ -21,6 +25,7 @@ from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import ( from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import ( run_lifecycle_probes, ) +from tests.e2e.online_serving.run_minicpmo_realtime_duplex_server_vad import run_server_vad_interrupt from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import ( run_soft_interrupt, ) @@ -31,11 +36,11 @@ pytestmark = [pytest.mark.full_model, pytest.mark.omni] @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) -def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_path: Path) -> None: +def test_duplex_admission_and_expiry_reaper(omni_server, tmp_path: Path) -> None: args = multi_session_args( omni_server=omni_server, input_wav=validated_input_wav(), - ref_audio=resolve_ref_audio(model_prefix), + ref_audio=resolve_ref_audio(), output_dir=tmp_path / "admission_expiry", response_required=False, ) @@ -43,7 +48,7 @@ def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_ args.disconnect_session_index = None args.takeover_session_index = None args.expire_session_index = 0 - args.verify_admission_limit = 2 + args.verify_admission_limit = deploy_max_sessions() result = asyncio.run(run_lifecycle_probes(args)) assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2) @@ -59,7 +64,7 @@ def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_ # so it never reaches a mid-stream decision point. @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) -def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -> None: +def test_duplex_soft_interrupt(omni_server, tmp_path: Path) -> None: input_wav = validated_soft_interrupt_wav() result = asyncio.run( run_soft_interrupt( @@ -67,7 +72,7 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) - url=realtime_url(omni_server), model=omni_server.model, input_wav=str(input_wav), - ref_audio=str(resolve_ref_audio(model_prefix)), + ref_audio=str(resolve_ref_audio()), output_dir=str(tmp_path / "soft_interrupt"), summary_output=None, chunk_ms=200, @@ -88,3 +93,22 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) - assert result["response_lifecycle_ok"] is True assert result["response_audio_contract_ok"] is True assert result["followup_response_transcript_ok"] is True + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True) +def test_duplex_server_vad_hard_interrupt(omni_server) -> None: + result = asyncio.run( + run_server_vad_interrupt( + SimpleNamespace( + url=realtime_url(omni_server), + model=omni_server.model, + input_wav=str(validated_input_wav()), + interrupt_wav=str(validated_soft_interrupt_wav()), + ref_audio=str(resolve_ref_audio()), + chunk_ms=200, + timeout_s=180.0, + ) + ) + ) + assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2) diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py index e84985e3e..e7fa50c16 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py @@ -1,7 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """E2E online expansion tests for MiniCPM-o 4.5. These cover modality combinations, separate Code2Wav behavior, sequential -request isolation, longer audio output, and Chinese speech. +request isolation, longer audio output, long-form generation, and Chinese speech. Like Qwen3-Omni expansion, ``default`` exercises ``--no-async-chunk`` and ``async_chunk`` exercises ``--async-chunk``. @@ -75,6 +78,15 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) +# Close and emit <|tts_bos|> so Talker speaks the answer, not reasoning. +_TTS_EXTRA_BODY = { + "chat_template_kwargs": { + "use_tts_template": True, + "enable_thinking": False, + } +} + + @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_video_to_text_001(omni_server, openai_client) -> None: @@ -127,6 +139,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages_1, "stream": True, + "extra_body": _TTS_EXTRA_BODY, }, request_num=1, ) @@ -138,6 +151,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None: "messages": messages_2, "stream": True, "key_words": {"text": ["Beijing"]}, + "extra_body": _TTS_EXTRA_BODY, }, request_num=1, ) @@ -164,20 +178,36 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: "messages": messages, "modalities": ["text", "audio"], "stream": True, - # Delimit the assistant answer with the TTS template so the Talker - # does not spend its codec-token budget speaking hidden reasoning. - "extra_body": { - "chat_template_kwargs": { - "use_tts_template": True, - "enable_thinking": False, - } - }, + "extra_body": _TTS_EXTRA_BODY, "key_words": {"audio": ["Beijing"]}, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_audio_long_form_001(omni_server, openai_client) -> None: + """ + Input Modal: text only (long-form generation prompt). + Output Modal: text, audio (default ``modalities``); + Input Setting: stream=True + Datasets: single request + """ + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + content_text="帮我讲一个100字的故事.", + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "extra_body": _TTS_EXTRA_BODY, + } + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + + @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_chinese_text_to_audio(omni_server, openai_client) -> None: @@ -193,5 +223,6 @@ def test_chinese_text_to_audio(omni_server, openai_client) -> None: "messages": messages, "stream": True, "key_words": {"text": ["北京"]}, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config) diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index b2e40c58f..9569078d2 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 import importlib.util @@ -15,6 +18,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] DRIVER_DIR = Path(__file__).resolve().parent HELPERS_DIR = DRIVER_DIR / "helpers" DEMO_PATH = HELPERS_DIR / "minicpmo_realtime_duplex_scenarios.py" +EXAMPLE_DEMO_PATH = DRIVER_DIR.parents[2] / "examples/online_serving/minicpmo/realtime_duplex_demo.py" MULTI_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_multi_session.py" PAIR_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_demo_pair.py" SOFT_INTERRUPT_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_soft_interrupt.py" @@ -47,6 +51,18 @@ def _load_pair_demo_module(): return module +def _load_example_demo_module(): + spec = importlib.util.spec_from_file_location( + "minicpmo_realtime_duplex_example_demo_test", + EXAMPLE_DEMO_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + def _load_soft_interrupt_demo_module(): spec = importlib.util.spec_from_file_location( "minicpmo_realtime_duplex_soft_interrupt_test", @@ -524,6 +540,111 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp assert summary["response_summaries"][1]["transcript"] == "一加一等于二。" +def test_realtime_duplex_soft_interrupt_accepts_followup_done_after_commit(tmp_path): + """A follow-up response can drain a residual model unit past the final commit. + + Sequence: ``r1.done -> listen -> r2.created -> ... -> commit -> r2.done -> listen``. + + The globally-last ``response.done`` (r2) lands *after* the commit, so anchoring + ``listen_after_response_before_commit`` on it leaves the ``(last_done, commit)`` + interval empty and fails the contract even though the turn yielded the floor + correctly. The check must anchor on the last done strictly before the commit + (r1.done), with the listen at 3.1 s satisfying the sandwich. + """ + demo = _load_soft_interrupt_demo_module() + output = tmp_path / "soft_interrupt_residual_drain" + output.mkdir() + first_response_id = "resp-first" + second_response_id = "resp-second" + events = [ + {"type": "response.listen", "_client_received_at_s": 1.0}, + {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, + {"type": "response.speak", "response_id": first_response_id, "_client_received_at_s": 2.0}, + { + "type": "response.audio.delta", + "response_id": first_response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.audio_transcript.delta", + "response_id": first_response_id, + "delta": "中国古代四大发明", + "_client_received_at_s": 2.1, + }, + { + "type": "response.audio.delta", + "response_id": first_response_id, + "delta": "AAAA", + "_client_received_at_s": 2.6, + }, + { + "type": "response.audio_transcript.delta", + "response_id": first_response_id, + "delta": "是造纸术。", + "_client_received_at_s": 2.6, + }, + {"type": "response.done", "response_id": first_response_id, "_client_received_at_s": 3.0}, + {"type": "response.listen", "_client_received_at_s": 3.1}, + {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 3.5}, + {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 3.5}, + { + "type": "response.audio.delta", + "response_id": second_response_id, + "delta": "AAAA", + "_client_received_at_s": 3.6, + }, + { + "type": "response.audio_transcript.delta", + "response_id": second_response_id, + "delta": "一加一等于", + "_client_received_at_s": 3.6, + }, + { + "type": "response.audio.delta", + "response_id": second_response_id, + "delta": "AAAA", + "_client_received_at_s": 3.7, + }, + { + "type": "response.audio_transcript.delta", + "response_id": second_response_id, + "delta": "二。", + "_client_received_at_s": 3.7, + }, + # The final commit lands while the follow-up response is still draining. + {"type": "input_audio_buffer.committed", "_client_received_at_s": 4.0}, + {"type": "response.done", "response_id": second_response_id, "_client_received_at_s": 4.5}, + {"type": "response.listen", "_client_received_at_s": 4.6}, + ] + (output / "events.jsonl").write_text( + "".join(demo.json.dumps(event) + "\n" for event in events), + encoding="utf-8", + ) + (output / "result.json").write_text( + demo.json.dumps({"ok": True, "response_ids": [first_response_id, second_response_id]}), + encoding="utf-8", + ) + + summary = demo.summarize_artifacts( + output_dir=output, + validation_mode="response-required", + min_responses=2, + min_audio_deltas_per_response=2, + expect_followup_response_substring="一加一等于二", + ) + + assert summary["ok"] is True + # Anchored on r1.done (the last done before the commit), not the late r2.done. + assert summary["listen_after_response_before_commit"] is True + # listen_after_last_done still keys off the globally-last done, hence the final listen. + assert summary["listen_after_last_done"] is True + assert summary["final_listen_after_commit"] is True + assert summary["second_response_before_final_commit"] is True + assert summary["followup_response_transcript_ok"] is True + assert summary["response_summaries"][1]["transcript"] == "一加一等于二。" + + def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp_path): demo = _load_soft_interrupt_demo_module() output = tmp_path / "model_policy" @@ -758,6 +879,34 @@ def test_realtime_duplex_multi_session_rejects_cross_session_response_identity() ) +def test_realtime_duplex_multi_session_rejects_overlapping_expected_tokens(monkeypatch, capsys): + demo = _load_multi_demo_module() + monkeypatch.setattr( + demo.sys, + "argv", + [ + "multi.py", + "--sessions", + "2", + "--input-wav", + "fallback.wav", + "--session-input-wav", + "a.wav", + "--session-input-wav", + "b.wav", + "--session-expected-token", + "cat", + "--session-expected-token", + "cater", + ], + ) + + with pytest.raises(SystemExit): + demo.parse_args() + + assert "must not overlap after normalization" in capsys.readouterr().err + + def test_realtime_duplex_multi_session_reads_nested_terminal_identity(): demo = _load_multi_demo_module() @@ -768,6 +917,67 @@ def test_realtime_duplex_multi_session_reads_nested_terminal_identity(): ) +def test_realtime_duplex_synchronized_start_gate_releases_all_participants(): + demo = _load_multi_demo_module() + + async def exercise(): + gate = demo._SynchronizedStartGate(3, timeout_s=0.5) + released = [] + + async def participant(index): + await gate.wait() + released.append(index) + + await asyncio.wait_for( + asyncio.gather(*(participant(index) for index in range(3))), + timeout=1.0, + ) + return released + + assert set(asyncio.run(exercise())) == {0, 1, 2} + + +def test_realtime_duplex_synchronized_start_gate_aborts_waiters_on_failure(monkeypatch): + demo = _load_multi_demo_module() + + async def fake_run_demo(args): + if args.fail: + raise RuntimeError("session creation failed") + await args.start_barrier.wait() + return {"ok": True} + + monkeypatch.setattr(demo, "run_demo", fake_run_demo) + + async def exercise(): + gate = demo._SynchronizedStartGate(2, timeout_s=0.5) + return await asyncio.wait_for( + asyncio.gather( + demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=True)), + demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=False)), + return_exceptions=True, + ), + timeout=1.0, + ) + + failed, aborted = asyncio.run(exercise()) + assert isinstance(failed, RuntimeError) + assert str(failed) == "session creation failed" + assert isinstance(aborted, RuntimeError) + assert "synchronized start aborted" in str(aborted) + assert "session creation failed" in str(aborted) + + +def test_realtime_duplex_synchronized_start_gate_has_bounded_wait(): + demo = _load_multi_demo_module() + + async def exercise(): + gate = demo._SynchronizedStartGate(2, timeout_s=0.01) + await gate.wait() + + with pytest.raises(TimeoutError, match=r"1/2 sessions ready"): + asyncio.run(exercise()) + + def test_realtime_duplex_demo_resolves_distinct_turn_inputs(): demo = _load_demo_module() primary = Path("first.wav") @@ -1721,6 +1931,241 @@ def test_realtime_duplex_demo_waits_at_each_model_unit_and_stops_after_speak(): assert len(ws.messages) == 10 +def _sent_video_frames(messages: list[dict[str, object]]) -> list[list[str]]: + frames: list[list[str]] = [] + for message in messages: + raw = message.get("video_frames") + if not isinstance(raw, list): + continue + frames.append([frame for frame in raw if isinstance(frame, str)]) + return frames + + +def _decode_jpeg_size(frame_b64: str) -> tuple[int, int]: + import io + + from PIL import Image + + return Image.open(io.BytesIO(base64.b64decode(frame_b64))).size + + +class _RecordingWebSocket: + def __init__(self, demo): + self._demo = demo + self.messages: list[dict[str, object]] = [] + + async def send(self, payload): + self.messages.append(self._demo.json.loads(payload)) + + +def _send_seconds_of_audio(demo, *, seconds: int, frames: list[str], stacked: list[str | None] | None = None): + ws = _RecordingWebSocket(demo) + asyncio.run( + demo._send_pcm16( + ws, + b"\x01\x00" * (demo.PCM16_SAMPLE_RATE * seconds), + chunk_ms=200, + realtime_delay=False, + frames_b64=frames, + stacked_frames_b64=stacked, + ) + ) + return ws.messages + + +def test_realtime_duplex_demo_streams_one_video_frame_per_model_unit(): + demo = _load_demo_module() + + messages = _send_seconds_of_audio(demo, seconds=3, frames=["f0", "f1", "f2", "f3"]) + + # Frame k rides the append that closes model unit k, so Stage0 can bind it + # to that unit's audio. 3 s of audio closes units at 1030 ms and 2030 ms and + # leaves a residual too short for a third, so only two frames go out. + assert _sent_video_frames(messages) == [["f0"], ["f1"]] + frame_indices = [index for index, message in enumerate(messages) if "video_frames" in message] + assert frame_indices == [5, 10] + # The carrying appends are exactly the ones that reach a unit boundary. + assert [messages[index]["audio_end_ms"] for index in frame_indices] == [1200, 2200] + + +def test_realtime_duplex_demo_never_sends_a_frame_before_its_unit_can_close(): + """Regression: frame 0 used to ride the very first 200 ms append. + + Stage0 cannot close a unit until 1030 ms of audio has arrived, and it does + not carry frames across appends, so a frame sent that early was silently + dropped and every later frame ended up one unit ahead of its audio. + """ + demo = _load_demo_module() + + messages = _send_seconds_of_audio(demo, seconds=3, frames=["f0", "f1", "f2"]) + + carried = [message for message in messages if "video_frames" in message] + assert carried, "expected at least one frame to be sent" + for message in carried: + assert message["audio_end_ms"] >= demo.duplex_unit_boundary_ms(0) + + +def test_realtime_duplex_demo_holds_the_last_video_frame_when_audio_outlives_the_clip(): + demo = _load_demo_module() + + # 4 s of audio closes three units (1030/2030/3030 ms); the two-frame clip + # holds its last frame for the third. + assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=4, frames=["a", "b"])) == [["a"], ["b"], ["b"]] + # A still image is a one-element clip and therefore repeats every unit. + assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=3, frames=["still"])) == [["still"], ["still"]] + + +def test_realtime_duplex_demo_sends_each_units_stacked_composite_next_to_its_base_frame(): + demo = _load_demo_module() + + messages = _send_seconds_of_audio( + demo, + seconds=4, + frames=["f0", "f1", "f2"], + stacked=["s0", "s1", None], + ) + + # Official pairing is frame_list=[base, composite of that unit's interior], + # so the composite belongs to the same unit as the base beside it -- never + # to the previous one. A unit without interior sub-frames sends base alone. + assert _sent_video_frames(messages) == [["f0", "s0"], ["f1", "s1"], ["f2"]] + + +def test_minicpmo_duplex_camera_fixture_returns_flat_base_frame_track(tmp_path, monkeypatch): + from tests.e2e.online_serving.helpers import minicpmo_4_5_duplex as fixtures + from tests.e2e.online_serving.helpers import minicpmo_realtime_duplex_scenarios as demo + from tests.helpers import media + + video_path = tmp_path / "camera.mp4" + monkeypatch.setattr( + media, + "generate_synthetic_video", + lambda *args, **kwargs: {"file_path": str(video_path)}, + ) + monkeypatch.setattr( + demo, + "_video_frames_from_file", + lambda path: (["frame-0", "frame-1"], [None, None]), + ) + + frames = fixtures.duplex_camera_frames(seconds=2, cache_dir=tmp_path) + + assert frames == ["frame-0", "frame-1"] + + +def test_realtime_duplex_demo_decodes_a_video_file_into_one_frame_per_second(tmp_path): + pytest.importorskip("cv2") + pytest.importorskip("imageio") + from tests.helpers.media import generate_synthetic_video + + demo = _load_demo_module() + # 30 fps synthetic clip: 90 frames is 3 s of camera input. + video = generate_synthetic_video(64, 64, 90, cache_dir=tmp_path) + + frames, stacked = demo._resolve_video_frames(SimpleNamespace(input_video=video["file_path"])) + + assert len(frames) == 3 + assert all(base64.b64decode(frame, validate=True)[:3] == b"\xff\xd8\xff" for frame in frames) + assert len(set(frames)) == 3 + # stack_frames defaults to 1, so there is nothing to stack. + assert stacked == [None, None, None] + + +def test_realtime_duplex_demo_stacks_a_units_interior_subframes_into_one_composite(tmp_path): + pytest.importorskip("cv2") + pytest.importorskip("imageio") + from tests.helpers.media import generate_synthetic_video + + demo = _load_demo_module() + video = generate_synthetic_video(64, 64, 90, cache_dir=tmp_path) + + frames, stacked = demo._resolve_video_frames(SimpleNamespace(input_video=video["file_path"], stack_frames=5)) + + # stack_frames=5 raises the visual refresh rate to 5 fps, but the four + # interior sub-frames of each second collapse into a single composite, so + # the wire still carries 2 images per unit and the audio cadence is + # untouched. + assert len(frames) == 3 + assert len(stacked) == 3 + assert all(frame is not None for frame in stacked) + assert all(base64.b64decode(frame, validate=True)[:3] == b"\xff\xd8\xff" for frame in stacked) + composite = _decode_jpeg_size(stacked[0]) + assert composite != _decode_jpeg_size(frames[0]) + + +def test_example_demo_uses_external_wav_duration_and_skips_video_soundtrack(tmp_path, monkeypatch): + demo = _load_example_demo_module() + wav = tmp_path / "external.wav" + demo.write_pcm16_wav(wav, b"\x00\x00" * (demo.PCM16_SAMPLE_RATE * 2), sample_rate_hz=demo.PCM16_SAMPLE_RATE) + soundtrack_calls: list[object] = [] + + def fake_soundtrack(*args, **kwargs): + soundtrack_calls.append(1) + raise AssertionError("video soundtrack must not be loaded when --input-wav is set") + + def fake_frames(video_path, work_dir, *, fps, max_side, duration_s, stack_frames): + assert video_path == Path("/tmp/silent.mp4") + assert fps == 1.0 + assert duration_s == pytest.approx(2.0) + return ["f0", "f1"], [None, None] + + monkeypatch.setattr(demo, "_extract_video_soundtrack", fake_soundtrack) + monkeypatch.setattr(demo, "_extract_video_frames", fake_frames) + + wav_out, frames, stacked = demo._resolve_duplex_av_inputs( + input_wav=str(wav), + input_video="/tmp/silent.mp4", + work_dir=tmp_path / "video_input", + fps=1.0, + max_side=0, + ) + + assert wav_out == str(wav) + assert frames == ["f0", "f1"] + assert stacked == [None, None] + assert soundtrack_calls == [] + + +def test_example_demo_extracts_video_soundtrack_only_when_wav_is_absent(tmp_path, monkeypatch): + demo = _load_example_demo_module() + extracted = tmp_path / "from_video.wav" + demo.write_pcm16_wav(extracted, b"\x00\x00" * demo.PCM16_SAMPLE_RATE, sample_rate_hz=demo.PCM16_SAMPLE_RATE) + + def fake_extract(video_path, work_dir, *, fps, max_side, stack_frames): + assert video_path == Path("/tmp/clip.mp4") + return extracted, ["v0"], ["s0"] + + monkeypatch.setattr(demo, "_extract_video_input", fake_extract) + + wav_out, frames, stacked = demo._resolve_duplex_av_inputs( + input_wav=None, + input_video="/tmp/clip.mp4", + work_dir=tmp_path / "video_input", + fps=1.0, + max_side=0, + ) + + assert wav_out == str(extracted) + assert frames == ["v0"] + assert stacked == ["s0"] + + +def test_realtime_duplex_demo_prefers_a_video_over_a_still_frame_image(tmp_path): + demo = _load_demo_module() + still = tmp_path / "frame.jpg" + still.write_bytes(b"\xff\xd8\xff-still") + + # A still has no interior sub-frames, so it never stacks. + assert demo._resolve_video_frames(SimpleNamespace(frame_image=str(still))) == ( + [base64.b64encode(still.read_bytes()).decode("ascii")], + [None], + ) + assert demo._resolve_video_frames( + SimpleNamespace(frame_image=str(still), video_frames_b64=["preset"]), + ) == (["preset"], [None]) + assert demo._resolve_video_frames(SimpleNamespace()) == ([], []) + + def test_realtime_duplex_demo_listen_only_overlap_accepts_silence_unit_before_first_done(monkeypatch): demo = _load_demo_module() state = demo.DemoState() diff --git a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py new file mode 100644 index 000000000..148520ff0 --- /dev/null +++ b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""MiniMax-H3 DLO + DP2 online-serving smoke test.""" + +from __future__ import annotations + +import concurrent.futures +import io +import json +import os + +import av +import pytest +import requests + +from tests.helpers.assertions import assert_video_valid +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = os.environ.get("VLLM_TEST_MINIMAX_H3_MODEL", "MiniMaxAI/MiniMax-H3") +WIDTH = 1344 +HEIGHT = 768 +FPS = 24 +NUM_INFERENCE_STEPS = 4 +REQUEST_TIMEOUT_SECONDS = 1800 +H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) + +SERVER_ARGS = [ + "--trust-remote-code", + "--task-type", + "fl2va", + "--num-gpus", + "2", + "--tensor-parallel-size", + "1", + "--data-parallel-size", + "2", + "--request-batch-max-wait-ms", + "500", + "--usp", + "1", + "--ring", + "1", + "--text-encoder-tp-size", + "1", + "--vae-patch-parallel-size", + "1", + "--vae-parallel-mode", + "tile", + "--vae-use-tiling", + "--enable-distributed-layerwise-offload", +] + + +def _assert_audio_stream_present(video: bytes) -> None: + """Assert that the generated MP4 contains decodable audio samples.""" + with av.open(io.BytesIO(video)) as container: + audio_streams = [stream for stream in container.streams if stream.type == "audio"] + assert audio_streams, "MiniMax-H3 MP4 has no audio stream" + audio_frame = next(container.decode(audio=0), None) + assert audio_frame is not None and audio_frame.samples > 0, "MiniMax-H3 MP4 audio stream is empty" + + +def _run_t2va_request(client: OpenAIClientHandler, seed: int) -> bytes: + """Submit one synchronous T2VA request and return its MP4 body.""" + request_data = { + "model": MODEL, + "prompt": "In a snowy blue-purple forest, a traveler walks past a sleeping giant; footsteps crunch in the snow while the creature softly breathes.", + "width": str(WIDTH), + "height": str(HEIGHT), + "fps": str(FPS), + "num_inference_steps": str(NUM_INFERENCE_STEPS), + "flow_shift": "12", + "seed": str(seed), + "extra_params": json.dumps( + { + "task": "t2va", + "duration": 4.0, + "aspect_ratio": "16:9", + "audio_flow_shift": 3.0, + }, + separators=(",", ":"), + ), + } + response = requests.post( + f"{client.base_url.rstrip('/')}/v1/videos/sync", + data=request_data, + headers={"Accept": "video/mp4"}, + timeout=REQUEST_TIMEOUT_SECONDS, + ) + response.raise_for_status() + assert response.headers.get("content-type", "").startswith("video/mp4") + assert response.content, "MiniMax-H3 returned an empty video body" + return response.content + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.slow +@pytest.mark.parametrize( + "omni_server", + [ + pytest.param( + OmniServerParams( + model=MODEL, + server_args=SERVER_ARGS, + stage_init_timeout=1800, + init_timeout=1800, + ), + id="minimax_h3_dlo_dp2_t2va", + marks=H100_TWO_CARD_MARKS, + ) + ], + indirect=True, +) +def test_minimax_h3_dlo_dp2_t2va(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + """Validate one complete DLO all-gather DP2 wave with two T2VA jobs.""" + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(_run_t2va_request, openai_client, seed) for seed in (1101, 1102)] + videos = [future.result() for future in futures] + + for video in videos: + assert_video_valid(video, width=WIDTH, height=HEIGHT, fps=FPS) + _assert_audio_stream_present(video) diff --git a/tests/e2e/online_serving/test_minimax_music3.py b/tests/e2e/online_serving/test_minimax_music3.py new file mode 100644 index 000000000..3d8718e70 --- /dev/null +++ b/tests/e2e/online_serving/test_minimax_music3.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E online tests for MiniMax Music 3 text-to-music. + +This is a music model on the speech endpoint, so the request shape differs +from every other TTS model here: ``input`` carries the lyrics, ``instructions`` +carries the caption that decides genre and arrangement, and both are required. + +Weekly rather than per-PR: the checkpoint is 27 GB across five components and a +single request occupies two decode rows, because classifier-free guidance is +mandatory for this model. Rejection cases live in +``tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py``. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "MiniMaxAI/MiniMax-Music3" +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 600.0 + +# 25 frames per second, so 250 frames is ten seconds and 750 is thirty. The +# byte floors sit below the true payload (32 kHz stereo, 16-bit) to leave room +# for the model ending a song before the cap. +_FRAMES_10S = 250 +_FRAMES_30S = 750 +_MIN_BYTES_10S = 500_000 +_MIN_BYTES_30S = 1_500_000 + +LYRICS = ( + "[Verse]\nWalking down the empty street at midnight\n" + "Streetlights flicker like a broken dream\n" + "[Chorus]\nAnd I keep on walking\nTill the morning finds me" +) + +CAPTION = ( + "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, " + "soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. " + "Intimate bedroom production, gentle tape saturation, no bright cymbals." +) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("minimax_music3.yaml"), + server_args=["--trust-remote-code", "--disable-log-stats"], + ), + id="minimax_music3", + ) +] + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_music_001(omni_server, openai_client) -> None: + """Baseline: lyrics plus caption in, one ten-second song out. + + Deploy Setting: minimax_music3.yaml + Input Modal: text (lyrics) + text (caption) + Output Modal: audio, 32 kHz stereo + Input Setting: stream=False, seed pinned, 250 frames + """ + request_config = { + "model": omni_server.model, + "input": LYRICS, + "instructions": CAPTION, + "seed": 1, + "max_new_tokens": _FRAMES_10S, + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_BYTES_10S, + } + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_music_multi_window_002(omni_server, openai_client) -> None: + """A song long enough to span several acoustic windows. + + The acoustic stage decodes overlapping 200-frame windows on a 100-frame + hop and crops the overlap back off, so a request that fits in one window + never exercises the stitching. 750 frames covers seven of them. + + Deploy Setting: minimax_music3.yaml + Input Setting: stream=False, seed pinned, 750 frames + """ + request_config = { + "model": omni_server.model, + "input": LYRICS, + "instructions": CAPTION, + "seed": 1, + "max_new_tokens": _FRAMES_30S, + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_BYTES_30S, + } + openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py index 2d048fe52..598326a1e 100644 --- a/tests/e2e/online_serving/test_moss_tts_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E online tests for the full MOSS-TTS family via /v1/audio/speech. @@ -30,10 +30,14 @@ from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +_SKIP_ISSUE_6417 = pytest.mark.skip( + reason="https://github.com/vllm-project/vllm-omni/issues/6417", +) + pytestmark = [ pytest.mark.slow, pytest.mark.tts, - pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), + _SKIP_ISSUE_6417, ] MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" @@ -92,7 +96,7 @@ tts_server_params = [ @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> None: +def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None: """ Realtime voice_clone via /v1/audio/speech, non-streaming. Deploy Setting: moss_tts_realtime.yaml @@ -114,12 +118,12 @@ def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> No "min_audio_bytes": _MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_002_streaming(omni_server, openai_client, ref_audio_data_url) -> None: +def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_url) -> None: """ Realtime voice_clone via /v1/audio/speech, streaming PCM. Deploy Setting: moss_tts_realtime.yaml @@ -143,12 +147,12 @@ def test_text_to_audio_002_streaming(omni_server, openai_client, ref_audio_data_ "min_hnr_db": -5.0, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_003_chinese(omni_server, openai_client, ref_audio_data_url) -> None: +def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_url) -> None: """ Realtime voice_clone via /v1/audio/speech, Chinese input. Deploy Setting: moss_tts_realtime.yaml @@ -168,4 +172,4 @@ def test_text_to_audio_003_chinese(omni_server, openai_client, ref_audio_data_ur "min_audio_bytes": _MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py index 3468a9656..3ea8c91ea 100644 --- a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for MOSS-TTS-Nano model with text input and audio output. @@ -86,7 +86,7 @@ tts_server_params = [ @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> None: +def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None: """ Test voice_clone mode (no ref_text) via /v1/audio/speech. Deploy Setting: default yaml @@ -103,12 +103,12 @@ def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> No "ref_audio": ref_audio_data_url, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_002(omni_server, openai_client, ref_audio_data_url) -> None: +def test_text_to_audio_002(omni_server, online_client, ref_audio_data_url) -> None: """ Test streaming voice_clone via /v1/audio/speech. Deploy Setting: default yaml @@ -135,12 +135,12 @@ def test_text_to_audio_002(omni_server, openai_client, ref_audio_data_url) -> No "min_hnr_db": -5.0, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_003(omni_server, openai_client, ref_audio_data_url) -> None: +def test_text_to_audio_003(omni_server, online_client, ref_audio_data_url) -> None: """ Test Chinese voice_clone via /v1/audio/speech. Deploy Setting: default yaml @@ -157,12 +157,12 @@ def test_text_to_audio_003(omni_server, openai_client, ref_audio_data_url) -> No "ref_audio": ref_audio_data_url, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_004_ref_text_ignored(omni_server, openai_client, ref_audio_data_url) -> None: +def test_text_to_audio_004_ref_text_ignored(omni_server, online_client, ref_audio_data_url) -> None: """ Sending ``ref_text`` alongside ``ref_audio`` is accepted but ignored. @@ -180,4 +180,4 @@ def test_text_to_audio_004_ref_text_ignored(omni_server, openai_client, ref_audi "ref_text": REF_AUDIO_TRANSCRIPT, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py b/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py new file mode 100644 index 000000000..acdffd20d --- /dev/null +++ b/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path + +import pytest +from huggingface_hub import snapshot_download + +from tests.e2e.online_serving.nemotron_voicechat_realtime_duplex import parse_args, run +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams, get_model_prefix +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_MODEL", "nvidia/NVIDIA-NemotronLabs-VoiceChat-11B") +DEPLOY_CONFIG = get_deploy_config_path("nemotron_labs_voicechat_duplex.yaml") +TOKENIZER = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_LLM_PATH") + +# The client gives up after CLIENT_TIMEOUT_S. The server-side stage-input +# deadline must fire first, otherwise a stalled stream surfaces as a bare +# client TimeoutError with no server-side error at all. See #6816. +CLIENT_TIMEOUT_S = 300 +SERVER_INPUT_WAIT_TIMEOUT_S = 240 + +_SERVER_ENV: dict[str, str] = {"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S": str(SERVER_INPUT_WAIT_TIMEOUT_S)} +if TOKENIZER: + _SERVER_ENV["NEMOTRON_VOICECHAT_LLM_PATH"] = TOKENIZER + +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize( + "omni_server", + [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=DEPLOY_CONFIG, + env_dict=_SERVER_ENV, + ), + id="native-duplex", + ) + ], + indirect=True, +) +def test_native_duplex_turn_taking_streams_model_audio(omni_server, tmp_path: Path) -> None: + model_prefix = get_model_prefix() + root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL) + if not root.is_dir(): + root = Path(snapshot_download(MODEL, local_files_only=True)) + args = parse_args( + ( + f"--url ws://{omni_server.host}:{omni_server.port}/v1/realtime --model {omni_server.model} " + f"--input-wav {root / 'turn_taking.wav'} --input-channel 0 --max-frames 190 " + f"--minimum-audio-chunks 48 --minimum-audio-rms 0.001 --no-realtime --timeout-s {CLIENT_TIMEOUT_S} " + f"--output-dir {tmp_path / 'native_duplex'}" + ).split() + ) + result = asyncio.run(run(args)) + assert result["ok"] is True + assert result["input_frames"] == 190 + assert result["event_counts"]["response.speak"] > 0 + assert result["audio_bytes"] >= result["input_frames"] * 1764 * 2 // 4 diff --git a/tests/e2e/online_serving/test_nextstep_expansion.py b/tests/e2e/online_serving/test_nextstep_expansion.py index 381079320..3e0bc53b5 100644 --- a/tests/e2e/online_serving/test_nextstep_expansion.py +++ b/tests/e2e/online_serving/test_nextstep_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Online serving E2E for NextStep-1.1 text-to-image (tensor parallel). """ @@ -10,7 +13,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) @@ -66,7 +69,7 @@ def _get_diffusion_feature_cases(model: str): _get_diffusion_feature_cases(model=os.environ.get("VLLM_TEST_NEXTSTEP_MODEL", _DEFAULT_MODEL)), indirect=True, ) -def test_nextstep_11(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_nextstep_11(omni_server: OmniServer, online_client: OnlineOmniClient): messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT) request_config = { "model": omni_server.model, @@ -82,4 +85,4 @@ def test_nextstep_11(omni_server: OmniServer, openai_client: OpenAIClientHandler }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py index ab35b74ea..3a042082e 100644 --- a/tests/e2e/online_serving/test_omnivoice_expansion.py +++ b/tests/e2e/online_serving/test_omnivoice_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for OmniVoice TTS model via /v1/audio/speech endpoint. @@ -14,9 +14,10 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.entrypoints.openai.serving_speech import _DEFAULT_VOICE_NAME try: from transformers import HiggsAudioV2TokenizerModel # noqa: F401 @@ -46,7 +47,7 @@ TEST_PARAMS = [ _DEFAULT_MIN_AUDIO_BYTES = 5000 -REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." @@ -62,7 +63,7 @@ class TestOmniVoiceTTS: """E2E tests for OmniVoice TTS model.""" @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_speech_auto_voice(self, omni_server, openai_client) -> None: + def test_speech_auto_voice(self, omni_server, online_client) -> None: """Test auto voice TTS generation (text only, no reference audio).""" request_config = { "model": omni_server.model, @@ -71,7 +72,7 @@ class TestOmniVoiceTTS: "timeout": 180.0, "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) @@ -79,7 +80,7 @@ class TestOmniVoiceSeed: """E2E tests for OmniVoice Seed Params.""" @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_speech_auto_voice_seed_deterministic(self, omni_server, openai_client) -> None: + def test_speech_auto_voice_seed_deterministic(self, omni_server, online_client) -> None: cfg = { "model": omni_server.model, "input": get_prompt("text"), @@ -88,12 +89,12 @@ class TestOmniVoiceSeed: "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, } - r1 = openai_client.send_audio_speech_request(cfg)[0] - r2 = openai_client.send_audio_speech_request(cfg)[0] + r1 = online_client.send_audio_speech_request(cfg)[0] + r2 = online_client.send_audio_speech_request(cfg)[0] assert r1.audio_bytes == r2.audio_bytes @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_speech_auto_voice_seed_non_deterministic(self, omni_server, openai_client) -> None: + def test_speech_auto_voice_seed_non_deterministic(self, omni_server, online_client) -> None: cfg = { "model": omni_server.model, "input": get_prompt("text"), @@ -104,8 +105,8 @@ class TestOmniVoiceSeed: cfg1 = {**cfg, "seed": 42} cfg2 = {**cfg, "seed": 43} - r1 = openai_client.send_audio_speech_request(cfg1)[0] - r2 = openai_client.send_audio_speech_request(cfg2)[0] + r1 = online_client.send_audio_speech_request(cfg1)[0] + r2 = online_client.send_audio_speech_request(cfg2)[0] assert r1.audio_bytes != r2.audio_bytes @@ -115,7 +116,7 @@ class TestOmniVoiceVoiceCloning: """E2E tests for OmniVoice voice cloning functionality.""" @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_voice_clone_ref_audio_only(self, omni_server, openai_client) -> None: + def test_voice_clone_ref_audio_only(self, omni_server, online_client) -> None: """Test voice cloning with ref_audio only (x_vector mode).""" request_config = { "model": omni_server.model, @@ -125,10 +126,10 @@ class TestOmniVoiceVoiceCloning: "timeout": 180.0, "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) - def test_voice_clone_ref_audio_and_text(self, omni_server, openai_client) -> None: + def test_voice_clone_ref_audio_and_text(self, omni_server, online_client) -> None: """Test voice cloning with ref_audio and ref_text (in-context mode).""" request_config = { "model": omni_server.model, @@ -139,4 +140,108 @@ class TestOmniVoiceVoiceCloning: "timeout": 180.0, "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +@pytest.mark.parametrize("voice", [_DEFAULT_VOICE_NAME, {"id": _DEFAULT_VOICE_NAME}]) +def test_speech_with_voice_param_accepted(omni_server, online_client, voice) -> None: + """The default voice should be accepted as both a plain string and a + VoiceID dict, even when the model has no uploaded speakers.""" + request_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "voice": voice, + "response_format": "wav", + "timeout": 180.0, + "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + } + online_client.send_audio_speech_request(request_config) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +def test_voice_request_before_and_after_clone_registration(omni_server, online_client) -> None: + """The default voice should always be listed and usable, even after + a different cloned voice is registered.""" + speech_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "voice": _DEFAULT_VOICE_NAME, + } + + # 1) "default" should appear in the voices list before any uploads + voices = online_client.send_audio_voices_list_http_request()[0] + assert _DEFAULT_VOICE_NAME in voices.json_body["voices"] + + # 2) Speech with the default voice should succeed + online_client.send_audio_speech_request(speech_config) + + # 3) Register a cloned voice under a different name + register_config = { + "data": {"name": "bar", "consent": "test-consent"}, + "files": { + "audio_sample": ( + "clone_2.wav", + get_asset_path("qwen3_tts/clone_2.wav").read_bytes(), + "audio/wav", + ), + }, + } + online_client.send_audio_voices_create_http_request(register_config) + + # 4) "default" should still be listed alongside the new voice + voices = online_client.send_audio_voices_list_http_request()[0] + assert _DEFAULT_VOICE_NAME in voices.json_body["voices"] + assert "bar" in voices.json_body["voices"] + + # 5) Speech with the default voice should still succeed + online_client.send_audio_speech_request(speech_config) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) +def test_registered_default_voice_overrides_placeholder(omni_server, online_client) -> None: + """When a real voice is uploaded under the default name, it should be + used as a cloned voice; after deletion, the placeholder behavior resumes.""" + speech_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "voice": _DEFAULT_VOICE_NAME, + } + + # 1) Placeholder default works before any registration + online_client.send_audio_speech_request(speech_config) + + # 2) Register a real voice under the default name + register_config = { + "data": {"name": _DEFAULT_VOICE_NAME, "consent": "test-consent"}, + "files": { + "audio_sample": ( + "clone_2.wav", + get_asset_path("qwen3_tts/clone_2.wav").read_bytes(), + "audio/wav", + ), + }, + } + online_client.send_audio_voices_create_http_request(register_config) + + # 3) After creating the voice with the default name, it's visible as an uploaded_voice + voices = online_client.send_audio_voices_list_http_request()[0] + assert _DEFAULT_VOICE_NAME in voices.json_body["voices"] + uploaded_names = [v["name"] for v in voices.json_body["uploaded_voices"]] + assert _DEFAULT_VOICE_NAME in uploaded_names + online_client.send_audio_speech_request(speech_config) + + # 4) Delete the uploaded default voice + online_client.send_audio_voices_delete_http_request( + {"name": _DEFAULT_VOICE_NAME}, + ) + + # 5) Ensure the placeholder behavior is consistent, but it's not an uploaded voice now + voices = online_client.send_audio_voices_list_http_request()[0] + assert _DEFAULT_VOICE_NAME in voices.json_body["voices"] + uploaded_names = [v["name"] for v in voices.json_body["uploaded_voices"]] + assert _DEFAULT_VOICE_NAME not in uploaded_names + online_client.send_audio_speech_request(speech_config) diff --git a/tests/e2e/online_serving/test_pi0_expansion.py b/tests/e2e/online_serving/test_pi0_expansion.py new file mode 100644 index 000000000..3a9c3b177 --- /dev/null +++ b/tests/e2e/online_serving/test_pi0_expansion.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E online serving test for π0 (Pi-Zero) OpenPI websocket serving. + +Boots ``vllm serve --omni --deploy-config pi0.yaml`` and drives the real OpenPI +websocket (``/v1/realtime/robot/openpi``) — the same wire path a robot uses +(handshake metadata → send observation → receive action chunk). Mirrors +``tests/e2e/online_serving/test_dreamzero_expansion.py``. Needs a GPU + a +pi0_base checkpoint; skipped in CI unless explicitly run. + +The in-process LeRobot parity oracle lives separately in +``tests/diffusion/models/pi0/test_pi0_parity.py``. +""" + +from __future__ import annotations + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import ( + OmniServerParams, + get_open_port, + pi0_openpi_require_dependencies, + pi0_openpi_run_policy_session, + pi0_openpi_validate_session_result, +) + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +MODEL = "lerobot/pi0_base" + + +test_params = [ + OmniServerParams( + model=MODEL, + port=8092, + server_args=[ + "--deploy-config", + "vllm_omni/deploy/pi0.yaml", + "--served-model-name", + "pi0", + "--enforce-eager", + "--disable-log-stats", + ], + env_dict={ + "ATTENTION_BACKEND": "torch", + "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA", + "VLLM_DISABLE_COMPILE_CACHE": "1", + "MASTER_PORT": str(get_open_port()), + }, + ) +] + + +@pytest.mark.full_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_pi0_openpi_online(omni_server): + try: + pi0_openpi_require_dependencies() + except ModuleNotFoundError as exc: + pytest.skip(str(exc)) + + result = pi0_openpi_run_policy_session( + host=omni_server.host, + port=omni_server.port, + prompt="pick up the red block and place it in the bin", + session_id="pi0-online-e2e", + num_steps=2, + ) + + # Asserts every returned chunk is [50, 32] + finite, and the handshake + # metadata matches pi0.yaml's policy_server_config. + pi0_openpi_validate_session_result(result) + + metadata = result["metadata"] + assert tuple(metadata["image_resolution"]) == (224, 224) + assert metadata["needs_wrist_camera"] is True + assert metadata["needs_session_id"] is False + assert metadata["action_space"] == "joint_position" diff --git a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py index bbf9ca2b6..1e5e59aed 100644 --- a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ E2E Online tests for Qwen2.5-Omni model with video input and audio output. """ @@ -59,7 +62,7 @@ def get_max_batch_size(size_type="few"): @pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_mix_to_text_audio_001(omni_server, openai_client) -> None: +def test_mix_to_text_audio_001(omni_server, online_client) -> None: """ Test multi-modal input processing and text/audio output generation via OpenAI API. Deploy Setting: default yaml @@ -90,13 +93,13 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: } # Test single completion - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @pytest.mark.slow @pytest.mark.omni @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_001(omni_server, openai_client) -> None: +def test_text_to_text_001(omni_server, online_client) -> None: """ Test text input processing and text/audio output generation via OpenAI API. Deploy Setting: default yaml @@ -114,4 +117,4 @@ def test_text_to_text_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index a54d2fe8a..f6ede8944 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ E2E Online tests for Qwen3-Omni model with video input and audio output. """ @@ -82,7 +85,7 @@ def get_max_batch_size(size_type="few"): @pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.") @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_mix_to_text_audio_001(omni_server, openai_client) -> None: +def test_mix_to_text_audio_001(omni_server, online_client) -> None: """ Test multi-modal input processing and text/audio output generation via OpenAI API. Deploy Setting: default yaml @@ -113,7 +116,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: } # Test single completion - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @pytest.mark.advanced_model @@ -122,7 +125,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: @pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.") @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_001(omni_server, openai_client) -> None: +def test_text_to_text_001(omni_server, online_client) -> None: """ Test text input processing and text/audio output generation via OpenAI API. Deploy Setting: default yaml @@ -140,18 +143,27 @@ def test_text_to_text_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) -def _run_prefix_cache_check(openai_client, request_config: dict): +def _run_prefix_cache_check(online_client, request_config: dict): """Make two requests given a request config, and validate that: 1. The second request actually had cached tokens 2. The number of cached tokens is divisible by the block size used in test_params, because currently upstream vLLM does not cache partial blocks. + 3. The first request also reports prompt token details; details must be + present whenever --enable-prompt-tokens-details is on, even when zero + tokens are cached. (An exact cached_tokens == 0 check would be + order-dependent here: the module-scoped server may have already cached + the shared system-prompt blocks from earlier tests.) + + Returns the (first, second) responses for additional per-test assertions. """ - openai_client.send_omni_request(request_config, request_num=1)[0] - cached_response = openai_client.send_omni_request(request_config, request_num=1)[0] + first_response = online_client.send_omni_request(request_config, request_num=1)[0] + cached_response = online_client.send_omni_request(request_config, request_num=1)[0] + + assert first_response.cached_tokens is not None # Ensure that we have a prefix cache hit on the second request and that only the last # partial block is uncached (since currently we don't cache partial blocks). @@ -163,13 +175,15 @@ def _run_prefix_cache_check(openai_client, request_config: dict): assert num_cached_tokens % BLOCK_SIZE == 0 assert (num_cached_tokens + num_uncached_tokens) == num_prompt_tokens + return first_response, cached_response + @pytest.mark.advanced_model @pytest.mark.core_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None: +def test_thinker_prefix_caching_text_output(omni_server, online_client) -> None: """ Test thinker prefix caching by sending identical requests with an image (i.e., a large shared prefix) and verifying that the second request uses cached tokens @@ -195,7 +209,10 @@ def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None: "stream": False, "modalities": ["text"], } - _run_prefix_cache_check(openai_client, request_config) + first_response, _ = _run_prefix_cache_check(online_client, request_config) + + assert first_response.multimodal_tokens is not None + assert first_response.multimodal_tokens.get("image", 0) > 0 @pytest.mark.advanced_model @@ -203,7 +220,7 @@ def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None: +def test_thinker_prefix_caching_audio_output(omni_server, online_client) -> None: """ Verify that thinker prefix caching does not hang when the request produces audio output (text + audio modalities). Sends two identical @@ -225,7 +242,7 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None }, } - _run_prefix_cache_check(openai_client, request_config) + _run_prefix_cache_check(online_client, request_config) @pytest.mark.advanced_model @@ -233,13 +250,13 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_completions_rejected_for_thinker_talker(omni_server, openai_client) -> None: +def test_completions_rejected_for_thinker_talker(omni_server, online_client) -> None: """Ensure Thinker-talker models reject /v1/completions; we do this because the thinker-talker handoff implementations currently use ChatML <|im_start|> and <|im_end|> markers to segment the input sequence; when we don't have them, the talker does not get any embeddings, which breaks the server. """ - responses = openai_client.send_completions_http_request( + responses = online_client.send_completions_http_request( { "json": { "model": omni_server.model, diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 1ba94c011..10c5dd8cb 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for Qwen3-Omni model. """ @@ -7,7 +7,6 @@ E2E Online tests for Qwen3-Omni model. import os import pytest -from openai import BadRequestError from tests.helpers.mark import hardware_test from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video @@ -89,7 +88,7 @@ test_params = [ ), ] -test_token_params = [ +reduced_token_params = [ pytest.param( OmniServerParams( model=model, @@ -97,7 +96,7 @@ test_token_params = [ use_stage_cli=True, server_args=["--async-chunk"], ), - id="batch_token_64", + id="batch_token_2048", ) ] @@ -138,30 +137,8 @@ def get_max_batch_size(size_type="few"): @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_audio_001(omni_server, openai_client) -> None: - """ - Input Modal: text - Output Modal: audio - Input Setting: stream=True - Datasets: single request - """ - messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()) - - request_config = { - "model": omni_server.model, - "messages": messages, - "modalities": ["audio"], - "stream": True, - "key_words": {"text": ["beijing"]}, - } - - openai_client.send_omni_request(request_config) - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) -def test_text_to_text_audio_001(omni_server, openai_client) -> None: +@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True) +def test_text_to_text_audio_001(omni_server, online_client) -> None: """ Input Modal: text Output Modal: text, audio @@ -176,12 +153,12 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_video_to_text_001(omni_server, openai_client) -> None: +def test_text_video_to_text_001(omni_server, online_client) -> None: """ Input Modal: long synthetic video (120s @ 30fps, LONG_VIDEO_FRAMES frames) Output Modal: text @@ -200,40 +177,17 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None: "key_words": {"video": VIDEO_KEY}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) -def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None: - """ - Input Modal: text, audio - Output Modal: text, audio - Input Setting: stream=False - Datasets: single request - """ - audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" - messages = dummy_messages_from_mix_data( - audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio") - ) - - request_config = { - "model": omni_server.model, - "messages": messages, - "key_words": {"audio": AUDIO_KEY}, - } - - openai_client.send_omni_request(request_config) - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) -def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None: +@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True) +def test_text_audio_to_text_audio_002(omni_server, online_client) -> None: """ Input Modal: text, long-duration audio (~LONG_AUDIO_DURATION_SEC s WAV) Output Modal: text, audio Input Setting: stream=False - Datasets: single request + Datasets: few requests """ audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(LONG_AUDIO_DURATION_SEC, 1)['base64']}" messages = dummy_messages_from_mix_data( @@ -248,41 +202,17 @@ def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None: "key_words": {"audio": AUDIO_KEY}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) -def test_text_image_to_text_audio_001(omni_server, openai_client) -> None: - """ - Input Modal: text, image - Output Modal: text, audio - Input Setting: stream=False - Datasets: single request - """ - image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" - - messages = dummy_messages_from_mix_data( - image_data_url=image_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_image") - ) - - request_config = { - "model": omni_server.model, - "messages": messages, - "key_words": {"image": IMAGE_KEY}, - } - - openai_client.send_omni_request(request_config) - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) -def test_large_image_to_text_audio_001(omni_server, openai_client) -> None: +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_large_image_to_text_audio_001(omni_server, online_client) -> None: """ Input Modal: text, high-resolution image (1080p-class JPEG) Output Modal: text, audio Input Setting: stream=False - Datasets: single request + Datasets: few requests """ image_data_url = ( f"data:image/jpeg;base64,{generate_synthetic_image(LARGE_IMAGE_WIDTH, LARGE_IMAGE_HEIGHT)['base64']}" @@ -300,12 +230,12 @@ def test_large_image_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"image": IMAGE_KEY}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) -def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_video_to_text_audio_001(omni_server, online_client) -> None: """ Input Modal: text, video Output Modal: text, audio @@ -325,42 +255,12 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"video": VIDEO_KEY}, } - openai_client.send_omni_request(request_config) - - -@pytest.mark.skip(reason="There is a known issue with shape mismatch error.") -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) -def test_mix_to_text_audio_001(omni_server, openai_client) -> None: - """ - Input Modal: text, audio, image, video - Output Modal: text, audio - Input Setting: stream=True - Datasets: few requests - """ - video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" - image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" - audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}" - messages = dummy_messages_from_mix_data( - system_prompt=get_system_prompt(), - video_data_url=video_data_url, - image_data_url=image_data_url, - audio_data_url=audio_data_url, - content_text=get_prompt("mix"), - ) - - request_config = { - "model": omni_server.model, - "messages": messages, - "stream": True, - "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY}, - } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_audio_in_video_001(omni_server, openai_client) -> None: +def test_audio_in_video_001(omni_server, online_client) -> None: """ Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video). Output Modal: text, audio @@ -381,12 +281,12 @@ def test_audio_in_video_001(omni_server, openai_client) -> None: "use_audio_in_video": True, "key_words": {"video": VIDEO_KEY}, } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_audio_in_video_002(omni_server, openai_client) -> None: +def test_audio_in_video_002(omni_server, online_client) -> None: """ Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video). Output Modal: text, audio @@ -408,12 +308,12 @@ def test_audio_in_video_002(omni_server, openai_client) -> None: "key_words": {"video": VIDEO_KEY}, } - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None: +def test_audio_in_video_default_loader_sampling_regression(omni_server, online_client) -> None: """ Regression: ``use_audio_in_video`` with default video loader sampling (no ``media_io_kwargs.video.fps``) after a prior request that pinned fps. @@ -438,7 +338,7 @@ def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_c } # Prime mm-cache with explicit loader fps (historical failure mode). - openai_client.send_omni_request( + online_client.send_omni_request( { **base_config, "extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}}, @@ -446,17 +346,37 @@ def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_c ) # Default loader sampling: must succeed without explicit media_io fps. - openai_client.send_omni_request(dict(base_config)) + online_client.send_omni_request(dict(base_config)) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_one_word_prompt_001(omni_server, openai_client) -> None: - """ +def test_one_word_prompt_001(omni_server, online_client) -> None: + """Catastrophic-regression gate on one-word pronunciation. + + Whether a one-word answer is pronounced intelligibly is sampled, not guaranteed: the + talker runs at temperature 0.9 (the upstream default, see qwen3_omni_moe.yaml stage 1), + so asserting it per request gates on a random variable, which is what the old 10-retry + loop was compensating for. + + Sends a fixed 40 requests and requires at least 29 to transcribe as the expected + word. 85.6% is a predeclared minimum acceptable success probability, taken from a + measured HF reference run (77/90); vLLM-Omni itself measured 92.2%. Assuming 40 + independent Bernoulli trials, >=29/40 has a nominal rejection probability of 0.9% + at p=.856, with 92.9% power at p=.60 and 55.9% at p=.70. + + Those are conditional on p, not a measured false-failure rate for this CI: 77/90 is + one point estimate (Wilson 76.8%-91.4%), and at the low end of that interval the gate + would reject about 20% of runs. Pinning down a real rate would take repeated + report-only runs on the target hardware. + + A drop to 70% is caught only about half the time, so this is a smoke gate for gross + pipeline breakage, not a quality tracker -- resolving moderate drift needs ~80+ + samples and belongs in an accuracy benchmark. See #5395. + Input Modal: text only (one-word answer constraint). Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text. Input Setting: stream=True - Datasets: single request """ messages = dummy_messages_from_mix_data( system_prompt=get_system_prompt(), @@ -474,47 +394,43 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None: "transcript_language": "en", } - # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/helpers/assertions.py). - _similarity_assert_msg = "The audio content is not same as the text" - _max_retries = 10 - for attempt in range(_max_retries): - try: - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) - break - except AssertionError as e: - if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1: - raise - print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}") + # The judge runs per request so mismatches can be counted, which also transcribes + # concurrently where the previous request_num=5 path judged serially -- worth watching + # on tighter cards, since each judge loads whisper on the same device as the talker and + # code2wav stages. + online_client.send_omni_request( + request_config, request_num=40, min_successes=29, max_concurrency=get_max_batch_size() + ) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_speaker_001(omni_server, openai_client) -> None: +def test_speaker_001(omni_server, online_client) -> None: """ - Input Modal: text only (one-word answer constraint). - Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text. + Input Modal: Chinese text. + Output Modal: Chinese text and audio using an uppercase preset speaker name. Input Setting: stream=True Datasets: single request """ messages = dummy_messages_from_mix_data( system_prompt=get_system_prompt(), - content_text=get_prompt("text"), + content_text=get_prompt("text_chinese"), ) request_config = { "model": omni_server.model, "messages": messages, "stream": True, - "speaker": "Chelsie", - "key_words": {"text": ["beijing"]}, + "speaker": "CHELSIE", + "key_words": {"text": ["北京"]}, } - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_speaker_002(omni_server, openai_client) -> None: +def test_speaker_002(omni_server, online_client) -> None: """ Input Modal: text only (one-word answer constraint). Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text. @@ -539,7 +455,7 @@ def test_speaker_002(omni_server, openai_client) -> None: _max_retries = 10 for attempt in range(_max_retries): try: - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + online_client.send_omni_request(request_config, request_num=get_max_batch_size()) break except AssertionError as e: if _gender_assert_substr not in str(e) or attempt == _max_retries - 1: @@ -549,61 +465,12 @@ def test_speaker_002(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_speaker_003(omni_server, openai_client) -> None: - """ - Input Modal: text only (one-word answer constraint). - Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text. - Input Setting: stream=True - Datasets: single request - """ - messages = dummy_messages_from_mix_data( - system_prompt=get_system_prompt(), - content_text=get_prompt("text"), - ) - - request_config = { - "model": omni_server.model, - "messages": messages, - "stream": True, - "speaker": "CHELSIE", - "key_words": {"text": ["beijing"]}, - } - - openai_client.send_omni_request(request_config) - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_language_001(omni_server, openai_client) -> None: - """ - Input Modal: text only (one-word answer constraint). - Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text. - Input Setting: stream=True - Datasets: single request - """ - messages = dummy_messages_from_mix_data( - system_prompt=get_system_prompt(), - content_text=get_prompt("text_chinese"), - ) - - request_config = { - "model": omni_server.model, - "messages": messages, - "stream": True, - "key_words": {"text": ["北京"]}, - } - - openai_client.send_omni_request(request_config) - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: +def test_text_to_audio_long_output_001(omni_server, online_client) -> None: """ Input Modal: text only (long-form generation prompt). Output Modal: text, audio (default ``modalities``); Input Setting: stream=True - Datasets: single request + Datasets: few requests """ messages = dummy_messages_from_mix_data( system_prompt=get_system_prompt(), @@ -611,25 +478,7 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: ) request_config = {"model": omni_server.model, "messages": messages, "stream": True} - responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + responses = online_client.send_omni_request(request_config, request_num=get_max_batch_size()) text = responses[0].text_content if responses else "" word_count = len(text.split()) assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}" - - -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) -@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True) -def test_invalid_audio_format_rejected(omni_server, openai_client) -> None: - """Ensure invalid audio format is rejected with 400 before streaming starts.""" - messages = dummy_messages_from_mix_data( - system_prompt=get_system_prompt(), - content_text=get_prompt(), - ) - with pytest.raises(BadRequestError, match="audio format"): - openai_client.client.chat.completions.create( - model=omni_server.model, - messages=messages, - modalities=["text", "audio"], - audio={"voice": "alloy", "format": "aac"}, - stream=True, - ) diff --git a/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py index f269836f3..1a27a5449 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py +++ b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Core-model CI guard for Qwen3-Omni multi-replica stage-pool routing on 4 GPUs. """ @@ -90,7 +93,7 @@ def _assert_audio_outputs(responses: list[OmniResponse], *, expect_text: bool) - @pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) -> None: +def test_text_only_batch_uses_multi_replica_talker(omni_server, online_client) -> None: request_config = { "model": omni_server.model, "messages": _text_messages(), @@ -98,7 +101,7 @@ def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) - "modalities": ["text"], } - responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS) + responses = online_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS) _assert_batch_size(responses, ROUTE_STRESS_REQUESTS) _assert_text_outputs(responses) @@ -107,7 +110,7 @@ def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) - @pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, openai_client) -> None: +def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, online_client) -> None: request_config = { "model": omni_server.model, "messages": _text_messages(), @@ -115,7 +118,7 @@ def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, open "modalities": ["audio"], } - responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS) + responses = online_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS) _assert_batch_size(responses, ROUTE_STRESS_REQUESTS) _assert_audio_outputs(responses, expect_text=False) @@ -124,13 +127,13 @@ def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, open @pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards=4) @pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_mixed_modal_stream_batch_generates_text_and_audio(omni_server, openai_client) -> None: +def test_mixed_modal_stream_batch_generates_text_and_audio(omni_server, online_client) -> None: request_config = { "model": omni_server.model, "messages": _mixed_messages(), "stream": True, } - responses = openai_client.send_omni_request(request_config, request_num=MIXED_MODAL_REQUESTS) + responses = online_client.send_omni_request(request_config, request_num=MIXED_MODAL_REQUESTS) _assert_batch_size(responses, MIXED_MODAL_REQUESTS) _assert_audio_outputs(responses, expect_text=True) diff --git a/tests/e2e/online_serving/test_qwen3_tts_base.py b/tests/e2e/online_serving/test_qwen3_tts_base.py index 60caf8f82..dd25520f9 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for Qwen3-TTS model with text input and audio output. @@ -14,7 +14,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path @@ -24,7 +24,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0 # Vendored under tests/assets/qwen3_tts/clone_2.wav so the server does not need # to fetch the reference audio over HTTPS at request time (CI runners in some # regions cannot reach the original Aliyun OSS host; see issue #3263). -REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." @@ -55,11 +55,10 @@ tts_server_params = [ @pytest.mark.advanced_model -@pytest.mark.core_model @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_server, openai_client) -> None: +def test_text_to_audio_001(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -80,14 +79,14 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few")) + online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few")) @pytest.mark.advanced_model @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_002(omni_server, openai_client) -> None: +def test_text_to_audio_002(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -108,4 +107,4 @@ def test_text_to_audio_002(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py index 0edaa4668..e126877db 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for Qwen3-TTS model with text input and audio output. @@ -12,7 +12,7 @@ import os import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path @@ -23,7 +23,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" # See tests/e2e/online_serving/test_qwen3_tts_base.py for the vendored-asset rationale. -REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." @@ -41,7 +41,7 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -tts_server_params = [ +tts_async_chunk_server_params = [ pytest.param( OmniServerParams( model=MODEL, @@ -50,9 +50,12 @@ tts_server_params = [ ), id="async_chunk", ), - # Synchronous (no async-chunk) variant — ``--no-async-chunk`` alone - # flips the deploy yaml's bool and the pipeline dispatches to the - # end-to-end codec processor. No variant yaml / pipeline needed. +] + +# Synchronous (no async-chunk) variant — ``--no-async-chunk`` alone +# flips the deploy yaml's bool and the pipeline dispatches to the +# end-to-end codec processor. No variant yaml / pipeline needed. +tts_no_async_chunk_server_params = [ pytest.param( OmniServerParams( model=MODEL, @@ -63,10 +66,16 @@ tts_server_params = [ ), ] +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0 + @hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_clone_streaming_001(omni_server, openai_client) -> None: +@pytest.mark.parametrize( + "omni_server", + tts_async_chunk_server_params + tts_no_async_chunk_server_params, + indirect=True, +) +def test_voice_clone_streaming_001(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -87,12 +96,16 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few")) + online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few")) @hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_response_format_001(omni_server, openai_client) -> None: +@pytest.mark.parametrize( + "omni_server", + tts_async_chunk_server_params + tts_no_async_chunk_server_params, + indirect=True, +) +def test_response_format_001(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -111,4 +124,31 @@ def test_response_format_001(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True) +def test_xvector_then_icl_same_ref_audio_keeps_engine_alive(omni_server, online_client) -> None: + """Regression for #5049: x-vector-only then ICL on the same ``ref_audio``. + + Pre-fix, readiness was mode-agnostic: an x-vector request marked the + artifact ready (embedding only, no ``ref_code``), so the following ICL + request took the artifact-only path and killed EngineCore. After the fix + (readiness keyed by ``(artifact_key, x_vector_only)``), both modes succeed + and the server stays up. Only runs on the ``async_chunk`` server param. + """ + base_request = { + "model": omni_server.model, + "input": get_prompt(), + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "task_type": "Base", + "ref_audio": REF_AUDIO_URL, + "min_audio_bytes": 1, + } + + online_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True}) + online_client.send_audio_speech_request({**base_request, "x_vector_only_mode": False, "ref_text": REF_TEXT}) + online_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True}) diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py index 7ebe08e1a..6192f8410 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for Qwen3-TTS model with text input and audio output. @@ -20,7 +20,6 @@ from tests.helpers.stage_config import ( get_deploy_config_stage, modify_stage_config, ) -from vllm_omni.platforms import current_omni_platform MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" @@ -92,13 +91,9 @@ def test_default_cuda_graph_startup(omni_server) -> None: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.tts -@pytest.mark.skipif( - current_omni_platform.is_npu(), - reason="#issue 5479", -) @hardware_test(res={"cuda": "L4", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_server, openai_client) -> None: +def test_text_to_audio_001(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -116,19 +111,15 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: "voice": "vivian", } - openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size()) + online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size()) @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.tts -@pytest.mark.skipif( - current_omni_platform.is_npu(), - reason="#issue 5479", -) @hardware_test(res={"cuda": "L4", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_002(omni_server, openai_client) -> None: +def test_text_to_audio_002(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -147,4 +138,4 @@ def test_text_to_audio_002(omni_server, openai_client) -> None: "voice": "vivian", } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py index 9c941cad8..07808af27 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for Qwen3-TTS model with text input and audio output. @@ -62,7 +62,7 @@ tts_server_params = [ @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_001(omni_server, openai_client) -> None: +def test_voice_001(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -84,7 +84,7 @@ def test_voice_001(omni_server, openai_client) -> None: # Retry once on assertion failures from transcript similarity / gender checks (flaky ASR or estimators). for attempt in range(2): try: - openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size()) + online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size()) break except AssertionError: if attempt == 0: @@ -94,7 +94,7 @@ def test_voice_001(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_002(omni_server, openai_client) -> None: +def test_voice_002(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -112,12 +112,12 @@ def test_voice_002(omni_server, openai_client) -> None: "voice": "Serena", } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_003(omni_server, openai_client) -> None: +def test_voice_003(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -135,12 +135,12 @@ def test_voice_003(omni_server, openai_client) -> None: "voice": "SERENA", } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_language_001(omni_server, openai_client) -> None: +def test_language_001(omni_server, online_client) -> None: """ Test text input processing and audio output via OpenAI API. Deploy Setting: default yaml @@ -158,4 +158,4 @@ def test_language_001(omni_server, openai_client) -> None: "voice": "vivian", } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py b/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py deleted file mode 100644 index e6c396932..000000000 --- a/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py +++ /dev/null @@ -1,88 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E regression for Qwen3-TTS ref_audio artifact mode collision (#5049). - -Reproduces the issue's exact sequence on a live server: an -``x_vector_only_mode=true`` request followed by an ICL -(``x_vector_only_mode=false`` + ``ref_text``) request using the **same** -``ref_audio``. - -Before the fix, the ref_audio artifact readiness was mode-agnostic: the -x-vector-only request (which caches a speaker embedding but no ``ref_code``) -marked the artifact "ready", so the ICL request took the artifact-only serving -path (``ref_audio`` stripped before dispatch), found no ``ref_code`` and no -audio to recompute from, and raised in ``build_prompt_embeds`` -> -``EngineDeadError`` — taking the whole server down rather than failing the one -request. - -After the fix (readiness keyed by ``(artifact_key, x_vector_only)``), the ICL -request no longer matches the x-vector artifact, so it sends ``ref_audio`` and -the worker computes ``ref_code``. Both requests succeed and the server stays up. -""" - -import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - -import pytest - -from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url -from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path - -MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" -DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0 - -# Vendored under tests/assets/qwen3_tts/clone_2.wav (see test_qwen3_tts_base.py). -REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") -REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." -INPUT_TEXT = "The weather is nice today, perfect for a walk in the park." - -tts_server_params = [ - pytest.param( - OmniServerParams( - model=MODEL, - stage_config_path=get_deploy_config_path("qwen3_tts.yaml"), - server_args=["--trust-remote-code"], - ), - id="async_chunk", - ) -] - - -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.tts -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_xvector_then_icl_same_ref_audio_keeps_engine_alive(omni_server, openai_client) -> None: - """Regression for #5049: x-vector-only then ICL on the same ref_audio. - - Both requests must return audio and the engine must stay alive (pre-fix the - second request killed EngineCore). - """ - base_request = { - "model": omni_server.model, - "input": INPUT_TEXT, - "stream": False, - "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, - "response_format": "wav", - "task_type": "Base", - "ref_audio": REF_AUDIO_URL, - # Assert non-empty audio came back (not just HTTP 200). - "min_audio_bytes": 1, - } - - # 1) x-vector-only request: caches a speaker embedding (no ref_code) and, on - # completion, marks the ref_audio artifact "ready". - openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True}) - - # 2) ICL request with the SAME ref_audio. Pre-fix this reused the x-vector - # artifact via the artifact-only path, hit the missing ref_code, and - # killed EngineCore (#5049). It must now succeed with audio. - openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": False, "ref_text": REF_TEXT}) - - # 3) The server must still be serving after the ICL request (pre-fix this - # would fail with EngineDeadError / connection refused). - openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True}) diff --git a/tests/e2e/online_serving/test_qwen_image.py b/tests/e2e/online_serving/test_qwen_image.py index c40cfbda8..aa037d9b8 100644 --- a/tests/e2e/online_serving/test_qwen_image.py +++ b/tests/e2e/online_serving/test_qwen_image.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Online serving tests for ``Qwen/Qwen-Image`` (text-to-image). @@ -24,7 +24,7 @@ import os import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -61,7 +61,7 @@ def _get_default_case(model: str): @pytest.mark.advanced_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True) -def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: +def test_text_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: """Default Qwen-Image T2I smoke (single ``default`` server config).""" messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT) request_config = { @@ -76,14 +76,14 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH "seed": 42, }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True) -def test_batch_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: +def test_batch_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: """Concurrent T2I: one ``request_config`` dict per prompt (``send_diffusion_request`` list mode).""" request_config = [ { @@ -100,4 +100,4 @@ def test_batch_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) } for prompt in TEST_PROMPTS ] - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen_image_edit.py b/tests/e2e/online_serving/test_qwen_image_edit.py index 738a2e50c..dc1f60521 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit.py +++ b/tests/e2e/online_serving/test_qwen_image_edit.py @@ -1,21 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ -Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions). +Online serving tests for ``Qwen/Qwen-Image-Edit-2511`` (image-to-image via chat completions). -- ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512. -- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512. -- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input - resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one - scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``). +- ``test_single_image_to_image_001``: one reference image, fixed 512×512. +- ``test_multi_images_to_image_001``: two reference images, fixed 512×512. +- ``test_different_sizes_001``: ``slow`` — mixed input resolutions; ``extra_body`` uses + per-output ``width``/``height`` lists; the test client sends one scalar-size request + per list index in parallel and merges images (see ``OnlineOmniClient.send_diffusion_request``). -``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row per model. +``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row. From ``tests/``:: - pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "core_model and diffusion" --run-level=core_model - pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "advanced_model and diffusion" --run-level=advanced_model + pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "slow and diffusion" --run-level=full_model """ import pytest @@ -25,10 +24,11 @@ from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) +MODEL = "Qwen/Qwen-Image-Edit-2511" EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." MULTI_EDIT_PROMPT = ( "Transform the first image into a Dadaism collage art. " @@ -39,7 +39,7 @@ NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) -def _get_diffusion_feature_cases(model: str): +def _get_diffusion_feature_cases(model: str = MODEL): """Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``).""" return [ pytest.param( @@ -52,15 +52,15 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"), + _get_diffusion_feature_cases(), indirect=True, ) -def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Single-reference edit smoke for ``Qwen/Qwen-Image-Edit`` (CFG when negative prompt + ``true_cfg_scale`` > 1).""" +def test_single_image_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient): + """Single-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511`` (CFG when negative prompt + ``true_cfg_scale`` > 1).""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) @@ -79,18 +79,17 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"), + _get_diffusion_feature_cases(), indirect=True, ) -def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_multi_images_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient): """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511``.""" image_data_url_list = [f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" for _ in range(2)] @@ -111,17 +110,17 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"), + _get_diffusion_feature_cases(), indirect=True, ) -def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_different_sizes_001(omni_server: OmniServer, online_client: OnlineOmniClient): """Multi-reference edit with distinct input resolutions; per-output ``width``/``height`` as lists (client splits into concurrent scalar-size calls).""" image_data_url_list = [ f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}", @@ -144,4 +143,4 @@ def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClien }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index b3d6f8b19..47bc2809b 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py @@ -1,20 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Comprehensive tests of diffusion features that are available in online serving mode -and are supported by the following models: -- Qwen-Image-Edit: single image input -- Qwen-Image-Edit-2511: single image input and two image inputs +and are supported by Qwen-Image-Edit-2511 (single image input and two image inputs). """ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] -EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." -SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting." +MODEL = "Qwen/Qwen-Image-Edit-2511" +SINGLE_EDIT_PROMPT = "Restyle this image into a Vincent van Gogh style impressionist painting." MULTI_EDIT_PROMPT = ( "Transform the first image into a Dadaism collage art. " "Transform the second image into a Vincent van Gogh style painting. " @@ -25,9 +26,7 @@ SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) -# This test file targets two models, so I write a helper function. -# If a similar test only involves one model, one can just define a global list variable. -def _get_diffusion_feature_cases(model: str): +def _get_diffusion_feature_cases(model: str = MODEL): return [ pytest.param( OmniServerParams( @@ -123,38 +122,10 @@ def _get_diffusion_feature_cases(model: str): @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"), + _get_diffusion_feature_cases(), indirect=True, ) -def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Test all diffusion features with Qwen-Image-Edit in regular end-user scenarios.""" - image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" - - messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) - - # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present - request_config = { - "model": omni_server.model, - "messages": messages, - "extra_body": { - "height": 512, - "width": 512, - "num_inference_steps": 2, - "negative_prompt": NEGATIVE_PROMPT, - "true_cfg_scale": 4.0, - "seed": 42, - }, - } - - openai_client.send_diffusion_request(request_config) - - -@pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"), - indirect=True, -) -def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, online_client: OnlineOmniClient): """Test Qwen-Image-Edit-2511 with a single image input. Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache @@ -166,7 +137,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien """ image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" - messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511) + messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT) request_config = { "model": omni_server.model, @@ -181,15 +152,15 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"), + _get_diffusion_feature_cases(), indirect=True, ) -def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, online_client: OnlineOmniClient): """Test Qwen-Image-Edit-2511 with two image inputs.""" image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" @@ -212,4 +183,4 @@ def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index c900ca4ce..890b077d2 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Comprehensive tests of diffusion features that are available in online serving mode and are supported by the following text-to-image models: @@ -5,15 +8,14 @@ and are supported by the following text-to-image models: - Qwen-Image-2512 One feature per test case, matching the Test Plan in PR #1682 (Qwen-Image-Edit). -Supported features for Qwen-Image series: TeaCache, Cache-DiT, Ulysses-SP, Ring-Attention, -CFG-Parallel, Tensor-Parallel, VAE-Patch-Parallel, CPU offload (layerwise). -See docs/user_guide/diffusion_acceleration.md. +Nightly covers each feature once, alternating models (5× Qwen-Image, 4× Qwen-Image-2512). +Ulysses stays on Qwen-Image. See docs/user_guide/diffusion_acceleration.md. """ import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] @@ -22,118 +24,116 @@ NEGATIVE_PROMPT = "blurry, low quality" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +MODEL_IMAGE = "Qwen/Qwen-Image" +MODEL_2512 = "Qwen/Qwen-Image-2512" -def _get_diffusion_feature_cases(model: str): - return [ - pytest.param( - OmniServerParams( - model=model, - server_args=["--enable-cpu-offload"], - ), - id="cpu_offload", - marks=SINGLE_CARD_FEATURE_MARKS, +# One server per feature. Alternate models; keep feature pytest ids unchanged. +FEATURE_CASES = [ + pytest.param( + OmniServerParams( + model=MODEL_IMAGE, + server_args=["--enable-cpu-offload"], ), - pytest.param( - OmniServerParams(model=model, server_args=["--step-execution"]), - id="step_execution", - marks=SINGLE_CARD_FEATURE_MARKS, + id="cpu_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams(model=MODEL_2512, server_args=["--step-execution"]), + id="step_execution", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams(model=MODEL_IMAGE, server_args=["--cache-backend", "tea_cache"]), + id="cache_tea_cache", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=MODEL_2512, + server_args=[ + "--cache-backend", + "cache_dit", + "--enable-layerwise-offload", + ], ), - pytest.param( - OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]), - id="cache_tea_cache", - marks=SINGLE_CARD_FEATURE_MARKS, + id="layerwise_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=MODEL_IMAGE, + server_args=[ + "--cache-backend", + "cache_dit", + "--ulysses-degree", + "2", + ], ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--cache-backend", - "cache_dit", - "--enable-layerwise-offload", - ], - ), - id="layerwise_offload", - marks=SINGLE_CARD_FEATURE_MARKS, + id="ulysses_2", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=MODEL_2512, + server_args=[ + "--cache-backend", + "cache_dit", + "--ring", + "2", + ], ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--cache-backend", - "cache_dit", - "--ulysses-degree", - "2", - ], - ), - id="ulysses_2", - marks=PARALLEL_FEATURE_MARKS, + id="ring_2", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=MODEL_IMAGE, + server_args=[ + "--cache-backend", + "tea_cache", + "--cfg-parallel-size", + "2", + ], ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--cache-backend", - "cache_dit", - "--ring", - "2", - ], - ), - id="ring_2", - marks=PARALLEL_FEATURE_MARKS, + id="cfg_parallel_2", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=MODEL_2512, + server_args=[ + "--cache-backend", + "cache_dit", + "--tensor-parallel-size", + "2", + "--vae-patch-parallel-size", + "2", + "--vae-use-tiling", + "--diffusion-quantization-config", + '{"method":"fp8"}', + ], ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--cache-backend", - "tea_cache", - "--cfg-parallel-size", - "2", - ], - ), - id="cfg_parallel_2", - marks=PARALLEL_FEATURE_MARKS, + id="vae_patch_parallel_2", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=MODEL_IMAGE, + server_args=[ + "--use-hsdp", + "--hsdp-shard-size", + "2", + ], ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--cache-backend", - "cache_dit", - "--tensor-parallel-size", - "2", - "--vae-patch-parallel-size", - "2", - "--vae-use-tiling", - "--diffusion-quantization-config", - '{"method":"fp8"}', - ], - ), - id="vae_patch_parallel_2", - marks=PARALLEL_FEATURE_MARKS, - ), - pytest.param( - OmniServerParams( - model=model, - server_args=[ - "--use-hsdp", - "--hsdp-shard-size", - "2", - ], - ), - id="parallel_hsdp", - marks=PARALLEL_FEATURE_MARKS, - ), - ] + id="parallel_hsdp", + marks=PARALLEL_FEATURE_MARKS, + ), +] -@pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image"), - indirect=True, -) -def test_qwen_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Test each diffusion feature with Qwen-Image (text-to-image), one feature per case.""" +@pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True) +def test_qwen_image(omni_server: OmniServer, online_client: OnlineOmniClient): + """One diffusion feature per case; model is chosen in FEATURE_CASES.""" messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT) request_config = { "model": omni_server.model, @@ -147,27 +147,4 @@ def test_qwen_image(omni_server: OmniServer, openai_client: OpenAIClientHandler) "seed": 42, }, } - openai_client.send_diffusion_request(request_config) - - -@pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases("Qwen/Qwen-Image-2512"), - indirect=True, -) -def test_qwen_image_2512(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Test each diffusion feature with Qwen-Image-2512 (text-to-image), one feature per case.""" - messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT) - request_config = { - "model": omni_server.model, - "messages": messages, - "extra_body": { - "height": 512, - "width": 512, - "num_inference_steps": 2, - "negative_prompt": NEGATIVE_PROMPT, - "true_cfg_scale": 4.0, - "seed": 42, - }, - } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen_image_layered.py b/tests/e2e/online_serving/test_qwen_image_layered.py index b29609351..298f4ca9e 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered.py +++ b/tests/e2e/online_serving/test_qwen_image_layered.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Online serving tests for ``Qwen/Qwen-Image-Layered`` (layer decomposition). @@ -18,7 +18,7 @@ from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) @@ -44,11 +44,10 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_single_image_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient): """Default Qwen-Image-Layered smoke (single ``default`` server config).""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" @@ -65,4 +64,4 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index 315fba6ef..bc36c9fa8 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Comprehensive tests of diffusion features that are available in online serving mode and are supported by Qwen-Image-Layered model. @@ -12,13 +12,18 @@ Kept cases (maximum feature coverage, no cpu-offload): Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload. """ +import base64 +import io +from typing import Any + import pytest +from PIL import Image from tests.helpers.mark import hardware_marks -from tests.helpers.media import decode_b64_image, generate_synthetic_image -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "Qwen/Qwen-Image-Layered" EDIT_PROMPT = "Decompose this image into layers." @@ -82,7 +87,7 @@ FEATURE_CASES = [ @pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True) -def test_feature(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_feature(omni_server: OmniServer, online_client: OnlineOmniClient): """Test feature combinations with Qwen-Image-Layered.""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" @@ -99,7 +104,7 @@ def test_feature(omni_server: OmniServer, openai_client: OpenAIClientHandler): }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) # ── Guard: Issue #1969 – layered output must return `layers` images ───────── @@ -123,19 +128,19 @@ LAYERS_GUARD_CASES = [ ] -def _collect_image_url_items(openai_client: OpenAIClientHandler, request_config: dict): +def _collect_image_url_items(online_client: OnlineOmniClient, request_config: dict): """Send a request and return all image_url content items from the API response. Handles both object-style (Pydantic) and dict-style (raw JSON) content items, because the OpenAI SDK may return either depending on how the server serializes the multimodal content list. """ - chat_completion = openai_client.client.chat.completions.create( + chat_completion = online_client.client.chat.completions.create( model=request_config["model"], messages=request_config["messages"], extra_body=request_config.get("extra_body"), ) - image_items = [] + image_items: list[Any] = [] for choice in chat_completion.choices: content = getattr(choice.message, "content", None) assert content is not None, "API response content is None" @@ -164,7 +169,7 @@ def _collect_image_url_items(openai_client: OpenAIClientHandler, request_config: ) def test_layered_output_image_count( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, expected_layers: int, ): """Guard for https://github.com/vllm-project/vllm-omni/issues/1969 @@ -189,7 +194,7 @@ def test_layered_output_image_count( }, } - image_items = _collect_image_url_items(openai_client, request_config) + image_items = _collect_image_url_items(online_client, request_config) # Core assertion: the API must NOT collapse multiple layers into 1 image assert len(image_items) != 1 or expected_layers == 1, ( @@ -209,7 +214,8 @@ def test_layered_output_image_count( url = item.image_url.url assert url.startswith("data:image"), f"image_url item {i} is not a data URI: {url[:80]}" b64 = url.split(",", 1)[1] - img = decode_b64_image(b64) + img = Image.open(io.BytesIO(base64.b64decode(b64))) + img.load() assert img is not None, f"Failed to decode image at index {i}" @@ -231,7 +237,7 @@ PROMPT_CASES = [ @pytest.mark.parametrize("omni_server", PROMPT_CASES, indirect=True) -def test_empty_prompt(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_empty_prompt(omni_server: OmniServer, online_client: OnlineOmniClient): """Test feature combinations with Qwen-Image-Layered.""" image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" @@ -248,4 +254,4 @@ def test_empty_prompt(omni_server: OmniServer, openai_client: OpenAIClientHandle }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_sana_video_expansion.py b/tests/e2e/online_serving/test_sana_video_expansion.py new file mode 100644 index 000000000..31e0f8552 --- /dev/null +++ b/tests/e2e/online_serving/test_sana_video_expansion.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""SANA-Video 2B online smoke tests for the native and Diffusers backends.""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL_480P = os.environ.get( + "SANA_VIDEO_480P_MODEL", + "Efficient-Large-Model/SANA-Video_2B_480p_diffusers", +) +MODEL_720P = os.environ.get( + "SANA_VIDEO_720P_MODEL", + "Efficient-Large-Model/SANA-Video_2B_720p_diffusers", +) +PROMPT = "A cat walking on grass toward the camera. motion score: 30." +NEGATIVE_PROMPT = "blurry, low quality, temporal artifacts" + +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _t2v_backend_cases(): + cases = [] + for variant, model in (("480p", MODEL_480P), ("720p", MODEL_720P)): + cases.extend( + [ + pytest.param( + OmniServerParams(model=model, server_args=["--model-class-name", "SanaVideoPipeline"]), + id=f"native-{variant}", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--diffusion-load-format", + "diffusers", + "--diffusion-attention-backend", + "TORCH_SDPA", + ], + ), + id=f"diffusers-adapter-{variant}", + marks=SINGLE_CARD_MARKS, + ), + ] + ) + return cases + + +@pytest.mark.parametrize("omni_server", _t2v_backend_cases(), indirect=True) +def test_sana_video_t2v_backends(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": 192, + "width": 320, + "num_frames": 9, + "fps": 8, + "num_inference_steps": 2, + "guidance_scale": 4.0, + "seed": 42, + }, + } + openai_client.send_video_diffusion_request(request_config) + + +def _i2v_cases(): + cases = [] + for variant, model in (("480p", MODEL_480P), ("720p", MODEL_720P)): + cases.extend( + [ + pytest.param( + OmniServerParams(model=model, server_args=["--model-class-name", "SanaImageToVideoPipeline"]), + id=f"native-{variant}", + marks=SINGLE_CARD_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--model-class-name", + "SanaImageToVideoPipeline", + "--diffusion-load-format", + "diffusers", + "--diffusion-attention-backend", + "TORCH_SDPA", + ], + ), + id=f"diffusers-adapter-{variant}", + marks=SINGLE_CARD_MARKS, + ), + ] + ) + return cases + + +@pytest.mark.parametrize("omni_server", _i2v_cases(), indirect=True) +def test_sana_video_i2v_variants(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + request_config = { + "model": omni_server.model, + "form_data": { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "height": 192, + "width": 320, + "num_frames": 9, + "fps": 8, + "num_inference_steps": 2, + "guidance_scale": 4.0, + "seed": 42, + }, + "image_reference": (f"data:image/jpeg;base64,{generate_synthetic_image(320, 192, seed=42)['base64']}"), + } + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_sana_wm.py b/tests/e2e/online_serving/test_sana_wm.py new file mode 100644 index 000000000..883e4c367 --- /dev/null +++ b/tests/e2e/online_serving/test_sana_wm.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Online serving smoke for SANA-WM image-to-video (first-frame I2V via ``/v1/videos``). + +Boots the server straight from the model id — the Stage-1 repo's +``model_index.json`` names ``SanaWmPipeline``, so the class resolves on its own +— submits one async ``/v1/videos`` job with a first-frame reference image, and +asserts the job completes and returns video bytes. + +From ``tests/``:: + + pytest -s -v e2e/online_serving/test_sana_wm.py -m "advanced_model and diffusion" --run-level=advanced_model +""" + +import base64 +import json +import os +from io import BytesIO + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from vllm_omni.diffusion.models.sana_wm import ( + SANA_WM_MODEL_ID, + SANA_WM_OUTPUT_HEIGHT, + SANA_WM_OUTPUT_WIDTH, +) + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = SANA_WM_MODEL_ID +PROMPT = "A slow forward camera move through a quiet city street." +NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark" + +# Smoke knobs: a handful of frames / steps so the job finishes quickly. +SMOKE_NUM_FRAMES = 9 +SMOKE_NUM_INFERENCE_STEPS = 2 +# First-frame I2V needs a camera trajectory; "w-" is the forward-move action DSL. +# Passed via the ``sana_wm`` form field (a JSON string parsed server-side); explicit +# intrinsics avoid the optional Pi3X camera-calibration dependency. +SANA_WM_PARAMS = { + "action": f"w-{SMOKE_NUM_FRAMES - 1}", + "translation_speed": 0.055, + "rotation_speed_deg": 1.2, + "intrinsics": { + "fx": SANA_WM_OUTPUT_WIDTH / 2, + "fy": SANA_WM_OUTPUT_WIDTH / 2, + "cx": SANA_WM_OUTPUT_WIDTH / 2, + "cy": SANA_WM_OUTPUT_HEIGHT / 2, + }, +} + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _first_frame_data_url() -> str: + """A solid-color first frame encoded as a PNG ``data:`` URL for I2V.""" + from PIL import Image + + image = Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160)) + buffer = BytesIO() + image.save(buffer, format="PNG") + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + return f"data:image/png;base64,{encoded}" + + +def _get_diffusion_feature_cases(model: str): + """Single default server row.""" + return [ + pytest.param( + OmniServerParams(model=model), + id="default", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_image_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + """Default SANA-WM I2V smoke: async ``/v1/videos`` job completes and returns video bytes.""" + request_config = { + "model": omni_server.model, + "image_reference": _first_frame_data_url(), + "form_data": { + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + # The camera payload rides inside extra_params; there is no + # top-level sana_wm form field on the video endpoints, so sending + # one is silently dropped and the request fails in the + # preprocessor for missing camera control. + "extra_params": json.dumps({"sana_wm": SANA_WM_PARAMS}), + "height": SANA_WM_OUTPUT_HEIGHT, + "width": SANA_WM_OUTPUT_WIDTH, + "num_frames": SMOKE_NUM_FRAMES, + "fps": 16, + "num_inference_steps": SMOKE_NUM_INFERENCE_STEPS, + "guidance_scale": 5.0, + "seed": 42, + }, + } + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_sd3_expansion.py b/tests/e2e/online_serving/test_sd3_expansion.py index f881a212b..9793d8a4b 100644 --- a/tests/e2e/online_serving/test_sd3_expansion.py +++ b/tests/e2e/online_serving/test_sd3_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Tests for Stable Diffusion 3.5 medium model. """ @@ -8,7 +11,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) @@ -51,10 +54,7 @@ def _get_diffusion_feature_cases(model: str): "2", ], ), - marks=[ - *FOUR_CARD_FEATURE_MARKS, - pytest.mark.skip(reason="#3432"), - ], + marks=FOUR_CARD_FEATURE_MARKS, ), ] @@ -66,7 +66,7 @@ def _get_diffusion_feature_cases(model: str): ), indirect=True, ) -def test_sd3_medium(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_sd3_medium(omni_server: OmniServer, online_client: OnlineOmniClient): messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT) request_config = { "model": omni_server.model, @@ -81,4 +81,4 @@ def test_sd3_medium(omni_server: OmniServer, openai_client: OpenAIClientHandler) }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_sdxl.py b/tests/e2e/online_serving/test_sdxl.py index 4d7fc34d5..052d9e57e 100644 --- a/tests/e2e/online_serving/test_sdxl.py +++ b/tests/e2e/online_serving/test_sdxl.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Tests for Stable Diffusion XL Base 1.0 model. @@ -10,7 +13,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import ( OmniServer, OmniServerParams, - OpenAIClientHandler, + OnlineOmniClient, dummy_messages_from_mix_data, ) @@ -41,7 +44,7 @@ def _get_diffusion_feature_cases(model: str): ), indirect=True, ) -def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_text_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient): messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT) request_config = { "model": omni_server.model, @@ -56,4 +59,4 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_stable_audio_online_expansion.py b/tests/e2e/online_serving/test_stable_audio_online_expansion.py new file mode 100644 index 000000000..fbd6bdbfd --- /dev/null +++ b/tests/e2e/online_serving/test_stable_audio_online_expansion.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online serving test for Stable Audio Open text-to-audio diffusion. + +Stable Audio Open is served through the OpenAI-compatible +`POST /v1/audio/generate` endpoint (JSON in, binary WAV out). +""" + +import os + +import pytest + +from tests.helpers import skip_if_gated_repo_inaccessible +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +# Stable Audio Open is gated; override to a local bundle locally if needed. +STABLE_AUDIO_TEST_MODEL = os.environ.get("STABLE_AUDIO_TEST_MODEL", "stabilityai/stable-audio-open-1.0") +T2A_PROMPT = "A piano playing a gentle melody with soft room ambience." + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) + + +def _stable_audio_server_cases(model: str): + return [ + pytest.param( + OmniServerParams(model=model), + id="t2a", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.fixture +def _require_stable_audio_access() -> None: + """Skip cleanly (before the server boots) if the gated checkpoint is inaccessible.""" + skip_if_gated_repo_inaccessible(STABLE_AUDIO_TEST_MODEL) + + +@pytest.mark.slow +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _stable_audio_server_cases(STABLE_AUDIO_TEST_MODEL), indirect=True) +def test_stable_audio_t2a_online( + _require_stable_audio_access: None, + omni_server: OmniServer, + openai_client: OpenAIClientHandler, +) -> None: + """Stable Audio Open text-to-audio: `/v1/audio/generate` returns a non-empty WAV. + + Uses tiny steps / short duration to keep CI light, matching the offline smoke + test in `tests/e2e/offline_inference/test_stable_audio_expansion.py`. + """ + request_config = { + "json": { + "model": omni_server.model, + "input": T2A_PROMPT, + "audio_length": 2.0, + "num_inference_steps": 4, + "guidance_scale": 7.0, + "negative_prompt": "Low quality.", + "seed": 42, + "response_format": "wav", + }, + "timeout": 300, + } + responses = openai_client.send_audio_generate_http_request(request_config) + assert responses, "no response from /v1/audio/generate" + resp = responses[0] + assert resp.success, f"audio generate failed: {resp.status_code} {resp.error_message}" + assert resp.status_code == 200 diff --git a/tests/e2e/online_serving/test_step_audio2_expansion.py b/tests/e2e/online_serving/test_step_audio2_expansion.py index 9c8280b59..5542d81df 100644 --- a/tests/e2e/online_serving/test_step_audio2_expansion.py +++ b/tests/e2e/online_serving/test_step_audio2_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E expansion tests for Step-Audio2 online serving (weekly CI). @@ -94,22 +94,22 @@ def _build_audio_to_text_request_config(omni_server, audio_base64: str) -> dict: @hardware_test(res={"cuda": "H100"}, num_cards=1) -def test_single_audio_to_text_request(omni_server, openai_client) -> None: +def test_single_audio_to_text_request(omni_server, online_client) -> None: """Test a single audio-to-text request via the chat completions API.""" request_config = _build_audio_to_text_request_config( omni_server, _synthetic_audio_base64(duration_sec=2), ) - openai_client.send_omni_request(request_config) + online_client.send_omni_request(request_config) @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("request_num", [2, 4]) -def test_concurrent_audio_to_text_requests(omni_server, openai_client, request_num: int) -> None: +def test_concurrent_audio_to_text_requests(omni_server, online_client, request_num: int) -> None: """Test concurrent audio-to-text requests.""" request_config = _build_audio_to_text_request_config( omni_server, _synthetic_audio_base64(duration_sec=2), ) - responses = openai_client.send_omni_request(request_config, request_num=request_num) + responses = online_client.send_omni_request(request_config, request_num=request_num) assert len(responses) == request_num diff --git a/tests/e2e/online_serving/test_voxcpm2_tts.py b/tests/e2e/online_serving/test_voxcpm2_tts.py index 19ddc0401..0d0364138 100644 --- a/tests/e2e/online_serving/test_voxcpm2_tts.py +++ b/tests/e2e/online_serving/test_voxcpm2_tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online tests for VoxCPM2 with text input and audio output. @@ -46,12 +46,11 @@ tts_server_params = [ ] -@pytest.mark.advanced_model -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_001(omni_server, openai_client) -> None: +def test_text_to_audio_001(omni_server, online_client) -> None: """ Test zero-shot text-to-audio via OpenAI API. Deploy Setting: voxcpm2.yaml @@ -69,14 +68,14 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: "voice": "default", "min_audio_bytes": _MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) + online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_text_to_audio_002(omni_server, openai_client) -> None: +def test_text_to_audio_002(omni_server, online_client) -> None: """ Test zero-shot streaming text-to-audio via OpenAI API. Deploy Setting: voxcpm2.yaml @@ -95,4 +94,4 @@ def test_text_to_audio_002(omni_server, openai_client) -> None: "voice": "default", "min_audio_bytes": _MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py b/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py index 996585a1d..d630c8627 100644 --- a/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py +++ b/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E Online expansion tests for VoxCPM2 (voice clone, PCM, concurrency). @@ -12,11 +12,11 @@ import os import pytest from tests.helpers.mark import hardware_test -from tests.helpers.media import load_test_audio_data_url +from tests.helpers.media import get_asset_path from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -pytestmark = [pytest.mark.full_model, pytest.mark.tts] +pytestmark = [pytest.mark.slow, pytest.mark.tts] os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -25,7 +25,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0 _MIN_AUDIO_BYTES = 40_000 MAX_CONCURRENT = 4 -REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav") +REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) def get_prompt(prompt_type="text"): @@ -49,7 +49,7 @@ tts_server_params = [ @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_voice_clone_streaming_001(omni_server, openai_client) -> None: +def test_voice_clone_streaming_001(omni_server, online_client) -> None: """ Test voice-clone text-to-audio via OpenAI API. Deploy Setting: voxcpm2.yaml @@ -69,12 +69,12 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None: "ref_audio": REF_AUDIO_URL, "min_audio_bytes": _MIN_AUDIO_BYTES, } - openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) + online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -def test_response_format_001(omni_server, openai_client) -> None: +def test_response_format_001(omni_server, online_client) -> None: """ Test voice-clone non-stream PCM output via OpenAI API. Deploy Setting: voxcpm2.yaml @@ -94,4 +94,4 @@ def test_response_format_001(omni_server, openai_client) -> None: "min_audio_bytes": _MIN_AUDIO_BYTES, "min_hnr_db": -2.0, } - openai_client.send_audio_speech_request(request_config) + online_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_voxtral_tts_expansion.py b/tests/e2e/online_serving/test_voxtral_tts_expansion.py index dc8db9d73..278e563e8 100644 --- a/tests/e2e/online_serving/test_voxtral_tts_expansion.py +++ b/tests/e2e/online_serving/test_voxtral_tts_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E tests for Voxtral TTS online serving. @@ -31,9 +31,9 @@ pytestmark = [ @hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1) -def test_speech_english_basic(omni_server, openai_client) -> None: +def test_speech_english_basic(omni_server, online_client) -> None: """Test basic English TTS generation.""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "how are you", @@ -47,9 +47,9 @@ def test_speech_english_basic(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1) -def test_speech_english_streaming(omni_server, openai_client) -> None: +def test_speech_english_streaming(omni_server, online_client) -> None: """Test basic streaming English TTS generation (PCM via streaming API).""" - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Hello, how are you?", @@ -65,11 +65,11 @@ def test_speech_english_streaming(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_speech_different_voices(omni_server, openai_client) -> None: +def test_speech_different_voices(omni_server, online_client) -> None: """Test TTS with different voice presets.""" voices = ["casual_female", "neutral_male"] for voice in voices: - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Testing voice selection.", @@ -82,11 +82,11 @@ def test_speech_different_voices(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_speech_speed(omni_server, openai_client) -> None: +def test_speech_speed(omni_server, online_client) -> None: """Request with speed parameters.""" speeds = [0.5, 1, 1.5, 2, 2.5] for speed in speeds: - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "The boy was there when the sun rose.", @@ -100,7 +100,7 @@ def test_speech_speed(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_speech_instructions(omni_server, openai_client) -> None: +def test_speech_instructions(omni_server, online_client) -> None: """Request with instructions parameters.""" instructions = [ "Speak formally", @@ -109,7 +109,7 @@ def test_speech_instructions(omni_server, openai_client) -> None: "Speak with a chirpy happy voice", ] for instruction in instructions: - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "The boy was there when the sun rose.", @@ -134,11 +134,11 @@ def test_speech_instructions(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_speech_response_formats(omni_server, openai_client) -> None: +def test_speech_response_formats(omni_server, online_client) -> None: """Test TTS with different response formats.""" response_formats = ["wav", "mp3"] for response_format in response_formats: - openai_client.send_audio_speech_request( + online_client.send_audio_speech_request( { "model": omni_server.model, "input": "Testing various response formats.", @@ -151,7 +151,7 @@ def test_speech_response_formats(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_speech_batches(omni_server, openai_client) -> None: +def test_speech_batches(omni_server, online_client) -> None: """Test TTS batches.""" items = [ {"input": "The birch canoe slid on the smooth planks."}, @@ -161,7 +161,7 @@ def test_speech_batches(omni_server, openai_client) -> None: {"input": "Rice is often served in round bowls."}, ] - openai_client.send_audio_speech_batch_http_request( + online_client.send_audio_speech_batch_http_request( { "json": { "model": omni_server.model, diff --git a/tests/e2e/online_serving/test_wan22_expansion.py b/tests/e2e/online_serving/test_wan22_expansion.py index 85eee128a..322cde3cd 100644 --- a/tests/e2e/online_serving/test_wan22_expansion.py +++ b/tests/e2e/online_serving/test_wan22_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Comprehensive tests of diffusion features that are available in online serving mode. @@ -19,7 +22,7 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] @@ -140,7 +143,7 @@ def _get_wan22_feature_cases(): ) def test_wan22_diffusion_features( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): model_path = omni_server.model is_i2v_or_ti2v = any(kw in model_path for kw in ["I2V", "TI2V"]) @@ -176,4 +179,4 @@ def test_wan22_diffusion_features( if is_i2v_or_ti2v: request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" - openai_client.send_video_diffusion_request(request_config) + online_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_wan22_t2v.py b/tests/e2e/online_serving/test_wan22_t2v.py index e6fe38086..f3f21ae7d 100644 --- a/tests/e2e/online_serving/test_wan22_t2v.py +++ b/tests/e2e/online_serving/test_wan22_t2v.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Online serving smoke for ``Wan-AI/Wan2.2-T2V-A14B-Diffusers`` (text-to-video via ``/v1/videos``). @@ -18,7 +18,7 @@ import os import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -44,7 +44,7 @@ def _get_diffusion_feature_cases(model: str): @pytest.mark.advanced_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: +def test_text_to_video_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None: """Default Wan2.2 T2V smoke: async ``/v1/videos`` job completes and returns video bytes.""" request_config = { "model": omni_server.model, @@ -60,4 +60,4 @@ def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientH "seed": 42, }, } - openai_client.send_video_diffusion_request(request_config) + online_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py index d58bb9e32..cf24e34b0 100644 --- a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py +++ b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Comprehensive e2e tests of diffusion features for Wan2.1-VACE in online serving mode. @@ -24,9 +24,9 @@ Coverage: import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers" PROMPT = "A cat walking slowly across a sunlit garden path" @@ -150,9 +150,9 @@ def _get_vace_feature_cases(): _get_vace_feature_cases(), indirect=True, ) -def test_wan_2_1_vace(omni_server: OmniServer, openai_client: OpenAIClientHandler): +def test_wan_2_1_vace(omni_server: OmniServer, online_client: OnlineOmniClient): """Test VACE T2V generation with all supported diffusion acceleration features.""" - openai_client.send_video_diffusion_request( + online_client.send_video_diffusion_request( { "model": MODEL, "form_data": { diff --git a/tests/e2e/online_serving/test_z_image.py b/tests/e2e/online_serving/test_z_image.py index 2ef78ce0f..7e3af1132 100644 --- a/tests/e2e/online_serving/test_z_image.py +++ b/tests/e2e/online_serving/test_z_image.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Tests of common diffusion feature combinations in online serving mode for Z-Image. @@ -11,7 +14,7 @@ exercise Z-Image's supported parallel feature combinations: import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data MODEL = "Tongyi-MAI/Z-Image-Turbo" PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard." @@ -32,7 +35,7 @@ def _get_diffusion_feature_cases(): ] -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", @@ -41,7 +44,7 @@ def _get_diffusion_feature_cases(): ) def test_basic_001( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """Exercise supported Z-Image diffusion features in minimal CI cases.""" messages = dummy_messages_from_mix_data(content_text=PROMPT) @@ -57,10 +60,10 @@ def test_basic_001( }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", @@ -69,7 +72,7 @@ def test_basic_001( ) def test_different_sizes_001( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """Exercise supported Z-Image diffusion features in minimal CI cases.""" messages = dummy_messages_from_mix_data(content_text=PROMPT) @@ -85,4 +88,4 @@ def test_different_sizes_001( }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_zimage_expansion.py b/tests/e2e/online_serving/test_zimage_expansion.py index ac27267e5..c37a947f4 100644 --- a/tests/e2e/online_serving/test_zimage_expansion.py +++ b/tests/e2e/online_serving/test_zimage_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Tests of common diffusion feature combinations in online serving mode for Z-Image. @@ -14,7 +17,7 @@ exercise Z-Image's supported feature combinations: import pytest from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] @@ -110,7 +113,7 @@ def _get_diffusion_feature_cases(): ) def test_zimage( omni_server: OmniServer, - openai_client: OpenAIClientHandler, + online_client: OnlineOmniClient, ): """Exercise supported Z-Image diffusion features in minimal CI cases.""" request_config = { @@ -125,4 +128,4 @@ def test_zimage( }, } - openai_client.send_diffusion_request(request_config) + online_client.send_diffusion_request(request_config) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 055866bd5..fb8231a75 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -4,11 +4,12 @@ invariant to the specific attributes of vLLM config except in cases where we explicitly patch values that differ from vLLM. """ -import argparse import inspect import json import os import shutil +import sys +import types from types import SimpleNamespace from unittest.mock import Mock @@ -20,8 +21,9 @@ from vllm.engine.arg_utils import EngineArgs from vllm_omni.config.model import OmniModelConfig from vllm_omni.engine.arg_utils import OmniEngineArgs -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import build_engine_args_dict +from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -32,6 +34,11 @@ def test_sync_config_is_omni(): assert isinstance(cfg, OmniModelConfig) +def test_session_mode_reaches_omni_model_config(): + assert OmniEngineArgs().create_model_config().session_mode == "turn" + assert OmniEngineArgs(session_mode="duplex").create_model_config().session_mode == "duplex" + + def test_default_stage_id_is_concrete_int(): """Ensure `stage_id` stays safe for downstream arithmetic/indexing.""" engine_args = OmniEngineArgs() @@ -44,7 +51,102 @@ def test_default_stage_id_is_concrete_int(): assert cfg.stage_id == 0 -def test_multimodal_kwarg_overrides(mocker): +def test_full_payload_capability_reaches_omni_model_config(monkeypatch): + """Keep the pipeline capability intact through the engine adapter.""" + captured: dict[str, object] = {} + baseline_config = Mock() + worker_module = types.ModuleType("test_connector_worker") + + class ConnectorRunner(OmniConnectorModelRunnerMixin): + pass + + class ConnectorWorker: + model_runner_cls = ConnectorRunner + + worker_module.ConnectorWorker = ConnectorWorker + monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) + + monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda _self: None) + monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config) + + def capture_omni_config(_cls, model_config, **omni_kwargs): + assert model_config is baseline_config + captured.update(omni_kwargs) + return baseline_config + + monkeypatch.setattr( + OmniModelConfig, + "from_vllm_model_config", + classmethod(capture_omni_config), + ) + + OmniEngineArgs( + stage_id=1, + requires_full_payload_input=True, + worker_cls="test_connector_worker.ConnectorWorker", + ).create_model_config() + + assert captured["requires_full_payload_input"] is True + + +def test_stage_without_connector_configuration_accepts_plain_runner(monkeypatch): + worker_module = types.ModuleType("test_async_scheduler_worker") + + class WorkerWithoutConnector: + model_runner_cls = object + + worker_module.WorkerWithoutConnector = WorkerWithoutConnector + monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) + + args = OmniEngineArgs( + stage_id=1, + async_chunk=True, + worker_cls="test_async_scheduler_worker.WorkerWithoutConnector", + ) + + assert args.requires_full_payload_input is False + + +def test_full_payload_capability_requires_selected_worker_connector(monkeypatch): + worker_module = types.ModuleType("test_worker_without_connector") + + class WorkerWithoutConnector: + model_runner_cls = object + + worker_module.WorkerWithoutConnector = WorkerWithoutConnector + monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) + + with pytest.raises(ValueError, match="does not provide an Omni connector model runner"): + OmniEngineArgs( + stage_id=1, + requires_full_payload_input=True, + worker_cls="test_worker_without_connector.WorkerWithoutConnector", + ) + + +def test_full_payload_capability_validates_platform_selected_worker(monkeypatch): + worker_module = types.ModuleType("test_platform_worker_without_connector") + + class WorkerWithoutConnector: + model_runner_cls = object + + worker_module.WorkerWithoutConnector = WorkerWithoutConnector + monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) + monkeypatch.setattr( + current_omni_platform, + "get_omni_ar_worker_cls", + lambda: "test_platform_worker_without_connector.WorkerWithoutConnector", + ) + + with pytest.raises(ValueError, match="does not provide an Omni connector model runner"): + OmniEngineArgs( + stage_id=1, + requires_full_payload_input=True, + worker_type="ar", + ) + + +def test_multimodal_kwarg_overrides(monkeypatch): """Ensure that overrides in the multimodal config are preserved.""" sig = inspect.signature(OmniEngineArgs) default_mm_cache = sig.parameters["mm_processor_cache_gb"].default @@ -58,8 +160,8 @@ def test_multimodal_kwarg_overrides(mocker): assert self.mm_processor_cache_gb == override_val return fake_model_config - mocker.patch.object(EngineArgs, "create_model_config", _fake_parent_create_model_config) - mocker.patch.object(OmniModelConfig, "from_vllm_model_config", side_effect=lambda model_config, **_: model_config) + monkeypatch.setattr(EngineArgs, "create_model_config", _fake_parent_create_model_config) + monkeypatch.setattr(OmniModelConfig, "from_vllm_model_config", lambda model_config, **_: model_config) cfg = OmniEngineArgs( model="Qwen/Qwen2-VL-2B-Instruct", @@ -149,19 +251,6 @@ def test_qwen3_tts_startup_task_type_is_validated(): ) -def test_from_cli_args_picks_up_stage_configs_path(): - """from_cli_args should pick up stage_configs_path from namespace.""" - ns = argparse.Namespace( - model="facebook/opt-125m", - stage_configs_path="/some/path.yaml", - custom_pipeline_args=None, - ) - - args = OmniEngineArgs.from_cli_args(ns) - assert args.stage_configs_path == "/some/path.yaml" - assert args.custom_pipeline_args is None - - def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch): """Ensure Code2Wav mirrors stage max_model_len into nested HF overrides. @@ -197,6 +286,28 @@ def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch): } +def test_remote_tokenizer_subfolder_download_does_not_report_failure(tmp_path, monkeypatch, mocker): + """A successful remote tokenizer download must not enter the error path.""" + tokenizer_dir = tmp_path / "CosyVoice-BlankEN" + tokenizer_dir.mkdir() + baseline_config = Mock() + warning = mocker.patch("vllm_omni.engine.arg_utils.logger.warning") + + monkeypatch.setattr("huggingface_hub.snapshot_download", lambda *args, **kwargs: str(tmp_path)) + monkeypatch.setattr(OmniEngineArgs, "_patch_empty_hf_config", lambda *args, **kwargs: None) + monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config) + monkeypatch.setattr( + OmniModelConfig, + "from_vllm_model_config", + classmethod(lambda cls, model_config, **omni_kwargs: model_config), + ) + + args = OmniEngineArgs(model="remote/cosyvoice3", model_arch="CosyVoice3Model") + assert args.create_model_config() is baseline_config + assert args.tokenizer == str(tokenizer_dir) + warning.assert_not_called() + + def test_patch_missing_local_hf_config(tmp_path): """Native model bundles may only contain a checkpoints/ directory.""" (tmp_path / "checkpoints").mkdir() @@ -418,130 +529,6 @@ def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size(): assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size -def test_stage_configs_path_field(): - """OmniEngineArgs with stage_configs_path should construct without error.""" - args = OmniEngineArgs(stage_configs_path="/some/path.yaml") - assert args.stage_configs_path == "/some/path.yaml" - - -def test_strip_single_engine_args(): - """_strip_single_engine_args should remove EngineArgs fields but keep omni fields.""" - kwargs = { - # Parent EngineArgs fields — stripped unless explicitly allowlisted - "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}', - "tensor_parallel_size": 4, - "gpu_memory_utilization": 0.9, - "model": "some/model", - # Parent field that should be kept (allowlisted) - "worker_extension_cls": "some.Extension", - # OmniEngineArgs-only / non-engine fields — should pass through - "stage_configs_path": "/path/to/yaml", - "custom_pipeline_args": {"pipeline_class": "my.Pipeline"}, - "mode": "text-to-image", - "lora_path": "/some/lora", - } - - filtered = AsyncOmniEngine._strip_single_engine_args(kwargs) - - # Stripped — parent EngineArgs fields - assert "compilation_config" not in filtered - assert filtered["tensor_parallel_size"] == 4 - assert "gpu_memory_utilization" not in filtered - assert "model" not in filtered - - # Stripped — orchestrator-level OmniEngineArgs field - assert "stage_configs_path" not in filtered - - # Kept - assert filtered["worker_extension_cls"] == "some.Extension" - assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"} - assert filtered["mode"] == "text-to-image" - assert filtered["lora_path"] == "/some/lora" - - -def test_strip_single_engine_args_model_does_not_trigger_warning(mocker): - """model is always in kwargs (callers set it via from_cli_args/asdict), - so it should not cause the override warning by itself or appear in it.""" - mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning") - - # Typical caller kwargs: model is always present, no other parent - # EngineArgs fields are explicitly overridden. - AsyncOmniEngine._strip_single_engine_args( - { - "model": "some/model", - "custom_pipeline_args": {"pipeline_class": "my.Pipeline"}, - } - ) - mock_warn.assert_not_called() - - # When there *are* genuinely surprising overrides alongside model, - # the warning should mention them but not model. Keep-listed fields such as - # tensor_parallel_size are intentionally passed through and should not warn. - AsyncOmniEngine._strip_single_engine_args( - { - "model": "some/model", - "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}', - "tensor_parallel_size": 4, - "custom_pipeline_args": {"pipeline_class": "my.Pipeline"}, - } - ) - mock_warn.assert_called_once() - warned_args = mock_warn.call_args[0][-1] # the formatted arg list - assert "compilation_config" in warned_args - assert "tensor_parallel_size" not in warned_args - assert "model" not in warned_args - - -def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path): - """New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs. - - Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping - those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni - accuracy from ~78% to ~70%. - """ - from vllm_omni.entrypoints.utils import is_new_format_deploy_config - - yaml_path = tmp_path / "minicpm_deploy.yaml" - yaml_path.write_text( - """ -pipeline: minicpmo_4_5 -async_chunk: true -stages: - - stage_id: 0 - devices: "0" -""".strip(), - encoding="utf-8", - ) - path = str(yaml_path) - assert is_new_format_deploy_config(path) - - kwargs = { - "stage_configs_path": path, - "trust_remote_code": True, - "interleave_mm_strings": True, - "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}, - "gpu_memory_utilization": 0.55, - } - # Mirror AsyncOmniEngine._resolve_stage_configs strip decision. - if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]): - base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs)) - else: - base = dict(kwargs) - assert base.get("interleave_mm_strings") is True - assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}} - # Parent EngineArgs fields also remain available for deploy merge. - assert base.get("gpu_memory_utilization") == 0.55 - - legacy = tmp_path / "legacy.yaml" - legacy.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8") - assert not is_new_format_deploy_config(str(legacy)) - legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy)) - stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs) - assert "interleave_mm_strings" not in stripped - assert "media_io_kwargs" not in stripped - assert "gpu_memory_utilization" not in stripped - - # For https://github.com/vllm-project/vllm-omni/issues/3293 def test_tensor_parallel_size_none_is_handled(): """Ensure the tensor parallel size of None isn't forwarded.""" diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py index 40f89df84..56b25704d 100644 --- a/tests/engine/test_async_omni_engine_abort.py +++ b/tests/engine/test_async_omni_engine_abort.py @@ -68,7 +68,7 @@ async def test_abort(): with ExitStack() as after: engine = AsyncOmni( model=model, - stage_configs_path=stage_config, + deploy_config=stage_config, ) after.callback(engine.shutdown) diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py new file mode 100644 index 000000000..86276ec66 --- /dev/null +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for acknowledged abort on AsyncOmniEngine.""" + +from __future__ import annotations + +import asyncio +import queue +import threading +from types import SimpleNamespace + +import pytest +from pytest_mock import MockerFixture + +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.messages import AbortRequestMessage, AbortResultMessage +from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_engine(request_q: queue.Queue, rpc_q: queue.Queue) -> AsyncOmniEngine: + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = SimpleNamespace(sync_q=request_q) + engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) + engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) + return engine + + +def test_abort_result_message_correlation_key(): + msg = AbortResultMessage(rpc_id="abc", success=True) + assert msg.rpc_correlation_key == ("abort", "abc") + + +def test_sync_abort_is_fire_and_forget(): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + + engine.abort(["req-1"]) + + msg = request_q.get_nowait() + assert isinstance(msg, AbortRequestMessage) + assert msg.request_ids == ["req-1"] + assert msg.rpc_id is None + assert rpc_q.empty() + engine._correlated_rpc_client.close() + + +def test_abort_async_waits_for_ack(mocker: MockerFixture): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="abort-rpc-1"), + ) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-a", "req-b"], timeout=1)) + msg = await asyncio.to_thread(request_q.get, True, 1) + assert isinstance(msg, AbortRequestMessage) + assert msg.request_ids == ["req-a", "req-b"] + assert msg.rpc_id == "abort-rpc-1" + assert not task.done() + + rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True)) + result = await task + assert result == [] + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="abort-rpc-err"), + ) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-x"], timeout=1)) + await asyncio.to_thread(request_q.get, True, 1) + rpc_q.put( + AbortResultMessage( + rpc_id="abort-rpc-err", + success=False, + error="stage abort failed", + ) + ) + with pytest.raises(RuntimeError, match="stage abort failed"): + await task + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_times_out_without_result(mocker: MockerFixture): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="abort-rpc-timeout"), + ) + + async def _run() -> None: + with pytest.raises(TimeoutError, match="abort timed out"): + await engine.abort_async(["req-timeout"], timeout=0.05) + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture): + class SignallingQueue(queue.Queue): + def __init__(self) -> None: + super().__init__(maxsize=1) + self.put_attempted = threading.Event() + + def put(self, item, block=True, timeout=None): + self.put_attempted.set() + return super().put(item, block=block, timeout=timeout) + + request_q = SignallingQueue() + request_q.put("queue-is-full") + request_q.put_attempted.clear() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_omni_engine.uuid.uuid4", + return_value=SimpleNamespace(hex="blocked-abort"), + ) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-blocked"], timeout=1)) + assert await asyncio.to_thread(request_q.put_attempted.wait, 1) + assert not task.done() + + assert request_q.get(timeout=1) == "queue-is-full" + msg = await asyncio.to_thread(request_q.get, True, 1) + assert isinstance(msg, AbortRequestMessage) + assert msg.rpc_id == "blocked-abort" + rpc_q.put(AbortResultMessage(rpc_id="blocked-abort", success=True)) + await task + + try: + asyncio.run(_run()) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_is_noop_after_shutdown_starts(): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = True + engine.request_queue = None + engine._correlated_rpc_client = None + + asyncio.run(engine.abort_async(["req-late"])) + + +def test_abort_async_is_noop_for_empty_request_ids(): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = None + engine._correlated_rpc_client = None + + asyncio.run(engine.abort_async([])) + + +def test_abort_async_tolerates_request_queue_close_during_shutdown( + mocker: MockerFixture, +): + class SyntheticSyncQueueShutDownError(Exception): + pass + + rpc_q: queue.Queue = queue.Queue() + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + + class ClosingRequestQueue: + def put(self, item) -> None: + del item + engine._shutdown_called = True + raise SyntheticSyncQueueShutDownError + + request_q = ClosingRequestQueue() + engine.request_queue = SimpleNamespace(sync_q=request_q) + engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + SyntheticSyncQueueShutDownError, + ) + + try: + asyncio.run(engine.abort_async(["req-submit-race"], timeout=1)) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_surfaces_request_queue_close_while_engine_is_live( + mocker: MockerFixture, +): + class SyntheticSyncQueueShutDownError(Exception): + pass + + class ClosedRequestQueue: + def put(self, item) -> None: + del item + raise SyntheticSyncQueueShutDownError + + request_q = ClosedRequestQueue() + rpc_q: queue.Queue = queue.Queue() + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = SimpleNamespace(sync_q=request_q) + engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + SyntheticSyncQueueShutDownError, + ) + + try: + with pytest.raises(SyntheticSyncQueueShutDownError): + asyncio.run(engine.abort_async(["req-live-submit"], timeout=1)) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_tolerates_rpc_router_close_during_shutdown(): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-ack-race"], timeout=1)) + await asyncio.to_thread(request_q.get, True, 1) + engine._shutdown_called = True + engine._correlated_rpc_client.close() + await task + + asyncio.run(_run()) + + +def test_abort_async_surfaces_rpc_router_close_while_engine_is_live(): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-live"], timeout=1)) + await asyncio.to_thread(request_q.get, True, 1) + engine._correlated_rpc_client.close() + with pytest.raises(RuntimeError, match="RPC result router closed"): + await task + + asyncio.run(_run()) diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 740064173..7ad111e25 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest import torch from pytest_mock import MockerFixture @@ -14,6 +17,104 @@ from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +class _SyntheticSyncQueueShutDownError(Exception): + pass + + +def test_abort_is_noop_for_empty_request_ids(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + + engine.abort([]) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_abort_is_noop_after_shutdown_starts(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = True + engine.request_queue = mocker.Mock() + + engine.abort(["req-1"]) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_abort_tolerates_queue_close_race_during_shutdown(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + _SyntheticSyncQueueShutDownError, + ) + + def close_queue_during_abort(*args, **kwargs): + del args, kwargs + engine._shutdown_called = True + raise _SyntheticSyncQueueShutDownError + + engine.request_queue.sync_q.put.side_effect = close_queue_during_abort + + engine.abort(["req-1"]) + + +def test_abort_tolerates_legacy_janus_close_error_during_shutdown(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + None, + ) + + def close_queue_during_abort(*args, **kwargs): + del args, kwargs + engine._shutdown_called = True + raise RuntimeError("Operation on the closed queue is forbidden") + + engine.request_queue.sync_q.put.side_effect = close_queue_during_abort + + engine.abort(["req-1"]) + + +def test_abort_surfaces_unrelated_legacy_runtime_error_during_shutdown( + mocker: MockerFixture, +): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + None, + ) + + def fail_during_abort(*args, **kwargs): + del args, kwargs + engine._shutdown_called = True + raise RuntimeError("unrelated queue failure") + + engine.request_queue.sync_q.put.side_effect = fail_during_abort + + with pytest.raises(RuntimeError, match="unrelated queue failure"): + engine.abort(["req-1"]) + + +def test_abort_surfaces_unexpected_closed_queue(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + _SyntheticSyncQueueShutDownError, + ) + engine.request_queue.sync_q.put.side_effect = _SyntheticSyncQueueShutDownError + + with pytest.raises(_SyntheticSyncQueueShutDownError): + engine.abort(["req-1"]) + + def _make_engine_core_request(request_id: str = "req-1") -> EngineCoreRequest: return EngineCoreRequest( request_id=request_id, @@ -67,6 +168,39 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock output_processor.add_request.assert_not_called() +def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.default_sampling_params_list = [params] + engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")] + engine.supported_tasks = ("speech",) + + input_processor = mocker.Mock() + input_processor.process_inputs.return_value = _make_engine_core_request() + engine.input_processor = input_processor + engine.output_processors = [mocker.Mock()] + + seen_global_ids: list[list[str]] = [] + + def transform(prompt, _sampling_params): + seen_global_ids.append(prompt["additional_information"]["global_request_id"]) + return prompt + + engine.prompt_transform_func = transform + prompt = {"prompt_token_ids": [1, 1, 1]} + + message = engine._build_add_request_message( + request_id="req-1", + prompt=prompt, + sampling_params_list=[params], + final_stage_id=0, + arrival_time=0.0, + ) + + assert seen_global_ids == [["req-1"]] + assert message.original_prompt["additional_information"]["global_request_id"] == ["req-1"] + + def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: MockerFixture): engine = object.__new__(AsyncOmniEngine) params = SamplingParams(max_tokens=8) @@ -371,3 +505,106 @@ def test_build_add_request_message_releases_preselected_replica_on_preprocess_er ) assert stage_pool.get_bound_replica_id("req-error") is None + + +def test_cfg_companion_build_failure_admits_nothing(mocker: MockerFixture): + """A guided request is all-or-nothing. + + Admitting the parent before its companion is built leaves an orphan: a + model whose guidance is mandatory never completes the pair, so the request + holds scheduler and KV capacity for the scheduler's whole hold budget and + then produces no audio. The build has to raise before anything is enqueued. + """ + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = mocker.Mock(side_effect=ValueError("cannot build the null twin")) + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.request_queue = mocker.Mock() + + with pytest.raises(ValueError, match="null twin"): + engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_cfg_companion_processing_failure_admits_nothing(mocker: MockerFixture): + """The same holds when expansion succeeds but input processing does not.""" + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [ + ExpandedPrompt( + prompt={"prompt": "negative"}, + role="cfg_text", + request_id_suffix="__cfg_text", + ) + ] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.input_processor.process_inputs.side_effect = RuntimeError("tokenizer rejected the prompt") + engine.request_queue = mocker.Mock() + + with pytest.raises(RuntimeError, match="tokenizer rejected"): + engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_cfg_companion_build_returns_messages_without_enqueueing(mocker: MockerFixture): + """Construction and admission are separate steps.""" + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [ + ExpandedPrompt( + prompt={"prompt": "negative"}, + role="cfg_text", + request_id_suffix="__cfg_text", + ) + ] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text") + engine.request_queue = mocker.Mock() + + companions = engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + assert len(companions) == 1 + assert companions[0].companion_id == "req__cfg_text" + assert companions[0].parent_id == "req" + engine.request_queue.sync_q.put.assert_not_called() + + +def test_cfg_expansion_returning_nothing_is_not_an_error(mocker: MockerFixture): + """An unguided request expands to no companions and admits normally.""" + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.request_queue = mocker.Mock() + + assert ( + engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + == [] + ) + engine.request_queue.sync_q.put.assert_not_called() diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 27aee8706..c1e56b7e1 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1,6 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import concurrent.futures import contextlib import importlib +import json import os import time import types @@ -15,6 +19,7 @@ from vllm_omni.engine.stage_init_utils import ( ReplicaInitPlan, build_stage0_input_processor, compute_replica_layout, + split_devices_for_replicas, ) from vllm_omni.engine.stage_runtime import StageRuntime @@ -161,6 +166,16 @@ def _make_diffusion_plan( ) +def _make_stage_runtime() -> StageRuntime: + return StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + async_chunk=False, + ) + + def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): """Regression test for forward references in make_async_mp_client.""" import vllm_omni.engine.stage_engine_core_client as client_mod @@ -180,7 +195,6 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp engine.model = "dummy-model" engine.config_path = "dummy-config" engine.single_stage_mode = False - engine.diffusion_batch_size = 1 engine.async_chunk = False engine.tokenizer = None engine._single_stage_id_filter = None @@ -221,6 +235,75 @@ def test_compute_replica_layout_splits_diffusion_devices_by_world_size(): assert replica_devices_map == {0: ["0,1", "2,3"]} +@pytest.mark.parametrize("num_replicas", [1, 2, 4]) +def test_split_devices_for_replicas_returns_one_slot_per_replica_when_devices_unset(num_replicas): + """``devices`` is optional, and the result is indexed by replica id. + + A stage that declares ``num_replicas`` without ``devices`` lets every + replica inherit the launcher's CUDA_VISIBLE_DEVICES, so each one still + needs its own (empty) slot. + """ + assert split_devices_for_replicas(None, num_replicas, 1, 0) == [None] * num_replicas + + +def test_compute_replica_layout_covers_every_replica_when_devices_unset(): + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={}, + runtime=types.SimpleNamespace(num_replicas=3), + ) + + replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg]) + + assert replicas_per_stage == [3] + assert replica_devices_map == {0: [None, None, None]} + + +def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatch): + """Regression: a multi-replica stage with no ``devices`` must still plan. + + ``_build_logical_stage_init_plans`` indexes ``replica_devices_map`` by + replica id, so a short list raised ``IndexError`` for every replica after + the first. + """ + import vllm_omni.engine.stage_runtime as runtime_mod + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={}, + runtime=types.SimpleNamespace(num_replicas=3), + ) + runtime = StageRuntime( + stage_configs=[stage_cfg], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + async_chunk=False, + ) + + monkeypatch.setattr( + runtime_mod, + "extract_legacy_stage_metadata", + lambda cfg: _make_llm_metadata(cfg.stage_id), + ) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "build_vllm_config", lambda *_args, **_kwargs: (types.SimpleNamespace(), object)) + + replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg]) + stage_plans = runtime._build_logical_stage_init_plans( + omni_transfer_config=None, + replicas_per_stage=replicas_per_stage, + replica_devices_map=replica_devices_map, + ) + + assert [replica.replica_id for replica in stage_plans[0].replicas] == [0, 1, 2] + assert [replica.stage_cfg.runtime.devices for replica in stage_plans[0].replicas] == [None, None, None] + + def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): shared = types.SimpleNamespace(name="shared") extra = types.SimpleNamespace(name="extra") @@ -233,6 +316,48 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): assert cleanup_clients == [shared, extra] +def test_initialize_local_diffusion_replica_scopes_runtime_env(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.platforms import current_omni_platform + + runtime = _make_stage_runtime() + plan = _make_diffusion_plan(0, stage_id=0).replicas[0] + + runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV" + device_env_var = current_omni_platform.device_control_env_var + runtime._init_visible_devices_baseline = "0,1" + plan.metadata.runtime_cfg = { + "devices": "0", + "env": {runtime_env_var: "stage-value"}, + } + monkeypatch.delenv(runtime_env_var, raising=False) + monkeypatch.setenv(device_env_var, "0,1") + + captured: dict[str, str | None] = {} + monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None) + + def _capture_launch_diffusion_stage_replica(**_kwargs): + captured["runtime_env"] = os.environ.get(runtime_env_var) + captured["device_env"] = os.environ.get(device_env_var) + raise RuntimeError("stop after capturing launch environment") + + monkeypatch.setattr( + runtime_mod, + "launch_diffusion_stage_replica", + _capture_launch_diffusion_stage_replica, + ) + + with pytest.raises(RuntimeError, match="stop after capturing launch environment"): + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) + + assert captured == { + "runtime_env": "stage-value", + "device_env": "0", + } + assert runtime_env_var not in os.environ + assert os.environ[device_env_var] == "0,1" + + def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.engine.stage_engine_startup import StageReplicaResources @@ -243,7 +368,6 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -271,16 +395,23 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc os.environ[env_var] = old_env -def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch): +@pytest.mark.parametrize( + ("num_stages", "expected_inline"), + [(1, True), (2, False)], +) +def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag( + monkeypatch, + num_stages, + expected_inline, +): import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.engine.stage_engine_startup import StageReplicaResources runtime = StageRuntime( - stage_configs=[types.SimpleNamespace()], + stage_configs=[types.SimpleNamespace()] * num_stages, model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=4, async_chunk=False, ) @@ -293,7 +424,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline def _capture_launch_diffusion_stage_replica(**kwargs): captured["stage_id"] = kwargs["metadata"].stage_id captured["stage_init_timeout"] = kwargs["stage_init_timeout"] - captured["batch_size"] = kwargs["batch_size"] captured["use_inline"] = kwargs["use_inline"] captured["omni_master_server"] = kwargs["omni_master_server"] return types.SimpleNamespace(), StageReplicaResources() @@ -305,12 +435,238 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline assert captured == { "stage_id": 0, "stage_init_timeout": 302, - "batch_size": 4, - "use_inline": True, + "use_inline": expected_inline, "omni_master_server": None, } +def test_initialize_local_diffusion_replica_uses_explicit_inline_opt_in(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.engine.stage_engine_startup import StageReplicaResources + + runtime = _make_stage_runtime() + plan = _make_diffusion_plan(0, stage_id=1).replicas[0] + plan.stage_cfg.engine_args["inline_diffusion"] = True + captured: dict[str, object] = {} + + monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None) + monkeypatch.setattr( + runtime_mod, + "launch_diffusion_stage_replica", + lambda **kwargs: ( + captured.update(use_inline=kwargs["use_inline"]), + (types.SimpleNamespace(), StageReplicaResources()), + )[1], + ) + + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) + + assert captured["use_inline"] is True + + +def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = _make_stage_runtime() + plan = _make_llm_plan(0, stage_id=0, vllm_config=object()).replicas[0] + plan.engine_args_dict = {} + + runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV" + runtime._init_visible_devices_baseline = "0" + plan.metadata.runtime_cfg = { + "devices": "0", + "env": {runtime_env_var: "stage-value"}, + } + monkeypatch.setenv(runtime_env_var, "parent-value") + # Pin the logical->physical device mapping so the assertion below is + # deterministic regardless of the ambient CUDA_VISIBLE_DEVICES (which on a + # rebase host can be e.g. "3,4", mapping logical device 0 -> physical 3). + runtime._init_visible_devices_baseline = "0" + + captured: dict[str, str | None] = {} + + @contextlib.contextmanager + def _capture_launch_stage_replica(*, stage_visible_devices, **_kwargs): + captured["runtime_env"] = os.environ.get(runtime_env_var) + captured["stage_visible_devices"] = stage_visible_devices + yield None + + monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica) + + with pytest.raises(RuntimeError, match="launcher returned no resources"): + runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) + + assert captured == { + "runtime_env": "stage-value", + "stage_visible_devices": "0", + } + assert os.environ[runtime_env_var] == "parent-value" + + +def test_initialize_diffusion_stage_preserves_configured_max_num_seqs(monkeypatch): + import vllm_omni.diffusion.stage_diffusion_client as client_mod + import vllm_omni.engine.stage_init_utils as init_mod + + od_config = types.SimpleNamespace(max_num_seqs=4) + captured: dict[str, object] = {} + monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config) + + def _capture_client(model, config, metadata, stage_init_timeout, use_inline): + captured.update( + model=model, + config=config, + metadata=metadata, + stage_init_timeout=stage_init_timeout, + use_inline=use_inline, + ) + return object() + + monkeypatch.setattr(client_mod, "create_diffusion_client", _capture_client) + metadata = _make_diffusion_metadata(0) + + init_mod.initialize_diffusion_stage( + 0, + "dummy-model", + types.SimpleNamespace(), + metadata, + stage_init_timeout=12, + use_inline=True, + ) + + assert od_config.max_num_seqs == 4 + assert captured == { + "model": "dummy-model", + "config": od_config, + "metadata": metadata, + "stage_init_timeout": 12, + "use_inline": True, + } + + +def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkeypatch): + import vllm_omni.diffusion.stage_diffusion_client as client_mod + import vllm_omni.diffusion.stage_diffusion_proc as proc_mod + import vllm_omni.engine.stage_engine_startup as startup_mod + + od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1)) + monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config) + monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: []) + monkeypatch.setattr( + startup_mod, + "register_stage_with_omni_master", + lambda **kwargs: types.SimpleNamespace( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + ), + ) + + omni_master_server = types.SimpleNamespace( + address="127.0.0.1", + port=25000, + release_route_port_reservations=lambda *args, **kwargs: None, + ) + proc_manager = types.SimpleNamespace( + addresses=types.SimpleNamespace( + inputs=["tcp://127.0.0.1:26002"], + outputs=["tcp://127.0.0.1:26003"], + ) + ) + monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager) + sentinel_client = object() + monkeypatch.setattr( + client_mod.StageDiffusionClient, + "from_addresses", + lambda metadata, **kwargs: sentinel_client, + ) + + result, resources = startup_mod.launch_diffusion_stage_replica( + model="dummy-model", + stage_config=types.SimpleNamespace(), + metadata=types.SimpleNamespace(stage_id=0), + stage_init_timeout=12, + use_inline=False, + omni_master_server=omni_master_server, + ) + + assert result is sentinel_client + assert od_config.max_num_seqs == 4 + assert resources.manager is proc_manager + + +def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch): + import vllm_omni.diffusion.stage_diffusion_client as client_mod + import vllm_omni.engine.stage_init_utils as init_mod + + od_config = types.SimpleNamespace(max_num_seqs=8, step_execution=True) + monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config) + monkeypatch.setattr(client_mod, "create_diffusion_client", lambda *args: object()) + metadata = _make_diffusion_metadata(0) + + init_mod.initialize_diffusion_stage( + 0, + "dummy-model", + types.SimpleNamespace(), + metadata, + stage_init_timeout=12, + use_inline=True, + ) + + assert od_config.max_num_seqs == 8 + + +def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(monkeypatch): + import vllm_omni.diffusion.stage_diffusion_client as client_mod + import vllm_omni.diffusion.stage_diffusion_proc as proc_mod + import vllm_omni.engine.stage_engine_startup as startup_mod + + od_config = types.SimpleNamespace( + max_num_seqs=8, + step_execution=True, + parallel_config=types.SimpleNamespace(world_size=1), + ) + monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config) + monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: []) + monkeypatch.setattr( + startup_mod, + "register_stage_with_omni_master", + lambda **kwargs: types.SimpleNamespace( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + ), + ) + omni_master_server = types.SimpleNamespace( + address="127.0.0.1", + port=25000, + release_route_port_reservations=lambda *args, **kwargs: None, + ) + proc_manager = types.SimpleNamespace( + addresses=types.SimpleNamespace( + inputs=["tcp://127.0.0.1:26002"], + outputs=["tcp://127.0.0.1:26003"], + ) + ) + monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager) + monkeypatch.setattr( + client_mod.StageDiffusionClient, + "from_addresses", + lambda metadata, **kwargs: object(), + ) + + startup_mod.launch_diffusion_stage_replica( + model="dummy-model", + stage_config=types.SimpleNamespace(), + metadata=types.SimpleNamespace(stage_id=0), + stage_init_timeout=12, + use_inline=False, + omni_master_server=omni_master_server, + ) + + assert od_config.max_num_seqs == 8 + + def test_stage_runtime_initializes_stage_pools(monkeypatch): import vllm_omni.engine.stage_runtime as runtime_mod @@ -319,7 +675,6 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch): model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -385,7 +740,6 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch): model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, log_stats=True, ) @@ -409,11 +763,12 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch): monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: []) monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica) - monkeypatch.setattr( - runtime_mod.StageEngineCoreClientBase, - "make_async_mp_client", - lambda **kwargs: captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client, - ) + + def _make_async_mp_client(**kwargs): + captured["client_log_stats"] = kwargs["log_stats"] + return stage_client + + monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", _make_async_mp_client) assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client assert captured["log_stats"] is True @@ -428,7 +783,6 @@ def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch): model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, log_stats=True, ) @@ -481,7 +835,6 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -522,7 +875,6 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk model="dummy-model", config_path="dummy-config", stage_init_timeout=123, - diffusion_batch_size=1, async_chunk=False, ) @@ -560,7 +912,6 @@ def test_remote_replicas_use_distinct_init_group_keys(): model="dummy-model", config_path="dummy-config", stage_init_timeout=123, - diffusion_batch_size=1, async_chunk=False, ) plan = _make_llm_plan( @@ -586,7 +937,6 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -623,7 +973,6 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_ model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) @@ -635,7 +984,7 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_ def _initialize_stage_replicas(_stage_plans, _stage_init_timeout): exc = RuntimeError("replica launch failed") - exc._initialized_clients_by_stage = {0: [None, initialized_client]} + setattr(exc, "_initialized_clients_by_stage", {0: [None, initialized_client]}) raise exc monkeypatch.setattr(runtime, "_initialize_stage_replicas", _initialize_stage_replicas) @@ -662,7 +1011,6 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag model="dummy-model", config_path="dummy-config", stage_init_timeout=1, - diffusion_batch_size=1, async_chunk=False, ) stage_init_timeout = 302 @@ -732,6 +1080,8 @@ def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_p default_sampling_params={}, ) + (tmp_path / "llm").mkdir() + engine_args = build_engine_args_dict( stage_cfg, str(tmp_path), @@ -793,6 +1143,9 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path): default_sampling_params={}, ) + (tmp_path / "llm").mkdir() + (tmp_path / "tokenizer").mkdir() + engine_args = build_engine_args_dict( stage_cfg, str(tmp_path), @@ -803,6 +1156,445 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path): assert engine_args["tokenizer"] == os.path.join(str(tmp_path), "tokenizer") +def test_build_engine_args_rejects_missing_subdir_in_local_model_dir(tmp_path): + """A local model directory without the stage subfolder must fail closed. + + Forwarding the joined path as ``EngineArgs.model`` sends a non-directory + string to HuggingFace, which reports it as a malformed repo id. + """ + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + (tmp_path / "tokenizer").mkdir() + + with pytest.raises(RuntimeError, match="language_model"): + build_engine_args_dict(stage_cfg, str(tmp_path)) + + +# A real snapshot subfolder always carries artifacts; empty directories would +# assert the broken exists-means-complete predicate this suite regresses. +_SUBDIR_ARTIFACT = {"language_model": "model.safetensors", "tokenizer": "tokenizer.json"} + + +def _make_snapshot(root, subdirs): + for subdir in subdirs: + folder = root / subdir + folder.mkdir(parents=True, exist_ok=True) + (folder / _SUBDIR_ARTIFACT.get(subdir, "data.bin")).write_text("x") + return str(root) + + +def test_build_engine_args_pulls_stage_subdirs_missing_from_cached_snapshot(monkeypatch, tmp_path): + """Regression for issue #6638. + + ``snapshot_download(local_files_only=True)`` hands back a snapshot root as + soon as the repo is partly cached, even when this stage's subfolders were + never materialized. Stage init must fetch exactly those subfolders instead + of joining onto a path that exists nowhere. + """ + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "deadbeef" + _make_snapshot(snapshot, ["tokenizer"]) + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if kwargs.get("local_files_only"): + return str(snapshot) + _make_snapshot(snapshot, ["language_model"]) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert engine_args["model"] == os.path.join(str(snapshot), "language_model") + assert engine_args["tokenizer"] == os.path.join(str(snapshot), "tokenizer") + assert os.path.isdir(engine_args["model"]) + # Only the stage's own subfolders are pulled, not the whole repo. + assert calls[-1]["allow_patterns"] == ["language_model/*", "tokenizer/*"] + + +def test_build_engine_args_skips_hub_call_when_cached_snapshot_is_complete(monkeypatch, tmp_path): + """A warm cache stays offline-friendly: no outgoing Hub request.""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "deadbeef" + _make_snapshot(snapshot, ["language_model", "tokenizer"]) + + def fake_snapshot_download(self, repo_id, **kwargs): + assert kwargs.get("local_files_only"), "warm cache must not reach the Hub" + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert engine_args["model"] == os.path.join(str(snapshot), "language_model") + + +def test_build_engine_args_redownloads_a_partial_subdir(monkeypatch, tmp_path): + """A subfolder holding only config.json is an interrupted download. + + The warm-cache early return must not accept it: converting the Hub ID into + that local path leaves vLLM's loader without any fallback for the missing + weights. + """ + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "deadbeef" + partial = snapshot / "language_model" + partial.mkdir(parents=True) + (partial / "config.json").write_text("{}") + _make_snapshot(snapshot, ["tokenizer"]) + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if kwargs.get("local_files_only"): + return str(snapshot) + _make_snapshot(snapshot, ["language_model"]) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert any(not call.get("local_files_only") for call in calls), "a partial subdir must fall through to the Hub" + assert engine_args["model"] == os.path.join(str(snapshot), "language_model") + + +def test_build_engine_args_redownloads_when_index_lists_missing_shards(monkeypatch, tmp_path): + """An index plus a subset of its shards is an interrupted download. + + The index downloads early; accepting it as proof of weights converts the + Hub ID into a local path vLLM cannot fetch the remaining shards for. + """ + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "deadbeef" + partial = snapshot / "language_model" + partial.mkdir(parents=True) + (partial / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "a.weight": "model-00001-of-00002.safetensors", + "b.weight": "model-00002-of-00002.safetensors", + } + } + ) + ) + (partial / "model-00001-of-00002.safetensors").write_text("x") + _make_snapshot(snapshot, ["tokenizer"]) + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if not kwargs.get("local_files_only"): + (partial / "model-00002-of-00002.safetensors").write_text("x") + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert any(not call.get("local_files_only") for call in calls), "missing shards must fall through to the Hub" + assert engine_args["model"] == os.path.join(str(snapshot), "language_model") + + +def test_build_engine_args_redownloads_shards_without_their_index(monkeypatch, tmp_path): + """Shard-named weights always ship an index; one without it is partial.""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "deadbeef" + partial = snapshot / "language_model" + partial.mkdir(parents=True) + (partial / "model-00001-of-00004.safetensors").write_text("x") + _make_snapshot(snapshot, ["tokenizer"]) + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if not kwargs.get("local_files_only"): + shards = {f"w{i}.weight": f"model-0000{i}-of-00004.safetensors" for i in range(1, 5)} + (partial / "model.safetensors.index.json").write_text(json.dumps({"weight_map": shards})) + for shard in shards.values(): + (partial / shard).write_text("x") + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert any(not call.get("local_files_only") for call in calls) + + +def test_build_engine_args_redownloads_a_tokenizer_folder_without_vocabulary(monkeypatch, tmp_path): + """Templates and configs download first; alone they are not a tokenizer.""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "deadbeef" + _make_snapshot(snapshot, ["language_model"]) + partial = snapshot / "tokenizer" + partial.mkdir(parents=True) + (partial / "chat_template.jinja").write_text("{{ messages }}") + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if not kwargs.get("local_files_only"): + _make_snapshot(snapshot, ["tokenizer"]) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert any(not call.get("local_files_only") for call in calls) + assert engine_args["tokenizer"] == os.path.join(str(snapshot), "tokenizer") + + +def test_build_engine_args_forwards_revision_and_download_dir(monkeypatch, tmp_path): + """revision/download_dir must shape snapshot selection (they cannot be + corrected downstream once the repo ID is a local path).""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "pinned" + _make_snapshot(snapshot, ["language_model", "tokenizer"]) + seen = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + seen.append(kwargs) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={ + "model_subdir": "language_model", + "tokenizer_subdir": "tokenizer", + "revision": "v2.0", + "download_dir": str(tmp_path / "cache"), + }, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert engine_args["model"] == os.path.join(str(snapshot), "language_model") + assert seen, "snapshot resolution must have been consulted" + for call in seen: + assert call.get("revision") == "v2.0" + assert call.get("cache_dir") == str(tmp_path / "cache") + + +def test_build_engine_args_resolves_tokenizer_revision_separately(monkeypatch, tmp_path): + """A tokenizer pinned to a different revision resolves against its own snapshot.""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + model_snapshot = tmp_path / "snapshots" / "model-rev" + _make_snapshot(model_snapshot, ["language_model"]) + tokenizer_snapshot = tmp_path / "snapshots" / "tok-rev" + _make_snapshot(tokenizer_snapshot, ["tokenizer"]) + + def fake_snapshot_download(self, repo_id, **kwargs): + if kwargs.get("revision") == "tok-rev": + return str(tokenizer_snapshot) + return str(model_snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={ + "model_subdir": "language_model", + "tokenizer_subdir": "tokenizer", + "revision": "model-rev", + "tokenizer_revision": "tok-rev", + }, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert engine_args["model"] == os.path.join(str(model_snapshot), "language_model") + assert engine_args["tokenizer"] == os.path.join(str(tokenizer_snapshot), "tokenizer") + + +def test_build_engine_args_resolves_root_tokenizer_revision_separately(monkeypatch, tmp_path): + """An empty tokenizer_subdir targets the snapshot root and still honors + its own revision.""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + model_snapshot = tmp_path / "snapshots" / "model-rev" + _make_snapshot(model_snapshot, ["language_model"]) + # An empty tokenizer_subdir means the tokenizer IS the snapshot root, so + # the vocabulary artifact must sit at the root. Building it under + # ``tokenizer/`` returned a root that holds no loadable tokenizer, which + # the root-tokenizer validation now correctly rejects. + tokenizer_snapshot = tmp_path / "snapshots" / "tok-rev" + tokenizer_snapshot.mkdir(parents=True, exist_ok=True) + (tokenizer_snapshot / _SUBDIR_ARTIFACT["tokenizer"]).write_text("x") + + def fake_snapshot_download(self, repo_id, **kwargs): + if kwargs.get("revision") == "tok-rev": + return str(tokenizer_snapshot) + return str(model_snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={ + "model_subdir": "language_model", + "tokenizer_subdir": "", + "revision": "model-rev", + "tokenizer_revision": "tok-rev", + }, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + assert engine_args["model"] == os.path.join(str(model_snapshot), "language_model") + assert engine_args["tokenizer"] == str(tokenizer_snapshot) + + +def test_build_engine_args_fails_closed_when_subdir_cannot_be_downloaded(monkeypatch, tmp_path): + """An undownloadable subfolder raises here instead of reaching HuggingFace.""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + snapshot = tmp_path / "snapshots" / "deadbeef" + _make_snapshot(snapshot, ["tokenizer"]) + + def fake_snapshot_download(self, repo_id, **kwargs): + if kwargs.get("local_files_only"): + return str(snapshot) + raise OSError("offline") + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + with pytest.raises(RuntimeError, match="language_model"): + build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + +def test_build_engine_args_fails_closed_on_cold_cache_instead_of_joining_repo_id(monkeypatch): + """With nothing cached, the stage must not join a subdir onto the repo id.""" + from huggingface_hub import HfApi + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + def fake_snapshot_download(self, repo_id, **kwargs): + raise OSError("offline") + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + with pytest.raises(RuntimeError, match="nothing is cached locally"): + build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3") + + +def test_model_path_resolver_is_generic_and_model_owned(tmp_path): + from vllm_omni.engine.stage_init_utils import _resolve_model_path + + engine_args = { + "model_path_resolver": ("vllm_omni.model_executor.models.minimax_h3.checkpoint.resolve_minimax_h3_model_root"), + "revision": None, + "task_type": "ref2va", + } + + resolved = _resolve_model_path(str(tmp_path), engine_args) + + assert resolved == str(tmp_path / "Ref2VA" / "text_encoder") + assert "model_path_resolver" not in engine_args + + def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch): import vllm_omni.engine.stage_init_utils as init_mod diff --git a/tests/engine/test_data_entry_keys.py b/tests/engine/test_data_entry_keys.py index 4720cd8ea..df0d7b2ed 100644 --- a/tests/engine/test_data_entry_keys.py +++ b/tests/engine/test_data_entry_keys.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for data_entry_keys.""" import msgspec @@ -71,6 +74,25 @@ class TestOmniPayloadStruct: for top, sub in original.items(): assert sorted(d[top].keys()) == sorted(sub.keys()) + def test_text_conditioning_uses_shared_payload_fields(self): + hidden = torch.zeros(4, 5120) + token_role_ids = torch.tensor([[1], [1], [0], [0]]) + + payload = to_dict( + to_struct( + { + "hidden_states": {"output": hidden}, + "meta": {"token_role_ids": token_role_ids}, + } + ) + ) + + assert torch.equal(payload["hidden_states"]["output"], hidden) + assert torch.equal(payload["meta"]["token_role_ids"], token_role_ids) + for legacy_key in ("encoder_hidden_states", "token_tags"): + with pytest.raises(msgspec.ValidationError, match="unknown field"): + to_struct({legacy_key: torch.zeros(1)}) + def test_to_dict_drops_unset_fields(self): s = OmniPayloadStruct(meta=MetaStruct(left_context_size=10)) d = to_dict(s) @@ -372,12 +394,12 @@ class TestSerializeDeserializePayload: assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0])) def test_tensor_dtype_preserved(self): - # bfloat16 excluded: numpy() doesn't support it; callers must cast before serializing. - for dtype in [torch.float16, torch.float32, torch.int64, torch.int32, torch.bool]: + for dtype in [torch.float16, torch.bfloat16, torch.float32, torch.int64, torch.int32, torch.bool]: original: OmniPayload = {"codes": {"audio": torch.tensor([1], dtype=dtype)}} wire = serialize_payload(original) restored = deserialize_payload(wire) assert restored["codes"]["audio"].dtype == dtype, f"dtype mismatch for {dtype}" + assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"]) def test_tensor_shape_preserved(self): t = torch.randn(3, 4, 5) @@ -394,3 +416,12 @@ class TestSerializeDeserializePayload: original: OmniPayload = {"meta": {"finished": None}} wire = serialize_payload(original) assert wire is None + + def test_minimax_h3_prepared_reference_video_descriptor_round_trip(self): + descriptor = '{"artifact_dir":"/tmp/h3","videos":[]}' + + wire = serialize_payload({"meta": {"minimax_h3_prepared_reference_videos": descriptor}}) + assert wire is not None + restored = deserialize_payload(wire) + + assert restored["meta"]["minimax_h3_prepared_reference_videos"] == descriptor diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py index 6eae2e698..d04ed2e70 100644 --- a/tests/engine/test_duplex_import_boundary.py +++ b/tests/engine/test_duplex_import_boundary.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os import subprocess import sys from pathlib import Path @@ -14,9 +15,18 @@ REPO_ROOT = Path(__file__).resolve().parents[2] def _assert_isolated_import_succeeds(script: str) -> None: + # These tests validate Python import/package boundaries and do not use a + # GPU. Other tests in the combined entrypoints/engine suite may temporarily + # configure only one of the CUDA/ROCm visibility variables. Do not let an + # unrelated, inconsistent device-selection environment prevent the child + # interpreter from reaching the boundary assertion. + env = os.environ.copy() + env.pop("CUDA_VISIBLE_DEVICES", None) + env.pop("HIP_VISIBLE_DEVICES", None) result = subprocess.run( [sys.executable, "-c", script], cwd=REPO_ROOT, + env=env, capture_output=True, text=True, timeout=60, @@ -25,6 +35,14 @@ def _assert_isolated_import_succeeds(script: str) -> None: assert result.returncode == 0, result.stderr +def test_isolated_import_ignores_conflicting_gpu_visibility(monkeypatch: pytest.MonkeyPatch) -> None: + """Import-boundary checks must not depend on GPU state leaked by prior tests.""" + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + + _assert_isolated_import_succeeds("import vllm_omni.outputs") + + def test_stable_engine_imports_do_not_load_experimental_duplex() -> None: _assert_isolated_import_succeeds(""" import sys diff --git a/tests/engine/test_membership_controller.py b/tests/engine/test_membership_controller.py index 1b8bf2adb..5f0416377 100644 --- a/tests/engine/test_membership_controller.py +++ b/tests/engine/test_membership_controller.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -30,6 +30,14 @@ class FakeHub: self.closed = True +async def _wait_until(predicate, timeout=1): + async def _poll(): + while not predicate(): + await asyncio.sleep(0) + + await asyncio.wait_for(_poll(), timeout=timeout) + + class FakePool: def __init__(self, stage_id: int): self.stage_id = stage_id @@ -39,6 +47,7 @@ class FakePool: self.invalidated = [] self.hub = None self.lb = None + self.replica_ids = {} def attach_hub(self, hub): self.hub = hub @@ -49,7 +58,9 @@ class FakePool: def add_client(self, input_addr, client, *, replica_id=None): self.added.append((input_addr, client, replica_id)) self.clients.append(client) - return len(self.clients) - 1 if replica_id is None else replica_id + resolved_replica_id = len(self.clients) - 1 if replica_id is None else replica_id + self.replica_ids[input_addr] = resolved_replica_id + return resolved_replica_id def invalidate_addr(self, input_addr): self.invalidated.append(input_addr) @@ -59,6 +70,12 @@ class FakePool: self.removed.append(input_addr) return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown")) + def get_client_by_addr(self, input_addr): + return SimpleNamespace() if input_addr in self.replica_ids else None + + def get_replica_id_by_addr(self, input_addr): + return self.replica_ids.get(input_addr) + def _snapshot(*replicas): return SimpleNamespace(replicas=list(replicas)) @@ -159,6 +176,29 @@ async def test_do_register_offloads_remote_factory(monkeypatch): assert pool.added[0] == ("tcp://stage-0-replica-1", pool.clients[0], 1) +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["missing-address", "add-client"]) +async def test_do_register_shuts_down_rejected_client(monkeypatch, failure): + shutdown_calls = 0 + + class Client: + client_addresses = {} if failure == "missing-address" else {"input_address": "tcp://rejected"} + + def shutdown(self): + nonlocal shutdown_calls + shutdown_calls += 1 + + pool = FakePool(stage_id=0) + if failure == "add-client": + monkeypatch.setattr(pool, "add_client", lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("reject"))) + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=lambda *_args: Client()) + + with pytest.raises(RuntimeError): + await controller._do_register(0, 1) + + assert shutdown_calls == 1 + + @pytest.mark.asyncio async def test_unregister_then_register_restores_coordinator_replica_slot(monkeypatch): class Client: @@ -178,6 +218,16 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey return replacement controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory) + removed_replicas = [] + + async def cleanup(_request_ids): + return None + + controller.install_unregister_handlers( + output_queue=asyncio.Queue(), + cleanup_callback=cleanup, + replica_removed_callback=lambda stage_id, replica_id: removed_replicas.append((stage_id, replica_id)), + ) await controller.handle_unregister(0, "tcp://old-0") await controller.handle_unregister(0, "tcp://old-1") @@ -188,3 +238,368 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey assert pool.get_replica_id_by_addr("tcp://new-1") == 1 assert pool.get_replica_id_by_addr("tcp://old-1") is None assert pool.available_replica_ids() == [1] + assert removed_replicas == [(0, 0), (0, 1)] + + +@pytest.mark.asyncio +async def test_handle_register_deduplicates_inflight_and_attached_replica(monkeypatch): + pool = FakePool(stage_id=0) + factory_started = threading.Event() + release_factory = threading.Event() + factory_calls = [] + + def _factory(stage_id, replica_id): + factory_calls.append((stage_id, replica_id)) + factory_started.set() + assert release_factory.wait(timeout=1) + return SimpleNamespace( + client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}, + shutdown=lambda: None, + ) + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory) + + await asyncio.gather( + controller.handle_register(0, 2), + controller.handle_register(0, 2), + ) + assert await asyncio.to_thread(factory_started.wait, 1) + + await controller.handle_register(0, 2) + release_factory.set() + await controller.drain_tasks(timeout=1) + + await controller.handle_register(0, 2) + await controller.drain_tasks(timeout=1) + + assert factory_calls == [(0, 2)] + assert len(pool.added) == 1 + + +@pytest.mark.asyncio +async def test_handle_register_allows_retry_after_factory_failure(monkeypatch): + pool = FakePool(stage_id=0) + factory_calls = 0 + + def _factory(stage_id, replica_id): + nonlocal factory_calls + factory_calls += 1 + if factory_calls == 1: + raise RuntimeError("injected registration failure") + return SimpleNamespace( + client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}, + shutdown=lambda: None, + ) + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory) + + await controller.handle_register(0, 3) + await controller.drain_tasks(timeout=1) + await controller.handle_register(0, 3) + await controller.drain_tasks(timeout=1) + + assert factory_calls == 2 + assert len(pool.added) == 1 + + +@pytest.mark.asyncio +async def test_unregister_allows_same_replica_to_register_again(monkeypatch): + pool = FakePool(stage_id=0) + factory_calls = 0 + + def _factory(stage_id, replica_id): + nonlocal factory_calls + factory_calls += 1 + return SimpleNamespace( + client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}, + shutdown=lambda: None, + ) + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory) + + await controller.handle_register(0, 4) + await controller.drain_tasks(timeout=1) + await controller.handle_unregister(0, "tcp://stage-0-replica-4") + await controller.handle_register(0, 4) + await controller.drain_tasks(timeout=1) + + assert factory_calls == 2 + assert len(pool.added) == 2 + + +@pytest.mark.asyncio +async def test_handle_register_ignores_registration_after_shutdown(monkeypatch): + pool = FakePool(stage_id=0) + factory_calls = 0 + + def _factory(stage_id, replica_id): + nonlocal factory_calls + factory_calls += 1 + return SimpleNamespace( + client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}, + shutdown=lambda: None, + ) + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory) + controller.shutdown() + + await controller.handle_register(0, 5) + await controller.drain_tasks(timeout=1) + + assert factory_calls == 0 + assert pool.added == [] + + +@pytest.mark.asyncio +async def test_shutdown_discards_client_from_inflight_registration(monkeypatch): + pool = FakePool(stage_id=0) + factory_started = threading.Event() + release_factory = threading.Event() + client_shutdown = threading.Event() + + def _factory(stage_id, replica_id): + factory_started.set() + assert release_factory.wait(timeout=1) + return SimpleNamespace( + client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}, + shutdown=client_shutdown.set, + ) + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory) + + await controller.handle_register(0, 6) + assert await asyncio.to_thread(factory_started.wait, 1) + controller.shutdown() + release_factory.set() + await controller.drain_tasks(timeout=1) + + assert client_shutdown.is_set() + assert pool.added == [] + + +@pytest.mark.asyncio +async def test_watcher_reconciles_replica_that_disappears_during_registration(monkeypatch): + input_addr = "tcp://stage-0-replica-7" + pool = StagePool(0, []) + hub = FakeHub( + snapshots=[ + _snapshot(_replica(0, input_addr)), + _snapshot(), + ] + ) + factory_started = threading.Event() + release_factory = threading.Event() + client_shutdown = threading.Event() + + def _factory(stage_id, replica_id): + assert (stage_id, replica_id) == (0, 7) + factory_started.set() + assert release_factory.wait(timeout=1) + return SimpleNamespace( + client_addresses={"input_address": input_addr}, + shutdown=client_shutdown.set, + ) + + controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory) + controller.WATCH_INTERVAL_S = 0 + original_handle_unregister = controller.handle_unregister + first_unregister = asyncio.Event() + second_unregister = asyncio.Event() + unregister_calls = 0 + + async def _track_unregister(stage_id, addr, *args, **kwargs): + nonlocal unregister_calls + await original_handle_unregister(stage_id, addr, *args, **kwargs) + unregister_calls += 1 + if unregister_calls == 1: + first_unregister.set() + elif unregister_calls == 2: + second_unregister.set() + + controller.handle_unregister = _track_unregister # type: ignore[method-assign] + + await controller.handle_register(0, 7) + assert await asyncio.to_thread(factory_started.wait, 1) + watcher = controller.start() + + try: + await asyncio.wait_for(first_unregister.wait(), timeout=1) + assert pool.live_num_replicas == 0 + + release_factory.set() + await asyncio.wait_for(second_unregister.wait(), timeout=1) + + assert client_shutdown.is_set() + assert pool.live_num_replicas == 0 + assert controller._attached_remote_replicas == set() + assert controller._remote_replica_keys_by_input_addr == {} + finally: + release_factory.set() + controller.shutdown() + try: + await watcher + except asyncio.CancelledError: + pass + await controller.drain_tasks(timeout=1) + + +@pytest.mark.asyncio +async def test_attached_replica_survives_until_hub_observes_it_up(monkeypatch): + class EmptyHub: + def __init__(self): + self.polls = 0 + + def get_replica_list(self): + self.polls += 1 + return SimpleNamespace(replicas=[], timestamp=0.0) + + def close(self): + pass + + input_addr = "tcp://stage-0-replica-8" + pool = StagePool(0, []) + shutdown_calls = [] + hub = EmptyHub() + + def _factory(stage_id, replica_id): + return SimpleNamespace( + client_addresses={"input_address": input_addr}, + shutdown=lambda: shutdown_calls.append((stage_id, replica_id)), + ) + + controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory) + controller.WATCH_INTERVAL_S = 0.001 + + await controller.handle_register(0, 8) + await controller.drain_tasks(timeout=1) + assert pool.live_num_replicas == 1 + + watcher = controller.start() + try: + await _wait_until(lambda: hub.polls >= 3) + + assert shutdown_calls == [] + assert pool.live_num_replicas == 1 + finally: + controller.shutdown() + try: + await watcher + except asyncio.CancelledError: + pass + await controller.drain_tasks(timeout=1) + + +@pytest.mark.asyncio +async def test_up_observed_before_attach_does_not_qualify_attachment(monkeypatch): + import vllm_omni.engine.membership_controller as membership_mod + + input_addr = "tcp://stage-0-replica-10" + pool = StagePool(0, []) + hub = FakeHub( + snapshots=[ + _snapshot(_replica(0, input_addr)), + _snapshot(), + ] + ) + shutdown_calls = [] + sleep_gates = asyncio.Queue() + + async def _controlled_sleep(_delay): + gate = asyncio.Event() + await sleep_gates.put(gate) + await gate.wait() + + monkeypatch.setattr(membership_mod.asyncio, "sleep", _controlled_sleep) + + def _factory(stage_id, replica_id): + return SimpleNamespace( + client_addresses={"input_address": input_addr}, + shutdown=lambda: shutdown_calls.append((stage_id, replica_id)), + ) + + controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory) + watcher = controller.start() + + try: + first_sleep = await asyncio.wait_for(sleep_gates.get(), timeout=1) + await controller.handle_register(0, 10) + await controller.drain_tasks(timeout=1) + assert pool.live_num_replicas == 1 + + first_sleep.set() + await asyncio.wait_for(sleep_gates.get(), timeout=1) + await controller.drain_tasks(timeout=1) + + assert shutdown_calls == [] + assert pool.live_num_replicas == 1 + finally: + controller.shutdown() + try: + await watcher + except asyncio.CancelledError: + pass + await controller.drain_tasks(timeout=1) + + +@pytest.mark.asyncio +async def test_reattached_address_requires_a_fresh_up_observation(monkeypatch): + class MutableHub: + def __init__(self, input_addr): + self.input_addr = input_addr + self.up = True + self.polls = 0 + + def get_replica_list(self): + self.polls += 1 + replicas = [_replica(0, self.input_addr)] if self.up else [] + return _snapshot(*replicas) + + def close(self): + pass + + class Client: + def __init__(self, input_addr): + self.client_addresses = {"input_address": input_addr} + self.shutdown_calls = 0 + + def shutdown(self): + self.shutdown_calls += 1 + + input_addr = "tcp://stage-0-replica-9" + pool = StagePool(0, []) + hub = MutableHub(input_addr) + clients = [] + + def _factory(_stage_id, _replica_id): + client = Client(input_addr) + clients.append(client) + return client + + controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory) + controller.WATCH_INTERVAL_S = 0.001 + + await controller.handle_register(0, 9) + await controller.drain_tasks(timeout=1) + watcher = controller.start() + + try: + await _wait_until(lambda: hub.polls >= 2) + hub.up = False + await _wait_until(lambda: clients[0].shutdown_calls == 1) + assert pool.live_num_replicas == 0 + + await controller.handle_register(0, 9) + await controller.drain_tasks(timeout=1) + polls_after_reattach = hub.polls + await _wait_until(lambda: hub.polls >= polls_after_reattach + 3) + + assert len(clients) == 2 + assert clients[1].shutdown_calls == 0 + assert pool.live_num_replicas == 1 + finally: + controller.shutdown() + try: + await watcher + except asyncio.CancelledError: + pass + await controller.drain_tasks(timeout=1) diff --git a/tests/engine/test_omni_request_output.py b/tests/engine/test_omni_request_output.py index fa25bba9e..eb1bdfd41 100644 --- a/tests/engine/test_omni_request_output.py +++ b/tests/engine/test_omni_request_output.py @@ -270,3 +270,56 @@ class TestMsgpackRoundTrip: # Content merged from legacy nested key assert result.prompt == "legacy prompt" assert result.prompt_token_ids == [1, 2] + + +class TestRequestOutputFieldParity: + """``OmniRequestOutput`` must expose everything ``RequestOutput`` sets. + + ``RequestOutput`` is a plain class, so the dataclass-generated ``__init__`` + replaces its ``__init__`` and the inherited attributes only exist because + they are redeclared as fields. Any vLLM field that is not redeclared simply + does not exist on an omni output, and vLLM's own serving code reads those + attributes unconditionally — ``/v1/completions`` returned HTTP 500 with + ``AttributeError: 'OmniRequestOutput' object has no attribute + 'ec_transfer_params'`` for exactly that reason. This test fails on the next + vLLM bump that adds a field, instead of leaving it for an endpoint to hit. + """ + + def test_carries_every_attribute_request_output_sets(self) -> None: + reference = _make_text_request_output() + omni = OmniRequestOutput(request_id="req-parity") + + missing = sorted(name for name in vars(reference) if not name.startswith("_") and not hasattr(omni, name)) + assert not missing, ( + f"OmniRequestOutput is missing RequestOutput attributes {missing}; " + "redeclare them as dataclass fields in vllm_omni/outputs/__init__.py" + ) + + def test_transfer_param_fields_default_to_none(self) -> None: + omni = OmniRequestOutput(request_id="req-defaults") + assert omni.kv_transfer_params is None + assert omni.ec_transfer_params is None + assert omni.num_cached_tokens is None + assert omni.num_cache_creation_tokens is None + + def test_transfer_param_fields_are_copied_from_the_stage_output(self) -> None: + """Declaring the fields is not enough — they must also be copied. + + ``from_stage_output`` only carries the names listed in + ``_REQUEST_OUTPUT_CONTENT_ATTRS``. With the fields declared but absent + from that list the endpoint stops raising ``AttributeError`` while + silently dropping external-connector metadata and cache-creation + accounting, which is harder to notice than the crash it replaced. + """ + source = _make_text_request_output() + source.num_cache_creation_tokens = 7 + source.ec_transfer_params = {"connector": "shm", "handle": "abc"} + source.num_cached_tokens = 3 + source.kv_transfer_params = {"role": "producer"} + + omni = OmniRequestOutput.from_stage_output(source, stage_id=1) + + assert omni.num_cache_creation_tokens == 7 + assert omni.ec_transfer_params == {"connector": "shm", "handle": "abc"} + assert omni.num_cached_tokens == 3 + assert omni.kv_transfer_params == {"role": "producer"} diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 3d63943a7..9dcf066cd 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -14,11 +17,13 @@ import janus import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams +from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats from vllm_omni.engine.messages import ( AbortRequestMessage, + AbortResultMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, @@ -30,6 +35,7 @@ from vllm_omni.engine.messages import ( from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, + StreamingSegmentState, _build_terminal_empty_output, ) from vllm_omni.engine.stage_pool import StagePool @@ -102,6 +108,13 @@ class OrchestratorFixture: result_future: concurrent.futures.Future[None] +@dataclass +class FakePromptRequest: + request_id: str + prompt_token_ids: list[int] + resumable: bool = True + + class FakeStageClient: def __init__( self, @@ -244,8 +257,32 @@ class FakeOutputProcessor: ) def abort_requests(self, request_ids, internal: bool = False): - self.abort_calls.append(request_ids) - return request_ids + aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal) + return aborted_ids + + def abort_requests_collecting_outputs(self, request_ids, *, internal: bool = False, commit_state: bool = True): + """Mirror MultimodalOutputProcessor collecting for AR abort-prefix tests.""" + del internal, commit_state + ids = list(request_ids) + self.abort_calls.append(ids) + outputs: list[RequestOutput] = [] + for rid in ids: + seeded = next( + (ro for ro in self.request_outputs if getattr(ro, "request_id", None) == rid), + None, + ) + token_ids = list(seeded.outputs[0].token_ids) if seeded is not None and seeded.outputs else [1, 2] + # Intentionally stamp an internal-looking id on the RequestOutput so + # StagePool must re-key by the orchestrator id it passed in. + outputs.append( + _build_request_output( + f"engine-internal-{rid}", + token_ids=token_ids, + finished=True, + finish_reason="abort", + ) + ) + return ids, outputs def update_scheduler_stats(self, _scheduler_stats) -> None: return None @@ -287,6 +324,20 @@ def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace: return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None) +def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs: + return EngineCoreOutputs( + outputs=[ + EngineCoreOutput( + request_id=request_id, + new_token_ids=[], + finish_reason=FinishReason.STOP, + ) + ], + timestamp=1.0, + finished_requests={request_id}, + ) + + def _build_request_output( request_id: str, *, @@ -294,6 +345,7 @@ def _build_request_output( prompt_token_ids: list[int] | None = None, finished: bool = True, text: str = "test", + finish_reason: str | None = None, ) -> RequestOutput: completion = CompletionOutput( index=0, @@ -301,7 +353,7 @@ def _build_request_output( token_ids=list(token_ids or [1, 2]), cumulative_logprob=0.0, logprobs=None, - finish_reason="stop" if finished else None, + finish_reason=(finish_reason if finish_reason is not None else ("stop" if finished else None)), stop_reason=None, ) return RequestOutput( @@ -760,6 +812,148 @@ async def test_run_async_chunk(orchestrator_factory) -> None: await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +async def test_async_chunk_raw_terminal_without_processed_output_finishes_request(orchestrator_factory) -> None: + stage0 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio") + orchestrator_fixture = orchestrator_factory( + [stage0, stage1], + output_processors=[FakeOutputProcessor(), FakeOutputProcessor()], + async_chunk=True, + ) + request = FakePromptRequest( + request_id="req-stream-terminal", + prompt_token_ids=[1, 2, 3, 4], + ) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id=request.request_id, + prompt=request, + original_prompt={"prompt": "stream audio"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + + await _wait_for(lambda: len(stage1.add_request_calls) == 1) + stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request.request_id)) + + output_msg = await _get_output_message(orchestrator_fixture) + + assert output_msg.request_id == request.request_id + assert output_msg.stage_id == 1 + assert output_msg.finished is True + assert output_msg.engine_outputs.finished is True + assert output_msg.engine_outputs.outputs[0].multimodal_output["audio"].numel() == 0 + await _wait_for(lambda: request.request_id not in orchestrator_fixture.orchestrator.request_states) + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_async_chunk_data_then_raw_terminal_finishes_once(orchestrator_factory) -> None: + request_id = "req-stream-data-terminal" + stage0 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio") + final_data = _build_request_output( + request_id, + token_ids=[20, 21], + finished=False, + text="last audio chunk", + ) + orchestrator_fixture = orchestrator_factory( + [stage0, stage1], + output_processors=[ + FakeOutputProcessor(), + FakeOutputProcessor(request_outputs=[final_data]), + ], + async_chunk=True, + ) + request = FakePromptRequest( + request_id=request_id, + prompt_token_ids=[1, 2, 3, 4], + ) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id=request_id, + prompt=request, + original_prompt={"prompt": "stream audio"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + + await _wait_for(lambda: len(stage1.add_request_calls) == 1) + stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request_id)) + + data_msg = await _get_output_message(orchestrator_fixture) + terminal_msg = await _get_output_message(orchestrator_fixture) + + assert data_msg.engine_outputs is final_data + assert data_msg.finished is False + assert terminal_msg.request_id == request_id + assert terminal_msg.stage_id == 1 + assert terminal_msg.finished is True + assert terminal_msg.engine_outputs.finished is True + await _wait_for(lambda: request_id not in orchestrator_fixture.orchestrator.request_states) + await asyncio.sleep(0.05) + with pytest.raises(queue.Empty): + orchestrator_fixture.output_sync_q.get_nowait() + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_async_chunk_processed_terminal_and_raw_terminal_finishes_once(orchestrator_factory) -> None: + request_id = "req-stream-processed-terminal" + stage0 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio") + processed_terminal = _build_request_output( + request_id, + token_ids=[20, 21], + finished=True, + text="last audio chunk", + ) + orchestrator_fixture = orchestrator_factory( + [stage0, stage1], + output_processors=[ + FakeOutputProcessor(), + FakeOutputProcessor(request_outputs=[processed_terminal]), + ], + async_chunk=True, + ) + request = FakePromptRequest( + request_id=request_id, + prompt_token_ids=[1, 2, 3, 4], + ) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id=request_id, + prompt=request, + original_prompt={"prompt": "stream audio"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + + await _wait_for(lambda: len(stage1.add_request_calls) == 1) + stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request_id)) + + terminal_msg = await _get_output_message(orchestrator_fixture) + + assert terminal_msg.engine_outputs is processed_terminal + assert terminal_msg.finished is True + await _wait_for(lambda: request_id not in orchestrator_fixture.orchestrator.request_states) + await asyncio.sleep(0.05) + with pytest.raises(queue.Empty): + orchestrator_fixture.output_sync_q.get_nowait() + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + @pytest.mark.asyncio async def test_run_shutdown(orchestrator_factory) -> None: stages = [ @@ -805,10 +999,125 @@ async def test_run_abort(orchestrator_factory) -> None: assert stages[0].abort_calls == [["req-abort"]] assert stages[1].abort_calls == [] assert "req-abort" not in orchestrator_fixture.orchestrator.request_states + # Fire-and-forget abort must not emit an RPC result. + assert orchestrator_fixture.queues[2].sync_q.empty() finally: await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> None: + stages = [ + FakeStageClient(stage_type="llm", final_output=False), + FakeStageClient( + stage_type="llm", + final_output=True, + next_inputs=[{"prompt_token_ids": [7, 8, 9]}], + ), + ] + processors = [ + FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]), + FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[2], finished=True)]), + ] + orchestrator_fixture = orchestrator_factory(stages, output_processors=processors) + request = SimpleNamespace(request_id="req-ack", prompt_token_ids=[1, 2, 3]) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-ack", + prompt=request, + original_prompt={"prompt": "cancel with ack"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + await _wait_for(lambda: len(stages[0].add_request_calls) == 1) + # Abort outputs come from the final AR stage's output processor, and + # StagePool.abort_requests only collects for requests with a live + # replica binding. Forward off stage-0 first so stage-1 is bound. + stages[0].push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0)) + await _wait_for(lambda: len(stages[1].add_request_calls) == 1) + + orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["req-ack"], rpc_id="abort-1")) + result = await _get_rpc_message(orchestrator_fixture) + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-1" + assert result.success is True + assert result.error is None + assert result.rpc_correlation_key == ("abort", "abort-1") + assert stages[0].abort_calls == [["req-ack"]] + assert stages[1].abort_calls == [["req-ack"]] + assert "req-ack" not in orchestrator_fixture.orchestrator.request_states + assert result.abort_outputs is not None + assert len(result.abort_outputs) == 1 + abort_msg = result.abort_outputs[0] + assert abort_msg.request_id == "req-ack" + assert abort_msg.finished is True + assert abort_msg.stage_id == 1 + assert list(abort_msg.engine_outputs.outputs[0].token_ids) == [2] + assert abort_msg.engine_outputs.outputs[0].finish_reason == "abort" + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +@pytest.mark.asyncio +async def test_abort_marks_only_last_final_stage_output_finished() -> None: + """A request with two final AR stages must keep earlier abort outputs unfinished.""" + stages = [ + FakeStageClient(stage_type="llm", final_output=True), + FakeStageClient(stage_type="llm", final_output=True), + ] + processors = [ + FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[1], finished=True)]), + FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[2], finished=True)]), + ] + pools = _build_stage_pools([[stages[0]], [stages[1]]], output_processors=processors) + pools[0]._request_bindings["req-two"] = 0 + pools[1]._request_bindings["req-two"] = 0 + + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=pools, + ) + orchestrator.request_states["req-two"] = OrchestratorRequestState( + request_id="req-two", + final_stage_id=1, + final_output_stage_ids={0, 1}, + ) + + outputs = await orchestrator._abort_request_ids(["req-two"]) + assert [(msg.stage_id, msg.finished) for msg in outputs] == [(0, False), (1, True)] + assert list(outputs[0].engine_outputs.outputs[0].token_ids) == [1] + assert list(outputs[1].engine_outputs.outputs[0].token_ids) == [2] + + +@pytest.mark.asyncio +async def test_handle_abort_emits_error_result_on_failure() -> None: + rpc_queue: asyncio.Queue = asyncio.Queue() + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=rpc_queue, + stage_pools=[], + ) + + async def boom(_request_ids, *, abort=False): + del _request_ids, abort + raise RuntimeError("cleanup exploded") + + orchestrator._cleanup_request_ids = boom # type: ignore[method-assign] + + await orchestrator._handle_abort(AbortRequestMessage(request_ids=["req-fail"], rpc_id="abort-err")) + + result = rpc_queue.get_nowait() + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-err" + assert result.success is False + assert "cleanup exploded" in (result.error or "") + + def _duplex_open_message( session_id: str, *, @@ -1951,7 +2260,7 @@ async def test_resumable_segment_boundary_builds_stage_metrics() -> None: final_stage_id=0, ) req_state.streaming.enabled = True - req_state.streaming.segment_finished = True + req_state.streaming.segments[0] = StreamingSegmentState(finished=True) req_state.stage_submit_ts[0] = time.time() orchestrator.request_states = {"req-stream": req_state} orchestrator.stage_pools = [pool] @@ -1984,7 +2293,12 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None: ) output = SimpleNamespace( request_id="req-stream", - outputs=[SimpleNamespace(cumulative_token_ids=list(range(11)))], + outputs=[ + SimpleNamespace( + cumulative_token_ids=list(range(11)), + finish_reason=FinishReason.LENGTH, + ) + ], ) metrics = pool.build_stage_metrics( @@ -1996,6 +2310,29 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None: assert metrics.num_tokens_out == 3 assert metrics.output_unit_count == 3 + assert metrics.finish_reason == "length" + + +def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None: + stage = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image") + pool = StagePool( + 1, + [stage], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + output = SimpleNamespace(request_id="req-image", _custom_output={"image": "non-empty"}) + pool.record_output_timestamps([output], output_ts=132.0) + + metrics = pool.build_stage_metrics( + [output], + submit_ts=130.0, + request_timestamp=120.0, + replica_id=0, + ) + + assert metrics.serving_time_to_first_output_ms == pytest.approx(12000.0) + assert metrics.image_time_to_first_output_ms == pytest.approx(2000.0) @pytest.mark.asyncio diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 780c5a17b..90db55edb 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for error propagation paths within the Orchestrator. Covers: @@ -16,7 +16,9 @@ from types import SimpleNamespace import janus import pytest +from vllm.sampling_params import SamplingParams from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.serial_utils import MsgpackEncoder from vllm_omni.engine.messages import ( AddCompanionRequestMessage, @@ -25,7 +27,11 @@ from vllm_omni.engine.messages import ( ShutdownRequestMessage, StageSubmissionMessage, ) -from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + cleanup_request_artifact_dirs, +) from vllm_omni.engine.stage_pool import StageUnavailableError from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -42,12 +48,22 @@ from .test_orchestrator import ( _wait_for, ) + +def test_request_artifact_cleanup_is_idempotent(tmp_path): + artifact_dir = tmp_path / "request-artifacts" + artifact_dir.mkdir() + (artifact_dir / "prepared.mp4").touch() + + cleanup_request_artifact_dirs([str(artifact_dir)]) + cleanup_request_artifact_dirs([str(artifact_dir)]) + + assert not artifact_dir.exists() + + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _sampling_params(max_tokens: int = 4): - from vllm.sampling_params import SamplingParams - return SamplingParams(max_tokens=max_tokens) @@ -188,6 +204,45 @@ class _UnavailableDiffusionStageClient(FakeStageClient): raise StageUnavailableError(f"stage {self.stage_id} has no live replica") +class _OverflowOnEncodeStageClient(FakeStageClient): + """Msgpack-encodes the request's sampling params on dispatch, as the real + StageEngineCoreClient does. We use this to make sure overflow sampling params + do not break the orchestrator.""" + + async def add_request_async(self, request, *args, **kwargs) -> None: + MsgpackEncoder().encode(request.sampling_params) + await super().add_request_async(request, *args, **kwargs) # unreachable + + +@pytest.mark.asyncio +async def test_encode_overflow_should_fail_request_not_loop(orchestrator_factory) -> None: + """Regression test for over-range seed handling.""" + stage0 = _OverflowOnEncodeStageClient(stage_type="llm", final_output=True) + orchestrator_fixture = orchestrator_factory([stage0]) + + # seed > 2**63-1 cannot be msgpack-serialized -> the malicious input under test. + overflow_params = SamplingParams(seed=2**70) + + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-overflow", + prompt=SimpleNamespace( + request_id="req-overflow", + prompt_token_ids=[1, 2], + sampling_params=overflow_params, + ), + original_prompt={"prompt": "hi"}, + sampling_params_list=[overflow_params], + final_stage_id=0, + ) + + error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-overflow") + assert error_msg.fatal is False + assert error_msg.status_code == 400 + assert orchestrator_fixture.thread.is_alive() + assert "req-overflow" not in orchestrator_fixture.orchestrator.request_states + + @pytest.mark.asyncio async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None: """Multi-replica stage: one replica dying must fail only the requests bound @@ -411,6 +466,51 @@ async def test_async_chunk_prewarm_failure_reports_failing_downstream_stage(orch orchestrator_fixture.thread.join(timeout=5) +@pytest.mark.asyncio +async def test_async_chunk_prewarm_without_prompt_token_ids_fails_only_that_request( + orchestrator_factory, +) -> None: + """R1.3 of #4855: an async-chunk request whose stage-0 prompt carries no + token ids cannot prewarm anything, so it must fail with a client error + instead of stalling. + + Non-fatal is the point: ``fatal=True`` reads as "the engine died" and the + entrypoints act on it (``OmniLLM.generate`` raises out of its batch loop, + the serving layer calls ``terminate_if_errored``). One client's embeds-only + prompt must not reach other requests. + """ + stage0 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="llm", final_output=True) + orchestrator_fixture = orchestrator_factory([stage0, stage1], async_chunk=True) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-embeds", + # No prompt_token_ids: what an embeds-only prompt looks like here. + prompt=SimpleNamespace(request_id="req-embeds"), + original_prompt={"prompt": "embeds only"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + + error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-embeds") + assert error_msg.fatal is False + assert error_msg.status_code == 400 + assert error_msg.error_type == "BadRequestError" + assert error_msg.stage_id == 0 + assert "prompt_token_ids" in error_msg.error + + # Nothing was prewarmed downstream, and the request is gone. + assert stage1.add_request_calls == [] + await _wait_for(lambda: "req-embeds" not in orchestrator_fixture.orchestrator.request_states) + assert orchestrator_fixture.thread.is_alive() + finally: + if orchestrator_fixture.thread.is_alive(): + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + # ───────── Direct unit tests for the fault-isolation helpers ───────── @@ -835,6 +935,36 @@ async def test_diffusion_client_error_output_propagates_status_code(orchestrator orchestrator_fixture.thread.join(timeout=5) +@pytest.mark.asyncio +async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None: + class FailingStage(FakeStageClient): + def process_engine_inputs(self, *_args, **_kwargs): + raise ValueError("No latent or hidden_states found in thinker output") + + fixture = orchestrator_factory( + [FakeStageClient(stage_type="llm"), FailingStage(stage_type="llm", final_output=True)] + ) + state = OrchestratorRequestState( + request_id="bad", + prompt=SimpleNamespace(request_id="bad", prompt_token_ids=[1]), + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + duplex_identity=SimpleNamespace(), + ) + + async def no_cleanup(*_args, **_kwargs): + pass + + monkeypatch.setattr(fixture.orchestrator, "_cleanup_request_ids", no_cleanup) + try: + await fixture.orchestrator._forward_to_next_stage_unguarded("bad", 0, _build_request_output("raw"), state) + error = await _wait_for_error_message(fixture, request_id="bad") + assert error.fatal is False + finally: + fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + fixture.thread.join(timeout=5) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("status_code", "error_type"), diff --git a/tests/engine/test_orchestrator_event_driven.py b/tests/engine/test_orchestrator_event_driven.py new file mode 100644 index 000000000..284dde55d --- /dev/null +++ b/tests/engine/test_orchestrator_event_driven.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Event-driven orchestration loop (``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``) tests. + +Parity suite: re-runs the legacy orchestration scenarios from +``test_orchestrator.py`` / ``test_orchestrator_error_handling.py`` with the +event-driven loop selected, so both loops are held to the same behavior. Plus +event-driven-specific coverage: reader reconcile on client swap, the blocking +final-output drain, and flag parsing. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import janus +import pytest + +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.messages import ShutdownRequestMessage +from vllm_omni.engine.orchestrator import _event_driven_orch_enabled + +from . import test_orchestrator as legacy +from . import test_orchestrator_error_handling as legacy_errors +from .test_orchestrator import ( + FakeOutputProcessor, + FakeStageClient, + OrchestratorFixture, + _build_harness, + _build_request_output, + _engine_core_outputs, + _enqueue_add_request, + _get_output_message, + _sampling_params, + _shutdown_orchestrator, + _wait_for, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def orchestrator_factory(monkeypatch): + """Flag-setting clone of the legacy harness fixture. + + Sets ``VLLM_OMNI_EVENT_DRIVEN_ORCH=1`` before any Orchestrator is + constructed and asserts the flag actually took effect, so the parity tests + cannot silently exercise the legacy poll loop. + """ + monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", "1") + fixtures: list[OrchestratorFixture] = [] + + def _factory(*args, **kwargs) -> OrchestratorFixture: + fixture = _build_harness(*args, **kwargs) + assert fixture.orchestrator._event_driven_orch is True + fixtures.append(fixture) + return fixture + + yield _factory + + for fixture in fixtures: + if fixture.thread.is_alive(): + fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + fixture.thread.join(timeout=5) + for q in fixture.queues: + q.close() + + +# --------------------------------------------------------------------------- +# Parity: the legacy scenario matrix, re-run through the event-driven loop +# --------------------------------------------------------------------------- + +_PARITY_TESTS = [ + legacy.test_run_two_stage_llm, + legacy.test_run_single_stage_diffusion, + legacy.test_run_single_stage_diffusion_streaming_forwards_intermediate_chunks, + legacy.test_run_llm_to_diffusion, + legacy.test_run_async_chunk, + legacy.test_run_shutdown, + legacy.test_run_abort, + legacy.test_multi_replica_round_robin_distribution, + legacy.test_multi_replica_abort_broadcasts_to_all_replicas, + legacy.test_multi_replica_shutdown_all_replicas, + legacy.test_multi_replica_cfg_companion_inherits_parent_affinity, + # Stats plumbing. The scheduler-stats cases matter most here: a batch with + # no request outputs still carries SchedulerStats on throttled ticks, and a + # reader that drops every output-less batch silently stops reporting + # KV/queue gauges under the event-driven loop. + legacy.test_orchestrator_records_iteration_stats_without_scheduler_stats, + legacy.test_orchestrator_records_scheduler_stats_without_outputs, + legacy.test_orchestrator_does_not_build_iteration_stats_for_finished_only_batch, + legacy.test_orchestrator_does_not_build_iteration_stats_without_stat_logger, + # Per-replica fault isolation (#4285): a dead replica must be evicted and + # the server kept up, on both the LLM reader path and the diffusion poller. + legacy_errors.test_engine_dead_error_evicts_replica_and_keeps_running, + legacy_errors.test_engine_dead_error_fails_only_dead_replica_requests, + legacy_errors.test_forward_to_dead_downstream_stage_fails_request_not_server, + legacy_errors.test_add_request_to_dead_stage_fails_request_not_server, + legacy_errors.test_diffusion_replica_death_on_poll_keeps_server, + legacy_errors.test_diffusion_error_output_routed_as_finished, + legacy_errors.test_diffusion_client_error_output_propagates_status_code, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("legacy_test", _PARITY_TESTS, ids=lambda f: f.__name__) +async def test_event_driven_parity(legacy_test, orchestrator_factory) -> None: + await legacy_test(orchestrator_factory) + + +# --------------------------------------------------------------------------- +# Event-driven-specific behavior +# --------------------------------------------------------------------------- + + +def test_flag_parsing(monkeypatch) -> None: + monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False) + assert _event_driven_orch_enabled() is False + for value in ("1", "true", "True", "YES", "on"): + monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", value) + assert _event_driven_orch_enabled() is True + for value in ("0", "false", "off", ""): + monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", value) + assert _event_driven_orch_enabled() is False + + +def test_default_is_legacy_loop(monkeypatch) -> None: + """Without the env flag, the harness runs the legacy poll loop.""" + monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False) + fixture = _build_harness([FakeStageClient(stage_type="llm", final_output=True)]) + try: + assert fixture.orchestrator._event_driven_orch is False + finally: + fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + fixture.thread.join(timeout=5) + for q in fixture.queues: + q.close() + + +@pytest.mark.asyncio +async def test_reader_reconcile_picks_up_swapped_client(orchestrator_factory) -> None: + """Outputs from a replica whose client object was replaced still flow. + + The event-driven loop binds one reader task per client object; the + periodic reconcile must respawn the reader when ``pool.clients[replica]`` + is swapped (replica replacement), otherwise the new client's outputs + would never be drained. + """ + stage0 = FakeStageClient(stage_type="llm", final_output=True) + processor = FakeOutputProcessor(request_outputs=[_build_request_output("req-swap", token_ids=[3], finished=True)]) + orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor]) + request = SimpleNamespace(request_id="req-swap", prompt_token_ids=[1, 2]) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-swap", + prompt=request, + original_prompt={"prompt": "swap"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + # Swap in a fresh client for the same replica slot; keep the pool's + # other wiring intact. The reconcile tick (0.5 s) must respawn the + # reader bound to the new client object. + pool = orchestrator_fixture.orchestrator.stage_pools[0] + replacement = FakeStageClient(stage_type="llm", final_output=True) + replacement.stage_id = stage0.stage_id + replacement.replica_id = stage0.replica_id + pool.clients[0] = replacement + + replacement.push_engine_core_outputs(_engine_core_outputs("swapped-raw", 1.0)) + + output_msg = await _get_output_message(orchestrator_fixture, timeout=5.0) + assert output_msg.request_id == "req-swap" + assert output_msg.finished is True + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + +# --------------------------------------------------------------------------- +# Blocking final-output drain (AsyncOmniEngine.get_output_blocking_async) +# --------------------------------------------------------------------------- + + +def _drain_engine(alive: bool = True) -> AsyncOmniEngine: + engine = object.__new__(AsyncOmniEngine) + engine.output_queue = janus.Queue() + engine.orchestrator_thread = SimpleNamespace(is_alive=lambda: alive) + return engine + + +def _drain_cleanup(engine: AsyncOmniEngine) -> None: + if engine._output_drain_executor is not None: + engine._output_drain_executor.shutdown(wait=False) + engine._output_drain_executor = None + engine.output_queue.close() + + +@pytest.mark.asyncio +async def test_blocking_drain_returns_queued_message() -> None: + engine = _drain_engine() + try: + engine.output_queue.sync_q.put_nowait("msg-1") + assert await engine.get_output_blocking_async(timeout=1.0) == "msg-1" + finally: + _drain_cleanup(engine) + + +@pytest.mark.asyncio +async def test_blocking_drain_wakes_on_late_message() -> None: + """A message put after the wait starts wakes the drain, no polling.""" + engine = _drain_engine() + try: + + async def _delayed_put() -> None: + await asyncio.sleep(0.05) + engine.output_queue.sync_q.put_nowait("late-msg") + + put_task = asyncio.create_task(_delayed_put()) + msg = await engine.get_output_blocking_async(timeout=5.0) + await put_task + assert msg == "late-msg" + finally: + _drain_cleanup(engine) + + +@pytest.mark.asyncio +async def test_blocking_drain_timeout_returns_none_when_alive() -> None: + engine = _drain_engine(alive=True) + try: + assert await engine.get_output_blocking_async(timeout=0.05) is None + finally: + _drain_cleanup(engine) + + +@pytest.mark.asyncio +async def test_blocking_drain_raises_when_orchestrator_dead() -> None: + engine = _drain_engine(alive=False) + try: + with pytest.raises(RuntimeError, match="Orchestrator died"): + await engine.get_output_blocking_async(timeout=0.05) + finally: + _drain_cleanup(engine) diff --git a/tests/engine/test_orchestrator_segment_stage_keying.py b/tests/engine/test_orchestrator_segment_stage_keying.py new file mode 100644 index 000000000..02bb9a338 --- /dev/null +++ b/tests/engine/test_orchestrator_segment_stage_keying.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Streaming segment-boundary state must be keyed by stage. + +``Orchestrator._orchestration_loop`` polls every stage into the same +``OrchestratorRequestState``. While the segment-boundary fields lived in a single +flat slot per request, whichever stage polled most recently overwrote the others, +and the per-stage consumers then read back a different stage's boundary. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + StreamingSegmentState, +) +from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +THINKER_STAGE = 0 +TALKER_STAGE = 1 + + +@dataclass +class _StageMetrics: + pipeline_timings: dict[str, float] = field(default_factory=dict) + + +@dataclass +class _StageOutput: + """Stand-in for a processed stage output, with the fields the router reads.""" + + request_id: str + finished: bool = False + error: object | None = None + + +class _RecordingPool: + """Stage pool that records the requests it was asked to build metrics for.""" + + def __init__(self) -> None: + self.built_for: list[str] = [] + + def build_stage_metrics(self, outputs: list[_StageOutput], **_kwargs: object) -> _StageMetrics: + self.built_for.extend(output.request_id for output in outputs) + return _StageMetrics() + + +def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState: + req_state = OrchestratorRequestState( + request_id=request_id, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=TALKER_STAGE, + ) + req_state.streaming.enabled = True + req_state.duplex_identity = DuplexRequestIdentity( + session_id="session-1", + fence=DuplexFence(session_id="session-1"), + ) + return req_state + + +def test_duplex_output_context_reads_only_its_own_stage_segment() -> None: + req_state = _duplex_req_state() + req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState( + finished=True, + token_ids=[11, 22], + output_metadata={"stage": "thinker"}, + ) + # The talker polls after the thinker and is still mid-segment. + req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False) + + thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE) + talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE) + + assert thinker is not None + assert talker is not None + + # These are the assertions that fail on the flat field: the talker is written + # second, so the thinker loses its boundary to whichever stage polled last. + assert thinker.segment_finished is True + assert thinker.segment_token_ids == (11, 22) + assert thinker.segment_output_metadata == {"stage": "thinker"} + + assert talker.segment_finished is False + assert talker.segment_token_ids == () + assert talker.segment_output_metadata == {} + + +@pytest.mark.asyncio +async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None: + req_state = _duplex_req_state() + req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True) + + talker_pool = _RecordingPool() + orchestrator = object.__new__(Orchestrator) + orchestrator.request_states = {req_state.request_id: req_state} + orchestrator.stage_pools = [_RecordingPool(), talker_pool] + + routed: list[_StageMetrics | None] = [] + + async def record_route( + _stage: int, + _replica: int, + _output: _StageOutput, + _state: OrchestratorRequestState, + metrics: _StageMetrics | None, + ) -> None: + routed.append(metrics) + + orchestrator._route_output = record_route + + # The talker's own output is neither finished nor at a segment boundary, so + # it must not be treated as one just because the thinker reached one. + await orchestrator._handle_processed_outputs( + TALKER_STAGE, + 0, + [_StageOutput(request_id=req_state.request_id)], + ) + + assert talker_pool.built_for == [] + assert routed == [None] diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index cf347e49f..52cee4657 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, + StreamingSegmentState, _OrchestratorDuplexStagePort, ) from vllm_omni.engine.stage_pool import StagePool @@ -160,6 +161,26 @@ def _duplex_stage_port_submission(): return port, stage_pools, request_states, prewarm, submission +def test_duplex_bridge_state_catches_up_after_serving_turn_boundary() -> None: + port, _, request_states, _, submission = _duplex_stage_port_submission() + request_state = request_states[submission.context.request_id] + bridge_state = request_state.streaming.bridge_states["duplex"] + bridge_state["model_turn_id"] = 0 + + next_context = DuplexStageRequestContext( + request_id=submission.context.request_id, + session_id=submission.context.session_id, + fence=DuplexFence(submission.context.session_id, turn_id=1), + stage_id=submission.context.stage_id, + final_stage_id=submission.context.final_stage_id, + config_generation=submission.context.config_generation, + sampling_params=submission.context.sampling_params, + ) + port.ensure_request(next_context) + + assert bridge_state["model_turn_id"] == 1 + + def _request_output(request_id: str) -> RequestOutput: completion = CompletionOutput( index=0, @@ -256,12 +277,13 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None: duplex_identity=SimpleNamespace(), ) - await orchestrator._prewarm_async_chunk_stages( + prewarmed = await orchestrator._prewarm_async_chunk_stages( "req-prewarm", SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), req_state, ) + assert prewarmed is True assert stage1.submitted == [] assert len(stage2.submitted) == 1 assert 1 not in req_state.stage_submit_ts @@ -277,6 +299,7 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None: @pytest.mark.asyncio async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None: port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() + prewarm.return_value = True result = await port.submit(submission) @@ -285,6 +308,30 @@ async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None: prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"]) +@pytest.mark.asyncio +async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None: + """A failed prewarm has already aborted the request and popped its state. + + Returning a success result here would hand the control plane a replica for a + request that no longer exists, and the trailing bookkeeping would re-register + a running counter the cleanup just released -- a leak that never decrements. + """ + port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission() + counter = MagicMock() + port._running_counter = counter + prewarm.return_value = False + request_state = request_states["req-duplex"] + + with pytest.raises(RuntimeError, match="prewarm failed"): + await port.submit(submission) + + prewarm.assert_awaited_once() + assert request_state.duplex_stage_fences == {} + assert request_state.stage_submit_ts == {} + assert request_state.running_counter_registered is False + counter.increment.assert_not_called() + + @pytest.mark.asyncio async def test_async_route_forwards_to_outgoing_only_stage() -> None: orchestrator = object.__new__(Orchestrator) @@ -332,7 +379,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: final_output_stage_ids={0}, ) req_state.streaming.enabled = True - req_state.streaming.segment_finished = True + req_state.streaming.segments[0] = StreamingSegmentState(finished=True) output = SimpleNamespace( request_id=req_state.request_id, finished=True, diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 23ac601b3..62ed8334b 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard.""" from types import SimpleNamespace @@ -581,18 +581,18 @@ def test_mm_only_terminal_finish_removes_request_state(monkeypatch): def test_no_detokenizer_make_request_output_with_routed_experts(): """make_request_output accepts the routed_experts arg that the multimodal - output channel (_process_mm_only_outputs) passes positionally, and attaches - it to the completion output. + output channel (_process_mm_only_outputs) passes, and attaches it to the + completion output. - Regression: the call site passes 6 positional args - (..., kv_transfer_params, routed_experts) but the override previously took - only 5, raising TypeError on every generation-stage output. + routed_experts is omni-specific keyword-only: the 6th positional now + matches upstream's ec_transfer_params so that super().process_outputs() + calls do not misroute it. """ s = _make_no_detok_state(RequestOutputKind.CUMULATIVE) s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO) routed_experts = np.zeros((2, 3), dtype=np.int32) # Mirror the exact call shape of _process_mm_only_outputs. - result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts) + result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts=routed_experts) assert result is not None assert not isinstance(result, PoolingRequestOutput) assert result.outputs[0].routed_experts is routed_experts @@ -663,3 +663,99 @@ def test_mm_only_outputs_update_iteration_stats(): assert finished.finish_reason == FinishReason.STOP assert finished.num_prompt_tokens == state.prompt_len assert finished.num_generation_tokens == 2 + + +def test_ec_transfer_params_survive_both_construction_paths(): + """v0.28 contract: encoder-cache transfer metadata must reach + RequestOutput.ec_transfer_params through BOTH omni construction paths — + the upstream-delegating one (logprobs processor present) and the + no-detokenizer generation-stage one (direct RequestOutput build). + Regression: the override accepted the argument but dropped it.""" + ec = {"remote_handle": "enc-cache-1"} + + state = _make_state(RequestOutputKind.CUMULATIVE) + out = state.make_request_output([1], None, None, None, {"kv": 1}, ec) + assert out is not None and out.ec_transfer_params == ec + + kwargs = dict(_DEFAULT_STATE_KWARGS) + kwargs.update(logprobs_processor=None, detokenizer=None) + gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE) + completion = gen_state._new_completion_output([1], None, None) + direct = gen_state._new_request_output("r", [completion], False, {"kv": 1}, ec) + assert direct.ec_transfer_params == ec + + +def test_num_cache_creation_tokens_reaches_direct_request_output(): + """v0.28 contract: prefix-cache creation usage must reach + RequestOutput.num_cache_creation_tokens through the no-detokenizer + direct build, which previously passed only num_cached_tokens.""" + kwargs = dict(_DEFAULT_STATE_KWARGS) + kwargs.update(logprobs_processor=None, detokenizer=None) + gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE) + gen_state.num_cached_tokens = 3 + gen_state.num_cache_creation_tokens = 2 + + completion = gen_state._new_completion_output([1], None, None) + direct = gen_state._new_request_output("r", [completion], False, None, None) + + assert direct.num_cached_tokens == 3 + assert direct.num_cache_creation_tokens == 2 + + +def _abort_processor_with_parent(child_ids: list[str]): + processor = object.__new__(MultimodalOutputProcessor) + processor.request_states = {} + processor.external_req_ids = {} + processor.parent_requests = {} + processor._native_text_metrics_by_request = {} + processor.lora_states = SimpleNamespace(request_finished=lambda *_args, **_kwargs: None) + parent = SimpleNamespace(request_id="parent", child_requests=set(child_ids)) + processor.parent_requests["parent"] = parent + for child_id in child_ids: + req_state = SimpleNamespace( + parent_req=parent, + lora_name=None, + output_kind=RequestOutputKind.CUMULATIVE, + detokenizer=SimpleNamespace(output_token_ids=[7, 8]), + queue=None, + external_req_id="parent", + make_request_output=MagicMock(return_value=SimpleNamespace(request_id=child_id)), + ) + processor.request_states[child_id] = req_state + return processor, parent + + +def test_abort_last_parallel_child_drops_parent_request(): + processor, parent = _abort_processor_with_parent(["0_parent", "1_parent"]) + + aborted, _outputs = processor.abort_requests_collecting_outputs(["0_parent"], internal=True) + assert aborted == ["0_parent"] + assert "parent" in processor.parent_requests + assert parent.child_requests == {"1_parent"} + + aborted, _outputs = processor.abort_requests_collecting_outputs(["1_parent"], internal=True) + assert aborted == ["1_parent"] + assert "parent" not in processor.parent_requests + assert not parent.child_requests + + +def test_abort_snapshot_leaves_state_until_commit(): + processor, parent = _abort_processor_with_parent(["0_parent"]) + processor.external_req_ids["parent"] = ["0_parent"] + + aborted, outputs = processor.abort_requests_collecting_outputs( + ["parent"], + internal=False, + commit_state=False, + ) + assert aborted == ["0_parent"] + assert outputs + assert "0_parent" in processor.request_states + assert processor.external_req_ids["parent"] == ["0_parent"] + assert "parent" in processor.parent_requests + assert parent.child_requests == {"0_parent"} + + processor.commit_aborted_request_state(["parent"], internal=False) + assert "0_parent" not in processor.request_states + assert "parent" not in processor.parent_requests + assert "parent" not in processor.external_req_ids diff --git a/tests/engine/test_rpc_result_router.py b/tests/engine/test_rpc_result_router.py index ec983392b..1b3159617 100644 --- a/tests/engine/test_rpc_result_router.py +++ b/tests/engine/test_rpc_result_router.py @@ -4,6 +4,7 @@ from typing import Literal import pytest from vllm_omni.engine.messages import ( + AbortResultMessage, CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, @@ -93,6 +94,19 @@ def test_rpc_result_router_does_not_broadcast_uncorrelated_nonfatal_errors(): router.close() +def test_rpc_result_router_routes_abort_results_by_correlation_id(): + source: queue.Queue = queue.Queue() + router = RpcResultRouter(source) + waiter = router.register(("abort", "abort-42")) + + source.put(AbortResultMessage(rpc_id="abort-42", success=True)) + + result = waiter.get(timeout=1) + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-42" + router.close() + + def test_rpc_result_router_close_unblocks_waiters_and_stops_consumer(): source: queue.Queue = queue.Queue() router = RpcResultRouter(source) diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index bde8ba490..cacda1f60 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for AsyncOmniEngine single-stage mode and OmniMasterServer.""" from __future__ import annotations @@ -428,7 +431,7 @@ class TestSingleStageModeDetection: assert engine.single_stage_mode is False assert engine._single_stage_id_filter is None - def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture): + def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture): duplex_session = DuplexSessionRuntimeConfig(max_sessions=2) get_pipeline_config = mocker.patch( "vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config", @@ -441,7 +444,7 @@ class TestSingleStageModeDetection: engine = self._make_engine_no_thread( mocker, - stage_configs_path="/fake/duplex.yaml", + deploy_config="/fake/duplex.yaml", ) get_pipeline_config.assert_called_once_with( @@ -572,7 +575,6 @@ class TestSingleStageInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=stage_id_filter, omni_master_address="127.0.0.1", @@ -804,7 +806,6 @@ class TestSingleStageInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, ) mocker.patch.object(runtime_mod, "prepare_engine_environment") @@ -857,7 +858,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -925,7 +925,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -947,7 +946,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -999,7 +997,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=2, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1065,7 +1062,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=4, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1117,7 +1113,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=4, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1138,7 +1133,8 @@ class TestSingleStageReplicaInitialization: runtime._init_visible_devices_baseline = "0" mocker.patch.object(runtime_mod, "inject_kv_stage_info") - mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1)) + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config) mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", return_value=StageRegistrationResponse( @@ -1174,6 +1170,7 @@ class TestSingleStageReplicaInitialization: os.environ[device_env_var] = prev_device_env assert result is sentinel_client + assert od_config.max_num_seqs == 4 mock_register.assert_called_once_with( omni_master_address="127.0.0.1", omni_master_port=25000, @@ -1183,7 +1180,7 @@ class TestSingleStageReplicaInitialization: ) mock_manager.assert_called_once_with( model="fake-model", - od_config="diffusion-config", + od_config=od_config, stage_init_timeout=60, handshake_address="tcp://127.0.0.1:26001", addresses=mocker.ANY, @@ -1198,10 +1195,50 @@ class TestSingleStageReplicaInitialization: request_address="tcp://127.0.0.1:26002", response_address="tcp://127.0.0.1:26003", proc_manager=mocker.ANY, - batch_size=4, ) assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc + def test_initialize_local_custom_diffusion_replica_uses_inline_client(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.platforms import current_omni_platform + + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._coordinator_runtime = None + plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0] + plan.stage_cfg.engine_args = {"custom_pipeline_args": {"pipeline_class": "test.CustomPipeline"}} + sentinel_client = SimpleNamespace() + mock_launch = mocker.patch.object( + runtime_mod, + "launch_diffusion_stage_replica", + return_value=(sentinel_client, StageReplicaResources()), + ) + mocker.patch.object(runtime_mod, "inject_kv_stage_info") + + device_env_var = current_omni_platform.device_control_env_var + prev_device_env = os.environ.get(device_env_var) + os.environ[device_env_var] = "0" + runtime._init_visible_devices_baseline = "0" + try: + result = runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60) + finally: + if prev_device_env is None: + os.environ.pop(device_env_var, None) + else: + os.environ[device_env_var] = prev_device_env + + assert result is sentinel_client + assert mock_launch.call_args.kwargs["use_inline"] is True + def test_initialize_local_diffusion_replica_failure_terminates_proc(self, mocker: MockerFixture): import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform @@ -1211,7 +1248,6 @@ class TestSingleStageReplicaInitialization: model="fake-model", config_path="/fake/stages.yaml", stage_init_timeout=60, - diffusion_batch_size=4, async_chunk=False, single_stage_id_filter=None, omni_master_address="127.0.0.1", @@ -1230,7 +1266,8 @@ class TestSingleStageReplicaInitialization: runtime._init_visible_devices_baseline = "0" mocker.patch.object(runtime_mod, "inject_kv_stage_info") - mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1)) + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config) mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", return_value=StageRegistrationResponse( @@ -1371,7 +1408,7 @@ class TestConnectRemoteEngineCoresCoordinator: omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2) omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2) mock_wait.assert_called_once() - _, _, core_engines, parallel_config, *_ = mock_wait.call_args.args + _, core_engines, parallel_config, *_ = mock_wait.call_args.args assert core_engines[0].local is False assert parallel_config.data_parallel_size_local == 0 diff --git a/tests/engine/test_stage_device_layout.py b/tests/engine/test_stage_device_layout.py new file mode 100644 index 000000000..0473dc341 --- /dev/null +++ b/tests/engine/test_stage_device_layout.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Regression tests for issue #5003: a per-stage world size that its assigned +``devices`` cannot satisfy must fail early in ``build_vllm_config`` with a clear +message, rather than surfacing as an opaque worker-side ``local rank ... out of +bounds`` assertion. + +Root cause: a top-level ``--tensor-parallel-size`` is broadcast to every stage, +but each stage's ``devices`` is not adjusted, so a stage can end up with e.g. +tensor_parallel_size=4 while still holding a single-GPU deploy default. Without +``--strategy-config`` the strategy-path device check never runs. +""" + +import json +import re +import types +from unittest import mock + +import pytest + +from vllm_omni.engine import stage_init_utils +from vllm_omni.engine.stage_init_utils import ( + _check_stage_device_layout, + build_vllm_config, + compute_replica_layout, + get_stage_devices_per_replica, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _stage(stage_id, devices, num_replicas=1, engine_args=None): + return types.SimpleNamespace( + stage_id=stage_id, + stage_type="llm", + engine_args=engine_args or {}, + runtime=types.SimpleNamespace(devices=devices, num_replicas=num_replicas), + ) + + +def test_tp_broadcast_without_devices_fails_early(): + """stage0 gets tensor_parallel_size=4 (broadcast) but only 1 device -> clear error.""" + stage = _stage(0, devices="0") + engine_args = { + "tensor_parallel_size": 4, + "data_parallel_size": 1, + "pipeline_parallel_size": 1, + } + with pytest.raises(ValueError) as excinfo: + _check_stage_device_layout(stage, engine_args) + msg = str(excinfo.value) + # Message names the stage, the mismatch, and the actionable workaround. + assert "Stage 0" in msg + assert "1 device" in msg and "4" in msg + assert "--stage-overrides" in msg + + +def test_workaround_example_sets_tp_on_every_stage(): + """The JSON example in the error message must not reproduce the bug: every + stage entry it suggests has to set ``tensor_parallel_size`` (a single-GPU + stage that only overrides ``devices`` would still inherit the broadcast + tp and crash again — the original #5003 failure mode).""" + stage = _stage(0, devices="0") + with pytest.raises(ValueError) as excinfo: + _check_stage_device_layout(stage, {"tensor_parallel_size": 4}) + msg = str(excinfo.value) + + match = re.search(r"\{.*\}", msg, re.DOTALL) + assert match, f"no JSON example found in error message: {msg}" + example = json.loads(match.group(0)) + + assert example, "example must contain at least one stage" + for stage_id, overrides in example.items(): + assert "tensor_parallel_size" in overrides, ( + f"stage {stage_id} in the suggested workaround omits " + f"tensor_parallel_size, which reproduces #5003: {overrides}" + ) + + +def test_consistent_tp_and_devices_pass(): + """tensor_parallel_size=4 with 4 assigned devices is valid.""" + stage = _stage(0, devices="0,1,2,3") + _check_stage_device_layout( + stage, + {"tensor_parallel_size": 4, "data_parallel_size": 1, "pipeline_parallel_size": 1}, + ) + + +def test_single_gpu_stage_passes(): + """A TP=1 stage on a single device (talker/code2wav default) is valid.""" + stage = _stage(1, devices="1") + _check_stage_device_layout(stage, {"tensor_parallel_size": 1}) + + +def test_missing_devices_is_skipped(): + """No explicit devices -> vLLM assigns them; nothing to validate here.""" + stage = _stage(0, devices=None) + _check_stage_device_layout(stage, {"tensor_parallel_size": 4}) + + +def test_replica_pool_layout_passes(): + """Pool mode: num_replicas=2 x tp=2 => 4 devices is a valid pool shape.""" + stage = _stage(0, devices="0,1,2,3", num_replicas=2) + _check_stage_device_layout(stage, {"tensor_parallel_size": 2}) + + +def test_multinode_dp_uses_local_width_for_device_validation(): + """One local DP engine needs one device even when global DP spans four nodes.""" + stage = _stage(0, devices="0") + _check_stage_device_layout( + stage, + { + "tensor_parallel_size": 1, + "data_parallel_size": 4, + "data_parallel_size_local": 1, + "pipeline_parallel_size": 1, + }, + ) + + +def test_zero_local_dp_skips_local_device_validation(): + """A head process with no local engines does not own a local device layout.""" + stage = _stage(0, devices="0") + _check_stage_device_layout( + stage, + { + "tensor_parallel_size": 4, + "data_parallel_size": 4, + "data_parallel_size_local": 0, + "pipeline_parallel_size": 1, + }, + ) + + +@pytest.mark.parametrize( + ("devices", "expected"), + [ + ("0,1", ["0,1", "2,3"]), + ("0,1,2,3", ["0,1", "2,3"]), + ], +) +def test_replica_guard_and_splitter_share_local_world_size(devices, expected): + """PP and local DP contribute to the same per-replica size in both paths.""" + engine_args = { + "tensor_parallel_size": 1, + "data_parallel_size": 4, + "data_parallel_size_local": 1, + "pipeline_parallel_size": 2, + } + stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args) + + _check_stage_device_layout(stage, engine_args) + assert get_stage_devices_per_replica(stage) == 2 + _, replica_devices_map = compute_replica_layout([stage]) + assert replica_devices_map == {0: expected} + + +def test_non_tp_layout_error_uses_generic_guidance(): + """A PP mismatch must not be reported as a top-level TP broadcast.""" + stage = _stage(0, devices="0") + with pytest.raises(ValueError) as excinfo: + _check_stage_device_layout( + stage, + { + "tensor_parallel_size": 1, + "data_parallel_size": 1, + "pipeline_parallel_size": 2, + }, + ) + + msg = str(excinfo.value) + assert "local world size" in msg + assert "top-level --tensor-parallel-size" not in msg + + +def test_build_vllm_config_fails_before_engine_config_on_mismatch(): + """Pin the production wiring: ``build_vllm_config`` must run the device-layout + guard *before* ``create_engine_config``. If the guard call is dropped or moved + after engine-config creation, #5003 regresses while the unit tests above stay + green — so assert here that a mismatched stage raises and that neither + ``create_engine_config`` nor ``Executor.get_class`` is reached.""" + stage = _stage(0, devices="0") + with ( + mock.patch.object(stage_init_utils.OmniEngineArgs, "create_engine_config") as create_engine_config, + mock.patch.object(stage_init_utils.Executor, "get_class") as get_class, + ): + with pytest.raises(ValueError) as excinfo: + build_vllm_config( + stage, + model="dummy-model", + engine_args_dict={ + "tensor_parallel_size": 4, + "data_parallel_size": 1, + "pipeline_parallel_size": 1, + }, + ) + msg = str(excinfo.value) + assert "Stage 0" in msg + assert "--stage-overrides" in msg + # Guard fired early: worker/config construction was never reached. + create_engine_config.assert_not_called() + get_class.assert_not_called() + + +def test_build_vllm_config_proceeds_on_consistent_layout(): + """The guard must not false-positive: a consistent single-GPU stage flows past + it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual.""" + stage = _stage(1, devices="1") + fake_config = types.SimpleNamespace( + quant_config=None, + model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()), + ) + sentinel_executor = object() + with ( + mock.patch.object( + stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config + ) as create_engine_config, + mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor), + mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant), + ): + vllm_config, executor_class = build_vllm_config( + stage, + model="dummy-model", + engine_args_dict={"tensor_parallel_size": 1}, + ) + create_engine_config.assert_called_once() + assert vllm_config is fake_config + assert executor_class is sentinel_executor diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 75660cbc7..0a54e2e80 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import copy @@ -8,8 +11,23 @@ from pathlib import Path import pytest from pydantic.fields import FieldInfo +from vllm.config import CacheConfig as VllmCacheConfig +from vllm.config import CompilationConfig as VllmCompilationConfig +from vllm.config import LoadConfig as VllmLoadConfig +from vllm.config import ParallelConfig as VllmParallelConfig +from vllm.config import ProfilerConfig as VllmProfilerConfig +from vllm.config import SchedulerConfig as VllmSchedulerConfig +from vllm.engine.arg_utils import EngineArgs -from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.omni_config import ( + _LLM_STAGE_ENGINE_FIELDS, + OmniStageCacheConfig, + OmniStageLoadConfig, + OmniStageParallelConfig, + OmniStageSchedulerConfig, + VllmOmniARStageConfig, + VllmOmniConfig, +) from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config from vllm_omni.config.stage_config import ( DeployConfig, @@ -30,6 +48,7 @@ from vllm_omni.engine.stage_init_utils import ( build_engine_args_dict_from_omni_stage_config, build_legacy_engine_args_dict, ) +from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -53,6 +72,32 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, _LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs)) _DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig)) +_TOPOLOGY_ONLY_ENGINE_ARGS = frozenset({"inline_diffusion"}) +_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset( + { + "active_stream_window", + "codec_frame_rate_hz", + "custom_voice_dir", + "devices", + "disable_autocast", + "enable_multithread_weight_load", + "env", + "has_sampling_extra_args", + "log_level", + "log_stats", + "model_arch", + "model_subdir", + "num_gpus", + "num_replicas", + "num_weight_load_threads", + "omni_kv_config", + "parallel_config", + "silence_ban_frames", + "subtalker_sampling_params", + "task_type", + "tokenizer_subdir", + } +) @pytest.fixture(autouse=True) @@ -80,6 +125,17 @@ def _stable_engine_arg_environment(monkeypatch, tmp_path): monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", _resolve_test_worker) + class _TestConnectorRunner(OmniConnectorModelRunnerMixin): + pass + + class _TestWorker: + model_runner_cls = _TestConnectorRunner + + for worker_type in ("ar", "generation"): + worker_module = types.ModuleType(f"test.{worker_type}") + worker_module.Worker = _TestWorker + monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) + def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]: parent_model = tmp_path / "parent-model" @@ -139,6 +195,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st engine_extras={ "model": str(stage_model), "attention_backend": "FLASHINFER", + "attention_config": {"use_trtllm_attention": False}, "hf_overrides": {"rope_scaling": {"factor": 2}}, "limit_mm_per_prompt": {"audio": 1}, "logits_processors": ["test.custom.LogitsProcessor"], @@ -196,6 +253,100 @@ def _legacy_and_typed_stages( return legacy_stages, omni_config +def test_llm_stage_engine_field_schema_tracks_upstream_engine_args(): + upstream_engine_fields = frozenset(field.name for field in fields(EngineArgs)) + + assert _LLM_STAGE_ENGINE_FIELDS - upstream_engine_fields == _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS + + +def test_typed_llm_projection_discovers_explicit_upstream_config_fields(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + load_config=OmniStageLoadConfig(ignore_patterns=["*.bin"]), + cache_config=OmniStageCacheConfig(block_size=32, cache_dtype="fp8"), + scheduler_config=OmniStageSchedulerConfig( + prefill_schedule_interval=2, + policy="priority", + ), + parallel_config=OmniStageParallelConfig( + data_parallel_master_ip="10.0.0.1", + enable_dbo=True, + ), + ) + + engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config) + + assert engine_args["ignore_patterns"] == ["*.bin"] + assert engine_args["block_size"] == 32 + assert engine_args["kv_cache_dtype"] == "fp8" + assert engine_args["prefill_schedule_interval"] == 2 + assert engine_args["scheduling_policy"] == "priority" + assert engine_args["data_parallel_address"] == "10.0.0.1" + assert engine_args["enable_dbo"] is True + + +def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + ) + + engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config) + + inherited_defaults = { + "ignore_patterns", + "block_size", + "kv_cache_dtype", + "prefill_schedule_interval", + "scheduling_policy", + "data_parallel_address", + "data_parallel_rank", + "data_parallel_size_local", + "data_parallel_rpc_port", + "enable_dbo", + } + assert inherited_defaults.isdisjoint(engine_args) + + +def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary(): + stage_config = VllmOmniARStageConfig( + stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"), + scheduler_config=OmniStageSchedulerConfig(scheduler_cls="test.CustomScheduler"), + ) + + with pytest.raises(ValueError, match=r"no EngineArgs projection: scheduler_cls"): + stage_init_utils._project_omni_stage_engine_args(stage_config) + + +@pytest.mark.parametrize("stage_id", [0, 1], ids=["ar", "generation"]) +def test_typed_llm_engine_args_preserve_upstream_config_objects(tmp_path, stage_id): + pipeline, deploy, model = _engine_arg_inputs(tmp_path) + compilation_config = VllmCompilationConfig(backend="eager") + profiler_config = VllmProfilerConfig(profiler="cuda") + stage_prefix = f"stage_{stage_id}_" + omni_config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=copy.deepcopy(deploy), + cli_overrides={ + "model": model, + f"{stage_prefix}compilation_config": compilation_config, + f"{stage_prefix}profiler_config": profiler_config, + }, + ) + + stage_config = omni_config.stage_by_id(stage_id) + typed_args = build_engine_args_dict_from_omni_stage_config(stage_config, model) + + assert isinstance(stage_config.compilation_config, VllmCompilationConfig) + assert isinstance(typed_args["compilation_config"], VllmCompilationConfig) + assert typed_args["compilation_config"] is not stage_config.compilation_config + assert typed_args["compilation_config"].backend == "eager" + + assert isinstance(stage_config.profiler_config, VllmProfilerConfig) + assert isinstance(typed_args["profiler_config"], VllmProfilerConfig) + assert typed_args["profiler_config"] is not stage_config.profiler_config + assert typed_args["profiler_config"].profiler == "cuda" + + def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path): pipeline, deploy, model = _engine_arg_inputs(tmp_path) legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model) @@ -218,9 +369,18 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path): ) typed_args_by_stage[stage_id] = typed_args - assert {name: typed_args[name] for name in legacy_args} == legacy_args + expected_args = {name: value for name, value in legacy_args.items() if name not in _TOPOLOGY_ONLY_ENGINE_ARGS} + assert {name: typed_args[name] for name in expected_args} == expected_args + assert _TOPOLOGY_ONLY_ENGINE_ARGS.isdisjoint(typed_args) thinker_args = typed_args_by_stage[0] + inherited_vllm_fields = { + field.name + for config_cls in (VllmLoadConfig, VllmCacheConfig, VllmSchedulerConfig, VllmParallelConfig) + for field in fields(config_cls) + } + topology_projected_fields = {"scheduler_cls"} + assert (inherited_vllm_fields - _LLM_STAGE_ENGINE_FIELDS - topology_projected_fields).isdisjoint(thinker_args) assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model") assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer") assert thinker_args["stage_id"] == 0 @@ -229,6 +389,7 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path): assert thinker_args["omni_kv_config"] == {"need_send_cache": True} assert thinker_args["worker_cls"] == "test.custom.Worker" assert thinker_args["attention_backend"] == "FLASHINFER" + assert thinker_args["attention_config"] == {"use_trtllm_attention": False} assert thinker_args["trust_remote_code"] is True assert thinker_args["dtype"] == "float16" assert thinker_args["distributed_executor_backend"] == "mp" @@ -469,8 +630,24 @@ def test_typed_ming_image_engine_args_defer_diffusion_batch_default(): assert OmniDiffusionConfig(**typed_backend_args).max_num_seqs == 1 +def _fake_model_root(pipeline, tmp_path): + """A model directory carrying every subfolder the pipeline's stages declare. + + Stage init fails closed when a declared ``model_subdir``/``tokenizer_subdir`` + is not a real directory, because the joined path would otherwise reach + HuggingFace as a malformed repo id (issue #6638). + """ + root = tmp_path / "model" + root.mkdir(exist_ok=True) + for stage in pipeline.stages: + for subdir in (stage.model_subdir, stage.tokenizer_subdir): + if subdir: + (root / subdir).mkdir(parents=True, exist_ok=True) + return str(root) + + @pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) -def test_typed_engine_args_match_current_registry_backend_semantics(model_type): +def test_typed_engine_args_match_current_registry_backend_semantics(model_type, tmp_path): pipeline = resolve_pipeline_config(model_type) if pipeline is None: pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve") @@ -480,18 +657,19 @@ def test_typed_engine_args_match_current_registry_backend_semantics(model_type): if pipeline.default_deploy_config_name is not None else DeployConfig() ) + model = _fake_model_root(pipeline, tmp_path) legacy_stages, omni_config = _legacy_and_typed_stages( pipeline, deploy, - model="/tmp", + model=model, ) for legacy_stage in legacy_stages: stage_id = legacy_stage.stage_id - legacy_args = build_legacy_engine_args_dict(legacy_stage, model="/tmp") + legacy_args = build_legacy_engine_args_dict(legacy_stage, model=model) typed_args = build_engine_args_dict_from_omni_stage_config( omni_config.stage_by_id(stage_id), - model="/tmp", + model=model, ) if legacy_stage.stage_type == StageType.DIFFUSION: backend_fields = _DIFFUSION_BACKEND_FIELDS @@ -500,6 +678,7 @@ def test_typed_engine_args_match_current_registry_backend_semantics(model_type): backend_fields = _LLM_BACKEND_FIELDS backend_config_cls = OmniEngineArgs + backend_fields -= _TOPOLOGY_ONLY_ENGINE_ARGS missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys() assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}" diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py new file mode 100644 index 000000000..e8a228827 --- /dev/null +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -0,0 +1,124 @@ +"""Unit tests for StagePool.collective_rpc EngineCore control dispatch.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from vllm_omni.engine.stage_pool import StagePool + +pytestmark = [pytest.mark.core_model] + + +def _make_pool(*, stage_type: str = "llm", **client_methods: AsyncMock) -> tuple[StagePool, SimpleNamespace]: + client = SimpleNamespace(stage_type=stage_type, **client_methods) + if "collective_rpc_async" not in client_methods: + client.collective_rpc_async = AsyncMock(return_value={"via": "collective"}) + pool = StagePool(0, [client]) # type: ignore[arg-type] + return pool, client + + +@pytest.mark.cpu +def test_collective_rpc_normalizes_none_args_on_control_helper(): + async def run() -> None: + pause = AsyncMock(return_value="paused") + pool, client = _make_pool(pause_scheduler_async=pause) + + result = await pool.collective_rpc(0, "pause_scheduler", args=None, kwargs={"mode": "abort"}) + + assert result == "paused" + pause.assert_awaited_once_with(mode="abort") + client.collective_rpc_async.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_unrelated_async_helper_uses_collective_path(): + async def run() -> None: + other = AsyncMock(return_value="should-not-run") + pool, client = _make_pool(reset_prefix_cache_async=other) + + result = await pool.collective_rpc(0, "reset_prefix_cache", timeout=1.5, args=("x",), kwargs={"k": 1}) + + assert result == {"via": "collective"} + other.assert_not_awaited() + client.collective_rpc_async.assert_awaited_once_with( + method="reset_prefix_cache", + timeout=1.5, + args=("x",), + kwargs={"k": 1}, + ) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_control_helper_honors_timeout(): + async def run() -> None: + async def slow_sleep(*_args, **_kwargs): + await asyncio.sleep(1.0) + return "slept" + + pool, _client = _make_pool(sleep_async=slow_sleep) + with pytest.raises(asyncio.TimeoutError): + await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,)) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_control_method_reraises_worker_error(): + async def run() -> None: + pool, _client = _make_pool(wake_up_async=AsyncMock(side_effect=RuntimeError("worker died"))) + with pytest.raises(RuntimeError, match="worker died"): + await pool.collective_rpc(0, "wake_up") + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_collective_rpc_non_control_method_still_returns_error_dict(): + async def run() -> None: + pool, _client = _make_pool(collective_rpc_async=AsyncMock(side_effect=RuntimeError("probe failed"))) + result = await pool.collective_rpc(0, "reset_prefix_cache") + assert result["supported"] is False + assert "probe failed" in result["error"] + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails(): + async def run() -> None: + class RecordingOutputProcessor: + def __init__(self) -> None: + self.collected = False + self.committed = False + + def abort_requests_collecting_outputs(self, request_ids, *, internal=False, commit_state=True): + del internal + self.collected = True + assert commit_state is False + return list(request_ids), [SimpleNamespace(request_id=request_ids[0])] + + def commit_aborted_request_state(self, request_ids, *, internal=False): + del request_ids, internal + self.committed = True + + abort = AsyncMock(side_effect=RuntimeError("engine abort failed")) + output_processor = RecordingOutputProcessor() + client = SimpleNamespace(stage_type="llm", abort_requests_async=abort) + pool = StagePool(0, [client], output_processor=output_processor) # type: ignore[arg-type] + pool._request_bindings["req-1"] = 0 + + with pytest.raises(RuntimeError, match="engine abort failed"): + await pool.abort_requests(["req-1"]) + assert output_processor.collected is True + assert output_processor.committed is False + abort.assert_awaited_once() + + asyncio.run(run()) diff --git a/tests/engine/test_stage_request_token_budget.py b/tests/engine/test_stage_request_token_budget.py new file mode 100644 index 000000000..045dfd31f --- /dev/null +++ b/tests/engine/test_stage_request_token_budget.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Token-budget clamping when the orchestrator builds a stage request.""" + +from __future__ import annotations + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _StageModelConfig: + """Only the field ``build_engine_core_request_from_tokens`` reads.""" + + max_model_len = 4096 + + +class TestStageRequestTokenBudget: + def test_min_tokens_is_clamped_to_the_remaining_context(self): + # A min_tokens wider than the leftover window makes ``check_stop`` + # return early forever: the request can neither stop nor be granted + # more tokens, so it stalls instead of being length-capped. + params = SamplingParams(max_tokens=4096, min_tokens=50) + + request = build_engine_core_request_from_tokens( + request_id="req-clamped", + prompt={"prompt_token_ids": [0] * 4050}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.min_tokens == 46 + assert params.min_tokens == 50 + + def test_min_tokens_survives_when_the_prompt_leaves_room(self): + params = SamplingParams(max_tokens=4096, min_tokens=50) + + request = build_engine_core_request_from_tokens( + request_id="req-untouched", + prompt={"prompt_token_ids": [0] * 68}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.min_tokens == 50 + assert request.sampling_params.max_tokens == 4096 + + def test_min_tokens_floors_at_zero_when_the_prompt_fills_the_context(self): + params = SamplingParams(max_tokens=16, min_tokens=8) + + request = build_engine_core_request_from_tokens( + request_id="req-full", + prompt={"prompt_token_ids": [0] * 4096}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.min_tokens == 0 + + def test_absent_max_tokens_still_defaults_to_the_remaining_context(self): + params = SamplingParams(min_tokens=0) + params.max_tokens = None + + request = build_engine_core_request_from_tokens( + request_id="req-default", + prompt={"prompt_token_ids": [0] * 96}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.max_tokens == 4000 diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py index 3c3666026..9fd61095b 100644 --- a/tests/entrypoints/openai/test_duplex_protocol.py +++ b/tests/entrypoints/openai/test_duplex_protocol.py @@ -452,7 +452,7 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease(): assert caps["supports_core_resumable_request"] is True assert caps["supports_stage_connector_handoff"] is True assert caps["supports_audio_truncate"] is True - assert caps["supports_barge_in"] is False + assert caps["supports_barge_in"] is True assert caps["target_barge_in_latency_ms"] is None assert caps["supports_multi_session"] is True assert caps["supports_multi_session_same_replica"] is True @@ -462,6 +462,15 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease(): assert caps["stage_handoff_transport"] == "scheduler_data_plane" +def test_duplex_overlap_policy_auto_falls_back_to_listen_only(): + from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexOverlapPolicy, + DuplexSessionConfig, + ) + + assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value + + def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment(): caps = DuplexCapabilities.minicpmo45_native(max_sessions=1).as_dict() diff --git a/tests/entrypoints/openai/test_profiler_endpoints.py b/tests/entrypoints/openai/test_profiler_endpoints.py new file mode 100644 index 000000000..02e48e19a --- /dev/null +++ b/tests/entrypoints/openai/test_profiler_endpoints.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm_omni.entrypoints.openai import api_server + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("path", "engine_method"), + [ + ("/start_profile", "start_profile"), + ("/stop_profile", "stop_profile"), + ], +) +def test_profiler_endpoint_forwards_selected_stages(path: str, engine_method: str) -> None: + engine_client = SimpleNamespace( + start_profile=AsyncMock(), + stop_profile=AsyncMock(), + ) + + app = FastAPI() + app.include_router(api_server.profiler_router) + app.state.engine_client = engine_client + + response = TestClient(app).post(path, json={"stages": [1]}) + + assert response.status_code == 200 + assert response.json() == {"status": "SUCCESS"} + getattr(engine_client, engine_method).assert_awaited_once_with(stages=[1]) diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py new file mode 100644 index 000000000..2b9ed6ec0 --- /dev/null +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -0,0 +1,882 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""API-server surface guards for the OpenAI Omni entrypoint. + +Purpose +------- +Freeze the *observable* server contract so helper moves (P0.2), route +extraction (P0.3), and later endpoint-family PRs cannot silently change: + +* which routes exist, +* how upstream routes are overridden at assembly time, +* unavailable / missing-handler envelopes, +* shared engine-error JSON fields, +* app-state keys wired by ``omni_init_app_state``. + +CPU-only (fakes / TestClient). No checked-in OpenAPI golden JSON. + +Common wrong changes these tests are meant to catch +---------------------------------------------------- +* Route accidentally removed, renamed, method-changed, or double-registered + after a move (census / assembled-app checks). +* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse). +* Assembly forgets to strip overridden upstream chat / models / profiler + routes, or also deletes unrelated upstream routes. +* Profiler / Omni router not mounted; storage not started; Omni engine + exception handlers not registered. +* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes. +* Missing-handler WS/HTTP responses change status, JSON shape, or wording + (clients treat these as capability probes). +* Engine error payloads lose ``request_id`` / ``error_stage_id``. +* Pure-diffusion vs multi-stage init drops an ``app.state`` key, or leaves a + handler as ``None`` while the route is still registered. + +Not covered here (add elsewhere if needed) +------------------------------------------ +* Deep per-endpoint success behavior (family UTs). +* Byte-identical helper bodies after a move (P0.2 AST integrity). +* Live GPU / speech byte-identity checks. +""" + +from __future__ import annotations + +import asyncio +import json +from contextlib import asynccontextmanager +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI, HTTPException +from fastapi.routing import APIRoute, APIWebSocketRoute +from fastapi.testclient import TestClient +from starlette.datastructures import State +from starlette.requests import Request +from starlette.websockets import WebSocketDisconnect +from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError + +from vllm_omni.entrypoints.openai import api_server + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here +# would invent ("HEAD", path) entries for every GET and fail the current census. +# TODO: extend this set and the expected census if @router.head / @router.options appear. +_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"} + +# Exact Omni router census on main. Update intentionally when routes change. +_EXPECTED_ROUTER_ROUTES = { + ("POST", "/v1/chat/completions"), + ("POST", "/v1/chat/completions/batch"), + ("POST", "/v1/audio/speech"), + ("POST", "/v1/audio/speech/batch"), + ("POST", "/v1/audio/generate"), + ("GET", "/v1/audio/voices"), + ("POST", "/v1/audio/voices"), + ("DELETE", "/v1/audio/voices/{name}"), + ("WEBSOCKET", "/v1/audio/speech/stream"), + ("WEBSOCKET", "/v1/video/chat/stream"), + ("WEBSOCKET", "/v1/realtime/video"), + ("WEBSOCKET", "/v1/realtime"), + ("WEBSOCKET", "/v1/realtime/robot/openpi"), + ("WEBSOCKET", "/v1/duplex"), + ("GET", "/health"), + ("GET", "/v1/models"), + ("POST", "/v1/images/generations"), + ("POST", "/v1/images/edits"), + ("POST", "/v1/videos"), + ("POST", "/v1/videos/sync"), + ("GET", "/v1/videos"), + ("GET", "/v1/videos/{video_id}"), + ("DELETE", "/v1/videos/{video_id}"), + ("GET", "/v1/videos/{video_id}/content"), + ("POST", "/v1/omni/sleep"), + ("POST", "/v1/omni/wakeup"), +} +_EXPECTED_PROFILER_ROUTES = { + ("POST", "/start_profile"), + ("POST", "/stop_profile"), +} + +# App-state keys that must remain available across refactor phases. +# Presence is not enough: "not wired" is written as ``state.X = None``. +_DIFFUSION_APP_STATE_KEYS = { + "engine_client", + "log_stats", + "args", + "sleeping_stages", + "stage_configs", + "vllm_config", + "diffusion_engine", + "openai_serving_models", + "serving_tokenization", + "openai_serving_chat", + "openai_serving_chat_batch", + "openai_serving_speech", + "openai_serving_audio_generate", + "openai_serving_video", + "openai_streaming_video_output", + "openai_serving_duplex", + "openai_streaming_speech", + "openai_streaming_video", + "openai_serving_realtime_robot", + "enable_server_load_tracking", + "server_load_metrics", +} +_DIFFUSION_MUST_BE_NONE = { + "vllm_config", + "serving_tokenization", + "openai_serving_duplex", + "openai_streaming_speech", + "openai_streaming_video", +} +_DIFFUSION_MUST_BE_WIRED = _DIFFUSION_APP_STATE_KEYS - _DIFFUSION_MUST_BE_NONE +_MULTISTAGE_APP_STATE_KEYS = { + "engine_client", + "log_stats", + "args", + "sleeping_stages", + "stage_configs", + "vllm_config", + "openai_serving_models", + "serving_tokenization", + "online_renderer", + "openai_serving_chat", + "openai_serving_chat_batch", + "openai_serving_speech", + "openai_serving_audio_generate", + "openai_streaming_speech", + "openai_streaming_video", + "openai_serving_duplex", + "openai_serving_realtime", + "openai_serving_video", + "openai_serving_realtime_robot", + "enable_server_load_tracking", + "server_load_metrics", +} +_MULTISTAGE_MUST_BE_NONE = { + "openai_serving_duplex", + "openai_serving_realtime_robot", +} +_MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE + + +def _route_entries(routes) -> list[tuple[str, str]]: + entries: list[tuple[str, str]] = [] + for route in routes: + if isinstance(route, APIRoute): + for method in sorted((route.methods or set()) & _HTTP_METHODS): + entries.append((method, route.path)) + elif isinstance(route, APIWebSocketRoute): + entries.append(("WEBSOCKET", route.path)) + return entries + + +def _route_manifest(routes) -> set[tuple[str, str]]: + return set(_route_entries(routes)) + + +def _assert_unique_route_entries(routes) -> None: + entries = _route_entries(routes) + duplicates = sorted({item for item in entries if entries.count(item) > 1}) + assert not duplicates, f"duplicate route registrations: {duplicates}" + + +def _matching_http_routes(routes, method: str, path: str) -> list[APIRoute]: + return [ + route + for route in routes + if isinstance(route, APIRoute) and route.path == path and method in (route.methods or set()) + ] + + +def _single_http_route(routes, method: str, path: str) -> APIRoute: + matches = _matching_http_routes(routes, method, path) + assert len(matches) == 1, f"expected exactly one {method} {path}, got {len(matches)}" + return matches[0] + + +def _assert_app_state_snapshot( + state, + *, + expected_keys: set[str], + must_be_wired: set[str], + must_be_none: set[str], +) -> None: + assert must_be_wired.isdisjoint(must_be_none) + assert must_be_wired | must_be_none <= expected_keys + present = {key for key in expected_keys if hasattr(state, key)} + assert present == expected_keys + not_wired = sorted(key for key in must_be_wired if getattr(state, key) is None) + assert not not_wired, f"app.state keys registered but not wired: {not_wired}" + unexpectedly_set = sorted(key for key in must_be_none if getattr(state, key) is not None) + assert not unexpectedly_set, f"app.state keys expected to stay None: {unexpectedly_set}" + + +def _minimal_args(**overrides) -> SimpleNamespace: + args = SimpleNamespace( + tool_parser_plugin=None, + reasoning_parser_plugin=None, + reasoning_parser=None, + structured_outputs_config=SimpleNamespace(reasoning_parser=None), + enable_ssl_refresh=False, + host="127.0.0.1", + port=8000, + uvicorn_log_level="info", + disable_uvicorn_access_log=True, + ssl_keyfile=None, + ssl_certfile=None, + ssl_ca_certs=None, + ssl_cert_reqs=0, + ssl_ciphers=None, + h11_max_incomplete_event_size=None, + h11_max_header_count=None, + log_config_file=None, + disable_access_log_for_endpoints=None, + disable_log_stats=True, + enable_log_requests=False, + max_log_len=None, + served_model_name=None, + model="demo-model", + enable_server_load_tracking=False, + log_error_stack=False, + chat_template=None, + chat_template_content_format="auto", + trust_request_chat_template=False, + tool_server=None, + lora_modules=None, + response_role="assistant", + default_chat_template_kwargs=None, + return_tokens_as_token_ids=False, + enable_auto_tool_choice=False, + exclude_tools_when_tool_choice_none=False, + tool_call_parser=None, + enable_prompt_tokens_details=False, + enable_force_include_usage=False, + enable_log_outputs=False, + enable_log_deltas=False, + tokens_only=False, + shutdown_timeout=0, + stage_configs_path=None, + deploy_config=None, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +def _request_for(app: FastAPI, *, method: str = "GET", path: str = "/") -> Request: + return Request( + { + "type": "http", + "method": method, + "path": path, + "headers": [], + "app": app, + "state": {}, + } + ) + + +class _FakeSocket: + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + +class _FakeEngineClient: + def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None: + self.stage_configs = stage_configs if stage_configs is not None else [] + self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {} + self.model_config = SimpleNamespace() + self.input_processor = object() + self.renderer = object() + self.errored = False + self.vllm_config = vllm_config + + async def get_vllm_config(self): + return self.vllm_config + + async def get_supported_tasks(self) -> tuple[str, ...]: + return ("generate",) + + async def get_tokenizer(self): + return SimpleNamespace(chat_template="dummy") + + async def check_health(self) -> None: + return None + + +def _marker(name: str): + return SimpleNamespace(name=name) + + +def test_router_route_manifest_is_explicit_and_complete() -> None: + """Lock the Omni + profiler router census. + + Fails if a route is added/removed/renamed, a method changes, or a + WebSocket path disappears during helper/router moves. Also fails if the + same public route is registered more than once, which a set manifest alone + would collapse. + """ + # TODO(P0.3): retarget this object-level census when routes move onto family + # routers. The assembled-app census is the client-facing lock that must survive. + _assert_unique_route_entries(api_server.router.routes) + _assert_unique_route_entries(api_server.profiler_router.routes) + assert _route_manifest(api_server.router.routes) == _EXPECTED_ROUTER_ROUTES + assert _route_manifest(api_server.profiler_router.routes) == _EXPECTED_PROFILER_ROUTES + + +def test_router_openapi_paths_cover_http_manifest() -> None: + """Lock that OpenAPI advertises every HTTP census route/method. + + Fails if a decorator still registers a route but OpenAPI loses the path + (or method), e.g. after moving route bodies / response models. + """ + app = FastAPI() + app.include_router(api_server.router) + app.include_router(api_server.profiler_router) + paths = app.openapi().get("paths", {}) + + http_routes = { + (method, path) for method, path in _EXPECTED_ROUTER_ROUTES | _EXPECTED_PROFILER_ROUTES if method != "WEBSOCKET" + } + for method, path in sorted(http_routes): + assert path in paths, f"missing OpenAPI path {path}" + assert method.lower() in paths[path], f"missing OpenAPI method {method} {path}" + + +@pytest.mark.asyncio +async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None: + """Lock worker assembly: override, mount, storage, handlers, full census. + + Fails if assembly stops replacing upstream chat/batch/models/profiler, + deletes unrelated upstream routes, forgets Omni/profiler mounts, skips + storage start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError`` + handlers. + """ + captured: dict[str, object] = {} + + def fake_build_openai_app(args, supported_tasks): + app = FastAPI() + + @app.post("/v1/chat/completions") + async def upstream_chat(): + return {"owner": "upstream"} + + @app.post("/v1/chat/completions/batch") + async def upstream_batch(): + return {"owner": "upstream"} + + @app.get("/v1/models") + async def upstream_models(): + return {"owner": "upstream"} + + @app.post("/start_profile") + async def upstream_start_profile(): + return {"owner": "upstream"} + + @app.post("/stop_profile") + async def upstream_stop_profile(): + return {"owner": "upstream"} + + @app.get("/v1/upstream-only") + async def upstream_only(): + return {"owner": "upstream"} + + captured["supported_tasks"] = supported_tasks + return app + + @asynccontextmanager + async def fake_build_async_omni(*_args, **_kwargs): + yield _FakeEngineClient( + stage_configs=[{"engine_args": {"profiler_config": {"profiler": "torch"}}}], + vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(_api_process_rank=0)), + ) + + async def fake_omni_init_app_state(engine_client, state, args): + state.engine_client = engine_client + state.initialized_by_omni = True + state.openai_serving_video = SimpleNamespace( + shutdown=lambda: captured.__setitem__("video_shutdown", True), + ) + state.openai_serving_speech = SimpleNamespace( + shutdown=lambda: captured.__setitem__("speech_shutdown", True), + ) + + async def fake_storage_start(): + captured["storage_started"] = True + + async def fake_serve_http(app, **_kwargs): + captured["served_app"] = app + task = asyncio.create_task(asyncio.sleep(0)) + await asyncio.sleep(0) + return task + + monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app) + monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni) + monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state) + monkeypatch.setattr( + api_server, + "shutdown_unsupported_routes", + lambda app, restrictions: captured.setdefault("restrictions", restrictions), + ) + monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=fake_storage_start)) + monkeypatch.setattr(api_server, "serve_http", fake_serve_http) + + sock = _FakeSocket() + await api_server.omni_run_server_worker("127.0.0.1:8000", sock, _minimal_args()) + + served_wrapper = captured["served_app"] + # Timestamp middleware wraps the FastAPI app. + assert hasattr(served_wrapper, "_inner") + served_app = served_wrapper._inner + routes = served_app.routes + _assert_unique_route_entries(routes) + manifest = _route_manifest(routes) + + assert captured["supported_tasks"] == ("generate",) + assert captured["storage_started"] is True + assert captured["restrictions"] == {} + assert captured["video_shutdown"] is True + assert captured["speech_shutdown"] is True + assert sock.closed is True + assert served_app.state.initialized_by_omni is True + + # Upstream override: unrelated route kept; Omni owns overridden pairs once. + # Endpoint identity ensures the surviving route is the Omni handler, not + # the fake upstream handler with the same method/path. + assert ("GET", "/v1/upstream-only") in manifest + assert _single_http_route(routes, "POST", "/v1/chat/completions").endpoint is api_server.create_chat_completion + assert ( + _single_http_route(routes, "POST", "/v1/chat/completions/batch").endpoint + is api_server.create_batch_chat_completion + ) + assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models + assert _single_http_route(routes, "POST", "/start_profile").endpoint is api_server.start_profile + assert _single_http_route(routes, "POST", "/stop_profile").endpoint is api_server.stop_profile + + # Assembled app includes Omni router HTTP/WS surface + profiler. + for item in _EXPECTED_ROUTER_ROUTES: + assert item in manifest, f"assembled app missing {item}" + assert ("POST", "/start_profile") in manifest + assert ("POST", "/stop_profile") in manifest + + assert served_app.exception_handlers[EngineGenerateError] + assert served_app.exception_handlers[EngineDeadError] + + # Light OpenAPI check on the assembled FastAPI app (not the ASGI wrapper). + openapi_paths = served_app.openapi().get("paths", {}) + assert "/v1/chat/completions" in openapi_paths + assert "/v1/models" in openapi_paths + assert "/health" in openapi_paths + + +@pytest.mark.asyncio +async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch) -> None: + """Lock outermost timestamp middleware behavior. + + Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes + start getting stamped / mutated incorrectly. + """ + captured: dict[str, object] = {} + + def fake_build_openai_app(args, supported_tasks): + return FastAPI() + + @asynccontextmanager + async def fake_build_async_omni(*_args, **_kwargs): + yield _FakeEngineClient() + + async def fake_omni_init_app_state(engine_client, state, args): + state.engine_client = engine_client + + async def fake_serve_http(app, **_kwargs): + captured["served_app"] = app + task = asyncio.create_task(asyncio.sleep(0)) + await asyncio.sleep(0) + return task + + monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app) + monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni) + monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state) + monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_a, **_k: None) + monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=lambda: asyncio.sleep(0))) + monkeypatch.setattr(api_server, "serve_http", fake_serve_http) + + await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args()) + middleware = captured["served_app"] + + http_scope = {"type": "http", "state": {}} + ws_scope = {"type": "websocket"} + + async def _receive(): + return {"type": "http.disconnect"} + + async def _send(_message): + return None + + called = {"http": False, "websocket": False} + + async def inner(scope, receive, send): + called[scope["type"]] = True + + middleware._inner = inner + await middleware(http_scope, _receive, _send) + await middleware(ws_scope, _receive, _send) + + assert called["http"] is True + assert called["websocket"] is True + assert "request_timestamp" in http_scope["state"] + assert isinstance(http_scope["state"]["request_timestamp"], float) + assert "request_timestamp" not in ws_scope.get("state", {}) + + +@pytest.mark.parametrize( + ("path", "payload"), + [ + ("/v1/audio/speech/stream", {"type": "error", "message": "Streaming speech is not available"}), + ("/v1/video/chat/stream", {"type": "error", "message": "Streaming video chat is not available"}), + ("/v1/realtime/video", {"type": "error", "message": "Streaming video generation is not available"}), + ("/v1/realtime", {"type": "error", "error": "Realtime API is not available", "code": "unsupported"}), + ( + "/v1/realtime/robot/openpi", + {"type": "error", "error": "Robot policy not available", "code": "unsupported"}, + ), + ("/v1/duplex", {"type": "error", "error": "Duplex API is not available", "code": "unsupported"}), + ], +) +def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, payload: dict[str, str]) -> None: + """Lock missing-handler WebSocket capability probes. + + Fails if unavailable JSON changes (keys/wording/code) or the socket stays + open after the error frame — common when WS wrappers move between modules. + """ + app = FastAPI() + app.include_router(api_server.router) + + with TestClient(app) as client: + with client.websocket_connect(path) as websocket: + assert websocket.receive_json() == payload + with pytest.raises(WebSocketDisconnect): + websocket.receive_text() + + +def test_health_without_engine_returns_stable_unhealthy_response() -> None: + """Lock ``/health`` with no engine initialized. + + Fails if status drifts from 503 or the unhealthy payload shape/reason + string changes during health-route extraction. + """ + app = FastAPI() + response = asyncio.run(api_server.health(_request_for(app, path="/health"))) + + assert response.status_code == 503 + assert json.loads(response.body) == { + "status": "unhealthy", + "reason": "No engine initialized", + } + + +def test_models_without_handler_returns_empty_openai_list() -> None: + """Lock ``/v1/models`` when ``openai_serving_models`` is unset. + + Fails if the empty listing becomes an error, or the OpenAI list envelope + changes (``object`` / ``data``). + """ + app = FastAPI() + response = asyncio.run(api_server.show_available_models(_request_for(app, path="/v1/models"))) + + assert response.status_code == 200 + assert json.loads(response.body) == {"object": "list", "data": []} + + +def test_images_generation_without_engine_preserves_service_unavailable_error() -> None: + """Lock images generation when the multi-stage engine is missing. + + Fails if the 503 detail string changes — clients/tests use it as a + capability signal when bootstrap failed. + """ + app = FastAPI() + raw_request = _request_for(app, method="POST", path="/v1/images/generations") + request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model") + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(api_server.generate_images(request, raw_request)) + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model." + + +def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None: + """Lock images generation when chat serving was not wired. + + Fails if multi-stage images stop requiring ``openai_serving_chat``, or the + 503 detail changes after images/chat ownership splits. + """ + app = FastAPI() + app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()]) + app.state.stage_configs = app.state.engine_client.stage_configs + app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")]) + app.state.openai_serving_chat = None + app.state.args = SimpleNamespace(max_generated_image_size=None) + + raw_request = _request_for(app, method="POST", path="/v1/images/generations") + request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model") + monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion") + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(api_server.generate_images(request, raw_request)) + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == "openai_serving_chat is not initialized for multi-stage image generation." + + +def test_speech_without_handler_preserves_not_found_http_error() -> None: + """Lock speech HTTP when ``openai_serving_speech`` is unset. + + Fails if missing speech becomes a different status/detail after TTS route + wrappers move out of ``api_server``. + """ + app = FastAPI() + app.state.openai_serving_speech = None + raw_request = _request_for(app, method="POST", path="/v1/audio/speech") + request = api_server.OpenAICreateSpeechRequest(input="hello", model="demo-model", voice="alloy") + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(api_server.create_speech(request, raw_request)) + + assert exc_info.value.status_code == 404 + assert exc_info.value.detail == "The model does not support Speech API" + + +def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None: + """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers. + + Fails if ``request_id`` / ``error_stage_id`` are dropped when error helpers + move (P0.2) or exception-handler wiring changes. + + Exercises the envelope through the registered handler, not the private + ``_create_engine_error_json_response`` helper, so relocating that helper + in P0.2 is not a false red. + """ + app = FastAPI() + # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers`` + # leaves api_server (red here is a rename, not a surface regression). + api_server._register_omni_exception_handlers(app) + app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False)) + app.state.server = object() + app.state.args = SimpleNamespace(log_error_stack=False) + + req = _request_for(app, method="POST", path="/v1/chat/completions") + req.state.request_metadata = SimpleNamespace(request_id="req-123") + + monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None) + + exc = EngineGenerateError("boom") + exc.error_stage_id = 2 # type: ignore[attr-defined] + handler = app.exception_handlers[EngineGenerateError] + response = asyncio.run(handler(req, exc)) + payload = json.loads(response.body) + + assert "error" in payload + assert payload["error"]["request_id"] == "req-123" + assert payload["error"]["error_stage_id"] == 2 + assert response.status_code >= 400 + + +def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None: + """Lock Omni ``EngineDeadError`` handler registration + JSON envelope. + + Fails if registration is skipped or the handler stops returning OpenAI-style + JSON with ``request_id``. + """ + app = FastAPI() + # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers`` + # leaves api_server (red here is a rename, not a surface regression). + api_server._register_omni_exception_handlers(app) + app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False)) + app.state.server = object() + app.state.args = SimpleNamespace(log_error_stack=False) + + monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None) + + handler = app.exception_handlers[EngineDeadError] + req = _request_for(app) + req.state.request_metadata = SimpleNamespace(request_id="dead-1") + response = asyncio.run(handler(req, EngineDeadError())) + payload = json.loads(response.body) + assert payload["error"]["request_id"] == "dead-1" + assert response.status_code >= 400 + + +@pytest.mark.asyncio +async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None: + """Lock pure-diffusion ``app.state`` keys after init, including live vs None. + + Fails if bootstrap drops keys route owners still read (e.g. video/speech/ + models), wires a handler that must stay None, or leaves a required + handler as None. + """ + stage = SimpleNamespace(engine_args={}) + engine = _FakeEngineClient(stage_configs=[stage]) + + def _for_diffusion_factory(label: str): + @classmethod + def _factory(cls, *args, **kwargs): + return _marker(label) + + return _factory + + monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion") + monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat")) + monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch")) + monkeypatch.setattr( + api_server.OmniOpenAIServingAudioGenerate, + "for_diffusion", + _for_diffusion_factory("audio_generate"), + ) + monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video")) + monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None) + monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _for_diffusion_factory("speech")) + monkeypatch.setattr( + api_server.ServingRealtimeRobotOpenPI, + "create_policy_server", + classmethod(lambda cls, *a, **k: _marker("openpi")), + ) + + state = State() + await api_server.omni_init_app_state( + engine, + state, + _minimal_args(), + ) + + _assert_app_state_snapshot( + state, + expected_keys=_DIFFUSION_APP_STATE_KEYS, + must_be_wired=_DIFFUSION_MUST_BE_WIRED, + must_be_none=_DIFFUSION_MUST_BE_NONE, + ) + assert state.diffusion_engine is engine + + +@pytest.mark.asyncio +async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> None: + stage = SimpleNamespace(engine_args={}) + engine = _FakeEngineClient(stage_configs=[stage]) + speech_kwargs = {} + + def _for_diffusion_factory(label: str): + @classmethod + def _factory(cls, *args, **kwargs): + return _marker(label) + + return _factory + + @classmethod + def _speech_factory(cls, *args, **kwargs): + speech_kwargs.update(kwargs) + return _marker("speech") + + monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion") + monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat")) + monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch")) + monkeypatch.setattr( + api_server.OmniOpenAIServingAudioGenerate, + "for_diffusion", + _for_diffusion_factory("audio_generate"), + ) + monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video")) + monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None) + monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _speech_factory) + monkeypatch.setattr( + api_server.ServingRealtimeRobotOpenPI, + "create_policy_server", + classmethod(lambda cls, *a, **k: _marker("openpi")), + ) + + state = State() + await api_server.omni_init_app_state( + engine, + state, + _minimal_args( + allowed_local_media_path="/allowed/media", + allowed_media_domains=["media.example.com"], + ), + ) + + assert speech_kwargs["allowed_local_media_path"] == "/allowed/media" + assert speech_kwargs["allowed_media_domains"] == ["media.example.com"] + + +@pytest.mark.asyncio +async def test_multistage_app_state_key_snapshot(monkeypatch) -> None: + """Lock multi-stage ``app.state`` keys after init, including live vs None. + + Fails if chat/speech/video/realtime/tokenization keys disappear or are + left as None after init refactors — classic “route still registered, + handler never wired” failure mode. + """ + engine = _FakeEngineClient( + stage_configs=[object(), object()], + vllm_config=SimpleNamespace( + lora_config=None, + model_config=SimpleNamespace(), + parallel_config=SimpleNamespace(_api_process_rank=0), + ), + ) + + class _FakeModels: + def __init__(self, *args, **kwargs): + self.base_model_paths = kwargs.get("base_model_paths") or [] + + async def init_static_loras(self): + return None + + class _FakeCtor: + def __init__(self, *args, **kwargs): + pass + + def warmup(self): + return None + + class _FakeSpeech(_FakeCtor): + async def warmup(self): + return None + + monkeypatch.setattr(api_server, "load_chat_template", lambda *_a, **_k: None) + monkeypatch.setattr(api_server, "process_lora_modules", lambda modules, _defaults: modules or []) + monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels) + monkeypatch.setattr(api_server, "OnlineRenderer", _FakeCtor) + monkeypatch.setattr(api_server, "OpenAIServingResponses", _FakeCtor) + monkeypatch.setattr(api_server, "OmniOpenAIServingChat", _FakeCtor) + monkeypatch.setattr(api_server, "OmniOpenAIServingChatBatch", _FakeCtor) + monkeypatch.setattr(api_server, "OpenAIServingCompletion", _FakeCtor) + monkeypatch.setattr(api_server, "ServingPooling", _FakeCtor) + monkeypatch.setattr(api_server, "OpenAIServingEmbedding", _FakeCtor) + monkeypatch.setattr(api_server, "ServingClassification", _FakeCtor) + monkeypatch.setattr(api_server, "ServingScores", _FakeCtor) + monkeypatch.setattr(api_server, "ServingTokenization", _FakeCtor) + monkeypatch.setattr(api_server, "OpenAIServingTranscription", _FakeCtor) + monkeypatch.setattr(api_server, "OpenAIServingTranslation", _FakeCtor) + monkeypatch.setattr(api_server, "AnthropicServingMessages", _FakeCtor) + monkeypatch.setattr(api_server, "ServingTokens", _FakeCtor) + monkeypatch.setattr(api_server, "OmniOpenAIServingSpeech", _FakeSpeech) + monkeypatch.setattr(api_server, "OmniOpenAIServingAudioGenerate", _FakeCtor) + monkeypatch.setattr(api_server, "OmniStreamingSpeechHandler", _FakeCtor) + monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video")) + monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor) + monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor) + monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False) + + state = State() + await api_server.omni_init_app_state(engine, state, _minimal_args()) + + _assert_app_state_snapshot( + state, + expected_keys=_MULTISTAGE_APP_STATE_KEYS, + must_be_wired=_MULTISTAGE_MUST_BE_WIRED, + must_be_none=_MULTISTAGE_MUST_BE_NONE, + ) diff --git a/tests/entrypoints/openai_api/test_audex_serving_guards.py b/tests/entrypoints/openai_api/test_audex_serving_guards.py index 68661865f..27ed11356 100644 --- a/tests/entrypoints/openai_api/test_audex_serving_guards.py +++ b/tests/entrypoints/openai_api/test_audex_serving_guards.py @@ -26,8 +26,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech: + from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter + from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) serving._tts_model_type = model_type + adapter_cls = resolve_adapter(model_type) + serving._adapter = adapter_cls(SpeechServingContext(server=serving)) if adapter_cls is not None else None serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None serving._discard_ref_audio_artifact_warmup = lambda request_id: None return serving @@ -154,7 +159,7 @@ class TestAudexCfgValidation: assert _adapter().validate(_speech_request(cfg_scale=None)) is None serving = TestAudexCfgInjection()._serving_with_fake_tokenizer() params = SimpleNamespace(extra_args={"cfg_scale": None}) - out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params) + out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params) assert out.extra_args == {} @pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt"]) @@ -177,21 +182,21 @@ class TestAudexCfgInjection: def test_no_extra_args_is_untouched(self): serving = self._serving_with_fake_tokenizer() params = SimpleNamespace(extra_args=None) - out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params) + out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params) assert out.extra_args is None assert params.extra_args is None def test_scale_one_leaves_params_byte_identical(self): serving = self._serving_with_fake_tokenizer() params = SimpleNamespace(extra_args={"cfg_scale": 1.0}) - out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params) + out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params) assert out.extra_args == {} def test_guided_request_gains_pair_contract(self): serving = self._serving_with_fake_tokenizer() params = SimpleNamespace(extra_args={"cfg_scale": 1.5}) cond = self._cond_prompt("Some transcription text here.") - out = serving._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params) + out = serving._adapter._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params) extra = out.extra_args assert extra["cfg_scale"] == 1.5 @@ -207,14 +212,14 @@ class TestAudexCfgInjection: serving = self._serving_with_fake_tokenizer() params = SimpleNamespace(extra_args={"cfg_scale": 1.5}) with pytest.raises(ValueError, match="adapter-built text prompt"): - serving._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params) + serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params) def test_injection_never_mutates_input_params(self): """The input may be the engine's SHARED default params (review P2): pair state must land only on the returned clone.""" serving = self._serving_with_fake_tokenizer() shared = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1) - out = serving._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared) + out = serving._adapter._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared) assert out is not shared assert shared.extra_args == {"cfg_scale": 1.5} assert shared.temperature == 0.1 @@ -252,11 +257,11 @@ def test_guided_injection_sets_guided_temperature(): serving = TestAudexCfgInjection()._serving_with_fake_tokenizer() cond = TestAudexCfgInjection()._cond_prompt() params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1) - out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params) + out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params) assert out.temperature == 0.05 unguided = SimpleNamespace(extra_args={"cfg_scale": 1.0}, temperature=0.1) - out = serving._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided) + out = serving._adapter._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided) assert out.temperature == 0.1 @@ -340,7 +345,7 @@ class TestAudexTTAInjection: def test_default_injection_applies_rvq_and_cfg3(self): serving = self._serving() params = SimpleNamespace(extra_args=None) - out = serving._inject_audex_tta_args("req-9", self._prompt(), params) + out = serving._adapter._inject_audex_tta_args("req-9", self._prompt(), params) extra = out.extra_args rvq = extra["tta_rvq"] @@ -355,15 +360,15 @@ class TestAudexTTAInjection: def test_explicit_scale_one_disables_cfg_but_keeps_rvq(self): serving = self._serving() params = SimpleNamespace(extra_args={"cfg_scale": 1.0}) - out = serving._inject_audex_tta_args("req-1", self._prompt(), params) + out = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), params) assert "tta_rvq" in out.extra_args assert "cfg_scale" not in out.extra_args assert "cfg_role" not in out.extra_args def test_rvq_contract_cached_per_process(self): serving = self._serving() - out_a = serving._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None)) - out_b = serving._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None)) + out_a = serving._adapter._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None)) + out_b = serving._adapter._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None)) assert out_a.extra_args["tta_rvq"] is out_b.extra_args["tta_rvq"] def test_injection_never_mutates_input_params(self): @@ -372,8 +377,8 @@ class TestAudexTTAInjection: per-request pair state must land only on the returned clone.""" serving = self._serving() shared = SimpleNamespace(extra_args=None) - out_1 = serving._inject_audex_tta_args("req-1", self._prompt(), shared) - out_2 = serving._inject_audex_tta_args("req-2", self._prompt(), shared) + out_1 = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), shared) + out_2 = serving._adapter._inject_audex_tta_args("req-2", self._prompt(), shared) assert shared.extra_args is None assert out_1 is not shared and out_2 is not shared assert out_1.extra_args["cfg_pair_id"] == "req-1" diff --git a/tests/entrypoints/openai_api/test_decode_audio_url.py b/tests/entrypoints/openai_api/test_decode_audio_url.py index 2a5c5be47..5dff3b052 100644 --- a/tests/entrypoints/openai_api/test_decode_audio_url.py +++ b/tests/entrypoints/openai_api/test_decode_audio_url.py @@ -5,6 +5,7 @@ import asyncio import base64 import os +import tempfile import httpx import pytest @@ -101,7 +102,7 @@ class TestDecodeAudioUrlSuffixSanitization: try: # ".." contains non-alnum dots, so falls back to .wav assert path.endswith(".wav") - assert os.path.dirname(path) == "/tmp" + assert os.path.dirname(path) == tempfile.gettempdir() finally: _cleanup(path) diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index a37d9d2b1..c75f1e599 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -30,6 +33,7 @@ from vllm_omni.experimental.fullduplex.minicpmo45.runtime import ( from vllm_omni.experimental.fullduplex.minicpmo45.session import ( MiniCPMO45ServingSessionState, ) +from vllm_omni.experimental.fullduplex.openai import vad as realtime_vad from vllm_omni.experimental.fullduplex.openai.protocol import ( DuplexCapabilities, DuplexOverlapPolicy, @@ -46,6 +50,9 @@ from vllm_omni.experimental.fullduplex.openai.serving import ( ) from vllm_omni.experimental.fullduplex.openai.websocket import DuplexWebSocketActor from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision +from vllm_omni.experimental.fullduplex.personaplex.serving_adapter import ( + PersonaPlexServingRuntimeAdapter, +) from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -105,16 +112,16 @@ class FakeEngineClient: self.close_result = close_result self.opened: list[str] = [] self.opened_fences: list[DuplexFence | None] = [] - self.appended: list[tuple[str, str, object, bool]] = [] + self.appended: list[tuple[str, str, Any, bool]] = [] self.append_operation_ids: list[str | None] = [] self.appended_fences: list[DuplexFence | None] = [] - self.opened_configs: list[dict[str, object]] = [] - self.opened_runtime_configs: list[dict[str, object]] = [] + self.opened_configs: list[dict[str, Any]] = [] + self.opened_runtime_configs: list[dict[str, Any]] = [] self.signals: list[tuple[str, str]] = [] self.signal_fences: list[DuplexFence | None] = [] self.signal_next_fences: list[DuplexFence | None] = [] - self.signal_session_configs: list[dict[str, object] | None] = [] - self.signal_runtime_configs: list[dict[str, object] | None] = [] + self.signal_session_configs: list[dict[str, Any] | None] = [] + self.signal_runtime_configs: list[dict[str, Any] | None] = [] self.closed: list[tuple[str, str]] = [] self.touched: list[tuple[str, DuplexLeaseActivity]] = [] self.touch_fences: list[DuplexFence] = [] @@ -136,7 +143,7 @@ class FakeEngineClient: runtime_config: dict[str, object] | None = None, timeout: float | None = None, fence: DuplexFence | None = None, - ) -> None: + ) -> Any: del timeout if self.fail_open: raise RuntimeError("open failed") @@ -157,7 +164,7 @@ class FakeEngineClient: timeout: float | None = None, collect_outputs: bool = True, fence: DuplexFence | None = None, - ) -> None: + ) -> Any: del timeout, collect_outputs self.appended.append((session_id, mode, payload, final)) self.append_operation_ids.append(operation_id) @@ -191,7 +198,7 @@ class FakeEngineClient: next_fence: DuplexFence | None = None, session_config: dict[str, object] | None = None, runtime_config: dict[str, object] | None = None, - ) -> None: + ) -> Any: del timeout if self.fail_signal or event in self.fail_signal_events: raise RuntimeError("signal failed") @@ -209,7 +216,7 @@ class FakeEngineClient: reason: str = "client_close", timeout: float | None = None, fence: DuplexFence | None = None, - ) -> None: + ) -> Any: del timeout if self.fail_close: raise RuntimeError("close failed") @@ -344,7 +351,7 @@ def _project_data_plane( result: object, *, session: DuplexSession | None = None, -) -> list[dict[str, object]]: +) -> list[dict[str, Any]]: return list(data_plane.project(result, context=_data_plane_context(session))) @@ -545,6 +552,54 @@ async def test_native_realtime_protocol_conversation_item_create_commits_user_te } +@pytest.mark.asyncio +async def test_function_output_is_recorded_only_after_runtime_delivery(): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + protocol.encode_outbound_event( + { + "type": "conversation.item.created", + "item": { + "id": "item-call", + "type": "function_call", + "call_id": "call-1", + "name": "weather", + "arguments": "{}", + "status": "completed", + }, + } + ) + event = { + "type": "conversation.item.create", + "item": { + "id": "item-output", + "type": "function_call_output", + "call_id": "call-1", + "output": '{"temperature":20}', + }, + } + + first = await protocol._to_duplex_event(event) + retry = await protocol._to_duplex_event(event) + + assert first is not None + assert retry is not None + assert "item-output" not in protocol._conversation_items + + delivered_item = first["payload"]["item"] + protocol.encode_outbound_event( + { + "type": "conversation.item.created", + "item": delivered_item, + } + ) + + assert protocol._conversation_items["item-output"] == delivered_item + assert await protocol._to_duplex_event(event) is None + assert ws.sent[-1]["error"]["code"] == "invalid_function_call_output" + + @pytest.mark.asyncio async def test_native_realtime_protocol_audio_commit_requires_non_empty_buffer(): ws = TimedWebSocket() @@ -652,7 +707,7 @@ async def test_native_realtime_protocol_preserves_input_turn_policy_hints(): @pytest.mark.asyncio -async def test_native_realtime_protocol_rejects_unimplemented_server_vad(): +async def test_native_realtime_protocol_rejects_server_vad_without_interrupt(): ws = TimedWebSocket() protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] protocol.bind_sender(ws.send_json) @@ -675,7 +730,11 @@ async def test_native_realtime_protocol_rejects_unimplemented_server_vad(): error = ws.sent[-1] assert error["type"] == "error" assert error["error"]["code"] == "unsupported_turn_detection" - assert "server_vad" in error["error"]["message"] + assert "interrupt_response=false" in error["error"]["message"] + threshold_error = protocol._validate_realtime_turn_detection( + {"turn_detection": {"type": "server_vad", "threshold": 0.15}} + ) + assert threshold_error == "turn_detection.threshold must be greater than 0.15 and at most 1" @pytest.mark.asyncio @@ -699,7 +758,7 @@ async def test_native_realtime_protocol_accepts_disabled_turn_detection(): @pytest.mark.asyncio -async def test_native_realtime_protocol_rejects_nested_vad_even_when_top_level_is_disabled(): +async def test_native_realtime_protocol_prefers_nested_vad_over_top_level(): ws = TimedWebSocket() protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] protocol.bind_sender(ws.send_json) @@ -719,8 +778,55 @@ async def test_native_realtime_protocol_rejects_nested_vad_even_when_top_level_i } ) - assert translated is None - assert ws.sent[-1]["error"]["code"] == "unsupported_turn_detection" + assert translated["session"]["extra_body"]["realtime_session_payload"]["overlap_policy"] == "barge_in_on_speech" + + +@pytest.mark.asyncio +async def test_realtime_vad_edge_order_and_tiny_chunk_resampling(monkeypatch): + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + protocol.bind_sender(ws.send_json) + protocol._server_vad = realtime_vad.SileroStreamingVAD(realtime_vad.SileroVADConfig()) + result = realtime_vad.StreamingVADResult(True, True, True, True) + monkeypatch.setattr(protocol._server_vad, "process_base64", lambda *args, **kwargs: result) + protocol._input_speech_started = True + protocol._hold_realtime_output_until_session_created = False + await protocol._to_duplex_event( + {"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(1), "format": "pcm_f32le"} + ) + assert [event["type"] for event in ws.sent] == [ + "input_audio_buffer.speech_stopped", + "input_audio_buffer.speech_started", + ] + ws.sent.clear() + protocol._input_speech_started = False + result = realtime_vad.StreamingVADResult(True, False, True, True) + await protocol._to_duplex_event({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"}) + assert [event["type"] for event in ws.sent] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + ] + detector = realtime_vad.SileroStreamingVAD(realtime_vad.SileroVADConfig(), frame_scorer=lambda _: 0.0) + for _ in range(1024): + detector.process_base64(_pcm_f32_b64(1), fmt="pcm_f32le", sample_rate_hz=48_000) + assert detector._processed_samples + detector._pending.size == 341 + + +@pytest.mark.asyncio +async def test_realtime_server_vad_runtime_failures_are_terminal(monkeypatch): + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + ws = TimedWebSocket() + ws.put(_server_vad_update("generic-vad", model="test-model", session_id="generic-vad")) + await handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({})) + assert ws.sent[-1]["error"]["code"] == "server_vad_requires_native_duplex" + monkeypatch.setattr(realtime_vad.SileroStreamingVAD, "process_base64", lambda *args, **kwargs: 1 / 0) + ws = TimedWebSocket() + create = _native_realtime_session_update("vad-inference-failure") + create["session"]["turn_detection"] = {"type": "server_vad"} + ws.put(create) + ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(1), "format": "pcm_f32le"}) + await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({})), timeout=2) + assert any(e.get("error", {}).get("code") == "realtime_input_failed" for e in ws.sent) def test_native_duplex_handler_has_no_fixed_session_admission_cap(): @@ -988,6 +1094,178 @@ async def test_minicpmo_native_session_update_requires_ref_audio_before_enabling assert ("sid-runtime-update-ref-required", "session.update") not in engine.signals +def test_personaplex_candidate_update_does_not_leak_minicpmo_ref_audio_check(): + handler = OmniDuplexSessionHandler( + chat_service=SimpleNamespace(engine_client=SimpleNamespace()), + serving_runtime_adapter=PersonaPlexServingRuntimeAdapter(lambda *_: None), + ) + session = DuplexSession( + session_id="sid-personaplex-voice-change", + config=DuplexSessionConfig(modalities=["audio"], voice="NATF2.pt"), + capabilities=PersonaPlexServingRuntimeAdapter.capabilities(max_sessions=1), + ) + candidate_config = DuplexSessionConfig(modalities=["audio"], voice="OtherVoice.pt") + + error = handler._runtime_session_candidate_update_error(session, candidate_config) + + assert error is None + + +@pytest.mark.asyncio +async def test_minicpmo_native_session_update_rejects_instructions_after_native_context_locked(): + engine = FakeEngineClient() + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + ws = TimedWebSocket() + ws.put(_native_session_create("sid-instructions-locked", modalities=["text"])) + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(16000, value=0.05), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + "duration_ms": 1000, + "is_speech": True, + } + ) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": {"instructions": "You are now a pirate."}, + } + ) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + error = next(message for message in ws.sent if message.get("type") == "error") + assert error["code"] == "instructions_update_unsupported" + assert "session.updated" not in ws.sent_types() + assert ("sid-instructions-locked", "session.update") not in engine.signals + + +@pytest.mark.asyncio +async def test_minicpmo_native_session_update_allows_unchanged_instructions_after_lock(): + engine = FakeEngineClient() + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + ws = TimedWebSocket() + ws.put(_native_session_create("sid-instructions-unchanged", modalities=["text"])) + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(16000, value=0.05), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + "duration_ms": 1000, + "is_speech": True, + } + ) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": {"instructions": "You are a concise assistant.", "temperature": 0.3}, + } + ) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + assert "session.updated" in ws.sent_types() + assert not any( + message.get("type") == "error" and message.get("code") == "instructions_update_unsupported" + for message in ws.sent + ) + + +@pytest.mark.asyncio +async def test_minicpmo_native_session_update_rejects_instructions_before_first_real_append(): + engine = FakeEngineClient() + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + ws = TimedWebSocket() + ws.put(_native_session_create("sid-buffered-before-lock", modalities=["text"])) + ws.put( + { + "type": "input_audio_buffer.append", + "audio": _pcm_f32_b64(3200, value=0.05), + "format": "pcm_f32le", + "sample_rate_hz": 16000, + "duration_ms": 200, + "is_speech": True, + } + ) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": { + "instructions": ( + "You are now a swashbuckling pirate captain who speaks entirely in " + "nautical slang, references buried treasure constantly, and never " + "breaks character no matter what the user asks." + ) + }, + } + ) + ws.put({"type": "input_audio_buffer.commit", "final": True, "response_create": True}) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + error = next(message for message in ws.sent if message.get("type") == "error") + assert error["code"] == "instructions_update_unsupported" + assert "session.updated" not in ws.sent_types() + assert engine.appended, "the commit must still flush the buffered chunk despite the rejected update" + + +@pytest.mark.asyncio +async def test_minicpmo_native_session_update_rejects_native_duplex_mode_flip(): + engine = FakeEngineClient() + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(engine), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + ws = TimedWebSocket() + ws.put(_native_session_create("sid-native-duplex-flag-flip", modalities=["text"])) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": {"extra_body": {"minicpmo45_native_duplex": False}}, + } + ) + ws.put( + { + "type": "turn.signal", + "event": "session.update", + "payload": {"instructions": "You are now a pirate."}, + } + ) + ws.put({"type": "session.close"}) + + await handler.handle_session(ws) + + errors = [message for message in ws.sent if message.get("type") == "error"] + assert [error["code"] for error in errors] == [ + "native_duplex_mode_update_unsupported", + "instructions_update_unsupported", + ] + assert "session.updated" not in ws.sent_types() + + def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): ws = TimedWebSocket() protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] @@ -1325,6 +1603,14 @@ def _native_realtime_session_update( } +def _server_vad_update(event_id: str, **session: object) -> dict[str, object]: + return { + "type": "session.update", + "event_id": event_id, + "session": {"turn_detection": {"type": "server_vad"}, **session}, + } + + def test_native_stage0_request_id_matches_engine_after_same_session_id_reopen(): handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) first = handler._registry.create(session_id="sid-reused") @@ -1525,6 +1811,27 @@ def test_auto_response_playback_overlap_keeps_model_owned_listen_speak_decision( assert handler._should_force_listen_for_auto_response_overlap(session, {}, speech_payload) is False +def test_legacy_auto_overlap_policy_falls_back_to_listen_only(): + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + session = DuplexSession( + session_id="sid-auto-policy-fallback", + config=DuplexSessionConfig( + overlap_policy=DuplexSessionConfig._normalize_overlap_policy("auto"), + ), + ) + session.capabilities = DuplexCapabilities.minicpmo45_native() + + decision = handler._overlap_decision( + session, + {"duration_ms": session.config.overlap_barge_in_ms, "is_speech": True}, + _native_audio_payload(is_speech=True), + ) + + assert session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert decision["action"] == "listen" + assert decision["reason"] == "policy_listen_only" + + def test_auto_response_overlap_silence_advances_model_unit_and_preserves_realtime_input(): handler, session = _auto_response_context( "sid-auto-silent-wire-buffer", @@ -2922,7 +3229,7 @@ async def test_minicpmo_pre_response_continuation_drops_after_model_turn_ends(): @pytest.mark.asyncio -async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap(): +async def test_minicpmo_auto_response_continuation_stops_at_large_safety_boundary(): request_id = "duplex-sid-long-response-e0-stage0" engine = FakeEngineClient() handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) @@ -2938,7 +3245,7 @@ async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap(): _install_direct_silence_scheduler(handler, session) ws = TimedWebSocket() - for _ in range(9): + for _ in range(handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS): await handler._maybe_continue_native_response( ws.send_json, session=session, @@ -2946,9 +3253,160 @@ async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap(): ) await asyncio.sleep(0.02) - assert len(engine.appended) == 9 + response_id = session.active_response_id + assert response_id is not None + assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS assert all(payload["duplex_turn_id"] == 0 for _, _, payload, _ in engine.appended) - assert all("force_speak" not in payload for _, _, payload, _ in engine.appended) + assert all(not {"force_speak", "force_listen"} & payload.keys() for _, _, payload, _ in engine.appended[:-1]) + assert engine.appended[-1][2]["force_listen"] is True + assert handler._native_response_continuations_remaining(session, response_id) is False + + await handler._maybe_continue_native_response( + ws.send_json, + session=session, + expected_epoch=session.epoch, + ) + + assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS + assert ws.sent_types()[-2:] == ["response.listen", "response.done"] + assert session.active_response_id is None + assert session.active_request_id == request_id + assert session.turn_id == 1 + + +@pytest.mark.asyncio +async def test_minicpmo_auto_response_boundary_listen_closes_response(): + request_id = "duplex-sid-boundary-listen-e0-stage0" + engine = FakeEngineClient() + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) + session = DuplexSession( + session_id="sid-boundary-listen", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.capabilities = DuplexCapabilities.minicpmo45_native() + response_id = session.begin_response(turn_id=0) + session.bind_request(request_id) + native = handler._minicpmo_session_state(session) + native.continuation_owner_id = f"response:{response_id}" + native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS + _install_direct_silence_scheduler(handler, session) + ws = TimedWebSocket() + + close_reason, emitted = await handler._send_one_native_duplex_event( + ws.send_json, + { + "supported": True, + "stage_role": "llm", + "is_listen": True, + "model_listen": True, + "data_plane_request_id": request_id, + "end_of_turn": False, + "uses_model_runner_scheduler": True, + "runner_kv_backed": True, + "runtime_impl": "scheduler_data_plane", + }, + session=session, + expected_epoch=session.epoch, + ) + + assert close_reason is None + assert emitted is True + assert engine.appended == [] + assert ws.sent_types() == ["response.listen", "response.done"] + assert session.active_response_id is None + assert session.active_request_id == request_id + assert session.turn_id == 1 + assert native.continuation_owner_id is None + + +@pytest.mark.asyncio +async def test_minicpmo_auto_response_boundary_allows_next_model_turn_response(): + request_id = "duplex-sid-boundary-next-turn-e0-stage0" + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + session = DuplexSession( + session_id="sid-boundary-next-turn", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.capabilities = DuplexCapabilities.minicpmo45_native() + session.begin_response(turn_id=0) + session.bind_request(request_id) + native = handler._minicpmo_session_state(session) + native.continuation_owner_id = f"response:{session.active_response_id}" + native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS + ws = TimedWebSocket() + + await handler._maybe_continue_native_response( + ws.send_json, + session=session, + expected_epoch=session.epoch, + ) + + assert session.active_response_id is None + assert session.turn_id == 1 + + close_reason, emitted = await handler._send_one_native_duplex_event( + ws.send_json, + { + "supported": True, + "stage_role": "tts", + "is_listen": False, + "data_plane_request_id": request_id, + "text": "next response", + "audio_data": "next-audio", + "audio_format": "pcm16", + "audio_duration_ms": 100, + "end_of_turn": False, + "model_turn_id": 1, + }, + session=session, + expected_epoch=session.epoch, + ) + + assert close_reason is None + assert emitted is True + assert ws.sent_types() == [ + "response.listen", + "response.done", + "response.created", + "response.speak", + "response.output_audio.delta", + ] + assert session.active_response_turn_id == 1 + + +@pytest.mark.asyncio +async def test_minicpmo_auto_response_boundary_does_not_close_new_epoch_response(): + request_id = "duplex-sid-boundary-race-e0-stage0" + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + session = DuplexSession( + session_id="sid-boundary-race", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.capabilities = DuplexCapabilities.minicpmo45_native() + old_response_id = session.begin_response(turn_id=0) + session.bind_request(request_id) + native = handler._minicpmo_session_state(session) + native.continuation_owner_id = f"response:{old_response_id}" + native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS + sent: list[dict[str, Any]] = [] + new_response_ids: list[str] = [] + + async def _barge_in_on_listen(payload: dict[str, Any]) -> None: + sent.append(payload) + if payload["type"] == "response.listen": + session.barge_in() + session.bind_request("duplex-sid-boundary-race-e1-stage0") + new_response_ids.append(session.begin_response(turn_id=1)) + + await handler._maybe_continue_native_response( + _barge_in_on_listen, + session=session, + expected_epoch=session.epoch, + ) + + assert [payload["type"] for payload in sent] == ["response.listen"] + assert len(new_response_ids) == 1 + assert session.active_response_id == new_response_ids[0] @pytest.mark.asyncio @@ -3641,6 +4099,55 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): ] +@pytest.mark.asyncio +async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot(): + class SnapshotChatService(FakeChatService): + async def create_chat_completion(self, request, raw_request=None): + del raw_request + self.seen_request_ids.append(request.request_id) + snapshots = [ + {"num_tokens_out": 1, "vllm_tpot_ms": 0.0, "vllm_itls_ms": []}, + {"num_tokens_out": 2, "vllm_tpot_ms": 10.0, "vllm_itls_ms": [10.0]}, + {"num_tokens_out": 3, "vllm_tpot_ms": 11.0, "vllm_itls_ms": [10.0, 12.0]}, + ] + + async def generate(): + for index, snapshot in enumerate(snapshots): + payload = { + "modality": "audio", + "choices": [{"delta": {"content": f"audio-{index}"}}], + "metrics": {"stage_metrics": {"0": snapshot}}, + } + yield f"data: {json.dumps(payload)}\n\n" + + return generate() + + handler = OmniDuplexSessionHandler( + chat_service=SnapshotChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-chat-snapshots", + config=DuplexSessionConfig(instructions="speak"), + ) + sent: list[dict[str, Any]] = [] + + async def send_json(data: dict[str, Any]) -> None: + sent.append(data) + + await handler._run_response(session, send_json) + + audio_deltas = [event for event in sent if event.get("type") == "response.output_audio.delta"] + done = next(event for event in sent if event.get("type") == "response.done") + assert all("vllm_omni" not in event for event in audio_deltas) + assert done["vllm_omni"]["stage_metrics"]["0"] == { + "num_tokens_out": 3, + "vllm_tpot_ms": 11.0, + "vllm_itls_ms": [10.0, 12.0], + } + + @pytest.mark.asyncio async def test_duplex_chat_rejects_unknown_output_modality(): handler = OmniDuplexSessionHandler( @@ -3840,6 +4347,385 @@ async def test_playback_ack_response_id_selects_matching_pending_history_item(): assert second_item_id not in session.pending_history_item_ids +@pytest.mark.asyncio +async def test_playback_ack_recovers_missing_response_history_registration(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-playback-missing-registration", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + response_id = session.begin_response() + session.append_assistant_text("response text") + session.mark_audio_sent(1000, text_chars=len("response text")) + session.end_response(commit_text=True, preserve_request=True) + assert not session.history_item_ids + assert f"item_{response_id}" in session.pending_history_item_ids + + ws = TimedWebSocket() + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") + assert ack["history_committed"] is True + assert f"item_{response_id}" in session.history_item_ids + assert session.history[-1] == {"role": "assistant", "content": "response text"} + + +@pytest.mark.asyncio +async def test_playback_ack_rejects_response_after_followup_user_commit(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-playback-after-followup-user", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + session.append_text("user A") + session.commit_user_input() + response_id = session.begin_response() + session.append_assistant_text("assistant A") + session.mark_audio_sent(1000, text_chars=len("assistant A")) + session.end_response(commit_text=False, preserve_request=True) + session.append_text("user B") + session.commit_user_input() + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + assert ws.sent_types() == ["error"] + assert ws.sent[0]["code"] == "playback_ack_too_late" + assert session.history == ( + {"role": "user", "content": "user A"}, + {"role": "user", "content": "user B"}, + ) + + +@pytest.mark.asyncio +async def test_precommit_zero_audio_ack_reserves_history_for_response_that_finishes_after_commit(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-zero-audio-checkpoint", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + session.mark_user_input_activity() + response_id = session.begin_response() + item_id = f"item_{response_id}" + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": item_id, + "played_ms": 0, + "committed_ms": 0, + }, + ws.send_json, + ) + session.commit_native_audio_input(transcript="later user input") + session.append_assistant_text("hello world") + session.mark_audio_sent(1000, text_chars=len("hello world")) + committed_message = session.end_response(commit_text=True, preserve_request=True) + session.register_history_item(item_id, committed_message) + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": item_id, + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + assert not any(message.get("type") == "error" for message in ws.sent) + assert session.history == ( + {"role": "assistant", "content": "hello world"}, + { + "role": "user", + "content": [ + { + "type": "audio_url", + "audio_url": {"url": "native-duplex:input-audio"}, + "transcript": "later user input", + } + ], + "transcript": "later user input", + }, + ) + + +@pytest.mark.asyncio +async def test_precommit_partial_ack_extends_history_after_response_finishes(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-partial-checkpoint", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + session.mark_user_input_activity() + response_id = session.begin_response() + item_id = f"item_{response_id}" + session.append_assistant_text("hello world") + session.mark_audio_sent(1000, text_chars=len("hello world")) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": item_id, + "played_ms": 500, + "committed_ms": 500, + }, + ws.send_json, + ) + session.commit_native_audio_input(transcript="later user input") + committed_message = session.end_response(commit_text=True, preserve_request=True) + session.register_history_item(item_id, committed_message) + + assert session.history[0] == {"role": "assistant", "content": "hello"} + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": item_id, + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + assert not any(message.get("type") == "error" for message in ws.sent) + assert session.history[0] == {"role": "assistant", "content": "hello world"} + assert session.history[1]["role"] == "user" + + +@pytest.mark.asyncio +async def test_explicit_truncate_seals_snapshot_against_later_playback_ack(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-hard-playback-truncate", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + response_id = session.begin_response() + item_id = f"item_{response_id}" + session.append_assistant_text("hello world") + session.mark_audio_sent(1000, text_chars=len("hello world")) + session.end_response(commit_text=False, preserve_request=True) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": item_id, + "played_ms": 500, + "committed_ms": 500, + "truncate": True, + }, + ws.send_json, + ) + assert session.history == ({"role": "assistant", "content": "hello"},) + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": item_id, + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + assert not any(message.get("type") == "error" for message in ws.sent) + assert session.history == ({"role": "assistant", "content": "hello"},) + + +@pytest.mark.asyncio +async def test_cancel_does_not_append_old_assistant_after_followup_user(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-cancel-after-followup-user", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + session.append_text("user A") + session.commit_user_input() + session.begin_response() + session.append_assistant_text("assistant A") + session.mark_audio_sent(1000, text_chars=len("assistant A")) + session.acknowledge_playback(500, 500) + session.append_text("user B") + session.commit_user_input() + ws = TimedWebSocket() + + cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="new_response") + + assert cancelled is True + assert session.history == ( + {"role": "user", "content": "user A"}, + {"role": "user", "content": "user B"}, + ) + + +@pytest.mark.asyncio +async def test_playback_ack_rejects_item_from_another_response(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-playback-mismatch", config=DuplexSessionConfig()) + first_response_id = session.begin_response() + session.append_assistant_text("first response") + session.mark_audio_sent(1000) + session.end_response(commit_text=False, preserve_request=True) + session.register_history_item(f"item_{first_response_id}", None) + second_response_id = session.begin_response() + session.mark_audio_sent(1000) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": second_response_id, + "item_id": f"item_{first_response_id}", + "played_ms": 500, + "committed_ms": 500, + }, + ws.send_json, + ) + + assert ws.sent_types() == ["error"] + assert ws.sent[0]["code"] == "playback_item_mismatch" + assert session.active_response_id == second_response_id + assert session.playback.as_dict() == { + "generated_ms": 1000, + "sent_ms": 1000, + "played_ms": 0, + "committed_ms": 0, + } + + +@pytest.mark.asyncio +async def test_playback_ack_rejects_user_history_item(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-playback-user-item", config=DuplexSessionConfig()) + user_message = {"role": "user", "content": "user input"} + session.append_history_message(user_message) + session.register_history_item("item_user", user_message) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": "user", + "item_id": "item_user", + "played_ms": 500, + "committed_ms": 500, + }, + ws.send_json, + ) + + assert ws.sent_types() == ["error"] + assert ws.sent[0]["code"] == "playback_item_not_found" + assert session.history == ({"role": "user", "content": "user input"},) + + +@pytest.mark.asyncio +async def test_late_playback_ack_recovers_response_local_history_after_next_response_begins(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-playback-response-local", config=DuplexSessionConfig()) + first_response_id = session.begin_response() + session.append_assistant_text("first response") + session.mark_audio_sent(1000, text_chars=len("first response")) + session.end_response(commit_text=False, preserve_request=True) + second_response_id = session.begin_response() + session.mark_audio_sent(2000, text_chars=len("second response")) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": first_response_id, + "item_id": f"item_{first_response_id}", + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") + assert ack["history_committed"] is True + assert session.history[-1] == {"role": "assistant", "content": "first response"} + assert session.active_response_id == second_response_id + assert session.playback.as_dict() == { + "generated_ms": 2000, + "sent_ms": 2000, + "played_ms": 0, + "committed_ms": 0, + } + + @pytest.mark.asyncio async def test_late_playback_ack_does_not_advance_new_response_cursor(): handler = OmniDuplexSessionHandler( @@ -4194,6 +5080,7 @@ async def test_realtime_resume_preserves_append_tail_order_across_connections(): await asyncio.sleep(0.01) else: pytest.fail("resumable session did not open") + assert created is not None session = handler._registry.get("sid-resume-append-order") assert session is not None native_state = handler._minicpmo_session_state(session) @@ -4323,6 +5210,7 @@ async def test_realtime_takeover_does_not_clear_new_attachment_pending_turns(): await asyncio.sleep(0.001) else: pytest.fail("first attachment did not create the session") + assert created is not None second = TimedWebSocket(receive_timeout_s=1.0) second.put( @@ -4386,6 +5274,7 @@ async def test_realtime_takeover_drops_input_queued_by_replaced_attachment(): await asyncio.sleep(0.001) else: pytest.fail("first attachment did not create the session") + assert created is not None second = TimedWebSocket(receive_timeout_s=1.0) second.put( @@ -4988,7 +5877,7 @@ async def test_audio_clear_preserves_later_wire_order_append(): @pytest.mark.asyncio -async def test_native_append_is_cancelled_by_later_wire_order_input_cancel(): +async def test_native_append_is_cancelled_by_later_wire_order_input_cancel(mocker): class SlowAppendEngine(FakeEngineClient): def __init__(self) -> None: super().__init__() @@ -5002,11 +5891,12 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel(): *, mode: str, payload: object, + operation_id: str | None = None, final: bool = False, timeout: float | None = None, collect_outputs: bool = True, fence: DuplexFence | None = None, - ): + ) -> Any: self.append_started.set() try: await self.release_append.wait() @@ -5017,12 +5907,14 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel(): session_id, mode=mode, payload=payload, + operation_id=operation_id, final=final, timeout=timeout, collect_outputs=collect_outputs, fence=fence, ) + rollback = mocker.spy(MiniCPMO45PcmAppendBuffer, "_rollback_reservation") engine = SlowAppendEngine() handler = OmniDuplexSessionHandler( chat_service=FakeChatService(engine), @@ -5053,6 +5945,7 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel(): await asyncio.wait_for(handler_task, timeout=1) assert engine.append_cancelled.is_set() + rollback.assert_called_once() assert engine.signal_fences[-1] == DuplexFence("sid-cancel-slow-append") assert engine.signal_next_fences[-1] == DuplexFence("sid-cancel-slow-append", epoch=1) @@ -5128,7 +6021,9 @@ async def test_duplex_handler_session_update_control_failure_rolls_back_config() error = next(m for m in ws.sent if m.get("type") == "error") assert error["code"] == "runtime_signal_failed" assert error["runtime_control"]["error_count"] == 1 - assert engine.signal_session_configs[-1]["instructions"] == "must roll back" + signaled_config = engine.signal_session_configs[-1] + assert signaled_config is not None + assert signaled_config["instructions"] == "must roll back" @pytest.mark.asyncio @@ -5146,33 +6041,59 @@ async def test_native_session_update_commits_config_only_after_runtime_ack(): return await super().signal_duplex_turn_async(session_id, **kwargs) engine = BlockingUpdateEngine() - handler = OmniDuplexSessionHandler( - chat_service=FakeChatService(engine), - config_timeout_s=0.1, - idle_timeout_s=1, - ) + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine)) ws = TimedWebSocket() - ws.put(_session_create("sid-update-two-phase")) - ws.put( - { - "type": "turn.signal", - "event": "session.update", - "payload": {"instructions": "candidate instructions"}, - } - ) + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + ws.put(_native_realtime_session_update("sid-update-two-phase")) + ws.put(_server_vad_update("update-two-phase")) - handler_task = asyncio.create_task(handler.handle_session(ws)) + handler_task = asyncio.create_task(handler.handle_session(ws, realtime_protocol=protocol)) await asyncio.wait_for(engine.update_started.wait(), timeout=1) live_session = handler._registry.get("sid-update-two-phase") assert live_session is not None - assert live_session.config.instructions != "candidate instructions" + assert live_session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert protocol._pending_turn_detection_update is not None + assert protocol._server_vad is None + + protocol.encode_outbound_event({"type": "error", "code": "unrelated", "error": "unrelated"}) + assert protocol._pending_turn_detection_update is not None engine.release_update.set() ws.put({"type": "session.close"}) await asyncio.wait_for(handler_task, timeout=2) - updated = next(message for message in ws.sent if message.get("type") == "session.updated") - assert updated["session"]["instructions"] == "candidate instructions" + updated = next(message for message in reversed(ws.sent) if message.get("type") == "session.updated") + assert updated["session"]["overlap_policy"] == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + assert isinstance(protocol._server_vad, realtime_vad.SileroStreamingVAD) + + +@pytest.mark.parametrize( + ("append_fails", "expected_code"), + [(True, "session_update_aborted"), (False, "instructions_update_unsupported")], +) +@pytest.mark.asyncio +async def test_realtime_vad_update_rejection_does_not_stall(append_fails, expected_code): + append_result = {"operation": "append", "ok": False, "error_count": 1} if append_fails else None + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient(append_result=append_result))) + ws = TimedWebSocket() + protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] + create = _native_realtime_session_update("sid-rejected-vad-update") + create["session"]["extra_body"]["auto_response"] = append_fails + ws.put(create) + if append_fails: + ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(16_000), "format": "pcm_f32le"}) + update = _server_vad_update("rejected-vad-update") + if not append_fails: + update["session"]["instructions"] = "You are now a pirate." + ws.put(update) + ws.put({"type": "session.close"}) + await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=protocol), timeout=2) + + errors = [message["error"] for message in ws.sent if isinstance(message.get("error"), dict)] + error = next(error for error in errors if error.get("code") == expected_code) + assert error["event_id"] == "rejected-vad-update" + assert protocol._pending_turn_detection_update is None + assert protocol._server_vad is None @pytest.mark.asyncio @@ -5296,7 +6217,7 @@ async def test_failed_native_append_prevents_queued_append_from_running(): @pytest.mark.asyncio -async def test_failed_final_append_keeps_committed_audio_for_response_retry(): +async def test_failed_final_append_discards_retained_audio_without_retry(): failed_result = { "operation": "append", "session_id": "sid-final-append-retry", @@ -5305,15 +6226,7 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry(): "error_count": 1, "stage_results": [], } - successful_result = { - "operation": "append", - "session_id": "sid-final-append-retry", - "ok": True, - "unsupported_count": 0, - "error_count": 0, - "stage_results": [], - } - engine = FakeEngineClient(append_results=[failed_result, successful_result]) + engine = FakeEngineClient(append_result=failed_result) handler = OmniDuplexSessionHandler( chat_service=FakeChatService(engine), config_timeout_s=0.1, @@ -5338,15 +6251,18 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry(): handler_task = asyncio.create_task(handler.handle_session(ws)) for _ in range(100): - if len(engine.appended) >= 2: + state = handler._serving_runtime_adapter.session_states.get("sid-final-append-retry") + if len(engine.appended) >= 1 and state is not None and state.committed_audio_payload is None: break await asyncio.sleep(0.01) + state = handler._serving_runtime_adapter.session_states["sid-final-append-retry"] + assert state.committed_audio_payload is None + assert state.committed_audio_reserved_bytes == 0 + assert handler._registry.get("sid-final-append-retry").pending_input_bytes == 0 ws.put({"type": "session.close"}) await asyncio.wait_for(handler_task, timeout=2) - assert len(engine.appended) == 2 - assert engine.appended[0][2]["audio"] == engine.appended[1][2]["audio"] - assert engine.append_operation_ids[0] == engine.append_operation_ids[1] + assert len(engine.appended) == 1 @pytest.mark.asyncio @@ -5429,7 +6345,7 @@ async def test_minicpmo_native_duplex_open_unsupported_fails_session_create(): @pytest.mark.asyncio async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract(): class LegacyEngineClient(FakeEngineClient): - async def open_duplex_session_async( + async def open_duplex_session_async( # type: ignore[override] self, session_id: str, *, @@ -5437,7 +6353,7 @@ async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract(): capabilities: dict[str, object] | None = None, session_config: dict[str, object] | None = None, timeout: float | None = None, - ) -> None: + ) -> Any: return await super().open_duplex_session_async( session_id, session_mode=session_mode, @@ -5534,6 +6450,9 @@ async def test_minicpmo_native_duplex_separates_public_and_runtime_config(monkey assert opened_runtime_config["duplex_stage_max_tokens"] == {"0": 20, "1": 8192} assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_k"] == 20 assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_p"] == 0.8 + # The Talker YAML carries upstream's min_new_token=50 for chat TTS; a duplex + # chunk is 26 codec samples, so check_stop would never release the request. + assert opened_runtime_config["duplex_stage_sampling_params"]["1"] == {"min_tokens": 0} assert base64.b64decode(opened_runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.25] * 1600)) @@ -5674,7 +6593,7 @@ async def test_minicpmo_native_duplex_rejects_ref_audio_path(): assert engine.opened == [] -def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen(): +def test_minicpmo_native_duplex_explicit_barge_in_request_interrupts(): engine = FakeEngineClient() handler = OmniDuplexSessionHandler( chat_service=FakeChatService(engine), @@ -5701,12 +6620,10 @@ def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen( ) assert decision == { - "action": "listen", - "reason": "barge_in_unsupported", + "action": "barge_in", + "reason": "client_force_barge_in", "duration_ms": 1000, - "overlap_speech_ms": 1000, "buffer_audio": True, - "defer_runtime_append": True, } @@ -5718,7 +6635,7 @@ def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen( {"type": "turn.signal", "event": "barge_in"}, ], ) -async def test_minicpmo_native_duplex_rejects_unadvertised_barge_in_control( +async def test_minicpmo_native_duplex_accepts_explicit_barge_in_control( barge_in_event: dict[str, object], ): engine = FakeEngineClient() @@ -5734,9 +6651,8 @@ async def test_minicpmo_native_duplex_rejects_unadvertised_barge_in_control( await handler.handle_session(ws) - error = next(message for message in ws.sent if message.get("code") == "barge_in_unsupported") - assert error["session_id"] == "sid-native-barge-control-disabled" - assert ("sid-native-barge-control-disabled", "barge_in") not in engine.signals + assert not any(message.get("code") == "barge_in_unsupported" for message in ws.sent) + assert ("sid-native-barge-control-disabled", "barge_in") in engine.signals @pytest.mark.asyncio diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 2136d84f8..b2eb2edfc 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -117,10 +117,23 @@ def test_encode_image_base64(): class MockGenerationResult: """Mock result object compatible with current diffusion output shape.""" - def __init__(self, images): + def __init__(self, images, stage_durations=None, peak_memory_mb=0.0): self.images = images - self.stage_durations = {} - self.peak_memory_mb = 0.0 + self.stage_durations = {} if stage_durations is None else stage_durations + self.peak_memory_mb = peak_memory_mb + + +def _stub_engine_generate_with_metrics(engine, *, stage_durations, peak_memory_mb): + """Stub engine.generate to return profiler metrics.""" + + async def generate(*args, **kwargs): + yield MockGenerationResult( + [Image.new("RGB", (16, 16), color="blue")], + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + engine.generate = generate class MockStageResult: @@ -972,6 +985,53 @@ def test_image_edits_streaming_rejects_single_stage_before_loading_url(test_clie assert "multi-stage" in response.json()["detail"] +def test_generate_images_returns_metrics_single_stage(test_client): + """Single-stage /v1/images/generations copies getattr metrics onto the response.""" + stage_durations = {"queue_wait_ms": 1.0, "stage_0_gen_ms": 2.0} + peak_memory_mb = 1024.0 + _stub_engine_generate_with_metrics( + test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = test_client.post( + "/v1/images/generations", + json={"prompt": "a cat", "n": 1, "size": "256x256"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + +def test_generate_images_returns_metrics_multistage(async_omni_test_client): + """Multi-stage /v1/images/generations unpacks generate_diffusion_images metrics.""" + stage_durations = { + "queue_wait_ms": 1.0, + "preprocess_ms": 2.0, + "stage_0_gen_ms": 3.0, + "stage_1_gen_ms": 4.0, + } + peak_memory_mb = 2048.0 + _stub_engine_generate_with_metrics( + async_omni_test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = async_omni_test_client.post( + "/v1/images/generations", + json={"prompt": "a cat", "n": 1, "size": "256x256"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + def test_generate_images_max_size_rejected(async_omni_test_client): """Test that a size exceeding max_generated_image_size returns 400.""" response = async_omni_test_client.post( @@ -1987,6 +2047,56 @@ def test_image_edit_with_seed_zero(async_omni_test_client): ) +def test_image_edits_returns_metrics_single_stage(test_client): + """Single-stage /v1/images/edits copies getattr metrics onto the response.""" + stage_durations = {"queue_wait_ms": 1.0, "stage_0_gen_ms": 2.0} + peak_memory_mb = 1024.0 + _stub_engine_generate_with_metrics( + test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = test_client.post( + "/v1/images/edits", + files=[("image", make_test_image_bytes((16, 16)))], + data={"prompt": "edit me"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + +def test_image_edits_returns_metrics_multistage(async_omni_test_client): + """Multi-stage /v1/images/edits unpacks generate_diffusion_images metrics.""" + stage_durations = { + "queue_wait_ms": 1.0, + "preprocess_ms": 2.0, + "ar2diffusion_ms": 3.0, + "stage_0_gen_ms": 4.0, + "stage_1_gen_ms": 5.0, + } + peak_memory_mb = 2048.0 + _stub_engine_generate_with_metrics( + async_omni_test_client.app.state.engine_client, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + ) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[("image", make_test_image_bytes((16, 16)))], + data={"prompt": "edit me"}, + ) + assert response.status_code == 200 + assert response.json()["metrics"] == { + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + } + + def test_image_edit_with_seed_zero_single_stage(test_client): """Test that seed=0 is correctly handled in image editing (single stage). diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 325daabdd..7664b2402 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out). @@ -38,6 +38,10 @@ MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" REALTIME_SYNTH_PHRASE_TEXT = ( "Translate into Chinese: Beijing is the Capital of China. It is the center of culture and politics" ) +ISSUE_6474_SYNTH_PHRASE_TEXT = ( + "Can you tell me the current temperature and weather conditions in New York City? " + "What about Los Angeles? Please compare them in detail using at least eight complete sentences." +) # Simulate realtime upload pacing (``openai_realtime_client.py --send-delay-ms``). SEND_DELAY_MS = 200 @@ -100,6 +104,7 @@ async def _run_realtime_audio_roundtrip( *, chunk_ms: int = 100, send_delay_ms: int = 0, + completion_timeout_s: float = 600, ) -> dict: uri = f"ws://{host}:{port}/v1/realtime" incremental: list[bytes] = [] @@ -131,7 +136,7 @@ async def _run_realtime_audio_roundtrip( await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) while True: - message = await asyncio.wait_for(ws.recv(), timeout=600) + message = await asyncio.wait_for(ws.recv(), timeout=completion_timeout_s) if isinstance(message, bytes): continue @@ -178,12 +183,16 @@ async def _run_realtime_audio_roundtrip( } -def _synthetic_pcm16_input() -> bytes: +def _synthetic_pcm16_input( + *, + phrase_text: str = REALTIME_SYNTH_PHRASE_TEXT, + duration_s: int = 10, +) -> bytes: syn = generate_synthetic_audio( - 10, + duration_s, 1, sample_rate=16000, - phrase_text=REALTIME_SYNTH_PHRASE_TEXT, + phrase_text=phrase_text, ) wav_bytes = base64.b64decode(syn["base64"]) return _pcm16_mono_16k_from_wav_bytes(wav_bytes) @@ -260,6 +269,33 @@ class TestQwen3OmniRealtimeWebSocket: _assert_realtime_smoke(result) _assert_realtime_accuracy(result) + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True) + def test_long_audio_response_completes_async_chunk(self, omni_server) -> None: + """Regression for #6474: long async-chunk audio must emit response.audio.done.""" + issue_6474_pcm16 = _synthetic_pcm16_input( + phrase_text=ISSUE_6474_SYNTH_PHRASE_TEXT, + ) + issue_6474_result = asyncio.run( + _run_realtime_audio_roundtrip( + omni_server.host, + omni_server.port, + omni_server.model, + issue_6474_pcm16, + chunk_ms=100, + send_delay_ms=SEND_DELAY_MS, + completion_timeout_s=180, + ) + ) + + _assert_realtime_smoke(issue_6474_result) + output_duration_s = len(issue_6474_result["output_pcm"]) / (2 * issue_6474_result["output_sample_rate"]) + assert output_duration_s > 5, ( + f"Expected an issue-like audio response longer than 5 seconds, got {output_duration_s:.2f}s" + ) + @pytest.mark.advanced_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py index 0647c40a3..2fdf9f646 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_chat_metrics.py @@ -2,9 +2,18 @@ """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field.""" from types import SimpleNamespace +from unittest.mock import MagicMock import pytest +from tests.helpers.serving_chat import ( + build_serving_chat, + collect_stream, + make_request, + make_text_omni_output, + parse_sse_chunks, +) + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -54,6 +63,65 @@ def test_omni_chat_completion_stream_response_metrics(): assert response.metrics == {"stage_latency": 0.5} +@pytest.mark.asyncio +async def test_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens(): + serving_chat = build_serving_chat() + serving_chat.enable_prompt_tokens_details = True + request = make_request(modalities=["text"], include_usage=True) + + async def result_generator(): + yield make_text_omni_output( + token_ids=[10], + finish_reason="stop", + num_cached_tokens=0, + ) + + raw_lines = await collect_stream( + serving_chat.chat_completion_stream_generator( + request=request, + result_generator=result_generator(), + request_id="test-req", + model_name="test-model", + conversation=[], + tokenizer=MagicMock(), + request_metadata=MagicMock(), + mm_token_counts={"image": 7}, + ) + ) + + chunks = parse_sse_chunks(raw_lines) + usage = next(chunk["usage"] for chunk in chunks if chunk.get("usage")) + + assert usage["prompt_tokens_details"] == { + "cached_tokens": 0, + "multimodal_tokens": {"image": 7}, + } + + +def test_non_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens(): + serving_chat = build_serving_chat() + serving_chat.enable_prompt_tokens_details = True + request = make_request(modalities=["text"], stream=False) + omni_output = make_text_omni_output( + token_ids=[10], + finish_reason="stop", + num_cached_tokens=0, + ) + + _, usage, _, _, _ = serving_chat._create_text_choice( + request=request, + omni_outputs=omni_output, + tokenizer=MagicMock(), + conversation=[], + role="assistant", + mm_token_counts={"audio": 5}, + ) + + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 0 + assert usage.prompt_tokens_details.multimodal_tokens == {"audio": 5} + + def test_create_image_choice_exposes_diffusion_metrics(): """Ensure image chat content exposes profiler metrics for clients.""" from PIL import Image diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py new file mode 100644 index 000000000..59d070a39 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression tests for MiniCPM-o 4.5 chat reference audio transport.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import cast + +import numpy as np +import pytest +import torch +from pytest_mock import MockerFixture +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.inputs import TokensPrompt + +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY +from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import ( + llm2tts, + tts2code2wav_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@dataclass +class _EngineClientStub: + stage_configs: list[dict[str, object]] + + +@dataclass +class _ModelConfigStub: + allowed_local_media_path: str = "" + allowed_media_domains: list[str] | None = None + + +@dataclass +class _ThinkerCompletionStub: + token_ids: list[int] + text: str + hidden_states: torch.Tensor + multimodal_output: dict[str, object] = field(default_factory=dict) + + +@dataclass +class _ThinkerRequestOutputStub: + request_id: str + prompt_token_ids: list[int] + outputs: list[_ThinkerCompletionStub] + multimodal_output: dict[str, object] | None = None + + +@dataclass +class _Code2WavRequestStub: + external_req_id: str + request_id: str + model_intermediate_buffer: dict[str, object] + + +@dataclass +class _TransferManagerStub: + request_payload: dict[str, object] = field(default_factory=dict) + + +@pytest.fixture +def serving_chat() -> OmniOpenAIServingChat: + return object.__new__(OmniOpenAIServingChat) + + +def _stage_config(model_arch: str) -> dict[str, object]: + return {"engine_args": {"model_arch": model_arch}} + + +def _chat_request(reference_audio_url: str) -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "target text"}], + ref_audio=reference_audio_url, + ) + + +@pytest.mark.asyncio +async def test_chat_reference_audio_reaches_code2wav_transport( + serving_chat: OmniOpenAIServingChat, + mocker: MockerFixture, +) -> None: + reference_audio_url = "data:audio/wav;base64,AAAA" + reference_waveform = np.array([0.25, -0.5], dtype=np.float32) + serving_chat.engine_client = _EngineClientStub( + [_stage_config("MiniCPMO45OmniForConditionalGeneration")], + ) + serving_chat.model_config = _ModelConfigStub() + engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]} + media_connector_class = mocker.patch( + "vllm_omni.entrypoints.openai.serving_chat.MediaConnector", + ) + media_connector = media_connector_class.return_value + media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000)) + + await serving_chat._attach_minicpmo45_reference_audio( + engine_prompt, + _chat_request(reference_audio_url), + ) + + media_connector_class.assert_called_once_with( + media_io_kwargs=None, + allowed_local_media_path="", + allowed_media_domains=None, + ) + media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url) + preserved_reference = cast(dict[str, object], engine_prompt)[MINICPMO45_REFERENCE_AUDIO_KEY] + assert isinstance(preserved_reference, tuple) + preserved_waveform, preserved_sample_rate = preserved_reference + assert isinstance(preserved_waveform, np.ndarray) + assert isinstance(preserved_sample_rate, int) + np.testing.assert_array_equal(preserved_waveform, reference_waveform) + assert preserved_sample_rate == 16000 + + thinker_output = _ThinkerRequestOutputStub( + request_id="request-0", + prompt_token_ids=[10], + outputs=[ + _ThinkerCompletionStub( + token_ids=[20], + text="target text", + hidden_states=torch.zeros((2, 4)), + ) + ], + ) + talker_prompt = llm2tts([thinker_output], prompt=engine_prompt)[0] + talker_request = _Code2WavRequestStub( + external_req_id="request-0", + request_id="request-0", + model_intermediate_buffer=talker_prompt["model_intermediate_buffer"], + ) + transfer_manager = _TransferManagerStub() + + code2wav_payload = tts2code2wav_async_chunk( + transfer_manager, + {"codes": {"audio": torch.arange(7).reshape(-1, 1)}}, + talker_request, + True, + ) + + assert code2wav_payload is not None + torch.testing.assert_close(code2wav_payload.codes.ref, torch.from_numpy(reference_waveform)) + assert code2wav_payload.meta.ref_audio_sr == 16000 + + +@pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"]) +@pytest.mark.asyncio +async def test_reference_audio_accepts_request_extra_mappings( + serving_chat: OmniOpenAIServingChat, + mocker: MockerFixture, + request_shape: str, +) -> None: + reference_audio_url = "data:audio/wav;base64,AAAA" + reference_waveform = np.array([0.25, -0.5], dtype=np.float32) + serving_chat.engine_client = _EngineClientStub( + [_stage_config("MiniCPMO45OmniForConditionalGeneration")], + ) + serving_chat.model_config = _ModelConfigStub() + engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]} + if request_shape == "extra_body": + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "target text"}], + extra_body={"ref_audio": reference_audio_url}, + ) + else: + request = SimpleNamespace(model_extra={"ref_audio": reference_audio_url}) + assert not hasattr(request, "ref_audio") + media_connector = mocker.patch( + "vllm_omni.entrypoints.openai.serving_chat.MediaConnector", + ).return_value + media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000)) + + await serving_chat._attach_minicpmo45_reference_audio(engine_prompt, request) + + media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url) + preserved_reference = cast(dict[str, object], engine_prompt)[MINICPMO45_REFERENCE_AUDIO_KEY] + assert isinstance(preserved_reference, tuple) + np.testing.assert_array_equal(preserved_reference[0], reference_waveform) + assert preserved_reference[1] == 16000 + + +@pytest.mark.asyncio +async def test_reference_audio_is_ignored_for_other_models( + serving_chat: OmniOpenAIServingChat, + mocker: MockerFixture, +) -> None: + reference_audio_url = "data:audio/wav;base64,AAAA" + serving_chat.engine_client = _EngineClientStub([_stage_config("SomeOtherOmniForConditionalGeneration")]) + engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]} + media_connector_class = mocker.patch( + "vllm_omni.entrypoints.openai.serving_chat.MediaConnector", + ) + + await serving_chat._attach_minicpmo45_reference_audio( + engine_prompt, + _chat_request(reference_audio_url), + ) + + media_connector_class.assert_not_called() + assert MINICPMO45_REFERENCE_AUDIO_KEY not in cast(dict[str, object], engine_prompt) diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index e7f014a3b..f81e61cf1 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -14,6 +14,8 @@ from pytest_mock import MockerFixture from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.sampling_params import SamplingParams +from tests.helpers.serving_chat import build_serving_chat + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -77,18 +79,8 @@ def mock_engine_client( return engine_client -@pytest.fixture -def serving_chat(mock_engine_client): - """Create OmniOpenAIServingChat instance with mocked dependencies.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - # Create instance without calling __init__ - instance = object.__new__(OmniOpenAIServingChat) - instance.engine_client = mock_engine_client - return instance - - -def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_chat): +def test_serving_boundary_normalizes_declared_root_and_nested_extras(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset({"cfg_text_scale", "negative_prompt"}) request = ChatCompletionRequest( model="test", @@ -118,7 +110,8 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha assert diffusion_request_args["modalities"] == ["image"] -def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat): +def test_unknown_root_extra_does_not_claim_canonical_extra(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset() request = ChatCompletionRequest( model="test", @@ -135,7 +128,8 @@ def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat): assert "pipeline_option" not in diffusion_request_args -def test_serving_boundary_rejects_invalid_quality(serving_chat): +def test_serving_boundary_rejects_invalid_quality(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset() request = ChatCompletionRequest( model="test", @@ -148,7 +142,8 @@ def test_serving_boundary_rejects_invalid_quality(serving_chat): serving_chat._normalize_diffusion_request_args(request) -def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat): +def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(mock_engine_client): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset() request = ChatCompletionRequest( model="test", @@ -170,7 +165,8 @@ def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat): ], ids=["registered-model-extra", "unregistered-common-alias"], ) -def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, request_kwargs): +def test_cfg_scale_owner_conflicts_are_rejected(mock_engine_client, registered, request_kwargs): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_extra_body_params = frozenset(registered) request = ChatCompletionRequest(model="test", messages=[], **request_kwargs) @@ -180,10 +176,11 @@ def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, reques @pytest.mark.parametrize("diffusion_mode", [True, False], ids=["pure", "mixed"]) def test_duplicate_extras_return_the_same_bad_request_before_dispatch( - serving_chat, + mock_engine_client, mocker: MockerFixture, diffusion_mode: bool, ): + serving_chat = build_serving_chat(engine_client=mock_engine_client) serving_chat._diffusion_mode = diffusion_mode serving_chat._diffusion_extra_body_params = frozenset({"sample_solver"}) serving_chat.engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] @@ -211,9 +208,10 @@ def test_duplicate_extras_return_the_same_bad_request_before_dispatch( pure_dispatch.assert_not_awaited() -def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, mocker: MockerFixture): +def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(mock_engine_client, mocker: MockerFixture): from vllm_omni.inputs.data import OmniDiffusionSamplingParams + serving_chat = build_serving_chat(engine_client=mock_engine_client) captured: dict[str, object] = {} async def generate(**kwargs): @@ -261,19 +259,20 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, assert response.error.message == "No output generated from AsyncOmni" -def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, mocker: MockerFixture): +def test_mixed_consumer_keeps_root_common_args_with_nested_extras(mock_engine_client, mocker: MockerFixture): from vllm_omni.inputs.data import OmniDiffusionSamplingParams - serving_chat.engine_client.stage_configs = [ + mock_engine_client.stage_configs = [ SimpleNamespace(stage_type="llm", is_comprehension=True), SimpleNamespace(stage_type="diffusion", is_comprehension=False), ] - serving_chat.engine_client.default_sampling_params_list = [ + mock_engine_client.default_sampling_params_list = [ SamplingParams(), OmniDiffusionSamplingParams(), ] - serving_chat.engine_client.output_modalities = ["image"] - serving_chat.engine_client.errored = False + mock_engine_client.output_modalities = ["image"] + mock_engine_client.errored = False + mock_engine_client.renderer = SimpleNamespace(get_tokenizer=lambda: object()) captured: dict[str, object] = {} @@ -285,21 +284,15 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, captured.update(kwargs) return results() - serving_chat.engine_client.generate = generate - serving_chat.__dict__.update( - _diffusion_mode=False, - _diffusion_extra_body_params=frozenset(), + mock_engine_client.generate = generate + serving_chat = build_serving_chat( + engine_client=mock_engine_client, models=SimpleNamespace(model_name=lambda _: "test"), - renderer=SimpleNamespace(get_tokenizer=lambda: object()), online_renderer=SimpleNamespace(validate_chat_template=lambda **_: None), - parser_cls=None, - use_harmony=False, - enable_auto_tools=False, - exclude_tools_when_tool_choice_none=False, trust_request_chat_template=True, - chat_template=None, - chat_template_content_format="auto", ) + serving_chat._diffusion_mode = False + serving_chat._diffusion_extra_body_params = frozenset() mocker.patch.multiple( serving_chat, _check_model=mocker.AsyncMock(return_value=None), @@ -386,8 +379,9 @@ def test_openai_sampling_fields_contains_expected_fields(): # ============================================================================= -def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_request): +def test_preserves_yaml_defaults_when_no_request_params(mock_engine_client, mock_request): """Test that YAML defaults are preserved when request has no params.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._build_sampling_params_list_from_request(mock_request) assert len(result) == 2 @@ -400,8 +394,9 @@ def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_reque assert comprehension_params.repetition_penalty == 1.05 # YAML custom param preserved -def test_request_temperature_overrides_yaml_default(serving_chat, mock_request): +def test_request_temperature_overrides_yaml_default(mock_engine_client, mock_request): """Test that request temperature overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.temperature = 0.8 mock_request.model_fields_set = {"temperature"} @@ -413,8 +408,9 @@ def test_request_temperature_overrides_yaml_default(serving_chat, mock_request): assert comprehension_params.top_k == 1 # YAML custom param preserved -def test_request_top_p_overrides_yaml_default(serving_chat, mock_request): +def test_request_top_p_overrides_yaml_default(mock_engine_client, mock_request): """Test that request top_p overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.top_p = 0.95 mock_request.model_fields_set = {"top_p"} @@ -425,8 +421,9 @@ def test_request_top_p_overrides_yaml_default(serving_chat, mock_request): assert comprehension_params.temperature == 0.4 # Preserved from YAML -def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request): +def test_request_max_tokens_overrides_yaml_default(mock_engine_client, mock_request): """Test that request max_tokens overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.max_tokens = 100 mock_request.model_fields_set = {"max_tokens"} @@ -435,15 +432,17 @@ def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request): assert result[0].max_tokens == 100 -def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_request): +def test_max_tokens_uses_yaml_default_when_not_specified(mock_engine_client, mock_request): """Test that max_tokens falls back to YAML default when not in request.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._build_sampling_params_list_from_request(mock_request) assert result[0].max_tokens == 4353 -def test_request_seed_overrides_yaml_default(serving_chat, mock_request): +def test_request_seed_overrides_yaml_default(mock_engine_client, mock_request): """Test that request seed overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.seed = 123 mock_request.model_fields_set = {"seed"} @@ -454,8 +453,9 @@ def test_request_seed_overrides_yaml_default(serving_chat, mock_request): assert comprehension_params.temperature == 0.4 # Preserved from YAML -def test_request_frequency_penalty_overrides(serving_chat, mock_request): +def test_request_frequency_penalty_overrides(mock_engine_client, mock_request): """Test that request frequency_penalty is applied.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.frequency_penalty = 0.5 mock_request.model_fields_set = {"frequency_penalty"} @@ -464,8 +464,9 @@ def test_request_frequency_penalty_overrides(serving_chat, mock_request): assert result[0].frequency_penalty == 0.5 -def test_request_presence_penalty_overrides(serving_chat, mock_request): +def test_request_presence_penalty_overrides(mock_engine_client, mock_request): """Test that request presence_penalty is applied.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.presence_penalty = 0.3 mock_request.model_fields_set = {"presence_penalty"} @@ -474,8 +475,9 @@ def test_request_presence_penalty_overrides(serving_chat, mock_request): assert result[0].presence_penalty == 0.3 -def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request): +def test_non_comprehension_stages_use_cloned_defaults(mock_engine_client, mock_request): """Test that non-comprehension stages always use cloned YAML defaults.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.max_tokens = 50 mock_request.temperature = 0.1 @@ -488,8 +490,9 @@ def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request assert other_params.seed == 42 # YAML default -def test_multiple_params_override_together(serving_chat, mock_request): +def test_multiple_params_override_together(mock_engine_client, mock_request): """Test that multiple request params can override together.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.max_tokens = 200 mock_request.temperature = 0.7 mock_request.top_p = 0.85 @@ -514,15 +517,17 @@ def test_multiple_params_override_together(serving_chat, mock_request): # ============================================================================= -def test_apply_request_overrides_clones_params(serving_chat, mock_request, default_comprehension_params): +def test_apply_request_overrides_clones_params(mock_engine_client, mock_request, default_comprehension_params): """Test that _apply_request_overrides returns a cloned object.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request) assert result is not default_comprehension_params # Different object -def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request, default_comprehension_params): +def test_apply_request_overrides_preserves_defaults(mock_engine_client, mock_request, default_comprehension_params): """Test that _apply_request_overrides preserves defaults when request has None.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request) assert result.temperature == 0.4 @@ -531,8 +536,9 @@ def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request, assert result.top_k == 1 # YAML custom param -def test_apply_request_overrides_applies_values(serving_chat, mock_request, default_comprehension_params): +def test_apply_request_overrides_applies_values(mock_engine_client, mock_request, default_comprehension_params): """Test that _apply_request_overrides applies non-None request values.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.temperature = 0.8 mock_request.seed = 123 mock_request.model_fields_set = {"temperature", "seed"} @@ -550,8 +556,9 @@ def test_apply_request_overrides_applies_values(serving_chat, mock_request, defa # ============================================================================= -def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker): +def test_apply_overrides_empty_stop_list_preserves_default(mock_engine_client, mocker): """Test that request.stop=[] does NOT override YAML default stop words.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"]) request = mocker.MagicMock() request.temperature = None @@ -573,8 +580,9 @@ def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker) assert result.stop == ["<|im_end|>"] # YAML default preserved -def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mocker): +def test_apply_overrides_nonempty_stop_list_overrides_default(mock_engine_client, mocker): """Test that request.stop=["\\n"] overrides YAML default stop words.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"]) request = mocker.MagicMock() request.temperature = None @@ -596,8 +604,9 @@ def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mock assert result.stop == ["\n"] # Overridden by request -def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mocker): +def test_apply_overrides_empty_stop_token_ids_preserves_default(mock_engine_client, mocker): """Test that request.stop_token_ids=[] does NOT override YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3]) request = mocker.MagicMock() request.temperature = None @@ -617,8 +626,9 @@ def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mo assert result.stop_token_ids == [2, 3] # YAML default preserved -def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat, mocker): +def test_apply_overrides_nonempty_stop_token_ids_overrides_default(mock_engine_client, mocker): """Test that request.stop_token_ids=[100] overrides YAML default.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3]) request = mocker.MagicMock() request.temperature = None @@ -640,8 +650,9 @@ def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat, assert result.stop_token_ids == [100] # Overridden by request -def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker): +def test_apply_overrides_mixed_empty_and_nonempty_lists(mock_engine_client, mocker): """Test mixing empty and non-empty list fields with scalar fields.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams( temperature=0.4, stop=["<|end|>"], @@ -669,8 +680,9 @@ def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker): assert result.stop_token_ids == [100, 200] # Non-empty list DID override -def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocker): +def test_apply_overrides_none_scalar_still_preserves_default(mock_engine_client, mocker): """Regression: ensure None scalar values still don't override defaults.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams(temperature=0.5, max_tokens=100, seed=42) request = mocker.MagicMock() request.temperature = None @@ -694,8 +706,9 @@ def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocke assert result.seed == 42 -def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocker): +def test_apply_overrides_both_lists_empty_preserves_defaults(mock_engine_client, mocker): """Test that both stop=[] and stop_token_ids=[] preserve YAML defaults.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) default_params = SamplingParams( temperature=0.5, stop=["<|end|>", "\\n"], @@ -722,9 +735,10 @@ def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocke assert result.stop_token_ids == [2, 32000] -def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock_request): +def test_build_sampling_params_list_empty_stop_preserves_yaml(mock_engine_client, mock_request): """Test that empty stop list in request preserves YAML defaults via _build_sampling_params_list_from_request.""" + serving_chat = build_serving_chat(engine_client=mock_engine_client) mock_request.stop = [] mock_request.stop_token_ids = [] @@ -736,21 +750,18 @@ def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock assert comprehension_params.stop_token_ids == [] -def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(): +def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(mocker: MockerFixture): """AURA callers may pass 3 semantic model params for a 4-stage engine pipeline.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) default_params = [ SamplingParams(max_tokens=10), SamplingParams(max_tokens=20), SamplingParams(max_tokens=30), SamplingParams(max_tokens=40), ] - instance.engine_client = SimpleNamespace( - stage_configs=[SimpleNamespace(stage_type="llm") for _ in range(4)], - default_sampling_params_list=default_params, - ) + engine_client = mocker.MagicMock() + engine_client.stage_configs = [SimpleNamespace(stage_type="llm") for _ in range(4)] + engine_client.default_sampling_params_list = default_params + instance = build_serving_chat(engine_client=engine_client) result = instance._to_sampling_params_list( [ @@ -772,44 +783,35 @@ def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(): def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client): """Test finding comprehension stage when it's at index 0.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) - instance.engine_client = mock_engine_client + instance = build_serving_chat(engine_client=mock_engine_client) assert instance._get_comprehension_stage_index() == 0 def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture): """Test finding comprehension stage when it's at index 1.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) - other = mocker.MagicMock() other.is_comprehension = False comprehension = mocker.MagicMock() comprehension.is_comprehension = True - instance.engine_client = mocker.MagicMock() - instance.engine_client.stage_configs = [other, comprehension] + engine_client = mocker.MagicMock() + engine_client.stage_configs = [other, comprehension] + instance = build_serving_chat(engine_client=engine_client) assert instance._get_comprehension_stage_index() == 1 def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixture): """Test that ValueError is raised when no comprehension stage exists.""" - from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat - - instance = object.__new__(OmniOpenAIServingChat) - stage1 = mocker.MagicMock() stage1.is_comprehension = False stage2 = mocker.MagicMock() stage2.is_comprehension = False - instance.engine_client = mocker.MagicMock() - instance.engine_client.stage_configs = [stage1, stage2] + engine_client = mocker.MagicMock() + engine_client.stage_configs = [stage1, stage2] + instance = build_serving_chat(engine_client=engine_client) with pytest.raises(ValueError, match="No comprehension stage"): instance._get_comprehension_stage_index() diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 94d84d041..19d383504 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # tests/entrypoints/openai/test_serving_speech.py import asyncio import base64 @@ -8,6 +11,7 @@ import logging import os import struct import wave +from dataclasses import FrozenInstanceError, replace from inspect import Signature, signature from pathlib import Path from types import SimpleNamespace @@ -33,15 +37,24 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( OpenAICreateAudioGenerateRequest, OpenAICreateSpeechRequest, SpeechBatchItem, + SpeechInputTokenDetails, + SpeechTokenUsage, StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import ( - _TTS_LANGUAGES, OmniOpenAIServingSpeech, _create_wav_header, ) -from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest, SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + DEFAULT_TTS_LANGUAGES, + PreparedRequest, + SpeechServingContext, + TTSGenerationError, +) +from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter +from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError +from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( FISH_TEXT_ONLY_SYSTEM_PROMPT, build_fish_voice_clone_prompt_ids, @@ -52,6 +65,7 @@ from vllm_omni.model_executor.models.ming_tts.constants import ( KEY_SPEAKER_WAVEFORM, KEY_SPEAKER_WAVEFORM_LENGTHS, SPEAKER_EMBEDDING_DIM, + TEXT_EOS_TOKEN_ID, ) from vllm_omni.outputs import OmniRequestOutput @@ -154,6 +168,7 @@ class TestAudioMixin: # Helper to create mock model output for endpoint tests def create_mock_audio_output_for_test( request_id: str = "speech-mock-123", + metrics: dict | None = None, ) -> OmniRequestOutput: class MockCompletionOutput: def __init__(self, index: int = 0): @@ -180,11 +195,16 @@ def create_mock_audio_output_for_test( audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples)) mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor) - return OmniRequestOutput.from_stage_output( + if metrics is None: + metrics = {"stage_metrics": {"0": {"num_tokens_in": 0, "num_tokens_out": 5}}} + + output = OmniRequestOutput.from_stage_output( mock_request_output, stage_id=0, final_output_type="audio", ) + output.metrics = metrics + return output def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None: @@ -209,7 +229,8 @@ def _wav_data_url(samples: np.ndarray, sample_rate: int) -> str: @pytest.fixture -def test_app(mocker: MockerFixture): +def test_app(mocker: MockerFixture, tmp_path, monkeypatch): + monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path)) # Mock the engine client mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False @@ -232,6 +253,12 @@ def test_app(mocker: MockerFixture): models=mock_models, request_logger=mock_request_logger, ) + # Voice endpoint tests exercise generic embedding parsing/persistence with + # multiple valid dimensions. Supply the adapter contract without imposing + # one real model's fixed embedding dimension on this shared fixture. + speech_server._adapter = mocker.MagicMock() + speech_server._adapter.validate_tts_embedding_dim.return_value = None + speech_server._tts_tokenizer = lambda text, padding=False: {"input_ids": text.split()} # Skip TTS validation in tests (mock doesn't set up supported_speakers) speech_server._validate_tts_request = mocker.MagicMock(return_value=None) @@ -253,31 +280,9 @@ def test_app(mocker: MockerFixture): speech_server.create_speech = awaitable_patched_create_speech app = FastAPI() + app.state.openai_serving_speech = speech_server app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) - - # Add list_voices endpoint - async def list_voices(): - speakers = sorted(speech_server.supported_speakers) if speech_server.supported_speakers else [] - uploaded_voices = [] - if hasattr(speech_server, "uploaded_speakers"): - for voice_name, info in speech_server.uploaded_speakers.items(): - voice_entry = { - "name": info.get("name", voice_name), - "consent": info.get("consent", ""), - "created_at": info.get("created_at", 0), - "file_size": info.get("file_size", 0), - "mime_type": info.get("mime_type", ""), - "embedding_source": info.get("embedding_source", "audio"), - "embedding_dim": info.get("embedding_dim"), - } - if info.get("ref_text"): - voice_entry["ref_text"] = info["ref_text"] - if info.get("speaker_description"): - voice_entry["speaker_description"] = info["speaker_description"] - uploaded_voices.append(voice_entry) - return {"voices": speakers, "uploaded_voices": uploaded_voices} - - app.add_api_route("/v1/audio/voices", list_voices, methods=["GET"]) + app.add_api_route("/v1/audio/voices", api_server_module.list_voices, methods=["GET"]) app.add_api_route("/v1/audio/speech/batch", speech_server.create_speech_batch, methods=["POST"]) # Add upload_voice endpoint @@ -369,6 +374,37 @@ class TestSpeechAPI: assert response.headers["content-type"] == "audio/wav" assert len(response.content) > 0 + def test_create_speech_includes_token_usage_and_detail_headers(self, client): + payload = { + "input": "Hello world", + "model": "tts-model", + "voice": "alloy", + "response_format": "wav", + } + response = client.post("/v1/audio/speech", json=payload) + assert response.status_code == 200 + assert response.headers["x-vllm-omni-input-tokens"] == "2" + assert response.headers["x-vllm-omni-output-tokens"] == "5" + assert response.headers["x-vllm-omni-total-tokens"] == "7" + assert response.headers["x-vllm-omni-input-text-tokens"] == "2" + assert response.headers["x-vllm-omni-input-audio-tokens"] == "0" + + def test_build_speech_usage_headers_uses_usage_and_detail_field_names(self): + usage = SpeechTokenUsage( + input_tokens=7, + output_tokens=11, + total_tokens=18, + input_token_details=SpeechInputTokenDetails(text_tokens=3, audio_tokens=4), + ) + headers = OmniOpenAIServingSpeech._build_speech_usage_headers(usage) + assert headers == { + "X-VLLM-OMNI-INPUT-TOKENS": "7", + "X-VLLM-OMNI-OUTPUT-TOKENS": "11", + "X-VLLM-OMNI-TOTAL-TOKENS": "18", + "X-VLLM-OMNI-INPUT-TEXT-TOKENS": "3", + "X-VLLM-OMNI-INPUT-AUDIO-TOKENS": "4", + } + def test_create_speech_mp3_format(self, client): payload = { "input": "Hello world", @@ -418,9 +454,31 @@ class TestSpeechAPI: assert audio_obj.speed == 2.5 def test_list_voices_endpoint(self, client): + handler = client.app.state.openai_serving_speech + handler._adapter.capabilities.supported_speakers = frozenset({"alice"}) + handler._adapter.capabilities.precomputed_speakers = {"bob": {"name": "Bob"}} + handler.uploaded_speakers = { + "carol": { + "name": "Carol", + "created_at": 1, + } + } + response = client.get("/v1/audio/voices") + assert response.status_code == 200 - assert "voices" in response.json() + assert response.json()["voices"] == ["alice", "bob", "carol", "default"] + assert response.json()["uploaded_voices"] == [ + { + "name": "Carol", + "consent": "", + "created_at": 1, + "file_size": 0, + "mime_type": "", + "embedding_source": "audio", + "embedding_dim": None, + } + ] def test_upload_voice_success(self, client, tmp_path): """Test successful voice upload without ref_text.""" @@ -678,6 +736,20 @@ class TestSpeechAPI: assert response.status_code == 400 assert "finite" in response.json()["detail"] + @pytest.mark.asyncio + async def test_diffusion_create_speech_with_unknown_voice(self, mocker: MockerFixture): + engine_client = mocker.MagicMock() + server = OmniOpenAIServingSpeech.for_diffusion( + diffusion_engine=engine_client, + model_name="test-model", + ) + assert server._adapter is None + + response = await server.create_speech(OpenAICreateSpeechRequest(input="test-input", voice="test-voice")) + + assert response.status_code == 400 + assert b"Invalid voice" in response.body + @pytest.mark.asyncio async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture): """Test public diffusion speech success and extra_params propagation.""" @@ -726,7 +798,8 @@ class TestTTSMethods: """Unit tests for TTS validation and parameter building.""" @pytest.fixture - def speech_server(self, mocker: MockerFixture): + def speech_server(self, mocker, tmp_path, monkeypatch): + monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path)) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.stage_configs = [] @@ -949,12 +1022,18 @@ class TestTTSMethods: def test_speaker_embedding_valid_base_task(self, speech_server): """speaker_embedding with Base task, x_vector_only_mode, and no ref_audio is accepted.""" + speech_server.engine_client.model_config = SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(hidden_size=1024)) + ) emb = [0.1] * 1024 req = OpenAICreateSpeechRequest(input="Hello", task_type="Base", speaker_embedding=emb, x_vector_only_mode=True) assert speech_server._validate_tts_request(req) is None def test_speaker_embedding_auto_sets_x_vector_only_mode(self, speech_server): """speaker_embedding auto-implies x_vector_only_mode, so validation passes.""" + speech_server.engine_client.model_config = SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(hidden_size=1024)) + ) emb = [0.1] * 1024 req = OpenAICreateSpeechRequest(input="Hello", task_type="Base", speaker_embedding=emb) result = speech_server._validate_tts_request(req) @@ -1026,6 +1105,7 @@ class TestTTSMethods: """Embedding uploads must match the loaded Qwen3-TTS model before being stored.""" speech_server._tts_model_type = "qwen3_tts" + speech_server._adapter = speech_server._get_tts_adapter() speech_server.engine_client.model_config = SimpleNamespace( hf_config=SimpleNamespace( talker_config=SimpleNamespace(hidden_size=2048), @@ -1044,6 +1124,7 @@ class TestTTSMethods: def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server): """Ming embedding uploads must match the model dimension before replacement.""" speech_server._tts_model_type = "ming_tts" + speech_server._adapter = speech_server._get_tts_adapter() invalid_dim = SPEAKER_EMBEDDING_DIM - 1 existing = { "name": "bad_emb_voice", @@ -1071,9 +1152,13 @@ class TestTTSMethods: def test_upload_ming_audio_voice_defers_speaker_extraction(self, speech_server, mocker: MockerFixture, tmp_path): """Ming audio uploads stay model-agnostic and defer extraction to use.""" speech_server._tts_model_type = "ming_tts" + speech_server._adapter = speech_server._get_tts_adapter() speech_server.uploaded_speakers_dir = tmp_path speech_server.uploaded_speakers = {} - speech_server.supported_speakers = set() + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + supported_speakers=frozenset(), + ) speech_server._speaker_cache.clear() left = np.full(32000, 0.5, dtype=np.float32) @@ -1161,23 +1246,421 @@ class TestTTSMethods: assert second[1] == 24000 assert first[0] is second[0] assert first[0][0] == pytest.approx(float(wav[0]), abs=1e-4) + cache_key = first[2] assert speech_server._get_resolved_ref_audio_artifact_key( - ref_audio + cache_key ) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000) + @pytest.mark.asyncio + async def test_diffusion_ref_audio_uses_media_access_config( + self, + mocker: MockerFixture, + monkeypatch, + tmp_path, + ): + monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path / "speakers")) + connector = mocker.MagicMock() + connector.fetch_audio_async = mocker.AsyncMock(return_value=(np.zeros(24000, dtype=np.float32), 24000)) + connector_cls = mocker.patch.object( + serving_speech_module, + "MediaConnector", + return_value=connector, + ) + media_dir = tmp_path / "media" + media_dir.mkdir() + ref_audio = media_dir / "reference.wav" + ref_audio.write_bytes(b"v1") + + server = OmniOpenAIServingSpeech.for_diffusion( + diffusion_engine=mocker.MagicMock(), + model_name="test-model", + allowed_local_media_path=str(media_dir), + allowed_media_domains=["media.example.com"], + ) + try: + first = await server._resolve_ref_audio(ref_audio.as_uri()) + ref_audio.write_bytes(b"v2-with-different-size") + second = await server._resolve_ref_audio(ref_audio.as_uri()) + finally: + server.shutdown() + + connector_cls.assert_called_once_with( + allowed_local_media_path=str(media_dir), + allowed_media_domains=["media.example.com"], + ) + assert first[2] != second[2] + assert connector.fetch_audio_async.await_count == 2 + + # ── ref-audio cache key tests (static helper) ── + + def test_local_file_cache_key_invalidation(self): + """Modified file must produce a different cache key (mtime change).""" + import pathlib + import tempfile + import time as _time + + with tempfile.TemporaryDirectory() as tmp_dir: + test_file = pathlib.Path(tmp_dir) / "test_audio.wav" + test_file.write_text("dummy audio data 1") + file_uri = f"file://{test_file}" + key1 = OmniOpenAIServingSpeech._get_ref_audio_cache_key( + file_uri, + allowed_local_media_path=tmp_dir, + ) + _time.sleep(0.01) + test_file.write_text("dummy audio data 2") + key2 = OmniOpenAIServingSpeech._get_ref_audio_cache_key( + file_uri, + allowed_local_media_path=tmp_dir, + ) + assert key1 != key2, "Cache key should change when file is modified!" + + def test_remote_url_cache_key(self): + """Remote URLs must produce a stable, repeatable cache key.""" + url = "https://example.com/audio.wav" + key1 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url) + key2 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url) + assert key1 == key2, "Cache key for URLs should remain constant!" + + def test_missing_file_cache_key(self, caplog): + """A missing file falls back to the URI-string hash and emits a log + so the stale-cache risk is visible.""" + file_uri = "file:///path/to/nonexistent/file.wav" + target_logger = logging.getLogger("vllm_omni.entrypoints.openai.serving_speech") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.DEBUG) + try: + key = OmniOpenAIServingSpeech._get_ref_audio_cache_key( + file_uri, + allowed_local_media_path="/path/to/nonexistent", + ) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) + expected_key = hashlib.sha1(file_uri.encode("utf-8")).hexdigest() + assert key == expected_key, "Missing files should fallback to the string hash." + assert any("stale cache" in r.getMessage() for r in caplog.records), ( + "A log message about stale cache must be emitted when os.stat fails" + ) + + def test_percent_encoded_file_uri(self): + """Percent-encoded file:// URIs must be decoded so os.stat hits the real file.""" + import pathlib + import tempfile + + with tempfile.TemporaryDirectory() as tmp_dir: + test_file = pathlib.Path(tmp_dir) / "my audio.wav" + test_file.write_text("dummy audio data") + encoded_uri = f"file://{str(test_file).replace(' ', '%20')}" + key = OmniOpenAIServingSpeech._get_ref_audio_cache_key( + encoded_uri, + allowed_local_media_path=tmp_dir, + ) + string_only_key = hashlib.sha1(encoded_uri.encode("utf-8")).hexdigest() + assert key != string_only_key, ( + "Percent-encoded URI should be decoded and stat'd, not fall back to string-only cache key!" + ) + + def test_no_allowlist_skips_stat_for_local_paths(self): + """When allowed_local_media_path is None, local paths must NOT be stat'd + — the key falls back to the string-only hash.""" + import pathlib + import tempfile + + with tempfile.TemporaryDirectory() as tmp_dir: + test_file = pathlib.Path(tmp_dir) / "test.wav" + test_file.write_text("data") + file_uri = f"file://{test_file}" + key = OmniOpenAIServingSpeech._get_ref_audio_cache_key( + file_uri, + allowed_local_media_path=None, + ) + # With no allowlist the stat is skipped, so the key is just sha1(uri) + expected = hashlib.sha1(file_uri.encode("utf-8")).hexdigest() + assert key == expected, "No allowlist should skip stat and use string-only key" + + def test_empty_string_allowlist_skips_stat_for_local_paths(self): + """vLLM's default allowed_local_media_path is '' not None; that must + also skip stat rather than treating cwd as the allowlist.""" + import pathlib + import tempfile + + with tempfile.TemporaryDirectory() as tmp_dir: + test_file = pathlib.Path(tmp_dir) / "test.wav" + test_file.write_text("data") + file_uri = f"file://{test_file}" + key = OmniOpenAIServingSpeech._get_ref_audio_cache_key( + file_uri, + allowed_local_media_path="", + ) + expected = hashlib.sha1(file_uri.encode("utf-8")).hexdigest() + assert key == expected, "Empty-string allowlist should skip stat and use string-only key" + + def test_uppercase_file_uri_uses_local_metadata(self): + """FILE:// must be treated as a local file, not a string-only locator.""" + import pathlib + import tempfile + + with tempfile.TemporaryDirectory() as tmp_dir: + test_file = pathlib.Path(tmp_dir) / "test.wav" + test_file.write_text("data") + upper_uri = f"FILE://{test_file}" + key = OmniOpenAIServingSpeech._get_ref_audio_cache_key( + upper_uri, + allowed_local_media_path=tmp_dir, + ) + string_only = hashlib.sha1(upper_uri.encode("utf-8")).hexdigest() + assert key != string_only + + def test_uppercase_http_uri_is_remote(self, speech_server): + url = "HTTP://example.com/audio.wav" + key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url) + assert key == hashlib.sha1(url.encode("utf-8")).hexdigest() + assert speech_server._validate_ref_audio_format(url) is None + assert speech_server._validate_ref_audio_format("FILE:///tmp/a.wav") is None + assert speech_server._validate_ref_audio_format("/tmp/a.wav") is not None + + @pytest.mark.asyncio + async def test_resolve_ref_audio_retries_when_file_changes_during_fetch(self, speech_server, mocker): + import pathlib + import tempfile + + wav = np.linspace(-0.5, 0.5, 48000, dtype=np.float32) + with tempfile.TemporaryDirectory() as tmp_dir: + path = pathlib.Path(tmp_dir) / "spk.wav" + path.write_bytes(b"v1") + uri = f"file://{path}" + speech_server.model_config.allowed_local_media_path = tmp_dir + speech_server.model_config.allowed_media_domains = None + speech_server._diffusion_mode = False + + fetches = {"n": 0} + + async def fake_fetch(_ref): + fetches["n"] += 1 + if fetches["n"] == 1: + path.write_bytes(b"v2-longer-content") + return wav, 24000 + + mock_connector = mocker.MagicMock() + mock_connector.fetch_audio_async = fake_fetch + mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.MediaConnector", + return_value=mock_connector, + ) + + wav_list, sr, cache_key = await speech_server._resolve_ref_audio(uri) + expected_key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(uri, allowed_local_media_path=tmp_dir) + assert sr == 24000 + assert len(wav_list) == 48000 + assert fetches["n"] == 2 + assert cache_key == expected_key + assert cache_key in speech_server._ref_audio_resolve_cache + + def test_conditioning_cache_salt_changes_with_ref_audio_cache_key(self): + """Prefix-cache salt must fold the content-aware resolve key so a + same-path file rewrite cannot reuse KV from the previous reference.""" + from vllm_omni.entrypoints.openai.serving_speech import _conditioning_cache_salt + + req = OpenAICreateSpeechRequest(input="hello", ref_audio="file:///data/spk.wav") + salt_a = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_aaa"}) + salt_b = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_bbb"}) + salt_locator_only = _conditioning_cache_salt(req, {}) + assert salt_a != salt_b + assert salt_a != salt_locator_only + + @pytest.mark.asyncio + async def test_higgs_v3_cache_salt_changes_with_ref_audio_cache_key(self, speech_server, mocker): + """Higgs v3 voice clone uses placeholder prompt ids + prefix caching. + The salt must move when the resolve key moves, or a same-path rewrite + reuses KV from the previous reference.""" + from unittest.mock import AsyncMock, MagicMock + + adapter = MagicMock() + adapter.build_prompt.return_value = [1, 1, 1, 1] + speech_server._resolve_higgs_audio_v3_adapter = AsyncMock(return_value=adapter) + codes = torch.arange(8, dtype=torch.long).view(8, 1) + speech_server._resolve_higgs_audio_v3_ref_codes = AsyncMock(return_value=(codes, False, False)) + speech_server._get_resolved_ref_audio_artifact_key = MagicMock(return_value="art") + + req = OpenAICreateSpeechRequest( + input="hello", + ref_audio="file:///data/spk.wav", + ref_text="transcript", + ) + speech_server._resolve_ref_audio = AsyncMock(return_value=([0.1] * 48000, 24000, "key_aaa")) + prompt_a = await speech_server._build_higgs_audio_v3_params(req) + speech_server._resolve_ref_audio = AsyncMock(return_value=([0.9] * 48000, 24000, "key_bbb")) + prompt_b = await speech_server._build_higgs_audio_v3_params(req) + + assert prompt_a["prompt_token_ids"] == prompt_b["prompt_token_ids"] + assert prompt_a["cache_salt"] != prompt_b["cache_salt"] + assert prompt_a["additional_information"]["ref_audio_cache_key"] == "key_aaa" + assert prompt_b["additional_information"]["ref_audio_cache_key"] == "key_bbb" + + # ── encode_reference_codes: speaker cache invalidation ── + + @pytest.mark.asyncio + async def test_encode_reference_codes_reencodes_on_key_change(self): + """Changing the resolve_cache_key must trigger a fresh encode, not serve + the stale cached tensor. This is the user-visible fix for delay-family + voice clones that were reproducing the old speaker.""" + from unittest.mock import MagicMock + + from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes + from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + + speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024) + + call_count = 0 + codes_v1 = torch.arange(10, dtype=torch.int64) + codes_v2 = torch.arange(10, 20, dtype=torch.int64) + + # Stub processor that tracks call count and returns different codes + processor = MagicMock() + + def fake_encode(wavs, sampling_rate, n_vq): + nonlocal call_count + call_count += 1 + return [codes_v1 if call_count == 1 else codes_v2] + + processor.encode_audios_from_wav = fake_encode + + # resolve_ref_audio returns different cache keys on successive calls + resolve_keys = iter(["key_aaa", "key_bbb"]) + + async def mock_resolve(ref_str): + return ([0.1, 0.2, 0.3], 24000, next(resolve_keys)) + + # First call: cold miss, encodes and stores + result1 = await encode_reference_codes( + "data:audio/wav;base64,fake", + processor=processor, + resolve_ref_audio=mock_resolve, + speaker_cache=speaker_cache, + variant="tts", + n_vq=32, + sr_target=24000, + ) + assert call_count == 1 + assert torch.equal(result1, codes_v1) + + # Second call: different resolve key → must re-encode, not serve cached + result2 = await encode_reference_codes( + "data:audio/wav;base64,fake", + processor=processor, + resolve_ref_audio=mock_resolve, + speaker_cache=speaker_cache, + variant="tts", + n_vq=32, + sr_target=24000, + ) + assert call_count == 2, "Different resolve key should trigger re-encode" + assert torch.equal(result2, codes_v2) + + @pytest.mark.asyncio + async def test_encode_reference_codes_named_voice_skips_resolve_on_hit(self): + """Named-voice cache hit must NOT call resolve_ref_audio.""" + from unittest.mock import MagicMock + + from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes + from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + + speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024) + + processor = MagicMock() + processor.encode_audios_from_wav = lambda wavs, sampling_rate, n_vq: [torch.zeros(5)] + + resolve_called = 0 + + async def mock_resolve(ref_str): + nonlocal resolve_called + resolve_called += 1 + return ([0.1], 24000, "some_key") + + # First call: cold miss, resolves + encodes + await encode_reference_codes( + "data:audio/wav;base64,fake", + processor=processor, + resolve_ref_audio=mock_resolve, + speaker_cache=speaker_cache, + variant="tts", + n_vq=32, + sr_target=24000, + voice_name="alice", + voice_created_at=100, + ) + assert resolve_called == 1 + + # Second call: warm hit, must NOT resolve + await encode_reference_codes( + "data:audio/wav;base64,fake", + processor=processor, + resolve_ref_audio=mock_resolve, + speaker_cache=speaker_cache, + variant="tts", + n_vq=32, + sr_target=24000, + voice_name="alice", + voice_created_at=100, + ) + assert resolve_called == 1, "Named-voice cache hit should skip resolve_ref_audio" + + @pytest.mark.asyncio + async def test_ttsd_second_reference_does_not_use_named_voice(self, speech_server, mocker): + """Uploaded voice names speaker 1 only; speaker 2 must stay anonymous.""" + speech_server._moss_variant = "ttsd" + speech_server.uploaded_speakers = {"alice": {}} + mocker.patch.object(speech_server, "_voice_created_at", return_value=42) + + proc = mocker.MagicMock() + proc.model_config.n_vq = 8 + proc.model_config.sampling_rate = 24000 + proc.build_user_message.return_value = "msg" + proc.return_value = {"input_ids": [torch.zeros((4, 9), dtype=torch.int64)]} + mocker.patch.object(speech_server, "_get_moss_processor", return_value=proc) + + seen: list[tuple[str, str | None]] = [] + + async def fake_encode(ref_str, **kwargs): + seen.append((ref_str, kwargs.get("voice_name"))) + return torch.zeros(3, dtype=torch.int64) + + mocker.patch( + "vllm_omni.model_executor.models.moss_tts.reference_encoder.encode_reference_codes", + fake_encode, + ) + + req = OpenAICreateSpeechRequest( + input="hello", + voice="alice", + ref_audio="data:audio/wav;base64,aaa", + ref_audio_2="data:audio/wav;base64,bbb", + ) + await speech_server._build_moss_tts_params(req, has_inline_ref_audio=False) + assert seen == [ + ("data:audio/wav;base64,aaa", "alice"), + ("data:audio/wav;base64,bbb", None), + ] + def test_precomputed_qwen3_voice_infers_base_without_ref_audio(self, speech_server): """Precomputed Qwen3 voices are reusable by name without per-request ref_audio.""" speech_server._tts_model_type = "qwen3_tts" - speech_server.precomputed_speakers = { - "alice": { - "name": "Alice", - "model_type": "qwen3_tts", - "mode": "icl", - "ref_text": "reference transcript", - "ref_code_length": 3, - } - } - speech_server.supported_speakers = {"alice"} + speech_server._adapter = speech_server._get_tts_adapter() + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers={ + "alice": { + "name": "Alice", + "model_type": "qwen3_tts", + "mode": "icl", + "ref_text": "reference transcript", + "ref_code_length": 3, + } + }, + supported_speakers=frozenset({"alice"}), + ) req = OpenAICreateSpeechRequest(input="Hello", voice="Alice") assert speech_server._validate_tts_request(req) is None @@ -1198,6 +1681,7 @@ class TestTTSMethods: uploaded_path = tmp_path / "alice.safetensors" save_file({"speaker_embedding": torch.tensor([0.1] * 4)}, str(uploaded_path)) speech_server._tts_model_type = "qwen3_tts" + speech_server._adapter = speech_server._get_tts_adapter() speech_server.uploaded_speakers_dir = tmp_path speech_server.uploaded_speakers = { "alice": { @@ -1207,15 +1691,18 @@ class TestTTSMethods: "embedding_source": "direct", } } - speech_server.precomputed_speakers = { - "alice": { - "name": "Alice", - "model_type": "qwen3_tts", - "mode": "icl", - "ref_text": "precomputed transcript", - "ref_code_length": 3, - } - } + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers={ + "alice": { + "name": "Alice", + "model_type": "qwen3_tts", + "mode": "icl", + "ref_text": "precomputed transcript", + "ref_code_length": 3, + } + }, + ) req = OpenAICreateSpeechRequest(input="Hello", voice="Alice") assert speech_server._validate_tts_request(req) is None @@ -1242,13 +1729,16 @@ class TestTTSMethods: talker_config=SimpleNamespace(hidden_size=4), ) ) + speech_server._adapter = speech_server._get_tts_adapter() - profiles = speech_server._load_precomputed_speakers() + profiles = speech_server._adapter._load_precomputed_speakers() assert profiles == {} - speech_server.precomputed_speakers = profiles - speech_server.supported_speakers = set(profiles) - assert "alice" not in speech_server.supported_speakers + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers=profiles, + ) + assert "alice" not in speech_server._adapter.capabilities.precomputed_speakers req = OpenAICreateSpeechRequest(input="Hello", voice="Alice") assert speech_server._validate_tts_request(req) is not None @@ -1276,13 +1766,16 @@ class TestTTSMethods: talker_config=SimpleNamespace(hidden_size=4), ) ) + speech_server._adapter = speech_server._get_tts_adapter() - profiles = speech_server._load_precomputed_speakers() + profiles = speech_server._adapter._load_precomputed_speakers() assert profiles == {} - speech_server.precomputed_speakers = profiles - speech_server.supported_speakers = set(profiles) - assert "alice" not in speech_server.supported_speakers + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers=profiles, + ) + assert "alice" not in speech_server._adapter.capabilities.precomputed_speakers req = OpenAICreateSpeechRequest(input="Hello", voice="Alice") assert speech_server._validate_tts_request(req) is not None @@ -1297,13 +1790,16 @@ class TestTTSMethods: speech_server.engine_client.model_config = SimpleNamespace( hf_config=SimpleNamespace(custom_voice_dir=str(tmp_path)) ) + speech_server._adapter = speech_server._get_tts_adapter() - profiles = speech_server._load_precomputed_speakers() + profiles = speech_server._adapter._load_precomputed_speakers() assert profiles == {} - speech_server.precomputed_speakers = profiles - speech_server.supported_speakers = set(profiles) - assert "bob" not in speech_server.supported_speakers + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers=profiles, + ) + assert "bob" not in speech_server._adapter.capabilities.precomputed_speakers req = OpenAICreateSpeechRequest(input="Hello", voice="Bob") assert speech_server._validate_tts_request(req) is not None assert speech_server._validate_tts_request(OpenAICreateSpeechRequest(input="Hello")) is None @@ -1311,9 +1807,12 @@ class TestTTSMethods: def test_prepare_voxcpm2_rejects_supported_speaker_without_voice_profile(self, speech_server, mocker): """VoxCPM2 named voices must be uploaded or precomputed on the real request path.""" speech_server._tts_model_type = "voxcpm2" - speech_server.supported_speakers = {"bob", "default"} + speech_server._adapter = speech_server._get_tts_adapter() speech_server.uploaded_speakers = {} - speech_server.precomputed_speakers = {} + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers={}, + ) speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] speech_server.engine_client.generate = mocker.MagicMock(return_value="generator") speech_server._build_voxcpm2_prompt = mocker.AsyncMock( @@ -1329,9 +1828,12 @@ class TestTTSMethods: def test_prepare_voxcpm2_accepts_default_voice(self, speech_server, mocker): """VoxCPM2 default voice preserves the built-in zero-shot request path.""" speech_server._tts_model_type = "voxcpm2" - speech_server.supported_speakers = {"default"} + speech_server._adapter = speech_server._get_tts_adapter() speech_server.uploaded_speakers = {} - speech_server.precomputed_speakers = {} + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers={}, + ) speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(())) speech_server._build_voxcpm2_prompt = mocker.AsyncMock( @@ -1346,23 +1848,26 @@ class TestTTSMethods: def test_prepare_voxcpm2_precomputed_voice_sets_model_cache_key(self, speech_server, mocker): """VoxCPM2 precomputed voices must carry voice metadata to the model cache lookup.""" speech_server._tts_model_type = "voxcpm2" - speech_server.supported_speakers = {"alice"} + speech_server._adapter = speech_server._get_tts_adapter() speech_server.uploaded_speakers = {} - speech_server.precomputed_speakers = { - "alice": { - "name": "Alice", - "model_type": "voxcpm2", - "mode": "reference", - "ref_audio_feat_len": 2, - } - } + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers={ + "alice": { + "name": "Alice", + "model_type": "voxcpm2", + "mode": "reference", + "ref_audio_feat_len": 2, + } + }, + ) speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(())) speech_server._build_voxcpm2_prompt = mocker.AsyncMock( return_value={ "prompt_token_ids": [1], "additional_information": { - "voice_profile": speech_server.precomputed_speakers["alice"], + "voice_profile": speech_server._adapter.capabilities.precomputed_speakers["alice"], }, } ) @@ -1441,6 +1946,10 @@ class TestTTSMethods: mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True + mock_stage = mocker.MagicMock() + mock_stage.engine_args.model_stage = "qwen3_tts" + mock_engine_client.stage_configs = [mock_stage] + server = OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, @@ -1448,7 +1957,10 @@ class TestTTSMethods: ) # Verify speakers are normalized to lowercase - assert server.supported_speakers == {"ryan", "vivian", "aiden"} + assert server._adapter.capabilities.supported_speakers == {"ryan", "vivian", "aiden"} + assert isinstance(server._adapter.capabilities.supported_speakers, frozenset) + with pytest.raises(FrozenInstanceError): + setattr(server._adapter.capabilities, "codec_frame_rate", 25.0) def test_load_supported_speakers_skips_non_tts_omni_model( self, @@ -1461,7 +1973,7 @@ class TestTTSMethods: ) assert speech_server._is_tts is False - assert speech_server._load_supported_speakers() == set() + assert load_supported_speakers(speech_server.engine_client) == set() warning.assert_not_called() def test_load_supported_languages_from_config(self, speech_server): @@ -1474,7 +1986,8 @@ class TestTTSMethods: ) ) ) - assert speech_server._load_supported_languages() == { + adapter = speech_server._get_tts_adapter() + assert adapter._load_supported_languages() == { "Chinese", "English", "Beijing_Dialect", @@ -1487,11 +2000,17 @@ class TestTTSMethods: speech_server.engine_client.model_config = SimpleNamespace( hf_config=SimpleNamespace(talker_config={"codec_language_id": {"chinese": 2055, "english": 2050}}) ) - assert speech_server._load_supported_languages() == {"Chinese", "English", "Auto"} + adapter = speech_server._get_tts_adapter() + assert adapter._load_supported_languages() == {"Chinese", "English", "Auto"} def test_validate_language_custom_dialect_accepted(self, speech_server): """A language present in the model config passes validation, case-insensitively.""" - speech_server.supported_languages = {"Chinese", "English", "Beijing_Dialect", "Auto"} + speech_server._tts_model_type = "qwen3_tts" + speech_server._adapter = speech_server._get_tts_adapter() + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + supported_languages=frozenset({"Chinese", "English", "Beijing_Dialect", "Auto"}), + ) for language in ("Beijing_Dialect", "beijing_dialect", "English", "english", "Auto", "AUTO"): req = OpenAICreateSpeechRequest(input="Hello", language=language) result = speech_server._validate_tts_request(req) @@ -1501,7 +2020,12 @@ class TestTTSMethods: def test_validate_language_unknown_rejected(self, speech_server): """A language not in the configured set is rejected.""" - speech_server.supported_languages = {"Chinese", "English", "Auto"} + speech_server._tts_model_type = "qwen3_tts" + speech_server._adapter = speech_server._get_tts_adapter() + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + supported_languages=frozenset({"Chinese", "English", "Auto"}), + ) for language in ("Klingon", "klingon"): req = OpenAICreateSpeechRequest(input="Hello", language=language) assert "Invalid language" in speech_server._validate_tts_request(req) @@ -1512,18 +2036,21 @@ class TestTTSMethods: speech_server.engine_client.model_config = SimpleNamespace( hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_language_id={})) ) - assert speech_server._load_supported_languages() == _TTS_LANGUAGES + adapter = speech_server._get_tts_adapter() + assert adapter._load_supported_languages() == DEFAULT_TTS_LANGUAGES def test_load_supported_languages_default_on_config_error(self, speech_server): """If the model config cannot be read, fall back to the default list.""" speech_server._tts_model_type = "qwen3_tts" speech_server.engine_client = SimpleNamespace() # no model_config -> AttributeError - assert speech_server._load_supported_languages() == _TTS_LANGUAGES + speech_server._adapter = None + adapter = speech_server._get_tts_adapter() + assert adapter._load_supported_languages() == DEFAULT_TTS_LANGUAGES def test_load_supported_languages_default_for_non_qwen(self, speech_server): """Non-qwen3_tts model types get the default language set.""" - speech_server._tts_model_type = None - assert speech_server._load_supported_languages() == _TTS_LANGUAGES + adapter = MingTTSAdapter(SpeechServingContext(server=speech_server, engine_client=speech_server.engine_client)) + assert adapter._load_supported_languages() == DEFAULT_TTS_LANGUAGES def test_build_tts_params_with_uploaded_voice(self, speech_server, mocker: MockerFixture): """Test _build_tts_params auto-sets ref_audio for uploaded voices (x_vector only).""" @@ -1536,7 +2063,6 @@ class TestTTSMethods: "created_at": 1711234567, } } - speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"} mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data") mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv" @@ -1560,7 +2086,6 @@ class TestTTSMethods: "created_at": 1711234567, } } - speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"} mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data") mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv" @@ -1577,7 +2102,6 @@ class TestTTSMethods: """Test _build_tts_params does not auto-set ref_audio for non-uploaded voices.""" # No uploaded speakers speech_server.uploaded_speakers = {} - speech_server.supported_speakers = {"ryan", "vivian"} req = OpenAICreateSpeechRequest(input="Hello", voice="ryan", task_type="Base") @@ -1597,7 +2121,6 @@ class TestTTSMethods: "mime_type": "audio/wav", } } - speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"} req = OpenAICreateSpeechRequest( input="Hello", voice="custom_voice", task_type="Base", ref_audio="data:audio/wav;base64,ZXhwbGljaXQ=" @@ -1703,7 +2226,6 @@ class TestTTSMethods: "cache_file": "/tmp/voice_samples/emb_voice.safetensors", } } - speech_server.supported_speakers = {"ryan", "vivian", "emb_voice"} fake_embedding = [0.1] * 1024 mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding") @@ -1777,7 +2299,7 @@ class TestTTSMethods: mocker.patch.object( speech_server, "_resolve_ref_audio", - new=mocker.AsyncMock(return_value=ref_audio_data), + new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")), ) mock_prompt = mocker.patch.object( speech_server, @@ -1894,7 +2416,7 @@ class TestTTSMethods: mocker.patch.object( speech_server, "_resolve_ref_audio", - new=mocker.AsyncMock(return_value=ref_audio_data), + new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")), ) mock_prompt = mocker.patch.object( speech_server, @@ -2162,6 +2684,28 @@ class TestTTSMethods: assert request.ref_audio is not None assert call.kwargs["has_inline_ref_audio"] is False + @pytest.mark.asyncio + async def test_warmup_with_no_adapter(self, speech_server): + assert speech_server._adapter is None + await speech_server.warmup() + + @pytest.mark.asyncio + async def test_voxcpm2_warmup_successfully(self, speech_server, mocker): + speech_server._tts_model_type = "voxcpm2" + speech_server._adapter = speech_server._get_tts_adapter() + speech_server.uploaded_speakers = {} + speech_server._adapter.capabilities = replace( + speech_server._adapter.capabilities, + precomputed_speakers={}, + ) + speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] + speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(())) + speech_server._build_voxcpm2_prompt = mocker.AsyncMock( + return_value={"prompt_token_ids": [1], "additional_information": {}} + ) + + await speech_server.warmup() + class TestFileValidationFunctions: """Unit tests for file validation helper functions.""" @@ -2390,6 +2934,40 @@ class TestStreamingResponse: assert "text/event-stream" not in response.headers["content-type"] assert len(response.content) > 0 + def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monkeypatch): + """The raw route must not bypass the shared Qwen3-TTS EOS guard.""" + speech_server = streaming_app.state.speech_server + finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]} + captured: dict = {} + + async def prepare(_request, request_id=None): + return request_id, object(), finalized_tts_params + + async def generate_chunks( + _generator, + _request_id, + _response_format="pcm", + raw_request=None, + request_start_s=None, + include_sample_rate=False, + usage_acc=None, + tts_params=None, + collect=None, + ): + captured["tts_params"] = tts_params + yield b"\x00\x00" + + monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare) + monkeypatch.setattr(speech_server, "_generate_audio_chunks", generate_chunks) + + response = TestClient(streaming_app).post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"}, + ) + + assert response.status_code == 200 + assert captured["tts_params"] is finalized_tts_params + def test_sse_streaming(self, streaming_app): """stream_format=sse without stream=True returns audio deltas as SSE.""" client = TestClient(streaming_app) @@ -2932,6 +3510,39 @@ def test_api_server_upload_voice_exception_returns_500(mocker: MockerFixture): ) +def test_api_server_upload_voice_with_no_adapter( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path)) + _patch_api_server_base(mocker) + engine_client = mocker.MagicMock() + engine_client.errored = False + engine_client.stage_configs = [] + models = mocker.MagicMock() + models.is_base_model.return_value = True + handler = OmniOpenAIServingSpeech( + engine_client=engine_client, + models=models, + request_logger=mocker.MagicMock(), + ) + assert handler._adapter is None + raw_request = _make_api_server_request(handler) + + response = asyncio.run( + api_server_module.upload_voice( + raw_request, + audio_sample=None, + speaker_embedding=json.dumps([0.1] * 1024), + consent="cons_test", + name="probe", + ) + ) + + assert response.status_code == 200 + + def test_api_server_delete_voice_without_speech_handler_returns_404(mocker: MockerFixture): _patch_api_server_base(mocker) raw_request = _make_api_server_request(None, method="DELETE", path="/v1/audio/voices/probe") @@ -3261,8 +3872,14 @@ class _FakeFishTokenizer: @pytest.fixture def fish_speech_server(mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False @@ -3532,8 +4149,14 @@ class TestWAVStreaming: @pytest.fixture def cosyvoice3_server(mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False @@ -3607,18 +4230,31 @@ class TestCosyVoice3Serving: error = cosyvoice3_server._validate_tts_request(request) assert error is None - def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server): - request = OpenAICreateSpeechRequest( - input="Hello", - ref_audio="data:audio/wav;base64,abc", - ref_text="hello", - max_new_tokens=0, - ) - error = cosyvoice3_server._validate_tts_request(request) - assert error is not None - assert "max_new_tokens" in error + def test_validate_cosyvoice3_max_new_tokens_range(self): + """Ensure max_new_tokens below the minimum is rejected during request validation.""" + with pytest.raises(ValidationError, match="max_new_tokens"): + OpenAICreateSpeechRequest( + input="Hello", + ref_audio="data:audio/wav;base64,abc", + ref_text="hello", + max_new_tokens=0, + ) - def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server, mocker: MockerFixture): + @pytest.mark.parametrize( + ("max_new_tokens", "expected_min_tokens", "expected_max_tokens"), + [ + (None, 10, 2000), + (5, 5, 5), + ], + ) + def test_prepare_speech_generation_cosyvoice3( + self, + cosyvoice3_server, + mocker: MockerFixture, + max_new_tokens: int | None, + expected_min_tokens: int, + expected_max_tokens: int, + ): cosyvoice3_server._build_cosyvoice3_prompt = mocker.AsyncMock( return_value={ "prompt": "Hello", @@ -3626,15 +4262,23 @@ class TestCosyVoice3Serving: "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000}, } ) - cosyvoice3_server._apply_cosyvoice3_dynamic_tokens = mocker.MagicMock(side_effect=lambda spl, req: spl) + cosyvoice3_server.model_config.hf_config = SimpleNamespace( + min_token_text_ratio=1, max_token_text_ratio=200, allowed_special=True + ) + cosyvoice3_server._cosyvoice3_tokenizer = mocker.MagicMock() + mocker.patch("vllm_omni.model_executor.models.cosyvoice3.utils.extract_text_token", return_value=[None, 10]) request = OpenAICreateSpeechRequest( input="Hello", ref_audio="data:audio/wav;base64,abc", ref_text="Reference text", + max_new_tokens=max_new_tokens, ) request_id, generator, tts_params = asyncio.run(cosyvoice3_server._prepare_speech_generation(request)) + sampling_params = cosyvoice3_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0] + assert sampling_params.min_tokens == expected_min_tokens + assert sampling_params.max_tokens == expected_max_tokens assert request_id.startswith("speech-") assert generator == "generator" assert tts_params == {} @@ -3646,8 +4290,14 @@ class TestCosyVoice3Serving: @pytest.fixture def glm_tts_server(mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False @@ -3669,7 +4319,6 @@ def glm_tts_server(mocker: MockerFixture): mock_models = mocker.MagicMock() mock_models.is_base_model.return_value = True - return OmniOpenAIServingSpeech( engine_client=mock_engine_client, models=mock_models, @@ -3718,11 +4367,107 @@ class TestGLMTTSServing: assert text_token_len == 3 load_tokenizer.assert_called_once() + def test_prepare_speech_generation_glm_tts(self, glm_tts_server, mocker: MockerFixture): + glm_tts_server._build_glm_tts_prompt = mocker.AsyncMock( + return_value={ + "prompt": "Hello", + "multi_modal_data": {"audio": (np.zeros(24000), 24000)}, + "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000}, + } + ) + request = OpenAICreateSpeechRequest( + input="Hello", + ref_audio="data:audio/wav;base64,abc", + ref_text="Reference text", + ) + glm_tts_server.model_config.hf_config = SimpleNamespace(min_token_text_ratio=1, max_token_text_ratio=200) + glm_tts_server._estimate_glm_tts_text_token_len = mocker.MagicMock(return_value=10) + request_id, generator, tts_params = asyncio.run(glm_tts_server._prepare_speech_generation(request)) + sampling_params = glm_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0] + + assert sampling_params.min_tokens == 10 + assert sampling_params.max_tokens == 2000 + assert request_id.startswith("speech-") + assert generator == "generator" + assert tts_params == {} + glm_tts_server._build_glm_tts_prompt.assert_awaited_once() + + +@pytest.fixture +def ming_tts_server(mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) + + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + mock_engine_client.model_config = mocker.MagicMock( + model="inclusionAI/Ming-omni-tts-0.5B", + hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20), + ) + mock_engine_client.default_sampling_params_list = [ + SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None) + ] + mock_engine_client.tts_batch_max_items = 32 + mock_engine_client.generate = mocker.MagicMock(return_value="generator") + mock_engine_client.stage_configs = [ + SimpleNamespace( + engine_args=SimpleNamespace( + model_stage="llm", model_arch="MingTTSForConditionalGeneration", worker_type="ar" + ), + tts_args={}, + ) + ] + + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + return OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + + +class TestMingTTSServing: + def test_prepare_speech_generation_ming(self, ming_tts_server, mocker: MockerFixture): + ming_tts_server._adapter.build = mocker.AsyncMock( + return_value=PreparedRequest( + prompt={"prompt": "hello"}, + tts_params={}, + model_type="ming_tts", + ) + ) + + request = OpenAICreateSpeechRequest( + input="Hello", + ref_audio="data:audio/wav;base64,abc", + ref_text="Reference text", + max_new_tokens=7, + ) + asyncio.run(ming_tts_server._prepare_speech_generation(request)) + + assert ming_tts_server._tts_model_type == "ming_tts" + ming_tts_server._adapter.build.assert_awaited_once() + sampling_params = ming_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0] + assert sampling_params.stop_token_ids == [TEXT_EOS_TOKEN_ID] + assert sampling_params.max_tokens == 8 + @pytest.fixture def ming_flash_omni_tts_server(mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False @@ -3785,8 +4530,14 @@ class TestTTSAsyncOffloading: @pytest.fixture def voxtral_server(self, mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.voxtral.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.model_config = mocker.MagicMock(model="mistralai/Voxtral") @@ -3811,8 +4562,14 @@ class TestTTSAsyncOffloading: @pytest.fixture def qwen3_tts_server(self, mocker: MockerFixture): - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.model_config = mocker.MagicMock(model="Qwen/Qwen3-TTS", hf_config=mocker.MagicMock()) @@ -3836,6 +4593,32 @@ class TestTTSAsyncOffloading: yield server server.shutdown() + def test_voxtral_loads_supported_speakers(self, mocker: MockerFixture): + engine_client = mocker.MagicMock() + engine_client.model_config.hf_config.audio_config = { + "speaker_id": { + "Alice": 0, + "Bob": 1, + } + } + + server = mocker.MagicMock() + adapter = VoxtralTTSAdapter( + SpeechServingContext( + server=server, + engine_client=engine_client, + ) + ) + + adapter.load_capabilities() + + assert adapter.capabilities.supported_speakers == frozenset( + { + "alice", + "bob", + } + ) + def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server, mocker: MockerFixture): """Voxtral path in _prepare_speech_generation should call the async wrapper.""" voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock( @@ -3901,6 +4684,9 @@ class TestTTSAsyncOffloading: model_type=adapter_model_type, ) + def apply_sampling_overrides(self, sampling_params_list, request, prompt, request_id=None): + return sampling_params_list + voxtral_server._tts_model_type = legacy_tts_model_type mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter()) log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info") @@ -3909,9 +4695,7 @@ class TestTTSAsyncOffloading: assert adapter_model_type != legacy_tts_model_type assert any( - call.args - and call.args[0] == "TTS speech request %s: text=%r, model=%s" - and call.args[3] == adapter_model_type + call.args and call.args[0] == "TTS speech request %s: model=%s" and call.args[2] == adapter_model_type for call in log_info.call_args_list ) @@ -4034,7 +4818,7 @@ class TestTTSAsyncOffloading: ): """Base explicit true should reach the model prompt additional_information.""" qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) - qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000)) + qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000, "fake_cache_key")) qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None) qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) @@ -4066,7 +4850,10 @@ class TestTTSAsyncOffloading: artifact_key, ) qwen3_tts_server._ref_audio_model_artifact_ready.add((artifact_key, False)) - qwen3_tts_server._codec_frame_rate = 25.0 + qwen3_tts_server._adapter.capabilities = replace( + qwen3_tts_server._adapter.capabilities, + codec_frame_rate=25.0, + ) qwen3_tts_server._tts_tokenizer = lambda _text, padding=False: {"input_ids": list(range(10))} qwen3_tts_server.engine_client.model_config.hf_config.talker_config = SimpleNamespace( codec_language_id={}, @@ -4101,6 +4888,57 @@ class TestTTSAsyncOffloading: assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()} + @pytest.mark.asyncio + async def test_qwen3_tts_nonstream_retries_codec_limit_once(self, qwen3_tts_server, mocker: MockerFixture) -> None: + qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None) + qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock( + side_effect=[ + Qwen3TTSCodecLimitError("codec limit"), + (b"RIFF" + b"\x00" * 32, "audio/wav"), + ] + ) + request = OpenAICreateSpeechRequest(input="Hello", task_type="Base") + + response = await qwen3_tts_server.create_speech(request) + + assert response.status_code == 200 + assert qwen3_tts_server._generate_audio_bytes.await_count == 2 + first_call, retry_call = qwen3_tts_server._generate_audio_bytes.await_args_list + assert first_call.args[0] is request + assert first_call.kwargs["request_id"] != retry_call.kwargs["request_id"] + assert retry_call.args[0].seed is not None + assert request.seed is None + + @pytest.mark.asyncio + async def test_nonstream_does_not_retry_nonretryable_generation_error( + self, qwen3_tts_server, mocker: MockerFixture + ) -> None: + qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None) + qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=TTSGenerationError("invalid generation")) + + response = await qwen3_tts_server.create_speech(OpenAICreateSpeechRequest(input="Hello", task_type="Base")) + + assert response.status_code == 500 + qwen3_tts_server._generate_audio_bytes.assert_awaited_once() + + @pytest.mark.asyncio + @pytest.mark.parametrize("request_kwargs", [{"seed": 100423}, {"max_new_tokens": 192}]) + async def test_qwen3_tts_nonstream_does_not_retry_explicit_sampling_controls( + self, + qwen3_tts_server, + mocker: MockerFixture, + request_kwargs: dict[str, int], + ) -> None: + qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None) + qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=Qwen3TTSCodecLimitError("codec limit")) + + response = await qwen3_tts_server.create_speech( + OpenAICreateSpeechRequest(input="Hello", task_type="Base", **request_kwargs) + ) + + assert response.status_code == 500 + qwen3_tts_server._generate_audio_bytes.assert_awaited_once() + @pytest.mark.asyncio async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server): async def failing_generator(): @@ -4115,6 +4953,80 @@ class TestTTSAsyncOffloading: assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready + @pytest.mark.asyncio + async def test_generate_audio_chunks_rejects_qwen3_codec_limit_before_success(self, qwen3_tts_server): + async def length_limited_generator(): + yield SimpleNamespace( + multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000}, + metrics={ + "stage_metrics": { + "0": { + "num_tokens_out": 191, + "finish_reason": "length", + } + } + }, + ) + + chunks = qwen3_tts_server._generate_audio_chunks( + length_limited_generator(), + "req-codec-limit", + tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]}, + ) + with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"): + async for _ in chunks: + pass + + @pytest.mark.asyncio + async def test_sse_codec_limit_marks_partial_audio_for_discard(self, qwen3_tts_server): + async def length_limited_generator(): + yield SimpleNamespace( + multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000}, + metrics={ + "stage_metrics": { + "0": { + "num_tokens_out": 191, + "finish_reason": "length", + } + } + }, + ) + + events = [ + event + async for event in qwen3_tts_server._generate_audio_sse_events( + length_limited_generator(), + "req-sse-codec-limit", + tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]}, + ) + ] + + assert any("event: speech.audio.delta" in event for event in events) + assert not any("event: speech.audio.done" in event for event in events) + error_event = next(event for event in events if "event: speech.audio.error" in event) + payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:]) + assert payload["error"]["partial_audio"] is True + assert payload["error"]["action"] == "discard" + + @pytest.mark.asyncio + async def test_sse_error_before_audio_is_not_marked_partial(self, qwen3_tts_server): + async def failing_generator(): + raise RuntimeError("boom") + yield # pragma: no cover + + events = [ + event + async for event in qwen3_tts_server._generate_audio_sse_events( + failing_generator(), + "req-sse-no-audio", + ) + ] + + error_event = next(event for event in events if "event: speech.audio.error" in event) + payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:]) + assert "partial_audio" not in payload["error"] + assert "action" not in payload["error"] + @pytest.mark.asyncio async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server): async def pcm_generator(): @@ -4171,8 +5083,14 @@ class TestTTSAsyncOffloading: def test_shutdown_is_idempotent(self, mocker: MockerFixture): """Calling shutdown() twice should not raise.""" - mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) - mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers", + return_value=set(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate", + return_value=None, + ) mock_engine_client = mocker.MagicMock() mock_engine_client.errored = False mock_engine_client.stage_configs = [] diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 37e3509f5..7c553af8f 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 from types import SimpleNamespace @@ -44,7 +47,9 @@ def _build_test_app( speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {})) speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): + async def mock_generate_pcm_chunks( + _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + ): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield (chunk, 24000) if include_sample_rate else chunk @@ -244,6 +249,7 @@ class TestStreamingSpeechWebSocket: def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): captured_requests = [] + captured_tts_params = [] speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) @@ -253,11 +259,12 @@ class TestStreamingSpeechWebSocket: async def mock_prepare_speech_generation(request): captured_requests.append(request) - return "req-stream", object(), {} + return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]} speech_service._prepare_speech_generation = mock_prepare_speech_generation - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + captured_tts_params.append(tts_params) for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"): yield (chunk, 24000) if include_sample_rate else chunk @@ -302,6 +309,7 @@ class TestStreamingSpeechWebSocket: assert captured_requests[0].stream is True assert captured_requests[0].response_format == "pcm" assert captured_requests[0].initial_codec_chunk_frames == 12 + assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}] assert speech_service._generate_audio_bytes.await_count == 0 def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture): @@ -344,7 +352,9 @@ class TestStreamingSpeechWebSocket: # The forced-aligner stage rides the same generator: its pooling output # is surfaced via the ``collect`` channel once the audio has streamed. - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): + async def mock_generate_pcm_chunks( + _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + ): for chunk in (first_chunk, second_chunk): yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: @@ -427,7 +437,9 @@ class TestStreamingSpeechWebSocket: speech_service.forced_aligner_enabled = True speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): + async def mock_generate_pcm_chunks( + _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + ): chunk = b"\x01" * 1000 yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: @@ -608,7 +620,7 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): yield b"\x01\x02" raise RuntimeError("stream boom") @@ -633,6 +645,8 @@ class TestStreamingSpeechWebSocket: assert ws.receive_json() == { "type": "error", "message": "Generation failed for utterance 0, sentence 0: stream boom", + "partial_audio": True, + "action": "discard", } assert ws.receive_json() == { "type": "audio.done", @@ -702,7 +716,7 @@ class TestStreamingSpeechWebSocket: speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): yield b"\x01\x02" speech_service._generate_pcm_chunks = mock_generate_pcm_chunks diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py index 27f1bed1e..e57e3e279 100644 --- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py @@ -47,7 +47,13 @@ def _build_test_app( engine_client.abort = mocker.AsyncMock() engine_client.submit_interaction_async = mocker.AsyncMock() engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()] - engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] + engine_client.stage_configs = [ + SimpleNamespace( + stage_type="diffusion", + final_output=True, + final_output_type="video", + ) + ] engine_client.generate = mock_generate class FakeStreamingVideoEncoder: @@ -72,10 +78,6 @@ def _build_test_app( "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", encoder_factory, ) - mocker.patch( - "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", - return_value="diffusion", - ) mocker.patch( "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", return_value=[OmniDiffusionSamplingParams()], @@ -398,7 +400,13 @@ def _build_handler_for_async_tests( engine_client.abort = mocker.AsyncMock() engine_client.submit_interaction_async = mocker.AsyncMock() engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()] - engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] + engine_client.stage_configs = [ + SimpleNamespace( + stage_type="diffusion", + final_output=True, + final_output_type="video", + ) + ] engine_client.generate = mock_generate chunks = streaming_chunks or [(b"mp4-chunk-0", True)] @@ -421,10 +429,6 @@ def _build_handler_for_async_tests( "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", encoder_factory, ) - mocker.patch( - "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", - return_value="diffusion", - ) mocker.patch( "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", return_value=[OmniDiffusionSamplingParams()], diff --git a/tests/entrypoints/openai_api/test_speech_usage.py b/tests/entrypoints/openai_api/test_speech_usage.py index 2fba10516..56b74c05b 100644 --- a/tests/entrypoints/openai_api/test_speech_usage.py +++ b/tests/entrypoints/openai_api/test_speech_usage.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for speech token-usage accounting (issue #4646). These exercise the pure usage logic without a model/server: input-token @@ -205,6 +206,14 @@ def test_output_tokens_from_stage_metrics(): assert acc.total() == 77 +def test_output_tokens_preserves_stage0_finish_reason(): + acc = SpeechOutputTokenCounter() + res = _final_res(stage0_out=128) + res.metrics["stage_metrics"]["0"]["finish_reason"] = "length" + acc.observe(res) + assert acc.stage0_finish_reason == "length" + + def test_output_tokens_takes_max_across_stages(): acc = SpeechOutputTokenCounter() acc.observe(_final_res(stage0_out=42, stage1_out=0)) diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 9ab744c83..03faacfcc 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the TTS serving adapter registry (RFC #4327). Pure-Python registry/resolution logic; no model or GPU resources are loaded. @@ -6,23 +7,36 @@ Pure-Python registry/resolution logic; no model or GPU resources are loaded. import asyncio from types import SimpleNamespace +from typing import Any import pytest +from vllm.sampling_params import SamplingParams +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest from vllm_omni.entrypoints.openai.tts_adapters import ( TTS_ADAPTER_REGISTRY, ARTTSAdapter, DiffusionTTSAdapter, + SpeechServingContext, all_tts_model_types, detect_tts_model_type, resolve_adapter, ) +from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import ( IndexTTS2Adapter, IndexTTS25Adapter, indextts2_conditioning_cache_salt, ) -from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter +from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import ( + MossTTSAdapter, + MossTTSNanoAdapter, +) +from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import ( + QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY, + Qwen3TTSAdapter, + Qwen3TTSCodecLimitError, +) from vllm_omni.model_executor.models.indextts2 import prompt_utils from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import ( INDEXTTS25_TOKENIZER_FILE, @@ -94,11 +108,114 @@ def test_all_adapters_are_ar_or_diffusion(): assert cls.backend in ("ar", "diffusion") +@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter]) +def test_moss_tts_applies_request_max_new_tokens(adapter_cls): + adapter = adapter_cls(SimpleNamespace(server=object())) + stage_defaults = [SimpleNamespace(max_tokens=4096)] + + overridden = adapter.apply_sampling_overrides( + stage_defaults, + SimpleNamespace(max_new_tokens=512), + ) + + assert overridden[0].max_tokens == 512 + assert stage_defaults[0].max_tokens == 4096 + + +def _build_moss_tts_request(adapter_cls, mocker, *, request_seed): + server = mocker.Mock() + server._build_moss_tts_params = mocker.AsyncMock(return_value={}) + adapter = adapter_cls(SpeechServingContext(server=server)) + request = OpenAICreateSpeechRequest(input="hello", seed=request_seed) + + return asyncio.run( + adapter.build( + request, + [SamplingParams(seed=42)], + has_inline_ref_audio=False, + ) + ) + + +# Full-family coverage pins the adapter contract; only Nano consumes this seed end to end today. +@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter]) +@pytest.mark.parametrize("request_seed", [0, 1234]) +def test_moss_tts_request_seed_overrides_stage_default(adapter_cls, request_seed, mocker): + prepared = _build_moss_tts_request( + adapter_cls, + mocker, + request_seed=request_seed, + ) + + assert prepared.tts_params["seed"] == [request_seed] + + +@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter]) +def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker): + prepared = _build_moss_tts_request( + adapter_cls, + mocker, + request_seed=None, + ) + + assert prepared.tts_params["seed"] == [42] + + def test_qwen3_tts_metadata(): assert Qwen3TTSAdapter.backend == "ar" assert issubclass(Qwen3TTSAdapter, ARTTSAdapter) +@pytest.mark.parametrize( + ("task_type", "text_tokens", "request_cap", "expected_cap"), + [ + ("Base", 0, None, 4096), + ("Base", 10, None, 192), + ("Base", 23, None, 276), + ("Base", 23, 128, 128), + ("Base", 23, 512, 512), + ("Base", 400, None, 4096), + ("CustomVoice", 10, None, 4096), + ("CustomVoice", 10, 256, 256), + ], +) +def test_qwen3_tts_applies_text_scaled_codec_safety_limit(task_type, text_tokens, request_cap, expected_cap, mocker): + server = mocker.Mock() + server._count_usage_text_tokens.return_value = text_tokens + adapter = Qwen3TTSAdapter(SpeechServingContext(server=server)) + stage_defaults = [SamplingParams(max_tokens=4096, min_tokens=2)] + request = OpenAICreateSpeechRequest( + input="test text", + task_type=task_type, + max_new_tokens=request_cap, + ) + prompt: dict[str, Any] = {"additional_information": {}} + + overridden = adapter.apply_sampling_overrides(stage_defaults, request, prompt) + + assert overridden[0].max_tokens == expected_cap + assert prompt["additional_information"][QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] == [expected_cap] + assert stage_defaults[0].max_tokens == 4096 + + +def test_qwen3_tts_rejects_only_length_finished_base_audio(mocker): + adapter = Qwen3TTSAdapter(SpeechServingContext(server=mocker.Mock())) + params = { + "task_type": ["Base"], + QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY: [192], + } + + # EOS at the exact budget is valid; the token count alone is not a + # sufficient failure signal. + adapter.validate_generation(params, stage0_finish_reason="stop", output_tokens=192) + adapter.validate_generation(params, stage0_finish_reason=None, output_tokens=192) + + # Some frontends expose 191 decoded frames for max_new_tokens=192. The + # engine terminal reason still makes this an unambiguous limit failure. + with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"): + adapter.validate_generation(params, stage0_finish_reason="length", output_tokens=191) + + def test_indextts_adapters_are_versioned(): assert resolve_adapter("indextts2") is IndexTTS2Adapter assert resolve_adapter("indextts2_5") is IndexTTS25Adapter @@ -162,6 +279,13 @@ def test_indextts25_speed_does_not_change_conditioning_cache_salt(): ) == indextts2_conditioning_cache_salt(slow_request, slow_params) +def test_indextts2_conditioning_cache_salt_changes_with_ref_audio_cache_key(): + request = SimpleNamespace(input="hello", ref_audio="file:///data/spk.wav") + salt_a = indextts2_conditioning_cache_salt(request, {"ref_audio_cache_key": ["key_aaa"]}) + salt_b = indextts2_conditioning_cache_salt(request, {"ref_audio_cache_key": ["key_bbb"]}) + assert salt_a != salt_b + + @pytest.mark.parametrize( ("hf_config", "expected_tokenizer_file"), [ @@ -216,5 +340,58 @@ def test_diffusion_adapter_extra_body_params_fallback(): assert _DiffAdapter.extra_body_params() == frozenset() +def _higgs_v2_adapter() -> HiggsAudioV2Adapter: + server = SimpleNamespace( + _apply_uploaded_speaker=lambda request: None, + uploaded_speakers={}, + ) + return HiggsAudioV2Adapter(SimpleNamespace(server=server)) + + +def _higgs_v2_request(**overrides: Any) -> SimpleNamespace: + fields: dict[str, Any] = { + "input": "Hello world.", + "ref_audio": None, + "ref_text": None, + "voice": None, + "x_vector_only_mode": None, + "speaker_embedding": None, + "instructions": None, + "task_type": None, + "language": None, + "speed": None, + "max_new_tokens": None, + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +@pytest.mark.parametrize( + "overrides, err_substr", + [ + pytest.param({"ref_audio": "data:audio/wav;base64,AA=="}, "ref_text", id="ref_audio_without_ref_text"), + pytest.param({"ref_text": "some transcript"}, "ref_audio", id="ref_text_without_ref_audio"), + pytest.param({"task_type": "Base"}, "task_type", id="task_type"), + pytest.param({"language": "Chinese"}, "language", id="language_override"), + pytest.param({"input": "[SPEAKER0] hi"}, "multi-speaker", id="multi_speaker_tag"), + pytest.param({"input": " "}, "empty", id="input_whitespace_only"), + ], +) +def test_higgs_audio_v2_validate_rejects_out_of_scope_fields(overrides: dict[str, object], err_substr: str) -> None: + """Adapter-only policy checks formerly covered by invalid_param e2e on a live V2 server.""" + adapter = _higgs_v2_adapter() + err = adapter.validate(_higgs_v2_request(**overrides)) + assert err is not None + assert err_substr.lower() in err.lower() + + +def test_higgs_audio_v2_validate_accepts_plain_text_and_paired_clone() -> None: + adapter = _higgs_v2_adapter() + assert adapter.validate(_higgs_v2_request()) is None + assert ( + adapter.validate(_higgs_v2_request(ref_audio="data:audio/wav;base64,AA==", ref_text="some transcript")) is None + ) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index c63e5c8e4..6b8b97849 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -14,7 +14,6 @@ import itertools import pytest from vllm_omni.entrypoints.openai.tts_adapters import ( - LEGACY_TTS_DETECTORS, TTS_ADAPTER_REGISTRY, all_tts_stage_keys, detect_tts_model_type, @@ -154,6 +153,7 @@ _PIPELINE_STAGES = [ "indextts2_talker", "latent_generator", "llm", + "minimax_music3_ar", "ming_tts", "moss_tts", "moss_tts_codec", @@ -282,27 +282,17 @@ def test_entry_stage_archs_is_ming_only(): assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"}) -def test_every_detected_type_is_adapter_backed_or_declared_legacy(): - """Detection may only name a model that has an adapter or an explicit - :data:`LEGACY_TTS_DETECTORS` entry — never an undeclared string.""" - legacy_names = {d.name for d in LEGACY_TTS_DETECTORS} +def test_every_detected_type_is_adapter_backed(): + """Detection may only name a model that has a registered adapter.""" for stage, arch in itertools.product(_STAGES, _ARCHS): # full cross product on purpose detected = detect_tts_model_type(stage, arch) if detected is None: continue - assert detected in TTS_ADAPTER_REGISTRY or detected in legacy_names, ( + assert detected in TTS_ADAPTER_REGISTRY, ( f"detection produced undeclared model type {detected!r} for stage={stage!r} arch={arch!r}" ) -def test_legacy_detectors_have_no_adapter(): - """A legacy entry that gained an adapter must be deleted from the list.""" - for detector in LEGACY_TTS_DETECTORS: - assert resolve_adapter(detector.name) is None, ( - f"{detector.name!r} now has an adapter; remove it from LEGACY_TTS_DETECTORS" - ) - - def test_no_ambiguous_detectors_at_equal_priority(): """Two detectors that can match the same deployment must be ordered. diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 08d223c68..9f643d8ca 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -1,17 +1,127 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for OpenAI-compatible video API encoding helpers.""" +import base64 +import threading +from io import BytesIO +from typing import Any + +import av +import httpx import numpy as np import pytest import torch +from PIL import Image +from vllm import envs from vllm_omni.diffusion.postprocess import rife_interpolator +from vllm_omni.diffusion.utils import media_utils from vllm_omni.entrypoints.openai import video_api_utils +from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _png_bytes() -> bytes: + buffer = BytesIO() + Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG") + return buffer.getvalue() + + +def _install_http_transport(monkeypatch, handler, client_kwargs): + async_client = httpx.AsyncClient + + def _client_factory(*args, **kwargs): + client_kwargs.append(kwargs.copy()) + return async_client(*args, transport=httpx.MockTransport(handler), **kwargs) + + monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory) + + +@pytest.mark.asyncio +async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch): + requested_paths = [] + client_kwargs: list[dict[str, object]] = [] + + def _handler(request): + requested_paths.append(request.url.path) + if request.url.path == "/redirect.png": + return httpx.Response(302, headers={"location": "/image.png"}) + return httpx.Response(200, content=_png_bytes(), headers={"content-type": "image/png"}) + + monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True) + _install_http_transport(monkeypatch, _handler, client_kwargs) + + image = await video_api_utils.decode_image_url("https://example.com/redirect.png") + + assert image.size == (2, 1) + assert image.mode == "RGB" + assert requested_paths == ["/redirect.png", "/image.png"] + assert client_kwargs == [{"timeout": 60, "follow_redirects": True}] + + +@pytest.mark.asyncio +async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch): + requested_paths = [] + client_kwargs: list[dict[str, object]] = [] + + def _handler(request): + requested_paths.append(request.url.path) + return httpx.Response(302, headers={"location": "/image.png"}) + + monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", False) + _install_http_transport(monkeypatch, _handler, client_kwargs) + + with pytest.raises(InvalidInputReferenceError, match="redirect.*VLLM_MEDIA_URL_ALLOW_REDIRECTS"): + await video_api_utils.decode_image_url("https://example.com/redirect.png") + + assert requested_paths == ["/redirect.png"] + assert client_kwargs == [{"timeout": 60, "follow_redirects": False}] + + +@pytest.mark.asyncio +async def test_decode_image_url_reports_http_status(monkeypatch): + client_kwargs: list[dict[str, object]] = [] + + def _handler(request): + return httpx.Response(404) + + monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True) + _install_http_transport(monkeypatch, _handler, client_kwargs) + + with pytest.raises(InvalidInputReferenceError, match="server returned HTTP 404"): + await video_api_utils.decode_image_url("https://example.com/missing.png") + + +@pytest.mark.asyncio +async def test_decode_image_url_reports_connection_failure(monkeypatch): + client_kwargs: list[dict[str, object]] = [] + + def _handler(request): + raise httpx.ConnectError("connection refused", request=request) + + monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True) + _install_http_transport(monkeypatch, _handler, client_kwargs) + + with pytest.raises(InvalidInputReferenceError, match="failed to download image"): + await video_api_utils.decode_image_url("https://example.com/unreachable.png") + + +@pytest.mark.asyncio +async def test_decode_image_url_keeps_data_urls_local(monkeypatch): + def _unexpected_http_client(*args, **kwargs): + pytest.fail("data URLs must not create an HTTP client") + + monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _unexpected_http_client) + encoded_image = base64.b64encode(_png_bytes()).decode() + + image = await video_api_utils.decode_image_url(f"data:image/png;base64,{encoded_image}") + + assert image.size == (2, 1) + assert image.mode == "RGB" + + def _install_fake_video_mux(monkeypatch, mux_calls): def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None): mux_calls.append( @@ -32,7 +142,7 @@ def _install_fake_video_mux(monkeypatch, mux_calls): def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): - mux_calls = [] + mux_calls: list[dict[str, Any]] = [] _install_fake_video_mux(monkeypatch, mux_calls) frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)] @@ -113,6 +223,686 @@ def test_channel_first_video_tensor_is_converted_to_channel_last(): np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy()) +def test_contiguous_uint8_video_reaches_mux_without_conversion(monkeypatch): + video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3) + + def fail_prepare(*args, **kwargs): + raise AssertionError("ready-to-encode uint8 video must not be normalized") + + monkeypatch.setattr(video_api_utils, "_prepare_video_frames", fail_prepare) + + frames = video_api_utils._coerce_video_to_uint8_frames(video) + + assert frames is video + + +def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch): + calls = [] + + def fake_planar_mux(frames, width, height, audio_waveform, **kwargs): + calls.append((list(frames), width, height, audio_waveform, kwargs)) + return b"planar-video" + + def fail_compat_mux(*args, **kwargs): + raise AssertionError("CFHW-backed FHWC frames should use the planar path") + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux) + monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fail_compat_mux) + + video = torch.linspace(-1, 1, 3 * 2 * 4 * 6, dtype=torch.float32).reshape(3, 2, 4, 6) + encoded = video_api_utils._encode_video_bytes( + video, + fps=12, + ) + + assert encoded == b"planar-video" + assert len(calls) == 1 + frames, width, height, audio, options = calls[0] + assert (len(frames), width, height) == (2, 6, 4) + assert all(frame.format.name == "gbrp" for frame in frames) + assert audio is None + assert options["fps"] == 12.0 + + +def test_planar_converter_serial_and_parallel_outputs_match(): + frames = [np.full((2, 3, 3), fill_value=value / 10, dtype=np.float32) for value in (1, 4, 7)] + + serial_converter = video_api_utils._PlanarFrameConverter(max_workers=1) + parallel_converter = video_api_utils._PlanarFrameConverter(max_workers=2) + try: + serial = [ + frame.to_ndarray(format="rgb24") for frame in serial_converter.iter_frames(frames, np.dtype(np.float32)) + ] + parallel = [ + frame.to_ndarray(format="rgb24") for frame in parallel_converter.iter_frames(frames, np.dtype(np.float32)) + ] + finally: + serial_converter.shutdown() + parallel_converter.shutdown() + + np.testing.assert_array_equal(np.stack(serial), np.stack(parallel)) + + +def test_planar_converter_preserves_fifo_order_and_bounds_pending_futures(monkeypatch): + converter = video_api_utils._PlanarFrameConverter(max_workers=2) + started = threading.Event() + release = threading.Event() + four_submitted = threading.Event() + outstanding = 0 + peak_outstanding = 0 + lock = threading.Lock() + real_submit = converter._executor.submit + + def fake_build(frame, _common_dtype): + started.set() + release.wait(timeout=2) + return int(frame[0, 0, 0]) + + def submit(*args, **kwargs): + nonlocal outstanding, peak_outstanding + with lock: + outstanding += 1 + peak_outstanding = max(peak_outstanding, outstanding) + if peak_outstanding >= 4: + four_submitted.set() + future = real_submit(*args, **kwargs) + + def done(_future): + nonlocal outstanding + with lock: + outstanding -= 1 + + future.add_done_callback(done) + return future + + monkeypatch.setattr(converter, "_build_frame", fake_build) + monkeypatch.setattr(converter._executor, "submit", submit) + frames = [np.full((1, 1, 3), value, dtype=np.uint8) for value in range(8)] + result: list[int] = [] + iterator = converter.iter_frames(frames, np.dtype(np.uint8)) + consumer = threading.Thread(target=lambda: result.extend(iterator)) + consumer.start() + try: + assert started.wait(timeout=2) + assert four_submitted.wait(timeout=2) + release.set() + consumer.join(timeout=2) + assert not consumer.is_alive() + finally: + release.set() + iterator.close() + converter.shutdown() + + assert result == list(range(8)) + assert peak_outstanding <= 4 + + +def test_planar_converter_reuses_executor_and_shutdowns_after_cancel(): + converter = video_api_utils._PlanarFrameConverter(max_workers=2) + executor = converter._executor + frames = [np.zeros((2, 2, 3), dtype=np.uint8) for _ in range(4)] + try: + list(converter.iter_frames(frames, np.dtype(np.uint8))) + iterator = converter.iter_frames(frames, np.dtype(np.uint8)) + next(iterator) + iterator.close() + assert converter._executor is executor + finally: + converter.shutdown() + + assert executor._shutdown + + +@pytest.mark.parametrize( + ("video", "expected_path", "expected_reason", "expected_mux"), + [ + ( + np.transpose(np.zeros((3, 2, 4, 4), dtype=np.float32), (1, 2, 3, 0)), + "direct_planar", + None, + "direct", + ), + ( + np.zeros((2, 4, 6, 3), dtype=np.float32), + "legacy_fallback", + "non_contiguous_rgb_planes", + "legacy", + ), + ], + ids=["automatic-direct", "interleaved-legacy-fallback"], +) +def test_audio_sample_rate_defaults_and_path_log(monkeypatch, video, expected_path, expected_reason, expected_mux): + mux_calls = [] + path_logs = [] + default_rate_logs = [] + monkeypatch.setattr( + video_api_utils.logger, + "info", + lambda message, *args, **kwargs: path_logs.append(message % args), + ) + monkeypatch.setattr( + video_api_utils.logger, + "info_once", + lambda message, *args, **kwargs: default_rate_logs.append(message % args), + ) + + def fake_planar_mux(frames, width, height, audio_waveform, **kwargs): + mux_calls.append({"kind": "direct", "audio_sample_rate": kwargs["audio_sample_rate"]}) + return b"planar-video" + + def fake_compat_mux(frames, audio, **kwargs): + mux_calls.append({"kind": "legacy", "audio_sample_rate": kwargs["audio_sample_rate"]}) + return b"legacy-video" + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux) + monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux) + + encoded = video_api_utils._encode_video_bytes( + video, + fps=12, + audio=np.zeros((1, 8), dtype=np.float32), + ) + + assert encoded == (b"planar-video" if expected_mux == "direct" else b"legacy-video") + assert mux_calls == [{"kind": expected_mux, "audio_sample_rate": 24000}] + assert len(default_rate_logs) == 1 + assert "default sample rate of 24000 Hz" in default_rate_logs[0] + assert len(path_logs) == 1 + assert f"selected_path={expected_path}" in path_logs[0] + if expected_reason is not None: + assert f"reason={expected_reason}" in path_logs[0] + assert "frames=2" in path_logs[0] + expected_shape = "(4, 4, 3)" if expected_mux == "direct" else "(4, 6, 3)" + assert f"frame_shape={expected_shape}" in path_logs[0] + assert "dtype=float32" in path_logs[0] + assert "fps=12" in path_logs[0] + assert "audio_present=True" in path_logs[0] + assert "effective_audio_sample_rate=24000" in path_logs[0] + + +def test_explicit_audio_sample_rate_bypasses_default_log(monkeypatch): + log_messages = [] + monkeypatch.setattr( + video_api_utils.logger, + "info_once", + lambda message, *args, **kwargs: log_messages.append(message % args), + ) + + audio = [[0.0, 1.0], [1.0, 0.0]] + assert video_api_utils._resolve_audio_sample_rate(audio, 32000) == 32000 + assert log_messages == [] + + +def test_video_only_does_not_log_default_audio_sample_rate(monkeypatch): + log_messages = [] + monkeypatch.setattr( + video_api_utils.logger, + "info_once", + lambda message, *args, **kwargs: log_messages.append(message % args), + ) + monkeypatch.setattr( + media_utils, + "mux_video_audio_bytes", + lambda frames, audio, **kwargs: b"legacy-video", + ) + + video_api_utils._encode_video_bytes(np.zeros((1, 2, 2, 3), dtype=np.uint8), fps=12) + + assert log_messages == [] + + +def test_planar_mux_accepts_none_rate_and_defaults_audio_to_44100(): + def make_frame(): + return next( + video_api_utils._iter_planar_video_frames( + [np.zeros((2, 2, 3), dtype=np.uint8)], + np.dtype(np.uint8), + ) + ) + + video_only = media_utils.mux_av_video_audio_bytes( + [make_frame()], + width=2, + height=2, + audio_sample_rate=None, + ) + with av.open(BytesIO(video_only), mode="r", format="mp4") as container: + assert not container.streams.audio + + with_audio = media_utils.mux_av_video_audio_bytes( + [make_frame()], + width=2, + height=2, + audio_waveform=np.zeros((1, 1024), dtype=np.float32), + audio_sample_rate=None, + ) + with av.open(BytesIO(with_audio), mode="r", format="mp4") as container: + assert container.streams.audio[0].rate == 44100 + + +def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch): + calls = [] + path_logs = [] + coerce_calls = [] + + real_coerce_video_to_frames = video_api_utils._coerce_video_to_frames + + def count_coerce_video_to_frames(video): + coerce_calls.append(video) + return real_coerce_video_to_frames(video) + + def fail_planar_mux(*args, **kwargs): + raise AssertionError("interleaved FHWC frames must use the legacy path") + + def fake_compat_mux(frames, audio, **kwargs): + calls.append((frames, audio, kwargs)) + return b"legacy-video" + + monkeypatch.setattr( + video_api_utils.logger, "info", lambda message, *args, **kwargs: path_logs.append(message % args) + ) + monkeypatch.setattr(video_api_utils, "_coerce_video_to_frames", count_coerce_video_to_frames) + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fail_planar_mux) + monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux) + + video = np.zeros((2, 4, 6, 3), dtype=np.float32) + encoded = video_api_utils._encode_video_bytes(video, fps=12) + + assert encoded == b"legacy-video" + assert calls[0][0].shape == (2, 4, 6, 3) + assert calls[0][0].dtype == np.uint8 + assert len(path_logs) == 1 + assert "selected_path=legacy_fallback" in path_logs[0] + assert "reason=non_contiguous_rgb_planes" in path_logs[0] + assert "audio_present=False" in path_logs[0] + assert "effective_audio_sample_rate=None" in path_logs[0] + assert "effective_frame_conversion_workers=0" in path_logs[0] + assert len(coerce_calls) == 1 + + +def test_uint8_frames_reach_the_muxer_without_a_float_round_trip(): + """uint8 is the muxer's own dtype, so normalisation must leave it alone. + + MiniMax-H3 quantises on the accelerator to keep the response small across + its two process hops. Re-normalising to float here would undo that at the + last step and cost a full-size conversion in each direction. + """ + video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3) + + frames, frame_shape, common_dtype = video_api_utils._prepare_video_frames(video) + + assert common_dtype == np.dtype(np.uint8) + assert frame_shape == (4, 5, 3) + np.testing.assert_array_equal(np.stack(frames), video) + + +def test_non_uint8_integer_frames_are_still_normalised(): + """Only uint8 is the muxer's dtype; other integer payloads still scale.""" + video = np.full((2, 4, 5, 3), 255, dtype=np.int32) + + frames, _, common_dtype = video_api_utils._prepare_video_frames(video) + + assert np.issubdtype(common_dtype, np.floating) + np.testing.assert_allclose(frames[0], 1.0) + + +@pytest.mark.parametrize( + "video", + [ + np.array( + [[[[0, 1, 127], [128, 254, 255]], [[255, 127, 1], [2, 3, 4]]]], + dtype=np.uint8, + ), + np.array( + [[[[-1.0, -0.5, 0.0], [0.5, 1.0, 2.0]], [[1.5, -2.0, 0.25], [0.75, 0.125, 0.875]]]], + dtype=np.float32, + ), + np.array( + [[[[-512, -1, 0], [128, 255, 512]], [[1024, -256, 64], [127, 300, -300]]]], + dtype=np.int16, + ), + ], + ids=["uint8", "float-bounded-and-out-of-range", "signed-int-out-of-range"], +) +def test_prepared_automatic_fallback_preserves_legacy_quantization(monkeypatch, video): + mux_inputs = [] + + def fail_planar_mux(*args, **kwargs): + raise AssertionError("interleaved frames must use the legacy mux") + + def fake_compat_mux(frames, audio, **kwargs): + mux_inputs.append(frames.copy()) + return b"legacy-video" + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fail_planar_mux) + monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux) + + prepared_frames, frame_shape, _ = video_api_utils._prepare_video_frames(video) + stacked = np.stack([frame[..., :3] for frame in prepared_frames]) + # uint8 frames are the muxer's own dtype and reach it unchanged; everything + # else is still quantised out of the normalised [0, 1] range. + reference = stacked if stacked.dtype == np.uint8 else np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8) + + assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video" + assert len(mux_inputs) == 1 + assert mux_inputs[0].dtype == np.uint8 + assert mux_inputs[0].shape == (len(prepared_frames), *frame_shape[:-1], 3) + np.testing.assert_array_equal(mux_inputs[0], reference) + + +@pytest.mark.parametrize( + ("worker_count", "expected_path", "expected_reason"), + [ + (None, "legacy_fallback", "non_contiguous_rgb_planes"), + (1, "legacy_fallback", "non_contiguous_rgb_planes"), + (8, "direct_planar", None), + ], +) +def test_interleaved_video_routing(monkeypatch, worker_count, expected_path, expected_reason): + routes = [] + + def fake_planar_mux(*args, **kwargs): + return b"planar-video" + + def fake_compat_mux(frames, audio, **kwargs): + return b"legacy-video" + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux) + monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux) + monkeypatch.setattr(video_api_utils, "_log_video_encoding_path", lambda **kwargs: routes.append(kwargs)) + + converter = None if worker_count is None else video_api_utils._PlanarFrameConverter(max_workers=worker_count) + try: + encoded = video_api_utils._encode_video_bytes( + np.zeros((2, 4, 6, 3), dtype=np.float32), + fps=12, + frame_converter=converter, + ) + finally: + if converter is not None: + converter.shutdown() + + assert encoded == (b"planar-video" if expected_path == "direct_planar" else b"legacy-video") + assert len(routes) == 1 + assert routes[0]["selected_path"] == expected_path + assert routes[0].get("reason") == expected_reason + + +@pytest.mark.parametrize( + ("video", "reason"), + [ + (np.zeros((2, 4, 6), dtype=np.float32), "unsupported_shape"), + (np.empty((2, 0, 6, 3), dtype=np.float32), "unsupported_shape"), + (np.empty((2, 4, 0, 3), dtype=np.float32), "unsupported_shape"), + (np.zeros((2, 2, 2, 3), dtype=np.complex64), "unsupported_dtype"), + (np.zeros((2, 2, 2, 3), dtype=object), "unsupported_dtype"), + (np.zeros((2, 2, 2, 3), dtype="U1"), "unsupported_dtype"), + (np.zeros((2, 2, 2, 3), dtype="datetime64[D]"), "unsupported_dtype"), + ], +) +def test_unsupported_direct_capability_uses_legacy_with_reason(monkeypatch, video, reason): + events = [] + + def fake_legacy(*args, **kwargs): + events.append("legacy") + return b"legacy-video" + + monkeypatch.setattr(video_api_utils, "_encode_prepared_video_bytes_legacy", fake_legacy) + monkeypatch.setattr( + video_api_utils.logger, + "info", + lambda message, *args, **kwargs: events.append(message % args), + ) + + assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video" + assert events[-1] == "legacy" + assert f"selected_path=legacy_fallback reason={reason}" in events[0] + assert "frames=" in events[0] + assert "fps=12" in events[0] + assert "audio_present=False" in events[0] + assert "effective_audio_sample_rate=None" in events[0] + + +def test_non_contiguous_fallback_is_logged_before_legacy_mux(monkeypatch): + events = [] + video = np.zeros((2, 4, 6, 3), dtype=np.float32) + + def fake_legacy(*args, **kwargs): + events.append("legacy") + return b"legacy-video" + + monkeypatch.setattr(video_api_utils, "_encode_prepared_video_bytes_legacy", fake_legacy) + monkeypatch.setattr( + video_api_utils.logger, + "info", + lambda message, *args, **kwargs: events.append(message % args), + ) + + assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video" + assert events[0].startswith("Video response encoding route selected: selected_path=legacy_fallback") + assert events[1] == "legacy" + + +def test_direct_mux_error_is_propagated_without_legacy_retry(monkeypatch): + legacy_calls = [] + + def fail_direct(*args, **kwargs): + raise RuntimeError("direct mux failed") + + def fail_legacy(*args, **kwargs): + legacy_calls.append((args, kwargs)) + raise AssertionError("direct failures must not retry through legacy") + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fail_direct) + monkeypatch.setattr(video_api_utils, "_encode_prepared_video_bytes_legacy", fail_legacy) + planar = np.transpose(np.zeros((3, 2, 4, 4), dtype=np.float32), (1, 2, 3, 0)) + + with pytest.raises(RuntimeError, match="direct mux failed"): + video_api_utils._encode_video_bytes(planar, fps=12) + + assert legacy_calls == [] + + +@pytest.mark.parametrize("dtype", [np.float16, np.float32]) +@pytest.mark.parametrize("channels", [3, 4]) +def test_planar_frames_decode_with_bounded_rgb_quantization(dtype, channels): + channel_values = np.array( + [ + [[-0.1, 0.0, 0.1, 0.5, 0.9, 1.0, 1.1]], + [[1.0, 0.9, 0.5, 0.1, 0.0, -0.1, 1.1]], + [[0.25, 0.75, 0.501, 0.499, 0.125, 0.875, 0.5]], + [[0.2, 0.4, 0.6, 0.8, 1.0, 0.0, 0.5]], + ], + dtype=dtype, + )[:channels] + fhwc = np.transpose(channel_values[:, None, :, :], (1, 2, 3, 0)) + frames = list(video_api_utils._iter_planar_video_frames(list(fhwc), np.dtype(dtype))) + + decoded = frames[0].to_ndarray(format="rgb24") + expected = np.rint(np.clip(fhwc[0, ..., :3], 0.0, 1.0) * 255.0).astype(np.uint8) + np.testing.assert_array_equal(decoded, expected) + + +def test_planar_frame_writes_gbr_planes_and_clears_padding(): + width = 17 + planar = np.stack( + [ + np.full((2, width), 11, dtype=np.uint8), + np.full((2, width), 22, dtype=np.uint8), + np.full((2, width), 33, dtype=np.uint8), + ] + ) + fhwc = np.transpose(planar[:, None, :, :], (1, 2, 3, 0)) + + frame = next(video_api_utils._iter_planar_video_frames(list(fhwc), np.dtype(np.uint8))) + + for plane, expected in zip(frame.planes, (22, 33, 11)): + plane_data = np.frombuffer(plane, dtype=np.uint8).reshape(plane.height, plane.line_size) + np.testing.assert_array_equal(plane_data[:2, :width], expected) + np.testing.assert_array_equal(plane_data[:2, width:], 0) + np.testing.assert_array_equal(plane_data[2:], 0) + + +def test_planar_bool_frames_match_bounded_compatible_output(): + planar = np.array( + [ + [[[False, True], [True, False]]], + [[[True, False], [True, False]]], + [[[False, False], [True, True]]], + ], + dtype=np.bool_, + ) + direct_video = np.transpose(planar, (1, 2, 3, 0)) + compatible_video = np.ascontiguousarray(direct_video) + + prepared_frames, _, common_dtype = video_api_utils._prepare_video_frames(direct_video) + direct_frames = list(video_api_utils._iter_planar_video_frames(prepared_frames, common_dtype)) + decoded_direct = np.stack([frame.to_ndarray(format="rgb24") for frame in direct_frames]) + bounded_compatible = video_api_utils._coerce_video_to_uint8_frames(compatible_video) + + np.testing.assert_array_equal(decoded_direct, bounded_compatible) + + +@pytest.mark.parametrize( + ("plane_height", "line_size"), + [(1, 2), (2, 1)], +) +def test_planar_frame_rejects_undersized_plane(monkeypatch, plane_height, line_size): + class FakePlane: + height: int + line_size: int + + plane = FakePlane() + plane.height = plane_height + plane.line_size = line_size + + class FakeFrame: + planes = [plane] + + monkeypatch.setattr(av, "VideoFrame", lambda *args, **kwargs: FakeFrame()) + planar = np.zeros((3, 1, 2, 2), dtype=np.uint8) + fhwc = np.transpose(planar, (1, 2, 3, 0)) + + with pytest.raises(ValueError, match="smaller than"): + next(video_api_utils._iter_planar_video_frames(list(fhwc), np.dtype(np.uint8))) + + +def test_unequal_frame_shapes_fail_before_output_allocation(monkeypatch): + def fail_empty(*args, **kwargs): + raise AssertionError("output allocation must happen after shape validation") + + monkeypatch.setattr(video_api_utils.np, "empty", fail_empty) + + with pytest.raises(ValueError, match="same shape"): + video_api_utils._encode_video_bytes( + [np.zeros((2, 3, 3), dtype=np.float32), np.zeros((3, 3, 3), dtype=np.float32)], + fps=12, + ) + + +@pytest.mark.parametrize("with_audio", [False, True]) +def test_direct_and_compatible_paths_produce_equivalent_mp4(with_audio): + rng = np.random.default_rng(7) + planar = rng.random((3, 3, 32, 32), dtype=np.float32) + direct_video = np.transpose(planar, (1, 2, 3, 0)) + compatible_video = np.ascontiguousarray(direct_video) + audio = np.zeros((2, 2400), dtype=np.float32) if with_audio else None + + direct_bytes = video_api_utils._encode_video_bytes( + direct_video, + fps=12, + audio=audio, + audio_sample_rate=24000, + video_codec_options={"preset": "ultrafast", "threads": "1"}, + ) + compatible_bytes = video_api_utils._encode_video_bytes( + compatible_video, + fps=12, + audio=audio, + audio_sample_rate=24000, + video_codec_options={"preset": "ultrafast", "threads": "1"}, + ) + + def decode(mp4_bytes): + with av.open(BytesIO(mp4_bytes), mode="r", format="mp4") as container: + stream_types = [stream.type for stream in container.streams] + video_stream = container.streams.video[0] + frames = np.stack([frame.to_ndarray(format="rgb24") for frame in container.decode(video_stream)]) + return stream_types, frames + + direct_streams, direct_frames = decode(direct_bytes) + compatible_streams, compatible_frames = decode(compatible_bytes) + assert direct_bytes == compatible_bytes + assert direct_streams == compatible_streams + assert ("audio" in direct_streams) is with_audio + np.testing.assert_array_equal(direct_frames, compatible_frames) + + +@pytest.mark.parametrize("with_audio", [False, True]) +def test_parallel_interleaved_and_legacy_paths_produce_equivalent_mp4(with_audio): + rng = np.random.default_rng(7) + video = rng.random((3, 32, 32, 3), dtype=np.float32) + audio = np.zeros((2, 2400), dtype=np.float32) if with_audio else None + kwargs = { + "fps": 12, + "audio": audio, + "audio_sample_rate": 24000, + "video_codec_options": {"preset": "ultrafast", "threads": "1"}, + } + + converter = video_api_utils._PlanarFrameConverter(max_workers=8) + try: + parallel_bytes = video_api_utils._encode_video_bytes( + video, + **kwargs, + frame_converter=converter, + ) + finally: + converter.shutdown() + legacy_bytes = video_api_utils._encode_video_bytes_legacy(video, **kwargs) + + assert parallel_bytes == legacy_bytes + + +def test_mux_closes_container_and_preserves_generator_error(monkeypatch): + class GeneratorError(RuntimeError): + pass + + class FakeStream: + options = None + + def encode(self, frame=None): + return [] + + class FakeContainer: + def __init__(self): + self.closed = False + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.closed = True + return False + + def add_stream(self, codec, rate): + return FakeStream() + + def mux(self, packet): + pass + + container = FakeContainer() + monkeypatch.setattr(media_utils.av, "open", lambda *args, **kwargs: container) + + def failing_frames(): + yield object() + raise GeneratorError("frame generation failed") + + with pytest.raises(GeneratorError, match="frame generation failed"): + media_utils.mux_av_video_audio_bytes(failing_frames(), width=2, height=2) + assert container.closed + + def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch): muxers = [] diff --git a/tests/entrypoints/openai_api/test_video_pipeline_capability.py b/tests/entrypoints/openai_api/test_video_pipeline_capability.py new file mode 100644 index 000000000..cfa11f96d --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_pipeline_capability.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.diffusion.io_support import get_diffusion_output_type +from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_video_pipeline_requires_declared_final_video_stage(): + assert is_video_generation_pipeline( + [ + SimpleNamespace( + stage_type="llm", + final_output=False, + final_output_type=None, + ), + SimpleNamespace( + stage_type="diffusion", + final_output=True, + final_output_type="video", + ), + ] + ) + + +@pytest.mark.parametrize( + "stage_configs", + [ + [SimpleNamespace(stage_type="diffusion")], + [ + SimpleNamespace( + stage_type="diffusion", + final_output=True, + final_output_type="image", + ) + ], + [ + { + "stage_type": "diffusion", + "final_output": False, + "final_output_type": "video", + } + ], + ], +) +def test_video_pipeline_rejects_non_video_final_outputs(stage_configs): + assert not is_video_generation_pipeline(stage_configs) + + +@pytest.mark.parametrize( + "model_class_name", + [ + "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "WanDMDPipeline", + "LingBotWorldCausalDMDPipeline", + "LongCatVideoAvatarPipeline", + ], +) +def test_registered_video_aliases_declare_video_output(model_class_name): + assert get_diffusion_output_type(model_class_name) == "video" diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index bcda12e2d..6d8977c55 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -16,15 +16,17 @@ from pathlib import Path from types import SimpleNamespace import av +import httpx import numpy as np import pytest from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from PIL import Image from pytest_mock import MockerFixture +from vllm import envs from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes -from vllm_omni.entrypoints.openai import api_server +from vllm_omni.entrypoints.openai import api_server, video_api_utils from vllm_omni.entrypoints.openai.api_server import router from vllm_omni.entrypoints.openai.protocol.videos import ( VideoGenerationRequest, @@ -32,7 +34,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoParams, VideoResponse, ) -from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo +from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage from vllm_omni.entrypoints.openai.storage import LocalStorageManager from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry from vllm_omni.errors import GuardrailViolationError @@ -62,7 +64,13 @@ class MockVideoResult: class FakeAsyncOmni: def __init__(self): - self.stage_configs = [SimpleNamespace(stage_type="diffusion")] + self.stage_configs = [ + SimpleNamespace( + stage_type="diffusion", + final_output=True, + final_output_type="video", + ) + ] self.default_sampling_params_list = [OmniDiffusionSamplingParams()] self.model_class_name = "WanPipeline" self.captured_prompt = None @@ -87,6 +95,58 @@ class FakeAsyncOmni: yield MockVideoResult(videos) +def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixture): + engine = FakeAsyncOmni() + handler = OmniOpenAIServingVideo.for_diffusion( + engine, + model_name="test-model", + ) + assert handler._video_frame_converter.max_workers == 8 + raw_encoder = mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"encoded-video", + ) + base64_encoder = mocker.patch( + "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", + return_value="encoded-video", + ) + + async def _generate_both_response_types(): + request = VideoGenerationRequest(prompt="test prompt") + await handler.generate_video_bytes(request, "raw-request") + await handler.generate_videos(request, "base64-request") + + asyncio.run(_generate_both_response_types()) + + assert raw_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter + assert base64_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter + handler.shutdown() + + +def test_resolve_diffusion_od_config_falls_back_to_attribute(): + od_config = SimpleNamespace(model_class_name="WanPipeline") + handler = OmniOpenAIServingVideo.for_diffusion( + SimpleNamespace(od_config=od_config), + model_name="test-model", + ) + + assert handler._resolve_diffusion_od_config() is od_config + + +def test_resolve_diffusion_od_config_prefers_getter_over_attribute(): + attribute_config = SimpleNamespace(model_class_name="WanPipeline") + getter_config = SimpleNamespace(model_class_name="MiniMaxH3Pipeline") + handler = OmniOpenAIServingVideo.for_diffusion( + SimpleNamespace( + od_config=attribute_config, + get_diffusion_od_config=lambda: getter_config, + ), + model_name="test-model", + ) + + assert handler._resolve_diffusion_od_config() is getter_config + + class BlockingVideoHandler: def __init__(self): self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" @@ -461,6 +521,51 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client, assert input_image.size == (96, 64) +def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): + engine = FakeAsyncOmni() + engine.get_diffusion_od_config = lambda: SimpleNamespace( + model="org/model", + model_class_name="ExamplePipeline", + revision="pinned-revision", + ) + captured = {} + + def fake_policy(model_class_name, *, model, revision=None): + captured.update( + model_class_name=model_class_name, + model=model, + revision=revision, + ) + return True + + monkeypatch.setattr( + "vllm_omni.entrypoints.openai.serving_video.should_preserve_reference_image_size", + fake_policy, + ) + handler = OmniOpenAIServingVideo.for_diffusion( + diffusion_engine=engine, + model_name="fallback/model", + ) + image = Image.new("RGB", (48, 32)) + + asyncio.run( + handler._run_and_extract( + VideoGenerationRequest(prompt="A bear playing with yarn.", width=96, height=64), + "pipeline-owned-resize", + reference_image=ReferenceImage(image), + ) + ) + + input_image = engine.captured_prompt["multi_modal_data"]["image"] + assert isinstance(input_image, Image.Image) + assert input_image.size == (48, 32) + assert captured == { + "model_class_name": "ExamplePipeline", + "model": "org/model", + "revision": "pinned-revision", + } + + def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture): image_bytes = _make_test_image_bytes((48, 48)) mocker.patch( @@ -554,6 +659,45 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc assert input_image.size == (40, 24) +def test_i2v_video_generation_follows_allowed_image_redirect(test_client, mocker: MockerFixture, monkeypatch): + requested_paths = [] + async_client = httpx.AsyncClient + + def _handler(request): + requested_paths.append(request.url.path) + if request.url.path == "/redirect.png": + return httpx.Response(302, headers={"location": "/image.png"}) + return httpx.Response(200, content=_make_test_image_bytes((40, 24))) + + def _client_factory(*args, **kwargs): + return async_client(*args, transport=httpx.MockTransport(_handler), **kwargs) + + monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True) + monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + + response = test_client.post( + "/v1/videos", + data={ + "prompt": "A fox running through snow.", + "image_reference": json.dumps({"image_url": "https://example.com/redirect.png"}), + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + assert requested_paths == ["/redirect.png", "/image.png"] + + engine = test_client.app.state.openai_serving_video._engine_client + input_image = engine.captured_prompt["multi_modal_data"]["image"] + assert isinstance(input_image, Image.Image) + assert input_image.size == (40, 24) + + def test_v2v_video_generation_form(test_client, mocker: MockerFixture): video_bytes = _make_test_video_bytes((32, 24), num_frames=3) @@ -692,6 +836,18 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation( assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0 +def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_false(test_client): + handler = test_client.app.state.openai_serving_video + handler._engine_client.model_class_name = None + handler._engine_client.stage_configs = [ + SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3TextEncoder"}), + SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3Pipeline"}), + ] + handler._stage_configs = handler._engine_client.stage_configs + + assert handler.supports_mixed_reference_inputs + + def test_decode_video_bytes_can_keep_first_frames(): from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes @@ -1071,8 +1227,15 @@ def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture): audio_sample_rates = [] - def _fake_encode(video, fps, audio=None, audio_sample_rate=None, video_codec_options=None): - del video, fps, audio, video_codec_options + def _fake_encode( + video, + fps, + audio=None, + audio_sample_rate=None, + video_codec_options=None, + frame_converter=None, + ): + del video, fps, audio, video_codec_options, frame_converter audio_sample_rates.append(audio_sample_rate) return b"fake-video" @@ -2082,6 +2245,31 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture): assert captured.quality == "high" +def test_sync_sana_wm_extra_params_payload_passes_to_engine_prompt(test_client, mocker: MockerFixture): + _mock_encode_video_bytes(mocker) + response = test_client.post( + "/v1/videos/sync", + data={ + "prompt": "drive forward", + "extra_params": json.dumps( + { + "sana_wm": {"action": "d-4", "rotation_speed_deg": 1.5}, + "sana_wm_native_max_tokens": 30000, + } + ), + }, + ) + + assert response.status_code == 200 + engine = test_client.app.state.openai_serving_video._engine_client + captured_params = engine.captured_sampling_params_list[0] + # The camera block reaches the model through extra_args; the Sana preprocess + # hook is what lifts it onto the prompt. + assert captured_params.extra_args["sana_wm"]["action"] == "d-4" + assert captured_params.extra_args["sana_wm"]["rotation_speed_deg"] == 1.5 + assert captured_params.extra_args["sana_wm_native_max_tokens"] == 30000 + + def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture): """Frame interpolation parameters should be forwarded on the sync path.""" encode_mock = _mock_encode_video_bytes(mocker) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 6120d5ea4..2549889b3 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import re from types import SimpleNamespace @@ -8,6 +11,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.outputs import OmniRequestOutput @@ -77,7 +81,7 @@ def test_generate_submits_randomized_id_to_engine(): """Ensure the engine receives a UUID-suffixed ID, not the raw request ID""" async def run(): - submitted_ids = [] + submitted_ids: list[str] = [] omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids)) req_id = "my-req-1" @@ -106,8 +110,8 @@ def test_generate_forwards_lora_request_to_engine(): """ async def run(): - submitted_ids = [] - submitted_loras = [] + submitted_ids: list[str] = [] + submitted_loras: list[LoRARequest | None] = [] omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids, submitted_loras)) lora = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake") @@ -143,7 +147,7 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi user provided request ID will be aborted.""" async def run(): - aborted_batches = [] + aborted_batches: list[list[str]] = [] omni = get_async_omni_instance( fake_abort_request=get_fake_abort(aborted_batches), ) @@ -176,7 +180,6 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid) assert len(aborted_ids) == expected_cancel_count assert len(set(aborted_ids)) == expected_cancel_count - assert len(omni.request_states) == len(req_ids) - expected_cancel_count for t in tasks: t.cancel() @@ -185,11 +188,152 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi asyncio.run(run()) +@pytest.mark.cpu +def test_abort_keeps_request_states_until_generate_cleanup(): + """Frontend abort must not drop request_states; generate() owns cleanup.""" + + async def run(): + release = asyncio.Event() + seen_during_wait: list[int] = [] + + async def slow_abort_async(request_ids): + del request_ids + seen_during_wait.append(len(omni.request_states)) + await release.wait() + + omni = get_async_omni_instance(fake_abort_request=slow_abort_async) + omni.request_states["req-1-aaaa"] = SimpleNamespace( + request_id="req-1-aaaa", + external_request_id="req-1", + input_stream_task=None, + ) + + task = asyncio.create_task(omni.abort("req-1")) + await asyncio.sleep(0) + assert not task.done() + assert "req-1-aaaa" in omni.request_states + assert seen_during_wait == [1] + + release.set() + await task + assert "req-1-aaaa" in omni.request_states + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_propagates_engine_errors_without_popping_state(): + async def run(): + async def failing_abort_async(request_ids): + del request_ids + raise RuntimeError("orchestrator abort failed") + + omni = get_async_omni_instance(fake_abort_request=failing_abort_async) + omni.request_states["req-1-bbbb"] = SimpleNamespace( + request_id="req-1-bbbb", + external_request_id="req-1", + input_stream_task=None, + ) + + with pytest.raises(RuntimeError, match="orchestrator abort failed"): + await omni.abort("req-1") + assert "req-1-bbbb" in omni.request_states + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_enqueues_prefix_tokens_from_engine(): + """AR abort must deliver the generated prefix on the request queue (CPU).""" + + async def run(): + from vllm.outputs import CompletionOutput + + from vllm_omni.engine.messages import OutputMessage + + queue: asyncio.Queue = asyncio.Queue() + + async def abort_with_prefix(request_ids): + rid = request_ids[0] + engine_output = OmniRequestOutput( + request_id=rid, + finished=True, + stage_id=0, + final_output_type="text", + outputs=[ + CompletionOutput( + index=0, + text="partial", + token_ids=[7, 8, 9], + cumulative_logprob=None, + logprobs=None, + finish_reason="abort", + stop_reason=None, + ) + ], + ) + return [ + OutputMessage( + request_id=rid, + stage_id=0, + replica_id=None, + engine_outputs=engine_output, + metrics=None, + finished=True, + stage_submit_ts=None, + ) + ] + + omni = get_async_omni_instance(fake_abort_request=abort_with_prefix) + omni.request_states["req-1-cccc"] = SimpleNamespace( + request_id="req-1-cccc", + external_request_id="req-1", + input_stream_task=None, + queue=queue, + ) + await omni._abort(["req-1-cccc"]) + msg = queue.get_nowait() + assert msg.finished is True + assert list(msg.engine_outputs.outputs[0].token_ids) == [7, 8, 9] + assert msg.engine_outputs.outputs[0].finish_reason == "abort" + assert "req-1-cccc" in omni.request_states + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_enqueues_synthetic_finished_when_engine_returns_empty(): + """Empty abort_async must still unblock generate() with finish_reason=abort.""" + + async def run(): + queue: asyncio.Queue = asyncio.Queue() + + async def empty_abort_async(request_ids): + del request_ids + return [] + + omni = get_async_omni_instance(fake_abort_request=empty_abort_async) + omni.request_states["req-1-dddd"] = SimpleNamespace( + request_id="req-1-dddd", + external_request_id="req-1", + input_stream_task=None, + queue=queue, + ) + await omni._abort(["req-1-dddd"]) + msg = queue.get_nowait() + assert msg.finished is True + assert msg.request_id == "req-1-dddd" + assert msg.engine_outputs.outputs[0].finish_reason == "abort" + assert "req-1-dddd" in omni.request_states + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_accepts_request_after_repeated_cancellations(): async def run_test(): - submitted_request_ids = [] - aborted_request_batches = [] + submitted_request_ids: list[str] = [] + aborted_request_batches: list[list[str]] = [] async def collect_outputs(request_id): outputs = [] @@ -234,6 +378,43 @@ def test_generate_accepts_request_after_repeated_cancellations(): asyncio.run(run_test()) +@pytest.mark.cpu +def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker): + async def run_test(): + submitted_request_ids: list[str] = [] + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = True + engine.request_queue = mocker.Mock() + engine.num_stages = 1 + engine.add_request_async = get_fake_add_request(submitted_request_ids) + + omni = get_async_omni_instance() + omni.engine = engine + + async def collect_outputs(): + async for _ in omni.generate( + prompt={"prompt": "prompt"}, + request_id="cancel-shutdown", + sampling_params_list=[SimpleNamespace()], + output_modalities=["image"], + ): + pass + + task = asyncio.create_task(collect_outputs()) + await asyncio.sleep(0) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert len(submitted_request_ids) == 1 + assert submitted_request_ids[0].startswith("cancel-shutdown-") + assert omni.request_states == {} + engine.request_queue.sync_q.put.assert_not_called() + + asyncio.run(run_test()) + + @pytest.mark.cpu def test_generate_yields_streaming_diffusion_chunks_before_final(): """AsyncOmni.generate yields every intermediate diffusion chunk before the final one.""" @@ -364,7 +545,7 @@ async def test_omni_generate_request_id(): Text modality only; asserts caller-visible ids are preserved across the multi-stage orchestrator path on H100. """ - engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG) + engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG) try: for req_id in _OMNI_REQ_IDS: async for output in engine.generate( diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 2fa8c92bd..71a6e00a5 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace @@ -33,6 +33,77 @@ def test_default_stage_config_includes_cache_backend(): assert engine_args["model_stage"] == "diffusion" +def test_default_stage_config_preserves_ulysses_a2a_permute() -> None: + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "ulysses_degree": 4, + "ulysses_a2a_permute": True, + } + )[0] + + parallel_config = stage_cfg["engine_args"]["parallel_config"] + assert parallel_config.ulysses_degree == 4 + assert parallel_config.ulysses_a2a_permute is True + + +def test_default_stage_config_preserves_model_extras(): + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0] + + assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True + + +def test_default_stage_config_preserves_and_overrides_promoted_extras(): + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "extras": { + "auxiliary_text_encoder": "/models/extras-llama", + "default_llama_model_id": "extras/default-llama", + }, + "auxiliary_text_encoder": None, + } + )[0] + + extras = stage_cfg["engine_args"]["extras"] + assert extras["auxiliary_text_encoder"] == "/models/extras-llama" + assert extras["default_llama_model_id"] == "extras/default-llama" + + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "extras": { + "auxiliary_text_encoder": "/models/extras-llama", + "default_llama_model_id": "extras/default-llama", + }, + "auxiliary_text_encoder": "/models/top-level-llama", + "default_llama_model_id": "top-level/default-llama", + } + )[0] + + extras = stage_cfg["engine_args"]["extras"] + assert extras["auxiliary_text_encoder"] == "/models/top-level-llama" + assert extras["default_llama_model_id"] == "top-level/default-llama" + + +def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker): + """Local/unregistered Diffusers checkpoints still honor stage-0 extras.""" + + def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs): + return None, default_stage_cfg_factory(), None + + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + side_effect=resolve_with_default, + ) + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + + _, stage_configs = engine._resolve_stage_configs( + "/models/LTX-2.5-Diffusers", + {"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'}, + trust_remote_code=False, + ) + + assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True + + def test_default_cache_config_used_when_missing(): """Ensure default cache_config is synthesized when only backend is given.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( @@ -58,6 +129,23 @@ def test_default_stage_devices_from_sequence_parallel(): assert stage_cfg["runtime"]["devices"] == "0,1,2,3" +def test_default_stage_devices_and_dp_from_num_gpus(): + """Resolve omitted DP before deriving devices from an explicit GPU count.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "num_gpus": 8, + "tensor_parallel_size": 2, + "ulysses_degree": 2, + } + )[0] + + parallel_config = stage_cfg["engine_args"]["parallel_config"] + assert parallel_config.data_parallel_size == 2 + assert parallel_config.world_size == 8 + assert stage_cfg["engine_args"]["num_gpus"] == 8 + assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7" + + def test_default_stage_config_uses_parallel_size_kwargs(): """Ensure default diffusion parallel_config uses CLI/API parallel sizes.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( @@ -76,8 +164,8 @@ def test_default_stage_config_uses_parallel_size_kwargs(): assert parallel_config.enable_expert_parallel is True -def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): - """Ensure nullified diffusion parallel-size kwargs fall back to defaults.""" +def test_default_stage_config_preserves_omitted_dp_for_runtime_inference(): + """Keep omitted DP unresolved until the runtime WORLD size is known.""" stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( { "pipeline_parallel_size": None, @@ -92,7 +180,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs(): parallel_config = stage_cfg["engine_args"]["parallel_config"] assert parallel_config.pipeline_parallel_size == 1 - assert parallel_config.data_parallel_size == 1 + assert parallel_config.data_parallel_size is None assert parallel_config.tensor_parallel_size == 1 assert parallel_config.enable_expert_parallel is False assert stage_cfg["engine_args"]["enforce_eager"] is False @@ -208,7 +296,7 @@ def test_default_stage_config_whitelist_none_fallback(): engine_args = stage_cfg["engine_args"] assert engine_args["trust_remote_code"] is False - assert engine_args["distributed_executor_backend"] == "mp" + assert engine_args["distributed_executor_backend"] is None assert engine_args["dtype"] == "auto" assert engine_args["enforce_eager"] is False @@ -240,6 +328,29 @@ def test_serve_cli_accepts_ulysses_mode(): assert parallel_config.ulysses_mode == "advanced_uaa" +def test_serve_cli_accepts_text_encoder_tp_size(): + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "MiniMaxAI/MiniMax-H3", + "--omni", + "--text-encoder-tp-size", + "4", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + parallel_config = stage_cfg["engine_args"]["parallel_config"] + + assert args.text_encoder_tp_size == 4 + assert parallel_config.text_encoder_tp_size == 4 + + def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage(): parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="command") @@ -345,6 +456,39 @@ def test_serve_cli_forwards_distributed_offload_residency(): assert engine_args["dlo_resident_layers"] == 20 +def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "MiniMaxAI/MiniMax-H3", + "--omni", + "--enable-distributed-layerwise-offload", + "--dlo-no-use-allgather", + "--host-weight-runtime-mode", + "preferred", + "--host-weight-runtime-root", + "/var/cache/vllm-omni/hwr", + "--dlo-host-registration-limit-gib", + "80", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + engine_args = stage_cfg["engine_args"] + + assert explicit_kwargs["host_weight_runtime_mode"] == "preferred" + assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" + assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80 + assert engine_args["host_weight_runtime_mode"] == "preferred" + assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" + assert engine_args["dlo_host_registration_limit_gib"] == 80 + + def test_serve_cli_accepts_diffusion_compile_controls(): """Ensure both compile controls reach the diffusion stage.""" parser = TrackingArgumentParser() @@ -383,7 +527,9 @@ def test_serve_cli_accepts_diffusion_attention_backend(): "Qwen/Qwen-Image", "--omni", "--diffusion-attention-backend", - "FLASH_ATTN", + "FASTVIDEO_VSA", + "--fastvideo-vsa-topk", + "96", ] ) @@ -391,10 +537,12 @@ def test_serve_cli_accepts_diffusion_attention_backend(): stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"] - assert args.diffusion_attention_backend == "FLASH_ATTN" + assert args.diffusion_attention_backend == "FASTVIDEO_VSA" + assert args.fastvideo_vsa_topk == 96 assert isinstance(diffusion_attention_config, AttentionConfig) assert diffusion_attention_config.default is not None - assert diffusion_attention_config.default.backend == "FLASH_ATTN" + assert diffusion_attention_config.default.backend == "FASTVIDEO_VSA" + assert diffusion_attention_config.default.backend_kwargs() == {"topk": 96} def test_serve_cli_accepts_request_batch_max_wait_ms(): @@ -462,6 +610,66 @@ def test_serve_cli_accepts_additional_config(): assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}} +def test_default_stage_resolves_video_output_from_checkpoint(mocker): + captured = {} + + def resolve_with_default(*args, default_stage_cfg_factory, **kwargs): + del args, kwargs + captured.update(default_stage_cfg_factory()[0]) + stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace()) + return ("", [stage], None) + + resolver = mocker.patch( + "vllm_omni.engine.async_omni_engine.resolve_model_class_name", + return_value="MiniMaxH3Pipeline", + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + side_effect=resolve_with_default, + ) + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + engine._strip_single_engine_args = lambda kwargs: kwargs + + engine._resolve_stage_configs( + "/models/MiniMax-H3/FL2VA", + {}, + trust_remote_code=False, + ) + + resolver.assert_called_once_with("/models/MiniMax-H3/FL2VA", "default") + assert captured["final_output_type"] == "video" + + +def test_default_diffusers_stage_preserves_video_model_identity(mocker): + captured = {} + + def resolve_with_default(*args, default_stage_cfg_factory, **kwargs): + del args, kwargs + captured.update(default_stage_cfg_factory()[0]) + stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace()) + return ("", [stage], None) + + mocker.patch( + "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index", + return_value={"_class_name": "WanImageToVideoPipeline"}, + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + side_effect=resolve_with_default, + ) + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + engine._strip_single_engine_args = lambda kwargs: kwargs + + engine._resolve_stage_configs( + "/models/Wan2.2-I2V", + {"diffusion_load_format": "diffusers"}, + trust_remote_code=False, + ) + + assert captured["engine_args"]["model_class_name"] == "WanImageToVideoPipeline" + assert captured["final_output_type"] == "video" + + def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker): """Ensure YAML/deploy stage resolution forwards top-level additional_config.""" fake_diffusion_stage = SimpleNamespace( @@ -478,12 +686,11 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) - engine._strip_single_engine_args = lambda kwargs: kwargs _, stage_configs = engine._resolve_stage_configs( "dummy-model", { - "stage_configs_path": "dummy.yaml", + "deploy_config": "dummy.yaml", "additional_config": {"torchair_graph_config": {"enabled": True}}, }, trust_remote_code=False, @@ -493,6 +700,24 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}} +@pytest.mark.parametrize( + ("legacy_arg", "value"), + [ + ("stage_configs_path", "legacy.yaml"), + ("stage_configs", [{"stage_id": 0}]), + ], +) +def test_resolve_stage_configs_rejects_legacy_config_arguments(legacy_arg, value): + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + + with pytest.raises(ValueError, match=rf"`{legacy_arg}`.*`deploy_config`"): + engine._resolve_stage_configs( + "dummy-model", + {legacy_arg: value}, + trust_remote_code=False, + ) + + def test_default_stage_config_includes_quantization_config(): """Ensure structured quantization_config survives default diffusion-stage creation.""" quantization_config = { @@ -516,12 +741,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage( ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) - engine._strip_single_engine_args = lambda kwargs: kwargs _, stage_configs = engine._resolve_stage_configs( "dummy-model", { - "stage_configs_path": "dummy.yaml", + "deploy_config": "dummy.yaml", "quantization_config": {"method": "bitsandbytes"}, }, trust_remote_code=False, diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py new file mode 100644 index 000000000..9600b990d --- /dev/null +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -0,0 +1,441 @@ +"""Unit tests for AR EngineCore vs diffusion worker pause/sleep routing.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, call + +import pytest + +from vllm_omni.diffusion.data import CuMemTag, OmniACK +from vllm_omni.entrypoints.async_omni import AsyncOmni + +pytestmark = [pytest.mark.core_model] + + +def _make_omni(*, stage_types: list[str]) -> AsyncOmni: + omni = object.__new__(AsyncOmni) + omni._pause_cond = asyncio.Condition() + omni._paused = False + omni._admitting = 0 + omni._hold_admission_until_resume = False + omni._sleeping_tags = set() + omni._stage_sleeping_tags = {} + omni._level2_sleeping = False + omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock()) + omni._final_output_handler = lambda: None + omni.reset_mm_cache = AsyncMock() + + stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] + omni.engine = SimpleNamespace( + stage_clients=stage_clients, + stage_vllm_configs=[None] * len(stage_types), + collective_rpc_async=AsyncMock(return_value=[True]), + ) + omni.collective_rpc = AsyncMock(return_value=[True]) + return omni + + +@pytest.mark.cpu +def test_split_stage_ids_by_type(): + omni = _make_omni(stage_types=["llm", "diffusion", "llm"]) + ar_ids, diff_ids = omni._split_stage_ids_by_type() + assert ar_ids == [0, 2] + assert diff_ids == [1] + + +@pytest.mark.cpu +def test_split_stage_ids_by_type_rejects_out_of_range(): + omni = _make_omni(stage_types=["llm", "diffusion"]) + with pytest.raises(ValueError, match=r"Invalid stage_ids \[2\].*0\.\.1"): + omni._split_stage_ids_by_type([0, 2]) + with pytest.raises(ValueError, match=r"Invalid stage_ids \[-1\]"): + omni._split_stage_ids_by_type([-1]) + + +@pytest.mark.cpu +def test_pause_generation_routes_ar_via_collective_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + + await omni.pause_generation(mode="abort", clear_cache=True) + + assert omni._paused is True + assert omni._hold_admission_until_resume is True + omni.collective_rpc.assert_awaited_once_with( + method="pause_scheduler", + args=(), + kwargs={"mode": "abort", "clear_cache": True}, + stage_ids=[0], + ) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_pause_generation_still_rpcs_when_already_paused(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + omni._paused = True + + await omni.pause_generation(mode="abort", clear_cache=True, stage_ids=[0]) + + omni.collective_rpc.assert_awaited_once_with( + method="pause_scheduler", + args=(), + kwargs={"mode": "abort", "clear_cache": True}, + stage_ids=[0], + ) + omni.reset_prefix_cache.assert_awaited_once() + omni.reset_mm_cache.assert_awaited_once() + omni.reset_encoder_cache.assert_awaited_once() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_resume_generation_resumes_ar_then_clears_frontend_pause(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni._paused = True + + await omni.resume_generation() + + omni.collective_rpc.assert_awaited_once_with( + method="resume_scheduler", + args=(), + kwargs=None, + stage_ids=[0], + ) + assert omni._paused is False + assert omni._hold_admission_until_resume is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0) + omni._sleep_diffusion = AsyncMock(return_value=[diffusion_ack]) + + acks = await omni.sleep(stage_ids=[0, 1], level=1, mode="abort") + + omni.collective_rpc.assert_awaited_once_with( + method="sleep", + args=(1, "abort"), + kwargs=None, + stage_ids=[0], + ) + omni._sleep_diffusion.assert_awaited_once_with([1], 1) + assert {ack.stage_id for ack in acks} == {0, 1} + assert any(ack.metadata.get("path") == "engine_core" for ack in acks if ack.stage_id == 0) + assert CuMemTag.WEIGHTS.value in omni._sleeping_tags + assert CuMemTag.KV_CACHE.value in omni._sleeping_tags + # Sleep gates frontend admission for the trainer resume contract. + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_blocks_admission_before_engine_core_rpc(): + """Problem 3: _paused must be set before awaiting sleep RPC.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + paused_at_rpc: list[bool] = [] + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + paused_at_rpc.append(omni._paused) + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + + assert omni._paused is False + await omni.sleep(level=1, mode="abort") + + assert paused_at_rpc == [True] + assert omni._paused is True + assert omni._hold_admission_until_resume is True + # Sleep must not also call pause_scheduler (EngineCore.sleep pauses). + assert all(c.kwargs.get("method") != "pause_scheduler" for c in omni.collective_rpc.await_args_list) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_waits_for_in_flight_generate_admission(): + """Sleep must not offload while generate() is still in add_request.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + omni._admitting = 1 + rpc_started = asyncio.Event() + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + rpc_started.set() + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort")) + await asyncio.sleep(0.05) + assert not sleep_task.done() + assert not rpc_started.is_set() + await omni._release_generate_admission() + await asyncio.wait_for(sleep_task, timeout=1) + assert rpc_started.is_set() + assert omni._admitting == 0 + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni._paused = True + omni._hold_admission_until_resume = True + omni._sleeping_tags = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value} + diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0) + omni._wake_diffusion = AsyncMock(return_value=[diffusion_ack]) + + acks = await omni.wake_up(stage_ids=[0, 1]) + + omni.collective_rpc.assert_awaited_once() + wake_kwargs = omni.collective_rpc.await_args.kwargs + assert wake_kwargs["method"] == "wake_up" + assert wake_kwargs["stage_ids"] == [0] + assert set(wake_kwargs["kwargs"]["tags"]) == { + CuMemTag.WEIGHTS.value, + CuMemTag.KV_CACHE.value, + } + omni._wake_diffusion.assert_awaited_once() + assert {ack.stage_id for ack in acks} == {0, 1} + assert not omni._sleeping_tags + # Mixed/AR wake restores memory but does not resume frontend admission. + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_wake_up_does_not_resume_frontend_admission(): + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + await omni.sleep(level=1, mode="abort") + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + await omni.wake_up() + + assert omni._paused is True + assert omni._hold_admission_until_resume is True + assert not omni._sleeping_tags + # Explicit resume is required after AR sleep/wake. + await omni.resume_generation() + assert omni._paused is False + assert omni._hold_admission_until_resume is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_level1_wake_without_tags_clears_all_sleeping_tags(): + """sleep(level=1) tracks WEIGHTS+KV; untagged wake_up must clear both.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + + await omni.sleep(level=1, mode="abort") + assert CuMemTag.WEIGHTS.value in omni._sleeping_tags + assert CuMemTag.KV_CACHE.value in omni._sleeping_tags + + await omni.wake_up() + assert not omni._sleeping_tags + assert omni._paused is True + assert omni._hold_admission_until_resume is True + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_sleep_diffusion_only_skips_engine_core_collective_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)]) + + await omni.sleep(level=1) + + # AR EngineCore sleep path must not run for diffusion-only. + assert call(method="sleep", args=(1, "abort"), kwargs=None, stage_ids=[0]) not in ( + omni.collective_rpc.await_args_list + ) + omni._sleep_diffusion.assert_awaited_once_with([0], 1) + assert omni._paused is True + assert omni._hold_admission_until_resume is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_streaming_wait_for_first_chunk_does_not_block_sleep(): + """Waiting for the next client chunk must not hold an admission slot.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + rpc_started = asyncio.Event() + first_chunk = asyncio.Event() + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + rpc_started.set() + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + + async def wait_then_submit() -> None: + await first_chunk.wait() + await omni._submit_with_admission(asyncio.sleep(0)) + + wait_task = asyncio.create_task(wait_then_submit()) + sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort")) + await asyncio.sleep(0.05) + assert not wait_task.done() + assert omni._admitting == 0 + await asyncio.wait_for(sleep_task, timeout=1) + assert rpc_started.is_set() + + wait_task.cancel() + with pytest.raises(asyncio.CancelledError): + await wait_task + assert omni._admitting == 0 + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_streaming_in_flight_add_blocks_sleep(): + """Sleep must wait while a streaming ADD/update holds an admission slot.""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm"]) + add_started = asyncio.Event() + add_release = asyncio.Event() + rpc_started = asyncio.Event() + + async def rpc_side_effect(**kwargs): + if kwargs.get("method") == "sleep": + rpc_started.set() + return [True] + + omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect) + + async def slow_add() -> None: + add_started.set() + await add_release.wait() + + add_task = asyncio.create_task(omni._submit_with_admission(slow_add())) + await add_started.wait() + sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort")) + await asyncio.sleep(0.05) + assert not sleep_task.done() + assert not rpc_started.is_set() + assert omni._admitting == 1 + + add_release.set() + await asyncio.wait_for(add_task, timeout=1) + await asyncio.wait_for(sleep_task, timeout=1) + assert rpc_started.is_set() + assert omni._admitting == 0 + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_wake_up_restores_admission_for_diffusion_only(): + """Pure diffusion must keep sleep → wake → generate (no resume).""" + + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)]) + omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)]) + + await omni.sleep(level=1) + assert omni._paused is True + + async def wait_for_admission() -> None: + async with omni._pause_cond: + await omni._pause_cond.wait_for(lambda: not omni._paused) + + waiter = asyncio.create_task(wait_for_admission()) + await omni.wake_up() + await asyncio.wait_for(waiter, timeout=1.0) + assert omni._paused is False + assert omni._hold_admission_until_resume is False + assert not omni._sleeping_tags + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion(): + """Explicit pause_generation still requires resume after diffusion wake.""" + + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)]) + omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)]) + + await omni.pause_generation() + await omni.sleep(level=1) + await omni.wake_up() + + assert omni._paused is True + assert omni._hold_admission_until_resume is True + await omni.resume_generation() + assert omni._paused is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_partial_wake_does_not_skip_remaining_sleeping_stage(): + """sleep(stage_ids=[0]) then wake([0]) must not skip a later wake([1]).""" + + async def run() -> None: + omni = _make_omni(stage_types=["llm", "llm"]) + await omni.sleep(stage_ids=[0, 1], level=1, mode="abort") + assert omni._stage_sleeping_tags.keys() == {0, 1} + + await omni.wake_up(stage_ids=[0]) + assert 0 not in omni._stage_sleeping_tags + assert 1 in omni._stage_sleeping_tags + assert CuMemTag.WEIGHTS.value in omni._sleeping_tags + + omni.collective_rpc.reset_mock() + acks = await omni.wake_up(stage_ids=[1]) + assert acks + omni.collective_rpc.assert_awaited_once() + wake_kwargs = omni.collective_rpc.await_args.kwargs + assert wake_kwargs["method"] == "wake_up" + assert wake_kwargs["stage_ids"] == [1] + assert not omni._sleeping_tags + assert not omni._stage_sleeping_tags + + asyncio.run(run()) diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py index 3f793e89e..827ab67b6 100644 --- a/tests/entrypoints/test_omni_base_profiler.py +++ b/tests/entrypoints/test_omni_base_profiler.py @@ -46,6 +46,16 @@ class TestOmniBaseProfiler: stage_ids=None, ) + def test_removed_diffusion_batch_size_fails_loudly(self, mock_engine, mocker: MockerFixture): + """Legacy diffusion_batch_size callers receive an actionable error.""" + mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine) + mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x) + mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize") + from vllm_omni.entrypoints.omni_base import OmniBase + + with pytest.raises(TypeError, match="diffusion_batch_size.*max_num_seqs"): + OmniBase(model="test-model", diffusion_batch_size=8) + def test_start_profile_with_prefix(self, omni_base_instance, mock_engine): """Test that start_profile passes profile_prefix to collective_rpc.""" omni_base_instance.start_profile(profile_prefix="test_trace") diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 161b3fe2a..4f9fc592b 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -12,12 +12,13 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError +from vllm_omni.config.stage_config import StageConfig from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ErrorMessage, OutputMessage from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni import Omni -from vllm_omni.entrypoints.omni_base import OmniEngineDeadError +from vllm_omni.entrypoints.omni_base import OmniBase, OmniEngineDeadError from vllm_omni.errors import ( OmniClientError, client_error_from_metadata, @@ -188,14 +189,39 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) def _make_base(): - from vllm_omni.entrypoints.omni_base import OmniBase - obj = object.__new__(OmniBase) obj.engine = MagicMock() obj.request_states = {} return obj +def test_resolve_sampling_params_list_preserves_stage_constraints(): + base = _make_base() + base.engine.num_stages = 1 + base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])] + base.engine.stage_configs = [ + StageConfig( + stage_id=0, + model_stage="dummy-model", + sampling_constraints={"detokenize": False, "stop_token_ids": [42]}, + ).to_omegaconf() + ] + base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs) + assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}] + caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7]) + + resolved = base.resolve_sampling_params_list(caller_params) + + assert resolved[0] is not caller_params + assert resolved[0].seed == 1234 + assert resolved[0].max_tokens == 7 + assert resolved[0].detokenize is False + assert resolved[0].stop_token_ids == [42] + assert 42 in resolved[0]._all_stop_token_ids + assert caller_params.detokenize is True + assert caller_params.stop_token_ids == [7] + + def _stage_spec( stage_id: int, *, @@ -585,7 +611,13 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat # external ID to avoid collisions, so this also tests mapping external_req_id = "req-1" req_id = "req-1-12345678" + recorded_failures = [] try: + monkeypatch.setattr( + app, + "_record_request_failure_once", + lambda request_id, reason: recorded_failures.append((request_id, reason)), + ) app.request_states[req_id] = ClientRequestState( request_id=req_id, external_request_id=external_req_id, @@ -596,6 +628,7 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat assert engine.aborted == [[req_id]] assert external_req_id not in app.request_states + assert recorded_failures == [(req_id, "client_abort")] @pytest.mark.asyncio @@ -788,7 +821,13 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch): _patch_engine(monkeypatch, engine) app = Omni("dummy-model") + recorded_failures = [] try: + monkeypatch.setattr( + app, + "_record_request_failure_once", + lambda request_id, reason: recorded_failures.append((request_id, reason)), + ) app.request_states["req-1"] = object() app.abort("req-1") finally: @@ -796,6 +835,7 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch): assert engine.aborted == [["req-1"]] assert "req-1" not in app.request_states + assert recorded_failures == [("req-1", "client_abort")] def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch): diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index 0377e1a9a..9567cbc96 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Omni sleep mode: entrypoint-level VRAM/ACK tests plus multi-stage e2e. @@ -109,11 +109,15 @@ def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str: async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None: - """Best-effort full wake so the next shared-engine test starts active.""" + """Best-effort full wake + resume so the next shared-engine test starts active.""" try: await engine.wake_up(stage_ids=stage_ids) except Exception as e: logger.warning("ensure_awake failed (stage_ids=%s): %s", stage_ids, e) + try: + await engine.resume_generation(stage_ids=stage_ids) + except Exception as e: + logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e) def _build_llm_stages() -> tuple[list[dict], list[dict]]: @@ -148,13 +152,13 @@ def _build_llm_stages() -> tuple[list[dict], list[dict]]: @pytest.fixture(scope="module", autouse=True) def _module_device_cleanup(): """One device cleanup pass around the module (engines are class-scoped).""" - from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup + from tests.helpers.clean import cleanup_test_environment print("\n=== PRE-MODULE DEVICE CLEANUP (sleep_mode) ===") - run_pre_test_cleanup() + cleanup_test_environment() yield print("\n=== POST-MODULE DEVICE CLEANUP (sleep_mode) ===") - run_post_test_cleanup() + cleanup_test_environment() @pytest_asyncio.fixture(scope="class", loop_scope="class") @@ -240,6 +244,9 @@ class TestOmniLlmSleepMode: try: await llm_engine.sleep(stage_ids=[0], level=1) await llm_engine.wake_up(stage_ids=[0], tags=["weights"]) + # wake_up restores weights but keeps the AR admission hold. + # Resume so generate() can reach the sleeping-tags rejection. + await llm_engine.resume_generation(stage_ids=[0]) with pytest.raises(RuntimeError, match="partially or fully asleep"): async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)): pass @@ -364,6 +371,8 @@ class TestOmniDiffusionSleepMode: # --- wake + VRAM restore to initial (±3 GiB) --- logger.info("Waking up (Reloading Weights)...") await diffusion_engine.wake_up(stage_ids=[0]) + # AR stages keep the frontend admission gate held after wake. + await diffusion_engine.resume_generation(stage_ids=[0]) await asyncio.sleep(2.0) gc.collect() get_vram_info(device_id) @@ -490,6 +499,7 @@ async def test_multistage_llm_diffusion_sleep_wake(tp_size: int): assert len(acks) == 2 await engine.wake_up(stage_ids=[0, 1]) + await engine.resume_generation(stage_ids=[0, 1]) async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]): pass finally: @@ -522,6 +532,7 @@ async def test_pure_diffusion_sleep_wake(): try: await engine.sleep(level=1) await engine.wake_up() + await engine.resume_generation() async for _ in engine.generate( "test", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256), diff --git a/tests/entrypoints/test_omni_snapshot_download.py b/tests/entrypoints/test_omni_snapshot_download.py index d2c76a249..eae774178 100644 --- a/tests/entrypoints/test_omni_snapshot_download.py +++ b/tests/entrypoints/test_omni_snapshot_download.py @@ -1,17 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """``omni_snapshot_download`` must honor vLLM's ``VLLM_USE_MODELSCOPE`` semantics. -vLLM treats the flag as enabled only for the literal string ``"true"`` -(case-insensitive). Reading ``os.environ`` directly made every non-empty value -truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0`` still took the -ModelScope path. +vLLM treats the flag as enabled for the literal strings ``"1"`` or ``"true"`` +(case-insensitive; see vllm.envs). Reading ``os.environ`` directly made every +non-empty value truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0`` +still took the ModelScope path. """ import sys import types +from pathlib import Path import pytest +from pytest_mock import MockerFixture from vllm import envs from vllm_omni.entrypoints import omni_base @@ -21,6 +23,64 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] MODEL_ID = "Qwen/Qwen3-Omni-30B-A3B-Instruct" +@pytest.mark.parametrize( + "model_uri", + [ + "s3://bucket/model", + "gs://bucket/model", + "az://bucket/model", + ], +) +def test_omni_snapshot_download_preserves_object_storage_uri( + model_uri: str, + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) + hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific") + + assert omni_base.omni_snapshot_download(model_uri) == model_uri + hf_download.assert_not_called() + + +def test_omni_snapshot_download_preserves_existing_local_path( + tmp_path: Path, + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) + absolute_model_path = tmp_path / "absolute-model" + absolute_model_path.mkdir() + relative_model_path = "relative-model" + (tmp_path / relative_model_path).mkdir() + monkeypatch.chdir(tmp_path) + hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific") + + assert omni_base.omni_snapshot_download(str(absolute_model_path)) == str(absolute_model_path) + assert omni_base.omni_snapshot_download(relative_model_path) == relative_model_path + hf_download.assert_not_called() + + +def test_omni_snapshot_download_uses_hf_for_relative_repo_id( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False) + model_id = "org/model" + # Stub the modular-index probe so the test never reaches huggingface.co. + hf_file_probe = mocker.patch.object(omni_base, "file_or_path_exists", return_value=False) + hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific") + + assert omni_base.omni_snapshot_download(model_id) == model_id + hf_file_probe.assert_called_once_with(model_id, "modular_model_index.json", revision=None) + hf_download.assert_called_once_with( + model_name_or_path=model_id, + cache_dir=None, + allow_patterns=["*"], + require_all=True, + ) + + @pytest.fixture def download_backend(monkeypatch: pytest.MonkeyPatch): """Run ``omni_snapshot_download`` and report which backend it selected. @@ -67,7 +127,7 @@ def download_backend(monkeypatch: pytest.MonkeyPatch): return run -@pytest.mark.parametrize("value", ["0", "1", "False", "false", "no", "off"]) +@pytest.mark.parametrize("value", ["0", "False", "false", "no", "off"]) def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend, value): monkeypatch.setenv("VLLM_USE_MODELSCOPE", value) @@ -75,7 +135,7 @@ def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend, assert download_backend() == "huggingface" -@pytest.mark.parametrize("value", ["true", "True", "TRUE"]) +@pytest.mark.parametrize("value", ["1", "true", "True", "TRUE"]) def test_true_values_enable_modelscope(monkeypatch, download_backend, value): monkeypatch.setenv("VLLM_USE_MODELSCOPE", value) diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index cc459c269..660203710 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -49,7 +49,9 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch): def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch): monkeypatch.setattr( "vllm.transformers_utils.config.get_hf_file_to_dict", - lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}, + lambda path, _model, **_kwargs: ( + None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]} + ), ) od_config = OmniDiffusionConfig( diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 04a29edd1..e759e0f15 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for the Omni serve CLI helpers.""" from __future__ import annotations @@ -76,6 +79,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: assert args.cfg_parallel_size == 4 +def test_serve_parser_accepts_ulysses_a2a_permute() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args(["serve", "fake-model", "--omni", "--ulysses-a2a-permute"]) + + assert args.ulysses_a2a_permute is True + assert args.get_explicit_kwargs_dict()["ulysses_a2a_permute"] is True + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", @@ -193,7 +207,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) with pytest.raises(ValueError, match="No stage config found for stage_id=0"): run_headless(args) - assert captured["args"][1] is None + assert len(captured["args"]) == 2 assert captured["deploy_config_path"] == "/tmp/deploy.yaml" assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}} assert captured["strategy_config_path"] == "/tmp/strategy.yaml" diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index e2890797f..73702df3a 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -216,6 +216,20 @@ def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytes assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4" +@pytest.mark.core_model +@pytest.mark.cpu +def test_resolve_stage_physical_devices_captured_unset_ignores_live_env(monkeypatch: pytest.MonkeyPatch): + """Parallel init captures an unset CUDA_VISIBLE_DEVICES as None. + + A sibling thread may then narrow the live env for spawn. The captured + None must not fall back to that live value, or logical device 0 would + map onto the other stage's GPU. + """ + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7") + assert resolve_stage_physical_devices(0, "0", visible_baseline=None) == "0" + assert resolve_stage_physical_devices(1, "1", visible_baseline=None) == "1" + + @pytest.mark.core_model @pytest.mark.cpu def test_map_device_list_raises_on_non_numeric(): diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index b02936838..97ed2370b 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -16,19 +16,20 @@ Key invariants tested: - voice/speaker parameter compatibility in chat completions """ -import json import time from unittest.mock import MagicMock import pytest -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ChatCompletionResponseStreamChoice, -) -from vllm.entrypoints.openai.engine.protocol import DeltaMessage -from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.outputs import CompletionOutput, RequestOutput +from tests.helpers.serving_chat import ( + build_serving_chat, + collect_stream, + make_request, + make_text_omni_output, + parse_sse_chunks, +) from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.outputs import OmniRequestOutput @@ -40,43 +41,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # --------------------------------------------------------------------------- -def _make_text_omni_output( - request_id: str = "test-req", - text: str = "hello", - token_ids: list[int] | None = None, - finish_reason: str | None = None, - index: int = 0, - num_prompt_tokens: int = 3, -) -> OmniRequestOutput: - """Build an OmniRequestOutput wrapping a text RequestOutput.""" - if token_ids is None: - token_ids = [10, 11, 12] - res = RequestOutput( - request_id=request_id, - prompt="test", - prompt_token_ids=list(range(num_prompt_tokens)), - prompt_logprobs=None, - outputs=[ - CompletionOutput( - index=index, - text=text, - token_ids=token_ids, - cumulative_logprob=0.0, - logprobs=None, - finish_reason=finish_reason, - stop_reason=None, - ) - ], - finished=finish_reason is not None, - ) - return OmniRequestOutput.from_stage_output( - res, - request_id=request_id, - final_output_type="text", - finished=finish_reason is not None, - ) - - def _make_audio_omni_output( request_id: str = "test-req", index: int = 0, @@ -118,82 +82,6 @@ def _make_audio_omni_output( ) -def _mock_audio_choices(index: int = 0, role: str = "assistant"): - return [ - ChatCompletionResponseStreamChoice( - index=index, - delta=DeltaMessage(role=role, content="dGVzdA=="), - logprobs=None, - finish_reason="stop", - ) - ] - - -def _build_serving_chat(): - """Create a minimal OmniOpenAIServingChat for testing.""" - mock_engine = MagicMock() - mock_engine.errored = False - - models = OpenAIServingModels( - engine_client=mock_engine, - base_model_paths=[], - ) - mock_render = MagicMock() - - instance = OmniOpenAIServingChat( - engine_client=mock_engine, - models=models, - response_role="assistant", - online_renderer=mock_render, - request_logger=None, - chat_template=None, - chat_template_content_format="auto", - ) - instance._create_audio_choice = MagicMock( - side_effect=lambda omni_res, role, request, stream=False: _mock_audio_choices( - index=omni_res.outputs[0].index, - role=role, - ) - ) - return instance - - -def _make_request(modalities: list[str], n: int = 1) -> ChatCompletionRequest: - req = ChatCompletionRequest( - model="test-model", - messages=[{"role": "user", "content": "hello"}], - n=n, - stream=True, - ) - req.modalities = modalities # type: ignore[attr-defined] - return req - - -def _parse_sse_chunks(lines: list[str]) -> list[dict]: - """Parse SSE lines into JSON dicts.""" - prefix = "data: " - chunks = [] - for line in lines: - line = line.strip() - if not line.startswith(prefix): - continue - payload = line[len(prefix) :].strip() - if payload == "[DONE]": - continue - try: - chunks.append(json.loads(payload)) - except json.JSONDecodeError: - pass - return chunks - - -async def _collect_stream(gen): - result = [] - async for item in gen: - result.append(item) - return result - - # --------------------------------------------------------------------------- # Tests: finish_reason correctness # --------------------------------------------------------------------------- @@ -202,14 +90,14 @@ async def _collect_stream(gen): @pytest.mark.asyncio async def test_single_modality_text_only_one_stop(): """Text-only streaming: exactly one chunk has finish_reason='stop'.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text"]) async def result_generator(): - yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) - yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") + yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) + yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -221,7 +109,7 @@ async def test_single_modality_text_only_one_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [c["choices"][0]["finish_reason"] for c in chunks if c.get("choices")] assert finish_reasons[-1] == "stop" @@ -233,15 +121,15 @@ async def test_single_modality_text_only_one_stop(): @pytest.mark.asyncio async def test_multi_modal_text_audio_only_last_stop(): """text+audio: text finish sends finish_reason=null, audio sends stop.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) async def result_generator(): - yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) - yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") + yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None) + yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop") yield _make_audio_omni_output() - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -253,7 +141,7 @@ async def test_multi_modal_text_audio_only_last_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1 @@ -269,18 +157,18 @@ async def test_multi_modal_text_audio_only_last_stop(): @pytest.mark.asyncio async def test_multi_modal_n2_independent_per_choice(): """n=2 with text+audio: each choice gets exactly one stop, at the end.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"], n=2) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"], n=2) async def result_generator(): - yield _make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0) - yield _make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1) - yield _make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0) - yield _make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1) + yield make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0) + yield make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1) + yield make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0) + yield make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1) yield _make_audio_omni_output(index=0) yield _make_audio_omni_output(index=1) - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -292,7 +180,7 @@ async def test_multi_modal_n2_independent_per_choice(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) per_choice: dict[int, list] = {} for c in chunks: for ch in c.get("choices", []): @@ -306,13 +194,13 @@ async def test_multi_modal_n2_independent_per_choice(): @pytest.mark.asyncio async def test_single_modality_audio_only_one_stop(): """Audio-only streaming: the audio chunk carries finish_reason='stop'.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["audio"]) async def result_generator(): yield _make_audio_omni_output() - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -324,7 +212,7 @@ async def test_single_modality_audio_only_one_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1 @@ -349,8 +237,8 @@ async def test_streaming_audio_metrics_resolve_replica_id( import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod from vllm_omni.entrypoints.client_request_state import ClientRequestState - serving_chat = _build_serving_chat() - request = _make_request(modalities=["audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["audio"]) req_state = ClientRequestState( request_id="internal-req", external_request_id="test-req", @@ -377,7 +265,7 @@ async def test_streaming_audio_metrics_resolve_replica_id( async def result_generator(): yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400) - await _collect_stream( + await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -408,15 +296,15 @@ async def test_streaming_audio_metrics_resolve_replica_id( async def test_declared_modality_not_produced_emits_fallback_stop(): """If request.modalities declares ["text","audio"] but engine only produces text, a fallback stop chunk is emitted at stream end.""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) async def result_generator(): # Engine only produces text, no audio output at all - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) - yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) + yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -428,7 +316,7 @@ async def test_declared_modality_not_produced_emits_fallback_stop(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] # Text finish is suppressed (audio not seen yet), but fallback stop @@ -441,15 +329,15 @@ async def test_declared_modality_not_produced_emits_fallback_stop(): async def test_declared_modality_not_produced_text_finish_suppressed(): """When text finishes but audio (declared in modalities) never appears, the text finish chunk has finish_reason=null (suppressed).""" - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) async def result_generator(): - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) - yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) + yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") # No audio output — stream ends - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -461,7 +349,7 @@ async def test_declared_modality_not_produced_text_finish_suppressed(): ) ) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) # Find the text finish chunk (content "!") for c in chunks: @@ -485,8 +373,8 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): yields (field, value) tuples, so the loop raised ``AttributeError: 'tuple' object has no attribute 'finish_reason'``. """ - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) empty_audio = _make_audio_omni_output() empty_audio.outputs[0].multimodal_output = {"audio": []} @@ -497,11 +385,11 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): serving_chat._create_audio_choice = create_audio_choice async def result_generator(): - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) - yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None) + yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop") yield empty_audio - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -514,7 +402,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive(): ) assert not any("Error in chat completion stream generator" in line for line in raw_lines) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}" assert finish_reasons[-1] == "stop" @@ -525,18 +413,18 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices(): """Any ErrorResponse from the audio path is skipped, not iterated.""" from vllm.entrypoints.openai.engine.protocol import ErrorResponse - serving_chat = _build_serving_chat() - request = _make_request(modalities=["text", "audio"]) + serving_chat = build_serving_chat() + request = make_request(modalities=["text", "audio"]) serving_chat._create_audio_choice = MagicMock( side_effect=lambda omni_res, role, request, stream=False: serving_chat._create_error_response("boom") ) async def result_generator(): - yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop") + yield make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop") yield _make_audio_omni_output() - raw_lines = await _collect_stream( + raw_lines = await collect_stream( serving_chat.chat_completion_stream_generator( request=request, result_generator=result_generator(), @@ -550,7 +438,7 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices(): assert isinstance(serving_chat._create_error_response("boom"), ErrorResponse) assert not any("AttributeError" in line for line in raw_lines) - chunks = _parse_sse_chunks(raw_lines) + chunks = parse_sse_chunks(raw_lines) finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])] assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}" diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index c6d8e0354..1fc4779a7 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for vllm_omni.entrypoints.utils module.""" import os @@ -22,11 +25,12 @@ from vllm_omni.engine.stage_init_utils import ( from vllm_omni.entrypoints.utils import ( _convert_dataclasses_to_dict, _filter_dict_like_object, + _try_resolve_omni_model_type, coerce_param_message_types, filter_dataclass_kwargs, filter_stages, load_and_resolve_stage_configs, - load_stage_configs_from_yaml, + load_stage_configs_from_model, resolve_model_config_path, ) @@ -371,17 +375,60 @@ class TestResolveModelConfigPath: assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000 assert sender["extra"]["connector_get_max_wait"] == 300 + def test_object_storage_uri_resolves_via_materialized_configs( + self, + mocker: MockerFixture, + tmp_path, + ): + """An object-storage URI must not reach HF helpers; resolution reads the + locally materialized config files instead.""" + hf_config = SimpleNamespace( + model_type="qwen3_omni_moe", + architectures=["Qwen3OmniMoeForConditionalGeneration"], + ) + uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct" + materialize = mocker.patch( + "vllm_omni.entrypoints.utils._materialize_object_storage_configs", + return_value=str(tmp_path), + ) + get_config = mocker.patch( + "vllm_omni.entrypoints.utils.get_config", + return_value=hf_config, + ) + + result = resolve_model_config_path(uri) + + materialize.assert_called_once_with(uri) + get_config.assert_called_once() + assert get_config.call_args[0][0] == str(tmp_path) + assert result is not None + assert Path(result).as_posix().endswith("qwen3_omni_moe.yaml") + + +class TestTryResolveOmniModelType: + """Name matching must only scan the last path component.""" + + def test_uri_bucket_name_cannot_hijack_pipeline(self): + # "qwen3-tts-models" normalizes to "qwen3ttsmodels", which contains the + # registered key "qwen3tts" — only the basename may participate, so a + # misleadingly named bucket selects nothing. + assert _try_resolve_omni_model_type("s3://qwen3-tts-models/plain-checkpoint") is None + + def test_basename_still_matches_pipeline(self): + assert _try_resolve_omni_model_type("gs://any-bucket/Fun-CosyVoice3-0.5B-2512") == "cosyvoice3" + class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture): - """Ensure that dtype survives default stage creation.""" - kwargs = {"dtype": torch.float32} + """Ensure that kwargs survive default stage creation.""" + engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine" + kwargs = {"dtype": torch.float32, "engine_backend": engine_backend} + kwargs["revision"] = "pinned-revision" mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None) mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None)) config_path, stage_configs, _ = load_and_resolve_stage_configs( model="black-forest-labs/FLUX.2-klein-4B", - stage_configs_path=None, kwargs=kwargs, trust_remote_code=False, default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs), @@ -389,10 +436,10 @@ class TestLoadAndResolveStageConfigs: assert config_path is None assert len(stage_configs) == 1 assert "dtype" in stage_configs[0]["engine_args"] + assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend + assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision" - def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas( - self, tmp_path, mocker: MockerFixture - ): + def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture): deploy_path = tmp_path / "qwen3_multi.yaml" deploy_path.write_text( 'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n', @@ -413,18 +460,23 @@ class TestLoadAndResolveStageConfigs: "vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=(returned_stage_configs, None), ) + cli_overrides = { + "interleave_mm_strings": True, + "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}, + "gpu_memory_utilization": 0.55, + } config_path, stage_configs, _ = load_and_resolve_stage_configs( model="dummy-model", - stage_configs_path=str(deploy_path), - kwargs={}, + kwargs=cli_overrides, trust_remote_code=True, + deploy_config_path=str(deploy_path), ) load_stage_configs.assert_called_once_with( "dummy-model", trust_remote_code=True, - base_engine_args={}, + base_engine_args=cli_overrides, deploy_config_path=str(deploy_path), stage_overrides=None, strategy_config_path=None, @@ -473,71 +525,20 @@ class TestLoadAndResolveStageConfigs: assert filtered[0].final_output_type is None -class TestLoadStageConfigsFromYaml: - """Regression tests for stage-config loading and merging.""" - - def test_deep_merges_stage_engine_args(self, mocker: MockerFixture): - yaml_config = create_config( - { - "async_chunk": True, - "stage_args": [ - { - "stage_id": 0, - "runtime": {"device": 0}, - "engine_args": { - "parallel_config": {"tensor_parallel_size": 4}, - }, - } - ], - } - ) +class TestLoadStageConfigsFromModel: + def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog): mocker.patch( - "vllm_omni.entrypoints.utils.load_yaml_config", - return_value=yaml_config, + "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model", + return_value=(None, None), ) + resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path") - stages = load_stage_configs_from_yaml( - "fake.yaml", - base_engine_args={ - "parallel_config": { - "tensor_parallel_size": 1, - "pipeline_parallel_size": 2, - }, - "model": "base-model", - }, - ) + result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False) - merged_engine_args = stages[0]["engine_args"] - assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4 - assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2 - assert merged_engine_args["model"] == "base-model" - assert merged_engine_args["async_chunk"] is True - - def test_merges_nested_stage_engine_args(self, mocker: MockerFixture): - yaml_config = create_config( - { - "stage_args": [ - { - "stage_id": 0, - "engine_args": { - "nested": {"override": 2}, - }, - } - ], - } - ) - mocker.patch( - "vllm_omni.entrypoints.utils.load_yaml_config", - return_value=yaml_config, - ) - - stages = load_stage_configs_from_yaml( - "fake.yaml", - base_engine_args={"nested": {"base": 1}}, - ) - - assert stages[0]["engine_args"]["nested"]["base"] == 1 - assert stages[0]["engine_args"]["nested"]["override"] == 2 + assert result == ([], None) + resolve_path.assert_not_called() + assert "No registered PipelineConfig resolved" in caplog.text + assert "deploy_config" in caplog.text class TestCumulativeStreamingCoercion: diff --git a/tests/examples/offline_inference/test_image_task_prompts.py b/tests/examples/offline_inference/test_image_task_prompts.py index 57e8f1560..054756e33 100644 --- a/tests/examples/offline_inference/test_image_task_prompts.py +++ b/tests/examples/offline_inference/test_image_task_prompts.py @@ -1,17 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +import functools +import importlib.util + import pytest from PIL import Image -from examples.offline_inference.image_to_video.image_to_video import build_image_to_video_prompt -from examples.offline_inference.text_to_image.text_to_image import build_text_to_image_prompt +from tests.examples.helpers import EXAMPLES pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@functools.cache +def _load_example_module(relative_path: str, module_name: str): + path = EXAMPLES / relative_path + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + @pytest.mark.parametrize( ("negative_prompt", "expected_negative_prompt"), [(None, None), ("", ""), ("blurry", "blurry")], @@ -20,7 +32,11 @@ def test_text_to_image_builds_canonical_prompt( negative_prompt: str | None, expected_negative_prompt: str | None, ) -> None: - result = build_text_to_image_prompt("a red fox", negative_prompt) + mod = _load_example_module( + "offline_inference/text_to_image/text_to_image.py", + "text_to_image_example", + ) + result = mod.build_text_to_image_prompt("a red fox", negative_prompt) assert result["prompt"] == "a red fox" assert result["modalities"] == ["image"] @@ -31,9 +47,13 @@ def test_text_to_image_builds_canonical_prompt( def test_image_to_video_builds_canonical_prompt() -> None: + mod = _load_example_module( + "offline_inference/image_to_video/image_to_video.py", + "image_to_video_example", + ) image = Image.new("RGB", (32, 16), "red") - result = build_image_to_video_prompt( + result = mod.build_image_to_video_prompt( prompt="the fox turns toward the camera", negative_prompt="flicker", media_inputs={"image": image}, @@ -45,3 +65,38 @@ def test_image_to_video_builds_canonical_prompt() -> None: "multi_modal_data": {"image": image}, "negative_prompt": "flicker", } + + +@pytest.mark.parametrize( + ("relative_path", "module_name"), + [ + ("offline_inference/text_to_video/text_to_video.py", "text_to_video_output_example"), + ("offline_inference/image_to_video/image_to_video.py", "image_to_video_output_example"), + ], +) +@pytest.mark.parametrize("output_type", ["video", "image"]) +def test_video_examples_accept_semantic_and_legacy_output_types( + relative_path: str, + module_name: str, + output_type: str, +) -> None: + mod = _load_example_module(relative_path, module_name) + + mod._validate_video_output_type(output_type) + + +@pytest.mark.parametrize( + ("relative_path", "module_name"), + [ + ("offline_inference/text_to_video/text_to_video.py", "text_to_video_invalid_output_example"), + ("offline_inference/image_to_video/image_to_video.py", "image_to_video_invalid_output_example"), + ], +) +def test_video_examples_reject_unrelated_output_types( + relative_path: str, + module_name: str, +) -> None: + mod = _load_example_module(relative_path, module_name) + + with pytest.raises(ValueError, match="Unexpected output type 'audio'"): + mod._validate_video_output_type("audio") diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py index 9b5597c3e..5ee2a2fae 100644 --- a/tests/examples/offline_inference/test_text_to_image.py +++ b/tests/examples/offline_inference/test_text_to_image.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Offline inference tests: text-to-image. See examples/offline_inference/text_to_image/README.md @@ -26,9 +29,8 @@ EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_t2i" def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, str]: if h2_title == "Web UI Demo": return True, f"README section '{h2_title}' is intentionally excluded for examples tests" - # Krea-2 Turbo/Raw README snippets currently fail in CI (more_cli_examples_006/007). - if "krea/Krea-2-Turbo" in code or "krea/Krea-2-Raw" in code: - return True, "Krea-2 README snippets are temporarily skipped (CI failures)" + if "krea/Krea-2-Raw" in code: + return True, "krea/Krea-2-Raw is gated on Hugging Face Hub; CI has no access" return False, "" diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py index ff7244d65..c588b7a38 100644 --- a/tests/examples/offline_inference/test_text_to_video_prompts.py +++ b/tests/examples/offline_inference/test_text_to_video_prompts.py @@ -3,19 +3,27 @@ from __future__ import annotations +import functools +import importlib.util + import pytest import torch -from examples.offline_inference.text_to_video.text_to_video import ( - _MODEL_PRESETS, - _detect_preset, - _normalize_float_tensor, - build_text_to_video_prompt, -) +from tests.examples.helpers import EXAMPLES pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@functools.cache +def _load_text_to_video(): + path = EXAMPLES / "offline_inference" / "text_to_video" / "text_to_video.py" + spec = importlib.util.spec_from_file_location("text_to_video_example", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + @pytest.mark.parametrize( ("negative_prompt", "expected_negative_prompt"), [(None, None), ("", ""), ("blurry", "blurry")], @@ -24,7 +32,8 @@ def test_text_to_video_builds_canonical_prompt( negative_prompt: str | None, expected_negative_prompt: str | None, ) -> None: - result = build_text_to_video_prompt("a robot waves", negative_prompt) + mod = _load_text_to_video() + result = mod.build_text_to_video_prompt("a robot waves", negative_prompt) assert result["prompt"] == "a robot waves" assert result["modalities"] == ["video"] @@ -37,6 +46,9 @@ def test_text_to_video_builds_canonical_prompt( @pytest.mark.parametrize( ("model", "model_class_name", "preset_name"), [ + ("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", None, "sana_480p"), + ("Efficient-Large-Model/SANA-Video_2B_720p_diffusers", None, "sana_720p"), + ("/models/custom-checkpoint", "SanaVideoPipeline", "sana_480p"), ("robbyant/lingbot-video-dense-1.3b", None, "lingbot"), ("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"), ("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"), @@ -50,11 +62,13 @@ def test_detect_preset_is_scoped_to_model_family( model_class_name: str | None, preset_name: str, ) -> None: - assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name] + mod = _load_text_to_video() + assert mod._detect_preset(model, model_class_name) is mod._MODEL_PRESETS[preset_name] def test_lingbot_preset_matches_lingbot_defaults() -> None: - assert _MODEL_PRESETS["lingbot"] == { + mod = _load_text_to_video() + assert mod._MODEL_PRESETS["lingbot"] == { "model_class_name": "LingBotVideoPipeline", "height": 192, "width": 320, @@ -68,24 +82,27 @@ def test_lingbot_preset_matches_lingbot_defaults() -> None: def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None: + mod = _load_text_to_video() frames = torch.tensor([0.0, 0.25, 1.0, 1.5]) - assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0])) + assert torch.equal(mod._normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0])) def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None: + mod = _load_text_to_video() frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5]) assert torch.equal( - _normalize_float_tensor(frames, "negative_one_to_one"), + mod._normalize_float_tensor(frames, "negative_one_to_one"), torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]), ) def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None: + mod = _load_text_to_video() frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])] - normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames] + normalized = [mod._normalize_float_tensor(frame, "negative_one_to_one") for frame in frames] assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0])) assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0])) diff --git a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py index 93fc4db38..dfb7b32b1 100644 --- a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py +++ b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py @@ -1,3 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio import base64 import importlib.util import sys @@ -169,6 +173,87 @@ def test_open_streaming_response_requires_post_commit_drain(): ) +def test_final_input_commit_checkpoints_playback_first(): + demo = _load_demo_module() + client = demo.RealtimeDuplexClient("ws://unused") + client.events.add( + { + "type": "response.created", + "response": {"id": "resp-zero-audio"}, + } + ) + calls = [] + + async def send(event): + calls.append(event) + + client.send = send + commit_sent_at_s = asyncio.run(demo._commit_input_after_playback_checkpoint(client)) + + assert [event["type"] for event in calls] == ["playback.ack", "input_audio_buffer.commit"] + assert calls[0]["response_id"] == "resp-zero-audio" + assert calls[0]["played_ms"] == 0 + assert isinstance(commit_sent_at_s, float) + + +def _asymmetric_image(): + from PIL import Image + + image = Image.new("RGB", (4, 2), (0, 0, 0)) + image.putpixel((0, 0), (255, 0, 0)) + return image + + +@pytest.mark.parametrize( + "rotation,expected_corner", + [ + (0, (0, 0)), + (90, (0, 3)), + (180, (3, 1)), + (-90, (1, 0)), + (270, (1, 0)), + (-270, (0, 3)), + ], +) +def test_upright_applies_the_display_matrix_quarter_turn(rotation, expected_corner): + demo = _load_demo_module() + + rotated = demo._upright(_asymmetric_image(), rotation) + + assert rotated.getpixel(expected_corner) == (255, 0, 0) + + +def test_decoded_frames_apply_the_container_display_matrix(monkeypatch, tmp_path): + """A -90 phone clip must reach the model upright, not on its side.""" + demo = _load_demo_module() + source = _asymmetric_image() + + class _Frame: + rotation = -90 + + def to_image(self): + return source + + class _Container: + streams = type("_Streams", (), {"video": [type("_Stream", (), {})()]})() + + def decode(self, stream): + yield _Frame() + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + monkeypatch.setitem(sys.modules, "av", type("_Av", (), {"open": staticmethod(lambda path: _Container())})) + + (image,) = demo._decode_video_frames_rgb(tmp_path / "clip.mp4", [0]) + + assert image.size == (2, 4) + assert image.getpixel((1, 0)) == (255, 0, 0) + + def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, capsys): demo = _load_demo_module() writer = demo._StreamingOutputWriter(tmp_path) diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index 9929beab8..9fa313409 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -1,6 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Shared, importable test helper utilities. -Submodules (``assertions``, ``env``, ``media``, ``runtime``, …) are imported +Submodules (``assertions``, ``clean``, ``client``, ``media``, ``runtime``, …) are imported explicitly by callers. Unit tests for these helpers live under ``tests/``. Avoid star-importing everything here: that ran before refactor only inside the old monolithic ``conftest``; a greedy ``__init__`` changes import order and can @@ -11,11 +14,16 @@ affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based import pytest -def skip_if_gated_repo_inaccessible(repo_id: str) -> None: +def skip_if_gated_repo_inaccessible( + repo_id: str, + *, + revision: str | None = None, + filename: str = "config.json", +) -> None: """Skip the test if a gated HuggingFace repo is not accessible. - Tries to download the model's config.json via ``hf_hub_download``, - which performs an actual file-access check (unlike ``HfApi().model_info()`` + Tries to download ``filename`` via ``hf_hub_download``, which performs an + actual file-access check (unlike ``HfApi().model_info()`` that only checks metadata). If the token has metadata access but not file-download access, ``hf_hub_download`` will raise ``GatedRepoError`` and we skip cleanly. @@ -26,7 +34,7 @@ def skip_if_gated_repo_inaccessible(repo_id: str) -> None: except Exception: return try: - hf_hub_download(repo_id=repo_id, filename="config.json") + hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) except GatedRepoError as exc: pytest.skip( f"Skipping: gated HF repo {repo_id!r} inaccessible to the current " diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 3b6fbe137..755cf0a4e 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -1,16 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Assertion and response validation helpers for tests.""" +import base64 import io import json +import math import tempfile import threading import wave +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from io import BytesIO from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, TypeVar if TYPE_CHECKING: - from tests.helpers.runtime import DiffusionResponse + from tests.helpers.client import DiffusionResponse import av import numpy as np @@ -20,12 +28,112 @@ from PIL import Image from tests.helpers.media import ( convert_audio_bytes_to_text, cosine_similarity_text, - decode_b64_image, preprocess_text, ) +_ResponseT = TypeVar("_ResponseT") + _GENDER_PIPELINE = None _GENDER_PIPELINE_LOCK = threading.Lock() + +# Assertions that sample a random variable rather than testing whether serving works: +# whether TTS audio is intelligible enough to transcribe back, and whether a preset voice's +# timbre estimates as the expected gender. Both are decided by the model and the estimator, +# so a caller measuring a success rate has to tell them apart from a serving failure. The +# messages are constants because a caller matching the wording by hand would silently stop +# matching if it were reworded, reclassifying every quality failure as a hard failure. +AUDIO_MISMATCH_MESSAGE = "The audio content is not same as the text" +GENDER_MISMATCH_MESSAGE = "estimated gender is" +_QUALITY_FAILURE_MESSAGES = (AUDIO_MISMATCH_MESSAGE, GENDER_MISMATCH_MESSAGE) + + +@dataclass(frozen=True) +class SuccessRateGate: + """Gate a sampled quality check on how many requests pass, not on each one passing. + + Built by the send helpers from their own arguments, so a test states a policy and never + touches this class. See ``collect_at_success_rate``. + + Attributes: + min_successes: Fail below this many successes out of ``request_num``. Predeclare it + from a measured baseline; tuning it until CI passes defeats the gate. + concurrency: In-flight requests, in batches until the sample count is reached. + Defaults to sending all of them at once. + """ + + min_successes: int + concurrency: int | None = None + + +def is_quality_failure(exc: BaseException) -> bool: + """True when ``exc`` is one of the sampled quality assertions above. + + A caller absorbing a few of these into a success-rate budget must not absorb an HTTP + failure, missing audio or a text-keyword miss as well, or a broken server passes. + """ + if not isinstance(exc, AssertionError): + return False + return any(message in str(exc) for message in _QUALITY_FAILURE_MESSAGES) + + +def collect_at_success_rate( + sample: Callable[[], _ResponseT], + gate: SuccessRateGate, + *, + request_num: int, + report: Callable[[str], None] = print, +) -> list[_ResponseT]: + """Send ``request_num`` samples and gate on how many passed, returning the successes. + + ``sample`` sends one request and asserts it. Only the sampled quality assertions are + counted; anything else propagates immediately. Does not retry and does not exit early, + since both would bias the rate being measured. + """ + if gate.min_successes > request_num: + raise ValueError(f"min_successes={gate.min_successes} exceeds request_num={request_num}, gate can never pass") + concurrency = min(gate.concurrency or request_num, request_num) + responses: list[_ResponseT] = [] + failures: list[str] = [] + lock = threading.Lock() + + def _one_sample(_: int) -> None: + try: + response = sample() + except AssertionError as exc: + if not is_quality_failure(exc): + raise + with lock: + failures.append(str(exc)) + return + with lock: + responses.append(response) + + for start in range(0, request_num, concurrency): + size = min(concurrency, request_num - start) + with ThreadPoolExecutor(max_workers=size) as executor: + for _ in executor.map(_one_sample, range(size)): + pass + + successes = request_num - len(failures) + report(f"success rate: {successes}/{request_num} succeeded") + assert successes >= gate.min_successes, ( + f"success rate {successes}/{request_num} is below the " + f"{gate.min_successes}/{request_num} gate " + f"(95% CI {wilson_interval(successes, request_num)}). Failures:\n" + + "\n".join(f" - {failure}" for failure in failures) + ) + return responses + + +def wilson_interval(successes: int, total: int, z: float = 1.96) -> str: + """Format a 95% Wilson score interval, so a rate failure shows how marginal it is.""" + p = successes / total + denom = 1 + z * z / total + centre = (p + z * z / (2 * total)) / denom + half = z / denom * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total)) + return f"{max(0.0, centre - half):.1%}-{min(1.0, centre + half):.1%}" + + # Transcript gates default to whisper ``small`` for speed. ``small`` mishears a # short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a # leading SFX token), which flakes the deterministic similarity gate. Short @@ -66,7 +174,7 @@ def _short_transcript_contains_expected(transcript: str, expected: str) -> bool: def assert_image_diffusion_response( response: "DiffusionResponse", request_config: dict[str, Any], - run_level: str = None, + run_level: str | None = None, ) -> None: """ Validate image diffusion response. @@ -138,14 +246,15 @@ def assert_images_generations_response( assert isinstance(item, dict), "Image generation data entries must be objects" b64_json = item.get("b64_json") assert isinstance(b64_json, str) and b64_json, "Image generation response is missing b64_json" - image = decode_b64_image(b64_json) + image = Image.open(io.BytesIO(base64.b64decode(b64_json))) + image.load() assert_image_valid(image, width=width, height=height) def assert_video_diffusion_response( response: "DiffusionResponse", request_config: dict[str, Any], - run_level: str = None, + run_level: str | None = None, ) -> None: """ Validate video diffusion response. @@ -214,7 +323,7 @@ def assert_video_first_frame_matches( def assert_audio_diffusion_response( response: "DiffusionResponse", request_config: dict[str, Any], - run_level: str = None, + run_level: str | None = None, ) -> None: """ Validate audio diffusion response. @@ -474,7 +583,7 @@ def _assert_preset_voice_gender_from_audio( print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}") if estimated_gender != "unknown": assert estimated_gender == expected_gender, ( - f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}" + f"{voice_name!r} is expected {expected_gender}, but {GENDER_MISMATCH_MESSAGE} {estimated_gender!r}" ) @@ -529,7 +638,7 @@ def _response_has_audio_output(response: Any) -> bool: return bool(audio_data) -def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool: +def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str | None) -> bool: if run_level not in {"advanced_model", "full_model"}: return False modalities = request_config.get("modalities", ["text", "audio"]) @@ -551,18 +660,18 @@ def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_l return "text" in modalities -def _speech_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool: +def _speech_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str | None) -> bool: if run_level not in {"advanced_model", "full_model"}: return False if request_config.get("response_format") == "pcm": return False - return bool(request_config.get("input")) + return bool(request_config.get("transcript_expected_text", request_config.get("input"))) def _resolve_audio_transcript( response: Any, request_config: dict[str, Any], - run_level: str, + run_level: str | None, *, speech_api: bool, ) -> str | None: @@ -655,7 +764,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve shorter_clean = _re.sub(r"[^\w\s]", "", shorter).strip() longer_clean = _re.sub(r"[^\w\s]", "", longer).strip() assert shorter_clean and (shorter_clean in longer_clean), ( - f"The audio content is not same as the text " + f"{AUDIO_MISMATCH_MESSAGE} " f"(short-text containment check failed: " f"text={text_output!r}, transcript={transcript!r})" ) @@ -666,7 +775,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve text_output.lower(), ) print(f"similarity is: {similarity}") - assert similarity > similarity_threshold, "The audio content is not same as the text" + assert similarity > similarity_threshold, AUDIO_MISMATCH_MESSAGE if audio_ref_text: assert transcript is not None, "No audio transcript for reference-text validation" audio_similarity = cosine_similarity_text( @@ -737,28 +846,13 @@ def _assert_transcript_matches( ) -def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None: - """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`. +def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str | None = None) -> None: + """Validate speech API results from :class:`~tests.helpers.client.OmniResponse`. - When ``request_config`` carries ``status_code`` and/or ``err_message``, the - request is expected to be rejected: assert it failed and that the HTTP status - / error text match. Otherwise the normal success-path checks run. + Success-path checks only. Negative / contract cases belong on + :meth:`~tests.helpers.client.OnlineOmniClient.send_audio_speech_http_request` + with :func:`assert_http_error`. """ - expected_status = request_config.get("status_code") - expected_err = request_config.get("err_message") - if expected_status is not None or expected_err is not None: - assert not response.success, "Expected an error response, but the request succeeded." - if expected_status is not None: - allowed = expected_status if isinstance(expected_status, (list, tuple)) else (expected_status,) - assert response.status_code in allowed, f"Expected HTTP status in {allowed}, got {response.status_code}" - if expected_err is not None: - alternatives = expected_err if isinstance(expected_err, (list, tuple)) else (expected_err,) - error_text = response.error_message or "" - assert any(alt in error_text for alt in alternatives), ( - f"Expected one of {alternatives} in error text, got: {error_text!r}" - ) - return - assert response.success, "The request failed." # Optional floor on decoded audio size (models with very short clips may use a lower value). @@ -786,7 +880,7 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True) if transcript is not None: - expected_text = request_config.get("input") + expected_text = request_config.get("transcript_expected_text", request_config.get("input")) if expected_text: print(f"audio content is: {transcript}") print(f"input text is: {expected_text}") @@ -805,7 +899,9 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], ) -def assert_diffusion_response(response: "DiffusionResponse", request_config: dict[str, Any], run_level: str = None): +def assert_diffusion_response( + response: "DiffusionResponse", request_config: dict[str, Any], run_level: str | None = None +): assert response.success, "The request failed." has_any_content = any(content is not None for content in (response.images, response.videos, response.audios)) assert has_any_content, "Response contains no images, videos, or audios" @@ -871,9 +967,9 @@ def assert_http_error( err_message: str | tuple[str, ...] | list[str] | None = None, websocket_json_message: bool = False, ) -> dict[str, Any] | None: - """Validate a raw-HTTP :class:`~tests.helpers.runtime.HttpResponse`-like object. + """Validate a raw-HTTP :class:`~tests.helpers.client.HttpResponse`-like object. - Used by :class:`~tests.helpers.runtime.OpenAIClientHandler` ``send_*_http_request`` helpers when + Used by :class:`~tests.helpers.client.OnlineOmniClient` ``send_*_http_request`` helpers when ``request_config`` contains optional ``err_code`` and/or ``err_message``. When ``websocket_json_message=True``, only ``json_body`` is checked (first JSON WebSocket text frame). diff --git a/tests/helpers/clean.py b/tests/helpers/clean.py new file mode 100644 index 000000000..17fef825d --- /dev/null +++ b/tests/helpers/clean.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Test cleanup helpers (device memory, distributed state, and memory monitoring). + +``vllm_omni.platforms`` is imported only inside functions that need it so importing this module +at pytest plugin load does not run before session autouse fixtures. +""" + +from __future__ import annotations + +import gc +import logging +import os +import subprocess +import time + +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + + +def get_physical_device_indices(devices): + """Map logical device indices to physical IDs via the platform visibility env. + + Uses ``current_omni_platform.device_control_env_var`` (e.g. ``CUDA_VISIBLE_DEVICES``, + ``ASCEND_RT_VISIBLE_DEVICES``) rather than hard-coding CUDA. + """ + visible_devices = os.environ.get(current_omni_platform.device_control_env_var) + if visible_devices is None: + return devices + visible_indices = [int(x) for x in visible_devices.split(",") if x.strip() != ""] + index_mapping = {i: physical for i, physical in enumerate(visible_indices)} + return [index_mapping[i] for i in devices if i in index_mapping] + + +def pick_least_used_device_indices(num_devices: int) -> list[int]: + """Return physical device indices for the least-used accelerators. + + Args: + num_devices: Number of device indices to return. + + Returns: + Physical device indices sorted from least to most used memory. + """ + if num_devices <= 0: + raise ValueError(f"num_devices must be positive, got {num_devices}.") + + logical_count = current_omni_platform.get_device_count() + if logical_count < num_devices: + raise RuntimeError(f"Need {num_devices} devices, but only {logical_count} are available.") + + device_usage: list[tuple[int, int]] = [] + for device_index in range(logical_count): + with current_omni_platform.device(device_index): + free_bytes, total_bytes = current_omni_platform.get_device_memory() + device_usage.append((total_bytes - free_bytes, device_index)) + device_usage.sort() + logical_indices = [device_index for _, device_index in device_usage[:num_devices]] + return get_physical_device_indices(logical_indices) + + +def wait_for_gpu_memory_to_clear( + *, + devices: list[int], + threshold_bytes: int | None = None, + threshold_ratio: float | None = None, + timeout_s: float = 120, +) -> None: + assert threshold_bytes is not None or threshold_ratio is not None + devices = get_physical_device_indices(devices) + start_time = time.time() + + device_list = ", ".join(str(d) for d in devices) + if threshold_bytes is not None: + condition_str = f"Memory usage ≤ {threshold_bytes / 2**30:.2f} GiB" + + def is_free(used, total): + return used <= threshold_bytes / 2**30 + else: + ratio = threshold_ratio + assert ratio is not None + condition_str = f"Memory usage ratio ≤ {ratio * 100:.1f}%" + + def is_free(used, total): + return used / total <= ratio + + print(f"[Device Memory Monitor] Waiting for device(s) {device_list} to free memory, Condition: {condition_str}") + + def get_mem_gib(device: int) -> tuple[float, float]: + with current_omni_platform.device(device): + free_bytes, total_bytes = current_omni_platform.mem_get_info() + return (total_bytes - free_bytes) / 2**30, total_bytes / 2**30 + + while True: + output_raw = {d: get_mem_gib(d) for d in devices} + output = { + d: f"{used:.1f}GiB/{total:.1f}GiB ({(used / total) * 100 if total > 0 else 0:.1f}%)" + for d, (used, total) in output_raw.items() + } + + print("[Device Memory Status] Current usage:") + for device_id, mem_info in output.items(): + print(f" Device {device_id}: {mem_info}") + + dur_s = time.time() - start_time + if all(is_free(used, total) for used, total in output_raw.values()): + print(f"[Device Memory Freed] Device(s) {device_list} meet memory condition") + print(f" Condition: {condition_str}") + print(f" Wait time: {dur_s:.1f} seconds ({dur_s / 60:.1f} minutes)") + break + + if dur_s >= timeout_s: + raise ValueError( + f"[Device Memory Timeout] Device(s) {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n" + f"Condition: {condition_str}\n" + f"Current status:\n" + "\n".join(f" Device {d}: {output[d]}" for d in devices) + ) + + gc.collect() + current_omni_platform.empty_cache() + time.sleep(5) + + +def _run_smi(label: str, cmd: list[str], head_lines: int, timeout: float = 5) -> None: + print("\n" + "=" * 80) + print(label) + print("=" * 80) + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) + if result.returncode == 0 and result.stdout.strip(): + lines = result.stdout.strip().split("\n") + for line in lines[:head_lines]: + print(line) + if len(lines) > head_lines: + print(f"... (showing first {head_lines} of {len(lines)} lines)") + else: + print(f"{cmd[0]} command failed or produced no output") + except (subprocess.TimeoutExpired, FileNotFoundError): + print(f"{cmd[0]} not available or timed out") + except Exception as e: + print(f"Error running {' '.join(cmd)}: {e}") + + +def _print_device_processes() -> None: + """Print device information via platform SMI tools.""" + if current_omni_platform.is_cuda(): + _run_smi("NVIDIA GPU Information (nvidia-smi)", ["nvidia-smi"], 20) + _run_smi("Detailed GPU Processes (nvidia-smi pmon)", ["nvidia-smi", "pmon", "-c", "1"], 100, timeout=3) + elif current_omni_platform.is_npu(): + _run_smi("Ascend NPU Information (npu-smi info)", ["npu-smi", "info"], 40) + elif current_omni_platform.is_rocm(): + # amd-smi can enter uninterruptible sleep in the amdgpu CPER ioctl + # (amdgpu_cper_ring_write). In that state subprocess.run's timeout + # cannot reap it, causing this diagnostic hook to hang the entire test. + # rocm-smi provides the information needed here without that ioctl. + _run_smi( + "AMD GPU Information (rocm-smi)", + ["rocm-smi", "--showuse", "--showmeminfo", "vram"], + 60, + ) + _run_smi("Detailed AMD GPU Processes (rocm-smi)", ["rocm-smi", "--showpids"], 100, timeout=3) + elif current_omni_platform.is_xpu(): + _run_smi("Intel XPU Information (xpu-smi discovery)", ["xpu-smi", "discovery"], 40) + elif current_omni_platform.is_musa(): + _run_smi("Moore Threads GPU Information (mthreads-gmi)", ["mthreads-gmi"], 30) + else: + print("\n" + "=" * 80) + print("WARNING: No supported device platform detected") + print("=" * 80) + + +def _cleanup_stale_device_locks() -> None: + """Remove stale device-initialization lock files whose recorded PID is dead. + + Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a + crashed / killed test run and block subsequent orchestrator startups. + """ + import glob as _glob + + for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"): + try: + with open(lock_file) as fh: + content = fh.read().strip() + if not content: + continue + pid = int(content) + except (OSError, ValueError): + continue + + try: + os.kill(pid, 0) + except ProcessLookupError: + print(f" Removing stale device lock {lock_file} (PID {pid} is dead)") + try: + os.unlink(lock_file) + except OSError: + pass + except PermissionError: + pass + + +def cleanup_test_environment(*, shutdown_ray: bool = False) -> None: + """Tear down distributed state and reset device memory for tests.""" + from vllm import envs + from vllm.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + ) + + envs.disable_envs_cache() + + if current_omni_platform.is_rocm(): + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.refresh_env_variables() + + gc.unfreeze() + destroy_model_parallel() + destroy_distributed_environment() + if shutdown_ray: + import ray + + ray.shutdown() + + print("Pre-test device status:") + _cleanup_stale_device_locks() + + num_devices = current_omni_platform.device_count() + if num_devices > 0: + try: + wait_for_gpu_memory_to_clear( + devices=list(range(num_devices)), + threshold_ratio=0.05, + timeout_s=60, + ) + except Exception as e: + print(f"Device cleanup note: {e}") + + gc.collect() + if not current_omni_platform.is_cpu(): + current_omni_platform.empty_cache() + try: + import torch + + torch._C._host_emptyCache() + except AttributeError: + logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5") + + if current_omni_platform.is_available(): + print("Post-test device status:") + _print_device_processes() + + +__all__ = [ + "cleanup_test_environment", + "get_physical_device_indices", + "pick_least_used_device_indices", + "wait_for_gpu_memory_to_clear", +] diff --git a/tests/helpers/client.py b/tests/helpers/client.py new file mode 100644 index 000000000..9ef4a7a0c --- /dev/null +++ b/tests/helpers/client.py @@ -0,0 +1,2148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Online/offline request clients and response types for tests.""" + +from __future__ import annotations + +import asyncio +import base64 +import concurrent.futures +import copy +import io +import json +import time +from dataclasses import asdict, dataclass +from io import BytesIO +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast +from urllib.parse import quote + +import numpy as np +import requests +import soundfile as sf +import torch +from openai import OpenAI, omit +from PIL import Image +from vllm.logger import init_logger + +from tests.helpers.assertions import ( + SuccessRateGate, + assert_audio_speech_response, + assert_diffusion_response, + assert_http_error, + assert_images_generations_response, + assert_omni_response, + collect_at_success_rate, +) +from tests.helpers.media import ( + _merge_base64_audio_to_segment, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.outputs import OmniRequestOutput + +if TYPE_CHECKING: + from tests.helpers.runtime import OmniRunner + +logger = init_logger(__name__) + + +def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None: + try: + data = r.json() + if isinstance(data, (dict, list)): + return data + except Exception: + pass + return None + + +def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None: + """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``).""" + eb = cfg.get("extra_body") + if not eb: + return None + h, w = eb.get("height"), eb.get("width") + if (isinstance(h, (list, tuple)) or isinstance(w, (list, tuple))) and not ( + isinstance(h, (list, tuple)) and isinstance(w, (list, tuple)) + ): + raise ValueError("extra_body height and width must both be lists or both be scalars") + if not (isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))): + return None + if len(h) != len(w): + raise ValueError(f"height and width lists must have equal length; got {len(h)=} {len(w)=}") + n = len(h) + n_out = eb.get("num_outputs_per_prompt") + if n_out is not None: + n_out = int(n_out) + if n_out != n: + raise ValueError( + "When height/width are lists, num_outputs_per_prompt must equal their length; " + f"got num_outputs_per_prompt={n_out}, len(lists)={n}" + ) + splits: list[dict[str, Any]] = [] + for i in range(n): + sub = copy.deepcopy(cfg) + sub_eb = dict(sub.get("extra_body") or {}) + sub_eb["height"] = int(h[i]) + sub_eb["width"] = int(w[i]) + sub_eb["num_outputs_per_prompt"] = 1 + sub["extra_body"] = sub_eb + splits.append(sub) + return splits + + +@dataclass +class OmniResponse: + """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``).""" + + text_content: str | None = None + audio_data: list[str] | None = None + audio_content: str | None = None + audio_format: str | None = None + audio_bytes: bytes | None = None + #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the + #: OpenAI client call through response parsing and local post-process (e.g. audio decode). + e2e_latency: float | None = None + success: bool = False + prompt_tokens: int | None = None + cached_tokens: int | None = None + multimodal_tokens: dict[str, int] | None = None + logprobs: list | None = None + + +@dataclass +class DiffusionResponse: + """Decoded diffusion output from chat completions or offline runner (not raw ``requests``).""" + + text_content: str | None = None + images: list[Image.Image] | None = None + audios: list[Any] | None = None + videos: list[Any] | None = None + #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before + #: ``chat.completions.create`` through local image / audio decode. + e2e_latency: float | None = None + success: bool = False + + +@dataclass +class HttpResponse: + """Normalized view of a ``requests`` response from :class:`OnlineOmniClient` HTTP helpers.""" + + status_code: int + success: bool + error_message: str | None = None + json_body: dict[str, Any] | list[Any] | None = None + + +@dataclass +class WebSocketJsonResponse: + """First JSON object delivered as a text WebSocket frame (streaming endpoints).""" + + json_body: dict[str, Any] + + +@dataclass +class OpenPIWebSocketResponse: + """Msgpack WebSocket session against ``/v1/realtime/robot/openpi``.""" + + server_metadata: dict[str, Any] + operation_responses: list[Any] + actions: dict[str, np.ndarray] | None = None + action_tensors: list[np.ndarray] | None = None + + +def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]: + """Build a minimal DROID-style observation payload for the OpenPI robot endpoint.""" + identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32) + identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) + return { + "session_id": session_id, + "video": { + "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), + "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), + }, + "state": { + "eef_9d": identity_eef_9d, + "gripper_position": np.zeros((1, 1, 1), dtype=np.float32), + "joint_position": np.zeros((1, 1, 7), dtype=np.float32), + }, + "language": {"annotation.language.language_instruction": [["pick up the object"]]}, + } + + +DREAMZERO_DEFAULT_PROMPT = ( + "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan" +) +DREAMZERO_ACTION_HORIZON = 24 +DREAMZERO_ACTION_DIM = 8 +DREAMZERO_CAMERA_FILES = { + "observation/exterior_image_0_left": "exterior_image_1_left.mp4", + "observation/exterior_image_1_left": "exterior_image_2_left.mp4", + "observation/wrist_image_left": "wrist_image_left.mp4", +} + + +def _require_opencv() -> Any: + try: + import cv2 + except ImportError as exc: # pragma: no cover - optional e2e dependency + raise ModuleNotFoundError("DreamZero OpenPI test dependencies are missing: opencv-python") from exc + return cv2 + + +def load_dreamzero_camera_frames(video_dir: Path) -> dict[str, np.ndarray]: + cv2 = _require_opencv() + camera_frames: dict[str, np.ndarray] = {} + for camera_key, file_name in DREAMZERO_CAMERA_FILES.items(): + video_path = video_dir / file_name + if not video_path.exists(): + raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}") + cap = cv2.VideoCapture(str(video_path)) + frames = [] + while True: + ok, frame = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + cap.release() + if not frames: + raise RuntimeError(f"No frames loaded from {video_path}") + camera_frames[camera_key] = np.stack(frames, axis=0) + return camera_frames + + +def build_dreamzero_demo_observations( + camera_frames: dict[str, np.ndarray], + *, + prompt: str, + session_id: str, + num_chunks: int = 2, +) -> list[dict[str, Any]]: + if num_chunks < 1: + raise ValueError("num_chunks must be at least 1") + + relative_offsets = [-23, -16, -8, 0] + total_frames = min(frames.shape[0] for frames in camera_frames.values()) + frame_schedules = [[0]] + current_frame = 23 + for _ in range(num_chunks - 1): + indices = [max(current_frame + offset, 0) for offset in relative_offsets] + if indices[-1] >= total_frames: + break + frame_schedules.append(indices) + current_frame += DREAMZERO_ACTION_HORIZON + + observations: list[dict[str, Any]] = [] + for frame_indices in frame_schedules: + obs: dict[str, Any] = {} + for camera_key, all_frames in camera_frames.items(): + selected = all_frames[frame_indices] + obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + observations.append(obs) + return observations + + +class OpenPIWebSocketSession: + """Persistent msgpack session for ``/v1/realtime/robot/openpi``.""" + + DEFAULT_PING_INTERVAL_SECS = 300 + DEFAULT_PING_TIMEOUT_SECS = 3600 + + @staticmethod + def _require_dependencies() -> tuple[Any, Any]: + try: + import websockets.sync.client as websockets_client + except ImportError as exc: # pragma: no cover - optional e2e dependency + raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: websockets") from exc + try: + from openpi_client import msgpack_numpy + except ImportError as exc: # pragma: no cover - optional e2e dependency + raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: openpi-client") from exc + return websockets_client, msgpack_numpy + + def __init__( + self, + uri: str, + *, + ping_interval: float = DEFAULT_PING_INTERVAL_SECS, + ping_timeout: float = DEFAULT_PING_TIMEOUT_SECS, + open_timeout: float = 120.0, + close_timeout: float = 120.0, + ) -> None: + websockets_client, msgpack_numpy = self._require_dependencies() + self._msgpack_numpy = msgpack_numpy + self._packer = msgpack_numpy.Packer() + self._conn = websockets_client.connect( + uri, + compression=None, + max_size=None, + ping_interval=ping_interval, + ping_timeout=ping_timeout, + open_timeout=open_timeout, + close_timeout=close_timeout, + ) + server_metadata = msgpack_numpy.unpackb(self._conn.recv()) + if not isinstance(server_metadata, dict): + raise TypeError(f"Expected dict metadata from {uri}, got {type(server_metadata)!r}") + self._server_metadata = server_metadata + + def get_server_metadata(self) -> dict[str, Any]: + return dict(self._server_metadata) + + def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]: + response = self._send_operation("infer", obs) + if not isinstance(response, dict): + raise TypeError(f"Expected dict infer response, got {type(response)!r}") + return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()} + + def reset(self, reset_info: dict[str, Any] | None = None) -> str: + response = self._send_operation("reset", dict(reset_info or {})) + return str(response["status"]) + + def close(self) -> None: + self._conn.close() + + def _send_operation(self, endpoint: str, payload: dict[str, Any]) -> Any: + body = dict(payload) + body["endpoint"] = endpoint + self._conn.send(self._packer.pack(body)) + raw = self._conn.recv() + if isinstance(raw, str): + raise RuntimeError(f"OpenPI {endpoint!r} failed: {raw}") + decoded = self._msgpack_numpy.unpackb(raw) + if isinstance(decoded, dict) and decoded.get("type") == "error": + raise RuntimeError(f"OpenPI {endpoint!r} failed: {decoded.get('message')}") + if endpoint == "reset": + if not isinstance(decoded, dict) or decoded.get("status") != "reset successful": + raise RuntimeError(f"Unexpected OpenPI reset response: {decoded!r}") + return decoded + return decoded + + +def _merge_http_expectation_kwargs( + base: dict[str, Any] | None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, +) -> dict[str, Any]: + cfg = dict(base or {}) + if err_code is not None: + cfg["err_code"] = err_code + if err_message is not None: + cfg["err_message"] = err_message + return cfg + + +def _merge_ws_expectation_kwargs( + base: dict[str, Any] | None, + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, +) -> dict[str, Any]: + cfg = dict(base or {}) + if err_message is not None: + cfg["err_message"] = err_message + if ws_json_type is not None: + cfg["ws_json_type"] = ws_json_type + if ws_error_code is not None: + cfg["ws_error_code"] = ws_error_code + return cfg + + +def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None: + jb = resp.json_body + want_type = cfg.get("ws_json_type") + if want_type is not None: + assert jb.get("type") == want_type, (jb, want_type) + want_code = cfg.get("ws_error_code") + if want_code is not None: + assert jb.get("code") == want_code, (jb, want_code) + err_message = cfg.get("err_message") + if err_message is not None: + assert_http_error(resp, err_message=err_message, websocket_json_message=True) + + +def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse: + """Concatenate images in order; ``e2e_latency`` is wall-clock of the batch (set by caller) or max of parts.""" + merged = DiffusionResponse() + merged.success = all(p.success for p in parts) and len(parts) > 0 + imgs: list[Image.Image] = [] + for p in parts: + if p.images: + imgs.extend(p.images) + merged.images = imgs if imgs else None + latencies = [p.e2e_latency for p in parts if p.e2e_latency is not None] + merged.e2e_latency = max(latencies) if latencies else None + return merged + + +class OnlineOmniClient: + def __init__( + self, + host: str = "127.0.0.1", + port: int | None = None, + api_key: str = "EMPTY", + run_level: str | None = None, + *, + log_stats: bool = True, + ): + if port is None: + from tests.helpers.runtime import get_open_port + + port = get_open_port() + self.base_url = f"http://{host}:{port}" + self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key) + self.run_level = run_level + self.log_stats = log_stats + + def _print_client_stat(self, message: str) -> None: + if self.log_stats: + print(message, flush=True) + + def _process_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse: + """Wall clock from *before* ``chat.completions.create`` through stream drain + local decode.""" + result = OmniResponse() + try: + text_content = "" + audio_data = [] + for chunk in chat_completion: + for choice in chunk.choices: + content = getattr(getattr(choice, "delta", None), "content", None) + modality = getattr(chunk, "modality", None) + if modality == "audio" and content: + audio_data.append(content) + elif modality == "text" and content: + text_content += content + # Usage is yielded after the last token + if chunk.usage: + result.prompt_tokens = chunk.usage.prompt_tokens + if details := getattr(chunk.usage, "prompt_tokens_details", None): + result.cached_tokens = details.cached_tokens + result.multimodal_tokens = getattr(details, "multimodal_tokens", None) + + if audio_data: + merged_seg = _merge_base64_audio_to_segment(audio_data) + wav_buf = BytesIO() + merged_seg.export(wav_buf, format="wav") + result.audio_bytes = wav_buf.getvalue() + result.text_content = text_content + result.audio_data = audio_data + result.e2e_latency = time.perf_counter() - wall_start + result.success = True + except Exception as e: + msg = f"Stream processing error: {str(e)}" + print(f"Error: {msg}") + return result + + def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse: + """Wall clock from *before* ``chat.completions.create`` through response parse + local decode.""" + result = OmniResponse() + try: + audio_data = None + text_content = None + for choice in chat_completion.choices: + if hasattr(choice.message, "audio") and choice.message.audio is not None: + audio_data = choice.message.audio.data + if hasattr(choice.message, "content") and choice.message.content is not None: + text_content = choice.message.content + # Extract cached & prompt token counts for prefix caching tests + usage = getattr(chat_completion, "usage", None) + if usage: + result.prompt_tokens = usage.prompt_tokens + if details := getattr(usage, "prompt_tokens_details", None): + result.cached_tokens = details.cached_tokens + result.multimodal_tokens = getattr(details, "multimodal_tokens", None) + if audio_data: + result.audio_bytes = base64.b64decode(audio_data) + result.text_content = text_content + result.e2e_latency = time.perf_counter() - wall_start + if chat_completion.choices and chat_completion.choices[0].logprobs is not None: + result.logprobs = chat_completion.choices[0].logprobs.content + result.success = True + except Exception as e: + msg = f"Non-stream processing error: {str(e)}" + print(f"Error: {msg}") + return result + + def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse: + """Wall clock from *before* ``chat.completions.create`` through image decode.""" + result = DiffusionResponse() + try: + images = [] + audios = [] + for choice in chat_completion.choices: + content = getattr(choice.message, "content", None) + if isinstance(content, list): + for item in content: + image_url = None + if isinstance(item, dict): + image_url = item.get("image_url", {}).get("url") + else: + image_url_obj = getattr(item, "image_url", None) + image_url = getattr(image_url_obj, "url", None) if image_url_obj else None + if image_url and image_url.startswith("data:image"): + b64_data = image_url.split(",", 1)[1] + img = Image.open(io.BytesIO(base64.b64decode(b64_data))) + img.load() + images.append(img) + + # OpenAI audio responses (e.g. AudioX text-to-audio) populate `message.audio`. + audio_obj = getattr(choice.message, "audio", None) + audio_b64 = getattr(audio_obj, "data", None) if audio_obj is not None else None + if audio_b64: + audios.append( + { + "wav_bytes": base64.b64decode(audio_b64), + "id": getattr(audio_obj, "id", None), + "expires_at": getattr(audio_obj, "expires_at", None), + } + ) + result.images = images if images else None + result.audios = audios if audios else None + result.e2e_latency = time.perf_counter() - wall_start + result.success = True + except Exception as e: + msg = f"Diffusion response processing error: {str(e)}" + print(f"Error: {msg}") + return result + + def _http_response_from_requests(self, r: requests.Response) -> HttpResponse: + payload = _parse_response_json(r) + ok = 200 <= r.status_code < 300 + return HttpResponse( + status_code=r.status_code, + success=ok, + error_message=None if ok else (r.text[:8000] if r.text else None), + json_body=payload, + ) + + def send_health_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/health`` (raw ``requests``). + + ``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for + :func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args). + """ + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0))) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_models_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`).""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get( + self._build_url("/v1/models"), + headers={"Accept": "application/json"}, + timeout=float(cfg.get("timeout", 120.0)), + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_chat_completions_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests).""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_batched_chat_completions_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """Post to batched chat completions.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_completions_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/completions`` with ``json`` or ``raw_body``.""" + # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_omni_sleep_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_omni_wakeup_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/omni/wakeup``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_voices_list_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/audio/voices``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get( + self._build_url("/v1/audio/voices"), + headers={"Accept": "application/json"}, + timeout=float(cfg.get("timeout", 120.0)), + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_voices_create_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_voices_delete_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + name = cfg["name"] + timeout = float(cfg.get("timeout", 120.0)) + path = f"/v1/audio/voices/{quote(str(name), safe='')}" + r = requests.delete( + self._build_url(path), + headers={"Accept": "application/json"}, + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_speech_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/speech`` with ``json`` or ``raw_body``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_speech_batch_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/speech/batch``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_audio_generate_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/audio/generate``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_images_generations_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + if cfg.get("err_code") is None: + assert resp.success, resp.error_message + payload = resp.json_body + assert isinstance(payload, dict) + assert_images_generations_response(payload, cfg, run_level=self.run_level) + return [resp] + + def send_images_edits_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_videos_create_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_videos_sync_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """POST ``/v1/videos/sync``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_videos_list_http_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/videos`` — optional ``params``, ``timeout``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + r = requests.get( + self._build_url("/v1/videos"), + params=cfg.get("params"), + headers={"Accept": "application/json"}, + timeout=float(cfg.get("timeout", 120.0)), + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_video_retrieve_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/videos/{video_id}``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + video_id = cfg["video_id"] + timeout = float(cfg.get("timeout", 120.0)) + r = requests.get( + self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"), + headers={"Accept": "application/json"}, + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_video_delete_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """DELETE ``/v1/videos/{video_id}``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + video_id = cfg["video_id"] + timeout = float(cfg.get("timeout", 120.0)) + r = requests.delete( + self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"), + headers={"Accept": "application/json"}, + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def send_video_content_http_request( + self, + request_config: dict[str, Any], + *, + err_code: int | tuple[int, ...] | list[int] | None = None, + err_message: str | tuple[str, ...] | list[str] | None = None, + ) -> list[HttpResponse]: + """GET ``/v1/videos/{video_id}/content``.""" + cfg = _merge_http_expectation_kwargs( + request_config, + err_code=err_code, + err_message=err_message, + ) + video_id = cfg["video_id"] + timeout = float(cfg.get("timeout", 120.0)) + r = requests.get( + self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"), + timeout=timeout, + ) + resp = self._http_response_from_requests(r) + assert_http_error( + resp, + err_code=cfg.get("err_code"), + err_message=cfg.get("err_message"), + ) + return [resp] + + def _build_ws_url(self, path: str) -> str: + """Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers.""" + base = self.base_url.rstrip("/") + suffix = "/" + path.lstrip("/") + if base.startswith("http://"): + return "ws://" + base.removeprefix("http://") + suffix + if base.startswith("https://"): + return "wss://" + base.removeprefix("https://") + suffix + raise ValueError(f"Unsupported base_url for WebSocket: {base!r}") + + def _send_websocket_first_json_request( + self, + path: str, + cfg: dict[str, Any], + ) -> list[WebSocketJsonResponse]: + """Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`. + + ``request_config`` keys: + + - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server + speaks first, e.g. ``/v1/realtime`` rejection path). + - ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame + (e.g. ``["session.created"]`` on ``/v1/realtime``). + - ``timeout``: seconds to wait for the first inbound text frame (default ``120``). + - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present. + """ + send_frames_raw = cfg.get("send_frames") + if send_frames_raw is None: + frames: list[str] = [] + elif isinstance(send_frames_raw, str): + frames = [send_frames_raw] + else: + frames = list(send_frames_raw) + + timeout = float(cfg.get("timeout", 120.0)) + uri = self._build_ws_url(path) + skip_types = set(cfg.get("ws_skip_types") or []) + + connect_kw: dict[str, Any] = {} + if "ws_max_size" in cfg: + connect_kw["max_size"] = cfg["ws_max_size"] + + async def _recv_first_json_object() -> WebSocketJsonResponse: + import websockets + + async with websockets.connect(uri, **connect_kw) as ws: + for frame in frames: + await ws.send(frame) + while True: + raw = await asyncio.wait_for(ws.recv(), timeout=timeout) + if not isinstance(raw, str): + raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}") + try: + data = json.loads(raw) + except json.JSONDecodeError as exc: + raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc + if not isinstance(data, dict): + raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}") + if skip_types and data.get("type") in skip_types: + continue + return WebSocketJsonResponse(json_body=data) + + resp = asyncio.run(_recv_first_json_object()) + _run_ws_expectations_from_request_config(cfg, resp) + return [resp] + + def send_audio_speech_stream_ws_request( + self, + request_config: dict[str, Any], + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, + ) -> list[WebSocketJsonResponse]: + """WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame.""" + cfg = _merge_ws_expectation_kwargs( + request_config, + err_message=err_message, + ws_json_type=ws_json_type, + ws_error_code=ws_error_code, + ) + return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg) + + def send_video_chat_stream_ws_request( + self, + request_config: dict[str, Any], + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, + ) -> list[WebSocketJsonResponse]: + """WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame.""" + cfg = _merge_ws_expectation_kwargs( + request_config, + err_message=err_message, + ws_json_type=ws_json_type, + ws_error_code=ws_error_code, + ) + return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg) + + def send_realtime_ws_request( + self, + request_config: dict[str, Any] | None = None, + *, + err_message: str | tuple[str, ...] | list[str] | None = None, + ws_json_type: str | None = None, + ws_error_code: str | None = None, + ) -> list[WebSocketJsonResponse]: + """WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated).""" + cfg = _merge_ws_expectation_kwargs( + request_config, + err_message=err_message, + ws_json_type=ws_json_type, + ws_error_code=ws_error_code, + ) + return self._send_websocket_first_json_request("/v1/realtime", cfg) + + def send_robot_openpi_ws_request( + self, + request_config: dict[str, Any] | None = None, + ) -> list[OpenPIWebSocketResponse]: + """WebSocket ``/v1/realtime/robot/openpi`` — msgpack metadata plus optional infer/reset ops. + + ``request_config`` keys: + + - ``operations``: optional sequence of ``{"endpoint": "infer"|"reset", "payload": {...}}``. + - ``run_default_policy_session``: when true and ``operations`` is omitted, run infer then reset + using :func:`build_openpi_droid_observation`. + - ``run_dreamzero_policy_session``: when true and ``operations`` is omitted, run the DreamZero + infer/reset/infer sequence using :func:`build_dreamzero_demo_observations`. + - ``video_dir``: required for ``run_dreamzero_policy_session``. + - ``prompt``: language instruction for DreamZero observations (optional). + - ``session_id``: used by the default policy session observation builder. + - ``num_chunks``: DreamZero chunk count (default ``2``). + - ``timeout``: seconds to wait for each inbound msgpack frame (default ``120``). + - ``ping_interval`` / ``ping_timeout``: forwarded to :func:`websockets.sync.client.connect`. + """ + cfg = dict(request_config or {}) + operations_cfg = cfg.get("operations") + if operations_cfg is not None: + operations = [dict(op) for op in operations_cfg] + elif cfg.get("run_dreamzero_policy_session"): + video_dir = cfg.get("video_dir") + if video_dir is None: + raise ValueError("run_dreamzero_policy_session requires video_dir") + session_id = str(cfg.get("session_id", "dreamzero-smoke")) + prompt = str(cfg.get("prompt", DREAMZERO_DEFAULT_PROMPT)) + num_chunks = int(cfg.get("num_chunks", 2)) + observations = build_dreamzero_demo_observations( + load_dreamzero_camera_frames(Path(video_dir)), + prompt=prompt, + session_id=session_id, + num_chunks=num_chunks, + ) + operations = [{"endpoint": "infer", "payload": obs} for obs in observations] + operations.append({"endpoint": "reset", "payload": {}}) + operations.append({"endpoint": "infer", "payload": observations[0]}) + elif cfg.get("run_default_policy_session"): + session_id = str(cfg.get("session_id", "gr00t-smoke")) + operations = [ + {"endpoint": "infer", "payload": build_openpi_droid_observation(session_id=session_id)}, + {"endpoint": "reset", "payload": {}}, + ] + else: + operations = [] + timeout = float(cfg.get("timeout", 120.0)) + ping_interval = float(cfg.get("ping_interval", OpenPIWebSocketSession.DEFAULT_PING_INTERVAL_SECS)) + ping_timeout = float(cfg.get("ping_timeout", OpenPIWebSocketSession.DEFAULT_PING_TIMEOUT_SECS)) + + uri = self._build_ws_url("/v1/realtime/robot/openpi") + session = OpenPIWebSocketSession( + uri, + ping_interval=ping_interval, + ping_timeout=ping_timeout, + open_timeout=timeout, + close_timeout=timeout, + ) + try: + operation_responses: list[Any] = [] + actions: dict[str, np.ndarray] | None = None + action_tensors: list[np.ndarray] = [] + for operation in operations: + endpoint = str(operation["endpoint"]) + payload = dict(operation.get("payload") or {}) + response = session._send_operation(endpoint, payload) + operation_responses.append(response) + if endpoint == "infer": + if isinstance(response, dict): + actions = {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()} + else: + action_tensors.append(np.asarray(response, dtype=np.float32)) + finally: + session.close() + + return [ + OpenPIWebSocketResponse( + server_metadata=session.get_server_metadata(), + operation_responses=operation_responses, + actions=actions, + action_tensors=action_tensors or None, + ) + ] + + def send_omni_request( + self, + request_config: dict[str, Any], + request_num: int = 1, + *, + min_successes: int | None = None, + max_concurrency: int | None = None, + ) -> list[OmniResponse]: + """Chat completions via the OpenAI Python SDK (not raw HTTP). + + ``min_successes`` gates the batch on how many of the ``request_num`` requests pass + rather than requiring each one to, and returns only the successes; ``max_concurrency`` + caps how many are in flight at once. See ``SuccessRateGate``. + """ + responses: list[OmniResponse] = [] + stream = request_config.get("stream", False) + modalities = request_config.get("modalities", ["text", "audio"]) + extra_body: dict[str, Any] = {} + if "speaker" in request_config: + extra_body["speaker"] = request_config["speaker"] + if request_config.get("use_audio_in_video"): + mm = dict(extra_body.get("mm_processor_kwargs") or {}) + mm["use_audio_in_video"] = True + extra_body["mm_processor_kwargs"] = mm + if "sampling_params_list" in request_config: + extra_body["sampling_params_list"] = request_config["sampling_params_list"] + if request_config.get("extra_body"): + extra_body.update(request_config["extra_body"]) + + create_kwargs: dict[str, Any] = { + "model": request_config.get("model"), + "messages": request_config.get("messages"), + "stream": stream, + "modalities": modalities, + } + if "logprobs" in request_config: + create_kwargs["logprobs"] = request_config["logprobs"] + if "top_logprobs" in request_config: + create_kwargs["top_logprobs"] = request_config["top_logprobs"] + if "stream_options" in request_config: + create_kwargs["stream_options"] = request_config["stream_options"] + if extra_body: + create_kwargs["extra_body"] = extra_body + + def _one(): + wall_start = time.perf_counter() + chat_completion = self.client.chat.completions.create(**create_kwargs) + return ( + self._process_stream_omni_response(chat_completion, wall_start=wall_start) + if stream + else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start) + ) + + if min_successes is not None: + gate = SuccessRateGate(min_successes=min_successes, concurrency=max_concurrency) + + def _sample() -> OmniResponse: + resp = _one() + assert_omni_response(resp, request_config, run_level=self.run_level) + return resp + + # Not _print_client_stat: the count is the gate's verdict, so it has to survive + # log_stats=False. + return collect_at_success_rate( + _sample, + gate, + request_num=request_num, + report=lambda line: print(f"[omni] {line}", flush=True), + ) + + if request_num == 1: + resp = _one() + assert_omni_response(resp, request_config, run_level=self.run_level) + if resp.e2e_latency is not None: + self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s") + else: + self._print_client_stat("[omni] request#1 completed") + responses.append(resp) + return responses + + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = {executor.submit(_one): i + 1 for i in range(request_num)} + for future in concurrent.futures.as_completed(futures): + request_idx = futures[future] + resp = future.result() + assert_omni_response(resp, request_config, run_level=self.run_level) + if resp.e2e_latency is not None: + self._print_client_stat(f"[omni] request#{request_idx} success in {resp.e2e_latency:.3f}s") + else: + self._print_client_stat(f"[omni] request#{request_idx} completed") + responses.append(resp) + return responses + + def _process_stream_audio_speech_response( + self, response, *, response_format: str | None = None, wall_start: float + ) -> OmniResponse: + """ + Process streaming /v1/audio/speech responses into an OmniResponse. + + This mirrors _process_stream_omni_response but operates on low-level + audio bytes. Whisper transcription runs in assert_audio_speech_response + when the run_level requires it. + """ + result = OmniResponse() + + try: + # Aggregate all audio bytes from the streaming response. + data = bytearray() + + # Preferred OpenAI helper. + if hasattr(response, "iter_bytes") and callable(getattr(response, "iter_bytes")): + for chunk in response.iter_bytes(): + if chunk: + data.extend(chunk) + else: + # Generic iterable-of-bytes fallback (e.g., generator or list of chunks). + try: + iterator = iter(response) + except TypeError: + iterator = None + + if iterator is not None: + for chunk in iterator: + if not chunk: + continue + if isinstance(chunk, (bytes, bytearray)): + data.extend(chunk) + elif hasattr(chunk, "data"): + data.extend(chunk.data) # type: ignore[arg-type] + elif hasattr(chunk, "content"): + data.extend(chunk.content) # type: ignore[arg-type] + else: + raise TypeError(f"Unsupported stream chunk type: {type(chunk)}") + else: + raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}") + + raw_bytes = bytes(data) + + # Populate OmniResponse. + result.audio_bytes = raw_bytes + result.e2e_latency = time.perf_counter() - wall_start + result.success = True + result.audio_format = getattr(response, "response", None) + if result.audio_format is not None: + result.audio_format = result.audio_format.headers.get("content-type", "") + + except Exception as e: + msg = f"Audio speech stream processing error: {str(e)}" + print(f"Error: {msg}") + + return result + + def _process_non_stream_audio_speech_response( + self, response, *, response_format: str | None = None, wall_start: float + ) -> OmniResponse: + """ + Process non-streaming /v1/audio/speech responses into an OmniResponse. + + This mirrors _process_non_stream_omni_response but for the binary + audio payload returned by audio.speech.create. + """ + result = OmniResponse() + + try: + # OpenAI non-streaming audio.speech.create returns HttpxBinaryResponseContent (.read() or .content) + if hasattr(response, "read") and callable(getattr(response, "read")): + raw_bytes = response.read() + elif hasattr(response, "content"): + raw_bytes = response.content # type: ignore[assignment] + else: + raise TypeError(f"Unsupported audio speech response type: {type(response)}") + + result.audio_bytes = raw_bytes + result.e2e_latency = time.perf_counter() - wall_start + result.success = True + result.audio_format = getattr(response, "response", None) + if result.audio_format is not None: + result.audio_format = result.audio_format.headers.get("content-type", "") + + except Exception as e: + msg = f"Audio speech non-stream processing error: {str(e)}" + print(f"Error: {msg}") + + return result + + def send_audio_speech_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + """ + Call the /v1/audio/speech endpoint using the same configuration-dict + style as send_omni_request, but via the OpenAI Python client's + audio.speech APIs. + + Expected keys in request_config: + - model: model name/path (required) + - input: text to synthesize (required) + - response_format: audio format such as "wav" or "pcm" (optional) + - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body) + - min_audio_bytes: optional minimum ``len(audio_bytes)`` checked in ``assert_audio_speech_response`` + - transcript_expected_text: local expected spoken text; defaults to ``input`` + - timeout: request timeout in seconds (float, optional, default 120.0) + - stream: whether to use streaming API (bool, optional, default False) + + For negative / contract checks (expected 4xx), use + :meth:`send_audio_speech_http_request` with ``err_code`` / ``err_message``. + """ + timeout = float(request_config.get("timeout", 120.0)) + + model = request_config["model"] + text_input = request_config["input"] + stream = bool(request_config.get("stream", False)) + voice = request_config.get("voice", None) + + # Standard OpenAI param: use omit when not provided to keep default behavior. + response_format = request_config.get("response_format", omit) + + # Qwen3-TTS custom fields, forwarded via extra_body. + extra_body: dict[str, Any] = {} + # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params. + for key in ( + "task_type", + "ref_text", + "ref_audio", + "language", + "max_new_tokens", + "seed", + "instructions", + "speed", + "stream_format", + "x_vector_only_mode", + ): + if key in request_config: + extra_body[key] = request_config[key] + + responses: list[OmniResponse] = [] + + speech_fmt: str | None = None if response_format is omit else str(response_format).lower() + + print(f"[audio.speech] start model={model}, stream={stream}, request_num={request_num}, timeout={timeout:.1f}s") + + if request_num == 1: + if stream: + # Use streaming response helper. + wall_start = time.perf_counter() + with self.client.audio.speech.with_streaming_response.create( + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) as resp: + omni_resp = self._process_stream_audio_speech_response( + resp, response_format=speech_fmt, wall_start=wall_start + ) + else: + # Non-streaming response. + wall_start = time.perf_counter() + resp = self.client.audio.speech.create( + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) + omni_resp = self._process_non_stream_audio_speech_response( + resp, response_format=speech_fmt, wall_start=wall_start + ) + + assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) + if omni_resp.e2e_latency is not None: + self._print_client_stat(f"[audio.speech] request#1 success in {omni_resp.e2e_latency:.3f}s") + else: + self._print_client_stat("[audio.speech] request#1 completed") + responses.append(omni_resp) + return responses + else: + # request_num > 1: concurrent requests (use same params as single-request path) + + if stream: + + def _stream_task(request_idx: int): + wall_start = time.perf_counter() + with self.client.audio.speech.with_streaming_response.create( + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) as resp: + result = self._process_stream_audio_speech_response( + resp, response_format=speech_fmt, wall_start=wall_start + ) + if result.e2e_latency is not None: + self._print_client_stat( + f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s" + ) + else: + self._print_client_stat(f"[audio.speech] request#{request_idx} completed") + return result + + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = {executor.submit(_stream_task, i + 1): i + 1 for i in range(request_num)} + for future in concurrent.futures.as_completed(futures): + request_idx = futures[future] + try: + omni_resp = future.result() + except Exception as e: + print( + f"[audio.speech] request#{request_idx} failed " + f"(stream={stream}, timeout={timeout:.1f}s): {e!r}" + ) + raise + assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) + responses.append(omni_resp) + else: + + def _non_stream_task(request_idx: int): + wall_start = time.perf_counter() + r = self.client.audio.speech.create( + model=model, + input=text_input, + response_format=response_format, + extra_body=extra_body or None, + timeout=timeout, + voice=voice, + ) + result = self._process_non_stream_audio_speech_response( + r, response_format=speech_fmt, wall_start=wall_start + ) + if result.e2e_latency is not None: + self._print_client_stat( + f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s" + ) + else: + self._print_client_stat(f"[audio.speech] request#{request_idx} completed") + return result + + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = {executor.submit(_non_stream_task, i + 1): i + 1 for i in range(request_num)} + for future in concurrent.futures.as_completed(futures): + request_idx = futures[future] + try: + omni_resp = future.result() + except Exception as e: + print( + f"[audio.speech] request#{request_idx} failed " + f"(stream={stream}, timeout={timeout:.1f}s): {e!r}" + ) + raise + assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) + responses.append(omni_resp) + + return responses + + def send_diffusion_request( + self, request_config: dict[str, Any] | list[dict[str, Any]], request_num: int = 1 + ) -> list[DiffusionResponse]: + """ + Send OpenAI requests for diffusion models. + If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel + (scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order. + + Args: + request_config: A single request configuration dict, or a list of + request configuration dicts (one request per element) + request_num: Number of requests to send concurrently, defaults to 1 (single request) + Returns: + list[DiffusionResponse]: List of DiffusionResponse objects containing the response data + """ + responses: list[DiffusionResponse] = [] + + def _create_from_config(cfg: dict[str, Any]) -> tuple[Any, float]: + stream = cfg.get("stream", False) + if stream: + raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test") + modalities = cfg.get("modalities", omit) # Most diffusion models don't require modalities param + eb = cfg.get("extra_body") + extra = copy.deepcopy(eb) if eb else None + wall_start = time.perf_counter() + chat_completion = self.client.chat.completions.create( + model=cfg.get("model"), + messages=cfg.get("messages"), + extra_body=extra, + modalities=modalities, + ) + return chat_completion, wall_start + + if isinstance(request_config, list): + if not request_config: + raise ValueError("request_config list must not be empty") + if request_num != 1: + raise ValueError("request_num is not supported when request_config is a list") + with concurrent.futures.ThreadPoolExecutor(max_workers=len(request_config)) as executor: + futures = { + executor.submit(_create_from_config, cfg): (i + 1, cfg) for i, cfg in enumerate(request_config) + } + for future in concurrent.futures.as_completed(futures): + request_idx, cfg = futures[future] + chat_completion, wall_start = future.result() + response = self._process_diffusion_response(chat_completion, wall_start=wall_start) + assert_diffusion_response(response, cfg, run_level=self.run_level) + if response.e2e_latency is not None: + self._print_client_stat( + f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s" + ) + else: + self._print_client_stat(f"[diffusion] request#{request_idx} completed") + responses.append(response) + return responses + + size_splits = _split_request_config_by_per_output_sizes(request_config) + if size_splits is not None: + if request_num != 1: + raise ValueError( + "request_num must be 1 when extra_body height/width are lists (split into concurrent per-size calls)" + ) + t0 = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=len(size_splits)) as executor: + split_futures = [executor.submit(_create_from_config, sub) for sub in size_splits] + chat_completions = [f.result() for f in split_futures] + parts = [self._process_diffusion_response(cc, wall_start=ws) for cc, ws in chat_completions] + merged = _merge_diffusion_responses(parts) + merged.e2e_latency = time.perf_counter() - t0 + assert_diffusion_response(merged, request_config, run_level=self.run_level) + if merged.e2e_latency is not None: + self._print_client_stat(f"[diffusion] request#1 success in {merged.e2e_latency:.3f}s") + else: + self._print_client_stat("[diffusion] request#1 completed") + return [merged] + + if request_num == 1: + # Send single request + chat_completion, wall_start = _create_from_config(request_config) + response = self._process_diffusion_response(chat_completion, wall_start=wall_start) + assert_diffusion_response(response, request_config, run_level=self.run_level) + if response.e2e_latency is not None: + self._print_client_stat(f"[diffusion] request#1 success in {response.e2e_latency:.3f}s") + else: + self._print_client_stat("[diffusion] request#1 completed") + responses.append(response) + return responses + + # Send concurrent requests for the same request_config + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + idx_by_future = {executor.submit(_create_from_config, request_config): i + 1 for i in range(request_num)} + for future in concurrent.futures.as_completed(idx_by_future): + request_idx = idx_by_future[future] + chat_completion, wall_start = future.result() + response = self._process_diffusion_response(chat_completion, wall_start=wall_start) + assert_diffusion_response(response, request_config, run_level=self.run_level) + if response.e2e_latency is not None: + self._print_client_stat(f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s") + else: + self._print_client_stat(f"[diffusion] request#{request_idx} completed") + responses.append(response) + return responses + + def send_video_diffusion_request( + self, request_config: dict[str, Any], request_num: int = 1 + ) -> list[DiffusionResponse]: + """ + Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content. + + For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc. + """ + if request_num != 1: + raise NotImplementedError("Concurrent video diffusion requests are not currently implemented") + + form_data = request_config.get("form_data") + if not isinstance(form_data, dict): + raise ValueError("Video request_config must contain 'form_data'") + normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None} + files: dict[str, tuple[str, BytesIO, str]] = {} + image_reference = request_config.get("image_reference") + video_reference = request_config.get("video_reference") + if image_reference and video_reference: + raise ValueError("Only one of image_reference or video_reference can be provided") + if image_reference: + if image_reference.startswith("data:image"): + header, encoded = image_reference.split(",", 1) + content_type = header.split(";")[0].removeprefix("data:") + extension = content_type.split("/")[-1] + file_data = base64.b64decode(encoded) + files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type) + else: + normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference}) + if video_reference: + if video_reference.startswith("data:video"): + header, encoded = video_reference.split(",", 1) + content_type = header.split(";")[0].removeprefix("data:") + extension = content_type.split("/")[-1] + file_data = base64.b64decode(encoded) + files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type) + else: + normalized_form_data["video_reference"] = json.dumps({"video_url": video_reference}) + + result = DiffusionResponse() + create_url = self._build_url("/v1/videos") + response = requests.post( + create_url, + data=normalized_form_data, + files=files, + headers={"Accept": "application/json"}, + timeout=60, + ) + start_time = time.perf_counter() + response.raise_for_status() + job_data = response.json() + video_id = job_data["id"] + self._wait_until_video_completed(video_id) + end_time = time.perf_counter() + video_content = self._download_video_content(video_id) + result.success = True + result.videos = [video_content] + result.e2e_latency = end_time - start_time + assert_diffusion_response(result, request_config, run_level=self.run_level) + if result.e2e_latency is not None: + self._print_client_stat(f"[diffusion] request#1 success in {result.e2e_latency:.3f}s") + else: + self._print_client_stat("[diffusion] request#1 completed") + return [result] + + def _post_json_endpoint( + self, + path: str, + request_config: dict[str, Any], + *, + default_timeout: float, + ) -> requests.Response: + url = self._build_url(path) + timeout = float(request_config.get("timeout", default_timeout)) + if "raw_body" in request_config: + raw = request_config["raw_body"] + payload = raw.encode("utf-8") if isinstance(raw, str) else raw + return requests.post( + url, + data=payload, + headers={"Content-Type": "application/json", "Accept": "application/json"}, + timeout=timeout, + ) + if "json" not in request_config: + raise ValueError(f"{path} request_config must include 'json' or 'raw_body'") + return requests.post( + url, + json=request_config["json"], + headers={"Accept": "application/json", "Content-Type": "application/json"}, + timeout=timeout, + ) + + def _post_form_endpoint( + self, + path: str, + request_config: dict[str, Any], + *, + default_timeout: float = 120.0, + ) -> requests.Response: + url = self._build_url(path) + timeout = float(request_config.get("timeout", default_timeout)) + data = request_config.get("data") + files = request_config.get("files") + if data is None and not files: + data = {} + return requests.post( + url, + data=data, + files=files, + headers={"Accept": "application/json"} if not files else {"Accept": "application/json"}, + timeout=timeout, + ) + + def send_streaming_video_diffusion_request( + self, + request_config: dict[str, Any], + request_num: int = 1, + *, + timeout_seconds: float = 600.0, + ) -> list[DiffusionResponse]: + """ + Send a native ``/v1/realtime/video`` WebSocket request and return one + finalized MP4 artifact assembled from the streamed binary fragments. + """ + if request_num != 1: + raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented") + + response = asyncio.run( + self._send_streaming_video_diffusion_request_once( + request_config, + timeout_seconds=timeout_seconds, + ) + ) + assert_diffusion_response(response, request_config, run_level=self.run_level) + if response.e2e_latency is not None: + self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s") + else: + self._print_client_stat("[diffusion.stream] request#1 completed") + return [response] + + async def _send_streaming_video_diffusion_request_once( + self, + request_config: dict[str, Any], + *, + timeout_seconds: float, + ) -> DiffusionResponse: + form_data = request_config.get("form_data") + if not isinstance(form_data, dict): + raise ValueError("Video request_config must contain 'form_data'") + payload: dict[str, Any] = { + "type": "session.start", + **{key: value for key, value in form_data.items() if value is not None}, + } + model = request_config.get("model") + if model is not None: + payload["model"] = model + payload.setdefault("format", "m4s") + + fps = float(payload.get("fps") or 16) + stream_format = payload["format"] + url = self._build_ws_url("/v1/realtime/video") + + result = DiffusionResponse() + chunks: list[bytes] = [] + start_time = time.perf_counter() + deadline = start_time + timeout_seconds + + import websockets + + async with websockets.connect(url, max_size=None) as websocket: + await websocket.send(json.dumps(payload)) + + while True: + remaining = deadline - time.perf_counter() + if remaining <= 0: + raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s") + + message = await asyncio.wait_for(websocket.recv(), timeout=remaining) + if isinstance(message, bytes): + chunks.append(message) + continue + + msg = json.loads(message) + msg_type = msg.get("type") + if msg_type == "video.start": + stream_format = msg.get("format") or stream_format + continue + if msg_type == "session.done": + break + if msg_type == "error": + raise RuntimeError(str(msg.get("message", msg))) + + from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes + from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat + + streamed_bytes = b"".join(chunks) + if not streamed_bytes: + raise RuntimeError("Streaming video request completed without binary video chunks") + result.videos = [ + finalize_streaming_video_bytes( + streamed_bytes, + input_format=cast(StreamingVideoFormat, stream_format), + fps=fps, + ) + ] + result.e2e_latency = time.perf_counter() - start_time + result.success = True + return result + + def _wait_until_video_completed( + self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300 + ) -> None: + status_url = self._build_url(f"/v1/videos/{video_id}") + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + status_resp = requests.get(status_url, headers={"Accept": "application/json"}, timeout=30) + status_resp.raise_for_status() + status_data = status_resp.json() + current_status = status_data["status"] + if current_status == "completed": + return + if current_status == "failed": + error_msg = status_data.get("last_error", "Unknown error") + raise RuntimeError(f"Job failed: {error_msg}") + time.sleep(poll_interval_seconds) + raise TimeoutError(f"Video job {video_id} did not complete within {timeout_seconds}s") + + def _download_video_content(self, video_id: str) -> bytes: + download_url = self._build_url(f"/v1/videos/{video_id}/content") + video_resp = requests.get(download_url, stream=True, timeout=60) + video_resp.raise_for_status() + video_bytes = BytesIO() + for chunk in video_resp.iter_content(chunk_size=8192): + if chunk: + video_bytes.write(chunk) + return video_bytes.getvalue() + + def _build_url(self, path: str) -> str: + return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}" + + +class OfflineOmniClient: + def __init__(self, omni_runner: OmniRunner): + self.runner = omni_runner + + def _process_omni_output(self, outputs: list[OmniRequestOutput]) -> OmniResponse: + result = OmniResponse() + try: + text_content = None + audio_content = None + for stage_output in outputs: + if getattr(stage_output, "final_output_type", None) == "text": + text_content = stage_output.outputs[0].text + if getattr(stage_output, "final_output_type", None) == "audio": + audio_content = stage_output.outputs[0].multimodal_output["audio"] + result.audio_content = audio_content + result.text_content = text_content + result.success = True + except Exception as e: + msg = f"Output processing error: {str(e)}" + result.success = False + print(f"Error: {msg}") + return result + + def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse: + result = DiffusionResponse() + output = outputs[0] + if isinstance(output.images[0], list): + # Returning frames of images as a video + result.videos = output.images + else: + # Returning actual images + result.images = output.images + # [TODO] Add audio processing when tests are introduced + result.success = True + return result + + def send_omni_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse: + if request_config is None: + request_config = {} + prompts_value = request_config.get("prompts") + prompts: str | list[str] = prompts_value if isinstance(prompts_value, (str, list)) else "" + videos = request_config.get("videos") + images = request_config.get("images") + audios = request_config.get("audios") + modalities = request_config.get("modalities", ["text", "audio"]) + outputs = self.runner.generate_multimodal( + prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities + ) + response = self._process_omni_output(outputs) + assert_omni_response(response, request_config, run_level="core_model") + return response + + def send_diffusion_request(self, request_config: dict[str, Any]) -> DiffusionResponse: + prompt = request_config.get("prompt") + if prompt is None: + prompts = request_config.get("prompts") + if not prompts: + raise ValueError("request_config must contain a prompt") + if len(prompts) > 1: + raise ValueError( + "In the current internal data structure, " + "only one prompt is supported for diffusion requests. " + "Because one prompt can contain multiple images or videos, " + "the current internal data structure is ambiguous when multiple prompts are provided." + ) + prompt = prompts[0] + + # Sync previous `extra_body` field with sampling params object + sampling_params: OmniDiffusionSamplingParams | None = request_config.get("sampling_params") + if sampling_params is None: + extra_body = request_config.get("extra_body", {}) + if extra_body: + sampling_params = OmniDiffusionSamplingParams(**extra_body) + else: + extra_body = asdict(sampling_params) + request_config["extra_body"] = extra_body + if not extra_body: + logger.warning("No sampling params provided in request_config, will skip output assertion") + + negative_prompt = extra_body.get("negative_prompt") or request_config.get("negative_prompt") + videos = request_config.get("videos") + images = request_config.get("images") + audios = request_config.get("audios") + # Full dict (e.g. image + mask_image for inpainting) or partial; merged with top-level image/video keys. + extra_multi_modal = request_config.get("multi_modal_data") + modalities = request_config.get("modalities") # only used by limited models. Do not add default value here + + prompt_object = OmniTextPrompt(prompt=prompt) + if negative_prompt: + prompt_object["negative_prompt"] = negative_prompt + multi_modal: dict = {} + if extra_multi_modal is not None: + multi_modal.update(dict(extra_multi_modal)) + if videos is not None: + multi_modal["video"] = videos + if images is not None: + multi_modal["image"] = images + if audios is not None: + multi_modal["audio"] = audios + if multi_modal: + prompt_object["multi_modal_data"] = multi_modal + if modalities: + prompt_object["modalities"] = modalities # pyright: ignore[reportGeneralTypeIssues] + + start_time = time.perf_counter() + raw_outputs = self.runner.generate([prompt_object], [sampling_params] if sampling_params else None) + end_time = time.perf_counter() + if not isinstance(raw_outputs, list): + raw_outputs = list(raw_outputs) + + response = self._process_diffusion_output(raw_outputs) + response.e2e_latency = end_time - start_time + assert_diffusion_response(response, request_config, run_level="core_model") + return response + + def send_audio_speech_request(self, request_config: dict[str, Any]) -> OmniResponse: + """ + Offline TTS: text -> audio via generate_multimodal, then validate with assert_audio_speech_response. + + request_config must contain: + - 'input' or 'prompts': text to synthesize. + Optional keys: + - 'voice' -> speaker (CustomVoice) + - 'task_type' -> task_type in additional_information (default: "CustomVoice") + - 'language' -> language in additional_information (default: "Auto") + - 'max_new_tokens' -> max_new_tokens in additional_information (default: 2048) + - 'response_format' -> desired audio format (used only for assertion) + """ + input_text = request_config.get("input") or request_config.get("prompts") + if input_text is None: + raise ValueError("request_config must contain 'input' or 'prompts' for TTS") + if isinstance(input_text, list): + input_text = input_text[0] if input_text else "" + + mm_processor_kwargs: dict[str, Any] = {} + if "voice" in request_config: + mm_processor_kwargs["speaker"] = request_config["voice"] + if "task_type" in request_config: + mm_processor_kwargs["task_type"] = request_config["task_type"] + if "ref_audio" in request_config: + mm_processor_kwargs["ref_audio"] = request_config["ref_audio"] + if "ref_text" in request_config: + mm_processor_kwargs["ref_text"] = request_config["ref_text"] + if "language" in request_config: + mm_processor_kwargs["language"] = request_config["language"] + if "max_new_tokens" in request_config: + mm_processor_kwargs["max_new_tokens"] = request_config["max_new_tokens"] + + outputs = self.runner.generate_multimodal( + prompts=input_text, + modalities=["audio"], + mm_processor_kwargs=mm_processor_kwargs or None, + ) + mm_out: dict[str, Any] | None = None + for stage_out in outputs: + if getattr(stage_out, "final_output_type", None) == "audio": + mm_out = stage_out.outputs[0].multimodal_output + break + if mm_out is None: + raise AssertionError("No audio output from pipeline") + + audio_data = mm_out.get("audio") + if audio_data is None: + raise AssertionError("No audio tensor in multimodal output") + + sr_raw = mm_out.get("sr") + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sr = int(sr_val.item() if hasattr(sr_val, "item") else sr_val) + wav_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data + wav_buf = io.BytesIO() + sf.write( + wav_buf, + wav_tensor.float().cpu().numpy().reshape(-1), + samplerate=sr, + format="WAV", + subtype="PCM_16", + ) + result = OmniResponse(success=True, audio_bytes=wav_buf.getvalue(), audio_format="audio/wav") + assert_audio_speech_response(result, request_config, run_level="core_model") + return result + + def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]: + return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages) + + def stop_profile(self, stages: list[int] | None = None) -> list[Any]: + return self.runner.stop_profile(stages=stages) + + +# Backward-compatible aliases (historical test fixture / type names). +OpenAIClientHandler = OnlineOmniClient +OmniRunnerHandler = OfflineOmniClient + + +__all__ = [ + "DiffusionResponse", + "HttpResponse", + "WebSocketJsonResponse", + "OpenPIWebSocketResponse", + "build_openpi_droid_observation", + "build_dreamzero_demo_observations", + "DREAMZERO_ACTION_DIM", + "DREAMZERO_ACTION_HORIZON", + "DREAMZERO_CAMERA_FILES", + "DREAMZERO_DEFAULT_PROMPT", + "load_dreamzero_camera_frames", + "OpenPIWebSocketSession", + "OmniResponse", + "OfflineOmniClient", + "OnlineOmniClient", + "OpenAIClientHandler", + "OmniRunnerHandler", +] diff --git a/tests/helpers/env.py b/tests/helpers/env.py deleted file mode 100644 index 39cec30e7..000000000 --- a/tests/helpers/env.py +++ /dev/null @@ -1,297 +0,0 @@ -"""Test environment / lifecycle helpers (device cleanup hooks and memory monitoring for tests). - -``vllm_omni.platforms`` is imported only inside functions that need it so importing this module -at pytest plugin load does not run before session autouse fixtures. -""" - -from __future__ import annotations - -import gc -import os -import subprocess -import threading -import time -from contextlib import contextmanager - -import torch - -from vllm_omni.platforms import current_omni_platform - - -def get_physical_device_indices(devices): - visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") - if visible_devices is None: - return devices - visible_indices = [int(x) for x in visible_devices.split(",")] - index_mapping = {i: physical for i, physical in enumerate(visible_indices)} - return [index_mapping[i] for i in devices if i in index_mapping] - - -def pick_least_used_device_indices(num_devices: int) -> list[int]: - """Return physical device indices for the least-used accelerators. - - Args: - num_devices: Number of device indices to return. - - Returns: - Physical device indices sorted from least to most used memory. - """ - if num_devices <= 0: - raise ValueError(f"num_devices must be positive, got {num_devices}.") - - logical_count = current_omni_platform.get_device_count() - if logical_count < num_devices: - raise RuntimeError(f"Need {num_devices} devices, but only {logical_count} are available.") - - device_usage: list[tuple[int, int]] = [] - for device_index in range(logical_count): - with current_omni_platform.device(device_index): - free_bytes, total_bytes = current_omni_platform.get_device_memory() - device_usage.append((total_bytes - free_bytes, device_index)) - device_usage.sort() - logical_indices = [device_index for _, device_index in device_usage[:num_devices]] - return get_physical_device_indices(logical_indices) - - -def wait_for_gpu_memory_to_clear( - *, - devices: list[int], - threshold_bytes: int | None = None, - threshold_ratio: float | None = None, - timeout_s: float = 120, -) -> None: - assert threshold_bytes is not None or threshold_ratio is not None - devices = get_physical_device_indices(devices) - start_time = time.time() - - device_list = ", ".join(str(d) for d in devices) - if threshold_bytes is not None: - condition_str = f"Memory usage ≤ {threshold_bytes / 2**30:.2f} GiB" - - def is_free(used, total): - return used <= threshold_bytes / 2**30 - else: - condition_str = f"Memory usage ratio ≤ {threshold_ratio * 100:.1f}%" - - def is_free(used, total): - return used / total <= threshold_ratio - - print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}") - - @contextmanager - def smi_scope(): - if current_omni_platform.is_rocm(): - from amdsmi import amdsmi_init, amdsmi_shut_down - - amdsmi_init() - try: - yield - finally: - amdsmi_shut_down() - elif current_omni_platform.is_cuda(): - from vllm.third_party.pynvml import nvmlInit, nvmlShutdown - - nvmlInit() - try: - yield - finally: - nvmlShutdown() - else: - yield - - def get_mem_gib(device: int) -> tuple[float, float]: - if current_omni_platform.is_rocm(): - from amdsmi import amdsmi_get_gpu_vram_usage, amdsmi_get_processor_handles - - info = amdsmi_get_gpu_vram_usage(amdsmi_get_processor_handles()[device]) - return info["vram_used"] / 2**10, info["vram_total"] / 2**10 - if current_omni_platform.is_npu(): - free_bytes, total_bytes = torch.npu.mem_get_info(device) - return (total_bytes - free_bytes) / 2**30, total_bytes / 2**30 - from vllm.third_party.pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo - - info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(device)) - return info.used / 2**30, info.total / 2**30 - - with smi_scope(): - while True: - output_raw = {d: get_mem_gib(d) for d in devices} - output = { - d: f"{used:.1f}GiB/{total:.1f}GiB ({(used / total) * 100 if total > 0 else 0:.1f}%)" - for d, (used, total) in output_raw.items() - } - - print("[GPU Memory Status] Current usage:") - for device_id, mem_info in output.items(): - print(f" GPU {device_id}: {mem_info}") - - dur_s = time.time() - start_time - if all(is_free(used, total) for used, total in output_raw.values()): - print(f"[GPU Memory Freed] Devices {device_list} meet memory condition") - print(f" Condition: {condition_str}") - print(f" Wait time: {dur_s:.1f} seconds ({dur_s / 60:.1f} minutes)") - break - - if dur_s >= timeout_s: - raise ValueError( - f"[GPU Memory Timeout] Devices {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n" - f"Condition: {condition_str}\n" - f"Current status:\n" + "\n".join(f" GPU {d}: {output[d]}" for d in devices) - ) - - gc.collect() - current_omni_platform.empty_cache() - time.sleep(5) - - -def _run_smi(label: str, cmd: list[str], head_lines: int, timeout: float = 5) -> None: - print("\n" + "=" * 80) - print(label) - print("=" * 80) - try: - result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) - if result.returncode == 0 and result.stdout.strip(): - lines = result.stdout.strip().split("\n") - for line in lines[:head_lines]: - print(line) - if len(lines) > head_lines: - print(f"... (showing first {head_lines} of {len(lines)} lines)") - else: - print(f"{cmd[0]} command failed or produced no output") - except (subprocess.TimeoutExpired, FileNotFoundError): - print(f"{cmd[0]} not available or timed out") - except Exception as e: - print(f"Error running {' '.join(cmd)}: {e}") - - -def _print_gpu_processes() -> None: - """Print device information via nvidia-smi/npu-smi/amd-smi.""" - from vllm_omni.platforms import current_omni_platform - - if current_omni_platform.is_cuda(): - _run_smi("NVIDIA GPU Information (nvidia-smi)", ["nvidia-smi"], 20) - _run_smi("Detailed GPU Processes (nvidia-smi pmon)", ["nvidia-smi", "pmon", "-c", "1"], 100, timeout=3) - elif current_omni_platform.is_npu(): - _run_smi("Ascend NPU Information (npu-smi info)", ["npu-smi", "info"], 40) - elif current_omni_platform.is_rocm(): - _run_smi("AMD GPU Information (amd-smi)", ["amd-smi"], 30) - _run_smi("Detailed AMD GPU Processes (amd-smi process)", ["amd-smi", "process"], 100, timeout=3) - else: - print("\n" + "=" * 80) - print("WARNING: No supported device platform detected") - print("=" * 80) - - print("\n" + "=" * 80) - print("System Processes with GPU keywords") - print("=" * 80) - - -def _cleanup_stale_device_locks() -> None: - """Remove stale device-initialization lock files whose recorded PID is dead. - - Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a - crashed / killed test run and block subsequent orchestrator startups. - """ - import glob as _glob - - for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"): - try: - with open(lock_file) as fh: - content = fh.read().strip() - if not content: - continue - pid = int(content) - except (OSError, ValueError): - continue - - try: - os.kill(pid, 0) - except ProcessLookupError: - print(f" Removing stale device lock {lock_file} (PID {pid} is dead)") - try: - os.unlink(lock_file) - except OSError: - pass - except PermissionError: - pass - - -def run_pre_test_cleanup() -> None: - print("Pre-test GPU status:") - - _cleanup_stale_device_locks() - - num_gpus = current_omni_platform.device_count() - if num_gpus > 0: - try: - wait_for_gpu_memory_to_clear( - devices=list(range(num_gpus)), - threshold_ratio=0.05, - timeout_s=60, - ) - except Exception as e: - print(f"Pre-test cleanup note: {e}") - - -def run_post_test_cleanup() -> None: - if current_omni_platform.is_available(): - gc.collect() - current_omni_platform.empty_cache() - - print("Post-test GPU status:") - _print_gpu_processes() - - -class DeviceMemoryMonitor: - """Poll global device memory usage.""" - - def __init__(self, device_index: int, interval: float = 0.05): - self.device_index = device_index - self.interval = interval - self._peak_used_mb = 0.0 - self._stop_event = threading.Event() - self._thread: threading.Thread | None = None - - def start(self) -> None: - from vllm_omni.platforms import current_omni_platform - - def monitor_loop() -> None: - while not self._stop_event.is_set(): - try: - with current_omni_platform.device(self.device_index): - free_bytes, total_bytes = current_omni_platform.mem_get_info() - used_mb = (total_bytes - free_bytes) / (1024**2) - self._peak_used_mb = max(self._peak_used_mb, used_mb) - except Exception: - pass - time.sleep(self.interval) - - self._thread = threading.Thread(target=monitor_loop, daemon=False) - self._thread.start() - - def stop(self) -> None: - if self._thread is None: - return - self._stop_event.set() - self._thread.join(timeout=2.0) - - @property - def peak_used_mb(self) -> float: - from vllm_omni.platforms import current_omni_platform - - fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2) - fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2) - return max(self._peak_used_mb, fallback_alloc, fallback_reserved) - - def __del__(self): - self.stop() - - -__all__ = [ - "DeviceMemoryMonitor", - "get_physical_device_indices", - "pick_least_used_device_indices", - "run_post_test_cleanup", - "run_pre_test_cleanup", - "wait_for_gpu_memory_to_clear", -] diff --git a/tests/helpers/fixtures/clean.py b/tests/helpers/fixtures/clean.py new file mode 100644 index 000000000..c07f563bb --- /dev/null +++ b/tests/helpers/fixtures/clean.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Opt-in cleanup fixtures for GPU memory, pipeline registry, and speaker cache.""" + +from __future__ import annotations + +import pytest + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES + + +@pytest.fixture +def clean_gpu_memory_between_tests(): + """Opt-in GPU pre/post hooks for a test (no environment-variable gate). + + Use as a test parameter or ``@pytest.mark.usefixtures("clean_gpu_memory_between_tests")``. + """ + from tests.helpers.clean import cleanup_test_environment + + print("\n=== PRE-TEST DEVICE CLEANUP ===") + cleanup_test_environment() + yield + cleanup_test_environment() + + +@pytest.fixture +def clean_pipeline_registry(): + """Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it.""" + snapshot = dict(OMNI_PIPELINES) + yield + OMNI_PIPELINES.clear() + OMNI_PIPELINES.update(snapshot) + + +@pytest.fixture +def clean_speaker_cache(): + """Reset the process-wide speaker cache singleton before and after the test.""" + import vllm_omni.utils.speaker_cache as _sc + + def _reset(): + with _sc._SINGLETON_LOCK: + if _sc._SINGLETON is not None: + _sc._SINGLETON.clear() + _sc._SINGLETON = None + + _reset() + yield + _reset() diff --git a/tests/helpers/fixtures/config.py b/tests/helpers/fixtures/config.py index 2e309b672..2f46004a6 100644 --- a/tests/helpers/fixtures/config.py +++ b/tests/helpers/fixtures/config.py @@ -1,14 +1,47 @@ -"""Fixtures config creation & out of tree registry management.""" +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Session-scoped config and environment fixtures.""" + +from __future__ import annotations + +import os import pytest - -from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +import torch +from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config -@pytest.fixture -def clean_pipeline_registry(): - """Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it.""" - snapshot = dict(OMNI_PIPELINES) +@pytest.fixture(scope="session", autouse=True) +def default_env(): + # Keep behavior but avoid import-time side effects (RFC #2299). + keys = ("VLLM_WORKER_MULTIPROC_METHOD", "VLLM_TARGET_DEVICE") + previous = {key: os.environ.get(key) for key in keys} + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn" + os.environ["VLLM_TARGET_DEVICE"] = previous["VLLM_TARGET_DEVICE"] or ( + "cuda" if torch.cuda.is_available() and torch.accelerator.device_count() > 0 else "cpu" + ) yield - OMNI_PIPELINES.clear() - OMNI_PIPELINES.update(snapshot) + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +@pytest.fixture(scope="session", autouse=True) +def default_vllm_config(): + """Set a default VllmConfig for the whole test session. + + Session scope ensures module-scoped fixtures (e.g. ``omni_runner``) and + deferred imports of ``tests.helpers.runtime`` both see the same context. + Function-scoped autouse ran too late for ``OmniRunner`` setup and could + desynchronize vLLM init vs request preprocessing (e.g. renderer state). + """ + # Use CPU device if no GPU is available (e.g., in CI environments) + has_gpu = torch.cuda.is_available() and torch.accelerator.device_count() > 0 + device = "cuda" if has_gpu else "cpu" + device_config = DeviceConfig(device=device) + + with set_current_vllm_config(VllmConfig(device_config=device_config)): + yield diff --git a/tests/helpers/fixtures/env.py b/tests/helpers/fixtures/env.py deleted file mode 100644 index 2d38e4907..000000000 --- a/tests/helpers/fixtures/env.py +++ /dev/null @@ -1,74 +0,0 @@ -import os - -import pytest -import torch -from vllm.distributed import parallel_state - - -@pytest.fixture(scope="session", autouse=True) -def default_env(): - # Keep behavior but avoid import-time side effects (RFC #2299). - keys = ("VLLM_WORKER_MULTIPROC_METHOD", "VLLM_TARGET_DEVICE") - previous = {key: os.environ.get(key) for key in keys} - os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn" - os.environ["VLLM_TARGET_DEVICE"] = previous["VLLM_TARGET_DEVICE"] or ( - "cuda" if torch.cuda.is_available() and torch.accelerator.device_count() > 0 else "cpu" - ) - yield - for key, value in previous.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - - -@pytest.fixture(scope="session") -def model_prefix() -> str: - prefix = os.environ.get("MODEL_PREFIX", "") - return f"{prefix.rstrip('/')}/" if prefix else "" - - -@pytest.fixture -def clean_gpu_memory_between_tests(): - """Opt-in GPU pre/post hooks for a test (no environment-variable gate). - - Use as a test parameter or ``@pytest.mark.usefixtures("clean_gpu_memory_between_tests")``. - """ - from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup - - print("\n=== PRE-TEST GPU CLEANUP ===") - run_pre_test_cleanup() - yield - run_post_test_cleanup() - - -@pytest.fixture(scope="session", autouse=True) -def default_vllm_config(): - """Set a default VllmConfig for the whole test session. - - Session scope ensures module-scoped fixtures (e.g. ``omni_runner``) and - deferred imports of ``tests.helpers.runtime`` both see the same context. - Function-scoped autouse ran too late for ``OmniRunner`` setup and could - desynchronize vLLM init vs request preprocessing (e.g. renderer state). - """ - from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config - - # Use CPU device if no GPU is available (e.g., in CI environments) - has_gpu = torch.cuda.is_available() and torch.accelerator.device_count() > 0 - device = "cuda" if has_gpu else "cpu" - device_config = DeviceConfig(device=device) - - with set_current_vllm_config(VllmConfig(device_config=device_config)): - yield - - -@pytest.fixture() -def init_fake_tp_group(mocker): - """Provide a fake TP group so vllm linear layers can be instantiated.""" - mock_tp = mocker.MagicMock() - mock_tp.world_size = 1 - mock_tp.rank_in_group = 0 - old = parallel_state._TP - parallel_state._TP = mock_tp - yield - parallel_state._TP = old diff --git a/tests/helpers/fixtures/media.py b/tests/helpers/fixtures/media.py new file mode 100644 index 000000000..c34de093e --- /dev/null +++ b/tests/helpers/fixtures/media.py @@ -0,0 +1,24 @@ +"""Fixtures for the reused Whisper transcription worker.""" + +import sys + +import pytest + + +@pytest.fixture(autouse=True, scope="module") +def release_audio_transcriber_after_module(): + """Free the transcription worker's device memory once a test module is done. + + A backstop. The server and runner fixtures already release the worker when + each of their instances tears down, which is what keeps Whisper off the + device while the next parametrized server initializes. This covers the + remaining case: a module that transcribes without using those fixtures. + + Reached through ``sys.modules`` because ``tests/conftest.py`` deliberately + keeps ``tests.helpers.media`` lazy: a module that never transcribed should + not pay to import it. + """ + yield + media = sys.modules.get("tests.helpers.media") + if media is not None: + media.release_audio_transcriber() diff --git a/tests/helpers/fixtures/pytest_collection.py b/tests/helpers/fixtures/pytest_collection.py new file mode 100644 index 000000000..912749242 --- /dev/null +++ b/tests/helpers/fixtures/pytest_collection.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Pytest collection-phase hooks. + +All ``pytest_collection_modifyitems`` policies live here so their order is +explicit. Add a named helper and call it from the hook; do not register a +second ``pytest_collection_modifyitems`` in another plugin. +""" + +from __future__ import annotations + +import pytest + + +def pytest_collection_modifyitems(config, items): + _pin_xdist_serial(config, items) + + +def _pin_xdist_serial(config, items): + """Pin unmarked tests to one xdist worker. + + Only tests explicitly marked ``pytest.mark.xdist`` run concurrently. + Others are pinned to a single xdist worker. + """ + if not config.pluginmanager.hasplugin("xdist"): + return + for item in items: + if not item.get_closest_marker("xdist"): + item.add_marker(pytest.mark.xdist_group(name="serial")) diff --git a/tests/helpers/fixtures/run_args.py b/tests/helpers/fixtures/pytest_run_args.py similarity index 68% rename from tests/helpers/fixtures/run_args.py rename to tests/helpers/fixtures/pytest_run_args.py index 975584d20..3766e9c4e 100644 --- a/tests/helpers/fixtures/run_args.py +++ b/tests/helpers/fixtures/pytest_run_args.py @@ -1,3 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Pytest CLI options and the fixtures that expose them.""" + +from __future__ import annotations + import pytest diff --git a/tests/helpers/fixtures/runtime.py b/tests/helpers/fixtures/runtime.py index 3bffa475d..5c4d0c055 100644 --- a/tests/helpers/fixtures/runtime.py +++ b/tests/helpers/fixtures/runtime.py @@ -1,10 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Runtime fixtures (OmniRunner / OmniServer). Imports are deferred to fixture time. Loading ``tests.helpers.runtime`` at plugin import time (before session fixtures) pulls in vLLM/vllm_omni too early and breaks initialization order vs the legacy -monolithic conftest. Defer imports until fixtures run so ``default_env`` / -``default_vllm_config`` run first. Implementation helpers live in -``tests.helpers.runtime`` (``iter_omni_server`` / ``iter_omni_runner``). +monolithic conftest. Defer imports until fixtures run so session fixtures in +``tests.helpers.fixtures.config`` run first. Server/runner helpers live in +``tests.helpers.runtime``; request clients live in ``tests.helpers.client``. """ from __future__ import annotations @@ -21,19 +24,30 @@ if TYPE_CHECKING: omni_fixture_lock = threading.Lock() +@pytest.fixture +def _normalized_hardware_marks(request: pytest.FixtureRequest): + """Dummy carrier for mixed-``num_cards`` ``@hardware_test`` mark normalization. + + Pytest has no API to clone an item and only change marks. ``@hardware_test`` + therefore parametrizes this fixture (``indirect=True``) so each collected + item gets one platform's SKU and ``cards_{n}`` without changing test + signatures. The returned platform id is unused by test bodies. + """ + return request.param + + @pytest.fixture(scope="function") def omni_server_function( request: pytest.FixtureRequest, run_level: str, - model_prefix: str, ) -> Generator[OmniServer, Any, None]: from tests.helpers.runtime import iter_omni_server - yield from iter_omni_server(request, run_level, model_prefix, omni_fixture_lock) + yield from iter_omni_server(request, run_level, omni_fixture_lock) @pytest.fixture(scope="module") -def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]: +def omni_server(request: pytest.FixtureRequest, run_level: str) -> Generator[OmniServer, Any, None]: """Start vLLM-Omni through the standard or stage-CLI launcher. The fixture stays module-scoped because multi-stage initialization is costly. @@ -42,16 +56,16 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st """ from tests.helpers.runtime import iter_omni_server - yield from iter_omni_server(request, run_level, model_prefix, omni_fixture_lock) + yield from iter_omni_server(request, run_level, omni_fixture_lock) @pytest.fixture -def openai_client(request: pytest.FixtureRequest, run_level: str): +def online_client(request: pytest.FixtureRequest, run_level: str): """Resolve ``omni_server`` lazily so parametrized server fixtures work like upstream.""" - from tests.helpers.runtime import OpenAIClientHandler + from tests.helpers.client import OnlineOmniClient server = request.getfixturevalue("omni_server") - return OpenAIClientHandler( + return OnlineOmniClient( host=server.host, port=server.port, api_key="EMPTY", @@ -61,12 +75,18 @@ def openai_client(request: pytest.FixtureRequest, run_level: str): @pytest.fixture -def openai_client_function(request: pytest.FixtureRequest, run_level: str): +def openai_client(online_client): + """Historical alias for :func:`online_client`.""" + return online_client + + +@pytest.fixture +def online_client_function(request: pytest.FixtureRequest, run_level: str): """Resolve ``omni_server_function`` lazily for function-scoped reliability tests.""" - from tests.helpers.runtime import OpenAIClientHandler + from tests.helpers.client import OnlineOmniClient server = request.getfixturevalue("omni_server_function") - return OpenAIClientHandler( + return OnlineOmniClient( host=server.host, port=server.port, api_key="EMPTY", @@ -75,10 +95,15 @@ def openai_client_function(request: pytest.FixtureRequest, run_level: str): ) +@pytest.fixture +def openai_client_function(online_client_function): + """Historical alias for :func:`online_client_function`.""" + return online_client_function + + @pytest.fixture(scope="function") def omni_runner_function( request: pytest.FixtureRequest, - model_prefix: str, run_level: str, ) -> Generator[OmniRunner, Any, None]: """Function-scoped :class:`~tests.helpers.runtime.OmniRunner` (cf. :func:`omni_server_function`). @@ -88,30 +113,42 @@ def omni_runner_function( """ from tests.helpers.runtime import iter_omni_runner - yield from iter_omni_runner(request, model_prefix, run_level, omni_fixture_lock) + yield from iter_omni_runner(request, run_level, omni_fixture_lock) @pytest.fixture(scope="module") -def omni_runner(request: pytest.FixtureRequest, model_prefix: str, run_level: str) -> Generator[OmniRunner, Any, None]: +def omni_runner(request: pytest.FixtureRequest, run_level: str) -> Generator[OmniRunner, Any, None]: """Module-scoped :class:`~tests.helpers.runtime.OmniRunner` (cf. :func:`omni_server`). Reuses one runner for the whole module to amortize multi-stage init cost. """ from tests.helpers.runtime import iter_omni_runner - yield from iter_omni_runner(request, model_prefix, run_level, omni_fixture_lock) + yield from iter_omni_runner(request, run_level, omni_fixture_lock) @pytest.fixture -def omni_runner_handler_function(omni_runner_function: OmniRunner): - """Resolve :class:`~tests.helpers.runtime.OmniRunnerHandler` for :func:`omni_runner_function`.""" - from tests.helpers.runtime import OmniRunnerHandler +def offline_client_function(omni_runner_function: OmniRunner): + """Resolve :class:`~tests.helpers.client.OfflineOmniClient` for :func:`omni_runner_function`.""" + from tests.helpers.client import OfflineOmniClient - return OmniRunnerHandler(omni_runner_function) + return OfflineOmniClient(omni_runner_function) @pytest.fixture -def omni_runner_handler(omni_runner: OmniRunner): - from tests.helpers.runtime import OmniRunnerHandler +def omni_runner_handler_function(offline_client_function): + """Historical alias for :func:`offline_client_function`.""" + return offline_client_function - return OmniRunnerHandler(omni_runner) + +@pytest.fixture +def offline_client(omni_runner: OmniRunner): + from tests.helpers.client import OfflineOmniClient + + return OfflineOmniClient(omni_runner) + + +@pytest.fixture +def omni_runner_handler(offline_client): + """Historical alias for :func:`offline_client`.""" + return offline_client diff --git a/tests/helpers/fixtures/speaker_cache.py b/tests/helpers/fixtures/speaker_cache.py deleted file mode 100644 index cf7a62122..000000000 --- a/tests/helpers/fixtures/speaker_cache.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Fixtures for the process-wide speaker cache singleton.""" - -from __future__ import annotations - -import pytest - - -@pytest.fixture -def fresh_speaker_cache(): - """Reset the process-wide speaker cache singleton before and after the test.""" - import vllm_omni.utils.speaker_cache as _sc - - def _reset(): - with _sc._SINGLETON_LOCK: - if _sc._SINGLETON is not None: - _sc._SINGLETON.clear() - _sc._SINGLETON = None - - _reset() - yield - _reset() diff --git a/tests/helpers/mark.py b/tests/helpers/mark.py index 676732a8e..18bdbd394 100644 --- a/tests/helpers/mark.py +++ b/tests/helpers/mark.py @@ -1,19 +1,33 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …).""" +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …). + +Module import is stdlib-only so ``tools/pre_commit/check_test_marks.py`` can load +this file without pytest or vllm. Decorator helpers import pytest lazily. +""" from __future__ import annotations import re from functools import lru_cache from pathlib import Path +from typing import TYPE_CHECKING -import pytest -from vllm.platforms import current_platform +if TYPE_CHECKING: + import pytest -# Marker description tag in ``pyproject.toml`` ``tool.pytest.ini_options.markers``. -# Example: ``"H100: [hardware-resource] Tests that require H100 GPU"``. +# Marker description tags in ``pyproject.toml`` ``tool.pytest.ini_options.markers``. +# Example: ``"H100: [hardware-resource] [cuda] Tests that require H100 GPU"``. _HARDWARE_RESOURCE_MARKER_TAG = "[hardware-resource]" +_HARDWARE_CARDS_MARKER_TAG = "[hardware-cards]" +_HARDWARE_PLATFORM_MARKER_TAG = "[hardware-platform]" +_CI_LEVEL_MARKER_TAG = "[ci-level]" +# On ``[hardware-platform]`` lines (not SKUs): also attach ``pytest.mark.gpu``. +_GPU_DEVICE_CLASS_TAG = "[gpu]" +_CARDS_MARK_NAME_RE = re.compile(r"^cards_(\d+)$") +# Indirect fixture in ``tests.helpers.fixtures.runtime``: mixed-count ``@hardware_test`` +# parametrizes it so each collected item carries one platform's SKU and ``cards_{n}``. +HARDWARE_MARK_NORMALIZATION_FIXTURE = "_normalized_hardware_marks" def _repo_root() -> Path: @@ -34,23 +48,16 @@ def _load_pyproject() -> dict: return {} -@lru_cache(maxsize=1) -def get_hardware_mark_list() -> frozenset[str]: - """Return hardware SKU marker names from ``pyproject.toml``. - - A marker is treated as a hardware resource marker when its description - contains the ``[hardware-resource]`` tag, e.g.:: - - "H100: [hardware-resource] Tests that require H100 GPU" - """ +@lru_cache(maxsize=16) +def _marker_names_with_tag(tag: str) -> frozenset[str]: + """Return pytest marker names whose description contains ``tag``.""" data = _load_pyproject() entries = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("markers", []) if not entries: - # Fallback when TOML parsing fails: scan tagged marker lines. text = (_repo_root() / "pyproject.toml").read_text(encoding="utf-8") return frozenset( re.findall( - rf'^\s*"([A-Za-z0-9_]+)\s*:\s*{re.escape(_HARDWARE_RESOURCE_MARKER_TAG)}', + rf'^\s*"([A-Za-z0-9_]+)\s*:.*{re.escape(tag)}', text, flags=re.M, ) @@ -59,7 +66,7 @@ def get_hardware_mark_list() -> frozenset[str]: names: set[str] = set() for entry in entries: text = str(entry) - if _HARDWARE_RESOURCE_MARKER_TAG not in text: + if tag not in text: continue name = text.split(":", 1)[0].strip() if name: @@ -67,138 +74,256 @@ def get_hardware_mark_list() -> frozenset[str]: return frozenset(names) -def cuda_marks(*, res: str, num_cards: int): +@lru_cache(maxsize=1) +def get_hardware_mark_list() -> frozenset[str]: + """Return hardware SKU marker names tagged ``[hardware-resource]``.""" + return _marker_names_with_tag(_HARDWARE_RESOURCE_MARKER_TAG) + + +@lru_cache(maxsize=1) +def get_supported_card_counts() -> frozenset[int]: + """Return ``n`` for each ``cards_{n}`` marker tagged ``[hardware-cards]``.""" + counts: set[int] = set() + for name in _marker_names_with_tag(_HARDWARE_CARDS_MARKER_TAG): + matched = _CARDS_MARK_NAME_RE.fullmatch(name) + if matched: + counts.add(int(matched.group(1))) + return frozenset(counts) + + +@lru_cache(maxsize=1) +def get_supported_platforms() -> frozenset[str]: + """Return platform marker names tagged ``[hardware-platform]``. + + Includes ``cpu`` / ``gpu`` (device-class marks) and accelerator keys. + ``hardware_marks`` / ``hardware_test`` ``res`` dicts reject ``cpu`` / ``gpu``. + """ + return _marker_names_with_tag(_HARDWARE_PLATFORM_MARKER_TAG) + + +@lru_cache(maxsize=1) +def get_level_markers() -> frozenset[str]: + """Return CI level marker names tagged ``[ci-level]`` (``core_model``, ``slow``, …).""" + return _marker_names_with_tag(_CI_LEVEL_MARKER_TAG) + + +@lru_cache(maxsize=8) +def _skus_for_platform(platform: str) -> frozenset[str]: + """SKU names tagged ``[hardware-resource]`` and ``[platform]`` (e.g. ``[cuda]``).""" + return get_hardware_mark_list() & _marker_names_with_tag(f"[{platform}]") + + +def _require_sku(platform: str, res: str) -> None: + allowed = _skus_for_platform(platform) + if res not in allowed: + supported = ", ".join(sorted(allowed)) or f"(none tagged [hardware-resource] [{platform}])" + raise ValueError(f"Invalid {platform} resource type: {res}. Supported: {supported}") + + +def _cards_mark(num_cards: int) -> pytest.MarkDecorator: + """Return ``cards_{n}`` so collection can filter with ``-m cards_2`` (etc.). + + Multi-card selection is ``not cards_1`` (or an explicit ``cards_2`` / ``cards_4``). + """ + import pytest + + if not isinstance(num_cards, int) or isinstance(num_cards, bool) or num_cards < 1: + raise ValueError(f"num_cards must be a positive int, got {num_cards!r}") + if num_cards not in get_supported_card_counts(): + counts = get_supported_card_counts() + supported = ", ".join(f"cards_{n}" for n in sorted(counts)) + supported = supported or "(none tagged [hardware-cards])" + raise ValueError(f"num_cards={num_cards} has no registered marker. Supported: {supported}.") + return getattr(pytest.mark, f"cards_{num_cards}") + + +def _cuda_marks(*, res: str, num_cards: int): + import pytest + from vllm.platforms import current_platform + + _require_sku("cuda", res) + test_resource = getattr(pytest.mark, res) test_platform_detail = pytest.mark.cuda - if res == "L4": - test_resource = pytest.mark.L4 - elif res == "H100": - test_resource = pytest.mark.H100 - else: - raise ValueError(f"Invalid CUDA resource type: {res}. Supported: L4, H100") - marks = [test_resource, test_platform_detail] - if num_cards == 1: + marks = [test_resource, test_platform_detail, _cards_mark(num_cards)] + if num_cards == 1 or not current_platform.is_cuda(): return marks - test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards) - - if not current_platform.is_cuda(): - return marks + [test_distributed] - test_skipif = pytest.mark.skipif( - current_platform.device_count() < num_cards, - reason=f"Need at least {num_cards} CUDA GPUs to run the test.", - ) - return marks + [test_distributed, test_skipif] + return marks + [ + pytest.mark.skipif( + current_platform.device_count() < num_cards, + reason=f"Need at least {num_cards} CUDA GPUs to run the test.", + ) + ] -def rocm_marks(*, res: str, num_cards: int): - test_platform_detail = pytest.mark.rocm - if res == "MI325": - test_resource = pytest.mark.MI325 - else: - raise ValueError(f"Invalid ROCm resource type: {res}. Supported: MI325") - marks = [test_resource, test_platform_detail] - if num_cards == 1: +def _rocm_marks(*, res: str, num_cards: int): + import pytest + + _require_sku("rocm", res) + return [getattr(pytest.mark, res), pytest.mark.rocm, _cards_mark(num_cards)] + + +def _xpu_marks(*, res: str, num_cards: int): + import pytest + from vllm.platforms import current_platform + + _require_sku("xpu", res) + marks = [getattr(pytest.mark, res), pytest.mark.xpu, _cards_mark(num_cards)] + if num_cards == 1 or not current_platform.is_xpu(): return marks - test_distributed = pytest.mark.distributed_rocm(num_cards=num_cards) - return marks + [test_distributed] + return marks + [ + pytest.mark.skipif( + current_platform.device_count() < num_cards, + reason=f"Need at least {num_cards} XPUs to run the test.", + ) + ] -def xpu_marks(*, res: str, num_cards: int): - test_platform_detail = pytest.mark.xpu - if res == "B60": - test_resource = pytest.mark.B60 - else: - raise ValueError(f"Invalid XPU resource type: {res}. Supported: B60") - marks = [test_resource, test_platform_detail] - if num_cards == 1: - return marks - test_distributed = pytest.mark.distributed_xpu(num_cards=num_cards) +def _musa_marks(*, res: str, num_cards: int): + import pytest - if not current_platform.is_xpu(): - return marks + [test_distributed] - test_skipif = pytest.mark.skipif( - current_platform.device_count() < num_cards, - reason=f"Need at least {num_cards} XPUs to run the test.", - ) - return marks + [test_distributed, test_skipif] + _require_sku("musa", res) + return [getattr(pytest.mark, res), pytest.mark.musa, _cards_mark(num_cards)] -def musa_marks(*, res: str, num_cards: int): - test_platform_detail = pytest.mark.musa - if res == "S5000": - test_resource = pytest.mark.S5000 - else: - raise ValueError(f"Invalid MUSA resource type: {res}. Supported: S5000") - marks = [test_resource, test_platform_detail] - if num_cards == 1: - return marks - test_distributed = pytest.mark.distributed_musa(num_cards=num_cards) - return marks + [test_distributed] +@lru_cache(maxsize=1) +def _gpu_res_platforms() -> frozenset[str]: + """Accelerator platforms tagged ``[hardware-platform] [gpu]`` (not the ``gpu`` mark).""" + return (get_supported_platforms() & _marker_names_with_tag(_GPU_DEVICE_CLASS_TAG)) - {"gpu"} -def gpu_marks(*, res: str, num_cards: int): - test_platform = pytest.mark.gpu - if res in ("L4", "H100"): - return [test_platform] + cuda_marks(res=res, num_cards=num_cards) - if res == "MI325": - return [test_platform] + rocm_marks(res=res, num_cards=num_cards) - if res == "B60": - return [test_platform] + xpu_marks(res=res, num_cards=num_cards) - if res == "S5000": - return [test_platform] + musa_marks(res=res, num_cards=num_cards) - raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325, B60, S5000") +def _npu_marks(*, res: str, num_cards: int): + import pytest + + _require_sku("npu", res) + return [pytest.mark.npu, getattr(pytest.mark, res), _cards_mark(num_cards)] -def npu_marks(*, res: str, num_cards: int): - test_platform = pytest.mark.npu - if res == "A2": - test_resource = pytest.mark.A2 - elif res == "A3": - test_resource = pytest.mark.A3 - else: - test_resource = None - if num_cards == 1: - return [mark for mark in [test_platform, test_resource] if mark is not None] - test_distributed = pytest.mark.distributed_npu(num_cards=num_cards) - return [mark for mark in [test_platform, test_resource, test_distributed] if mark is not None] +def _res_platforms() -> frozenset[str]: + """Accelerator keys allowed in ``hardware_test(res=...)`` / ``hardware_marks(res=...)``. + + A ``[hardware-platform]`` name is a ``res`` key only if some SKU is tagged + with that platform (``[cuda]``, ``[npu]``, …). ``cpu`` / ``gpu`` have no SKUs. + """ + return frozenset(p for p in get_supported_platforms() if _skus_for_platform(p)) + + +def _normalize_num_cards(res: dict[str, str], num_cards: int | dict[str, int]) -> dict[str, int]: + allowed = _res_platforms() + device_class = get_supported_platforms() - allowed + for platform in res: + if platform in device_class: + raise ValueError( + f"{platform!r} is not a res dict key; use pytest.mark.{platform} " + f"or an accelerator key ({', '.join(sorted(allowed))})." + ) + if platform not in allowed: + supported = ", ".join(sorted(allowed)) or "(none)" + raise ValueError(f"Unsupported platform: {platform}. res keys: {supported}.") + if isinstance(num_cards, int) and not isinstance(num_cards, bool): + return {platform: num_cards for platform in res} + if not isinstance(num_cards, dict): + raise ValueError(f"num_cards must be a positive int or dict, got {num_cards!r}") + num_cards_dict = dict(num_cards) + for platform in num_cards_dict: + if platform not in res: + raise ValueError(f"Platform '{platform}' in num_cards but not in res.") + for platform in res: + num_cards_dict.setdefault(platform, 1) + return num_cards_dict + + +def _marks_for_platform(platform: str, resource: str, num_cards: int) -> list[pytest.MarkDecorator]: + import pytest + + builders = { + "cuda": _cuda_marks, + "rocm": _rocm_marks, + "xpu": _xpu_marks, + "musa": _musa_marks, + "npu": _npu_marks, + } + builder = builders.get(platform) + if builder is None: + raise ValueError(f"Unsupported platform: {platform}") + marks = builder(res=resource, num_cards=num_cards) + if platform in _gpu_res_platforms(): + return [pytest.mark.gpu] + marks + return marks + + +def _skipif_not_platform(platform: str) -> pytest.MarkDecorator: + import pytest + + from vllm_omni.platforms import current_omni_platform + + checkers = { + "cuda": current_omni_platform.is_cuda, + "rocm": current_omni_platform.is_rocm, + "xpu": current_omni_platform.is_xpu, + "npu": current_omni_platform.is_npu, + "musa": current_omni_platform.is_musa, + } + return pytest.mark.skipif(not checkers[platform](), reason=f"Requires {platform} platform") + + +def _apply_marks(func, marks: list[pytest.MarkDecorator]): + for mark in reversed(marks): + func = mark(func) + return func def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): - for platform, _ in res.items(): - if platform not in ("cuda", "rocm", "xpu", "npu", "musa"): - raise ValueError(f"Unsupported platform: {platform}") - if isinstance(num_cards, int): - num_cards_dict = {platform: num_cards for platform in res.keys()} - else: - num_cards_dict = num_cards - for platform in num_cards_dict.keys(): - if platform not in res: - raise ValueError(f"Platform '{platform}' in num_cards but not in res.") - for platform in res.keys(): - if platform not in num_cards_dict: - num_cards_dict[platform] = 1 + """Return marks for a **single** pytest item. + + Different per-platform ``num_cards`` cannot share one item: ``cards_2`` and + ``cards_4`` would both match ``-m cards_2`` / ``-m cards_4``. Use + ``@hardware_test`` (one collected variant per platform) or call this once + per platform and attach each list to its own ``pytest.param``. + """ + num_cards_dict = _normalize_num_cards(res, num_cards) + unique_counts = set(num_cards_dict.values()) + if len(res) > 1 and len(unique_counts) > 1: + raise ValueError( + "hardware_marks() cannot attach different per-platform num_cards to one pytest item " + f"(got {num_cards_dict}). Use @hardware_test, which collects one variant per platform, " + "or call hardware_marks() once per platform." + ) all_marks: list[pytest.MarkDecorator] = [] for platform, resource in res.items(): - cards = num_cards_dict[platform] - if platform in ("cuda", "rocm", "xpu"): - marks = gpu_marks(res=resource, num_cards=cards) - elif platform == "musa": - marks = musa_marks(res=resource, num_cards=cards) - elif platform == "npu": - marks = npu_marks(res=resource, num_cards=cards) - else: - raise ValueError(f"Unsupported platform: {platform}") - all_marks.extend(marks) + all_marks.extend(_marks_for_platform(platform, resource, num_cards_dict[platform])) return all_marks def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): - all_marks = hardware_marks(res=res, num_cards=num_cards) + """Apply hardware marks; split collected items when card counts differ by platform. - def wrapper(f): - func = f - for mark in reversed(all_marks): - func = mark(func) - return func + A single item cannot carry both ``cards_4`` (CUDA) and ``cards_2`` (ROCm): + ``-m "H100 and cards_4 and cuda"`` would otherwise also match + ``-m "H100 and cards_2 and cuda"``. When ``num_cards`` differs across + platforms, this decorator parametrizes one variant per platform so each + item only has that platform's SKU and ``cards_{n}``. + """ + import pytest - return wrapper + num_cards_dict = _normalize_num_cards(res, num_cards) + mixed_counts = len(res) > 1 and len(set(num_cards_dict.values())) > 1 + if not mixed_counts: + all_marks = hardware_marks(res=res, num_cards=num_cards) + + def apply_union_marks(f): + return _apply_marks(f, all_marks) + + return apply_union_marks + + params = [] + for platform, resource in res.items(): + marks = _marks_for_platform(platform, resource, num_cards_dict[platform]) + marks.append(_skipif_not_platform(platform)) + params.append(pytest.param(platform, marks=marks, id=platform)) + + def apply_split_params(f): + func = pytest.mark.usefixtures(HARDWARE_MARK_NORMALIZATION_FIXTURE)(f) + return pytest.mark.parametrize(HARDWARE_MARK_NORMALIZATION_FIXTURE, params, indirect=True)(func) + + return apply_split_params diff --git a/tests/helpers/media.py b/tests/helpers/media.py index 8eb82fc1f..623c218ad 100644 --- a/tests/helpers/media.py +++ b/tests/helpers/media.py @@ -1,23 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Synthetic media generation and media/text utilities for tests.""" import atexit import base64 import concurrent.futures -import gc import hashlib import io import logging import math +import mimetypes import multiprocessing import os import random import re import subprocess +import sys import tempfile +import threading import time +from concurrent.futures.process import BrokenProcessPool from contextlib import contextmanager from pathlib import Path -from typing import Any +from typing import Any, Literal, TypedDict, overload import numpy as np import soundfile as sf @@ -391,8 +397,16 @@ def generate_synthetic_video( import cv2 import imageio + class _BouncingBall(TypedDict): + x: float + y: float + vx: float + vy: float + radius: int + color_bgr: tuple[int, int, int] + num_balls = random.randint(3, 8) - balls = [] + balls: list[_BouncingBall] = [] for _ in range(num_balls): radius = min(width, height) // 8 if radius < 1: @@ -404,7 +418,16 @@ def generate_synthetic_video( vx = speed * math.cos(angle) vy = speed * math.sin(angle) color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255)) - balls.append({"x": x, "y": y, "vx": vx, "vy": vy, "radius": radius, "color_bgr": color_bgr}) + balls.append( + { + "x": float(x), + "y": float(y), + "vx": float(vx), + "vy": float(vy), + "radius": radius, + "color_bgr": color_bgr, + } + ) video_frames = [] for _ in range(num_frames): @@ -524,27 +547,34 @@ _AUDIO_MIME_BY_SUFFIX = { } -def get_asset_path(relative_path: str | os.PathLike) -> Path: - """Resolve a path under ``tests/assets/`` to its absolute on-disk location.""" - return _TEST_ASSETS_ROOT / Path(relative_path) +def _asset_mime_type(path: Path) -> str: + mime = _AUDIO_MIME_BY_SUFFIX.get(path.suffix.lower()) + if mime is not None: + return mime + guessed, _ = mimetypes.guess_type(path.name) + return guessed or "application/octet-stream" -def load_test_audio_data_url(relative_path: str | os.PathLike) -> str: - """Load a vendored test audio file under ``tests/assets/`` as a base64 data URL. +@overload +def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: Literal[False] = False) -> Path: ... - Used by tests that need real reference audio (e.g. voice cloning) without - relying on the server's ability to fetch external URLs at request time. + +@overload +def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: Literal[True]) -> str: ... + + +def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: bool = False) -> Path | str: + """Resolve a path under ``tests/assets/``. + + When ``as_data_url`` is true, read the file and return a base64 data URL + suitable for embedding reference media in API requests without external URLs. """ - path = get_asset_path(relative_path) - mime = _AUDIO_MIME_BY_SUFFIX.get(path.suffix.lower(), "application/octet-stream") + path = _TEST_ASSETS_ROOT / Path(relative_path) + if not as_data_url: + return path + encoded = base64.b64encode(path.read_bytes()).decode("ascii") - return f"data:{mime};base64,{encoded}" - - -def decode_b64_image(b64: str): - img = Image.open(io.BytesIO(base64.b64decode(b64))) - img.load() - return img + return f"data:{_asset_mime_type(path)};base64,{encoded}" def concat_audio(audio_val) -> np.ndarray: @@ -657,6 +687,10 @@ def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer: @contextmanager def _serialize_whisper_model_download(model_size: str = "small"): """Serialize Whisper cache writes across processes (Linux/Unix), per model.""" + if sys.platform == "win32": + yield + return + import fcntl lock_path = Path.home() / ".cache" / "whisper" / f".{model_size}_model_download.lock" @@ -670,59 +704,152 @@ def _serialize_whisper_model_download(model_size: str = "small"): f.close() -def _whisper_transcribe_in_current_process( - output_path: str, model_size: str = "small", language: str | None = None -) -> str: - import whisper - +def _select_whisper_device() -> str: device_index = None from vllm_omni.platforms import current_omni_platform if current_omni_platform.is_available(): n = current_omni_platform.get_device_count() # Single-GPU runners (e.g. the L4 nightly): the model server already - # occupies device 0. Loading Whisper there, once per concurrent - # request, competes for VRAM and OOMs. Only borrow an accelerator - # when a spare device exists; otherwise validate on CPU. + # occupies device 0, and a Whisper model resident there competes with + # it for VRAM and OOMs. Only borrow an accelerator when a spare device + # exists; otherwise validate on CPU. if n > 1: device_index = n - 1 - if device_index is not None: - torch_device = current_omni_platform.get_torch_device(device_index) - current_omni_platform.set_device(torch_device) - device = str(torch_device) - use_accelerator = True - else: - use_accelerator = False - device = "cpu" + if device_index is None: + return "cpu" - with _serialize_whisper_model_download(model_size): - model = whisper.load_model(model_size, device=device) - try: - text = model.transcribe( - output_path, - temperature=0.0, - word_timestamps=True, - condition_on_previous_text=False, - # None keeps whisper's auto-detection. Do not default this to a - # language: callers include non-English audio tests. - language=language, - )["text"] - finally: - del model - gc.collect() - if use_accelerator: - current_omni_platform.synchronize() - current_omni_platform.empty_cache() + torch_device = current_omni_platform.get_torch_device(device_index) + current_omni_platform.set_device(torch_device) + return str(torch_device) + + +# Populated in the transcription worker, not in the pytest process. +_WHISPER_MODELS: dict[str, Any] = {} + + +def _get_whisper_model(model_size: str) -> Any: + model = _WHISPER_MODELS.get(model_size) + if model is None: + import whisper + + # The device is picked on first load and the model stays on it for the + # worker's lifetime: the current server or runner fixture instance, or + # the test module for callers that transcribe without those fixtures. + device = _select_whisper_device() + with _serialize_whisper_model_download(model_size): + model = whisper.load_model(model_size, device=device) + _WHISPER_MODELS[model_size] = model + return model + + +def _whisper_transcribe_in_current_process( + output_path: str, model_size: str = "small", language: str | None = None +) -> str: + model = _get_whisper_model(model_size) + text = model.transcribe( + output_path, + temperature=0.0, + word_timestamps=True, + condition_on_previous_text=False, + # None keeps whisper's auto-detection. Do not default this to a + # language: callers include non-English audio tests. + language=language, + )["text"] return text or "" +# Serializes a whole submit->result->cleanup on the parent side, so at most one +# call is ever in flight on the single-worker pool. The child already runs +# transcriptions serially, so this costs no throughput; it lets a failed call +# discard the worker without disrupting another caller's in-flight future. +_TRANSCRIBER_CALL_LOCK = threading.Lock() +# Guards the _TRANSCRIBER pointer itself. +_TRANSCRIBER_LOCK = threading.Lock() +_TRANSCRIBER: concurrent.futures.ProcessPoolExecutor | None = None + + +def _get_transcriber() -> concurrent.futures.ProcessPoolExecutor: + global _TRANSCRIBER + with _TRANSCRIBER_LOCK: + if _TRANSCRIBER is None: + ctx = multiprocessing.get_context("spawn") + _TRANSCRIBER = concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) + return _TRANSCRIBER + + +def _discard_transcriber(executor: concurrent.futures.ProcessPoolExecutor) -> None: + """Drop ``executor``, but only while it is still the current one. + + Identity-checked so a stale reference can never shut down a newer worker that + was installed after ``executor`` was replaced. + """ + global _TRANSCRIBER + with _TRANSCRIBER_LOCK: + if _TRANSCRIBER is not executor: + return + _TRANSCRIBER = None + # Joining the worker can block; do it outside the lock. + executor.shutdown(wait=True) + + +def release_audio_transcriber() -> None: + """Shut the transcription worker down, freeing the device memory its models hold. + + Called when a server or runner fixture instance tears down, so that the next + one -- including the next parametrization inside the same test module -- does + not initialize its model while Whisper still occupies the device. A module + teardown fixture repeats it for tests that transcribe without those fixtures. + + Takes the call lock, so it waits for any in-flight transcription to finish + rather than shutting the worker down underneath it. + """ + global _TRANSCRIBER + with _TRANSCRIBER_CALL_LOCK: + with _TRANSCRIBER_LOCK: + executor, _TRANSCRIBER = _TRANSCRIBER, None + if executor is not None: + executor.shutdown(wait=True) + + def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str: - """Convert an audio file to text in an isolated subprocess.""" - ctx = multiprocessing.get_context("spawn") - with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor: - future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size, language) - return future.result() + """Convert an audio file to text in a reused, isolated subprocess. + + The worker outlives the call so its Whisper model is loaded once rather than + once per transcription. The call lock serializes callers onto the single + worker (the child already transcribes serially, so this costs no throughput), + which lets a failed call tear the worker down without racing another caller. + The worker caches one model per size it is asked for, so a run that escalates + to a stronger ASR keeps both models resident until release. + + Any failure discards the worker, restoring the failure isolation of the old + one-process-per-call design: a failure (a ``torch`` OOM, or a + ``KeyboardInterrupt``/``SystemExit`` transported through the future or + raised while ``result()`` blocks) propagates unchanged after the worker -- + and its resident model -- is torn down, and a dead worker + (``BrokenProcessPool``) is additionally retried once. + """ + with _TRANSCRIBER_CALL_LOCK: + for attempt in range(2): + executor = _get_transcriber() + try: + return executor.submit( + _whisper_transcribe_in_current_process, output_path, model_size, language + ).result() + except BrokenProcessPool: + _discard_transcriber(executor) + if attempt == 1: + raise + except BaseException: + # Any other failure -- a task exception (a torch OOM included), or + # a KeyboardInterrupt/SystemExit transported through the future or + # interrupting result() -- leaves the worker and its model + # resident; drop it so it cannot contaminate later calls, matching + # the old per-call teardown. Do not retry. + _discard_transcriber(executor) + raise + raise AssertionError("unreachable") def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small", language: str | None = None) -> str: @@ -742,11 +869,10 @@ __all__ = [ "convert_audio_bytes_to_text", "convert_audio_file_to_text", "cosine_similarity_text", - "decode_b64_image", "generate_synthetic_audio", "generate_synthetic_image", "generate_synthetic_video", "get_asset_path", - "load_test_audio_data_url", "preprocess_text", + "release_audio_transcriber", ] diff --git a/tests/helpers/monitor.py b/tests/helpers/monitor.py new file mode 100644 index 000000000..3402bd186 --- /dev/null +++ b/tests/helpers/monitor.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Device memory monitoring helpers for tests.""" + +from __future__ import annotations + +import threading +import time + +from vllm_omni.platforms import current_omni_platform + + +class DeviceMemoryMonitor: + """Poll global device memory usage.""" + + def __init__(self, device_index: int, interval: float = 0.05): + self.device_index = device_index + self.interval = interval + self._peak_used_mb = 0.0 + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + + def start(self) -> None: + def monitor_loop() -> None: + while not self._stop_event.is_set(): + try: + with current_omni_platform.device(self.device_index): + free_bytes, total_bytes = current_omni_platform.mem_get_info() + used_mb = (total_bytes - free_bytes) / (1024**2) + self._peak_used_mb = max(self._peak_used_mb, used_mb) + except Exception: + pass + time.sleep(self.interval) + + self._thread = threading.Thread(target=monitor_loop, daemon=False) + self._thread.start() + + def stop(self) -> None: + if self._thread is None: + return + self._stop_event.set() + self._thread.join(timeout=2.0) + + @property + def peak_used_mb(self) -> float: + fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2) + fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2) + return max(self._peak_used_mb, fallback_alloc, fallback_reserved) + + def __del__(self): + self.stop() + + +__all__ = ["DeviceMemoryMonitor"] diff --git a/tests/helpers/process.py b/tests/helpers/process.py deleted file mode 100644 index 094de9652..000000000 --- a/tests/helpers/process.py +++ /dev/null @@ -1,134 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import functools -import os -import signal -import subprocess -import sys -import tempfile -from collections.abc import Callable -from contextlib import ExitStack, suppress -from typing import Any, Literal - -import cloudpickle -from typing_extensions import ParamSpec -from vllm.platforms import current_platform - -_P = ParamSpec("_P") - - -def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]: - """Decorator to fork a new process for each test function.""" - - @functools.wraps(func) - def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None: - os.setpgrp() - from _pytest.outcomes import Skipped - - with ( - tempfile.NamedTemporaryFile( - delete=False, mode="w+b", prefix=f"vllm_test_{func.__name__}_{os.getpid()}_", suffix=".exc" - ) as exc_file, - ExitStack() as delete_after, - ): - exc_file_path = exc_file.name - delete_after.callback(os.remove, exc_file_path) - pid = os.fork() - if pid == 0: - delete_after.pop_all() - try: - func(*args, **kwargs) - except Skipped as e: - print(str(e)) - os._exit(0) - except Exception as e: - import traceback - - tb_string = traceback.format_exc() - exc_to_serialize: dict[str, Any] - try: - exc_to_serialize = {"pickled_exception": e} - cloudpickle.dumps(exc_to_serialize) - except (Exception, KeyboardInterrupt): - exc_to_serialize = { - "exception_type": type(e).__name__, - "exception_msg": str(e), - "traceback": tb_string, - } - try: - with open(exc_file_path, "wb") as f: - cloudpickle.dump(exc_to_serialize, f) - except Exception: - print(tb_string) - os._exit(1) - else: - os._exit(0) - else: - pgid = os.getpgid(pid) - _pid, _exitcode = os.waitpid(pid, 0) - old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN) - os.killpg(pgid, signal.SIGTERM) - signal.signal(signal.SIGTERM, old_signal_handler) - if _exitcode != 0: - exc_info = {} - if os.path.exists(exc_file_path): - with suppress(Exception), open(exc_file_path, "rb") as f: - exc_info = cloudpickle.load(f) - if (original_exception := exc_info.get("pickled_exception")) is not None: - assert isinstance(original_exception, Exception) - raise original_exception - if (original_tb := exc_info.get("traceback")) is not None: - raise AssertionError( - f"Test {func.__name__} failed when called with args {args} and kwargs {kwargs}" - f" (exit code: {_exitcode}):\n{original_tb}" - ) from None - raise AssertionError( - f"function {func.__name__} failed when called with args {args} and kwargs {kwargs}" - f" (exit code: {_exitcode})" - ) from None - - return wrapper - - -def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]: - """Decorator to spawn a new process for each test function.""" - - @functools.wraps(f) - def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None: - if os.environ.get("RUNNING_IN_SUBPROCESS") == "1": - return f(*args, **kwargs) - - import torch.multiprocessing as mp - - with suppress(RuntimeError): - mp.set_start_method("spawn") - - module_name = f.__module__ - env = os.environ.copy() - env["RUNNING_IN_SUBPROCESS"] = "1" - - with tempfile.TemporaryDirectory() as tempdir: - output_filepath = os.path.join(tempdir, "new_process.tmp") - input_bytes = cloudpickle.dumps((f, output_filepath)) - cmd = [sys.executable, "-m", f"{module_name}"] - returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env) - try: - returned.check_returncode() - except Exception as e: - raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e - - return wrapper - - -def create_new_process_for_each_test( - method: Literal["spawn", "fork"] | None = None, -) -> Callable[[Callable[_P, None]], Callable[_P, None]]: - """Creates a decorator that runs each test function in a new process.""" - if method is None: - use_spawn = current_platform.is_xpu() - method = "spawn" if use_spawn else "fork" - assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'" - if method == "fork": - return fork_new_process_for_each_test - return spawn_new_process_for_each_test diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 82fcc308d..db4cba893 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -1,13 +1,13 @@ -"""Server/client/runner runtime primitives for tests.""" +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Server and offline runner runtime primitives for tests. + +Request clients and response types live in ``tests.helpers.client``; +this module re-exports them for backward-compatible imports. +""" -import asyncio -import base64 -import concurrent.futures -import copy import errno -import gc -import io -import json import os import socket import subprocess @@ -16,126 +16,27 @@ import tempfile import threading import time from collections.abc import Generator -from dataclasses import asdict, dataclass -from io import BytesIO +from contextlib import contextmanager from pathlib import Path -from typing import Any, NamedTuple, cast -from urllib.parse import quote +from typing import Any, NamedTuple import numpy as np import psutil -import requests -import soundfile as sf -import torch import yaml -from openai import APIError, OpenAI, omit -from PIL import Image -from vllm import TextPrompt, envs -from vllm.distributed.parallel_state import ( - destroy_distributed_environment, - destroy_model_parallel, -) +from vllm import TextPrompt from vllm.logger import init_logger -from tests.helpers.assertions import ( - assert_audio_speech_response, - assert_diffusion_response, - assert_http_error, - assert_images_generations_response, - assert_omni_response, -) -from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.clean import cleanup_test_environment from tests.helpers.media import ( - _merge_base64_audio_to_segment, - decode_b64_image, + release_audio_transcriber, ) +from tests.helpers.stage_config import load_stage_ids, load_stage_replica_counts from tests.model_tests.diffusion.utils import resolve_tiny_model_path from vllm_omni.config.stage_config import resolve_deploy_yaml -from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) - -def cleanup_dist_env_and_memory(shutdown_ray: bool = False): - # Reset environment variable cache - envs.disable_envs_cache() - - # Reset rocm_aiter_ops class variables to match current os.environ. - # These are class-level attributes that persist across tests and are - # NOT restored by monkeypatch (which only restores os.environ). - from vllm_omni.platforms import current_omni_platform - - if current_omni_platform.is_rocm(): - from vllm._aiter_ops import rocm_aiter_ops - - rocm_aiter_ops.refresh_env_variables() - - # Ensure all objects are not frozen before cleanup - gc.unfreeze() - - destroy_model_parallel() - destroy_distributed_environment() - if shutdown_ray: - import ray # Lazy import Ray - - ray.shutdown() - gc.collect() - - if not current_omni_platform.is_cpu(): - current_omni_platform.empty_cache() - try: - torch._C._host_emptyCache() - except AttributeError: - logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5") - - -def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None: - try: - data = r.json() - if isinstance(data, (dict, list)): - return data - except Exception: - pass - return None - - -def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None: - """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``).""" - eb = cfg.get("extra_body") - if not eb: - return None - h, w = eb.get("height"), eb.get("width") - if (isinstance(h, (list, tuple)) or isinstance(w, (list, tuple))) and not ( - isinstance(h, (list, tuple)) and isinstance(w, (list, tuple)) - ): - raise ValueError("extra_body height and width must both be lists or both be scalars") - if not (isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))): - return None - if len(h) != len(w): - raise ValueError(f"height and width lists must have equal length; got {len(h)=} {len(w)=}") - n = len(h) - n_out = eb.get("num_outputs_per_prompt") - if n_out is not None: - n_out = int(n_out) - if n_out != n: - raise ValueError( - "When height/width are lists, num_outputs_per_prompt must equal their length; " - f"got num_outputs_per_prompt={n_out}, len(lists)={n}" - ) - splits: list[dict[str, Any]] = [] - for i in range(n): - sub = copy.deepcopy(cfg) - sub_eb = dict(sub.get("extra_body") or {}) - sub_eb["height"] = int(h[i]) - sub_eb["width"] = int(w[i]) - sub_eb["num_outputs_per_prompt"] = 1 - sub["extra_body"] = sub_eb - splits.append(sub) - return splits - - PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None PromptImageInput = list[Any] | Any | None PromptVideoInput = list[Any] | Any | None @@ -174,16 +75,30 @@ def get_open_port(host: str = "127.0.0.1", *, max_attempts: int = 128) -> int: ) from last_exc +def get_distributed_init_method(prefix: str = "torch_dist_init_") -> str: + """Return a ``file://`` init_method for a ``torch.distributed`` process group. + + Args: + prefix: Prefix for the temporary rendezvous filename. Defaults to + ``torch_dist_init_``. + + Returns: + An init_method string for ``torch.distributed.init_process_group``. + """ + with tempfile.NamedTemporaryFile(prefix=prefix) as f: + return f"file://{f.name}" + + def dummy_messages_from_mix_data( - system_prompt: dict[str, Any] = None, + system_prompt: dict[str, Any] | None = None, video_data_url: Any = None, audio_data_url: Any = None, image_data_url: Any = None, - content_text: str = None, + content_text: str | None = None, ): """Create messages with video、image、audio data URL for OpenAI API.""" if content_text is not None: - content = [{"type": "text", "text": content_text}] + content: list[dict[str, Any]] = [{"type": "text", "text": content_text}] else: content = [] @@ -246,9 +161,7 @@ class OmniServer: env_dict: dict[str, str] | None = None, use_omni: bool = True, ) -> None: - run_pre_test_cleanup() - run_post_test_cleanup() - cleanup_dist_env_and_memory() + cleanup_test_environment() self.model = model args = list(serve_args) self.serve_args = args @@ -261,7 +174,6 @@ class OmniServer: def _start_server(self) -> None: env = os.environ.copy() - env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") if self.env_dict is not None: env.update(self.env_dict) @@ -322,7 +234,7 @@ class OmniServer: if proc.status() != psutil.STATUS_ZOMBIE: return False try: - os.waitpid(proc.pid, os.WNOHANG) + os.waitpid(proc.pid, getattr(os, "WNOHANG", 0)) except ChildProcessError: pass return True @@ -446,15 +358,19 @@ class OmniServer: pass def __enter__(self): - self._start_server() + try: + self._start_server() + except BaseException: + # ``__exit__`` is not invoked when ``__enter__`` raises; roll back + # any processes / log handles already launched. + self.__exit__(None, None, None) + raise return self def __exit__(self, exc_type, exc_val, exc_tb): if self.proc: self._kill_process_tree(self.proc.pid) - run_pre_test_cleanup() - run_post_test_cleanup() - cleanup_dist_env_and_memory() + cleanup_test_environment() class OmniServerStageCli(OmniServer): @@ -473,7 +389,6 @@ class OmniServerStageCli(OmniServer): super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True) self.stage_config_path = stage_config_path self.master_port = get_open_port() - self.visible_device_list = self._load_visible_device_list(env_dict) resolved_cfg = resolve_deploy_yaml(stage_config_path) # Dump the resolved deploy config so CI logs show each stage's # gpu_memory_utilization / max_model_len / max_num_seqs after @@ -484,113 +399,15 @@ class OmniServerStageCli(OmniServer): f"{yaml.safe_dump(resolved_cfg, sort_keys=False, default_flow_style=False)}", flush=True, ) - self.stage_runtime_devices = self._load_stage_runtime_devices(resolved_cfg) - self.stage_ids = stage_ids or self._load_stage_ids(resolved_cfg) + self.stage_ids = stage_ids or load_stage_ids(resolved_cfg) if 0 not in self.stage_ids: raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}") - self.stage_replica_counts = self._load_stage_replica_counts(resolved_cfg) + self.stage_replica_counts = load_stage_replica_counts(resolved_cfg) self.stage_procs: dict[tuple[int, int], subprocess.Popen] = {} + self._stage_log_paths: dict[tuple[int, int], Path] = {} + self._stage_log_files: dict[tuple[int, int], Any] = {} self.proc = None - @staticmethod - def _stage_entries(cfg: dict) -> list[dict]: - """Return the list of stage entries from either legacy (``stage_args``) - or new-schema (``stages``) deploy YAMLs.""" - return cfg.get("stage_args") or cfg.get("stages") or [] - - @staticmethod - def _load_stage_ids(resolved_config: dict) -> list[int]: - stage_ids = [ - stage["stage_id"] for stage in OmniServerStageCli._stage_entries(resolved_config) if "stage_id" in stage - ] - if not stage_ids: - raise ValueError("No stage IDs found in resolved config") - return stage_ids - - @staticmethod - def _load_stage_runtime_devices(resolved_config: dict) -> dict[int, str]: - runtime_devices: dict[int, str] = {} - for stage in OmniServerStageCli._stage_entries(resolved_config): - stage_id = stage.get("stage_id") - # New schema: stage.devices is flat at stage level. - # Legacy schema: stage.runtime.devices is nested. - devices = stage.get("devices") or stage.get("runtime", {}).get("devices") - if stage_id is not None and devices: - runtime_devices[int(stage_id)] = str(devices) - return runtime_devices - - @staticmethod - def _load_stage_replica_counts(resolved_config: dict) -> dict[int, int]: - replica_counts: dict[int, int] = {} - for stage in OmniServerStageCli._stage_entries(resolved_config): - stage_id = stage.get("stage_id") - if stage_id is None: - continue - replica_counts[int(stage_id)] = max( - 1, - int(stage.get("num_replicas") or stage.get("runtime", {}).get("num_replicas", 1)), - ) - return replica_counts - - @classmethod - def _parse_device_list(cls, devices: str | int) -> list[str]: - if isinstance(devices, int): - if devices < 0: - raise ValueError("Device IDs must be non-negative integers") - return [str(devices)] - return [token.strip() for token in str(devices).split(",") if token.strip()] - - @classmethod - def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None: - env = os.environ.copy() - if env_dict is not None: - env.update(env_dict) - - env_var = getattr(current_omni_platform, "device_control_env_var", None) - if env_var and env_var in env: - return [token.strip() for token in env[env_var].split(",") if token.strip()] - return None - - @classmethod - def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str: - device_list = cls._parse_device_list(devices) - - if visible_device_list is None: - return ",".join(device_list) - - if not all(device.isdigit() for device in device_list): - raise ValueError("Logical devices must be non-negative integers") - - logical_ids = [int(device) for device in device_list] - if logical_ids and max(logical_ids) >= len(visible_device_list): - raise ValueError( - f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices" - ) - - return ",".join(visible_device_list[idx] for idx in logical_ids) - - def _devices_for_replica(self, stage_id: int, devices: str, replica_id: int) -> str: - replica_count = self.stage_replica_counts.get(stage_id, 1) - if replica_count == 1: - return devices - - device_list = self._parse_device_list(devices) - if len(device_list) % replica_count != 0: - raise ValueError( - f"Stage {stage_id} has {len(device_list)} device(s) for {replica_count} replica(s); " - "device count must be divisible by replica count" - ) - devices_per_replica = len(device_list) // replica_count - start = replica_id * devices_per_replica - return ",".join(device_list[start : start + devices_per_replica]) - - def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str, replica_id: int = 0) -> None: - replica_devices = self._devices_for_replica(stage_id, devices, replica_id) - mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, replica_devices) - env_var = getattr(current_omni_platform, "device_control_env_var", None) - if env_var: - env[env_var] = mapped_devices - def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> list[str]: cmd = [ sys.executable, @@ -621,7 +438,6 @@ class OmniServerStageCli(OmniServer): def _launch_stage(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> None: env = os.environ.copy() - env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" if self.env_dict is not None: env.update(self.env_dict) @@ -631,11 +447,9 @@ class OmniServerStageCli(OmniServer): # debugging "Stage N exited before API server ready" doesn't rely on # guessing; the file is surfaced in the RuntimeError message. log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_replica_{replica_id}_{self.master_port}.log" - self._stage_log_paths = getattr(self, "_stage_log_paths", {}) stage_key = (stage_id, replica_id) self._stage_log_paths[stage_key] = log_path - log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in __exit__ - self._stage_log_files = getattr(self, "_stage_log_files", {}) + log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in ``__exit__`` self._stage_log_files[stage_key] = log_fh proc = subprocess.Popen( cmd, @@ -655,7 +469,7 @@ class OmniServerStageCli(OmniServer): for (stage_id, replica_id), proc in self.stage_procs.items(): ret = proc.poll() if ret is not None: - log_path = getattr(self, "_stage_log_paths", {}).get((stage_id, replica_id)) + log_path = self._stage_log_paths.get((stage_id, replica_id)) tail = "" if log_path and log_path.exists(): try: @@ -712,7 +526,7 @@ class OmniServerStageCli(OmniServer): captures them post-run. Head covers engine init; tail covers whatever state the stage was in when it was torn down. """ - log_paths = getattr(self, "_stage_log_paths", {}) or {} + log_paths = self._stage_log_paths for stage_id, replica_id in sorted(log_paths): log_path = log_paths[(stage_id, replica_id)] if not log_path or not log_path.exists(): @@ -746,1877 +560,30 @@ class OmniServerStageCli(OmniServer): proc = self.stage_procs[stage_key] if proc.poll() is None: self._kill_process_tree(proc.pid) - run_pre_test_cleanup() - run_post_test_cleanup() - cleanup_dist_env_and_memory() - -@dataclass -class OmniResponse: - """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``).""" - - text_content: str | None = None - audio_data: list[str] | None = None - audio_content: str | None = None - audio_format: str | None = None - audio_bytes: bytes | None = None - #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the - #: OpenAI client call through response parsing and local post-process (e.g. audio decode). - e2e_latency: float | None = None - success: bool = False - prompt_tokens: int | None = None - cached_tokens: int | None = None - logprobs: list | None = None - #: HTTP status + error text for the error-handling path (e.g. validator - #: rejections); populated when the OpenAI client raises an APIError. - status_code: int | None = None - error_message: str | None = None - - -@dataclass -class DiffusionResponse: - """Decoded diffusion output from chat completions or offline runner (not raw ``requests``).""" - - text_content: str | None = None - images: list[Image.Image] | None = None - audios: list[Any] | None = None - videos: list[Any] | None = None - #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before - #: ``chat.completions.create`` through local image / audio decode. - e2e_latency: float | None = None - success: bool = False - - -@dataclass -class HttpResponse: - """Normalized view of a ``requests`` response from :class:`OpenAIClientHandler` HTTP helpers.""" - - status_code: int - success: bool - error_message: str | None = None - json_body: dict[str, Any] | list[Any] | None = None - - -@dataclass -class WebSocketJsonResponse: - """First JSON object delivered as a text WebSocket frame (streaming endpoints).""" - - json_body: dict[str, Any] - - -@dataclass -class OpenPIWebSocketResponse: - """Msgpack WebSocket session against ``/v1/realtime/robot/openpi``.""" - - server_metadata: dict[str, Any] - operation_responses: list[Any] - actions: dict[str, np.ndarray] | None = None - action_tensors: list[np.ndarray] | None = None - - -def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]: - """Build a minimal DROID-style observation payload for the OpenPI robot endpoint.""" - identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32) - identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) - return { - "session_id": session_id, - "video": { - "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), - "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), - }, - "state": { - "eef_9d": identity_eef_9d, - "gripper_position": np.zeros((1, 1, 1), dtype=np.float32), - "joint_position": np.zeros((1, 1, 7), dtype=np.float32), - }, - "language": {"annotation.language.language_instruction": [["pick up the object"]]}, - } - - -DREAMZERO_DEFAULT_PROMPT = ( - "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan" -) -DREAMZERO_ACTION_HORIZON = 24 -DREAMZERO_ACTION_DIM = 8 -DREAMZERO_CAMERA_FILES = { - "observation/exterior_image_0_left": "exterior_image_1_left.mp4", - "observation/exterior_image_1_left": "exterior_image_2_left.mp4", - "observation/wrist_image_left": "wrist_image_left.mp4", -} - - -def _require_opencv() -> Any: - try: - import cv2 - except ImportError as exc: # pragma: no cover - optional e2e dependency - raise ModuleNotFoundError("DreamZero OpenPI test dependencies are missing: opencv-python") from exc - return cv2 - - -def load_dreamzero_camera_frames(video_dir: Path) -> dict[str, np.ndarray]: - cv2 = _require_opencv() - camera_frames: dict[str, np.ndarray] = {} - for camera_key, file_name in DREAMZERO_CAMERA_FILES.items(): - video_path = video_dir / file_name - if not video_path.exists(): - raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}") - cap = cv2.VideoCapture(str(video_path)) - frames = [] - while True: - ok, frame = cap.read() - if not ok: - break - frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) - cap.release() - if not frames: - raise RuntimeError(f"No frames loaded from {video_path}") - camera_frames[camera_key] = np.stack(frames, axis=0) - return camera_frames - - -def build_dreamzero_demo_observations( - camera_frames: dict[str, np.ndarray], - *, - prompt: str, - session_id: str, - num_chunks: int = 2, -) -> list[dict[str, Any]]: - if num_chunks < 1: - raise ValueError("num_chunks must be at least 1") - - relative_offsets = [-23, -16, -8, 0] - total_frames = min(frames.shape[0] for frames in camera_frames.values()) - frame_schedules = [[0]] - current_frame = 23 - for _ in range(num_chunks - 1): - indices = [max(current_frame + offset, 0) for offset in relative_offsets] - if indices[-1] >= total_frames: - break - frame_schedules.append(indices) - current_frame += DREAMZERO_ACTION_HORIZON - - observations: list[dict[str, Any]] = [] - for frame_indices in frame_schedules: - obs: dict[str, Any] = {} - for camera_key, all_frames in camera_frames.items(): - selected = all_frames[frame_indices] - obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected - obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) - obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) - obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) - obs["prompt"] = prompt - obs["session_id"] = session_id - observations.append(obs) - return observations - - -class OpenPIWebSocketSession: - """Persistent msgpack session for ``/v1/realtime/robot/openpi``.""" - - DEFAULT_PING_INTERVAL_SECS = 300 - DEFAULT_PING_TIMEOUT_SECS = 3600 - - @staticmethod - def _require_dependencies() -> tuple[Any, Any]: - try: - import websockets.sync.client as websockets_client - except ImportError as exc: # pragma: no cover - optional e2e dependency - raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: websockets") from exc - try: - from openpi_client import msgpack_numpy - except ImportError as exc: # pragma: no cover - optional e2e dependency - raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: openpi-client") from exc - return websockets_client, msgpack_numpy - - def __init__( - self, - uri: str, - *, - ping_interval: float = DEFAULT_PING_INTERVAL_SECS, - ping_timeout: float = DEFAULT_PING_TIMEOUT_SECS, - open_timeout: float = 120.0, - close_timeout: float = 120.0, - ) -> None: - websockets_client, msgpack_numpy = self._require_dependencies() - self._msgpack_numpy = msgpack_numpy - self._packer = msgpack_numpy.Packer() - self._conn = websockets_client.connect( - uri, - compression=None, - max_size=None, - ping_interval=ping_interval, - ping_timeout=ping_timeout, - open_timeout=open_timeout, - close_timeout=close_timeout, - ) - server_metadata = msgpack_numpy.unpackb(self._conn.recv()) - if not isinstance(server_metadata, dict): - raise TypeError(f"Expected dict metadata from {uri}, got {type(server_metadata)!r}") - self._server_metadata = server_metadata - - def get_server_metadata(self) -> dict[str, Any]: - return dict(self._server_metadata) - - def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]: - response = self._send_operation("infer", obs) - if not isinstance(response, dict): - raise TypeError(f"Expected dict infer response, got {type(response)!r}") - return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()} - - def reset(self, reset_info: dict[str, Any] | None = None) -> str: - response = self._send_operation("reset", dict(reset_info or {})) - return str(response["status"]) - - def close(self) -> None: - self._conn.close() - - def _send_operation(self, endpoint: str, payload: dict[str, Any]) -> Any: - body = dict(payload) - body["endpoint"] = endpoint - self._conn.send(self._packer.pack(body)) - raw = self._conn.recv() - if isinstance(raw, str): - raise RuntimeError(f"OpenPI {endpoint!r} failed: {raw}") - decoded = self._msgpack_numpy.unpackb(raw) - if isinstance(decoded, dict) and decoded.get("type") == "error": - raise RuntimeError(f"OpenPI {endpoint!r} failed: {decoded.get('message')}") - if endpoint == "reset": - if not isinstance(decoded, dict) or decoded.get("status") != "reset successful": - raise RuntimeError(f"Unexpected OpenPI reset response: {decoded!r}") - return decoded - return decoded - - -def _merge_http_expectation_kwargs( - base: dict[str, Any] | None, - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, -) -> dict[str, Any]: - cfg = dict(base or {}) - if err_code is not None: - cfg["err_code"] = err_code - if err_message is not None: - cfg["err_message"] = err_message - return cfg - - -def _merge_ws_expectation_kwargs( - base: dict[str, Any] | None, - *, - err_message: str | tuple[str, ...] | list[str] | None = None, - ws_json_type: str | None = None, - ws_error_code: str | None = None, -) -> dict[str, Any]: - cfg = dict(base or {}) - if err_message is not None: - cfg["err_message"] = err_message - if ws_json_type is not None: - cfg["ws_json_type"] = ws_json_type - if ws_error_code is not None: - cfg["ws_error_code"] = ws_error_code - return cfg - - -def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None: - jb = resp.json_body - want_type = cfg.get("ws_json_type") - if want_type is not None: - assert jb.get("type") == want_type, (jb, want_type) - want_code = cfg.get("ws_error_code") - if want_code is not None: - assert jb.get("code") == want_code, (jb, want_code) - err_message = cfg.get("err_message") - if err_message is not None: - assert_http_error(resp, err_message=err_message, websocket_json_message=True) - - -def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse: - """Concatenate images in order; ``e2e_latency`` is wall-clock of the batch (set by caller) or max of parts.""" - merged = DiffusionResponse() - merged.success = all(p.success for p in parts) and len(parts) > 0 - imgs: list[Image.Image] = [] - for p in parts: - if p.images: - imgs.extend(p.images) - merged.images = imgs if imgs else None - latencies = [p.e2e_latency for p in parts if p.e2e_latency is not None] - merged.e2e_latency = max(latencies) if latencies else None - return merged - - -class OpenAIClientHandler: - def __init__( - self, - host: str = "127.0.0.1", - port: int = None, - api_key: str = "EMPTY", - run_level: str = None, - *, - log_stats: bool = True, - ): - if port is None: - port = get_open_port() - self.base_url = f"http://{host}:{port}" - self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key) - self.run_level = run_level - self.log_stats = log_stats - - def _print_client_stat(self, message: str) -> None: - if self.log_stats: - print(message, flush=True) - - def _process_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse: - """Wall clock from *before* ``chat.completions.create`` through stream drain + local decode.""" - result = OmniResponse() - try: - text_content = "" - audio_data = [] - for chunk in chat_completion: - for choice in chunk.choices: - content = getattr(getattr(choice, "delta", None), "content", None) - modality = getattr(chunk, "modality", None) - if modality == "audio" and content: - audio_data.append(content) - elif modality == "text" and content: - text_content += content - # Usage is yielded after the last token - if chunk.usage: - result.prompt_tokens = chunk.usage.prompt_tokens - if details := getattr(chunk.usage, "prompt_tokens_details", None): - result.cached_tokens = details.cached_tokens - - if audio_data: - merged_seg = _merge_base64_audio_to_segment(audio_data) - wav_buf = BytesIO() - merged_seg.export(wav_buf, format="wav") - result.audio_bytes = wav_buf.getvalue() - result.text_content = text_content - result.audio_data = audio_data - result.e2e_latency = time.perf_counter() - wall_start - result.success = True - except Exception as e: - msg = f"Stream processing error: {str(e)}" - print(f"Error: {msg}") - return result - - def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse: - """Wall clock from *before* ``chat.completions.create`` through response parse + local decode.""" - result = OmniResponse() - try: - audio_data = None - text_content = None - for choice in chat_completion.choices: - if hasattr(choice.message, "audio") and choice.message.audio is not None: - audio_data = choice.message.audio.data - if hasattr(choice.message, "content") and choice.message.content is not None: - text_content = choice.message.content - # Extract cached & prompt token counts for prefix caching tests - usage = getattr(chat_completion, "usage", None) - if usage: - result.prompt_tokens = usage.prompt_tokens - if details := getattr(usage, "prompt_tokens_details", None): - result.cached_tokens = details.cached_tokens - if audio_data: - result.audio_bytes = base64.b64decode(audio_data) - result.text_content = text_content - result.e2e_latency = time.perf_counter() - wall_start - if chat_completion.choices and chat_completion.choices[0].logprobs is not None: - result.logprobs = chat_completion.choices[0].logprobs.content - result.success = True - except Exception as e: - msg = f"Non-stream processing error: {str(e)}" - print(f"Error: {msg}") - return result - - def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse: - """Wall clock from *before* ``chat.completions.create`` through image decode.""" - result = DiffusionResponse() - try: - images = [] - audios = [] - for choice in chat_completion.choices: - content = getattr(choice.message, "content", None) - if isinstance(content, list): - for item in content: - image_url = None - if isinstance(item, dict): - image_url = item.get("image_url", {}).get("url") - else: - image_url_obj = getattr(item, "image_url", None) - image_url = getattr(image_url_obj, "url", None) if image_url_obj else None - if image_url and image_url.startswith("data:image"): - b64_data = image_url.split(",", 1)[1] - images.append(decode_b64_image(b64_data)) - - # OpenAI audio responses (e.g. AudioX text-to-audio) populate `message.audio`. - audio_obj = getattr(choice.message, "audio", None) - audio_b64 = getattr(audio_obj, "data", None) if audio_obj is not None else None - if audio_b64: - audios.append( - { - "wav_bytes": base64.b64decode(audio_b64), - "id": getattr(audio_obj, "id", None), - "expires_at": getattr(audio_obj, "expires_at", None), - } - ) - result.images = images if images else None - result.audios = audios if audios else None - result.e2e_latency = time.perf_counter() - wall_start - result.success = True - except Exception as e: - msg = f"Diffusion response processing error: {str(e)}" - print(f"Error: {msg}") - return result - - def _http_response_from_requests(self, r: requests.Response) -> HttpResponse: - payload = _parse_response_json(r) - ok = 200 <= r.status_code < 300 - return HttpResponse( - status_code=r.status_code, - success=ok, - error_message=None if ok else (r.text[:8000] if r.text else None), - json_body=payload, - ) - - def send_health_http_request( - self, - request_config: dict[str, Any] | None = None, - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """GET ``/health`` (raw ``requests``). - - ``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for - :func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args). - """ - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0))) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_models_http_request( - self, - request_config: dict[str, Any] | None = None, - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`).""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = requests.get( - self._build_url("/v1/models"), - headers={"Accept": "application/json"}, - timeout=float(cfg.get("timeout", 120.0)), - ) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_chat_completions_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests).""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_batched_chat_completions_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """Post to batched chat completions.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_completions_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/completions`` with ``json`` or ``raw_body``.""" - # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_omni_sleep_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_omni_wakeup_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/omni/wakeup``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_audio_voices_list_http_request( - self, - request_config: dict[str, Any] | None = None, - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """GET ``/v1/audio/voices``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = requests.get( - self._build_url("/v1/audio/voices"), - headers={"Accept": "application/json"}, - timeout=float(cfg.get("timeout", 120.0)), - ) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_audio_voices_create_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_audio_voices_delete_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - name = cfg["name"] - timeout = float(cfg.get("timeout", 120.0)) - path = f"/v1/audio/voices/{quote(str(name), safe='')}" - r = requests.delete( - self._build_url(path), - headers={"Accept": "application/json"}, - timeout=timeout, - ) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_audio_speech_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/audio/speech`` with ``json`` or ``raw_body``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_audio_speech_batch_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/audio/speech/batch``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_audio_generate_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/audio/generate``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_images_generations_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - if cfg.get("err_code") is None: - assert resp.success, resp.error_message - payload = resp.json_body - assert isinstance(payload, dict) - assert_images_generations_response(payload, cfg, run_level=self.run_level) - return [resp] - - def send_images_edits_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_videos_create_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_videos_sync_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """POST ``/v1/videos/sync``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_videos_list_http_request( - self, - request_config: dict[str, Any] | None = None, - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """GET ``/v1/videos`` — optional ``params``, ``timeout``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - r = requests.get( - self._build_url("/v1/videos"), - params=cfg.get("params"), - headers={"Accept": "application/json"}, - timeout=float(cfg.get("timeout", 120.0)), - ) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_video_retrieve_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """GET ``/v1/videos/{video_id}``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - video_id = cfg["video_id"] - timeout = float(cfg.get("timeout", 120.0)) - r = requests.get( - self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"), - headers={"Accept": "application/json"}, - timeout=timeout, - ) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_video_delete_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """DELETE ``/v1/videos/{video_id}``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - video_id = cfg["video_id"] - timeout = float(cfg.get("timeout", 120.0)) - r = requests.delete( - self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"), - headers={"Accept": "application/json"}, - timeout=timeout, - ) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def send_video_content_http_request( - self, - request_config: dict[str, Any], - *, - err_code: int | tuple[int, ...] | list[int] | None = None, - err_message: str | tuple[str, ...] | list[str] | None = None, - ) -> list[HttpResponse]: - """GET ``/v1/videos/{video_id}/content``.""" - cfg = _merge_http_expectation_kwargs( - request_config, - err_code=err_code, - err_message=err_message, - ) - video_id = cfg["video_id"] - timeout = float(cfg.get("timeout", 120.0)) - r = requests.get( - self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"), - timeout=timeout, - ) - resp = self._http_response_from_requests(r) - assert_http_error( - resp, - err_code=cfg.get("err_code"), - err_message=cfg.get("err_message"), - ) - return [resp] - - def _build_ws_url(self, path: str) -> str: - """Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers.""" - base = self.base_url.rstrip("/") - suffix = "/" + path.lstrip("/") - if base.startswith("http://"): - return "ws://" + base.removeprefix("http://") + suffix - if base.startswith("https://"): - return "wss://" + base.removeprefix("https://") + suffix - raise ValueError(f"Unsupported base_url for WebSocket: {base!r}") - - def _send_websocket_first_json_request( - self, - path: str, - cfg: dict[str, Any], - ) -> list[WebSocketJsonResponse]: - """Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`. - - ``request_config`` keys: - - - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server - speaks first, e.g. ``/v1/realtime`` rejection path). - - ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame - (e.g. ``["session.created"]`` on ``/v1/realtime``). - - ``timeout``: seconds to wait for the first inbound text frame (default ``120``). - - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present. - """ - send_frames_raw = cfg.get("send_frames") - if send_frames_raw is None: - frames: list[str] = [] - elif isinstance(send_frames_raw, str): - frames = [send_frames_raw] - else: - frames = list(send_frames_raw) - - timeout = float(cfg.get("timeout", 120.0)) - uri = self._build_ws_url(path) - skip_types = set(cfg.get("ws_skip_types") or []) - - connect_kw: dict[str, Any] = {} - if "ws_max_size" in cfg: - connect_kw["max_size"] = cfg["ws_max_size"] - - async def _recv_first_json_object() -> WebSocketJsonResponse: - import websockets - - async with websockets.connect(uri, **connect_kw) as ws: - for frame in frames: - await ws.send(frame) - while True: - raw = await asyncio.wait_for(ws.recv(), timeout=timeout) - if not isinstance(raw, str): - raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}") + # Close per-stage log handles; delete temp files unless VLLM_OMNI_KEEP_LOG is set. + keep_logs = os.environ.get("VLLM_OMNI_KEEP_LOG", "").lower() in ("1", "true", "yes") + log_files = self._stage_log_files + log_paths = self._stage_log_paths + for stage_key in set(log_files) | set(log_paths): + log_fh = log_files.get(stage_key) + if log_fh is not None: + try: + log_fh.close() + except Exception: + pass + if not keep_logs: + log_path = log_paths.get(stage_key) + if log_path is not None: try: - data = json.loads(raw) - except json.JSONDecodeError as exc: - raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc - if not isinstance(data, dict): - raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}") - if skip_types and data.get("type") in skip_types: - continue - return WebSocketJsonResponse(json_body=data) - - resp = asyncio.run(_recv_first_json_object()) - _run_ws_expectations_from_request_config(cfg, resp) - return [resp] - - def send_audio_speech_stream_ws_request( - self, - request_config: dict[str, Any], - *, - err_message: str | tuple[str, ...] | list[str] | None = None, - ws_json_type: str | None = None, - ws_error_code: str | None = None, - ) -> list[WebSocketJsonResponse]: - """WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame.""" - cfg = _merge_ws_expectation_kwargs( - request_config, - err_message=err_message, - ws_json_type=ws_json_type, - ws_error_code=ws_error_code, - ) - return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg) - - def send_video_chat_stream_ws_request( - self, - request_config: dict[str, Any], - *, - err_message: str | tuple[str, ...] | list[str] | None = None, - ws_json_type: str | None = None, - ws_error_code: str | None = None, - ) -> list[WebSocketJsonResponse]: - """WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame.""" - cfg = _merge_ws_expectation_kwargs( - request_config, - err_message=err_message, - ws_json_type=ws_json_type, - ws_error_code=ws_error_code, - ) - return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg) - - def send_realtime_ws_request( - self, - request_config: dict[str, Any] | None = None, - *, - err_message: str | tuple[str, ...] | list[str] | None = None, - ws_json_type: str | None = None, - ws_error_code: str | None = None, - ) -> list[WebSocketJsonResponse]: - """WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated).""" - cfg = _merge_ws_expectation_kwargs( - request_config, - err_message=err_message, - ws_json_type=ws_json_type, - ws_error_code=ws_error_code, - ) - return self._send_websocket_first_json_request("/v1/realtime", cfg) - - def send_robot_openpi_ws_request( - self, - request_config: dict[str, Any] | None = None, - ) -> list[OpenPIWebSocketResponse]: - """WebSocket ``/v1/realtime/robot/openpi`` — msgpack metadata plus optional infer/reset ops. - - ``request_config`` keys: - - - ``operations``: optional sequence of ``{"endpoint": "infer"|"reset", "payload": {...}}``. - - ``run_default_policy_session``: when true and ``operations`` is omitted, run infer then reset - using :func:`build_openpi_droid_observation`. - - ``run_dreamzero_policy_session``: when true and ``operations`` is omitted, run the DreamZero - infer/reset/infer sequence using :func:`build_dreamzero_demo_observations`. - - ``video_dir``: required for ``run_dreamzero_policy_session``. - - ``prompt``: language instruction for DreamZero observations (optional). - - ``session_id``: used by the default policy session observation builder. - - ``num_chunks``: DreamZero chunk count (default ``2``). - - ``timeout``: seconds to wait for each inbound msgpack frame (default ``120``). - - ``ping_interval`` / ``ping_timeout``: forwarded to :func:`websockets.sync.client.connect`. - """ - cfg = dict(request_config or {}) - operations_cfg = cfg.get("operations") - if operations_cfg is not None: - operations = [dict(op) for op in operations_cfg] - elif cfg.get("run_dreamzero_policy_session"): - video_dir = cfg.get("video_dir") - if video_dir is None: - raise ValueError("run_dreamzero_policy_session requires video_dir") - session_id = str(cfg.get("session_id", "dreamzero-smoke")) - prompt = str(cfg.get("prompt", DREAMZERO_DEFAULT_PROMPT)) - num_chunks = int(cfg.get("num_chunks", 2)) - observations = build_dreamzero_demo_observations( - load_dreamzero_camera_frames(Path(video_dir)), - prompt=prompt, - session_id=session_id, - num_chunks=num_chunks, - ) - operations = [{"endpoint": "infer", "payload": obs} for obs in observations] - operations.append({"endpoint": "reset", "payload": {}}) - operations.append({"endpoint": "infer", "payload": observations[0]}) - elif cfg.get("run_default_policy_session"): - session_id = str(cfg.get("session_id", "gr00t-smoke")) - operations = [ - {"endpoint": "infer", "payload": build_openpi_droid_observation(session_id=session_id)}, - {"endpoint": "reset", "payload": {}}, - ] - else: - operations = [] - timeout = float(cfg.get("timeout", 120.0)) - ping_interval = float(cfg.get("ping_interval", OpenPIWebSocketSession.DEFAULT_PING_INTERVAL_SECS)) - ping_timeout = float(cfg.get("ping_timeout", OpenPIWebSocketSession.DEFAULT_PING_TIMEOUT_SECS)) - - uri = self._build_ws_url("/v1/realtime/robot/openpi") - session = OpenPIWebSocketSession( - uri, - ping_interval=ping_interval, - ping_timeout=ping_timeout, - open_timeout=timeout, - close_timeout=timeout, - ) - try: - operation_responses: list[Any] = [] - actions: dict[str, np.ndarray] | None = None - action_tensors: list[np.ndarray] = [] - for operation in operations: - endpoint = str(operation["endpoint"]) - payload = dict(operation.get("payload") or {}) - response = session._send_operation(endpoint, payload) - operation_responses.append(response) - if endpoint == "infer": - if isinstance(response, dict): - actions = {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()} - else: - action_tensors.append(np.asarray(response, dtype=np.float32)) - finally: - session.close() - - return [ - OpenPIWebSocketResponse( - server_metadata=session.get_server_metadata(), - operation_responses=operation_responses, - actions=actions, - action_tensors=action_tensors or None, - ) - ] - - def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """Chat completions via the OpenAI Python SDK (not raw HTTP).""" - responses: list[OmniResponse] = [] - stream = request_config.get("stream", False) - modalities = request_config.get("modalities", ["text", "audio"]) - extra_body: dict[str, Any] = {} - if "speaker" in request_config: - extra_body["speaker"] = request_config["speaker"] - if request_config.get("use_audio_in_video"): - mm = dict(extra_body.get("mm_processor_kwargs") or {}) - mm["use_audio_in_video"] = True - extra_body["mm_processor_kwargs"] = mm - if "sampling_params_list" in request_config: - extra_body["sampling_params_list"] = request_config["sampling_params_list"] - if request_config.get("extra_body"): - extra_body.update(request_config["extra_body"]) - - create_kwargs: dict[str, Any] = { - "model": request_config.get("model"), - "messages": request_config.get("messages"), - "stream": stream, - "modalities": modalities, - } - if "logprobs" in request_config: - create_kwargs["logprobs"] = request_config["logprobs"] - if "top_logprobs" in request_config: - create_kwargs["top_logprobs"] = request_config["top_logprobs"] - if "stream_options" in request_config: - create_kwargs["stream_options"] = request_config["stream_options"] - if extra_body: - create_kwargs["extra_body"] = extra_body - - if request_num == 1: - wall_start = time.perf_counter() - chat_completion = self.client.chat.completions.create(**create_kwargs) - resp = ( - self._process_stream_omni_response(chat_completion, wall_start=wall_start) - if stream - else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start) - ) - assert_omni_response(resp, request_config, run_level=self.run_level) - if resp.e2e_latency is not None: - self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s") - else: - self._print_client_stat("[omni] request#1 completed") - responses.append(resp) - return responses - - def _one(): - wall_start = time.perf_counter() - chat_completion = self.client.chat.completions.create(**create_kwargs) - return ( - self._process_stream_omni_response(chat_completion, wall_start=wall_start) - if stream - else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start) - ) - - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = {executor.submit(_one): i + 1 for i in range(request_num)} - for future in concurrent.futures.as_completed(futures): - request_idx = futures[future] - resp = future.result() - assert_omni_response(resp, request_config, run_level=self.run_level) - if resp.e2e_latency is not None: - self._print_client_stat(f"[omni] request#{request_idx} success in {resp.e2e_latency:.3f}s") - else: - self._print_client_stat(f"[omni] request#{request_idx} completed") - responses.append(resp) - return responses - - def _process_stream_audio_speech_response( - self, response, *, response_format: str | None = None, wall_start: float - ) -> OmniResponse: - """ - Process streaming /v1/audio/speech responses into an OmniResponse. - - This mirrors _process_stream_omni_response but operates on low-level - audio bytes. Whisper transcription runs in assert_audio_speech_response - when the run_level requires it. - """ - result = OmniResponse() - - try: - # Aggregate all audio bytes from the streaming response. - data = bytearray() - - # Preferred OpenAI helper. - if hasattr(response, "iter_bytes") and callable(getattr(response, "iter_bytes")): - for chunk in response.iter_bytes(): - if chunk: - data.extend(chunk) - else: - # Generic iterable-of-bytes fallback (e.g., generator or list of chunks). - try: - iterator = iter(response) - except TypeError: - iterator = None - - if iterator is not None: - for chunk in iterator: - if not chunk: - continue - if isinstance(chunk, (bytes, bytearray)): - data.extend(chunk) - elif hasattr(chunk, "data"): - data.extend(chunk.data) # type: ignore[arg-type] - elif hasattr(chunk, "content"): - data.extend(chunk.content) # type: ignore[arg-type] - else: - raise TypeError(f"Unsupported stream chunk type: {type(chunk)}") - else: - raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}") - - raw_bytes = bytes(data) - - # Populate OmniResponse. - result.audio_bytes = raw_bytes - result.e2e_latency = time.perf_counter() - wall_start - result.success = True - result.audio_format = getattr(response, "response", None) - if result.audio_format is not None: - result.audio_format = result.audio_format.headers.get("content-type", "") - - except Exception as e: - msg = f"Audio speech stream processing error: {str(e)}" - print(f"Error: {msg}") - - return result - - def _process_non_stream_audio_speech_response( - self, response, *, response_format: str | None = None, wall_start: float - ) -> OmniResponse: - """ - Process non-streaming /v1/audio/speech responses into an OmniResponse. - - This mirrors _process_non_stream_omni_response but for the binary - audio payload returned by audio.speech.create. - """ - result = OmniResponse() - - try: - # OpenAI non-streaming audio.speech.create returns HttpxBinaryResponseContent (.read() or .content) - if hasattr(response, "read") and callable(getattr(response, "read")): - raw_bytes = response.read() - elif hasattr(response, "content"): - raw_bytes = response.content # type: ignore[assignment] - else: - raise TypeError(f"Unsupported audio speech response type: {type(response)}") - - result.audio_bytes = raw_bytes - result.e2e_latency = time.perf_counter() - wall_start - result.success = True - result.audio_format = getattr(response, "response", None) - if result.audio_format is not None: - result.audio_format = result.audio_format.headers.get("content-type", "") - - except Exception as e: - msg = f"Audio speech non-stream processing error: {str(e)}" - print(f"Error: {msg}") - - return result - - def send_audio_speech_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: - """ - Call the /v1/audio/speech endpoint using the same configuration-dict - style as send_omni_request, but via the OpenAI Python client's - audio.speech APIs. - - Expected keys in request_config: - - model: model name/path (required) - - input: text to synthesize (required) - - response_format: audio format such as "wav" or "pcm" (optional) - - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body) - - min_audio_bytes: optional minimum ``len(audio_bytes)`` checked in ``assert_audio_speech_response`` - - status_code: if set, HTTP status is asserted (int or e.g. ``(400, 422)``); uses APIError handling - - err_message: optional substring(s) to match against error text (``str`` or list/tuple of alternatives; - see ``assert_audio_speech_response``). If set, uses the same APIError path as ``status_code``. - - When both ``status_code`` and ``err_message`` are absent (or each is ``None``), the normal request path - is used (no try/except around ``APIError``). - - timeout: request timeout in seconds (float, optional, default 120.0) - - stream: whether to use streaming API (bool, optional, default False) - """ - timeout = float(request_config.get("timeout", 120.0)) - - model = request_config["model"] - text_input = request_config["input"] - stream = bool(request_config.get("stream", False)) - voice = request_config.get("voice", None) - - # Standard OpenAI param: use omit when not provided to keep default behavior. - response_format = request_config.get("response_format", omit) - - # Qwen3-TTS custom fields, forwarded via extra_body. - extra_body: dict[str, Any] = {} - # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params. - for key in ( - "task_type", - "ref_text", - "ref_audio", - "language", - "max_new_tokens", - "seed", - "instructions", - "speed", - "stream_format", - "x_vector_only_mode", - ): - if key in request_config: - extra_body[key] = request_config[key] - - responses: list[OmniResponse] = [] - - speech_fmt: str | None = None if response_format is omit else str(response_format).lower() - - print(f"[audio.speech] start model={model}, stream={stream}, request_num={request_num}, timeout={timeout:.1f}s") - - # Error validation path: only when at least one of these is set to a non-``None`` value. - expect_error_handling = (request_config.get("status_code") is not None) or ( - request_config.get("err_message") is not None - ) - if expect_error_handling and request_num != 1: - raise ValueError( - "request_config error validation (status_code / err_message) is only supported when request_num=1" - ) - - if request_num == 1: - if expect_error_handling: - # ``status`` and/or ``err_message`` requested: catch APIError (4xx) and (optionally) assert body text; - # HTTP 200 with JSON error body is handled in ``assert_audio_speech_response``. - try: - if stream: - wall_start = time.perf_counter() - with self.client.audio.speech.with_streaming_response.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) as resp: - omni_resp = self._process_stream_audio_speech_response( - resp, response_format=speech_fmt, wall_start=wall_start - ) - else: - wall_start = time.perf_counter() - resp = self.client.audio.speech.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) - omni_resp = self._process_non_stream_audio_speech_response( - resp, response_format=speech_fmt, wall_start=wall_start - ) - except APIError as e: - sc = getattr(e, "status_code", None) - if sc is None: - raise - omni_resp = OmniResponse( - success=False, - status_code=sc, - error_message=str(e), - ) - else: - if getattr(omni_resp, "status_code", None) is None: - omni_resp.status_code = 200 - elif stream: - # Use streaming response helper. - wall_start = time.perf_counter() - with self.client.audio.speech.with_streaming_response.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) as resp: - omni_resp = self._process_stream_audio_speech_response( - resp, response_format=speech_fmt, wall_start=wall_start - ) - else: - # Non-streaming response. - wall_start = time.perf_counter() - resp = self.client.audio.speech.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) - omni_resp = self._process_non_stream_audio_speech_response( - resp, response_format=speech_fmt, wall_start=wall_start - ) - - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - if omni_resp.e2e_latency is not None: - self._print_client_stat(f"[audio.speech] request#1 success in {omni_resp.e2e_latency:.3f}s") - else: - self._print_client_stat("[audio.speech] request#1 completed") - responses.append(omni_resp) - return responses - else: - # request_num > 1: concurrent requests (use same params as single-request path) - - if stream: - - def _stream_task(request_idx: int): - wall_start = time.perf_counter() - with self.client.audio.speech.with_streaming_response.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) as resp: - result = self._process_stream_audio_speech_response( - resp, response_format=speech_fmt, wall_start=wall_start - ) - if result.e2e_latency is not None: - self._print_client_stat( - f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s" - ) - else: - self._print_client_stat(f"[audio.speech] request#{request_idx} completed") - return result - - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = {executor.submit(_stream_task, i + 1): i + 1 for i in range(request_num)} - for future in concurrent.futures.as_completed(futures): - request_idx = futures[future] - try: - omni_resp = future.result() - except Exception as e: - print( - f"[audio.speech] request#{request_idx} failed " - f"(stream={stream}, timeout={timeout:.1f}s): {e!r}" - ) - raise - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - else: - - def _non_stream_task(request_idx: int): - wall_start = time.perf_counter() - r = self.client.audio.speech.create( - model=model, - input=text_input, - response_format=response_format, - extra_body=extra_body or None, - timeout=timeout, - voice=voice, - ) - result = self._process_non_stream_audio_speech_response( - r, response_format=speech_fmt, wall_start=wall_start - ) - if result.e2e_latency is not None: - self._print_client_stat( - f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s" - ) - else: - self._print_client_stat(f"[audio.speech] request#{request_idx} completed") - return result - - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = {executor.submit(_non_stream_task, i + 1): i + 1 for i in range(request_num)} - for future in concurrent.futures.as_completed(futures): - request_idx = futures[future] - try: - omni_resp = future.result() - except Exception as e: - print( - f"[audio.speech] request#{request_idx} failed " - f"(stream={stream}, timeout={timeout:.1f}s): {e!r}" - ) - raise - assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level) - responses.append(omni_resp) - - return responses - - def send_diffusion_request( - self, request_config: dict[str, Any] | list[dict[str, Any]], request_num: int = 1 - ) -> list[DiffusionResponse]: - """ - Send OpenAI requests for diffusion models. - If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel - (scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order. - - Args: - request_config: A single request configuration dict, or a list of - request configuration dicts (one request per element) - request_num: Number of requests to send concurrently, defaults to 1 (single request) - Returns: - list[DiffusionResponse]: List of DiffusionResponse objects containing the response data - """ - responses: list[DiffusionResponse] = [] - - def _create_from_config(cfg: dict[str, Any]) -> tuple[Any, float]: - stream = cfg.get("stream", False) - if stream: - raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test") - modalities = cfg.get("modalities", omit) # Most diffusion models don't require modalities param - eb = cfg.get("extra_body") - extra = copy.deepcopy(eb) if eb else None - wall_start = time.perf_counter() - chat_completion = self.client.chat.completions.create( - model=cfg.get("model"), - messages=cfg.get("messages"), - extra_body=extra, - modalities=modalities, - ) - return chat_completion, wall_start - - if isinstance(request_config, list): - if not request_config: - raise ValueError("request_config list must not be empty") - if request_num != 1: - raise ValueError("request_num is not supported when request_config is a list") - with concurrent.futures.ThreadPoolExecutor(max_workers=len(request_config)) as executor: - futures = { - executor.submit(_create_from_config, cfg): (i + 1, cfg) for i, cfg in enumerate(request_config) - } - for future in concurrent.futures.as_completed(futures): - request_idx, cfg = futures[future] - chat_completion, wall_start = future.result() - response = self._process_diffusion_response(chat_completion, wall_start=wall_start) - assert_diffusion_response(response, cfg, run_level=self.run_level) - if response.e2e_latency is not None: - self._print_client_stat( - f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s" - ) - else: - self._print_client_stat(f"[diffusion] request#{request_idx} completed") - responses.append(response) - return responses - - size_splits = _split_request_config_by_per_output_sizes(request_config) - if size_splits is not None: - if request_num != 1: - raise ValueError( - "request_num must be 1 when extra_body height/width are lists (split into concurrent per-size calls)" - ) - t0 = time.perf_counter() - with concurrent.futures.ThreadPoolExecutor(max_workers=len(size_splits)) as executor: - futures = [executor.submit(_create_from_config, sub) for sub in size_splits] - chat_completions = [f.result() for f in futures] - parts = [self._process_diffusion_response(cc, wall_start=ws) for cc, ws in chat_completions] - merged = _merge_diffusion_responses(parts) - merged.e2e_latency = time.perf_counter() - t0 - assert_diffusion_response(merged, request_config, run_level=self.run_level) - if merged.e2e_latency is not None: - self._print_client_stat(f"[diffusion] request#1 success in {merged.e2e_latency:.3f}s") - else: - self._print_client_stat("[diffusion] request#1 completed") - return [merged] - - if request_num == 1: - # Send single request - chat_completion, wall_start = _create_from_config(request_config) - response = self._process_diffusion_response(chat_completion, wall_start=wall_start) - assert_diffusion_response(response, request_config, run_level=self.run_level) - if response.e2e_latency is not None: - self._print_client_stat(f"[diffusion] request#1 success in {response.e2e_latency:.3f}s") - else: - self._print_client_stat("[diffusion] request#1 completed") - responses.append(response) - return responses - - # Send concurrent requests for the same request_config - with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: - futures = {executor.submit(_create_from_config, request_config): i + 1 for i in range(request_num)} - for future in concurrent.futures.as_completed(futures): - request_idx = futures[future] - chat_completion, wall_start = future.result() - response = self._process_diffusion_response(chat_completion, wall_start=wall_start) - assert_diffusion_response(response, request_config, run_level=self.run_level) - if response.e2e_latency is not None: - self._print_client_stat(f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s") - else: - self._print_client_stat(f"[diffusion] request#{request_idx} completed") - responses.append(response) - return responses - - def send_video_diffusion_request( - self, request_config: dict[str, Any], request_num: int = 1 - ) -> list[DiffusionResponse]: - """ - Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content. - - For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc. - """ - if request_num != 1: - raise NotImplementedError("Concurrent video diffusion requests are not currently implemented") - - form_data = request_config.get("form_data") - if not isinstance(form_data, dict): - raise ValueError("Video request_config must contain 'form_data'") - normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None} - files: dict[str, tuple[str, BytesIO, str]] = {} - image_reference = request_config.get("image_reference") - video_reference = request_config.get("video_reference") - if image_reference and video_reference: - raise ValueError("Only one of image_reference or video_reference can be provided") - if image_reference: - if image_reference.startswith("data:image"): - header, encoded = image_reference.split(",", 1) - content_type = header.split(";")[0].removeprefix("data:") - extension = content_type.split("/")[-1] - file_data = base64.b64decode(encoded) - files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type) - else: - normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference}) - if video_reference: - if video_reference.startswith("data:video"): - header, encoded = video_reference.split(",", 1) - content_type = header.split(";")[0].removeprefix("data:") - extension = content_type.split("/")[-1] - file_data = base64.b64decode(encoded) - files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type) - else: - normalized_form_data["video_reference"] = json.dumps({"video_url": video_reference}) - - result = DiffusionResponse() - create_url = self._build_url("/v1/videos") - response = requests.post( - create_url, - data=normalized_form_data, - files=files, - headers={"Accept": "application/json"}, - timeout=60, - ) - start_time = time.perf_counter() - response.raise_for_status() - job_data = response.json() - video_id = job_data["id"] - self._wait_until_video_completed(video_id) - end_time = time.perf_counter() - video_content = self._download_video_content(video_id) - result.success = True - result.videos = [video_content] - result.e2e_latency = end_time - start_time - assert_diffusion_response(result, request_config, run_level=self.run_level) - if result.e2e_latency is not None: - self._print_client_stat(f"[diffusion] request#1 success in {result.e2e_latency:.3f}s") - else: - self._print_client_stat("[diffusion] request#1 completed") - return [result] - - def _post_json_endpoint( - self, - path: str, - request_config: dict[str, Any], - *, - default_timeout: float, - ) -> requests.Response: - url = self._build_url(path) - timeout = float(request_config.get("timeout", default_timeout)) - if "raw_body" in request_config: - raw = request_config["raw_body"] - payload = raw.encode("utf-8") if isinstance(raw, str) else raw - return requests.post( - url, - data=payload, - headers={"Content-Type": "application/json", "Accept": "application/json"}, - timeout=timeout, - ) - if "json" not in request_config: - raise ValueError(f"{path} request_config must include 'json' or 'raw_body'") - return requests.post( - url, - json=request_config["json"], - headers={"Accept": "application/json", "Content-Type": "application/json"}, - timeout=timeout, - ) - - def _post_form_endpoint( - self, - path: str, - request_config: dict[str, Any], - *, - default_timeout: float = 120.0, - ) -> requests.Response: - url = self._build_url(path) - timeout = float(request_config.get("timeout", default_timeout)) - data = request_config.get("data") - files = request_config.get("files") - if data is None and not files: - data = {} - return requests.post( - url, - data=data, - files=files, - headers={"Accept": "application/json"} if not files else {"Accept": "application/json"}, - timeout=timeout, - ) - - def send_streaming_video_diffusion_request( - self, - request_config: dict[str, Any], - request_num: int = 1, - *, - timeout_seconds: float = 600.0, - ) -> list[DiffusionResponse]: - """ - Send a native ``/v1/realtime/video`` WebSocket request and return one - finalized MP4 artifact assembled from the streamed binary fragments. - """ - if request_num != 1: - raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented") - - response = asyncio.run( - self._send_streaming_video_diffusion_request_once( - request_config, - timeout_seconds=timeout_seconds, - ) - ) - assert_diffusion_response(response, request_config, run_level=self.run_level) - if response.e2e_latency is not None: - self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s") - else: - self._print_client_stat("[diffusion.stream] request#1 completed") - return [response] - - async def _send_streaming_video_diffusion_request_once( - self, - request_config: dict[str, Any], - *, - timeout_seconds: float, - ) -> DiffusionResponse: - form_data = request_config.get("form_data") - if not isinstance(form_data, dict): - raise ValueError("Video request_config must contain 'form_data'") - payload: dict[str, Any] = { - "type": "session.start", - **{key: value for key, value in form_data.items() if value is not None}, - } - model = request_config.get("model") - if model is not None: - payload["model"] = model - payload.setdefault("format", "m4s") - - fps = float(payload.get("fps") or 16) - stream_format = payload["format"] - url = self._build_ws_url("/v1/realtime/video") - - result = DiffusionResponse() - chunks: list[bytes] = [] - start_time = time.perf_counter() - deadline = start_time + timeout_seconds - - import websockets - - async with websockets.connect(url, max_size=None) as websocket: - await websocket.send(json.dumps(payload)) - - while True: - remaining = deadline - time.perf_counter() - if remaining <= 0: - raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s") - - message = await asyncio.wait_for(websocket.recv(), timeout=remaining) - if isinstance(message, bytes): - chunks.append(message) - continue - - msg = json.loads(message) - msg_type = msg.get("type") - if msg_type == "video.start": - stream_format = msg.get("format") or stream_format - continue - if msg_type == "session.done": - break - if msg_type == "error": - raise RuntimeError(str(msg.get("message", msg))) - - from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes - from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat - - streamed_bytes = b"".join(chunks) - if not streamed_bytes: - raise RuntimeError("Streaming video request completed without binary video chunks") - result.videos = [ - finalize_streaming_video_bytes( - streamed_bytes, - input_format=cast(StreamingVideoFormat, stream_format), - fps=fps, - ) - ] - result.e2e_latency = time.perf_counter() - start_time - result.success = True - return result - - def _wait_until_video_completed( - self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300 - ) -> None: - status_url = self._build_url(f"/v1/videos/{video_id}") - deadline = time.monotonic() + timeout_seconds - while time.monotonic() < deadline: - status_resp = requests.get(status_url, headers={"Accept": "application/json"}, timeout=30) - status_resp.raise_for_status() - status_data = status_resp.json() - current_status = status_data["status"] - if current_status == "completed": - return - if current_status == "failed": - error_msg = status_data.get("last_error", "Unknown error") - raise RuntimeError(f"Job failed: {error_msg}") - time.sleep(poll_interval_seconds) - raise TimeoutError(f"Video job {video_id} did not complete within {timeout_seconds}s") - - def _download_video_content(self, video_id: str) -> bytes: - download_url = self._build_url(f"/v1/videos/{video_id}/content") - video_resp = requests.get(download_url, stream=True, timeout=60) - video_resp.raise_for_status() - video_bytes = BytesIO() - for chunk in video_resp.iter_content(chunk_size=8192): - if chunk: - video_bytes.write(chunk) - return video_bytes.getvalue() - - def _build_url(self, path: str) -> str: - return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}" + Path(log_path).unlink(missing_ok=True) + except Exception: + pass + self._stage_log_files = {} + if not keep_logs: + self._stage_log_paths = {} + + cleanup_test_environment() class OmniRunner: @@ -2625,7 +592,6 @@ class OmniRunner: model_name: str, seed: int = 42, stage_init_timeout: int = 600, - batch_timeout: int = 10, # Bumped from 900s -> 1800s to give CI cold-cache loads of large # diffusion models enough headroom (Buildkite #8418 hit a 6-second # overrun loading Tongyi-MAI/Z-Image-Turbo: weights alone took 690s, @@ -2635,30 +601,34 @@ class OmniRunner: # affects the test runner wrapper. init_timeout: int = 1800, log_stats: bool = False, - stage_configs_path: str | None = None, + deploy_config: str | None = None, **kwargs, ) -> None: startup_t0 = time.perf_counter() - cleanup_dist_env_and_memory() - run_pre_test_cleanup() - run_post_test_cleanup() + cleanup_test_environment() self.model_name = model_name self.seed = seed self._prompt_len_estimate_cache: dict[str, Any] = {} - from vllm_omni.entrypoints.omni import Omni + self.omni: Any = None + try: + from vllm_omni.entrypoints.omni import Omni - self.omni = Omni( - model=model_name, - log_stats=log_stats, - stage_init_timeout=stage_init_timeout, - batch_timeout=batch_timeout, - init_timeout=init_timeout, - stage_configs_path=stage_configs_path, - **kwargs, - ) - startup_s = time.perf_counter() - startup_t0 - if log_stats: - print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True) + self.omni = Omni( + model=model_name, + log_stats=log_stats, + stage_init_timeout=stage_init_timeout, + init_timeout=init_timeout, + deploy_config=deploy_config, + **kwargs, + ) + startup_s = time.perf_counter() - startup_t0 + if log_stats: + print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True) + except BaseException: + # ``with OmniRunner(...)`` never reaches ``__enter__``/``__exit__`` + # when construction fails after worker processes have started. + self.__exit__(None, None, None) + raise def get_default_sampling_params_list(self) -> list[Any]: if not hasattr(self.omni, "default_sampling_params_list"): @@ -2833,14 +803,14 @@ class OmniRunner: f"<|im_start|>user\n{user_content}<|im_end|>\n" f"<|im_start|>assistant\n" ) - input_dict: dict[str, Any] = {"prompt": full_prompt} + prompt_dict: dict[str, Any] = {"prompt": full_prompt} if multi_modal_data: - input_dict["multi_modal_data"] = multi_modal_data + prompt_dict["multi_modal_data"] = multi_modal_data if modalities: - input_dict["modalities"] = modalities + prompt_dict["modalities"] = modalities if mm_processor_kwargs: - input_dict["mm_processor_kwargs"] = mm_processor_kwargs - omni_inputs.append(input_dict) + prompt_dict["mm_processor_kwargs"] = mm_processor_kwargs + omni_inputs.append(prompt_dict) return omni_inputs def generate( @@ -2919,197 +889,11 @@ class OmniRunner: return self def __exit__(self, exc_type, exc_val, exc_tb): - if hasattr(self.omni, "close"): - self.omni.close() + omni = getattr(self, "omni", None) + if omni is not None and hasattr(omni, "close"): + omni.close() self._cleanup_process() - run_pre_test_cleanup() - run_post_test_cleanup() - cleanup_dist_env_and_memory() - - -class OmniRunnerHandler: - def __init__(self, omni_runner: OmniRunner): - self.runner = omni_runner - - def _process_omni_output(self, outputs: list[OmniRequestOutput]) -> OmniResponse: - result = OmniResponse() - try: - text_content = None - audio_content = None - for stage_output in outputs: - if getattr(stage_output, "final_output_type", None) == "text": - text_content = stage_output.outputs[0].text - if getattr(stage_output, "final_output_type", None) == "audio": - audio_content = stage_output.outputs[0].multimodal_output["audio"] - result.audio_content = audio_content - result.text_content = text_content - result.success = True - except Exception as e: - msg = f"Output processing error: {str(e)}" - result.success = False - print(f"Error: {msg}") - return result - - def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse: - result = DiffusionResponse() - output = outputs[0] - if isinstance(output.images[0], list): - # Returning frames of images as a video - result.videos = output.images - else: - # Returning actual images - result.images = output.images - # [TODO] Add audio processing when tests are introduced - result.success = True - return result - - def send_omni_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse: - if request_config is None: - request_config = {} - prompts = request_config.get("prompts") - videos = request_config.get("videos") - images = request_config.get("images") - audios = request_config.get("audios") - modalities = request_config.get("modalities", ["text", "audio"]) - outputs = self.runner.generate_multimodal( - prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities - ) - response = self._process_omni_output(outputs) - assert_omni_response(response, request_config, run_level="core_model") - return response - - def send_diffusion_request(self, request_config: dict[str, Any]) -> DiffusionResponse: - prompt = request_config.get("prompt") - if prompt is None: - prompts = request_config.get("prompts") - if not prompts: - raise ValueError("request_config must contain a prompt") - if len(prompts) > 1: - raise ValueError( - "In the current internal data structure, " - "only one prompt is supported for diffusion requests. " - "Because one prompt can contain multiple images or videos, " - "the current internal data structure is ambiguous when multiple prompts are provided." - ) - prompt = prompts[0] - - # Sync previous `extra_body` field with sampling params object - sampling_params: OmniDiffusionSamplingParams | None = request_config.get("sampling_params") - if sampling_params is None: - extra_body = request_config.get("extra_body", {}) - if extra_body: - sampling_params = OmniDiffusionSamplingParams(**extra_body) - else: - extra_body = asdict(sampling_params) - request_config["extra_body"] = extra_body - if not extra_body: - logger.warning("No sampling params provided in request_config, will skip output assertion") - - negative_prompt = extra_body.get("negative_prompt") or request_config.get("negative_prompt") - videos = request_config.get("videos") - images = request_config.get("images") - audios = request_config.get("audios") - # Full dict (e.g. image + mask_image for inpainting) or partial; merged with top-level image/video keys. - extra_multi_modal = request_config.get("multi_modal_data") - modalities = request_config.get("modalities") # only used by limited models. Do not add default value here - - prompt_object = OmniTextPrompt(prompt=prompt) - if negative_prompt: - prompt_object["negative_prompt"] = negative_prompt - multi_modal: dict = {} - if extra_multi_modal is not None: - multi_modal.update(dict(extra_multi_modal)) - if videos is not None: - multi_modal["video"] = videos - if images is not None: - multi_modal["image"] = images - if audios is not None: - multi_modal["audio"] = audios - if multi_modal: - prompt_object["multi_modal_data"] = multi_modal - if modalities: - prompt_object["modalities"] = modalities # pyright: ignore[reportGeneralTypeIssues] - - start_time = time.perf_counter() - response = self.runner.generate([prompt_object], [sampling_params] if sampling_params else None) - end_time = time.perf_counter() - - response = self._process_diffusion_output(response) - response.e2e_latency = end_time - start_time - assert_diffusion_response(response, request_config, run_level="core_model") - return response - - def send_audio_speech_request(self, request_config: dict[str, Any]) -> OmniResponse: - """ - Offline TTS: text -> audio via generate_multimodal, then validate with assert_audio_speech_response. - - request_config must contain: - - 'input' or 'prompts': text to synthesize. - Optional keys: - - 'voice' -> speaker (CustomVoice) - - 'task_type' -> task_type in additional_information (default: "CustomVoice") - - 'language' -> language in additional_information (default: "Auto") - - 'max_new_tokens' -> max_new_tokens in additional_information (default: 2048) - - 'response_format' -> desired audio format (used only for assertion) - """ - input_text = request_config.get("input") or request_config.get("prompts") - if input_text is None: - raise ValueError("request_config must contain 'input' or 'prompts' for TTS") - if isinstance(input_text, list): - input_text = input_text[0] if input_text else "" - - mm_processor_kwargs: dict[str, Any] = {} - if "voice" in request_config: - mm_processor_kwargs["speaker"] = request_config["voice"] - if "task_type" in request_config: - mm_processor_kwargs["task_type"] = request_config["task_type"] - if "ref_audio" in request_config: - mm_processor_kwargs["ref_audio"] = request_config["ref_audio"] - if "ref_text" in request_config: - mm_processor_kwargs["ref_text"] = request_config["ref_text"] - if "language" in request_config: - mm_processor_kwargs["language"] = request_config["language"] - if "max_new_tokens" in request_config: - mm_processor_kwargs["max_new_tokens"] = request_config["max_new_tokens"] - - outputs = self.runner.generate_multimodal( - prompts=input_text, - modalities=["audio"], - mm_processor_kwargs=mm_processor_kwargs or None, - ) - mm_out: dict[str, Any] | None = None - for stage_out in outputs: - if getattr(stage_out, "final_output_type", None) == "audio": - mm_out = stage_out.outputs[0].multimodal_output - break - if mm_out is None: - raise AssertionError("No audio output from pipeline") - - audio_data = mm_out.get("audio") - if audio_data is None: - raise AssertionError("No audio tensor in multimodal output") - - sr_raw = mm_out.get("sr") - sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw - sr = int(sr_val.item() if hasattr(sr_val, "item") else sr_val) - wav_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data - wav_buf = io.BytesIO() - sf.write( - wav_buf, - wav_tensor.float().cpu().numpy().reshape(-1), - samplerate=sr, - format="WAV", - subtype="PCM_16", - ) - result = OmniResponse(success=True, audio_bytes=wav_buf.getvalue(), audio_format="audio/wav") - assert_audio_speech_response(result, request_config, run_level="core_model") - return result - - def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]: - return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages) - - def stop_profile(self, stages: list[int] | None = None) -> list[Any]: - return self.runner.stop_profile(stages=stages) + cleanup_test_environment() # --------------------------------------------------------------------------- @@ -3118,16 +902,37 @@ class OmniRunnerHandler: # --------------------------------------------------------------------------- +@contextmanager +def _whisper_device_free_around(): + """Keep the Whisper judge off the accelerator around a server/runner lifecycle. + + Released on entry so the next instance -- e.g. the next parametrization in + the same module -- initializes on a clean device, and again on exit + (including when a test raises) so it does not linger for the instance after. + """ + release_audio_transcriber() + try: + yield + finally: + release_audio_transcriber() + + +def get_model_prefix() -> str: + """Return ``MODEL_PREFIX`` with a trailing slash, or ``""`` when unset.""" + prefix = os.environ.get("MODEL_PREFIX", "") + return f"{prefix.rstrip('/')}/" if prefix else "" + + def iter_omni_server( request: Any, run_level: str, - model_prefix: str, omni_fixture_lock: threading.Lock, ) -> Generator[Any, Any, None]: """Start/stop an Omni HTTP server; used by ``omni_server`` / ``omni_server_function`` fixtures.""" from tests.helpers.stage_config import stage_config_path_for_run_level - with omni_fixture_lock: + model_prefix = get_model_prefix() + with omni_fixture_lock, _whisper_device_free_around(): params: OmniServerParams = request.param # For now, when a tiny model is substituted, we preserve the original model # name via --served-model-name (so that the server still accepts requests with @@ -3179,12 +984,7 @@ def iter_omni_server( print("OmniServer stopping...") else: if stage_config_path is not None: - from vllm_omni.entrypoints.utils import is_new_format_deploy_config - - if is_new_format_deploy_config(stage_config_path): - server_args += ["--deploy-config", stage_config_path] - else: - server_args += ["--stage-configs-path", stage_config_path] + server_args += ["--deploy-config", stage_config_path] with ( OmniServer( @@ -3213,14 +1013,14 @@ def iter_omni_server( def iter_omni_runner( request: Any, - model_prefix: str, run_level: str, omni_fixture_lock: threading.Lock, ) -> Generator[Any, None, None]: """Yield an :class:`OmniRunner`; used by ``omni_runner`` / ``omni_runner_function`` fixtures.""" from tests.helpers.stage_config import stage_config_path_for_run_level - with omni_fixture_lock: + model_prefix = get_model_prefix() + with omni_fixture_lock, _whisper_device_free_around(): param = request.param if not isinstance(param, (tuple, list)) or len(param) not in (2, 3): raise ValueError( @@ -3237,7 +1037,7 @@ def iter_omni_runner( model = model_prefix + model if run_level == "core_model" and request.node.get_closest_marker("diffusion"): model = resolve_tiny_model_path(model) - with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner: + with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner: print("OmniRunner started successfully") yield runner print("OmniRunner stopping...") @@ -3245,6 +1045,166 @@ def iter_omni_runner( print("OmniRunner stopped") +# ───────────────────────────────────────────────────────────────────── +# π0 (Pi-Zero) OpenPI websocket policy client +# +# Minimal client for the OpenPI realtime robot protocol +# (``/v1/realtime/robot/openpi``): connect → read handshake metadata → send +# observation → receive an action chunk. Used by the π0 online-serving e2e +# (tests/e2e/online_serving/test_pi0_expansion.py). π0's observation is 3 RGB +# cameras + proprioceptive state + a language prompt, built as synthetic numpy +# frames here (no video assets). +# ───────────────────────────────────────────────────────────────────── +PI0_OPENPI_PATH = "/v1/realtime/robot/openpi" +PI0_OPENPI_DEFAULT_PROMPT = "pick up the red block and place it in the bin" +# π0 / pi0_base camera identities (must match the server's image_feature_keys, +# i.e. vllm_omni/deploy/pi0.yaml, whose image_key_map is empty). +PI0_CAMERA_KEYS = ( + "observation.images.base_0_rgb", + "observation.images.left_wrist_0_rgb", + "observation.images.right_wrist_0_rgb", +) +PI0_ACTION_HORIZON = 50 +PI0_ACTION_DIM = 32 +PI0_STATE_DIM = 32 +PI0_IMAGE_SIZE = 224 + + +def pi0_openpi_require_dependencies() -> None: + """Raise ModuleNotFoundError if the OpenPI websocket client deps are missing.""" + missing = [] + try: + import websockets.sync.client # noqa: F401 + except ImportError: + missing.append("websockets") + try: + from openpi_client import msgpack_numpy # noqa: F401 + except ImportError: + missing.append("openpi-client") + if missing: + raise ModuleNotFoundError(f"π0 OpenPI test dependencies are missing: {', '.join(missing)}") + + +def _pi0_decode_action_response(response: bytes | str) -> np.ndarray: + from openpi_client import msgpack_numpy + + if isinstance(response, str): + raise RuntimeError(f"Inference failed: {response}") + decoded = msgpack_numpy.unpackb(response) + if isinstance(decoded, dict) and decoded.get("type") == "error": + raise RuntimeError(f"Inference failed: {decoded.get('message', decoded)}") + return np.asarray(decoded, dtype=np.float32) + + +def pi0_make_dummy_obs(*, prompt: str, session_id: str, image_size: int = PI0_IMAGE_SIZE) -> dict[str, Any]: + """A single π0 observation: 3 blank cameras (HWC uint8) + zero state + prompt.""" + obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in PI0_CAMERA_KEYS} + obs["state"] = np.zeros(PI0_STATE_DIM, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + return obs + + +def pi0_openpi_run_policy_session( + *, + host: str = "127.0.0.1", + port: int = 8000, + path: str = PI0_OPENPI_PATH, + prompt: str = PI0_OPENPI_DEFAULT_PROMPT, + session_id: str | None = None, + num_steps: int = 2, +) -> dict[str, Any]: + """Connect, read handshake metadata, send ``num_steps`` observations.""" + import uuid + + import websockets.sync.client as websockets_client + from openpi_client import msgpack_numpy + + session_id = session_id or str(uuid.uuid4()) + uri = f"ws://{host}:{port}{path}" + packer = msgpack_numpy.Packer() + conn = websockets_client.connect(uri, compression=None, max_size=None, ping_interval=300, ping_timeout=3600) + try: + metadata = msgpack_numpy.unpackb(conn.recv()) + if not isinstance(metadata, dict): + raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}") + actions = [] + for _ in range(num_steps): + payload = pi0_make_dummy_obs(prompt=prompt, session_id=session_id) + payload["endpoint"] = "infer" + conn.send(packer.pack(payload)) + actions.append(_pi0_decode_action_response(conn.recv())) + return {"metadata": dict(metadata), "actions": actions, "session_id": session_id} + finally: + conn.close() + + +def pi0_openpi_validate_session_result( + result: dict[str, Any], + *, + expected_action_horizon: int = PI0_ACTION_HORIZON, + expected_action_dim: int = PI0_ACTION_DIM, +) -> None: + """Assert the handshake metadata + every returned action chunk for π0.""" + metadata = result["metadata"] + required_keys = ( + "image_resolution", + "needs_wrist_camera", + "needs_session_id", + "action_space", + "action_horizon", + "action_dim", + ) + missing = [key for key in required_keys if key not in metadata] + if missing: + raise AssertionError(f"Missing π0 metadata keys: {missing}") + + if tuple(metadata["image_resolution"]) != (PI0_IMAGE_SIZE, PI0_IMAGE_SIZE): + raise AssertionError(f"Unexpected image_resolution: {metadata['image_resolution']!r}") + if not metadata["needs_wrist_camera"]: + raise AssertionError("π0 test expects needs_wrist_camera=True") + if metadata["needs_session_id"]: + raise AssertionError("π0 is stateless; needs_session_id must be False") + if metadata["action_space"] != "joint_position": + raise AssertionError(f"Unexpected action_space: {metadata['action_space']!r}") + if int(metadata["action_horizon"]) != expected_action_horizon: + raise AssertionError(f"Unexpected action_horizon: {metadata['action_horizon']!r}") + if int(metadata["action_dim"]) != expected_action_dim: + raise AssertionError(f"Unexpected action_dim: {metadata['action_dim']!r}") + + actions = result["actions"] + if not actions: + raise AssertionError("No actions returned from the server") + for index, action in enumerate(actions): + if action.shape != (expected_action_horizon, expected_action_dim): + raise AssertionError( + f"Action {index} shape mismatch: expected " + f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}" + ) + if not np.isfinite(action).all(): + raise AssertionError(f"Action {index} contains non-finite values") + + +# Re-export client types for backward-compatible imports. +from tests.helpers.client import ( # noqa: E402 + DREAMZERO_ACTION_DIM, + DREAMZERO_ACTION_HORIZON, + DREAMZERO_CAMERA_FILES, + DiffusionResponse, + HttpResponse, + OfflineOmniClient, + OmniResponse, + OmniRunnerHandler, + OnlineOmniClient, + OpenAIClientHandler, + OpenPIWebSocketResponse, + OpenPIWebSocketSession, + WebSocketJsonResponse, + build_dreamzero_demo_observations, + build_openpi_droid_observation, + load_dreamzero_camera_frames, +) + __all__ = [ "DiffusionResponse", "HttpResponse", @@ -3259,11 +1219,18 @@ __all__ = [ "OpenPIWebSocketSession", "OmniResponse", "OmniRunner", - "OmniRunnerHandler", + "OfflineOmniClient", "OmniServer", "OmniServerParams", "OmniServerStageCli", + "OnlineOmniClient", "OpenAIClientHandler", + "OmniRunnerHandler", + "get_model_prefix", "get_open_port", "dummy_messages_from_mix_data", + "pi0_openpi_require_dependencies", + "pi0_openpi_run_policy_session", + "pi0_openpi_validate_session_result", + "pi0_make_dummy_obs", ] diff --git a/tests/helpers/serving_chat.py b/tests/helpers/serving_chat.py new file mode 100644 index 000000000..f1f2710b5 --- /dev/null +++ b/tests/helpers/serving_chat.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared scaffolding for OmniOpenAIServingChat unit tests. + +Provides factories for a minimal serving-chat instance, chat requests, and +OmniRequestOutput objects, plus SSE stream collection/parsing utilities. +""" + +import json +from unittest.mock import MagicMock + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponseStreamChoice, +) +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.outputs import CompletionOutput, RequestOutput + +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.outputs import OmniRequestOutput + + +def make_text_omni_output( + request_id: str = "test-req", + text: str = "hello", + token_ids: list[int] | None = None, + finish_reason: str | None = None, + index: int = 0, + num_prompt_tokens: int = 3, + num_cached_tokens: int | None = None, +) -> OmniRequestOutput: + """Build an OmniRequestOutput wrapping a text RequestOutput.""" + if token_ids is None: + token_ids = [10, 11, 12] + res = RequestOutput( + request_id=request_id, + prompt="test", + prompt_token_ids=list(range(num_prompt_tokens)), + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=index, + text=text, + token_ids=token_ids, + cumulative_logprob=0.0, + logprobs=None, + finish_reason=finish_reason, + stop_reason=None, + ) + ], + finished=finish_reason is not None, + num_cached_tokens=num_cached_tokens, + ) + return OmniRequestOutput.from_stage_output( + res, + request_id=request_id, + final_output_type="text", + finished=finish_reason is not None, + ) + + +def mock_audio_choices(index: int = 0, role: str = "assistant"): + return [ + ChatCompletionResponseStreamChoice( + index=index, + delta=DeltaMessage(role=role, content="dGVzdA=="), + logprobs=None, + finish_reason="stop", + ) + ] + + +def build_serving_chat( + *, + engine_client=None, + models=None, + online_renderer=None, + response_role: str = "assistant", + request_logger=None, + chat_template=None, + chat_template_content_format="auto", + **init_overrides, +) -> OmniOpenAIServingChat: + """Create a real OmniOpenAIServingChat via its actual __init__. + + engine_client/models/online_renderer default to minimal mocks; pass your + own to wire in test-specific behavior (e.g. engine_client.stage_configs). + Any other OmniOpenAIServingChat.__init__ keyword — enable_auto_tools, + enable_prompt_tokens_details, trust_request_chat_template, etc. — can be + passed through **init_overrides instead of set on the instance after + construction. + """ + if engine_client is None: + engine_client = MagicMock() + engine_client.errored = False + if models is None: + models = OpenAIServingModels( + engine_client=engine_client, + base_model_paths=[], + ) + if online_renderer is None: + online_renderer = MagicMock() + + instance = OmniOpenAIServingChat( + engine_client=engine_client, + models=models, + response_role=response_role, + online_renderer=online_renderer, + request_logger=request_logger, + chat_template=chat_template, + chat_template_content_format=chat_template_content_format, + **init_overrides, + ) + instance._create_audio_choice = MagicMock( + side_effect=lambda omni_res, role, request, stream=False: mock_audio_choices( + index=omni_res.outputs[0].index, + role=role, + ) + ) + return instance + + +def make_request( + modalities: list[str], + n: int = 1, + *, + stream: bool = True, + include_usage: bool = False, +) -> ChatCompletionRequest: + req = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + n=n, + stream=stream, + stream_options={"include_usage": True} if include_usage else None, + ) + req.modalities = modalities # type: ignore[attr-defined] + return req + + +def parse_sse_chunks(lines: list[str]) -> list[dict]: + """Parse SSE lines into JSON dicts.""" + prefix = "data: " + chunks = [] + for line in lines: + line = line.strip() + if not line.startswith(prefix): + continue + payload = line[len(prefix) :].strip() + if payload == "[DONE]": + continue + chunks.append(json.loads(payload)) + return chunks + + +async def collect_stream(gen): + result = [] + async for item in gen: + result.append(item) + return result diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index 8706bb1bc..36d8c406a 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Config/message construction helpers used by tests.""" import atexit @@ -8,11 +11,13 @@ from typing import Any import yaml +from vllm_omni.config.stage_config import load_deploy_config + def modify_stage_config( yaml_path: str, - updates: dict[str, Any] = None, - deletes: dict[str, Any] = None, + updates: dict[str, Any] | None = None, + deletes: dict[str, Any] | None = None, ) -> str: """ Modify configurations in a YAML file, supporting both top-level and stage-specific modifications, @@ -23,17 +28,17 @@ def modify_stage_config( updates: Dictionary containing both top-level and stage-specific modifications to add or update. Format: { 'async_chunk': True, - 'stage_args': { - 0: {'engine_args.max_model_len': 5800}, - 1: {'engine_args.max_num_seqs': 2} + 'stages': { + 0: {'max_model_len': 5800}, + 1: {'max_num_seqs': 2} } } deletes: Dictionary containing configurations to delete. Format: { 'old_config': None, # Delete entire key - 'stage_args': { - 0: ['engine_args.old_param'], - 1: ['runtime.unused_setting'] + 'stages': { + 0: ['old_param'], + 1: ['unused_setting'] } } @@ -164,24 +169,20 @@ def modify_stage_config( else: print(f"Path {path} does not exist") - _stage_key = "stages" if "stages" in config else "stage_args" + stages = get_stage_entries(config) # Apply deletions first if deletes: for key, value in deletes.items(): - if key in ("stage_args", "stages"): + if key == "stages": if value and isinstance(value, dict): - stage_args = config.get(_stage_key, []) - if not stage_args: - raise ValueError("stage_args does not exist in config") - for stage_id, delete_paths in value.items(): if not delete_paths: continue # Find stage by ID target_stage = None - for stage in stage_args: + for stage in stages: if stage.get("stage_id") == int(stage_id): target_stage = stage break @@ -204,28 +205,24 @@ def modify_stage_config( # Apply updates if updates: for key, value in updates.items(): - if key in ("stage_args", "stages"): + if key == "stages": if value and isinstance(value, dict): - stage_args = config.get(_stage_key, []) - if not stage_args: - raise ValueError("stage_args does not exist in config") - for stage_id, stage_updates in value.items(): # Find stage by ID target_stage = None - for stage in stage_args: + for stage in stages: if stage.get("stage_id") == int(stage_id): target_stage = stage break if target_stage is None: - available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s] + available_ids = [s.get("stage_id") for s in stages if "stage_id" in s] raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}") # Apply updates to this stage for update_path, val in stage_updates.items(): # Check if this is a simple key (not dot-separated) - # Example: 'engine_input_source' vs 'engine_args.max_model_len' + # Example: 'engine_input_source' vs nested 'default_sampling_params.max_tokens' if "." not in update_path: # Direct key assignment (e.g., updating a list value) target_stage[update_path] = val @@ -624,6 +621,36 @@ def _materialize_ci_overlay(model_type: str) -> Path: return out +def get_stage_entries(cfg: dict[str, Any]) -> list[dict[str, Any]]: + """Return the ``stages`` list from a new-schema deploy config dict. + + Raises ``ValueError`` when the top-level ``stages`` key is missing. + """ + stages = cfg.get("stages") + if stages is None: + raise ValueError("Deploy config missing top-level 'stages'") + return list(stages) + + +def load_stage_ids(resolved_config: dict[str, Any]) -> list[int]: + """Return ``stage_id`` values from a resolved deploy config dict.""" + stage_ids = [stage["stage_id"] for stage in get_stage_entries(resolved_config) if "stage_id" in stage] + if not stage_ids: + raise ValueError("No stage IDs found in resolved config") + return stage_ids + + +def load_stage_replica_counts(resolved_config: dict[str, Any]) -> dict[int, int]: + """Return ``{stage_id: num_replicas}`` from flat stage entries (default 1).""" + replica_counts: dict[int, int] = {} + for stage in get_stage_entries(resolved_config): + stage_id = stage.get("stage_id") + if stage_id is None: + continue + replica_counts[int(stage_id)] = max(1, int(stage.get("num_replicas", 1))) + return replica_counts + + def get_deploy_config_path(rel_path: str) -> str: """Resolve a deploy yaml; ``ci/.yaml`` materializes from ``_CI_OVERLAYS``.""" if rel_path.startswith("ci/") and rel_path.endswith(".yaml"): @@ -638,33 +665,30 @@ def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]: with open(get_deploy_config_path(rel_path), encoding="utf-8") as f: cfg = yaml.safe_load(f) or {} - stage_key = "stages" if "stages" in cfg else "stage_args" - for stage in cfg.get(stage_key, []): + for stage in get_stage_entries(cfg): if stage.get("stage_id") == stage_id: return stage raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}") -def _get_config_value_by_path(config_dict: dict, path: str) -> Any: - """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``).""" - current: Any = config_dict - for key in path.split("."): - if not isinstance(current, dict) or key not in current: - return None - current = current[key] - return current +def get_deploy_duplex_max_sessions(rel_path: str) -> int: + """Return the duplex session capacity a deploy yaml admits. + + Loads through ``load_deploy_config`` so ``base_config`` merging and + ``DuplexSessionRuntimeConfig`` defaults stay in lockstep with the server + instead of being re-implemented here. + """ + return load_deploy_config(get_deploy_config_path(rel_path)).duplex_session.max_sessions -def _stage_load_format_paths(stage_config_path: str) -> tuple[str, str, list[int]]: - """Return ``(stage_key, load_format_field_path, stage_ids)`` for a deploy YAML.""" +def _stage_ids_from_deploy_yaml(stage_config_path: str) -> list[int]: + """Return ``stage_id`` values from a new-schema deploy YAML (``stages``).""" with open(stage_config_path, encoding="utf-8") as f: cfg = yaml.safe_load(f) or {} - new_schema_stages = cfg.get("stages") - stage_key = "stages" if new_schema_stages is not None else "stage_args" - load_format_path = "load_format" if new_schema_stages is not None else "engine_args.load_format" - stage_entries = cfg.get(stage_key, []) - stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage] - return stage_key, load_format_path, stage_ids + try: + return load_stage_ids(cfg) + except ValueError as exc: + raise ValueError(f"Deploy YAML missing top-level 'stages': {stage_config_path}") from exc def _add_dummy_load_format( @@ -674,10 +698,10 @@ def _add_dummy_load_format( """For ``core_model`` runs, patch every stage in the deploy YAML to ``load_format: dummy``.""" if run_level != "core_model" or stage_config_path is None: return stage_config_path - stage_key, load_format_path, stage_ids = _stage_load_format_paths(stage_config_path) + stage_ids = _stage_ids_from_deploy_yaml(stage_config_path) return modify_stage_config( stage_config_path, - updates={stage_key: {stage_id: {load_format_path: "dummy"} for stage_id in stage_ids}}, + updates={"stages": {stage_id: {"load_format": "dummy"} for stage_id in stage_ids}}, ) @@ -688,25 +712,27 @@ def _delete_dummy_load_format( """For ``advanced_model`` / ``full_model``, strip ``load_format: dummy`` so real weights load.""" if run_level not in {"advanced_model", "full_model"} or stage_config_path is None: return stage_config_path - stage_key, load_format_path, _stage_ids = _stage_load_format_paths(stage_config_path) with open(stage_config_path, encoding="utf-8") as f: cfg = yaml.safe_load(f) or {} - stage_entries = cfg.get(stage_key, []) + try: + stages = get_stage_entries(cfg) + except ValueError as exc: + raise ValueError(f"Deploy YAML missing top-level 'stages': {stage_config_path}") from exc deletes: dict[int, list[str]] = {} - for stage in stage_entries: + for stage in stages: stage_id = stage.get("stage_id") if stage_id is None: continue - if _get_config_value_by_path(stage, load_format_path) == "dummy": - deletes[stage_id] = [load_format_path] + if stage.get("load_format") == "dummy": + deletes[stage_id] = ["load_format"] if not deletes: return stage_config_path return modify_stage_config( stage_config_path, - deletes={stage_key: deletes}, + deletes={"stages": deletes}, ) @@ -720,6 +746,10 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st __all__ = [ "get_deploy_config_path", "get_deploy_config_stage", + "get_deploy_duplex_max_sessions", + "get_stage_entries", + "load_stage_ids", + "load_stage_replica_counts", "modify_stage_config", "stage_config_path_for_run_level", ] diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py index 9cb1560b9..98deaa511 100644 --- a/tests/helpers/tests/test_assertions.py +++ b/tests/helpers/tests/test_assertions.py @@ -1,12 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace import pytest from tests.helpers import assertions -from tests.helpers.assertions import _assert_transcript_matches, _resolve_audio_transcript +from tests.helpers.assertions import ( + _assert_transcript_matches, + _resolve_audio_transcript, + assert_audio_speech_response, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -71,6 +75,40 @@ def test_resolve_transcript_leaves_language_unset_by_default(monkeypatch): assert captured["model_size"] == "small" +@pytest.mark.parametrize( + ("request_config", "expected_text"), + [ + ({"input": "spoken words"}, "spoken words"), + ( + { + "input": "<|style:whispering|>spoken words", + "transcript_expected_text": "spoken words", + }, + "spoken words", + ), + ], +) +def test_speech_transcript_expected_text(monkeypatch, request_config, expected_text): + captured: dict = {} + + monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "spoken words") + + def capture_match(_transcript, _audio_bytes, expected_text, **_kwargs): + captured["expected_text"] = expected_text + + monkeypatch.setattr(assertions, "_assert_transcript_matches", capture_match) + response = SimpleNamespace(success=True, audio_bytes=b"fake-wav", audio_format="audio/wav") + request_config["response_format"] = "wav" + + assert_audio_speech_response( + response, + request_config, + "advanced_model", + ) + + assert captured["expected_text"] == expected_text + + def test_escalated_transcript_keeps_declared_language(monkeypatch): # The escalated pass must honour the same language, otherwise a request that # pins one silently falls back to auto-detection on retry. diff --git a/tests/helpers/tests/test_mark.py b/tests/helpers/tests/test_mark.py new file mode 100644 index 000000000..d2fed4b6d --- /dev/null +++ b/tests/helpers/tests/test_mark.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from tests.helpers.mark import ( + _gpu_res_platforms, + hardware_marks, + hardware_test, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _mark_names(marks) -> set[str]: + return {mark.name for mark in marks} + + +def test_gpu_platforms_come_from_pyproject_gpu_tag(): + assert "npu" not in _gpu_res_platforms() + assert "gpu" not in _gpu_res_platforms() + + +def test_default_num_cards_adds_cards_1(): + names = _mark_names(hardware_marks(res={"cuda": "H100"})) + assert "cards_1" in names + assert "H100" in names + assert "cards_2" not in names + + +def test_h100_two_cards_is_selectable_with_cards_2(): + names = _mark_names(hardware_marks(res={"cuda": "H100"}, num_cards=2)) + assert names >= {"H100", "cuda", "gpu", "cards_2"} + assert "cards_1" not in names + assert "cards_4" not in names + + +def test_h100_four_cards_is_selectable_with_cards_4(): + names = _mark_names(hardware_marks(res={"cuda": "H100"}, num_cards=4)) + assert "cards_4" in names + assert "cards_2" not in names + + +def test_npu_card_count_mark_is_independent_of_sku(): + names = _mark_names(hardware_marks(res={"npu": "A2"}, num_cards=4)) + assert "cards_4" in names + assert "npu" in names + assert "gpu" not in names + + +def test_hardware_marks_rejects_mixed_counts_on_one_item(): + with pytest.raises(ValueError, match="different per-platform num_cards"): + hardware_marks( + res={"cuda": "H100", "rocm": "MI325"}, + num_cards={"cuda": 4, "rocm": 2}, + ) + + +def test_same_count_multi_platform_stays_one_item(): + names = _mark_names(hardware_marks(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)) + assert names >= {"H100", "cuda", "MI325", "rocm", "cards_2"} + + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + def _probe(): + return None + + names = {mark.name for mark in _probe.pytestmark} + assert "parametrize" not in names + assert "cards_2" in names + assert "H100" in names + assert "MI325" in names + + +def test_hardware_test_decorator_applies_cards_mark(): + @hardware_test(res={"cuda": "H100"}, num_cards=2) + def _probe(): + return None + + names = {mark.name for mark in _probe.pytestmark} + assert "cards_2" in names + assert "H100" in names + + +def _hardware_platform_params(func): + marks = func.pytestmark if isinstance(func.pytestmark, list) else [func.pytestmark] + for mark in marks: + if mark.name != "parametrize": + continue + argnames = mark.args[0] + if argnames == "_normalized_hardware_marks": + return mark.args[1] + raise AssertionError("expected @hardware_test to parametrize _normalized_hardware_marks") + + +def test_mixed_card_counts_split_so_cuda_filter_does_not_see_rocm_count(): + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) + def _probe(): + return None + + variants = {param.id: {mark.name for mark in param.marks} for param in _hardware_platform_params(_probe)} + cuda_names = variants["cuda"] + rocm_names = variants["rocm"] + + assert cuda_names >= {"H100", "cuda", "cards_4"} + assert "cards_2" not in cuda_names + assert "rocm" not in cuda_names + assert "MI325" not in cuda_names + + assert rocm_names >= {"MI325", "rocm", "cards_2"} + assert "cards_4" not in rocm_names + assert "cuda" not in rocm_names + assert "H100" not in rocm_names + + +def test_unsupported_card_count_raises(): + with pytest.raises(ValueError, match="num_cards=16"): + hardware_marks(res={"cuda": "H100"}, num_cards=16) + + +@pytest.mark.parametrize("platform", ["cpu", "gpu"]) +def test_res_rejects_cpu_and_gpu(platform): + with pytest.raises(ValueError, match="not a res dict key"): + hardware_marks(res={platform: "H100"}) + + +def test_npu_rejects_unsupported_sku(): + with pytest.raises(ValueError, match="Invalid npu resource type"): + hardware_marks(res={"npu": "A100"}) + + +@pytest.mark.parametrize("sku", ["H200", "B200", "H800"]) +def test_cuda_registered_skus_attach_sku_and_cards(sku): + names = _mark_names(hardware_marks(res={"cuda": sku}, num_cards=2)) + assert names >= {sku, "cuda", "gpu", "cards_2"} + + +def test_cuda_rejects_sku_from_another_platform(): + with pytest.raises(ValueError, match="Invalid cuda resource type"): + hardware_marks(res={"cuda": "MI325"}) diff --git a/tests/helpers/tests/test_media.py b/tests/helpers/tests/test_media.py index 4a8e5393d..8045f7a69 100644 --- a/tests/helpers/tests/test_media.py +++ b/tests/helpers/tests/test_media.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import sys +from concurrent.futures.process import BrokenProcessPool +from contextlib import nullcontext from types import SimpleNamespace import numpy as np @@ -22,7 +24,11 @@ def _stub_transcribe(monkeypatch) -> dict: captured.update(kwargs) return {"text": "London"} - monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=lambda size, device=None: FakeModel())) + monkeypatch.setitem( + sys.modules, + "whisper", + SimpleNamespace(load_model=lambda size, device=None: FakeModel()), + ) # Keep the unit test off the accelerator probe so it stays hermetic. monkeypatch.setitem( sys.modules, @@ -50,32 +56,183 @@ def test_transcribe_defaults_to_auto_language(monkeypatch): assert captured.get("language") is None +class _FakeExecutor: + """Stand-in for ProcessPoolExecutor recording what was submitted and shut down.""" + + def __init__(self, outcomes=()): + self._outcomes = list(outcomes) + self.submitted: list[tuple] = [] + self.shutdown_calls = 0 + + def submit(self, fn, *args): + self.submitted.append((fn, args)) + outcome = self._outcomes.pop(0) if self._outcomes else "London" + + def _result(): + if isinstance(outcome, BaseException): + raise outcome + return outcome + + return SimpleNamespace(result=_result) + + def shutdown(self, wait=True): + self.shutdown_calls += 1 + + +def _patch_executors(monkeypatch, outcomes_per_executor=()) -> list[_FakeExecutor]: + """Replace the pool constructor and hand back the fakes it produces, in creation order.""" + created: list[_FakeExecutor] = [] + + def factory(*args, **kwargs): + index = len(created) + outcomes = outcomes_per_executor[index] if index < len(outcomes_per_executor) else () + executor = _FakeExecutor(outcomes) + executor.init_kwargs = kwargs + created.append(executor) + return executor + + monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", factory) + return created + + +@pytest.fixture(autouse=True) +def _reset_transcriber_singletons(): + """The worker and its model cache are module-level state; do not leak them across tests.""" + media.release_audio_transcriber() + media._WHISPER_MODELS.clear() + yield + media.release_audio_transcriber() + media._WHISPER_MODELS.clear() + + def test_bytes_entrypoint_forwards_language_to_subprocess(monkeypatch, tmp_path): """Cover the two hops the tests above skip: bytes -> file -> executor.submit. The real call crosses a spawn ProcessPoolExecutor, so capture what gets submitted rather than what the worker eventually does. """ - submitted: dict = {} - - class FakeExecutor: - def __enter__(self): - return self - - def __exit__(self, *exc): - return False - - def submit(self, fn, *args): - submitted["fn"] = fn - submitted["args"] = args - return SimpleNamespace(result=lambda: "London") - - monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", lambda *a, **kw: FakeExecutor()) + created = _patch_executors(monkeypatch) wav = tmp_path / "clip.wav" sf.write(wav, np.zeros(2400, dtype=np.float32), 24000) assert media.convert_audio_bytes_to_text(wav.read_bytes(), "small", "en") == "London" - assert submitted["fn"] is media._whisper_transcribe_in_current_process - assert submitted["args"][1:] == ("small", "en") + fn, args = created[0].submitted[0] + assert fn is media._whisper_transcribe_in_current_process + assert args[1:] == ("small", "en") + + +def test_parent_reuses_one_spawned_worker_across_calls(monkeypatch): + """Two calls share one spawned worker. + + Pins the parent-side contract: the pool is constructed once, with + ``spawn`` (CUDA is not fork-safe) and ``max_workers=1`` (one resident + model). Reuse of the worker's *loaded model* is pinned by the worker-side + cache test below, which is where the startup cost actually lives. + """ + created = _patch_executors(monkeypatch) + + media.convert_audio_file_to_text("/tmp/a.wav") + media.convert_audio_file_to_text("/tmp/b.wav") + + assert len(created) == 1 + assert len(created[0].submitted) == 2 + assert created[0].init_kwargs["max_workers"] == 1 + assert created[0].init_kwargs["mp_context"].get_start_method() == "spawn" + + +def test_release_forces_a_new_process_pool(monkeypatch): + created = _patch_executors(monkeypatch) + + media.convert_audio_file_to_text("/tmp/a.wav") + media.release_audio_transcriber() + media.convert_audio_file_to_text("/tmp/b.wav") + + assert len(created) == 2 + assert created[0].shutdown_calls == 1 + + +def test_worker_loads_each_model_size_once(monkeypatch): + """The worker-side cache is what removes the repeated whisper.load_model cost.""" + loaded: list[str] = [] + + class FakeModel: + def transcribe(self, path, **kwargs): + return {"text": "London"} + + def load_model(size, device=None): + loaded.append(size) + return FakeModel() + + monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=load_model)) + monkeypatch.setitem( + sys.modules, + "vllm_omni.platforms", + SimpleNamespace(current_omni_platform=SimpleNamespace(is_available=lambda: False)), + ) + monkeypatch.setattr(media, "_serialize_whisper_model_download", lambda model_size: nullcontext()) + + media._whisper_transcribe_in_current_process("/tmp/a.wav", "small") + media._whisper_transcribe_in_current_process("/tmp/b.wav", "small") + + assert loaded == ["small"] + + # The ASR escalation path asks the same worker for a stronger model. + media._whisper_transcribe_in_current_process("/tmp/c.wav", "large-v3") + + assert loaded == ["small", "large-v3"] + + +@pytest.mark.parametrize( + ("second_outcome", "recovers"), + [(["London"], True), ([BrokenProcessPool("second")], False)], +) +def test_broken_pool_is_discarded_and_retried_once(monkeypatch, second_outcome, recovers): + """A dead worker is discarded and the call retried exactly once. + + The retry either succeeds, or -- if the replacement also dies -- fails with + both pools discarded and no executor left installed. + """ + created = _patch_executors( + monkeypatch, + outcomes_per_executor=([BrokenProcessPool("first")], second_outcome), + ) + + if recovers: + assert media.convert_audio_file_to_text("/tmp/a.wav") == "London" + else: + with pytest.raises(BrokenProcessPool): + media.convert_audio_file_to_text("/tmp/a.wav") + assert created[1].shutdown_calls == 1 + assert media._TRANSCRIBER is None + + assert len(created) == 2 # discarded the first, retried on a fresh one + assert created[0].shutdown_calls == 1 + + +@pytest.mark.parametrize( + "error", + [RuntimeError("CUDA out of memory"), KeyboardInterrupt()], + ids=["Exception", "BaseException"], +) +def test_transcription_failure_propagates_and_discards_the_worker(monkeypatch, error): + """Any failure -- an Exception or a BaseException like KeyboardInterrupt -- + propagates AND tears the worker down. + + It arrives as the task's own exception, not ``BrokenProcessPool``, but the + worker -- and its resident model -- is still discarded, restoring the old + one-process-per-call isolation. It is not retried; the next call builds a + fresh worker. + """ + created = _patch_executors(monkeypatch, outcomes_per_executor=([error],)) + + with pytest.raises(type(error)): + media.convert_audio_file_to_text("/tmp/a.wav") + + assert len(created) == 1 # not retried + assert created[0].shutdown_calls == 1 + assert media._TRANSCRIBER is None + + assert media.convert_audio_file_to_text("/tmp/b.wav") == "London" + assert len(created) == 2 # a fresh worker for the next call diff --git a/tests/helpers/tests/test_runtime_omni_runner.py b/tests/helpers/tests/test_runtime_omni_runner.py new file mode 100644 index 000000000..aaada20f2 --- /dev/null +++ b/tests/helpers/tests/test_runtime_omni_runner.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Tests for OmniRunner startup rollback.""" + +from __future__ import annotations + +import pytest + +from tests.helpers.runtime import OmniRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_runner_init_rolls_back_on_omni_startup_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """``__exit__`` is skipped when construction raises; rollback must still run.""" + cleaned: list[str] = [] + monkeypatch.setattr( + "tests.helpers.runtime.cleanup_test_environment", + lambda: cleaned.append("env"), + ) + monkeypatch.setattr( + OmniRunner, + "_cleanup_process", + lambda self: cleaned.append("process"), + ) + + class _BoomOmni: + def __init__(self, *args, **kwargs) -> None: + raise RuntimeError("Orchestrator initialization failed") + + monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _BoomOmni) + + with pytest.raises(RuntimeError, match="Orchestrator initialization failed"): + with OmniRunner("fake-model"): + raise AssertionError("context body must not run after a failed constructor") + + assert cleaned.count("process") == 1 + # Once at the start of ``__init__``, once from the constructor rollback. + assert cleaned.count("env") == 2 diff --git a/tests/helpers/tests/test_runtime_stage_cli.py b/tests/helpers/tests/test_runtime_stage_cli.py index 5a2ab5d68..564f12cf1 100644 --- a/tests/helpers/tests/test_runtime_stage_cli.py +++ b/tests/helpers/tests/test_runtime_stage_cli.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for OmniServerStageCli process planning helpers.""" from __future__ import annotations from pathlib import Path +from types import SimpleNamespace import pytest @@ -38,23 +42,48 @@ def test_stage_cli_builds_headless_replica_cmd(tmp_path: Path) -> None: assert cmd[cmd.index("--deploy-config") + 1] == server.stage_config_path -def test_stage_cli_splits_devices_per_replica(tmp_path: Path) -> None: +def test_stage_cli_loads_stage_ids_and_replica_counts(tmp_path: Path) -> None: server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), []) - assert server._devices_for_replica(1, "1,2,3", 0) == "1" - assert server._devices_for_replica(1, "1,2,3", 1) == "2" - assert server._devices_for_replica(1, "1,2,3", 2) == "3" + assert server.stage_ids == [0, 1] + assert server.stage_replica_counts == {0: 1, 1: 3} -def test_stage_cli_maps_visible_devices_after_replica_split(tmp_path: Path) -> None: - server = OmniServerStageCli( - "fake-model", - _write_stage_config(tmp_path), - [], - env_dict={"CUDA_VISIBLE_DEVICES": "4,5,6,7"}, - ) - env: dict[str, str] = {} +def test_stage_cli_enter_rolls_back_launched_stages_on_startup_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """``__exit__`` is skipped when ``__enter__`` raises; rollback must still run.""" + monkeypatch.setattr("tests.helpers.runtime.time.sleep", lambda *_args, **_kwargs: None) + monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None) - server._set_stage_device_env(1, env, "1,2,3", replica_id=1) + server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), ["--disable-log-stats"]) + log_path = tmp_path / "stage0.log" + log_fh = open(log_path, "w", encoding="utf-8") + log_fh.write("partial launch\n") + log_fh.flush() - assert env["CUDA_VISIBLE_DEVICES"] == "6" + proc = SimpleNamespace(pid=4242, poll=lambda: None) + killed: list[int] = [] + monkeypatch.setattr(server, "_kill_process_tree", lambda pid: killed.append(pid)) + + def _fake_launch(stage_id: int, *, headless: bool, replica_id: int = 0) -> None: + stage_key = (stage_id, replica_id) + server.stage_procs[stage_key] = proc # type: ignore[assignment] + server._stage_log_paths[stage_key] = log_path + server._stage_log_files[stage_key] = log_fh + if stage_id == 0 and replica_id == 0: + server.proc = proc # type: ignore[assignment] + + def _fail_alive() -> None: + raise RuntimeError("Stage 0 replica 0 exited with code 1") + + monkeypatch.setattr(server, "_launch_stage", _fake_launch) + monkeypatch.setattr(server, "_ensure_stage_processes_alive", _fail_alive) + + with pytest.raises(RuntimeError, match="exited with code 1"): + with server: + raise AssertionError("context body must not run after a failed enter") + + assert killed == [4242] + assert log_fh.closed + assert not log_path.exists() diff --git a/tests/helpers/tests/test_runtime_transcriber_release.py b/tests/helpers/tests/test_runtime_transcriber_release.py new file mode 100644 index 000000000..5e7f29ebe --- /dev/null +++ b/tests/helpers/tests/test_runtime_transcriber_release.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""The server/runner fixtures free the Whisper judge around each instance. + +iter_omni_server and iter_omni_runner both wrap their body in the same +_whisper_device_free_around() context manager, so one representative path +exercises its enter/exit and exception behavior. +""" + +import threading +from types import SimpleNamespace + +import pytest + +from tests.helpers import runtime +from tests.helpers.runtime import OmniServerParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def events(monkeypatch) -> list[str]: + """Record construction and release in order.""" + calls: list[str] = [] + + class FakeServer: + def __init__(self, *args, **kwargs): + calls.append("construct") + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + monkeypatch.setattr(runtime, "OmniServer", FakeServer) + monkeypatch.setattr(runtime, "release_audio_transcriber", lambda: calls.append("release")) + monkeypatch.setattr( + "tests.helpers.stage_config.stage_config_path_for_run_level", + lambda path, run_level: None, + ) + return calls + + +def _server_generator(): + request = SimpleNamespace( + param=OmniServerParams(model="fake-model"), + node=SimpleNamespace(get_closest_marker=lambda name: None), + ) + return runtime.iter_omni_server(request, "core_model", threading.Lock()) + + +def test_releases_before_construction_and_after_teardown(events): + generator = _server_generator() + generator.send(None) + + assert events == ["release", "construct"], "the device must be clear before the model loads" + + with pytest.raises(StopIteration): + generator.send(None) + + assert events == ["release", "construct", "release"] + + +def test_failing_test_still_releases(events): + """A test that raises must not strand Whisper on the device.""" + generator = _server_generator() + generator.send(None) + + with pytest.raises(RuntimeError, match="assertion blew up"): + generator.throw(RuntimeError("assertion blew up")) + + assert events[-1] == "release" diff --git a/tests/helpers/tests/test_success_rate_helpers.py b/tests/helpers/tests/test_success_rate_helpers.py new file mode 100644 index 000000000..c1b8f4417 --- /dev/null +++ b/tests/helpers/tests/test_success_rate_helpers.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +from types import SimpleNamespace + +import pytest + +from tests.helpers import assertions +from tests.helpers.assertions import ( + AUDIO_MISMATCH_MESSAGE, + SuccessRateGate, + assert_omni_response, + collect_at_success_rate, + is_quality_failure, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_GATE = SuccessRateGate(min_successes=7, concurrency=2) + + +def _fake_response(text: str) -> SimpleNamespace: + return SimpleNamespace(success=True, text_content=text, audio_content=None, audio_bytes=b"fake-wav") + + +def _omni_response_failure(monkeypatch, *, text: str, transcript: str) -> AssertionError: + """Drive the real assert_omni_response to failure and hand back what it raised.""" + monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *a, **kw: transcript) + request_config = {"modalities": ["text", "audio"], "key_words": {"text": ["london"]}} + with pytest.raises(AssertionError) as excinfo: + assert_omni_response(_fake_response(text), request_config, run_level="full_model") + return excinfo.value + + +def test_real_audio_mismatch_is_tolerated(monkeypatch): + # Ties the classifier to the actual producer: if assert_omni_response's wording drifts, + # this fails instead of the gate silently turning every quality failure into a hard one. + assert is_quality_failure(_omni_response_failure(monkeypatch, text="London", transcript="런던")) + + +def test_real_keyword_failure_is_not_tolerated(monkeypatch): + assert not is_quality_failure(_omni_response_failure(monkeypatch, text="Paris", transcript="London")) + + +@pytest.mark.parametrize("exc", [RuntimeError(AUDIO_MISMATCH_MESSAGE), ValueError("boom")]) +def test_non_assertion_failures_are_not_tolerated(exc): + # A transport error carrying the same text is still a serving failure, not a sample. + assert not is_quality_failure(exc) + + +class _ScriptedSampler: + """One scripted outcome per call: raised if it is an exception, else returned.""" + + def __init__(self, outcomes: list) -> None: + self._outcomes = outcomes + self._lock = threading.Lock() + self.calls = 0 + + def __call__(self): + with self._lock: + outcome = self._outcomes[self.calls] + self.calls += 1 + if isinstance(outcome, BaseException): + raise outcome + return outcome + + +def _mismatch(detail: str) -> AssertionError: + return AssertionError(f"{AUDIO_MISMATCH_MESSAGE} (transcript={detail!r})") + + +def test_tolerated_failures_are_counted_and_every_sample_is_sent(): + sampler = _ScriptedSampler(["ok"] * 7 + [_mismatch("a"), _mismatch("b"), _mismatch("c")]) + + responses = collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None) + + # No early exit once the gate is already satisfied: it would bias the measured rate. + assert sampler.calls == 10 + assert responses == ["ok"] * 7 + + +def test_gate_failure_reports_the_count_and_every_failing_sample(): + outcomes = ["ok"] * 6 + [_mismatch("first"), _mismatch("second"), _mismatch("third"), _mismatch("fourth")] + + with pytest.raises(AssertionError) as excinfo: + collect_at_success_rate(_ScriptedSampler(outcomes), _GATE, request_num=10, report=lambda _: None) + + message = str(excinfo.value) + assert "6/10" in message and "7/10 gate" in message + assert all(detail in message for detail in ("first", "second", "third", "fourth")) + + +def test_serving_failure_is_not_absorbed_into_the_budget(): + sampler = _ScriptedSampler(["ok", AssertionError("No audio output is generated"), *["ok"] * 8]) + + with pytest.raises(AssertionError, match="No audio output"): + collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None) + + # Stopped on the serving failure rather than spending the remaining samples. + assert sampler.calls < 10 + + +def test_unreachable_gate_is_rejected_before_sending(): + sampler = _ScriptedSampler(["ok"] * 3) + + with pytest.raises(ValueError, match="can never pass"): + collect_at_success_rate(sampler, _GATE, request_num=3, report=lambda _: None) + + assert sampler.calls == 0 diff --git a/tests/host_weight_runtime/test_contracts.py b/tests/host_weight_runtime/test_contracts.py new file mode 100644 index 000000000..5c7036c8e --- /dev/null +++ b/tests/host_weight_runtime/test_contracts.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Contract tests for semantic identity and immutable artifact schemas.""" + +from __future__ import annotations + +import ast +import math +from pathlib import Path +from typing import get_type_hints + +import pytest + +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + ComponentIdentity, + FileManifestEntry, + IdentityValidationError, + ManifestValidationError, + ProducerIdentity, + ProductionPolicy, + PublicationMarker, + RuntimeWeightLayout, + TensorKind, + TensorManifestEntry, + WeightArtifactIdentity, + WeightManifest, + WeightRepresentation, + WeightRestorePlan, + WeightRestorer, + WeightSourceIdentity, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_SHA_A = "a" * 64 +_SHA_B = "b" * 64 +_SHA_C = "c" * 64 + + +def _identity( + *, + tp_size: int = 1, + tp_rank: int = 0, + sp_size: int = 1, + adaptation: AdaptationIdentity | None = None, + source_metadata: CanonicalJson | None = None, +) -> WeightArtifactIdentity: + return WeightArtifactIdentity( + schema_version=1, + source=WeightSourceIdentity( + model_id="test-org/test-model", + revision="0123456789abcdef", + fingerprint="source-files-sha256", + metadata=source_metadata or CanonicalJson.empty(), + ), + component=ComponentIdentity(name="transformer", ownership="diffusion.dit"), + representation=WeightRepresentation(name="runtime-bf16", dtype="torch.bfloat16"), + layout=RuntimeWeightLayout( + name="final-module-layout", + tensor_parallel_size=tp_size, + tensor_parallel_rank=tp_rank, + sequence_parallel_size=sp_size, + sequence_parallel_backend="ulysses" if sp_size > 1 else "none", + ), + adaptation=adaptation or AdaptationIdentity(), + producer=ProducerIdentity( + producer_id="test.final-layout", + version="1", + implementation_fingerprint="implementation-sha256", + manifest_schema="test-producer-v1", + restorer_schema="test-restorer-v1", + ), + ) + + +def test_canonical_metadata_is_a_stable_immutable_snapshot() -> None: + items: list[object] = [1, {"value": 2}] + mutable = {"z": items, "a": True} + snapshot = CanonicalJson.from_value(mutable) + items.append(3) + + assert snapshot.encoded == b'{"a":true,"z":[1,{"value":2}]}' + assert snapshot.to_value() == {"a": True, "z": [1, {"value": 2}]} + assert CanonicalJson.from_value({"a": True, "z": [1, {"value": 2}]}) == snapshot + + +@pytest.mark.parametrize("value", [math.inf, math.nan, object(), {1: "non-string-key"}]) +def test_canonical_metadata_rejects_process_specific_or_non_json_values(value: object) -> None: + with pytest.raises(IdentityValidationError): + CanonicalJson.from_value(value) + + +@pytest.mark.parametrize("encoded", [b"not-json", b'{"b":1, "a":2}', b'{"value":NaN}']) +def test_canonical_metadata_cannot_bypass_stable_encoding(encoded: bytes) -> None: + with pytest.raises(IdentityValidationError): + CanonicalJson(encoded) + + +def test_identity_round_trip_and_hash_are_deterministic() -> None: + first = _identity(source_metadata=CanonicalJson.from_value({"files": ["b", "a"], "size": 3})) + second = _identity(source_metadata=CanonicalJson.from_value({"size": 3, "files": ["b", "a"]})) + + assert first == second + assert first.key == second.key + assert WeightArtifactIdentity.from_dict(first.to_dict()) == first + assert len(first.key) == 64 + + +def test_identity_separates_semantic_parallel_layout_and_static_adaptation() -> None: + base = _identity() + tp0 = _identity(tp_size=2, tp_rank=0) + tp1 = _identity(tp_size=2, tp_rank=1) + sp2 = _identity(sp_size=2) + merged_lora = _identity(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256")) + + assert len({item.key for item in (base, tp0, tp1, sp2, merged_lora)}) == 5 + assert "data_parallel" not in base.layout.to_dict() + assert "sequence_parallel_rank" not in sp2.layout.to_dict() + + +def test_identity_schema_rejects_unknown_fields() -> None: + document = _identity().to_dict() + document["transfer_mode"] = "registered" + + with pytest.raises(IdentityValidationError, match="unexpected"): + WeightArtifactIdentity.from_dict(document) + + +def test_identity_schema_does_not_coerce_field_types() -> None: + document = _identity().to_dict() + document["schema_version"] = "1" + + with pytest.raises(IdentityValidationError, match="must be an integer"): + WeightArtifactIdentity.from_dict(document) + + +def test_restoration_contract_has_a_validation_only_one_shot_commit_boundary() -> None: + assert get_type_hints(WeightRestorePlan.commit)["return"] is type(None) + assert "one-shot" in (WeightRestorePlan.__doc__ or "") + assert "without mutating" in (WeightRestorer.__doc__ or "") + + +def test_post_load_publication_policy_is_explicitly_enabled() -> None: + policy = ProductionPolicy(allow_post_load_publish=True) + + assert policy.allow_local_build + assert policy.allow_post_load_publish + + +def test_manifest_and_publication_marker_round_trip_exactly() -> None: + identity = _identity() + manifest = WeightManifest( + schema_version=1, + artifact_key=identity.key, + identity=identity, + producer_schema="test-producer-v1", + restorer_schema="test-restorer-v1", + tensors=( + TensorManifestEntry( + name="layer.weight", + kind=TensorKind.PARAMETER, + role="weight", + file_name="weights.safetensors", + storage_key="layer.weight", + shape=(2, 4), + dtype="torch.bfloat16", + stride=(4, 1), + nbytes=16, + sha256=_SHA_A, + ), + ), + files=(FileManifestEntry(file_name="weights.safetensors", size=128, sha256=_SHA_B),), + artifact_content_sha256=_SHA_C, + format_metadata=CanonicalJson.from_value({"quantization": "none"}), + ) + marker = PublicationMarker( + schema_version=1, + publication_id="publication-id", + manifest_sha256=_SHA_A, + artifact_content_sha256=_SHA_C, + integrity_mode="full_checksum", + ) + + assert WeightManifest.from_dict(manifest.to_dict()) == manifest + assert PublicationMarker.from_dict(marker.to_dict()) == marker + + document = manifest.to_dict() + document["unknown"] = True + with pytest.raises(ManifestValidationError, match="unexpected"): + WeightManifest.from_dict(document) + + malformed = manifest.to_dict() + tensors = malformed["tensors"] + assert isinstance(tensors, list) and isinstance(tensors[0], dict) + tensors[0]["nbytes"] = "16" + with pytest.raises(ManifestValidationError, match="must be an integer"): + WeightManifest.from_dict(malformed) + + +@pytest.mark.parametrize("file_name", ["", ".", "..", "../weights.safetensors", "nested/weights.safetensors"]) +def test_manifest_rejects_unsafe_payload_names(file_name: str) -> None: + with pytest.raises(ManifestValidationError, match="safe relative basename"): + FileManifestEntry(file_name=file_name, size=1, sha256=_SHA_A) + + +def test_package_source_keeps_the_neutral_dependency_boundary() -> None: + package_root = Path(__file__).resolve().parents[2] / "vllm_omni" / "host_weight_runtime" + assert package_root.is_dir() + forbidden = ( + "vllm_omni.diffusion", + "vllm_omni.distributed", + "torch.cuda", + "torch.distributed", + ) + violations: list[str] = [] + for source_path in package_root.rglob("*.py"): + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + violations.extend( + f"{source_path.name}:{node.lineno}:{alias.name}" + for alias in node.names + if alias.name.startswith(forbidden) + ) + elif isinstance(node, ast.ImportFrom) and node.level == 0: + module = node.module or "" + if module.startswith(forbidden) or ( + module == "torch" and any(alias.name in {"cuda", "distributed"} for alias in node.names) + ): + violations.append(f"{source_path.name}:{node.lineno}:{module}") + + assert violations == [] diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py new file mode 100644 index 000000000..7fe716c5d --- /dev/null +++ b/tests/host_weight_runtime/test_filesystem_store.py @@ -0,0 +1,2581 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Lifecycle and concurrency tests for the local filesystem store.""" + +from __future__ import annotations + +import contextlib +import errno +import hashlib +import json +import multiprocessing as mp +import os +import shutil +import stat +import tempfile +import threading +import time +import warnings +from collections.abc import Sequence +from dataclasses import replace +from multiprocessing.process import BaseProcess +from pathlib import Path + +import pytest +import torch + +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + ArtifactInventoryState, + BuildRequest, + CapacityPolicy, + ComponentIdentity, + CoordinationScope, + FailureCode, + HostWeightError, + HostWeightLease, + IntegrityPolicy, + ProducerIdentity, + ProductionMetadata, + ProductionSourceMode, + ResolutionStage, + RuntimeWeightLayout, + StorageClass, + StorageDomainPolicy, + StorageScope, + StoreResult, + StoreStatus, + TensorKind, + TensorWriteSpec, + ValidationLevel, + WeightArtifactIdentity, + WeightProductionSpec, + WeightRepresentation, + WeightSourceIdentity, +) +from vllm_omni.host_weight_runtime.filesystem import ( + FilesystemHostWeightStore, + detect_filesystem_type, + detect_storage_class, +) +from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module +from vllm_omni.host_weight_runtime.filesystem import writer as filesystem_writer_module +from vllm_omni.host_weight_runtime.filesystem.locks import FileLock +from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Spawning imports the complete vLLM package and exceeded 10 seconds on a loaded +# CI worker. Keep process-bootstrap bounds separate from the short store +# deadlines exercised by these tests. +_PROCESS_START_TIMEOUT_SECONDS = 60.0 +_PROCESS_COMPLETION_TIMEOUT_SECONDS = 60.0 +_PROCESS_CLEANUP_TIMEOUT_SECONDS = 5.0 + + +def _identity(*, tp_rank: int = 0) -> WeightArtifactIdentity: + return WeightArtifactIdentity( + schema_version=1, + source=WeightSourceIdentity("test-org/test-model", "0123456789abcdef", "source-files-sha256"), + component=ComponentIdentity("transformer", "diffusion.dit"), + representation=WeightRepresentation("runtime-bf16", "torch.bfloat16"), + layout=RuntimeWeightLayout( + "final-module-layout", + tensor_parallel_size=2, + tensor_parallel_rank=tp_rank, + ), + adaptation=AdaptationIdentity(), + producer=ProducerIdentity( + "test.final-layout", + "1", + "implementation-sha256", + "test-producer-v1", + "test-restorer-v1", + ), + ) + + +def _weight(offset: int = 0) -> torch.Tensor: + return torch.arange(offset, offset + 12, dtype=torch.float32).to(torch.bfloat16).reshape(4, 3) + + +def _specs() -> tuple[TensorWriteSpec, ...]: + return ( + TensorWriteSpec("layer.weight", (4, 3), torch.bfloat16), + TensorWriteSpec("layer.scale", (2,), torch.float32, kind=TensorKind.BUFFER, role="scale"), + ) + + +def _open_fd_targets() -> tuple[str, ...]: + targets = [] + for path in Path("/proc/self/fd").iterdir(): + with contextlib.suppress(OSError): + targets.append(os.readlink(path)) + return tuple(targets) + + +def _direct_writer(tmp_path: Path) -> FilesystemArtifactWriter: + staging = tmp_path / "staging" + staging.mkdir() + return FilesystemArtifactWriter(staging, capacity_check=lambda _projected, _additional: None) + + +def test_ordered_writer_hashes_payload_during_writes( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + weight = _weight() + scale = torch.tensor([1.0, 2.0]) + + def fail_rescan(_path: Path, _specs: tuple[TensorWriteSpec, ...]) -> object: + raise AssertionError("ordered payload must not be rescanned") + + monkeypatch.setattr(filesystem_writer_module, "_scan_payload", fail_rescan) + with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output: + # The writer's ordered contract follows canonical storage-key order. + output.write_tensor("layer.scale", scale) + output.write_rows("layer.weight", 0, weight[:2]) + output.write_rows("layer.weight", 2, weight[2:]) + + manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + payload = writer.staging_dir / "weights.safetensors" + assert manifest.files[0].sha256 == hashlib.sha256(payload.read_bytes()).hexdigest() + expected_tensor_digests = { + "layer.scale": hashlib.sha256(memoryview(scale.view(torch.uint8).numpy())).hexdigest(), + "layer.weight": hashlib.sha256(memoryview(weight.contiguous().view(torch.uint8).numpy())).hexdigest(), + } + assert {entry.name: entry.sha256 for entry in manifest.tensors} == expected_tensor_digests + + +def test_ordered_writer_rejects_out_of_order_payload(tmp_path: Path) -> None: + writer = _direct_writer(tmp_path) + output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) + + with pytest.raises(ArtifactWriterError, match="ordered tensor payload expected offset"): + output.write_tensor("layer.weight", _weight()) + + output.abort() + assert not list(writer.staging_dir.iterdir()) + + +def test_payload_fsync_overlaps_producer_work( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + started = threading.Event() + release = threading.Event() + closed = threading.Event() + original_fsync = filesystem_writer_module._fsync_payload + + def blocking_fsync(fd: int) -> None: + started.set() + if not release.wait(timeout=5): + raise TimeoutError("test did not release payload fsync") + original_fsync(fd) + + monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", blocking_fsync) + output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) + output.write_tensor("layer.scale", torch.tensor([1.0, 2.0])) + output.write_tensor("layer.weight", _weight()) + close_errors: list[BaseException] = [] + + def close_output() -> None: + try: + output.close() + except BaseException as exc: + close_errors.append(exc) + finally: + closed.set() + + close_thread = threading.Thread(target=close_output) + close_thread.start() + try: + assert started.wait(timeout=2) + assert closed.wait(timeout=1), "payload close waited for fsync instead of overlapping it" + finally: + release.set() + close_thread.join(timeout=5) + + assert not close_errors + writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + +def test_payload_fsync_failure_prevents_publication( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + + def failing_fsync(fd: int) -> None: + try: + raise OSError(errno.EIO, "injected fsync failure") + finally: + os.close(fd) + + monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", failing_fsync) + with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output: + output.write_tensor("layer.scale", torch.tensor([1.0, 2.0])) + output.write_tensor("layer.weight", _weight()) + + with pytest.raises(OSError, match="injected fsync failure"): + writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + assert not (writer.staging_dir / "manifest.json").exists() + assert not (writer.staging_dir / "READY.json").exists() + + +def test_unordered_payload_fallback_hashes_files_in_parallel( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + specs = ( + TensorWriteSpec("a", (2,), torch.float32), + TensorWriteSpec("b", (2,), torch.float32), + ) + for name, spec in zip(("a.safetensors", "b.safetensors"), specs, strict=True): + with writer.open_tensor_file(name, (spec,)) as output: + output.write_tensor(spec.name, torch.tensor([1.0, 2.0])) + + original_scan = filesystem_writer_module._scan_payload + barrier = threading.Barrier(2) + thread_ids: set[int] = set() + thread_ids_lock = threading.Lock() + + def synchronized_scan(path: Path, file_specs: tuple[TensorWriteSpec, ...]) -> object: + with thread_ids_lock: + thread_ids.add(threading.get_ident()) + barrier.wait(timeout=5) + return original_scan(path, file_specs) + + monkeypatch.setattr(filesystem_writer_module, "_scan_payload", synchronized_scan) + manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + assert len(manifest.files) == 2 + assert len(thread_ids) == 2 + + +class FakeProducer: + def __init__( + self, + identity: WeightArtifactIdentity, + *, + counter_path: Path | None = None, + delay_seconds: float = 0.0, + started: object | None = None, + release: object | None = None, + write_mode: str = "rows", + weight_offset: int = 0, + coordination_scope: CoordinationScope = CoordinationScope.SINGLE_PROCESS, + ) -> None: + self._spec = WeightProductionSpec( + producer_id="test.final-layout", + outputs=(identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + coordination_scope=coordination_scope, + ) + self.counter_path = counter_path + self.delay_seconds = delay_seconds + self.started = started + self.release = release + self.write_mode = write_mode + self.weight_offset = weight_offset + + @property + def spec(self) -> WeightProductionSpec: + return self._spec + + def produce(self, writer: object) -> ProductionMetadata: + if self.counter_path is not None: + fd = os.open(self.counter_path, os.O_CREAT | os.O_APPEND | os.O_WRONLY | os.O_CLOEXEC, 0o600) + try: + os.write(fd, f"{os.getpid()}\n".encode()) + os.fsync(fd) + finally: + os.close(fd) + if self.started is not None: + self.started.set() # type: ignore[attr-defined] + if self.release is not None and not self.release.wait(timeout=10): # type: ignore[attr-defined] + raise TimeoutError("test producer was not released") + if self.delay_seconds: + time.sleep(self.delay_seconds) + + with writer.open_tensor_file("weights.safetensors", _specs()) as output: # type: ignore[attr-defined] + if self.write_mode == "rows": + output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:2]) + output.write_rows("layer.weight", 2, _weight(self.weight_offset)[2:]) + elif self.write_mode == "complete": + output.write_tensor("layer.weight", _weight(self.weight_offset)) + elif self.write_mode == "incomplete": + output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:2]) + elif self.write_mode == "overlap": + output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:3]) + output.write_rows("layer.weight", 2, _weight(self.weight_offset)[2:]) + else: # pragma: no cover - test helper guard + raise AssertionError(f"unknown write mode {self.write_mode}") + output.write_tensor("layer.scale", torch.tensor([1.0, 2.0])) + return ProductionMetadata("test-producer-v1", "test-restorer-v1") + + +class CrashProducer(FakeProducer): + def produce(self, writer: object) -> ProductionMetadata: + output = writer.open_tensor_file("weights.safetensors", _specs()) # type: ignore[attr-defined] + output.write_rows("layer.weight", 0, _weight()[:2]) + os._exit(17) + + +class NonCooperativeStore(FilesystemHostWeightStore): + """Test double representing a publisher that does not join the build lock.""" + + def _build_lock_path(self, key: str) -> Path: + return self.locks_dir / f"{key}.noncooperative-build.lock" + + def _remove_stale_temps_locked(self, _key: str) -> None: + return + + +class InitialLookupBarrierStore(FilesystemHostWeightStore): + """Make both test workers observe the initial miss before either builds.""" + + def __init__( + self, + domain: StorageDomainPolicy, + capacity: CapacityPolicy, + integrity: IntegrityPolicy, + *, + initial_lookup_barrier: object, + ) -> None: + super().__init__(domain, capacity, integrity) + self._initial_lookup_barrier = initial_lookup_barrier + self._lookup_count = 0 + + def lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None = None, + ) -> StoreResult: + result = super().lookup(identity, validation=validation, deadline=deadline) + self._lookup_count += 1 + if self._lookup_count == 1: + self._initial_lookup_barrier.wait(timeout=10) # type: ignore[attr-defined] + return result + + +class PausingOpenLeaseStore(FilesystemHostWeightStore): + """Pause after opening a lease so another lookup can publish a deny marker.""" + + def __init__( + self, + domain: StorageDomainPolicy, + capacity: CapacityPolicy, + integrity: IntegrityPolicy, + *, + lease_opened: threading.Event, + resume_lookup: threading.Event, + ) -> None: + super().__init__(domain, capacity, integrity) + self._lease_opened = lease_opened + self._resume_lookup = resume_lookup + + def _open_lease( + self, + identity: WeightArtifactIdentity, + entry_path: Path, + artifact_lock: FileLock, + validation: ValidationLevel, + ) -> HostWeightLease: + lease = super()._open_lease(identity, entry_path, artifact_lock, validation) + self._lease_opened.set() + if not self._resume_lookup.wait(5): + lease.close() + raise TimeoutError("test did not resume paused lookup") + return lease + + +def _make_store( + root: Path, + *, + capacity: CapacityPolicy | None = None, + integrity: IntegrityPolicy | None = None, +) -> FilesystemHostWeightStore: + root.mkdir(parents=True, exist_ok=True, mode=0o700) + domain = StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)) + return FilesystemHostWeightStore(domain, capacity or CapacityPolicy(), integrity or IntegrityPolicy()) + + +def _publish_test_artifact( + store: FilesystemHostWeightStore, + identity: WeightArtifactIdentity | None = None, +) -> tuple[WeightArtifactIdentity, Path]: + identity = identity or _identity() + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert result.status is StoreStatus.BUILT + assert result.lease is not None + result.lease.close() + return identity, store.artifacts_dir / identity.key + + +def _corrupt_test_payload(artifact: Path) -> None: + payload = artifact / "weights.safetensors" + os.chmod(artifact, 0o755) + os.chmod(payload, 0o644) + with payload.open("r+b") as handle: + handle.seek(-1, os.SEEK_END) + value = handle.read(1) + handle.seek(-1, os.SEEK_END) + handle.write(bytes([value[0] ^ 0xFF])) + handle.flush() + os.fsync(handle.fileno()) + os.chmod(payload, 0o444) + os.chmod(artifact, 0o555) + + +def _build_process(root: str, counter: str, initial_lookup_barrier: object, result_queue: object) -> None: + try: + identity = _identity() + root_path = Path(root) + root_path.mkdir(parents=True, exist_ok=True, mode=0o700) + store = InitialLookupBarrierStore( + StorageDomainPolicy(root=root_path, storage_class=detect_storage_class(root_path)), + CapacityPolicy(), + IntegrityPolicy(), + initial_lookup_barrier=initial_lookup_barrier, + ) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, counter_path=Path(counter), delay_seconds=0.25), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + fingerprint = result.lease.provenance.backing_fingerprint if result.lease is not None else None + if result.lease is not None: + result.lease.close() + result_queue.put((result.status.value, fingerprint, None)) # type: ignore[attr-defined] + except BaseException as exc: + result_queue.put(("error", None, repr(exc))) # type: ignore[attr-defined] + + +def _controlled_build_process(root: str, started: object, release: object, result_queue: object) -> None: + try: + identity = _identity() + store = _make_store(Path(root)) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, started=started, release=release), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + if result.lease is not None: + result.lease.close() + result_queue.put((result.status.value, None)) # type: ignore[attr-defined] + except BaseException as exc: + result_queue.put(("error", repr(exc))) # type: ignore[attr-defined] + + +def _pausing_weak_lookup_process(root: str, opened: object, resume: object, result_queue: object) -> None: + try: + policy = StorageDomainPolicy(root=Path(root), storage_class=detect_storage_class(Path(root))) + store = PausingOpenLeaseStore( + policy, + CapacityPolicy(), + IntegrityPolicy(), + lease_opened=opened, # type: ignore[arg-type] + resume_lookup=resume, # type: ignore[arg-type] + ) + result = store.lookup(_identity(), validation=ValidationLevel.MANIFEST_AND_METADATA) + if result.lease is not None: + result.lease.close() + code = result.failure.code.value if result.failure is not None else None + result_queue.put((result.status.value, code, None)) # type: ignore[attr-defined] + except BaseException as exc: + result_queue.put(("error", None, repr(exc))) # type: ignore[attr-defined] + + +def _crash_build_process(root: str) -> None: + identity = _identity() + store = _make_store(Path(root)) + store.get_or_build( + BuildRequest(identity), + CrashProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + + +def _crash_cleanup_process(root: str) -> None: + identity = _identity() + store = _make_store(Path(root)) + + def crash_before_removal(path: Path) -> None: + if path.parent == store.quarantine_dir and ".cleanup." in path.name: + os._exit(23) + FilesystemHostWeightStore._remove_tree(path) + + store._remove_tree = crash_before_removal # type: ignore[method-assign] + store.cleanup(identity) + raise AssertionError("cleanup did not reach the injected crash point") + + +def _crash_cleanup_before_rename_process(root: str) -> None: + identity = _identity() + store = _make_store(Path(root)) + artifact = store.artifacts_dir / identity.key + original_replace = os.replace + + def crash_before_rename( + source: str | os.PathLike[str], + destination: str | os.PathLike[str], + ) -> None: + if Path(source) == artifact and ".cleanup." in Path(destination).name: + os._exit(31) + original_replace(source, destination) + + filesystem_store_module.os.replace = crash_before_rename + store.cleanup(identity) + raise AssertionError("cleanup did not reach the pre-rename crash point") + + +def _crash_quarantine_after_rename_process(root: str) -> None: + identity = _identity() + store = _make_store(Path(root)) + original_harden = filesystem_store_module._harden_open_directory + + def crash_before_destination_hardening(fd: int, path: Path, mode: int) -> None: + if path.parent == store.quarantine_dir: + os._exit(37) + original_harden(fd, path, mode) + + filesystem_store_module._harden_open_directory = crash_before_destination_hardening + with FileLock(store._artifact_lock_path(identity.key), exclusive=True, deadline=None): + store._quarantine_locked(identity.key, reason="validation_failure") + raise AssertionError("quarantine did not reach the post-rename crash point") + + +def _crash_after_artifacts_fsync_process(root: str) -> None: + identity = _identity() + store = _make_store(Path(root)) + original_fsync_directory = filesystem_store_module._fsync_directory + + def crash_after_artifacts_fsync(path: Path) -> None: + original_fsync_directory(path) + if path == store.artifacts_dir: + os._exit(29) + + filesystem_store_module._fsync_directory = crash_after_artifacts_fsync + store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + raise AssertionError("publication did not reach the injected crash point") + + +def _crash_after_publication_hardening_process(root: str) -> None: + identity = _identity() + store = _make_store(Path(root)) + artifact = store.artifacts_dir / identity.key + original_chmod = os.chmod + + def crash_after_hardening( + path: str | os.PathLike[str], + mode: int, + *args: object, + **kwargs: object, + ) -> None: + original_chmod(path, mode, *args, **kwargs) # type: ignore[arg-type] + if Path(path) == artifact and mode == 0o555: + os._exit(41) + + filesystem_store_module.os.chmod = crash_after_hardening + store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 15, + ) + raise AssertionError("publication did not reach the post-hardening crash point") + + +def _join_processes( + processes: Sequence[BaseProcess], + timeout: float = _PROCESS_COMPLETION_TIMEOUT_SECONDS, +) -> None: + deadline = time.monotonic() + timeout + for process in processes: + process.join(max(0.0, deadline - time.monotonic())) + timed_out = [process for process in processes if process.is_alive()] + if not timed_out: + return + timed_out_workers = ", ".join(f"{process.name}(pid={process.pid})" for process in timed_out) + for process in timed_out: + process.terminate() + for process in timed_out: + process.join(_PROCESS_CLEANUP_TIMEOUT_SECONDS) + survivors = [process for process in timed_out if process.is_alive()] + for process in survivors: + process.kill() + process.join(_PROCESS_CLEANUP_TIMEOUT_SECONDS) + unstoppable = [process for process in survivors if process.is_alive()] + assert not unstoppable, "timed-out test worker processes could not be stopped" + pytest.fail(f"test worker processes exceeded the {timeout:g}-second group deadline: {timed_out_workers}") + + +def test_build_lookup_share_one_published_backing_and_lease_blocks_cleanup(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.status is StoreStatus.BUILT + assert built.lease is not None + assert torch.equal(built.lease.tensors["layer.weight"], _weight()) + assert torch.equal(built.lease.tensors["layer.scale"], torch.tensor([1.0, 2.0])) + assert built.lease.mapped_regions + assert all(region.address > 0 and region.length > 0 and region.read_only for region in built.lease.mapped_regions) + with pytest.raises(TypeError, match="process-local"): + built.lease.__reduce_ex__(4) + + second = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert second.status is StoreStatus.HIT + assert second.lease is not None + assert second.lease.provenance.publication_id == built.lease.provenance.publication_id + assert second.lease.provenance.backing_fingerprint == built.lease.provenance.backing_fingerprint + payload = store.artifacts_dir / identity.key / "weights.safetensors" + assert any(str(payload) in target for target in _open_fd_targets()) + + artifact = store.inspect_artifact(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert artifact.status is StoreStatus.HIT + assert artifact.manifest is not None and artifact.manifest.identity == identity + assert artifact.publication_id == built.lease.provenance.publication_id + assert artifact.artifact_lock_active + domain = store.inspect_domain() + assert domain.domain_uuid == built.lease.provenance.domain_uuid + assert domain.filesystem_type == detect_filesystem_type(store.root) + assert any( + item.artifact_key == identity.key and item.state is ArtifactInventoryState.READY and item.size_bytes > 0 + for item in domain.inventory + ) + + active = store.cleanup(identity) + assert active is not None and active.code is FailureCode.ACTIVE_LEASE + second.lease.close() + built.lease.close() + built.lease.close() + assert not built.lease.tensors + assert not any(str(payload) in target for target in _open_fd_targets()) + + assert store.cleanup(identity) is None + assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS + + +def test_publication_hardens_artifact_after_atomic_rename(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert built.status is StoreStatus.BUILT + assert built.lease is not None + built.lease.close() + artifact_dir = store.artifacts_dir / identity.key + assert artifact_dir.stat().st_mode & 0o777 == 0o555 + + +@pytest.mark.parametrize( + ("expected_filesystem", "filesystem_root"), + [("overlay", Path("/tmp")), ("tmpfs", Path("/dev/shm"))], +) +def test_hardened_artifact_cleanup_on_supported_filesystem( + expected_filesystem: str, + filesystem_root: Path, +) -> None: + if not filesystem_root.is_dir(): + pytest.skip(f"{filesystem_root} is unavailable") + try: + filesystem_type = detect_filesystem_type(filesystem_root) + except OSError as exc: + pytest.skip(f"cannot inspect {filesystem_root} filesystem type: {exc}") + if filesystem_type != expected_filesystem: + pytest.skip(f"{filesystem_root} is {filesystem_type!r}, not {expected_filesystem}") + if not os.access(filesystem_root, os.W_OK | os.X_OK): + pytest.skip(f"{filesystem_root} is not writable") + + with tempfile.TemporaryDirectory(prefix="vllm-omni-hwr-", dir=filesystem_root) as directory: + store = _make_store(Path(directory) / "store") + identity, artifact = _publish_test_artifact(store) + assert artifact.stat().st_mode & 0o777 == 0o555 + + assert store.cleanup(identity) is None + assert not artifact.exists() + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +@pytest.mark.parametrize("failure_point", ["chmod", "rename"]) +def test_failed_cleanup_move_preserves_ready_immutability( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_point: str, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + artifact_inode = artifact.stat().st_ino + + if failure_point == "chmod": + original_fchmod = os.fchmod + + def fail_make_writable(fd: int, mode: int) -> None: + if os.fstat(fd).st_ino == artifact_inode and mode & stat.S_IWUSR: + raise OSError(errno.EIO, "injected artifact chmod failure") + original_fchmod(fd, mode) + + monkeypatch.setattr(os, "fchmod", fail_make_writable) + else: + original_replace = os.replace + + def fail_artifact_rename( + source: str | os.PathLike[str], + destination: str | os.PathLike[str], + ) -> None: + if Path(source) == artifact: + raise OSError(errno.EIO, "injected artifact rename failure") + original_replace(source, destination) + + monkeypatch.setattr(os, "replace", fail_artifact_rename) + + failure = store.cleanup(identity) + + assert failure is not None + assert failure.code is FailureCode.QUARANTINE_FAILED + assert failure.retryable + assert artifact.stat().st_mode & 0o777 == 0o555 + assert not list(store.quarantine_dir.iterdir()) + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + + +def test_move_artifact_one_shot_restore_failure_still_hardens_destination( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + destination = store.quarantine_dir / f"{identity.key}.direct-test" + artifact_inode = artifact.stat().st_ino + original_fchmod = os.fchmod + injected = False + + def fail_first_hardening(fd: int, mode: int) -> None: + nonlocal injected + if os.fstat(fd).st_ino == artifact_inode and destination.exists() and not mode & 0o222 and not injected: + injected = True + raise OSError(errno.EIO, "injected artifact hardening failure") + original_fchmod(fd, mode) + + monkeypatch.setattr(os, "fchmod", fail_first_hardening) + with FileLock(store._artifact_lock_path(identity.key), exclusive=True, deadline=None): + filesystem_store_module._move_artifact_locked(artifact, destination) + + assert injected + assert not artifact.exists() + assert destination.is_dir() + assert not destination.stat().st_mode & 0o222 + + +@pytest.mark.parametrize("failing_parent", ["artifacts", "quarantine"]) +def test_cleanup_move_fsync_failure_leaves_retryable_immutable_tombstone( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failing_parent: str, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + original_fsync_directory = filesystem_store_module._fsync_directory + failed_path = store.artifacts_dir if failing_parent == "artifacts" else store.quarantine_dir + + def fail_parent_fsync(path: Path) -> None: + if path == failed_path: + raise OSError(errno.EIO, "injected cleanup fsync failure") + original_fsync_directory(path) + + monkeypatch.setattr(filesystem_store_module, "_fsync_directory", fail_parent_fsync) + failure = store.cleanup(identity) + + assert failure is not None + assert failure.code is FailureCode.QUARANTINE_FAILED + assert failure.retryable + assert not artifact.exists() + tombstones = list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + assert len(tombstones) == 1 + assert tombstones[0].stat().st_mode & 0o777 == 0o555 + assert tombstones[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + + monkeypatch.setattr(filesystem_store_module, "_fsync_directory", original_fsync_directory) + assert store.cleanup(identity) is None + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +def test_cleanup_retries_removal_failure_tombstone( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + original_rmtree = shutil.rmtree + + def fail_cleanup_removal( + path: str | os.PathLike[str], + ) -> None: + if Path(path).parent == store.quarantine_dir and ".cleanup." in Path(path).name: + raise OSError(errno.EIO, "injected cleanup removal failure") + original_rmtree(path) + + monkeypatch.setattr(shutil, "rmtree", fail_cleanup_removal) + failure = store.cleanup(identity) + + assert failure is not None + assert failure.code is FailureCode.QUARANTINE_FAILED + assert failure.retryable + assert not artifact.exists() + tombstones = list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + assert len(tombstones) == 1 + + monkeypatch.setattr(shutil, "rmtree", original_rmtree) + assert store.cleanup(identity) is None + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +@pytest.mark.parallel +def test_cleanup_after_move_crash_tombstone_is_removed_by_next_build(tmp_path: Path) -> None: + root = tmp_path / "store" + store = _make_store(root) + identity, artifact = _publish_test_artifact(store) + process = mp.get_context("spawn").Process(target=_crash_cleanup_process, args=(str(root),)) + + process.start() + _join_processes([process]) + + assert process.exitcode == 23 + assert not artifact.exists() + tombstones = list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + assert len(tombstones) == 1 + assert not tombstones[0].stat().st_mode & 0o222 + + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt.lease.close() + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + assert not (store.deny_dir / f"{identity.key}.json").exists() + + +@pytest.mark.parallel +def test_pre_rename_cleanup_crash_leaves_persistently_denied_artifact(tmp_path: Path) -> None: + root = tmp_path / "store" + store = _make_store(root) + identity, artifact = _publish_test_artifact(store) + process = mp.get_context("spawn").Process( + target=_crash_cleanup_before_rename_process, + args=(str(root),), + ) + + process.start() + _join_processes([process]) + + assert process.exitcode == 31 + assert artifact.is_dir() + assert artifact.stat().st_mode & stat.S_IWUSR + assert artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + weak_store = _make_store( + root, + integrity=IntegrityPolicy(require_trusted_permissions=False), + ) + denied = weak_store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + assert denied.status is StoreStatus.INVALID + assert denied.failure is not None and denied.failure.code is FailureCode.ARTIFACT_DENIED + + rebuilt = weak_store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt.lease.close() + quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*")) + assert len(quarantined) == 1 + assert not quarantined[0].stat().st_mode & 0o222 + assert not quarantined[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert not list(store.tmp_dir.iterdir()) + + +@pytest.mark.parallel +def test_post_rename_quarantine_crash_is_rehardened_by_next_build(tmp_path: Path) -> None: + root = tmp_path / "store" + store = _make_store(root) + identity, artifact = _publish_test_artifact(store) + process = mp.get_context("spawn").Process( + target=_crash_quarantine_after_rename_process, + args=(str(root),), + ) + + process.start() + _join_processes([process]) + + assert process.exitcode == 37 + assert not artifact.exists() + quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*")) + assert len(quarantined) == 1 + assert quarantined[0].stat().st_mode & stat.S_IWUSR + assert quarantined[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt.lease.close() + assert not quarantined[0].stat().st_mode & 0o222 + assert not quarantined[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert not list(store.tmp_dir.iterdir()) + + +@pytest.mark.parametrize("failure_point", ["chmod", "artifact_inode_fsync", "artifacts_fsync"]) +def test_post_rename_failure_quarantines_non_authoritative_publication( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_point: str, +) -> None: + store = _make_store(tmp_path / "store") + identity = _identity() + original_fsync_directory = filesystem_store_module._fsync_directory + injected = False + + if failure_point == "chmod": + original_chmod = os.chmod + + def fail_publication_chmod( + path: str | os.PathLike[str], + mode: int, + *args: object, + **kwargs: object, + ) -> None: + nonlocal injected + if Path(path) == store.artifacts_dir / identity.key and mode == 0o555 and not injected: + injected = True + raise OSError(errno.EIO, "injected publication chmod failure") + original_chmod(path, mode, *args, **kwargs) # type: ignore[arg-type] + + monkeypatch.setattr(os, "chmod", fail_publication_chmod) + else: + failed_path = ( + store.artifacts_dir / identity.key if failure_point == "artifact_inode_fsync" else store.artifacts_dir + ) + + def fail_publication_fsync(path: Path) -> None: + nonlocal injected + if path == failed_path and not injected: + injected = True + raise OSError(errno.EIO, f"injected {failure_point} publication failure") + original_fsync_directory(path) + + monkeypatch.setattr(filesystem_store_module, "_fsync_directory", fail_publication_fsync) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert injected + assert result.status is StoreStatus.FAILED + assert result.failure is not None + assert result.failure.stage is ResolutionStage.LIFECYCLE + assert result.failure.code is FailureCode.PUBLICATION_FAILED + assert not result.failure.retryable + assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS + quarantined = list(store.quarantine_dir.iterdir()) + assert len(quarantined) == 1 + assert ".publication_failure." in quarantined[0].name + assert quarantined[0].stat().st_mode & 0o777 == 0o555 + assert not list(store.tmp_dir.iterdir()) + + monkeypatch.setattr(filesystem_store_module, "_fsync_directory", original_fsync_directory) + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt.lease.close() + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert not list(store.tmp_dir.iterdir()) + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +def test_uncontained_publication_failure_is_retryable_and_outcome_uncertain( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _make_store(tmp_path / "store") + identity = _identity() + artifact = store.artifacts_dir / identity.key + original_fsync_directory = filesystem_store_module._fsync_directory + original_open = os.open + publication_failed = False + + def fail_artifact_inode_fsync(path: Path) -> None: + nonlocal publication_failed + if path == artifact: + publication_failed = True + raise OSError(errno.EIO, "injected publication inode fsync failure") + original_fsync_directory(path) + + def fail_containment_open( + path: str | bytes | os.PathLike[str] | os.PathLike[bytes], + flags: int, + mode: int = 0o777, + *, + dir_fd: int | None = None, + ) -> int: + if publication_failed and Path(os.fsdecode(path)) == artifact: + raise OSError(errno.EIO, "injected persistent containment failure") + return original_open(path, flags, mode, dir_fd=dir_fd) + + with monkeypatch.context() as fault: + fault.setattr(filesystem_store_module, "_fsync_directory", fail_artifact_inode_fsync) + fault.setattr(os, "open", fail_containment_open) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert publication_failed + assert result.status is StoreStatus.FAILED + assert result.failure is not None + assert result.failure.stage is ResolutionStage.DOMAIN + assert result.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert result.failure.retryable + details = result.failure.details.to_value() + assert isinstance(details, dict) + assert details["outcome_uncertain"] is True + assert artifact.is_dir() + assert not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + assert not list(store.quarantine_dir.iterdir()) + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + + +@pytest.mark.parallel +def test_crash_after_publication_hardening_leaves_logically_committed_hit(tmp_path: Path) -> None: + root = tmp_path / "store" + store = _make_store(root) + identity = _identity() + process = mp.get_context("spawn").Process( + target=_crash_after_publication_hardening_process, + args=(str(root),), + ) + + process.start() + _join_processes([process]) + + assert process.exitcode == 41 + artifact = store.artifacts_dir / identity.key + assert artifact.is_dir() + assert not artifact.stat().st_mode & 0o222 + assert not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert not list(store.tmp_dir.iterdir()) + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + assert not list(store.quarantine_dir.iterdir()) + + +@pytest.mark.parallel +def test_crash_after_final_artifacts_fsync_leaves_authoritative_hit(tmp_path: Path) -> None: + root = tmp_path / "store" + store = _make_store(root) + identity = _identity() + process = mp.get_context("spawn").Process(target=_crash_after_artifacts_fsync_process, args=(str(root),)) + + process.start() + _join_processes([process]) + + assert process.exitcode == 29 + artifact = store.artifacts_dir / identity.key + assert artifact.is_dir() + assert not artifact.stat().st_mode & 0o222 + assert not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert not list(store.tmp_dir.iterdir()) + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + + joined = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert joined.status is StoreStatus.HIT + assert joined.lease is not None + joined.lease.close() + assert not list(store.quarantine_dir.iterdir()) + + +@pytest.mark.parametrize("lock_kind", ["artifact", "build"]) +def test_post_commit_lock_release_error_keeps_authoritative_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + lock_kind: str, +) -> None: + store = _make_store(tmp_path / "store") + identity = _identity() + lock_path = ( + store._artifact_lock_path(identity.key) if lock_kind == "artifact" else store._build_lock_path(identity.key) + ) + original_fsync_directory = filesystem_store_module._fsync_directory + original_close = FileLock.close + committed = False + injected = False + + def observe_publication_commit(path: Path) -> None: + nonlocal committed + original_fsync_directory(path) + if path == store.artifacts_dir: + committed = True + + def fail_post_commit_close(lock: FileLock) -> None: + nonlocal injected + original_close(lock) + if committed and lock.path == lock_path and not injected: + injected = True + raise OSError(errno.EIO, f"injected post-commit {lock_kind} lock release failure") + + monkeypatch.setattr(filesystem_store_module, "_fsync_directory", observe_publication_commit) + monkeypatch.setattr(FileLock, "close", fail_post_commit_close) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert committed + assert injected + assert result.status is StoreStatus.BUILT + assert result.lease is not None + result.lease.close() + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + assert not list(store.quarantine_dir.iterdir()) + assert not list(store.tmp_dir.iterdir()) + + +def test_joined_build_lock_release_error_keeps_authoritative_outcome( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _make_store(tmp_path / "store") + identity, _ = _publish_test_artifact(store) + build_lock_path = store._build_lock_path(identity.key) + producer_count = tmp_path / "producer-count" + original_lookup = store.lookup + original_close = FileLock.close + lookup_count = 0 + injected = False + + def miss_once( + lookup_identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None = None, + ) -> StoreResult: + nonlocal lookup_count + lookup_count += 1 + if lookup_count == 1: + return StoreResult(StoreStatus.MISS) + return original_lookup(lookup_identity, validation=validation, deadline=deadline) + + def fail_joined_build_lock_close(lock: FileLock) -> None: + nonlocal injected + original_close(lock) + if lock.path == build_lock_path and not injected: + injected = True + raise OSError(errno.EIO, "injected joined build-lock release failure") + + monkeypatch.setattr(store, "lookup", miss_once) + monkeypatch.setattr(FileLock, "close", fail_joined_build_lock_close) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, counter_path=producer_count), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert lookup_count == 2 + assert injected + assert result.status is StoreStatus.JOINED + assert result.lease is not None + result.lease.close() + assert not producer_count.exists() + + +def test_concurrent_lease_close_waits_for_resource_teardown(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + lease = built.lease + teardown_started = threading.Event() + allow_teardown = threading.Event() + second_started = threading.Event() + second_returned = threading.Event() + errors: list[BaseException] = [] + + def blocking_closer() -> None: + teardown_started.set() + if not allow_teardown.wait(5): + raise TimeoutError("test did not release lease teardown") + + def close_lease(*, started: threading.Event | None = None, returned: threading.Event | None = None) -> None: + if started is not None: + started.set() + try: + lease.close() + except BaseException as exc: + errors.append(exc) + finally: + if returned is not None: + returned.set() + + lease._resource_closers = (*lease._resource_closers, blocking_closer) + first = threading.Thread(target=close_lease) + second = threading.Thread(target=close_lease, kwargs={"started": second_started, "returned": second_returned}) + first.start() + assert teardown_started.wait(5) + second.start() + assert second_started.wait(5) + assert not second_returned.wait(0.2) + + allow_teardown.set() + first.join(5) + second.join(5) + assert not first.is_alive() and not second.is_alive() + assert second_returned.is_set() + assert not errors + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork-specific flock ownership regression") +def test_forked_child_close_does_not_release_parent_lease_lock(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + child_pid = os.fork() + if child_pid == 0: # pragma: no cover - assertions run in the parent + try: + built.lease.close() + except BaseException: + os._exit(1) + os._exit(0) + _, child_status = os.waitpid(child_pid, 0) + assert os.waitstatus_to_exitcode(child_status) == 0 + + active = store.cleanup(identity) + assert active is not None and active.code is FailureCode.ACTIVE_LEASE + built.lease.close() + assert store.cleanup(identity) is None + + +@pytest.mark.parametrize("write_mode", ["incomplete", "overlap"]) +def test_failed_streaming_writes_never_publish_partial_artifacts(tmp_path: Path, write_mode: str) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, write_mode=write_mode), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.FAILED + assert result.failure is not None and result.failure.code is FailureCode.PRODUCER_FAILED + assert not list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + +@pytest.mark.parametrize( + "dtype", + [ + dtype + for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz") + if (dtype := getattr(torch, name, None)) is not None + ], +) +def test_available_float8_formats_round_trip_through_streaming_store(tmp_path: Path, dtype: torch.dtype) -> None: + identity = replace( + _identity(), + representation=WeightRepresentation(name="runtime-fp8", dtype=str(dtype)), + ) + expected = torch.zeros((2, 4), dtype=dtype) + + class Float8Producer: + spec = WeightProductionSpec( + producer_id="test.final-layout", + outputs=(identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + ) + + def produce(self, writer: object) -> ProductionMetadata: + spec = TensorWriteSpec("weight", tuple(expected.shape), dtype) + with writer.open_tensor_file("weights.safetensors", (spec,)) as output: # type: ignore[attr-defined] + output.write_tensor("weight", expected) + return ProductionMetadata("test-producer-v1", "test-restorer-v1") + + store = _make_store(tmp_path / "store") + result = store.get_or_build( + BuildRequest(identity), + Float8Producer(), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.BUILT + assert result.lease is not None + actual = result.lease.tensors["weight"] + assert actual.dtype is dtype + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + result.lease.close() + + +def test_lookup_treats_orphan_deny_marker_as_miss(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + deny = store.deny_dir / f"{identity.key}.json" + deny.write_text("{}", encoding="utf-8") + + result = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert result.status is StoreStatus.MISS + + +def test_overlapping_weaker_lookup_cannot_return_after_artifact_is_denied(tmp_path: Path) -> None: + identity = _identity() + root = tmp_path / "store" + store = _make_store(root) + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + payload = artifact_dir / "weights.safetensors" + os.chmod(artifact_dir, 0o755) + os.chmod(payload, 0o644) + with payload.open("r+b") as handle: + handle.seek(-1, os.SEEK_END) + original = handle.read(1) + handle.seek(-1, os.SEEK_END) + handle.write(bytes([original[0] ^ 0xFF])) + handle.flush() + os.fsync(handle.fileno()) + os.chmod(payload, 0o444) + os.chmod(artifact_dir, 0o555) + + lease_opened = threading.Event() + resume_lookup = threading.Event() + overlapping = PausingOpenLeaseStore( + StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)), + CapacityPolicy(), + IntegrityPolicy(), + lease_opened=lease_opened, + resume_lookup=resume_lookup, + ) + weak_results: list[StoreResult] = [] + + def weak_lookup() -> None: + weak_results.append(overlapping.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA)) + + thread = threading.Thread(target=weak_lookup) + thread.start() + try: + assert lease_opened.wait(5) + strong = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert strong.status is StoreStatus.INVALID + assert (store.deny_dir / f"{identity.key}.json").is_file() + finally: + resume_lookup.set() + thread.join(5) + + assert not thread.is_alive() + assert len(weak_results) == 1 + weak = weak_results[0] + assert weak.status is StoreStatus.INVALID + assert weak.failure is not None and weak.failure.code is FailureCode.ARTIFACT_DENIED + + +@pytest.mark.parametrize( + ("failure_errno", "expected_stage", "expected_code"), + [ + (errno.ENOSPC, ResolutionStage.CAPACITY, FailureCode.ENOSPC), + (errno.EDQUOT, ResolutionStage.CAPACITY, FailureCode.ENOSPC), + (errno.EIO, ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE), + ], +) +@pytest.mark.skipif( + not Path("/proc/self/fd").is_dir(), + reason="fault targeting requires Linux /proc/self/fd", +) +def test_deny_write_failure_surfaces_as_retryable_storage_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_errno: int, + expected_stage: ResolutionStage, + expected_code: FailureCode, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + _corrupt_test_payload(artifact) + deny_path = store.deny_dir / f"{identity.key}.json" + original_write = os.write + + def fail_deny_write(fd: int, data: bytes) -> int: + if Path(os.readlink(f"/proc/self/fd/{fd}")) == deny_path: + raise OSError(failure_errno, "injected deny write failure") + return original_write(fd, data) + + with monkeypatch.context() as fault: + fault.setattr(os, "write", fail_deny_write) + failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None + assert failed.failure.stage is expected_stage + assert failed.failure.code is expected_code + assert failed.failure.retryable + assert deny_path.exists() + assert not artifact.exists() + assert len(list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))) == 1 + + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt.lease.close() + assert not deny_path.exists() + + +def test_deny_open_failure_is_not_reported_as_successful_invalidation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + _corrupt_test_payload(artifact) + deny_path = store.deny_dir / f"{identity.key}.json" + original_open = os.open + + def fail_deny_open( + path: str | bytes | os.PathLike[str] | os.PathLike[bytes], + flags: int, + mode: int = 0o777, + *, + dir_fd: int | None = None, + ) -> int: + if Path(os.fsdecode(path)) == deny_path and flags & os.O_CREAT: + raise OSError(errno.EIO, "injected deny open failure") + return original_open(path, flags, mode, dir_fd=dir_fd) + + with monkeypatch.context() as fault: + fault.setattr(os, "open", fail_deny_open) + failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None + assert failed.failure.stage is ResolutionStage.DOMAIN + assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failed.failure.retryable + assert not deny_path.exists() + assert not artifact.exists() + quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*")) + assert len(quarantined) == 1 + assert not quarantined[0].stat().st_mode & 0o222 + + assert store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA).status is StoreStatus.MISS + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt.lease.close() + assert not deny_path.exists() + + +def test_deny_method_eio_synchronously_excludes_artifact_from_weaker_lookup( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + _corrupt_test_payload(artifact) + + def fail_deny(_key: str, _failure: object) -> None: + raise OSError(errno.EIO, "injected deny publication failure") + + with monkeypatch.context() as fault: + fault.setattr(store, "_write_deny", fail_deny) + failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None + assert failed.failure.stage is ResolutionStage.DOMAIN + assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failed.failure.retryable + assert "injected deny publication failure" in failed.failure.message + assert not artifact.exists() + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA).status is StoreStatus.MISS + quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*")) + assert len(quarantined) == 1 + assert not quarantined[0].stat().st_mode & 0o222 + + +def test_deny_failure_containment_preserves_a_replacement_publication(tmp_path: Path) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + old_inode = (artifact.stat().st_dev, artifact.stat().st_ino) + _corrupt_test_payload(artifact) + shared_lock = FileLock( + store._artifact_lock_path(identity.key), + exclusive=False, + deadline=time.monotonic() + 10, + ) + rebuilt_publications: list[tuple[str, tuple[int, int]]] = [] + + class RebuildOnSharedLockRelease: + def close(self) -> None: + shared_lock.close() + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt_inode = (artifact.stat().st_dev, artifact.stat().st_ino) + rebuilt_publications.append((rebuilt.lease.provenance.publication_id, rebuilt_inode)) + rebuilt.lease.close() + + containment = store._contain_invalid_after_deny_failure( + identity.key, + RebuildOnSharedLockRelease(), # type: ignore[arg-type] + deadline=time.monotonic() + 10, + ) + + assert containment["invalidation_marker"] == "durable" + assert containment["quarantine"] == "already_replaced" + assert len(rebuilt_publications) == 1 + publication_id, rebuilt_inode = rebuilt_publications[0] + assert rebuilt_inode != old_inode + assert (artifact.stat().st_dev, artifact.stat().st_ino) == rebuilt_inode + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + assert hit.lease.provenance.publication_id == publication_id + hit.lease.close() + assert len(list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))) == 1 + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert not list(store.tmp_dir.iterdir()) + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +@pytest.mark.parallel +def test_deny_failure_blocks_overlapping_weaker_lookup_across_processes( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + ctx = mp.get_context("spawn") + root = tmp_path / "store" + store = _make_store(root) + identity, artifact = _publish_test_artifact(store) + _corrupt_test_payload(artifact) + lease_opened = ctx.Event() + resume_lookup = ctx.Event() + result_queue = ctx.Queue() + process = ctx.Process( + target=_pausing_weak_lookup_process, + args=(str(root), lease_opened, resume_lookup, result_queue), + ) + strong_results: list[StoreResult] = [] + + def fail_deny(_key: str, _failure: object) -> None: + raise OSError(errno.EIO, "injected deny publication failure") + + def strong_lookup() -> None: + strong_results.append( + store.lookup( + identity, + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + ) + + process.start() + strong_thread = threading.Thread(target=strong_lookup) + strong_started = False + try: + assert lease_opened.wait(timeout=_PROCESS_START_TIMEOUT_SECONDS) + monkeypatch.setattr(store, "_write_deny", fail_deny) + strong_thread.start() + strong_started = True + marker_deadline = time.monotonic() + 5 + while not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE: + assert time.monotonic() < marker_deadline, "invalidation marker was not published" + time.sleep(0.01) + finally: + resume_lookup.set() + if strong_started: + strong_thread.join(timeout=10) + _join_processes([process]) + + assert not strong_thread.is_alive() + assert process.exitcode == 0 + assert result_queue.get(timeout=3) == ( + StoreStatus.INVALID.value, + FailureCode.ARTIFACT_DENIED.value, + None, + ) + assert len(strong_results) == 1 + strong = strong_results[0] + assert strong.status is StoreStatus.FAILED + assert strong.failure is not None and strong.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert not artifact.exists() + assert len(list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))) == 1 + + +@pytest.mark.parametrize("failure_point", ["entry_open", "ready_read", "payload_hash"]) +def test_operational_lookup_eio_is_retryable_without_denying_valid_artifact( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_point: str, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + + with monkeypatch.context() as fault: + if failure_point == "entry_open": + original_open = os.open + + def fail_entry_open( + path: str | bytes | os.PathLike[str] | os.PathLike[bytes], + flags: int, + mode: int = 0o777, + *, + dir_fd: int | None = None, + ) -> int: + if Path(os.fsdecode(path)) == artifact: + raise OSError(errno.EIO, "injected artifact open failure") + return original_open(path, flags, mode, dir_fd=dir_fd) + + fault.setattr(os, "open", fail_entry_open) + elif failure_point == "ready_read": + original_read_json_at = filesystem_store_module._read_json_at + + def fail_ready_read( + directory_fd: int, + name: str, + *, + require_read_only: bool = False, + ) -> dict[str, object]: + if name == "READY.json": + raise OSError(errno.EIO, "injected READY metadata read failure") + return original_read_json_at( + directory_fd, + name, + require_read_only=require_read_only, + ) + + fault.setattr(filesystem_store_module, "_read_json_at", fail_ready_read) + else: + + def fail_hash(_fd: int, _size: int) -> str: + raise OSError(errno.EIO, "injected hash failure") + + fault.setattr(filesystem_store_module, "_sha256_fd", fail_hash) + failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None + assert failed.failure.stage is ResolutionStage.DOMAIN + assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failed.failure.retryable + assert not (store.deny_dir / f"{identity.key}.json").exists() + + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + + +@pytest.mark.parametrize("failure_point", ["entry_stat", "deny_stat"]) +@pytest.mark.skipif( + not Path("/proc/self/fd").is_dir(), + reason="lock descriptor assertions require Linux /proc/self/fd", +) +def test_lookup_state_probe_eio_is_typed_and_releases_artifact_lock( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_point: str, +) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + deny_path = store.deny_dir / f"{identity.key}.json" + target = artifact if failure_point == "entry_stat" else deny_path + artifact_lock_target = str(store._artifact_lock_path(identity.key)) + lock_fds_before = _open_fd_targets().count(artifact_lock_target) + original_lstat = Path.lstat + injected = False + + def fail_state_probe(path: Path) -> os.stat_result: + nonlocal injected + if path == target and not injected: + injected = True + raise OSError(errno.EIO, f"injected {failure_point} failure") + return original_lstat(path) + + with monkeypatch.context() as fault: + fault.setattr(Path, "lstat", fail_state_probe) + failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert injected + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None + assert failed.failure.stage is ResolutionStage.DOMAIN + assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failed.failure.retryable + assert _open_fd_targets().count(artifact_lock_target) == lock_fds_before + assert not deny_path.exists() + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + + +def test_full_checksum_denies_quarantines_and_rebuilds_corrupt_payload(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + rebuild_count = tmp_path / "rebuild-count" + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + _corrupt_test_payload(artifact_dir) + + metadata_only = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + assert metadata_only.status is StoreStatus.HIT + assert metadata_only.lease is not None + metadata_only.lease.close() + + inspected = store.inspect_artifact(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert inspected.status is StoreStatus.INVALID + assert inspected.failure is not None and inspected.failure.code is FailureCode.CHECKSUM_MISMATCH + assert not (store.deny_dir / f"{identity.key}.json").exists() + + invalid = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert invalid.status is StoreStatus.INVALID + assert invalid.failure is not None and invalid.failure.code is FailureCode.CHECKSUM_MISMATCH + assert (store.deny_dir / f"{identity.key}.json").is_file() + + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, counter_path=rebuild_count), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + assert torch.equal(rebuilt.lease.tensors["layer.weight"], _weight()) + rebuilt.lease.close() + assert rebuild_count.read_text(encoding="utf-8").count("\n") == 1 + assert not (store.deny_dir / f"{identity.key}.json").exists() + quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*")) + assert len(quarantined) == 1 + assert not quarantined[0].stat().st_mode & 0o222 + assert not list(store.tmp_dir.iterdir()) + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +def test_pre_writable_crash_state_is_quarantined_immutably_and_rebuilt(tmp_path: Path) -> None: + store = _make_store(tmp_path / "store") + identity, artifact = _publish_test_artifact(store) + rebuild_count = tmp_path / "rebuild-count" + os.chmod(artifact, 0o755) + + invalid = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + + assert invalid.status is StoreStatus.INVALID + assert invalid.failure is not None and invalid.failure.code is FailureCode.MANIFEST_CORRUPT + assert (store.deny_dir / f"{identity.key}.json").is_file() + + rebuilt = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, counter_path=rebuild_count), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert rebuilt.status is StoreStatus.BUILT + assert rebuilt.lease is not None + rebuilt.lease.close() + assert rebuild_count.read_text(encoding="utf-8").count("\n") == 1 + assert not (store.artifacts_dir / identity.key).stat().st_mode & 0o222 + quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*")) + assert len(quarantined) == 1 + assert not quarantined[0].stat().st_mode & 0o222 + assert not (store.deny_dir / f"{identity.key}.json").exists() + assert not list(store.tmp_dir.iterdir()) + assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*")) + + +def test_malformed_safetensors_is_a_typed_invalid_artifact(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + payload = artifact_dir / "weights.safetensors" + os.chmod(artifact_dir, 0o755) + os.chmod(payload, 0o644) + with payload.open("r+b") as handle: + handle.write(b"\xff" * 8) + handle.flush() + os.fsync(handle.fileno()) + os.chmod(payload, 0o444) + os.chmod(artifact_dir, 0o555) + + result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + + assert result.status is StoreStatus.INVALID + assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_CORRUPT + assert (store.deny_dir / f"{identity.key}.json").is_file() + + +@pytest.mark.parametrize("schema_field", ["producer_schema", "restorer_schema"]) +def test_lookup_rejects_manifest_schema_that_differs_from_identity( + tmp_path: Path, + schema_field: str, +) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + manifest_path = artifact_dir / "manifest.json" + marker_path = artifact_dir / "READY.json" + manifest_document = json.loads(manifest_path.read_bytes()) + marker_document = json.loads(marker_path.read_bytes()) + manifest_document[schema_field] = "incompatible-schema" + manifest_bytes = json.dumps( + manifest_document, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + marker_document["manifest_sha256"] = hashlib.sha256(manifest_bytes).hexdigest() + marker_bytes = json.dumps( + marker_document, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + + os.chmod(artifact_dir, 0o755) + for path, contents in ((manifest_path, manifest_bytes), (marker_path, marker_bytes)): + os.chmod(path, 0o644) + path.write_bytes(contents) + os.chmod(path, 0o444) + os.chmod(artifact_dir, 0o555) + + result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + + assert result.status is StoreStatus.INVALID + assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_INCOMPATIBLE + assert (store.deny_dir / f"{identity.key}.json").is_file() + + +def test_noncooperative_conflicting_publication_is_reported_without_replacement(tmp_path: Path) -> None: + identity = _identity() + root = tmp_path / "store" + store = _make_store(root) + started = threading.Event() + release = threading.Event() + thread_result: list[StoreResult] = [] + + def build_expected_content() -> None: + thread_result.append( + store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, started=started, release=release), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + ) + + builder = threading.Thread(target=build_expected_content) + builder.start() + try: + assert started.wait(timeout=5) + competing = NonCooperativeStore(store.domain_policy, store.capacity_policy, store.integrity_policy) + published = competing.get_or_build( + BuildRequest(identity), + FakeProducer(identity, weight_offset=100), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert published.status is StoreStatus.BUILT + assert published.lease is not None + published.lease.close() + finally: + release.set() + builder.join(timeout=10) + assert not builder.is_alive() + assert len(thread_result) == 1 + collision = thread_result[0] + assert collision.status is StoreStatus.FAILED + assert collision.failure is not None and collision.failure.code is FailureCode.IDENTITY_COLLISION + + winner = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert winner.status is StoreStatus.HIT + assert winner.lease is not None + assert torch.equal(winner.lease.tensors["layer.weight"], _weight(100)) + winner.lease.close() + assert len(list(store.quarantine_dir.iterdir())) == 1 + + +def test_same_semantic_competing_publication_remains_authoritative_on_staging_cleanup_eio( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + root = tmp_path / "store" + store = _make_store(root) + injected = False + original_remove_tree = FilesystemHostWeightStore._remove_tree + + class CompetingProducer(FakeProducer): + def produce(self, writer: object) -> ProductionMetadata: + competing = NonCooperativeStore(store.domain_policy, store.capacity_policy, store.integrity_policy) + published = competing.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert published.status is StoreStatus.BUILT + assert published.lease is not None + published.lease.close() + return super().produce(writer) + + def fail_primary_staging_cleanup_once(path: Path) -> None: + nonlocal injected + if path.parent == store.tmp_dir and not injected: + injected = True + raise OSError(errno.EIO, "injected same-semantic staging cleanup failure") + original_remove_tree(path) + + monkeypatch.setattr(store, "_remove_tree", fail_primary_staging_cleanup_once) + result = store.get_or_build( + BuildRequest(identity), + CompetingProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert injected + assert result.status is StoreStatus.BUILT + assert result.lease is not None + result.lease.close() + assert not list(store.tmp_dir.iterdir()) + assert not list(store.deny_dir.iterdir()) + assert not list(store.quarantine_dir.iterdir()) + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + + +def test_competing_publication_manifest_eio_is_retryable_and_preserves_entry( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + root = tmp_path / "store" + store = _make_store(root) + artifact = store.artifacts_dir / identity.key + manifest_path = artifact / "manifest.json" + original_read_json_file = filesystem_store_module._read_json_file + injected = False + + class CompetingProducer(FakeProducer): + def produce(self, writer: object) -> ProductionMetadata: + competing = NonCooperativeStore(store.domain_policy, store.capacity_policy, store.integrity_policy) + published = competing.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert published.status is StoreStatus.BUILT + assert published.lease is not None + published.lease.close() + return super().produce(writer) + + def fail_competing_manifest_read_once(path: Path) -> dict[str, object]: + nonlocal injected + if path == manifest_path and not injected: + injected = True + raise OSError(errno.EIO, "injected competing manifest read failure") + return original_read_json_file(path) + + monkeypatch.setattr(filesystem_store_module, "_read_json_file", fail_competing_manifest_read_once) + result = store.get_or_build( + BuildRequest(identity), + CompetingProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert injected + assert result.status is StoreStatus.FAILED + assert result.failure is not None + assert result.failure.stage is ResolutionStage.DOMAIN + assert result.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert result.failure.retryable + details = result.failure.details.to_value() + assert isinstance(details, dict) + assert details["outcome_uncertain"] is True + assert details["competing_entry"] == "preserved" + assert artifact.is_dir() + assert not list(store.quarantine_dir.iterdir()) + assert not list(store.tmp_dir.iterdir()) + hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM) + assert hit.status is StoreStatus.HIT + assert hit.lease is not None + hit.lease.close() + + +def test_payload_symlink_is_rejected_without_following_it(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + built = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert built.lease is not None + built.lease.close() + + artifact_dir = store.artifacts_dir / identity.key + payload = artifact_dir / "weights.safetensors" + outside = tmp_path / "outside.safetensors" + os.chmod(artifact_dir, 0o755) + payload.rename(outside) + payload.symlink_to(outside) + os.chmod(artifact_dir, 0o555) + + result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA) + assert result.status is StoreStatus.INVALID + assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_CORRUPT + + +def test_coordination_deadline_does_not_claim_to_cancel_an_active_producer(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + started = time.monotonic() + + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, delay_seconds=0.15), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 0.02, + ) + + assert time.monotonic() - started >= 0.1 + assert result.status is StoreStatus.BUILT + assert result.lease is not None + result.lease.close() + + +@pytest.mark.parallel +def test_multiprocess_contention_runs_one_builder_and_joins_same_artifact(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + root = tmp_path / "store" + _make_store(root) + counter = tmp_path / "producer-count" + initial_lookup_barrier = ctx.Barrier(2) + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=_build_process, + args=(str(root), str(counter), initial_lookup_barrier, result_queue), + ) + for _ in range(2) + ] + for process in processes: + process.start() + _join_processes(processes) + + results = [result_queue.get(timeout=3) for _ in processes] + assert all(process.exitcode == 0 for process in processes) + assert sorted(item[0] for item in results) == [StoreStatus.BUILT.value, StoreStatus.JOINED.value] + assert all(item[2] is None for item in results) + assert len({item[1] for item in results}) == 1 + assert counter.read_text(encoding="utf-8").count("\n") == 1 + + +@pytest.mark.parallel +def test_reader_sees_miss_during_build_and_waiter_has_bounded_timeout(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + root = tmp_path / "store" + store = _make_store(root) + identity = _identity() + started = ctx.Event() + release = ctx.Event() + result_queue = ctx.Queue() + process = ctx.Process( + target=_controlled_build_process, + args=(str(root), started, release, result_queue), + ) + process.start() + try: + assert started.wait(timeout=_PROCESS_START_TIMEOUT_SECONDS) + assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS + timed_out = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 0.1, + ) + assert timed_out.status is StoreStatus.TIMEOUT + assert timed_out.failure is not None and timed_out.failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT + finally: + release.set() + _join_processes([process]) + + status, error = result_queue.get(timeout=3) + assert process.exitcode == 0 + assert (status, error) == (StoreStatus.BUILT.value, None) + + +@pytest.mark.parallel +def test_crashed_builder_leaves_only_invisible_temp_and_next_builder_recovers(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + root = tmp_path / "store" + store = _make_store(root) + process = ctx.Process(target=_crash_build_process, args=(str(root),)) + process.start() + _join_processes([process]) + + assert process.exitcode == 17 + assert list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + identity = _identity() + recovered = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert recovered.status is StoreStatus.BUILT + assert recovered.lease is not None + recovered.lease.close() + assert not list(store.tmp_dir.iterdir()) + + +def test_capacity_and_enospc_failures_are_typed_and_clean_staging( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + identity = _identity() + limited = _make_store(tmp_path / "limited", capacity=CapacityPolicy(max_artifact_bytes=64)) + too_large = limited.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert too_large.status is StoreStatus.FAILED + assert too_large.failure is not None and too_large.failure.code is FailureCode.ARTIFACT_TOO_LARGE + assert not list(limited.tmp_dir.iterdir()) + + no_space = _make_store(tmp_path / "enospc") + + def raise_enospc(_fd: int, _offset: int, _length: int) -> None: + raise OSError(errno.ENOSPC, "injected no space") + + monkeypatch.setattr(os, "posix_fallocate", raise_enospc) + failed = no_space.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None and failed.failure.code is FailureCode.ENOSPC + assert not list(no_space.tmp_dir.iterdir()) + assert not list(no_space.artifacts_dir.iterdir()) + + +@pytest.mark.parametrize( + ("failure_errno", "expected_stage", "expected_code"), + [ + (errno.ENOSPC, ResolutionStage.CAPACITY, FailureCode.ENOSPC), + (errno.EDQUOT, ResolutionStage.CAPACITY, FailureCode.ENOSPC), + (errno.EACCES, ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE), + (errno.EIO, ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE), + ], +) +def test_staging_directory_failures_are_typed_and_clean( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_errno: int, + expected_stage: ResolutionStage, + expected_code: FailureCode, +) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + original_mkdir = Path.mkdir + + def fail_staging_mkdir( + path: Path, + mode: int = 0o777, + parents: bool = False, + exist_ok: bool = False, + ) -> None: + if path.parent == store.tmp_dir: + raise OSError(failure_errno, "injected staging directory failure") + original_mkdir(path, mode=mode, parents=parents, exist_ok=exist_ok) + + monkeypatch.setattr(Path, "mkdir", fail_staging_mkdir) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.FAILED + assert result.failure is not None + assert result.failure.stage is expected_stage + assert result.failure.code is expected_code + assert result.failure.retryable + assert not list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + +def test_post_producer_staging_io_failure_is_retryable_domain_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + original_fsync_directory = filesystem_store_module._fsync_directory + + def fail_staging_fsync(path: Path) -> None: + if path.parent == store.tmp_dir: + raise OSError(errno.EIO, "injected post-producer staging fsync failure") + original_fsync_directory(path) + + monkeypatch.setattr(filesystem_store_module, "_fsync_directory", fail_staging_fsync) + result = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.FAILED + assert result.failure is not None + assert result.failure.stage is ResolutionStage.DOMAIN + assert result.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert result.failure.retryable + assert not list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + +def test_async_payload_fsync_eio_is_retryable_domain_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + injection_count = 0 + + def fail_payload_fsync(fd: int) -> None: + nonlocal injection_count + injection_count += 1 + try: + raise OSError(errno.EIO, "injected payload-fsync EIO") + finally: + os.close(fd) + + with monkeypatch.context() as fault: + fault.setattr(filesystem_writer_module, "_fsync_payload", fail_payload_fsync) + failed = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert injection_count == 1 + assert failed.status is StoreStatus.FAILED + assert failed.failure is not None + assert failed.failure.stage is ResolutionStage.DOMAIN + assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failed.failure.retryable + assert "injected payload-fsync EIO" in failed.failure.message + assert not list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + recovered = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert recovered.status is StoreStatus.BUILT + assert recovered.lease is not None + recovered.lease.close() + + +def test_producer_raised_eio_remains_nonretryable_producer_failure(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + class FailingProducer(FakeProducer): + def produce(self, writer: object) -> ProductionMetadata: + del writer + raise OSError(errno.EIO, "injected producer failure") + + result = store.get_or_build( + BuildRequest(identity), + FailingProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + + assert result.status is StoreStatus.FAILED + assert result.failure is not None + assert result.failure.stage is ResolutionStage.PRODUCTION + assert result.failure.code is FailureCode.PRODUCER_FAILED + assert not result.failure.retryable + assert not list(store.tmp_dir.iterdir()) + assert not list(store.artifacts_dir.iterdir()) + + +def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path: Path) -> None: + identity = _identity() + store_limited = _make_store(tmp_path / "store-limit", capacity=CapacityPolicy(max_store_bytes=1)) + store_result = store_limited.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert store_result.status is StoreStatus.FAILED + assert store_result.failure is not None and store_result.failure.code is FailureCode.STORE_LIMIT_EXCEEDED + assert not list(store_limited.artifacts_dir.iterdir()) + + free_root = tmp_path / "free-limit" + free_root.mkdir(mode=0o700) + # Leave enough margin for filesystem accounting to move between the + # preflight measurement and the capacity check. + unavailable_reserve = shutil.disk_usage(free_root).free + (1 << 30) + free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve)) + free_result = free_limited.get_or_build( + BuildRequest(identity), + FakeProducer(identity), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert free_result.status is StoreStatus.FAILED + assert free_result.failure is not None and free_result.failure.code is FailureCode.INSUFFICIENT_CAPACITY + assert not list(free_limited.artifacts_dir.iterdir()) + + +def test_domain_capacity_policy_is_authoritative(tmp_path: Path) -> None: + root = tmp_path / "store" + _make_store(root, capacity=CapacityPolicy(max_store_bytes=4096)) + + with pytest.raises(HostWeightError, match="capacity policy differs"): + _make_store(root, capacity=CapacityPolicy(max_store_bytes=8192)) + + +@pytest.mark.parametrize("filesystem_type", ["nfs4", "cifs", "lustre", "ceph", "unknownfs"]) +def test_remote_and_unknown_filesystems_cannot_back_the_local_store( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + filesystem_type: str, +) -> None: + monkeypatch.setattr(filesystem_store_module, "detect_filesystem_type", lambda _path: filesystem_type) + root = tmp_path / "store" + + with pytest.raises(ValueError, match="not node-local|not recognized"): + detect_storage_class(root) + with pytest.raises(HostWeightError) as error: + FilesystemHostWeightStore( + StorageDomainPolicy(root=root, storage_class=StorageClass.DISK), + CapacityPolicy(), + IntegrityPolicy(), + ) + + assert error.value.failure.code is FailureCode.DOMAIN_POLICY_MISMATCH + assert not root.exists() + + +def test_trusted_store_rejects_group_or_world_writable_root(tmp_path: Path) -> None: + root = tmp_path / "untrusted" + root.mkdir(mode=0o700) + root.chmod(0o777) + + with pytest.raises(HostWeightError) as error: + FilesystemHostWeightStore( + StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)), + CapacityPolicy(), + IntegrityPolicy(), + ) + + assert error.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE + + +def test_unimplemented_coordination_and_fs_verity_return_typed_failures(tmp_path: Path) -> None: + identity = _identity() + store = _make_store(tmp_path / "store") + + unsupported_producer = store.get_or_build( + BuildRequest(identity), + FakeProducer(identity, coordination_scope=CoordinationScope.GROUP_COHORT), + validation=ValidationLevel.FULL_CHECKSUM, + deadline=time.monotonic() + 10, + ) + assert unsupported_producer.status is StoreStatus.FAILED + assert unsupported_producer.failure is not None + assert unsupported_producer.failure.code is FailureCode.PRODUCER_UNSUPPORTED + + fs_verity = store.lookup(identity, validation=ValidationLevel.FS_VERITY) + assert fs_verity.status is StoreStatus.FAILED + assert fs_verity.failure is not None and fs_verity.failure.code is FailureCode.UNSUPPORTED + + +def test_unimplemented_numa_enforcement_and_background_scrub_fail_explicitly(tmp_path: Path) -> None: + numa_root = tmp_path / "numa" + numa_root.mkdir(mode=0o700) + with pytest.raises(HostWeightError) as numa_error: + FilesystemHostWeightStore( + StorageDomainPolicy( + root=numa_root, + scope=StorageScope.NUMA, + storage_class=detect_storage_class(numa_root), + numa_prefault=True, + ), + CapacityPolicy(), + IntegrityPolicy(), + ) + assert numa_error.value.failure.code is FailureCode.UNSUPPORTED + + with pytest.raises(HostWeightError) as scrub_error: + _make_store(tmp_path / "scrub", integrity=IntegrityPolicy(background_scrub=True)) + assert scrub_error.value.failure.code is FailureCode.UNSUPPORTED diff --git a/tests/host_weight_runtime/test_runtime_resolution.py b/tests/host_weight_runtime/test_runtime_resolution.py new file mode 100644 index 000000000..21bab37de --- /dev/null +++ b/tests/host_weight_runtime/test_runtime_resolution.py @@ -0,0 +1,632 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Runtime resolution and publication-policy tests for Host Weight Runtime.""" + +from __future__ import annotations + +from pathlib import Path +from typing import cast + +import pytest +import torch + +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + ComponentIdentity, + FailureCode, + HostWeightFailure, + HostWeightLease, + HostWeightLeaseCarrier, + HostWeightRuntime, + HostWeightRuntimeConfig, + IntegrityPolicy, + LookupPhase, + PostLoadPublicationOutcome, + PostLoadPublicationReport, + ProducerIdentity, + ProductionMetadata, + ProductionPolicy, + ProductionSourceMode, + RemoteImportPolicy, + RemoteOnMiss, + ResolutionOutcome, + ResolutionReport, + ResolutionStage, + RuntimeMode, + RuntimeWeightLayout, + StorageDomainPolicy, + StoreResult, + StoreStatus, + TensorWriteSpec, + ValidationLevel, + WeightArtifactIdentity, + WeightProductionSpec, + WeightRepresentation, + WeightSourceIdentity, +) +from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _identity() -> WeightArtifactIdentity: + return WeightArtifactIdentity( + schema_version=1, + source=WeightSourceIdentity("test-org/test-model", "0123456789abcdef", "source-files-sha256"), + component=ComponentIdentity("transformer", "diffusion.dit"), + representation=WeightRepresentation("runtime-bf16", "torch.bfloat16"), + layout=RuntimeWeightLayout("final-module-layout"), + adaptation=AdaptationIdentity(), + producer=ProducerIdentity( + "test.final-layout", + "1", + "implementation-sha256", + "test-producer-v1", + "test-restorer-v1", + ), + ) + + +class CountingProducer: + def __init__( + self, + identity: WeightArtifactIdentity, + *, + lookup_phase: LookupPhase = LookupPhase.PRE_LOAD_SAFE, + ) -> None: + self.calls = 0 + self._spec = WeightProductionSpec( + producer_id="test.final-layout", + outputs=(identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + lookup_phase=lookup_phase, + ) + + @property + def spec(self) -> WeightProductionSpec: + return self._spec + + def produce(self, writer: object) -> ProductionMetadata: + self.calls += 1 + spec = TensorWriteSpec("weight", (2, 2), torch.bfloat16) + with writer.open_tensor_file("weights.safetensors", (spec,)) as output: # type: ignore[attr-defined] + output.write_tensor("weight", torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)) + return ProductionMetadata( + producer_schema="test-producer-v1", + restorer_schema="test-restorer-v1", + format_metadata=CanonicalJson.from_value({"source": "test"}), + ) + + +def _domain(root: Path) -> StorageDomainPolicy: + return StorageDomainPolicy(root=root, storage_class=detect_storage_class(root.parent)) + + +def test_disabled_mode_does_not_probe_or_create_configured_root(tmp_path: Path) -> None: + root = tmp_path / "must-not-exist" + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.DISABLED, domain=_domain(root))) + + resolution = runtime.resolve() + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_DIRECT + assert resolution.report.attempts == () + assert resolution.lease is None + assert not root.exists() + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_local_miss_uses_configured_fallback_policy( + tmp_path: Path, + mode: RuntimeMode, + expected: ResolutionOutcome, +) -> None: + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path / mode.value))) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is expected + assert resolution.lease is None + assert [attempt.result.value for attempt in resolution.report.attempts] == ["miss", "skipped"] + expected_action = "canonical_fallback" if mode is RuntimeMode.PREFERRED else "fail_startup" + assert [attempt.action.value for attempt in resolution.report.attempts] == [expected_action, expected_action] + + +def test_local_production_then_exact_warm_hit_emits_one_terminal_report_each(tmp_path: Path) -> None: + reports: list[ResolutionReport] = [] + identity = _identity() + producer = CountingProducer(identity) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")), + observer=reports.append, + ) + + cold = runtime.resolve(identity, producer=producer) + assert cold.report.outcome is ResolutionOutcome.LOCAL_PRODUCTION + assert cold.lease is not None + assert torch.equal( + cold.lease.tensors["weight"], + torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2), + ) + cold.lease.close() + assert cold.report.attempts[0].action.value == "try_producer" + + warm = runtime.resolve(identity) + assert warm.report.outcome is ResolutionOutcome.LOCAL_HIT + assert warm.lease is not None + warm.lease.close() + + assert producer.calls == 1 + assert reports == [cold.report, warm.report] + assert all(report.resolution_id for report in reports) + + +def test_lease_carrier_is_single_take_and_process_local(tmp_path: Path) -> None: + identity = _identity() + producer = CountingProducer(identity) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")) + ) + resolution = runtime.resolve(identity, producer=producer) + assert resolution.lease is not None + + carrier = HostWeightLeaseCarrier(resolution.lease) + with pytest.raises(TypeError, match="cannot be serialized"): + carrier.__reduce_ex__(4) + + lease = carrier.take() + assert lease is resolution.lease + with pytest.raises(RuntimeError, match="already taken"): + carrier.take() + carrier.close() # ownership has moved; the carrier must not close twice. + assert not lease.closed + lease.close() + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.FAILED), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_remote_required_is_explicitly_unsupported_and_skips_local_production( + tmp_path: Path, + mode: RuntimeMode, + expected: ResolutionOutcome, +) -> None: + identity = _identity() + producer = CountingProducer(identity) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=mode, + domain=_domain(tmp_path / mode.value), + remote=RemoteImportPolicy( + on_local_miss=RemoteOnMiss.REQUIRE, + providers=("future-mooncake",), + ), + ) + ) + + resolution = runtime.resolve(identity, producer=producer) + + assert resolution.report.outcome is expected + assert producer.calls == 0 + assert [attempt.stage.value for attempt in resolution.report.attempts] == ["lookup", "remote"] + assert resolution.report.attempts[-1].failure is not None + assert resolution.report.attempts[-1].failure.code.value == "unsupported" + assert resolution.report.attempts[-1].action.value == "fail_startup" + + +def test_observer_failure_cannot_change_resolution_outcome(tmp_path: Path) -> None: + def broken_observer(_report: object) -> None: + raise RuntimeError("metrics backend unavailable") + + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")), + observer=broken_observer, + ) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + + +def test_preload_resolution_does_not_run_postload_only_producer(tmp_path: Path) -> None: + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")) + ) + + resolution = runtime.resolve(identity, producer=producer) + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + assert producer.calls == 0 + assert resolution.report.attempts[0].action.value == "canonical_fallback" + + +def test_postload_publication_closes_lease_and_only_warms_future_resolution( + tmp_path: Path, +) -> None: + publications: list[PostLoadPublicationReport] = [] + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + production=ProductionPolicy( + allow_local_build=False, + allow_post_load_publish=True, + ), + ), + publication_observer=publications.append, + ) + + preload = runtime.resolve(identity, producer=producer) + assert preload.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + assert producer.calls == 0 + + publication = runtime.publish_after_load(identity, producer=producer) + assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED + assert publication.identity_digest == identity.key + + already_present = runtime.publish_after_load(identity, producer=producer) + assert already_present.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT + + warm = runtime.resolve(identity) + assert warm.report.outcome is ResolutionOutcome.LOCAL_HIT + assert warm.lease is not None + assert torch.equal( + warm.lease.tensors["weight"], + torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2), + ) + warm.lease.close() + assert producer.calls == 1 + assert publications == [publication, already_present] + assert isinstance(runtime.store, FilesystemHostWeightStore) + assert runtime.store.cleanup(identity) is None + + +def test_postload_publication_policy_skip_does_not_run_producer(tmp_path: Path) -> None: + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")) + ) + + publication = runtime.publish_after_load(identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.POLICY_DISABLED + assert producer.calls == 0 + + +def test_postload_publication_reports_lease_cleanup_failure( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + production=ProductionPolicy(allow_post_load_publish=True), + ) + ) + assert runtime.store is not None + + class FailingLease: + def close(self) -> None: + raise RuntimeError("injected lease cleanup failure") + + def return_failing_lease(*_args: object, **_kwargs: object) -> StoreResult: + return StoreResult(StoreStatus.BUILT, lease=cast(HostWeightLease, FailingLease())) + + monkeypatch.setattr(runtime.store, "get_or_build", return_failing_lease) + + publication = runtime.publish_after_load(identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.FAILED + assert publication.failure is not None + assert publication.failure.stage is ResolutionStage.LIFECYCLE + assert publication.failure.code is FailureCode.PUBLICATION_FAILED + assert publication.failure.details.to_value() == { + "exception_type": "RuntimeError", + "store_outcome": "built", + } + + +def test_disabled_postload_publication_does_not_probe_configured_root(tmp_path: Path) -> None: + root = tmp_path / "must-not-exist" + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.DISABLED, + domain=_domain(root), + production=ProductionPolicy(allow_post_load_publish=True), + ) + ) + + publication = runtime.publish_after_load() + + assert publication.outcome is PostLoadPublicationOutcome.RUNTIME_DISABLED + assert not root.exists() + + +def test_postload_publication_rejects_preload_producer(tmp_path: Path) -> None: + identity = _identity() + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + production=ProductionPolicy(allow_post_load_publish=True), + ) + ) + + with pytest.raises(ValueError, match="POST_LOAD_ONLY"): + runtime.publish_after_load(identity, producer=CountingProducer(identity)) + + +def test_postload_publication_failure_is_reported_without_raising( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + production=ProductionPolicy(allow_post_load_publish=True), + ) + ) + assert runtime.store is not None + failure = HostWeightFailure( + stage=ResolutionStage.CAPACITY, + code=FailureCode.ENOSPC, + retryable=True, + message="injected post-load publication failure", + details=CanonicalJson.empty(), + ) + + def fail_publication(*_args: object, **_kwargs: object) -> StoreResult: + return StoreResult(StoreStatus.FAILED, failure=failure) + + monkeypatch.setattr(runtime.store, "get_or_build", fail_publication) + publication = runtime.publish_after_load(identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.FAILED + assert publication.failure is failure + assert producer.calls == 0 + + +def test_postload_publication_reports_store_initialization_failure(tmp_path: Path) -> None: + root = tmp_path / "not-a-directory" + root.write_text("occupied", encoding="utf-8") + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(root), + production=ProductionPolicy(allow_post_load_publish=True), + ) + ) + + publication = runtime.publish_after_load(identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.FAILED + assert publication.failure is not None + assert publication.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert publication.failure.retryable + assert producer.calls == 0 + + +def test_postload_publication_maps_joined_and_closes_lease( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + production=ProductionPolicy(allow_post_load_publish=True), + ) + ) + assert runtime.store is not None + + class ClosingLease: + closed = False + + def close(self) -> None: + self.closed = True + + lease = ClosingLease() + + def return_joined_lease(*_args: object, **_kwargs: object) -> StoreResult: + return StoreResult(StoreStatus.JOINED, lease=cast(HostWeightLease, lease)) + + monkeypatch.setattr(runtime.store, "get_or_build", return_joined_lease) + + publication = runtime.publish_after_load(identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.JOINED + assert lease.closed + + +def test_postload_publication_synthesizes_failure_for_unexpected_store_status( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY) + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + production=ProductionPolicy(allow_post_load_publish=True), + ) + ) + assert runtime.store is not None + + def return_unexpected_miss(*_args: object, **_kwargs: object) -> StoreResult: + return StoreResult(StoreStatus.MISS) + + monkeypatch.setattr(runtime.store, "get_or_build", return_unexpected_miss) + + publication = runtime.publish_after_load(identity, producer=producer) + + assert publication.outcome is PostLoadPublicationOutcome.FAILED + assert publication.failure is not None + assert publication.failure.stage is ResolutionStage.PRODUCTION + assert publication.failure.code is FailureCode.PRODUCER_FAILED + assert not publication.failure.retryable + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_retryable_store_initialization_failure_obeys_runtime_mode( + tmp_path: Path, + mode: RuntimeMode, + expected: ResolutionOutcome, +) -> None: + root = tmp_path / "not-a-directory" + root.write_text("occupied", encoding="utf-8") + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(root))) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is expected + assert resolution.lease is None + assert len(resolution.report.attempts) == 1 + assert resolution.report.attempts[0].failure is not None + assert resolution.report.attempts[0].failure.code.value == "domain_unavailable" + assert resolution.report.attempts[0].failure.retryable + + +@pytest.mark.parametrize( + ("stage", "code"), + [ + (ResolutionStage.PRODUCTION, FailureCode.PRODUCER_UNSUPPORTED), + (ResolutionStage.VALIDATION, FailureCode.IDENTITY_COLLISION), + (ResolutionStage.LIFECYCLE, FailureCode.PUBLICATION_FAILED), + ], +) +def test_nonretryable_post_initialization_failure_cannot_fall_back_in_preferred_mode( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + stage: ResolutionStage, + code: FailureCode, +) -> None: + identity = _identity() + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / code.value)) + ) + assert runtime.store is not None + + def fail_production( + _request: object, + _producer: object, + *, + validation: object, + deadline: float, + ) -> StoreResult: + del validation, deadline + return StoreResult( + StoreStatus.FAILED, + failure=HostWeightFailure( + stage=stage, + code=code, + retryable=False, + message="injected nonretryable production failure", + details=CanonicalJson.empty(), + ), + ) + + monkeypatch.setattr(runtime.store, "get_or_build", fail_production) + resolution = runtime.resolve(identity, producer=CountingProducer(identity)) + + assert resolution.report.outcome is ResolutionOutcome.FAILED + assert resolution.report.attempts[-1].action.value == "fail_startup" + assert resolution.report.attempts[-1].failure is not None + assert resolution.report.attempts[-1].failure.code is code + + +def test_retryable_post_initialization_failure_can_fall_back_in_preferred_mode( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + identity = _identity() + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "retryable")) + ) + assert runtime.store is not None + + def fail_production( + _request: object, + _producer: object, + *, + validation: object, + deadline: float, + ) -> StoreResult: + del validation, deadline + return StoreResult( + StoreStatus.FAILED, + failure=HostWeightFailure( + stage=ResolutionStage.CAPACITY, + code=FailureCode.ENOSPC, + retryable=True, + message="injected recoverable capacity failure", + details=CanonicalJson.empty(), + ), + ) + + monkeypatch.setattr(runtime.store, "get_or_build", fail_production) + resolution = runtime.resolve(identity, producer=CountingProducer(identity)) + + assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK + assert resolution.report.attempts[-1].action.value == "canonical_fallback" + + +def test_nonretryable_store_configuration_failure_remains_visible_in_preferred_mode(tmp_path: Path) -> None: + root = tmp_path / "untrusted" + root.mkdir(mode=0o700) + root.chmod(0o777) + runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(root))) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is ResolutionOutcome.FAILED + assert resolution.report.attempts[0].action.value == "fail_startup" + assert resolution.report.attempts[0].failure is not None + assert not resolution.report.attempts[0].failure.retryable + + +def test_unsupported_integrity_policy_fails_instead_of_disguising_configuration_error(tmp_path: Path) -> None: + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=RuntimeMode.PREFERRED, + domain=_domain(tmp_path / "store"), + integrity=IntegrityPolicy(local_lookup=ValidationLevel.FS_VERITY), + ) + ) + + resolution = runtime.resolve(_identity()) + + assert resolution.report.outcome is ResolutionOutcome.FAILED + assert resolution.report.attempts[0].failure is not None + assert resolution.report.attempts[0].failure.code.value == "unsupported" diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py new file mode 100644 index 000000000..b46f4080f --- /dev/null +++ b/tests/metrics/test_emit_calls.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Verify emit-call wiring in production code paths. + +Uses behavior and Prometheus exposition to verify production call semantics. +""" + +from __future__ import annotations + +import time +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +from vllm_omni.metrics import OmniPrometheusMetrics +from vllm_omni.metrics import definitions as defs + +if TYPE_CHECKING: + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODEL = "emit-calls-test" + + +# --------------------------------------------------------------------------- +# Behavioral pins — verify call semantics with mock OmniPrometheusMetrics. +# --------------------------------------------------------------------------- + + +def _make_omni_base_with_mock_prom(mocker): + """Build a minimal OmniBase shell wired with a mock prom_metrics. + + Mirrors the pattern in ``test_prometheus.py::TestRequestLifecycleGauges`` + — ``object.__new__`` skips ``__init__`` so we don't need a real engine. + """ + from vllm_omni.entrypoints.omni_base import OmniBase + + obj = object.__new__(OmniBase) + obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) + obj.request_states = {} + obj._consumed_metric_messages = {} + obj.log_stats = True + return obj, obj.prom_metrics + + +class TestFailureCounterWiring: + """All failure paths share one request-idempotent counter helper.""" + + def test_failure_reason_taxonomy_is_bounded(self) -> None: + from vllm_omni.metrics.utils import normalize_failure_reason + + assert normalize_failure_reason("client_abort") == "client_abort" + assert normalize_failure_reason("client_disconnect") == "client_disconnect" + assert normalize_failure_reason("stage_error") == "stage_error" + assert normalize_failure_reason("unknown") == "unknown" + assert normalize_failure_reason("timeout for request req-1") == "unknown" + assert normalize_failure_reason("") == "unknown" + assert normalize_failure_reason(None) == "unknown" + + def test_fire_failure_counter_passes_reason_through(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace(e2e_done=set()), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="client_disconnect") + + prom.request_failed.assert_called_once() + prom.inc_requests_failed.assert_called_once_with("client_disconnect") + + def test_fire_failure_counter_normalizes_unrecognized_reason(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace(e2e_done=set()), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="CUDA error for req-1") + + prom.inc_requests_failed.assert_called_once_with("unknown") + + def test_fire_failure_counter_skips_when_request_already_succeeded(self, mocker) -> None: + # When the request is in e2e_done (i.e. finalize already fired + # request_succeeded), the failure path must NOT double-count. + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace(e2e_done={"req-1"}), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="client_disconnect") + + prom.request_failed.assert_not_called() + prom.inc_requests_failed.assert_not_called() + + def test_fire_failure_counter_skips_when_request_state_missing(self, mocker) -> None: + # No request_states entry — already popped by abort path. Fail-safe + # must NOT raise. + obj, prom = _make_omni_base_with_mock_prom(mocker) + + obj._record_request_failure_once("missing-req", reason="oom") + + prom.request_failed.assert_not_called() + prom.inc_requests_failed.assert_not_called() + + def test_failure_counter_is_idempotent_across_abort_and_cleanup(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace( + e2e_done=set(), + build_and_log_summary=lambda: None, + ), + failure_recorded=False, + ) + + obj._record_request_failure_once("req-1", reason="client_disconnect") + obj._log_summary_and_cleanup("req-1", reason="stage_error") + + prom.request_failed.assert_called_once() + prom.inc_requests_failed.assert_called_once_with("client_disconnect") + + def test_log_summary_and_cleanup_default_reason(self, mocker) -> None: + obj, prom = _make_omni_base_with_mock_prom(mocker) + obj.request_states["req-1"] = SimpleNamespace( + metrics=SimpleNamespace( + e2e_done=set(), + build_and_log_summary=lambda: None, + ), + failure_recorded=False, + ) + + obj._log_summary_and_cleanup("req-1") + + prom.inc_requests_failed.assert_called_once_with("stage_error") + + +class TestEarlyReturnOnLogStatsOff: + """When --log-stats is off (default), observe methods must be silent no-ops. + + This is the gating contract — emit sites can call unconditionally and the + helper short-circuits. + """ + + def test_observe_methods_silent_when_log_stats_false(self) -> None: + prom = OmniPrometheusMetrics(model_name=_MODEL, log_stats=False) + # None of these should raise or write to the registry. + prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=1.5) + prom.observe_stage_in_queue(stage=0, in_queue_s=0.2) + prom.observe_queue_wait(queue_wait_s=0.5) + prom.set_stage_waiting_requests(stage=0, n_waiting=3) + prom.observe_num_inference_steps(n_steps=20) + prom.inc_image_count(n_images=1) + prom.observe_image_pixels(n_pixels=512 * 512) + prom.set_peak_memory(stage=0, peak_memory_mb=1024.0) + prom.inc_requests_failed(reason="oom") + prom.observe_kv_wait(connector_type="shm", kv_wait_s=0.01) + + +class TestStageGenerationLabels: + def test_stage_gen_time_exposes_bounded_stage_type(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + model = _MODEL + "-stage-type" + prom = OmniPrometheusMetrics(model_name=model, log_stats=True) + prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=0.25) + prom.observe_stage_gen_time(stage=1, stage_type="diffusion", gen_time_s=1.5) + + out = generate_latest(REGISTRY).decode() + assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="0",stage_type="llm"}} 1.0' in out + assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="1",stage_type="diffusion"}} 1.0' in out + + +class TestCounterExpositionNames: + def test_counter_families_have_exactly_one_total_suffix(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + model = _MODEL + "-counter-suffix" + prom = OmniPrometheusMetrics(model_name=model, log_stats=True) + prom.inc_image_count(2) + prom.inc_requests_failed("stage_error") + + out = generate_latest(REGISTRY).decode() + assert f'{defs.IMAGE_COUNT_METRIC}_total{{model_name="{model}"}} 2.0' in out + assert f'{defs.REQUESTS_FAILED}_total{{model_name="{model}",reason="stage_error"}} 1.0' in out + assert f"{defs.IMAGE_COUNT_METRIC}_total_total" not in out + assert f"{defs.REQUESTS_FAILED}_total_total" not in out + + +class TestQueueWaitExtraction: + def test_present_zero_is_a_valid_observation(self) -> None: + from vllm_omni.metrics.utils import extract_queue_wait_s + + assert extract_queue_wait_s({"queue_wait_ms": 0.0}) == 0.0 + + def test_missing_queue_wait_is_not_synthetic_zero(self) -> None: + from vllm_omni.metrics.utils import extract_queue_wait_s + + assert extract_queue_wait_s({"preprocess_ms": 2.0}) is None + assert extract_queue_wait_s(None) is None + + +class TestStageInQueueObservation: + def test_prometheus_records_valid_zero_wait(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + model = _MODEL + "-zero-stage-wait" + prom = OmniPrometheusMetrics(model_name=model, log_stats=True) + + prom.observe_stage_in_queue(stage=1, in_queue_s=0.0) + + out = generate_latest(REGISTRY).decode() + sample = f'{defs.STAGE_IN_QUEUE_S}_count{{model_name="{model}",stage="1"}} 1.0' + assert sample in out + + +class TestStageWorkloadMetricScope: + @staticmethod + def _observe(mocker, *, stage_type: str, output_unit_type: str): + from vllm_omni.metrics.utils import observe_stage_workload_metrics + + prom = mocker.Mock(spec=OmniPrometheusMetrics) + stage_metrics = SimpleNamespace( + num_inference_steps=50, + output_unit_type=output_unit_type, + image_pixels=1024 * 1024, + output_unit_count=2, + ) + observe_stage_workload_metrics( + prom, + stage_type=stage_type, + stage_metrics=stage_metrics, + ) + return prom + + def test_diffusion_image_observes_steps_and_image_workload(self, mocker) -> None: + prom = self._observe(mocker, stage_type="diffusion", output_unit_type="image") + + prom.observe_num_inference_steps.assert_called_once_with(50) + prom.observe_image_pixels.assert_called_once_with(1024 * 1024) + prom.inc_image_count.assert_called_once_with(2) + + def test_diffusion_video_observes_steps_only(self, mocker) -> None: + prom = self._observe(mocker, stage_type="diffusion", output_unit_type="video") + + prom.observe_num_inference_steps.assert_called_once_with(50) + prom.observe_image_pixels.assert_not_called() + prom.inc_image_count.assert_not_called() + + def test_llm_text_observes_no_workload_metrics(self, mocker) -> None: + prom = self._observe(mocker, stage_type="llm", output_unit_type="text") + + prom.observe_num_inference_steps.assert_not_called() + prom.observe_image_pixels.assert_not_called() + prom.inc_image_count.assert_not_called() + + def test_non_diffusion_image_observes_image_workload_only(self, mocker) -> None: + prom = self._observe(mocker, stage_type="llm", output_unit_type="image") + + prom.observe_num_inference_steps.assert_not_called() + prom.observe_image_pixels.assert_called_once_with(1024 * 1024) + prom.inc_image_count.assert_called_once_with(2) + + @pytest.mark.parametrize("peak_memory_mb", [2048.0, 0.0]) + def test_same_finished_image_message_is_observed_exactly_once(self, mocker, peak_memory_mb: float) -> None: + from vllm_omni.entrypoints.omni_base import OmniBase + + obj = object.__new__(OmniBase) + obj._enable_ar_profiler = False + obj._consumed_metric_messages = {} + obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) + obj.mod_metrics = mocker.Mock() + obj.engine = SimpleNamespace( + get_stage_metadata=lambda _stage_id: SimpleNamespace( + stage_type="diffusion", + final_output=False, + final_output_type="image", + ) + ) + stage_metrics = SimpleNamespace( + stage_gen_time_ms=1000.0, + diffusion_metrics={"diffusion_engine_exec_time_s": 0.8, "scheduler_queue_wait_s": 0.1}, + num_inference_steps=20, + output_unit_type="image", + image_pixels=1024 * 1024, + output_unit_count=1, + serving_time_to_first_output_ms=250.0, + image_time_to_first_output_ms=250.0, + pipeline_timings={}, + ) + result = SimpleNamespace( + request_id="req-replay", + stage_id=1, + replica_id=0, + stage_submit_ts=10.0, + metrics=stage_metrics, + engine_outputs=SimpleNamespace( + stage_durations={}, + peak_memory_mb=peak_memory_mb, + finished=True, + final_output_type="image", + ), + ) + aggregator = mocker.Mock() + aggregator.stage_events = {} + aggregator.stage_first_ts = [None, None] + aggregator.stage_last_ts = [None, None] + + for _ in range(2): + assert obj._process_single_result(result, 1, aggregator, {}, 0.0, 1) is None + + aggregator.on_stage_metrics.assert_called_once_with(1, "req-replay", stage_metrics, "image") + obj.prom_metrics.inc_image_count.assert_called_once_with(1) + obj.prom_metrics.observe_image_pixels.assert_called_once_with(1024 * 1024) + obj.prom_metrics.observe_num_inference_steps.assert_called_once_with(20) + obj.prom_metrics.observe_stage_in_queue.assert_called_once_with(1, 0.1) + if peak_memory_mb > 0: + obj.prom_metrics.set_peak_memory.assert_called_once_with(1, peak_memory_mb) + else: + obj.prom_metrics.set_peak_memory.assert_not_called() + obj.mod_metrics.observe_image_ttfp.assert_called_once_with("1", "0", 0.25) + + +class TestStageWaitingAggregation: + @staticmethod + def _make_orchestrator(mocker): + from vllm_omni.engine.orchestrator import Orchestrator + + orchestrator = object.__new__(Orchestrator) + orchestrator._prom_metrics = mocker.Mock() + orchestrator._stage_replica_waiting = {} + return orchestrator + + def test_multiple_replicas_are_summed(self, mocker) -> None: + orchestrator = self._make_orchestrator(mocker) + + orchestrator._update_stage_replica_waiting(1, 0, 3) + orchestrator._update_stage_replica_waiting(1, 1, 5) + + orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 8) + + def test_zero_update_and_stage_isolation(self, mocker) -> None: + orchestrator = self._make_orchestrator(mocker) + + orchestrator._update_stage_replica_waiting(1, 0, 3) + orchestrator._update_stage_replica_waiting(2, 0, 7) + orchestrator._update_stage_replica_waiting(1, 0, 0) + + orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 0) + assert orchestrator._stage_replica_waiting[(2, 0)] == 7 + + def test_dead_replica_snapshot_is_removed(self, mocker) -> None: + orchestrator = self._make_orchestrator(mocker) + orchestrator._update_stage_replica_waiting(1, 0, 3) + orchestrator._update_stage_replica_waiting(1, 1, 5) + + orchestrator._remove_stage_replica_waiting(1, 1) + + orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 3) + + +# --------------------------------------------------------------------------- +# Behavioral pin: call sites now scope workload observations by stage/output +# type. Keep the Prometheus wrapper's <= 0 guard as a second line of defense. +# --------------------------------------------------------------------------- + + +class TestZeroGuardContract: + """Zero-valued observations must not bump histogram ``_count``.""" + + def test_zero_image_pixels_not_observed(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-guard") + prom.observe_image_pixels(n_pixels=0) + out = generate_latest(REGISTRY).decode() + # ``.labels()`` in OmniPrometheusMetrics.__init__ already creates the + # child sample, so the ``_count`` line exists with value 0.0 even + # without observations. The guard prevents ``_count`` from bumping + # to 1.0 — assert the value stays at 0. + needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-zero-guard"}}' + count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)] + assert count_lines, "expected image_pixels_count line to exist after OmniPrometheusMetrics construction" + assert float(count_lines[0].split()[-1]) == 0.0, ( + "zero-pixel observation leaked to registry — guard should early-return on <= 0" + ) + + def test_zero_num_inference_steps_not_observed(self) -> None: + from prometheus_client import REGISTRY, generate_latest + + prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-steps") + prom.observe_num_inference_steps(n_steps=0) + out = generate_latest(REGISTRY).decode() + needle = f'vllm_omni:num_inference_steps_count{{model_name="{_MODEL}-zero-steps"}}' + count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)] + assert count_lines, "expected num_inference_steps_count line to exist after construction" + assert float(count_lines[0].split()[-1]) == 0.0, ( + "zero-step observation leaked to registry — guard should early-return on <= 0" + ) + + def test_positive_image_pixels_observed(self) -> None: + # Sanity check: positive values DO get observed. + from prometheus_client import REGISTRY, generate_latest + + prom = OmniPrometheusMetrics(model_name=_MODEL + "-pos-guard") + prom.observe_image_pixels(n_pixels=512) + out = generate_latest(REGISTRY).decode() + needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-pos-guard"}}' + count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)] + assert count_lines, "expected image_pixels_count line after positive observation" + assert float(count_lines[0].split()[-1]) == 1.0, "positive-pixel observation did not increment count to 1" + + +# --------------------------------------------------------------------------- +# KV-wait emit wiring — scheduler ENTER/EXIT lifecycle + orchestrator dispatch. +# --------------------------------------------------------------------------- + + +def _make_scheduler_shell() -> OmniARScheduler: + """Minimal OmniARScheduler shell — skips upstream __init__ via object.__new__.""" + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + obj = object.__new__(OmniARScheduler) + obj.requests = {} + obj.running = [] + obj.waiting = [] + obj._kv_wait_start_ts = {} + obj._omni_kv_config = None + return obj + + +class TestKvWaitSchedulerEmit: + """Scheduler-side half of kv_wait_s: pops start ts, carries the wait across.""" + + def test_skips_when_no_start_ts_recorded(self) -> None: + sched = _make_scheduler_shell() + outputs: dict[int, list] = {} + + sched._emit_kv_wait_output(outputs, "req-no-wait", req=SimpleNamespace(client_index=0)) + + assert outputs == {}, "emit must not append output when no start ts recorded" + assert "req-no-wait" not in sched._kv_wait_start_ts + + def test_emits_wait_duration_and_connector_type(self) -> None: + sched = _make_scheduler_shell() + + sched._kv_wait_start_ts["req-1"] = time.monotonic() - 0.25 + outputs: dict[int, list] = {} + live_req = SimpleNamespace(client_index=3) + + sched._emit_kv_wait_output(outputs, "req-1", req=live_req) + + assert "req-1" not in sched._kv_wait_start_ts, "start ts must be popped after emit" + assert 3 in outputs and len(outputs[3]) == 1 + params = outputs[3][0].kv_transfer_params + assert "kv_wait_s" in params and "connector_type" in params + assert 0.0 < params["kv_wait_s"] < 1.0 + assert params["connector_type"] == "unknown" + + def test_connector_type_resolved_from_omni_kv_config(self) -> None: + sched = _make_scheduler_shell() + sched._omni_kv_config = {"connector_config": {"type": "SharedMemoryConnector"}} + + sched._kv_wait_start_ts["req-3"] = time.monotonic() - 0.01 + outputs: dict[int, list] = {} + + sched._emit_kv_wait_output(outputs, "req-3", req=SimpleNamespace(client_index=0)) + + eco = outputs[0][0] + assert eco.kv_transfer_params["connector_type"] == "SharedMemoryConnector" + + +class TestKvWaitTerminalCleanup: + def test_finish_requests_clears_requested_wait_only(self) -> None: + from vllm.v1.request import RequestStatus + + scheduler = _make_scheduler_shell() + scheduler._kv_wait_start_ts = {"req-1": 10.0, "req-2": 20.0} + + finished = scheduler.finish_requests( + (request_id for request_id in ["req-1"]), + RequestStatus.FINISHED_ABORTED, + ) + + assert finished == [] + assert "req-1" not in scheduler._kv_wait_start_ts + assert scheduler._kv_wait_start_ts["req-2"] == 20.0 + + def test_cleanup_is_idempotent_when_timestamp_is_already_gone(self) -> None: + from vllm.v1.request import RequestStatus + + scheduler = _make_scheduler_shell() + + scheduler.finish_requests("req-already-emitted", RequestStatus.FINISHED_ABORTED) + + assert scheduler._kv_wait_start_ts == {} diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py index 315594233..694808a31 100644 --- a/tests/metrics/test_modality.py +++ b/tests/metrics/test_modality.py @@ -1,5 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations +from types import SimpleNamespace + import pytest from prometheus_client import REGISTRY, generate_latest @@ -10,6 +15,7 @@ from vllm_omni.metrics.modality import ( observe_audio_streaming_finalize, observe_modality_at_finalize, ) +from vllm_omni.metrics.stats import OrchestratorAggregator pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -46,6 +52,11 @@ _EXPECTED_FAMILIES = [ defs.DIFFUSION_EXEC_PER_STEP_S, defs.DIFFUSION_PREPROCESS_S, defs.DIFFUSION_POSTPROCESS_S, + defs.VAE_DECODE_S, + defs.DIFFUSION_FORWARD_S, + defs.DIFFUSION_KV_LOAD_S, + defs.IMAGE_TTFP_S, + defs.DENOISE_STEP_LATENCY_S, ] @@ -63,6 +74,11 @@ class TestRegistration: mod.observe_diffusion_exec_per_step("s", "r", 0.01) mod.observe_diffusion_preprocess("s", "r", 0.01) mod.observe_diffusion_postprocess("s", "r", 0.2) + mod.observe_vae_decode("s", "r", 0.3) + mod.observe_diffusion_forward("s", "r", 0.8) + mod.observe_diffusion_kv_load("s", "r", 0.04) + mod.observe_image_ttfp("s", "r", 2.5) + mod.observe_denoise_step_latency("s", "r", 0.05) out = generate_latest(REGISTRY).decode() for name in _EXPECTED_FAMILIES: @@ -247,6 +263,21 @@ class _StubModMetrics: def observe_diffusion_postprocess(self, s, r, v): self.calls.append(("observe_diffusion_postprocess", s, r, v)) + def observe_vae_decode(self, s, r, seconds): + self.calls.append(("observe_vae_decode", s, r, seconds)) + + def observe_diffusion_forward(self, s, r, seconds): + self.calls.append(("observe_diffusion_forward", s, r, seconds)) + + def observe_diffusion_kv_load(self, s, r, seconds): + self.calls.append(("observe_diffusion_kv_load", s, r, seconds)) + + def observe_image_ttfp(self, s, r, seconds): + self.calls.append(("observe_image_ttfp", s, r, seconds)) + + def observe_denoise_step_latency(self, s, r, seconds): + self.calls.append(("observe_denoise_step_latency", s, r, seconds)) + class _Bag: def __init__(self, **kw): @@ -306,11 +337,13 @@ class TestObserveModalityAtFinalize: def test_diffusion_path_full(self): stub = _StubModMetrics() stage_metrics = _Bag( + output_unit_type="image", stage_gen_time_ms=1000.0, diffusion_metrics={ "preprocess_time_s": 0.05, "diffusion_engine_exec_time_s": 1.2, "postprocess_time_s": 0.3, + "forward_time_s": 1.0, "image_num": 1, "resolution": 1024, "num_inference_steps": 50, @@ -328,6 +361,71 @@ class TestObserveModalityAtFinalize: assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls + assert ("observe_denoise_step_latency", "2", "0", pytest.approx(0.02)) in stub.calls + + def test_denoise_step_latency_requires_forward_time(self): + stub = _StubModMetrics() + stage_metrics = _Bag( + output_unit_type="image", + num_inference_steps=20, + diffusion_metrics={"diffusion_engine_exec_time_s": 1.2}, + ) + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=stage_metrics, + engine_outputs=_Bag(), + ) + assert not any(c[0] == "observe_denoise_step_latency" for c in stub.calls) + + def test_stage_snapshot_keeps_forward_based_denoise_latency(self): + aggregator = OrchestratorAggregator( + num_stages=1, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=0, + ) + aggregator.diffusion_metrics["req-image"]["forward_time_s"] = 0.8 + stage_metrics = _Bag( + stage_id=None, + request_id=None, + final_output_type=None, + num_tokens_in=0, + num_tokens_out=0, + stage_gen_time_ms=1200.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=_Bag(), + diffusion_metrics=None, + num_inference_steps=20, + output_unit_type="image", + denoise_step_latency_ms=0.0, + ) + + stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image") + + assert stats.denoise_step_latency_ms == pytest.approx(40.0) + + def test_stage_snapshot_falls_back_to_stage_time_without_profiler(self): + aggregator = object.__new__(OrchestratorAggregator) + stage_metrics = _Bag( + stage_id=None, + request_id=None, + final_output_type=None, + stage_gen_time_ms=1200.0, + diffusion_metrics=None, + num_inference_steps=20, + output_unit_type="image", + denoise_step_latency_ms=0.0, + ) + + stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image") + + assert stats.diffusion_metrics is None + assert stats.denoise_step_latency_ms == pytest.approx(60.0) def test_diffusion_per_step_skipped_without_num_inference_steps(self): stub = _StubModMetrics() @@ -422,6 +520,80 @@ class TestObserveModalityAtFinalize: assert stub.calls == [] +class TestDiffusionUnitConversion: + def test_accumulator_converts_all_engine_milliseconds_to_seconds(self): + aggregator = OrchestratorAggregator( + num_stages=1, + log_stats=False, + wall_start_ts=0.0, + final_stage_id_for_e2e=0, + ) + engine_output = SimpleNamespace( + metrics={ + "preprocess_time_ms": 50.0, + "diffusion_engine_exec_time_ms": 1500.0, + "postprocess_time_ms": 20.0, + "vae_decode_time_ms": 300.0, + "forward_time_ms": 800.0, + "scheduler_queue_wait_ms": 125.0, + "kv_recv_time_ms": 40.0, + } + ) + + aggregator.accumulate_diffusion_metrics("diffusion", "req-units", engine_output) + + assert dict(aggregator.diffusion_metrics["req-units"]) == { + "preprocess_time_s": pytest.approx(0.05), + "diffusion_engine_exec_time_s": pytest.approx(1.5), + "postprocess_time_s": pytest.approx(0.02), + "vae_decode_time_s": pytest.approx(0.3), + "forward_time_s": pytest.approx(0.8), + "scheduler_queue_wait_s": pytest.approx(0.125), + "kv_recv_time_s": pytest.approx(0.04), + } + + +class TestObserveDiffusionFinalize: + @pytest.mark.parametrize( + ("metric_key", "call_name", "value"), + [ + ("vae_decode_time_s", "observe_vae_decode", 0.3), + ("forward_time_s", "observe_diffusion_forward", 0.8), + ("kv_recv_time_s", "observe_diffusion_kv_load", 0.04), + ], + ) + def test_emits_optional_breakdown_metric(self, metric_key, call_name, value): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="image", + stage_id=2, + replica_id=0, + stage_metrics=_Bag(diffusion_metrics={metric_key: value}), + engine_outputs=_Bag(), + ) + + assert (call_name, "2", "0", value) in stub.calls + + def test_audio_diffusion_stage_fires_both_paths(self): + stub = _StubModMetrics() + observe_modality_at_finalize( + stub, + output_type="audio", + stage_id=1, + replica_id=0, + stage_metrics=_Bag( + stage_gen_time_ms=1500.0, + audio_generated_frames=24000, + diffusion_metrics={"diffusion_engine_exec_time_s": 1.2}, + ), + engine_outputs=_Bag(multimodal_output={"audio_sample_rate": 24000}), + ) + + assert ("observe_diffusion_exec", "1", "0", 1.2) in stub.calls + assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls + + class TestObserveAudioFirstPacket: def test_observes_with_valid_inputs(self): stub = _StubModMetrics() @@ -541,3 +713,39 @@ class TestBucketSelection: # SECONDS_FAST_BUCKETS includes le=0.001 which is absent from SECONDS_BUCKETS. fast_marker = f'{defs.AUDIO_UNDERRUN_S}_bucket{{le="0.001"' assert fast_marker in out, "audio_underrun_s should use SECONDS_FAST_BUCKETS containing le=0.001" + + def test_vae_decode_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None: + # VAE decode can stretch into seconds for high-resolution latents; + # stays on the wide bucket so the 300 s ceiling covers video VAE. + stage, replica = "vae_decode_sec", "0" + mod.observe_vae_decode(stage, replica, 1.5) + out = generate_latest(REGISTRY).decode() + sec_marker = f'{defs.VAE_DECODE_S}_bucket{{le="300.0"' + assert sec_marker in out, "vae_decode_s should use SECONDS_BUCKETS containing le=300.0" + + def test_diffusion_forward_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None: + # Denoise loop can exceed the 60s fast-bucket ceiling for video gen. + stage, replica = "diff_forward_sec", "0" + mod.observe_diffusion_forward(stage, replica, 2.0) + out = generate_latest(REGISTRY).decode() + sec_marker = f'{defs.DIFFUSION_FORWARD_S}_bucket{{le="300.0"' + assert sec_marker in out, "diffusion_forward_s should use SECONDS_BUCKETS containing le=300.0" + + def test_denoise_step_latency_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None: + # Per-step denoise is sub-second; stays on the fast bucket for ms-level + # resolution. + stage, replica = "denoise_fast", "0" + mod.observe_denoise_step_latency(stage, replica, 0.075) + out = generate_latest(REGISTRY).decode() + fast_marker = f'{defs.DENOISE_STEP_LATENCY_S}_bucket{{le="0.001"' + assert fast_marker in out, "denoise_step_latency_s should use SECONDS_FAST_BUCKETS containing le=0.001" + + +class TestDenoiseStepLatencyGuard: + @pytest.mark.parametrize(("value", "expected_count"), [(0.0, None), (-0.1, None), (0.05, 1.0)]) + def test_nonpositive_values_are_skipped(self, mod: OmniModalityMetrics, value, expected_count) -> None: + stage, replica = f"denoise_{value}", "0" + mod.observe_denoise_step_latency(stage, replica, value) + out = generate_latest(REGISTRY).decode() + prefix = f'{defs.DENOISE_STEP_LATENCY_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}' + assert _sample_value(out, prefix) == expected_count diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py index f01933dbc..e45365cf1 100644 --- a/tests/metrics/test_prometheus.py +++ b/tests/metrics/test_prometheus.py @@ -198,3 +198,32 @@ class TestRequestLifecycleGauges: out = generate_latest(registry).decode() assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test-2"}') == 1.0 assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0 + + def test_engine_queued_requests_count_as_waiting(self, registry: CollectorRegistry) -> None: + from types import SimpleNamespace + + from vllm_omni.entrypoints.omni_base import OmniBase + from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter + + obj = object.__new__(OmniBase) + obj.engine = SimpleNamespace( + _running_counter=OmniRequestCounter(), + _engines_waiting_counter=OmniRequestCounter(), + ) + obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-3") + obj.request_states = {} + obj.log_stats = False + + # Three requests dispatched to a stage engine whose scheduler admits + # one and queues two (e.g. a diffusion engine with + # max_num_running_reqs=1). The orchestrator reports the queued count + # via _engines_waiting_counter; they must show as waiting, not running. + for _ in range(3): + obj.engine._running_counter.increment() + obj.engine._engines_waiting_counter.value = 2 + + obj._publish_request_gauges(3) + + out = generate_latest(registry).decode() + assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test-3"}') == 1.0 + assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test-3"}') == 2.0 diff --git a/tests/model_executor/model_loader/test_resolve_model_path.py b/tests/model_executor/model_loader/test_resolve_model_path.py new file mode 100644 index 000000000..996313b06 --- /dev/null +++ b/tests/model_executor/model_loader/test_resolve_model_path.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for resolve_model_to_local_path.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import huggingface_hub +import pytest +from huggingface_hub.errors import LocalEntryNotFoundError + +from vllm_omni.model_executor.model_loader import weight_utils +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Both resolution paths are patched in every test, so nothing is ever fetched +REPO_ID = "test-org/test-model" +RESOLVED = "/cache/snapshots/deadbeef" + + +@pytest.fixture +def paths(mocker): + """The resolver's two ways out, both patched.""" + return SimpleNamespace( + cache_only=mocker.patch.object(huggingface_hub, "snapshot_download", return_value=RESOLVED), + download=mocker.patch.object(weight_utils, "download_weights_from_hf_specific", return_value=RESOLVED), + ) + + +@pytest.mark.parametrize("allow_download", [False, True]) +def test_local_directory_short_circuits(tmp_path, paths, allow_download): + """A local dir resolves to itself without either path running.""" + assert resolve_model_to_local_path(str(tmp_path), allow_download=allow_download) == str(tmp_path) + + paths.cache_only.assert_not_called() + paths.download.assert_not_called() + + +@pytest.mark.parametrize( + ("kwargs", "allow_patterns", "cache_dir"), + [ + ({}, None, None), + ({"allow_patterns": ["subdir/*"], "cache_dir": "/cache"}, ["subdir/*"], "/cache"), + ], +) +def test_repo_id_is_cache_only_by_default(paths, kwargs, allow_patterns, cache_dir): + """Without `allow_download` nothing is fetched, and narrowing is forwarded.""" + assert resolve_model_to_local_path(REPO_ID, **kwargs) == RESOLVED + + paths.cache_only.assert_called_once_with( + REPO_ID, + cache_dir=cache_dir, + allow_patterns=allow_patterns, + local_files_only=True, + ) + paths.download.assert_not_called() + + +def test_cache_miss_propagates(paths): + """A failed resolution raises instead of degrading into the raw repo id.""" + paths.cache_only.side_effect = LocalEntryNotFoundError("cold cache") + + with pytest.raises(LocalEntryNotFoundError): + resolve_model_to_local_path(REPO_ID) + + paths.download.assert_not_called() + + +@pytest.mark.parametrize( + ("kwargs", "allow_patterns"), + [ + ({}, ["*"]), + ({"allow_patterns": ["subdir/model*.safetensors"]}, ["subdir/model*.safetensors"]), + ], +) +def test_download_routes_through_the_locked_helper(paths, kwargs, allow_patterns): + """Downloads go through the vLLM helper holding the node-wide lock. + + That helper is also what honors `VLLM_USE_MODELSCOPE`; + calling `snapshot_download` directly would drop both. + """ + assert resolve_model_to_local_path(REPO_ID, allow_download=True, cache_dir="/cache", **kwargs) == RESOLVED + + paths.download.assert_called_once_with( + model_name_or_path=REPO_ID, + cache_dir="/cache", + allow_patterns=allow_patterns, + require_all=True, + ) + paths.cache_only.assert_not_called() diff --git a/tests/model_executor/models/common/__init__.py b/tests/model_executor/models/common/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py b/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py new file mode 100644 index 000000000..e202321dc --- /dev/null +++ b/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py @@ -0,0 +1,292 @@ +"""Unit tests for fused_adaptive_group_norm_silu operator. + +Tests numeric correctness against PyTorch native implementation: + F.silu(F.group_norm(x, num_groups, weight, bias, eps) * (1 + scale) + shift) +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.common.ops import fused_adaptive_group_norm_silu + +# Skip tests if CUDA not available +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Triton kernels"), +] + + +def _reference(x, weight, bias, scale, shift, num_groups, eps): + """AdaGN reference, computed in fp32. + + As with fused_group_norm_silu, a same-dtype reference is unusable: the + GroupNorm affine can cancel to near zero, and rounding its O(1) terms to + bf16 wipes out the remaining significant digits. + """ + B, C = x.shape[:2] + normed = F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps) + scale_b = scale.float().view(B, C, *([1] * (x.ndim - 2))) + shift_b = shift.float().view(B, C, *([1] * (x.ndim - 2))) + return F.silu(normed * (1 + scale_b) + shift_b) + + +def _tolerance(dtype): + if dtype == torch.float32: + return 1e-5, 1e-5 + if dtype == torch.float16: + return 2e-3, 2e-3 + return 1e-2, 1e-2 # bfloat16 + + +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("channels", [32, 64, 128]) +@pytest.mark.parametrize("spatial_size", [(16, 16), (32, 32)]) +@pytest.mark.parametrize("num_groups", [8, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fused_adaptive_group_norm_silu_correctness(batch_size, channels, spatial_size, num_groups, dtype): + """Test numeric correctness against PyTorch native ops.""" + if channels % num_groups != 0: + pytest.skip(f"channels {channels} not divisible by num_groups {num_groups}") + + H, W = spatial_size + eps = 1e-6 + torch.manual_seed(42) + kw = dict(device="cuda", dtype=dtype) + + x = torch.randn(batch_size, channels, H, W, **kw) + weight = torch.randn(channels, **kw) + bias = torch.randn(channels, **kw) + scale = torch.randn(batch_size, channels, **kw) + shift = torch.randn(batch_size, channels, **kw) + + ref_out = _reference(x, weight, bias, scale, shift, num_groups, eps).to(dtype) + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, num_groups, eps) + + rtol, atol = _tolerance(dtype) + torch.testing.assert_close( + fused_out, + ref_out, + rtol=rtol, + atol=atol, + msg=f"Mismatch for batch={batch_size}, C={channels}, HW={spatial_size}, groups={num_groups}, dtype={dtype}", + ) + + +def test_broadcast_scale_shift_shapes(): + """scale/shift may arrive as (B, C) or already broadcast to (B, C, 1, 1).""" + B, C, H, W = 2, 64, 16, 16 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(B, C, H, W, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + flat = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + expanded = fused_adaptive_group_norm_silu(x, weight, bias, scale.view(B, C, 1, 1), shift.view(B, C, 1, 1), 32, 1e-6) + torch.testing.assert_close(flat, expanded, rtol=0, atol=0) + + +def test_output_dtype_matches_eager(): + """Output dtype must match what eager GroupNorm would produce.""" + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + B, C = 2, 64 + kw = dict(device="cuda", dtype=dtype) + x = torch.randn(B, C, 16, 16, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + eager_dtype = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)).dtype + out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + + assert out.dtype == eager_dtype, ( + f"Output dtype {out.dtype} != eager dtype {eager_dtype} for input dtype {dtype}" + ) + assert out.shape == x.shape + + +@pytest.mark.parametrize( + "autocast_dtype, input_dtype", + [ + (torch.float16, torch.float16), + (torch.bfloat16, torch.bfloat16), + (torch.float16, torch.float32), + (torch.bfloat16, torch.float32), + ], +) +def test_autocast_matches_eager(autocast_dtype, input_dtype): + """Inside autocast the fused op must stay a drop-in replacement. + + Also pins the fp32-activation corner so the inlined ``is_autocast_enabled + -> fp32`` branch cannot silently downcast an already-fp32 activation. + """ + B, C = 2, 64 + x = torch.randn(B, C, 16, 16, device="cuda", dtype=input_dtype) + weight = torch.randn(C, device="cuda", dtype=torch.float32) + bias = torch.randn(C, device="cuda", dtype=torch.float32) + scale = torch.randn(B, C, device="cuda", dtype=torch.float32) + shift = torch.randn(B, C, device="cuda", dtype=torch.float32) + + with torch.autocast("cuda", dtype=autocast_dtype): + eager_out = F.silu( + F.group_norm(x, 32, weight, bias, 1e-6) * (1 + scale.view(B, C, 1, 1)) + shift.view(B, C, 1, 1) + ) + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + + assert eager_out.dtype == torch.float32, "sanity check: eager group_norm is expected to return fp32 under autocast" + assert fused_out.dtype == eager_out.dtype, ( + f"Output dtype {fused_out.dtype} != eager dtype {eager_out.dtype} " + f"under autocast({autocast_dtype}) for input {input_dtype}" + ) + torch.testing.assert_close(fused_out, eager_out, rtol=1e-5, atol=1e-5) + + +def test_invalid_channels_not_divisible(): + """Test that assertion fires when channels not divisible by num_groups.""" + B, C = 1, 33 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(B, C, 16, 16, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + with pytest.raises(AssertionError, match="must be divisible by num_groups"): + fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(2, 64, 3, 16, 16), (1, 32, 8, 8), (2, 64, 128)]) +def test_non_2d_spatial_matches_eager(dtype, shape): + """Any spatial rank must work; the wrapper collapses the spatial axes. + + GroupNorm reduces over the whole channel group and every spatial position, + so flattening is exact rather than an approximation. + """ + torch.manual_seed(0) + B, C = shape[:2] + kw = dict(device="cuda", dtype=dtype) + x = torch.randn(*shape, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + ref_out = _reference(x, weight, bias, scale, shift, 32, 1e-6).to(dtype) + + assert fused_out.shape == x.shape + rtol, atol = _tolerance(dtype) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("batch_size", [1, 2, 3]) +def test_scale_shift_from_chunk_are_not_contiguous(batch_size): + """Regression: scale/shift arrive as halves of a torch.chunk. + + ResBlock builds them as ``torch.chunk(emb_out, 2, dim=1)`` where emb_out is + (B, 2C, 1, 1). For B > 1 each half is a strided view, so the wrapper must + ``reshape`` (which copies) instead of ``view`` (which raises). + """ + C, H, W = 64, 16, 16 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(batch_size, C, H, W, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + + emb_out = torch.randn(batch_size, 2 * C, 1, 1, **kw) + scale, shift = torch.chunk(emb_out, 2, dim=1) + if batch_size > 1: + assert not scale.is_contiguous(), "test premise: chunk half is strided" + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6) * (1.0 + scale) + shift) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-5) + + +def test_non_contiguous_input_matches_eager(): + """A channels-last / permuted activation must still be handled correctly.""" + B, C, H, W = 2, 64, 16, 16 + kw = dict(device="cuda", dtype=torch.float32) + x = torch.randn(B, H, W, C, **kw).permute(0, 3, 1, 2) + assert not x.is_contiguous(), "test premise: input is strided" + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6) + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6) * (1.0 + scale.view(B, C, 1, 1)) + shift.view(B, C, 1, 1)) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize("spatial", [32, 64]) +def test_dit_sized_activation(spatial): + """Shapes the HunyuanImage3 DiT actually sees. + + The earlier kernel walked the spatial axis one position at a time with + ``spatial_size`` as a ``tl.constexpr``; at 64x64 that is 4096 serial steps + per pass. This test pins the realistic size so a regression to that layout + shows up as a timeout rather than silently costing throughput. + """ + B, C = 2, 256 + kw = dict(device="cuda", dtype=torch.bfloat16) + x = torch.randn(B, C, spatial, spatial, **kw) + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + scale = torch.randn(B, C, **kw) + shift = torch.randn(B, C, **kw) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-5) + ref_out = _reference(x, weight, bias, scale, shift, 32, 1e-5).to(torch.bfloat16) + + rtol, atol = _tolerance(torch.bfloat16) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_large_offset_small_variance(dtype): + """A large constant offset with a small variance must stay stable. + + Mirrors the fused_group_norm_silu case: ``x ~ 10000 ± 0.1`` has mean ~1e4 + and variance ~1e-2, so ``E[x^2] - E[x]^2`` cancels catastrophically in fp32 + (rounding the variance to ~8.0). The Welford reduction keeps it accurate. + fp32 tolerance is 0.1 (not 1e-6) because ``x``'s fp32 representation noise + (~0.001 ULP at 1e4) is amplified by ``rstd ~ 10`` to ~0.01-0.02 — the same + floor PyTorch's own reference hits; the old naive implementation would be + off by ~1.0 here. fp16/bf16 quantize the perturbation away, so for those + dtypes the meaningful assertion is just "no NaN". + """ + torch.manual_seed(0) + B, C, H, W = 2, 64, 16, 16 + num_groups = 32 + eps = 1e-6 + device = torch.device("cuda") + + x = torch.full((B, C, H, W), 10000.0, device=device, dtype=dtype) + x = x + torch.randn(B, C, H, W, device=device, dtype=dtype) * 0.1 + weight = torch.randn(C, device=device, dtype=dtype) + bias = torch.randn(C, device=device, dtype=dtype) + scale = torch.randn(B, C, device=device, dtype=dtype) + shift = torch.randn(B, C, device=device, dtype=dtype) + + fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, num_groups, eps) + ref_out = _reference(x, weight, bias, scale, shift, num_groups, eps).to(dtype) + + assert not torch.isnan(fused_out).any(), "fused output must not contain NaN" + + if dtype == torch.float32: + # ~0.02 is the representation-noise floor at offset 1e4; old impl ~1.0. + torch.testing.assert_close(fused_out, ref_out, rtol=0.05, atol=0.1) + # fp16/bf16: perturbation is swallowed by the dtype — no meaningful close, + # the no-NaN assertion above is the regression guard. + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/model_executor/models/common/test_fused_group_norm_silu.py b/tests/model_executor/models/common/test_fused_group_norm_silu.py new file mode 100644 index 000000000..3e9d055af --- /dev/null +++ b/tests/model_executor/models/common/test_fused_group_norm_silu.py @@ -0,0 +1,338 @@ +"""Unit tests for fused_group_norm_silu operator. + +Tests numeric correctness against PyTorch native implementation: + F.silu(F.group_norm(x, num_groups, weight, bias, eps)) +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.common.ops import fused_group_norm_silu + +# Skip tests if CUDA not available +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Triton kernels"), +] + + +def _tolerance(dtype): + """Tolerance against an fp32 reference, rounded once to ``dtype``.""" + if dtype == torch.float32: + return 1e-5, 1e-6 + if dtype == torch.float16: + return 2e-3, 2e-3 + return 1e-2, 1e-2 # bfloat16 + + +@pytest.mark.parametrize("batch_size", [1, 2, 4]) +@pytest.mark.parametrize("channels", [32, 64, 128]) +@pytest.mark.parametrize("spatial_size", [(16, 16), (32, 32), (64, 64)]) +@pytest.mark.parametrize("num_groups", [8, 16, 32]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fused_group_norm_silu_correctness(batch_size, channels, spatial_size, num_groups, dtype): + """Test numeric correctness against PyTorch native ops.""" + # Skip invalid configurations + if channels % num_groups != 0: + pytest.skip(f"channels {channels} not divisible by num_groups {num_groups}") + + H, W = spatial_size + device = torch.device("cuda") + eps = 1e-6 + + # Create inputs + torch.manual_seed(42) + x = torch.randn(batch_size, channels, H, W, device=device, dtype=dtype) + weight = torch.randn(channels, device=device, dtype=dtype) + bias = torch.randn(channels, device=device, dtype=dtype) + + # Reference implementation, computed in fp32. + # + # A same-dtype reference is NOT usable here: GroupNorm's affine step + # ``(x - mean) * rstd * weight + bias`` adds two O(1) terms that can cancel + # down to ~1e-3. Rounding those terms to bf16 (ULP ~8e-3 near 1.0) destroys + # every remaining significant digit -- measured cases flip sign, e.g. a true + # +0.006136 comes out as -0.000471. The low-precision "reference" is then + # far less accurate than the kernel it is supposed to validate: the fused + # output matches the correctly-rounded fp32 result for 99.997% of elements, + # a bf16 reference for only 65.2%. + ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype) + + # Fused implementation + fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps) + + # Check correctness + # Tolerances are set from the measured worst case across this parameter + # sweep, leaving roughly 2x headroom for each dtype. + if dtype == torch.float32: + rtol, atol = 1e-5, 1e-6 + elif dtype == torch.float16: + rtol, atol = 2e-3, 2e-3 + else: # bfloat16 + rtol, atol = 1e-2, 1e-2 + + torch.testing.assert_close( + fused_out, + ref_out, + rtol=rtol, + atol=atol, + msg=f"Mismatch for batch={batch_size}, C={channels}, HW={spatial_size}, groups={num_groups}, dtype={dtype}", + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_fp32_accumulation_precision(dtype): + """Test that fp32 accumulation matches PyTorch's behavior. + + GroupNorm uses fp32 for mean/variance computation regardless of input dtype. + Our kernel should match this behavior. + """ + batch_size, channels, H, W = 2, 64, 32, 32 + num_groups = 32 + eps = 1e-6 + device = torch.device("cuda") + + # Create inputs with values that expose precision differences + torch.manual_seed(123) + x = torch.randn(batch_size, channels, H, W, device=device, dtype=dtype) * 100 + weight = torch.ones(channels, device=device, dtype=dtype) + bias = torch.zeros(channels, device=device, dtype=dtype) + + # Reference: fp32 accumulation, then a single round back to ``dtype``. + # (Same discipline as the main correctness test -- a same-dtype reference + # wouldn't isolate the accumulation path this test is named after.) + ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype) + + # Fused + fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps) + + # Should match closely even with large values + if dtype == torch.float32: + rtol, atol = 1e-5, 1e-5 + else: + rtol, atol = 1e-2, 1e-2 + + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +def test_edge_case_zero_input(): + """Test with zero input.""" + x = torch.zeros(1, 32, 16, 16, device="cuda") + weight = torch.ones(32, device="cuda") + bias = torch.zeros(32, device="cuda") + + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-6) + + +def test_edge_case_large_eps(): + """Test with large epsilon value.""" + x = torch.randn(1, 32, 16, 16, device="cuda") + weight = torch.randn(32, device="cuda") + bias = torch.randn(32, device="cuda") + eps = 1e-2 # Larger than typical + + ref_out = F.silu(F.group_norm(x, 32, weight, bias, eps)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, eps) + + torch.testing.assert_close(fused_out, ref_out, rtol=1e-4, atol=1e-5) + + +def test_output_dtype_matches_eager(): + """Output dtype must match what eager GroupNorm+SiLU would produce. + + Outside autocast that means the input dtype; inside autocast it means fp32, + because group_norm is on autocast's fp32 cast policy. HunyuanImage3's VAE + runs fp32 weights under fp16 autocast, so getting this wrong would silently + downcast the post-norm activation. + """ + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + x = torch.randn(2, 64, 32, 32, device="cuda", dtype=dtype) + weight = torch.randn(64, device="cuda", dtype=dtype) + bias = torch.randn(64, device="cuda", dtype=dtype) + + eager_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + assert out.dtype == eager_out.dtype, ( + f"Output dtype {out.dtype} != eager dtype {eager_out.dtype} for input dtype {dtype}" + ) + assert out.shape == x.shape, f"Output shape {out.shape} != input shape {x.shape}" + + +@pytest.mark.parametrize( + "autocast_dtype, input_dtype", + [ + (torch.float16, torch.float16), + (torch.bfloat16, torch.bfloat16), + (torch.float16, torch.float32), + (torch.bfloat16, torch.float32), + ], +) +def test_autocast_matches_eager(autocast_dtype, input_dtype): + """Inside autocast the fused op must stay a drop-in replacement. + + Mirrors the HunyuanImage3 VAE setup: fp32 module weights, activations + arriving as fp16/bf16 from an upstream conv, all under autocast. Also pins + the fp32-activation corner (some VAE stages) so the inlined + ``is_autocast_enabled -> fp32`` branch cannot silently downcast it. + """ + x = torch.randn(2, 64, 32, 32, device="cuda", dtype=input_dtype) + weight = torch.randn(64, device="cuda", dtype=torch.float32) + bias = torch.randn(64, device="cuda", dtype=torch.float32) + + with torch.autocast("cuda", dtype=autocast_dtype): + eager_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + assert eager_out.dtype == torch.float32, "sanity check: eager group_norm is expected to return fp32 under autocast" + assert fused_out.dtype == eager_out.dtype, ( + f"Output dtype {fused_out.dtype} != eager dtype {eager_out.dtype} " + f"under autocast({autocast_dtype}) for input {input_dtype}" + ) + torch.testing.assert_close(fused_out, eager_out, rtol=1e-5, atol=1e-5) + + +def test_invalid_channels_not_divisible(): + """Test that assertion fires when channels not divisible by num_groups.""" + x = torch.randn(1, 33, 16, 16, device="cuda") # 33 not divisible by 32 + weight = torch.randn(33, device="cuda") + bias = torch.randn(33, device="cuda") + + with pytest.raises(AssertionError, match="must be divisible by num_groups"): + fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + + +def test_invalid_weight_shape(): + """Test that assertion fires with wrong weight shape.""" + x = torch.randn(1, 64, 16, 16, device="cuda") + weight = torch.randn(32, device="cuda") # Wrong: should be 64 + bias = torch.randn(64, device="cuda") + + with pytest.raises(AssertionError, match="doesn't match channels"): + fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(2, 64, 3, 16, 16), (1, 32, 1, 8, 8), (2, 64, 128)]) +def test_non_2d_spatial_matches_eager(dtype, shape): + """Spatial ranks other than 2 must work: the 3D VAE feeds (N, C, T, H, W). + + GroupNorm reduces over the whole channel group and every spatial position, + so flattening the spatial axes inside the wrapper has to be exact. + """ + torch.manual_seed(0) + C = shape[1] + x = torch.randn(*shape, device="cuda", dtype=dtype) + weight = torch.randn(C, device="cuda", dtype=dtype) + bias = torch.randn(C, device="cuda", dtype=dtype) + + fused_out = fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + ref_out = F.silu(F.group_norm(x.float(), 32, weight.float(), bias.float(), 1e-6)).to(dtype) + + assert fused_out.shape == x.shape + rtol, atol = _tolerance(dtype) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_permuted_input_matches_eager(dtype): + """A permuted (channel-stride-1) activation must still be correct. + + This is not a hypothetical layout: HunyuanImage3's ``UNetUp`` calls this op + on the output of ``rearrange(x, "b (h w) c -> b c h w")``. An earlier + version of the kernel indexed such inputs through their strides, which was + correct but uncoalesced enough to run at 0.44x of eager, so the wrapper now + materializes them first. + """ + torch.manual_seed(0) + B, C, H, W = 2, 64, 16, 16 + kw = dict(device="cuda", dtype=dtype) + x = torch.randn(B, H * W, C, **kw).permute(0, 2, 1).reshape(B, C, H, W) + assert not x.is_contiguous(), "test premise: input is strided" + + weight = torch.randn(C, **kw) + bias = torch.randn(C, **kw) + + fused_out = fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6) + ref_out = F.silu(F.group_norm(x.float(), 32, weight.float(), bias.float(), 1e-6)).to(dtype) + + rtol, atol = _tolerance(dtype) + torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("channels", [64, 128]) +def test_hunyuan_vae_config(batch_size, channels): + """Test with HunyuanImage3 VAE actual configuration. + + HunyuanImage3 ResnetBlock uses: + - num_groups=32 + - eps=1e-6 + - Typical spatial sizes: 16x16 to 128x128 (after tiling) + """ + spatial_sizes = [(16, 16), (32, 32), (64, 64), (128, 128)] + + for H, W in spatial_sizes: + x = torch.randn(batch_size, channels, H, W, device="cuda", dtype=torch.float32) + weight = torch.randn(channels, device="cuda") + bias = torch.randn(channels, device="cuda") + + ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)) + fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6) + + torch.testing.assert_close( + fused_out, + ref_out, + rtol=1e-5, + atol=1e-6, + msg=f"HunyuanVAE config: batch={batch_size}, C={channels}, HW=({H},{W})", + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_large_offset_small_variance(dtype): + """A large constant offset with a small variance must stay stable. + + ``x ~ 10000 ± 0.1`` has mean ~1e4 and variance ~1e-2, so the naive + ``E[x^2] - E[x]^2`` subtracts two ~1e8 quantities and rounds the variance to + ~8.0 (the ULP of 1e8) — catastrophic cancellation, producing garbage + normalization. The Welford reduction keeps the variance accurate. + + Why fp32 tolerance is 0.1, not the usual 1e-6: at this offset, ``x`` is stored + in fp32 with ULP ~0.001, so ``x - mean`` carries ~0.001 of representation + noise, and ``rstd ~ 10`` amplifies it to ~0.01-0.02. Even PyTorch's own + GroupNorm reference hits this floor (the old naive implementation would be + off by ~1.0 here). fp16/bf16 quantize the 0.1 perturbation away entirely, so + for those dtypes the meaningful assertion is just "no NaN". + """ + torch.manual_seed(0) + B, C, H, W = 2, 64, 16, 16 + num_groups = 32 + eps = 1e-6 + device = torch.device("cuda") + + x = torch.full((B, C, H, W), 10000.0, device=device, dtype=dtype) + x = x + torch.randn(B, C, H, W, device=device, dtype=dtype) * 0.1 + weight = torch.randn(C, device=device, dtype=dtype) + bias = torch.randn(C, device=device, dtype=dtype) + + fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps) + # fp32 reference = PyTorch's Welford path, the ground truth. + ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype) + + assert not torch.isnan(fused_out).any(), "fused output must not contain NaN" + + if dtype == torch.float32: + # ~0.02 is the representation-noise floor at offset 1e4; old impl ~1.0. + torch.testing.assert_close(fused_out, ref_out, rtol=0.05, atol=0.1) + # fp16/bf16: perturbation is swallowed by the dtype — no meaningful close, + # the no-NaN assertion above is the regression guard. + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 2a4cb2f1e..906b27919 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for CosyVoice3 components.""" from types import SimpleNamespace @@ -9,10 +9,41 @@ import torch import torch.nn as nn from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( + CausalHiFTGenerator, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.fixture +def causal_hift(): + return CausalHiFTGenerator( + base_channels=32, + upsample_rates=[2, 2], + upsample_kernel_sizes=[4, 4], + source_resblock_kernel_sizes=[3, 3], + source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]], + ) + + +def test_causal_hift_moves_stft_window_with_model(causal_hift): + assert causal_hift.get_buffer("stft_window") is causal_hift.stft_window + assert "stft_window" not in causal_hift.state_dict() + causal_hift.to(dtype=torch.float64) + assert causal_hift.stft_window.dtype == torch.float64 + + +def test_causal_hift_stft_moves_window_to_input_device(causal_hift): + waveform = torch.empty((1, 64), device="meta") + + real, imag = causal_hift._stft(waveform) + + assert real.device == waveform.device + assert imag.device == waveform.device + assert causal_hift.stft_window.device == waveform.device + + class TestPreLookaheadLayer: """Tests for PreLookaheadLayer.""" @@ -250,6 +281,153 @@ class TestCFM: assert out.shape == mu.shape + @pytest.mark.core_model + @pytest.mark.cpu + def test_trt_estimator_uses_stream_dependencies(self, monkeypatch): + from omegaconf import DictConfig + + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ConditionalCFM + + class FakeStream: + def __init__(self, name, state): + self.name = name + self.state = state + self.cuda_stream = hash(name) + self.synchronize_calls = 0 + self.waited_on = [] + + def synchronize(self): + self.synchronize_calls += 1 + + def wait_stream(self, stream): + self.waited_on.append(stream) + + class FakeStreamContext: + def __init__(self, stream): + self.stream = stream + + def __enter__(self): + self.previous = self.stream.state.current + self.stream.state.current = self.stream + return self.stream + + def __exit__(self, exc_type, exc, traceback): + self.stream.state.current = self.previous + + state = SimpleNamespace(current=None) + caller_stream = FakeStream("caller", state) + estimator_stream = FakeStream("estimator", state) + state.current = caller_stream + monkeypatch.setattr(torch.cuda, "current_stream", lambda *args, **kwargs: state.current) + stream_contexts = [] + + def stream_context(stream): + stream_contexts.append(stream) + return FakeStreamContext(stream) + + monkeypatch.setattr(torch.cuda, "stream", stream_context) + + class FakeContext: + def __init__(self): + self.execute_stream = None + + def set_input_shape(self, name, shape): + pass + + def set_tensor_address(self, name, address): + pass + + def execute_async_v3(self, stream): + self.execute_stream = stream + return True + + class FakeEngine: + @staticmethod + def get_tensor_name(index): + return f"tensor_{index}" + + context = FakeContext() + + class FakeEstimatorPool: + io_dtype = torch.float32 + + def __init__(self): + self.released = [] + + def acquire_estimator(self): + return [context, estimator_stream], FakeEngine() + + def release_estimator(self, released_context, released_stream): + self.released.append((released_context, released_stream)) + + estimator_pool = FakeEstimatorPool() + cfm = ConditionalCFM( + in_channels=80, + cfm_params=DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ), + n_spks=1, + spk_emb_dim=80, + estimator=estimator_pool, + ) + + x = torch.randn(2, 80, 4) + mask = torch.ones(2, 1, 4) + mu = torch.randn(2, 80, 4) + timestep = torch.randn(2) + speakers = torch.randn(2, 80) + condition = torch.randn(2, 80, 4) + + output = cfm.forward_estimator(x, mask, mu, timestep, speakers, condition) + + assert output.shape == x.shape + assert caller_stream.synchronize_calls == 0 + assert estimator_stream.synchronize_calls == 0 + assert estimator_stream.waited_on == [caller_stream] + assert caller_stream.waited_on == [estimator_stream] + assert stream_contexts == [estimator_stream] + assert context.execute_stream == estimator_stream.cuda_stream + assert estimator_pool.released == [(context, estimator_stream)] + + @pytest.mark.core_model + @pytest.mark.cpu + def test_trt_context_pool_stores_cuda_stream(self, monkeypatch): + from vllm_omni.model_executor.models.cosyvoice3.flow_estimator_trt import TrtContextWrapper + + execution_context = object() + + class FakeEngine: + @staticmethod + def create_execution_context(): + return execution_context + + cuda_stream = object() + monkeypatch.setattr(torch.cuda, "Stream", lambda device: cuda_stream) + + def reject_stream_context(stream): + pytest.fail("TrtContextWrapper must store the CUDA stream, not a StreamContext") + + monkeypatch.setattr(torch.cuda, "stream", reject_stream_context) + + engine = FakeEngine() + wrapper = TrtContextWrapper(engine, device="cuda:0") + [context, stream], acquired_engine = wrapper.acquire_estimator() + + assert context is execution_context + assert stream is cuda_stream + assert acquired_engine is engine + + wrapper.release_estimator(context, stream) + [reused_context, reused_stream], _ = wrapper.acquire_estimator() + assert reused_context is context + assert reused_stream is stream + class TestSDPAFallback: """Test SDPA fallback for float32 inputs.""" diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py new file mode 100644 index 000000000..7a5492d55 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.config.model import OmniModelArchConfigConvertor +from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_attention_geometry_matches_qwen_backbone() -> None: + config = CosyVoice3Config() + model_arch_config = OmniModelArchConfigConvertor(config, config.get_text_config()).convert() + + assert model_arch_config.total_num_attention_heads == 14 + assert model_arch_config.total_num_kv_heads == 2 + assert model_arch_config.head_size == 64 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index a2c3142c7..e16927096 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -401,6 +401,52 @@ def test_sample_tolerates_padded_rows_without_history(): assert out.sampled_token_ids.shape == (2, 1) +def test_sample_excludes_non_finite_logits(): + model = _make_talker_model() + metadata = _make_sampling_metadata(output_token_ids=[[]]) + metadata.temperature.fill_(0.5) + logits = torch.tensor([[float("nan"), 1.0, float("inf"), float("-inf")]], dtype=torch.bfloat16) + + out = model.sample(logits, metadata) + + assert out is not None + assert out.sampled_token_ids.tolist() == [[1]] + + +def test_sample_preserves_allowed_token_mask_with_invalid_logits(): + model = _make_talker_model() + metadata = _make_sampling_metadata(output_token_ids=[[]]) + metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]]) + logits = torch.tensor([[float("nan"), 1.0, float("inf")]], dtype=torch.float32) + + out = model.sample(logits, metadata) + + assert out is not None + assert out.sampled_token_ids.tolist() == [[1]] + + +def test_sample_rejects_rows_without_finite_logits(): + model = _make_talker_model() + metadata = _make_sampling_metadata(output_token_ids=[[]]) + metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]]) + logits = torch.tensor([[0.0, float("nan"), 0.0]], dtype=torch.float32) + + with pytest.raises(ValueError, match="no finite logits"): + model.sample(logits, metadata) + + +def test_sample_keeps_only_finite_token_after_ras_rejection(): + model = _make_talker_model() + metadata = _make_sampling_metadata(output_token_ids=[[1] * 10]) + metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]]) + logits = torch.tensor([[0.0, 1.0, 0.0]], dtype=torch.float32) + + out = model.sample(logits, metadata) + + assert out is not None + assert out.sampled_token_ids.tolist() == [[1]] + + def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in(): metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]]) expected = SamplerOutput( @@ -423,9 +469,14 @@ def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in(): _, GPUARModelRunner = _cosyvoice3_model_and_runner() runner = object.__new__(GPUARModelRunner) runner.input_batch = _DummyInputBatch() + + def model_sample(logits, sampling_metadata): + calls.append(logits.clone()) + return expected + runner.model = SimpleNamespace( prefer_model_sampler=True, - sample=lambda logits, sampling_metadata: calls.append(logits.clone()) or expected, + sample=model_sample, ) runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used")) @@ -459,12 +510,14 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler(): _, GPUARModelRunner = _cosyvoice3_model_and_runner() runner = object.__new__(GPUARModelRunner) runner.input_batch = _DummyInputBatch() + + def model_sample(logits, sampling_metadata): + seen_histories.append([list(x) for x in sampling_metadata.output_token_ids]) + return SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None) + runner.model = SimpleNamespace( prefer_model_sampler=True, - sample=lambda logits, sampling_metadata: ( - seen_histories.append([list(x) for x in sampling_metadata.output_token_ids]) - or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None) - ), + sample=model_sample, ) runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used")) @@ -504,12 +557,14 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler(): _, GPUARModelRunner = _cosyvoice3_model_and_runner() runner = object.__new__(GPUARModelRunner) runner.input_batch = _DummyInputBatch() + + def model_sample(logits, sampling_metadata): + seen_histories.append([list(x) for x in sampling_metadata.output_token_ids]) + return SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None) + runner.model = SimpleNamespace( prefer_model_sampler=True, - sample=lambda logits, sampling_metadata: ( - seen_histories.append([list(x) for x in sampling_metadata.output_token_ids]) - or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None) - ), + sample=model_sample, ) runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used")) diff --git a/tests/model_executor/models/dots_tts/__init__.py b/tests/model_executor/models/dots_tts/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/tests/model_executor/models/dots_tts/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/model_executor/models/dots_tts/test_talker.py b/tests/model_executor/models/dots_tts/test_talker.py new file mode 100644 index 000000000..ed3c6bcaa --- /dev/null +++ b/tests/model_executor/models/dots_tts/test_talker.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the dots.tts talker's stop-signal pairing and +per-request state lifecycle. + +Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes +the methods under test read) — no weights, no GPU, no engine. Mirrors +tests/model_executor/models/voxcpm2/test_talker_state_eviction.py. +""" + +from __future__ import annotations + +import functools +from types import SimpleNamespace + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +torch = pytest.importorskip("torch") + +_VOCAB_SIZE = 151672 + + +@functools.lru_cache(maxsize=1) +def _dots_tts_talker_mod(): + """Defer talker import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import ( + DotsTTSForConditionalGeneration, + _RequestState, + ) + + return DotsTTSForConditionalGeneration, _RequestState + + +def _make_bare_talker(): + DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod() + talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration) + talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE) + talker._active_states = {} + talker._pending_requests = [] + talker._results_queue = [] + talker._audio_queue = [] + talker._deferred_cleanup_ids = set() + talker._beta_trace = False + return talker + + +def _add_state(talker, req_id: str, *, is_stopping: bool = False): + _, _RequestState = _dots_tts_talker_mod() + state = _RequestState(request_id=req_id) + state.prefill_completed = True + state.is_stopping = is_stopping + talker._active_states[req_id] = state + return state + + +def _hidden(bsz: int) -> torch.Tensor: + return torch.zeros(bsz, 1536, dtype=torch.float32) + + +class TestComputeLogitsPairing: + """compute_logits pairs _results_queue entries with batch rows + positionally, so every request must push exactly one entry per step — + including prefills, which push a (req_id, None) placeholder.""" + + def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None: + talker = _make_bare_talker() + _add_state(talker, "prefill-req") + decode_state = _add_state(talker, "decode-req", is_stopping=True) + decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]]) + + # Row 0 is the prefill (None placeholder), row 1 the stopping decode. + talker._results_queue = [ + ("prefill-req", None), + ("decode-req", torch.tensor([[0.1, 0.9]])), + ] + logits = talker.compute_logits(_hidden(2)) + + assert logits.shape == (2, _VOCAB_SIZE) + # Prefill row: forced continue. + assert logits[0, 0].item() == 1.0 + assert logits[0, 1].item() == float("-inf") + # Stopping decode row: stop slot wins. + assert logits[1, 1].item() == 1.0 + assert logits[1, 0].item() == 0.0 + # Everything else stays masked. + assert logits[0, 2].item() == float("-inf") + assert logits[1, 2].item() == float("-inf") + # Queue drained. + assert talker._results_queue == [] + + def test_below_threshold_stop_logits_force_continue(self) -> None: + """Raw (continue, stop) softmax must not reach the sampler: a + prob_stop in (0.5, 0.8] would flip greedy argmax even though + is_stopping stayed False (the 0.8-threshold contract).""" + talker = _make_bare_talker() + _add_state(talker, "req", is_stopping=False) + talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))] + + logits = talker.compute_logits(_hidden(1)) + + assert logits[0, 0].item() == 1.0 + assert logits[0, 1].item() == float("-inf") + + def test_empty_queue_defaults_to_all_continue(self) -> None: + talker = _make_bare_talker() + logits = talker.compute_logits(_hidden(3)) + + assert torch.all(logits[:, 0] == 1.0) + assert torch.all(logits[:, 1] == float("-inf")) + + +class TestStateEviction: + """on_requests_finished fires BEFORE forward(), so eviction must be + deferred to _flush_deferred_cleanup at the end of the step — the + finishing request's audio still drains through the current forward.""" + + def test_finished_request_survives_until_flush(self) -> None: + talker = _make_bare_talker() + _add_state(talker, "req-a") + _add_state(talker, "req-b") + + talker.on_requests_finished(["req-a"]) + assert "req-a" in talker._active_states # not evicted yet + + talker._flush_deferred_cleanup() + assert "req-a" not in talker._active_states + assert "req-b" in talker._active_states + assert talker._deferred_cleanup_ids == set() + + def test_unknown_request_id_is_ignored(self) -> None: + talker = _make_bare_talker() + talker.on_requests_finished(["never-seen"]) + talker._flush_deferred_cleanup() + assert talker._active_states == {} diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 6d81f300a..f5a7094ef 100755 --- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Deterministic unit tests for higgs-audio v3. These tests verify AC-1 (config), AC-3 (prompt), AC-4 (fused modules), @@ -343,7 +344,22 @@ class TestSamplerMethods: return t @staticmethod - def _sampling_metadata(temperature=1.0, top_k=50, top_p=0.95): + def _make_state_tracking_talker(num_rows): + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + talker_cls = mod.HiggsAudioV3TalkerForConditionalGeneration + talker = talker_cls.__new__(talker_cls) + torch.nn.Module.__init__(talker) + talker.num_codebooks = 8 + talker._decode_last_codes = torch.zeros(num_rows, 8, dtype=torch.long) + talker._decode_has_codes = torch.zeros(num_rows, dtype=torch.bool) + talker._decode_delay_count = torch.zeros(num_rows, dtype=torch.long) + talker._decode_eoc_countdown = torch.full((num_rows,), -1, dtype=torch.long) + talker._decode_generation_done = torch.zeros(num_rows, dtype=torch.bool) + return talker + + @staticmethod + def _sampling_metadata(temperature=1.0, top_k=50, top_p=0.95, generators=None): return type( "SamplingMetadata", (), @@ -352,9 +368,16 @@ class TestSamplerMethods: "top_k": top_k, "top_p": top_p, "all_greedy": temperature is None, + "generators": generators or {}, }, )() + @staticmethod + def _seeded_generator(seed): + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + return generator + def test_sample_respects_mask(self): """Tokens masked to -inf must never be sampled.""" t = self._make_batched_sampler_talker(num_rows=1) @@ -442,6 +465,51 @@ class TestSamplerMethods: assert torch.count_nonzero(seen[0][0]).item() == 1 assert torch.count_nonzero(seen[0][2]).item() > 1 + def test_seeded_sampling_is_batch_independent(self): + t = self._make_minimal_talker() + num_codebooks = 8 + torch.manual_seed(123) + first_logits = torch.randn(num_codebooks, 64) + second_logits = torch.randn(num_codebooks, 64) + filler_logits = torch.randn(num_codebooks, 64) + reference_generator = self._seeded_generator(42) + batched_generators = { + 1: self._seeded_generator(42), + 2: self._seeded_generator(42), + } + + reference_first = t._sample_audio_codes( + first_logits, + self._sampling_metadata(generators={0: reference_generator}), + num_codebooks=num_codebooks, + ) + batched_first = t._sample_audio_codes( + torch.cat([filler_logits, first_logits, first_logits]), + self._sampling_metadata(generators=batched_generators), + num_codebooks=num_codebooks, + ) + assert torch.equal(reference_first, batched_first[num_codebooks : 2 * num_codebooks]) + assert torch.equal(reference_first, batched_first[2 * num_codebooks :]) + + reference_second = t._sample_audio_codes( + second_logits, + self._sampling_metadata(generators={0: reference_generator}), + num_codebooks=num_codebooks, + ) + condensed_second = t._sample_audio_codes( + torch.cat([second_logits, second_logits]), + self._sampling_metadata( + generators={ + 0: batched_generators[1], + 1: batched_generators[2], + } + ), + num_codebooks=num_codebooks, + ) + + assert torch.equal(reference_second, condensed_second[:num_codebooks]) + assert torch.equal(reference_second, condensed_second[num_codebooks:]) + def test_delay_masking_forces_boc_during_delay(self): """During delay phase, codebooks beyond delay_count must have only BOC allowed.""" from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID @@ -678,6 +746,58 @@ class TestSamplerMethods: assert not t._decode_generation_done[2:5].any() assert t._decode_eoc_countdown[2:5].eq(-1).all() + def test_decode_state_follows_request_after_batch_condensation(self): + talker = self._make_state_tracking_talker(3) + + talker._sync_decode_state_with_batch(["A", "B", "C"]) + c_codes = torch.arange(300, 308) + talker._decode_last_codes[:] = torch.stack((torch.arange(100, 108), torch.zeros(8, dtype=torch.long), c_codes)) + talker._decode_has_codes[:] = torch.tensor([True, False, True]) + talker._decode_delay_count[:] = torch.tensor([3, 0, 8]) + + talker.on_requests_finished({"B"}) + talker._sync_decode_state_with_batch(["A", "C"]) + + assert talker._decode_has_codes[1].item() is True + assert talker._decode_delay_count[1].item() == 8 + assert torch.equal(talker._decode_last_codes[1], c_codes) + + def test_finished_request_state_is_cleared_before_pool_slot_reuse(self): + talker = self._make_state_tracking_talker(1) + + talker._sync_decode_state_with_batch(["finished"]) + talker._decode_has_codes[0] = True + talker._decode_last_codes[0] = 7 + talker._decode_delay_count[0] = 5 + talker._decode_eoc_countdown[0] = 0 + talker._decode_generation_done[0] = True + talker.on_requests_finished({"finished"}) + talker._sync_decode_state_with_batch(["new"]) + + assert talker._decode_has_codes[0].item() is False + assert talker._decode_last_codes[0].eq(0).all() + assert talker._decode_delay_count[0].item() == 0 + assert talker._decode_eoc_countdown[0].item() == -1 + assert talker._decode_generation_done[0].item() is False + + def test_finished_request_state_is_cleared_when_request_id_is_reused(self): + talker = self._make_state_tracking_talker(1) + + talker._sync_decode_state_with_batch(["request"]) + talker._decode_has_codes[0] = True + talker._decode_last_codes[0] = 7 + talker._decode_delay_count[0] = 5 + talker._decode_eoc_countdown[0] = 0 + talker._decode_generation_done[0] = True + talker.on_requests_finished({"request"}) + talker._sync_decode_state_with_batch(["request"]) + + assert talker._decode_has_codes[0].item() is False + assert talker._decode_last_codes[0].eq(0).all() + assert talker._decode_delay_count[0].item() == 0 + assert talker._decode_eoc_countdown[0].item() == -1 + assert talker._decode_generation_done[0].item() is False + def test_mixed_batch_prefill_mask_targets_request_rows(self): """Mixed prefill/decode must reset only prefill request rows.""" t = self._make_batched_sampler_talker(num_rows=3) @@ -1326,7 +1446,9 @@ class TestRegistry: assert "higgs_multimodal_qwen3" in OMNI_PIPELINES - def test_deploy_yaml_exists(self): + def test_deploy_yaml_uses_native_flashinfer(self): + from vllm_omni.config.stage_config import load_deploy_config + for name in ( "higgs_multimodal_qwen3.yaml", "higgs_multimodal_qwen3_high_throughput.yaml", @@ -1334,6 +1456,10 @@ class TestRegistry: ): yaml_path = get_deploy_config_path(name) assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}" + deploy_config = load_deploy_config(yaml_path) + stage_0 = next(stage for stage in deploy_config.stages if stage.stage_id == 0) + assert stage_0.engine_extras["attention_backend"] == "FLASHINFER" + assert stage_0.engine_extras["attention_config"] == {"use_trtllm_attention": False} # ---- AC-3: Prompt Builder ---- diff --git a/tests/test_indextts2_optional_dependencies.py b/tests/model_executor/models/indextts2/test_optional_dependencies.py similarity index 99% rename from tests/test_indextts2_optional_dependencies.py rename to tests/model_executor/models/indextts2/test_optional_dependencies.py index cc3c46b63..24c6131fd 100644 --- a/tests/test_indextts2_optional_dependencies.py +++ b/tests/model_executor/models/indextts2/test_optional_dependencies.py @@ -14,7 +14,7 @@ from packaging.requirements import Requirement pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -ROOT = Path(__file__).parents[1] +ROOT = Path(__file__).parents[4] TEXT_PROCESSING_PATH = ROOT / "vllm_omni" / "model_executor" / "models" / "indextts2" / "text_processing_v2_5.py" INDEXTTS2_FRONTEND_PACKAGES = { "fugashi", diff --git a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py index 7a0ad68de..21c8dfd14 100644 --- a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py +++ b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py @@ -7,8 +7,14 @@ import pytest import torch from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig -from vllm_omni.model_executor.models.common.ming.fm import Solver +from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler +from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT +from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss +from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, get_epss_timesteps +from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor +from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor from vllm_omni.model_executor.models.ming_tts.constants import ( AGGREGATOR_HIDDEN_SIZE, HISTORY_PATCH_SIZE, @@ -24,7 +30,6 @@ from vllm_omni.model_executor.models.ming_tts.constants import ( SPEAKER_EMBEDDING_DIM, VAE_PATCH_SIZE, ) -from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel from vllm_omni.model_executor.models.ming_tts.speaker_extractor import _resolve_speaker_embeddings from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config @@ -44,7 +49,87 @@ def test_ming_tts_audio_vae_uses_common_config(): def test_ming_tts_cfm_solver_uses_common_implementation(): """Ming dense imports the shared solver implementation directly.""" - assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm" + assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.cfm_solver" + + +@pytest.mark.parametrize( + ("symbol", "expected_module"), + [ + (Aggregator, "vllm_omni.model_executor.models.common.ming.aggregator"), + (DiT, "vllm_omni.model_executor.models.common.ming.cfm_dit"), + (FlowLoss, "vllm_omni.model_executor.models.common.ming.cfm_head"), + (CFMGraphExecutor, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"), + (CFMSampler, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"), + (SpeakerEmbeddingExtractor, "vllm_omni.model_executor.models.common.ming.speaker_extractor"), + (SpkembExtractor, "vllm_omni.model_executor.models.common.ming.spk_embedding"), + ], +) +def test_ming_tts_reusable_modules_live_in_common(symbol, expected_module): + """The pieces a second Ming-family TTS model reuses live under common/ming. + + Guards the promotion: if an implementation drifts back into `ming_tts` + (or into another model package), a second model would have to reach across + packages to get it. + """ + assert symbol.__module__ == expected_module + + +def test_ming_flash_talker_uses_the_shared_aggregator(): + """Ming flash no longer carries its own copy of the aggregator. + + Both Ming packages ran near-identical aggregators one directory apart, + which is the duplication this module layout exists to remove. + """ + from vllm_omni.model_executor.models.ming_flash_omni import talker_module + + assert talker_module.Aggregator is Aggregator + assert not any( + isinstance(obj, type) and obj.__module__ == talker_module.__name__ and name == "Aggregator" + for name, obj in vars(talker_module).items() + ) + + +def test_shared_aggregator_matches_the_ming_flash_construction(): + """The shared aggregator reproduces Ming flash's former default geometry. + + `talker_module.Aggregator` defaulted to `in_channels=64` while the + shared class defaults to 4, so the call site now passes `latent_dim` + explicitly. This pins the resulting shapes rather than the default. + """ + latent_dim, llm_hidden = 64, 32 + aggregator = Aggregator( + in_channels=latent_dim, + llm_input_dim=llm_hidden, + hidden_size=16, + depth=1, + num_heads=2, + ) + + assert aggregator.x_embedder.in_features == latent_dim + assert aggregator.final_layer.linear.out_features == llm_hidden + + out = aggregator(torch.randn(1, 3, latent_dim)) + # Keeps only the CLS projection: [B, T, D] -> [B, 1, llm_hidden]. + assert out.shape == (1, 1, llm_hidden) + + +def test_shared_aggregator_rejects_mismatched_latent_width(): + """Wrong-width input fails loudly instead of silently projecting incorrect.""" + aggregator = Aggregator(in_channels=64, llm_input_dim=32, hidden_size=16, depth=1, num_heads=2) + + with pytest.raises(ValueError, match="feature dim mismatch"): + aggregator(torch.randn(1, 3, 4)) + + +def test_cfm_solver_owns_the_timestep_schedule(): + """get_epss_timesteps sits with the solver, not with the DiT blocks.""" + from vllm_omni.model_executor.models.common.ming import dit_blocks + + assert not hasattr(dit_blocks, "get_epss_timesteps") + + t = get_epss_timesteps(10, device="cpu", dtype=torch.float32) + assert t.shape == (11,) + assert float(t[0]) == 0.0 and float(t[-1]) == 1.0 def test_ming_tts_flowloss_preserves_checkpoint_prefix(): diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py new file mode 100644 index 000000000..af71c07ca --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""What is allowed to keep the CFM wrapper from capturing a graph. + +These run without a GPU: ``_capture`` is mocked, so nothing here touches a +CUDA context. The point is the decision made *before* capture, not the capture. +""" + +from typing import NamedTuple +from unittest.mock import Mock + +import pytest +import torch + +from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + CFMGraphWrapper, + _format_memory_delta, + _memory_snapshot, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _InputSignature(NamedTuple): + """The three properties ``replay`` reads from an input before capturing. + + Everything ``replay`` does ahead of the capture decision -- the device + check and the cache key -- comes from these, and ``_capture`` is mocked + out, so the tests need no CUDA tensors and no GPU. + """ + + shape: tuple[int, ...] + dtype: torch.dtype + device: torch.device + + +def _cfm_input_signatures(chunk_size: int) -> tuple[_InputSignature, ...]: + device = torch.device("cuda:0") + shapes = ((2, 16, chunk_size), (2, 1, 8), (2, 2, 8, 2), (2, 2, 1, 4, 8), (2, 2, 8, 2), (2, 2, 1, 12, 8)) + return tuple(_InputSignature(shape, torch.float32, device) for shape in shapes) + + +def _wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 4) -> CFMGraphWrapper: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = object.__new__(CFMGraphWrapper) + wrapper.max_graphs = max_graphs + wrapper.enabled = True + wrapper.graph_fn = Mock(return_value=torch.tensor([42.0])) + wrapper.device = torch.device("cuda:0") + wrapper._cache = {} + wrapper._unsupported = set() + wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0} + wrapper._capture = Mock(return_value=None) + return wrapper + + +def test_capture_is_not_gated_on_free_device_memory(monkeypatch: pytest.MonkeyPatch) -> None: + """Whether to capture must not be decided from free device memory. + + A capture is served by the caching allocator out of its private graph + pool, and the allocator falls back on memory it has already reserved. Free + device memory is therefore neither an upper nor a lower bound on what a + capture can draw, and a threshold on it decides the wrong question -- one + whose answer also silently varies with card size and with whatever else + shares the device. Anything consulting it here would fail this test. + """ + import vllm_omni.platforms as platforms_module + + def _saturated(device: torch.device | None = None) -> int: + raise AssertionError("capture must not depend on device-level free memory") + + # Setting the attribute on the module shadows the lazy ``__getattr__``, so + # this neither builds nor requires a real platform singleton. + monkeypatch.setattr( + platforms_module, + "current_omni_platform", + Mock(get_free_memory=_saturated), + raising=False, + ) + + wrapper = _wrapper(monkeypatch) + wrapper.replay(*_cfm_input_signatures(10)) + + wrapper._capture.assert_called_once() + + +def test_memory_reporting_is_inert_off_cuda() -> None: + """Capture logging must not break the platforms that have no such counters.""" + assert _memory_snapshot(torch.device("cpu")) is None + assert _format_memory_delta(None, (1, 2)) == "" + assert _format_memory_delta((1, 2), None) == "" + assert "allocated" in _format_memory_delta((0, 0), (1 << 20, 2 << 20)) diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index ba3c7292b..a0e7ed4db 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -1,3 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import subprocess +import sys +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -5,8 +11,12 @@ import pytest import torch import torch.nn as nn +import vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav as batched_token2wav_module from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( BatchedToken2Wav, + _token2wav_sdpa_context, + plan_token2wav_encode_slices, + relpos_encode_token_budget, ) from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( MiniCPMO45Code2Wav, @@ -50,6 +60,9 @@ class _FakeEstimator(nn.Module): self.blocks = [_FakeBlock()] self.cfg_batches: list[int] = [] self.speaker_order: list[list[float]] = [] + self.attention_cache_dtypes: list[torch.dtype | None] = [] + self.register_buffer("att_cache_buffer", torch.ones(1), persistent=False) + self.register_buffer("cnn_cache_buffer", torch.ones(1), persistent=False) def t_embedder(self, time): return time[:, None] @@ -64,7 +77,8 @@ class _FakeEstimator(nn.Module): cnn_out, att_out, ): - del time, mask, cnn_cache, att_cache + del time, mask, cnn_cache + self.attention_cache_dtypes.append(att_cache.dtype if isinstance(att_cache, torch.Tensor) else None) self.cfg_batches.append(inputs.shape[0]) self.speaker_order.append(inputs[:, 2, 0].tolist()) marker = inputs[:, 1, 0] @@ -98,16 +112,17 @@ class _FakeHiFT(nn.Module): super().__init__() self.calls: list[int] = [] - def forward(self, mel, source): + def forward(self, batch, device): + mel = batch["speech_feat"].transpose(1, 2).to(device) + return self.inference(mel, mel.new_zeros((mel.shape[0], 1, 0))) + + def inference(self, mel, source): del source self.calls.append(mel.shape[0]) speech = mel[:, 0].repeat_interleave(3, dim=1) generated_source = speech[:, None] return speech, generated_source - def inference(self, mel, source): - return self.forward(mel, source) - class _FakeToken2Wav: def __init__(self): @@ -138,13 +153,14 @@ class _FakeToken2Wav: raise AssertionError("sequential __call__ fallback must never be called") -def _config(minimum: int = 1): +def _config(minimum: int = 1, initial: int = 0): return SimpleNamespace( model_config=SimpleNamespace( model="/fake/model", stage_connector_config={ "extra": { "code2wav_min_batch_size": minimum, + "code2wav_initial_batch_size": initial, "prompt_cache_id": "shared", "prompt_wav": "/fake/prompt.wav", } @@ -153,14 +169,45 @@ def _config(minimum: int = 1): ) -def _model(): +def _model(initial: int = 0, minimum: int = 1): token2wav = _FakeToken2Wav() backend = BatchedToken2Wav(token2wav) - model = MiniCPMO45Code2Wav(vllm_config=_config()) + _enable_fake_ragged_kernel(backend) + model = MiniCPMO45Code2Wav(vllm_config=_config(minimum=minimum, initial=initial)) model.backend = backend return model, token2wav +def _enable_fake_ragged_kernel(adapter: BatchedToken2Wav) -> None: + """Keep fake-model tests focused on ragged grouping and row mapping.""" + estimator = adapter.flow.decoder.estimator + estimator.in_proj = nn.Identity() + + def fake_ragged_kernel( + estimator, + estimator_input, + time_embedding, + attn_mask, + cnn_cache, + att_cache, + cnn_cache_buffer, + att_cache_buffer, + valid_lengths, + ): + del valid_lengths + return estimator.blocks_forward_chunk( + estimator_input, + time_embedding, + attn_mask, + cnn_cache, + att_cache, + cnn_cache_buffer, + att_cache_buffer, + ) + + adapter._blocks_forward_chunk_ragged = fake_ragged_kernel # type: ignore[method-assign] + + def _info( request_id: str, chunk_seq: int, @@ -217,6 +264,437 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches(): assert cache1[0, 0, 0, 0, 0].item() == 20 +def test_relpos_encode_token_budget_leaves_room_for_upsample_and_cache(): + # The NPU crash was 6968 vs 985: 3333 tokens * 2 + cache overflowed max_pos=5000. + assert relpos_encode_token_budget(max_pos=5000, stride=2, cache_offset=150, lookahead=3) <= 1024 + assert relpos_encode_token_budget(max_pos=5000, stride=2, cache_offset=150, lookahead=3, cap=3000) == 2346 + assert plan_token2wav_encode_slices(3333, max_frames=1024, min_nonfinal=4, last_chunk=True) == [ + (0, 1024), + (1024, 2048), + (2048, 3072), + (3072, 3333), + ] + assert plan_token2wav_encode_slices(50, max_frames=40, min_nonfinal=20, last_chunk=False) == [ + (0, 30), + (30, 50), + ] + + +def test_decode_batch_splits_overlong_prefill_inside_relpos_budget(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 1) + setup_encodes = len(token2wav.flow.encoder.calls) + adapter._max_encode_token_frames = lambda _states: 4 + + audios, _ = adapter.decode_batch( + torch.arange(10, dtype=torch.long).reshape(1, -1), + prompt, + states, + last_chunk=True, + ) + + assert token2wav.flow.encoder.calls[setup_encodes:] == [1, 1, 1] + assert token2wav.flow.encoder.last_chunk_calls[setup_encodes:] == [False, False, True] + assert len(audios) == 1 + assert audios[0].numel() > 0 + + +def test_ragged_overlong_prefill_uses_relpos_safe_exact_slices(): + token2wav = _FakeToken2Wav() + token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3) + adapter = BatchedToken2Wav(token2wav) + _enable_fake_ragged_kernel(adapter) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 2) + setup_encodes = len(token2wav.flow.encoder.calls) + adapter._max_encode_token_frames = lambda _states: 4 + + tokens = [ + torch.arange(10, dtype=torch.long), + torch.arange(4, dtype=torch.long), + ] + last_chunks = [True, False] + audios, next_states = adapter.decode_ragged_batch( + tokens, + prompt, + states, + last_chunks=last_chunks, + ) + + assert token2wav.flow.encoder.calls[setup_encodes:] == [1, 1, 1, 1] + assert token2wav.flow.encoder.last_chunk_calls[setup_encodes:] == [False, False, True, False] + assert len(audios) == len(next_states) == 2 + assert all(audio.numel() > 0 for audio in audios) + + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + reference_token2wav = _FakeToken2Wav() + reference_token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3) + reference_adapter = BatchedToken2Wav(reference_token2wav) + reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav") + reference_states = reference_adapter.setup_batch(reference_prompt, 1) + reference_adapter._max_encode_token_frames = lambda _states: 4 + reference_audio, reference_state = reference_adapter.decode_batch( + row_tokens.unsqueeze(0), + reference_prompt, + reference_states, + last_chunk=last_chunk, + ) + torch.testing.assert_close(audios[row], reference_audio[0]) + for cache_name, expected in reference_state[0].flow_cache.items(): + torch.testing.assert_close( + next_states[row].flow_cache[cache_name], + expected, + ) + for cache_name, expected in reference_state[0].hift_cache.items(): + torch.testing.assert_close( + next_states[row].hift_cache[cache_name], + expected, + ) + + +def test_ragged_diffusion_batches_eight_rows_and_preserves_exact_results(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav) + _enable_fake_ragged_kernel(adapter) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 8) + token2wav.flow.encoder.calls.clear() + token2wav.flow.decoder.estimator.cfg_batches.clear() + token2wav.hift.calls.clear() + tokens = [ + *(torch.tensor([10 + row, 11 + row, 12 + row]) for row in range(6)), + torch.tensor([20, 21]), + torch.tensor([30, 31]), + ] + last_chunks = [False] * 6 + [True, True] + + audios, next_states = adapter.decode_ragged_batch( + tokens, + prompt, + states, + last_chunks=last_chunks, + ) + + assert token2wav.flow.encoder.calls == [6, 2] + assert token2wav.flow.decoder.estimator.cfg_batches == [16, 16] + assert token2wav.hift.calls == [6, 2] + assert len(audios) == len(next_states) == 8 + + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + reference_adapter = BatchedToken2Wav(_FakeToken2Wav()) + reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav") + reference_states = reference_adapter.setup_batch(reference_prompt, 1) + reference_audio, reference_state = reference_adapter.decode_batch( + row_tokens.unsqueeze(0), + reference_prompt, + reference_states, + last_chunk=last_chunk, + ) + torch.testing.assert_close(audios[row], reference_audio[0]) + for cache_name, expected in reference_state[0].flow_cache.items(): + torch.testing.assert_close( + next_states[row].flow_cache[cache_name], + expected, + ) + for cache_name, expected in reference_state[0].hift_cache.items(): + torch.testing.assert_close( + next_states[row].hift_cache[cache_name], + expected, + ) + + +def test_npu_patch_preserves_ragged_estimator_contract(): + script = r""" +import torch + +from tests.model_executor.models.minicpmo_4_5.test_code2wav_batching import ( + _FakeToken2Wav, + _enable_fake_ragged_kernel, +) +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav +import vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav as npu_patch + + +class _FailGraphRunner: + def run(self, *args, **kwargs): + raise AssertionError("ragged decode must bypass the exact-shape NPU Graph") + + +npu_patch.apply_minicpmo_4_5_code2wav_patch() +adapter = BatchedToken2Wav(_FakeToken2Wav()) +_enable_fake_ragged_kernel(adapter) +prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") +states = adapter.setup_batch(prompt, 2) +npu_patch._backend_graph_runners[adapter] = _FailGraphRunner() +audios, next_states = adapter.decode_ragged_batch( + [ + torch.tensor([10, 11, 12]), + torch.tensor([20, 21]), + ], + prompt, + states, + last_chunks=[False, True], +) +assert len(audios) == len(next_states) == 2 +""" + completed = subprocess.run( + [sys.executable, "-c", script], + cwd=Path(__file__).resolve().parents[4], + capture_output=True, + text=True, + ) + + assert completed.returncode == 0, completed.stderr + + +def test_decode_cfm_enters_platform_sdpa_context(monkeypatch): + entered: list[str] = [] + + @contextmanager + def recording_context(): + entered.append("enter") + yield + entered.append("exit") + + monkeypatch.setattr( + batched_token2wav_module, + "_token2wav_sdpa_context", + lambda _device: recording_context(), + ) + adapter = BatchedToken2Wav(_FakeToken2Wav()) + _enable_fake_ragged_kernel(adapter) + adapter._decode_cfm( + torch.ones((2, 1, 2)), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + valid_lengths=[2, 1], + ) + + assert entered == ["enter", "exit"] + + +def test_ragged_decode_bypasses_exact_shape_accelerators(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + _enable_fake_ragged_kernel(adapter) + + class ForbiddenAccelerator: + def replay(self, *args): + pytest.fail("ragged decode must not use the exact-shape CFM graph") + + def step(self, **kwargs): + pytest.fail("ragged decode must not use the exact-shape TensorRT stepper") + + accelerator = ForbiddenAccelerator() + adapter._cfm_graph_wrapper = accelerator + adapter._trt_stepper = accelerator + adapter._decode_cfm( + torch.ones((2, 1, 2)), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + valid_lengths=[2, 1], + ) + + +def test_ragged_decode_fails_when_estimator_has_no_ragged_kernel(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + + with pytest.raises( + RuntimeError, + match=r'"reason":"ragged_kernel_unavailable"', + ): + adapter._decode_cfm( + torch.ones((2, 1, 2)), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + valid_lengths=[2, 1], + ) + + +def test_cfm_graph_receives_bfloat16_cache_in_compute_dtype(): + adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True) + seen_dtypes: list[torch.dtype] = [] + + class RecordingGraph: + def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out): + del time_embedding, cnn_cache + seen_dtypes.append(att_cache.dtype) + return estimator_input[:, :1], cnn_out, att_out + + adapter._cfm_graph_wrapper = RecordingGraph() + adapter._estimator_step( + adapter.flow.decoder.estimator, + x=torch.ones((2, 1, 2)), + mu=torch.ones((2, 1, 2)), + time=torch.zeros(2), + speakers=torch.ones((2, 1)), + cond=torch.zeros((2, 1, 2)), + cnn_cache=torch.zeros((1, 2, 2, 1)), + att_cache=torch.zeros((1, 2, 1, 1, 2), dtype=torch.bfloat16), + ) + + assert seen_dtypes == [torch.float32] + + +def test_npu_sdpa_context_uses_platform_patch(monkeypatch): + entered: list[str] = [] + + @contextmanager + def recording_context(): + entered.append("enter") + yield + entered.append("exit") + + monkeypatch.setitem( + sys.modules, + "vllm_omni.platforms.npu.models.step_audio2_token2wav", + SimpleNamespace(npu_token2wav_sdpa_context=recording_context), + ) + with _token2wav_sdpa_context(SimpleNamespace(type="npu")): + assert entered == ["enter"] + + assert entered == ["enter", "exit"] + + +def test_ragged_outputs_fail_closed_on_middle_hole(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 3) + audios: list[torch.Tensor | None] = [ + torch.zeros(1), + None, + torch.ones(1), + ] + next_states = [states[0], None, states[2]] + + with pytest.raises( + RuntimeError, + match=r'"reason":"incomplete_ragged_output","audio_rows":\[1\],"state_rows":\[1\]', + ): + adapter._require_complete_ragged_outputs(audios, next_states) + + +@pytest.mark.parametrize( + "bfloat16_attention_cache", + [False, True], + ids=["fp32-cache", "bf16-cache"], +) +def test_ragged_dit_two_step_decode_matches_per_row_exact_decode( + bfloat16_attention_cache: bool, +): + from cosyvoice2.flow.decoder_dit import DiT + + torch.manual_seed(17) + template = DiT( + in_channels=4, + out_channels=1, + depth=2, + num_heads=2, + head_dim=2, + hidden_size=4, + ).eval() + template_state = {name: value.detach().clone() for name, value in template.state_dict().items()} + rand_noise = torch.randn(1, 1, 32) + + def make_adapter() -> BatchedToken2Wav: + token2wav = _FakeToken2Wav() + estimator = DiT( + in_channels=4, + out_channels=1, + depth=2, + num_heads=2, + head_dim=2, + hidden_size=4, + ).eval() + estimator.load_state_dict(template_state) + token2wav.flow.decoder.estimator = estimator + token2wav.flow.decoder.rand_noise = rand_noise.clone() + return BatchedToken2Wav( + token2wav, + bfloat16_attention_cache=bfloat16_attention_cache, + ) + + ragged_adapter = make_adapter() + ragged_prompt = ragged_adapter.prepare_prompt("shared", "/fake/prompt.wav") + ragged_states = ragged_adapter.setup_batch(ragged_prompt, 2) + + exact_adapters = [make_adapter(), make_adapter()] + exact_prompts = [adapter.prepare_prompt("shared", "/fake/prompt.wav") for adapter in exact_adapters] + exact_states = [ + adapter.setup_batch(prompt, 1) for adapter, prompt in zip(exact_adapters, exact_prompts, strict=True) + ] + tolerance = (2e-2, 2e-3) if bfloat16_attention_cache else (1e-5, 1e-6) + steps = [ + ( + [torch.tensor([10, 11, 12]), torch.tensor([20, 21, 22])], + [False, True], + ), + ( + [torch.tensor([13, 14]), torch.tensor([22, 23, 24])], + [True, True], + ), + ] + first_step_history_width: int | None = None + + for step, (tokens, last_chunks) in enumerate(steps): + if step == 1: + assert first_step_history_width is not None + assert all( + int(state.flow_cache["estimator_att_cache"].shape[4]) == first_step_history_width + for state in ragged_states + ) + ragged_audios, ragged_states = ragged_adapter.decode_ragged_batch( + tokens, + ragged_prompt, + ragged_states, + last_chunks=last_chunks, + ) + for row, (adapter, prompt) in enumerate(zip(exact_adapters, exact_prompts, strict=True)): + exact_audios, exact_states[row] = adapter.decode_batch( + tokens[row].unsqueeze(0), + prompt, + exact_states[row], + last_chunk=last_chunks[row], + ) + torch.testing.assert_close( + ragged_audios[row], + exact_audios[0], + rtol=tolerance[0], + atol=tolerance[1], + ) + for cache_name, expected in exact_states[row][0].flow_cache.items(): + actual = ragged_states[row].flow_cache[cache_name] + assert actual.dtype == expected.dtype + torch.testing.assert_close( + actual, + expected, + rtol=tolerance[0], + atol=tolerance[1], + ) + for cache_name, expected in exact_states[row][0].hift_cache.items(): + torch.testing.assert_close( + ragged_states[row].hift_cache[cache_name], + expected, + rtol=tolerance[0], + atol=tolerance[1], + ) + + estimator_att = ragged_states[0].flow_cache["estimator_att_cache"] + assert estimator_att.shape[4] > 0 + if step == 0: + first_step_history_width = int(estimator_att.shape[4]) + else: + assert first_step_history_width is not None + assert int(estimator_att.shape[4]) > first_step_history_width + + def test_fade_in_out_limits_overlap_to_available_previous_audio(): speech = torch.arange(6, dtype=torch.float32).reshape(1, -1) previous = torch.full((1, 3), 2.0) @@ -256,6 +734,69 @@ def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows(): ) +def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav, bfloat16_attention_cache=True) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 2) + + assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in states) + assert ( + states[0].flow_cache["estimator_att_cache"].data_ptr() != states[1].flow_cache["estimator_att_cache"].data_ptr() + ) + + stacked = adapter._stack_flow_cache(states) + assert stacked["estimator_att_cache"].dtype == torch.bfloat16 + restored = adapter._split_flow_cache(stacked, 2) + for original, round_tripped in zip(states, restored, strict=True): + assert round_tripped["estimator_att_cache"].dtype == torch.bfloat16 + torch.testing.assert_close( + round_tripped["estimator_att_cache"], + original.flow_cache["estimator_att_cache"], + ) + + _, next_states = adapter.decode_batch( + torch.tensor([[10, 11], [20, 21]]), + prompt, + states, + last_chunk=False, + ) + assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in next_states) + assert token2wav.flow.decoder.estimator.attention_cache_dtypes[-2:] == [torch.float32, torch.float32] + + +def test_bfloat16_estimator_attention_cache_materializes_each_timestep(): + adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True) + + _, estimator_cnn, estimator_att = adapter._decode_cfm( + torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + ) + + assert estimator_cnn.dtype == torch.float32 + assert estimator_att.dtype == torch.bfloat16 + assert estimator_att.shape[:3] == (2, 1, 4) + + +def test_float32_estimator_attention_cache_materializes_each_timestep(): + adapter = BatchedToken2Wav(_FakeToken2Wav()) + + _, estimator_cnn, estimator_att = adapter._decode_cfm( + torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]), + torch.ones((2, 1)), + torch.zeros((2, 1, 2)), + cnn_cache=None, + att_cache=None, + ) + + assert estimator_cnn.dtype == torch.float32 + assert estimator_att.dtype == torch.float32 + assert estimator_att.shape[:3] == (2, 1, 4) + + def test_model_preserves_output_slots_and_prefers_runtime_codes(): model, token2wav = _model() output = _forward( @@ -347,6 +888,28 @@ def test_initial_empty_segment_marker_initializes_stream_without_audio(): assert "duplex" in model._states +def test_initial_empty_segment_markers_respect_initial_batch_limit(): + model, token2wav = _model(initial=2) + names = ["a", "b", "c", "d", "e"] + boundaries = [] + for name in names: + boundary = _info(name, 0, []) + boundary["meta"].update( + { + "code_flat_numel": 0, + "tts_is_last_chunk": True, + "turn_end": False, + } + ) + boundaries.append(boundary) + + output = _forward(model, boundaries) + + assert token2wav.flow.encoder.calls == [2, 2, 1] + assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0] * len(names) + assert set(model._states) == set(names) + + def test_shared_runtime_prompt_recreates_missing_file_before_second_owner(tmp_path, monkeypatch): monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path)) model, _ = _model() @@ -498,6 +1061,15 @@ def test_empty_final_ignores_generation_scheduler_placeholder_token(): # The runner injects the engine request id on every step (GPU # _preprocess, NPU _gather_runtime_additional_information)... {"request_id": "a", "meta": {"request_id": "a"}}, + # Runtime snapshot bookkeeping alone is not a Talker payload. + { + "request_id": "a", + "meta": { + "request_id": "a", + "num_processed_tokens": 0, + "resumable": True, + }, + }, # ...but a pre-warm step can also reach the model with nothing at all. {}, ], @@ -599,11 +1171,56 @@ def test_singleton_and_mixed_shape_buckets_use_same_batched_backend_without_fall assert token2wav.hift.calls[-2:] == [1, 1] +def test_initial_batch_limit_allows_later_full_batch(): + model, token2wav = _model(initial=4) + names = [f"request-{index}" for index in range(8)] + + for chunk_seq in (0, 1): + infos = [_info(name, chunk_seq, [1, 2]) for name in names] + _forward(model, infos) + assert token2wav.hift.calls[-2:] == [4, 4] + + token2wav.hift.calls.clear() + _forward(model, [_info(name, 2, [1, 2]) for name in names]) + assert token2wav.hift.calls == [8] + + +def test_initial_batch_partition_preserves_minimum(): + model, token2wav = _model(initial=4, minimum=3) + names = [f"request-{index}" for index in range(6)] + + _forward(model, [_info(name, 0, [1, 2]) for name in names]) + + assert token2wav.hift.calls == [3, 3] + + +def test_initial_batch_partition_rejects_impossible_remainder(): + model, _ = _model(initial=4, minimum=3) + names = [f"request-{index}" for index in range(5)] + + with pytest.raises(RuntimeError, match="initial_batch_partition_below_minimum"): + _forward(model, [_info(name, 0, [1, 2]) for name in names]) + + +@pytest.mark.parametrize("chunk_seqs", [[0, 2, 2, 2], [0, 0, 0, 2]]) +def test_initial_steady_partition_rejects_undersized_wave(chunk_seqs): + model, _ = _model(initial=4, minimum=3) + bucket = [SimpleNamespace(chunk_seq=chunk_seq) for chunk_seq in chunk_seqs] + + with pytest.raises(RuntimeError, match="decode_wave_below_minimum"): + list(model._iter_decode_batches([bucket])) + + def test_backend_failure_does_not_commit_any_request_state(monkeypatch): model, _ = _model() _forward( model, - [_info(name, 0, [index + 1, index + 2]) for index, name in enumerate(("a", "b", "c", "d"))], + [ + _info("a", 0, [1, 2], cache_epoch=0), + _info("b", 0, [2, 3], cache_epoch=0), + _info("c", 0, [3, 4], cache_epoch=1), + _info("d", 0, [4, 5], cache_epoch=1), + ], ) before = dict(model._states) original = model.backend.decode_batch @@ -623,8 +1240,8 @@ def test_backend_failure_does_not_commit_any_request_state(monkeypatch): [ _info("a", 1, [5, 6]), _info("b", 1, [7, 8]), - _info("c", 1, [9, 10, 11]), - _info("d", 1, [12, 13, 14]), + _info("c", 1, [9, 10], cache_epoch=1), + _info("d", 1, [12, 13], cache_epoch=1), ], ) assert call_count == 2 diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 7ad35b86d..26daeff7b 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -from functools import lru_cache +import sys from types import SimpleNamespace from unittest.mock import Mock @@ -10,6 +10,7 @@ import torch import torch.nn as nn from vllm.platforms import current_platform +import vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper as wrapper_module from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( HiFTGenerator, ) @@ -248,50 +249,203 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes( torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5) torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5) - wrapper._capture_graph.cache_clear() + wrapper._flush() + + +def test_importing_wrapper_does_not_resolve_platform() -> None: + """Importing this module must not build the OmniPlatform singleton. + + The NPU platform's ``__init__`` patches Code2Wav, which imports + ``batched_token2wav`` -> ``cuda_graph_wrapper``. Resolving + ``current_omni_platform`` at module scope re-enters + ``platforms.__getattr__`` while the singleton is still under construction + and the import graph deadlocks. The module reaches for nothing on + ``vllm_omni.platforms`` today; this keeps it that way. + """ + import importlib + + import vllm_omni.platforms as platforms_module + + name = "vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper" + importlib.import_module(name) + saved = platforms_module._current_omni_platform + try: + del sys.modules[name] + platforms_module._current_omni_platform = None + importlib.import_module(name) + assert platforms_module._current_omni_platform is None + finally: + platforms_module._current_omni_platform = saved def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) wrapper = object.__new__(CFMGraphWrapper) wrapper.max_graphs = max_graphs + wrapper.enabled = True wrapper.graph_fn = Mock(return_value=torch.tensor([42.0])) wrapper.device = torch.device("cuda") - - @lru_cache(maxsize=max_graphs) - def _capture_graph(key: tuple): - return wrapper._do_capture(key) - - wrapper._capture_graph = _capture_graph - wrapper._do_capture = Mock(return_value=None) + wrapper._cache = {} + wrapper._unsupported = set() + wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0} + wrapper._capture = Mock(return_value=None) return wrapper def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None: wrapper = _cfm_mock_wrapper(monkeypatch) - wrapper._capture_graph = Mock(return_value=None) inputs = _cfm_inputs(2, 10, 0) wrapper.replay(*inputs) - wrapper._capture_graph.assert_called_once() + wrapper._capture.assert_called_once() wrapper.graph_fn.assert_called_once() -def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None: +def test_cfm_returning_no_entry_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None: + """A capture that yields no entry must still serve the request eagerly. + + ``_capture`` is mocked here, so this says nothing about ``_disable``; see + ``test_cfm_capture_failure_disables_further_capture`` for that. + """ wrapper = _cfm_mock_wrapper(monkeypatch) - wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None) inputs = _cfm_inputs(2, 10, 0) result = wrapper.replay(*inputs) wrapper.graph_fn.assert_called_once() assert result[0] is wrapper.graph_fn.return_value + assert wrapper._stats["eager"] == 1 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None: - """Verify that LRU cache evicts oldest entries when full.""" +def test_cfm_capture_failure_disables_further_capture(monkeypatch: pytest.MonkeyPatch) -> None: + """A real capture failure must stop the wrapper capturing for good. + + A failed capture can leave the capture stream current and the allocator + still routing into the graph pool, so the next shape must not try again. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4) + + def _explode(*args: object, **kwargs: object) -> None: + raise RuntimeError("capture failed") + + monkeypatch.setattr(torch.cuda, "graph", _explode) + + with torch.inference_mode(): + wrapper.replay(*_cfm_inputs(2, 10, 0)) + assert wrapper.enabled is False + assert wrapper._cache == {} + + captures_after_failure = wrapper._stats["captures"] + wrapper.replay(*_cfm_inputs(2, 12, 0)) + + assert wrapper._stats["captures"] == captures_after_failure + assert wrapper._stats["eager"] == 2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.MonkeyPatch) -> None: + """A key that cannot round-trip is a property of one shape, not of the GPU. + + It must not disable the wrapper or retire the generation the way a capture + failure does. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _MiniDiT().eval().cuda() + wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4) + + unbuildable_width = 10 + real_tensors_from_key = wrapper_module._tensors_from_key + + def _reject_one_shape(key: tuple) -> tuple: + if key[1][0][2] == unbuildable_width: + raise KeyError("torch.int64") + return real_tensors_from_key(key) + + monkeypatch.setattr(wrapper_module, "_tensors_from_key", _reject_one_shape) + + with torch.inference_mode(): + wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0)) + assert wrapper.enabled is True + assert wrapper._stats["captures"] == 0 + + # a capturable shape still gets a graph + wrapper.replay(*_cfm_inputs(2, 12, 0)) + assert wrapper._stats["captures"] == 1 + + # and the rejected shape stays eager without retrying the capture + wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0)) + assert wrapper._stats["captures"] == 1 + assert wrapper._stats["eager"] == 2 + + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_hift_replay_survives_a_cfm_generation_flush(monkeypatch: pytest.MonkeyPatch) -> None: + """Retiring a CFM generation must not disturb the vocoder. + + Both wrappers capture into ``get_global_graph_pool()``, so the HiFT graphs + are exactly the live graphs a CFM flush could strand. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + hift = _small_hift() + token2wav = SimpleNamespace( + hift=hift, + flow=SimpleNamespace( + encoder=SimpleNamespace(pre_lookahead_layer=SimpleNamespace(pre_lookahead_len=3)), + token_mel_ratio=2, + ), + mel_cache_len=2, + source_cache_len=960, + ) + hift_wrapper = HiFTGraphWrapper( + token2wav, + connector_config={"codec_chunk_frames": 2, "codec_left_context_frames": 3}, + capture_batch_sizes=[1], + ) + hift_wrapper.capture() + + estimator = _MiniDiT().eval().cuda() + cfm = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2) + + speech_feat = torch.randn(1, 80, 4, device="cuda") + cache_source = torch.zeros(1, 1, 0, device="cuda") + + with torch.inference_mode(): + expected_speech, expected_source = hift.inference(speech_feat, cache_source) + + for width in (10, 12, 14, 16): + cfm.replay(*_cfm_inputs(2, width, 0)) + assert cfm._stats["flushes"] >= 1, "cache never overflowed; the test proves nothing" + + actual_speech, actual_source = hift_wrapper.replay(speech_feat, cache_source) + + torch.testing.assert_close(actual_speech, expected_speech, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(actual_source, expected_source, rtol=1e-4, atol=1e-5) + + cfm._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cfm_cache_flushes_whole_generation(monkeypatch: pytest.MonkeyPatch) -> None: + """A full cache is retired all at once, never one graph at a time. + + Destroying a single graph while its peers stay live strands them, so the + cache must never hold a graph that outlived one of its generation-mates. + """ pool = torch.cuda.graph_pool_handle() monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) @@ -302,26 +456,32 @@ def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None: with torch.inference_mode(): inputs_a = _cfm_inputs(2, 10, 0) wrapper.replay(*inputs_a) - assert wrapper._capture_graph.cache_info().currsize == 1 + assert len(wrapper._cache) == 1 inputs_b = _cfm_inputs(2, 12, 0) wrapper.replay(*inputs_b) - assert wrapper._capture_graph.cache_info().currsize == 2 + assert len(wrapper._cache) == 2 + # A hit must not grow the cache or trigger a flush. wrapper.replay(*inputs_a) - assert wrapper._capture_graph.cache_info().currsize == 2 + assert len(wrapper._cache) == 2 + assert wrapper._stats["flushes"] == 0 + assert wrapper._stats["hits"] == 1 + # The third distinct shape flushes the generation, then captures alone. inputs_c = _cfm_inputs(2, 14, 0) wrapper.replay(*inputs_c) - assert wrapper._capture_graph.cache_info().currsize == 2 + assert wrapper._stats["flushes"] == 1 + assert len(wrapper._cache) == 1 - wrapper.replay(*inputs_b) - hits_before = wrapper._capture_graph.cache_info().hits - wrapper.replay(*inputs_b) - hits_after = wrapper._capture_graph.cache_info().hits - assert hits_after == hits_before + 1 + # The flushed shapes are gone, so they capture again rather than hit. + hits_before = wrapper._stats["hits"] + wrapper.replay(*inputs_a) + assert wrapper._stats["hits"] == hits_before + wrapper.replay(*inputs_a) + assert wrapper._stats["hits"] == hits_before + 1 - wrapper._capture_graph.cache_clear() + wrapper._flush() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") diff --git a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py index 9c4b50967..3fd4867cb 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py +++ b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the MiniCPM-o 4.5 stage 0 -> stage 1 bridge. Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2tts``: @@ -25,6 +25,7 @@ import torch from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, ) +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import llm2tts pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -313,6 +314,44 @@ class TestPromptAndMultiModal: ) assert out[0]["multi_modal_data"] == mm + def test_reference_audio_is_added_to_stage_handoff(self) -> None: + hidden_states = torch.zeros((2, _HIDDEN_DIM)) + waveform = torch.tensor([0.25, -0.5], dtype=torch.float32) + + stage_outputs = llm2tts( + [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)], + prompt={"multi_modal_data": {"audio": (waveform, 16000)}}, + ) + + intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"] + assert intermediate_buffer["codes"]["ref"] == [0.25, -0.5] + assert intermediate_buffer["meta"]["ref_audio_sr"] == 16000 + + def test_serving_reference_audio_is_added_to_stage_handoff(self) -> None: + hidden_states = torch.zeros((2, _HIDDEN_DIM)) + waveform = torch.tensor([0.125, -0.25], dtype=torch.float32) + + stage_outputs = llm2tts( + [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)], + prompt={MINICPMO45_REFERENCE_AUDIO_KEY: (waveform, 24000)}, + ) + + intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"] + assert intermediate_buffer["codes"]["ref"] == [0.125, -0.25] + assert intermediate_buffer["meta"]["ref_audio_sr"] == 24000 + + def test_missing_reference_audio_does_not_add_stage_handoff_fields(self) -> None: + hidden_states = torch.zeros((2, _HIDDEN_DIM)) + + stage_outputs = llm2tts( + [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)], + prompt={"multi_modal_data": {}}, + ) + + intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"] + assert "codes" not in intermediate_buffer + assert "ref_audio_sr" not in intermediate_buffer.get("meta", {}) + def test_internal_streaming_context_is_accepted(self) -> None: hidden = torch.zeros((2, _HIDDEN_DIM)) out = llm2tts( diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 148f28b16..a05e4076f 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the MiniCPM-o 4.5 pipeline registration. Covers: @@ -19,11 +19,15 @@ from pathlib import Path from types import SimpleNamespace import pytest +from transformers import PretrainedConfig +from vllm.config import ModelConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, + StagePipelineConfig, + _apply_platform_overrides, load_deploy_config, merge_pipeline_deploy, ) @@ -62,7 +66,9 @@ class TestRegistryDeclaration: assert not hasattr(pipeline, "max_native_duplex_sessions") def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None: - pipeline = PipelineConfig(model_type="ordinary") + pipeline = PipelineConfig( + model_type="ordinary", stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),) + ) assert pipeline.duplex_control_enabled is False assert pipeline.duplex_serving_adapter is None assert not hasattr(pipeline, "max_native_duplex_sessions") @@ -77,10 +83,6 @@ class TestPipelineTopology: assert len(pipeline.stages) == 3 assert [s.stage_id for s in pipeline.stages] == [0, 1, 2] - def test_topology_validates(self, pipeline: PipelineConfig) -> None: - # ``validate`` returns a list of structural errors; empty == valid. - assert pipeline.validate() == [] - def test_thinker_stage(self, pipeline: PipelineConfig) -> None: thinker = pipeline.get_stage(0) assert thinker is not None @@ -104,6 +106,10 @@ class TestPipelineTopology: assert talker.engine_output_type == "latent" # scope KV cache / mrope sizing to talker sub-config assert talker.hf_config_name == "tts_config" + assert talker.sampling_constraints == { + "detokenize": False, + "stop_token_ids": [6561], + } assert talker.custom_process_next_stage_input_func == ( "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload" ) @@ -162,8 +168,11 @@ class TestDeployTopology: assert connector["extra"]["connector_get_max_wait"] == 300 expected_processor = "tts2code2wav_async_chunk" if deploy.async_chunk else "tts2code2wav_full_payload" assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) + if filename != "minicpmo_4_5.yaml": + assert "hf_overrides" not in stages[1].yaml_engine_args if filename == "minicpmo_4_5.yaml": assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] + assert stages[1].yaml_engine_args["hf_overrides"] == {"tts_config": {"attention_type": "sliding_recompute"}} memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] assert memory_utilizations == [ 0.55, @@ -184,6 +193,42 @@ class TestDeployTopology: 0.35, ] + def test_duplex_talker_hf_override_resolves_nested_attention_policy(self) -> None: + deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5.yaml") + stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy) + overrides = stages[1].yaml_engine_args["hf_overrides"] + hf_config = PretrainedConfig() + hf_config.tts_config = PretrainedConfig(attention_type="full_attention") + + model_config = object.__new__(ModelConfig) + model_config._apply_dict_overrides(hf_config, overrides) + + assert hf_config.tts_config.attention_type == "sliding_recompute" + + @pytest.mark.parametrize( + "filename", + [ + "minicpmo_4_5.yaml", + "minicpmo_4_5_2gpu.yaml", + "minicpmo_4_5_3gpu.yaml", + ], + ) + def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None: + deploy = load_deploy_config(_DEPLOY_DIR / filename) + connector_extra = deploy.connectors["connector_of_shared_memory"]["extra"] + assert "code2wav_enable_npu_graph" not in connector_extra + assert "code2wav_max_npu_graphs" not in connector_extra + + deploy = _apply_platform_overrides(deploy, platform="npu") + stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy) + assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE" + assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE" + assert stages[2].yaml_engine_args["enforce_eager"] is True + assert stages[2].yaml_engine_args["additional_config"] == { + "code2wav_enable_npu_graph": True, + "code2wav_max_npu_graphs": 32, + } + def test_pipeline_exposes_full_and_async_payload_hooks(self) -> None: pipeline = OMNI_PIPELINES[_PIPELINE_KEY] talker = pipeline.get_stage(1) @@ -200,6 +245,14 @@ class TestDeployTopology: "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_token_only" ) + def test_no_async_chunk_selects_full_payload_processor(self) -> None: + deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5.yaml") + deploy.async_chunk = False + stages = merge_pipeline_deploy(OMNI_PIPELINES[_PIPELINE_KEY], deploy) + + assert stages[1].yaml_engine_args["async_chunk"] is False + assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(".tts2code2wav_full_payload") + def test_code2wav_model_is_lazily_registered() -> None: assert _OMNI_MODELS["MiniCPMO45Code2Wav"] == ( diff --git a/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py new file mode 100644 index 000000000..e867e0cab --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MiniCPM-o 4.5 streaming audio cache compatibility.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from transformers.models.whisper.modeling_whisper import WhisperConfig + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + WHISPER_ATTENTION_CLASSES, + MiniCPMWhisperEncoder, + MiniCPMWhisperEncoderLayer, + _get_audio_cache_length, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _LegacyWhisperAttention(nn.Module): + def __init__(self, **kwargs) -> None: + super().__init__() + self.seen_cache = None + + def forward(self, hidden_states, past_key_value=None, **kwargs): + self.seen_cache = past_key_value + return hidden_states, None, past_key_value + + +class _CurrentWhisperAttention(nn.Module): + def __init__(self, **kwargs) -> None: + super().__init__() + self.seen_cache = None + + def forward(self, hidden_states, past_key_values=None, **kwargs): + self.seen_cache = past_key_values + return hidden_states, None + + +@pytest.mark.parametrize( + ("implementation", "attention_cls"), + [ + ("test_legacy_cache", _LegacyWhisperAttention), + ("test_current_cache", _CurrentWhisperAttention), + ], +) +def test_whisper_attention_preserves_streaming_cache( + monkeypatch: pytest.MonkeyPatch, + implementation: str, + attention_cls: type[nn.Module], +) -> None: + monkeypatch.setitem(WHISPER_ATTENTION_CLASSES, implementation, attention_cls) + config = WhisperConfig( + d_model=4, + encoder_attention_heads=1, + encoder_ffn_dim=8, + encoder_layers=1, + ) + config._attn_implementation = implementation + layer = MiniCPMWhisperEncoderLayer(config, layer_idx=0) + cache = object() + + outputs = layer( + torch.randn(1, 2, config.d_model), + attention_mask=None, + layer_head_mask=None, + past_key_values=cache, + use_cache=True, + ) + + assert layer.self_attn.seen_cache is cache + assert outputs[-1] is cache + + +def test_audio_cache_length_uses_current_cache_api() -> None: + dynamic_cache = SimpleNamespace(get_seq_length=lambda: 7) + encoder_decoder_cache = SimpleNamespace(self_attention_cache=dynamic_cache) + + assert _get_audio_cache_length(encoder_decoder_cache) == 7 + + +def test_audio_cache_length_supports_legacy_cache() -> None: + key = torch.zeros(1, 1, 5, 2) + value = torch.zeros_like(key) + + assert _get_audio_cache_length(((key, value),)) == 5 + + +def _tiny_encoder_config(attn_implementation: str) -> WhisperConfig: + config = WhisperConfig( + num_mel_bins=16, + d_model=32, + encoder_layers=2, + encoder_attention_heads=2, + encoder_ffn_dim=64, + max_source_positions=100, + dropout=0.0, + attention_dropout=0.0, + ) + config._attn_implementation = attn_implementation + return config + + +@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"]) +def test_encoder_reuses_cache_across_streaming_chunks(attn_implementation: str) -> None: + """Encode two chunks against the installed Whisper attention and cache classes. + + Unlike the stub-based tests above, this drives the real transformers cache + API that ``get_audio_embedding_streaming`` relies on, so it fails whenever a + removed cache method or renamed attention kwarg breaks streaming reuse. + """ + torch.manual_seed(0) + config = _tiny_encoder_config(attn_implementation) + encoder = MiniCPMWhisperEncoder(config).eval() + + first_frames, second_frames = 20, 14 + first_len = (first_frames - 1) // 2 + 1 + second_len = (second_frames - 1) // 2 + 1 + first_mel = torch.randn(1, config.num_mel_bins, first_frames) + second_mel = torch.randn(1, config.num_mel_bins, second_frames) + + with torch.no_grad(): + first = encoder( + first_mel, + attention_mask=torch.zeros(1, 1, first_len, first_len), + past_key_values=None, + use_cache=True, + output_hidden_states=True, + ) + cache = first.past_key_values + assert cache is not None + assert _get_audio_cache_length(cache) == first_len + + second = encoder( + second_mel, + attention_mask=torch.zeros(1, 1, second_len, first_len + second_len), + past_key_values=cache, + use_cache=True, + output_hidden_states=True, + ) + assert _get_audio_cache_length(second.past_key_values) == first_len + second_len + + stateless = encoder( + second_mel, + attention_mask=torch.zeros(1, 1, second_len, second_len), + past_key_values=None, + use_cache=True, + output_hidden_states=True, + ) + + assert torch.isfinite(second.last_hidden_state).all() + assert not torch.allclose(second.last_hidden_state, stateless.last_hidden_state) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py index 8a7593623..95d845686 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Request-alignment tests for MiniCPM-o 4.5's native Talker.""" from __future__ import annotations +from dataclasses import dataclass +from typing import Any + import pytest import torch import torch.nn as nn @@ -11,35 +14,72 @@ import torch.nn as nn from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, ) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + ConditionalChatTTSConfig, +) from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _CODEC_PENALTY_WINDOW, + _DUPLEX_CODEC_TOKENS_PER_CHUNK, + _OFFLINE_CODEC_MAX_NEW_TOKENS, _REPETITION_PENALTY_CHUNK_SIZE, MiniCPMO45OmniTTSForConditionalGeneration, _apply_batched_repetition_penalty, - _max_audio_tokens, + _native_duplex_chunk_budget, _restore_weight_norm_weight, + blank_scheduler_prompt_for_penalties, ) from vllm_omni.utils.mm_outputs import to_payload_element pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@dataclass +class _PenaltyMetadata: + """Stand-in for the penalty-relevant slice of vLLM's SamplingMetadata.""" + + prompt_token_ids: torch.Tensor | None + + +@dataclass +class _SamplerOutput: + """Stand-in for vLLM's SamplerOutput.""" + + sampled_token_ids: torch.Tensor + + +@dataclass +class _CodecSamplingMetadata: + """Stand-in for the codec-penalty slice of vLLM's SamplingMetadata.""" + + repetition_penalties: torch.Tensor + prompt_token_ids: torch.Tensor | None = None + no_penalties: bool = False + + class _FakeNativeTalker(nn.Module): has_preprocess = True def __init__(self) -> None: super().__init__() - self.forward_kwargs = None + self.forward_kwargs: dict[str, Any] | None = None - def forward(self, **kwargs): + def forward(self, **kwargs: Any) -> torch.Tensor: self.forward_kwargs = kwargs return torch.ones(2, 4) +def test_tts_config_exposes_codec_vocab_to_vllm() -> None: + cfg = ConditionalChatTTSConfig(num_audio_tokens=6562) + assert cfg.vocab_size == 6562 + assert cfg.eos_token_id == 6561 + + def test_wrapper_always_delegates_talker_to_native_ar_path() -> None: model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) nn.Module.__init__(model) model.model_stage = "tts" - model.talker = _FakeNativeTalker() + talker = _FakeNativeTalker() + model.talker = talker output = model( input_ids=torch.tensor([1, 2]), @@ -48,19 +88,26 @@ def test_wrapper_always_delegates_talker_to_native_ar_path() -> None: ) assert output.shape == (2, 4) - assert model.talker.forward_kwargs["model_intermediate_buffer"][0]["request_id"] == "req" + assert talker.forward_kwargs is not None + assert talker.forward_kwargs["model_intermediate_buffer"][0]["request_id"] == "req" def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: talker = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) nn.Module.__init__(talker) talker._num_audio_tokens = 8 - talker._batch_stop_logits = None - talker._request_generators = {} + talker._codec_eos_id = 7 + talker._force_eos_rows = None + talker._mask_eos_rows = None + talker._pending_force_eos_rows = None + talker._penalty_histories = None talker._request_audio_states = {} + talker._request_condition_states = {} talker._deferred_cleanup_ids = set() - talker._codec_min_tokens = 50 - talker._codec_seed = 42 + talker._tts_config = ConditionalChatTTSConfig() + talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) + with torch.no_grad(): + talker.head_code[0].weight.copy_(torch.eye(8, 2)) return talker @@ -90,46 +137,6 @@ def _reference_repetition_penalty( return torch.where(logits < 0, logits * alpha, logits / alpha) -def _make_sampling_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: - talker = _make_talker() - talker._codec_temperature = 0.8 - talker._codec_top_k = 5 - talker._codec_top_p = 0.85 - talker._codec_repetition_penalty = 1.05 - talker._request_audio_states = { - "req-a": {"min_tokens": 2}, - "req-b": {"min_tokens": 0}, - } - talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) - with torch.no_grad(): - talker.head_code[0].weight.copy_( - torch.tensor( - [ - [0.1, -0.3], - [0.2, 0.4], - [-0.5, 0.2], - [0.7, -0.1], - [-0.2, 0.8], - [0.6, 0.3], - [-0.4, -0.7], - [0.3, 0.5], - ] - ) - ) - return talker - - -@pytest.mark.parametrize( - ("condition_tokens", "expected"), - [(3, 64), (100, 1000), (1000, 2048)], -) -def test_audio_token_limit_scales_with_condition_length( - condition_tokens: int, - expected: int, -) -> None: - assert _max_audio_tokens(condition_tokens) == expected - - def test_weight_norm_restore_matches_checkpoint_parametrization_in_bfloat16() -> None: generator = torch.Generator().manual_seed(42) weight_v = torch.randn(8, 16, generator=generator, dtype=torch.bfloat16) @@ -220,133 +227,454 @@ def test_batched_repetition_penalty_matches_rows_across_chunks(mocker) -> None: ] -def test_batched_codec_sampling_is_request_order_independent() -> None: - talker = _make_sampling_talker() - histories = [torch.tensor([1, 1, 2]), torch.tensor([3, 4, 4])] - hidden = torch.tensor([[0.25, -0.75], [-0.5, 0.5]]) +def test_scheduler_prompt_is_fully_blanked_for_penalties() -> None: + # The handoff path fills the prompt with 0s and the non-handoff path passes + # thinker token ids; neither is codec history, so every position is padded. + prompt = torch.tensor([[0, 0, 0], [12, 0, 340]]) - forward_order = talker._sample_audio_codes( - hidden, - histories, - ["req-a", "req-b"], - [0, 0], - ) - talker._request_generators = {} - reverse_order = talker._sample_audio_codes( - hidden.flip(0), - list(reversed(histories)), - ["req-b", "req-a"], - [0, 0], + blanked = blank_scheduler_prompt_for_penalties(prompt, vocab_size=8) + + assert blanked.tolist() == [[8, 8, 8], [8, 8, 8]] + assert prompt.tolist() == [[0, 0, 0], [12, 0, 340]] + + +def test_blanked_prompt_leaves_every_codec_logit_unpenalized() -> None: + from vllm.model_executor.layers.utils import apply_penalties + + vocab_size = 8 + logits = torch.arange(1.0, vocab_size + 1.0).unsqueeze(0) + prompt = blank_scheduler_prompt_for_penalties(torch.tensor([[0, 0, 3]]), vocab_size) + + penalized = apply_penalties( + logits.clone(), + prompt, + torch.full((1, 1), vocab_size, dtype=torch.long), + torch.zeros(1), + torch.zeros(1), + torch.full((1,), 1.02), ) - assert torch.equal(forward_order[0], reverse_order[1]) - assert torch.equal(forward_order[1], reverse_order[0]) + torch.testing.assert_close(penalized, logits) -def test_batched_codec_sampling_preserves_request_rng_after_compaction() -> None: - talker = _make_sampling_talker() - history_a = torch.tensor([1, 1, 2]) - history_b = torch.tensor([3, 4, 4]) - hidden_a = torch.tensor([[0.25, -0.75]]) - hidden_b = torch.tensor([[-0.5, 0.5]]) - - first_batch = talker._sample_audio_codes( - torch.cat([hidden_a, hidden_b]), - [history_a, history_b], - ["req-a", "req-b"], - [0, 0], - ) - compacted_b = talker._sample_audio_codes( - hidden_b, - [history_b], - ["req-b"], - [1], - ) - - talker._request_generators = {} - standalone_b_first = talker._sample_audio_codes( - hidden_b, - [history_b], - ["req-b"], - [0], - )[0] - standalone_b_second = talker._sample_audio_codes( - hidden_b, - [history_b], - ["req-b"], - [1], - )[0] - - assert torch.equal(first_batch[1], standalone_b_first) - assert torch.equal(compacted_b[0], standalone_b_second) - - -def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None: +def test_talker_sample_blanks_penalty_prompt_without_touching_runner_state(mocker) -> None: talker = _make_talker() - seen: list[tuple[str, list[float], list[int]]] = [] + metadata = _PenaltyMetadata(prompt_token_ids=torch.tensor([[0, 2, 0]])) + captured: dict[str, torch.Tensor] = {} - def sample(hidden, histories, request_ids, steps): - assert steps == [0, 0] - seen.extend( - (request_id, hidden[row].reshape(-1).tolist(), histories[row].tolist()) - for row, request_id in enumerate(request_ids) + class _Recorder: + def __call__(self, logits, sampling_metadata): + prompt = sampling_metadata.prompt_token_ids + assert prompt is not None + captured["prompt"] = prompt + return None + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=_Recorder(), + ) + talker.sample(torch.zeros(1, 8), metadata) + + assert captured["prompt"].tolist() == [[8, 8, 8]] + assert metadata.prompt_token_ids is not None + assert metadata.prompt_token_ids.tolist() == [[0, 2, 0]] + + +def test_compute_logits_returns_codec_vocab() -> None: + talker = _make_talker() + hidden = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + + logits = talker.compute_logits(hidden) + + assert logits.shape == (2, 8) + assert torch.allclose(logits[0], talker.head_code[0](hidden[0])) + + +def test_compute_logits_forces_eos_for_empty_speech() -> None: + talker = _make_talker() + talker._request_audio_states["req-empty"] = {"finished": True} + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-empty", "codes": {"audio": torch.empty(0)}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert output.multimodal_outputs["meta"]["finished"][0].item() is True + assert int(logits.argmax(dim=-1)) == 7 + + +def _capture_sampler_logits(mocker) -> dict[str, Any]: + captured: dict[str, Any] = {} + + def _record(logits, sampling_metadata): + captured["logits"] = logits.clone() + captured["metadata"] = sampling_metadata + return _SamplerOutput(sampled_token_ids=logits.argmax(dim=-1, keepdim=True)) + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=_record, + ) + return captured + + +def test_sample_scores_the_upstream_windowed_codec_penalty(mocker) -> None: + talker = _make_talker() + talker._request_audio_states["req-pen"] = {"finished": False, "step": 0} + # ``head_code`` is eye(8, 2) here, so only codes 0 and 1 carry a non-zero + # logit and are the only ones a penalty can visibly move. + hidden = torch.tensor([[1.0, -2.0]]) + raw = talker.head_code[0](hidden).float() + + output = talker.make_omni_output( + hidden, + model_intermediate_buffer=[{"request_id": "req-pen", "codes": {"audio": torch.tensor([[1]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + captured = _capture_sampler_logits(mocker) + talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.05]))) + + expected = _reference_repetition_penalty(raw, torch.tensor([1]), penalty=1.05, window_size=_CODEC_PENALTY_WINDOW) + torch.testing.assert_close(captured["logits"], expected) + # A negative logit is pushed further down, so the row really did change. + assert captured["logits"][0, 1].item() < raw[0, 1].item() + # Scored once: the Sampler's own whole-stream pass is neutralized. + assert captured["metadata"].repetition_penalties.tolist() == [1.0] + # Consumed once, so a later step cannot rescore a stale row. + assert talker._penalty_histories is None + + +def test_codec_penalty_forgets_codes_older_than_the_upstream_window(mocker) -> None: + # vLLM's sampler penalty spans the whole stream, so every code ever sampled + # stays taxed and a long tail drifts toward never-sampled codes. Upstream + # scores only the last ``past_window`` frames. + talker = _make_talker() + talker._request_audio_states["req-window"] = {"finished": False, "step": 0} + hidden = torch.tensor([[1.0, -2.0]]) + raw = talker.head_code[0](hidden).float() + captured = _capture_sampler_logits(mocker) + + def emit(code: int) -> torch.Tensor: + output = talker.make_omni_output( + hidden, + model_intermediate_buffer=[{"request_id": "req-window", "codes": {"audio": torch.tensor([[code]])}}], + request_token_spans=[(0, 1)], ) - return torch.tensor([2, 3]) + logits = talker.compute_logits(output.text_hidden_states) + talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.05]))) + return captured["logits"] - batched_sample = mocker.patch.object(talker, "_sample_audio_codes", side_effect=sample) + assert emit(1)[0, 1].item() < raw[0, 1].item() + # Push code 1 out of the window with a full window of code 0. + for _ in range(_CODEC_PENALTY_WINDOW): + scored = emit(0) + + assert talker._request_audio_states["req-window"]["recent_codes"] == [0] * _CODEC_PENALTY_WINDOW + torch.testing.assert_close(scored[0, 1], raw[0, 1]) + + +def test_sample_honours_a_request_that_disabled_penalties(mocker) -> None: + talker = _make_talker() + talker._request_audio_states["req-none"] = {"finished": False, "step": 0} + hidden = torch.tensor([[1.0, -2.0]]) + raw = talker.head_code[0](hidden).float() + + output = talker.make_omni_output( + hidden, + model_intermediate_buffer=[{"request_id": "req-none", "codes": {"audio": torch.tensor([[1]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + captured = _capture_sampler_logits(mocker) + talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.0]), no_penalties=True)) + + torch.testing.assert_close(captured["logits"], raw) + + +def test_sample_restores_forced_eos_that_min_tokens_masked_away(mocker) -> None: + # ``min_tokens`` masks every stage stop_token_id, including the codec EOS + # compute_logits just forced, so the sampler returns an arbitrary id for a + # request the model already finished. + talker = _make_talker() + talker._request_audio_states["req-empty"] = {"finished": True} + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-empty", "codes": {"audio": torch.empty(0)}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + logits[0, talker._codec_eos_id] = float("-inf") + + class _MinTokensSampler: + def __call__(self, logits, sampling_metadata): + del sampling_metadata + return _SamplerOutput(sampled_token_ids=logits.argmax(dim=-1, keepdim=True)) + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=_MinTokensSampler(), + ) + sampled = talker.sample(logits, _PenaltyMetadata(prompt_token_ids=None)) + + assert sampled.sampled_token_ids.tolist() == [[talker._codec_eos_id]] + + +def test_sample_leaves_unforced_rows_to_the_sampler(mocker) -> None: + talker = _make_talker() + talker._request_audio_states["req-a"] = {"finished": True} + talker._request_audio_states["req-b"] = {"finished": False, "step": 0} + + output = talker.make_omni_output( + torch.ones(2, 2), + model_intermediate_buffer=[ + {"request_id": "req-a", "codes": {"audio": torch.empty(0)}}, + {"request_id": "req-b", "codes": {"audio": torch.tensor([[3]])}}, + ], + request_token_spans=[(0, 1), (1, 2)], + ) + talker.compute_logits(output.text_hidden_states) + + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=lambda logits, sampling_metadata: _SamplerOutput(sampled_token_ids=torch.tensor([[0], [4]])), + ) + sampled = talker.sample(torch.zeros(2, 8), _PenaltyMetadata(prompt_token_ids=None)) + + assert sampled.sampled_token_ids.tolist() == [[talker._codec_eos_id], [4]] + # A single compute_logits decision must not leak into the next sample(). + assert talker._pending_force_eos_rows is None + + +def test_native_duplex_chunk_budget_masks_eos_until_generate_chunk_is_full() -> None: + assert _native_duplex_chunk_budget({}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, _DUPLEX_CODEC_TOKENS_PER_CHUNK) + assert _native_duplex_chunk_budget({"turn_start": True}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, 0) + assert _native_duplex_chunk_budget({"turn_end": True}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, 0) + + talker = _make_talker() + talker._request_audio_states["req-chunk"] = { + "finished": False, + "step": 0, + "max_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + "min_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-chunk", "codes": {"audio": torch.empty(0)}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert logits[0, 7].item() == float("-inf") + assert torch.isfinite(logits[0, 0]) + assert output.multimodal_outputs["meta"]["finished"][0].item() is False + + +def test_native_duplex_chunk_budget_forces_eos_after_twenty_five_frames() -> None: + talker = _make_talker() + talker._request_audio_states["req-full"] = { + "finished": False, + "step": 24, + "max_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + "min_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-full", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-full"]["step"] == 25 + assert talker._request_audio_states["req-full"]["finished"] is True + assert int(logits.argmax(dim=-1)) == 7 + assert output.multimodal_outputs["meta"]["finished"][0].item() is True + + +def test_offline_prefill_caps_codec_budget_to_the_remaining_talker_context() -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + talker._text_eos_id = 0 + talker._tts_bos_id = 0 + # Long enough that the context, not upstream's max_new_token, is the binding + # limit; the Talker must never plan past its own position embeddings. + prompt_len = int(talker._tts_config.max_position_embeddings) - _OFFLINE_CODEC_MAX_NEW_TOKENS + 100 + + _, _, updates = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_prompt_len=prompt_len, + request_id="req-offline-cap", + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + + remaining = int(talker._tts_config.max_position_embeddings) - prompt_len + assert remaining < _OFFLINE_CODEC_MAX_NEW_TOKENS + assert updates["audio_state"]["max_tokens"] == remaining + assert updates["audio_state"]["min_tokens"] is None + + +def test_talker_masks_eos_until_recorded_min_tokens() -> None: + talker = _make_talker() + talker._request_audio_states["req-early"] = { + "finished": False, + "step": 0, + "max_tokens": 8, + "min_tokens": 4, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-early", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-early"]["step"] == 1 + assert logits[0, 7].item() == float("-inf") + assert torch.isfinite(logits[0, 0]) + assert output.multimodal_outputs["meta"]["finished"][0].item() is False + + +def test_offline_talker_forces_eos_at_max_new_token() -> None: + talker = _make_talker() + max_tokens = 8 + talker._request_audio_states["req-cap"] = { + "finished": False, + "step": max_tokens - 2, + "max_tokens": max_tokens, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-cap", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-cap"]["step"] == max_tokens - 1 + assert talker._request_audio_states["req-cap"]["finished"] is True + assert int(logits.argmax(dim=-1)) == 7 + assert output.multimodal_outputs["meta"]["finished"][0].item() is True + + +def test_talker_without_a_recorded_budget_neither_forces_nor_masks_eos() -> None: + # A state rebuilt from a transported ``audio_state`` can lack budget keys; + # such rows must fall through to the raw codec head. + talker = _make_talker() + talker._request_audio_states["req-nobudget"] = {"finished": False, "step": 0} + + output = talker.make_omni_output( + torch.tensor([[1.0, 0.0]]), + model_intermediate_buffer=[{"request_id": "req-nobudget", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-nobudget"]["step"] == 1 + assert torch.allclose(logits[0], talker.head_code[0](output.text_hidden_states[0]).float()) + + +def test_offline_prefill_records_upstream_max_new_token_budget() -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + talker._text_eos_id = 0 + talker._tts_bos_id = 0 + + _, _, updates = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + request_id="req-offline", + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + + # A 2-token boundary prompt leaves far more context than upstream's + # max_new_token, so MiniCPMTTS.generate's own ceiling is what binds. + assert int(talker._tts_config.max_position_embeddings) - 2 > _OFFLINE_CODEC_MAX_NEW_TOKENS + assert updates["audio_state"]["max_tokens"] == _OFFLINE_CODEC_MAX_NEW_TOKENS + assert updates["audio_state"]["min_tokens"] is None + + +def test_native_duplex_prefill_records_generate_chunk_budget() -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + talker._text_eos_id = 0 + talker._tts_bos_id = 0 + + _, _, mid_turn = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + request_id="req-mid", + native_duplex=True, + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + _, _, boundary = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + request_id="req-end", + native_duplex=True, + meta={"turn_end": True}, + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + + assert mid_turn["audio_state"]["max_tokens"] == _DUPLEX_CODEC_TOKENS_PER_CHUNK + assert mid_turn["audio_state"]["min_tokens"] == _DUPLEX_CODEC_TOKENS_PER_CHUNK + assert boundary["audio_state"]["min_tokens"] == 0 + + +def test_make_omni_output_packs_the_previous_codec_id() -> None: + talker = _make_talker() infos = [ - {"request_id": "req-a", "audio_codes": {"accumulated": torch.tensor([1])}}, - {"request_id": "req-b", "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}}, + {"request_id": "req-a", "codes": {"audio": torch.tensor([[2]])}}, + {"request_id": "req-b", "codes": {"audio": torch.tensor([[3]])}}, ] output = talker.make_omni_output( - torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]), + torch.ones(2, 2), model_intermediate_buffer=infos, - request_token_spans=[(0, 2), (2, 3)], + request_token_spans=[(0, 1), (1, 2)], ) - assert seen == [ - ("req-a", [2.0, 0.0], [1]), - ("req-b", [3.0, 0.0], []), - ] - batched_sample.assert_called_once() - assert infos[0]["audio_codes"]["accumulated"].tolist() == [1, 2] - assert infos[1]["audio_codes"]["accumulated"].tolist() == [3] - assert set(output.multimodal_outputs) == {"codes", "meta"} - assert "model_outputs" not in output.multimodal_outputs - assert "sr" not in output.multimodal_outputs assert _routed(output, 0)["codes"]["audio"].tolist() == [[2]] assert _routed(output, 1)["codes"]["audio"].tolist() == [[3]] assert _routed(output, 0)["meta"]["finished"].item() is False - assert set(output.multimodal_outputs["meta"]) == {"finished"} - assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [0, 0] -def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: +def test_talker_projects_request_aligned_duplex_metadata() -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2, 2])) infos = [ { "request_id": "req-a", "native_duplex": True, "duplex": {"epoch": 3, "turn_id": 7}, "ids": {"tts": [41]}, - "meta": { - "native_duplex_segment_text": "first", - "turn_eos_token_id": 99, - }, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, + "meta": {"native_duplex_segment_text": "first", "turn_eos_token_id": 99}, + "codes": {"audio": torch.empty(0)}, }, { "request_id": "req-b", "native_duplex": True, "duplex": {"epoch": 4, "turn_id": 8}, "ids": {"tts": [42, 99]}, - "meta": { - "native_duplex_segment_text": "second", - "turn_eos_token_id": 99, - }, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, + "meta": {"native_duplex_segment_text": "second", "turn_eos_token_id": 99}, + "codes": {"audio": torch.empty(0)}, }, ] @@ -360,7 +688,6 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: assert [value.item() for value in meta["native_duplex"]] == [True, True] assert [value.item() for value in meta["duplex_epoch"]] == [3, 4] assert [value.item() for value in meta["duplex_turn_id"]] == [7, 8] - assert "native_duplex_segment_text" not in meta assert [bytes(value.tolist()).decode("utf-8") for value in meta["llm_output_text_utf8"]] == [ "first", "second", @@ -368,14 +695,9 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: assert [value.item() for value in meta["turn_end"]] == [False, True] -def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None: +def test_talker_rejects_native_duplex_without_fence_identity() -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2])) - info = { - "request_id": "req-missing-fence", - "native_duplex": True, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, - } + info = {"request_id": "req-missing-fence", "native_duplex": True, "codes": {"audio": torch.empty(0)}} with pytest.raises(RuntimeError, match="requires non-negative integer epoch and turn_id"): talker.make_omni_output( @@ -385,113 +707,35 @@ def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None: ) -def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None: +def test_decode_preprocess_embeds_the_sampled_codec_id() -> None: talker = _make_talker() - sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2])) - infos = [ - { - "request_id": "req-prefill", - "audio_state": {"step": 0}, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, - }, - { - "request_id": "req-decode", - "audio_state": {"step": 4}, - "audio_codes": {"accumulated": torch.tensor([1])}, - }, - ] + talker.emb_code = nn.ModuleList([nn.Embedding(8, 4)]) + with torch.no_grad(): + talker.emb_code[0].weight.copy_(torch.arange(32, dtype=torch.float32).reshape(8, 4)) - output = talker.make_omni_output( - torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]), - model_intermediate_buffer=infos, - request_token_spans=[(0, 2), (2, 3)], - request_sample_eligible=[False, True], + _, embeds, updates = talker.preprocess( + torch.tensor([3]), + None, + request_id="req-decode", + audio_state={"finished": False}, ) - sample.assert_called_once() - assert sample.call_args.args[2] == ["req-decode"] - assert infos[0]["audio_state"]["step"] == 0 - assert infos[0]["audio_codes"]["accumulated"].numel() == 0 - assert infos[1]["audio_state"]["step"] == 5 - assert _routed(output, 0)["codes"]["audio"].shape == (0, 1) - assert _routed(output, 1)["codes"]["audio"].tolist() == [[2]] + assert torch.equal(embeds, talker.emb_code[0](torch.tensor([3]))) + assert updates["codes"]["audio"].tolist() == [[3]] -def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None: +def test_decode_preprocess_drops_codec_eos() -> None: talker = _make_talker() - sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([7])) - info = { - "request_id": "req-stop", - "audio_state": {"step": 3}, - "audio_codes": {"accumulated": torch.tensor([4, 5])}, - } + talker.emb_code = nn.ModuleList([nn.Embedding(8, 4)]) - first = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - first_logits = talker.compute_logits(first.text_hidden_states) - second = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], + _, _, updates = talker.preprocess( + torch.tensor([7]), + None, + request_id="req-eos", + audio_state={"finished": False}, ) - sample.assert_called_once() - assert info["audio_codes"]["accumulated"].tolist() == [4, 5] - assert first.multimodal_outputs["codes"]["audio"][0].shape == (0, 1) - assert first.multimodal_outputs["meta"]["finished"][0].item() is True - assert second.multimodal_outputs["meta"]["finished"][0].item() is False - assert first_logits.argmax(dim=-1).tolist() == [1] - assert talker.compute_logits(second.text_hidden_states).argmax(dim=-1).tolist() == [1] - - -def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None: - talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3])) - info = { - "request_id": "req-limit", - "audio_state": {"step": 1, "max_tokens": 2}, - "audio_codes": {"accumulated": torch.tensor([4, 5])}, - } - - output = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - - # MiniCPMTTS.generate_chunk samples once at the max-token boundary to - # advance RNG state, but the sampled code is not fed into KV or returned. - assert info["audio_codes"]["accumulated"].tolist() == [4, 5] - assert output.multimodal_outputs["codes"]["audio"][0].shape == (0, 1) - assert output.multimodal_outputs["meta"]["finished"][0].item() is True - assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [1] - - -def test_request_local_state_survives_missing_runner_buffer_update(mocker) -> None: - talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3])) - first_info = { - "request_id": "req-local-state", - "audio_state": {"step": 1, "max_tokens": 3}, - "audio_codes": {"accumulated": torch.tensor([4])}, - } - - talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[first_info], - request_token_spans=[(0, 1)], - ) - second = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[{"request_id": "req-local-state"}], - request_token_spans=[(0, 1)], - ) - - assert second.multimodal_outputs["meta"]["finished"][0].item() is True - assert talker._request_audio_states["req-local-state"]["step"] == 3 + assert updates["codes"]["audio"].numel() == 0 def test_missing_conditioning_fails_clearly() -> None: @@ -524,18 +768,19 @@ def test_empty_speech_segment_finishes_without_sampling_codes() -> None: assert torch.equal(embeds, talker.emb_text(torch.tensor([5, 6]))) assert updates["audio_state"]["finished"] is True + assert updates["codes"]["audio"].numel() == 0 - # Stage 1's sampling min_tokens keeps scheduling decode steps until the stop - # token becomes eligible, and those steps have no previous code to embed. - _, decode_embeds, _ = talker.preprocess( + # min_tokens can keep scheduling leftover decode rows after an empty + # speech segment. Those rows have no previous codec id to embed. + _, decode_embeds, decode_updates = talker.preprocess( torch.zeros(1, dtype=torch.long), None, request_id="req-empty", audio_state=updates["audio_state"], - audio_codes=updates["audio_codes"], ) assert decode_embeds.shape == (1, 4) + assert decode_updates["codes"]["audio"].numel() == 0 def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> None: @@ -544,7 +789,7 @@ def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> Non condition = torch.arange(18, dtype=torch.float32).reshape(9, 2) mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) - _, embeds, _ = talker.preprocess( + _, embeds, updates = talker.preprocess( torch.zeros(9, dtype=torch.long), None, _omni_is_prefill=True, @@ -556,46 +801,9 @@ def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> Non ) assert torch.equal(embeds, condition) + assert updates["codes"]["audio"].numel() == 0 state = talker._request_audio_states["req-chunked-prefill"] - assert state["min_tokens"] == 50 - assert state["max_tokens"] == 64 - - -@pytest.mark.parametrize( - ("meta", "expected_min_tokens"), - [ - ({"turn_start": True}, 0), - ({}, 26), - ({"turn_end": True}, 0), - ], -) -def test_native_duplex_prefill_uses_official_chunk_limits( - mocker, - meta, - expected_min_tokens, -) -> None: - talker = _make_talker() - talker.emb_text = nn.Embedding(1, 2) - mocker.patch.object( - talker, - "_build_condition_embeddings", - return_value=torch.ones(3, 2), - ) - - talker.preprocess( - torch.zeros(3, dtype=torch.long), - None, - _omni_is_prefill=True, - request_id="req-duplex-chunk", - native_duplex=True, - meta=meta, - tts_token_ids=torch.tensor([1]), - tts_hidden_states=torch.ones(1, 2), - ) - - state = talker._request_audio_states["req-duplex-chunk"] - assert state["min_tokens"] == expected_min_tokens - assert state["max_tokens"] == 26 + assert state["finished"] is False def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None: @@ -626,13 +834,164 @@ def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None: assert condition.shape[0] == token_ids.shape[0] + 1 -def test_request_cleanup_evicts_ar_rng_and_decode_state() -> None: +def test_native_duplex_rollover_matches_official_sliding_recompute(mocker) -> None: talker = _make_talker() - talker._request_generators["req-done"] = torch.Generator() - talker._request_audio_states["req-done"] = {"step": 1} + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + with torch.no_grad(): + talker.emb_code[0].weight.copy_(torch.arange(16, dtype=torch.float32).reshape(8, 2)) + + previous_condition = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + current_condition = torch.tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]) + build_condition = mocker.patch.object( + talker, + "_build_condition_embeddings", + side_effect=[previous_condition, current_condition, current_condition], + ) + + talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_prompt_len=2, + request_id="req-window", + native_duplex=True, + meta={"turn_start": True, "streaming_condition_seq": 7}, + tts_token_ids=torch.tensor([1]), + tts_hidden_states=torch.ones(1, 2), + ) + + older_codes = [3, 4, 5] + talker._request_condition_states["req-window"]["base_recent_codes"] = tuple(older_codes) + previous_codes = [1, 2] + recompute_meta = { + "streaming_condition_seq": 8, + "streaming_prompt_recompute": True, + } + expected = torch.cat( + [ + previous_condition, + talker.emb_code[0](torch.tensor(previous_codes)), + current_condition, + ], + dim=0, + ) + _, full_embeds, _ = talker.preprocess( + torch.zeros(expected.shape[0], dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_prompt_len=expected.shape[0], + request_id="req-window", + native_duplex=True, + ids={"streaming_prompt_previous_codes": previous_codes}, + meta=recompute_meta, + tts_token_ids=torch.tensor([2]), + tts_hidden_states=torch.ones(1, 2), + ) + + assert torch.equal(full_embeds, expected) + assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes + + # Chunked-prefill/preemption retries reuse the immutable recipe for the + # same condition instead of advancing the one-chunk window twice. + _, retry_slice, _ = talker.preprocess( + torch.zeros(3, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_num_computed_tokens=2, + _omni_prompt_len=expected.shape[0], + request_id="req-window", + native_duplex=True, + ids={"streaming_prompt_previous_codes": previous_codes}, + meta=recompute_meta, + tts_token_ids=torch.tensor([2]), + tts_hidden_states=torch.ones(1, 2), + ) + assert torch.equal(retry_slice, expected[2:5]) + assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes + assert build_condition.call_count == 3 + + +def test_native_duplex_condition_advance_requires_recompute_marker(mocker) -> None: + talker = _make_talker() + condition = torch.ones(2, 2) + mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) + common = { + "_omni_is_prefill": True, + "_omni_prompt_len": 2, + "request_id": "req-history", + "native_duplex": True, + "tts_token_ids": torch.tensor([1]), + "tts_hidden_states": torch.ones(1, 2), + } + + talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 0}, **common) + with pytest.raises(ValueError, match="advanced without its streaming recompute marker"): + talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 1}, **common) + + +@pytest.mark.parametrize( + ("state", "previous_codes", "error"), + [ + (None, [1], "missing the previous Talker condition"), + ( + {"condition_seq": 0, "condition": torch.ones(2, 2), "base_recent_codes": ()}, + [7], + "invalid or terminal token", + ), + ], +) +def test_native_duplex_rollover_rejects_unverifiable_history(state, previous_codes, error) -> None: + talker = _make_talker() + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + if state is not None: + talker._request_condition_states["req-invalid-window"] = state + + with pytest.raises(ValueError, match=error): + talker._build_streaming_recompute_embeddings( + torch.ones(3, 2), + request_id="req-invalid-window", + info_dict={"ids": {"streaming_prompt_previous_codes": previous_codes}}, + meta={ + "streaming_condition_seq": 1, + "streaming_prompt_recompute": True, + }, + ) + + +def test_native_duplex_turn_start_resets_rollover_history() -> None: + talker = _make_talker() + talker._request_condition_states["req-turn"] = { + "condition_seq": 4, + "condition": torch.full((2, 2), 9.0), + "active_embeddings": torch.full((5, 2), 9.0), + } + current = torch.ones(3, 2) + + result = talker._build_streaming_recompute_embeddings( + current, + request_id="req-turn", + info_dict={}, + meta={"turn_start": True, "streaming_condition_seq": 5}, + ) + + assert torch.equal(result, current) + state = talker._request_condition_states["req-turn"] + assert state["condition_seq"] == 5 + assert torch.equal(state["condition"], current) + assert "active_embeddings" not in state + + +def test_request_cleanup_evicts_decode_state() -> None: + talker = _make_talker() + talker._request_audio_states["req-done"] = {"finished": False} + talker._request_condition_states["req-done"] = { + "condition_seq": 1, + "condition": torch.ones(2, 2), + } talker.on_requests_finished(["req-done"]) talker._flush_deferred_cleanup() - assert "req-done" not in talker._request_generators assert "req-done" not in talker._request_audio_states + assert "req-done" not in talker._request_condition_states diff --git a/tests/model_executor/models/moss_tts/test_compact_topk.py b/tests/model_executor/models/moss_tts/test_compact_topk.py new file mode 100644 index 000000000..e2b06c824 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_compact_topk.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU regression tests for MOSS-TTS compact top-k sampling.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts] + + +def test_compact_topk_only_samples_retained_candidates() -> None: + logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(7)) + + sampled = modeling_moss_tts_local._sample_token( + logits, + temperature=1.7, + top_k=25, + top_p=0.8, + do_sample=True, + generator=torch.Generator().manual_seed(42), + ) + + retained = logits.topk(25, dim=-1).indices + assert torch.all((retained == sampled.unsqueeze(-1)).any(dim=-1)) + + +def test_compact_topk_tiny_nucleus_selects_argmax() -> None: + logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(17)) + + sampled = modeling_moss_tts_local._sample_token( + logits, + temperature=1.7, + top_k=25, + top_p=1e-9, + do_sample=True, + generator=torch.Generator().manual_seed(42), + ) + + torch.testing.assert_close(sampled, logits.argmax(dim=-1)) diff --git a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py new file mode 100644 index 000000000..90c386ee0 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU parity tests for MOSS-TTS Local Depth RoPE caching.""" + +import pytest +import torch +from transformers import GPT2Config + +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( + MossTTSLocalDepthTransformer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts] + + +def _test_config() -> GPT2Config: + return GPT2Config( + n_embd=32, + n_head=4, + n_inner=64, + layer_norm_epsilon=1e-6, + rope_base=10000.0, + ) + + +def test_rope_cache_matches_reference_and_reuses_storage() -> None: + model = MossTTSLocalDepthTransformer(_test_config()).eval() + attn = model.h[0].attn + device = torch.device("cpu") + + attn.prepare_rope_cache(12, device, torch.float32) + cache_ptr = attn._rope_cos_cache.data_ptr() + cos, sin = attn._rope_cos_sin(4, device, torch.float32) + + positions = torch.arange(4, dtype=torch.float32) + freqs = torch.einsum("s,d->sd", positions, attn.inv_freq.float()) + expected_cos = freqs.cos().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim) + expected_sin = freqs.sin().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim) + torch.testing.assert_close(cos, expected_cos) + torch.testing.assert_close(sin, expected_sin) + + # Smaller slices retain the fixed allocation used by CUDA Graph. + attn.prepare_rope_cache(6, device, torch.float32) + assert attn._rope_cos_cache.data_ptr() == cache_ptr diff --git a/tests/model_executor/models/moss_tts/test_moss_fused_load.py b/tests/model_executor/models/moss_tts/test_moss_fused_load.py new file mode 100644 index 000000000..60039c693 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_moss_fused_load.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for MOSS-TTS-Realtime fused code-predictor loading. + +The MOSS-TTS-Realtime depth transformer shares ``CodePredictorBaseModel``, whose +q/k/v and gate/up projections are fused. Before the fix, the talker loaded those +shards with a per-tensor ``default_weight_loader``, which silently dropped every +q/k/v/gate/up shard and left the fused params randomly initialized -- no +exception, just a corrupted model. + +These tests assert the property that failure mode violates: after +``load_weights``, every fused parameter is both reported loaded and numerically +equal to the concatenation of the shards that fed it. +""" + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( + MossTTSLocalTransformerConfig, +) +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( + MossTTSRealtimeLocalTransformer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _tiny_cfg(): + return MossTTSLocalTransformerConfig( + head_dim=8, + hidden_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + intermediate_size=64, + max_position_embeddings=8, + ) + + +def _hf_shards(cfg, *, drop: set[str] | None = None): + """Synthetic HF-style checkpoint tensors for the depth transformer.""" + drop = drop or set() + h, inter = cfg.hidden_size, cfg.intermediate_size + q_out = cfg.num_attention_heads * cfg.head_dim + kv_out = cfg.num_key_value_heads * cfg.head_dim + out = {} + for layer in range(cfg.num_hidden_layers): + p = f"model.layers.{layer}" + for nm, rows in (("q_proj", q_out), ("k_proj", kv_out), ("v_proj", kv_out)): + key = f"{p}.self_attn.{nm}.weight" + if nm not in drop: + out[key] = torch.randn(rows, h) + out[f"{p}.self_attn.o_proj.weight"] = torch.randn(h, q_out) + for nm in ("gate_proj", "up_proj"): + key = f"{p}.mlp.{nm}.weight" + if nm not in drop: + out[key] = torch.randn(inter, h) + out[f"{p}.mlp.down_proj.weight"] = torch.randn(h, inter) + out[f"{p}.input_layernorm.weight"] = torch.randn(h) + out[f"{p}.post_attention_layernorm.weight"] = torch.randn(h) + for nm in ("q_norm", "k_norm"): + out[f"{p}.self_attn.{nm}.weight"] = torch.randn(cfg.head_dim) + out["model.norm.weight"] = torch.randn(h) + return out + + +def test_every_fused_param_is_loaded_and_correctly_packed(): + cfg = _tiny_cfg() + lt = MossTTSRealtimeLocalTransformer(cfg) + shards = _hf_shards(cfg) + + loaded = lt.load_weights(iter([(k, v) for k, v in shards.items()])) + + params = dict(lt.named_parameters()) + fused = [n for n in params if ".self_attn.qkv_proj." in n or ".mlp.gate_up_proj." in n] + assert fused, "model declares no fused params -- test would be vacuous" + + # 1. Reported as loaded (this is what the old path silently failed). + unreported = [n for n in fused if n not in loaded] + assert not unreported, f"fused params not reported loaded: {unreported}" + + # 2. Numerically equal to the concatenation of their shards, in order. + for layer in range(cfg.num_hidden_layers): + p = f"model.layers.{layer}" + want_qkv = torch.cat([shards[f"{p}.self_attn.{n}.weight"] for n in ("q_proj", "k_proj", "v_proj")], dim=0) + got_qkv = params[f"{p}.self_attn.qkv_proj.weight"].detach() + assert torch.equal(got_qkv, want_qkv), f"layer {layer}: qkv packing mismatch" + + want_gu = torch.cat([shards[f"{p}.mlp.{n}.weight"] for n in ("gate_proj", "up_proj")], dim=0) + got_gu = params[f"{p}.mlp.gate_up_proj.weight"].detach() + assert torch.equal(got_gu, want_gu), f"layer {layer}: gate_up packing mismatch" + + +def test_incomplete_shard_group_raises_instead_of_corrupting(): + """A checkpoint missing v_proj must fail loudly, not load a random qkv.""" + cfg = _tiny_cfg() + lt = MossTTSRealtimeLocalTransformer(cfg) + shards = _hf_shards(cfg, drop={"v_proj"}) + + with pytest.raises(RuntimeError, match="incomplete fused shards"): + lt.load_weights(iter([(k, v) for k, v in shards.items()])) + + +def test_absent_fused_group_raises(): + """No q/k/v at all must also raise, via the missing-fused-param guard.""" + cfg = _tiny_cfg() + lt = MossTTSRealtimeLocalTransformer(cfg) + shards = _hf_shards(cfg, drop={"q_proj", "k_proj", "v_proj"}) + + with pytest.raises(RuntimeError, match="missing fused parameters"): + lt.load_weights(iter([(k, v) for k, v in shards.items()])) + + +def test_returned_names_keep_the_model_prefix(): + """Returned names must be ``model.*`` for talker re-prefixing.""" + cfg = _tiny_cfg() + lt = MossTTSRealtimeLocalTransformer(cfg) + loaded = lt.load_weights(iter([(k, v) for k, v in _hf_shards(cfg).items()])) + assert loaded, "nothing reported loaded" + assert all(n.startswith("model.") for n in loaded), sorted(n for n in loaded if not n.startswith("model.")) diff --git a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py index 35c309979..b5866781c 100644 --- a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py +++ b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py @@ -23,6 +23,11 @@ class _RecordingLogitsProcessor: return torch.zeros((1, 4)) +class _EmbeddingBackbone(nn.Module): + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + return input_ids.to(dtype=torch.float32).unsqueeze(-1).repeat(1, 4) + + def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None: from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( MossTTSDelayTalkerForGeneration, @@ -41,3 +46,74 @@ def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> N assert logits.shape == (1, 4) assert model.logits_processor.args == (model.text_lm_head, hidden_states) assert model.logits_processor.kwargs == {} + + +def test_moss_tts_local_keeps_fixed_shape_rows_and_routes_stop_logits() -> None: + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSLocalTalkerForGeneration, + ) + + model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration) + nn.Module.__init__(model) + model.audio_pad_token_id = 1024 + model.audio_assistant_slot_token_id = 2 + model.im_end_token_id = 3 + model.text_vocab_size = 8 + model.n_vq = 4 + model._batch_state = None + model._batch_state_spans = None + model._batch_should_continue = None + + valid = torch.tensor([[1, 2, 3, 4]], dtype=torch.long) + stopped = torch.full((1, 4), model.audio_pad_token_id, dtype=torch.long) + infos = [ + {"audio_state": {}, "audio_codes": {"current": valid}}, + {"audio_state": {}, "audio_codes": {"current": stopped}}, + ] + spans = [(0, 3), (3, 4)] + + output = model.make_omni_output( + torch.randn(4, 8), + model_intermediate_buffer=infos, + request_token_spans=spans, + ) + + audio = output.multimodal_outputs["codes"]["audio"] + assert len(audio) == 2 + torch.testing.assert_close(audio[0], valid) + # Keeping this row avoids CUDA boolean-index shape discovery; the + # talker2codec raw async-chunk processor drops all-pad rows on CPU. + torch.testing.assert_close(audio[1], stopped) + assert model._batch_should_continue.tolist() == [True, False] + + logits = model.compute_logits(torch.randn(4, 8)) + + assert logits is not None + assert logits.argmax(dim=-1).tolist() == [2, 2, 2, 3] + + +def test_moss_tts_local_single_token_prefill_does_not_advance_audio() -> None: + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSLocalTalkerForGeneration, + ) + + model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration) + nn.Module.__init__(model) + model.model = _EmbeddingBackbone() + model.n_vq = 4 + + input_ids = torch.tensor([7], dtype=torch.long) + _, embeds, update = model.preprocess( + input_ids, + None, + _omni_is_prefill=True, + audio_state={"is_stopping": True}, + ref_offset=5, + ) + + torch.testing.assert_close(embeds, torch.full((1, 4), 7.0)) + assert update == { + "audio_state": {"is_stopping": False}, + "ref_offset": 6, + } + assert "mtp_inputs" not in update diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_quantization.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_quantization.py new file mode 100644 index 000000000..382e2f4d1 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_quantization.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression tests for nested Qwen3-Omni quantization config mapping.""" + +from dataclasses import dataclass + +import pytest + +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + _OUTER_MAPPING_APPLIED, + apply_nested_quant_config_mapping, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@dataclass +class _QuantConfig: + packed_modules_mapping: dict[str, list[str]] + + +@dataclass +class _NestedModel: + quant_config: _QuantConfig | None + + +@dataclass +class _NestedModelWithPackedModules(_NestedModel): + packed_modules_mapping: dict[str, list[str]] + + +def _mapped_quant_config(**packed_modules_mapping): + quant_config = _QuantConfig(packed_modules_mapping=packed_modules_mapping) + setattr(quant_config, _OUTER_MAPPING_APPLIED, True) + return quant_config + + +def test_nested_mapping_ignores_missing_quant_config(): + apply_nested_quant_config_mapping(_NestedModel(quant_config=None)) + + +def test_nested_mapping_allows_model_without_packed_modules_mapping(): + quant_config = _mapped_quant_config(existing=["module"]) + + apply_nested_quant_config_mapping(_NestedModel(quant_config=quant_config)) + + assert quant_config.packed_modules_mapping == {"existing": ["module"]} + + +def test_nested_mapping_merges_packed_modules_mapping(): + quant_config = _mapped_quant_config(existing=["module"]) + model = _NestedModelWithPackedModules( + quant_config=quant_config, + packed_modules_mapping={"qkv_proj": ["q_proj", "k_proj", "v_proj"]}, + ) + + apply_nested_quant_config_mapping(model) + + assert quant_config.packed_modules_mapping == { + "existing": ["module"], + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + } diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 20f8eb2f0..619f2169a 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -697,3 +697,481 @@ class TestGumbelMaxSampling: empirical = torch.bincount(samples, minlength=vocab).float() / n_draws # L1 distance under 0.03 with 20k draws is comfortable for vocab=16. assert (empirical - target).abs().sum().item() < 0.05 + + +class TestCodePredictorFusedProjections: + """Reference tests for the fused qkv_proj / gate_up_proj paths. + + These CPU tests confirm that the packed projections implement the same + linear algebra as the unfused reference under this setup. GPU executions + may still diverge at the bit level when kernel choices change. + """ + + def test_fused_qkv_matches_separate(self, loaded_target_classes) -> None: + torch.manual_seed(0) + hidden = 32 + num_heads, num_kv_heads, head_dim = 4, 2, 8 + q_size = num_heads * head_dim + kv_size = num_kv_heads * head_dim + x = torch.randn(2, 5, hidden) + + q = torch.nn.Linear(hidden, q_size, bias=False) + k = torch.nn.Linear(hidden, kv_size, bias=False) + v = torch.nn.Linear(hidden, kv_size, bias=False) + fused = torch.nn.Linear(hidden, q_size + 2 * kv_size, bias=False) + with torch.no_grad(): + fused.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0)) + + ref_q, ref_k, ref_v = q(x), k(x), v(x) + out_q, out_k, out_v = fused(x).split([q_size, kv_size, kv_size], dim=-1) + assert torch.equal(ref_q, out_q) + assert torch.equal(ref_k, out_k) + assert torch.equal(ref_v, out_v) + + def test_attention_split_redensifies_kv_slices(self, loaded_target_classes) -> None: + cp_config, _ = _make_tiny_config(loaded_target_classes) + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + attn = common_mod.CodePredictorAttention(cp_config) + + qkv = torch.randn(2, 5, attn._q_size + 2 * attn._kv_size) + q_ref, k_ref, v_ref = qkv.split([attn._q_size, attn._kv_size, attn._kv_size], dim=-1) + q_raw, k_raw, v_raw = attn._split_qkv(qkv) + + assert torch.equal(q_raw, q_ref) + assert torch.equal(k_raw, k_ref) + assert torch.equal(v_raw, v_ref) + assert k_raw.is_contiguous() + assert v_raw.is_contiguous() + + def test_fused_gate_up_matches_separate(self, loaded_target_classes) -> None: + torch.manual_seed(1) + hidden, intermediate = 32, 64 + x = torch.randn(2, 5, hidden) + + g = torch.nn.Linear(hidden, intermediate, bias=False) + u = torch.nn.Linear(hidden, intermediate, bias=False) + fused = torch.nn.Linear(hidden, 2 * intermediate, bias=False) + with torch.no_grad(): + fused.weight.copy_(torch.cat([g.weight, u.weight], dim=0)) + + ref_g, ref_u = g(x), u(x) + out_g, out_u = fused(x).split([intermediate, intermediate], dim=-1) + assert torch.equal(ref_g, out_g) + assert torch.equal(ref_u, out_u) + + def test_load_weights_repacks_hf_shards(self, mocker: MockerFixture, loaded_target_classes) -> None: + """``load_weights`` must transparently re-pack HF q/k/v and gate/up + shards into the fused ``qkv_proj`` and ``gate_up_proj`` weights.""" + _, _, _, code_predictor_model, _ = loaded_target_classes + cp_config, _ = _make_tiny_config(loaded_target_classes) + model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) + + # The test harness mocks ``default_weight_loader`` as a no-op; install a + # real copier so this test verifies both the returned names and the + # packed parameter values/order. + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object( + common_mod, + "default_weight_loader", + lambda param, weight: param.data.copy_(weight), + ) + + # Build a synthetic "HF checkpoint" that exposes separate q/k/v + # and gate/up shards, mirroring what HuggingFace ships. + param_dict = dict(model.named_parameters(remove_duplicate=False)) + weights: list[tuple[str, torch.Tensor]] = [] + expected_fused: dict[str, torch.Tensor] = {} + torch.manual_seed(42) + for name, param in param_dict.items(): + if name.endswith(".self_attn.qkv_proj.weight"): + prefix = name[: -len(".qkv_proj.weight")] + num_heads = cp_config.num_attention_heads + num_kv_heads = cp_config.num_key_value_heads + head_dim = cp_config.head_dim + q_size = num_heads * head_dim + kv_size = num_kv_heads * head_dim + hidden = cp_config.hidden_size + q_w = torch.randn(q_size, hidden) + k_w = torch.randn(kv_size, hidden) + v_w = torch.randn(kv_size, hidden) + weights.append((f"{prefix}.q_proj.weight", q_w)) + weights.append((f"{prefix}.k_proj.weight", k_w)) + weights.append((f"{prefix}.v_proj.weight", v_w)) + expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0) + elif name.endswith(".mlp.gate_up_proj.weight"): + prefix = name[: -len(".gate_up_proj.weight")] + inter = cp_config.intermediate_size + hidden = cp_config.hidden_size + gate_w = torch.randn(inter, hidden) + up_w = torch.randn(inter, hidden) + weights.append((f"{prefix}.gate_proj.weight", gate_w)) + weights.append((f"{prefix}.up_proj.weight", up_w)) + expected_fused[name] = torch.cat([gate_w, up_w], dim=0) + else: + # Pass-through parameters (norms, embeddings, o_proj, down_proj). + weights.append((name, torch.randn_like(param))) + + loaded = model.load_weights(weights) + after = dict(model.named_parameters(remove_duplicate=False)) + assert expected_fused, "test config must contain at least one fused projection layer" + for fname, expected in expected_fused.items(): + assert fname in loaded, f"{fname} was not assembled by load_weights" + assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order" + + def test_load_weights_repacks_bias_shards(self, mocker: MockerFixture, loaded_target_classes) -> None: + """When ``attention_bias=True``, the q/k/v ``.bias`` shards must be + concatenated into ``qkv_proj.bias`` instead of being dropped (which + would leave the fused bias randomly initialized).""" + _, _, _, code_predictor_model, _ = loaded_target_classes + cp_config, _ = _make_tiny_config(loaded_target_classes) + cp_config.attention_bias = True + model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) + + # The test harness mocks ``default_weight_loader`` as a no-op; install a + # real copier so we can assert the packed values land on the parameter. + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object( + common_mod, + "default_weight_loader", + lambda param, weight: param.data.copy_(weight), + ) + + param_dict = dict(model.named_parameters(remove_duplicate=False)) + weights: list[tuple[str, torch.Tensor]] = [] + expected_bias: dict[str, torch.Tensor] = {} + torch.manual_seed(11) + for name, param in param_dict.items(): + if name.endswith(".self_attn.qkv_proj.weight"): + prefix = name[: -len(".qkv_proj.weight")] + q_size = cp_config.num_attention_heads * cp_config.head_dim + kv_size = cp_config.num_key_value_heads * cp_config.head_dim + hidden = cp_config.hidden_size + q_w, k_w, v_w = (torch.randn(s, hidden) for s in (q_size, kv_size, kv_size)) + q_b, k_b, v_b = (torch.randn(s) for s in (q_size, kv_size, kv_size)) + weights.append((f"{prefix}.q_proj.weight", q_w)) + weights.append((f"{prefix}.k_proj.weight", k_w)) + weights.append((f"{prefix}.v_proj.weight", v_w)) + weights.append((f"{prefix}.q_proj.bias", q_b)) + weights.append((f"{prefix}.k_proj.bias", k_b)) + weights.append((f"{prefix}.v_proj.bias", v_b)) + expected_bias[f"{prefix}.qkv_proj.bias"] = torch.cat([q_b, k_b, v_b], dim=0) + elif name.endswith(".mlp.gate_up_proj.weight"): + prefix = name[: -len(".gate_up_proj.weight")] + inter, hidden = cp_config.intermediate_size, cp_config.hidden_size + weights.append((f"{prefix}.gate_proj.weight", torch.randn(inter, hidden))) + weights.append((f"{prefix}.up_proj.weight", torch.randn(inter, hidden))) + elif name.endswith(".qkv_proj.bias"): + # Fused target is assembled from the q/k/v bias shards above; a + # real HF checkpoint never ships a qkv_proj.bias tensor directly. + continue + else: + weights.append((name, torch.randn_like(param))) + + loaded = model.load_weights(weights) + after = dict(model.named_parameters()) + assert expected_bias, "test config must contain at least one qkv layer" + for bname, exp in expected_bias.items(): + assert bname in loaded, f"{bname} bias was not assembled by load_weights" + assert torch.equal(after[bname], exp), f"{bname} does not match the packed q/k/v bias" + + def test_wrapper_load_weights_handles_model_prefixed_shards( + self, + mocker: MockerFixture, + loaded_target_classes, + ) -> None: + _, _, code_predictor_wrapper, _, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker) + predictor = code_predictor_wrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object( + common_mod, + "default_weight_loader", + lambda param, weight: param.data.copy_(weight), + ) + + param_dict = dict(predictor.named_parameters(remove_duplicate=False)) + weights: list[tuple[str, torch.Tensor]] = [] + expected_fused: dict[str, torch.Tensor] = {} + torch.manual_seed(7) + for name, param in param_dict.items(): + if name.endswith(".self_attn.qkv_proj.weight"): + prefix = name[: -len(".qkv_proj.weight")] + q_size = cp_config.num_attention_heads * cp_config.head_dim + kv_size = cp_config.num_key_value_heads * cp_config.head_dim + hidden = cp_config.hidden_size + q_w = torch.randn(q_size, hidden) + k_w = torch.randn(kv_size, hidden) + v_w = torch.randn(kv_size, hidden) + weights.append((f"{prefix}.q_proj.weight", q_w)) + weights.append((f"{prefix}.k_proj.weight", k_w)) + weights.append((f"{prefix}.v_proj.weight", v_w)) + expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0) + elif name.endswith(".mlp.gate_up_proj.weight"): + prefix = name[: -len(".gate_up_proj.weight")] + inter = cp_config.intermediate_size + hidden = cp_config.hidden_size + gate_w = torch.randn(inter, hidden) + up_w = torch.randn(inter, hidden) + weights.append((f"{prefix}.gate_proj.weight", gate_w)) + weights.append((f"{prefix}.up_proj.weight", up_w)) + expected_fused[name] = torch.cat([gate_w, up_w], dim=0) + else: + weights.append((name, torch.randn_like(param))) + + loaded = predictor.load_weights(weights) + after = dict(predictor.named_parameters(remove_duplicate=False)) + assert expected_fused, "test config must contain at least one fused projection layer" + for fname, expected in expected_fused.items(): + assert fname in loaded, f"{fname} was not loaded through the outer wrapper" + assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order" + + def test_wrapper_load_weights_handles_legacy_unprefixed_model_shards( + self, + mocker: MockerFixture, + loaded_target_classes, + ) -> None: + """Qwen3-Omni checkpoints may expose code_predictor body tensors as + ``layers.*`` / ``codec_embedding.*`` instead of ``model.*``. The + wrapper must still route those shards through the fused inner loader.""" + _, _, code_predictor_wrapper, _, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker) + predictor = code_predictor_wrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object( + common_mod, + "default_weight_loader", + lambda param, weight: param.data.copy_(weight), + ) + + param_dict = dict(predictor.named_parameters(remove_duplicate=False)) + weights: list[tuple[str, torch.Tensor]] = [] + expected_fused: dict[str, torch.Tensor] = {} + torch.manual_seed(17) + for name, param in param_dict.items(): + if name.endswith(".self_attn.qkv_proj.weight"): + prefix = name[: -len(".qkv_proj.weight")] + legacy_prefix = prefix.removeprefix("model.") + q_size = cp_config.num_attention_heads * cp_config.head_dim + kv_size = cp_config.num_key_value_heads * cp_config.head_dim + hidden = cp_config.hidden_size + q_w = torch.randn(q_size, hidden) + k_w = torch.randn(kv_size, hidden) + v_w = torch.randn(kv_size, hidden) + weights.append((f"{legacy_prefix}.q_proj.weight", q_w)) + weights.append((f"{legacy_prefix}.k_proj.weight", k_w)) + weights.append((f"{legacy_prefix}.v_proj.weight", v_w)) + expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0) + elif name.endswith(".mlp.gate_up_proj.weight"): + prefix = name[: -len(".gate_up_proj.weight")] + legacy_prefix = prefix.removeprefix("model.") + inter = cp_config.intermediate_size + hidden = cp_config.hidden_size + gate_w = torch.randn(inter, hidden) + up_w = torch.randn(inter, hidden) + weights.append((f"{legacy_prefix}.gate_proj.weight", gate_w)) + weights.append((f"{legacy_prefix}.up_proj.weight", up_w)) + expected_fused[name] = torch.cat([gate_w, up_w], dim=0) + elif name.startswith("model."): + weights.append((name.removeprefix("model."), torch.randn_like(param))) + else: + weights.append((name, torch.randn_like(param))) + + loaded = predictor.load_weights(weights) + after = dict(predictor.named_parameters(remove_duplicate=False)) + assert expected_fused, "test config must contain at least one fused projection layer" + for fname, expected in expected_fused.items(): + assert fname in loaded, f"{fname} was not loaded from legacy unprefixed shards" + assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order" + + def test_wrapper_load_weights_handles_nested_code_predictor_prefixed_shards( + self, + mocker: MockerFixture, + loaded_target_classes, + ) -> None: + """Realtime Qwen3-Omni loaders may still prepend ``talker.code_predictor.`` + before the inner ``model.*`` body; the wrapper must normalize that layer + too before delegating to the fused inner loader.""" + _, _, code_predictor_wrapper, _, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker) + predictor = code_predictor_wrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + prefix="talker.code_predictor", + ) + + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object( + common_mod, + "default_weight_loader", + lambda param, weight: param.data.copy_(weight), + ) + + param_dict = dict(predictor.named_parameters(remove_duplicate=False)) + weights: list[tuple[str, torch.Tensor]] = [] + expected_fused: dict[str, torch.Tensor] = {} + torch.manual_seed(23) + wrapper_prefix = "talker.code_predictor." + for name, param in param_dict.items(): + wrapped_name = f"{wrapper_prefix}{name}" + if name.endswith(".self_attn.qkv_proj.weight"): + prefix = wrapped_name[: -len(".qkv_proj.weight")] + q_size = cp_config.num_attention_heads * cp_config.head_dim + kv_size = cp_config.num_key_value_heads * cp_config.head_dim + hidden = cp_config.hidden_size + q_w = torch.randn(q_size, hidden) + k_w = torch.randn(kv_size, hidden) + v_w = torch.randn(kv_size, hidden) + weights.append((f"{prefix}.q_proj.weight", q_w)) + weights.append((f"{prefix}.k_proj.weight", k_w)) + weights.append((f"{prefix}.v_proj.weight", v_w)) + expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0) + elif name.endswith(".mlp.gate_up_proj.weight"): + prefix = wrapped_name[: -len(".gate_up_proj.weight")] + inter = cp_config.intermediate_size + hidden = cp_config.hidden_size + gate_w = torch.randn(inter, hidden) + up_w = torch.randn(inter, hidden) + weights.append((f"{prefix}.gate_proj.weight", gate_w)) + weights.append((f"{prefix}.up_proj.weight", up_w)) + expected_fused[name] = torch.cat([gate_w, up_w], dim=0) + else: + weights.append((wrapped_name, torch.randn_like(param))) + + loaded = predictor.load_weights(weights) + after = dict(predictor.named_parameters(remove_duplicate=False)) + assert expected_fused, "test config must contain at least one fused projection layer" + for fname, expected in expected_fused.items(): + assert fname in loaded, f"{fname} was not loaded from nested code_predictor shards" + assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order" + + def test_wrapper_load_weights_accepts_incremental_non_layer_shard( + self, + mocker: MockerFixture, + loaded_target_classes, + ) -> None: + """Qwen3-Omni splits code-predictor weights across checkpoint shards. + + Its first shard contains the transformer layers, while the next shard + re-enters the wrapper with only codec embeddings and LM heads. Loading + that second shard must not require every fused layer to appear again. + """ + _, _, code_predictor_wrapper, _, _ = loaded_target_classes + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker) + predictor = code_predictor_wrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object( + common_mod, + "default_weight_loader", + lambda param, weight: param.data.copy_(weight), + ) + + first_shard: list[tuple[str, torch.Tensor]] = [] + second_shard: list[tuple[str, torch.Tensor]] = [] + for name, param in predictor.named_parameters(remove_duplicate=False): + if name.startswith("model.codec_embedding.") or name.startswith("lm_head."): + second_shard.append((name, torch.randn_like(param))) + elif name.endswith(".self_attn.qkv_proj.weight"): + prefix = name[: -len(".qkv_proj.weight")] + q_size = cp_config.num_attention_heads * cp_config.head_dim + kv_size = cp_config.num_key_value_heads * cp_config.head_dim + hidden = cp_config.hidden_size + first_shard.extend( + [ + (f"{prefix}.q_proj.weight", torch.randn(q_size, hidden)), + (f"{prefix}.k_proj.weight", torch.randn(kv_size, hidden)), + (f"{prefix}.v_proj.weight", torch.randn(kv_size, hidden)), + ] + ) + elif name.endswith(".mlp.gate_up_proj.weight"): + prefix = name[: -len(".gate_up_proj.weight")] + inter = cp_config.intermediate_size + hidden = cp_config.hidden_size + first_shard.extend( + [ + (f"{prefix}.gate_proj.weight", torch.randn(inter, hidden)), + (f"{prefix}.up_proj.weight", torch.randn(inter, hidden)), + ] + ) + else: + first_shard.append((name, torch.randn_like(param))) + + assert first_shard and second_shard + common_mod.CodePredictorWrapper.load_weights(predictor, first_shard) + loaded_second = common_mod.CodePredictorWrapper.load_weights(predictor, second_shard) + + assert {name for name, _ in second_shard} <= loaded_second + + def test_load_weights_raises_on_incomplete_shards(self, loaded_target_classes) -> None: + """Missing one of q/k/v (or gate/up) must surface a clear error. + + ``CodePredictorBaseModel`` is the inner transformer (its parameter + paths look like ``layers..self_attn.qkv_proj.weight`` -- without + the ``model.`` prefix that the outer ``CodePredictorWrapper`` strips + before forwarding). The test mimics that contract. + """ + _, _, _, code_predictor_model, _ = loaded_target_classes + cp_config, _ = _make_tiny_config(loaded_target_classes) + model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) + + # Find any actual self_attn layer prefix, to make the test robust + # to renaming of the inner model's parameter scheme. + attn_prefix = next( + ( + name[: -len(".qkv_proj.weight")] + for name in dict(model.named_parameters()).keys() + if name.endswith(".self_attn.qkv_proj.weight") + ), + None, + ) + assert attn_prefix is not None, "test config must contain at least one self_attn layer" + + num_heads = cp_config.num_attention_heads + num_kv_heads = cp_config.num_key_value_heads + head_dim = cp_config.head_dim + hidden = cp_config.hidden_size + + # Provide only q + k (no v). The fused param exists in the model, + # so load_weights must refuse to silently leave it uninitialized. + bad_weights: list[tuple[str, torch.Tensor]] = [ + (f"{attn_prefix}.q_proj.weight", torch.randn(num_heads * head_dim, hidden)), + (f"{attn_prefix}.k_proj.weight", torch.randn(num_kv_heads * head_dim, hidden)), + # v_proj missing on purpose. + ] + with pytest.raises(RuntimeError, match="incomplete fused shards"): + model.load_weights(bad_weights) + + def test_load_weights_raises_when_fused_shards_absent(self, loaded_target_classes) -> None: + """Entirely absent q/k/v or gate/up groups must not leave fused params + randomly initialized.""" + _, _, _, code_predictor_model, _ = loaded_target_classes + cp_config, _ = _make_tiny_config(loaded_target_classes) + model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size) + + weights = [ + (name, torch.randn_like(param)) + for name, param in model.named_parameters(remove_duplicate=False) + if not (name.endswith(".qkv_proj.weight") or name.endswith(".gate_up_proj.weight")) + ] + + with pytest.raises(RuntimeError, match="missing fused parameters"): + model.load_weights(weights) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py new file mode 100644 index 000000000..3552c5356 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py @@ -0,0 +1,322 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the Qwen3-TTS leading-silence ban (#4966, #5048). + +``compute_logits`` suppresses the checkpoint-derived silence codec vocabulary +for the first ``silence_ban_frames`` decode frames of a request, and only for +requests running Base voice-clone in x-vector-only mode. ICL onsets are +legitimate audio and must not be masked. + +Two per-request gates guard the mask, and both index by batch position: + +* the decode step, from ``sampling_metadata.output_token_ids[i]`` -- the + request's generated-token history, so the window closes on its own and a + preempted request cannot re-enter it; +* the generation mode, recorded at prefill in ``preprocess_batch`` and looked + up through the batch req_id order the runner supplies via + ``set_batch_req_ids``. + +These tests drive the methods directly against a talker built with ``__new__``: +no config, no checkpoint, no device. Only the attributes the masking path reads +are populated, and ``logits_processor`` is stubbed to return a fixed tensor. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import resolve_x_vector_only +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( + Qwen3TTSTalkerForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_VOCAB = 32 +_SILENCE_IDS = (3, 7, 11) + + +def _make_talker(*, ban_frames: int, logits: torch.Tensor, modes=None): + """Talker stub exposing only what ``compute_logits`` touches. + + ``modes`` is one entry per batch row: ``True`` for x-vector-only, ``False`` + for ICL, ``None`` for a request whose mode was never recorded. Defaults to + every row in scope, so a test that says nothing about mode exercises the + step gate alone. + """ + talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + + silence_mask = torch.zeros((_VOCAB,), dtype=torch.bool) + for token in _SILENCE_IDS: + silence_mask[token] = True + + if modes is None: + modes = [True] * logits.shape[0] + req_ids = [f"req{i}" for i in range(len(modes))] + + object.__setattr__(talker, "_silence_ban_frames", ban_frames) + object.__setattr__(talker, "_silence_mask", silence_mask) + object.__setattr__(talker, "_codec_disallowed_mask", torch.zeros((_VOCAB,), dtype=torch.bool)) + object.__setattr__(talker, "_batch_req_ids", req_ids) + object.__setattr__( + talker, + "_req_x_vector_only", + {rid: mode for rid, mode in zip(req_ids, modes, strict=True) if mode is not None}, + ) + object.__setattr__(talker, "lm_head", None) + object.__setattr__(talker, "logits_processor", lambda _head, _hidden: logits.clone()) + return talker + + +def _metadata(*steps: int): + """sampling_metadata whose output_token_ids encode one step per request.""" + return SimpleNamespace(output_token_ids=[[0] * step for step in steps]) + + +def _banned(row: torch.Tensor) -> set[int]: + return {int(i) for i in torch.nonzero(torch.isinf(row) & (row < 0)).flatten()} + + +# -------------------- the step gate -------------------- + + +def test_ban_disabled_leaves_logits_unchanged() -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=0, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(0)) + assert _banned(out[0]) == set(), "ban_frames=0 must be a no-op" + + +@pytest.mark.parametrize("step", [0, 1, 2]) +def test_early_steps_mask_only_the_silence_vocabulary(step: int) -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(step)) + assert _banned(out[0]) == set(_SILENCE_IDS), f"step {step} must mask exactly the silence ids" + + +@pytest.mark.parametrize("step", [3, 4, 50]) +def test_steps_past_the_window_are_untouched(step: int) -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(step)) + assert _banned(out[0]) == set(), f"step {step} is past the window; nothing may be masked" + + +def test_mixed_step_batch_masks_only_the_early_request() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 3)) + assert _banned(out[0]) == set(_SILENCE_IDS), "the step-0 request must be masked" + assert _banned(out[1]) == set(), "the step-3 request must not be masked" + + +# -------------------- the mode gate -------------------- + + +def test_icl_request_is_never_masked() -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits, modes=[False]) + out = talker.compute_logits(torch.zeros((1, 4)), _metadata(0)) + assert _banned(out[0]) == set(), "an ICL onset is real audio and must not be suppressed" + + +def test_mixed_mode_batch_masks_only_the_x_vector_request() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits, modes=[True, False]) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0)) + assert _banned(out[0]) == set(_SILENCE_IDS), "the x-vector-only request must be masked" + assert _banned(out[1]) == set(), "the ICL request in the same batch must not be" + + +def test_unrecorded_mode_is_left_alone() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits, modes=[True, None]) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0)) + assert _banned(out[0]) == set(_SILENCE_IDS) + assert _banned(out[1]) == set(), "a request that cannot be classified must not be masked" + + +def test_stale_batch_req_ids_skip_the_ban() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + # One req_id for a two-row batch: the recorded order no longer indexes the + # logits, so masking by position would hit the wrong request. + object.__setattr__(talker, "_batch_req_ids", ["req0"]) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0)) + assert _banned(out[0]) == set(), "a req_id/logits length mismatch must disable the ban" + assert _banned(out[1]) == set() + + +# -------------------- metadata plumbing -------------------- + + +def test_missing_sampling_metadata_does_not_crash() -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), None) + assert _banned(out[0]) == set() + + +@pytest.mark.parametrize( + "metadata", + [SimpleNamespace(output_token_ids=None), SimpleNamespace(), SimpleNamespace(output_token_ids=[])], +) +def test_absent_output_token_ids_warns_without_masking(metadata) -> None: + logits = torch.zeros((1, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((1, 4)), metadata) + assert _banned(out[0]) == set(), "the ban must skip, not crash, when history is unavailable" + + +def test_batch_length_mismatch_skips_the_ban() -> None: + logits = torch.zeros((2, _VOCAB)) + talker = _make_talker(ban_frames=3, logits=logits) + out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0)) + assert _banned(out[0]) == set(), "a runner shape mismatch must disable the ban, not misindex it" + assert _banned(out[1]) == set() + + +# -------------------- mode recording and eviction -------------------- + + +def _mode_talker(ban_frames: int = 3): + talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + object.__setattr__(talker, "_silence_ban_frames", ban_frames) + object.__setattr__(talker, "_req_x_vector_only", {}) + object.__setattr__(talker, "_batch_req_ids", []) + object.__setattr__(talker, "_prompt_builder", SimpleNamespace(preprocess_batch=lambda **_kwargs: None)) + return talker + + +def _info(*, task_type: str = "Base", xvec: bool = True, icl: bool | None = None) -> dict: + info: dict = {"task_type": [task_type], "x_vector_only_mode": [xvec]} + if icl is not None: + info["voice_clone_prompt"] = [{"icl_mode": icl}] + return info + + +def test_preprocess_batch_records_the_generation_mode() -> None: + talker = _mode_talker() + buffer = { + "req0": _info(xvec=True), + "req1": _info(xvec=False), + "req2": _info(task_type="CustomVoice"), + } + talker.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu")) + assert talker._req_x_vector_only == {"req0": True, "req1": False, "req2": None} + + +def test_recorded_mode_matches_the_prompt_builder_resolution() -> None: + """The recording path must agree with the resolution the prefill path uses; + a divergence here would mask requests the builder treats as in-context.""" + infos = [ + _info(xvec=True), + _info(xvec=False), + _info(xvec=True, icl=True), + _info(xvec=False, icl=False), + _info(task_type="VoiceDesign"), + ] + talker = _mode_talker() + buffer = {f"req{i}": info for i, info in enumerate(infos)} + talker.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu")) + for req_id, info in buffer.items(): + assert talker._req_x_vector_only[req_id] == resolve_x_vector_only(info) + + +def test_preprocess_batch_records_nothing_when_the_ban_is_off() -> None: + talker = _mode_talker(ban_frames=0) + buffer = {"req0": _info(xvec=True)} + talker.preprocess_batch(req_ids=["req0"], model_intermediate_buffer=buffer, device=torch.device("cpu")) + assert talker._req_x_vector_only == {} + + +def test_set_batch_req_ids_evicts_departed_requests() -> None: + talker = _mode_talker() + object.__setattr__(talker, "_req_x_vector_only", {"req0": True, "req1": False, "req2": True}) + talker.set_batch_req_ids(["req2", "req0"]) + assert talker._batch_req_ids == ["req2", "req0"], "batch order must be preserved, not sorted" + assert talker._req_x_vector_only == {"req0": True, "req2": True}, "req1 has left the batch" + + +# -------------------- _init_silence_mask sanity checks -------------------- +# +# The silence vocabulary is derived from the checkpoint at load time rather than +# hardcoded. A derivation that comes back empty, oversized, or out of range means +# the encode did not produce silence codes, so the ban disables itself instead of +# masking a guess. + +_TALKER_VOCAB = 2048 + + +def _make_init_talker(*, codes, ban_frames: int = 3, vocab: int = _TALKER_VOCAB): + """Talker stub exposing only what ``_init_silence_mask`` touches. + + ``codes`` is either the list ``_encode_ref_audio_batch`` should return or an + exception instance it should raise. + """ + talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + + def _encode(_wavs, _sr, *, device): + if isinstance(codes, Exception): + raise codes + return codes + + def _parameters(): + return iter([torch.zeros(1)]) + + object.__setattr__(talker, "_silence_ban_frames", ban_frames) + object.__setattr__(talker, "_silence_mask", torch.zeros((vocab,), dtype=torch.bool)) + object.__setattr__(talker, "talker_config", SimpleNamespace(vocab_size=vocab)) + object.__setattr__(talker, "_encode_ref_audio_batch", _encode) + object.__setattr__(talker, "parameters", _parameters) + return talker + + +def _codes_from(ids) -> list[torch.Tensor]: + """One [T, Q] codec tensor whose codebook-0 column holds ``ids``.""" + cb0 = torch.tensor(list(ids), dtype=torch.long).unsqueeze(1) + return [torch.cat([cb0, torch.zeros_like(cb0)], dim=1)] + + +def test_plausible_derivation_populates_the_mask() -> None: + talker = _make_init_talker(codes=_codes_from([4, 4, 9, 17])) + talker._init_silence_mask() + assert talker._silence_ban_frames == 3 + assert set(torch.nonzero(talker._silence_mask).flatten().tolist()) == {4, 9, 17} + + +def test_empty_derivation_disables_the_ban() -> None: + talker = _make_init_talker(codes=[]) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0, "an all-False mask must not stay enabled" + assert not talker._silence_mask.any() + + +def test_oversized_derivation_disables_the_ban() -> None: + limit = max(64, _TALKER_VOCAB // 20) + talker = _make_init_talker(codes=_codes_from(range(limit + 1))) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0, "a speech-sized vocabulary must not be masked" + assert not talker._silence_mask.any() + + +def test_out_of_range_derivation_disables_the_ban() -> None: + talker = _make_init_talker(codes=_codes_from([4, 9, _TALKER_VOCAB + 1])) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0, "a vocabulary-size disagreement must not be half-applied" + assert not talker._silence_mask.any() + + +def test_encoder_failure_disables_the_ban() -> None: + talker = _make_init_talker(codes=RuntimeError("encoder unavailable")) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0 + assert not talker._silence_mask.any() + + +def test_disabled_ban_skips_derivation_entirely() -> None: + talker = _make_init_talker(codes=RuntimeError("must not be called"), ban_frames=0) + talker._init_silence_mask() + assert talker._silence_ban_frames == 0 diff --git a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py index 050a5b724..53a90452e 100644 --- a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py +++ b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -76,8 +77,8 @@ def test_step_audio2_token2wav_async_chunk_batch_guard(): model._forward_async_chunk( input_ids=torch.tensor([1, 2, 3]), runtime_additional_information=[ - {"left_context_size": 0}, - {"left_context_size": 1}, + {"meta": {"left_context_size": 0}}, + {"meta": {"left_context_size": 1}}, ], ) @@ -116,7 +117,7 @@ def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch): out = model.forward( input_ids=torch.tensor([10, 11, 12]), positions=torch.tensor([0, 1, 2]), - runtime_additional_information=[{"left_context_size": 1}], + runtime_additional_information=[{"meta": {"left_context_size": 1}}], ) audio_list = out.multimodal_outputs["model_outputs"] @@ -126,12 +127,24 @@ def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch): assert calls == {"setup": 1, "stream": 1, "reset": 1} +def test_step_audio2_token2wav_routes_non_final_chunk_to_async_forward(): + model = _build_model() + model._forward_async_chunk = MagicMock() + + model.forward( + input_ids=torch.tensor([10, 11, 12]), + positions=torch.tensor([0, 1, 2]), + runtime_additional_information=[{"meta": {"left_context_size": 0}}], + ) + model._forward_async_chunk.assert_called_once() + + def test_step_audio2_token2wav_async_chunk_empty_eof_returns_zero_chunk(): model = _build_model() out = model._forward_async_chunk( input_ids=torch.tensor([], dtype=torch.int64), - runtime_additional_information=[{"left_context_size": 1}], + runtime_additional_information=[{"meta": {"left_context_size": 1}}], ) audio_list = out.multimodal_outputs["model_outputs"] diff --git a/tests/model_executor/models/test_encoder_quant_config.py b/tests/model_executor/models/test_encoder_quant_config.py index 28a1efa71..694e92a02 100644 --- a/tests/model_executor/models/test_encoder_quant_config.py +++ b/tests/model_executor/models/test_encoder_quant_config.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression test for #2686: pre-quantized methods must not apply quant config to vision / audio encoders. -For modelopt FP8/FP4/MXFP8 checkpoints the Thinker LM is the only -quantized component. Vision and audio encoder weights are BF16 with no -FP8 scale tensors — passing quant_config to them causes FP8 kernels to -run on BF16 weights, producing garbage embeddings. +For serialized ModelOpt and SVDQuant checkpoints the target model component is +the only quantized component. Vision and audio encoder weights are BF16 with no +serialized scale or correction tensors, so the same config must not be applied +to them. **Import note (pytest + ``tests/model_executor`` only):** :mod:`vllm_omni.quantization` :file:`__init__` pulls ``factory``/``inc`` and reproduces vLLM CustomOp double @@ -58,8 +58,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @pytest.mark.parametrize("method", sorted(PRE_QUANTIZED_METHODS)) def test_pre_quantized_returns_none(method: str) -> None: """visual_quant_config and audio_quant_config must be None for - pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8, - modelopt_mixed).""" + pre-quantized methods (ModelOpt variants and SVDQuant).""" mock_config = MagicMock() mock_config.get_name.return_value = method @@ -101,5 +100,5 @@ def test_component_config_passed_through() -> None: def test_pre_quantized_methods_contains_expected() -> None: """Guard against accidental removal of a known pre-quantized method.""" - expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"} + expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed", "svdquant"} assert PRE_QUANTIZED_METHODS == expected diff --git a/tests/model_executor/models/test_gepard_window.py b/tests/model_executor/models/test_gepard_window.py new file mode 100644 index 000000000..67c61c7ab --- /dev/null +++ b/tests/model_executor/models/test_gepard_window.py @@ -0,0 +1,363 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU coverage for Gepard's prompt layout, streaming-decode window and config. + +The e2e test needs a GPU, NeMo and the checkpoint, so it only runs nightly. +These pin the parts that are pure logic and break silently: which frames each +decode covers, which samples survive the lookback trim, the prompt layout, and +that the packaged deploy config sits where the example looks for it. + +The codec is replaced by a shift-invariant stand-in whose output depends only +on the codes, which makes "chunked emissions concatenate to one whole-history +decode" an exact assertion — the contract lookback exists to provide. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +import vllm_omni +from vllm_omni.config.stage_config import load_deploy_config +from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig +from vllm_omni.model_executor.models.gepard.gepard_talker import ( + CHUNK_FRAMES, + FIRST_CHUNK_FRAMES, + GepardTalkerForConditionalGeneration, + _GepardState, +) +from vllm_omni.model_executor.models.gepard.nanocodec import ( + LOOKBACK_FRAMES, + NanoCodec, + apply_end_of_speech_tail, +) +from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SAMPLES_PER_FRAME = 1024 +NUM_HEADS = 32 +SAMPLE_RATE = 22050 +# Text-token count at or above which the layout is left unrepeated, read from +# the checkpoint defaults so the boundary cases follow the config. +APPLY_BELOW = GepardConfig().text_repetition_apply_below + + +class _ShiftInvariantCodec: + """Frame tagged ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits.""" + + def __init__(self) -> None: + self.windows: list[int] = [] # frame count of every decode call + + def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor): + assert codes.shape[0] == 1, "streaming decode is per-request" + assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels" + tags = codes[0, 0, :].to(torch.int64) # (T,) + self.windows.append(int(tags.shape[0])) + offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64) + audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1) + return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME + + +def _make_codec() -> tuple[NanoCodec, _ShiftInvariantCodec]: + codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE) + fake = _ShiftInvariantCodec() + codec._codec = fake # bypasses load(); NeMo is never imported + assert codec.is_loaded + return codec, fake + + +def _make_talker(codec: NanoCodec) -> GepardTalkerForConditionalGeneration: + """A real talker with only the audio-path attributes populated. + + ``__init__`` needs a VllmConfig and builds the backbone; the methods under + test touch only these four attributes. + """ + talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration) + nn.Module.__init__(talker) # sets up _parameters/_buffers/_modules only + talker._codec = codec + talker._audio_queue = [] + talker._active_states = {} + talker._deferred_cleanup_ids = set() + return talker + + +def _frame(tag: int) -> torch.Tensor: + """One committed frame: 32 per-dimension codes, all carrying the tag.""" + return torch.full((NUM_HEADS,), tag, dtype=torch.long) + + +def _expected_samples(num_frames: int) -> torch.Tensor: + offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64) + tags = torch.arange(num_frames, dtype=torch.int64).unsqueeze(1) + return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32) + + +def _drain(talker) -> torch.Tensor: + """Concatenate and clear the queued deltas, as make_omni_output does.""" + if not talker._audio_queue: + return torch.empty(0) + out = torch.cat([a.reshape(-1) for _, a in talker._audio_queue]) + talker._audio_queue.clear() + return out + + +def _generate(talker, state, num_frames: int, *, stop_at_end: bool) -> None: + """Drive the commit loop the way forward() does, one frame per step.""" + for tag in range(num_frames): + state.frames.append(_frame(tag)) + state.frame_count += 1 + state.past_first_step = True + talker._emit_audio(state, is_final=stop_at_end and tag == num_frames - 1) + + +# Counts straddling both cadence thresholds, where an off-by-one hides. +@pytest.mark.parametrize( + "num_frames", + [ + 1, + 2, + FIRST_CHUNK_FRAMES - 1, + FIRST_CHUNK_FRAMES, + FIRST_CHUNK_FRAMES + 1, + CHUNK_FRAMES - 1, + CHUNK_FRAMES, + CHUNK_FRAMES + 1, + FIRST_CHUNK_FRAMES + CHUNK_FRAMES, + 79, + 200, + ], +) +def test_chunked_emissions_equal_a_whole_history_decode(num_frames: int) -> None: + codec, _fake = _make_codec() + talker = _make_talker(codec) + state = _GepardState(request_id="req-0") + + _generate(talker, state, num_frames, stop_at_end=True) + + torch.testing.assert_close(_drain(talker), _expected_samples(num_frames)) + assert state.emitted_frames == num_frames + + +def test_every_frame_is_emitted_exactly_once() -> None: + """The trim must not drop or duplicate samples at a chunk boundary.""" + codec, _fake = _make_codec() + talker = _make_talker(codec) + state = _GepardState(request_id="req-0") + num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES + 3 + + _generate(talker, state, num_frames, stop_at_end=True) + audio = _drain(talker) + + assert audio.numel() == num_frames * SAMPLES_PER_FRAME + # Tags are globally unique, so a duplicated or dropped frame shows up here. + assert len(torch.unique(audio)) == audio.numel() + + +def test_decode_windows_carry_lookback_context() -> None: + """Every decode after the first must re-decode LOOKBACK_FRAMES of context.""" + codec, fake = _make_codec() + talker = _make_talker(codec) + state = _GepardState(request_id="req-0") + num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES + + _generate(talker, state, num_frames, stop_at_end=True) + + assert len(fake.windows) >= 3, "expected a first chunk plus steady-state chunks" + assert fake.windows[0] == FIRST_CHUNK_FRAMES, "the first chunk has no left context" + for window in fake.windows[1:]: + assert window == CHUNK_FRAMES + LOOKBACK_FRAMES + + +def test_codec_runs_once_per_chunk_not_once_per_frame() -> None: + codec, fake = _make_codec() + talker = _make_talker(codec) + state = _GepardState(request_id="req-0") + num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES + + _generate(talker, state, num_frames, stop_at_end=True) + + assert len(fake.windows) == 4 + assert sum(fake.windows) < num_frames * 2 # far below one decode per frame + + +def test_final_emit_drains_a_partial_chunk() -> None: + """A request can end mid-chunk (max_tokens, STOP); the frames committed + since the last cadence boundary must still reach the queue. + + Arithmetic only: this drives ``_emit_audio`` directly, so it says nothing + about *whether* the engine reaches it on the request's last step. That + wiring lives in ``test_gepard_wiring.py`` — a bug there leaves this green. + """ + codec, _fake = _make_codec() + talker = _make_talker(codec) + state = _GepardState(request_id="req-0") + num_frames = FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5 + + _generate(talker, state, num_frames, stop_at_end=False) + assert state.emitted_frames < num_frames, "a partial chunk should still be pending" + + talker._emit_audio(state, is_final=True) + + torch.testing.assert_close(_drain(talker), _expected_samples(num_frames)) + assert state.emitted_frames == num_frames + + +def test_emitting_with_nothing_pending_is_a_no_op() -> None: + codec, fake = _make_codec() + talker = _make_talker(codec) + state = _GepardState(request_id="req-0") + + _generate(talker, state, 4, stop_at_end=True) + calls_after_flush = len(fake.windows) + talker._audio_queue.clear() + + talker._emit_audio(state, is_final=True) + + assert len(fake.windows) == calls_after_flush, "codec ran with no new frames" + assert talker._audio_queue == [] + + +def test_concurrent_requests_keep_separate_windows() -> None: + codec, _fake = _make_codec() + talker = _make_talker(codec) + states = [_GepardState(request_id=f"req-{i}") for i in range(2)] + + for state in states: + _generate(talker, state, FIRST_CHUNK_FRAMES, stop_at_end=True) + + by_req: dict[str, list[torch.Tensor]] = {} + for req_id, audio in talker._audio_queue: + by_req.setdefault(req_id, []).append(audio.reshape(-1)) + assert set(by_req) == {"req-0", "req-1"} + for chunks in by_req.values(): + torch.testing.assert_close(torch.cat(chunks), _expected_samples(FIRST_CHUNK_FRAMES)) + + +def test_end_of_speech_tail_fades_and_pads() -> None: + audio = torch.ones(SAMPLE_RATE) # 1 s of full-scale signal + out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=10.0, silence_ms=20.0) + + fade = int(SAMPLE_RATE * 10.0 / 1000.0) + pad = int(SAMPLE_RATE * 20.0 / 1000.0) + assert out.numel() == SAMPLE_RATE + pad + assert torch.all(out[-pad:] == 0.0), "trailing pad must be silent" + assert out[SAMPLE_RATE - fade - 1] == 1.0, "the fade must not reach back past fade_ms" + # Monotonically ramped to zero — the pad alone would leave the click. + ramp = out[SAMPLE_RATE - fade : SAMPLE_RATE] + assert torch.all(ramp[:-1] >= ramp[1:]) + assert ramp[-1] == 0.0 + + +def test_end_of_speech_tail_is_a_no_op_by_default() -> None: + audio = torch.ones(128) + out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=0.0, silence_ms=0.0) + assert out is audio + + +def _repeat_count(ids: list[int], cfg: GepardConfig) -> int: + """How many ``[SOT ... EOT]`` copies the layout carries.""" + return ids.count(cfg.start_of_text) + + +def test_gepard_prompt_layout() -> None: + """The prompt carries the speaker slots and exactly one SOS, at the end.""" + cfg = GepardConfig() + ids = build_gepard_prompt_ids([1, 2, 3], config=cfg) + + slots = [cfg.speaker_token_base + i for i in range(cfg.num_speaker_prefix)] + assert ids[: cfg.num_speaker_prefix] == slots + assert ids.count(cfg.start_of_speech) == 1, "only the canonical copy may carry SOS" + assert ids[-1] == cfg.start_of_speech, "the first frame is sampled from the SOS position" + + +def test_short_text_repeats_to_reach_the_target_token_mass() -> None: + """A short text is repeated until its region holds ~target_text_tokens.""" + cfg = GepardConfig() + text = [1, 2, 3, 4] # target 16 / 4 -> 4 copies + ids = build_gepard_prompt_ids(text, config=cfg) + + assert _repeat_count(ids, cfg) == 4 + assert ids.count(cfg.start_of_speech) == 1, "only the last copy opens the speech region" + assert ids[-1] == cfg.start_of_speech + + +@pytest.mark.parametrize( + ("num_text_tokens", "repeated"), + [ + (APPLY_BELOW - 1, True), + (APPLY_BELOW, False), + (APPLY_BELOW + 1, False), + ], +) +def test_apply_below_is_the_repetition_boundary(num_text_tokens: int, repeated: bool) -> None: + """Both sides of the threshold: at apply_below the text is left alone. + + The layout has to match training here or WER collapses, and the boundary + is where an off-by-one would be invisible in a single happy-path clip. + """ + cfg = GepardConfig() + ids = build_gepard_prompt_ids(list(range(1, num_text_tokens + 1)), config=cfg) + + assert (_repeat_count(ids, cfg) > 1) is repeated + + +def test_max_repeats_caps_a_one_token_text() -> None: + """target/1 would ask for 16 copies; the cap keeps the prefill bounded.""" + cfg = GepardConfig() + ids = build_gepard_prompt_ids([1], config=cfg) + + assert _repeat_count(ids, cfg) == cfg.text_repetition_max_repeats + + +def test_repetition_thresholds_are_read_from_the_config() -> None: + """The numbers come from the checkpoint sidecar, not from literals. + + A checkpoint whose text_repetition block differs must change the layout; + duplicating the values in the builder would silently ignore it. + """ + cfg = GepardConfig(text_repetition={"enabled": True, "target_text_tokens": 9, "apply_below": 8, "max_repeats": 2}) + assert _repeat_count(build_gepard_prompt_ids([1, 2, 3], config=cfg), cfg) == 2 # ceil(9/3)=3, capped at 2 + assert _repeat_count(build_gepard_prompt_ids(list(range(8)), config=cfg), cfg) == 1 # at apply_below + + +def test_repetition_can_be_disabled() -> None: + cfg = GepardConfig(text_repetition={"enabled": False}) + ids = build_gepard_prompt_ids([1], config=cfg) + + assert _repeat_count(ids, cfg) == 1 + assert ids[-1] == cfg.start_of_speech + + +def test_empty_text_is_rejected() -> None: + """An empty text would still assemble a valid layout and be voiced. + + ``[slots, SOT, EOT, SOS]`` parses, prefills and generates: the failure is + audible, not raised, which is the worst shape for a caller to debug. + """ + with pytest.raises(ValueError, match="at least one text token"): + build_gepard_prompt_ids([], config=GepardConfig()) + + +def test_packaged_deploy_config_is_where_the_example_looks_for_it() -> None: + """The example resolves its default off ``vllm_omni.__file__``; tests reach + the same YAML through ``tests.helpers.stage_config``, so nothing covered the + packaged location until this. + + ``pipeline: gepard`` is the load-bearing key: the checkpoint self-identifies + as qwen3_5_text, so without the pin the architectures fallback routes it to + the diffusion registry. + + This asserts the location, not the example's constant — importing the + example would pull ``Omni`` and NeMo into a CPU test. Running the quick + start verbatim stays a release check. + """ + path = Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml" + assert path.is_file(), f"packaged deploy config missing at {path}" + + deploy = load_deploy_config(path) + assert deploy.pipeline == "gepard" + assert deploy.stages, "deploy config declares no stages" diff --git a/tests/model_executor/models/test_gepard_wiring.py b/tests/model_executor/models/test_gepard_wiring.py new file mode 100644 index 000000000..9a08ddf0a --- /dev/null +++ b/tests/model_executor/models/test_gepard_wiring.py @@ -0,0 +1,732 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU coverage for the seam between Gepard's talker and the AR runner. + +``test_gepard_window.py`` pins the decode arithmetic *below* this layer, and +the e2e test exercises the layer *above* it but only along the path where the +stop head fires. Neither covers the wiring itself: whether anything calls the +final emit on a request's last step, and whether the payload it produces +survives the runner's sparse-audio routing. + +These drive the real ``preprocess`` -> ``forward`` -> ``make_omni_output`` +sequence through a stand-in for ``GPUARModelRunner.execute_model``, and use the +shared ``sparse_audio.resolve_sparse_mm_routing`` helper used by the runner to +decide what actually reaches the caller. Only the LM backbone and the frame +sampler are stubbed out; every method under test is the real one, so a test +here goes red when the engine +path breaks — which a test that calls the flush by hand cannot do. +""" + +from __future__ import annotations + +import logging + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig +from vllm_omni.model_executor.models.gepard.gepard_talker import ( + CHUNK_FRAMES, + FIRST_CHUNK_FRAMES, + GepardTalkerForConditionalGeneration, + _ReqInfo, +) +from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec +from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids +from vllm_omni.worker import sparse_audio + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SAMPLES_PER_FRAME = 1024 +NUM_HEADS = 32 +SAMPLE_RATE = 22050 +LOGGER_NAME = "vllm_omni.model_executor.models.gepard.gepard_talker" + +# Each request tags its frames from a distinct base, so a sample that reaches +# the wrong request is visible in the value, not just the length. +_TAG_STRIDE = 100_000 + + +class _TagCodec: + """Frame carrying tag ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits. + + Shift-invariant, so "what the caller received" can be compared against one + whole-history decode exactly. The tag rides channel 1: channel 0 is head0, + which the talker keeps inside its real (small) vocab. + """ + + def __init__(self) -> None: + self.decode_calls = 0 + + def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor): + assert codes.shape[0] == 1, "streaming decode is per-request" + assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels" + self.decode_calls += 1 + tags = codes[0, 1, :].to(torch.int64) # (T,) + offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64) + audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1) + return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME + + +class _StubBackbone(nn.Module): + """Stands in for the Qwen3.5 backbone; only shapes matter to this seam.""" + + def __init__(self, hidden: int) -> None: + super().__init__() + self.hidden = hidden + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros(int(input_ids.shape[0]), self.hidden) + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + src = inputs_embeds if inputs_embeds is not None else input_ids + return torch.zeros(int(src.shape[0]), self.hidden) + + +def _expected_samples(tag_base: int, num_frames: int) -> torch.Tensor: + offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64) + tags = (torch.arange(num_frames, dtype=torch.int64) + tag_base).unsqueeze(1) + return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32) + + +class _MiniEngine: + """Drives the talker the way ``GPUARModelRunner.execute_model`` does. + + Reproduces the three orderings the dropped-tail bug lived in: + + * ``on_requests_finished`` fires at the START of the step AFTER the one + that finished the request (it reads ``scheduler_output.finished_req_ids``); + * a finished request is not scheduled again, so it is absent from + ``req_ids_output_copy`` and from that step's routing; + * with nothing left to schedule there is no step at all, so no further + ``forward`` runs; + * a stop-token finish gets one more scheduled step, a budget finish does + not -- under async scheduling only the latter is guarded against it. + """ + + def __init__(self, talker: GepardTalkerForConditionalGeneration, *, max_model_len: int) -> None: + self.talker = talker + self.max_model_len = max_model_len + self.reqs: dict[str, dict] = {} + self.order: list[str] = [] + self._finished_last_step: list[str] = [] + self.delivered: dict[str, list[torch.Tensor]] = {} + self.committed_frames: dict[str, int] = {} + self.steps = 0 + self._tags_handed_out = 0 + + def add_request( + self, + req_id: str, + *, + max_tokens: int, + stop_after: int | None = None, + text_len: int = 3, + seed: int | None = None, + ) -> None: + """Submit a request. An id already seen is a resubmit of that id. + + Tags come from a counter, so a resubmitted id gets frames the previous + request could not have produced. + """ + prompt_ids = build_gepard_prompt_ids(list(range(10, 10 + text_len)), config=self.talker.config) + self._tags_handed_out += 1 + self.reqs[req_id] = { + "prompt_ids": torch.tensor(prompt_ids, dtype=torch.long), + "prompt_len": len(prompt_ids), + "max_tokens": max_tokens, + "stop_after": stop_after, + "stop_fired": False, + "seed": seed, + "computed": 0, + "generated": 0, + "finished": False, + "ghost_step_pending": False, + "next_token": None, + "tag_base": _TAG_STRIDE * self._tags_handed_out, + } + if req_id not in self.order: + self.order.append(req_id) + self.delivered[req_id] = [] + self.committed_frames[req_id] = 0 + + def abort(self, req_id: str) -> None: + """Kill a request the way an aborted client does: no STOP, no budget.""" + self.reqs[req_id]["finished"] = True + self._finished_last_step.append(req_id) + + def _live(self) -> list[str]: + return [r for r in self.order if not self.reqs[r]["finished"]] + + def step(self) -> bool: + """Run one engine step. False when there was nothing to schedule.""" + live = self._live() + if not live: + return False + talker = self.talker + self.steps += 1 + + # --- runner: notify the model of requests finished since last step --- + if self._finished_last_step: + talker.on_requests_finished(set(self._finished_last_step)) + self._finished_last_step = [] + + # --- runner: per-request preprocess, in input_batch.req_ids order --- + embeds_parts, ids_parts = [], [] + for req_id in live: + r = self.reqs[req_id] + is_prefill = r["computed"] < r["prompt_len"] + if is_prefill: + ids = r["prompt_ids"] + else: + ids = torch.tensor([r["next_token"] or 0], dtype=torch.long) + span = int(ids.shape[0]) + out_ids, out_embeds, _ = talker.preprocess( + input_ids=ids, + input_embeds=None, + request_id=req_id, + _omni_prompt_len=r["prompt_len"], + _omni_num_computed_tokens=r["computed"], + _omni_is_prefill=is_prefill, + _omni_max_tokens=r["max_tokens"], + _omni_seed=r["seed"], + ) + ids_parts.append(out_ids) + embeds_parts.append(out_embeds) + r["span"] = span + + inputs_embeds = torch.cat(embeds_parts, dim=0) + input_ids = torch.cat(ids_parts, dim=0) + positions = torch.arange(int(input_ids.shape[0]), dtype=torch.long) + + # --- runner: forward, then build the multimodal payload --- + hidden = talker.forward(input_ids, positions, inputs_embeds=inputs_embeds) + results = list(talker._results_queue) + omni_out = talker.make_omni_output(hidden) + mm = omni_out.multimodal_outputs + # compute_logits drains this; it does not touch the audio path. + talker._results_queue.clear() + + for req_id in live: + state = talker._active_states.get(req_id) + if state is not None: + self.committed_frames[req_id] = state.frame_count + + # --- runner: sparse-audio routing decides what reaches the caller --- + downstream, sparse_index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=list(live), + downstream_req_ids=list(live), + multimodal_outputs=mm, + ) + assert is_sparse, "gepard always marks meta.sparse_audio" + for rid in downstream: + self.delivered[rid].append(mm["model_outputs"][sparse_index[rid]].reshape(-1)) + + # --- runner/scheduler bookkeeping: advance and detect finishes --- + stopped = {rid for rid, _head0, do_stop in results if do_stop} + for req_id in live: + r = self.reqs[req_id] + r["computed"] += r["span"] + r["generated"] = r["computed"] - r["prompt_len"] + 1 + r["next_token"] = 0 + length_cap = min(r["max_tokens"], self.max_model_len - r["prompt_len"]) + if r["ghost_step_pending"]: + # This step WAS the extra one. The engine finished the request + # on the stop token last step and discards whatever came back + # here, whether or not the model chose to stop again. + r["finished"] = True + self._finished_last_step.append(req_id) + elif req_id in stopped: + # Ends on the stop token: one more step is already scheduled. + r["ghost_step_pending"] = True + elif r["generated"] >= length_cap: + r["finished"] = True + self._finished_last_step.append(req_id) + return True + + def run(self, max_steps: int = 4000) -> None: + for _ in range(max_steps): + if not self.step(): + return + raise AssertionError("engine did not drain") + + def audio_for(self, req_id: str) -> torch.Tensor: + chunks = self.delivered[req_id] + return torch.cat(chunks) if chunks else torch.empty(0) + + +def _make_talker(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _TagCodec]: + """A real talker with only the audio-path collaborators populated. + + ``__init__`` needs a VllmConfig and builds the Qwen3.5 backbone; the seam + under test needs none of that, so the backbone and the frame sampler are + stubbed and everything else is the class's own code. + """ + cfg = GepardConfig() + # Fixed, not cfg.get_text_config().hidden_size: nothing in this seam reads + # the real width, and a small one keeps the test off the backbone config. + hidden = 64 + codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE) + fake = _TagCodec() + codec._codec = fake # bypasses load(); NeMo is never imported + assert codec.is_loaded + + talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration) + nn.Module.__init__(talker) + talker.config = cfg + talker.model = _StubBackbone(hidden) + # requires_grad=False: preprocess writes the prefix into inputs_embeds + # in place, and production runs that under inference_mode. + talker.null_prefix = nn.Parameter(torch.zeros(cfg.num_speaker_prefix, hidden), requires_grad=False) + talker._codec = codec + talker._sample_rate = SAMPLE_RATE + talker._max_model_len = max_model_len + talker._active_states = {} + talker._pending_requests = [] + talker._results_queue = [] + talker._audio_queue = [] + talker._deferred_cleanup_ids = set() + return talker, fake + + +def _install_scripted_sampling(talker: GepardTalkerForConditionalGeneration, engine: _MiniEngine) -> None: + """Replace the 32-head sampler with a scripted one. + + Frames carry a per-request, per-step tag so the caller's audio is + self-describing. STOP fires once, on the request's configured frame: on the + extra step the engine then runs, re-firing is likely but not guaranteed, + and the model may not assume it. + """ + head0_vocab = talker.config.head0_vocab_size + + # Signature mirrors the real ``_sample_frame`` minus ``self``; the scripted + # sampler ignores ``generators`` because seeding is decided in + # ``_build_with_real_sampler``, which keeps the model's own sampler. + def _sample_frame(rows: torch.Tensor, generators: list[torch.Generator | None] | None = None): + n = int(rows.shape[0]) + head0 = torch.zeros(n, dtype=torch.long) + heads = torch.zeros(n, NUM_HEADS - 1, dtype=torch.long) + stop = torch.zeros(n, dtype=torch.bool) + for i, (req_id, _span, _samples_frame, *_rest) in enumerate(talker._pending_requests): + state = talker._active_states.get(req_id) + frame_index = state.frame_count if state is not None else 0 + tag = engine.reqs[req_id]["tag_base"] + frame_index + head0[i] = tag % head0_vocab + heads[i] = tag + r = engine.reqs[req_id] + if r["stop_after"] is not None and frame_index >= r["stop_after"] and not r["stop_fired"]: + r["stop_fired"] = True + stop[i] = True + return head0, heads, stop + + talker._sample_frame = _sample_frame + talker._audio_frame_embed = lambda h0, h: torch.zeros(int(h0.shape[0]), talker.model.hidden) + + +def _build(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _TagCodec, _MiniEngine]: + talker, codec = _make_talker(max_model_len=max_model_len) + engine = _MiniEngine(talker, max_model_len=max_model_len) + _install_scripted_sampling(talker, engine) + return talker, codec, engine + + +def _build_with_real_sampler(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _MiniEngine]: + """Same seam, but ``_sample_frame`` is the model's own 32-head sampler. + + The stub backbone hands every row the same hidden state, so the sampled + codes — and therefore the audio the codec tags them into — vary with + nothing but the Gumbel draws. That makes this the one place where "did the + caller's seed reach the audio" is a decidable question. + """ + talker, _codec = _make_talker(max_model_len=max_model_len) + engine = _MiniEngine(talker, max_model_len=max_model_len) + cfg = talker.config + talker.vocab_sizes = list(cfg.audio_head_levels) + talker.num_heads = cfg.num_audio_heads + talker.head0_vocab = cfg.head0_vocab_size + talker.stop_token = cfg.stop_token + talker.stop_threshold = cfg.stop_threshold + talker.temperature = cfg.temperature + talker._greedy = False + gather_idx, mask = GepardTalkerForConditionalGeneration._build_gather_mask(talker.vocab_sizes) + talker._cb_gather_idx = gather_idx + talker._cb_mask = mask + # Fixed weights: two engine runs must differ only by their RNG. + torch.manual_seed(0) + talker.fused_codebook_head = nn.Linear(talker.model.hidden, sum(talker.vocab_sizes)) + talker.stop_head = nn.Linear(talker.model.hidden, 1) + # Never stop: these requests end on their token budget, so the frame count + # is fixed and only the sampled codes can differ. + torch.nn.init.constant_(talker.stop_head.bias, -20.0) + talker._audio_frame_embed = lambda h0, h: torch.zeros(int(h0.shape[0]), talker.model.hidden) + return talker, engine + + +# -------------------------------------------------------------------------- +# The tail must reach the caller, whichever way the request ends. +# -------------------------------------------------------------------------- + + +def test_last_in_flight_request_delivers_every_committed_frame() -> None: + """The only request runs out of budget: there is no later step to flush in. + + This is the case a deferred flush cannot serve at all — nothing runs after + the final forward. + """ + _talker, _codec, engine = _build() + engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5) + engine.run() + + frames = engine.committed_frames["req-0"] + assert frames > 0 + torch.testing.assert_close( + engine.audio_for("req-0"), + _expected_samples(engine.reqs["req-0"]["tag_base"], frames), + ) + + +@pytest.mark.parametrize( + "max_tokens", + [ + 1, + FIRST_CHUNK_FRAMES, + FIRST_CHUNK_FRAMES + 1, + 30, + FIRST_CHUNK_FRAMES + CHUNK_FRAMES, + 100, + 150, + ], +) +def test_budget_truncated_request_delivers_every_committed_frame(max_tokens: int) -> None: + """Lowering max_tokens must shorten the clip, never silently cut it back + to the last cadence boundary.""" + _talker, _codec, engine = _build() + engine.add_request("req-0", max_tokens=max_tokens) + engine.run() + + frames = engine.committed_frames["req-0"] + assert frames == max_tokens, "every sampled step commits one frame when STOP never fires" + torch.testing.assert_close( + engine.audio_for("req-0"), + _expected_samples(engine.reqs["req-0"]["tag_base"], frames), + ) + + +def test_stop_head_path_still_delivers_every_committed_frame() -> None: + """Regression guard: the path that already worked must keep working.""" + _talker, _codec, engine = _build() + engine.add_request("req-0", max_tokens=4000, stop_after=FIRST_CHUNK_FRAMES + 7) + engine.run() + + frames = engine.committed_frames["req-0"] + assert frames == FIRST_CHUNK_FRAMES + 7, "the STOP frame itself is not committed" + torch.testing.assert_close( + engine.audio_for("req-0"), + _expected_samples(engine.reqs["req-0"]["tag_base"], frames), + ) + + +def test_the_step_after_a_stop_token_commits_nothing(caplog: pytest.LogCaptureFixture) -> None: + """Async scheduling runs one more step on a request that already stopped. + + Committing its frame would strand it — nothing ships after the stop step — + and would make the truncation warning fire on a request that delivered + everything it produced, which is the one thing that warning must not do. + """ + talker, _codec, engine = _build() + stop_after = FIRST_CHUNK_FRAMES + 6 + engine.add_request("req-0", max_tokens=4000, stop_after=stop_after) + + target = logging.getLogger(LOGGER_NAME) + target.addHandler(caplog.handler) + prev = target.level + target.setLevel(logging.WARNING) + try: + engine.run() + # The cleanup that would warn runs on the step after the request + # leaves, and here there is none. Drive it the way the runner would. + talker.on_requests_finished({"req-0"}) + talker._flush_deferred_cleanup() + finally: + target.removeHandler(caplog.handler) + target.setLevel(prev) + + assert engine.reqs["req-0"]["ghost_step_pending"], "the harness did not run the extra step" + frames = engine.committed_frames["req-0"] + assert frames == stop_after, f"the extra step committed a frame: {frames} > {stop_after}" + torch.testing.assert_close( + engine.audio_for("req-0"), + _expected_samples(engine.reqs["req-0"]["tag_base"], frames), + ) + assert not [m for m in caplog.messages if "undelivered frame" in m], ( + f"a fully delivered request was reported as truncated; saw {caplog.messages}" + ) + + +def test_a_resubmitted_request_id_keeps_its_new_state(caplog: pytest.LogCaptureFixture) -> None: + """An aborted id can come back before its cleanup has run. + + The runner treats the two as distinct requests, so the pending free must + not take the second one's state: without that, the state is popped at the + end of the very forward that created it and every later decode falls back + to the zero-embedding path — the request keeps producing frames, so only + the audio is wrong. + """ + _talker, _codec, engine = _build() + engine.add_request("req-0", max_tokens=4000) + for _ in range(FIRST_CHUNK_FRAMES + 2): + engine.step() + engine.abort("req-0") + + # Same id, resubmitted before the model has been told about the abort. + engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5) + + target = logging.getLogger(LOGGER_NAME) + target.addHandler(caplog.handler) + prev = target.level + target.setLevel(logging.WARNING) + try: + engine.run() + finally: + target.removeHandler(caplog.handler) + target.setLevel(prev) + + frames = engine.committed_frames["req-0"] + assert frames == engine.reqs["req-0"]["max_tokens"], "the second request lost frames" + torch.testing.assert_close( + engine.audio_for("req-0"), + _expected_samples(engine.reqs["req-0"]["tag_base"], frames), + msg="the second request did not receive its own audio", + ) + # The abort still gets reported: the cleanup that would have said so is the + # one this cancelled. + assert any("resubmitted" in m for m in caplog.messages), ( + f"the aborted request's truncation went unreported; saw {caplog.messages}" + ) + + +def test_a_preempted_request_is_reported_rather_than_resumed(caplog: pytest.LogCaptureFixture) -> None: + """Recompute preemption re-prefills a request that was mid-generation. + + This model cannot survive that, and the point of the test is that it says + so. The recomputed ids are the prompt plus the head0 codes sampled so far, + which go through the text embedding table, and each frame's other 31 codes + are never in the token stream at all — there is nothing to stitch the two + halves back together with, so generation restarts from frame 0. + + The harness re-prefills with the bare prompt, which understates the damage. + What is under test is the report and the reset, not the corruption, which + needs a real backbone to observe. + """ + _talker, _codec, engine = _build() + engine.add_request("req-0", max_tokens=4000) + for _ in range(FIRST_CHUNK_FRAMES + 2): + engine.step() + assert engine.committed_frames["req-0"] > 0 + + # Preemption the way the scheduler does it: computed tokens back to 0, so + # the next step re-prefills, and no id reaches on_requests_finished. + engine.reqs["req-0"]["computed"] = 0 + + target = logging.getLogger(LOGGER_NAME) + target.addHandler(caplog.handler) + prev = target.level + target.setLevel(logging.WARNING) + try: + engine.step() + finally: + target.removeHandler(caplog.handler) + target.setLevel(prev) + + assert any("preempted" in m for m in caplog.messages), ( + f"a preempted request restarted its audio silently; saw {caplog.messages}" + ) + assert not any("resubmitted" in m for m in caplog.messages), ( + "a preemption was reported as an id resubmit; the two need different words" + ) + assert engine.committed_frames["req-0"] == 1, "the re-prefill must start the frame history over" + + +def test_concurrent_requests_each_receive_their_own_full_audio() -> None: + """One request finishes while another stays live. + + The earlier finisher is the case where a deferred flush *does* run but the + payload is filtered out again, because its id has left the output batch. + """ + _talker, _codec, engine = _build() + engine.add_request("req-short", max_tokens=FIRST_CHUNK_FRAMES + 3) + engine.add_request("req-long", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 11) + engine.run() + + for req_id in ("req-short", "req-long"): + frames = engine.committed_frames[req_id] + assert frames == engine.reqs[req_id]["max_tokens"] + torch.testing.assert_close( + engine.audio_for(req_id), + _expected_samples(engine.reqs[req_id]["tag_base"], frames), + msg=f"{req_id} lost or gained samples", + ) + + +def test_codec_still_runs_once_per_chunk_not_once_per_frame() -> None: + """Emitting on the final step must not degrade into per-frame decoding.""" + _talker, codec, engine = _build() + num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES + engine.add_request("req-0", max_tokens=num_frames) + engine.run() + + assert codec.decode_calls == 4, f"expected one decode per chunk, got {codec.decode_calls}" + + +# -------------------------------------------------------------------------- +# Deferred cleanup is cleanup: it must not try to produce output. +# -------------------------------------------------------------------------- + + +def test_deferred_cleanup_never_queues_audio() -> None: + """Anything queued from the cleanup hook is unroutable by construction. + + By the time it runs the id has left ``req_ids_output_copy``, so a payload + under it is dropped; queueing there only risks leaking one request's audio + into the next step's payload for a different request. + """ + talker, _codec, engine = _build() + engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + 4) + engine.add_request("req-1", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES) + while engine.step(): + if engine.reqs["req-0"]["finished"] and not engine.reqs["req-1"]["finished"]: + break + + talker._audio_queue.clear() + talker.on_requests_finished({"req-0"}) + talker._flush_deferred_cleanup() + + assert talker._audio_queue == [], "cleanup emitted audio that routing cannot deliver" + assert "req-0" not in talker._active_states, "cleanup must still free the state" + + +def test_cleanup_warns_when_frames_are_left_undelivered(caplog: pytest.LogCaptureFixture) -> None: + """An abort is the one case with no last step to emit on. Say so.""" + _talker, _codec, engine = _build() + engine.add_request("req-0", max_tokens=4000) + engine.add_request("req-1", max_tokens=4000) + for _ in range(FIRST_CHUNK_FRAMES + 3): + engine.step() + engine.abort("req-0") + + target = logging.getLogger(LOGGER_NAME) + target.addHandler(caplog.handler) + prev = target.level + target.setLevel(logging.WARNING) + try: + engine.step() # cleanup for req-0 runs at the start of this step + finally: + target.removeHandler(caplog.handler) + target.setLevel(prev) + + assert any("undelivered frame" in m for m in caplog.messages), ( + f"aborted request dropped audio without a warning; saw {caplog.messages}" + ) + + +# -------------------------------------------------------------------------- +# A request's seed must reach the in-model sampler. +# -------------------------------------------------------------------------- + +_SEEDED_FRAMES = FIRST_CHUNK_FRAMES + 5 + + +def _run_seeded(seed: int | None, *, neighbour: bool = False, neighbour_seed: int | None = None) -> torch.Tensor: + """Audio delivered to ``req-a``, optionally sharing the batch with another.""" + _talker, engine = _build_with_real_sampler() + engine.add_request("req-a", max_tokens=_SEEDED_FRAMES, seed=seed) + if neighbour: + engine.add_request("req-b", max_tokens=_SEEDED_FRAMES, seed=neighbour_seed) + engine.run() + return engine.audio_for("req-a") + + +def test_a_seeded_request_is_reproducible_within_one_process() -> None: + """The bug this pins: the 32 heads are sampled in ``forward``, so a seed + that only reaches vLLM's sampler leaves the audio at the mercy of the + global RNG — two identical requests in one process then diverge.""" + torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234)) + + +def test_different_seeds_produce_different_audio() -> None: + """Guards the opposite failure: a seed that is threaded but never used.""" + assert not torch.equal(_run_seeded(1234), _run_seeded(4321)) + + +def test_a_seeded_request_is_unaffected_by_its_batch() -> None: + """Each seeded request draws its own noise, so the noise it sees cannot + depend on who else was scheduled alongside it. + + The stub backbone is what makes that checkable. A real one is a bf16 + batched matmul whose rows are not reproducible across batch shapes, so + end-to-end batch invariance is not attainable there; this pins the half + the model actually controls. + """ + torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234, neighbour=True, neighbour_seed=999)) + + +def test_an_unseeded_neighbour_does_not_disturb_a_seeded_request() -> None: + """The mixed case: one row draws from the global RNG, the other must not.""" + torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234, neighbour=True)) + + +# -------------------------------------------------------------------------- +# Budget arithmetic, straight through the helper. +# -------------------------------------------------------------------------- + + +def _info(**kwargs) -> _ReqInfo: + base = { + "request_id": "req-0", + "is_prefill": False, + "prompt_len": 20, + "num_computed_tokens": 20, + "max_tokens": 10, + "seed": None, + } + base.update(kwargs) + return _ReqInfo(**base) + + +def test_prefill_step_is_the_first_output_token() -> None: + talker, _codec, _engine = _build() + first = _info(is_prefill=True, num_computed_tokens=0, max_tokens=1) + assert talker._is_last_output_token(first, span=20, samples_frame=True) is True + more = _info(is_prefill=True, num_computed_tokens=0, max_tokens=2) + assert talker._is_last_output_token(more, span=20, samples_frame=True) is False + + +def test_partial_prefill_chunk_does_not_consume_budget() -> None: + """vLLM discards that row's token, so it must not count against max_tokens.""" + talker, _codec, _engine = _build() + info = _info(is_prefill=True, num_computed_tokens=0, prompt_len=40, max_tokens=1) + assert talker._is_last_output_token(info, span=10, samples_frame=False) is False + + +def test_budget_is_capped_by_max_model_len() -> None: + """A max_tokens larger than the context window never gets reached.""" + talker, _codec, _engine = _build(max_model_len=30) + # prompt 20 + 10 generated fills the window, whatever max_tokens says. + at_cap = _info(num_computed_tokens=28, max_tokens=1000) + assert talker._is_last_output_token(at_cap, span=1, samples_frame=True) is True + below_cap = _info(num_computed_tokens=27, max_tokens=1000) + assert talker._is_last_output_token(below_cap, span=1, samples_frame=True) is False + + +def test_missing_budget_is_not_treated_as_exhausted() -> None: + """A runner that does not thread max_tokens must not trigger a fade.""" + talker, _codec, _engine = _build(max_model_len=0) + assert talker._is_last_output_token(_info(max_tokens=None), span=1, samples_frame=True) is False diff --git a/tests/model_executor/models/test_minimax_h3_token_tags.py b/tests/model_executor/models/test_minimax_h3_token_tags.py new file mode 100644 index 000000000..69ee7cb84 --- /dev/null +++ b/tests/model_executor/models/test_minimax_h3_token_tags.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration + +from vllm_omni.model_executor.models.minimax_h3.text_encoder import MiniMaxH3TextEncoder + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _model_without_weights() -> MiniMaxH3TextEncoder: + model = MiniMaxH3TextEncoder.__new__(MiniMaxH3TextEncoder) + model.config = SimpleNamespace( + vision_start_token_id=1, + vision_end_token_id=2, + image_token_id=3, + video_token_id=4, + ) + model._token_tags = None + return model + + +def test_embed_failure_clears_token_tags(): + model = _model_without_weights() + model._token_tags = torch.tensor([0]) + + with ( + patch.object(Qwen3VLForConditionalGeneration, "embed_input_ids", side_effect=RuntimeError("failed")), + pytest.raises(RuntimeError, match="failed"), + ): + model.embed_input_ids(torch.tensor([5, 3])) + + assert model._token_tags is None + + +def test_invalid_output_length_consumes_token_tags(): + model = _model_without_weights() + model._token_tags = torch.tensor([1, 0]) + + with pytest.raises(RuntimeError, match="token tags exceed"): + model.make_omni_output(torch.zeros(1, 8)) + + assert model._token_tags is None + + +def test_make_omni_output_uses_shared_stage_payload_fields(): + model = _model_without_weights() + model._token_tags = torch.tensor([1, 0]) + hidden = torch.zeros(2, 5120) + + output = model.make_omni_output(hidden) + + assert output.multimodal_outputs is not None + assert torch.equal(output.multimodal_outputs["hidden_states"]["output"], hidden) + assert torch.equal(output.multimodal_outputs["meta"]["token_role_ids"], torch.tensor([[1], [0]])) + assert "encoder_hidden_states" not in output.multimodal_outputs + assert "token_tags" not in output.multimodal_outputs diff --git a/tests/model_executor/models/test_minimax_music3_batch_payload.py b/tests/model_executor/models/test_minimax_music3_batch_payload.py new file mode 100644 index 000000000..531500cd4 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_batch_payload.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MiniMax Music 3 batched payload partitioning. + +The stage-0 payload is split per request by indexing list values: a list is +indexed by batch row, anything else is handed to every request unchanged. So +per-request metadata has to be carried as lists. Merging it into one dict made +every request in a batch receive the last request's seed, window offset and +terminal flag, which crosses seeds between songs, misplaces acoustic windows, +and lets one request's terminal flag end another request's stream early. + +These drive ``make_omni_output`` directly with a hand-built batch, so no GPU +and no engine are needed. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.minimax_music3.talker import ( + MiniMaxMusic3TalkerForConditionalGeneration as Talker, +) +from vllm_omni.utils.mm_outputs import to_payload_element + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_WIDTH = 4 + + +def _info(*, seed: int, start: int, frames: int, last: bool) -> dict: + """One batch row's intermediate buffer, as ``postprocess`` leaves it.""" + return { + "latent": torch.full((frames, _WIDTH), float(seed)), + "codes": {"audio": torch.full((frames, 8), seed)}, + "meta": { + "audio_seed": seed, + "num_processed_tokens": start, + "codec_chunk_frames": frames, + "last_chunk": last, + }, + } + + +def _payload_for(output, row: int) -> dict: + """Partition the batch payload down to one request, as the runner does.""" + flat = {} + for key, value in output.multimodal_outputs.items(): + if isinstance(value, dict): + for qual, val in value.items(): + flat[f"{key}.{qual}"] = val + else: + flat[key] = value + return {key: to_payload_element(value, row, 1, 1) for key, value in flat.items()} + + +class TestBatchedMetadataPartitioning: + """Two requests in one batch must keep their own metadata.""" + + @staticmethod + def _two_request_batch(): + infos = [ + _info(seed=111, start=100, frames=3, last=False), + _info(seed=222, start=200, frames=5, last=True), + ] + talker = Talker.__new__(Talker) + return talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos) + + def test_seeds_are_not_crossed(self) -> None: + """A song decoded with another request's seed is a different song.""" + output = self._two_request_batch() + assert _payload_for(output, 0)["meta.audio_seed"] == 111 + assert _payload_for(output, 1)["meta.audio_seed"] == 222 + + def test_window_offsets_stay_with_their_request(self) -> None: + """The start frame decides where a window lands in the song.""" + output = self._two_request_batch() + assert _payload_for(output, 0)["meta.num_processed_tokens"] == 100 + assert _payload_for(output, 1)["meta.num_processed_tokens"] == 200 + + def test_terminal_flag_does_not_leak_across_requests(self) -> None: + """One request finishing must not end another request's stream.""" + output = self._two_request_batch() + assert _payload_for(output, 0)["meta.last_chunk"] is False + assert _payload_for(output, 1)["meta.last_chunk"] is True + + def test_span_length_matches_its_own_latent(self) -> None: + output = self._two_request_batch() + for row, frames in ((0, 3), (1, 5)): + entry = _payload_for(output, row) + assert entry["meta.codec_chunk_frames"] == frames + assert entry["latent"].shape[0] == frames + + def test_latent_and_codes_follow_the_same_row(self) -> None: + """The tensors were already per-request; pin that they stay aligned.""" + output = self._two_request_batch() + for row, seed in ((0, 111), (1, 222)): + entry = _payload_for(output, row) + assert entry["latent"][0, 0].item() == pytest.approx(float(seed)) + assert int(entry["codes.audio"][0, 0]) == seed + + +class TestPartialBatches: + """Only some rows carry a span on any given step.""" + + def test_row_without_a_span_does_not_take_another_row_metadata(self) -> None: + infos = [{}, _info(seed=777, start=300, frames=2, last=True)] + talker = Talker.__new__(Talker) + output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos) + + idle = _payload_for(output, 0) + assert idle["latent"].numel() == 0 + assert idle["meta.audio_seed"] is None + assert idle["meta.last_chunk"] is None + + active = _payload_for(output, 1) + assert active["meta.audio_seed"] == 777 + assert active["meta.last_chunk"] is True + + def test_batch_with_no_spans_emits_no_payload(self) -> None: + talker = Talker.__new__(Talker) + output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=[{}, {}]) + assert output.multimodal_outputs is None + + def test_consumed_entries_are_cleared_from_the_buffer(self) -> None: + """The runner never clears this buffer, so a span left behind would be + re-emitted on every later step.""" + infos = [_info(seed=1, start=0, frames=2, last=False)] + talker = Talker.__new__(Talker) + talker.make_omni_output(torch.zeros(1, _WIDTH), model_intermediate_buffer=infos) + assert "latent" not in infos[0] + assert "meta" not in infos[0] diff --git a/tests/model_executor/models/test_minimax_music3_frame_state.py b/tests/model_executor/models/test_minimax_music3_frame_state.py new file mode 100644 index 000000000..73dc9215b --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_frame_state.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MiniMax Music 3 per-request AR frame bookkeeping. + +The talker emits overlapping windows while dropping frames it can no longer +need, and flushes whatever is left as one span when the song ends. Getting the +eviction boundary wrong by one frame silently corrupts the conditioning handed +to the acoustic stage from the second window onward, which is audible but hard +to attribute, so the lifecycle is simulated here frame by frame. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.minimax_music3.chunking import chunk_windows +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + NUM_CODEBOOKS, +) +from vllm_omni.model_executor.models.minimax_music3.frame_state import ( + FrameBuffer, + RequestARState, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# One scalar per frame instead of 32768 dims: the bookkeeping is what is under +# test, and a frame's value doubles as its identity. +_WIDTH = 4 + + +def _frame(index: int) -> torch.Tensor: + return torch.full((1, _WIDTH), float(index)) + + +def _identities(chunk: torch.Tensor) -> list[int]: + return [int(row[0].item()) for row in chunk] + + +class TestFrameBuffer: + def test_slice_returns_the_requested_frames(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + assert _identities(buffer.slice(3, 7)) == [3, 4, 5, 6] + + def test_slice_is_relative_to_the_live_base(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + buffer.discard_before(4) + assert buffer.base_frame == 4 + assert _identities(buffer.slice(4, 8)) == [4, 5, 6, 7] + + def test_slicing_a_discarded_range_raises(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + buffer.discard_before(5) + with pytest.raises(RuntimeError, match="unavailable"): + buffer.slice(3, 6) + + def test_slicing_past_the_end_raises(self) -> None: + buffer = FrameBuffer() + for index in range(4): + buffer.append(_frame(index)) + with pytest.raises(RuntimeError, match="unavailable"): + buffer.slice(2, 9) + + def test_discard_is_clamped_at_both_ends(self) -> None: + buffer = FrameBuffer() + for index in range(5): + buffer.append(_frame(index)) + buffer.discard_before(-3) + assert buffer.base_frame == 0 + buffer.discard_before(99) + assert buffer.base_frame == 5 + assert buffer.end_frame == 5 + + def test_peak_tracks_the_high_water_mark(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + buffer.discard_before(8) + assert buffer.peak_frames == 10 + + +class TestWindowLifecycle: + """Drive the state exactly as the talker does, one frame per step.""" + + @staticmethod + def _run(total_frames: int) -> tuple[list[list[int]], int]: + """Generate ``total_frames`` frames, returning the emitted spans. + + Returns: + The frame identities of every emitted span, and the peak buffer + occupancy reached along the way. + """ + state = RequestARState(seed=1, max_frames=total_frames) + spans: list[list[int]] = [] + for index in range(total_frames): + state.frames.append(_frame(index)) + state.codes.append(torch.full((1, NUM_CODEBOOKS), index)) + state.generated_frames += 1 + window = state.ready_window() + if window is not None: + spans.append(_identities(state.take_window(window, is_final=False))) + + # Terminal flush: everything not yet emitted, as one contiguous span. + start = state.emitted_windows * AR_CHUNK_HOP_FRAMES + end = state.frames.end_frame + if end > start: + spans.append(_identities(state.frames.slice(start, end))) + return spans, state.frames.peak_frames + + def test_short_song_is_one_span(self) -> None: + spans, _ = self._run(50) + assert spans == [list(range(50))] + + def test_every_frame_is_delivered_exactly_once_in_order(self) -> None: + """Spans overlap by design, but their union must be the whole song. + + Reconstruct by taking each span's non-overlapping remainder, which is + what the acoustic stage does after cropping. + """ + for total in (250, 437, 600, 901): + spans, _ = self._run(total) + rebuilt: list[int] = [] + for span in spans: + rebuilt.extend(frame for frame in span if frame >= len(rebuilt)) + assert rebuilt == list(range(total)), f"gap or duplicate at {total} frames" + + def test_windows_advance_by_the_hop(self) -> None: + spans, _ = self._run(600) + for index, span in enumerate(spans[:-1]): + assert span[0] == index * AR_CHUNK_HOP_FRAMES + assert len(span) == AR_CHUNK_FRAMES + + def test_mid_stream_windows_are_never_the_final_frame(self) -> None: + """A window is held back until a further frame proves it is not last.""" + total = 600 + spans, _ = self._run(total) + for span in spans[:-1]: + assert span[-1] < total - 1 + + def test_terminal_span_covers_the_tail(self) -> None: + total = 437 + spans, _ = self._run(total) + assert spans[-1][-1] == total - 1 + + def test_terminal_span_is_bounded(self) -> None: + """The tail is at most one full window plus one hop.""" + for total in (250, 301, 437, 600, 901, 1000): + spans, _ = self._run(total) + assert len(spans[-1]) <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES + + def test_buffer_stays_bounded_over_a_long_song(self) -> None: + """Peak occupancy must not grow with song length; that is the point + of evicting retired frames.""" + _, peak_short = self._run(600) + _, peak_long = self._run(3000) + assert peak_long == peak_short + assert peak_long <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES + + def test_emitted_window_count_matches_the_window_plan(self) -> None: + """What the stream emitted must reconcile with chunk_windows().""" + for total in (250, 437, 600): + spans, _ = self._run(total) + planned = chunk_windows(total) + # Every window is accounted for: the pre-terminal spans are one + # window each, and the terminal span carries whatever remains. + assert len(spans) <= len(planned) + assert spans[0][0] == planned[0].start + + +class TestSamplingPositions: + def test_each_frame_reserves_one_slot_per_codebook(self) -> None: + """Codebook streams must never collide, or codes repeat across frames.""" + state = RequestARState(seed=1, max_frames=10) + positions = [state.advance_position() for _ in range(5)] + assert positions == [0, 8, 16, 24, 32] + + def test_draw_slots_are_globally_unique(self) -> None: + state = RequestARState(seed=1, max_frames=64) + slots: set[int] = set() + for _ in range(64): + base = state.advance_position() + for codebook in range(NUM_CODEBOOKS): + slot = base + codebook + assert slot not in slots + slots.add(slot) + assert len(slots) == 64 * NUM_CODEBOOKS + + +class TestCodeRetention: + def test_codes_span_matches_the_frame_span(self) -> None: + state = RequestARState(seed=1, max_frames=10) + for index in range(10): + state.codes.append(torch.full((1, NUM_CODEBOOKS), index)) + codes = state.take_codes(2, 6) + assert codes is not None + assert codes.shape == (4, NUM_CODEBOOKS) + assert [int(row[0]) for row in codes] == [2, 3, 4, 5] + + def test_out_of_range_span_is_dropped_not_raised(self) -> None: + """Codes are diagnostic; the audio path must not fail without them.""" + state = RequestARState(seed=1, max_frames=10) + for index in range(4): + state.codes.append(torch.full((1, NUM_CODEBOOKS), index)) + assert state.take_codes(2, 99) is None + assert state.take_codes(0, 0) is None diff --git a/tests/model_executor/models/test_minimax_music3_prompt.py b/tests/model_executor/models/test_minimax_music3_prompt.py new file mode 100644 index 000000000..8cd9e41b0 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_prompt.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic. + +The prompt transformations are not cosmetic: they decide token positions, and +token positions decide the audio. These tests pin the documented behaviours +that a well-meaning cleanup would otherwise "fix". +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.minimax_music3.chunking import ( + chunk_windows, + crop_sample_bounds, + overlap_mel_length, +) +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + BLEND_MEL_FRAMES, + DAV_HOP_SAMPLES, + HOP_MEL_FRAMES, +) +from vllm_omni.model_executor.models.minimax_music3.prompt import ( + SPECIAL_TOKEN_IDS, + build_cfg_null_token_ids, + build_prompt, + clean_caption, + normalize_lyrics, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestNormalizeLyrics: + def test_prepends_start_tag(self) -> None: + assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n") + + def test_lowercases_structure_tags(self) -> None: + assert "[verse]" in normalize_lyrics("[Verse]\nhello") + assert "[Verse]" not in normalize_lyrics("[Verse]\nhello") + + def test_text_after_a_leading_tag_is_dropped(self) -> None: + """Documented contract: a tag must sit on its own line. + + Lyrics written next to a tag are silently lost, which is why the + serving docs tell callers to put tags on their own line. + """ + out = normalize_lyrics("[Verse] Walking down the street") + assert "Walking" not in out + assert "[verse]" in out + + def test_text_on_its_own_line_survives(self) -> None: + out = normalize_lyrics("[Verse]\nWalking down the street") + assert "Walking down the street" in out + + def test_bracket_space_becomes_newline(self) -> None: + assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]" + + +class TestCleanCaption: + def test_special_token_form_is_rewritten_not_passed_through(self) -> None: + """A caption must not be able to inject a real special token.""" + assert clean_caption("<|bpm 92|>") == "bpm is 92" + assert "<|" not in clean_caption("<|bpm 92|>") + + def test_lone_tag_keeps_its_inner_text(self) -> None: + assert clean_caption("<|instrumental|>") == "instrumental" + + def test_markdown_is_stripped(self) -> None: + assert clean_caption("## Heading") == "Heading" + assert clean_caption("**bold**") == "bold" + assert clean_caption("- item") == "item" + + def test_blank_lines_collapse_at_two_not_three(self) -> None: + """The source collapses runs of two or more, not three or more. + + This shifts blank-line token positions and therefore the audio, so it + is pinned deliberately. + """ + assert clean_caption("a\n\nb") == "a\nb" + assert clean_caption("a\n\n\n\nb") == "a\nb" + + +class TestBuildPrompt: + def test_framing_tokens_are_present_and_ordered(self) -> None: + prompt = build_prompt("a bright pop song", "[Verse]\nhello") + assert prompt.startswith("<|im_start|><|caption_start|>") + assert prompt.endswith("<|im_end|><|audio_start|>") + assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>") + + +class TestBuildCfgNullTokenIds: + @staticmethod + def _prompt_ids(body_len: int) -> list[int]: + return ( + [SPECIAL_TOKEN_IDS["<|im_start|>"]] + + list(range(1000, 1000 + body_len)) + + [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]] + ) + + def test_length_is_preserved(self) -> None: + """Equal length is what lets the two rows decode in lockstep.""" + ids = self._prompt_ids(12) + assert len(build_cfg_null_token_ids(ids)) == len(ids) + + def test_body_is_replaced_and_framing_survives(self) -> None: + ids = self._prompt_ids(12) + null = build_cfg_null_token_ids(ids) + assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"] + assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"] + assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"] + assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]} + + def test_input_is_not_mutated(self) -> None: + ids = self._prompt_ids(6) + before = list(ids) + build_cfg_null_token_ids(ids) + assert ids == before + + def test_minimal_prompt_with_empty_body(self) -> None: + ids = [ + SPECIAL_TOKEN_IDS["<|im_start|>"], + SPECIAL_TOKEN_IDS["<|audio_cfg|>"], + SPECIAL_TOKEN_IDS["<|im_end|>"], + SPECIAL_TOKEN_IDS["<|audio_start|>"], + ] + assert build_cfg_null_token_ids(ids) == ids + + def test_too_short_is_rejected(self) -> None: + with pytest.raises(ValueError, match="too short"): + build_cfg_null_token_ids([1, 2, 3]) + + def test_wrong_framing_is_rejected(self) -> None: + with pytest.raises(ValueError, match="framing"): + build_cfg_null_token_ids([9, 9, 9, 9, 9]) + + +class TestChunkWindows: + def test_no_frames_yields_nothing(self) -> None: + assert chunk_windows(0) == [] + + def test_short_song_is_one_final_window(self) -> None: + windows = chunk_windows(50) + assert len(windows) == 1 + assert (windows[0].start, windows[0].end) == (0, 50) + assert windows[0].is_first and windows[0].is_last + + def test_exactly_one_window(self) -> None: + windows = chunk_windows(AR_CHUNK_FRAMES) + assert len(windows) == 1 + assert windows[0].is_last + + def test_windows_advance_by_the_hop(self) -> None: + windows = chunk_windows(450) + starts = [w.start for w in windows] + assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES)) + assert windows[-1].is_last + assert not any(w.is_last for w in windows[:-1]) + + def test_windows_cover_every_frame(self) -> None: + frames = 437 + windows = chunk_windows(frames) + covered: set[int] = set() + for window in windows: + covered.update(range(window.start, window.end)) + assert covered == set(range(frames)) + + def test_final_window_is_clamped_to_the_song(self) -> None: + windows = chunk_windows(437) + assert windows[-1].end == 437 + + def test_only_the_first_window_is_first(self) -> None: + windows = chunk_windows(600) + assert windows[0].is_first + assert not any(w.is_first for w in windows[1:]) + + +class TestCropBounds: + def test_single_window_song_is_not_cropped(self) -> None: + window = chunk_windows(50)[0] + assert crop_sample_bounds(window) == (0, 0) + + def test_first_window_keeps_its_head(self) -> None: + left, right = crop_sample_bounds(chunk_windows(600)[0]) + assert left == 0 + assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES + + def test_last_window_keeps_its_tail(self) -> None: + left, right = crop_sample_bounds(chunk_windows(600)[-1]) + assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES + assert right == 0 + + def test_middle_window_is_cropped_on_both_sides(self) -> None: + windows = chunk_windows(900) + middle = windows[len(windows) // 2] + left, right = crop_sample_bounds(middle) + assert left > 0 and right > 0 + + def test_overlap_is_half_a_window(self) -> None: + assert overlap_mel_length() == HOP_MEL_FRAMES // 2 diff --git a/tests/model_executor/models/test_minimax_music3_repo_root.py b/tests/model_executor/models/test_minimax_music3_repo_root.py new file mode 100644 index 000000000..182635ab7 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_repo_root.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Repository-root resolution for MiniMax Music 3's multi-component checkpoint. + +The acoustic stage declares no ``model_subdir`` because it reads the repo root +itself, so on a hub-id deployment its model path is ``MiniMaxAI/MiniMax-Music3`` +rather than a directory. Resolving that as a relative path silently points at +the server's working directory (issue #6638). +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.minimax_music3.weights import ( + _COMPONENT_WEIGHT_NAME, + _REQUIRED_COMPONENTS, + _ROOT_MARKERS, + resolve_repo_root, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_root(path): + # A real snapshot carries weights for every component a stage loads, not + # just the folders that identify the root. Building only ``_ROOT_MARKERS`` + # here would let a snapshot missing ``condition_encoder`` pass as complete, + # which is the gap this suite regresses. + for component in _REQUIRED_COMPONENTS: + folder = path / component + folder.mkdir(parents=True, exist_ok=True) + (folder / _COMPONENT_WEIGHT_NAME).write_bytes(b"x") + return path + + +def test_resolve_repo_root_accepts_the_root_itself(tmp_path): + root = _make_root(tmp_path / "snapshot") + assert resolve_repo_root(str(root)) == root + + +def test_resolve_repo_root_walks_up_from_a_model_subdir(tmp_path): + root = _make_root(tmp_path / "snapshot") + (root / "language_model").mkdir() + assert resolve_repo_root(str(root / "language_model")) == root + + +def test_resolve_repo_root_resolves_a_hub_id_to_a_local_snapshot(monkeypatch, tmp_path): + """A hub id must resolve through the cache, not against the working directory.""" + from huggingface_hub import HfApi + + root = _make_root(tmp_path / "snapshots" / "deadbeef") + seen = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + seen.append(kwargs) + return str(root) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == root + # Cache-first, and only the component folders. + assert seen[0]["local_files_only"] is True + assert "transformer/*" in seen[0]["allow_patterns"] + assert "qwen_7B/*" not in seen[0]["allow_patterns"] + + +def test_resolve_repo_root_falls_back_to_the_hub_when_the_cache_is_incomplete(monkeypatch, tmp_path): + from huggingface_hub import HfApi + + root = _make_root(tmp_path / "snapshots" / "deadbeef") + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if kwargs.get("local_files_only"): + raise OSError("incomplete snapshot") + return str(root) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == root + assert len(calls) == 2 + + +def test_resolve_repo_root_retries_online_when_the_local_lookup_returns_a_partial_root(monkeypatch, tmp_path): + """Hub clients differ on partial caches: some raise, some return the root. + + A returned partial root must fall through to the online call all the same, + or startup fails without ever consulting the Hub. + """ + from huggingface_hub import HfApi + + snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef" + (snapshot / "language_model").mkdir(parents=True) + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if not kwargs.get("local_files_only"): + _make_root(snapshot) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == snapshot + assert [bool(call.get("local_files_only")) for call in calls] == [True, False] + + +def test_resolve_repo_root_redownloads_weightless_marker_dirs(monkeypatch, tmp_path): + """Marker folders without their component weights are an interrupted download.""" + from huggingface_hub import HfApi + + snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef" + for marker in _ROOT_MARKERS: + (snapshot / marker).mkdir(parents=True) + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + if not kwargs.get("local_files_only"): + _make_root(snapshot) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == snapshot + assert any(not call.get("local_files_only") for call in calls) + + +def test_resolve_repo_root_reports_the_original_reference_when_unresolvable(monkeypatch): + from huggingface_hub import HfApi + + def fake_snapshot_download(self, repo_id, **kwargs): + raise OSError("offline") + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + with pytest.raises(FileNotFoundError, match="MiniMaxAI/MiniMax-Music3"): + resolve_repo_root("MiniMaxAI/MiniMax-Music3") + + +def test_resolve_repo_root_downloads_components_for_a_cold_cache_model_subdir(monkeypatch, tmp_path): + """Cold-cache stage 0: the talker gets an EXISTING ``language_model`` dir. + + Stage init pre-downloads only ``language_model/`` and ``tokenizer/``, so + the marker walk fails although the repo id is recoverable from the cache + layout and one snapshot_download away from working. + """ + from huggingface_hub import HfApi + + snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef" + (snapshot / "language_model").mkdir(parents=True) + (snapshot / "tokenizer").mkdir() + seen = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + seen.append(repo_id) + _make_root(snapshot) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root(str(snapshot / "language_model")) == snapshot + assert seen == ["MiniMaxAI/MiniMax-Music3"] + + +def test_resolve_repo_root_does_not_guess_a_repo_id_for_plain_directories(monkeypatch, tmp_path): + """A directory outside the HF cache has no recoverable repo id.""" + from huggingface_hub import HfApi + + def fake_snapshot_download(self, repo_id, **kwargs): + raise AssertionError("must not reach the Hub for a plain directory") + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + plain = tmp_path / "language_model" + plain.mkdir() + with pytest.raises(FileNotFoundError, match="expected sibling folders"): + resolve_repo_root(str(plain)) + + +def test_resolve_repo_root_rejects_a_snapshot_missing_condition_encoder(monkeypatch, tmp_path): + """``condition_encoder`` is loaded by the acoustic stage but is not a root marker. + + A cache holding only the marker folders satisfies the name probe while the + acoustic stage would still fail, so completeness must cover every required + component rather than the identification markers alone. + """ + from huggingface_hub import HfApi + + snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef" + for marker in _ROOT_MARKERS: + folder = snapshot / marker + folder.mkdir(parents=True) + (folder / _COMPONENT_WEIGHT_NAME).write_bytes(b"x") + repaired = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + if not kwargs.get("local_files_only"): + repaired.append(repo_id) + _make_root(snapshot) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root(str(snapshot)) == snapshot + assert repaired == ["MiniMaxAI/MiniMax-Music3"] + + +def test_resolve_repo_root_rejects_a_partial_shard_run(monkeypatch, tmp_path): + """One shard of a multi-shard component satisfies a bare glob but cannot load.""" + from huggingface_hub import HfApi + + stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors") + snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef" + for component in _REQUIRED_COMPONENTS: + (snapshot / component).mkdir(parents=True) + (snapshot / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x") + # Replace one component with an interrupted 1-of-3 shard run. + partial = snapshot / _ROOT_MARKERS[0] + (partial / _COMPONENT_WEIGHT_NAME).unlink() + (partial / f"{stem}-00001-of-00003.safetensors").write_bytes(b"x") + repaired = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + if not kwargs.get("local_files_only"): + repaired.append(repo_id) + for index in (2, 3): + (partial / f"{stem}-0000{index}-of-00003.safetensors").write_bytes(b"x") + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root(str(snapshot)) == snapshot + assert repaired == ["MiniMaxAI/MiniMax-Music3"] + + +def test_resolve_repo_root_threads_revision_and_download_dir(monkeypatch, tmp_path): + """A pinned revision must not be dropped when fetching the components. + + Resolving components against the default branch while the AR backbone is + pinned elsewhere would combine weights from two different commits. + """ + from huggingface_hub import HfApi + + snapshot = tmp_path / "snapshot" + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + _make_root(snapshot) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root("MiniMaxAI/MiniMax-Music3", revision="abc123", download_dir="/tmp/hub-cache") == snapshot + assert calls + assert all(call.get("revision") == "abc123" for call in calls) + assert all(call.get("cache_dir") == "/tmp/hub-cache" for call in calls) + + +def test_resolve_repo_root_reuses_the_snapshot_revision_from_the_cache_path(monkeypatch, tmp_path): + """The commit is recoverable from ``snapshots/`` and must be reused. + + Re-resolving the repo id alone would fetch the default branch rather than + the snapshot the caller is already sitting in. + """ + from huggingface_hub import HfApi + + snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "cafebabe" + (snapshot / "language_model").mkdir(parents=True) + calls = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + calls.append(kwargs) + _make_root(snapshot) + return str(snapshot) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root(str(snapshot / "language_model")) == snapshot + assert calls + assert all(call.get("revision") == "cafebabe" for call in calls) + + +def test_resolve_repo_root_accepts_shard_names_the_loader_can_read(monkeypatch, tmp_path): + """The completeness predicate must not be stricter than ``load_component_state``. + + The loader globs ``-*.safetensors`` and reads whatever it finds, so a + component whose weights do not use the ``--of-`` naming is + loadable. Judging it incomplete would force a needless re-download and, with + no network, fail on weights that are present and readable. + """ + from huggingface_hub import HfApi + + stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors") + root = tmp_path / "snapshot" + for component in _REQUIRED_COMPONENTS: + (root / component).mkdir(parents=True) + (root / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x") + unsharded = root / _ROOT_MARKERS[0] + (unsharded / _COMPONENT_WEIGHT_NAME).unlink() + (unsharded / f"{stem}-fp16.safetensors").write_bytes(b"x") + + def fake_snapshot_download(self, repo_id, **kwargs): + raise AssertionError("must not reach the Hub for readable weights") + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + assert resolve_repo_root(str(root)) == root + + +def test_resolve_repo_root_still_rejects_a_partial_run_beside_an_extra_file(monkeypatch, tmp_path): + """An unparsable extra must not mask a genuinely short ``-of-`` run.""" + from huggingface_hub import HfApi + + stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors") + root = tmp_path / "snapshot" + for component in _REQUIRED_COMPONENTS: + (root / component).mkdir(parents=True) + (root / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x") + partial = root / _ROOT_MARKERS[0] + (partial / _COMPONENT_WEIGHT_NAME).unlink() + (partial / f"{stem}-00001-of-00003.safetensors").write_bytes(b"x") + (partial / f"{stem}-extra.safetensors").write_bytes(b"x") + repaired = [] + + def fake_snapshot_download(self, repo_id, **kwargs): + repaired.append(repo_id) + for index in (2, 3): + (partial / f"{stem}-0000{index}-of-00003.safetensors").write_bytes(b"x") + return str(root) + + monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download) + + with pytest.raises(FileNotFoundError, match="components are incomplete"): + resolve_repo_root(str(root)) + assert repaired == [] diff --git a/tests/model_executor/models/test_minimax_music3_row_binding.py b/tests/model_executor/models/test_minimax_music3_row_binding.py new file mode 100644 index 000000000..94ff30199 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_row_binding.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MiniMax Music 3 row-to-request binding. + +vLLM compacts and reorders the persistent batch whenever a request finishes or +is preempted, so a batch row number is only meaningful within the step that +produced it. Every failure pinned here was a real defect: state keyed by row +survived into a step where the row belonged to someone else, or state keyed by +one spelling of the request id was released under another and leaked. + +These exercise the pure helpers and the binding arithmetic directly, so they +need no GPU and no engine. +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.minimax_music3.constants import ( + DEFAULT_MAX_AUDIO_FRAMES, +) +from vllm_omni.model_executor.models.minimax_music3.talker import ( + _CFG_UNCOND_SUFFIX, + _first_int, + _pair_id_from_request, + _request_key, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestRequestKey: + def test_engine_list_form_is_unwrapped(self) -> None: + """The engine injects global_request_id as a single-element list. + + Keying state on the list's ``str()`` would produce "['req-1']", which + never matches the bare id that ``on_requests_finished`` delivers, so + every finished song would leak its frame buffer. + """ + assert _request_key(["req-1"]) == "req-1" + + def test_bare_string_is_unchanged(self) -> None: + assert _request_key("req-1") == "req-1" + + def test_list_and_bare_forms_agree(self) -> None: + """Allocation and release must land on the same key.""" + assert _request_key(["req-1"]) == _request_key("req-1") + + def test_empty_and_missing_are_stable(self) -> None: + assert _request_key(None) == "" + assert _request_key([]) == "" + + def test_tuple_form(self) -> None: + assert _request_key(("req-9",)) == "req-9" + + +class TestPairIdDerivation: + def test_conditioned_request_is_its_own_pair(self) -> None: + assert _pair_id_from_request("req-1") == "req-1" + + def test_companion_maps_to_its_partner(self) -> None: + assert _pair_id_from_request(f"req-1{_CFG_UNCOND_SUFFIX}") == "req-1" + + def test_pair_members_agree_on_the_pair_id(self) -> None: + cond = "req-42" + uncond = f"req-42{_CFG_UNCOND_SUFFIX}" + assert _pair_id_from_request(cond) == _pair_id_from_request(uncond) + + def test_a_request_ending_in_a_similar_string_is_not_confused(self) -> None: + assert _pair_id_from_request("req-cfg_uncond") == "req-cfg_uncond" + + +class TestFirstInt: + def test_serving_layer_list_wrapping(self) -> None: + """tts_params values arrive wrapped in a single-element list.""" + assert _first_int([750], DEFAULT_MAX_AUDIO_FRAMES) == 750 + + def test_bare_value(self) -> None: + assert _first_int(750, DEFAULT_MAX_AUDIO_FRAMES) == 750 + + def test_missing_falls_back(self) -> None: + assert _first_int(None, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + assert _first_int([], DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + + def test_nonsense_falls_back_rather_than_raising(self) -> None: + assert _first_int("abc", DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + assert _first_int(0, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + assert _first_int(-5, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + + +def _feedback_plan( + request_ids: list[str], + computed: list[int], + scheduled: list[int], + prompt_tokens: dict[str, int], + has_codes: set[str], +) -> list[int]: + """Mirror of the talker's feedback row selection. + + Returns the flat token offsets whose embedding gets replaced by an audio + frame. Kept as a standalone function so the arithmetic can be checked + without constructing a model. + """ + offsets: list[int] = [] + running = 0 + for count in scheduled: + offsets.append(running) + running += count + + rows: list[int] = [] + for row, request_id in enumerate(request_ids): + if request_id not in has_codes: + continue + if scheduled[row] != 1 or computed[row] < prompt_tokens.get(request_id, 0): + continue + rows.append(offsets[row]) + return rows + + +class TestFeedbackRowSelection: + """The guard that decides which token embeddings become audio frames.""" + + def test_decode_only_step_replaces_every_row(self) -> None: + ids = ["a", "a__cfg_uncond", "b", "b__cfg_uncond"] + rows = _feedback_plan( + ids, + computed=[120, 120, 90, 90], + scheduled=[1, 1, 1, 1], + prompt_tokens={i: 50 for i in ids}, + has_codes=set(ids), + ) + assert rows == [0, 1, 2, 3] + + def test_prefill_row_is_never_overwritten(self) -> None: + """A prompt token must keep its real embedding. + + The earlier guard compared the flat token count against a row budget, + so a short prompt slipped through and had its first embeddings + replaced with a previous request's audio frame. + """ + rows = _feedback_plan( + ["a", "b"], + computed=[0, 300], + scheduled=[8, 1], + prompt_tokens={"a": 8, "b": 40}, + has_codes={"b"}, + ) + # Row 0 occupies flat tokens 0..7 and must be untouched; row 1's single + # decode token sits at flat offset 8. + assert rows == [8] + + def test_short_prefill_below_the_row_budget_is_still_safe(self) -> None: + """The exact shape that defeated the old token-count guard.""" + rows = _feedback_plan( + ["a"], + computed=[0], + scheduled=[10], + prompt_tokens={"a": 10}, + has_codes={"a"}, + ) + assert rows == [] + + def test_chunked_prefill_midway_is_not_fed_back(self) -> None: + rows = _feedback_plan( + ["a"], + computed=[2048], + scheduled=[2048], + prompt_tokens={"a": 5000}, + has_codes={"a"}, + ) + assert rows == [] + + def test_final_prefill_chunk_is_not_fed_back(self) -> None: + """The step that completes a prompt still carries prompt tokens.""" + rows = _feedback_plan( + ["a"], + computed=[4096], + scheduled=[904], + prompt_tokens={"a": 5000}, + has_codes={"a"}, + ) + assert rows == [] + + def test_mixed_batch_offsets_follow_the_flat_layout(self) -> None: + """Tokens are a flat concatenation, so offsets are a running sum.""" + ids = ["p", "d1", "d2"] + rows = _feedback_plan( + ids, + computed=[0, 500, 700], + scheduled=[64, 1, 1], + prompt_tokens={"p": 64, "d1": 40, "d2": 40}, + has_codes={"d1", "d2"}, + ) + assert rows == [64, 65] + + def test_request_without_codes_yet_is_skipped(self) -> None: + """The first decode step after prefill has no previous frame.""" + rows = _feedback_plan( + ["a"], + computed=[50], + scheduled=[1], + prompt_tokens={"a": 50}, + has_codes=set(), + ) + assert rows == [] + + def test_reordered_batch_binds_by_request_not_position(self) -> None: + """The same two requests in swapped rows get swapped offsets. + + This is the compaction case: holding codes in a row-indexed buffer + would have handed each request the other's frame. + """ + prompt_tokens = {"a": 50, "b": 50} + first = _feedback_plan(["a", "b"], [100, 200], [1, 1], prompt_tokens, {"a"}) + swapped = _feedback_plan(["b", "a"], [200, 100], [1, 1], prompt_tokens, {"a"}) + assert first == [0] + assert swapped == [1] diff --git a/tests/model_executor/models/test_nemotron_voicechat_registration.py b/tests/model_executor/models/test_nemotron_voicechat_registration.py index 09d85b271..935ffe96d 100644 --- a/tests/model_executor/models/test_nemotron_voicechat_registration.py +++ b/tests/model_executor/models/test_nemotron_voicechat_registration.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights).""" import subprocess @@ -59,7 +59,6 @@ def test_pipeline_registered_with_three_stages() -> None: assert code2wav.execution_type == StageExecutionType.LLM_GENERATION assert code2wav.final_output and code2wav.final_output_type == "audio" assert code2wav.input_sources == (1,) - assert pipeline.validate() == [] def test_architectures_registered() -> None: @@ -67,15 +66,6 @@ def test_architectures_registered() -> None: assert arch in _OMNI_MODELS, f"{arch} missing from _OMNI_MODELS" -def test_full_payload_whitelists_cover_code2wav_only() -> None: - from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES - - assert ("NemotronVoiceChatCode2Wav", "code2wav") in _FULL_PAYLOAD_INPUT_STAGES - # The thinker->talker hop is token-path only: the talker must NOT wait on a - # connector payload or stage 1 hangs. - assert ("NemotronVoiceChatTalkerForConditionalGeneration", "talker") not in _FULL_PAYLOAD_INPUT_STAGES - - def test_config_parses_nemo_layout_without_hub_access(monkeypatch) -> None: # Provide the thinker text config explicitly so the test never touches HF. cfg = NemotronVoiceChatConfig( diff --git a/tests/model_executor/models/test_nemotron_voicechat_shapes.py b/tests/model_executor/models/test_nemotron_voicechat_shapes.py index 8bd279bb8..9b3c01559 100644 --- a/tests/model_executor/models/test_nemotron_voicechat_shapes.py +++ b/tests/model_executor/models/test_nemotron_voicechat_shapes.py @@ -225,3 +225,39 @@ def test_torchaudio_mel_filterbank_matches_librosa() -> None: reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney") assert ta.shape == reference.shape assert abs(ta - reference).max() < 1e-6 + + +def test_code2wav_reuses_and_clears_per_request_streaming_cache() -> None: + import torch + from torch import nn + + from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import ( + NemotronVoiceChatCode2Wav, + ) + + class Codec(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.caches = [] + + def decode(self, codes, lengths, *, cache): + self.caches.append(cache) + wav = torch.zeros((1, int(codes.shape[1]) * 1764), device=codes.device) + return wav, lengths * 1764 + + model = NemotronVoiceChatCode2Wav.__new__(NemotronVoiceChatCode2Wav) + nn.Module.__init__(model) + model._sample_rate, model._num_quantizers, model._codebook_size, model._wav_per_frame = 22050, 31, 1024, 1764 + model._duplex_codec_caches = {} + model.audio_codec = Codec() + + def info(codes): + return [{"meta": {"codec_streaming": True, "request_id": "req"}, "codes": {"audio": codes}}] + + model.forward(runtime_additional_information=info(torch.zeros((1, 31), dtype=torch.long))) + model.forward(runtime_additional_information=info(torch.zeros((2, 31), dtype=torch.long))) + assert model.audio_codec.caches[0] is model.audio_codec.caches[1] + assert "req" in model._duplex_codec_caches + model.on_requests_finished({"req"}) + assert "req" not in model._duplex_codec_caches diff --git a/tests/model_executor/models/voxcpm2/test_talker_output_marker.py b/tests/model_executor/models/voxcpm2/test_talker_output_marker.py new file mode 100644 index 000000000..1a76525fe --- /dev/null +++ b/tests/model_executor/models/voxcpm2/test_talker_output_marker.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""VoxCPM2 dense-mode audio outputs must carry the sparse-alignment marker. + +Dense mode (`_uses_sparse_audio_outputs()` False) still yields a strict +SUBSET of the batch whenever some requests emit no audio in a step (e.g. +prefill phase). The runner routes per-request `model_outputs` lists by batch +position unless `meta.sparse_audio` declares `meta.req_id` alignment — so a +subset without the marker misroutes audio across requests (RFC #5450 C3 +review follow-up). +""" + +from __future__ import annotations + +import functools + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +torch = pytest.importorskip("torch") + + +@functools.lru_cache(maxsize=1) +def _talker_cls(): + """Defer talker import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + VoxCPM2TalkerForConditionalGeneration, + ) + + return VoxCPM2TalkerForConditionalGeneration + + +def _make_dense_talker(): + cls = _talker_cls() + talker = cls.__new__(cls) + # Dense mode: emit every step, no delayed copy, no batched VAE decode. + talker._audio_emit_every = 1 + talker._vae_decode_every = 1 + talker._enable_delayed_audio_copy = False + talker._coalesce_audio_d2h = False + talker._sample_rate = 16000 + talker._audio_queue = [] + return talker + + +def test_dense_subset_audio_carries_sparse_alignment_marker(): + talker = _make_dense_talker() + # Batch had [r0, r1, r2]; only r1/r2 produced audio this step. + talker._audio_queue = [("r1", torch.zeros(4)), ("r2", torch.ones(4))] + + out = talker.make_omni_output(torch.zeros(1)) + + mm = out.multimodal_outputs + assert mm["meta"] == {"req_id": ["r1", "r2"], "sparse_audio": ["1"]} + assert len(mm["model_outputs"]) == 2 + assert len(mm["sr"]) == 2 + # (The coalesce-D2H sibling branch carries the same marker; it needs + # device-resident chunks, so its marker contract is pinned here via the + # shared dense-subset shape rather than a GPU-only test.) + + +def test_dense_full_batch_audio_still_carries_marker_and_order(): + # Full coverage is the degenerate subset: the marker stays truthful and + # `meta.req_id` order matches the emitted list order. + talker = _make_dense_talker() + talker._audio_queue = [("r0", torch.zeros(2)), ("r1", torch.ones(2))] + + out = talker.make_omni_output(torch.zeros(1)) + + mm = out.multimodal_outputs + assert mm["meta"]["req_id"] == ["r0", "r1"] + assert mm["meta"]["sparse_audio"] == ["1"] diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py index 8ff67b5ff..c7686ca31 100644 --- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py +++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for CUDA Graph acoustic transformer wrapper numerical equivalence. @@ -18,11 +18,12 @@ import pytest import torch import torch.nn as nn +from tests.helpers.mark import hardware_marks + pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), pytest.mark.core_model, - pytest.mark.cuda, - pytest.mark.L4, + *hardware_marks(res={"cuda": "L4"}), ] DEVICE = torch.device("cuda:0") diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py index d5971efe7..08090d954 100644 --- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py +++ b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections import defaultdict from types import SimpleNamespace @@ -7,6 +7,7 @@ from types import SimpleNamespace import pytest import torch +from vllm_omni.data_entry_keys import serialize_payload from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import ( talker2code2wav_async_chunk, text2flow_full_payload, @@ -153,6 +154,36 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset(): assert payload.embed.embedding is not None +def test_talker2code2wav_async_chunk_decodes_bfloat16_conditioning(): + speech_feat = torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.bfloat16) + additional_information = serialize_payload( + { + "embed": { + "speech_token": torch.tensor([[11]], dtype=torch.long), + "speech_feat": speech_feat, + "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.bfloat16), + } + } + ) + request = SimpleNamespace( + external_req_id="rid-bfloat16", + output_token_ids=[1], + additional_information=additional_information, + is_finished=lambda: True, + ) + + payload = talker2code2wav_async_chunk( + transfer_manager=_transfer_manager(), + multimodal_output=None, + request=request, + is_finished=True, + ) + + assert payload is not None + assert payload.embed.speech_feat.dtype == torch.bfloat16 + assert torch.equal(payload.embed.speech_feat, speech_feat) + + def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes(): transfer_manager = _transfer_manager(chunk_frames=25) request = SimpleNamespace( diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py index f0da872b9..afa7b46c4 100644 --- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections import defaultdict from types import SimpleNamespace @@ -72,6 +72,16 @@ def _codes(payload) -> list[int]: return payload.codes.audio.tolist() +def test_empty_full_payload_releases_consumer_wait_gate() -> None: + payload = tts2code2wav_full_payload(_manager(), None, _request("req")) + + assert _codes(payload) == [] + assert payload.meta.code_flat_numel == 0 + assert payload.meta.left_context_size == 0 + assert payload.meta.last_chunk is True + assert payload.meta.finished.item() is True + + @pytest.mark.parametrize(("count", "emitted"), [(24, False), (25, True), (26, True)]) def test_first_chunk_threshold_is_25_generated_codes(count: int, emitted: bool) -> None: manager = _manager() @@ -158,10 +168,14 @@ def test_duplex_turn_end_waits_for_terminal_codec_flush() -> None: def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None: manager = _manager() request = _request("req") - request.additional_information = { + request.model_intermediate_buffer = { "codes": {"ref": [0.1, -0.1]}, "meta": {"ref_audio_sr": 16000}, } + request.additional_information = { + "codes": {"ref": [0.9]}, + "meta": {"ref_audio_sr": 8000}, + } payload = tts2code2wav_async_chunk( manager, @@ -183,10 +197,33 @@ def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None: assert payload.meta.turn_end is True +def test_first_chunk_falls_back_to_legacy_reference_fields() -> None: + manager = _manager() + request = _request("req") + request.model_intermediate_buffer = { + "meta": {"ref_audio_sr": 16000}, + } + request.additional_information = { + "codes": {"ref": [0.9]}, + "meta": {"ref_audio_sr": 8000}, + } + + payload = tts2code2wav_async_chunk( + manager, + _duplex_delta(*range(7), turn_end=True), + request, + True, + ) + + assert payload is not None + assert payload.codes.ref.tolist() == pytest.approx([0.9]) + assert payload.meta.ref_audio_sr == 16000 + + def test_full_payload_forwards_all_codes_and_request_metadata() -> None: manager = _manager() request = _request("req") - request.additional_information = { + request.model_intermediate_buffer = { "codes": {"ref": [0.1, -0.1]}, "meta": { "ref_audio_sr": 16000, @@ -423,6 +460,12 @@ def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None: request, True, ) + duplicate_boundary = tts2code2wav_async_chunk( + manager, + _duplex_delta(turn_id=7, turn_end=True), + request, + True, + ) next_turn = tts2code2wav_async_chunk( manager, _duplex_delta(20, turn_id=8), @@ -435,6 +478,10 @@ def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None: assert turn_end.meta.last_chunk is True assert turn_end.meta.turn_end is True assert turn_end.meta.is_segment_finished.item() is True + assert turn_end.meta.replace_runtime_additional_information is True + assert duplicate_boundary is not None + assert duplicate_boundary.meta.is_segment_finished.item() is True + assert duplicate_boundary.meta.replace_runtime_additional_information is True assert next_turn.meta.cache_epoch == turn_end.meta.cache_epoch + 1 assert next_turn.meta.chunk_seq == 0 assert next_turn.meta.last_chunk is False diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py index 652215a08..795047f1c 100644 --- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace import pytest @@ -131,6 +134,7 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None: assert converted["prompt_token_ids"] == [0, 0, 0] assert info["meta"]["replace_streaming_prompt"] is True assert info["meta"]["next_stage_prompt_len"] == 3 + assert info["meta"]["next_stage_generation_tokens"] == 26 assert info["meta"]["turn_start"] is True assert info["meta"]["segment_end"] is True assert info["duplex"]["epoch"] == 3 @@ -203,6 +207,7 @@ def test_native_duplex_continuation_appends_only_new_talker_condition() -> None: assert "replace_streaming_prompt" not in second_input["model_intermediate_buffer"]["meta"] assert third_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True assert second_input["model_intermediate_buffer"]["meta"]["next_stage_prompt_len"] == 3 + assert second_input["model_intermediate_buffer"]["meta"]["next_stage_generation_tokens"] == 26 assert second_input["prompt_token_ids"] == [0, 0, 0] diff --git a/tests/model_executor/stage_input_processors/test_minimax_h3.py b/tests/model_executor/stage_input_processors/test_minimax_h3.py new file mode 100644 index 000000000..599be81f2 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_minimax_h3.py @@ -0,0 +1,387 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Regression tests for MiniMax H3's disaggregated text-encoder contract.""" + +import os +import shutil +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from PIL import Image + +from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _load_audio, + resolve_minimax_h3_diffusion_model_path, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( + resolve_minimax_h3_model_root, + resolve_minimax_h3_partition, +) +from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MINIMAX_H3_CONDITION_LABELS_KEY, + MINIMAX_H3_PRESENTATION_TASK_KEY, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + minimax_h3_ref2va_presentation, + minimax_h3_ref2va_video_presentation, +) +from vllm_omni.model_executor.models.minimax_h3.reference_video import ( + MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, + deserialize_prepared_reference_videos, +) +from vllm_omni.model_executor.models.minimax_h3.text_encoder import ( + MiniMaxH3MultiModalProcessor, + _build_minimax_h3_presentation, +) +from vllm_omni.model_executor.stage_input_processors.minimax_h3 import ( + _audio_items, + _diffusion_sampling_params, + prepare_text_encoder_prompt, + text_encoder2diffusion, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _SegmentTokenizer: + _special_ids = { + "<|vision_start|>": 1, + "<|vision_end|>": 2, + "<|image_pad|>": 3, + "<|video_pad|>": 4, + } + + def __call__(self, text, *, add_special_tokens): + assert add_special_tokens is False + return {"input_ids": [100 + len(text), 1000 + sum(text.encode())]} + + def convert_tokens_to_ids(self, token): + return self._special_ids[token] + + +def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(mocker): + processor = object.__new__(MiniMaxH3MultiModalProcessor) + sentinel = ([1, 2, 3], object(), True) + apply_processor = mocker.patch.object( + MiniMaxH3MultiModalProcessor, + "_apply_hf_processor", + return_value=sentinel, + ) + inputs = object() + timing_ctx = object() + + result = processor._cached_apply_hf_processor(inputs, timing_ctx) + + assert result is sentinel + apply_processor.assert_called_once_with(inputs, timing_ctx) + + +@pytest.mark.parametrize( + ("value", "expected_count"), + [ + ((torch.zeros(16), 16_000), 1), + ([np.zeros(16), 16_000], 1), + ([(torch.zeros(16), 16_000), (torch.ones(16), 24_000)], 2), + (["first.wav", "second.wav"], 2), + ], +) +def test_audio_items_preserves_waveform_pairs(value, expected_count): + assert len(_audio_items(value)) == expected_count + + +def test_h3_selects_the_single_explicit_diffusion_stage_params() -> None: + stage_zero = SimpleNamespace(extra_args={"task": "t2va"}) + diffusion = OmniDiffusionSamplingParams(extra_args={"task": "ref2va"}) + + assert _diffusion_sampling_params([stage_zero, diffusion]) is diffusion + + +def test_h3_rejects_missing_or_ambiguous_diffusion_stage_params() -> None: + with pytest.raises(RuntimeError, match="exactly one OmniDiffusionSamplingParams"): + _diffusion_sampling_params([SimpleNamespace(extra_args={"task": "t2va"})]) + with pytest.raises(RuntimeError, match="got 2"): + _diffusion_sampling_params([OmniDiffusionSamplingParams(), OmniDiffusionSamplingParams()]) + + +def test_fused_audio_loader_accepts_list_waveform_pair(): + waveform, sample_rate = _load_audio([[0.0, 0.5, -0.5], 16_000]) + assert sample_rate == 16_000 + torch.testing.assert_close(waveform, torch.tensor([0.0, 0.5, -0.5])) + + +def test_prepare_ref2va_keeps_original_text_and_exact_condition_order(): + prompt = { + "prompt": "hello", + "multi_modal_data": { + "image": Image.new("RGB", (256, 256)), + "audio": [np.zeros(16), 16_000], + }, + } + sampling = OmniDiffusionSamplingParams( + height=256, + width=448, + extra_args={"task": "ref2va"}, + ) + + transformed = prepare_text_encoder_prompt(prompt, [sampling]) + + assert transformed["prompt"] == "hello" + assert len(transformed["multi_modal_data"]["image"]) == 1 + assert "audio" not in transformed["multi_modal_data"] + assert transformed["mm_processor_kwargs"][MINIMAX_H3_PRESENTATION_TASK_KEY] == "ref2va" + assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [ + ("image", 1), + ("audio", 1), + ] + + +def test_text_encoder_prompt_rejects_injected_prepared_video_descriptor(): + prompt = { + "prompt": "hello", + "additional_information": {"meta": {MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY: '{"artifact_dir":"/tmp/user"}'}}, + } + sampling = OmniDiffusionSamplingParams(height=256, width=448, extra_args={"task": "t2va"}) + + transformed = prepare_text_encoder_prompt(prompt, [sampling]) + + assert MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY not in transformed["additional_information"]["meta"] + + +def test_prepare_ref2va_video_uses_shared_frame_sampler_once(monkeypatch): + prepare_videos = Mock() + + def prepare(_videos, *, target_frame_count, workdir, start_time_seconds): + prepare_videos(_videos, target_frame_count, workdir, start_time_seconds) + prepared_path = os.path.join(workdir, "prepared.mp4") + open(prepared_path, "wb").close() + return [ + { + "original_path": "/tmp/input.mp4", + "prepared_path": prepared_path, + "input_has_audio": True, + "width": 448, + "height": 256, + "start_time_seconds": 0.0, + "duration_seconds": 5.2, + "audio_duration_seconds": 5.2, + } + ] + + sample_frames = Mock(return_value={"frames": [np.zeros((4, 4, 3), dtype=np.uint8)]}) + monkeypatch.setattr("vllm_omni.model_executor.stage_input_processors.minimax_h3.prepare_reference_videos", prepare) + monkeypatch.setattr( + "vllm_omni.model_executor.stage_input_processors.minimax_h3.sample_reference_video_frames", + sample_frames, + ) + prompt = { + "prompt": "hello", + "multi_modal_data": {"video": "/tmp/input.mp4"}, + } + sampling = OmniDiffusionSamplingParams(height=256, width=448, num_frames=124, extra_args={"task": "ref2va"}) + + transformed = prepare_text_encoder_prompt(prompt, [sampling]) + + prepare_videos.assert_called_once() + _, target_frame_count, artifact_dir, start_time_seconds = prepare_videos.call_args.args + assert target_frame_count == 124 + assert start_time_seconds is None + descriptor = transformed["additional_information"]["meta"][MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY] + described_dir, described_videos = deserialize_prepared_reference_videos(descriptor) + assert described_dir == artifact_dir + assert os.path.isfile(described_videos[0]["prepared_path"]) + sample_frames.assert_called_once_with(described_videos[0]["prepared_path"]) + assert prompt["multi_modal_data"] == {"video": "/tmp/input.mp4"} + assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [ + ("audio", 1), + ("video", 1), + ] + shutil.rmtree(artifact_dir) + + +def test_stage_wire_rejects_multiple_text_encoder_sources(): + with pytest.raises(RuntimeError, match="exactly one text-encoder source"): + text_encoder2diffusion([SimpleNamespace(), SimpleNamespace()], prompt={"prompt": "hello"}) + + +def test_stage_wire_rejects_request_id_mismatch(): + source = SimpleNamespace(request_id="other", outputs=[SimpleNamespace()]) + prompt = { + "prompt": "hello", + "additional_information": {"global_request_id": ["req-1"]}, + } + + with pytest.raises(RuntimeError, match="request ID does not match"): + text_encoder2diffusion([source], prompt=prompt) + + +def test_stage_wire_rejects_multiple_completions(): + source = SimpleNamespace(request_id="req-1", outputs=[SimpleNamespace(), SimpleNamespace()]) + + with pytest.raises(RuntimeError, match="exactly one completion"): + text_encoder2diffusion([source], prompt={"prompt": "hello"}) + + +def _source_output(payload: dict) -> SimpleNamespace: + completion = SimpleNamespace(multimodal_output=payload) + return SimpleNamespace(request_id="request-1", outputs=[completion]) + + +def _diffusion_prompt() -> dict: + return { + "request_id": "request-1", + "prompt": "test prompt", + "additional_information": {}, + } + + +def test_text_encoder2diffusion_reads_hidden_states_output_and_token_role_ids() -> None: + hidden = torch.randn(4, 5120) + token_role_ids = torch.tensor([[1], [1], [0], [0]]) + + result = text_encoder2diffusion( + [ + _source_output( + { + "hidden_states": {"output": hidden}, + "meta": {"token_role_ids": token_role_ids}, + } + ) + ], + _diffusion_prompt(), + ) + + conditioning = result["additional_information"]["text_encoder_output"] + torch.testing.assert_close(conditioning["hidden_states"], hidden) + torch.testing.assert_close(conditioning["token_tags"], token_role_ids.squeeze(-1)) + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ({"meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states payload"), + ({"hidden_states": {}, "meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states.output tensor"), + ({"hidden_states": {"output": torch.randn(4, 5120)}}, "no conditioning metadata"), + ({"hidden_states": {"output": torch.randn(4, 5120)}, "meta": {}}, "no token_role_ids tensor"), + ( + { + "hidden_states": {"output": torch.randn(4, 5120)}, + "meta": {"token_role_ids": torch.ones(4)}, + }, + r"must have shape \[tokens, 1\]", + ), + ], +) +def test_text_encoder2diffusion_rejects_invalid_conditioning_payload(payload: dict, message: str) -> None: + with pytest.raises(RuntimeError, match=message): + text_encoder2diffusion([_source_output(payload)], _diffusion_prompt()) + + +def test_ref2va_one_image_tokens_and_tags_match_fused_presentation(): + tokenizer = _SegmentTokenizer() + labels = [("image", 1), ("audio", 1)] + image_grid = torch.tensor([[1, 4, 4]]) + + actual = _build_minimax_h3_presentation( + tokenizer, + prompt="hello", + task="ref2va", + condition_labels=labels, + image_grid_thw=image_grid, + video_grid_thw=None, + video_timestamps=None, + merge_size=2, + ) + expected = minimax_h3_ref2va_presentation( + tokenizer, + prompt="hello", + condition_labels=labels, + image_token_count=[4], + ) + + torch.testing.assert_close(actual[0], expected[0]) + torch.testing.assert_close(actual[1], expected[1]) + + +def test_ref2va_video_tokens_and_tags_match_fused_without_outer_markers(): + tokenizer = _SegmentTokenizer() + labels = [("audio", 1), ("video", 1)] + video_grid = torch.tensor([[2, 4, 4]]) + timestamps = [[0.2, 0.4]] + + actual = _build_minimax_h3_presentation( + tokenizer, + prompt="hello", + task="ref2va", + condition_labels=labels, + image_grid_thw=None, + video_grid_thw=video_grid, + video_timestamps=timestamps, + merge_size=2, + ) + expected = minimax_h3_ref2va_video_presentation( + tokenizer, + prompt="hello", + condition_labels=labels, + image_token_count=None, + video_block_token_counts=[[4, 4]], + video_block_timestamps=timestamps, + ) + + torch.testing.assert_close(actual[0], expected[0]) + torch.testing.assert_close(actual[1], expected[1]) + assert int((actual[0] == tokenizer._special_ids["<|vision_start|>"]).sum()) == 2 + assert int((actual[0] == tokenizer._special_ids["<|vision_end|>"]).sum()) == 2 + + +def test_checkpoint_resolver_selects_local_partition(tmp_path): + root = tmp_path / "MiniMax-H3" + (root / "FL2VA" / "text_encoder").mkdir(parents=True) + (root / "Ref2VA" / "text_encoder").mkdir(parents=True) + + assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder") + assert resolve_minimax_h3_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder") + assert resolve_minimax_h3_model_root(str(root / "Ref2VA"), None, None) == str(root / "Ref2VA" / "text_encoder") + assert resolve_minimax_h3_model_root(str(root / "FL2VA"), None, "ref2va") == str(root / "Ref2VA" / "text_encoder") + + +def test_checkpoint_resolver_rejects_unknown_task(tmp_path): + with pytest.raises(ValueError, match="task_type must be one of"): + resolve_minimax_h3_model_root(str(tmp_path), None, "unknown") + + +def test_partition_resolver_preserves_consumer_auto_default(tmp_path): + root = tmp_path / "MiniMax-H3" + ref2va = root / "Ref2VA" + ref2va.mkdir(parents=True) + + assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="fl2va") == "fl2va" + assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="combined") == "combined" + assert resolve_minimax_h3_partition(str(ref2va), "auto", auto_partition="combined") == "ref2va" + + +def test_diffusion_resolver_selects_startup_partition(tmp_path): + root = tmp_path / "MiniMax-H3" + fl2va = root / "FL2VA" + ref2va = root / "Ref2VA" + fl2va.mkdir(parents=True) + ref2va.mkdir() + (fl2va / "model_index.json").write_text("{}") + (ref2va / "model_index.json").write_text("{}") + + assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(fl2va) + assert resolve_minimax_h3_diffusion_model_path(str(root), None, "ref2va") == str(ref2va) + assert resolve_minimax_h3_diffusion_model_path(str(root), None, None) == str(fl2va) + assert resolve_minimax_h3_diffusion_model_path(str(root), None, "combined") == str(root) + assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, None) == str(ref2va) + + +def test_diffusion_resolver_normalizes_partial_partition_directory(tmp_path): + root = tmp_path / "MiniMax-H3" + ref2va = root / "Ref2VA" + (ref2va / "text_encoder").mkdir(parents=True) + + assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, "ref2va") == str(ref2va) diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py index dcdff5bbb..172faf9fb 100644 --- a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py +++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace @@ -285,9 +285,46 @@ def test_talker2code2wav_async_chunk_requires_prompt_len() -> None: talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req) -# ========================= -# Talker drain-per-wake (async scheduling review P1/P2) -# ========================= +def test_talker2code2wav_duplex_keeps_cumulative_history_across_wakes() -> None: + from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import ( + talker2code2wav_async_chunk, + ) + + tm = _fake_transfer_manager(codec_chunk_frames=2) + prompt_len = 3 + req = _streaming_request( + prompt_len=prompt_len, + generated=[], + info={"meta": {"codec_streaming": torch.tensor([True])}, "nvc_logical_prompt_len": prompt_len}, + ) + prompt_rows = torch.zeros((prompt_len - 1, 31), dtype=torch.long) + first = torch.full((1, 31), 101, dtype=torch.long) + second = torch.full((1, 31), 202, dtype=torch.long) + + for rows, finished in ((prompt_rows, False), (torch.cat([prompt_rows, first]), False)): + payload = talker2code2wav_async_chunk( + tm, + {"codes": {"audio": rows}, "meta": {"codec_streaming": torch.tensor([True])}}, + req, + is_finished=finished, + ) + assert payload is not None and not bool(payload.meta.finished) + + payload = talker2code2wav_async_chunk( + tm, + { + "codes": {"audio": torch.cat([prompt_rows, first, second])}, + "meta": {"codec_streaming": torch.tensor([True])}, + }, + req, + is_finished=True, + ) + + assert payload is not None + assert not bool(payload.meta.finished) + assert not bool(payload.meta.is_segment_finished) + assert torch.equal(payload.codes.audio, torch.cat([first, second])) + assert tm.request_payload["req-0"]["nvc_frames_seen"] == 2 class _FakeTTSModelOutput: @@ -342,7 +379,14 @@ def _bare_talker(): return talker -def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict: +def _async_info( + timeline: list[int], + prompt_len: int, + finished: bool, + request_id: str = "req-0", + *, + codec_streaming: bool = False, +) -> dict: return { "request_id": request_id, "additional_information": { @@ -351,6 +395,7 @@ def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id "num_processed_tokens": prompt_len, "finished": torch.tensor(finished), "request_id": request_id, + "codec_streaming": codec_streaming, }, }, } @@ -369,11 +414,13 @@ def test_talker_drains_all_received_positions_per_wake() -> None: session = talker._sessions["req-0"] assert session["step"] == 1 and not session["sync_mode"] - # First decode wake: drains ALL prompt-region steps (t=1..P-1) at once. - _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info) + # First resumed wake: drains ALL prompt-region steps (t=1..P-1) at once. + # vLLM 0.28 reports each resumed one-token segment as prefill; the model + # session must still advance instead of restarting its EAR-TTS warmup. + _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=True, **info) assert session["step"] == prompt_len # t advanced 1 -> P in one wake assert update["codes"]["audio"].shape == (prompt_len - 1, 31) - assert float(embeds[0, 0]) == 0.0 # not finished + assert float(embeds[0, 0]) == 0.0 # offline stream waits for upstream terminal # Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save # thread merged two thinker steps). Both must be consumed by this wake. @@ -383,14 +430,13 @@ def test_talker_drains_all_received_positions_per_wake() -> None: assert update["codes"]["audio"].shape == (prompt_len + 1, 31) assert float(embeds[0, 0]) == 0.0 - # Zero-progress wake (no new positions, not finished): plain CONTINUE, no - # fabricated frames, no codes update. + # Zero-progress wake: never fabricate frames or end the offline stream + # before the thinker sends its terminal marker. _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2) assert session["step"] == prompt_len + 2 assert update == {} assert float(embeds[0, 0]) == 0.0 - # Final chunk: one more token + finished marker -> drain + stop flag. info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True) _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3) assert session["step"] == prompt_len + 3 @@ -398,6 +444,18 @@ def test_talker_drains_all_received_positions_per_wake() -> None: assert float(embeds[0, 0]) == 1.0 # stop +def test_talker_native_codec_stream_ends_each_received_segment() -> None: + talker = _bare_talker() + ids = torch.zeros(1, dtype=torch.long) + info = _async_info([12, 12, 7], 2, finished=False, codec_streaming=True) + talker.preprocess(ids, None, _omni_is_prefill=True, **info) + + _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info) + + assert update["codes"]["audio"].shape == (2, 31) + assert float(embeds[0, 0]) == 1.0 + + def test_talker_sync_mode_steps_once_per_wake() -> None: """The verified sync parity path keeps its one-TTS-step-per-engine-step cadence.""" talker = _bare_talker() diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 013c0b99e..996486d49 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -8,9 +8,12 @@ import pytest import torch from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( + AdaptiveChunkController, + compute_adaptive_emit, compute_dynamic_initial_chunk_size, compute_ramp_emit, max_ic_for_chunk_size, + parse_adaptive_config, parse_chunk_ramp, ramp_chunk_size, ramp_cumulative, @@ -41,7 +44,17 @@ def _req(rid, *, finished, initial_codec_chunk_frames=None): ) -def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frames=0, chunk_ramp=None): +def _tm( + *, + chunk_frames=25, + left_context=25, + max_num_seqs=1, + initial_chunk_frames=0, + chunk_ramp=None, + adaptive=False, + adaptive_min=2, + adaptive_margin=200.0, +): extra = { "codec_chunk_frames": chunk_frames, "codec_left_context_frames": left_context, @@ -49,11 +62,16 @@ def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frame } if chunk_ramp is not None: extra["codec_chunk_ramp"] = chunk_ramp + if adaptive: + extra["codec_chunk_adaptive"] = True + extra["codec_chunk_min_frames"] = adaptive_min + extra["codec_chunk_safety_margin_ms"] = adaptive_margin return SimpleNamespace( code_prompt_token_ids=defaultdict(list), scheduler_max_num_seqs=max_num_seqs, put_req_chunk=defaultdict(int), ramp_chunk_count=defaultdict(int), + _adaptive_states={}, request_payload={}, connector=SimpleNamespace(config={"extra": extra}), ) @@ -605,18 +623,32 @@ class TestRampHelpers: def test_parse_ramp_warns_on_tail_mismatch(self, caplog): import logging - with caplog.at_level(logging.WARNING): + target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.WARNING) + try: result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) assert result == [4, 4, 8] - assert any("reintroduces" in r.message for r in caplog.records) + assert "reintroduces" in caplog.text def test_parse_ramp_no_warn_on_tail_match(self, caplog): import logging - with caplog.at_level(logging.WARNING): + target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils") + target_logger.addHandler(caplog.handler) + prev_level = target_logger.level + target_logger.setLevel(logging.WARNING) + try: result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25) + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(prev_level) assert result == [4, 4, 8, 16, 25] - assert not any("reintroduces" in r.message for r in caplog.records) + assert "reintroduces" not in caplog.text @pytest.mark.parametrize( "index,ramp,steady,expected", @@ -941,3 +973,488 @@ class TestChunkRampEmission: p_seg2_0 = self._emit(tm, rid, 4) assert p_seg2_0 is not None assert p_seg2_0.meta.left_context_size == 0 + + +class TestAdaptiveController: + def test_compute_next_chunk_size_basic(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=20.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=10, + ) + now = 0.5 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=100.0) + assert 2 <= size <= 25 + + def test_compute_next_chunk_size_clamp_min(self): + """min_frames floor holds when ramp_floor (last_target + delta) is + below min_frames. Uses a small EWMA so delta=2 and last_target=0, + i.e. ramp_floor=2 < min_frames=4 -> greedy=min_frames wins.""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=1, + ) + # buffer = 80 - 70 = 10ms (positive but < margin=200) + # greedy: available=0, max_n=0, greedy=min_frames=4 + # ramp_floor: last_target=0 + delta=max(2,int(10/15))=2 -> 2 < 4 + now = 0.07 + size = ctrl.compute_next_chunk_size(now, min_frames=4, max_frames=25, safety_margin_ms=200.0) + assert size == 4 # min_frames floor dominates ramp_floor=2 + + def test_compute_next_chunk_size_clamp_max(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=100, + ) + now = 1.0 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert size == 25 + + def test_hysteresis_locks_steady(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=50, + ) + now = 0.1 + size1 = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert ctrl.steady_state_reached + assert size1 == 25 + size2 = ctrl.compute_next_chunk_size(now + 10.0, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert size2 == 25 + + def test_ramp_floor_overrides_greedy_when_buffer_low(self): + """When buffer is negative (overloaded), ramp_floor still forces + gradual growth to avoid stagnating at min_frames (which would cause + too many Code2Wav decode passes and RTF explosion).""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ) + # buffer = 160 - 200 = -40ms (negative) + # greedy: available=0, max_n=0, greedy=2 + # ramp_floor: last_target=4 + delta=max(2,int(50/15))=3 -> 7 + now = 0.2 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert not ctrl.steady_state_reached + assert size == 7 # ramp_floor (unconditional) dominates greedy=2 + + def test_ramp_floor_dominates_when_last_target_set(self): + """Once last_target is set (post-emit), the unconditional ramp_floor + (last_target + ewma-scaled delta) dominates greedy whenever greedy is + smaller. This is intended: chunk size grows monotonically toward max + to minimize Code2Wav decode passes (pre-W5; see W5 coupling).""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=4, + last_target=4, + ) + # buffer = 320 - 100 = 220ms (positive) + # greedy: available=170, max_n=int(170/50)=3, greedy=3 + # ramp_floor: 4 + delta=max(2,int(50/15))=3 -> 7 + now = 0.1 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert size == 7 # ramp_floor (7) dominates greedy (3) + + def test_greedy_dominates_when_ramp_floor_low(self): + """When last_target=0 (fresh, no prior emit) and buffer is positive + but below the hysteresis threshold, greedy selects an intermediate + size larger than ramp_floor and greedy dominates. This is the only + regime where greedy (not ramp_floor or hysteresis) sets the size.""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=8, + last_target=0, + ) + # buffer = 640 - 100 = 540ms (positive, < hysteresis 1300ms) + # greedy: available=490, max_n=int(490/50)=9, greedy=9 + # ramp_floor: 0 + delta=max(2,int(50/15))=3 -> 3 + now = 0.1 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert not ctrl.steady_state_reached + assert size == 9 # greedy (9) dominates ramp_floor (3) + + def test_record_emit_ewma_update(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=20.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=4, + ) + ctrl.record_emit(now=0.5, n_frames=8, accumulated=12, target_size=8) + assert ctrl.frames_emitted_this_segment == 12 + assert ctrl.accumulated_at_last_emit == 12 + assert ctrl.last_emit_time == 0.5 + assert ctrl.last_target == 8 # D6: stores target, not actual + measured = (0.5 - 0.0) * 1000.0 / 8 + expected = 0.3 * measured + 0.7 * 20.0 + assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01 + + def test_record_emit_first_then_second(self): + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=0.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=0, + ) + ctrl.record_emit(now=0.4, n_frames=4, accumulated=4, target_size=4) + assert ctrl.frame_time_ewma_ms == 0.0 + assert ctrl.frames_emitted_this_segment == 4 + assert ctrl.last_target == 4 + + ctrl.record_emit(now=0.8, n_frames=4, accumulated=8, target_size=4) + expected = (0.8 - 0.4) * 1000.0 / 4 + assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01 + + def test_record_emit_collects_underrun(self): + """record_emit mirrors stream_client.py's playback-cursor underrun + model: per_chunk_gap_ms records each chunk's lateness, total_gap_s + accumulates gaps > 5ms. Deterministic given explicit `now`.""" + ctrl = AdaptiveChunkController( + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=0, + ) + # chunk0: now=0.0 (= first_emit_time), n=2 -> 0.16s audio + ctrl.record_emit(now=0.0, n_frames=2, accumulated=2, target_size=2) + assert ctrl.chunk_sizes == [2] + assert ctrl.per_chunk_gap_ms == [0.0] + assert ctrl.total_gap_s == 0.0 + assert abs(ctrl.playback_cursor_s - 0.16) < 1e-9 + # chunk1: now=0.2 (200ms wall) -> arrival 0.2, late by 0.2-0.16=0.04s + ctrl.record_emit(now=0.2, n_frames=2, accumulated=4, target_size=2) + assert ctrl.chunk_sizes == [2, 2] + assert abs(ctrl.per_chunk_gap_ms[1] - 40.0) < 1e-6 + assert abs(ctrl.total_gap_s - 0.04) < 1e-9 + assert abs(ctrl.playback_cursor_s - 0.36) < 1e-9 + + def test_ewma_zero_skips_hysteresis(self): + """When EWMA=0, hysteresis must not trigger (steady_gen_ms=0 + would make any positive buffer satisfy the condition).""" + ctrl = AdaptiveChunkController( + frame_time_ewma_ms=0.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=10, + ) + now = 0.001 + size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0) + assert not ctrl.steady_state_reached + assert size == 2 # min_frames (EWMA=0 → max_n=min_frames) + + def test_ramp_divisor_configurable(self): + """ramp_divisor/ramp_delta_min are controller fields wired from config. + Defaults (15/2) preserve the original formula; changing them changes + delta. Lets a deploy tune the ramp via yaml without code changes, while + the unit tests (which use defaults) stay green.""" + # default divisor=15: ewma=50 -> delta=max(2,int(50/15))=3 -> floor=4+3=7 + c1 = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ) + assert c1.compute_next_chunk_size(0.2, 2, 25, 50.0) == 7 + # divisor=10: delta=max(2,int(50/10))=5 -> floor=4+5=9 + c2 = AdaptiveChunkController( + frame_time_ewma_ms=50.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ramp_divisor=10, + ) + assert c2.compute_next_chunk_size(0.2, 2, 25, 50.0) == 9 + # delta_min=4 + ewma=10 (int(10/15)=0): delta=max(4,0)=4 -> floor=4+4=8 + c3 = AdaptiveChunkController( + frame_time_ewma_ms=10.0, + first_emit_time=0.0, + last_emit_time=0.0, + frames_emitted_this_segment=2, + last_target=4, + ramp_delta_min=4, + ) + assert c3.compute_next_chunk_size(0.2, 2, 25, 50.0) == 8 + + +class TestAdaptiveEmitHelper: + @pytest.mark.parametrize( + "length,accumulated,target,finished,expected_emit,expected_ctx", + [ + (4, 0, 4, False, True, 4), + (3, 0, 4, False, False, 0), + (12, 4, 8, False, True, 8), + (10, 4, 8, False, False, 0), + # D3: finished flushes ALL remaining frames, not just target_size + (6, 4, 8, True, True, 2), + (4, 4, 8, True, True, 0), + (20, 4, 8, True, True, 16), # new_frames=16 > target=8 → flush 16 + (30, 0, 25, True, True, 30), # finished, new_frames=30 > target=25 → flush 30 + # B1: non-finished path also ships ALL accumulated frames (new_frames), + # not just target_size. When new_frames > target_size (e.g. scheduler + # stall dropped target below queued frames), returning target_size + # would cause the [-context_length:] slice to skip surplus frames. + (10, 0, 4, False, True, 10), # new_frames=10 > target=4 → ship 10 + (8, 2, 4, False, True, 6), # new_frames=6 > target=4 → ship 6 + (100, 0, 25, False, True, 100), # new_frames=100 > target=25 → ship 100 + ], + ) + def test_compute_adaptive_emit(self, length, accumulated, target, finished, expected_emit, expected_ctx): + emit, ctx = compute_adaptive_emit(length, accumulated, target, finished) + assert emit == expected_emit + assert ctx == expected_ctx + + +class TestAdaptiveConfigParsing: + def test_disabled_by_default(self): + enabled, min_f, margin, divisor, delta_min = parse_adaptive_config({}) + assert enabled is False + assert min_f == 2 + assert margin == 50.0 + assert divisor == 15 + assert delta_min == 2 + + def test_enabled(self): + cfg = {"codec_chunk_adaptive": True} + enabled, *_ = parse_adaptive_config(cfg) + assert enabled is True + + def test_custom_values(self): + cfg = { + "codec_chunk_adaptive": True, + "codec_chunk_min_frames": 4, + "codec_chunk_safety_margin_ms": 300.0, + "codec_chunk_ramp_divisor": 10, + "codec_chunk_ramp_delta_min": 3, + } + enabled, min_f, margin, divisor, delta_min = parse_adaptive_config(cfg) + assert enabled is True + assert min_f == 4 + assert margin == 300.0 + assert divisor == 10 + assert delta_min == 3 + + def test_min_frames_clamped(self): + cfg = {"codec_chunk_adaptive": True, "codec_chunk_min_frames": 0} + _, min_f, _, _, _ = parse_adaptive_config(cfg) + assert min_f == 1 + + def test_ramp_divisor_clamped(self): + cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_divisor": 0} + _, _, _, divisor, _ = parse_adaptive_config(cfg) + assert divisor == 1 + + def test_ramp_delta_min_clamped(self): + cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_delta_min": 0} + _, _, _, _, delta_min = parse_adaptive_config(cfg) + assert delta_min == 1 + + +class TestAdaptiveEmission: + def _emit(self, tm, rid, n_frames, finished=False): + tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)] + return talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req(rid, finished=finished), + is_finished=finished, + ) + + def test_adaptive_chunk0_uses_dynamic_ic(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-ic" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 2 + + def test_adaptive_chunk0_holds_until_ic(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-hold" + p = self._emit(tm, rid, 1) + assert p is None + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 2 + + def test_adaptive_finished_flush(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-flush" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p_fin = self._emit(tm, rid, 5, finished=True) + assert p_fin is not None + assert p_fin.meta.finished.item() is True + + def test_adaptive_finished_no_new_frames(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-no-new" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + + p_fin = self._emit(tm, rid, 2, finished=True) + assert p_fin is not None + assert p_fin.meta.finished.item() is True + assert p_fin.codes.audio.numel() == 0 + + def test_adaptive_precedence_over_fixed_ramp(self): + tm = _tm(chunk_ramp=[4, 4, 8, 16, 25], adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-precedence" + + p = self._emit(tm, rid, 1) + assert p is None + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 2 + + def test_adaptive_backward_compat(self): + tm = _tm(initial_chunk_frames=1) + rid = "no-adaptive" + + p0 = self._emit(tm, rid, 1) + assert p0 is not None + assert len(p0.codes.audio) == _Q * 1 + + def test_adaptive_segment_reset(self): + tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8) + rid = "adaptive-segment" + + p0 = self._emit(tm, rid, 2) + assert p0 is not None + tm.ramp_chunk_count[rid] = 1 + tm.put_req_chunk[rid] = 1 + + tm.code_prompt_token_ids[rid] = [] + tm.ramp_chunk_count.pop(rid, None) + tm._adaptive_states.pop(rid, None) + + p_seg2 = self._emit(tm, rid, 2) + assert p_seg2 is not None + assert p_seg2.meta.left_context_size == 0 + + def test_adaptive_with_ref_code_first_chunk(self): + tm = _tm(adaptive=True, adaptive_min=2, left_context=25, max_num_seqs=8) + rid = "adaptive-ref" + tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(2)] + ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) + + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}}, + request=_req(rid, finished=False), + is_finished=False, + ) + + assert payload is not None + assert payload.meta.ref_context_size == 2 + assert payload.meta.ref_context_included is True + assert payload.meta.left_context_size == 2 + + +class TestAdaptiveAccumulationPath: + """Regression test for the frame-skipping bug (PR #6001 review B1). + + Appends one frame per call with a controlled clock, asserting that + the concatenation of every emitted chunk equals the accumulated + frame sequence, in order, with no gaps. Each frame has a distinct + value so the assertion can detect a skipped frame rather than just + a short total. + """ + + def test_no_frames_skipped_when_target_drops(self, monkeypatch): + """When greedy target_size is high (large buffer) and then drops + (clock jumps), new_frames > target_size at the emit point. + compute_adaptive_emit now returns new_frames (not target_size), + so all accumulated frames are shipped and the [-context_length:] + slice stays contiguous with the previous emit. last_target is set + to target_size (not n_frames) to prevent the ramp floor from + ratcheting up on transient overshoot. + + Clock schedule (max_num_seqs=4 gives IC=2 via dynamic IC): + t=0.0 chunk 0 (IC=2), emit 2. ewma=0. + t=0.05 emit target=4 (ramp_floor). ewma=12.5 (direct assign). + t=0.05 greedy dominates (target=18), 17 frames accumulate (holds). + t=100.0 clock jumps -> greedy collapses, target drops to + ramp_floor=6. new_frames=18 > 6 -> emit 18 (all new frames). + last_target=6 (not 18, preventing floor ratchet). + """ + tm = _tm( + adaptive=True, + adaptive_min=2, + max_num_seqs=4, + ) + rid = "adaptive-accum" + tm.code_prompt_token_ids[rid] = [] + + clock = [0.0] + monkeypatch.setattr( + "vllm_omni.model_executor.stage_input_processors.qwen3_tts.time.monotonic", + lambda: clock[0], + ) + + emitted_indices: list[int] = [] + total_frames = 40 + + def append_and_emit(frame_idx, finished=False): + tm.code_prompt_token_ids[rid].append([frame_idx] * _Q) + p = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req(rid, finished=finished), + is_finished=finished, + ) + if p is not None: + if p.codes.audio.numel() > 0: + n = p.codes.audio.numel() // _Q + emitted_indices.extend(p.codes.audio[:n].tolist()) + tm.ramp_chunk_count[rid] += 1 + return p + + # Phase 1 (t=0.0): chunk 0, IC=2 + for i in range(2): + append_and_emit(i) + + # Phase 2 (t=0.05): build EWMA, emit target=4 at length=6 + clock[0] = 0.05 + for i in range(2, 6): + append_and_emit(i) + + # Phase 3 (t=0.05): greedy target=18 (large buffer), accumulate + # 17 frames while holding (new_frames 1..17 < 18) + for i in range(6, 23): + append_and_emit(i) + + # Phase 4 (t=100.0): clock jumps, greedy collapses, target drops + # to ramp_floor=6. new_frames=18 > 6 -> bug trigger point. + clock[0] = 100.0 + for i in range(23, total_frames): + append_and_emit(i, finished=(i == total_frames - 1)) + + # Assert: all frames 0..39 emitted, in order, no gaps. + expected = list(range(total_frames)) + assert emitted_indices == expected, ( + f"Frame loss detected! " + f"Missing: {sorted(set(expected) - set(emitted_indices))}, " + f"Extra: {sorted(set(emitted_indices) - set(expected))}" + ) diff --git a/tests/model_executor/test_cfg_pairing.py b/tests/model_executor/test_cfg_pairing.py new file mode 100644 index 000000000..c00313a2a --- /dev/null +++ b/tests/model_executor/test_cfg_pairing.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the model-neutral CFG request-pairing scheduler patches. + +The patches keep a cond/uncond request pair step-locked inside one engine: +admitted together, kept at equal ``num_computed_tokens``, finished together. +They are installed onto a lightweight fake scheduler here; vLLM's real +``Scheduler`` needs a full engine (and a GPU) to construct. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.model_executor.models.common import cfg_pairing +from vllm_omni.model_executor.models.common.cfg_pairing import ( + _equalize_cfg_pair_progress, + _patch_scheduler, + cfg_uncond_suffixes, + register_cfg_uncond_suffix, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# ───────── fakes ───────── + + +class _FakeRequest: + """Minimal stand-in for ``vllm.v1.request.Request`` (identity equality).""" + + def __init__(self, request_id: str, extra_args: dict[str, Any] | None = None, step_tokens: int = 1) -> None: + self.request_id = request_id + self.sampling_params = SimpleNamespace(extra_args=extra_args) + self.num_computed_tokens = 0 + self.step_tokens = step_tokens + self.skip_reading_prefix_cache = False + + +class _FakeWaitingQueue(deque): + """FCFS waiting queue: vLLM's is also a ``deque`` subclass.""" + + def remove_requests(self, requests: list[_FakeRequest]) -> None: + keep = [req for req in self if not any(req is dropped for dropped in requests)] + self.clear() + self.extend(keep) + + def prepend_request(self, request: _FakeRequest) -> None: + self.appendleft(request) + + +@dataclass +class _FakeSchedulerOutput: + num_scheduled_tokens: dict[str, int] = field(default_factory=dict) + total_num_scheduled_tokens: int = 0 + + +def _make_scheduler_cls() -> type: + """Fresh class per test so the in-place patches never stack.""" + + class FakeScheduler: + def __init__(self) -> None: + self.waiting = _FakeWaitingQueue() + self.running: list[_FakeRequest] = [] + self.requests: dict[str, _FakeRequest] = {} + self.finished: list[set[str]] = [] + self.observed_waiting: list[list[str]] = [] + self.max_num_scheduled_tokens = 512 + self.scheduler_config = SimpleNamespace(long_prefill_token_threshold=4096) + + def add_request(self, request: _FakeRequest) -> None: + self.requests[request.request_id] = request + self.waiting.append(request) + + def finish_requests(self, request_ids: Any, finished_status: Any) -> None: + if request_ids is None: + ids = set(self.requests) + elif isinstance(request_ids, str): + ids = {request_ids} + else: + ids = set(request_ids) + self.finished.append(ids) + self.running = [req for req in self.running if req.request_id not in ids] + for req_id in ids: + self.requests.pop(req_id, None) + + def schedule(self) -> _FakeSchedulerOutput: + self.observed_waiting.append([req.request_id for req in self.waiting]) + while self.waiting: + self.running.append(self.waiting.popleft()) + output = _FakeSchedulerOutput() + for req in self.running: + req.num_computed_tokens += req.step_tokens + output.num_scheduled_tokens[req.request_id] = req.step_tokens + output.total_num_scheduled_tokens += req.step_tokens + return output + + return FakeScheduler + + +@pytest.fixture +def scheduler() -> Any: + cls = _make_scheduler_cls() + _patch_scheduler(cls) + return cls() + + +@pytest.fixture(autouse=True) +def _restore_suffix_registry(): + original = set(cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES) + yield + cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.clear() + cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.update(original) + + +def _add_pair( + scheduler: Any, + pair_id: str | None = None, + base_id: str = "req-1", + suffix: str = "__cfg_uncond", +) -> tuple[_FakeRequest, _FakeRequest]: + cond_extra: dict[str, Any] = {"cfg_role": "cond"} + uncond_extra: dict[str, Any] = {"cfg_role": "uncond"} + if pair_id is not None: + cond_extra["cfg_pair_id"] = pair_id + uncond_extra["cfg_pair_id"] = pair_id + cond = _FakeRequest(base_id, cond_extra) + uncond = _FakeRequest(f"{base_id}{suffix}", uncond_extra) + scheduler.add_request(cond) + scheduler.add_request(uncond) + return cond, uncond + + +# ───────── pair-id derivation ───────── + + +def test_pair_id_is_derived_from_request_id_when_absent(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + + assert scheduler._cfg_pairs == {"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}} + assert scheduler._cfg_req_to_pair == {cond.request_id: "req-1", uncond.request_id: "req-1"} + + +def test_registered_companion_suffix_is_stripped(scheduler: Any) -> None: + register_cfg_uncond_suffix("::negative") + assert "::negative" in cfg_uncond_suffixes() + + cond, uncond = _add_pair(scheduler, base_id="music-7", suffix="::negative") + + assert scheduler._cfg_pairs == {"music-7": {"cond": cond.request_id, "uncond": uncond.request_id}} + + +def test_register_rejects_empty_suffix() -> None: + with pytest.raises(ValueError): + register_cfg_uncond_suffix("") + + +def test_uncond_without_known_suffix_keeps_its_own_request_id(scheduler: Any) -> None: + # Degenerate input: the companion never joins the cond pair, so the pair + # stays incomplete and the hold budget eventually releases it. + scheduler.add_request(_FakeRequest("solo-uncond", {"cfg_role": "uncond"})) + + assert scheduler._cfg_pairs == {"solo-uncond": {"uncond": "solo-uncond"}} + + +def test_explicit_pair_id_takes_precedence(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler, pair_id="pair-A") + + assert scheduler._cfg_pairs == {"pair-A": {"cond": cond.request_id, "uncond": uncond.request_id}} + assert cond.request_id not in scheduler._cfg_pairs + assert scheduler._cfg_req_to_pair[uncond.request_id] == "pair-A" + + +def test_cfg_requests_skip_prefix_cache_reads(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + + assert cond.skip_reading_prefix_cache is True + assert uncond.skip_reading_prefix_cache is True + + +# ───────── admission holding ───────── + + +def test_lone_member_is_held_then_released(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + + scheduler.schedule() + + assert scheduler.observed_waiting == [[]] # the base scheduler saw nothing + assert [req.request_id for req in scheduler.waiting] == ["req-1"] + assert scheduler.running == [] + + +def test_complete_pair_is_admitted_together(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + scheduler.schedule() + scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})) + + scheduler.schedule() + + assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond"] + assert [req.request_id for req in scheduler.running] == ["req-1", "req-1__cfg_uncond"] + assert list(scheduler.waiting) == [] + + +def test_hold_budget_expiry_drops_the_stale_pair(scheduler: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(cfg_pairing, "_MAX_PAIR_HOLD_STEPS", 2) + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + + scheduler.schedule() + scheduler.schedule() + assert scheduler.running == [] + assert scheduler._cfg_pairs # still hoping for the partner + + scheduler.schedule() + + assert scheduler._cfg_pairs == {} + assert scheduler._cfg_req_to_pair == {} + assert [req.request_id for req in scheduler.running] == ["req-1"] + + +def test_waiting_queue_reorder_puts_partners_adjacent(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + scheduler.add_request(_FakeRequest("other", None)) + scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})) + + scheduler.schedule() + + assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond", "other"] + + +# ───────── progress equalization ───────── + + +def _equalize_fixture( + cond_computed: int, + uncond_computed: int, + cond_sched: int, + uncond_sched: int, +) -> tuple[Any, _FakeSchedulerOutput, _FakeRequest, _FakeRequest]: + cond = _FakeRequest("req-1", {"cfg_role": "cond"}) + uncond = _FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}) + cond.num_computed_tokens = cond_computed + uncond.num_computed_tokens = uncond_computed + fake = SimpleNamespace( + _cfg_pairs={"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}}, + requests={cond.request_id: cond, uncond.request_id: uncond}, + ) + output = _FakeSchedulerOutput( + num_scheduled_tokens={cond.request_id: cond_sched, uncond.request_id: uncond_sched}, + total_num_scheduled_tokens=cond_sched + uncond_sched, + ) + return fake, output, cond, uncond + + +def test_equalize_pulls_the_ahead_member_back() -> None: + fake, output, cond, uncond = _equalize_fixture(10, 6, 5, 4) + + _equalize_cfg_pair_progress(fake, output) + + assert cond.num_computed_tokens == 6 + assert uncond.num_computed_tokens == 6 + assert output.num_scheduled_tokens == {cond.request_id: 1, uncond.request_id: 4} + assert output.total_num_scheduled_tokens == 5 + + +def test_equalize_is_a_noop_when_the_pair_is_already_aligned() -> None: + fake, output, cond, uncond = _equalize_fixture(8, 8, 4, 4) + + _equalize_cfg_pair_progress(fake, output) + + assert (cond.num_computed_tokens, uncond.num_computed_tokens) == (8, 8) + assert output.total_num_scheduled_tokens == 8 + + +def test_equalize_skips_an_infeasible_rollback() -> None: + # Clawing back 4 tokens would leave the cond row with 0 scheduled tokens, + # which the scheduler output cannot express; leave the pair to _drop_split_pairs. + fake, output, cond, uncond = _equalize_fixture(10, 6, 4, 4) + + _equalize_cfg_pair_progress(fake, output) + + assert cond.num_computed_tokens == 10 + assert output.total_num_scheduled_tokens == 8 + + +# ───────── finish ───────── + + +def test_finish_requests_finishes_the_partner(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + scheduler.schedule() + + scheduler.finish_requests(cond.request_id, "FINISHED_STOPPED") + + assert scheduler.finished == [{cond.request_id, uncond.request_id}] + assert scheduler._cfg_pairs == {} + assert scheduler._cfg_req_to_pair == {} + assert scheduler.running == [] + + +def test_finish_all_requests_clears_the_registry(scheduler: Any) -> None: + _add_pair(scheduler) + scheduler.schedule() + + scheduler.finish_requests(None, "FINISHED_ABORTED") + + assert scheduler._cfg_pairs == {} + assert scheduler._cfg_req_to_pair == {} + + +# ───────── non-CFG traffic ───────── + + +def test_plain_requests_are_untouched(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("plain-1", None)) + scheduler.add_request(_FakeRequest("plain-2", {"seed": 3})) + + output = scheduler.schedule() + + assert scheduler._cfg_pairs == {} + assert scheduler.observed_waiting == [["plain-1", "plain-2"]] + assert output.num_scheduled_tokens == {"plain-1": 1, "plain-2": 1} + assert output.total_num_scheduled_tokens == 2 + # No pair ever existed, so the hold bookkeeping is never allocated. + assert not hasattr(scheduler, "_cfg_hold_counts") + + +def test_default_suffix_matches_the_audex_stage_processor() -> None: + from vllm_omni.model_executor.stage_input_processors.audex import CFG_UNCOND_SUFFIX + + assert CFG_UNCOND_SUFFIX in cfg_uncond_suffixes() + + +def test_audex_cfg_module_still_exports_the_patch_entry_points() -> None: + from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches, cfg_patches_applied + + assert apply_cfg_patches is cfg_pairing.apply_cfg_patches + assert cfg_patches_applied is cfg_pairing.cfg_patches_applied + + +# ───────── install gate ───────── + + +def test_gate_triggers_on_a_cfg_logits_processor() -> None: + from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs + + class AudexCFGLogitsProcessor: + pass + + assert _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[AudexCFGLogitsProcessor])) + # vLLM keeps unresolved entries as dotted strings. + assert _stage_declares_cfg_pairs( + SimpleNamespace(logits_processors=["vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"]) + ) + + +def test_gate_triggers_on_a_declared_cfg_role_default() -> None: + from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs + + assert _stage_declares_cfg_pairs( + SimpleNamespace(logits_processors=None, sampling_extra_args_keys=("cfg_role", "cfg_scale")) + ) + + +def test_gate_stays_off_for_non_cfg_stages() -> None: + from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs + + assert not _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[], sampling_extra_args_keys=())) + assert not _stage_declares_cfg_pairs( + SimpleNamespace(logits_processors=["some.module:MinPLogitsProcessor"], sampling_extra_args_keys=("tta_rvq",)) + ) + assert not _stage_declares_cfg_pairs(SimpleNamespace()) + + +def test_sampling_extra_args_keys_are_derived_from_stage_defaults() -> None: + from vllm_omni.engine.stage_init_utils import _sampling_extra_args_keys + + assert _sampling_extra_args_keys({"extra_args": {"cfg_scale": 1.5, "cfg_role": "cond"}}) == ( + "cfg_role", + "cfg_scale", + ) + assert _sampling_extra_args_keys({"temperature": 0.1}) == () + assert _sampling_extra_args_keys(None) == () diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index d57671adc..463bedfc3 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -20,6 +20,7 @@ from vllm_omni.model_extras import ( get_output_tensor_range, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, + should_preserve_reference_image_size, ) @@ -150,6 +151,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na "max_sequence_length", "use_resolution_template", "use_duration_template", + "negative_metadata_mode", "use_system_prompt", "system_prompt", "negative_prompt", @@ -192,23 +194,6 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False -@pytest.mark.core_model -@pytest.mark.cpu -def test_magi_human_extra_registry_declares_request_and_response_params() -> None: - assert get_extra_body_params("MagiHumanPipeline") == frozenset( - { - "seconds", - "audio_path", - "image_path", - "sr_height", - "sr_width", - "sr_num_inference_steps", - } - ) - assert get_extra_output_params("MagiHumanPipeline") == frozenset() - assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False - - @pytest.mark.core_model @pytest.mark.cpu def test_ltx_extra_registry_declares_official_guidance_params() -> None: @@ -230,13 +215,60 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None: "audio_rescale_scale", "video_stg_blocks", "audio_stg_blocks", + "sigmas", + "stage_1_sigmas", + "stage_2_sigmas", + "image_crf", } ) - assert get_extra_body_params("LTX2Pipeline") == expected - assert get_extra_output_params("LTX2Pipeline") == frozenset() - assert get_extra_body_params("LTX2DistilledPipeline") == expected - assert get_extra_output_params("LTX2DistilledPipeline") == frozenset() + for pipeline_name in ( + "LTX2Pipeline", + "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", + "LTX2DistilledPipeline", + ): + assert get_extra_body_params(pipeline_name) == expected + assert get_extra_output_params(pipeline_name) == frozenset() + + +@pytest.mark.parametrize( + ("model_version", "expected"), + [("2", False), ("2.3", False), ("2.5", True)], +) +def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected: bool) -> None: + (tmp_path / "model_index.json").write_text( + '{"_class_name": "LTX2Pipeline", "model_version": "' + model_version + '"}' + ) + + assert ( + should_preserve_reference_image_size( + None, + model=str(tmp_path), + ) + is expected + ) + + +def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None: + captured = {} + + def fake_policy(*, model, revision=None): + captured.update(model=model, revision=revision) + return True + + monkeypatch.setattr( + "vllm_omni.diffusion.models.ltx2.ltx2_components.preserves_reference_image_size", + fake_policy, + ) + + assert should_preserve_reference_image_size( + "LTX2Pipeline", + model="org/model", + revision="pinned-revision", + ) + assert captured == {"model": "org/model", "revision": "pinned-revision"} @pytest.mark.core_model @@ -337,47 +369,6 @@ def test_ming_flash_omni_image_to_image_prompt_builder() -> None: assert "negative_prompt" not in result -@pytest.mark.core_model -@pytest.mark.cpu -def test_audiox_extra_registry_declares_request_and_response_params() -> None: - assert get_extra_body_params("AudioXPipeline") == frozenset( - { - "audiox_task", - "seconds_start", - "seconds_total", - "sigma_min", - "sigma_max", - "cfg_rescale", - "video_path", - "audio_path", - } - ) - assert get_extra_output_params("AudioXPipeline") == frozenset({"audiox_task"}) - assert should_init_extra_args_for_non_diffusion_stages("AudioXPipeline") is False - - -@pytest.mark.core_model -@pytest.mark.cpu -def test_audiox_declared_extra_args_route_into_sampling_params() -> None: - params = OmniDiffusionSamplingParams() - declared = get_extra_body_params("AudioXPipeline") - apply_declared_extra_args( - params, - declared, - { - "audiox_task": "t2a", - "seconds_total": 10.0, - "sigma_min": 0.03, - "unknown": "ignored", - }, - ) - assert params.extra_args == { - "audiox_task": "t2a", - "seconds_total": 10.0, - "sigma_min": 0.03, - } - - @pytest.mark.core_model @pytest.mark.cpu def test_helios_extra_registry_declares_request_and_response_params() -> None: @@ -421,6 +412,13 @@ def test_vace_extra_registry_has_no_pipeline_params() -> None: def test_unknown_pipeline_has_empty_extra_registry() -> None: assert get_extra_body_params("UnknownPipeline") == frozenset() assert get_extra_output_params("UnknownPipeline") == frozenset() + assert ( + should_preserve_reference_image_size( + "UnknownPipeline", + model="org/model", + ) + is False + ) assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False diff --git a/tests/model_tests/conftest.py b/tests/model_tests/conftest.py index ab71a02cd..7760fa25d 100644 --- a/tests/model_tests/conftest.py +++ b/tests/model_tests/conftest.py @@ -1,12 +1,14 @@ +import json from shutil import rmtree import pytest +from filelock import ReadWriteLock from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS @pytest.fixture(scope="session") -def tiny_model_paths(request, run_level): +def tiny_model_paths(request, run_level, tmp_path_factory): """Build or download the tiny models for the selected tests. At core_model level, builds tiny models via the builder function. @@ -15,27 +17,43 @@ def tiny_model_paths(request, run_level): NOTE: this is session scoped to avoid churn in tiny model creation, but will ensure all the tiny models you need are created for the selected tests before it starts to execute them.""" + base_path = tmp_path_factory.getbasetemp().parent / "tiny_models" + use_lock = ReadWriteLock(base_path.with_suffix(".use.lock")) + build_lock = ReadWriteLock(base_path.with_suffix(".build.lock")) + data_file = base_path.with_suffix(".json") + model_paths = {} built_paths = [] - print("Initializing tiny models...") - for item in request.session.items: - if not hasattr(item, "callspec"): - continue - model_name = item.callspec.params.get("model_name") - if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS: - continue - if model_name not in model_paths: - settings = DIFFUSION_TEST_SETTINGS[model_name] - if run_level == "core_model": - print(f"Calling tiny model builder for: {model_name}") - path = settings.builder() - built_paths.append(path) - else: - print(f"Run level is {run_level}; {model_name} will use the full model") - path = settings.model - model_paths[model_name] = path - yield model_paths - for path in built_paths: - print(f"Removing tiny model: {path}") - rmtree(path, ignore_errors=True) + with use_lock.read_lock(): + with build_lock.write_lock(): + if data_file.is_file(): + model_paths = json.loads(data_file.read_text()) + else: + print("Initializing tiny models...") + for item in request.session.items: + if not hasattr(item, "callspec"): + continue + model_name = item.callspec.params.get("model_name") + if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS: + continue + if model_name not in model_paths: + settings = DIFFUSION_TEST_SETTINGS[model_name] + if run_level == "core_model": + print(f"Calling tiny model builder for: {model_name}") + path = settings.builder() + built_paths.append(path) + else: + print(f"Run level is {run_level}; {model_name} will use the full model") + path = settings.model + model_paths[model_name] = path + data_file.write_text(json.dumps(model_paths)) + + yield model_paths + + with use_lock.write_lock(): + if built_paths: + data_file.unlink(missing_ok=True) + for path in built_paths: + print(f"Removing tiny model: {path}") + rmtree(path, ignore_errors=True) diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py index e47ea1895..7c1c9df67 100644 --- a/tests/model_tests/diffusion/diff_model_builders.py +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -1,7 +1,16 @@ from functools import partial from pathlib import Path -from tests.helpers.tiny_model import build_tiny_from_configs +import torch +from diffusers import ( + AutoencoderKLWan, + DPMSolverMultistepScheduler, + SanaVideoPipeline, + SanaVideoTransformer3DModel, +) +from transformers import Gemma2Config, Gemma2Model, GemmaTokenizerFast + +from tests.helpers.tiny_model import _get_tiny_model_path, build_tiny_from_configs TINY_CONFIGS_DIR = Path(__file__).parent / "tiny_configs" @@ -40,6 +49,72 @@ def tiny_ltx2_builder() -> str: return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline") +def tiny_sana_video_builder() -> str: + """Build a tiny 480p SANA-Video model without downloading model weights. + + The tokenizer is the only component loaded from the upstream repository. + All modules with weights are initialized locally from intentionally small + configs, and the scheduler is constructed from its weight-free config. + """ + model_id = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers" + model_dir = _get_tiny_model_path("SanaVideoPipeline") + + tokenizer = GemmaTokenizerFast.from_pretrained(model_id, subfolder="tokenizer") + scheduler = DPMSolverMultistepScheduler( + algorithm_type="dpmsolver++", + flow_shift=8.0, + prediction_type="flow_prediction", + use_flow_sigmas=True, + ) + text_encoder = Gemma2Model( + Gemma2Config( + vocab_size=256000, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=512, + sliding_window=256, + layer_types=["sliding_attention"], + ) + ) + transformer = SanaVideoTransformer3DModel( + in_channels=16, + out_channels=16, + num_attention_heads=2, + attention_head_dim=16, + num_layers=1, + num_cross_attention_heads=2, + cross_attention_head_dim=16, + cross_attention_dim=32, + caption_channels=32, + mlp_ratio=2.0, + sample_size=30, + patch_size=(1, 2, 2), + ) + vae = AutoencoderKLWan( + base_dim=8, + decoder_base_dim=8, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + latents_mean=[0.0] * 16, + latents_std=[1.0] * 16, + ) + pipeline = SanaVideoPipeline( + tokenizer=tokenizer, + text_encoder=text_encoder, + vae=vae, + transformer=transformer, + scheduler=scheduler, + ) + pipeline.to(dtype=torch.bfloat16).save_pretrained(model_dir) + return model_dir + + def _shrink_flux_clip_text_encoder(config: dict) -> dict: config["num_hidden_layers"] = 2 return config @@ -54,14 +129,18 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict: def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict: - text_config = config["text_config"] + """Shrink a Qwen2.5-VL text encoder (nested Qwen-Image or flat LongCat layout).""" + # Nested: text fields live under text_config, with top-level mirrors. + # Flat: text fields live at the top level. + text_config = config["text_config"] if "text_config" in config else config old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"] text_config["num_hidden_layers"] = 2 text_config["intermediate_size"] = 64 text_config["hidden_size"] = hidden_size text_config["num_attention_heads"] = 2 text_config["num_key_value_heads"] = 2 - text_config["layer_types"] = text_config["layer_types"][:2] + if "layer_types" in text_config: + text_config["layer_types"] = text_config["layer_types"][:2] factor = old_head_dim / (hidden_size / 2) mrope_section = text_config["rope_scaling"]["mrope_section"] text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section] @@ -105,6 +184,38 @@ def tiny_qwen_image_edit_plus_builder() -> str: ) +def tiny_longcat_image_builder() -> str: + return build_tiny_from_configs( + "LongCatImagePipeline", + "meituan-longcat/LongCat-Image", + transform={ + "text_encoder": _shrink_qwen_text_encoder_config, + "transformer": partial( + _shrink_dit_rope_config, + num_single_layers=2, + default_axes_dims_rope=[16, 56, 56], + joint_attention_dim=64, + ), + }, + ) + + +def tiny_longcat_image_edit_builder() -> str: + return build_tiny_from_configs( + "LongCatImageEditPipeline", + "meituan-longcat/LongCat-Image-Edit", + transform={ + "text_encoder": _shrink_qwen_text_encoder_config, + "transformer": partial( + _shrink_dit_rope_config, + num_single_layers=2, + default_axes_dims_rope=[16, 56, 56], + joint_attention_dim=64, + ), + }, + ) + + def tiny_flux_builder() -> str: return build_tiny_from_configs( "FluxPipeline", @@ -131,19 +242,18 @@ def tiny_flux_kontext_builder() -> str: def tiny_flux2_builder() -> str: def shrink_text_encoder(config: dict) -> dict: - # The real checkpoint ships language_model.lm_head.weight as its own - # tensor even though tie_word_embeddings defaults to True; keep it - # untied here too so transformers doesn't drop it on save. config["tie_word_embeddings"] = False - config["text_config"]["num_hidden_layers"] = 31 + config["text_config"]["num_hidden_layers"] = 3 + config["text_config"]["hidden_size"] = 32 config["text_config"]["intermediate_size"] = 64 - config["text_config"]["num_attention_heads"] = 4 - config["text_config"]["head_dim"] = 32 + config["text_config"]["num_attention_heads"] = 2 + config["text_config"]["head_dim"] = 16 config["text_config"]["num_key_value_heads"] = 2 + config["text_config"]["text_encoder_out_layers"] = [0, 1, 2] config["vision_config"]["num_hidden_layers"] = 2 config["vision_config"]["intermediate_size"] = 64 - config["vision_config"]["num_attention_heads"] = 4 - config["vision_config"]["head_dim"] = 32 + config["vision_config"]["num_attention_heads"] = 2 + config["vision_config"]["head_dim"] = 16 return config return build_tiny_from_configs( @@ -151,6 +261,6 @@ def tiny_flux2_builder() -> str: "black-forest-labs/FLUX.2-dev", transform={ "text_encoder": shrink_text_encoder, - "transformer": partial(_shrink_dit_rope_config, num_single_layers=2), + "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, joint_attention_dim=96), }, ) diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py index cbb37729a..f24564f62 100644 --- a/tests/model_tests/diffusion/model_settings.py +++ b/tests/model_tests/diffusion/model_settings.py @@ -49,11 +49,37 @@ DIFFUSION_TEST_SETTINGS = { [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD], ], ), + "SanaVideoPipeline": DiffusionModelTestOpts( + model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers", + builder=diff_model_builders.tiny_sana_video_builder, + supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO], + check_multi_output=False, + ), "QwenImagePipeline": DiffusionModelTestOpts( model="Qwen/Qwen-Image", builder=diff_model_builders.tiny_qwen_image_builder, supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE], ), + "LongCatImagePipeline": DiffusionModelTestOpts( + model="meituan-longcat/LongCat-Image", + builder=diff_model_builders.tiny_longcat_image_builder, + supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE], + extra_test_groups=[ + [DiffusionAccs.TEA_CACHE], + [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD], + [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD], + ], + ), + "LongCatImageEditPipeline": DiffusionModelTestOpts( + model="meituan-longcat/LongCat-Image-Edit", + builder=diff_model_builders.tiny_longcat_image_edit_builder, + supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE], + extra_test_groups=[ + [DiffusionAccs.TEA_CACHE], + [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD], + [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD], + ], + ), "FluxPipeline": DiffusionModelTestOpts( model="black-forest-labs/FLUX.1-schnell", builder=diff_model_builders.tiny_flux_builder, diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py index 2be2338c1..4fb31fb5c 100644 --- a/tests/model_tests/diffusion/task_runners.py +++ b/tests/model_tests/diffusion/task_runners.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Definitions running and validating individual tasks, e.g., text to image, image to image, and so on. These are called by the core test runner. @@ -6,12 +9,13 @@ image to image, and so on. These are called by the core test runner. import base64 import io from dataclasses import replace +from typing import Any import numpy as np import pytest from PIL import Image -from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data +from tests.helpers.runtime import DiffusionResponse, OmniServer, OnlineOmniClient, dummy_messages_from_mix_data from tests.model_tests.diffusion.config_types import DiffusionTasks from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -48,7 +52,7 @@ IMAGE_GEN_EXTRA_BODY = { } # Online form_data for video generation requests (multipart /v1/videos API) -VIDEO_GEN_FORM_DATA = { +VIDEO_GEN_FORM_DATA: dict[str, Any] = { "height": HEIGHT, "width": WIDTH, "num_inference_steps": 4, @@ -171,7 +175,7 @@ def _build_online_image_data_url() -> str: def _run_online_t2i( - server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None + server: OmniServer, client: OnlineOmniClient, extra_body: dict | None = None ) -> list[DiffusionResponse]: """Run a text to image request through the server.""" messages = dummy_messages_from_mix_data(content_text=PROMPT) @@ -184,7 +188,7 @@ def _run_online_t2i( def _run_online_i2i( - server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None + server: OmniServer, client: OnlineOmniClient, extra_body: dict | None = None ) -> list[DiffusionResponse]: """Run an image to image request through the server.""" image_data_url = _build_online_image_data_url() @@ -278,7 +282,7 @@ def run_and_validate_multi_output(omni: Omni, task_type: DiffusionTasks): def _run_online_t2v( - server: OmniServer, client: OpenAIClientHandler, form_data: dict | None = None + server: OmniServer, client: OnlineOmniClient, form_data: dict | None = None ) -> list[DiffusionResponse]: """Run a text to video request through the server's /v1/videos API.""" data = dict(form_data or VIDEO_GEN_FORM_DATA) @@ -287,7 +291,7 @@ def _run_online_t2v( return client.send_video_diffusion_request({"form_data": data}) -def _run_online_i2v(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]: +def _run_online_i2v(server: OmniServer, client: OnlineOmniClient) -> list[DiffusionResponse]: """Run an image to video request through the server's /v1/videos API.""" data = dict(VIDEO_GEN_FORM_DATA) data.setdefault("prompt", PROMPT) @@ -306,17 +310,17 @@ def _get_online_videos(responses: list[DiffusionResponse]) -> list: ### Online task runners -def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler): +def run_and_validate_online_text_to_image_request(server: OmniServer, client: OnlineOmniClient): """Run and validate a text to image request through the server.""" _validate_images(_get_online_images(_run_online_t2i(server, client))) -def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler): +def run_and_validate_online_image_to_image_request(server: OmniServer, client: OnlineOmniClient): """Run and validate an image to image request through the server.""" _validate_images(_get_online_images(_run_online_i2i(server, client))) -def run_and_validate_online_determinism(server: OmniServer, client: OpenAIClientHandler, task_type: DiffusionTasks): +def run_and_validate_online_determinism(server: OmniServer, client: OnlineOmniClient, task_type: DiffusionTasks): """Checks for determinism through the server, dispatching by task type. Video is skipped online only: send_video_diffusion_request returns encoded @@ -342,7 +346,7 @@ def run_and_validate_online_determinism(server: OmniServer, client: OpenAIClient raise ValueError(f"Task type {task_type} is not yet supported for determinism checks") -def run_and_validate_online_multi_output(server: OmniServer, client: OpenAIClientHandler, task_type: DiffusionTasks): +def run_and_validate_online_multi_output(server: OmniServer, client: OnlineOmniClient, task_type: DiffusionTasks): """Checks for multi-output through the server, dispatching by task type. Video is skipped online only: send_video_diffusion_request explicitly @@ -368,11 +372,11 @@ def run_and_validate_online_multi_output(server: OmniServer, client: OpenAIClien raise ValueError(f"Task type {task_type} is not yet supported for multi-output checks") -def run_and_validate_online_text_to_video_request(server: OmniServer, client: OpenAIClientHandler): +def run_and_validate_online_text_to_video_request(server: OmniServer, client: OnlineOmniClient): """Run and validate a text to video request through the server.""" _get_online_videos(_run_online_t2v(server, client)) -def run_and_validate_online_image_to_video_request(server: OmniServer, client: OpenAIClientHandler): +def run_and_validate_online_image_to_video_request(server: OmniServer, client: OnlineOmniClient): """Run and validate an image to video request through the server.""" _get_online_videos(_run_online_i2v(server, client)) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index be064f603..f9a5072f9 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -26,7 +26,11 @@ EXCLUDED_MODELS = [ "ZImagePipeline", "OvisImagePipeline", "WanPipeline", + "WanDMDPipeline", "WanVACEPipeline", + "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", + "LTX2DistilledTwoStagePipeline", "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", @@ -37,14 +41,12 @@ EXCLUDED_MODELS = [ "WanS2VPipeline", "WanT2VDMD2Pipeline", "WanI2VDMD2Pipeline", - "LongCatImagePipeline", "LongCatVideoAvatarPipeline", "BagelPipeline", "BooguImagePipeline", "LancePipeline", "MingImagePipeline", "InternVLAA1Pipeline", - "LongCatImageEditPipeline", "StableDiffusion3Pipeline", "HunyuanImage3ForCausalMM", "ErnieImagePipeline", @@ -55,24 +57,31 @@ EXCLUDED_MODELS = [ "OmniGen2Pipeline", "HeliosPipeline", "HeliosPyramidPipeline", - "DreamIDOmniPipeline", "SenseNovaU1Pipeline", - "AudioXPipeline", "HunyuanVideo15Pipeline", "HunyuanVideo15ImageToVideoPipeline", "LingBotVideoPipeline", - "MagiHumanPipeline", + # The common harness does not yet define an image-to-video task or supply + # an input image. SANA I2V is covered by + # tests/diffusion/models/sana_video/test_pipeline_sana_video.py and + # tests/e2e/online_serving/test_sana_video_expansion.py. SANA T2V is exercised through + # the common tiny-model setting in model_settings.py. + "SanaImageToVideoPipeline", + # Requires camera/action assets and persistent AR state; covered by its + # dedicated contract tests and real-checkpoint E2E smoke instead. + "LingBotWorldCausalDMDPipeline", "OmniVoicePipeline", "OmniVoice", "Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline", "DiffusersAdapterPipeline", "HiDreamImagePipeline", + "HiDreamO1ImagePipeline", "DreamZeroPipeline", "StableDiffusionXLPipeline", "Gr00tN1d7Pipeline", - "SoulXSingerPipeline", - "SoulXSingerSVCPipeline", + "Pi0Pipeline", + "SanaWmPipeline", ] diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py index b0eb26c9b..884710f2e 100644 --- a/tests/model_tests/diffusion/test_common_offline.py +++ b/tests/model_tests/diffusion/test_common_offline.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest from tests.model_tests.diffusion.case_filtering import get_parametrized_options @@ -17,7 +20,7 @@ from tests.model_tests.diffusion.task_runners import ( ) # NOTE : Hardware marks are added dynamically based on test requirements -pytestmark = [pytest.mark.diffusion] +pytestmark = [pytest.mark.diffusion, pytest.mark.xdist] @pytest.mark.parametrize( @@ -50,7 +53,7 @@ def test_pipeline_on_supported_tasks( ) try: for task_type in supported_tasks: - with subtests.test(msg=task_type): + with subtests.test(msg=task_type.value): if task_type == DiffusionTasks.TEXT_TO_IMAGE: run_and_validate_text_to_image_request(omni) elif task_type == DiffusionTasks.IMAGE_TO_IMAGE: diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py index 464adc13a..9c7ef8a0a 100644 --- a/tests/model_tests/diffusion/test_common_online.py +++ b/tests/model_tests/diffusion/test_common_online.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Analogous to test_common_offline, but for server tests. Validates the full online serving stack (CLI arg parsing, subprocess, API routing, response @@ -6,7 +9,7 @@ encoding) using tiny models. import pytest -from tests.helpers.runtime import OmniServer, OpenAIClientHandler +from tests.helpers.runtime import OmniServer, OnlineOmniClient from tests.model_tests.diffusion.case_filtering import get_parametrized_options from tests.model_tests.diffusion.config_types import ( DiffusionAccs, @@ -24,7 +27,7 @@ from tests.model_tests.diffusion.task_runners import ( ) # NOTE : Hardware marks are added dynamically based on test requirements -pytestmark = [pytest.mark.diffusion] +pytestmark = [pytest.mark.diffusion, pytest.mark.xdist] @pytest.mark.parametrize( @@ -50,7 +53,7 @@ def test_online_on_supported_tasks( # TODO: We may want to revisit run_level validation here, # because checks for things like image size etc should not # depend on whether or not the weights are real or random - client = OpenAIClientHandler( + client = OnlineOmniClient( host=server.host, port=server.port, api_key="EMPTY", @@ -58,7 +61,7 @@ def test_online_on_supported_tasks( log_stats=server.log_stats, ) for task_type in supported_tasks: - with subtests.test(msg=task_type): + with subtests.test(msg=task_type.value): if task_type == DiffusionTasks.TEXT_TO_IMAGE: run_and_validate_online_text_to_image_request(server, client) elif task_type == DiffusionTasks.IMAGE_TO_IMAGE: diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py index 99dbb4650..08cbb2adc 100644 --- a/tests/platforms/npu/test_310p_patches.py +++ b/tests/platforms/npu/test_310p_patches.py @@ -414,6 +414,108 @@ def test_qwen3_tts_code_predictor_forward_uses_projected_embedding_and_sampling( assert sample_calls == [{0: generator}] +def test_310p_attention_forward_runs_fused_qkv_runtime(monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise the mocked 310P fused-QKV path at runtime on CPU. + + The reviewer explicitly asked for more than source inspection: this test now + instantiates the patched attention, runs ``prepare_qkv_weights()``, verifies + the packed weight/bias cache, and drives the mocked 310P flash-attention op. + """ + module, _ = _load_qwen3_tts_patch(monkeypatch) + + trans_calls = [] + rotary_calls = [] + flash_calls = [] + + def maybe_trans_nz(weight: torch.Tensor) -> torch.Tensor: + trans_calls.append(weight.detach().clone()) + return weight + 1.0 + + def npu_rotary_mul(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + rotary_calls.append((hidden_states.detach().clone(), cos.detach().clone(), sin.detach().clone())) + return hidden_states + cos - sin + + def npu_flash_attention(**kwargs) -> None: + flash_calls.append( + {key: value.detach().clone() if torch.is_tensor(value) else value for key, value in kwargs.items()} + ) + kwargs["out"].copy_((kwargs["query"] + kwargs["key"] + kwargs["value"]).to(kwargs["out"].dtype)) + + monkeypatch.setattr(module, "maybe_trans_nz", maybe_trans_nz) + monkeypatch.setattr(module, "aligned_16", lambda tensor: tensor) + monkeypatch.setattr(module.torch_npu, "npu_rotary_mul", npu_rotary_mul, raising=False) + monkeypatch.setattr(module.torch_npu, "_npu_flash_attention", npu_flash_attention, raising=False) + + config = SimpleNamespace( + num_attention_heads=2, + num_key_value_heads=2, + head_dim=2, + hidden_size=4, + num_code_groups=2, + attention_bias=True, + rms_norm_eps=1e-6, + ) + attention = module._Qwen3CodePredictorAttention310P(config) + attention.num_heads = 2 + attention.num_kv_heads = 2 + attention.head_dim = 2 + attention.scaling = 0.5 + attention._q_size = attention.num_heads * attention.head_dim + attention._kv_size = attention.num_kv_heads * attention.head_dim + attention.q_norm = torch.nn.Identity() + attention.k_norm = torch.nn.Identity() + attention.o_proj = torch.nn.Identity() + attention.qkv_proj = torch.nn.Linear(4, attention._q_size + 2 * attention._kv_size, bias=True) + with torch.no_grad(): + attention.qkv_proj.weight.copy_(torch.arange(48, dtype=torch.float32).reshape(12, 4) / 10) + attention.qkv_proj.bias.copy_(torch.arange(12, dtype=torch.float32) / 10) + attention._split_qkv = lambda qkv: qkv.split([attention._q_size, attention._kv_size, attention._kv_size], dim=-1) + + attention.prepare_qkv_weights() + + assert len(trans_calls) == 1 + assert trans_calls[0].is_contiguous() + expected_fused_weight = attention.qkv_proj.weight.detach().clone().contiguous() + 1.0 + torch.testing.assert_close(attention._fused_qkv_weight, expected_fused_weight) + assert attention._fused_qkv_bias is attention.qkv_proj.bias + + hidden_states = (torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) - 3) / 10 + cos = torch.full((1, 3, 2), 0.25, dtype=torch.float32) + sin = torch.full((1, 3, 2), -0.5, dtype=torch.float32) + attention_mask = torch.arange(9, dtype=torch.float32).reshape(3, 3) + + output = attention(hidden_states, (cos, sin), attention_mask=attention_mask) + + expected_qkv = torch.nn.functional.linear(hidden_states, expected_fused_weight, attention.qkv_proj.bias) + expected_q_raw, expected_k_raw, expected_v_raw = expected_qkv.split( + [attention._q_size, attention._kv_size, attention._kv_size], + dim=-1, + ) + expected_q = expected_q_raw.view(1, 3, attention.num_heads, attention.head_dim).transpose(1, 2) + expected_k = expected_k_raw.view(1, 3, attention.num_kv_heads, attention.head_dim).transpose(1, 2) + expected_v = expected_v_raw.view(1, 3, attention.num_kv_heads, attention.head_dim).transpose(1, 2) + expected_cos = cos.unsqueeze(1) + expected_sin = sin.unsqueeze(1) + expected_q = expected_q + expected_cos - expected_sin + expected_k = expected_k + expected_cos - expected_sin + expected_q_f = expected_q.transpose(1, 2).reshape(3, attention.num_heads, attention.head_dim) + expected_k_f = expected_k.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim) + expected_v_f = expected_v.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim) + expected_flash = (expected_q_f + expected_k_f + expected_v_f).to(torch.float16).to(torch.float32) + + assert len(rotary_calls) == 2 + assert len(flash_calls) == 1 + torch.testing.assert_close(flash_calls[0]["query"], expected_q_f) + torch.testing.assert_close(flash_calls[0]["key"], expected_k_f) + torch.testing.assert_close(flash_calls[0]["value"], expected_v_f) + torch.testing.assert_close(flash_calls[0]["mask"], attention_mask) + torch.testing.assert_close(flash_calls[0]["seq_len"], torch.tensor([3], dtype=torch.int32)) + assert flash_calls[0]["scale_value"] == pytest.approx(0.5) + assert flash_calls[0]["num_heads"] == 2 + assert flash_calls[0]["num_kv_heads"] == 2 + torch.testing.assert_close(output, expected_flash.reshape(1, 3, 4)) + + def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.MonkeyPatch) -> None: module, _ = _load_qwen3_tts_patch(monkeypatch) talker = module._Qwen3TTSTalker310P(vllm_config=object()) diff --git a/tests/platforms/npu/test_diffusion_attn_backend_selector.py b/tests/platforms/npu/test_diffusion_attn_backend_selector.py new file mode 100644 index 000000000..1136ef2e3 --- /dev/null +++ b/tests/platforms/npu/test_diffusion_attn_backend_selector.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the NPU diffusion attention backend selector. + +Regression coverage for the eager ``mindiesd`` import in +``NPUOmniPlatform.get_diffusion_attn_backend_cls``: the import must run only +for backends that reach mindiesd kernels (FLASH_ATTN, and RAINFUSION_ATTN via +its dense FlashAttention fallback), so that a broken optional mindiesd +install cannot block backends that never use it (e.g. TORCH_SDPA). + +The platform module is loaded from source with fake vllm/vllm-ascend +dependencies, so the tests run on CPU without NPU or mindiesd installed. +""" + +from __future__ import annotations + +import builtins +import importlib.util +import sys +import types +from enum import Enum +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_BACKEND_NAMES = ( + "FLASH_ATTN", + "FLASH_ATTN_HUB", + "FLASH_ATTN_3_HUB", + "TORCH_SDPA", + "RAINFUSION_ATTN", +) + + +def _repo_root() -> Path: + marker = Path("vllm_omni") / "platforms" / "npu" / "platform.py" + for parent in Path(__file__).resolve().parents: + if (parent / marker).is_file(): + return parent + raise FileNotFoundError(f"could not locate repo root containing {marker}") + + +def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _make_fake_backend_enum(): + members = { + name: SimpleNamespace(get_path=lambda n=name: f"fake.backends.{n.lower()}.Backend") for name in _BACKEND_NAMES + } + + class _FakeEnumMeta(type): + def __getitem__(cls, name): + return members[name] + + def __getattr__(cls, name): + try: + return members[name] + except KeyError: + raise AttributeError(name) from None + + class FakeDiffusionAttentionBackendEnum(metaclass=_FakeEnumMeta): + pass + + return FakeDiffusionAttentionBackendEnum + + +class _FakeOmniPlatformEnum(Enum): + CUDA = "cuda" + NPU = "npu" + + +class _FakeNPUPlatform: + pass + + +class _FakeOmniPlatform: + _omni_enum = _FakeOmniPlatformEnum.NPU + + @classmethod + def validate_diffusion_attn_backend(cls, selected_backend: str) -> None: + pass + + +def _load_platform_module(monkeypatch: pytest.MonkeyPatch): + """Load vllm_omni/platforms/npu/platform.py with fake dependencies.""" + _install_fake_module(monkeypatch, "vllm") + _install_fake_module( + monkeypatch, + "vllm.logger", + init_logger=lambda name: SimpleNamespace( + debug=lambda *a, **k: None, + info=lambda *a, **k: None, + warning=lambda *a, **k: None, + error=lambda *a, **k: None, + ), + ) + _install_fake_module(monkeypatch, "vllm_ascend") + _install_fake_module(monkeypatch, "vllm_ascend.platform", NPUPlatform=_FakeNPUPlatform) + for pkg in ( + "vllm_omni", + "vllm_omni.diffusion", + "vllm_omni.diffusion.attention", + "vllm_omni.diffusion.attention.backends", + "vllm_omni.platforms", + ): + _install_fake_module(monkeypatch, pkg) + _install_fake_module( + monkeypatch, + "vllm_omni.diffusion.attention.backends.registry", + DiffusionAttentionBackendEnum=_make_fake_backend_enum(), + ) + _install_fake_module( + monkeypatch, + "vllm_omni.platforms.interface", + OmniPlatform=_FakeOmniPlatform, + OmniPlatformEnum=_FakeOmniPlatformEnum, + ) + + path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "platform.py" + spec = importlib.util.spec_from_file_location("npu_platform_under_test", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load import spec for {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.NPUOmniPlatform + + +@pytest.fixture +def mindiesd_tracker(monkeypatch: pytest.MonkeyPatch): + """Make mindiesd importable and record eager import attempts.""" + monkeypatch.delitem(sys.modules, "mindiesd", raising=False) + attempts = [] + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "mindiesd": + attempts.append(name) + return types.ModuleType("mindiesd") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + monkeypatch.setattr( + "importlib.util.find_spec", + lambda name, *a, **k: object() if name == "mindiesd" else None, + ) + return attempts + + +@pytest.mark.parametrize( + ("selected_backend", "expect_mindiesd_import"), + [ + ("FLASH_ATTN", True), + # HF-hub variants fall back to local FLASH_ATTN on NPU. + ("FLASH_ATTN_HUB", True), + ("FLASH_ATTN_3_HUB", True), + ("RAINFUSION_ATTN", True), + # Backends that never reach mindiesd must not trigger the import. + ("TORCH_SDPA", False), + ], +) +def test_eager_mindiesd_import_scoped_to_mindiesd_backends( + monkeypatch: pytest.MonkeyPatch, + mindiesd_tracker: list[str], + selected_backend: str, + expect_mindiesd_import: bool, +): + platform_cls = _load_platform_module(monkeypatch) + + path = platform_cls.get_diffusion_attn_backend_cls(selected_backend, head_size=64) + + assert (len(mindiesd_tracker) > 0) is expect_mindiesd_import + expected = "FLASH_ATTN" if selected_backend in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB") else selected_backend + assert path == f"fake.backends.{expected.lower()}.Backend" + + +def test_explicit_flash_attn_without_mindiesd_still_resolves( + monkeypatch: pytest.MonkeyPatch, +): + platform_cls = _load_platform_module(monkeypatch) + monkeypatch.delitem(sys.modules, "mindiesd", raising=False) + monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None) + + path = platform_cls.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64) + + assert path == "fake.backends.flash_attn.Backend" + + +def test_default_backend_prefers_flash_attn_when_mindiesd_available( + monkeypatch: pytest.MonkeyPatch, + mindiesd_tracker: list[str], +): + platform_cls = _load_platform_module(monkeypatch) + + path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64) + + assert path == "fake.backends.flash_attn.Backend" + assert len(mindiesd_tracker) > 0 + + +def test_default_backend_falls_back_to_sdpa_without_mindiesd( + monkeypatch: pytest.MonkeyPatch, +): + platform_cls = _load_platform_module(monkeypatch) + monkeypatch.delitem(sys.modules, "mindiesd", raising=False) + monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None) + + path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64) + + assert path == "fake.backends.torch_sdpa.Backend" diff --git a/tests/platforms/npu/test_diffusion_platform.py b/tests/platforms/npu/test_diffusion_platform.py new file mode 100644 index 000000000..10e43e034 --- /dev/null +++ b/tests/platforms/npu/test_diffusion_platform.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib.util +from types import SimpleNamespace + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("available", [False, True]) +def test_dense_flash_capability_tracks_mindiesd(monkeypatch, available): + pytest.importorskip("vllm_ascend") + from vllm_omni.platforms.npu.platform import NPUOmniPlatform + + monkeypatch.setattr( + importlib.util, + "find_spec", + lambda name: object() if name == "mindiesd" and available else None, + ) + + assert NPUOmniPlatform.supports_diffusion_dense_flash_attention() is available + + +def test_paged_config_uses_ascend_kernel_block_size() -> None: + pytest.importorskip("vllm_ascend") + from vllm_ascend.attention.attention_v1 import AscendAttentionBackend + + from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode + from vllm_omni.platforms.npu.platform import NPUOmniPlatform + + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16)) + NPUOmniPlatform.configure_diffusion_vllm_config( + vllm_config, + SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER), + ) + + assert vllm_config.cache_config.block_size == AscendAttentionBackend.get_supported_kernel_block_sizes()[0] + + +def test_strict_ulysses_paged_backend_bypasses_ascend_pcp() -> None: + pytest.importorskip("vllm_ascend") + from vllm_ascend.attention.attention_v1 import ( + AscendAttentionBackend, + AscendAttentionBackendImpl, + AscendAttentionMetadataBuilder, + ) + + from vllm_omni.platforms.npu.platform import NPUOmniPlatform + + assert ( + NPUOmniPlatform.get_diffusion_paged_kv_attn_backend( + AscendAttentionBackend, + ulysses_degree=1, + ) + is AscendAttentionBackend + ) + backend = NPUOmniPlatform.get_diffusion_paged_kv_attn_backend( + AscendAttentionBackend, + ulysses_degree=2, + ) + assert backend.get_impl_cls() is AscendAttentionBackendImpl + assert backend.get_builder_cls() is AscendAttentionMetadataBuilder diff --git a/tests/platforms/npu/test_graph_tools.py b/tests/platforms/npu/test_graph_tools.py new file mode 100644 index 000000000..703b8f630 --- /dev/null +++ b/tests/platforms/npu/test_graph_tools.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.platforms.npu.graph_tools import ( + CapturedDeviceGraph, + NPUExactGraphRunner, +) +from vllm_omni.platforms.npu.models import minicpmo_4_5_code2wav as code2wav_patch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_captured_graph_replay_clones_persistent_outputs(): + static_input = torch.zeros(1) + static_output = torch.zeros(1) + + class _Graph: + def replay(self): + static_output.copy_(static_input * 2) + + graph = CapturedDeviceGraph( + graph=_Graph(), + static_inputs=(static_input,), + static_outputs=(static_output,), + ) + + first = graph.replay((torch.tensor([2.0]),))[0] + second = graph.replay((torch.tensor([3.0]),))[0] + + torch.testing.assert_close(first, torch.tensor([4.0])) + torch.testing.assert_close(second, torch.tensor([6.0])) + assert first.data_ptr() != second.data_ptr() + + +def test_capture_uses_vllm_global_graph_pool(monkeypatch): + runner = NPUExactGraphRunner() + global_pool = object() + seen_pools = [] + synchronizations = 0 + + class _Graph: + def replay(self): + pass + + @contextmanager + def graph_context(graph, *, pool): + del graph + seen_pools.append(pool) + yield + + def synchronize(): + nonlocal synchronizations + synchronizations += 1 + + fake_npu = SimpleNamespace( + NPUGraph=_Graph, + graph=graph_context, + synchronize=synchronize, + ) + monkeypatch.setattr(torch, "npu", fake_npu, raising=False) + + import vllm.platforms + + monkeypatch.setattr( + vllm.platforms, + "current_platform", + SimpleNamespace(get_global_graph_pool=lambda: global_pool), + ) + + captured = runner.capture( + (torch.tensor([2.0]),), + lambda value: (value + 1,), + ) + + assert seen_pools == [global_pool] + assert synchronizations == 2 + torch.testing.assert_close(captured.static_outputs[0], torch.tensor([3.0])) + + +def test_exact_signature_dispatch_captures_then_replays(monkeypatch): + runner = NPUExactGraphRunner(max_graphs=2) + replay_count = 0 + + class _FunctionalGraph: + def __init__(self, compute): + self.compute = compute + + def replay(self, inputs): + nonlocal replay_count + replay_count += 1 + return tuple(output.detach().clone() for output in self.compute(*inputs)) + + monkeypatch.setattr(runner, "_eligible", lambda inputs: True) + monkeypatch.setattr( + runner, + "capture", + lambda inputs, compute: _FunctionalGraph(compute), + ) + + def compute(value): + return (value * 2,) + + first = runner.run("unit", (torch.tensor([2.0]),), (False,), compute) + second = runner.run("unit", (torch.tensor([3.0]),), (False,), compute) + + torch.testing.assert_close(first[0], torch.tensor([4.0])) + torch.testing.assert_close(second[0], torch.tensor([6.0])) + assert runner.stats == {"captures": 1, "failed": 0, "hits": 1} + assert replay_count == 1 + + +def test_capture_failure_is_fatal_for_stage_process(monkeypatch): + runner = NPUExactGraphRunner( + component_name="test component", + disable_config_hint="disable the test graph", + ) + captures = 0 + computes = 0 + + monkeypatch.setattr(runner, "_eligible", lambda inputs: True) + + def fail_capture(inputs, compute): + nonlocal captures + captures += 1 + raise RuntimeError("unsupported op") + + monkeypatch.setattr(runner, "capture", fail_capture) + + def compute(value): + nonlocal computes + computes += 1 + return (value + 1,) + + with pytest.raises(RuntimeError, match="restart the stage process"): + runner.run("unit", (torch.tensor([1.0]),), (False,), compute) + with pytest.raises(RuntimeError, match="cannot continue"): + runner.run( + "different-shape", + (torch.tensor([1.0, 2.0]),), + (True,), + compute, + ) + + assert captures == 1 + assert computes == 1 + assert runner.stats == {"captures": 0, "failed": 1, "hits": 0} + + +def test_code2wav_runtime_disables_internal_format_and_jit(monkeypatch): + monkeypatch.delenv("ASCEND_LAUNCH_BLOCKING", raising=False) + compile_modes = [] + fake_npu = SimpleNamespace( + config=SimpleNamespace(allow_internal_format=True), + set_compile_mode=lambda **kwargs: compile_modes.append(kwargs), + ) + monkeypatch.setattr(torch, "npu", fake_npu, raising=False) + + code2wav_patch.prepare_code2wav_graph_runtime() + + assert fake_npu.config.allow_internal_format is False + assert compile_modes == [{"jit_compile": False}] + + +def test_code2wav_runtime_rejects_launch_blocking(monkeypatch): + monkeypatch.setenv("ASCEND_LAUNCH_BLOCKING", "1") + + with pytest.raises(RuntimeError, match="ASCEND_LAUNCH_BLOCKING=1"): + code2wav_patch.prepare_code2wav_graph_runtime() + + +@pytest.mark.parametrize( + ("enabled", "max_graphs", "expected_prepared"), + [ + (True, 7, 1), + (False, 7, 0), + (True, 0, 0), + ], +) +def test_code2wav_patch_reads_stage_additional_config( + monkeypatch, + enabled, + max_graphs, + expected_prepared, +): + prepared = 0 + + class _Backend: + speech_window = torch.zeros(1) + flow = SimpleNamespace(training=False) + + model = SimpleNamespace( + backend=None, + _extra_config=lambda: {}, + vllm_config=SimpleNamespace( + additional_config={ + "code2wav_enable_npu_graph": enabled, + "code2wav_max_npu_graphs": max_graphs, + } + ), + ) + + def build_backend(instance): + instance.backend = _Backend() + + def prepare_runtime(): + nonlocal prepared + prepared += 1 + + monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend) + monkeypatch.setattr(code2wav_patch, "prepare_code2wav_graph_runtime", prepare_runtime) + code2wav_patch._patched_build_backend(model) + + assert prepared == expected_prepared + if expected_prepared: + graph_runner = code2wav_patch._backend_graph_runners[model.backend] + assert isinstance(graph_runner, NPUExactGraphRunner) + assert graph_runner.max_graphs == max_graphs + else: + assert model.backend not in code2wav_patch._backend_graph_runners + + +@pytest.mark.parametrize("value", [True, "true", "false", "0", "off"]) +def test_code2wav_patch_rejects_bfloat16_attention_cache(monkeypatch, value): + built = False + + model = SimpleNamespace( + backend=None, + _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value}, + ) + + def build_backend(_instance): + nonlocal built + built = True + + monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend) + + with pytest.raises(ValueError, match="currently supported only on CUDA"): + code2wav_patch._patched_build_backend(model) + + assert not built + + +@pytest.mark.parametrize("value", [False, None, 0, ""]) +def test_code2wav_patch_allows_disabled_bfloat16_attention_cache(monkeypatch, value): + built = False + + model = SimpleNamespace( + backend=None, + _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value}, + ) + + def build_backend(instance): + nonlocal built + built = True + instance.backend = object() + + monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend) + code2wav_patch._patched_build_backend(model) + + assert built + + +@pytest.mark.parametrize("with_cache", [False, True]) +def test_code2wav_estimator_graph_dispatch_is_platform_owned(monkeypatch, with_cache): + calls = [] + + class _Backend: + _trt_stepper = None + _cfm_graph_wrapper = None + + class _Runner: + def run(self, operation, inputs, constants, compute): + calls.append((operation, inputs, constants)) + return compute(*inputs) + + backend = _Backend() + runner = _Runner() + code2wav_patch._backend_graph_runners[backend] = runner + + def graphable_estimator_step( + instance, + estimator, + *, + x, + mu, + time_embedding, + speakers, + cond, + cnn_cache, + att_cache, + ): + del instance, estimator, time_embedding, speakers, cond + cache_marker = x.new_tensor(0 if cnn_cache is None else 1) + return x + mu, cache_marker, cache_marker.clone() + + monkeypatch.setattr(code2wav_patch, "_original_estimator_step", lambda *args, **kwargs: None) + monkeypatch.setattr(code2wav_patch, "_graphable_estimator_step", graphable_estimator_step) + value = torch.tensor([2.0]) + cache = torch.tensor([1.0]) if with_cache else None + estimator = SimpleNamespace(t_embedder=lambda time: time + 1) + outputs = code2wav_patch._patched_estimator_step( + backend, + estimator, + x=value, + mu=value, + time=value, + speakers=value, + cond=value, + cnn_cache=cache, + att_cache=cache, + ) + + assert calls[0][0] == "cfm_estimator" + assert calls[0][2] == (with_cache,) + assert len(calls[0][1]) == (7 if with_cache else 5) + torch.testing.assert_close(calls[0][1][2], torch.tensor([[3.0]])) + torch.testing.assert_close(outputs[0], torch.tensor([4.0])) + torch.testing.assert_close(outputs[1], torch.tensor(1.0 if with_cache else 0.0)) + + +def test_code2wav_platform_wraps_flow_execution_context(monkeypatch): + entered = [] + + class _Backend: + pass + + class _Context: + def __init__(self, device, require_math): + self.device = device + self.require_math = require_math + + def __enter__(self): + entered.append((self.device, self.require_math)) + + def __exit__(self, exc_type, exc, traceback): + return False + + backend = _Backend() + code2wav_patch._backend_graph_runners[backend] = NPUExactGraphRunner() + monkeypatch.setattr( + code2wav_patch, + "_flow_execution_context", + lambda device, *, require_math: _Context(device, require_math), + ) + monkeypatch.setattr( + code2wav_patch, + "_original_setup_batch", + lambda instance, features, batch_size: ("setup", batch_size), + ) + monkeypatch.setattr( + code2wav_patch, + "_original_decode_batch", + lambda instance, tokens, features, states, *, last_chunk, flush_encoder: ( + "decode", + last_chunk, + flush_encoder, + ), + ) + features = SimpleNamespace(speech_tokens=torch.zeros(1)) + + assert code2wav_patch._patched_setup_batch(backend, features, 3) == ("setup", 3) + assert code2wav_patch._patched_decode_batch( + backend, + torch.zeros(1), + features, + [], + last_chunk=True, + ) == ("decode", True, False) + assert entered == [(torch.device("cpu"), True), (torch.device("cpu"), True)] diff --git a/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py new file mode 100644 index 000000000..4ae544380 --- /dev/null +++ b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Ascend hardware validation for MiniCPM-o Code2Wav NPUGraph replay.""" + +from math import prod +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +pytest.importorskip("vllm_ascend") + +from tests.helpers.mark import hardware_marks +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( + BatchedToken2Wav, +) +from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner +from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import ( + _graphable_estimator_step, + prepare_code2wav_graph_runtime, +) +from vllm_omni.platforms.npu.models.step_audio2_token2wav import ( + npu_token2wav_sdpa_context, +) +from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.omni, + *hardware_marks(res={"npu": "A3"}, num_cards=1), +] + + +class _Flow(nn.Module): + def __init__(self, estimator: nn.Module): + super().__init__() + self.encoder = nn.Identity() + self.decoder = SimpleNamespace(estimator=estimator) + + +class _Token2Wav: + def __init__(self, estimator: nn.Module): + self.flow = _Flow(estimator).eval() + self.hift = nn.Identity().eval() + self.float16 = False + self.n_timesteps = 2 + self.mel_cache_len = 1 + self.source_cache_len = 2 + self.speech_window = torch.ones(4, device="npu") + + +def test_npu_runner_inherits_runtime_snapshot_replacement(): + assert issubclass(OmniNPUModelRunner, OmniGPUModelRunner) + assert OmniNPUModelRunner._replace_intermediate_buffer is OmniGPUModelRunner._replace_intermediate_buffer + runner = object.__new__(OmniNPUModelRunner) + request = SimpleNamespace(additional_information_cpu=None) + runner.requests = {"req": request} + runner.model_intermediate_buffer = { + "req": { + "codes": {"audio": torch.tensor([1, 2])}, + "meta": {"last_chunk": True}, + } + } + + runner._replace_intermediate_buffer("req", {"meta": {"is_segment_finished": True}}) + + assert runner.model_intermediate_buffer["req"] == {"meta": {"is_segment_finished": True}} + assert request.additional_information_cpu == runner.model_intermediate_buffer["req"] + + +def test_minicpmo_cfm_estimator_npugraph_matches_eager(): + decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit") + prepare_code2wav_graph_runtime() + + estimator = ( + decoder_dit.DiT( + in_channels=6, + out_channels=2, + depth=2, + num_heads=2, + head_dim=8, + hidden_size=16, + mlp_ratio=2.0, + ) + .npu() + .eval() + ) + graph_runner = NPUExactGraphRunner(max_graphs=4) + adapter = BatchedToken2Wav(_Token2Wav(estimator)) + + def inputs(seed: int): + def make(shape: tuple[int, ...], offset: float): + values = torch.arange( + prod(shape), + device="npu", + dtype=torch.float32, + ) + return torch.sin(values * 0.17 + seed + offset).reshape(shape) + + return ( + make((2, 2, 6), 0.0), + make((2, 2, 6), 0.5), + make((2, 1, 16), 1.0), + make((2, 1), 1.5), + make((2, 1, 6), 2.0), + ) + + def compute(*values, caches=None): + cnn_cache, att_cache = (None, None) if caches is None else caches + return _graphable_estimator_step( + adapter, + estimator, + x=values[0], + mu=values[1], + time_embedding=values[2], + speakers=values[3], + cond=values[4], + cnn_cache=cnn_cache, + att_cache=att_cache, + ) + + with torch.inference_mode(), npu_token2wav_sdpa_context(require_math=True): + warm_inputs = inputs(11) + warm_outputs = graph_runner.run( + "cfm_estimator", + warm_inputs, + (False,), + lambda *values: compute(*values), + ) + replay_inputs = inputs(13) + expected = compute(*replay_inputs) + actual = graph_runner.run( + "cfm_estimator", + replay_inputs, + (False,), + lambda *values: compute(*values), + ) + for eager, replayed in zip(expected, actual, strict=True): + torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3) + + cached_inputs = (*inputs(17), warm_outputs[1], warm_outputs[2]) + graph_runner.run( + "cfm_estimator", + cached_inputs, + (True,), + lambda *values: compute(*values[:5], caches=(values[5], values[6])), + ) + replay_cached_inputs = ( + *inputs(19), + warm_outputs[1] + 0.1, + warm_outputs[2] + 0.1, + ) + expected_cached = compute( + *replay_cached_inputs[:5], + caches=(replay_cached_inputs[5], replay_cached_inputs[6]), + ) + actual_cached = graph_runner.run( + "cfm_estimator", + replay_cached_inputs, + (True,), + lambda *values: compute(*values[:5], caches=(values[5], values[6])), + ) + for eager, replayed in zip(expected_cached, actual_cached, strict=True): + torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3) + assert graph_runner.stats == {"captures": 2, "failed": 0, "hits": 2} diff --git a/tests/platforms/npu/test_step_audio2_token2wav.py b/tests/platforms/npu/test_step_audio2_token2wav.py index ae86f92d9..23baa2c3e 100644 --- a/tests/platforms/npu/test_step_audio2_token2wav.py +++ b/tests/platforms/npu/test_step_audio2_token2wav.py @@ -7,6 +7,7 @@ from __future__ import annotations import importlib.util import sys import types +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -194,3 +195,71 @@ def test_hift_patch_reports_incompatible_layout(monkeypatch: pytest.MonkeyPatch) with pytest.raises(TypeError, match=r"m_source\.l_sin_gen\._f02sine"): module.patch_step_audio2_hift_for_npu(SimpleNamespace()) + + +def test_sdpa_context_preserves_exceptions_from_token2wav_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_patch_module(monkeypatch) + fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn") + fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None + + @contextmanager + def fake_math_context(*, require_available=False): + del require_available + yield + + fake_attention_patch.npu_math_sdpa_context = fake_math_context + monkeypatch.setitem( + sys.modules, + "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn", + fake_attention_patch, + ) + + with pytest.raises(RuntimeError, match="original graph failure"): + with module.npu_token2wav_sdpa_context(): + raise RuntimeError("original graph failure") + + +def test_sdpa_context_requires_math_for_npu_graph(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_patch_module(monkeypatch) + required = [] + fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn") + fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None + + @contextmanager + def fake_math_context(*, require_available=False): + required.append(require_available) + yield + + fake_attention_patch.npu_math_sdpa_context = fake_math_context + monkeypatch.setitem( + sys.modules, + "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn", + fake_attention_patch, + ) + + with module.npu_token2wav_sdpa_context(require_math=True): + pass + + assert required == [True] + + +def test_sdpa_context_does_not_hide_patch_setup_failures(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_patch_module(monkeypatch) + fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn") + + def fail_patch(): + raise RuntimeError("attention patch failed") + + fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = fail_patch + fake_attention_patch.npu_math_sdpa_context = contextmanager(lambda: iter((None,))) + monkeypatch.setitem( + sys.modules, + "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn", + fake_attention_patch, + ) + + with pytest.raises(RuntimeError, match="attention patch failed"): + with module.npu_token2wav_sdpa_context(): + pass diff --git a/tests/platforms/test_record_device_event.py b/tests/platforms/test_record_device_event.py index a979ca940..9b81f9cd4 100644 --- a/tests/platforms/test_record_device_event.py +++ b/tests/platforms/test_record_device_event.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for OmniPlatform.record_device_event implementations.""" import pytest +from tests.helpers.mark import hardware_test + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -38,6 +40,71 @@ class TestCudaOmniPlatformRecordDeviceEvent: assert result is None +class TestXPUOmniPlatformRecordDeviceEvent: + """XPU must record a device-agnostic ``torch.Event``, not ``torch.xpu.Event``. + + The async diffusion output thread gates its D2H copy with + ``torch.Stream.wait_event`` on a generic ``torch.Stream``, and that binding + silently drops the dependency for a ``torch.xpu.Event`` — the copy then races + the still-running forward and the host reads a partially written tensor + (garbage rows at the bottom of the image). + """ + + def test_records_device_agnostic_torch_event(self, mocker): + # vLLM's XPUPlatform imports vllm_xpu_kernels at module scope, and that + # package is absent on non-XPU images. + pytest.importorskip("vllm_xpu_kernels") + + from vllm_omni.platforms.xpu.platform import XPUOmniPlatform + + mock_event = mocker.MagicMock() + mocker.patch("torch.Event", return_value=mock_event) + xpu_event = mocker.patch("torch.xpu.Event") + + result = XPUOmniPlatform.record_device_event() + + assert result is mock_event + mock_event.record.assert_called_once() + xpu_event.assert_not_called() + + @hardware_test(res={"xpu": "B60"}) + def test_recorded_event_orders_side_stream_d2h(self): + """On-device: the event must make a side-stream D2H wait for compute.""" + import torch + + if not torch.xpu.is_available(): + pytest.skip("XPU not available") + + from vllm_omni.platforms.xpu.platform import XPUOmniPlatform + + device = torch.device("xpu", 0) + # Normalized so the matmul chain stays finite; the ``* 0.0`` term keeps + # every element exactly 7.0 while still depending on the whole chain. + base = torch.randn(4096, 4096, device=device) / 64.0 + chained = base.clone() + for _ in range(40): + chained = chained @ base + tensor = torch.full_like(chained, 7.0) + chained * 0.0 + + event = XPUOmniPlatform.record_device_event() + assert event is not None + + # Mirrors WorkerProc._async_output_loop. + d2h_stream = torch.Stream(device=device) + d2h_stream.wait_event(event) + previous = torch.accelerator.current_stream() + torch.accelerator.set_stream(d2h_stream) + try: + host = torch.empty(tensor.shape, dtype=tensor.dtype, pin_memory=True) + host.copy_(tensor, non_blocking=True) + finally: + torch.accelerator.set_stream(previous) + d2h_stream.synchronize() + + copied_early = int((host != 7.0).sum()) + assert copied_early == 0, f"{copied_early} elements copied before compute finished" + + class TestNPUOmniPlatformRecordDeviceEvent: """Test NPUOmniPlatform.record_device_event with mocked torch.npu.""" @@ -55,8 +122,8 @@ class TestNPUOmniPlatformRecordDeviceEvent: result = NPUOmniPlatform.record_device_event() assert result is None - def test_synchronizes_stream_then_records_event(self, mocker): - """NPU should sync stream first, then create and record an Event.""" + def test_synchronizes_stream_then_records_generic_event(self, mocker): + """NPU should return an event consumable by the generic torch.Stream.""" try: from vllm_omni.platforms.npu.platform import NPUOmniPlatform except ModuleNotFoundError: @@ -66,14 +133,16 @@ class TestNPUOmniPlatformRecordDeviceEvent: mock_stream = mocker.MagicMock() mock_torch = mocker.MagicMock() mock_torch.npu.current_stream.return_value = mock_stream - mock_torch.npu.Event.return_value = mock_event + mock_torch.Event.return_value = mock_event mocker.patch("vllm_omni.platforms.npu.platform.torch", mock_torch) result = NPUOmniPlatform.record_device_event() # Stream should be synced first (HCCL ordering) mock_stream.synchronize.assert_called_once() - # Then event should be created and recorded - mock_torch.npu.Event.assert_called_once() + # Then a backend-neutral event should be created and recorded. The + # worker waits on it through the public torch.Stream wrapper. + mock_torch.Event.assert_called_once() + mock_torch.npu.Event.assert_not_called() mock_event.record.assert_called_once() assert result is mock_event diff --git a/tests/tools/test_check_tts_adapter.py b/tests/tools/test_check_tts_adapter.py index 486a629d5..47c6ac89f 100644 --- a/tests/tools/test_check_tts_adapter.py +++ b/tests/tools/test_check_tts_adapter.py @@ -17,9 +17,7 @@ import pytest sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools" / "pre_commit")) from check_tts_adapter import ( # noqa: E402 - MAX_LEGACY_DETECTORS, MAX_MODEL_TYPE_BRANCHES, - _legacy_detectors, _model_type_branches, main, ) @@ -60,12 +58,6 @@ def test_dict_dispatch_is_a_known_gap(): assert _count('TABLE = {"foo": 1}\nvalue = TABLE[self._tts_model_type]') == 0 -def test_legacy_detector_counting(): - src = 'X = (LegacyDetector(name="a", stage_keys=frozenset({"s"})), LegacyDetector(name="b"))' - found = _legacy_detectors(ast.parse(src)) - assert [name for _, name in found] == ["a", "b"] - - def test_gate_passes_on_current_tree(): """The recorded budgets must match reality, or the ratchet is decorative.""" assert main([]) == 0 @@ -74,6 +66,4 @@ def test_gate_passes_on_current_tree(): def test_budgets_are_the_real_counts(): root = Path(__file__).resolve().parents[2] serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py" - adapters = root / "vllm_omni" / "entrypoints" / "openai" / "tts_adapters" / "__init__.py" assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES - assert len(_legacy_detectors(ast.parse(adapters.read_text()))) == MAX_LEGACY_DETECTORS diff --git a/tests/utils/test_speaker_cache.py b/tests/utils/test_speaker_cache.py index 373baacb4..ea50c5881 100644 --- a/tests/utils/test_speaker_cache.py +++ b/tests/utils/test_speaker_cache.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import threading import pytest @@ -144,7 +147,7 @@ class TestSpeakerEmbeddingCacheBehavior: class TestSingleton: - def test_singleton_identity(self, fresh_speaker_cache): + def test_singleton_identity(self, clean_speaker_cache): a = get_speaker_cache() b = get_speaker_cache() assert a is b diff --git a/tests/utils/test_speaker_cache_integration.py b/tests/utils/test_speaker_cache_integration.py index 8a24659f0..01b886725 100644 --- a/tests/utils/test_speaker_cache_integration.py +++ b/tests/utils/test_speaker_cache_integration.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Integration tests for the process-wide speaker cache across serving + models.""" import pytest @@ -26,7 +29,7 @@ class TestSpeakerCacheIntegration: assert cache.get(fish_key) is None assert cache.get(bob_key) is not None - def test_singleton_shared_across_call_sites(self, fresh_speaker_cache): + def test_singleton_shared_across_call_sites(self, clean_speaker_cache): cache_a = get_speaker_cache() cache_b = get_speaker_cache() assert cache_a is cache_b diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py index f00da2797..443ec0138 100644 --- a/tests/utils/test_tracking_parser.py +++ b/tests/utils/test_tracking_parser.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Tests for TrackingArgumentParser and related utilities.""" import argparse @@ -35,7 +38,7 @@ _TEST_PIPELINE = PipelineConfig( stages=( StagePipelineConfig(stage_id=0, model_stage="stage0"), StagePipelineConfig(stage_id=1, model_stage="stage1"), - StagePipelineConfig(stage_id=2, model_stage="stage2"), + StagePipelineConfig(stage_id=2, model_stage="stage2", final_output=True), ), ) diff --git a/tests/worker/test_capture_replay_contract.py b/tests/worker/test_capture_replay_contract.py index 504aaa7b0..a0ca095f3 100644 --- a/tests/worker/test_capture_replay_contract.py +++ b/tests/worker/test_capture_replay_contract.py @@ -14,9 +14,9 @@ cannot be masked by an adjacent ``else``/``elif`` that happens to read it. It matches by attribute chain (``self.input_ids.gpu``), so local variable renames do not break it. -Scope: the BASE runner (``OmniGPUModelRunner``) only. The generation runner is -KNOWN to lack the ``has_preprocess`` branch; extend this -check to the generation file when it adds the branch. +Scope: the BASE runner (``OmniGPUModelRunner``) and the generation runner's +``_dummy_run`` override (which gained the branch in RFC #5450 C6; its +``_preprocess`` is inherited from the base). """ import ast @@ -25,6 +25,7 @@ import textwrap import pytest +from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -69,11 +70,24 @@ def _has_preprocess_branch_reads_both_buffers(method) -> bool: return any(_REQUIRED_BUFFERS <= _attr_chains_in(node.body) for node in guarded) -@pytest.mark.parametrize("method_name", ["_dummy_run", "_preprocess"]) -def test_has_preprocess_branch_reads_shared_buffers(method_name): - method = getattr(OmniGPUModelRunner, method_name) +@pytest.mark.parametrize( + ("runner_cls", "method_name"), + [ + (OmniGPUModelRunner, "_dummy_run"), + (OmniGPUModelRunner, "_preprocess"), + (GPUGenerationModelRunner, "_dummy_run"), + ], +) +def test_has_preprocess_branch_reads_shared_buffers(runner_cls, method_name): + method = getattr(runner_cls, method_name) assert _has_preprocess_branch_reads_both_buffers(method), ( - f"{OmniGPUModelRunner.__name__}.{method_name}: the has_preprocess branch no " + f"{runner_cls.__name__}.{method_name}: the has_preprocess branch no " "longer reads BOTH self.input_ids.gpu and self.inputs_embeds.gpu — capture " "and replay would diverge." ) + + +def test_generation_dummy_run_overrides_with_branch(): + # Guard against the override being deleted (which would silently fall back + # to the base implementation and vacuously pass the parametrized check). + assert "_dummy_run" in GPUGenerationModelRunner.__dict__ diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 116b5ee4d..5e3cfe865 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -1,7 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import copy +import re from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier from types import SimpleNamespace +from typing import Any +from unittest.mock import patch import numpy as np import pytest @@ -14,12 +22,15 @@ from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechReques from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker import sparse_audio from vllm_omni.worker.gpu_ar_model_runner import ( ExecuteModelState, GPUARModelRunner, OmniAsyncGPUModelRunnerOutput, ) +from vllm_omni.worker.output import payload_build from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest +from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -200,6 +211,15 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc async def build(self, request, sampling_params_list, has_inline_ref_audio): return PreparedRequest(prompt={"prompt": request.input}, tts_params={}, model_type="higgs_audio_v3") + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + return sampling_params_list + engine = SimpleNamespace( errored=False, default_sampling_params_list=[SamplingParams(temperature=1.0, top_p=0.95, top_k=50)], @@ -291,11 +311,13 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests( def test_sparse_mm_req_ids_requires_sparse_audio_marker(): - assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None - assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None + assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None - assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"] - assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"sparse_audio": ["1"]}, "meta.req_id": ["r1"]}) == ["r1"] def test_runner_assisted_full_attention_metadata_request_is_opt_in(): @@ -875,6 +897,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk lambda *args, **kwargs: ( 0, None, + None, [], None, ["r1", "r2"], @@ -1188,3 +1211,643 @@ def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_dow assert torch.equal(output.inter_stage_outputs[0]["codes.ref"], ref_codes[0]) assert torch.equal(output.inter_stage_outputs[2]["codes.ref"], ref_codes[2]) + + +class TestPromptIdsClampPaddingConvention: + """Pin `sampling_utils.clamp_prompt_ids_to_penalty_padding` (the clamp + `sample_tokens` applies before penalties). + + The audit (RFC #5450 C2) originally flagged the clamp as an off-by-one + (OOB index). It is not: upstream penalties allocate `vocab_size + 1` bins + and drop the last column, so `vocab_size` is the designed padding value + (vllm/model_executor/layers/utils.py::get_token_bin_counts_and_mask). + These tests call the runner's own helper, so a future "fix" to the clamp + fails here. + """ + + LOGITS_VOCAB = 8 + BATCH_VOCAB = 10 # wider input_batch vocab; its pad value is BATCH_VOCAB + + def _padded_prompt(self): + # Two real tokens + two padding slots padded with the batch vocab size. + return torch.tensor([[1, 2, self.BATCH_VOCAB, self.BATCH_VOCAB]]) + + def test_clamp_to_logits_vocab_keeps_padding_inert(self): + from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask + + # The runner's own helper, not a re-implementation. + clamped = clamp_prompt_ids_to_penalty_padding(self._padded_prompt(), self.LOGITS_VOCAB) + + _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1) + + expected = torch.zeros(1, self.LOGITS_VOCAB, dtype=torch.bool) + expected[0, 1] = True + expected[0, 2] = True + # Pad ids landed in the dropped (vocab_size-th) bin: no penalty effect. + assert torch.equal(mask, expected) + + def test_clamp_one_lower_would_corrupt_last_token_penalties(self): + from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask + + # The "fix" the audit first proposed - shown here to be the actual bug: + # padding would count as occurrences of the last real vocab token. + clamped = self._padded_prompt().clamp(max=self.LOGITS_VOCAB - 1) + + _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1) + + assert bool(mask[0, self.LOGITS_VOCAB - 1]) + + def test_unclamped_padding_overflows_penalty_bins(self): + from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask + + # Why the clamp exists: the batch-level pad value (BATCH_VOCAB) is out + # of range even for the vocab_size+1 penalty bins of a narrower head. + with pytest.raises(RuntimeError): + get_token_bin_counts_and_mask(self._padded_prompt(), self.LOGITS_VOCAB, 1) + + +class TestSparseAudioMarkerRobustness: + """RFC #5450 C3: `sparse_audio.is_sparse_audio_marker` must not crash on tensor/array + markers, and the combined prefix-cache payload builder must never ship one + request's data as another's.""" + + def test_marker_list_and_str_forms(self): + marker = sparse_audio.is_sparse_audio_marker + assert marker(["1"]) is True + assert marker(["0"]) is False + assert marker("true") is True + assert marker("off") is False + + def test_marker_enforced_canonical_forms_no_error(self): + # The designed carrier (list-of-str; bare str accepted) and marker + # absence (None) are the ONLY silent forms. + marker = sparse_audio.is_sparse_audio_marker + with patch.object(sparse_audio.logger, "error") as err: + assert marker(["1"]) is True + assert marker(["0"]) is False + assert marker(["0", "1"]) is True + assert marker("true") is True + assert marker("off") is False + assert marker(None) is False + err.assert_not_called() + + def test_marker_illegal_carriers_raise_one_error_per_type(self): + # Enforcement, not normalization: any carrier outside the design + # contract is a producer bug - it raises InvalidSparseDeclarationError + # (caught inside the module by resolve_sparse_mm_routing, which + # fails closed) + # and logs at error level once per offending type. Interpreting + # deviants would need a second truth predicate, and predicate + # divergence is exactly how the np.array(["0"])-routes-as-sparse bug + # happened. + marker = sparse_audio.is_sparse_audio_marker + for value in (torch.tensor(1), torch.tensor([1, 1]), np.array(["1"]), 1, 1.0, True, ["1", 1]): + with ( + patch.object(sparse_audio, "_logged_offenders", set()), + patch.object(sparse_audio.logger, "error") as err, + ): + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(value) + assert err.call_count == 1, f"expected one error for {value!r}" + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(value) + assert err.call_count == 1 # deduped per type, not per step + + def test_marker_device_tensor_rejected_without_read(self): + # Deviant carriers are never read element-wise, so a device tensor + # can never introduce a D2H sync here. Meta tensors hold no data - + # any read (`.item()`/`bool()`) raises a RuntimeError, so seeing OUR + # exception type (not RuntimeError) proves the data is never touched. + marker = sparse_audio.is_sparse_audio_marker + with patch.object(sparse_audio, "_logged_offenders", set()): + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(torch.ones((), device="meta")) + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(torch.ones(3, device="meta")) + + def test_invalid_marker_fails_closed_at_routing(self): + # A present-but-invalid sparse declaration most likely belongs to a + # producer that INTENDED sparse output; falling back to dense would + # batch-index-assign its payloads to the WRONG requests. Routing + # must fail closed: no payload routed for the step (empty sparse + # set), never dense misrouting. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0", "r1"], + downstream_req_ids=["r0", "r1"], + multimodal_outputs={ + "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}, + "model_outputs": [torch.zeros(1)], + }, + ) + assert downstream == [] + assert index == {} + assert is_sparse is True + + def test_fail_closed_routing_error_remains_observable(self): + def route_invalid_marker(): + return sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={ + "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}, + "model_outputs": [torch.zeros(1)], + }, + ) + + with ( + patch.object(sparse_audio, "_logged_offenders", set()), + patch.object(sparse_audio, "_fail_closed_last_log_at", None), + patch.object(sparse_audio, "_fail_closed_suppressed", 0), + patch.object(sparse_audio.time, "monotonic", side_effect=[100.0, 120.0, 161.0]), + patch.object(sparse_audio.logger, "error") as err, + ): + for _ in range(3): + assert route_invalid_marker() == ([], {}, True) + + routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]] + assert len(routing_errors) == 2 + assert [call.args[1] for call in routing_errors] == [0, 1] + + def test_fail_closed_logging_is_thread_safe(self): + workers = 8 + ready = Barrier(workers) + + def route_invalid_marker(): + return sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={ + "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}, + }, + ) + + def route_concurrently(_): + ready.wait(timeout=10) + return route_invalid_marker() + + with ( + patch.object(sparse_audio, "_logged_offenders", set()), + patch.object(sparse_audio, "_fail_closed_last_log_at", None), + patch.object(sparse_audio, "_fail_closed_suppressed", 0), + patch.object(sparse_audio.time, "monotonic", return_value=100.0) as clock, + patch.object(sparse_audio.logger, "error") as err, + ): + with ThreadPoolExecutor(max_workers=workers) as executor: + results = list(executor.map(route_concurrently, range(workers))) + assert results == [([], {}, True)] * workers + + clock.return_value = 161.0 + assert route_invalid_marker() == ([], {}, True) + + routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]] + assert len(routing_errors) == 2 + assert [call.args[1] for call in routing_errors] == [0, workers - 1] + classifier_errors = [call for call in err.call_args_list if "Illegal sparse-audio marker" in call.args[0]] + assert len(classifier_errors) == 1 + + def test_duplicate_sparse_metadata_must_be_consistent(self): + equal = { + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": ["r1"], + "meta.sparse_audio": ["true"], + } + assert sparse_audio.resolve_sparse_mm_req_ids(equal) == ["r1"] + + invalid_declarations = ( + { + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": ["r2"], + }, + { + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": "r1", + }, + { + "meta": {"req_id": ["r1"], "sparse_audio": ["0"]}, + "meta.sparse_audio": ["1"], + }, + {"meta": {"req_id": ["r1", "r1"], "sparse_audio": ["1"]}}, + ) + with patch.object(sparse_audio, "_logged_offenders", set()): + for declaration in invalid_declarations: + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + sparse_audio.resolve_sparse_mm_req_ids(declaration) + + def test_conflicting_sparse_metadata_fails_closed_at_public_boundary(self): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r1", "r2"], + downstream_req_ids=["r1", "r2"], + multimodal_outputs={ + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": ["r2"], + }, + ) + assert downstream == [] + assert index == {} + assert is_sparse is True + + def test_flat_marker_with_nested_req_ids_is_not_erased(self): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0", "r1"], + downstream_req_ids=["r0", "r1"], + multimodal_outputs={"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]}, + ) + assert downstream == ["r1"] + assert index == {"r1": 0} + assert is_sparse is True + + def test_sparse_marker_with_unusable_req_ids_fails_closed(self): + # The marker says "sparse" but the alignment list is unusable - the + # declaration as a whole is out of contract; same fail-closed shape. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={"meta": {"req_id": "r0", "sparse_audio": ["1"]}}, + ) + assert downstream == [] + assert is_sparse is True + + def test_nested_marker_without_req_id_fails_closed(self): + # {"meta": {"sparse_audio": ["1"]}} - marker set, req_id missing. + # The flattened-encoding fallback must not ERASE the nested marker + # (which used to silently demote the declaration to dense); this is + # an unusable-req_id invalid declaration -> fail closed. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={"meta": {"sparse_audio": ["1"]}}, + ) + assert downstream == [] + assert is_sparse is True + + def test_invalid_marker_on_non_audio_pipeline_keeps_dense_routing(self): + # Non-audio pipelines never consult sparse routing; the carrier + # error is logged by the classifier but routing is unaffected. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="latent", + req_ids_output_copy=["r0", "r1"], + downstream_req_ids=["r0", "r1"], + multimodal_outputs={"meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}}, + ) + assert downstream == ["r0", "r1"] + assert is_sparse is False + + def test_legal_marker_still_routes_sparse(self): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0", "r1", "r2"], + downstream_req_ids=["r0", "r1", "r2"], + multimodal_outputs={"meta": {"req_id": ["r1", "r2"], "sparse_audio": ["1"]}}, + ) + assert downstream == ["r1", "r2"] + assert index == {"r1": 0, "r2": 1} + assert is_sparse is True + + def test_combined_payload_unwraps_aligned_list(self): + combined = {"codes.audio": {"req-1": [torch.zeros(1), torch.ones(1)]}} + payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=1) + assert torch.equal(payload["codes.audio"], torch.ones(1)) + + def test_combined_payload_broadcasts_qwen_request_invariant_singletons(self): + req_ids = ("r0", "r1") + tts_embeddings = { + "embed.tts_bos": torch.full((1, 1, 2), 1.0), + "embed.tts_eos": torch.full((1, 1, 2), 2.0), + "embed.tts_pad": torch.full((1, 1, 2), 3.0), + } + combined = {key: {rid: [value.clone()] for rid in req_ids} for key, value in tts_embeddings.items()} + + for idx, rid in enumerate(req_ids): + payload = payload_build.build_omni_mm_payload( + combined_multimodal_outputs=combined, + mm_cpu=None, + rid=rid, + idx=idx, + start=idx, + end=idx + 1, + audio_sparse_output=False, + sparse_mm_index={}, + hidden_seq_len=len(req_ids), + scheduled_seq_len=len(req_ids), + ) + assert payload.keys() == tts_embeddings.keys() + for key, expected in tts_embeddings.items(): + assert torch.equal(payload[key], expected) + + def test_combined_payload_drops_sparse_protocol_metadata(self): + combined = { + "meta.sparse_audio": {"req-1": ["1"]}, + "meta.req_id": {"req-1": ["req-1"]}, + "sr": {"req-1": 22050}, + } + payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=2) + assert payload == {"sr": 22050} + + def test_combined_payload_misalignment_drops_key_not_request_zero(self): + # A short multi-element list used to silently return `v[0]` - request + # 0's payload shipped under request `rid`. Never substitute another + # request's data; the key is dropped with an error log (same policy + # as the non-combined sparse mismatch) rather than raised - v1 + # runners don't catch per-request exceptions, so a raise would turn a + # per-request protocol break into a whole-stage outage. + combined = { + "codes.audio": {"req-3": [torch.zeros(1), torch.ones(1)]}, + "sr": {"req-3": 22050}, + } + payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-3", list_idx=2) + assert "codes.audio" not in payload + assert payload["sr"] == 22050 + + def test_direct_path_sparse_mismatch_drops_key_and_keeps_payload_shape(self): + # Mirror of the combined-path policy on the non-prefix-cache path: + # a sparse index beyond the per-request list drops that key (error + # log) and the rest of the request's payload survives. + payload = payload_build.build_omni_mm_payload( + combined_multimodal_outputs=None, + mm_cpu={"codes.audio": [torch.zeros(1)]}, + rid="r2", + idx=1, + start=0, + end=1, + audio_sparse_output=True, + sparse_mm_index={"r2": 1}, + hidden_seq_len=1, + scheduled_seq_len=1, + ) + assert payload == {} + + def test_combined_payload_uses_sparse_index_under_sparse_routing(self): + # Batch [r0, r1, r2], sparse outputs only for [r1, r2]: producers + # align per-request lists to the SPARSE order, so r1 (batch idx 1, + # sparse idx 0) must get element 0 and r2 (batch idx 2, sparse idx 1) + # element 1. Indexing by batch idx shipped r2's payload to r1 and + # went out of range for r2. The sparse->list_idx resolution lives in + # `_build_omni_mm_payload` (the combined helper itself is + # sparse-agnostic), so this exercises the caller. + sparse_mm_index = {"r1": 0, "r2": 1} + combined = { + "codes.audio": { + "r1": [torch.zeros(1), torch.ones(1)], + "r2": [torch.zeros(1), torch.ones(1)], + } + } + payloads = {} + for rid, idx in (("r1", 1), ("r2", 2)): + payloads[rid] = payload_build.build_omni_mm_payload( + combined_multimodal_outputs=combined, + mm_cpu=None, + rid=rid, + idx=idx, + start=0, + end=1, + audio_sparse_output=True, + sparse_mm_index=sparse_mm_index, + hidden_seq_len=1, + scheduled_seq_len=1, + ) + assert torch.equal(payloads["r1"]["codes.audio"], torch.zeros(1)) + assert torch.equal(payloads["r2"]["codes.audio"], torch.ones(1)) + + +class TestMergeModelKvTransferMetadata: + """RFC #5450 C4: model KV-transfer metadata must merge into a COPY of the + scheduler's finished-request data - `scheduler_output` can be shared with + the engine-core process and must never be mutated by the runner.""" + + def _runner_with_model(self, model): + runner = object.__new__(GPUARModelRunner) + runner.model = model + return runner + + def test_merge_does_not_mutate_scheduler_data(self): + class _Model: + def get_kv_transfer_metadata(self, req_id, num_computed_tokens): + return {"talker_codes": [7], "seen_tokens": num_computed_tokens} + + original_meta = {"a": 1} + original = {"r1": {"seq_len": 3, "block_ids": [4], "custom_metadata": original_meta}} + runner = self._runner_with_model(_Model()) + + merged = runner._merge_model_kv_transfer_metadata(original) + + # The manager sees the merged view... + assert merged["r1"]["custom_metadata"] == {"a": 1, "talker_codes": [7], "seen_tokens": 3} + assert merged["r1"]["seq_len"] == 3 + # ...while the engine-shared originals are untouched. + assert original["r1"]["custom_metadata"] is original_meta + assert original_meta == {"a": 1} + assert "talker_codes" not in original["r1"]["custom_metadata"] + + def test_merge_without_model_meta_passes_entry_through(self): + class _Model: + def get_kv_transfer_metadata(self, req_id, num_computed_tokens): + return None + + data = {"seq_len": 2, "block_ids": [1]} + runner = self._runner_with_model(_Model()) + + merged = runner._merge_model_kv_transfer_metadata({"r1": data}) + + assert merged["r1"] is data + + def test_merge_failure_raises_instead_of_shipping_partial_transfer(self): + # RFC #5450 C7: the bare `except Exception -> logger.warning` used to + # drop the model's custom metadata and let the KV transfer proceed + # incomplete; the failure then surfaced much later on the consumer + # stage. It must fail at the defect instead. + class _Model: + def get_kv_transfer_metadata(self, req_id, num_computed_tokens): + raise RuntimeError("boom") + + data = {"seq_len": 2, "block_ids": [1], "custom_metadata": {"a": 1}} + runner = self._runner_with_model(_Model()) + + with pytest.raises(RuntimeError, match="boom"): + runner._merge_model_kv_transfer_metadata({"r1": data}) + assert data["custom_metadata"] == {"a": 1} + + +class TestDownstreamPayloadMemoization: + """RFC #5450 C7: `_request_needs_downstream_stage_payload` must not + memoize before the deciding marker exists - the final-stage id arrives via + `model_intermediate_buffer`, which can be unpopulated on the first call.""" + + def _runner(self, stages): + runner = object.__new__(GPUARModelRunner) + runner._downstream_payload_cache = {} + runner._request_final_stage_id = lambda rid: stages[rid] + return runner + + def test_missing_marker_defaults_true_without_memoizing(self): + stages = {"r1": None} + runner = self._runner(stages) + + assert runner._request_needs_downstream_stage_payload("r1") is True + # Not cached: the answer must refresh once the marker arrives. + assert "r1" not in runner._downstream_payload_cache + + stages["r1"] = 0 # marker lands: this stage IS the final stage + assert runner._request_needs_downstream_stage_payload("r1") is False + assert runner._downstream_payload_cache["r1"] is False + + def test_memoizes_once_marker_known(self): + stages = {"r1": 2} + runner = self._runner(stages) + + assert runner._request_needs_downstream_stage_payload("r1") is True + assert runner._downstream_payload_cache["r1"] is True + # Cached value now authoritative for the request's lifetime. + stages["r1"] = 0 + assert runner._request_needs_downstream_stage_payload("r1") is True + + +# Matches a real ASSIGNMENT of the capability (optionally `self.`-qualified +# and/or annotated), not a bare mention: a comment, a docstring, or a read +# such as `getattr(self.model, "prefer_model_sampler", False)` must not count +# as declaring it. `=(?!=)` keeps `==` comparisons out. +_PREFER_MODEL_SAMPLER_ASSIGNMENT = re.compile( + r"^[ \t]*(?:self\.)?prefer_model_sampler[ \t]*(?::[^=\n]*)?=(?!=)[ \t]*(?P[^\n]*)", + re.MULTILINE, +) + + +def _declares_prefer_model_sampler(source: str) -> bool: + """Whether `source` opts INTO the model-sampler contract. + + An explicit `prefer_model_sampler = False` is an opt-OUT and does not + count; any other right-hand side does, including runtime-conditional + forms (minicpmo_4_5 assigns `self.model_stage in {...}`), which can + evaluate True. + """ + for match in _PREFER_MODEL_SAMPLER_ASSIGNMENT.finditer(source): + if match.group("rhs").split("#")[0].strip() != "False": + return True + return False + + +class TestPreferModelSamplerNoneFallback: + """RFC #5450 C7: a `prefer_model_sampler` model returning None from + `sample()` DECLINES the step and the runner falls back to the default + sampler. The fallback is load-bearing (cosyvoice3/glm_tts decline for + empty-logits / inapplicable-stage steps), so it must keep working - now + as a documented contract with a one-time visibility log.""" + + def _runner(self, model_sample_result, default_result): + runner = object.__new__(GPUARModelRunner) + runner.input_batch = SimpleNamespace( + sampling_metadata="smd", + update_async_output_token_ids=lambda: None, + ) + runner.model = SimpleNamespace( + prefer_model_sampler=True, + sample=lambda logits, metadata: model_sample_result, + ) + runner._sampling_metadata_for_model_sampler = lambda smd: smd + runner._resolve_duplex_sampling_hook = lambda: None + runner.sampler = lambda logits, sampling_metadata: default_result + return runner + + def test_none_falls_back_to_default_sampler_and_warns_once(self): + import logging + + from vllm_omni.worker import gpu_ar_model_runner as ar_module + + class _UniqueDecliningSampler: + prefer_model_sampler = True + + @staticmethod + def sample(logits, metadata): + return None + + sentinel = object() + runner = self._runner(model_sample_result=None, default_result=sentinel) + runner.model = _UniqueDecliningSampler() + + records: list[logging.LogRecord] = [] + + class _Capture(logging.Handler): + def emit(self, record): + records.append(record) + + handler = _Capture(level=logging.WARNING) + ar_module.logger.addHandler(handler) + try: + for _ in range(3): + out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None) + assert out is sentinel + finally: + ar_module.logger.removeHandler(handler) + + fallback_records = [r for r in records if "falling back to the default sampler" in r.getMessage()] + assert len(fallback_records) == 1 + assert "_UniqueDecliningSampler" in fallback_records[0].getMessage() + + def test_model_sampler_output_short_circuits_default(self): + model_out = object() + runner = self._runner(model_sample_result=model_out, default_result=object()) + + out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None) + + assert out is model_out + + def test_declaration_matcher_ignores_mentions_and_opt_outs(self): + # Guards the guard: the inventory below is only meaningful if + # "declares" means an actual opt-in assignment. + assert _declares_prefer_model_sampler(" prefer_model_sampler = True") + assert _declares_prefer_model_sampler(" prefer_model_sampler: bool = True") + assert _declares_prefer_model_sampler(' self.prefer_model_sampler = stage in {"llm"}') + assert not _declares_prefer_model_sampler(" prefer_model_sampler = False") + assert not _declares_prefer_model_sampler(" prefer_model_sampler = False # opted out") + assert not _declares_prefer_model_sampler(" # prefer_model_sampler = True (not anymore)") + assert not _declares_prefer_model_sampler(' """Does not use prefer_model_sampler."""') + assert not _declares_prefer_model_sampler(' x = getattr(model, "prefer_model_sampler", False)') + assert not _declares_prefer_model_sampler(" if prefer_model_sampler == True:") + + def test_in_tree_declarer_inventory_is_exact(self): + # The contract's audience: an exact inventory, so ADDING a declarer + # fails this test and the new model owner consciously signs up for the + # documented None-fallback semantics (update the set when they do). + import pathlib + + import vllm_omni + + models_dir = pathlib.Path(vllm_omni.__file__).parent / "model_executor" / "models" + declarers = { + p.parent.name + for p in models_dir.rglob("*.py") + if _declares_prefer_model_sampler(p.read_text(encoding="utf-8", errors="ignore")) + } + expected = { + "cosyvoice3", + "glm_tts", + "higgs_audio_v2", + "higgs_audio_v3", + "hunyuan_image3", + "minicpmo_4_5", + "minimax_music3", + "nemotron_voicechat", + } + assert declarers == expected, ( + "The set of models declaring `prefer_model_sampler` changed:\n" + f" added: {sorted(declarers - expected) or 'none'}\n" + f" removed: {sorted(expected - declarers) or 'none'}\n\n" + "This is a deliberate consent gate, not a broken test. Declaring " + "`prefer_model_sampler` opts a model into the runner's documented " + "None-fallback contract: `model.sample()` may return None to decline " + "a step, and the runner then falls back to the DEFAULT sampler " + "(see `GPUARModelRunner._sample`; RFC #5450 C7.3).\n\n" + "If you ADDED a declarer: confirm your model relies on -- or at " + "least tolerates -- that fallback, then add its directory name to " + "`expected` above. If you REMOVED one, drop its name." + ) diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index a6df4732f..b97168161 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -16,14 +16,14 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _DummyInputBatch: - def __init__(self): - self.req_ids = ["req-1"] - self.req_id_to_index = {"req-1": 0} - self.num_reqs = 1 + def __init__(self, num_reqs: int = 1): + self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)] + self.req_id_to_index = {rid: i for i, rid in enumerate(self.req_ids)} + self.num_reqs = num_reqs self.vocab_size = 10 -def _make_runner(multimodal_outputs): +def _make_runner(multimodal_outputs, num_reqs: int = 1): runner = object.__new__(GPUGenerationModelRunner) runner.execute_model_state = ExecuteModelState( None, @@ -40,7 +40,7 @@ def _make_runner(multimodal_outputs): None, ) runner.kv_connector_output = None - runner.input_batch = _DummyInputBatch() + runner.input_batch = _DummyInputBatch(num_reqs) runner.use_async_scheduling = False runner.device = torch.device("cpu") runner.supports_mm_inputs = False @@ -111,6 +111,12 @@ def _make_guard_runner(): runner.cache_config = CacheConfig() runner.input_batch = _DummyInputBatch() runner.model = object() + runner.parallel_config = SimpleNamespace( + distributed_executor_backend=None, + data_parallel_size=1, + ) + runner._dummy_run_calls = [] + runner._dummy_run = lambda n: runner._dummy_run_calls.append(n) runner._update_states = lambda scheduler_output: None runner.synchronize_input_prep = contextlib.nullcontext runner.attach_omni_connector_output = lambda result: result @@ -121,9 +127,112 @@ def _make_guard_runner(): def test_execute_model_skips_non_positive_scheduled_span(monkeypatch, total): """#5196: a negative span is truthy, so it used to reach `_prepare_inputs`, whose `assert total_num_scheduled_tokens > 0` killed the stage EngineCore.""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False) monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) runner = _make_guard_runner() output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total)) assert output is EMPTY_MODEL_RUNNER_OUTPUT + assert runner._dummy_run_calls == [] + + +@pytest.mark.parametrize("total", [-1, 0]) +def test_execute_model_zero_tokens_runs_dp_dummy_run(monkeypatch, total): + """A 0-token step under external_launcher + DP>1 must `_dummy_run(1)` so + `coordinate_batch_across_dp` stays in sync (it used to be shadowed by an + unconditional early return and never ran -> DP hang risk).""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False) + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) + runner = _make_guard_runner() + runner.parallel_config = SimpleNamespace( + distributed_executor_backend="external_launcher", + data_parallel_size=2, + ) + + output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total)) + + assert output is EMPTY_MODEL_RUNNER_OUTPUT + assert runner._dummy_run_calls == [1] + + +def test_execute_model_zero_tokens_kv_connector_no_forward(monkeypatch): + """With a KV-transfer group, a 0-token step must go through + `kv_connector_no_forward` (also shadowed by the old early return).""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False) + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: True) + runner = _make_guard_runner() + sentinel = object() + calls = [] + runner.kv_connector_no_forward = lambda scheduler_output, vllm_config: ( + calls.append(scheduler_output), + sentinel, + )[1] + runner.vllm_config = object() + + output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0)) + + assert output is sentinel + assert len(calls) == 1 + + +def test_execute_model_zero_tokens_ec_producer_takes_encoder_path(monkeypatch): + """AR lock-step: the EC-producer block sits ABOVE the zero-token block + (matching GPUARModelRunner), so a 0-token step on an EC producer returns + the empty ENCODER output. Before the C1 reorder, the unconditional early + return won and produced EMPTY_MODEL_RUNNER_OUTPUT instead — this is the + one cell whose behaviour the reorder changed.""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: True) + monkeypatch.setattr(gen_runner_module, "get_ec_transfer", lambda: SimpleNamespace(is_consumer=False)) + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) + sentinel = object() + monkeypatch.setattr(gen_runner_module, "make_empty_encoder_model_runner_output", lambda so: sentinel) + runner = _make_guard_runner() + runner.encoder_cache = {} + runner.maybe_get_ec_connector_output = lambda scheduler_output, encoder_cache: contextlib.nullcontext() + encoder_calls = [] + runner._execute_mm_encoder = lambda so: encoder_calls.append(so) + + output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0)) + + assert output is sentinel + assert len(encoder_calls) == 1 + assert runner._dummy_run_calls == [] + + +class TestSampleTokensBatched: + """RFC #5450 C5: `sample_tokens` must slice per request for num_reqs > 1. + The tensor branch's asserts used to force num_reqs == 1 and the list + branch built a length-1 payload misaligned with `req_ids`.""" + + def test_tensor_output_sliced_per_request(self): + multimodal_outputs = torch.arange(12, dtype=torch.float32).reshape(2, 2, 3) + runner = _make_runner(multimodal_outputs, num_reqs=2) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.multimodal_outputs) == 2 + assert torch.equal(output.multimodal_outputs[0]["model_outputs"], multimodal_outputs[0]) + assert torch.equal(output.multimodal_outputs[1]["model_outputs"], multimodal_outputs[1]) + + def test_list_output_one_entry_per_request(self): + entries = [torch.zeros(2, 1), torch.ones(2, 1)] + runner = _make_runner(list(entries), num_reqs=2) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.multimodal_outputs) == 2 + assert torch.equal(output.multimodal_outputs[0]["model_outputs"], entries[0]) + assert torch.equal(output.multimodal_outputs[1]["model_outputs"], entries[1]) + + def test_tensor_output_batch_mismatch_raises(self): + runner = _make_runner(torch.randn(1, 2, 3), num_reqs=2) + + with pytest.raises(ValueError, match="leading dim 1 .* 2 requests"): + GPUGenerationModelRunner.sample_tokens(runner) + + def test_list_output_batch_mismatch_raises(self): + runner = _make_runner([torch.randn(2, 1)], num_reqs=2) + + with pytest.raises(ValueError, match="length 1 .* 2 requests"): + GPUGenerationModelRunner.sample_tokens(runner) diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 383379202..a30dcdb58 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -1,5 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations +import ast import inspect from types import SimpleNamespace @@ -848,6 +852,161 @@ def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format(): assert result["prompt_suffix_len"] == 0 +def _stage0_vision_runtime(): + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + MiniCPMO45Stage0DuplexRuntime, + ) + + class _StageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Embedding(256, 2) + + def get_input_embeddings(self): + return self.embed + + def get_audio_hidden_states(self, _data): + return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)] + + runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime) + runtime.stage_model = _StageModel() + runtime.thinker = runtime.stage_model + runtime.tokenizer = SimpleNamespace( + unk_token_id=0, + convert_tokens_to_ids=lambda token: { + "": 1, + "": 2, + "<|listen|>": 3, + "<|speak|>": 4, + "<|tts_bos|>": 5, + "<|tts_eos|>": 6, + "<|tts_pad|>": 7, + "<|chunk_eos|>": 8, + "<|chunk_tts_eos|>": 9, + "<|turn_eos|>": 10, + "<|audio|>": 11, + "": 12, + "": 13, + "": 14, + "": 15, + }.get(token, 0), + encode=lambda text, add_special_tokens=False: [201, 5], + ) + runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4) + runtime.device = "cpu" + runtime._init_token_ids() + # One 2-row block per frame keeps the expected token_ids readable; the real + # resampler emits VISION_EMBEDS_PER_FRAME rows per slice. + runtime._stage_vision_embeddings = lambda frames: [ + [torch.full((2, 2), float(index), dtype=torch.float32)] for index in range(len(frames)) + ] + return runtime + + +def test_minicpmo_stage0_puts_every_frame_of_an_append_in_one_unit(): + """Official streaming_prefill feeds the whole frame_list into one unit. + + Stacked pairs used to be spread one frame per unit, so a two-frame append + that closed a single unit left a frame over, failed the prefill, and still + advanced the session state. + """ + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + runtime = _stage0_vision_runtime() + state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-stacked-frames") + + result = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + video_frames=["frame-prev", "frame-cur"], + seq=1, + ) + + assert result["success"] is True + # , then each frame as + embeds + , then the audio. + assert result["input_token_ids"] == [1, 12, 0, 0, 13, 12, 0, 0, 13, 11] + assert state.audio_chunk_idx == 1 + + +def test_minicpmo_stage0_wraps_hd_slices_after_the_source_image(): + """Official stacked pair is max_slice_nums=[2, 1]: source + slices, then audio.""" + import torch + + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + runtime = _stage0_vision_runtime() + state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-hd-slices") + runtime._stage_vision_embeddings = lambda frames: [ + [torch.ones((2, 2)), torch.full((2, 2), 2.0), torch.full((2, 2), 3.0)], + [torch.full((2, 2), 4.0)], + ] + + result = runtime._stage_prefill_embeddings_only( + state, + np.zeros(4, dtype=np.float32), + video_frames=["base", "stack"], + seq=1, + ) + + assert result["success"] is True + # src p0 p1 + # stack audio + assert result["input_token_ids"] == [ + 1, + 12, + 0, + 0, + 13, + 14, + 0, + 0, + 15, + 14, + 0, + 0, + 15, + 12, + 0, + 0, + 13, + 11, + ] + + +def test_minicpmo_stage0_reports_frames_that_close_no_unit_instead_of_dropping_them(): + """A frame only binds to the unit its append closes. + + Frames are not carried across appends, so one arriving early used to vanish + under a generic "audio not enough" buffering result. + """ + from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + + runtime = _stage0_vision_runtime() + state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-early-frame") + + assert runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)["success"] is True + + early = runtime._stage_prefill_embeddings_only( + state, + np.zeros(2, dtype=np.float32), + video_frames=["frame-cur"], + seq=2, + ) + + assert early["success"] is False + assert "closes no model unit" in early["reason"] + # The half-filled append must not advance the unit counter. + assert state.audio_chunk_idx == 1 + + def test_minicpmo_stage0_speech_append_sets_pending_context_once(): import torch @@ -2466,3 +2625,128 @@ def test_minicpmo_stage0_session_context_includes_resolved_ref_audio(): assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5] assert len(state.context_embeds) == 6 + + +# --------------------------------------------------------------------------- +# Cross-platform wiring: every AR model runner must reach the duplex hook. +# +# GPUARModelRunner and NPUARModelRunner are siblings -- neither inherits the +# other's _sample -- so the NPU runner once copied GPU _sample without the +# prepare_duplex_sampling call. Parse source instead of importing NPUARModelRunner: +# that class needs vllm_ascend, which CPU test images do not have. +# --------------------------------------------------------------------------- + +_REQUIRED_DUPLEX_CALLS = { + "__init__": "_init_duplex_sampling_state", + "load_model": "_resolve_duplex_sampling_hook", + "_update_states": "_update_duplex_sampling_states", + "_sample": "_apply_duplex_sampling", +} + +_MODEL_SAMPLER_CALL = "model_sample" + + +def _repo_root(): + from pathlib import Path + + return Path(__file__).parents[2] + + +def _ar_runner_sources(): + root = _repo_root() + paths = [root / "vllm_omni/worker/gpu_ar_model_runner.py"] + paths.extend(sorted(root.glob("vllm_omni/platforms/*/worker/*_ar_model_runner.py"))) + return paths + + +def _ar_runner_classdefs(): + root = _repo_root() + found = {} + for path in _ar_runner_sources(): + module = ast.parse(path.read_text()) + for node in module.body: + if isinstance(node, ast.ClassDef) and node.name.endswith("ARModelRunner"): + found[node.name] = (node, str(path.relative_to(root))) + assert found, "no AR model runners discovered; the glob in _ar_runner_sources() is stale" + return found + + +def _method_calls(method): + calls: dict[str, int] = {} + for node in ast.walk(method): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + calls.setdefault(node.func.attr, node.lineno) + return calls + + +def _wires_duplex_sampling_directly(classdef) -> bool: + bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)} + if "DuplexSamplingRunnerMixin" not in bases: + return False + methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)} + return all( + method_name in methods and required_call in _method_calls(methods[method_name]) + for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items() + ) + + +def test_every_ar_runner_reaches_the_duplex_sampling_hook(): + classdefs = _ar_runner_classdefs() + wired = {name for name, (classdef, _path) in classdefs.items() if _wires_duplex_sampling_directly(classdef)} + + for class_name, (classdef, path) in sorted(classdefs.items()): + if class_name in wired: + continue + base_names = {base.id for base in classdef.bases if isinstance(base, ast.Name)} + assert base_names & wired, ( + f"{path}::{class_name} neither wires the duplex sampling hook nor derives from a " + f"runner that does (wired runners: {sorted(wired)}). Inherit DuplexSamplingRunnerMixin " + f"and call {sorted(_REQUIRED_DUPLEX_CALLS.values())} from " + f"{sorted(_REQUIRED_DUPLEX_CALLS)} respectively." + ) + + +@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"]) +def test_ar_runner_calls_every_duplex_sampling_hook_site(class_name: str): + classdefs = _ar_runner_classdefs() + assert class_name in classdefs, f"{class_name} not found by _ar_runner_sources()" + classdef, path = classdefs[class_name] + + bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)} + assert "DuplexSamplingRunnerMixin" in bases, ( + f"{path}::{class_name} must inherit DuplexSamplingRunnerMixin so the duplex " + f"sampling hook is wired the same way on every platform" + ) + + methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)} + for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items(): + method = methods.get(method_name) + assert method is not None, f"{path}::{class_name}.{method_name} is missing; it must call {required_call}" + calls = _method_calls(method) + assert required_call in calls, f"{path}::{class_name}.{method_name} must call self.{required_call}" + + +@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"]) +def test_ar_runner_applies_duplex_sampling_before_the_model_sampler(class_name: str): + classdef, path = _ar_runner_classdefs()[class_name] + sample = next(node for node in classdef.body if isinstance(node, ast.FunctionDef) and node.name == "_sample") + + apply_lineno = _method_calls(sample).get("_apply_duplex_sampling") + sampler_lineno = next( + ( + node.lineno + for node in ast.walk(sample) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == _MODEL_SAMPLER_CALL + ), + None, + ) + assert apply_lineno is not None, f"{path}::{class_name}._sample must call self._apply_duplex_sampling" + assert sampler_lineno is not None, ( + f"{path}::{class_name}._sample no longer calls {_MODEL_SAMPLER_CALL}(); update _MODEL_SAMPLER_CALL " + f"so the ordering check keeps testing something" + ) + assert apply_lineno < sampler_lineno, ( + f"{path}::{class_name}._sample calls self._apply_duplex_sampling at line {apply_lineno}, after " + f"{_MODEL_SAMPLER_CALL}() at line {sampler_lineno}. The hook masks the logits and publishes the " + f"row -> session map the sampler reads, so running it afterwards is a silent no-op." + ) diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index ca282994a..81d2a6b8a 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -1,9 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from contextlib import contextmanager from types import SimpleNamespace import pytest import torch from vllm.v1.cudagraph_dispatcher import CUDAGraphMode +from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.gpu_model_runner import ( @@ -107,6 +111,23 @@ class DummyReqState: pass +def test_model_forward_passes_request_ids_to_decode_metadata(monkeypatch): + received = {} + model = SimpleNamespace( + supports_omni_decode_step_metadata=True, + update_decode_step_metadata=lambda **kwargs: received.update(kwargs), + ) + runner = object.__new__(OmniGPUModelRunner) + runner.model = model + runner.input_batch = DummyInputBatch(["request-a", "request-b"]) + runner._build_model_kwargs_extra = lambda: {} + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *_args, **_kwargs: torch.zeros(1)) + + OmniGPUModelRunner._model_forward(runner, input_ids=torch.ones(2, dtype=torch.long)) + + assert received["req_ids"] == ["request-a", "request-b"] + + class MiMoAudioForConditionalGeneration(torch.nn.Module): """Dummy model whose class name must exactly match the production check.""" @@ -543,6 +564,80 @@ def test_update_additional_information_deserializes_new_request_payload(): ) +def test_streaming_new_request_marker_replaces_terminal_chunk_snapshot(): + from vllm_omni.engine.serialization import serialize_additional_information + + runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4) + runner.model.replace_runtime_additional_information = True + terminal = { + "codes": {"audio": torch.tensor([1, 2])}, + "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True}, + } + peer = { + "codes": {"audio": torch.tensor([9])}, + "meta": {"cache_epoch": 3, "chunk_seq": 1, "last_chunk": False}, + } + runner.model_intermediate_buffer.update(r1=terminal, r2=peer) + marker = { + "meta": { + "finished": False, + "is_segment_finished": True, + "request_finished": False, + "replace_runtime_additional_information": True, + } + } + new_req = SimpleNamespace( + req_id="r1", + model_intermediate_buffer=marker, + additional_information=serialize_additional_information(terminal), + ) + + OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req, "r1") + OmniGPUModelRunner._update_additional_information( + runner, + SimpleNamespace( + scheduled_new_reqs=[new_req], + scheduled_cached_reqs=SimpleNamespace(), + ), + ) + + info = runner.model_intermediate_buffer["r1"] + assert "codes" not in info + assert info["meta"] == { + **marker["meta"], + "num_processed_tokens": 0, + "resumable": True, + } + assert runner.requests["r1"].additional_information_cpu == info + assert runner.model_intermediate_buffer["r2"] == peer + + +def test_cached_empty_marker_replaces_terminal_chunk_snapshot(): + runner = _make_runner(req_ids=("r1",), hidden_size=4) + runner.model.replace_runtime_additional_information = True + runner.model_intermediate_buffer["r1"] = { + "codes": {"audio": torch.tensor([1, 2])}, + "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True}, + } + marker = { + "meta": { + "is_segment_finished": torch.tensor(True, dtype=torch.bool), + "replace_runtime_additional_information": True, + } + } + + OmniGPUModelRunner._update_additional_information( + runner, + SimpleNamespace( + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace(additional_information={"r1": marker}), + ), + ) + + assert runner.model_intermediate_buffer["r1"] == marker + assert runner.requests["r1"].additional_information_cpu == marker + + def test_update_intermediate_buffer_skips_empty_update(): """Validate that an empty update dict is a no-op.""" runner = _make_runner(req_ids=("r1",), hidden_size=4) diff --git a/tests/worker/test_worker_base.py b/tests/worker/test_worker_base.py index 96e1ea5d3..305b0778f 100644 --- a/tests/worker/test_worker_base.py +++ b/tests/worker/test_worker_base.py @@ -199,7 +199,14 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags): worker = object.__new__(OmniGPUWorkerBase) worker.rank = 0 worker.device = "cuda:0" - monkeypatch.setattr(base, "current_omni_platform", _fake_platform()) + order: list[str] = [] + platform = SimpleNamespace( + get_current_memory_usage=lambda device: 0, + empty_cache=lambda: order.append("empty_cache"), + synchronize=lambda: order.append("sync"), + collect=lambda: None, + ) + monkeypatch.setattr(base, "current_omni_platform", platform) calls = {} @@ -209,12 +216,17 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags): return cls() def sleep(self, offload_tags): + order.append("allocator_sleep") calls["offload_tags"] = offload_tags monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator) assert OmniGPUWorkerBase.sleep(worker, level=level) is True assert calls["offload_tags"] == expected_tags + assert order[0] == "sync" + assert "allocator_sleep" in order + assert order.index("sync") < order.index("allocator_sleep") + assert "empty_cache" not in order def test_wake_up_forwards_tags_to_allocator(monkeypatch): diff --git a/tools/baselines/omni_hash_baseline.py b/tools/baselines/omni_hash_baseline.py index 1bb93099b..a1644aeb3 100644 --- a/tools/baselines/omni_hash_baseline.py +++ b/tools/baselines/omni_hash_baseline.py @@ -147,7 +147,7 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]: # (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time. omni = Omni( model, - stage_configs_path=yaml_path, + deploy_config=yaml_path, stage_init_timeout=1800, init_timeout=3600, async_chunk=async_mode, @@ -159,11 +159,10 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]: wav: bytes = b"" # Batch size 1: submit a single prompt so scheduling is deterministic. for step in omni.generate([_text_prompt(question)], _sampling_params(), py_generator=False): - out = step.request_output if step.final_output_type == "text": - text = out.outputs[0].text # exact text, no strip -> truly bit-identical + text = step.outputs[0].text # exact text, no strip -> truly bit-identical elif step.final_output_type == "audio": - wav = _audio_bytes(out.outputs[0].multimodal_output["audio"]) + wav = _audio_bytes(step.outputs[0].multimodal_output["audio"]) if text is None: raise RuntimeError(f"prompt_id={prompt_id} async={async_mode}: produced no text output") if not wav: diff --git a/tools/nightly/buildkite_testcase_statistics.py b/tools/nightly/buildkite_testcase_statistics.py index 740f896f5..06f09c791 100644 --- a/tools/nightly/buildkite_testcase_statistics.py +++ b/tools/nightly/buildkite_testcase_statistics.py @@ -1,8 +1,7 @@ #!/usr/bin/env python3 """ -Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml, -and test-weekly.yml; collect test cases (including parametrized) via pytest --collect-only -q -and produce an HTML report. +Parse pytest commands from Buildkite .buildkite/cuda/test-{ready,merge,nightly,weekly}.yml; +collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report. Leaf steps may live under ``group:`` blocks (nested ``steps``); those are flattened with the group title carried into the report. @@ -30,7 +29,7 @@ import yaml # Repo root (parent of the directory containing this script) REPO_ROOT = Path(__file__).resolve().parent.parent.parent -BUILDKITE_DIR = REPO_ROOT / ".buildkite" +BUILDKITE_DIR = REPO_ROOT / ".buildkite" / "cuda" PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml", "test-weekly.yml"] @@ -256,13 +255,34 @@ def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_l rel_dir = Path(target.rstrip("/")) dir_path = (repo_root / rel_dir).resolve() if dir_path.exists() and dir_path.is_dir(): - return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 120 + return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 300 if not extra: raise RuntimeError(f"Failed to parse pytest sidecar args from line: {raw_line!r}") if not (repo_root / "tests").exists(): raise FileNotFoundError("tests/ directory not found under repo root") - return ["tests/"], "tests/", 120 + # Broad ``tests/`` collection loads many conftests; allow more time on busy CI agents. + return ["tests/"], "tests/", 300 + + +def _format_pytest_collect_failure( + path_args: list[str], + extra: list[str], + result: subprocess.CompletedProcess[str], +) -> str: + """pytest often puts collection errors on stdout; include both streams + exit code.""" + stdout = (result.stdout or "").strip() + stderr = (result.stderr or "").strip() + parts = [ + f"pytest --collect-only failed for {path_args} with args {extra} (exit={result.returncode})", + ] + if stderr: + parts.append(f"stderr:\n{stderr}") + if stdout: + parts.append(f"stdout:\n{stdout}") + if not stderr and not stdout: + parts.append("(no stdout/stderr captured)") + return "\n".join(parts) def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]: @@ -273,6 +293,8 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str] flags ``--ignore``, ``--test-config-file``, ``-m``, ``-k``, ``--run-level`` are taken from the full ``raw_line`` and appended to the same invocation so collection matches CI (perf scripts parametrize from ``--test-config-file`` at import time). + + Pytest exit code 5 (no tests collected) is treated as an empty result, not a hard failure. """ extra = _pytest_collect_sidecar_args(raw_line) path_args, _fallback, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line) @@ -286,8 +308,15 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str] text=True, timeout=timeout_quiet, ) + # 5 == no tests collected under the given -m/-k filters. + if result.returncode == 5: + print( + f"pytest --collect-only: no tests matched for {path_args} with args {extra}", + file=sys.stderr, + ) + return [] if result.returncode != 0: - raise RuntimeError(f"pytest --collect-only failed for {path_args} with args {extra}: {result.stderr.strip()}") + raise RuntimeError(_format_pytest_collect_failure(path_args, extra, result)) print(f"pytest --collect-only success for {path_args} with args {extra}: {result.stdout.strip()}") return _parse_collect_only_stdout(result.stdout or "", stderr=result.stderr or "") @@ -962,7 +991,7 @@ def main() -> None: "--buildkite-dir", type=Path, default=BUILDKITE_DIR, - help="Path to .buildkite directory", + help="Path to Buildkite CUDA pipeline directory (default: .buildkite/cuda)", ) default_out = REPO_ROOT / "buildkite_testcase_statistics.html" parser.add_argument( @@ -1023,7 +1052,12 @@ def main() -> None: names = [] skipped_ids: set[str] = set() for target, raw in targets: - collected = collect_test_names(target, REPO_ROOT, raw) + try: + collected = collect_test_names(target, REPO_ROOT, raw) + except (RuntimeError, FileNotFoundError, subprocess.TimeoutExpired) as exc: + # Keep generating the report for other steps; surface the root cause. + print(f"WARNING: collect failed for {module_key} target={target!r}: {exc}", file=sys.stderr) + collected = [] skipped_ids.update(get_skip_status(target, raw, REPO_ROOT)) names.extend(collected) count = len(names) - len(skipped_ids) diff --git a/tools/nightly/run_nightly_jobs.sh b/tools/nightly/run_nightly_jobs.sh index a95eb8876..b5b6b92ce 100644 --- a/tools/nightly/run_nightly_jobs.sh +++ b/tools/nightly/run_nightly_jobs.sh @@ -20,6 +20,8 @@ # acc — label contains "Accuracy Test" (incl. GEBench / GEdit-Bench style) # function — label has neither "Perf Test" nor "Accuracy Test" (incl. Doc, Multi-Replica, etc.) # stability— fixed dfx stability scripts under tests/dfx/stability/scripts/ (see below) +# Uses --pytest-mark / PYTEST_MARK for pytest -m (hardware via JSON marks). +# When unset, -m defaults to H800 (CUDA stability fleet; use --pytest-mark A3 on NPU). # local — pytest -sv -m " and local_model" from repo root (no YAML step extract) # When LABEL_SUBSTR is set, also runs matching tests/**/test_*.py and perf JSON configs under # tests/dfx/perf/tests/*.json via run_benchmark.py / run_diffusion_benchmark.py. @@ -42,8 +44,10 @@ # whose filename contains the substring (benchmark runner chosen by JSON family). # # stability (when included in TEST_TYPE): -# From repo root: pytest -s -v --run-level full_model tests/dfx/stability/scripts/test_stability_*.py -# model_type: omni → qwen3_omni; tts → qwen3_tts + voxcpm2; diffusion → qwen_image + wan22 + hunyuan_image; all → all six +# From repo root: pytest -s -v --run-level full_model -m "" tests/dfx/stability/scripts/... +# -m comes from --pytest-mark / PYTEST_MARK when set (e.g. H800 or A3). +# Otherwise -m defaults to H800. PYTEST_MARK is rejected unless stability is enabled. +# model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22 + hunyuan_image; all → all five # LABEL_SUBSTR: if set, script path / job key / filename must contain it # # Requirements: bash, python3, PyYAML (pip install pyyaml) @@ -73,6 +77,7 @@ # MODEL_TYPE - comma-separated and/or repeated flags (default: all); see above # LABEL_SUBSTR - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename; # local: substring of test_*.py basename or tests/dfx/perf/tests/*.json filename +# PYTEST_MARK - pytest -m expression for stability jobs only (also --pytest-mark) # DRY_RUN=1 - print extracted commands only; do not write scripts or run pytest # RUN_JOB_TIMEOUT_KILL_AFTER - seconds after SIGTERM before SIGKILL on inline pytest # timeout (default: 60; baked into generated job scripts) @@ -82,6 +87,7 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" LABEL_SUBSTR="${LABEL_SUBSTR:-}" +PYTEST_MARK="${PYTEST_MARK:-}" TEST_TYPE_ENV="${TEST_TYPE:-all}" MODEL_TYPE_ENV="${MODEL_TYPE:-all}" TEST_TYPE_CLI_PARTS=() @@ -204,6 +210,10 @@ while [[ $# -gt 0 ]]; do LABEL_SUBSTR="$2" shift 2 ;; + --pytest-mark) + PYTEST_MARK="$2" + shift 2 + ;; --model-type) MODEL_TYPE_FROM_CLI=1 _split_append_csv_array MODEL_TYPE_CLI_PARTS "$2" @@ -339,7 +349,7 @@ fi mkdir -p "${LOG_DIR}/jobs" LOG_DIR="$(cd "${LOG_DIR}" && pwd)" -export REPO_ROOT LOG_DIR YML LABEL_SUBSTR TEST_TYPE MODEL_TYPE DRY_RUN +export REPO_ROOT LOG_DIR YML LABEL_SUBSTR TEST_TYPE MODEL_TYPE DRY_RUN PYTEST_MARK # Drop stale wrapper scripts, old tee logs, and perf manifest so a filtered run does not # mix with the previous outputs. @@ -588,7 +598,6 @@ def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) - STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [ ("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)), ("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)), - ("stability_voxcpm2", "tests/dfx/stability/scripts/test_stability_voxcpm2.py", ("tts",)), ("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)), ("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)), ("stability_hunyuan_image", "tests/dfx/stability/scripts/test_stability_hunyuan_image.py", ("diffusion",)), @@ -601,6 +610,26 @@ def _stability_model_matches(model_types: list[str], families: tuple[str, ...]) return bool(set(model_types) & set(families)) +STABILITY_DEFAULT_PYTEST_MARK = "H800" + + +def _shell_quote_pytest_mark(marker_expr: str) -> str: + """Shell-safe pytest ``-m`` value for generated bash job scripts.""" + expr = (marker_expr or "").strip() + if not expr: + raise ValueError("pytest -m expression must not be empty") + return shlex.quote(expr) + + +def _stability_pytest_mark(model_types: list[str]) -> str: + """Return explicit --pytest-mark or the default CUDA stability hardware marker.""" + del model_types # MODEL_TYPE filters scripts; default -m is hardware-specific. + override = (os.environ.get("PYTEST_MARK") or "").strip() + if override: + return override + return STABILITY_DEFAULT_PYTEST_MARK + + def _write_job_script(key: str, script_lines: list[str], jobs_dir: Path) -> None: body = "\n".join(script_lines) + "\n" if DRY_RUN: @@ -713,7 +742,7 @@ def run_local_mode( if rel_paths: paths_arg = " ".join(shlex.quote(p) for p in rel_paths) - pytest_line = f'pytest -sv -m "{marker_expr}" {paths_arg}' + pytest_line = f"pytest -sv -m {_shell_quote_pytest_mark(marker_expr)} {paths_arg}" slug = _local_job_slug(LABEL_SUBSTR) key = f"local_pytest_{slug}" header = ( @@ -733,7 +762,7 @@ def run_local_mode( for json_rel, runner_rel in perf_pairs: json_name = Path(json_rel).name pytest_line = ( - f'pytest -s -v -m "{marker_expr}" ' + f"pytest -s -v -m {_shell_quote_pytest_mark(marker_expr)} " f"{shlex.quote(runner_rel)} " f"--test-config-file {shlex.quote(json_rel)}" ) @@ -764,7 +793,7 @@ def run_local_mode( ) return matched - pytest_line = f'pytest -sv -m "{marker_expr}"' + pytest_line = f"pytest -sv -m {_shell_quote_pytest_mark(marker_expr)}" key = "local_pytest" header = f"# Local marker tests — MODEL_TYPE={','.join(model_types)}" script_lines = [ @@ -780,6 +809,8 @@ def run_local_mode( def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int: matched = 0 + marker_expr = _stability_pytest_mark(model_types) + mark_from_cli = bool((os.environ.get("PYTEST_MARK") or "").strip()) for key, rel_posix, families in STABILITY_CASES: if not _stability_model_matches(model_types, families): continue @@ -795,11 +826,17 @@ def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int: if not rel_path.is_file(): print(f"# skip (missing file): {rel_path}", file=sys.stderr) continue - pytest_line = f"pytest -s -v --run-level full_model {rel_posix}" + pytest_line = ( + f"pytest -s -v --run-level full_model " + f"-m {_shell_quote_pytest_mark(marker_expr)} " + f"{shlex.quote(rel_posix)}" + ) matched += 1 script_lines = [ "#!/usr/bin/env bash", f"# Stability: {key} — {rel_posix}", + f"# pytest -m {marker_expr!r}" + + (" (--pytest-mark)" if mark_from_cli else f" (default {STABILITY_DEFAULT_PYTEST_MARK})"), "set -euo pipefail", f'cd "{REPO_ROOT}"', pytest_line, @@ -817,6 +854,14 @@ def main() -> None: yaml_test_types = [t for t in test_types if t not in ("stability", "local")] needs_yaml = bool(yaml_test_types) + pytest_mark = (os.environ.get("PYTEST_MARK") or "").strip() + if pytest_mark and not want_stability: + print( + "PYTEST_MARK / --pytest-mark is only supported when --test-type includes stability", + file=sys.stderr, + ) + sys.exit(2) + jobs_dir = LOG_DIR / "jobs" if not DRY_RUN: jobs_dir.mkdir(parents=True, exist_ok=True) diff --git a/tools/pre_commit/check_buildkite.py b/tools/pre_commit/check_buildkite.py new file mode 100644 index 000000000..362253494 --- /dev/null +++ b/tools/pre_commit/check_buildkite.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Validate .buildkite YAML against the official Buildkite schema. + +Test pipelines (test-ready.yml, etc.) use Omni-only keys such as +``mirror_hardwares`` and ``source_file_dependencies``. Those are expanded or +stripped the same way ``upload_pipeline.py`` does before upload, then the +rendered YAML is checked with check-jsonschema. +""" + +from __future__ import annotations + +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + +import yaml + +ROOT = Path(__file__).resolve().parents[2] +SCRIPTS = ROOT / ".buildkite" / "common" / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +from upload_pipeline import _render_test_pipeline # noqa: E402 + +SKIP_FILES = { + ".buildkite/common/ci_mirror_hardwares.yml", + ".buildkite/npu/pipeline-npu-a3.yml", +} + +# AMD pipelines are uploaded by a different path and use list-valued +# mirror_hardwares plus agent_pool/grade, which CUDA's expander rejects. +AMD_STRIP_KEYS = frozenset({"mirror_hardwares", "agent_pool", "grade", "source_file_dependencies"}) + + +def _normalize(path: str) -> str: + return path.replace("\\", "/") + + +def _strip_keys(obj: Any, keys: frozenset[str]) -> Any: + if isinstance(obj, dict): + return {k: _strip_keys(v, keys) for k, v in obj.items() if k not in keys} + if isinstance(obj, list): + return [_strip_keys(item, keys) for item in obj] + return obj + + +def _prepare_pipeline(path: Path) -> dict[str, Any]: + rel = _normalize(str(path.relative_to(ROOT) if path.is_absolute() else path)) + doc = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(doc, dict): + raise ValueError(f"{rel}: expected a mapping (pipeline YAML), got {type(doc).__name__}") + + if "/amd/" in f"/{rel}": + return _strip_keys(doc, AMD_STRIP_KEYS) + + return _render_test_pipeline(doc, changed_files=None) + + +def _validate(path: Path, doc: dict[str, Any]) -> int: + with tempfile.NamedTemporaryFile( + mode="w", + suffix=".yml", + encoding="utf-8", + delete=False, + ) as handle: + yaml.safe_dump(doc, handle, sort_keys=False) + tmp = Path(handle.name) + + try: + result = subprocess.run( + [ + sys.executable, + "-m", + "check_jsonschema", + "--builtin-schema", + "vendor.buildkite", + str(tmp), + ], + capture_output=True, + text=True, + ) + finally: + tmp.unlink(missing_ok=True) + + if result.returncode == 0: + return 0 + + rel = _normalize(str(path)) + sys.stderr.write(f"{rel}: Buildkite schema validation failed\n") + sys.stderr.write(result.stdout or "") + sys.stderr.write(result.stderr or "") + return 1 + + +def main(argv: list[str]) -> int: + rc = 0 + for raw in argv: + path = Path(raw) + rel = _normalize(str(path)) + if rel in SKIP_FILES: + continue + if not path.is_file(): + path = ROOT / path + try: + doc = _prepare_pipeline(path) + except Exception as exc: # noqa: BLE001 - surface expander errors as hook failures + sys.stderr.write(f"{rel}: {exc}\n") + rc = 1 + continue + rc |= _validate(path, doc) + return rc + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py new file mode 100644 index 000000000..926e102fa --- /dev/null +++ b/tools/pre_commit/check_forbidden_imports.py @@ -0,0 +1,367 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Ban dangerous / policy-violating imports unless a file is allowlisted. + +Ported from vllm/tools/pre_commit/check_forbidden_imports.py. Existing +vllm_omni/ usages are grandfathered; add a path to the matching +allowed_files set only after a security/policy review. +""" + +import sys +from dataclasses import dataclass, field + +import regex as re + +# Hub entry points that must go through the vLLM-tagged repo_utils helpers. +_HF_NAMES = ( + r"HfApi|HfFileSystem|hf_hub_download|snapshot_download" + r"|list_repo_files|file_exists|try_to_load_from_cache" + r"|list_repo_refs|repo_exists" +) + +# Non-library trees: examples, tests, and tooling may keep stdlib imports. +_NON_LIBRARY_DIRS = { + "examples/", + "tests/", + "benchmarks/", + "apps/", + "docs/", + "tools/", + "scripts/", + ".buildkite/", + ".github/", +} + + +@dataclass +class ForbiddenImport: + pattern: str + tip: str + allowed_pattern: re.Pattern = re.compile(r"^$") # matches nothing by default + allowed_files: set[str] = field(default_factory=set) + allowed_dirs: set[str] = field(default_factory=set) + + +CHECK_IMPORTS = { + "pickle/cloudpickle": ForbiddenImport( + pattern=( + r"^\s*(import\s+(?:[\w.]+\s*(?:as\s+\w+)?\s*,\s*)*(pickle|cloudpickle)\b" + r"|from\s+(pickle|cloudpickle)\s+import\b)" + ), + tip=("Avoid using pickle or cloudpickle or add this file to tools/pre_commit/check_forbidden_imports.py."), + allowed_files={ + # STOP AND READ BEFORE YOU ADD ANYTHING ELSE TO THIS LIST: + # pickle/cloudpickle are unsafe when deserializing untrusted data. + "tests/diffusion/attention/test_attention_sp.py", + "tests/helpers/process.py", + "vllm_omni/diffusion/distributed/group_coordinator.py", + }, + ), + "base64": ForbiddenImport( + pattern=r"^\s*(?:import\s+base64(?:$|\s|,)|from\s+base64\s+import)", + tip=("Replace 'import base64' with 'import pybase64' or 'import pybase64 as base64'."), + allowed_pattern=re.compile(r"^\s*import\s+pybase64(\s*|\s+as\s+base64\s*)$"), + allowed_dirs=_NON_LIBRARY_DIRS, + allowed_files={ + "vllm_omni/benchmarks/data_modules/daily_omni_dataset.py", + "vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py", + "vllm_omni/benchmarks/data_modules/seed_tts_dataset.py", + "vllm_omni/entrypoints/openai/api_server.py", + "vllm_omni/entrypoints/openai/audio_utils_mixin.py", + "vllm_omni/entrypoints/openai/image_api_utils.py", + "vllm_omni/entrypoints/openai/protocol/images.py", + "vllm_omni/entrypoints/openai/realtime_connection.py", + "vllm_omni/entrypoints/openai/serving_chat.py", + "vllm_omni/entrypoints/openai/serving_speech.py", + "vllm_omni/entrypoints/openai/serving_speech_stream.py", + "vllm_omni/entrypoints/openai/video_api_utils.py", + "vllm_omni/entrypoints/openai/video_stream_base.py", + "vllm_omni/entrypoints/openai/video_stream_session.py", + "vllm_omni/experimental/fullduplex/client.py", + "vllm_omni/experimental/fullduplex/engine/contracts.py", + "vllm_omni/experimental/fullduplex/joyvl/memory/memory.py", + "vllm_omni/experimental/fullduplex/minicpmo45/adapter.py", + "vllm_omni/experimental/fullduplex/minicpmo45/input.py", + "vllm_omni/experimental/fullduplex/minicpmo45/runtime.py", + "vllm_omni/experimental/fullduplex/minicpmo45/stage0.py", + "vllm_omni/experimental/fullduplex/openai/audio.py", + "vllm_omni/experimental/fullduplex/openai/realtime_input.py", + "vllm_omni/experimental/fullduplex/openai/runtime_bridge.py", + "vllm_omni/experimental/fullduplex/openai/serving.py", + "vllm_omni/experimental/fullduplex/personaplex/input.py", + "vllm_omni/experimental/fullduplex/personaplex/runtime_extension.py", + "vllm_omni/experimental/fullduplex/personaplex/stage0.py", + "vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py", + "vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py", + "vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_tokenizer.py", + "vllm_omni/tokenizers/mammoth_moda2_tokenizer.py", + "vllm_omni/transformers_utils/processors/audiox.py", + }, + ), + "re": ForbiddenImport( + pattern=r"^\s*(?:import\s+re(?:$|\s|,)|from\s+re\s+import)", + tip="Replace 'import re' with 'import regex as re' or 'import regex'.", + allowed_pattern=re.compile(r"^\s*import\s+regex(\s*|\s+as\s+re\s*)$"), + allowed_dirs=_NON_LIBRARY_DIRS, + allowed_files={ + "setup.py", + "vllm_omni/benchmarks/data_modules/daily_omni_eval.py", + "vllm_omni/config/stage_config.py", + "vllm_omni/diffusion/layers/mot/ops/mot_gemm.py", + "vllm_omni/diffusion/model_loader/diffusers_loader.py", + "vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py", + "vllm_omni/diffusion/models/dreamid_omni/fusion.py", + "vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py", + "vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py", + "vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py", + "vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py", + "vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py", + "vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py", + "vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py", + "vllm_omni/diffusion/models/minimax_h3/encoder.py", + "vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py", + "vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py", + "vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py", + "vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py", + "vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py", + "vllm_omni/distributed/omni_connectors/utils/env.py", + "vllm_omni/entrypoints/cli/logo.py", + "vllm_omni/entrypoints/openai/serving_speech.py", + "vllm_omni/experimental/fullduplex/joyvl/memory/memory.py", + "vllm_omni/model_executor/models/common/whisper_vq.py", + "vllm_omni/model_executor/models/fish_speech/prompt_utils.py", + "vllm_omni/model_executor/models/glm_tts/text_frontend.py", + "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py", + "vllm_omni/model_executor/models/indextts2/indextts2_talker.py", + "vllm_omni/model_executor/models/indextts2/text_processing_v2_5.py", + "vllm_omni/model_executor/models/indextts2/utils/common.py", + "vllm_omni/model_executor/models/ming_flash_omni/text_processing.py", + "vllm_omni/model_executor/models/ming_flash_omni/vision_encoder.py", + "vllm_omni/model_executor/models/ming_tts/prompt_assembly.py", + "vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py", + "vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py", + "vllm_omni/outputs/output_modality.py", + "vllm_omni/quantization/mxfp4_config.py", + "vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py", + "vllm_omni/utils/custom_voice_io.py", + }, + ), + "triton": ForbiddenImport( + pattern=r"^(from|import)\s+triton(\s|\.|$)", + tip="Use 'from vllm.triton_utils import triton' instead.", + allowed_pattern=re.compile(r"from vllm.triton_utils import (triton|tl|tl, triton)"), + ), + "tilelang": ForbiddenImport( + pattern=r"^(from|import)\s+tilelang(\s|\.|$)", + tip="Use 'from vllm.tilelang_utils import tilelang, T' instead.", + allowed_pattern=re.compile( + r"from\s+vllm\.tilelang_utils\s+import\s+" + r"(tilelang|T|T, tilelang|tilelang, T)\b" + ), + ), + "huggingface_hub repo API": ForbiddenImport( + # Catch `from huggingface_hub import `, including parenthesized, + # multi-line imports. + pattern=( + r"^\s*from\s+huggingface_hub\s+import\s*\([^)]*\b(?:" + _HF_NAMES + r")\b" + r"|" + r"^\s*from\s+huggingface_hub\s+import\b[^\n]*\b(?:" + _HF_NAMES + r")\b" + ), + tip=( + "Use the shared, vLLM-tagged helpers from " + "vllm.transformers_utils.repo_utils (e.g. hf_api(), hf_fs(), " + "list_repo_files, file_exists) instead of calling " + "huggingface_hub directly." + ), + allowed_dirs=_NON_LIBRARY_DIRS, + allowed_files={ + "vllm_omni/benchmarks/data_modules/daily_omni_dataset.py", + "vllm_omni/benchmarks/data_modules/seed_tts_dataset.py", + "vllm_omni/benchmarks/data_modules/seed_tts_eval.py", + "vllm_omni/diffusion/lora/loader.py", + "vllm_omni/diffusion/model_loader/hub_prefetch.py", + "vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py", + "vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py", + "vllm_omni/diffusion/models/helios/pipeline_helios.py", + "vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py", + "vllm_omni/diffusion/models/ltx2/ltx2_components.py", + "vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py", + "vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py", + "vllm_omni/diffusion/models/pi0/pipeline_pi0.py", + "vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py", + "vllm_omni/diffusion/models/soulx_singer/utils.py", + "vllm_omni/diffusion/models/utils.py", + "vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py", + "vllm_omni/engine/arg_utils.py", + "vllm_omni/entrypoints/openai/api_server.py", + "vllm_omni/entrypoints/openai/serving_speech.py", + "vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py", + "vllm_omni/experimental/fullduplex/personaplex/runtime.py", + "vllm_omni/experimental/fullduplex/personaplex/serving/server.py", + "vllm_omni/model_executor/model_loader/weight_utils.py", + "vllm_omni/model_executor/models/audex/checkpoint.py", + "vllm_omni/model_executor/models/bagel/bagel.py", + "vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py", + "vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py", + "vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py", + "vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py", + "vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py", + "vllm_omni/model_executor/models/glm_tts/glm_tts.py", + "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py", + "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py", + "vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py", + "vllm_omni/model_executor/models/indextts2/preprocess_utils.py", + "vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py", + "vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py", + "vllm_omni/model_executor/models/ming_tts/speaker_extractor.py", + "vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py", + "vllm_omni/model_executor/models/omnivoice/omnivoice.py", + "vllm_omni/model_executor/models/personaplex/personaplex_mimi.py", + "vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py", + "vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py", + "vllm_omni/transformers_utils/configs/higgs_audio_v3.py", + }, + ), +} + + +def _normalize_path(path: str) -> str: + return path.replace("\\", "/") + + +def check_file(path: str) -> int: + path = _normalize_path(path) + with open(path, encoding="utf-8") as f: + content = f.read() + return_code = 0 + # Check all patterns in the whole file + for import_name, forbidden_import in CHECK_IMPORTS.items(): + # Skip files that are allowed for this import + if path in forbidden_import.allowed_files: + continue + # Skip directories that are allowed for this import + if any(path.startswith(prefix) for prefix in forbidden_import.allowed_dirs): + continue + # Search for forbidden imports + for match in re.finditer(forbidden_import.pattern, content, re.MULTILINE): + # Check if it's allowed + if forbidden_import.allowed_pattern.match(match.group()): + continue + # Skip matches inside a comment + line_start = content.rfind("\n", 0, match.start()) + 1 + if "#" in content[line_start : match.start()]: + continue + # Calculate line number from match position + line_num = content[: match.start() + 1].count("\n") + 1 + print( + f"{path}:{line_num}: " + "\033[91merror:\033[0m " # red color + f"Found forbidden import: {import_name}. {forbidden_import.tip}" + ) + return_code = 1 + return return_code + + +def main(): + returncode = 0 + for path in sys.argv[1:]: + returncode |= check_file(path) + return returncode + + +def test_regex(): + def matches(rule: str, content: str) -> bool: + return bool(re.search(CHECK_IMPORTS[rule].pattern, content, re.MULTILINE)) + + pickle_cases = [ + # Should match + ("import pickle", True), + ("import cloudpickle", True), + ("import pickle as pkl", True), + ("import cloudpickle as cpkl", True), + ("from pickle import *", True), + ("from cloudpickle import dumps", True), + ("from pickle import dumps, loads", True), + ("from cloudpickle import (dumps, loads)", True), + (" import pickle", True), + ("\timport cloudpickle", True), + ("from pickle import loads", True), + ("import pickle, os", True), + ("import os, pickle", True), + ("import os, pickle as pkl", True), + ("import pickle as pkl, os", True), + ("import os.path, cloudpickle", True), + ("import cloudpickle, json", True), + # Should not match + ("import somethingelse", False), + ("from somethingelse import pickle", False), + ("# import pickle", False), + ("print('import pickle')", False), + ("import pickleas as asdf", False), + ] + for i, (content, should_match) in enumerate(pickle_cases): + result = matches("pickle/cloudpickle", content) + assert result == should_match, f"pickle case {i} failed: {content!r} (expected {should_match}, got {result})" + + tilelang_cases = [ + # Should match + ("import tilelang", True), + ("import tilelang.language as T", True), + ("from tilelang.jit import JITImpl", True), + ("from tilelang.jit.kernel import JITKernel", True), + # Should not match: indented (local) imports are allowed, mirroring + # the "triton" rule, so mocked-module test imports are not flagged. + (" import tilelang", False), + (" from tilelang.jit import JITImpl", False), + ("from vllm.tilelang_utils import tilelang", False), + ("from vllm.tilelang_utils import T", False), + ("from vllm.tilelang_utils import T, tilelang", False), + ("from vllm.tilelang_utils import tilelang, T", False), + ("import tilelang_kernels", False), + ] + for i, (content, should_match) in enumerate(tilelang_cases): + result = matches("tilelang", content) + assert result == should_match, f"tilelang case {i} failed: {content!r} (expected {should_match}, got {result})" + + hf_cases = [ + # Should match + ("from huggingface_hub import snapshot_download", True), + ("from huggingface_hub import hf_hub_download", True), + ("from huggingface_hub import HfApi", True), + ("from huggingface_hub import HfFileSystem", True), + ("from huggingface_hub import list_repo_files", True), + ("from huggingface_hub import try_to_load_from_cache", True), + (" from huggingface_hub import snapshot_download", True), + ("from huggingface_hub import PyTorchModelHubMixin, hf_hub_download", True), + ("from huggingface_hub import (snapshot_download)", True), + # Parenthesized multi-line import must not bypass the hook + ("from huggingface_hub import (\n snapshot_download,\n)", True), + ( + "from huggingface_hub import (\n PyTorchModelHubMixin,\n HfApi,\n)", + True, + ), + # Should not match + ("import huggingface_hub", False), + ("import huggingface_hub as hf", False), + ("from huggingface_hub import PyTorchModelHubMixin", False), + ("from huggingface_hub.constants import HF_HUB_CACHE", False), + ("from huggingface_hub.utils import EntryNotFoundError", False), + ("from vllm.transformers_utils.repo_utils import hf_api", False), + ("from huggingface_hub import (\n PyTorchModelHubMixin,\n)", False), + ("# from huggingface_hub import snapshot_download", False), + ] + for i, (content, should_match) in enumerate(hf_cases): + result = matches("huggingface_hub repo API", content) + assert result == should_match, ( + f"huggingface_hub case {i} failed: {content!r} (expected {should_match}, got {result})" + ) + + print("All regex tests passed.") + + +if __name__ == "__main__": + if "--test-regex" in sys.argv: + test_regex() + else: + sys.exit(main()) diff --git a/tools/pre_commit/check_pickle_imports.py b/tools/pre_commit/check_pickle_imports.py deleted file mode 100644 index 819a7c8c3..000000000 --- a/tools/pre_commit/check_pickle_imports.py +++ /dev/null @@ -1,83 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import sys - -import regex as re - -# List of files (relative to repo root) that are allowed to import pickle or -# cloudpickle -# -# STOP AND READ BEFORE YOU ADD ANYTHING ELSE TO THIS LIST: -# The pickle and cloudpickle modules are known to be unsafe when deserializing -# data from potentially untrusted parties. They have resulted in multiple CVEs -# for vLLM and numerous vulnerabilities in the Python ecosystem more broadly. -# Before adding new uses of pickle/cloudpickle, please consider safer -# alternatives like msgpack or pydantic that are already in use in vLLM. Only -# add to this list if absolutely necessary and after careful security review. -ALLOWED_FILES = { - "tests/helpers/process.py", - "vllm_omni/diffusion/distributed/group_coordinator.py", - "tests/diffusion/attention/test_attention_sp.py", -} - -PICKLE_RE = re.compile( - r"^\s*(import\s+(pickle|cloudpickle)(\s|$|\sas)" - r"|from\s+(pickle|cloudpickle)\s+import\b)" -) - - -def scan_file(path: str) -> int: - with open(path, encoding="utf-8") as f: - for i, line in enumerate(f, 1): - if PICKLE_RE.match(line): - print( - f"{path}:{i}: " - "\033[91merror:\033[0m " # red color - "Found pickle/cloudpickle import" - ) - return 1 - return 0 - - -def main(): - returncode = 0 - for filename in sys.argv[1:]: - if filename in ALLOWED_FILES: - continue - returncode |= scan_file(filename) - return returncode - - -def test_regex(): - test_cases = [ - # Should match - ("import pickle", True), - ("import cloudpickle", True), - ("import pickle as pkl", True), - ("import cloudpickle as cpkl", True), - ("from pickle import *", True), - ("from cloudpickle import dumps", True), - ("from pickle import dumps, loads", True), - ("from cloudpickle import (dumps, loads)", True), - (" import pickle", True), - ("\timport cloudpickle", True), - ("from pickle import loads", True), - # Should not match - ("import somethingelse", False), - ("from somethingelse import pickle", False), - ("# import pickle", False), - ("print('import pickle')", False), - ("import pickleas as asdf", False), - ] - for i, (line, should_match) in enumerate(test_cases): - result = bool(PICKLE_RE.match(line)) - assert result == should_match, f"Test case {i} failed: '{line}' (expected {should_match}, got {result})" - print("All regex tests passed.") - - -if __name__ == "__main__": - if "--test-regex" in sys.argv: - test_regex() - else: - sys.exit(main()) diff --git a/tools/pre_commit/check_spdx_header.py b/tools/pre_commit/check_spdx_header.py new file mode 100644 index 000000000..5bfb26034 --- /dev/null +++ b/tools/pre_commit/check_spdx_header.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Require SPDX license headers on source files. + +Ported from vllm/tools/pre_commit/check_spdx_header.py. The copyright line +uses the vLLM-Omni project name rather than upstream vLLM. Shell scripts +(``.sh``) are included so wrappers such as ``shellcheck.sh`` cannot keep a +stale upstream copyright that Python/Rust/proto files would have rewritten. +``.pyi`` stubs use the same ``#`` header style as ``.py``. +""" + +import sys +from dataclasses import dataclass +from enum import Enum +from pathlib import Path + +# Older files still carry the upstream vLLM copyright line. Treat it as stale +# and rewrite it to the Omni line instead of inserting a duplicate header. +LEGACY_COPYRIGHT_TEXT = "SPDX-FileCopyrightText: Copyright contributors to the vLLM project" + + +@dataclass(frozen=True) +class HeaderStyle: + """Comment syntax and preamble handling for an SPDX header.""" + + comment_prefix: str + preserve_shebang: bool = False + + +class SPDXStatus(Enum): + """SPDX header status enumeration""" + + EMPTY = "empty" # empty __init__.py + COMPLETE = "complete" + MISSING_LICENSE = "missing_license" # Only has copyright line + MISSING_COPYRIGHT = "missing_copyright" # Only has license line + MISSING_BOTH = "missing_both" # Completely missing + + +LICENSE_TEXT = "SPDX-License-Identifier: Apache-2.0" +COPYRIGHT_TEXT = "SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project" +FILE_STYLES = { + ".py": HeaderStyle("#", preserve_shebang=True), + ".pyi": HeaderStyle("#"), + ".sh": HeaderStyle("#", preserve_shebang=True), + ".rs": HeaderStyle("//"), + ".proto": HeaderStyle("//"), +} + + +def file_style(file_path): + """Return the declared header style for a file.""" + suffix = Path(file_path).suffix + try: + return FILE_STYLES[suffix] + except KeyError: + raise ValueError(f"Unsupported file type: {file_path}") from None + + +def spdx_header(style): + """Return the SPDX header for a file style.""" + license_line = f"{style.comment_prefix} {LICENSE_TEXT}" + copyright_line = f"{style.comment_prefix} {COPYRIGHT_TEXT}" + return license_line, copyright_line + + +def legacy_copyright_line(style): + """Return the stale upstream vLLM copyright line for a file style.""" + return f"{style.comment_prefix} {LEGACY_COPYRIGHT_TEXT}" + + +def header_insertion_index(style, lines): + """Return the line index where a missing header should be inserted.""" + if style.preserve_shebang and lines and lines[0].startswith("#!"): + return 1 + return 0 + + +def check_spdx_header_status(file_path): + """Check SPDX header status of the file""" + style = file_style(file_path) + license_line, copyright_line = spdx_header(style) + with open(file_path, encoding="UTF-8") as file: + lines = file.readlines() + if not lines: + # Empty file + return SPDXStatus.EMPTY + + has_license = False + has_copyright = False + + # Check all lines for SPDX headers (not just the first two) + for raw_line in lines: + line = raw_line.strip() + if line == license_line: + has_license = True + elif line == copyright_line: + has_copyright = True + + # Determine status based on what we found + if has_license and has_copyright: + return SPDXStatus.COMPLETE + elif has_license and not has_copyright: + # Only has license line (or a stale vLLM copyright line) + return SPDXStatus.MISSING_COPYRIGHT + elif not has_license and has_copyright: + # Only has copyright line + return SPDXStatus.MISSING_LICENSE + else: + # Completely missing both lines + return SPDXStatus.MISSING_BOTH + + +def add_header(file_path, status): + """Add or supplement SPDX header based on status""" + style = file_style(file_path) + license_line, copyright_line = spdx_header(style) + stale_copyright_line = legacy_copyright_line(style) + full_spdx_header = f"{license_line}\n{copyright_line}" + with open(file_path, "r+", encoding="UTF-8") as file: + lines = file.readlines() + file.seek(0, 0) + file.truncate() + + if status == SPDXStatus.MISSING_BOTH: + # Completely missing, add complete header + insertion_index = header_insertion_index(style, lines) + file.writelines(lines[:insertion_index]) + file.write(full_spdx_header + "\n") + remaining_lines = lines[insertion_index:] + if remaining_lines and remaining_lines[0].strip(): + file.write("\n") + file.writelines(remaining_lines) + + elif status == SPDXStatus.MISSING_COPYRIGHT: + # Only has license line, need to add copyright line. + # Rewrite a stale upstream vLLM copyright line in place. + replaced = False + for i, line in enumerate(lines): + if line.strip() == stale_copyright_line: + lines[i] = f"{copyright_line}\n" + replaced = True + break + if not replaced: + for i, line in enumerate(lines): + if line.strip() == license_line: + lines.insert(i + 1, f"{copyright_line}\n") + break + + file.writelines(lines) + + elif status == SPDXStatus.MISSING_LICENSE: + # Only has copyright line, need to add license line + # Find the copyright line and add license line before it + for i, line in enumerate(lines): + if line.strip() == copyright_line: + # Insert license line before copyright line + lines.insert(i, f"{license_line}\n") + break + file.writelines(lines) + + +def main(): + """Main function""" + files_missing_both = [] + files_missing_copyright = [] + files_missing_license = [] + + for file_path in sys.argv[1:]: + if Path(file_path).suffix not in FILE_STYLES: + continue + status = check_spdx_header_status(file_path) + + if status == SPDXStatus.MISSING_BOTH: + files_missing_both.append(file_path) + elif status == SPDXStatus.MISSING_COPYRIGHT: + files_missing_copyright.append(file_path) + elif status == SPDXStatus.MISSING_LICENSE: + files_missing_license.append(file_path) + else: + continue + + # Collect all files that need fixing + all_files_to_fix = files_missing_both + files_missing_copyright + files_missing_license + if all_files_to_fix: + print("The following files are missing the SPDX header:") + if files_missing_both: + for file_path in files_missing_both: + print(f" {file_path}") + add_header(file_path, SPDXStatus.MISSING_BOTH) + + if files_missing_copyright: + for file_path in files_missing_copyright: + print(f" {file_path}") + add_header(file_path, SPDXStatus.MISSING_COPYRIGHT) + if files_missing_license: + for file_path in files_missing_license: + print(f" {file_path}") + add_header(file_path, SPDXStatus.MISSING_LICENSE) + + sys.exit(1 if all_files_to_fix else 0) + + +if __name__ == "__main__": + main() diff --git a/tools/pre_commit/check_test_marks.py b/tools/pre_commit/check_test_marks.py index 2afef8b6e..6af0ae825 100644 --- a/tools/pre_commit/check_test_marks.py +++ b/tools/pre_commit/check_test_marks.py @@ -1,52 +1,45 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Run a pre-commit hook that fails if test files are modified or added that (probably) never run in the CI. For now, this means that every tests file needs to have a CI level marker (e.g., core_model, advanced_model, full_model, local_model, slow, etc) and hardware mark / helper so that we ensure mutated tests will actually be selected as long as there are pytest commands pointing at the right paths. + +SKU markers (``H100``, ``L4``, … tagged ``[hardware-resource]`` in +``pyproject.toml``) must be applied via ``hardware_test(`` / ``hardware_marks(`` +so ``cards_{n}`` is attached. Direct ``pytest.mark.H100`` is rejected. + +Platform names allowed as ``pytest.mark.cpu`` / ``cuda`` come from +``get_supported_platforms()`` (``[hardware-platform]`` in ``pyproject.toml``). +CI level names come from ``get_level_markers()`` (``[ci-level]``). """ +from __future__ import annotations + +import importlib.util import os import re import sys +from functools import lru_cache +from pathlib import Path +from types import ModuleType -# CI level markers -LEVEL_MARKERS = ("core_model", "advanced_model", "full_model", "local_model", "slow") - -# Hardware markers. These are platforms and resources that CI filters on. -# NOTE: If new hardware marks are added to pyproject.toml etc, we need -# to add them here as well. -HARDWARE_MARKERS = ( - "cpu", - "gpu", - "cuda", - "rocm", - "xpu", - "npu", - "musa", - "H100", - "L4", - "MI325", - "B60", - "S5000", - "A2", - "A3", -) - -# Helpers from tests/helpers/mark.py that auto-apply hardware marks. +# Helpers from tests/helpers/mark.py that auto-apply hardware + cards_* marks. HARDWARE_HELPERS = ("hardware_test", "hardware_marks") +# The helper implementation is the only file allowed to write pytest.mark.. +_ALLOWED_DIRECT_SKU_FILES = frozenset({"tests/helpers/mark.py"}) + # Match mark.X since we could also do `from pytest import mark`. # \b prevents matching prefixes (e.g., mark.slow vs mark.slow_test). -LEVEL_RE = re.compile(r"mark\.(?:" + "|".join(LEVEL_MARKERS) + r")\b") -HARDWARE_RE = re.compile(r"mark\.(?:" + "|".join(HARDWARE_MARKERS) + r")\b") HELPER_RE = re.compile(r"(?:" + "|".join(HARDWARE_HELPERS) + r")\s*\(") MISSING_LEVEL_MARKER = "Level" MISSING_HARDWARE_MARKER = "Hardware" +DIRECT_SKU_MARKER = "Direct SKU" # Check if a file is located under tests/ and matches test_.py # or _test.py, since pytest technically collects on both. @@ -54,9 +47,68 @@ MISSING_HARDWARE_MARKER = "Hardware" TEST_FILE_RE = re.compile(r"^tests/(?:.*/)?(?:test_[^/]*\.py$|[^/]*_test\.py$)") +def _normalize_path(path: str) -> str: + return path.replace("\\", "/") + + +def _repo_root() -> Path: + return Path(__file__).resolve().parents[2] + + +@lru_cache(maxsize=1) +def _mark_module() -> ModuleType: + """Load ``mark.py`` by file path (no pytest/vllm; skip helpers ``__init__``).""" + path = _repo_root() / "tests" / "helpers" / "mark.py" + spec = importlib.util.spec_from_file_location("_vllm_omni_mark", path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load mark helpers from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _mark_name_re(names: tuple[str, ...]) -> re.Pattern[str]: + if not names: + return re.compile(r"(?!)") + return re.compile(r"mark\.(?:" + "|".join(re.escape(n) for n in names) + r")\b") + + +@lru_cache(maxsize=1) +def level_markers() -> tuple[str, ...]: + """CI level names tagged ``[ci-level]`` (``core_model``, ``slow``, …).""" + return tuple(sorted(_mark_module().get_level_markers())) + + +@lru_cache(maxsize=1) +def platform_markers() -> tuple[str, ...]: + """Names tests may apply as ``pytest.mark.cpu`` / ``cuda`` (not SKUs).""" + return tuple(sorted(_mark_module().get_supported_platforms())) + + +@lru_cache(maxsize=1) +def sku_markers() -> tuple[str, ...]: + """SKU marker names tagged ``[hardware-resource]`` in ``pyproject.toml``.""" + return tuple(sorted(_mark_module().get_hardware_mark_list())) + + +@lru_cache(maxsize=1) +def _level_re() -> re.Pattern[str]: + return _mark_name_re(level_markers()) + + +@lru_cache(maxsize=1) +def _platform_re() -> re.Pattern[str]: + return _mark_name_re(platform_markers()) + + +@lru_cache(maxsize=1) +def _sku_mark_re() -> re.Pattern[str]: + return _mark_name_re(sku_markers()) + + def is_test_file(path: str) -> bool: """Determine whether or not a path is pointing at a test file or not.""" - return bool(TEST_FILE_RE.match(path)) + return bool(TEST_FILE_RE.match(_normalize_path(path))) def read_test_file(path: str) -> str | None: @@ -69,22 +121,31 @@ def read_test_file(path: str) -> str | None: def has_level_marker(contents: str) -> bool: """Check if file contents contain at least one CI level marker.""" - return bool(LEVEL_RE.search(contents)) + return bool(_level_re().search(contents)) def has_hardware_marker(contents: str) -> bool: - """Check if file contents contain any hardware marker or helpers.""" - return bool(HARDWARE_RE.search(contents) or HELPER_RE.search(contents)) + """Check if file contents contain a platform marker or hardware helper.""" + return bool(_platform_re().search(contents) or HELPER_RE.search(contents)) + + +def has_direct_sku_marker(path: str, contents: str) -> bool: + """True when a test applies ``pytest.mark.`` instead of the helpers.""" + if _normalize_path(path) in _ALLOWED_DIRECT_SKU_FILES: + return False + return bool(_sku_mark_re().search(contents)) def get_files_missing_markers( staged_files: list[str], ) -> dict[str, list[str]]: - """Return a dict mapping file path to list of missing marker types.""" + """Return a dict mapping file path to list of missing / invalid marker types.""" results: dict[str, list[str]] = {} for path in staged_files: if is_test_file(path) and (contents := read_test_file(path)) is not None: missing = [] + if has_direct_sku_marker(path, contents): + missing.append(DIRECT_SKU_MARKER) if not has_level_marker(contents): missing.append(MISSING_LEVEL_MARKER) if not has_hardware_marker(contents): @@ -99,14 +160,17 @@ if __name__ == "__main__": if missing: file_lines = "\n".join(f" - {path} [{' and '.join(problems)}]" for path, problems in missing.items()) + sku = ", ".join(sku_markers()) print( "\033[91merror:\033[0m test files are missing pytest marks " - "required for Buildkite CI collection.\n\n" - f"Level marks, e.g.: {', '.join(LEVEL_MARKERS[:4])}\n" - f"Hardware marks, e.g.: {', '.join(HARDWARE_MARKERS[:4])}, ...\n" - f" or helpers: {', '.join(HARDWARE_HELPERS)}\n\n" + "required for Buildkite CI collection, or apply SKU marks directly.\n\n" + f"Level marks, e.g.: {', '.join(level_markers()[:4])}\n" + f"Hardware marks, e.g.: {', '.join(platform_markers()[:4])}, ...\n" + f" or helpers: {', '.join(HARDWARE_HELPERS)}\n" + f"Do not write pytest.mark. ({sku}). " + "Use hardware_test(...) / hardware_marks(...) so cards_* is attached.\n\n" "The following files are missing marks:\n" f"{file_lines}\n\n" - "To skip: SKIP=check-test-ci-coverage git commit ..." + "To skip: SKIP=check-mark git commit ..." ) sys.exit(1) diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py new file mode 100644 index 000000000..a298baae5 --- /dev/null +++ b/tools/pre_commit/check_torch_cuda.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Ban direct torch.cuda.* usage outside an allowlist. + +Modeled on vllm/tools/pre_commit/check_torch_cuda.py. Call sites should use +current_omni_platform / OmniPlatform instead of CUDA-only torch.cuda APIs. +""" + +import sys + +import regex as re + +# Match `torch.cuda.xxx` device/memory/seed helpers that have portable Omni +# platform equivalents. CUDAGraph / is_available / stream capture are not in +# this list (same split as upstream vLLM). +_TORCH_CUDA_PATTERNS = [ + r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|mem_get_info|set_device)\b", + # Match torch.cuda.device(...) including args. `device\()\b` misses + # `device()` (no word char after `)`) and `device(0)`. + r"\btorch\.cuda\.device\s*\(", + r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b", + r"\bwith\storch\.cuda\.device\b", + # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally + r"\bcuda_device_count_stateless\(\)\b", +] + +# Platform adapters must call native torch.cuda / torch.npu / torch.xpu APIs. +ALLOWED_PREFIXES = { + "vllm_omni/platforms/", +} + +# Grandfathered call sites. Remove a path after migrating it to +# current_omni_platform / OmniPlatform. Do not add new files without review. +ALLOWED_FILES = { + "benchmarks/kernels/mot_linear_benchmarks.py", + "examples/online_serving/qwen2_5_omni/gradio_demo.py", + "examples/online_serving/qwen3_omni/gradio_demo.py", + "tests/dfx/reliability/helpers.py", + "tests/diffusion/distributed/test_cfg_parallel.py", + "tests/diffusion/distributed/test_pipeline_parallel.py", + "tests/e2e/accuracy/wan22_i2v/run_wan22_i2v_diffusers_cp.py", + "tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py", + "tests/e2e/features/fullduplex/test_personaplex_elastic.py", + "tests/entrypoints/test_omni_sleep_mode.py", + "tools/configure_stage_memory.py", + "vllm_omni/diffusion/models/bagel/pipeline_bagel.py", + "vllm_omni/diffusion/models/lance/pipeline_lance.py", + "vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py", + "vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py", + "vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py", + "vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py", + "vllm_omni/experimental/ar_diffusion/runner.py", + "vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py", + "vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py", + "vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py", + "vllm_omni/worker/gpu_ar_model_runner.py", +} + +_GENERIC_TIP = ( + "Found torch.cuda API call. Use current_omni_platform / OmniPlatform " + "(empty_cache, synchronize, get_device_count, current_device, " + "get_device_memory, set_device) instead." +) +_SEED_TIP = ( + "Found CUDA manual-seed API call. Use " + "torch.Generator(device=current_omni_platform.device_type).manual_seed(...) " + "instead." +) + + +def _normalize_path(path: str) -> str: + return path.replace("\\", "/") + + +def _is_allowed(path: str) -> bool: + if any(path.startswith(prefix) for prefix in ALLOWED_PREFIXES): + return True + return path in ALLOWED_FILES + + +def scan_file(path: str) -> int: + with open(path, encoding="utf-8") as f: + content = f.read() + for pattern in _TORCH_CUDA_PATTERNS: + for match in re.finditer(pattern, content, re.MULTILINE): + line_start = content.rfind("\n", 0, match.start()) + 1 + if "#" in content[line_start : match.start()]: + continue + line_num = content[: match.start() + 1].count("\n") + 1 + matched_text = match.group(0) + tip = _SEED_TIP if "manual_seed" in matched_text else _GENERIC_TIP + print(f"{path}:{line_num}: \033[91merror:\033[0m {tip}") + return 1 + return 0 + + +def main(): + returncode = 0 + for filename in sys.argv[1:]: + path = _normalize_path(filename) + if _is_allowed(path): + continue + returncode |= scan_file(path) + return returncode + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/pre_commit/check_tts_adapter.py b/tools/pre_commit/check_tts_adapter.py index 626323b4f..d46928662 100644 --- a/tools/pre_commit/check_tts_adapter.py +++ b/tools/pre_commit/check_tts_adapter.py @@ -6,10 +6,8 @@ Per-model behaviour is being moved out of ``serving_speech.py`` and into nothing stopped a new model from adding one more ``self._tts_model_type == ...`` branch faster than the migration removed them. -This gate makes the two debt counters monotonically non-increasing: - -1. ``self._tts_model_type`` comparisons in ``serving_speech.py`` -2. entries in ``tts_adapters.LEGACY_TTS_DETECTORS`` (models with no adapter) +This gate makes the number of ``self._tts_model_type`` comparisons in +``serving_speech.py`` monotonically non-increasing. Lowering a budget after removing branches is expected and encouraged — the check insists on it, so the recorded numbers cannot silently drift upward. Raising one @@ -28,11 +26,9 @@ from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[2] SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py") -ADAPTERS_INIT = Path("vllm_omni/entrypoints/openai/tts_adapters/__init__.py") # Budgets. These may only go DOWN. See the module docstring. -MAX_MODEL_TYPE_BRANCHES = 27 -MAX_LEGACY_DETECTORS = 0 +MAX_MODEL_TYPE_BRANCHES = 9 _GUIDANCE = """ Move the per-model behaviour into its adapter under @@ -109,20 +105,6 @@ def _model_type_branches(tree: ast.AST) -> list[tuple[int, str]]: return found -def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]: - """Every ``LegacyDetector(...)`` construction.""" - found: list[tuple[int, str]] = [] - for node in ast.walk(tree): - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "LegacyDetector": - name = "?" - for kw in node.keywords: - if kw.arg == "name" and isinstance(kw.value, ast.Constant): - name = str(kw.value.value) - found.append((node.lineno, name)) - found.sort() - return found - - def _check( label: str, path: Path, @@ -151,21 +133,18 @@ def _check( def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) - # pre-commit passes the staged filenames; we always audit the same two files, + # pre-commit passes the staged filenames; we always audit the same file, # so they are accepted and ignored. parser.add_argument("filenames", nargs="*") parser.parse_args(argv) serving_path = REPO_ROOT / SERVING_SPEECH - adapters_path = REPO_ROOT / ADAPTERS_INIT - for path in (serving_path, adapters_path): - if not path.is_file(): - print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr) - return 1 + if not serving_path.is_file(): + print(f"check_tts_adapter: expected file not found: {serving_path}", file=sys.stderr) + return 1 notices: list[str] = [] branches = _model_type_branches(ast.parse(serving_path.read_text())) - detectors = _legacy_detectors(ast.parse(adapters_path.read_text())) errors = _check( "`_tts_model_type` comparisons", @@ -176,15 +155,6 @@ def main(argv: list[str] | None = None) -> int: "MAX_MODEL_TYPE_BRANCHES", notices, ) - errors += _check( - "unmigrated model types in LEGACY_TTS_DETECTORS", - ADAPTERS_INIT, - len(detectors), - MAX_LEGACY_DETECTORS, - [f"line {lineno}: {name}" for lineno, name in detectors], - "MAX_LEGACY_DETECTORS", - notices, - ) for notice in notices: print(f"check_tts_adapter: {notice}", file=sys.stderr) diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py new file mode 100644 index 000000000..e24f62444 --- /dev/null +++ b/tools/pre_commit/mypy.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Run mypy on changed files. + +Ported from vllm/tools/pre_commit/mypy.py. Group files so tests can use +``follow_imports=skip`` without pulling that setting into library code. + +Usage: + python tools/pre_commit/mypy.py + +Args: + python_version: Python version to use (e.g. "3.10") or "local" to use + the interpreter running this script. + changed_files: List of changed files to check. +""" + +from __future__ import annotations + +import subprocess +import sys + +import regex as re + +# tests/ uses follow_imports=skip so third-party and library errors do not +# leak into test type-checking. After a directory is clean under skip, it +# can join the default group (follow_imports from pyproject.toml). +SEPARATE_GROUPS = [ + "tests", +] + +# Model trees are large third-party ports. Drop a prefix from this list when +# that directory is ready to be type-checked (same ratchet as upstream vLLM). +EXCLUDE = [ + r"vllm_omni/model_executor/models/", + r"vllm_omni/diffusion/models/", +] + + +def _normalize_path(path: str) -> str: + return path.replace("\\", "/") + + +def group_files(changed_files: list[str]) -> dict[str, list[str]]: + """Group changed files into different mypy calls.""" + exclude_pattern = re.compile(f"^{'|'.join(EXCLUDE)}.*") + file_groups: dict[str, list[str]] = {"": []} + file_groups.update({k: [] for k in SEPARATE_GROUPS}) + # Longest path first so a sub-directory is not shadowed by its parent. + separate_groups = sorted(SEPARATE_GROUPS, key=len, reverse=True) + for changed_file in changed_files: + changed_file = _normalize_path(changed_file) + if exclude_pattern.match(changed_file): + continue + for directory in separate_groups: + if re.match(f"^{directory}.*", changed_file): + file_groups[directory].append(changed_file) + break + else: + if changed_file.startswith(("vllm_omni/", "tests/")): + file_groups[""].append(changed_file) + return file_groups + + +def mypy( + targets: list[str], + python_version: str | None, + follow_imports: str | None, + file_group: str, +) -> int: + """Run mypy on the given targets.""" + args = ["mypy"] + if python_version is not None: + args += ["--python-version", python_version] + if follow_imports is not None: + args += ["--follow-imports", follow_imports] + print(f"$ {' '.join(args)} {file_group}") + return subprocess.run(args + targets, check=False).returncode + + +def main() -> int: + python_version = sys.argv[1] + file_groups = group_files(sys.argv[2:]) + + if python_version == "local": + python_version = f"{sys.version_info.major}.{sys.version_info.minor}" + + returncode = 0 + for file_group, changed_files in file_groups.items(): + follow_imports = None if file_group == "" else "skip" + if changed_files: + returncode |= mypy(changed_files, python_version, follow_imports, file_group) + return returncode + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/pre_commit/shellcheck.sh b/tools/pre_commit/shellcheck.sh new file mode 100755 index 000000000..7794ffe51 --- /dev/null +++ b/tools/pre_commit/shellcheck.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# +# Lint bash scripts for undefined vars, quoting, and similar bugs. +# Do not download a binary: use a distro/package-manager install so the +# artifact is signed (apt/dnf/brew) or at least an explicit local install. +set -euo pipefail + +is_windows_exe() { + [[ "$1" == *.exe ]] +} + +is_git_bash() { + case "$(uname -s)" in + MINGW*|MSYS*|CYGWIN*) return 0 ;; + *) return 1 ;; + esac +} + +find_native_shellcheck() { + local cand + cand="$(command -v shellcheck 2>/dev/null || true)" + if [ -n "$cand" ] && ! is_windows_exe "$cand"; then + echo "$cand" + return 0 + fi + return 1 +} + +install_hint() { + echo "Please install shellcheck with your package manager, then re-run:" + echo " Debian/Ubuntu/WSL: sudo apt-get install shellcheck" + echo " Fedora: sudo dnf install ShellCheck" + echo " macOS: brew install shellcheck" + echo " Git Bash: scoop install shellcheck" + echo " https://github.com/koalaman/shellcheck?tab=readme-ov-file#installing" +} + +SHELLCHECK_BIN="" +if SHELLCHECK_BIN="$(find_native_shellcheck)"; then + : +elif is_git_bash && SHELLCHECK_BIN="$(command -v shellcheck.exe 2>/dev/null || true)" && [ -n "$SHELLCHECK_BIN" ]; then + : +else + install_hint + exit 1 +fi + +should_lint() { + local f="${1//\\//}" + f="${f#./}" + case "$f" in + *.sh) ;; + *) return 1 ;; + esac + git check-ignore -q "$f" && return 1 + return 0 +} + +run_shellcheck() { + local f + for f in "$@"; do + if should_lint "$f"; then + "$SHELLCHECK_BIN" -s bash "$f" + fi + done +} + +if [ "$#" -gt 0 ]; then + run_shellcheck "$@" + exit 0 +fi + +# Direct invocation with no args: lint every tracked *.sh. +while IFS= read -r -d '' f || [ -n "$f" ]; do + git check-ignore -q "$f" && continue + "$SHELLCHECK_BIN" -s bash "$f" +done < <(find . -path ./.git -prune -o -name "*.sh" -print0) diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py index 512e86538..abb02b1d5 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Daily-Omni Dataset loader for benchmark. Daily-Omni is an audio-visual reasoning benchmark with 684 videos @@ -7,7 +10,8 @@ Dataset source: https://huggingface.co/datasets/liarliar/Daily-Omni Supports loading QA metadata from: - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments -- HuggingFace datasets (``dataset_path``): legacy online mode +- HuggingFace Hub id (``dataset_path``): resolved with ``snapshot_download`` to the repo's + ``qa.json``, which reads the hub cache and therefore also works under ``HF_HUB_OFFLINE`` Video/audio files normally come from extracted ``Videos.tar``. When ``--daily-omni-video-dir`` is not set, the first request that needs on-disk media downloads that archive from the Hugging Face @@ -24,9 +28,9 @@ Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``? This class therefore inherits only ``BenchmarkDataset`` (minimal: ``dataset_path``, ``random_seed``, ``self.data``) and implements **two explicit loaders**: ``_load_from_local_json`` (default path for air-gapped runs) and ``_load_from_huggingface`` - (optional legacy path for users who prefer ``datasets`` + Hub cache). The latter is **not** + (Hub ids, resolved to the repo's ``qa.json`` through the hub cache). The latter is **not** inheritance; it is the same Hub rows as before, factored into a helper so one class can - serve both deployment modes without mandatory ``datasets`` when using ``qa_json_path``. + serve both deployment modes without mandatory ``datasets``. Usage: from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset @@ -38,7 +42,7 @@ Usage: random_seed=42, ) - # HuggingFace mode (legacy, requires network) + # Hub id mode (works offline once the repo is in the HF hub cache) dataset = DailyOmniDataset( dataset_path="liarliar/Daily-Omni", dataset_split="train", @@ -210,6 +214,68 @@ def _resolve_hf_cache_snapshot(path: Path) -> Path: return path +#: QA metadata pulled by :func:`ensure_daily_omni_hub_root`. ``Videos.tar`` is deliberately left +#: out: it is ~60GB and only :func:`ensure_daily_omni_hub_videos_dir` needs it, lazily, once a +#: request actually asks for media. +_DAILY_OMNI_QA_ALLOW_PATTERNS = ["qa.json", "QA.json"] + + +def ensure_daily_omni_hub_root(repo_id: str) -> Path: + """Resolve a Daily-Omni Hub id to its snapshot directory, downloading QA metadata if needed. + + Mirrors :func:`...seed_tts_dataset.resolve_seed_tts_root`: ``snapshot_download`` resolves the + repo through the **hub cache**, so an air-gapped run (``HF_HUB_OFFLINE=1``) keeps working as + long as the dataset was pre-staged there. ``datasets.load_dataset`` cannot do this — it only + consults the separate *datasets* cache, which a ``hf download`` never populates. + + Raises: + ImportError: if ``huggingface_hub`` is not installed. + ValueError: if ``repo_id`` is empty. + """ + rid = (repo_id or "").strip() + if not rid: + raise ValueError("repo_id is required to download Daily-Omni from Hugging Face") + + try: + from huggingface_hub import snapshot_download + except ImportError as e: + raise ImportError( + "Install huggingface_hub to load Daily-Omni from the Hub, or pass a local " + "--daily-omni-qa-json / --daily-omni-video-dir mirror." + ) from e + + cache = snapshot_download( + repo_id=rid, + repo_type="dataset", + allow_patterns=_DAILY_OMNI_QA_ALLOW_PATTERNS, + ) + root = _resolve_hf_cache_snapshot(Path(cache).resolve()) + logger.info("Daily-Omni Hub snapshot ready at %s (repo=%s)", root, rid) + return root + + +def _daily_omni_qa_load_error( + repo_id: str | None, + snapshot_error: Exception | None, + datasets_error: Exception, +) -> RuntimeError: + """Compose an actionable error after both QA-metadata sources failed. + + Offline CI hits this when the dataset is missing from the hub cache *or* staged under a + different repo id, so the message names both attempts and the flags that bypass them. + """ + parts = [f"Could not load Daily-Omni QA metadata for {repo_id!r}."] + if snapshot_error is not None: + parts.append(f"huggingface_hub snapshot_download failed: {type(snapshot_error).__name__}: {snapshot_error}") + parts.append(f"datasets.load_dataset failed: {type(datasets_error).__name__}: {datasets_error}") + parts.append( + "For offline runs, pre-stage the dataset in the HF hub cache " + "(`hf download --repo-type dataset `) and make sure the repo id matches, " + "or pass --daily-omni-qa-json / --daily-omni-video-dir pointing at a local mirror." + ) + return RuntimeError(" ".join(parts)) + + def resolve_daily_omni_local_root(dataset_path: str | None) -> Path | None: """Return the local Daily-Omni mirror root for ``dataset_path``, or ``None`` if not local. @@ -433,7 +499,8 @@ class DailyOmniDataset(BenchmarkDataset): QA metadata can be loaded from: - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments - - HuggingFace datasets (``dataset_path``): legacy online mode + - HuggingFace Hub id (``dataset_path``): the repo's ``qa.json`` via ``snapshot_download`` + (hub-cache backed, offline-safe); ``datasets.load_dataset`` only as a parquet fallback Video/audio files normally come from extracted ``Videos.tar``. When ``video_dir`` is not set, the first sample that needs on-disk media downloads that archive from the Hugging Face dataset @@ -444,7 +511,8 @@ class DailyOmniDataset(BenchmarkDataset): qa_json_path: Path to local qa.json file (offline mode, preferred). When provided, ``dataset_path`` and ``dataset_split`` are ignored. dataset_path: HuggingFace dataset path (e.g., "liarliar/Daily-Omni"). Used only if - ``qa_json_path`` is not provided (legacy online mode). + ``qa_json_path`` is not provided; resolved through the hub cache, so a pre-staged + repo also loads with ``HF_HUB_OFFLINE=1``. dataset_split: Dataset split to use (default: "train"). Used only in online mode. random_seed: Random seed for shuffling video_dir: Directory containing extracted video files (default: None; may be filled lazily @@ -562,10 +630,13 @@ class DailyOmniDataset(BenchmarkDataset): def _load_from_local_json(self) -> None: """Load QA data from local JSON file.""" - if not self.qa_json_path.exists(): - raise FileNotFoundError(f"QA JSON file not found: {self.qa_json_path}") + qa_json_path = self.qa_json_path + if qa_json_path is None: + raise RuntimeError("qa_json_path is required for local JSON loading") + if not qa_json_path.exists(): + raise FileNotFoundError(f"QA JSON file not found: {qa_json_path}") - with open(self.qa_json_path, encoding="utf-8") as f: + with open(qa_json_path, encoding="utf-8") as f: data = json.load(f) # Support both list format and dict with "train"/"test" splits @@ -602,21 +673,45 @@ class DailyOmniDataset(BenchmarkDataset): self.data = _ListDatasetIterator(data) def _load_from_huggingface(self) -> None: - """Load QA rows via ``datasets.load_dataset`` (legacy / convenience path). + """Load QA rows for a Hub id, preferring the repo's ``qa.json`` over ``datasets``. - Kept for backward compatibility: callers can still pass ``dataset_path=liarliar/Daily-Omni`` - and get the same parquet-backed rows as the Hub dataset card, with streaming (or - non-streaming if ``no_stream=True``) and shuffle. + The Hub repo ships ``qa.json`` (1,197 rows) next to ``Videos.tar``, so the primary path is + :func:`ensure_daily_omni_hub_root` → :meth:`_load_from_local_json`: ``snapshot_download`` + reads the hub cache and therefore works under ``HF_HUB_OFFLINE``, the same contract as + Seed-TTS. ``datasets.load_dataset`` remains only as a fallback for mirrors + that publish parquet rows instead of ``qa.json``; it consults the *datasets* cache rather + than the hub cache and streams shards over HTTP, so it cannot serve air-gapped runs. This is intentionally **not** implemented by subclassing ``HuggingFaceDataset``: that base always runs Hub ``load_dataset`` from its constructor and expects a Hub id as the primary API; Daily-Omni instead chooses the source in ``load_data()`` (JSON vs Hub) while sharing one ``sample()`` / request-building implementation for both. """ + snapshot_error: Exception | None = None + dataset_path = self.dataset_path + if dataset_path is None: + raise ValueError("dataset_path is required for HuggingFace loading") + try: + root = ensure_daily_omni_hub_root(dataset_path) + except Exception as e: # noqa: BLE001 - any failure just falls through to `datasets` + snapshot_error = e + else: + qa_json = daily_omni_local_qa_json(root) + if qa_json is not None: + self.qa_json_path = qa_json + logger.info("Loading Daily-Omni QA metadata from Hub snapshot: %s", qa_json) + self._load_from_local_json() + return + logger.info( + "Hub snapshot %s ships no qa.json; falling back to `datasets.load_dataset`", + root, + ) + if load_dataset is None: - raise ImportError( - "datasets library is required for HuggingFace mode. " - "Install with: pip install datasets, or use local JSON mode instead." + raise _daily_omni_qa_load_error( + self.dataset_path, + snapshot_error, + ImportError("datasets library is not installed (pip install datasets)"), ) load_kw: dict[str, Any] = { @@ -626,7 +721,10 @@ class DailyOmniDataset(BenchmarkDataset): } if self.dataset_subset is not None: load_kw["name"] = self.dataset_subset - ds = load_dataset(self.dataset_path, **load_kw) + try: + ds = load_dataset(self.dataset_path, **load_kw) + except Exception as e: + raise _daily_omni_qa_load_error(self.dataset_path, snapshot_error, e) from e if not getattr(self, "disable_shuffle", False): ds = ds.shuffle(seed=self.random_seed) self.data = ds @@ -1167,7 +1265,7 @@ class DailyOmniDataset(BenchmarkDataset): num_video_frames, avg_fps = _probe_video_frames_and_fps(video_path) duration = num_video_frames / avg_fps num_seconds = max(1, int(math.ceil(duration))) - second_timestamps = list(range(num_seconds)) + second_timestamps = [float(i) for i in range(num_seconds)] if duration > max_num_frames: timestamps = [round(i * 0.1, 1) for i in range(int(duration / 0.1))] diff --git a/vllm_omni/benchmarks/data_modules/omniinteract_dataset.py b/vllm_omni/benchmarks/data_modules/omniinteract_dataset.py new file mode 100644 index 000000000..04726eee6 --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/omniinteract_dataset.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniInteract dataset discovery for the standard serving benchmark.""" + +from __future__ import annotations + +import json +import logging +import os +import random +import shutil +import tarfile +import tempfile +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from vllm.benchmarks.datasets import BenchmarkDataset, SampleRequest +from vllm.tokenizers import TokenizerLike +from vllm.transformers_utils.repo_utils import hf_fs + +OMNIINTERACT_SUBSETS = ("1q1a", "1q1a_math", "1qna") +DEFAULT_OMNIINTERACT_REPO = "lucky-lance/OmniInteract" +DEFAULT_MAX_VIDEO_DURATION_S = 3600.0 +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class OmniInteractCase: + subset: str + video_rel: str + video_path: Path + annotation_path: Path + scene_type: str + + +@dataclass(frozen=True) +class OmniInteractSessionOptions: + output_root: Path + timeout_s: float + media_timeout_s: float + ref_audio: str + require_response: bool = False + max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S + + +@dataclass(frozen=True) +class OmniInteractPreparedInput: + duration_s: float + pcm16: bytes = field(repr=False) + video_frames: tuple[str | None, ...] = field(repr=False) + ref_audio_data_url: str = field(repr=False) + + +@dataclass +class OmniInteractSampleRequest(SampleRequest): + omniinteract_case: OmniInteractCase | None = None + omniinteract_options: OmniInteractSessionOptions | None = None + omniinteract_prepared_input: OmniInteractPreparedInput | None = field(default=None, repr=False) + + +class OmniInteractDataset(BenchmarkDataset): + def __init__( + self, + *, + data_root: str | None, + dataset_repo: str = DEFAULT_OMNIINTERACT_REPO, + subsets: Sequence[str] = OMNIINTERACT_SUBSETS, + random_seed: int = 0, + disable_shuffle: bool = False, + ) -> None: + super().__init__( + dataset_path=data_root or dataset_repo, + random_seed=random_seed, + disable_shuffle=disable_shuffle, + ) + self.root = resolve_omniinteract_root(data_root, dataset_repo) + self.subsets = tuple(subsets) + + def sample( + self, + tokenizer: TokenizerLike | None, + num_requests: int, + *, + request_id_prefix: str = "", + options: OmniInteractSessionOptions, + **_: Any, + ) -> list[SampleRequest]: + del tokenizer + cases = discover_omniinteract_cases( + self.root, + self.subsets, + num_prompts=num_requests, + seed=self.random_seed, + disable_shuffle=self.disable_shuffle, + ) + return [ + OmniInteractSampleRequest( + prompt="", + prompt_len=0, + expected_output_len=0, + multi_modal_data=None, + request_id=f"{request_id_prefix}{index}", + omniinteract_case=case, + omniinteract_options=options, + ) + for index, case in enumerate(cases) + ] + + +def _data_dir(root: Path) -> Path: + for candidate in (root, root / "data"): + if any((candidate / subset).is_dir() for subset in OMNIINTERACT_SUBSETS): + return candidate + raise FileNotFoundError(f"OmniInteract data not found under {root}") + + +def _confined_path(root: Path, relative: object, *, field: str) -> Path: + text, resolved_root = str(relative or ""), root.resolve() + path = Path(text) + destination = (root / path).resolve() + if not text or path.is_absolute() or ".." in path.parts or not destination.is_relative_to(resolved_root): + raise ValueError(f"Unsafe OmniInteract {field} path: {text!r}") + return destination + + +def _safe_extract(archive: tarfile.TarFile, target: Path) -> None: + root = target.resolve() + members = archive.getmembers() + for member in members: + path = Path(member.name) + if ( + path.is_absolute() + or ".." in path.parts + or not (root / path).resolve().is_relative_to(root) + or not (member.isdir() or member.isfile()) + ): + raise ValueError(f"Unsafe path in OmniInteract archive: {member.name!r}") + archive.extractall(target, members=members, filter="data") + + +def _extract_archive(archive: Path, target: Path) -> Path: + if target.is_symlink(): + raise ValueError(f"Refusing to extract through symlink: {target}") + stat = archive.stat() + published = target / f"{stat.st_size}-{stat.st_mtime_ns}" + if published.is_symlink(): + raise ValueError(f"Refusing to use symlinked extraction: {published}") + if published.is_dir(): + try: + return _data_dir(published) + except FileNotFoundError: + pass + target.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=".tmp-", dir=target)) + try: + with tarfile.open(archive, "r:*") as handle: + _safe_extract(handle, staging) + root = _data_dir(staging) + try: + staging.replace(published) + return published / root.relative_to(staging) + except OSError: + if not published.is_dir(): + raise + return _data_dir(published) + finally: + shutil.rmtree(staging, ignore_errors=True) + + +def resolve_omniinteract_root(data_root: str | None, dataset_repo: str = DEFAULT_OMNIINTERACT_REPO) -> Path: + """Resolve a local tree/archive or download the official archive.""" + + if data_root: + local = Path(data_root).expanduser().resolve() + if local.is_file(): + return _extract_archive(local, local.parent / f".{local.stem}.vllm_omni_extracted") + if not local.is_dir(): + raise FileNotFoundError(f"--dataset-path does not exist: {local}") + try: + return _data_dir(local) + except FileNotFoundError: + for name in ("data.tar.gz", "data.tar"): + if (archive := local / name).is_file(): + return _extract_archive(archive, local / ".vllm_omni_extracted") + raise + + cache = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) + target = cache / "vllm_omni" / "omniinteract" / dataset_repo.replace("/", "__") + downloads = target / ".downloads" + downloads.mkdir(parents=True, exist_ok=True) + filesystem, errors = hf_fs(), [] + for name in ("data.tar.gz", "data.tar"): + archive = downloads / name + try: + if not archive.is_file(): + descriptor, temporary_name = tempfile.mkstemp(prefix=f".{name}.", dir=downloads) + os.close(descriptor) + temporary = Path(temporary_name) + try: + filesystem.get_file(f"datasets/{dataset_repo}/{name}", str(temporary)) + if not temporary.stat().st_size: + raise ValueError("downloaded archive is empty") + temporary.replace(archive) + finally: + temporary.unlink(missing_ok=True) + return _extract_archive(archive, target) + except Exception as exc: # noqa: BLE001, PERF203 - Hub errors vary by version + errors.append(f"{name}: {exc}") + raise FileNotFoundError(f"Could not download OmniInteract from {dataset_repo!r}: {'; '.join(errors)}") + + +def _mapping_cases(root: Path, subset: str) -> list[OmniInteractCase]: + mapping_path = root / "video_json_map.json" + if not mapping_path.is_file(): + raise FileNotFoundError(f"Missing OmniInteract mapping: {mapping_path}") + mapping = json.loads(mapping_path.read_text()) + entries = mapping.get("entries") if isinstance(mapping, dict) else None + if not isinstance(entries, list): + raise ValueError(f"Invalid OmniInteract mapping: {mapping_path}") + cases = [] + for row in entries: + if not isinstance(row, dict): + raise ValueError(f"Invalid OmniInteract mapping row in {mapping_path}") + video_rel, annotation_rel = str(row.get("video") or ""), str(row.get("annotation") or "") + video = _confined_path(root, video_rel, field="video") + annotation = _confined_path(root, annotation_rel, field="annotation") + if not video.is_file() or not annotation.is_file(): + raise FileNotFoundError(f"OmniInteract mapping references missing files: {video_rel!r}, {annotation_rel!r}") + cases.append( + OmniInteractCase( + subset, + video_rel, + video, + annotation, + str(row.get("scene_type") or "multi_turn").lower(), + ) + ) + return cases + + +def _one_to_many_cases(root: Path) -> list[OmniInteractCase]: + videos, annotations = root / "videos_bench", root / "annotations" + if not videos.is_dir() or not annotations.is_dir(): + raise FileNotFoundError(f"Invalid OmniInteract 1qna layout under {root}") + cases = [] + for video in sorted(videos.rglob("*.mp4")): + if not video.resolve().is_relative_to(videos.resolve()): + raise ValueError(f"Unsafe OmniInteract video path: {video}") + relative = video.relative_to(videos) + annotation = (annotations / relative).with_suffix(".json").resolve() + if not annotation.is_relative_to(annotations.resolve()): + raise ValueError(f"Unsafe OmniInteract annotation path: {relative}") + if not annotation.is_file(): + raise FileNotFoundError(f"Missing OmniInteract annotation: {annotation}") + cases.append(OmniInteractCase("1qna", f"videos_bench/{relative.as_posix()}", video, annotation, "1qna")) + return cases + + +def discover_omniinteract_cases( + root: Path, + subsets: Sequence[str], + *, + num_prompts: int, + seed: int = 0, + disable_shuffle: bool = False, +) -> list[OmniInteractCase]: + invalid = set(subsets) - set(OMNIINTERACT_SUBSETS) + if invalid: + raise ValueError(f"Unsupported OmniInteract subsets: {sorted(invalid)}") + if not subsets: + raise ValueError("At least one OmniInteract subset is required") + if len(set(subsets)) != len(subsets): + raise ValueError("OmniInteract subsets must not contain duplicates") + if num_prompts < 0: + raise ValueError("num_prompts must be non-negative") + data_root, cases = _data_dir(root.resolve()), [] + for subset in subsets: + subset_root = data_root / subset + selected = _one_to_many_cases(subset_root) if subset == "1qna" else _mapping_cases(subset_root, subset) + if not selected: + raise ValueError(f"No OmniInteract sessions found for requested subset {subset!r}") + cases.extend(selected) + if len({case.video_path for case in cases}) != len(cases): + raise ValueError("OmniInteract dataset contains duplicate video paths") + if not disable_shuffle: + random.Random(seed).shuffle(cases) + if num_prompts: + if num_prompts > len(cases): + logger.warning( + "Requested %d OmniInteract prompts but only %d are available; using all cases", num_prompts, len(cases) + ) + cases = cases[:num_prompts] + return cases + + +def case_manifest(case: OmniInteractCase, output_dir: Path) -> dict[str, Any]: + return { + "sample_id": f"{case.subset}__{output_dir.name}", + "video": str(case.video_path), + "gt_json": str(case.annotation_path.resolve()), + "model_json": str((output_dir / "wav_transcript.json").resolve()), + "scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type, + } diff --git a/vllm_omni/benchmarks/duplex/__init__.py b/vllm_omni/benchmarks/duplex/__init__.py new file mode 100644 index 000000000..62d91dd86 --- /dev/null +++ b/vllm_omni/benchmarks/duplex/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CPU-friendly Omni-DuplexEval generation and scoring helpers.""" + +from .omni_duplex_eval_dataset import DuplexSample, load_samples +from .omni_duplex_eval_metrics import PROTOCOL_PIN + +__all__ = ["DuplexSample", "PROTOCOL_PIN", "load_samples"] diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py new file mode 100644 index 000000000..03997fc8e --- /dev/null +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Media-clock sentence extraction and response JSON normalization.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import regex as re + +_TERMINAL = re.compile(r"(.+?[.!?。!?]+)(?:\s+|$)", re.S) + + +@dataclass +class TimedSentence: + sentence: str + start: float + end: float + + def as_dict(self) -> dict[str, Any]: + return {"sentence": self.sentence, "start": round(self.start, 3), "end": round(self.end, 3)} + + +def split_text(text: str) -> list[str]: + text = str(text or "").strip() + if not text: + return [] + parts: list[str] = [] + position = 0 + for match in _TERMINAL.finditer(text): + parts.append(match.group(1).strip()) + position = match.end() + tail = text[position:].strip() + if tail: + parts.append(tail) + return [part for part in parts if part] + + +def extract_timed_sentences(events: Iterable[dict[str, Any]], *, clock: str = "media") -> list[TimedSentence]: + """Split output deltas while anchoring text to input ``audio_end_ms``.""" + if clock != "media": + raise ValueError("only clock=media is supported") + sentences: list[TimedSentence] = [] + pending = "" + start: float | None = None + last: float | None = None + for event in events: + event_type = event.get("type") + if event_type not in {"response.output_text.delta", "response.audio_transcript.delta", "response.text.delta"}: + continue + delta = event.get("delta") + if not isinstance(delta, str) or not delta.strip(): + continue + timestamp = event.get("audio_end_ms") + if timestamp is None: + timestamp = event.get("_media_clock_ms", event.get("current_time_ms", 0.0)) + try: + # Every accepted event field is explicitly millisecond-valued. + now = float(timestamp) / 1000.0 + except (TypeError, ValueError): + now = last or 0.0 + pending += delta + start = now if start is None else start + last = now + pieces = split_text(pending) + complete = pieces if pieces and re.search(r"[.!?。!?]$", pending.strip()) else pieces[:-1] + if complete: + for piece in complete: + sentences.append(TimedSentence(piece, start, max(start, now))) + consumed = " ".join(complete) + pending = pending[len(consumed) :].strip() + start = now if pending else None + if pending.strip(): + sentences.append(TimedSentence(pending.strip(), start or last or 0.0, last or start or 0.0)) + return sentences + + +def normalize_response_items(raw: Any) -> list[dict[str, Any]]: + if isinstance(raw, dict): + items = raw.get("sentences", raw.get("chunks")) + else: + items = raw + if not isinstance(items, list): + raise ValueError("response JSON must be a list or contain sentences/chunks") + normalized = [] + for item in items: + if not isinstance(item, dict): + continue + text = item.get("sentence", item.get("text", item.get("content", ""))) + if not isinstance(text, str) or not text.strip(): + continue + try: + start = float(item.get("start", item.get("current_time", item.get("time", 0.0))) or 0.0) + end = float(item.get("end", item.get("current_time", item.get("time", start))) or start) + except (TypeError, ValueError): + continue + normalized.append({"sentence": text.strip(), "start": start, "end": max(start, end)}) + return normalized + + +def validate_clock(meta: dict[str, Any], *, allow_invalid: bool = False) -> None: + if meta.get("clock") == "invalid" and not allow_invalid: + raise ValueError("generate artifact has clock=invalid; pass --allow-invalid-clock to score it") diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py new file mode 100644 index 000000000..9d6013e80 --- /dev/null +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Dataset and split normalization for Omni-DuplexEval. + +The loader accepts a Hugging Face dataset id, a JSON/JSONL manifest, or an +already materialized iterable. Media is deliberately resolved at use time so +the benchmark remains usable in air-gapped environments. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +DEFAULT_DATASET = "Hothan/Omni-DuplexEval" +RTD_SPLITS = frozenset( + { + "RTD_world_knowledge", + "RTD_counting", + "RTD_fine_grained_movement", + "RTD_interaction_relation", + "RTD_OCR", + "RTD_Omni", + } +) +PR_SPLITS = frozenset({"PR_correction", "PR_event_reminder", "PR_post_event_reminder"}) +_TASK_ALIASES = { + "correction": "correction", + "pr_correction": "correction", + "event_reminder": "proactive_reminder", + "pr_event_reminder": "proactive_reminder", + "proactive_reminder": "proactive_reminder", + "post_event_reminder": "post_event_reminder", + "pr_post_event_reminder": "post_event_reminder", +} + + +def canonical_task_type(value: str) -> str: + key = str(value).strip().lower().replace("-", "_").replace(" ", "_") + if key not in _TASK_ALIASES: + raise ValueError(f"unsupported Omni-DuplexEval task type: {value!r}") + return _TASK_ALIASES[key] + + +def task_type_for_split(split: str) -> str: + """Match the task routing used by the official HF batch evaluator.""" + key = str(split).strip().lower() + if "correction" in key: + return "correction" + if "post_event" in key: + return "post_event_reminder" + if key.startswith("pr"): + return "proactive_reminder" + raise ValueError(f"cannot infer proactive-reminder task from split {split!r}") + + +def family_for_split(split: str, task_type: str | None = None) -> str: + if split in RTD_SPLITS or str(split).upper().startswith("RTD"): + return "rtd" + if split in PR_SPLITS or str(split).upper().startswith("PR"): + return "pr" + if task_type: + return "rtd" if str(task_type).lower().startswith("rtd") else "pr" + raise ValueError(f"cannot infer benchmark family from split {split!r}") + + +def _value(row: dict[str, Any], *keys: str, default: Any = None) -> Any: + for key in keys: + if key in row and row[key] is not None: + return row[key] + return default + + +def _float(value: Any, default: float | None = None) -> float | None: + try: + return None if value is None else float(value) + except (TypeError, ValueError): + return default + + +@dataclass(frozen=True) +class DuplexSample: + id: str + split: str + family: str + task_type: str | None + video: Any + question_audio: Any = None + question_text: str = "" + answer1: str = "" + answer2: str = "" + reminder1: Any = None + reminder2: Any = None + video_duration: float | None = None + video_type: str = "" + raw: dict[str, Any] | None = None + + @classmethod + def from_row(cls, row: dict[str, Any], *, split: str | None = None, media_root: Path | None = None) -> DuplexSample: + chosen_split = str(split or _value(row, "split", "subset", "config", default="")) + task_value = _value(row, "task_type", "task", "type") + family = str(_value(row, "family", default="") or family_for_split(chosen_split, task_value)) + task = None + if family == "pr": + task = canonical_task_type(str(task_value)) if task_value else task_type_for_split(chosen_split) + sample_id = str(_value(row, "id", "sample_id", "uid", "name", default="")) + video = _value(row, "video", "video_path", "video_file") + audio = _value(row, "question_audio", "audio", "question_wav") + if media_root: + + def resolve(value: Any) -> Any: + if not isinstance(value, str): + return value + candidate = Path(value).expanduser() + return str(candidate if candidate.is_absolute() else media_root / candidate) + + video, audio = resolve(video), resolve(audio) + return cls( + id=sample_id, + split=chosen_split, + family=family, + task_type=task, + video=video, + question_audio=audio, + question_text=str(_value(row, "question_text", "question", "instruction", default="") or ""), + answer1=str(_value(row, "answer1", "answer", "ground_answer", default="") or ""), + answer2=str(_value(row, "answer2", default="") or ""), + reminder1=_value(row, "reminder1", "reminder_1"), + reminder2=_value(row, "reminder2", "reminder_2"), + video_duration=_float(_value(row, "video_duration", "duration")), + video_type=str(_value(row, "video_type", default="") or ""), + raw=dict(row), + ) + + +def _read_manifest(path: Path) -> list[dict[str, Any]]: + if path.suffix.lower() == ".jsonl": + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + payload = json.loads(path.read_text(encoding="utf-8")) + if isinstance(payload, dict): + payload = payload.get("data", payload.get("samples", [payload])) + if not isinstance(payload, list): + raise ValueError("manifest must contain a JSON list") + return payload + + +def load_samples( + dataset: str | Path | Iterable[dict[str, Any]] = DEFAULT_DATASET, + *, + split: str | None = None, + family: str = "all", + media_root: str | Path | None = None, + limit: int | None = None, + ids: Iterable[str] | None = None, +) -> list[DuplexSample]: + if isinstance(dataset, (str, Path)) and Path(str(dataset)).exists(): + rows = _read_manifest(Path(str(dataset))) + elif not isinstance(dataset, (str, Path)): + rows = list(dataset) + else: + try: + from datasets import load_dataset + except ImportError as exc: + raise RuntimeError("datasets is required for Hugging Face loading; use a local manifest instead") from exc + kwargs: dict[str, Any] = {} + if split and split != "all": + kwargs["split"] = split + loaded = load_dataset(str(dataset), **kwargs) + if isinstance(loaded, dict): + rows = [] + for name, table in loaded.items(): + rows.extend({**dict(row), "split": name} for row in table) + else: + rows = [dict(row) for row in loaded] + wanted = set(str(item) for item in ids) if ids else None + root = Path(media_root).expanduser() if media_root else None + result = [] + for row in rows: + sample = DuplexSample.from_row(row, split=None if split in (None, "all") else split, media_root=root) + if split and split != "all" and sample.split != split: + continue + if family != "all" and sample.family != family: + continue + if wanted is not None and sample.id not in wanted: + continue + result.append(sample) + if limit is not None: + result = result[: max(0, limit)] + return result diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_eval.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_eval.py new file mode 100644 index 000000000..a92376bd7 --- /dev/null +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_eval.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Evaluate and summarize generated Omni-DuplexEval artifacts.""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path +from typing import Any + +from .omni_duplex_eval_clock import normalize_response_items, validate_clock +from .omni_duplex_eval_dataset import DuplexSample +from .omni_duplex_eval_judge import DuplexJudge +from .omni_duplex_eval_media import extract_jpeg, materialize_media, video_duration +from .omni_duplex_eval_metrics import ( + PROTOCOL_PIN, + build_content_prompt, + build_reminder_prompt, + build_temporal_prompt, + parse_judge_json, + reminder_window, + summarize_pr_results, + summarize_temporal_results, + temporal_window, +) + + +def _read(path: Path) -> Any: + return json.loads(path.read_text(encoding="utf-8")) + + +def _write(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8") + + +def _text(items: list[dict[str, Any]]) -> str: + return " ".join(str(item.get("sentence", item.get("text", ""))).strip() for item in items).strip() + + +def evaluate_sample( + sample: DuplexSample, + response_path: Path, + score_path: Path, + judge: DuplexJudge, + *, + judge_fps: int = 2, + judge_video_mode: str = "video_url", + window_size: float = 10.0, + allow_invalid_clock: bool = False, +) -> dict[str, Any]: + raw = _read(response_path) + meta = ( + _read(response_path.with_name(response_path.stem + ".meta.json")) + if response_path.with_name(response_path.stem + ".meta.json").exists() + else {} + ) + validate_clock(meta, allow_invalid=allow_invalid_clock) + items = normalize_response_items(raw) + result: dict[str, Any] = { + "id": sample.id, + "family": sample.family, + "task_type": sample.task_type, + "protocol_pin": PROTOCOL_PIN, + "response_meta": meta, + "judge_model": getattr(judge, "model", None), + "judge_video_mode": judge_video_mode, + } + if sample.family == "rtd": + video_path = materialize_media(sample.video, score_path.parent / ".media", sample.id, ".mp4") + duration = sample.video_duration or video_duration(video_path) + temporal_rows = [] + for item in items: + window = temporal_window(item["start"], item["end"], duration) + if not window: + temporal_rows.append( + { + **item, + "sentence_start": item["start"], + "sentence_end": item["end"], + "sentence_duration": max(0.0, item["end"] - item["start"]), + "error": "No valid temporal window", + "temporal_score": 0, + "is_relevant": 0, + } + ) + continue + frames = _extract_frames(video_path, _times(*window, fps=judge_fps)) + parsed = parse_judge_json( + judge.temporal(build_temporal_prompt(*window, item["sentence"], sample.question_text), frames) + ) + temporal_rows.append( + { + **item, + "sentence_start": item["start"], + "sentence_end": item["end"], + "sentence_duration": max(0.0, item["end"] - item["start"]), + "window_start": window[0], + "window_end": window[1], + "error": None, + **parsed, + } + ) + content_frames = None + if judge_video_mode == "frame-sample": + content_frames = _extract_frames(video_path, _content_frame_times(duration)) + content = parse_judge_json( + judge.content( + build_content_prompt(_text(items), sample.question_text, [sample.answer1, sample.answer2]), + video_path, + frames=content_frames, + mode=judge_video_mode, + ) + ) + if content_frames is not None: + content["frame_count"] = len(content_frames) + result.update( + { + "temporal": {"sentences": temporal_rows, "summary": summarize_temporal_results(temporal_rows)}, + "content": content, + } + ) + else: + if sample.task_type is None: + raise ValueError(f"missing proactive-reminder task type for split {sample.split!r}") + task = sample.task_type + times = _reminder_times(sample) + event_rows = [] + if task == "correction" or not times: + response = _text(items) + parsed = parse_judge_json( + judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1)) + ) + event_rows.append({"response_segment": response, **parsed}) + else: + for start in times: + lo, hi = reminder_window(start, window_size) + response = " ".join(item["sentence"] for item in items if lo <= item["start"] <= hi) + if not response: + event_rows.append({"start_time": start, "response_segment": "", "success_score": 0}) + else: + parsed = parse_judge_json( + judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1)) + ) + event_rows.append({"start_time": start, "response_segment": response, **parsed}) + scores = [int(row.get("success_score", 0)) for row in event_rows] + result["pr"] = { + "events": event_rows, + "all_success": bool(scores) and all(score == 1 for score in scores), + "total_score": int(bool(scores) and all(score == 1 for score in scores)), + } + _write(score_path, result) + return result + + +def _times(start: float, end: float, *, fps: int) -> list[float]: + count = max(1, int((end - start) * fps)) + return ( + [start + (end - start) / 2] if count == 1 else [start + i * (end - start) / (count - 1) for i in range(count)] + ) + + +def _extract_frames(path: str | Path, timestamps: list[float]) -> list[bytes]: + """Extract only decodable JPEG frames; clips can end on non-keyframes.""" + frames = [] + for timestamp in timestamps: + try: + frame = extract_jpeg(path, timestamp=timestamp) + except (OSError, subprocess.CalledProcessError, ValueError): + continue + if frame.startswith(b"\xff\xd8"): + frames.append(frame) + return frames + + +def _content_frame_times(duration: float) -> list[float]: + """Sample the full video at approximately one frame every three seconds.""" + if duration <= 0: + return [] + times = [] + timestamp = 0.0 + while timestamp < duration: + times.append(min(timestamp, max(0.0, duration - 0.01))) + timestamp += 3.0 + # Qwen2.5-VL has a finite per-request image-token budget. Preserve + # coverage while bounding long clips to a safely portable frame count. + if len(times) > 16: + stride = (len(times) - 1) / 15 + times = [times[round(i * stride)] for i in range(16)] + return times + + +def _reminder_times(sample: DuplexSample) -> list[float]: + values = [] + for value in (sample.reminder1, sample.reminder2): + if isinstance(value, dict): + value = value.get("start", value.get("time")) + try: + if value is not None: + values.append(float(value)) + except (TypeError, ValueError): + pass + return values + + +def summarize_scores(score_root: str | Path) -> dict[str, Any]: + rows = [] + for path in sorted(Path(score_root).rglob("*.json")): + if not path.name.endswith("_summary.json"): + row = _read(path) + row.setdefault("split", path.parent.name) + rows.append(row) + rtd_temporal = [ + row["temporal"]["summary"] | {"content_score": row.get("content", {}).get("content_score", 0.0)} + for row in rows + if "temporal" in row + ] + pr = [{"task_type": row.get("task_type"), **row.get("pr", {})} for row in rows if "pr" in row] + result = {"protocol_pin": PROTOCOL_PIN, "samples": len(rows)} + if rtd_temporal: + by_task: dict[str, list[dict[str, float]]] = {} + for row in rows: + if "temporal" in row: + by_task.setdefault(str(row.get("split", row.get("task_type", "unknown"))), []).append( + { + "content": float(row.get("content", {}).get("content_score", 0.0)), + "temporal": float(row["temporal"]["summary"].get("avg_temporal_score", 0.0)), + } + ) + result["rtd"] = { + "mean_content_score": sum(item["content_score"] for item in rtd_temporal) / len(rtd_temporal), + "mean_avg_temporal_score": sum(item["avg_temporal_score"] for item in rtd_temporal) / len(rtd_temporal), + "by_task": { + task: { + "mean_content_score": sum(item["content"] for item in values) / len(values), + "mean_avg_temporal_score": sum(item["temporal"] for item in values) / len(values), + } + for task, values in by_task.items() + }, + } + if pr: + result["pr"] = summarize_pr_results(pr) + return result diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_judge.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_judge.py new file mode 100644 index 000000000..2fa4014d5 --- /dev/null +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_judge.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI-compatible local judge client.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pybase64 as base64 +import requests + + +def _build_openai_url(base_url: str, api_path: str) -> str: + base = base_url.rstrip("/") + normalized_path = api_path if api_path.startswith("/") else f"/{api_path}" + if base.endswith(normalized_path): + return base + if base.endswith("/v1"): + return f"{base}{normalized_path}" + return f"{base}/v1{normalized_path}" + + +class DuplexJudge: + def __init__(self, base_url: str, model: str, *, api_key: str = "EMPTY", timeout: int = 600) -> None: + self.base_url, self.model, self.api_key, self.timeout = base_url, model, api_key, timeout + + def chat(self, content: Any, *, system: str | None = None, max_tokens: int = 1200) -> str: + messages = ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": content}] + response = requests.post( + _build_openai_url(self.base_url, "/chat/completions"), + json={ + "model": self.model, + "messages": messages, + "temperature": 0.1, + "max_tokens": max_tokens, + }, + headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}, + timeout=self.timeout, + ) + if not response.ok: + raise requests.HTTPError( + f"{response.status_code} response from judge: {response.text[:500]}", response=response + ) + return str(response.json()["choices"][0]["message"]["content"]) + + def temporal(self, prompt: str, frames: list[bytes]) -> str: + content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + content.extend( + {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()}} + for frame in frames + ) + return self.chat(content, system="You are a careful evaluator for real-time multimodal systems.") + + def content( + self, + prompt: str, + video: str | Path | None = None, + frames: list[bytes] | None = None, + *, + mode: str = "video_url", + ) -> str: + content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + if mode == "video_url" and video is not None: + content.append({"type": "video_url", "video_url": {"url": "file://" + str(Path(video).resolve())}}) + else: + content.extend( + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()}, + } + for frame in (frames or []) + ) + return self.chat(content) diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py new file mode 100644 index 000000000..629d8e877 --- /dev/null +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Small ffmpeg-backed media helpers used by the duplex runner and judge.""" + +from __future__ import annotations + +import subprocess +import wave +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +from vllm_omni.experimental.fullduplex.client import PCM16_BYTES_PER_SAMPLE, PCM16_SAMPLE_RATE + + +def _run_ffmpeg(args: list[str]) -> bytes: + result = subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", *args], capture_output=True, check=True) + return result.stdout + + +def materialize_media(value: Any, output_dir: str | Path, stem: str, suffix: str) -> Path: + """Resolve common Hugging Face media values to a local file.""" + if isinstance(value, str | Path): + return Path(value) + path = value.get("path") if isinstance(value, dict) else getattr(value, "path", None) + if path: + return Path(path) + payload = value.get("bytes") if isinstance(value, dict) else None + if payload is None and isinstance(value, bytes | bytearray): + payload = value + if payload is None: + raise ValueError(f"cannot materialize media value for {stem!r}") + destination = Path(output_dir) / f"{stem}{suffix}" + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(bytes(payload)) + return destination + + +def video_duration(path: str | Path) -> float: + result = subprocess.run( + ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(path)], + capture_output=True, + text=True, + check=True, + ) + return max(0.0, float(result.stdout.strip())) + + +def extract_jpeg(path: str | Path, *, timestamp: float, quality: int = 3) -> bytes: + return _run_ffmpeg( + [ + "-ss", + f"{max(0.0, timestamp):.3f}", + "-i", + str(path), + "-frames:v", + "1", + "-q:v", + str(quality), + "-f", + "image2", + "pipe:1", + ] + ) + + +def iter_jpegs( + path: str | Path, *, fps: float = 1.0, duration: float | None = None, quality: int = 3 +) -> Iterator[tuple[float, bytes]]: + end = duration if duration is not None else video_duration(path) + step = 1.0 / fps + timestamp = 0.0 + while timestamp < end: + try: + frame = extract_jpeg(path, timestamp=timestamp, quality=quality) + except (OSError, subprocess.CalledProcessError): + break + if frame: + yield timestamp, frame + timestamp += step + + +def read_audio_pcm16(path: str | Path) -> bytes: + source = Path(path) + if source.suffix.lower() == ".wav": + try: + with wave.open(str(source), "rb") as wav_file: + if ( + wav_file.getnchannels() == 1 + and wav_file.getsampwidth() == 2 + and wav_file.getframerate() == PCM16_SAMPLE_RATE + and wav_file.getcomptype() == "NONE" + ): + return wav_file.readframes(wav_file.getnframes()) + except wave.Error: + pass + return _run_ffmpeg(["-i", str(source), "-f", "s16le", "-ac", "1", "-ar", str(PCM16_SAMPLE_RATE), "pipe:1"]) + + +def iter_av_units( + audio_pcm16: bytes, frames: Iterator[tuple[float, bytes]], *, unit_ms: int = 1000 +) -> Iterator[tuple[bytes, bytes | None]]: + unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * unit_ms // 1000 + frame_iter = iter(frames) + next_frame = next(frame_iter, None) + for offset in range(0, len(audio_pcm16), unit_bytes): + timestamp = offset / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE) + frame = None + if next_frame is not None and next_frame[0] <= timestamp + unit_ms / 1000: + frame = next_frame[1] + next_frame = next(frame_iter, None) + yield audio_pcm16[offset : offset + unit_bytes], frame diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_metrics.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_metrics.py new file mode 100644 index 000000000..a373be58d --- /dev/null +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_metrics.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Protocol prompts, windows, parsing, and aggregation. + +Prompt wording and window math are copied verbatim from the MIT-licensed +OpenBMB/Omni-DuplexEval evaluator at :data:`PROTOCOL_PIN`. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterable +from typing import Any + +import regex as re + +# Copied from OpenBMB/Omni-DuplexEval@PROTOCOL_PIN (MIT). +# Do not change prompt wording without bumping the pin and recording a score delta. +PROTOCOL_PIN = "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48" + + +def temporal_window(sentence_start: float, sentence_end: float, video_duration: float) -> tuple[float, float] | None: + start = max(0.0, float(sentence_start) - 2.0) + end = max(start + 0.5, float(sentence_end) - 2.0) + return (start, end) if end <= float(video_duration) and end - start >= 0.5 else None + + +def reminder_window(start_time: float, window_size: float = 10.0) -> tuple[float, float]: + return float(start_time), float(start_time) + float(window_size) + + +def parse_judge_json(text: str) -> dict[str, Any]: + for match in re.findall(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", str(text), re.S): + try: + value = json.loads(re.sub(r",\s*([}\]])", r"\1", match)) + except (TypeError, ValueError, json.JSONDecodeError): + continue + if isinstance(value, dict): + break + else: + value = {} + patterns = [ + (r'"temporal_score"\s*:\s*(\d+)', "temporal_score", int), + (r'"is_relevant"\s*:\s*(\d+)', "is_relevant", int), + (r'"success_score"\s*:\s*([01])', "success_score", int), + (r'"content_score"\s*:\s*(\d+(?:\.\d+)?)', "content_score", float), + ] + for pattern, key, caster in patterns: + if match := re.search(pattern, str(text)): + parsed = caster(match.group(1)) + value[key] = round(parsed, 2) if caster is float else parsed + for pattern, key in [ + (r'"temporal_reasoning"\s*:\s*"([^"]*)"', "temporal_reasoning"), + (r'"content_reasoning"\s*:\s*"([^"]*)"', "content_reasoning"), + (r'"reasoning"\s*:\s*"([^"]*)"', "reasoning"), + ]: + if match := re.search(pattern, str(text), re.S): + value[key] = match.group(1).strip() + break + for key in ("temporal_score", "is_relevant", "success_score"): + if key in value: + try: + value[key] = int(float(value[key])) + except (TypeError, ValueError): + value[key] = 0 + if "content_score" in value: + try: + value["content_score"] = round(float(value["content_score"]), 2) + except (TypeError, ValueError): + value["content_score"] = 0.0 + return value + + +def summarize_temporal_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]: + rows = list(results) + valid = [row for row in rows if not row.get("error")] + relevant = [row for row in valid if int(row.get("is_relevant", 0)) == 1] + irrelevant = [row for row in valid if int(row.get("is_relevant", 0)) == 0] + scores = [int(row.get("temporal_score", 0)) for row in relevant] + total_valid_duration = sum(float(row.get("sentence_duration", 0.0)) for row in valid) + irrelevant_duration = sum(float(row.get("sentence_duration", 0.0)) for row in irrelevant) + return { + "avg_temporal_score": round(sum(scores) / len(scores), 4) if scores else 0.0, + "total_sentences": len(rows), + "evaluated_sentences": len(valid), + "error_count": len(rows) - len(valid), + "relevant_sentences_count": len(relevant), + "irrelevant_sentences_count": len(irrelevant), + "score_distribution": { + "3_points": scores.count(3), + "2_points": scores.count(2), + "1_point": scores.count(1), + "0_points": scores.count(0), + }, + "relevance_stats": { + "first_relevant_time": min((row.get("sentence_start") for row in relevant), default=None), + "total_relevant_duration": round(sum(float(row.get("sentence_duration", 0.0)) for row in relevant), 4), + "total_irrelevant_duration": round(irrelevant_duration, 4), + "irrelevant_duration_ratio": round( + irrelevant_duration / total_valid_duration if total_valid_duration > 0 else 0.0, 4 + ), + }, + "excluded_irrelevant_sentences": [ + { + "sentence": row.get("sentence", ""), + "start": row.get("sentence_start"), + "end": row.get("sentence_end"), + "temporal_score": row.get("temporal_score", 0), + } + for row in irrelevant + ], + } + + +def summarize_pr_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]: + rows = list(results) + successes = [int(row.get("all_success", row.get("total_score", 0))) for row in rows] + by_task: dict[str, list[int]] = {} + for row, score in zip(rows, successes, strict=False): + by_task.setdefault(str(row.get("task_type", "unknown")), []).append(score) + return { + "samples": len(rows), + "mean_all_success": sum(successes) / len(successes) if successes else 0.0, + "by_task": {task: sum(values) / len(values) for task, values in by_task.items()}, + } + + +def build_temporal_prompt(start_time: float, end_time: float, response_text: str, question: str) -> str: + return f""" +You are evaluating a real-time video description system. + +Basic information: +- Video segment: {start_time:.2f}s to {end_time:.2f}s. +- User instruction: {question} +- Model response sentence: "{response_text}" + +Evaluation principles: +- Judge only the current video segment. +- Ignore response content that clearly belongs to earlier or later segments unless it contradicts the current segment. +- Language does not matter; evaluate semantic correctness and temporal alignment only. +- Decide whether the sentence is a substantive response or only a filler/polite phrase. + +Temporal sensitivity score: +- 3: Excellent temporal alignment. The sentence accurately describes the current segment and follows the instruction. +- 2: Mostly aligned. Minor inaccuracies, omissions, or harmless references to nearby segments are allowed. +- 1: Poor alignment. Major inaccuracies, wrong timing, or mostly irrelevant content. +- 0: No meaningful alignment with the current segment. + +Relevance label: +- 1: The sentence contains substantive task- or video-related information. +- 0: The sentence is only filler, acknowledgement, meta-commentary, or generic text without useful content. + +Return exactly this JSON object: +{{ + "temporal_score": <0, 1, 2, or 3>, + "temporal_reasoning": "", + "is_relevant": <0 or 1> +}} +""" + + +def build_content_prompt(response_text: str, question: str, references: list[str] | None = None) -> str: + reference_block = "" + if references: + reference_block = "Reference annotations, for reference only:\n" + for index, reference in enumerate(references, 1): + reference_block += f"{index}. {reference}\n" + + return f""" +You are a precise content-accuracy evaluator for video description. + +Basic information: +- User instruction: {question} +- Model response: "{response_text}" +{reference_block} + +Evaluation goal: +Judge whether the model response is factually accurate for the whole video and aligned with the instruction. +Consider object/action/color/count/spatial/event errors, hallucinations, omissions, and irrelevant content. +Reference annotations are optional guidance; the primary evidence is the video itself. + +Scoring: +- Use a decimal score from 0.00 to 3.00. +- Start from 3.00 and deduct for each error. +- Use 0.00 only when the response is empty, completely irrelevant, or contains no correct video facts. +- Output exactly two decimal places. + +Return exactly this JSON object: +{{ + "content_score": , + "content_reasoning": "" +}} +""" + + +def build_reminder_prompt(instruction: str, response: str, task_type: str, ground_answer: str = "") -> str: + if task_type == "correction": + return f""" +You are judging whether a model successfully completed a correction task. + +Task description: +The user instruction contains an incorrect statement about the video. +The model should identify the incorrect part and provide the correct information. + +Input: +- User instruction: {instruction} +- Reference correction: {ground_answer} +- Model response: "{response}" + +Success criteria: +1. Identify the error implied by the user instruction and the reference correction. +2. Check whether the model corrected all required error points. +3. The corrected content must preserve the correct context, including subject, object, action, and attributes. +4. Ignore extra information that is unrelated to both the instruction and the reference, unless it contradicts them. + +Scoring: +- 1 = all required error points are corrected and the context is consistent. +- 0 = at least one required correction is missing, wrong, or contextually inconsistent. + +Return exactly this JSON object: +{{ + "success_score": <0 or 1>, + "reasoning": "" +}} +""" + + if task_type == "post_event_reminder": + timing_note = "The response is evaluated immediately after the target event has occurred." + else: + timing_note = "The response is evaluated at the target event time or immediately after it." + + return f""" +You are judging whether a model successfully completed an event-reminder task. + +Task description: +The user gave an instruction asking the system to remind them when a specific event happens. +{timing_note} +The provided model text is the response segment generated in the evaluation window. + +Input: +- User instruction: {instruction} +- Model response segment: "{response}" + +Success criteria: +1. The response clearly refers to the target event in the instruction. +2. The response communicates a reminder, notification, or confirmation that the event has happened. +3. Vague narration, unrelated descriptions, or wrong-event reminders are failures. +4. The response does not need to match the instruction wording exactly if the meaning is clear. + +Scoring: +- 1 = successful reminder. +- 0 = unsuccessful reminder. + +Return exactly this JSON object: +{{ + "success_score": <0 or 1>, + "reasoning": "" +}} +""" diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py new file mode 100644 index 000000000..d5a28aa06 --- /dev/null +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Single-sample MiniCPM-o native-duplex generation.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +from pathlib import Path +from typing import Any + +import pybase64 as base64 + +from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, build_realtime_url, wait_for + +from .omni_duplex_eval_clock import extract_timed_sentences +from .omni_duplex_eval_media import iter_av_units, iter_jpegs, materialize_media, read_audio_pcm16, video_duration + + +def _ref_audio(path: str | Path) -> str: + value = Path(path).expanduser().read_bytes() + return "data:audio/wav;base64," + base64.b64encode(value).decode("ascii") + + +async def generate_sample( + sample: Any, + *, + url: str, + model: str, + ref_audio: str | Path, + output_root: str | Path, + fps: float = 1.0, + mix: str = "question", + pace: str = "realtime", + clock: str = "media", + overwrite: bool = False, + unit_ms: int = 1000, +) -> Path: + output = Path(output_root) / sample.split / f"{sample.id}.json" + meta_path = output.with_name(output.stem + ".meta.json") + if output.exists() and not overwrite: + return output + if mix != "question": + raise NotImplementedError("v1 supports mix=question; soundtrack mixing is reserved for P1") + media_dir = output.parent / ".media" + audio_path = materialize_media(sample.question_audio, media_dir, f"{sample.id}_question", ".wav") + video_path = materialize_media(sample.video, media_dir, sample.id, ".mp4") + pcm = read_audio_pcm16(audio_path) + duration = sample.video_duration or video_duration(video_path) + frames = iter_jpegs(video_path, fps=fps, duration=duration) + realtime = pace == "realtime" + if pace not in {"realtime", "as-fast-as-possible"}: + raise ValueError("pace must be realtime or as-fast-as-possible") + client = RealtimeDuplexClient(build_realtime_url(url, model, autostart=False)) + response_done = False + drain_timeout = None + close_timeout = None + async with client: + await client.configure(model, ref_audio=_ref_audio(ref_audio), instructions="Streaming Omni Conversation.") + ack_task = asyncio.create_task(_ack_playback(client)) + try: + await client.stream_av_units(iter_av_units(pcm, frames, unit_ms=unit_ms), realtime=realtime) + await client.commit() + try: + await wait_for( + lambda: client.events.count("response.done") > 0, + timeout_s=max(20.0, duration + 20.0), + label="response.done", + ) + response_done = True + except TimeoutError as exc: + drain_timeout = str(exc) + finally: + ack_task.cancel() + try: + await ack_task + except asyncio.CancelledError: + pass + await client.acknowledge_playback() + try: + await client.close_session(timeout_s=20.0) + except TimeoutError as exc: + close_timeout = str(exc) + events = list(client.events.events) + timed = [sentence.as_dict() for sentence in extract_timed_sentences(events, clock=clock)] + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(timed, ensure_ascii=False, indent=2), encoding="utf-8") + meta = { + "id": sample.id, + "split": sample.split, + "clock": clock if realtime else "invalid", + "pace": pace, + "mix": mix, + "fps": fps, + "unit_ms": unit_ms, + "model": model, + "response_done": response_done, + "drain_timeout": drain_timeout, + "close_timeout": close_timeout, + "ref_audio_sha256": hashlib.sha256(Path(ref_audio).read_bytes()).hexdigest(), + } + meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8") + return output + + +async def _ack_playback(client: RealtimeDuplexClient) -> None: + while True: + await asyncio.sleep(0.25) + await client.acknowledge_playback() diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 824b810f3..3d7c983f3 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -1,5 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import warnings from collections import defaultdict +from collections.abc import Sequence from dataclasses import field, make_dataclass import numpy as np @@ -15,6 +19,11 @@ _FLOAT_LIST_TYPE = list[float] _INT_LIST_TYPE = list[int] _MULTIMODAL_BENCHMARK_FIELDS = [ + ("num_ttft_samples", int, field(default=0)), + ("num_tpot_samples", int, field(default=0)), + ("num_itl_samples", int, field(default=0)), + ("num_audio_ttfp_samples", int, field(default=0)), + ("num_audio_rtf_samples", int, field(default=0)), (defs.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)), (defs.MEDIAN_AUDIO_TTFP_MS, float, field(default=0.0)), (defs.STD_AUDIO_TTFP_MS, float, field(default=0.0)), @@ -176,6 +185,17 @@ def _stage_modality_flags( ) +def has_metric_samples(metrics: object, metric_name: str) -> bool: + sample_count_attr = { + "ttft": "num_ttft_samples", + "tpot": "num_tpot_samples", + "itl": "num_itl_samples", + defs.AUDIO_TTFP: "num_audio_ttfp_samples", + defs.AUDIO_RTF: "num_audio_rtf_samples", + }.get(metric_name) + return sample_count_attr is None or getattr(metrics, sample_count_attr, 0) > 0 + + def print_metrics( task_type, selected_percentile_metrics, @@ -308,6 +328,9 @@ def process_one_metric( metric_attribute_name: str, metrics: MultiModalsBenchmarkMetrics, ): + if not has_metric_samples(metrics, metric_attribute_name): + return + metric_header_map = { "ttft": "Time to First Token", "tpot": "Time per Output Token (excl. 1st token)", @@ -718,10 +741,12 @@ def calculate_metrics( good_completed = 0 itls: list[float] = [] tpots: list[float] = [] - all_tpots: list[float] = [] + all_tpots: list[float | None] = [] ttfts: list[float] = [] + goodput_ttfts: list[float | None] = [] e2els: list[float] = [] audio_ttfps: list[float] = [] + goodput_audio_ttfps: list[float | None] = [] audio_rtfs: list[float] = [] audio_duration: list[float] = [] audio_frames: list[int] = [] @@ -748,27 +773,49 @@ def calculate_metrics( output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids) actual_output_lens.append(output_len) total_input += outputs[i].prompt_len - tpot = 0 + tpot: float | None = None if output_len > 1: if outputs[i].itl: # Use mean(ITL) directly so per-request TPOT == mean(ITL). # The ITL list records one entry per SSE chunk; server may # bundle multiple tokens per chunk, so len(itl)+1 != output_len. # Using mean(itl) keeps TPOT and ITL on the same footing. - tpot = sum(outputs[i].itl) / len(outputs[i].itl) - else: - try: - latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft - except Exception: - latency_minus_ttft = outputs[i].latency - outputs[i].ttft - tpot = latency_minus_ttft / (output_len - 1) - tpots.append(tpot) - # Note: if output_len <= 1, we regard tpot as 0 for goodput + measured_tpot = sum(outputs[i].itl) / len(outputs[i].itl) + if np.isfinite(measured_tpot) and measured_tpot > 0: + tpot = measured_tpot + elif getattr(outputs[i], "tpot_measured", True): + text_latency = getattr(outputs[i], "text_latency", None) + ttft = getattr(outputs[i], "ttft", None) + if ( + isinstance(text_latency, (int, float)) + and isinstance(ttft, (int, float)) + and np.isfinite(text_latency) + and np.isfinite(ttft) + and text_latency > ttft > 0 + ): + tpot = (text_latency - ttft) / (output_len - 1) + if tpot is not None: + tpots.append(tpot) + # TPOT is not measurable for a single-token response or when the + # client did not observe a positive decode interval. In + # particular, stage-level output token counts alone must not turn + # an unset text_latency into a negative TPOT sample. all_tpots.append(tpot) itls += outputs[i].itl - ttfts.append(outputs[i].ttft) - audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0)) - audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0)) + session_metrics = getattr(outputs[i], "duplex_session_metrics", None) + ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_ttft_ms") is not None + audio_ttfp_measured = ( + not isinstance(session_metrics, dict) or session_metrics.get("mean_ttfp_ms") is not None + ) + audio_rtf_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_rtf") is not None + if ttft_measured: + ttfts.append(outputs[i].ttft) + if audio_ttfp_measured: + audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0)) + if audio_rtf_measured: + audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0)) + goodput_ttfts.append(outputs[i].ttft if ttft_measured else None) + goodput_audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0) if audio_ttfp_measured else None) audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0)) audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0)) image_count = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0) @@ -789,14 +836,14 @@ def calculate_metrics( actual_output_lens.append(0) if goodput_config_dict: - valid_metrics = [] + valid_metrics: list[Sequence[float | None]] = [] slo_values = [] if "ttft" in goodput_config_dict: - valid_metrics.append(ttfts) + valid_metrics.append(goodput_ttfts) slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION) if "audio_ttft" in goodput_config_dict: - valid_metrics.append(audio_ttfps) + valid_metrics.append(goodput_audio_ttfps) slo_values.append(goodput_config_dict["audio_ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION) if "tpot" in goodput_config_dict: valid_metrics.append(all_tpots) @@ -806,7 +853,7 @@ def calculate_metrics( slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION) for req_metric in zip(*valid_metrics): - is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)]) + is_good_req = all(r is not None and s >= r for s, r in zip(slo_values, req_metric)) if is_good_req: good_completed += 1 @@ -822,6 +869,7 @@ def calculate_metrics( # Calculate max output tokens per second metric max_output_tokens_per_s = 0.0 max_concurrent_requests = 0 + token_timeline_available = True # Find the time range across all successful requests successful_outputs = [output for output in outputs if output.success] @@ -836,19 +884,26 @@ def calculate_metrics( concurrent_requests_per_second = np.zeros(duration_seconds) for i, output in enumerate(successful_outputs): - # Calculate token generation timestamp using - # start_time, ttft, and itl - token_times = [output.start_time + output.ttft] - current_time = token_times[0] - for itl_value in output.itl: - current_time += itl_value - token_times.append(current_time) + session_metrics = getattr(output, "duplex_session_metrics", None) + if isinstance(session_metrics, dict): + # Duplex TTFT/ITL samples are response-local. Without each + # response.created offset they cannot form a session-global + # token timeline, so do not publish a misleading peak rate. + token_timeline_available = False + else: + # Calculate token generation timestamp using + # start_time, ttft, and itl + token_times = [output.start_time + output.ttft] + current_time = token_times[0] + for itl_value in output.itl: + current_time += itl_value + token_times.append(current_time) - # Add tokens to second buckets - for token_time in token_times: - second_bucket = int(token_time - min_start_time) - if 0 <= second_bucket < duration_seconds: - tokens_per_second[second_bucket] += 1 + # Add tokens to second buckets + for token_time in token_times: + second_bucket = int(token_time - min_start_time) + if 0 <= second_bucket < duration_seconds: + tokens_per_second[second_bucket] += 1 # Track concurrent requests for each second this request was active request_start_second = int(output.start_time - min_start_time) @@ -860,18 +915,22 @@ def calculate_metrics( # Find the maximum tokens per second and corresponding # concurrent requests if len(tokens_per_second) > 0: - max_output_tokens_per_s = float(np.max(tokens_per_second)) + if token_timeline_available: + max_output_tokens_per_s = float(np.max(tokens_per_second)) + else: + max_output_tokens_per_s = float("nan") max_concurrent_requests = int(np.max(concurrent_requests_per_second)) if TERM_PLOTLIB_AVAILABLE: import termplotlib as tpl fig = tpl.figure() - fig.plot( - np.arange(len(tokens_per_second)), - tokens_per_second, - title="Output tokens per second", - ) + if token_timeline_available: + fig.plot( + np.arange(len(tokens_per_second)), + tokens_per_second, + title="Output tokens per second", + ) fig.plot( np.arange(len(concurrent_requests_per_second)), concurrent_requests_per_second, @@ -881,6 +940,10 @@ def calculate_metrics( else: print("tip: install termplotlib and gnuplot to plot the metrics") + duplex_metrics_present = any( + isinstance(getattr(output, "duplex_session_metrics", None), dict) for output in outputs + ) + missing_duplex_value = float("nan") if duplex_metrics_present else 0 metrics = MultiModalsBenchmarkMetrics( completed=completed, failed=len(failed_outputs), @@ -890,18 +953,18 @@ def calculate_metrics( request_goodput=good_completed / dur_s, output_throughput=sum(actual_output_lens) / dur_s, total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s, - mean_ttft_ms=np.mean(ttfts or 0) * 1000, # ttfts is empty if streaming is not supported by the endpoint - std_ttft_ms=np.std(ttfts or 0) * 1000, - median_ttft_ms=np.median(ttfts or 0) * 1000, - percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles], - mean_tpot_ms=np.mean(tpots or 0) * 1000, - std_tpot_ms=np.std(tpots or 0) * 1000, - median_tpot_ms=np.median(tpots or 0) * 1000, - percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles], - mean_itl_ms=np.mean(itls or 0) * 1000, - std_itl_ms=np.std(itls or 0) * 1000, - median_itl_ms=np.median(itls or 0) * 1000, - percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles], + mean_ttft_ms=np.mean(ttfts or missing_duplex_value) * 1000, + std_ttft_ms=np.std(ttfts or missing_duplex_value) * 1000, + median_ttft_ms=np.median(ttfts or missing_duplex_value) * 1000, + percentiles_ttft_ms=[(p, np.percentile(ttfts or missing_duplex_value, p) * 1000) for p in selected_percentiles], + mean_tpot_ms=np.mean(tpots or float("nan")) * 1000, + std_tpot_ms=np.std(tpots or float("nan")) * 1000, + median_tpot_ms=np.median(tpots or float("nan")) * 1000, + percentiles_tpot_ms=[(p, np.percentile(tpots or float("nan"), p) * 1000) for p in selected_percentiles], + mean_itl_ms=np.mean(itls or missing_duplex_value) * 1000, + std_itl_ms=np.std(itls or missing_duplex_value) * 1000, + median_itl_ms=np.median(itls or missing_duplex_value) * 1000, + percentiles_itl_ms=[(p, np.percentile(itls or missing_duplex_value, p) * 1000) for p in selected_percentiles], mean_e2el_ms=np.mean(e2els or 0) * 1000, std_e2el_ms=np.std(e2els or 0) * 1000, median_e2el_ms=np.median(e2els or 0) * 1000, @@ -910,11 +973,16 @@ def calculate_metrics( max_concurrent_requests=max_concurrent_requests, rtfx=input_audio_duration / dur_s, **{ - defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or 0) * 1000, - defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or 0) * 1000, - defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or 0) * 1000, + "num_ttft_samples": len(ttfts), + "num_tpot_samples": len(tpots), + "num_itl_samples": len(itls), + "num_audio_ttfp_samples": len(audio_ttfps), + "num_audio_rtf_samples": len(audio_rtfs), + defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or missing_duplex_value) * 1000, + defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or missing_duplex_value) * 1000, + defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or missing_duplex_value) * 1000, defs.PERCENTILES_AUDIO_TTFP_MS: [ - (p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles + (p, np.percentile(audio_ttfps or missing_duplex_value, p) * 1000) for p in selected_percentiles ], defs.MEAN_AUDIO_DURATION_S: np.mean(audio_duration or 0), defs.STD_AUDIO_DURATION_S: np.std(audio_duration or 0), @@ -925,10 +993,12 @@ def calculate_metrics( defs.TOTAL_AUDIO_DURATION_S: sum(audio_duration), defs.TOTAL_AUDIO_FRAMES: sum(audio_frames), defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s, - defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or 0), - defs.STD_AUDIO_RTF: np.std(audio_rtfs or 0), - defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or 0), - defs.PERCENTILES_AUDIO_RTF: [(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles], + defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or missing_duplex_value), + defs.STD_AUDIO_RTF: np.std(audio_rtfs or missing_duplex_value), + defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or missing_duplex_value), + defs.PERCENTILES_AUDIO_RTF: [ + (p, np.percentile(audio_rtfs or missing_duplex_value, p)) for p in selected_percentiles + ], defs.TOTAL_IMAGES: total_images, defs.IMAGE_THROUGHPUT: total_images / dur_s, defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0, diff --git a/vllm_omni/benchmarks/omniinteract.py b/vllm_omni/benchmarks/omniinteract.py new file mode 100644 index 000000000..46822b486 --- /dev/null +++ b/vllm_omni/benchmarks/omniinteract.py @@ -0,0 +1,933 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OmniInteract session execution and artifacts for serving benchmarks.""" + +from __future__ import annotations + +import asyncio +import binascii +import contextlib +import copy +import fcntl +import hashlib +import json +import logging +import math +import os +import subprocess +import tempfile +import time +from collections import Counter +from collections.abc import Callable, Iterator, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, cast +from urllib.parse import urljoin, urlsplit + +import pybase64 as base64 + +from vllm_omni.benchmarks.data_modules.omniinteract_dataset import ( + DEFAULT_MAX_VIDEO_DURATION_S, + OmniInteractCase, + OmniInteractPreparedInput, + case_manifest, +) +from vllm_omni.experimental.fullduplex.client import ( + PCM16_BYTES_PER_SAMPLE, + PCM16_SAMPLE_RATE, + RealtimeDuplexClient, + RealtimeEventCollector, + build_realtime_url, + chunk_period_ms, + has_residual_model_unit, + reference_audio_data_url, + summarize_session_request_metrics, + write_pcm16_wav, +) + +OUTPUT_SAMPLE_RATE = 24_000 +SUCCESS_ARTIFACTS = (".done", "output.wav", "wav_transcript.json", "events.json", "result.json") +BATCH_ARTIFACTS = ("batch_summary.json", "official_eval_manifest.jsonl") +ARTIFACT_LOCK_FILE = ".omniinteract.lock" +_INPUT_CHUNK_MS = 200 +VIDEO_FPS = 1.0 +_COMPLETION_SETTLE_S = 2.0 +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class OmniInteractBenchmarkConfig: + model: str + base_url: str = "http://127.0.0.1:8000" + endpoint: str = "/v1/realtime" + output_root: Path = Path("omniinteract-output") + timeout_s: float = 900.0 + media_timeout_s: float = 600.0 + max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S + ref_audio: str | None = None + require_response: bool = False + extra_headers: dict[str, str] | None = None + extra_body: dict[str, object] | None = None + + +@dataclass +class OmniInteractCaseResult: + subset: str + video: str + output_dir: str + success: bool = False + error: str = "" + session_id: str = "" + latency_s: float = 0.0 + input_audio_chunks: int = 0 + input_video_frames: int = 0 + pacing_mean_lag_s: float = 0.0 + pacing_max_lag_s: float = 0.0 + responses: int = 0 + audio_bytes: int = 0 + audio_clipped_bytes: int = 0 + transcript: str = "" + eligible_for_official_eval: bool = False + official_eval_ineligible_reasons: list[str] = field(default_factory=list) + artifact_warnings: list[str] = field(default_factory=list) + output_tokens: int = 0 + duplex_request_metrics: list[dict[str, object]] = field(default_factory=list) + duplex_session_metrics: dict[str, object] = field(default_factory=dict) + _artifact_context: _DeferredArtifactContext | None = field(default=None, repr=False, compare=False) + + def as_dict(self) -> dict[str, Any]: + return copy.deepcopy({key: value for key, value in vars(self).items() if not key.startswith("_")}) + + +@dataclass(frozen=True) +class _DeferredArtifactContext: + horizon_bytes: int + spans: tuple[tuple[int, bytes], ...] + chunks: list[dict[str, object]] + events: list[dict[str, object]] + + +def benchmark_summary(results: list[OmniInteractCaseResult]) -> dict[str, Any]: + succeeded = sum(result.success for result in results) + eligible = sum(result.success and result.eligible_for_official_eval for result in results) + return { + "total": len(results), + "success": succeeded, + "failed": len(results) - succeeded, + "eligible_for_official_eval": eligible, + "successful_but_ineligible": succeeded - eligible, + "audio_clipped_bytes": sum(result.audio_clipped_bytes for result in results), + "results": [result.as_dict() for result in results], + } + + +def _run_media_command( + command: list[str], + *, + timeout_s: float, + text: bool = False, +) -> subprocess.CompletedProcess[Any]: + try: + result = subprocess.run( + command, + capture_output=True, + check=False, + text=text, + timeout=timeout_s, + ) + except subprocess.TimeoutExpired as exc: + raise TimeoutError(f"Media command timed out after {timeout_s:g}s: {command[0]}") from exc + except FileNotFoundError as exc: + raise RuntimeError(f"Required media command is unavailable: {command[0]}") from exc + if result.returncode: + error = result.stderr if text else result.stderr.decode("utf-8", "ignore") + raise RuntimeError(f"{command[0]} failed: {error.strip()}") + return result + + +def prepare_media( + video: Path, + fps: float, + *, + timeout_s: float, + max_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S, +) -> tuple[float, bytes, list[str | None]]: + probe = [ + "ffprobe", + *"-v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1".split(), + str(video), + ] + duration = float(_run_media_command(probe, text=True, timeout_s=timeout_s).stdout.strip()) + if not math.isfinite(duration) or duration <= 0 or duration > max_duration_s: + raise ValueError(f"Invalid video duration for {video}: {duration!r}") + target_bytes = math.ceil(duration) * PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE + audio = [ + "ffmpeg", + *"-loglevel error -i".split(), + str(video), + "-t", + str(duration), + *f"-vn -f s16le -ac 1 -ar {PCM16_SAMPLE_RATE} -fs".split(), + str(target_bytes), + "pipe:1", + ] + pcm = _run_media_command(audio, timeout_s=timeout_s).stdout + pcm = (pcm + bytes(max(0, target_bytes - len(pcm))))[:target_bytes] + + frame_count = math.ceil(duration * fps) + frames: list[str | None] = [None] * frame_count + with tempfile.TemporaryDirectory(prefix="vllm-omni-frames-") as temp_dir: + output_pattern = Path(temp_dir) / "frame-%06d.jpg" + frame_filter = f"select=gte(t\\,(selected_n+0.5)/{fps}),scale=640:640:force_original_aspect_ratio=decrease" + frame_command = [ + "ffmpeg", + *"-loglevel error -i".split(), + str(video), + "-an", + "-vf", + frame_filter, + *f"-frames:v {frame_count} -vsync vfr -q:v 5".split(), + str(output_pattern), + ] + _run_media_command(frame_command, timeout_s=timeout_s) + for index, frame_path in enumerate(sorted(Path(temp_dir).glob("frame-*.jpg"))[:frame_count]): + frames[index] = base64.b64encode(frame_path.read_bytes()).decode("ascii") + return duration, pcm, frames + + +@dataclass(frozen=True) +class _AudioSegment: + event_index: int + response_id: str + start_s: float + end_s: float + pcm16: bytes + + +class _Playback: + def __init__(self) -> None: + self.cursor = 0 + self.end_s = 0.0 + self.segments: list[_AudioSegment] = [] + self.completed: set[str] = set() + self.completion_acked: set[str] = set() + self.warnings: list[str] = [] + self._total_samples: dict[str, int] = {} + self._response_end_s: dict[str, float] = {} + + def _warn_once(self, warning: str) -> None: + if warning not in self.warnings: + self.warnings.append(warning) + + def ingest(self, events: RealtimeEventCollector) -> None: + while self.cursor < len(events.events): + index, self.cursor = self.cursor, self.cursor + 1 + event = events.events[index] + response_id = events.response_id(event) + if event.get("type") == "response.done": + if response_id: + self.completed.add(response_id) + continue + if event.get("type") != "response.audio.delta": + continue + encoded = event.get("delta") or event.get("audio") + if not response_id: + raise ValueError("response audio has no response_id") + if not isinstance(encoded, str): + raise ValueError("response audio payload is missing") + if event.get("format") is None: + self._warn_once("response.audio.delta omitted format; assumed pcm16") + elif event.get("format") != "pcm16": + raise ValueError("OmniInteract output must be pcm16") + rate = event.get("sample_rate_hz") + if rate is None: + rate = events.output_sample_rate_hz or OUTPUT_SAMPLE_RATE + self._warn_once(f"response.audio.delta omitted sample_rate_hz; assumed {rate}") + if rate != OUTPUT_SAMPLE_RATE: + raise ValueError(f"OmniInteract output must use {OUTPUT_SAMPLE_RATE} Hz audio") + try: + raw = base64.b64decode(encoded, validate=True) + except (ValueError, binascii.Error) as exc: + raise ValueError("response audio is not valid base64") from exc + if not raw or len(raw) % PCM16_BYTES_PER_SAMPLE: + raise ValueError("response audio is empty or not PCM16 aligned") + start = max(events.event_received_at_s[index], self.end_s) + samples = len(raw) // PCM16_BYTES_PER_SAMPLE + segment = _AudioSegment(index, response_id, start, start + samples / rate, raw) + self.segments.append(segment) + self.end_s = segment.end_s + self._total_samples[response_id] = self._total_samples.get(response_id, 0) + samples + self._response_end_s[response_id] = segment.end_s + + async def acknowledge(self, client: RealtimeDuplexClient, now: float | None = None) -> None: + events = client.events + self.ingest(events) + now = time.monotonic() if now is None else now + for response_id in self.completed - self.completion_acked: + total_samples = self._total_samples.get(response_id, 0) + if not total_samples or now < self._response_end_s[response_id]: + continue + played_ms = total_samples * 1000 // OUTPUT_SAMPLE_RATE + await client.send_playback_ack(response_id, played_ms) + self.completion_acked.add(response_id) + + +async def stream_inputs( + client: RealtimeDuplexClient, + pcm: bytes, + frames: Sequence[str | None], + playback: _Playback, +) -> tuple[int, int, float, float]: + bytes_per_second = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE + chunk_bytes = bytes_per_second * _INPUT_CHUNK_MS // 1000 + started_at = time.monotonic() + frame_cursor = sent_frames = 0 + lags: list[float] = [] + for chunks, offset in enumerate(range(0, len(pcm), chunk_bytes), start=1): + end = min(offset + chunk_bytes, len(pcm)) + end_ms = end * 1000 // bytes_per_second + ready: list[str] = [] + while frame_cursor < len(frames) and end_ms >= (frame_cursor + 0.5) * 1000 / VIDEO_FPS: + if frames[frame_cursor]: + ready.append(frames[frame_cursor] or "") + frame_cursor += 1 + sent_frames += len(ready) + event: dict[str, object] = { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(pcm[offset:end]).decode("ascii"), + "input_audio_format": "pcm16", + "sample_rate_hz": PCM16_SAMPLE_RATE, + "duration_ms": (end - offset) * 1000 // bytes_per_second, + "audio_end_ms": end_ms, + } + if ready: + event["video_frames"] = ready + lags.append(max(0.0, time.monotonic() - started_at - offset / bytes_per_second)) + await client.send(event) + await playback.acknowledge(client) + await asyncio.sleep(max(0.0, started_at + end_ms / 1000 - time.monotonic())) + return chunks if pcm else 0, sent_frames, sum(lags) / len(lags) if lags else 0.0, max(lags, default=0.0) + + +def _response_status(event: dict[str, object]) -> str | None: + response = event.get("response") + if isinstance(response, dict) and response.get("status"): + return str(response["status"]) + status = event.get("status") + return str(status) if status else None + + +def response_ledger( + collector: RealtimeEventCollector, + *, + end: int | None = None, +) -> tuple[set[str], set[str]]: + created: set[str] = set() + done: set[str] = set() + for event in collector.events[:end]: + event_type = event.get("type") + if event_type not in {"response.created", "response.done"}: + continue + response_id = collector.response_id(event) + if not response_id: + raise ValueError(f"{event_type} has no response_id") + if event_type == "response.created": + if response_id in created: + raise ValueError(f"duplicate response.created for {response_id}") + created.add(response_id) + continue + if response_id not in created: + raise ValueError(f"response.done without response.created for {response_id}") + if response_id in done: + raise ValueError(f"duplicate response.done for {response_id}") + response = event.get("response") if isinstance(event.get("response"), dict) else {} + details = event.get("status_details") if isinstance(event.get("status_details"), dict) else {} + nested = response.get("status_details") if isinstance(response.get("status_details"), dict) else {} + if "failed" in {event.get("status"), response.get("status"), details.get("type"), nested.get("type")}: + raise ValueError(f"response.done reports failure for {response_id}") + done.add(response_id) + return created, done + + +def _is_model_listen_decision(event: dict[str, object]) -> bool: + response = event.get("response") + metadata = response.get("metadata") if isinstance(response, dict) else event.get("metadata") + return isinstance(metadata, dict) and metadata.get("buffering") is not True and metadata.get("model_listen") is True + + +def _has_post_commit_decision( + collector: RealtimeEventCollector, + events: list[dict[str, object]], + *, + prior_response_ids: set[str], +) -> bool: + pending_prior = set(prior_response_ids) + for event in events: + event_type = event.get("type") + if event_type == "response.done": + response_id = collector.response_id(event) + if response_id in pending_prior: + pending_prior.remove(response_id) + continue + if not pending_prior and _response_status(event) != "cancelled": + return True + elif event_type == "response.listen" and not pending_prior and _is_model_listen_decision(event): + return True + return False + + +def _raise_if_session_terminated( + collector: RealtimeEventCollector, + from_index: int, + *, + explicit_close_from: int | None = None, +) -> None: + errors = collector.errors() + if errors: + raise RuntimeError(str(errors[-1])) + for index, event in enumerate(collector.events[from_index:], start=from_index): + event_type = event.get("type") + if event_type not in {"session.expired", "session.closed"}: + continue + reason = event.get("reason") + nested = event.get("event") + if reason is None and isinstance(nested, dict): + reason = nested.get("reason") + expected = ( + event_type == "session.closed" + and explicit_close_from is not None + and index >= explicit_close_from + and reason is None + ) + if expected: + continue + detail = f": {reason}" if reason else "" + prefix = "Unexpected " if explicit_close_from is not None else "" + raise RuntimeError(f"{prefix}{event_type}{detail}") + + +def _ensure_final_commit_tail(pcm: bytes, events: list[dict[str, object]]) -> bytes: + """Reserve one PCM sample so an exact model unit is flushed by commit. + + A complete unit is otherwise emitted before commit, so its asynchronous + decision cannot be correlated with the accepted final input. Keeping an + almost-full unit buffered makes the server's final residual flush carry + that correlation; the server pads the missing sample back to a full unit. + """ + + period_ms = chunk_period_ms(events) + if len(pcm) >= PCM16_BYTES_PER_SAMPLE and not has_residual_model_unit(pcm, chunk_period_ms=period_ms): + return pcm[:-PCM16_BYTES_PER_SAMPLE] + return pcm + + +async def wait_for_session_completion( + client: RealtimeDuplexClient, + playback: _Playback, + *, + commit_from: int, + session_from: int | None = None, + timeout_s: float, + settle_s: float, +) -> int: + deadline = time.monotonic() + timeout_s + committed_index: int | None = None + prior_response_ids: set[str] = set() + last_event_count = len(client.events.events) + stable_since = time.monotonic() + while time.monotonic() < deadline: + client.raise_if_reader_stopped() + _raise_if_session_terminated(client.events, commit_from if session_from is None else session_from) + await playback.acknowledge(client) + if len(client.events.events) != last_event_count: + last_event_count = len(client.events.events) + stable_since = time.monotonic() + if committed_index is None: + committed_index = next( + ( + index + for index in range(commit_from, len(client.events.events)) + if client.events.events[index].get("type") == "input_audio_buffer.committed" + ), + None, + ) + if committed_index is not None: + committed_event = client.events.events[committed_index] + committed = committed_event.get("event") + if isinstance(committed, dict) and committed.get("overlap_deferred") is True: + prior_created, prior_done = response_ledger(client.events, end=committed_index) + prior_response_ids = prior_created - prior_done + stable_since = time.monotonic() + if committed_index is not None: + created, done = response_ledger(client.events) + post_commit = client.events.events[committed_index + 1 :] + decision = _has_post_commit_decision( + client.events, + post_commit, + prior_response_ids=prior_response_ids, + ) + if created == done and decision and time.monotonic() - stable_since >= settle_s: + return committed_index + await asyncio.sleep(0.05) + missing: set[str] = set() + with contextlib.suppress(ValueError): + created, done = response_ledger(client.events) + missing = created - done + raise TimeoutError( + "Timed out waiting for committed input and stable responses" + + (f"; unfinished response_ids={sorted(missing)}" if missing else "") + ) + + +def _output_dir(root: Path, case: OmniInteractCase) -> Path: + relative = case.video_rel.replace("\\", "/") + stem = Path(relative).with_suffix("").as_posix().replace("/", "__") + digest = hashlib.sha256(f"{case.subset}/{relative}".encode()).hexdigest()[:8] + return root / case.subset / f"{stem}--{digest}" + + +@contextlib.contextmanager +def omniinteract_output_lock(root: Path) -> Iterator[None]: + root.mkdir(parents=True, exist_ok=True) + with (root / ARTIFACT_LOCK_FILE).open("a+b") as lock_file: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + + +def _atomic_replace(path: Path, writer: Callable[[Path], None]) -> None: + descriptor, temporary_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + ) + os.close(descriptor) + temporary = Path(temporary_name) + try: + writer(temporary) + temporary.replace(path) + finally: + temporary.unlink(missing_ok=True) + + +def _atomic_write_text(path: Path, value: str) -> None: + _atomic_replace(path, lambda temporary: temporary.write_text(value, encoding="utf-8")) + + +def _atomic_write_json(path: Path, value: object) -> None: + _atomic_write_text(path, json.dumps(value, ensure_ascii=False, indent=2) + "\n") + + +def _atomic_write_wav(path: Path, pcm: bytes | bytearray, rate: int) -> None: + _atomic_replace( + path, + lambda temporary: write_pcm16_wav( + temporary, + cast(bytes, pcm), # wave.writeframes accepts bytearray. + sample_rate_hz=rate, + ), + ) + + +def _clear_artifacts(directory: Path, names: Sequence[str]) -> None: + directory.mkdir(parents=True, exist_ok=True) + for name in names: + (directory / name).unlink(missing_ok=True) + + +def _artifact_summary(case: OmniInteractCase, result: OmniInteractCaseResult) -> dict[str, Any]: + return { + **result.as_dict(), + "annotation": str(case.annotation_path), + "scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type, + } + + +def _publish_success_artifacts( + directory: Path, + result: OmniInteractCaseResult, + context: _DeferredArtifactContext, + summary: dict[str, Any], +) -> None: + pcm = bytearray(context.horizon_bytes) + for offset, chunk in context.spans: + pcm[offset : offset + len(chunk)] = chunk + try: + _clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json")) + _atomic_write_wav(directory / "output.wav", pcm, OUTPUT_SAMPLE_RATE) + _atomic_write_json( + directory / "wav_transcript.json", + { + "text": result.transcript, + "chunks": context.chunks, + "timestamp_semantics": "serialized playback queue time relative to input streaming start", + }, + ) + _atomic_write_json(directory / "events.json", context.events) + _atomic_write_json(directory / "result.json", result.as_dict()) + _atomic_write_json(directory / ".done", summary) + except Exception: + _clear_artifacts(directory, SUCCESS_ARTIFACTS) + result.success = result.eligible_for_official_eval = False + if "artifact_write_failed" not in result.official_eval_ineligible_reasons: + result.official_eval_ineligible_reasons.append("artifact_write_failed") + raise + + +def _collect_output( + collector: RealtimeEventCollector, + playback: _Playback, + *, + stream_start: float, + video_duration_s: float, + require_response: bool, + retain_events: bool, + result: OmniInteractCaseResult, +) -> _DeferredArtifactContext: + playback.ingest(collector) + created, done = response_ledger(collector) + if created != done: + raise ValueError(f"unfinished response_ids: {sorted(created - done)}") + done_events = [event for event in collector.events if event.get("type") == "response.done"] + cancelled = any(_response_status(event) == "cancelled" for event in done_events) + terminal_responses = { + response_id + for event in done_events + if _response_status(event) != "cancelled" and (response_id := collector.response_id(event)) + } + horizon_s = math.ceil(video_duration_s) + horizon_bytes = horizon_s * OUTPUT_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE + spans: list[tuple[int, bytes]] = [] + clipped_bytes = 0 + response_times: dict[str, list[float]] = {} + responses_with_audio: set[str] = set() + for segment in playback.segments: + response_id = segment.response_id + if response_id not in created: + raise ValueError("response audio has no matching response.created") + start_s = max(0.0, segment.start_s - stream_start) + end_s = max(start_s, segment.end_s - stream_start) + offset = round(start_s * OUTPUT_SAMPLE_RATE) * PCM16_BYTES_PER_SAMPLE + writable = max(0, min(len(segment.pcm16), horizon_bytes - offset)) + clipped_bytes += len(segment.pcm16) - writable + if writable: + clipped = segment.pcm16[:writable] + spans.append((offset, clipped)) + if any(clipped): + responses_with_audio.add(response_id) + timing = response_times.setdefault(response_id, [start_s, end_s]) + timing[0], timing[1] = min(timing[0], start_s), max(timing[1], end_s) + + text_event_types = {"response.audio_transcript.delta", "response.output_text.delta", "response.text.delta"} + text_ids = {collector.response_id(event) for event in collector.events if event.get("type") in text_event_types} + if None in text_ids or not text_ids <= created: + raise ValueError("response transcript has no matching response.created") + chunks: list[dict[str, object]] = [] + texts: list[str] = [] + for response_id in collector.response_ids: + text = collector.response_text(response_id).strip() + if not text: + continue + response_timing = response_times.get(response_id) + if response_timing is None or response_timing[0] >= horizon_s: + continue + texts.append(text) + chunks.append( + { + "response_id": response_id, + "text": text, + "timestamp": [round(response_timing[0], 6), round(min(response_timing[1], horizon_s), 6)], + } + ) + complete_outputs = terminal_responses & responses_with_audio & {str(chunk["response_id"]) for chunk in chunks} + if require_response and not complete_outputs: + raise ValueError("OmniInteract E2E requires a response with audio and transcript") + result.audio_bytes, result.audio_clipped_bytes = ( + sum(len(segment.pcm16) for segment in playback.segments), + clipped_bytes, + ) + result.transcript, result.responses, result.success = " ".join(texts).strip(), len(created), True + result.artifact_warnings = list(playback.warnings) + result.official_eval_ineligible_reasons = [] + if clipped_bytes: + result.official_eval_ineligible_reasons.append("audio_clipped") + if cancelled: + result.official_eval_ineligible_reasons.append("cancelled_response") + result.eligible_for_official_eval = not result.official_eval_ineligible_reasons + audio_sizes = {segment.event_index: len(segment.pcm16) for segment in playback.segments} + events = ( + [ + { + **{key: value for key, value in event.items() if key not in {"delta", "audio"}}, + "audio_bytes": audio_sizes[index], + } + if event.get("type") == "response.audio.delta" + else dict(event) + for index, event in enumerate(collector.events) + ] + if retain_events + else [] + ) + return _DeferredArtifactContext(horizon_bytes, tuple(spans), chunks, events) + + +def prepare_success_artifacts( + root: Path, + case: OmniInteractCase, + collector: RealtimeEventCollector, + *, + playback: _Playback | None = None, + stream_start: float, + video_duration_s: float, + require_response: bool, + result: OmniInteractCaseResult, + capture_artifacts: bool = True, +) -> dict[str, Any]: + directory = _output_dir(root, case) + context = _collect_output( + collector, + playback or _Playback(), + stream_start=stream_start, + video_duration_s=video_duration_s, + require_response=require_response, + retain_events=capture_artifacts, + result=result, + ) + result.output_dir = str(directory.resolve()) + summary = _artifact_summary(case, result) + if capture_artifacts: + result._artifact_context = context + return summary + + +def write_failure_artifacts(root: Path, case: OmniInteractCase, result: OmniInteractCaseResult) -> None: + directory = _output_dir(root, case) + result.output_dir = str(directory.resolve()) + result.success = result.eligible_for_official_eval = False + _clear_artifacts(directory, SUCCESS_ARTIFACTS) + _atomic_write_json(directory / ".failed.json", _artifact_summary(case, result)) + + +def publish_deferred_case_artifacts( + root: Path, + case: OmniInteractCase, + result: OmniInteractCaseResult, +) -> None: + context = result._artifact_context + result._artifact_context = None + if not result.success: + write_failure_artifacts(root, case, result) + return + if context is None: + raise RuntimeError("OmniInteract benchmark output lost its deferred artifact context") + directory = _output_dir(root, case) + _publish_success_artifacts(directory, result, context, _artifact_summary(case, result)) + + +def clear_case_artifacts(root: Path, case: OmniInteractCase) -> None: + directory = _output_dir(root, case) + _clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json")) + + +def clear_batch_artifacts(root: Path) -> None: + _clear_artifacts(root, BATCH_ARTIFACTS) + + +def write_batch_artifacts( + root: Path, + cases: list[OmniInteractCase], + results: list[OmniInteractCaseResult], +) -> None: + root.mkdir(parents=True, exist_ok=True) + rows = [ + case_manifest(case, _output_dir(root, case)) + for case, result in zip(cases, results, strict=True) + if result.success and result.eligible_for_official_eval + ] + ineligible = [result for result in results if result.success and not result.eligible_for_official_eval] + if ineligible: + reasons = Counter( + reason for result in ineligible for reason in (result.official_eval_ineligible_reasons or ["unspecified"]) + ) + logger.warning( + "%d successful OmniInteract cases were excluded from official evaluation: %s", + len(ineligible), + ", ".join(f"{reason}={count}" for reason, count in sorted(reasons.items())), + ) + _atomic_write_json(root / "batch_summary.json", benchmark_summary(results)) + _atomic_write_text( + root / "official_eval_manifest.jsonl", + "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows), + ) + + +def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str: + endpoint = ( + config.endpoint + if urlsplit(config.endpoint).scheme + else urljoin(config.base_url.rstrip("/") + "/", config.endpoint.lstrip("/")) + ) + return build_realtime_url( + endpoint, + config.model, + autostart=False, + native_duplex=True, + session_id=session_id, + ) + + +def _populate_response_metrics( + result: OmniInteractCaseResult, + collector: RealtimeEventCollector, + *, + stream_start: float, +) -> None: + measurement_origin = { + "ttft": "response.created client receive to first non-empty text delta", + "ttfp": "response.created client receive to first audio packet", + "rtf": "response.created client receive to last audio packet divided by emitted audio duration", + } + request_metrics: list[dict[str, object]] = [] + output_tokens = 0 + for request_index, response_id in enumerate(collector.response_ids): + timing = collector.timing_summary( + after_s=stream_start, + input_committed_at_s=None, + response_id=response_id, + measurement_origin=measurement_origin, + ) + raw_metric = timing.get("request_metrics") + stage0 = timing.get("stage0_tokens") + metric = { + "session_id": result.session_id, + "request_index": request_index, + "response_id": response_id, + **(raw_metric if isinstance(raw_metric, dict) else {}), + } + if isinstance(stage0, dict): + metric["stage0_tokens"] = dict(stage0) + output_tokens += int(stage0.get("output_token_count") or 0) + if isinstance(raw_metric, dict) or isinstance(stage0, dict): + request_metrics.append(metric) + result.output_tokens = output_tokens + result.duplex_request_metrics = request_metrics + result.duplex_session_metrics = summarize_session_request_metrics( + request_metrics, + session_id=result.session_id, + ) + + +async def run_omniinteract_case( + case: OmniInteractCase, + config: OmniInteractBenchmarkConfig, + *, + request_index: int | str, + capture_artifacts: bool = True, + prepared_input: OmniInteractPreparedInput | None = None, +) -> OmniInteractCaseResult: + if not config.ref_audio: + raise ValueError("ref_audio is required for MiniCPM-o native-duplex audio output") + for name, value in ( + ("timeout_s", config.timeout_s), + ("media_timeout_s", config.media_timeout_s), + ("max_video_duration_s", config.max_video_duration_s), + ): + if not math.isfinite(value) or value <= 0: + raise ValueError(f"{name} must be finite and positive") + session_id = f"omniinteract:{case.subset}:{request_index}:{time.monotonic_ns()}" + result = OmniInteractCaseResult( + case.subset, + str(case.video_path), + str(_output_dir(config.output_root, case).resolve()), + session_id=session_id, + ) + started_at = time.monotonic() + try: + if prepared_input is None: + reference_audio = reference_audio_data_url(config.ref_audio) + duration, pcm, frames = await asyncio.to_thread( + prepare_media, + case.video_path, + VIDEO_FPS, + timeout_s=config.media_timeout_s, + max_duration_s=config.max_video_duration_s, + ) + else: + reference_audio = prepared_input.ref_audio_data_url + duration = prepared_input.duration_s + pcm = prepared_input.pcm16 + frames = prepared_input.video_frames + if not any(frames): + raise ValueError(f"No video frames were decoded from {case.video_path}") + async with RealtimeDuplexClient( + _websocket_url(config, session_id), additional_headers=config.extra_headers + ) as client: + session_from = len(client.events.events) + await client.configure( + config.model, + ref_audio=reference_audio, + session_id=session_id, + extra_body=config.extra_body, + idle_timeout_s=config.timeout_s, + timeout_s=min(config.timeout_s, 20.0), + ) + pcm = _ensure_final_commit_tail(pcm, client.events.events) + playback = _Playback() + stream_start = time.monotonic() + try: + input_duration_s = len(pcm) / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE) + upload_timeout_s = config.timeout_s + input_duration_s + try: + chunks, frame_count, mean_lag, max_lag = await asyncio.wait_for( + stream_inputs(client, pcm, frames, playback), timeout=upload_timeout_s + ) + except asyncio.TimeoutError as exc: + raise TimeoutError(f"Realtime upload timed out after {upload_timeout_s:g}s") from exc + commit_from = len(client.events.events) + await client.commit() + await wait_for_session_completion( + client, + playback, + commit_from=commit_from, + session_from=session_from, + timeout_s=config.timeout_s, + settle_s=_COMPLETION_SETTLE_S, + ) + await playback.acknowledge(client, playback.end_s) + _raise_if_session_terminated(client.events, session_from) + close_from = len(client.events.events) + await client.close_session(timeout_s=min(config.timeout_s, 20.0)) + except Exception: + with contextlib.suppress(Exception): + await client.close_session(timeout_s=min(config.timeout_s, 20.0)) + raise + errors = client.events.errors() + if errors: + raise RuntimeError(str(errors[-1])) + _raise_if_session_terminated(client.events, session_from, explicit_close_from=close_from) + result.latency_s = time.monotonic() - started_at + result.input_audio_chunks, result.input_video_frames = chunks, frame_count + result.pacing_mean_lag_s, result.pacing_max_lag_s = mean_lag, max_lag + _populate_response_metrics(result, client.events, stream_start=stream_start) + await asyncio.to_thread( + prepare_success_artifacts, + config.output_root, + case, + client.events, + playback=playback, + stream_start=stream_start, + video_duration_s=duration, + require_response=config.require_response, + result=result, + capture_artifacts=capture_artifacts, + ) + except Exception as exc: + result.success = result.eligible_for_official_eval = False + if "case_failed" not in result.official_eval_ineligible_reasons: + result.official_eval_ineligible_reasons.append("case_failed") + result.error = str(exc) + result.latency_s = max(0.0, time.monotonic() - started_at) + return result diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 3f83a1b38..fd52c197b 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import contextlib import io @@ -16,7 +19,6 @@ from dataclasses import dataclass, replace from datetime import datetime from pathlib import Path from typing import Any, Literal -from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import aiohttp import numpy as np @@ -31,6 +33,7 @@ from vllm.benchmarks.lib.endpoint_request_func import ( RequestFuncOutput, StreamedResponseHandler, _get_chat_content, + _get_headers, _update_headers_common, _update_payload_common, _validate_api_url, @@ -46,6 +49,13 @@ from vllm_omni.benchmarks.data_modules.daily_omni_dataset import ( daily_omni_local_videos_dir, resolve_daily_omni_local_root, ) +from vllm_omni.benchmarks.data_modules.omniinteract_dataset import ( + DEFAULT_OMNIINTERACT_REPO, + OmniInteractDataset, + OmniInteractPreparedInput, + OmniInteractSampleRequest, + OmniInteractSessionOptions, +) from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT, @@ -56,8 +66,27 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( ) from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDataset from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset +from vllm_omni.benchmarks.omniinteract import ( + VIDEO_FPS, + OmniInteractBenchmarkConfig, + OmniInteractCaseResult, + clear_batch_artifacts, + clear_case_artifacts, + prepare_media, + publish_deferred_case_artifacts, + run_omniinteract_case, + write_failure_artifacts, +) +from vllm_omni.benchmarks.omniinteract import ( + benchmark_summary as omniinteract_benchmark_summary, +) +from vllm_omni.benchmarks.omniinteract import ( + write_batch_artifacts as write_omniinteract_batch_artifacts, +) from vllm_omni.experimental.fullduplex.client import ( RealtimeDuplexClient, + build_realtime_url, + reference_audio_data_url, summarize_session_request_metrics, wait_for, ) @@ -90,6 +119,17 @@ def should_request_stage_metrics(args: Any) -> bool: if backend in _IMAGE_STAGE_METRICS_BACKENDS: return True + percentile_metrics = getattr(args, "percentile_metrics", ()) or () + if isinstance(percentile_metrics, str): + percentile_metrics = percentile_metrics.split(",") + selected_metrics = {str(metric).strip().lower() for metric in percentile_metrics} + if backend == "openai-chat-omni" and selected_metrics.intersection({"tpot", "itl"}): + # Client receive timestamps normally provide these metrics. Under + # enough load, however, every text token can be coalesced into one + # event-loop read, leaving no positive interval. Request the engine's + # Stage 0 timings so that path still has an authoritative measurement. + return True + extra_body = getattr(args, "extra_body", None) or {} modalities = extra_body.get("modalities") if isinstance(extra_body, dict) else None return backend == "openai-chat-omni" and "image" in (modalities or []) @@ -221,6 +261,91 @@ def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFu rfi.extra_body = base +def _attach_omniinteract_to_request_func_input(sample: SampleRequest, rfi: RequestFuncInput) -> None: + if not isinstance(sample, OmniInteractSampleRequest): + return + setattr(rfi, "omniinteract_case", sample.omniinteract_case) + setattr(rfi, "omniinteract_options", sample.omniinteract_options) + setattr(rfi, "omniinteract_prepared_input", sample.omniinteract_prepared_input) + + +def _append_error(existing: str, message: str) -> str: + return f"{existing}\n{message}" if existing else message + + +def _finalize_omniinteract_batch( + input_requests: list[SampleRequest], + outputs: list[RequestFuncOutput], +) -> dict[str, object] | None: + rows = [ + (sample, output) + for sample, output in zip(input_requests, outputs, strict=True) + if isinstance(sample, OmniInteractSampleRequest) + ] + if not rows: + return None + options = rows[0][0].omniinteract_options + if not isinstance(options, OmniInteractSessionOptions): + raise RuntimeError("OmniInteract benchmark output lost its artifact options") + cases, results = [], [] + artifact_errors: list[str] = [] + for sample, output in rows: + case = sample.omniinteract_case + result = getattr(output, "omniinteract_case_result", None) + if case is None or not isinstance(result, OmniInteractCaseResult): + raise RuntimeError("OmniInteract benchmark output lost its dataset identity") + cases.append(case) + results.append(result) + try: + publish_deferred_case_artifacts(options.output_root, case, result) + except Exception as exc: # noqa: BLE001 - artifact backends raise heterogeneous errors + result.success = result.eligible_for_official_eval = False + if "artifact_write_failed" not in result.official_eval_ineligible_reasons: + result.official_eval_ineligible_reasons.append("artifact_write_failed") + artifact_error = f"Artifact publication failed: {exc}" + result.error = _append_error(result.error, artifact_error) + output.error = _append_error(output.error, artifact_error) + message = f"{case.subset}/{case.video_rel}: {artifact_error}" + artifact_errors.append(message) + logger.exception(message) + try: + write_failure_artifacts(options.output_root, case, result) + except Exception as recovery_exc: # noqa: BLE001 - preserve benchmark metrics after I/O failure + message = f"{case.subset}/{case.video_rel}: failure artifact publication failed: {recovery_exc}" + artifact_errors.append(message) + logger.exception(message) + try: + write_omniinteract_batch_artifacts(options.output_root, cases, results) + except Exception as exc: # noqa: BLE001 - preserve benchmark metrics after I/O failure + message = f"Batch artifact publication failed: {exc}" + artifact_errors.append(message) + logger.exception(message) + try: + clear_batch_artifacts(options.output_root) + except Exception as cleanup_exc: # noqa: BLE001 - best-effort removal of partial batch artifacts + message = f"Partial batch artifact cleanup failed: {cleanup_exc}" + artifact_errors.append(message) + logger.exception(message) + summary = omniinteract_benchmark_summary(results) + compact_summary = {key: value for key, value in summary.items() if key != "results"} + compact_summary["artifacts_complete"] = not artifact_errors + if artifact_errors: + compact_summary["artifact_errors"] = artifact_errors + return compact_summary + + +def _prepare_omniinteract_batch(input_requests: list[SampleRequest]) -> None: + roots: set[Path] = set() + for sample in input_requests: + if not isinstance(sample, OmniInteractSampleRequest): + continue + root = sample.omniinteract_options.output_root.resolve() + roots.add(root) + clear_case_artifacts(sample.omniinteract_options.output_root, sample.omniinteract_case) + for root in roots: + clear_batch_artifacts(root) + + def _daily_omni_repo_from_args(args) -> str | None: """Resolve HuggingFace repo id for Daily-Omni from CLI args. @@ -250,6 +375,7 @@ def get_samples(args, tokenizer): "ttsd", "sound-effect", ) + is_omniinteract = args.dataset_name == "omniinteract" # Check if we need to handle omni-related backends/datasets is_omni_backend = args.backend in [ @@ -258,12 +384,68 @@ def get_samples(args, tokenizer): "openai-realtime-duplex", "daily-omni", ] - is_omni_dataset = is_daily_omni or is_seed_tts or args.dataset_name == "random-mm" + is_omni_dataset = is_daily_omni or is_seed_tts or is_omniinteract or args.dataset_name == "random-mm" if not is_omni_backend and not is_omni_dataset: # Not an omni-related request, delegate to original implementation return get_samples_old(args, tokenizer) + if is_omniinteract: + dataset_path = getattr(args, "dataset_path", None) + data_root = None + dataset_repo = DEFAULT_OMNIINTERACT_REPO + if dataset_path: + candidate = Path(dataset_path).expanduser() + if candidate.exists() or candidate.is_absolute(): + data_root = str(candidate) + else: + dataset_repo = str(dataset_path) + dataset = OmniInteractDataset( + data_root=data_root, + dataset_repo=dataset_repo, + subsets=tuple(getattr(args, "omniinteract_subsets")), + random_seed=args.seed, + disable_shuffle=getattr(args, "disable_shuffle", False), + ) + options = OmniInteractSessionOptions( + output_root=Path(getattr(args, "omniinteract_output_dir")), + timeout_s=float(getattr(args, "omniinteract_timeout_s")), + media_timeout_s=float(getattr(args, "omniinteract_media_timeout_s")), + ref_audio=str(getattr(args, "omniinteract_ref_audio")), + require_response=bool(getattr(args, "omniinteract_require_response")), + max_video_duration_s=float(getattr(args, "omniinteract_max_video_duration_s")), + ) + requests = dataset.sample( + tokenizer, + args.num_prompts, + request_id_prefix=args.request_id_prefix, + options=options, + ) + if not requests: + raise ValueError("No OmniInteract sessions were selected") + encoded_ref_audio = reference_audio_data_url(options.ref_audio) + assert encoded_ref_audio is not None + for request in requests: + case = request.omniinteract_case + assert isinstance(request, OmniInteractSampleRequest) and case is not None + duration, pcm, frames = prepare_media( + case.video_path, + VIDEO_FPS, + timeout_s=options.media_timeout_s, + max_duration_s=options.max_video_duration_s, + ) + if not any(frames): + raise ValueError(f"No video frames were decoded from {case.video_path}") + request.omniinteract_prepared_input = OmniInteractPreparedInput( + duration_s=duration, + pcm16=pcm, + video_frames=tuple(frames), + ref_audio_data_url=encoded_ref_audio, + ) + # Replace OmniInteract's ``0`` (all) with the measured request count. + args.num_prompts = len(requests) + return requests + # Handle Daily-Omni dataset if is_daily_omni: # Support: @@ -476,6 +658,7 @@ class MixRequestFuncOutput(RequestFuncOutput): image_pixels: int = 0 denoise_step_latency_ms: float = 0.0 text_latency: float = 0.0 + tpot_measured: bool = True #: Worst-case streaming-audio underrun (wall-clock seconds the player #: would have been starved). Populated by the audio-speech backend; ``0.0`` #: for backends that do not run continuity analysis. @@ -682,6 +865,33 @@ def _update_output_stage_metrics_from_payload( output.stage_metrics.update(stage_snapshot) +def _apply_chat_stage0_token_timings(output: MixRequestFuncOutput) -> bool: + """Apply native Stage 0 timings from a chat response snapshot.""" + stage_metrics = output.stage_metrics + if not isinstance(stage_metrics, dict): + return False + stage0 = stage_metrics.get("0") + if not isinstance(stage0, dict): + return False + + output_tokens = coerce_positive_int_scalar(stage0.get(defs.NUM_TOKENS_OUT)) + expected_output_tokens = coerce_positive_int_scalar(output.output_tokens) + if output_tokens is None or expected_output_tokens is None or output_tokens != expected_output_tokens: + return False + + return _apply_stage0_token_timings( + output, + [ + { + "output_token_count": output_tokens, + "itls_ms": stage0.get(defs.VLLM_ITLS_MS), + "tpot_ms": stage0.get(defs.VLLM_TPOT_MS), + } + ], + expected_output_tokens=expected_output_tokens, + ) + + def _image_metrics_from_stage_metrics(metrics: dict[str, Any] | None) -> tuple[int, float, int, float]: if not isinstance(metrics, dict): return 0, 0.0, 0, 0.0 @@ -949,6 +1159,14 @@ async def async_request_openai_chat_omni_completions( output.latency = timestamp - st output.generated_text = generated_text + if output.output_tokens > 1 and not any( + isinstance(value, int | float) + and not isinstance(value, bool) + and np.isfinite(value) + and value > 0 + for value in output.itl + ): + _apply_chat_stage0_token_timings(output) if output.itl: # Align text_latency with ITL so TPOT formula and # mean(ITL) are consistent. Do NOT infer output_tokens @@ -1313,11 +1531,158 @@ async def async_request_openai_audio_speech( def _realtime_websocket_url(api_url: str) -> str: - parts = urlsplit(api_url) - scheme = "wss" if parts.scheme == "https" else "ws" - query = dict(parse_qsl(parts.query, keep_blank_values=True)) - query.setdefault("duplex", "1") - return urlunsplit((scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)) + return build_realtime_url(api_url, None, native_duplex=None) + + +def _nonnegative_number(value: object) -> bool: + return isinstance(value, int | float) and not isinstance(value, bool) and np.isfinite(value) and value >= 0 + + +def _apply_stage0_token_timings( + output: MixRequestFuncOutput, + raw_stages: list[object], + *, + expected_output_tokens: int, +) -> bool: + """Populate aggregate token timing from complete per-response Stage 0 metrics.""" + output.itl = [] + output.text_latency = output.ttft + output.tpot_measured = False + + stages: list[tuple[int, dict[str, object]]] = [] + for stage in raw_stages: + if not isinstance(stage, dict): + return False + token_count = stage.get("output_token_count") + if not isinstance(token_count, int) or isinstance(token_count, bool) or token_count < 0: + return False + stages.append((token_count, stage)) + + if not stages or sum(count for count, _ in stages) != expected_output_tokens: + return False + + timed = [(count - 1, stage) for count, stage in stages if count > 1] + if not timed: + return False + + measured_itls: list[float] = [] + exact_itls = True + for intervals, stage in timed: + values = stage.get("itls_ms") + if ( + not isinstance(values, list) + or len(values) != intervals + or not all(_nonnegative_number(value) for value in values) + ): + exact_itls = False + break + measured_itls.extend(float(value) / 1000.0 for value in values if isinstance(value, int | float)) + if exact_itls: + if sum(measured_itls) > 0: + output.itl = measured_itls + output.text_latency = output.ttft + sum(output.itl) + output.tpot_measured = True + return True + + weighted_tpot_ms = 0.0 + for intervals, stage in timed: + tpot_ms = stage.get("tpot_ms") + if not _nonnegative_number(tpot_ms) or not isinstance(tpot_ms, int | float) or tpot_ms <= 0: + break + weighted_tpot_ms += float(tpot_ms) * intervals + else: + decode_intervals = sum(intervals for intervals, _ in timed) + weighted_tpot = weighted_tpot_ms / decode_intervals / 1000.0 + # ``calculate_metrics`` consumes a request-shaped latency. Preserve + # the weighted per-response TPOT when projecting it onto the combined + # output token count; gaps and TTFTs between turns are intentionally + # excluded. + output.text_latency = output.ttft + weighted_tpot * (expected_output_tokens - 1) + output.tpot_measured = True + return True + + return False + + +async def _async_request_omniinteract( + request_func_input: RequestFuncInput, + *, + pbar: tqdm | None, +) -> MixRequestFuncOutput: + case = getattr(request_func_input, "omniinteract_case", None) + options = getattr(request_func_input, "omniinteract_options", None) + output = MixRequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + output.start_time = time.perf_counter() + try: + if case is None or not isinstance(options, OmniInteractSessionOptions): + raise ValueError("OmniInteract RequestFuncInput is missing its dataset session layout") + prepared_input = getattr(request_func_input, "omniinteract_prepared_input", None) + if not isinstance(prepared_input, OmniInteractPreparedInput): + raise ValueError("OmniInteract media must be prepared before benchmark request timing") + headers = _get_headers() + _update_headers_common(headers, request_func_input) + config = OmniInteractBenchmarkConfig( + endpoint=request_func_input.api_url, + model=request_func_input.model_name or request_func_input.model, + output_root=options.output_root, + timeout_s=options.timeout_s, + media_timeout_s=options.media_timeout_s, + max_video_duration_s=options.max_video_duration_s, + ref_audio=options.ref_audio, + require_response=options.require_response, + extra_headers=headers or None, + extra_body=dict(request_func_input.extra_body or {}) or None, + ) + case_result = await run_omniinteract_case( + case, + config, + request_index=request_func_input.request_id or uuid.uuid4().hex, + capture_artifacts=request_func_input.request_id is not None, + prepared_input=prepared_input, + ) + output.latency = case_result.latency_s + output.generated_text = case_result.transcript + output.output_tokens = case_result.output_tokens + output.audio_duration = case_result.audio_bytes / (24_000 * 2) + output.audio_frames = case_result.audio_bytes // 2 + session_metrics = case_result.duplex_session_metrics + output.ttft = float(session_metrics.get("mean_ttft_ms") or 0.0) / 1000.0 + output.audio_ttfp = float(session_metrics.get("mean_ttfp_ms") or 0.0) / 1000.0 + output.audio_rtf = float(session_metrics.get("mean_rtf") or 0.0) + token_timing_measured = _apply_stage0_token_timings( + output, + [request_metric.get("stage0_tokens") for request_metric in case_result.duplex_request_metrics], + expected_output_tokens=output.output_tokens, + ) + if not token_timing_measured: + if output.output_tokens > 1 or (output.output_tokens == 0 and output.generated_text): + logger.warning( + "OmniInteract session %s omitted complete engine token timing; standard TPOT/ITL are unavailable", + case_result.session_id, + ) + output.duplex_request_metrics = case_result.duplex_request_metrics + output.duplex_session_metrics = session_metrics + output.success = case_result.success + output.error = case_result.error + setattr(output, "omniinteract_case", case) + setattr(output, "omniinteract_case_result", case_result) + except Exception: + output.success = False + output.error = traceback.format_exc() + logger.error("OmniInteract Realtime request failed: %s", output.error) + if case is not None: + case_result = OmniInteractCaseResult( + subset=case.subset, + video=str(case.video_path), + output_dir="", + error=output.error, + ) + setattr(output, "omniinteract_case", case) + setattr(output, "omniinteract_case_result", case_result) + if pbar: + pbar.update(1) + return output async def async_request_openai_realtime_duplex( @@ -1325,14 +1690,18 @@ async def async_request_openai_realtime_duplex( session: aiohttp.ClientSession, pbar: tqdm | None = None, ) -> MixRequestFuncOutput: - """Run one or more Seed-TTS turns through one explicit Realtime TTS session.""" del session + if getattr(request_func_input, "omniinteract_case", None) is not None: + return await _async_request_omniinteract(request_func_input, pbar=pbar) output = MixRequestFuncOutput() output.prompt_len = request_func_input.prompt_len output.start_time = time.perf_counter() speech_extra = getattr(request_func_input, "seed_tts_speech_extra", None) if not isinstance(speech_extra, dict): speech_extra = {} + else: + speech_extra = dict(speech_extra) + speech_extra["return_stage_metrics"] = True configured_turns = getattr(request_func_input, "seed_tts_turns", ()) turn_prompts = [ ( @@ -1387,8 +1756,10 @@ async def async_request_openai_realtime_duplex( ) await client.send({"type": "response.create"}) await wait_for( - lambda: client.events.count("response.done") > done_before - or len(client.events.errors()) > errors_before, + lambda: ( + client.events.count("response.done") > done_before + or len(client.events.errors()) > errors_before + ), timeout_s=180.0, label=f"Seed-TTS Realtime TTS turn {request_index} response.done", ) @@ -1433,19 +1804,7 @@ async def async_request_openai_realtime_duplex( channels=1, ) ) - turn_transcripts.append( - "".join( - str(event.get("delta") or "") - for event in client.events.events - if client.events.response_id(event) == response_id - and event.get("type") - in { - "response.audio_transcript.delta", - "response.output_text.delta", - "response.text.delta", - } - ) - ) + turn_transcripts.append(client.events.response_text(response_id)) await client.acknowledge_playback() request_finished_at = time.perf_counter() session_metrics = summarize_session_request_metrics( @@ -1473,6 +1832,16 @@ async def async_request_openai_realtime_duplex( for timing in turn_timings if isinstance((stage0 := timing.get("stage0_tokens")), dict) ) + token_timing_measured = _apply_stage0_token_timings( + output, + [timing.get("stage0_tokens") for timing in turn_timings], + expected_output_tokens=output.output_tokens, + ) + if not token_timing_measured and output.output_tokens > 1: + logger.warning( + "Realtime TTS session %s omitted complete engine token timing; standard TPOT/ITL are unavailable", + session_id, + ) output.success = True except Exception: output.success = False @@ -1520,6 +1889,7 @@ from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, ge from vllm_omni.benchmarks.metrics.metrics import ( MultiModalsBenchmarkMetrics, calculate_metrics, + has_metric_samples, ) # ruff: noqa: E402 @@ -1626,6 +1996,7 @@ async def benchmark( ) _attach_daily_omni_to_request_func_input(input_requests[0], test_input) _attach_seed_tts_to_request_func_input(input_requests[0], test_input) + _attach_omniinteract_to_request_func_input(input_requests[0], test_input) if ready_check_timeout_sec > 0: test_output = await wait_for_endpoint( @@ -1741,6 +2112,7 @@ async def benchmark( print(f"Probe request rate: {probe_request_rate} req/s") probe_task = asyncio.create_task(probe_loop()) + _prepare_omniinteract_batch(input_requests) benchmark_start_time = time.perf_counter() tasks: list[asyncio.Task] = [] @@ -1801,6 +2173,7 @@ async def benchmark( ) _attach_daily_omni_to_request_func_input(request, request_func_input) _attach_seed_tts_to_request_func_input(request, request_func_input) + _attach_omniinteract_to_request_func_input(request, request_func_input) tasks.append( asyncio.create_task(limited_request_func(request_func_input=request_func_input, session=session, pbar=pbar)) ) @@ -1815,6 +2188,8 @@ async def benchmark( benchmark_duration = time.perf_counter() - benchmark_start_time + omniinteract_summary = _finalize_omniinteract_batch(input_requests, outputs) + if task_type == TaskType.GENERATION: metrics, actual_output_lens = calculate_metrics( input_requests=input_requests, @@ -1839,6 +2214,13 @@ async def benchmark( actual_output_lens = 0 if isinstance(metrics, MultiModalsBenchmarkMetrics): + + def measured_ttft(output: RequestFuncOutput) -> float | None: + session_metrics = getattr(output, "duplex_session_metrics", None) + if isinstance(session_metrics, dict) and session_metrics.get("mean_ttft_ms") is None: + return None + return output.ttft + result = { "duration": benchmark_duration, "completed": metrics.completed, @@ -1859,7 +2241,7 @@ async def benchmark( "input_lens": [output.prompt_len for output in outputs], "start_times": [output.start_time for output in outputs], "output_lens": actual_output_lens, - "ttfts": [output.ttft for output in outputs], + "ttfts": [measured_ttft(output) for output in outputs], "itls": [output.itl for output in outputs], "generated_texts": [output.generated_text for output in outputs], "errors": [output.error for output in outputs], @@ -1893,6 +2275,8 @@ async def benchmark( ] if duplex_session_metrics: result["duplex_session_metrics"] = duplex_session_metrics + if omniinteract_summary is not None: + result["omniinteract"] = omniinteract_summary from vllm_omni.benchmarks.data_modules.daily_omni_eval import ( compute_daily_omni_accuracy_metrics, @@ -1949,6 +2333,8 @@ async def benchmark( # metric. if metric_attribute_name not in result_percentile_metrics: return + if not has_metric_samples(metrics, metric_attribute_name): + return # No text tokens generated (e.g. pure TTS speech endpoint): per-token # latency metrics (ttft/tpot/itl) are undefined, so skip them. is_text_token_metric = not (metric_attribute_name == "e2el" or metric_attribute_name.startswith("audio")) diff --git a/vllm_omni/benchmarks/serve.py b/vllm_omni/benchmarks/serve.py index 8e3b19ea9..2c4f8d130 100644 --- a/vllm_omni/benchmarks/serve.py +++ b/vllm_omni/benchmarks/serve.py @@ -1,6 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + import argparse import asyncio +import contextlib import os +from pathlib import Path from typing import Any from vllm.benchmarks.serve import main_async @@ -9,6 +16,7 @@ from vllm.benchmarks.serve import main_async # This monkey-patches vllm.benchmarks.datasets.get_samples before it's used # Must be imported before any vllm.benchmarks module usage import vllm_omni.benchmarks.patch.patch # noqa: F401 +from vllm_omni.benchmarks.omniinteract import omniinteract_output_lock from vllm_omni.benchmarks.patch.patch import ( maybe_enable_stage_metrics, set_print_stage, @@ -28,4 +36,10 @@ def main(args: argparse.Namespace) -> dict[str, Any]: getattr(args, "extra_body", None), enabled=should_request_stage_metrics(args), ) - return asyncio.run(main_async(args)) + lock = ( + omniinteract_output_lock(Path(args.omniinteract_output_dir)) + if getattr(args, "dataset_name", None) == "omniinteract" + else contextlib.nullcontext() + ) + with lock: + return asyncio.run(main_async(args)) diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index 3069025be..f53a1ef74 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -2,7 +2,6 @@ Configuration module for vLLM-Omni. """ -from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.lora import LoRAConfig from vllm_omni.config.model import OmniModelConfig from vllm_omni.config.omni_config import ( @@ -22,7 +21,6 @@ from vllm_omni.config.omni_config import ( VllmOmniGenerationStageConfig, VllmOmniOrchestratorConfig, ) -from vllm_omni.config.pipeline_registry import register_pipeline from vllm_omni.config.stage_config import ( PIPELINE_WIDE_ENGINE_FIELDS, DeployConfig, @@ -42,6 +40,15 @@ from vllm_omni.config.yaml_util import ( to_dict, ) +# StageConfigFactory / register_pipeline pull pipeline_registry, which eagerly +# imports PI0_PIPELINE → diffusion.data. Keep those lazy so +# `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still +# loading) cannot close a circular import through DiffusionOutput. +_LAZY_ATTRS = { + "StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"), + "register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"), +} + __all__ = [ # Legacy model-level configs. "LoRAConfig", @@ -82,3 +89,19 @@ __all__ = [ "merge_configs", "to_dict", ] + + +def __getattr__(name: str): + spec = _LAZY_ATTRS.get(name) + if spec is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module_name, attr_name = spec + import importlib + + value = getattr(importlib.import_module(module_name), attr_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(__all__) | set(globals())) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index a1778eda3..b2eb64bca 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Config factories for vllm-omni, e.g., StageConfigFactory.""" from __future__ import annotations @@ -15,6 +15,7 @@ from transformers import PretrainedConfig from vllm.logger import init_logger from vllm.transformers_utils.config import get_config from vllm.transformers_utils.repo_utils import get_hf_file_to_dict +from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri from vllm_omni.config.endpoint_policy import EndpointRestriction from vllm_omni.config.omni_config import VllmOmniConfig @@ -28,8 +29,10 @@ from vllm_omni.config.stage_config import ( build_stage_runtime_overrides, load_deploy_config, merge_pipeline_deploy, + normalize_pipeline_cli_overrides, ) from vllm_omni.config.yaml_util import create_config +from vllm_omni.diffusion.io_support import get_diffusion_output_type from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero logger = init_logger(__name__) @@ -45,6 +48,39 @@ logger = init_logger(__name__) _DEFAULT_PARALLEL_DEGREE = 1 +@functools.cache +def _materialize_object_storage_configs(model: str) -> str: + """Materialize an object-storage model URI's config files locally. + + vLLM's Run:AI streamer keeps ``s3://``/``gs://``/``az://`` URIs opaque until + each stage builds its ``ModelConfig``; parent-process resolution (HF config + lookup, pipeline/pipeline-key matching) would instead hand the URI to + ``huggingface_hub`` helpers, which reject it with ``HFValidationError``. + Pull the lightweight files once into vLLM's deterministic + ``model_streamer/`` directory so config reads work here, and so the + stage processes' own pull lands in that same directory. + + Returns the input unchanged for non object-storage paths. + """ + if not is_runai_obj_uri(model): + return model + object_storage_model = ObjectStorageModel(url=model) + object_storage_model.pull_files(model, allow_pattern=["*.model", "*.py", "*.json"]) + logger.info("Materialized object-storage configs for %s at %s", model, object_storage_model.dir) + return object_storage_model.dir + + +def _name_match_candidate(model: str) -> str: + """Last path component of a model reference, used for name-based matching. + + Object-storage URIs and HF repo ids carry non-model segments (bucket name, + organization) that must not participate in substring matching; e.g. a + bucket named ``qwen3-tts-models`` holding a ``Qwen3-Omni`` checkpoint must + not resolve to the ``qwen3_tts`` pipeline. + """ + return model.rstrip("/").rsplit("/", 1)[-1] + + def with_trust_remote_code_override( overrides: Mapping[str, Any], trust_remote_code: bool | None, @@ -117,7 +153,7 @@ class StageConfigFactory: """ hf_config = None try: - return get_config(model, trust_remote_code=trust_remote_code) + return get_config(_materialize_object_storage_configs(model), trust_remote_code=trust_remote_code) except Exception as e: logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None") return hf_config @@ -163,10 +199,12 @@ class StageConfigFactory: if hf_config is not None: return hf_config.model_type + config_source = _materialize_object_storage_configs(model) + # Fallback: read config.json directly for custom model types that # are not registered with transformers (e.g. qwen3_tts). try: - config_dict = get_hf_file_to_dict("config.json", model, revision=None) + config_dict = get_hf_file_to_dict("config.json", config_source, revision=None) if config_dict: if "model_type" in config_dict: return config_dict["model_type"] @@ -183,14 +221,17 @@ class StageConfigFactory: # model_index.json with _class_name that maps to a pipeline key via # PipelineConfig.diffusers_class_name. try: - model_index = get_hf_file_to_dict("model_index.json", model, revision=None) + model_index = get_hf_file_to_dict("model_index.json", config_source, revision=None) if model_index and "_class_name" in model_index: class_name = model_index["_class_name"] for obj in OMNI_PIPELINES.values(): # If we have a resolver, call it with the optional hf_config # to get the default pipeline config for this key pipeline_cfg = obj(hf_config) if callable(obj) else obj - if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name: + if pipeline_cfg is not None and class_name in ( + pipeline_cfg.diffusers_class_name, + *pipeline_cfg.diffusers_class_aliases, + ): logger.info( "Detected pipeline %r from model_index.json (_class_name=%r)", pipeline_cfg.model_type, @@ -203,8 +244,10 @@ class StageConfigFactory: # Final fallback: some models (e.g. CosyVoice3) ship an empty # config.json and rely on naming conventions. Match the model path # basename against registered pipeline keys — longest match wins - # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). - model_lower = model.lower().replace("-", "").replace("_", "") + # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). Only + # the basename is scanned so URI segments such as the bucket name + # cannot select an unrelated pipeline. + model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "") best: str | None = None best_len = 0 for registered_key in OMNI_PIPELINES.keys(): @@ -395,6 +438,8 @@ class StageConfigFactory: load-balancer policy (``None`` when no strategy set one) travels with the stages instead of through a mutable out-param. """ + cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides) + deploy_cfg: DeployConfig | None if user_deploy_config is not None: deploy_cfg = user_deploy_config elif deploy_config_path is not None: @@ -402,9 +447,12 @@ class StageConfigFactory: assert deploy_cfg is not None elif pipeline_cfg.default_deploy_config_name is not None: deploy_cfg = load_deploy_config(_DEPLOY_DIR / pipeline_cfg.default_deploy_config_name) + assert deploy_cfg is not None else: deploy_cfg = DeployConfig() + assert deploy_cfg is not None + cli_async_chunk = cli_overrides.get("async_chunk") if cli_async_chunk is not None: deploy_cfg.async_chunk = bool(cli_async_chunk) @@ -583,6 +631,7 @@ class StageConfigFactory: engine_args["dtype"] = str(engine_args["dtype"]) engine_args.setdefault("max_num_seqs", 1) + model_class_name = engine_args.get("model_class_name") config_dict: dict[str, Any] = { "stage_id": 0, @@ -593,7 +642,7 @@ class StageConfigFactory: }, "engine_args": create_config(engine_args), "final_output": True, - "final_output_type": "image", + "final_output_type": get_diffusion_output_type(model_class_name), } return [config_dict] diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py index d22c0acc1..af53b8691 100644 --- a/vllm_omni/config/endpoint_policy.py +++ b/vllm_omni/config/endpoint_policy.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Endpoint restriction policy for omni pipelines.""" from dataclasses import dataclass @@ -8,7 +8,14 @@ from typing import NamedTuple from fastapi import FastAPI, Request from fastapi.responses import JSONResponse -from vllm.entrypoints.serve.utils.error_response import create_error_response + +# vLLM < 0.28 keeps create_error_response under serve.utils (it is not +# re-exported from vllm.entrypoints.serve); 0.28+ moved it under +# serve.exception_handling and re-exports it from the package root. +try: + from vllm.entrypoints.serve import create_error_response +except ImportError: + from vllm.entrypoints.serve.utils.error_response import create_error_response class RouteTarget(NamedTuple): diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py new file mode 100644 index 000000000..2a7385c1d --- /dev/null +++ b/vllm_omni/config/environment_variable_inventory.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reviewed inventory of environment variables used by vLLM-Omni. + +This module classifies environment-variable names; it does not read them or +make every classified name part of the public configuration contract. In +particular, model-specific variables remain transitional until their recorded +disposition is implemented. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from types import MappingProxyType + + +class EnvironmentVariableCategory(str, Enum): + """Ownership and documentation boundary for an environment variable.""" + + PUBLIC_OMNI = "public_omni" + INHERITED_VLLM = "inherited_vllm" + PLATFORM_EXTERNAL = "platform_external" + MODEL_SPECIFIC = "model_specific" + BENCHMARK_TRANSITIONAL = "benchmark_transitional" + INTERNAL = "internal" + + +class ModelEnvironmentVariableDisposition(str, Enum): + """Required destination for a model-specific environment variable.""" + + PROMOTE = "promote" + REQUEST_SCOPE = "request_scope" + EXTERNAL = "external" + INTERNALIZE = "internalize" + DEPRECATE_REMOVE = "deprecate_remove" + + +@dataclass(frozen=True) +class EnvironmentVariableClassification: + """Classification metadata for one statically identifiable name.""" + + name: str + category: EnvironmentVariableCategory + model_disposition: ModelEnvironmentVariableDisposition | None = None + redact_value: bool = False + + @property + def is_public_omni(self) -> bool: + return self.category is EnvironmentVariableCategory.PUBLIC_OMNI + + +# New entries must use the VLLM_OMNI_ prefix. The regression test explicitly +# grandfathers the older public names below. +_PUBLIC_OMNI = ( + "DIFFUSION_ATTENTION_BACKEND", + "DIFFUSION_CACHE_ADAPTER", + "DIFFUSION_CACHE_BACKEND", + "OMNI_DIFFUSION_PROMPT_EMBED_CACHE", + "OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", + "OMNI_DIFFUSION_SESSION_STATE_MANAGER", + "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS", + "SPEAKER_MAX_UPLOADED", + "SPEAKER_SAMPLES_DIR", + "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", + "VLLM_OMNI_EVENT_DRIVEN_ORCH", + "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", + "VLLM_OMNI_ORCH_MONITOR_PATH", + "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", + "VLLM_OMNI_SERVER_STORAGE__FILE_TTL", + "VLLM_OMNI_SERVER_STORAGE__PATH", + "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL", + "VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP", + "VLLM_OMNI_STORAGE_MAX_CONCURRENCY", + "VLLM_OMNI_STORAGE_PATH", + "VLLM_OMNI_USE_QUACK_FP8", + "VLLM_OMNI_VIDEO_SYNC_TIMEOUT", + "VLLM_VIDEO_ASYNC_CHUNK", + "VLLM_VIDEO_AUDIO_DELTA_MODE", +) + +_INHERITED_VLLM = ( + "CUDA_HOME", + "CUDA_VISIBLE_DEVICES", + "LOCAL_RANK", + "NO_COLOR", + "VLLM_ALLOW_LONG_MAX_MODEL_LEN", + "VLLM_BATCH_INVARIANT", + "VLLM_CACHE_ROOT", + "VLLM_DISABLE_LOG_LOGO", + "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", + "VLLM_HTTP_TIMEOUT_KEEP_ALIVE", + "VLLM_LOGGING_COLOR", + "VLLM_MOONCAKE_BOOTSTRAP_PORT", + "VLLM_PLUGINS", + "VLLM_ROCM_USE_AITER", + "VLLM_ROCM_USE_AITER_RMSNORM", + "VLLM_TUNED_CONFIG_FOLDER", + "VLLM_USE_MODELSCOPE", + "VLLM_USE_OINK_OPS", + "VLLM_WORKER_MULTIPROC_METHOD", + "VLLM_XPU_USE_SAMPLER_KERNEL", +) + +_PLATFORM_EXTERNAL = ( + "ASCEND_LAUNCH_BLOCKING", + "ASCEND_RT_VISIBLE_DEVICES", + "FLASHINFER_DISABLE_VERSION_CHECK", + "HF_HOME", + "HF_HUB_DISABLE_XET", + "HF_TOKEN", + "HIP_VISIBLE_DEVICES", + "HUGGINGFACE_HUB_TOKEN", + "MASTER_ADDR", + "MASTER_PORT", + "MINDIE_SD_FA_TYPE", + "MORI_RDMA_DEVICES", + "MUSA_VISIBLE_DEVICES", + "MV2_COMM_WORLD_LOCAL_RANK", + "NCCL_ASYNC_ERROR_HANDLING", + "OMPI_COMM_WORLD_LOCAL_RANK", + "OPENAI_API_KEY", + "PYTHONPATH", + "QUACK_CACHE_DIR", + "RANK", + "RANK_ID", + "RAY_RAYLET_PID", + "RDMA_DEVICE_NAME", + "SAGE_ATTN_TENSOR_LAYOUT", + "VLLM_LOCAL_RANK", + "WORLD_SIZE", + "XDG_CACHE_HOME", +) + +_MODEL_PROMOTE = ( + "AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", + "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH", + "COSMOS3_SOUND_TOKENIZER_PATH", + "COSYVOICE3_ESTIMATOR_ONNX", + "COSYVOICE3_TRT", + "COSYVOICE3_TRT_CACHE", + "DYNIN_CONFIG_PATH", + "DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY", + "DYNIN_MAGVIT_REMOTE_CODE_REPO_ID", + "DYNIN_MAGVIT_REMOTE_CODE_REVISION", + "DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY", + "DYNIN_REMOTE_CODE_REPO_ID", + "DYNIN_REMOTE_CODE_REVISION", + "HIGGS_AUDIO_TOKENIZER_PATH", + "INTERNVLA_A1_COSMOS_DECODER_PATH", + "INTERNVLA_A1_COSMOS_DIR", + "INTERNVLA_A1_COSMOS_ENCODER_PATH", + "INTERNVLA_A1_PROCESSOR_DIR", + "MIMO_AUDIO_TOKENIZER_CUDA_GRAPH", + "MIMO_AUDIO_TOKENIZER_DEVICE", + "MIMO_AUDIO_TOKENIZER_PATH", + "MING_CFM_CUDAGRAPH", + "MINICPMO_TOKEN2WAV_TRT", + "MINICPMO_TOKEN2WAV_TRT_CACHE", + "MINICPMO_TOKEN2WAV_TRT_DTYPE", + "NEMOTRON_VOICECHAT_LLM_PATH", + "OMNIVOICE_CUDA_GRAPH", + "OMNIVOICE_TF32", + "VLLM_GEPARD_CHUNK_FRAMES", + "VLLM_GEPARD_FIRST_CHUNK_FRAMES", + "VLLM_GEPARD_LOOKBACK_FRAMES", + "VLLM_OMNI_FISH_KVCACHE_ATTN", + "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER", +) + +_MODEL_REQUEST_SCOPE = ( + "DOTS_TTS_DIT_NUM_STEPS", + "GR00T_NOISE_SEED", + "STEP_AUDIO2_DEFAULT_PROMPT_WAV", + "VLLM_GEPARD_END_FADE_MS", + "VLLM_GEPARD_END_SILENCE_MS", + "VLLM_GEPARD_GREEDY", +) + +# No audited model variable is currently owned by a supported third-party +# contract. Keep the disposition explicit so future reviews cannot silently +# fold an externally owned name into "promote". +_MODEL_EXTERNAL: tuple[str, ...] = () + +_MODEL_INTERNALIZE = ( + "DOTS_TTS_BETA_TRACE", + "DYNIN_S2U_VENDOR_ROOT", + "DZ_PHASE_TIMING", + "MING_TTS_STAGE1_FINAL_LOG", + "MOSS_TTS_DEBUG_STOP", + "NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", + "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS", + "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS_LOG_EVERY", + "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS", + "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_FILE", + "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_LOG_EVERY", +) + +_MODEL_DEPRECATE_REMOVE = ( + "HIGGS_AUDIO_V2_TOKENIZER_PATH", + "VLLM_OMNI_LINGBOT_ACTION_ROOT", + "VLLM_OMNI_VOXCPM_CODE_PATH", + "VLLM_USE_FLASHINFER_MOE_FP16", + "XCODEC1_PATH", + "model_stage", +) + +_BENCHMARK_TRANSITIONAL = ( + "DAILY_OMNI_EXTRACT_MODE", + "DAILY_OMNI_SAVE_EVAL_ITEMS", + "MAX_NUM_FRAMES", + "SEED_TTS_EVAL_DEVICE", + "SEED_TTS_HF_WHISPER_MODEL", + "SEED_TTS_SIM_DEVICE", + "SEED_TTS_SIM_EVAL", + "SEED_TTS_UTMOS_EVAL", + "SEED_TTS_UTMOS_HF_REPO", + "SEED_TTS_UTMOS_JIT_FILE", + "SEED_TTS_WAVLM_MAX_SECONDS", + "SEED_TTS_WAVLM_MIN_SAMPLES", + "SEED_TTS_WAVLM_MODEL", + "SEED_TTS_WER_EVAL", + "SEED_TTS_WER_SAVE_AUDIO_DIR", + "SEED_TTS_WER_SAVE_ITEMS", + "VLLM_DAILY_OMNI_MEDIA_REPO", + "VLLM_OMNI_BENCH_AUDIO_CHANNELS", + "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S", + "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE", +) + +_INTERNAL = ( + "VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", + "VLLM_OMNI_REPLICA_ID", +) + +_REDACTED = frozenset( + { + "HF_TOKEN", + "HUGGINGFACE_HUB_TOKEN", + "OPENAI_API_KEY", + } +) + + +def _build_inventory() -> Mapping[str, EnvironmentVariableClassification]: + inventory: dict[str, EnvironmentVariableClassification] = {} + + def add( + names: tuple[str, ...], + category: EnvironmentVariableCategory, + disposition: ModelEnvironmentVariableDisposition | None = None, + ) -> None: + for name in names: + if name in inventory: + raise ValueError(f"Environment variable {name!r} is classified more than once") + inventory[name] = EnvironmentVariableClassification( + name=name, + category=category, + model_disposition=disposition, + redact_value=name in _REDACTED, + ) + + add(_PUBLIC_OMNI, EnvironmentVariableCategory.PUBLIC_OMNI) + add(_INHERITED_VLLM, EnvironmentVariableCategory.INHERITED_VLLM) + add(_PLATFORM_EXTERNAL, EnvironmentVariableCategory.PLATFORM_EXTERNAL) + add(_MODEL_PROMOTE, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.PROMOTE) + add( + _MODEL_REQUEST_SCOPE, + EnvironmentVariableCategory.MODEL_SPECIFIC, + ModelEnvironmentVariableDisposition.REQUEST_SCOPE, + ) + add(_MODEL_EXTERNAL, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.EXTERNAL) + add( + _MODEL_INTERNALIZE, + EnvironmentVariableCategory.MODEL_SPECIFIC, + ModelEnvironmentVariableDisposition.INTERNALIZE, + ) + add( + _MODEL_DEPRECATE_REMOVE, + EnvironmentVariableCategory.MODEL_SPECIFIC, + ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE, + ) + add(_BENCHMARK_TRANSITIONAL, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL) + add(_INTERNAL, EnvironmentVariableCategory.INTERNAL) + return MappingProxyType(dict(sorted(inventory.items()))) + + +# All reviewed environment-variable names, indexed by exact name. This is +# classification metadata, not the executable value registry used by +# ``vllm.envs``. +ENVIRONMENT_VARIABLE_INVENTORY: Mapping[str, EnvironmentVariableClassification] = _build_inventory() diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 9164e1352..7be569f47 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -93,6 +93,7 @@ class OmniModelConfig(ModelConfig): hf_text_config: The sub text_config of the model's hf_config (default: None) stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) async_chunk: If set to True, perform async chunk + session_mode: Request lifecycle mode, either turn-based or duplex model_stage: Stage type identifier, e.g., "thinker" or "talker" (default: "thinker") model_arch: Model architecture name @@ -121,6 +122,7 @@ class OmniModelConfig(ModelConfig): stage_id: int = 0 async_chunk: bool = False + session_mode: str = "turn" retains_state_across_chunks: bool = False # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin. active_stream_window: int = 0 @@ -141,11 +143,18 @@ class OmniModelConfig(ModelConfig): } ) subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 omni_kv_config: dict | None = None codec_frame_rate_hz: float | None = None task_type: str | None = None enable_sleep_mode: bool = False has_sampling_extra_args: bool = False + # Key names (not values) of the stage's default sampling ``extra_args``. + # Engine-core code runs before any request arrives, so this is the only + # place it can learn which request-shaping conventions a stage uses (e.g. + # ``cfg_role`` for classifier-free-guidance request pairs). + sampling_extra_args_keys: tuple[str, ...] = () + requires_full_payload_input: bool = False @property def registry(self): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 4b5c73a39..58dc55fb6 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Structured vLLM-Omni configuration classes. This module is additive for Phase 2 of RFC #4021. @@ -12,12 +12,22 @@ from __future__ import annotations import copy from collections.abc import Mapping -from dataclasses import dataclass, field, fields +from dataclasses import InitVar, dataclass, field, fields +from functools import wraps +from inspect import Parameter, signature from pathlib import Path from typing import Any, Literal, TypeAlias, TypedDict, cast -from pydantic import ConfigDict, Field +from pydantic import ConfigDict, Field, model_validator +from typing_extensions import Self +from vllm.config import CacheConfig as VllmCacheConfig +from vllm.config import CompilationConfig as VllmCompilationConfig +from vllm.config import LoadConfig as VllmLoadConfig +from vllm.config import ParallelConfig as VllmParallelConfig +from vllm.config import ProfilerConfig as VllmProfilerConfig +from vllm.config import SchedulerConfig as VllmSchedulerConfig from vllm.config.utils import config +from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.config.stage_config import ( @@ -36,6 +46,7 @@ from vllm_omni.config.stage_config import ( _select_processor_funcs, build_stage_runtime_overrides, load_deploy_config, + normalize_pipeline_cli_overrides, ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -53,7 +64,6 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset( "model", "omni", "output_modalities", - "stage_configs_path", "stage_id", "tokenizer", } @@ -78,6 +88,49 @@ class _QuantizationEngineOverrides(TypedDict, total=False): quantization: str +class _TrackExplicitConfigFields: + """Record constructor inputs without adding a serialized config field.""" + + @model_validator(mode="wrap") + @classmethod + def _record_explicit_config_fields(cls, value: Any, handler: Any) -> Any: + result = handler(value) + kwargs = getattr(value, "kwargs", None) + if kwargs is not None: + explicit_fields = frozenset(kwargs) + elif isinstance(value, cls): + explicit_fields = getattr(value, "_omni_explicit_fields", frozenset()) + else: + explicit_fields = frozenset() + object.__setattr__(result, "_omni_explicit_fields", explicit_fields) + return result + + +def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]: + """Make inherited Pydantic dataclass fields keyword-only as well.""" + generated_init = cls.__init__ + generated_signature = signature(cls) + keyword_only_signature = generated_signature.replace( + parameters=[ + parameter.replace(kind=Parameter.KEYWORD_ONLY) + if parameter.kind in {Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD} + else parameter + for parameter in generated_signature.parameters.values() + ] + ) + + @wraps(generated_init) + def keyword_only_init(self: Any, *args: Any, **kwargs: Any) -> None: + if args: + raise TypeError(f"{cls.__name__}() accepts keyword arguments only; got {len(args)} positional argument(s)") + generated_init(self, **kwargs) + + cls.__init__ = keyword_only_init + cls.__signature__ = keyword_only_signature + setattr(cls, "__match_args__", ()) + return cls + + class _ModelEngineOverrides(TypedDict, total=False): model: str model_arch: str @@ -91,6 +144,7 @@ class _ModelEngineOverrides(TypedDict, total=False): trust_remote_code: bool dtype: Any attention_backend: Any + attention_config: Any moe_backend: str hf_overrides: Any limit_mm_per_prompt: dict[str, Any] @@ -99,6 +153,7 @@ class _ModelEngineOverrides(TypedDict, total=False): active_stream_window: int enable_sleep_mode: bool subtalker_sampling_params: dict[str, Any] + silence_ban_frames: int has_sampling_extra_args: bool custom_voice_dir: str task_type: str @@ -106,7 +161,6 @@ class _ModelEngineOverrides(TypedDict, total=False): enforce_eager: bool max_cudagraph_capture_size: int enable_flashinfer_autotune: bool - compilation_config: dict[str, Any] enable_multithread_weight_load: bool num_weight_load_threads: int disable_autocast: bool @@ -148,7 +202,6 @@ class _RuntimeEngineOverrides(TypedDict, total=False): num_gpus: int log_level: str log_stats: bool - profiler_config: dict[str, Any] class _ParallelConfigEngineOverrides(TypedDict, total=False): @@ -160,6 +213,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False): ring_degree: int allgather_degree: int ulysses_mode: str + ulysses_a2a_permute: bool cfg_parallel_size: int vae_patch_parallel_size: int vae_parallel_mode: str @@ -192,6 +246,8 @@ class _StageEngineValues: runtime: _RuntimeEngineOverrides parallel: _ParallelEngineOverrides diffusion: _DiffusionEngineOverrides + compilation_config: Mapping[str, Any] | VllmCompilationConfig | None + profiler_config: Mapping[str, Any] | VllmProfilerConfig | None @dataclass(frozen=True) @@ -322,6 +378,7 @@ class OmniStageModelConfig: trust_remote_code: bool = False dtype: Any = "auto" attention_backend: Any = None + attention_config: Any = None moe_backend: str = "auto" hf_overrides: Any = None limit_mm_per_prompt: dict[str, Any] | None = None @@ -329,10 +386,12 @@ class OmniStageModelConfig: interleave_mm_strings: bool | None = None media_io_kwargs: dict[str, Any] | None = None active_stream_window: int = Field(default=0, ge=0) + session_mode: str = "turn" duplex_max_sessions: int = Field(default=1, ge=1) enable_sleep_mode: bool = False default_sampling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 has_sampling_extra_args: bool = False custom_voice_dir: str | None = None task_type: str | None = None @@ -340,7 +399,6 @@ class OmniStageModelConfig: enforce_eager: bool = False max_cudagraph_capture_size: int | None = Field(default=None, ge=0) enable_flashinfer_autotune: bool | None = None - compilation_config: dict[str, Any] | None = None enable_multithread_weight_load: bool = True num_weight_load_threads: int = Field(default=4, ge=1) disable_autocast: bool = False @@ -349,23 +407,24 @@ class OmniStageModelConfig: # StagePipelineConfig.model_subdir/tokenizer_subdir on the legacy path. model_subdir: str | None = None tokenizer_subdir: str | None = None + requires_full_payload_input: bool = False -@config -class OmniStageLoadConfig: - """Per-stage loading behavior and resolved tokenizer input.""" +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig): + """vLLM loading behavior plus Omni stage-specific tokenizer inputs.""" tokenizer: str | None = None - download_dir: str | None = None skip_tokenizer_init: bool = False - load_format: str = "auto" tokenizer_mode: str = "auto" config_format: str | None = None skip_mm_profiling: bool | None = None -@config -class OmniStageCacheConfig: +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageCacheConfig(_TrackExplicitConfigFields, VllmCacheConfig): """Per-stage engine cache and memory behavior. This is separate from ``_DiffusionConfigProjection.cache_config``, which configures @@ -382,17 +441,26 @@ class OmniStageCacheConfig: mamba_ssm_cache_dtype: str | None = None -@config -class OmniStageSchedulerConfig: +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig): """Per-stage request scheduling behavior.""" + # Upstream receives max_model_len only while materializing SchedulerConfig. + # Omni retains it as unresolved stage input until the owning engine process. max_num_seqs: int | None = Field(default=None, ge=1) max_num_batched_tokens: int | None = Field(default=None, ge=1) max_model_len: int | None = Field(default=None, ge=-1) + is_encoder_decoder: InitVar[bool] = False # type: ignore[assignment] enable_chunked_prefill: bool | None = None async_scheduling: bool | None = None - def __post_init__(self) -> None: + def __post_init__(self, is_encoder_decoder: bool = False) -> None: + # Upstream initializes these derived fields in its terminal post-init. + # Keep them serializable here without running model-dependent checks. + self.max_num_encoder_input_tokens = self.max_num_batched_tokens + self.encoder_cache_size = self.max_num_batched_tokens + if ( self.max_num_batched_tokens is not None and self.max_num_seqs is not None @@ -431,24 +499,73 @@ class OmniStageRuntimeConfig: num_gpus: int = Field(default=1, ge=1) log_level: str = "info" log_stats: bool = False - profiler_config: dict[str, Any] | None = None -@config -class OmniStageParallelConfig: +@_enforce_keyword_only_init +@config(kw_only=True) +class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig): """Common per-stage distributed parallelism behavior.""" - pipeline_parallel_size: int = Field(default=1, ge=1) - data_parallel_size: int = Field(default=1, ge=1) - tensor_parallel_size: int = Field(default=1, ge=1) - enable_expert_parallel: bool = False - world_size: int = Field(default=1, ge=1, init=False) + # EngineArgs intentionally leaves these unresolved. The upstream terminal + # ParallelConfig defaults (rank 0, local size 1, port 29550, worker + # ``"auto"``) must not turn into explicit engine inputs in the head + # process merely because this transport class inherits ParallelConfig. + data_parallel_size_local: int | None = Field(default=None, ge=0) + data_parallel_rank: int | None = Field(default=None, ge=0) + data_parallel_rpc_port: int | None = None + worker_cls: str | None = None + + @model_validator(mode="after") + def _validate_parallel_config(self) -> Self: + """Run upstream validation without resolving deferred engine inputs.""" + deferred_values = ( + self.data_parallel_size_local, + self.data_parallel_rank, + self.data_parallel_rpc_port, + self.worker_cls, + getattr(self, "all2all_backend", None), + getattr(self, "disable_custom_all_reduce", None), + ) + if self.data_parallel_size_local is not None and self.data_parallel_size_local > self.data_parallel_size: + raise ValueError( + f"data_parallel_size_local ({self.data_parallel_size_local}) " + f"must be <= data_parallel_size ({self.data_parallel_size})" + ) + if self.data_parallel_rank is not None and not 0 <= self.data_parallel_rank < self.data_parallel_size: + raise ValueError( + f"data_parallel_rank ({self.data_parallel_rank}) must be in the range [0, {self.data_parallel_size})" + ) + self.data_parallel_size_local = self.data_parallel_size + self.data_parallel_rank = 0 + self.data_parallel_rpc_port = VllmParallelConfig.data_parallel_rpc_port + self.worker_cls = VllmParallelConfig.worker_cls + try: + VllmParallelConfig._validate_parallel_config(self) + finally: + ( + self.data_parallel_size_local, + self.data_parallel_rank, + self.data_parallel_rpc_port, + self.worker_cls, + self.all2all_backend, + self.disable_custom_all_reduce, + ) = deferred_values + return self def __post_init__(self) -> None: + # Keep config construction transport-safe. Upstream runtime backend, + # rank, port, and platform resolution happens in the owning process. + self.data_parallel_index = self.data_parallel_rank self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size + @property + def world_size_across_dp(self) -> int: + # Omni's public world_size has historically included DP. + return self.world_size -@config + +@_enforce_keyword_only_init +@config(kw_only=True) class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): """Diffusion-stage distributed parallelism behavior.""" @@ -457,6 +574,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): ring_degree: int = Field(default=1, ge=1) allgather_degree: int = Field(default=1, ge=1) ulysses_mode: str = "strict" + ulysses_a2a_permute: bool = False cfg_parallel_size: int = Field(default=1, ge=1) vae_patch_parallel_size: int = Field(default=1, ge=1) text_encoder_tp_size: int = Field(default=1, ge=1) @@ -467,6 +585,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): hsdp_replicate_size: int = Field(default=1, ge=1) def __post_init__(self) -> None: + self.data_parallel_index = self.data_parallel_rank self.sequence_parallel_size = ( self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree ) @@ -488,12 +607,17 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): * self.cfg_parallel_size ) if self.use_hsdp: - if self.tensor_parallel_size > 1 or self.data_parallel_size > 1: - raise ValueError( - "HSDP (use_hsdp=True) cannot be used with TP or DP " - f"(tensor_parallel_size={self.tensor_parallel_size}, " - f"data_parallel_size={self.data_parallel_size})" - ) + incompatible = [] + if self.tensor_parallel_size > 1: + incompatible.append("TP") + if self.data_parallel_size > 1: + incompatible.append("DP") + if self.pipeline_parallel_size > 1: + incompatible.append("PP") + if self.enable_expert_parallel: + incompatible.append("EP") + if incompatible: + raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible)) if self.hsdp_shard_size == -1: if other_parallel_world_size == 1: raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1") @@ -544,7 +668,7 @@ class _DiffusionConfigProjection: dtype: Any = "auto" trust_remote_code: bool = False revision: str | None = None - distributed_executor_backend: str = "mp" + distributed_executor_backend: str | None = None dist_timeout: int | None = None nccl_port: int | None = None master_port: int | None = None @@ -559,6 +683,7 @@ class _DiffusionConfigProjection: cache_config: Any = field(default_factory=dict) enable_cache_dit_summary: bool = False diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY + diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True) enable_prompt_embed_cache: bool = False prompt_embed_cache_size: int = Field(default=32, ge=1) enable_session_state_manager: bool = False @@ -576,6 +701,9 @@ class _DiffusionConfigProjection: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = Field(default=0, ge=0) + host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled" + host_weight_runtime_root: str | None = None + dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0) pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) @@ -605,6 +733,7 @@ class _DiffusionConfigProjection: default_factory=lambda: { "transformer": True, "vae": True, + "text_encoder": True, } ) override_transformer_cls_name: str | None = None @@ -621,7 +750,7 @@ class _DiffusionConfigProjection: from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs) - valid_fields = {f.name for f in fields(cls)} + valid_fields = {f.name for f in fields(cast(Any, cls))} return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields}) def __post_init__(self) -> None: @@ -633,6 +762,8 @@ class _DiffusionConfigProjection: TransformerConfig, build_attention_config, parse_kv_cache_skip_selector, + validate_dlo_host_registration_options, + validate_host_weight_runtime_options, ) from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode from vllm_omni.quantization import build_quant_config @@ -697,6 +828,11 @@ class _DiffusionConfigProjection: ) self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) + if ( + self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + and self.diffusion_kv_max_rows_per_request is None + ): + raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set") self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps) self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers) @@ -705,6 +841,17 @@ class _DiffusionConfigProjection: elif self.max_cpu_loras < 1: raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + validate_host_weight_runtime_options( + mode=self.host_weight_runtime_mode, + root=self.host_weight_runtime_root, + ) + self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options( + limit_gib=self.dlo_host_registration_limit_gib, + enable_dlo=self.enable_distributed_layerwise_offload, + use_allgather=self.dlo_use_allgather, + hwr_mode=self.host_weight_runtime_mode, + ) + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): raise ValueError( "diffusers_load_kwargs and diffusers_call_kwargs are only " @@ -768,7 +915,7 @@ class _DiffusionConfigProjection: setattr(self, name, _copy_value(getattr(omni_diffusion_config, name))) -_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(_DiffusionConfigProjection)) +_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(cast(Any, _DiffusionConfigProjection))) # Current OmniDiffusionConfig still contains a flat mix of shared engine, # runtime, parallel, and diffusion-specific knobs. Keep this classification @@ -813,6 +960,10 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset( "omni_kv_config", "cfg_kv_collect_func", "max_num_seqs", + "kv_cache_memory_bytes", + "gpu_memory_utilization", + "max_num_batched_tokens", + "max_model_len", "enable_sleep_mode", "enforce_eager", "enable_multithread_weight_load", @@ -827,6 +978,7 @@ _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS) _DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset( { "diffusion_attention_backend", + "fastvideo_vsa_topk", "kv_cache_dtype", "kv_cache_skip_layers", "kv_cache_skip_steps", @@ -838,6 +990,58 @@ _DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMP "stage_id", } + +def _upstream_engine_field_map( + config_cls: type[Any], + *, + aliases: Mapping[str, str] = {}, + exclude: frozenset[str] = frozenset(), +) -> dict[str, str]: + """Map reusable upstream config fields to their EngineArgs inputs.""" + engine_fields = frozenset(config_field.name for config_field in fields(VllmEngineArgs)) + return { + config_field.name: engine_name + for config_field in fields(config_cls) + if config_field.init + and config_field.name not in exclude + and (engine_name := aliases.get(config_field.name, config_field.name)) in engine_fields + } + + +_LOAD_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(VllmLoadConfig) +_CACHE_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map( + VllmCacheConfig, + aliases={"cache_dtype": "kv_cache_dtype"}, +) +_SCHEDULER_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map( + VllmSchedulerConfig, + aliases={"policy": "scheduling_policy"}, + # ``scheduler_cls`` is selected by the immutable stage topology, while + # ``disable_hybrid_kv_cache_manager`` belongs to the cache concern in the + # Omni schema. Keep both out of the dynamic SchedulerConfig projection so + # one input cannot acquire two owners. + exclude=frozenset( + { + "scheduler_cls", + "disable_hybrid_kv_cache_manager", + } + ), +) +_PARALLEL_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map( + VllmParallelConfig, + aliases={"data_parallel_master_ip": "data_parallel_address"}, + # Runtime owns worker selection. The private API-process fields are + # terminal vLLM internals rather than per-stage user inputs. + exclude=frozenset( + { + "distributed_executor_backend", + "worker_cls", + "_api_process_count", + "_api_process_rank", + } + ), +) + _QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__) _MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__) _LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__) @@ -845,7 +1049,17 @@ _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__) _SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__) _CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__) _RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) -_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__) +_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"}) +_LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values()) +_LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values()) +_LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values()) +_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP.values()) +_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset( + f.name for f in fields(OmniStageDiffusionParallelConfig) +) & frozenset(_ParallelConfigEngineOverrides.__annotations__) +_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP = {name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS} +_LLM_PARALLEL_CONFIG_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP) +_PARALLEL_CONFIG_ENGINE_FIELDS = _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS _PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"} _COMMON_STAGE_ENGINE_FIELDS = ( _QUANTIZATION_ENGINE_FIELDS @@ -855,12 +1069,16 @@ _COMMON_STAGE_ENGINE_FIELDS = ( | _SCHEDULER_ENGINE_FIELDS | _CONNECTOR_ENGINE_FIELDS | _RUNTIME_ENGINE_FIELDS + | _DIRECT_VLLM_CONFIG_ENGINE_FIELDS +) +_LLM_STAGE_ENGINE_FIELDS = ( + _COMMON_STAGE_ENGINE_FIELDS + | _LLM_LOAD_ENGINE_FIELDS + | _LLM_CACHE_ENGINE_FIELDS + | _LLM_SCHEDULER_ENGINE_FIELDS + | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS + | {"parallel_config"} ) -_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageParallelConfig)) - {"world_size"} -_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageDiffusionParallelConfig)) - { - "world_size" -} -_LLM_STAGE_ENGINE_FIELDS = _COMMON_STAGE_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"} _DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = ( _COMMON_STAGE_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS @@ -881,6 +1099,27 @@ _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = { StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS, StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS, } +_PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = { + StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS, + StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS, + StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP), +} + +_LOAD_STAGE_ENGINE_FIELD_MAP = { + **{name: name for name in _LOAD_ENGINE_FIELDS}, + **_LOAD_CONFIG_ENGINE_FIELD_MAP, +} +_CACHE_STAGE_ENGINE_FIELD_MAP = { + **{name: name for name in _CACHE_ENGINE_FIELDS}, + **_CACHE_CONFIG_ENGINE_FIELD_MAP, +} +_SCHEDULER_STAGE_ENGINE_FIELD_MAP = { + **{name: name for name in _SCHEDULER_ENGINE_FIELDS}, + **_SCHEDULER_CONFIG_ENGINE_FIELD_MAP, +} +_DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP = {name: name for name in _LOAD_ENGINE_FIELDS} +_DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP = {name: name for name in _CACHE_ENGINE_FIELDS} +_DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {name: name for name in _SCHEDULER_ENGINE_FIELDS} def _validate_stage_engine_override_ownership( @@ -891,13 +1130,19 @@ def _validate_stage_engine_override_ownership( try: owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type] + parallel_config_fields = _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE[execution_type] except KeyError as exc: raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc unowned_fields = set(overrides) - owner_fields parallel_config = overrides.get("parallel_config") if isinstance(parallel_config, Mapping): - unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - parallel_owner_fields) + # Nested values traditionally use upstream config names, while flat + # EngineArgs uses aliases such as ``data_parallel_address``. Accept + # either spelling at the boundary; the builder canonicalizes aliases + # before instantiating the inherited config. + accepted_parallel_names = parallel_config_fields | parallel_owner_fields + unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - accepted_parallel_names) if unowned_fields: names = ", ".join(sorted(unowned_fields)) raise ValueError( @@ -918,7 +1163,7 @@ def _global_stage_cli_fields() -> frozenset[str]: ) externally_consumed = ( _NON_STAGE_ENGINE_CLI_FIELDS - | frozenset(f.name for f in fields(VllmOmniOrchestratorConfig)) + | frozenset(f.name for f in fields(cast(Any, VllmOmniOrchestratorConfig))) | (orchestrator_field_names() - _STAGE_ENGINE_FIELDS) ) return candidates - externally_consumed @@ -968,17 +1213,28 @@ def _stage_engine_values( topology.execution_type, engine, ) + if topology.execution_type in { + StageExecutionType.LLM_AR, + StageExecutionType.LLM_GENERATION, + }: + load_engine_fields = _LLM_LOAD_ENGINE_FIELDS + cache_engine_fields = _LLM_CACHE_ENGINE_FIELDS + scheduler_engine_fields = _LLM_SCHEDULER_ENGINE_FIELDS + else: + load_engine_fields = _LOAD_ENGINE_FIELDS + cache_engine_fields = _CACHE_ENGINE_FIELDS + scheduler_engine_fields = _SCHEDULER_ENGINE_FIELDS return _StageEngineValues( quantization=cast( _QuantizationEngineOverrides, _select_engine_overrides(engine, _QUANTIZATION_ENGINE_FIELDS), ), model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)), - load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, _LOAD_ENGINE_FIELDS)), - cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, _CACHE_ENGINE_FIELDS)), + load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)), + cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)), scheduler=cast( _SchedulerEngineOverrides, - _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS), + _select_engine_overrides(engine, scheduler_engine_fields), ), connector=cast( _ConnectorEngineOverrides, @@ -987,6 +1243,8 @@ def _stage_engine_values( runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)), parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)), diffusion=_DiffusionEngineOverrides.from_engine(engine), + compilation_config=_copy_value(engine.get("compilation_config")), + profiler_config=_copy_value(engine.get("profiler_config")), ) @@ -1003,7 +1261,7 @@ def _stage_sampling_params( def _orchestrator_cli_overrides(cli_overrides: Mapping[str, Any]) -> dict[str, Any]: overrides: dict[str, Any] = {} - for config_field in fields(VllmOmniOrchestratorConfig): + for config_field in fields(cast(Any, VllmOmniOrchestratorConfig)): name = config_field.name if name == "deploy_config_path": continue @@ -1041,6 +1299,8 @@ class BaseVllmOmniStageConfig: connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig) runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig) parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) + compilation_config: VllmCompilationConfig | None = None + profiler_config: VllmProfilerConfig | None = None quantization_config: _QuantizationConfigType = None @property @@ -1117,6 +1377,10 @@ class BaseVllmOmniStageConfig: def prompt_expand_func(self) -> str | None: return self.stage_pipeline_config.prompt_expand_func + @property + def prompt_transform_func(self) -> str | None: + return self.stage_pipeline_config.prompt_transform_func + @property def cfg_kv_collect_func(self) -> str | None: return self.stage_pipeline_config.cfg_kv_collect_func @@ -1192,6 +1456,8 @@ def _build_common_stage_config_kwargs( parallel_config, ), "parallel_config": parallel_config, + "compilation_config": _copy_value(engine.compilation_config), + "profiler_config": _copy_value(engine.profiler_config), "quantization_config": _copy_value(quantization_config), }, input_proc, @@ -1204,8 +1470,8 @@ def _with_resolved_processors( input_proc: str | None, next_stage_proc: str | None, ) -> StageConfigType: - stage_config._resolved_custom_process_input_func = input_proc - stage_config._resolved_custom_process_next_stage_input_func = next_stage_proc + setattr(stage_config, "_resolved_custom_process_input_func", input_proc) + setattr(stage_config, "_resolved_custom_process_next_stage_input_func", next_stage_proc) return stage_config @@ -1319,13 +1585,16 @@ def _build_stage_config( builder = _STAGE_CONFIG_BUILDERS[topology.execution_type] except KeyError as exc: raise ValueError(f"Unsupported stage execution type: {topology.execution_type!r}") from exc - return builder( - pipeline, - deploy, - topology, - stage_deploy, - engine, - model=model, + return cast( + StageConfigType, + builder( + pipeline, + deploy, + topology, + stage_deploy, + engine, + model=model, + ), ) @@ -1352,6 +1621,7 @@ def _build_model_config( ) -> OmniStageModelConfig: default_sampling_params = _stage_sampling_params(stage_deploy, topology) kwargs = _config_kwargs(engine) + kwargs["requires_full_payload_input"] = topology.requires_full_payload_input kwargs["model"] = _first_defined(kwargs.get("model"), model) if "model_arch" not in kwargs: kwargs["model_arch"] = topology.model_arch or pipeline.model_arch or None @@ -1369,8 +1639,9 @@ def _build_model_config( kwargs["model_subdir"] = topology.model_subdir if "tokenizer_subdir" not in kwargs and topology.tokenizer_subdir is not None: kwargs["tokenizer_subdir"] = topology.tokenizer_subdir - return OmniStageModelConfig( + return cast(Any, OmniStageModelConfig)( default_sampling_params=default_sampling_params, + session_mode=deploy.session_mode, duplex_max_sessions=duplex_max_sessions, **kwargs, ) @@ -1386,12 +1657,51 @@ def _build_load_config( return OmniStageLoadConfig(**kwargs) +def _config_kwargs_from_engine_args( + engine: Mapping[str, Any], + field_map: Mapping[str, str], +) -> dict[str, Any]: + """Convert EngineArgs names back to their upstream config field names.""" + return { + config_name: _copy_value(engine[engine_name]) + for config_name, engine_name in field_map.items() + if engine_name in engine and engine[engine_name] is not None + } + + +def _normalize_config_mapping( + values: Mapping[str, Any], + field_map: Mapping[str, str], +) -> dict[str, Any]: + """Canonicalize config and EngineArgs spellings to config field names.""" + engine_to_config = {engine_name: config_name for config_name, engine_name in field_map.items()} + normalized: dict[str, Any] = {} + source_names: dict[str, str] = {} + for name, value in values.items(): + if value is None: + continue + config_name = name if name in field_map else engine_to_config.get(name, name) + if config_name in normalized: + raise ValueError( + f"Config mapping specifies both {source_names[config_name]!r} and {name!r} " + f"for upstream field {config_name!r}" + ) + normalized[config_name] = _copy_value(value) + source_names[config_name] = name + return normalized + + +def _omni_config_kwargs(engine: Mapping[str, Any], fields_: frozenset[str]) -> dict[str, Any]: + return {name: _copy_value(engine[name]) for name in fields_ if name in engine and engine[name] is not None} + + def _build_cache_config( deploy: DeployConfig, engine: _CacheEngineOverrides, execution_type: StageExecutionType, ) -> OmniStageCacheConfig: - kwargs = _config_kwargs(engine) + kwargs = _config_kwargs_from_engine_args(engine, _CACHE_CONFIG_ENGINE_FIELD_MAP) + kwargs.update(_omni_config_kwargs(engine, _CACHE_ENGINE_FIELDS)) if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None: kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching) if "disable_hybrid_kv_cache_manager" not in kwargs and execution_type == StageExecutionType.LLM_GENERATION: @@ -1405,7 +1715,8 @@ def _build_scheduler_config( engine: _SchedulerEngineOverrides, execution_type: StageExecutionType, ) -> OmniStageSchedulerConfig: - kwargs = _config_kwargs(engine) + kwargs = _config_kwargs_from_engine_args(engine, _SCHEDULER_CONFIG_ENGINE_FIELD_MAP) + kwargs.update(_omni_config_kwargs(engine, _SCHEDULER_ENGINE_FIELDS)) if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None: kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill) if "async_scheduling" not in kwargs and execution_type == StageExecutionType.LLM_AR: @@ -1420,7 +1731,7 @@ def _build_connector_config( ) -> OmniStageConnectorConfig: output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None - return OmniStageConnectorConfig( + return cast(Any, OmniStageConnectorConfig)( async_chunk=bool(deploy.async_chunk), omni_kv_config=_copy_value(engine.get("omni_kv_config")), output_connectors=_copy_value(output_connectors) if output_connectors else None, @@ -1453,16 +1764,21 @@ def _build_parallel_config( config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig, ) -> OmniStageParallelConfig: parallel_config = _mapping_or_empty(engine.get("parallel_config")) - config_fields = { - config_field.name - for config_field in fields(config_cls) - if getattr(config_field.default, "init", None) is not False - } - kwargs = { - name: _copy_value(value) for name in _PARALLEL_CONFIG_ENGINE_FIELDS if (value := engine.get(name)) is not None - } + config_fields = frozenset(signature(config_cls).parameters) + is_diffusion = issubclass(config_cls, OmniStageDiffusionParallelConfig) + field_map = _DIFFUSION_PARALLEL_CONFIG_FIELD_MAP if is_diffusion else _PARALLEL_CONFIG_ENGINE_FIELD_MAP + owned_parallel_fields = frozenset(field_map) + kwargs = _config_kwargs(engine) if is_diffusion else _config_kwargs_from_engine_args(engine, field_map) + kwargs = {name: value for name, value in kwargs.items() if name in owned_parallel_fields} kwargs = {name: value for name, value in kwargs.items() if name in config_fields} - kwargs.update({name: value for name, value in _config_kwargs(parallel_config).items() if name in config_fields}) + nested_kwargs = _normalize_config_mapping(parallel_config, field_map) + kwargs.update( + { + name: value + for name, value in nested_kwargs.items() + if name in owned_parallel_fields and name in config_fields + } + ) if "pipeline_parallel_size" not in kwargs and deploy.pipeline_parallel_size is not None: kwargs["pipeline_parallel_size"] = _copy_value(deploy.pipeline_parallel_size) if "data_parallel_size" not in kwargs and deploy.data_parallel_size is not None: @@ -1528,6 +1844,7 @@ class VllmOmniConfig: """Create a structured config from a resolved pipeline and deploy YAML.""" if cli_overrides is None: cli_overrides = {} + cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides) deploy, loaded_deploy_config_path = _get_deploy_config( pipeline_cfg, @@ -1564,11 +1881,11 @@ class VllmOmniConfig: for topology in pipeline_cfg.stages ) - orchestrator_config = VllmOmniOrchestratorConfig( + orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)( deploy_config_path=loaded_deploy_config_path, **_orchestrator_cli_overrides(cli_overrides), ) - return cls( + return cast(Any, cls)( pipeline_config=pipeline_cfg, stage_configs=stage_configs, orchestrator_config=orchestrator_config, diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index f02a3578e..83f8de383 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Pipeline registry and factory for vllm-omni. ``OMNI_PIPELINES`` maps each ``model_type`` to either a ``PipelineConfig`` @@ -33,6 +33,7 @@ from vllm.logger import init_logger from vllm_omni.config.stage_config import ( PipelineConfig, ) +from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE from vllm_omni.model_executor.models.audex.pipeline import ( AUDEX_S2S_PIPELINE, AUDEX_THINKER_ONLY_PIPELINE, @@ -47,9 +48,11 @@ from vllm_omni.model_executor.models.bagel.pipeline import ( ) from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE +from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE +from vllm_omni.model_executor.models.gepard.pipeline import GEPARD_PIPELINE from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE from vllm_omni.model_executor.models.glm_tts.pipeline import GLM_TTS_PIPELINE from vllm_omni.model_executor.models.gr00t.pipeline import GR00T_N1D7_PIPELINE @@ -82,6 +85,8 @@ from vllm_omni.model_executor.models.ming_tts.pipeline import ( MING_TTS_PIPELINE, ) from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE +from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE +from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE from vllm_omni.model_executor.models.moss_tts.pipeline import ( MOSS_TTS_LOCAL_PIPELINE, MOSS_TTS_PIPELINE, @@ -97,12 +102,11 @@ from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import ( QWEN2_5_OMNI_PIPELINE, QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, ) -from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline -from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE -from vllm_omni.model_executor.models.soulx_singer.pipeline import ( - SOULXSINGER_SVC_PIPELINE, - SOULXSINGER_SVS_PIPELINE, +from vllm_omni.model_executor.models.qwen3_omni.pipeline import ( + QWEN3_OMNI_THINKER_ONLY_PIPELINE, + resolve_qwen3_omni_pipeline, ) +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE from vllm_omni.model_executor.models.step_audio2.pipeline import ( STEP_AUDIO2_ASR_PIPELINE, STEP_AUDIO2_PIPELINE, @@ -127,6 +131,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { # has no model_type key). "nemotron_labs_voicechat": NEMOTRON_VOICECHAT_PIPELINE, "qwen3_omni_moe": resolve_qwen3_omni_pipeline, + "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE, "qwen3_tts": QWEN3_TTS_PIPELINE, "step_audio_2": STEP_AUDIO2_PIPELINE, "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE, @@ -137,6 +142,8 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "lance": LANCE_PIPELINE, "dreamzero": DREAMZERO_PIPELINE, "Gr00tN1d7": GR00T_N1D7_PIPELINE, + "pi0": PI0_PIPELINE, + "gepard": GEPARD_PIPELINE, "glm_image": GLM_IMAGE_PIPELINE, "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE, "hunyuan_image3_ar": HUNYUAN_IMAGE3_AR_PIPELINE, @@ -144,6 +151,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE, "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, + "dots_tts": DOTS_TTS_PIPELINE, "cosyvoice3": COSYVOICE3_PIPELINE, "audex_tts": AUDEX_TTS_PIPELINE, "audex_tta": AUDEX_TTA_PIPELINE, @@ -164,6 +172,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE, "moss_tts_nano": MOSS_TTS_NANO_PIPELINE, + "minimax_h3_disaggregated": MINIMAX_H3_PIPELINE, "omnivoice": OMNIVOICE_PIPELINE, "mammoth_moda2": MAMMOTH_MODA2_PIPELINE, "mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE, @@ -171,13 +180,12 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE, "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE, "minicpmo_4_5": MINICPMO_4_5_PIPELINE, + "minimax_music3": MINIMAX_MUSIC3_PIPELINE, "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE, "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE, "dynin_omni": DYNIN_OMNI_PIPELINE, "indextts2": INDEXTTS2_PIPELINE, "indextts2_5": INDEXTTS25_PIPELINE, - "soulxsinger_svc": SOULXSINGER_SVC_PIPELINE, - "soulxsinger_svs": SOULXSINGER_SVS_PIPELINE, } @@ -188,18 +196,13 @@ def register_pipeline(pipeline: PipelineConfig | PipelineResolverFunc, model_typ since resolvers can return multiple different PipelineConfigs depending on the consumed config. """ - errors: list[str] = [] if isinstance(pipeline, PipelineConfig): - errors = pipeline.validate() model_type = model_type if model_type is not None else pipeline.model_type - else: - if model_type is None: - raise ValueError("Model type must be explicitly provided when registering a pipeline resolver") + elif model_type is None: + raise ValueError("Model type must be explicitly provided when registering a pipeline resolver") if model_type in OMNI_PIPELINES: - errors.append(f"Model type {model_type} is already registered; the old mapping will be clobbered") - if errors: - logger.warning("Registration for pipeline of type %s produced the following issues: %s", model_type, errors) + logger.warning(f"Model type {model_type} is already registered; the old mapping will be clobbered") OMNI_PIPELINES[model_type] = pipeline diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 496aea208..2c1dc670b 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage configuration system for vLLM-Omni.""" from __future__ import annotations -import dataclasses import functools import re import warnings @@ -56,7 +55,7 @@ def build_stage_runtime_overrides( ``internal_keys`` defaults to the union of ``arg_utils.internal_blacklist_keys()`` and ``arg_utils.SHARED_FIELDS`` so that neither orchestrator-only fields nor shared-pipeline fields - (``model`` / ``stage_configs_path`` / ``log_stats`` / ``stage_id``) leak + (``model`` / ``log_stats`` / ``stage_id``) leak into a stage's per-stage runtime overrides — the orchestrator sets those uniformly for every stage, they are not per-stage knobs. Callers can pass an explicit set for tests or specialized flows. @@ -90,44 +89,41 @@ def build_stage_runtime_overrides( return result -def strip_parent_engine_args( - kwargs: dict[str, Any], - *, - parent_fields: dict[str, dataclasses.Field], - keep_keys: set[str] | frozenset[str] = frozenset(), - strip_keys: set[str] | frozenset[str] = frozenset(), - no_warn_keys: set[str] | frozenset[str] = frozenset(), -) -> tuple[dict[str, Any], list[str]]: - """Strip parent ``EngineArgs`` fields before merging into stage YAML.""" - overridden: list[str] = [] - result: dict[str, Any] = {} - - for key, value in kwargs.items(): - if key in strip_keys: +def normalize_pipeline_cli_overrides( + pipeline: PipelineConfig, + cli_overrides: dict[str, Any], +) -> dict[str, Any]: + """Translate pipeline-owned global CLI aliases into stage-scoped overrides.""" + normalized = dict(cli_overrides) + for source, (stage_id, target) in pipeline.stage_cli_aliases.items(): + invalid_stage_keys = [ + key + for key, value in normalized.items() + if value is not None + and (match := _STAGE_OVERRIDE_PATTERN.match(key)) is not None + and match.group(2) == source + and int(match.group(1)) != stage_id + ] + if invalid_stage_keys: + invalid = ", ".join(sorted(invalid_stage_keys)) + raise ValueError( + f"{invalid} cannot be set for pipeline {pipeline.model_type!r}; " + f"{source} belongs to stage {stage_id} as {target}." + ) + value = normalized.pop(source, None) + if value is None: continue - - if key not in parent_fields or key in keep_keys: - result[key] = value + stage_key = f"stage_{stage_id}_{target}" + stage_value = normalized.get(stage_key) + if stage_value is not None and stage_value != value: + warnings.warn( + f"Ignoring {source}={value!r} because {stage_key}={stage_value!r} takes precedence.", + UserWarning, + stacklevel=2, + ) continue - - field_def = parent_fields[key] - if field_def.default is not dataclasses.MISSING: - default = field_def.default - elif field_def.default_factory is not dataclasses.MISSING: - default = field_def.default_factory() - else: - default = dataclasses.MISSING - - if default is dataclasses.MISSING or value is None: - continue - - if dataclasses.is_dataclass(default) and not isinstance(default, type): - default = asdict(default) - - if value != default and key not in no_warn_keys: - overridden.append(key) - - return result, sorted(overridden) + normalized[stage_key] = value + return normalized def _apply_diffusion_parallel_runtime_overrides( @@ -232,6 +228,10 @@ class StagePipelineConfig: # Alternates picked by ``merge_pipeline_deploy`` based on ``deploy.async_chunk``. async_chunk_process_next_stage_input_func: str | None = None sync_process_input_func: str | None = None + # Rewrites the Stage-0 view of a raw prompt before vLLM input processing. + # The callable receives ``(prompt, sampling_params_list)``; downstream + # stages continue to receive the original prompt. + prompt_transform_func: str | None = None prompt_expand_func: str | None = None cfg_kv_collect_func: str | None = None omni_kv_config: dict[str, Any] | None = None @@ -242,6 +242,16 @@ class StagePipelineConfig: # by ``stage_init_utils._resolve_model_tokenizer_paths``. model_subdir: str | None = None tokenizer_subdir: str | None = None + # Model-owned hook that resolves a pipeline root to this stage's checkpoint. + # Consumed and removed before backend engine args are constructed. + model_path_resolver: str | None = None + # Keep a single-replica diffusion stage in the orchestrator process. + # Disabled by default so existing multi-stage pipelines retain subprocess + # isolation unless their topology explicitly opts in. + inline_diffusion: bool = False + # Whether the non-async path waits for a complete upstream payload from + # the model-runner connector before scheduling this stage. + requires_full_payload_input: bool = False extras: dict[str, Any] = field(default_factory=dict) @@ -275,6 +285,7 @@ class PipelineConfig: # this value to auto-detect the pipeline. Only needed for diffusers-style # multi-component repos (e.g. GLM-Image). ``None`` = not a diffusers model. diffusers_class_name: str | None = None + diffusers_class_aliases: tuple[str, ...] = () endpoint_restrictions: tuple[EndpointRestriction, ...] = () # Optional model-owned duplex planner loaded by the stable engine runtime. duplex_runtime_extension: str | None = None @@ -288,6 +299,19 @@ class PipelineConfig: # Bundled deploy defaults for this concrete pipeline topology. The file is # loaded from vllm_omni/deploy; None uses DeployConfig defaults. default_deploy_config_name: str | None = None + # Global CLI spelling -> (stage id, stage-local spelling). + stage_cli_aliases: dict[str, tuple[int, str]] = field(default_factory=dict) + + def __post_init__(self) -> None: + errors = self.get_validation_errors() + + if errors: + # If we have multiple errors, put them on separate indented lines for readability + multi_sep = "\n\t" + error_str = multi_sep.join(errors) + if len(errors) > 1: + error_str = f"{multi_sep}{error_str}" + raise ValueError(f"PipelineConfig initialization failed with the following error(s): {error_str}") def get_stage(self, stage_id: int) -> StagePipelineConfig | None: """Look up a stage by its ID.""" @@ -296,7 +320,7 @@ class PipelineConfig: return stage return None - def validate(self) -> list[str]: + def get_validation_errors(self) -> list[str]: """Return list of topology errors (empty if valid).""" errors: list[str] = [] if not self.stages: @@ -314,6 +338,11 @@ class PipelineConfig: errors.append(f"Stage {stage.stage_id} references itself") if not any(not s.input_sources for s in self.stages): errors.append("No entry point (stage with empty input_sources)") + # Request completion and client-visible output are both driven by stages + # that declare ``final_output`` (see ``Orchestrator._route_output``). A + # pipeline without one can never emit a result, so every request hangs. + if not any(s.final_output for s in self.stages): + errors.append("No terminal stage (stage with final_output=True)") return errors @@ -342,6 +371,7 @@ class StageDeployConfig: default_sampling_params: dict[str, Any] | None = None default_pooling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 # === Generic stage engine fields === # Parallelism, scheduler, and memory-capacity controls. @@ -375,6 +405,7 @@ class StageDeployConfig: # Diffusion parallel_config deploy/runtime override fields. ulysses_degree: int | None = None ulysses_mode: str | None = None + ulysses_a2a_permute: bool | None = None ring_degree: int | None = None allgather_degree: int | None = None sequence_parallel_size: int | None = None @@ -396,6 +427,7 @@ class StageDeployConfig: diffusers_call_kwargs: dict[str, Any] | None = None diffusion_quantization_config: str | None = None diffusion_attention_backend: str | None = None + fastvideo_vsa_topk: int | None = None diffusion_attention_config: dict[str, Any] | None = None # Diffusion execution, cache, and VAE behavior. @@ -424,6 +456,9 @@ class StageDeployConfig: enable_distributed_layerwise_offload: bool | None = None dlo_use_allgather: bool | None = None dlo_resident_layers: int | None = None + host_weight_runtime_mode: str | None = None + host_weight_runtime_root: str | None = None + dlo_host_registration_limit_gib: float | None = None # Diffusion-specific debug and observability knobs. enable_diffusion_pipeline_profiler: bool | None = None @@ -573,7 +608,13 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: _DEEP_MERGE_KEYS = frozenset( - {"default_sampling_params", "default_pooling_params", "subtalker_sampling_params", "engine_extras", "engine_args"} + { + "default_sampling_params", + "default_pooling_params", + "subtalker_sampling_params", + "engine_extras", + "engine_args", + } ) @@ -673,6 +714,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]: def load_deploy_config(path: str | Path) -> DeployConfig: """Load a deploy YAML (with optional base_config inheritance).""" raw_dict = resolve_deploy_yaml(path) + if "stage_args" in raw_dict: + raise ValueError( + f"Deploy config {path} uses the removed `stage_args` schema; " + "define topology in PipelineConfig and deployment overrides under `stages`." + ) stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])] @@ -855,6 +901,9 @@ def _build_engine_args( engine_args["model_subdir"] = ps.model_subdir if ps.tokenizer_subdir: engine_args["tokenizer_subdir"] = ps.tokenizer_subdir + if ps.model_path_resolver: + engine_args["model_path_resolver"] = ps.model_path_resolver + engine_args["inline_diffusion"] = ps.inline_diffusion # Pipeline-wide top-level DeployConfig settings, applied to every stage. for name in _PIPELINE_WIDE_ENGINE_FIELDS: @@ -872,6 +921,7 @@ def _build_engine_args( # Materialize the resolved pipeline-wide async_chunk value into every # stage so explicit False overrides do not get lost downstream. engine_args["async_chunk"] = bool(deploy.async_chunk) + engine_args["session_mode"] = deploy.session_mode if deploy.session_mode == "duplex": # The engine admission limit is also the authoritative capacity for # model-owned streaming state. Propagate it to every stage instead of @@ -879,6 +929,7 @@ def _build_engine_args( engine_args["duplex_max_sessions"] = deploy.duplex_session.max_sessions if ps.omni_kv_config: engine_args["omni_kv_config"] = dict(ps.omni_kv_config) + engine_args["requires_full_payload_input"] = ps.requires_full_payload_input return engine_args @@ -900,6 +951,8 @@ def _build_extras( extras["output_connectors"] = dict(ds.output_connectors) if ds is not None and ds.input_connectors: extras["input_connectors"] = dict(ds.input_connectors) + if ps.prompt_transform_func: + extras["prompt_transform_func"] = ps.prompt_transform_func if ps.prompt_expand_func: extras["prompt_expand_func"] = ps.prompt_expand_func if ps.cfg_kv_collect_func: @@ -988,6 +1041,7 @@ def merge_pipeline_deploy( scheduler_cls=ps.scheduler_cls or _scheduler_path(sched_cls), hf_config_name=ps.hf_config_name, is_comprehension=ps.owns_tokenizer, + sampling_constraints=dict(ps.sampling_constraints), yaml_engine_args=engine_args, yaml_runtime=runtime, yaml_extras=extras, @@ -1015,6 +1069,7 @@ class StageConfig: scheduler_cls: str | None = None hf_config_name: str | None = None is_comprehension: bool = False + sampling_constraints: dict[str, Any] = field(default_factory=dict) yaml_engine_args: dict[str, Any] = field(default_factory=dict) yaml_runtime: dict[str, Any] = field(default_factory=dict) yaml_extras: dict[str, Any] = field(default_factory=dict) @@ -1075,6 +1130,7 @@ class StageConfig: "final_output": self.final_output, "final_output_type": self.final_output_type, "is_comprehension": self.is_comprehension, + "sampling_constraints": dict(self.sampling_constraints), } if self.custom_process_input_func: diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 66da231a6..ad7baf764 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -1,7 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations +import time from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Iterator from typing import Any import numpy as np @@ -104,6 +108,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Cache per-request flag to avoid repeated deserialization of additional_information self._omits_kv_transfer_cache: dict[str, bool] = {} + + # KV-wait start ts for the vllm_omni:kv_wait_s metric; see + # _emit_kv_wait_output for the engine-core → orchestrator carry. + self._kv_wait_start_ts: dict[str, float] = {} self._init_omni_io_scheduling_state() # Snapshot prompt length for each streaming input update self._new_prompt_len_snapshot: dict[str, int] = {} @@ -113,6 +121,70 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Output placeholders are zero when async scheduling isn't used return request.num_computed_tokens - request.num_output_placeholders + def _resolve_kv_connector_type(self) -> str: + """Connector backend name for the ``kv_wait_s`` label, or ``unknown``.""" + connector_cfg = self._get_omni_kv_config_value("connector_config") + if isinstance(connector_cfg, dict): + c_type = connector_cfg.get("type") + if isinstance(c_type, str) and c_type: + return c_type + return "unknown" + + def _emit_kv_wait_output( + self, + outputs: dict[int, list[OmniEngineCoreOutput]], + req_id: str, + req: Request, + ) -> None: + """Carry the KV-wait duration for ``req_id`` across the process boundary. + + Pops the start ts recorded at ENTER (``_free_request``); skips silently + when none was recorded. The wait rides ``kv_transfer_params`` to the + orchestrator, which calls ``observe_kv_wait``. + """ + start_ts = self._kv_wait_start_ts.pop(req_id, None) + if start_ts is None: + return + kv_wait_s = max(time.monotonic() - start_ts, 0.0) + outputs.setdefault(req.client_index, []).append( + OmniEngineCoreOutput( + request_id=req_id, + new_token_ids=[], + kv_transfer_params={ + "kv_wait_s": kv_wait_s, + "connector_type": self._resolve_kv_connector_type(), + }, + ) + ) + + def _clear_kv_wait_starts(self, request_ids: Iterable[str]) -> None: + """Drop incomplete KV-wait measurements for terminal requests.""" + wait_starts = getattr(self, "_kv_wait_start_ts", None) + if wait_starts is None: + return + for request_id in request_ids: + wait_starts.pop(request_id, None) + + def finish_requests( + self, + request_ids: str | Iterable[str] | None, + finished_status: RequestStatus, + ) -> list[Request]: + """Finish requests and discard any incomplete KV-wait timing.""" + if isinstance(request_ids, str): + cleanup_ids = (request_ids,) + finish_request_ids: str | tuple[str, ...] | None = request_ids + elif request_ids is None: + cleanup_ids = () + finish_request_ids = None + else: + finish_request_ids = list(request_ids) if isinstance(request_ids, Iterator) else request_ids + cleanup_ids = tuple(finish_request_ids) + + finished = super().finish_requests(finish_request_ids, finished_status) + self._clear_kv_wait_starts(cleanup_ids) + return finished + def _get_kv_transfer_criteria(self) -> dict | None: return self._get_omni_kv_config_value("kv_transfer_criteria") @@ -223,15 +295,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # them. Upstream vllm raises RuntimeError on this status; omni allows # async abort (e.g. client disconnect during TTS streaming) to leave # requests in the waiting/running queues temporarily. - for queue in (self.waiting, self.running): - for req in list(queue): - if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED: - queue.remove(req) + waiting = getattr(self, "waiting") + self._drop_aborted_queued_requests() self._process_pending_omni_inputs(model_mode="ar") + self._drop_aborted_queued_requests() + self._resync_streaming_input_counter() original_waiting = None if self._should_defer_waiting_admission(): - original_waiting = self.waiting + original_waiting = waiting self.waiting = create_request_queue(self.policy) try: @@ -270,8 +342,25 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): inter_stage_outputs = getattr(model_runner_output, "inter_stage_outputs", None) num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output + ec_connector_output = getattr(model_runner_output, "ec_connector_output", None) cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats + # Every GPU write enqueued by this and earlier steps has completed, so + # it is safe to return deferred-free blocks to the pool. This is the + # update-side half of the upstream v0.28 deferred-free fence; the schedule + # half advances sched_step_seq inside super().schedule(). + # getattr: __new__-constructed test schedulers carry no + # defer_block_free attribute. + if ( + getattr(self, "defer_block_free", False) + # getattr again: SchedulerOutput is mocked field-by-field in the + # scheduler unit tests, and dataclass fields without defaults are + # invisible to MagicMock(spec=...). + and getattr(scheduler_output, "total_num_scheduled_tokens", 0) > 0 + ): + self.processed_step_seq += 1 + self._drain_deferred_frees() + perf_stats: PerfStats | None = None if self.perf_metrics and self.perf_metrics.is_enabled(): perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) @@ -306,6 +395,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): kv_transfer_params={"kv_ready": True}, ) ) + # EXIT: extraction ack ends the wait — carry the duration. + # req is still alive here (the del happens in the cleanup + # loop below); skip if it was already freed upstream. + if req is not None: + self._emit_kv_wait_output(outputs, req_id, req) except Exception: init_logger(__name__).exception("Failed to pre-process KV extraction for %s", req_id) @@ -345,15 +439,24 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # request is aborted while the model is executing it (e.g., # in pipeline parallelism or async scheduling). continue - if output_is_stale: + req_index = model_runner_output.req_id_to_index[req_id] + generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] + + stale_async_tokens = int(getattr(request, "async_tokens_to_discard", 0) or 0) + async_output_is_stale = bool(generated_token_ids and stale_async_tokens > 0) + if async_output_is_stale: + # Drain this marker even when the same frame also belongs to + # the scheduled-token stale window below. Both accounting + # domains must consume the old frame before new output passes. + request.async_tokens_to_discard = max(0, stale_async_tokens - len(generated_token_ids)) + + if output_is_stale or async_output_is_stale: # Output of a step scheduled before the request's in-flight # tokens were discarded (segment stop / session replacement). # num_computed_tokens was rolled back at the discard site, so # this output must not be appended or emitted. continue - req_index = model_runner_output.req_id_to_index[req_id] - generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] status_before_stop = request.status new_logprobs = None logprob_validation_failed = False @@ -412,6 +515,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): mm_output = mm_outputs[req_index] if mm_outputs else None inter_stage_output = inter_stage_outputs[req_index] if inter_stage_outputs else None kv_transfer_params = None + ec_transfer_params = None finish_reason = None routed_experts = None @@ -463,13 +567,33 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): request.resumable = False stopped = True + # Finalize prefill stats BEFORE stop handling (upstream v0.28 + # order): _free_request below releases the KV blocks, after which + # estimate_cached_tokens(request) reports 0. kv_transfer_params is + # omitted from the emission predicate here because it only becomes + # non-None when stopped is already True. + prefill_stats = None + if new_token_ids or mm_output is not None or pooler_output is not None or stopped: + prefill_stats = request.take_prefill_stats() + if prefill_stats is not None: + prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request)) + + confirmed_num_computed_tokens = None + boundary_generation = None if stopped: + if self.chunk_transfer_adapter is not None: + confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request) if model_runner_output.routed_experts is not None: routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request) # Capture finish_reason BEFORE _handle_stopped_request, which may # reset the status to WAITING for streaming requests that continue. finish_reason = request.get_finished_reason() + if self.chunk_transfer_adapter is not None: + try: + boundary_generation = int(getattr(request, "_omni_segment_generation", 0) or 0) + except (TypeError, ValueError): + boundary_generation = 0 finished = self._handle_stopped_request(request) is_segment_finished = not finished if finished: @@ -478,9 +602,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # for streaming input request only if self.chunk_transfer_adapter: if self.vllm_config.model_config.stage_id != 0: - # Downstream async-chunk stages receive real payloads from the - # connector. This update only resumes polling for the next segment. - self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id) + # Only a connector-fed receiver in native duplex can + # poll the next segment without an external update. + # Sender-only and turn-mode stages remain parked. + self._resume_downstream_chunk_receiver(request) outstanding_async_tokens = request.num_output_placeholders # Always record the discard signal (0 when nothing is in # flight). Upstream a0c092ee72 removed the @@ -499,8 +624,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # valid new-segment frames or underflowed the drain # assert whenever placeholder counts diverged from # scheduled counts (spec drafts, in-flight prefill). + # Assign, never accumulate: _handle_stopped_request above + # may already have fenced the same in-flight tokens through + # a queued streaming update, and adding twice swallows the + # next duplex unit's listen/speak under async scheduling. if request.num_in_flight_tokens > 0: - request.num_stale_output_tokens += request.num_in_flight_tokens + request.num_stale_output_tokens = request.num_in_flight_tokens if outstanding_async_tokens > 0: # Discard only outputs that are already in flight and # roll back their optimistic computed-token accounting. @@ -509,7 +638,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): request.spec_token_ids = [] request._output_token_ids.clear() if finished: - kv_transfer_params, _ = self._free_request(request) + kv_transfer_params, ec_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK: @@ -537,8 +666,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): pooling_output=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, - prefill_stats=request.take_prefill_stats(), + prefill_stats=prefill_stats, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, is_segment_finished=is_segment_finished, @@ -551,10 +681,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter is not None and ( inter_stage_output is not None or is_segment_finished or finished ): + save_kwargs = { + "new_token_ids": new_token_ids, + "confirmed_num_computed_tokens": confirmed_num_computed_tokens, + } + if is_segment_finished: + save_kwargs["segment_generation"] = boundary_generation self.chunk_transfer_adapter.save_async( inter_stage_output, request, is_segment_finished, + **save_kwargs, ) self._remove_stopped_requests_from_queues( @@ -576,6 +713,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # EC Connector: update state from worker-side EC connector output. + # Use getattr for safety with test __new__/SimpleNamespace code paths. + if getattr(self, "ec_connector", None) is not None and ec_connector_output: + self.ec_connector.update_connector_output(ec_connector_output) + kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output) self._publish_kv_cache_events() @@ -635,10 +777,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # pre-replacement frame will drain, so the counter reaches exactly # zero and the new segment's frames are never swallowed. This also # covers an in-flight prefill chunk, which carries no placeholders - # but must still have its late output dropped. + # but must still have its late output dropped. Assign, never + # accumulate: update_from_output fences the same in-flight tokens when + # a resumable stop applies a queued update through this helper. in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0) if in_flight_tokens > 0: - session.num_stale_output_tokens += in_flight_tokens + session.num_stale_output_tokens = in_flight_tokens if outstanding_async_tokens > 0: # Async scheduling may already have sampled the previous # segment's next token. Drop that late token instead of @@ -648,9 +792,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): session.num_output_placeholders = 0 session.spec_token_ids = [] stage_id = self.vllm_config.model_config.stage_id + + update_infos = ( + getattr(update, "model_intermediate_buffer", None), + getattr(update, "additional_information", None), + ) if self.chunk_transfer_adapter and self.chunk_transfer_adapter.receives_chunks: self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None) if stage_id != 0: + session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 # Downstream async-chunk stages receive real payloads from the # connector. This update only resumes polling for the next segment. self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id) @@ -668,10 +818,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) return - update_infos = ( - getattr(update, "model_intermediate_buffer", None), - getattr(update, "additional_information", None), - ) replace_streaming_prompt = any( isinstance(info, dict) and isinstance(info.get("meta"), dict) @@ -679,8 +825,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): for info in update_infos ) if replace_streaming_prompt: + self._release_replaced_streaming_prompt_cache(session) self._replace_streaming_session(session, update) return + session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 super()._update_request_as_session(session, update) if hasattr(update, "model_intermediate_buffer"): session.model_intermediate_buffer = update.model_intermediate_buffer @@ -702,6 +850,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # 1. Standard cleanup parts from base _free_request connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) + # EC Connector: mirror the KV hook (upstream v0.28 _free_request). + # The contract requires firing before the encoder cache is freed so + # the connector can inspect per-request state (e.g. which mm_hashes + # it recorded during save_caches()) and emit ec_transfer_params for + # the response body. getattr: __new__-constructed test schedulers + # carry no ec_connector attribute. + ec_xfer_params: dict[str, Any] | None = None + if getattr(self, "ec_connector", None) is not None: + ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request) + connector_delay_free_blocks |= ec_delay_free + self.encoder_cache_manager.free(request) request_id = request.request_id self.finished_req_ids.add(request_id) @@ -725,11 +884,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # It triggered but hasn't finished yet. We MUST wait. logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.") self.waiting_for_transfer_free.add(request_id) + self._kv_wait_start_ts[request_id] = time.monotonic() kv_xfer_params = None - return kv_xfer_params, None + return kv_xfer_params, ec_xfer_params elif request_id in self.waiting_for_transfer_free: # Blocks held until KV extraction completes in a future step. - return None, None + return None, ec_xfer_params else: logger.debug( f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). " @@ -737,6 +897,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): ) else: self.waiting_for_transfer_free.add(request_id) + self._kv_wait_start_ts[request_id] = time.monotonic() confirmed_computed = self._get_confirmed_num_computed_tokens(request) self._mark_request_for_kv_transfer(request_id, confirmed_computed) # Return KV transfer metadata so it propagates to RequestOutput @@ -766,14 +927,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if isinstance(add_info, dict): add_info.update(kv_xfer_params) - return kv_xfer_params, None + return kv_xfer_params, ec_xfer_params # 3. Standard Freeing delay_free_blocks |= connector_delay_free_blocks if not delay_free_blocks: self._free_blocks(request) - return kv_xfer_params, None + return kv_xfer_params, ec_xfer_params finally: self._free_input_coordinator_request(request_id) # Normal completion runs through here, not finish_requests() diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index ae19a41a3..db0eedcf7 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import time @@ -98,7 +101,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Temporary queue: preserve waiting order while requests await input. skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 + self._drop_aborted_queued_requests() self._process_pending_omni_inputs(model_mode="generation") + self._drop_aborted_queued_requests() + self._resync_streaming_input_counter() # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled @@ -309,6 +315,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): ec_meta = self.ec_connector.build_connector_meta(scheduler_output) scheduler_output.ec_connector_metadata = ec_meta + # Advance the fence only for non-empty steps (those that actually + # write KV and have their output processed later in + # update_from_output). Must precede _update_after_schedule, which + # stamps request.last_sched_seq from the advanced value; the other + # half drains in update_from_output. The fallback path above gets + # both halves from super().schedule(). getattr: __new__-constructed + # test schedulers carry no defer_block_free attribute. + if getattr(self, "defer_block_free", False) and total_num_scheduled_tokens > 0: + self.sched_step_seq += 1 + # Update internal state (advance num_computed_tokens, free encoder inputs, # etc.) self._update_after_schedule(scheduler_output) @@ -345,8 +361,26 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): mm_outputs = getattr(model_runner_output, "multimodal_outputs", None) num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output + ec_connector_output = getattr(model_runner_output, "ec_connector_output", None) cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats + + # Every GPU write enqueued by this and earlier steps has completed, so + # it is safe to return deferred-free blocks to the pool. This is the + # update-side half of the upstream v0.28 deferred-free fence; the schedule + # half advances sched_step_seq in the fast path above. + # getattr: __new__-constructed test schedulers carry no + # defer_block_free attribute. + if ( + getattr(self, "defer_block_free", False) + # getattr again: SchedulerOutput is mocked field-by-field in the + # scheduler unit tests, and dataclass fields without defaults are + # invisible to MagicMock(spec=...). + and getattr(scheduler_output, "total_num_scheduled_tokens", 0) > 0 + ): + self.processed_step_seq += 1 + self._drain_deferred_frees() + perf_stats: PerfStats | None = None if self.perf_metrics and self.perf_metrics.is_enabled(): perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) @@ -438,6 +472,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): new_logprobs = None new_token_ids = generated_token_ids kv_transfer_params = None + ec_transfer_params = None pooler_output = pooler_outputs[req_index] if pooler_outputs else None mm_output = mm_outputs[req_index] if mm_outputs else None status_before_stop = request.status @@ -474,6 +509,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # (does not affect protocol) stopped = True + # Finalize prefill stats BEFORE stop handling (upstream v0.28 + # order): _free_request below releases the KV blocks, after which + # estimate_cached_tokens(request) reports 0. kv_transfer_params is + # omitted from the emission predicate here because it only becomes + # non-None when stopped is already True. + prefill_stats = None + if new_token_ids or mm_output is not None or pooler_output is not None or stopped: + prefill_stats = request.take_prefill_stats() + if prefill_stats is not None: + prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request)) + if stopped: if model_runner_output.routed_experts is not None: routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request) @@ -485,7 +531,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.segment_finished_requests.discard(req_id) if finished: - kv_transfer_params, _ = self._free_request(request) + kv_transfer_params, ec_transfer_params = self._free_request(request) if self.chunk_transfer_adapter is not None: self.chunk_transfer_adapter.cleanup( request.request_id, @@ -527,8 +573,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): pooling_output=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, - prefill_stats=request.take_prefill_stats(), + prefill_stats=prefill_stats, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, is_segment_finished=is_segment_finished, @@ -547,8 +594,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): finished = self._handle_stopped_request(request) is_segment_finished = not finished kv_transfer_params = None + ec_transfer_params = None if finished: - kv_transfer_params, _ = self._free_request(request) + kv_transfer_params, ec_transfer_params = self._free_request(request) if self.chunk_transfer_adapter is not None: self.chunk_transfer_adapter.cleanup( request.request_id, @@ -562,6 +610,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason=finish_reason, stop_reason=request.stop_reason, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, is_segment_finished=is_segment_finished, ) stopped_running_reqs.add(request) @@ -587,6 +636,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # EC Connector: update state from worker-side EC connector output. + # Use getattr for safety with test __new__/SimpleNamespace code paths. + if getattr(self, "ec_connector", None) is not None and ec_connector_output: + self.ec_connector.update_connector_output(ec_connector_output) + kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output) self._publish_kv_cache_events() diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index e0df774c0..58c3546cd 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations +import math import os import time -from collections.abc import Iterable +from collections.abc import Iterable, Iterator from typing import Any import torch @@ -42,25 +46,56 @@ logger = init_logger(__name__) _STATS_INTERVAL_S = 1.0 -# Upper bound on how long a request may sit in full-payload-input wait -# (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``) -# before the scheduler force-fails it. Defends against stuck consumer-side -# requests when the producer drops a full-payload, send fails, or recv -# never arrives. Override per-deployment via -# VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set <=0 to disable the safety net. +# Upper bound on how long a request may wait for stage input before the +# scheduler force-fails it. Defends against stuck consumer-side requests when +# the producer drops a payload, send fails, or recv never arrives. Override +# per-deployment via VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set exactly 0 to disable +# the safety net. Negative and non-finite values are rejected at startup. # -# Scope: this constant only covers the full-payload coordinator path -# (``input_coordinator``). The async-chunk path uses -# ``chunk_transfer_adapter`` and is not affected by this constant. +# Scope: both transports. The full-payload path measures time parked in +# WAITING_FOR_INPUT (``OmniSchedulingCoordinator._waiting_since``); the +# async-chunk path measures time stalled in WAITING_FOR_CHUNK +# (``OmniChunkTransferAdapter._waiting_since``, reset on each chunk arrival). +# One knob, because the question it answers -- "how long may a request wait for +# stage input?" -- does not depend on which transport carries it. +_INPUT_WAIT_TIMEOUT_DEFAULT = 600.0 _INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600") try: DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW) except ValueError: logger.warning( - "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.", + "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to %s seconds.", _INPUT_WAIT_TIMEOUT_RAW, + _INPUT_WAIT_TIMEOUT_DEFAULT, + ) + DEFAULT_INPUT_WAIT_TIMEOUT_S = _INPUT_WAIT_TIMEOUT_DEFAULT + +if not math.isfinite(DEFAULT_INPUT_WAIT_TIMEOUT_S): + # nan compares false against everything, so every deadline check silently + # passes; inf makes the deadline unreachable. Both disable the net while + # looking like a number, which is the failure this item exists to stop. + raise ValueError( + f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is not finite. " + "Use a positive number of seconds, or exactly 0 to disable the deadline." + ) +if DEFAULT_INPUT_WAIT_TIMEOUT_S < 0: + # Fail at startup rather than substituting a default. `-1` is a common idiom + # for "no limit", so silently turning it into 600 would give an operator who + # meant "never time out" mysterious failures ten minutes in. + raise ValueError( + f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is negative. " + "Use a positive number of seconds, or exactly 0 to disable the deadline." + ) +elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0: + # Zero stays a supported opt-out, but it is no longer silent: it disables + # the deadline on BOTH transports, so a request whose stage input never + # arrives waits forever. Say which nets just went away. + logger.warning( + "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=0: stage-input deadlines are DISABLED for " + "both the full-payload path (WAITING_FOR_INPUT) and the async-chunk path " + "(WAITING_FOR_CHUNK). A request whose input never arrives will wait " + "indefinitely instead of failing." ) - DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0 class OmniSchedulerMixin: @@ -114,9 +149,13 @@ class OmniSchedulerMixin: def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None: """Replace a downstream stage's placeholder with its next payload.""" + session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 adapter = getattr(self, "chunk_transfer_adapter", None) if adapter is not None: adapter.segment_finished_requests.discard(session.request_id) + watermark = getattr(adapter, "requests_num_chunks_sent", None) + if watermark is not None: + watermark.pop(session.external_req_id, None) session._output_token_ids.clear() session._all_token_ids.clear() # In-flight outputs from the previous segment were optimistically @@ -127,7 +166,10 @@ class OmniSchedulerMixin: # pre-replacement frame will drain, so the counter reaches exactly # zero; a placeholder-based seed swallowed valid new-segment frames # whenever placeholder counts diverged from scheduled counts. - session.num_stale_output_tokens += int(getattr(session, "num_in_flight_tokens", 0) or 0) + # num_in_flight_tokens already includes any undrained stale share. + # Assign instead of accumulating so callers that fenced the same + # rollover before entering this helper do not count it twice. + session.num_stale_output_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0) session.num_output_placeholders = 0 session.spec_token_ids = [] new_prompt = update.prompt_token_ids or () @@ -153,6 +195,44 @@ class OmniSchedulerMixin: if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) + def _release_replaced_streaming_prompt_cache(self, session: Request) -> None: + """Discard cache state that belongs to a replaced prompt.""" + # A prompt replacement is not a normal streaming extension: none of + # the old KV blocks or encoder state is valid for the new prompt. Use + # the scheduler's block-free path so an in-flight GPU step is fenced + # correctly before the blocks return to the pool. + self._free_request_blocks(session) + self.encoder_cache_manager.free(session) + getattr(self, "_inflight_prefills", set()).discard(session) + + def _reset_ready_async_chunk_replacements(self) -> None: + """Release stale cache state after an async-chunk prompt rollover.""" + adapter = getattr(self, "chunk_transfer_adapter", None) + if adapter is None: + return + replaced_ids = getattr(adapter, "replaced_streaming_prompt_ids", None) + ready_ids = getattr(adapter, "requests_with_ready_chunks", None) + if not replaced_ids or not ready_ids: + return + + for request_id in tuple(replaced_ids & ready_ids): + request = self.requests.get(request_id) + if request is None: + replaced_ids.discard(request_id) + continue + # The streaming update may already have fenced this same in-flight + # frame. Seed idempotently so the replacement does not count it twice. + request.num_stale_output_tokens = int(getattr(request, "num_in_flight_tokens", 0) or 0) + request.num_output_placeholders = 0 + request.spec_token_ids = [] + self._release_replaced_streaming_prompt_cache(request) + watermark = getattr(adapter, "requests_num_chunks_sent", None) + if watermark is not None: + watermark.pop(request.external_req_id, None) + # Consume this marker after the one-time cache reset. The separate + # ready-chunk marker remains until scheduler admission succeeds. + replaced_ids.discard(request_id) + def _consume_pending_connector_output(self, model_mode: str) -> None: """Drain ``self._latest_omni_connector_output`` into the coordinator. @@ -184,6 +264,29 @@ class OmniSchedulerMixin: self.running, scheduler_requests=self.requests, ) + # Prompt-window contracts are finalized on this scheduler thread + # and may surface a permanent receive failure. + self._process_chunk_receive_failures() + self._reset_ready_async_chunk_replacements() + self._process_pending_chunk_timeouts() + self._log_failed_chunk_sends() + + def _process_chunk_receive_failures(self) -> None: + """Fail requests whose consumed async chunk violated its contract.""" + adapter = getattr(self, "chunk_transfer_adapter", None) + collector = getattr(adapter, "collect_failed_receive_request_ids", None) + if collector is None: + return + failures = collector() + present_ids = {request_id for request_id in failures if request_id in self.requests} + if not present_ids: + return + logger.error( + "Marking %d request(s) as FINISHED_ERROR after invalid connector input: %s", + len(present_ids), + {request_id: failures[request_id] for request_id in sorted(present_ids)}, + ) + self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR) def _restore_omni_wait_queues(self) -> None: """Restore requests temporarily parked by Omni input gates.""" @@ -211,9 +314,9 @@ class OmniSchedulerMixin: Disabled when ``DEFAULT_INPUT_WAIT_TIMEOUT_S`` is <= 0. Scope: only covers ``input_coordinator`` (full-payload path). - Async-chunk requests park in ``chunk_transfer_adapter`` instead - and are not handled here -- if a similar safety net is needed - for the chunk path, it belongs in the chunk adapter. + Async-chunk requests park in ``chunk_transfer_adapter`` instead and + are handled by ``_process_pending_chunk_timeouts`` below, which shares + this timeout. """ if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0: return @@ -234,6 +337,84 @@ class OmniSchedulerMixin: ) self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR) + def _log_failed_chunk_sends(self) -> None: + """Surface chunks the sender gave up on (R1.2 of #4855). + + ``save_loop`` pops a task with ``popleft()`` and never re-queues it, and + ``connector.put`` can report failure without raising, so a dropped chunk + is a give-up rather than a retry. Until this change the only trace was a + warning that printed ``None`` for the request id, because the task dict + has no ``request_id`` key. + + This only reports. Failing the request from here does not work: the + producer is not the ``final_output`` stage, so an abort synthesized by + ``finish_requests`` is not turned into a client-visible termination by + ``Orchestrator._route_output`` (see the ``final_output`` guards there). + Making a producer-side drop client-visible needs an orchestrator-side + path like ``_fail_request_dead_stage``; the consumer's R1.1 deadline is + what actually ends the request today. + """ + adapter = getattr(self, "chunk_transfer_adapter", None) + if adapter is None: + return + failures = adapter.collect_failed_send_request_ids() + for request_id, reason in failures.items(): + if request_id in self.requests: + logger.error( + "[OmniScheduler] req=%s: chunk send gave up (%s); the consumer will " + "fail on the stage-input deadline", + request_id, + reason, + ) + else: + # Drained but no longer scheduled here: the request finished or + # was aborted between the failed send and this sweep. Still say + # so -- collect_* empties the map, so staying quiet would make + # this the one drop path in a change whose point is visibility. + logger.warning( + "[OmniScheduler] req=%s: chunk send gave up (%s) after the request " + "left this scheduler (finished or aborted)", + request_id, + reason, + ) + + def _process_pending_chunk_timeouts(self) -> None: + """Force-fail requests stalled in ``WAITING_FOR_CHUNK`` too long. + + The async-chunk counterpart of ``_process_pending_input_timeouts``, + called right after ``chunk_transfer_adapter.process_pending_chunks``. + Until now the chunk path had no safety net at all: the scheduler-side + comment on the full-payload net said a similar guard "belongs in the + chunk adapter", while the worker mixin assumed the full-payload net + already covered it. Only the first was true, so a dropped terminal + chunk or a producer that exhausted its send retries parked the request + forever (vllm-project/vllm-omni#3833). Since ``async_chunk: true`` is + the default for the flagship TTS pipelines, that was the default + streaming path. + + Shares ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S`` with the full-payload net: + one knob for "how long may a request wait for stage input", regardless + of which transport carries it. Disabled when the value is <= 0. + """ + if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0: + return + adapter = getattr(self, "chunk_transfer_adapter", None) + if adapter is None or not getattr(adapter, "receives_chunks", False): + return + timed_out_ids = adapter.collect_timed_out_request_ids(timeout_s=DEFAULT_INPUT_WAIT_TIMEOUT_S) + if not timed_out_ids: + return + present_ids = {req_id for req_id in timed_out_ids if req_id in self.requests} + if not present_ids: + return + logger.warning( + "Marking %d request(s) as FINISHED_ERROR after stalling > %.0fs waiting for a chunk: %s", + len(present_ids), + DEFAULT_INPUT_WAIT_TIMEOUT_S, + sorted(present_ids), + ) + self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR) + def _capture_omni_connector_output(self, model_runner_output: Any) -> None: """Stash the model runner's omni_connector_output for next schedule(). @@ -315,6 +496,7 @@ class OmniSchedulerMixin: stop_reason: Any = None, prefill_stats: Any = None, kv_transfer_params: Any = None, + ec_transfer_params: Any = None, routed_experts: Any = None, num_nans_in_logits: int = 0, is_segment_finished: bool | None = False, @@ -333,6 +515,7 @@ class OmniSchedulerMixin: events=request.take_events(), prefill_stats=prefill_stats, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, trace_headers=request.trace_headers, routed_experts=routed_experts, num_nans_in_logits=num_nans_in_logits, @@ -407,8 +590,41 @@ class OmniSchedulerMixin: if stopped_running_reqs: self.running = remove_all(self.running, stopped_running_reqs) if stopped_preempted_reqs: - self.waiting.remove_requests(stopped_preempted_reqs) - self.skipped_waiting.remove_requests(stopped_preempted_reqs) + # ``_handle_stopped_request`` re-enqueues a resumable segment as + # WAITING_FOR_STREAMING_REQ before this cleanup runs. A + # downstream async-chunk request can have entered the update with + # the older WAITING_FOR_CHUNK status, which puts it in + # ``stopped_preempted_reqs``; deleting it here would immediately + # undo the requeue and the receiver would never poll chunk 1. + removable = { + request + for request in stopped_preempted_reqs + if not ( + request.resumable + and request.status + in ( + RequestStatus.WAITING, + RequestStatus.WAITING_FOR_STREAMING_REQ, + ) + ) + } + self.waiting.remove_requests(removable) + self.skipped_waiting.remove_requests(removable) + + def _resume_downstream_chunk_receiver(self, request: Request) -> None: + """Resume duplex connector polling without an external update.""" + adapter = self.chunk_transfer_adapter + adapter.segment_finished_requests.discard(request.request_id) + if ( + not adapter.receives_chunks + or getattr(self.vllm_config.model_config, "session_mode", "turn") != "duplex" + or request.status != RequestStatus.WAITING_FOR_STREAMING_REQ + ): + return + self.num_waiting_for_streaming_input -= 1 + request.status = RequestStatus.WAITING + self.skipped_waiting.remove_requests((request,)) + self._enqueue_waiting_request(request) def _aggregate_kv_connector_stats( self, @@ -459,12 +675,36 @@ class OmniSchedulerMixin: finished_status: RequestStatus, ) -> list[Request]: """Finish requests and clean all Omni-owned queue/coordinator state.""" - if self.chunk_transfer_adapter: - self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + if isinstance(request_ids, str): + target_request_ids = {request_ids} + elif request_ids is None: + target_request_ids = set(self.requests) + else: + if isinstance(request_ids, Iterator): + request_ids = tuple(request_ids) + target_request_ids = set(request_ids) - self._realign_request_status_to_queues(request_ids) + skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())} + pre_adapter_streaming_wait_ids = { + rid + for rid in target_request_ids & skipped_waiting_ids + if (req := self.requests.get(rid)) is not None + and ( + req.status == RequestStatus.WAITING_FOR_STREAMING_REQ + or (getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED) + ) + } + + if chunk_transfer_adapter := getattr(self, "chunk_transfer_adapter", None): + chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + + self._realign_request_status_to_queues( + request_ids, + pre_adapter_streaming_wait_ids=pre_adapter_streaming_wait_ids, + ) finished = super().finish_requests(request_ids, finished_status) - self._purge_finished_from_running() + self._purge_finished_from_running(target_request_ids) + self._resync_streaming_input_counter() for request in finished: self._free_input_coordinator_request(request.request_id) @@ -480,6 +720,8 @@ class OmniSchedulerMixin: def _realign_request_status_to_queues( self, request_ids: str | Iterable[str] | None, + *, + pre_adapter_streaming_wait_ids: set[str] | None = None, ) -> None: """Realign ``request.status`` to actual queue membership. @@ -506,7 +748,13 @@ class OmniSchedulerMixin: Realign here: if a request lives in ``self.running`` but its status is not ``RUNNING``, set it to ``RUNNING``; symmetrically flip ``RUNNING → WAITING`` when the request is actually in - ``self.waiting``. This is a localized safety net for + ``self.waiting``. A resumable segment stop still held in + ``self.skipped_waiting`` is restored to + ``WAITING_FOR_STREAMING_REQ`` so upstream also balances its + paused-session counter. Because adapter cleanup may restore a + connector-owned request's prior status first, ``finish_requests`` + snapshots these counter-bearing rows before invoking the adapter. + This is a localized safety net for ``requests_origin_status`` staleness on the admit transition; it does not touch the adapter's invariants and is complementary to the chunk-transfer-adapter deque purge that already runs @@ -539,17 +787,36 @@ class OmniSchedulerMixin: running_ids = {r.request_id for r in self.running} waiting_ids = {r.request_id for r in self.waiting} + skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())} for rid in ids_to_align: req = self.requests.get(rid) - if req is None or req.is_finished(): + if req is None: continue - if rid in running_ids and req.status != RequestStatus.RUNNING: + # A persistent Session may be closed after its current segment + # reached FINISHED_STOPPED but while it is still owned by a live + # scheduler/connector queue. vLLM skips already-finished requests + # in finish_requests(), leaking the KV and worker slot. Only + # recover requests with positive queue ownership; an off-queue + # terminal can legitimately be waiting for deferred block free. + resumable_segment_stop = bool( + getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED + ) + streaming_wait = ( + rid in pre_adapter_streaming_wait_ids + if pre_adapter_streaming_wait_ids is not None + else resumable_segment_stop + ) + if req.is_finished() and not (resumable_segment_stop or streaming_wait): + continue + if rid in skipped_waiting_ids and streaming_wait: + req.status = RequestStatus.WAITING_FOR_STREAMING_REQ + elif rid in running_ids and req.status != RequestStatus.RUNNING: req.status = RequestStatus.RUNNING - elif rid in waiting_ids and req.status == RequestStatus.RUNNING: + elif rid in waiting_ids and (req.status == RequestStatus.RUNNING or resumable_segment_stop): req.status = RequestStatus.WAITING - def _purge_finished_from_running(self) -> None: + def _purge_finished_from_running(self, target_request_ids: set[str] | None = None) -> None: """Defensive post-finish sweep of ``self.running``. Belt-and-suspenders to ``_realign_request_status_to_queues``: @@ -566,18 +833,14 @@ class OmniSchedulerMixin: residue after ``super().finish_requests`` so any stale entries are reclaimed). - Scope of the predicate. ``is_finished()`` covers entries the - upstream ``finish_requests`` already drained from ``self.requests`` - but failed to remove from ``self.running``; the - ``request_id not in self.requests`` arm catches the same surface - from a different angle and is the post-cleanup mirror of the - deque purge ``_purge_untracked_chunk_requests`` already runs at - the chunk-transfer-adapter layer. It does **not** by itself make - arbitrary direct deletions of ``self.requests`` safe -- callers - that pop ``self.requests`` outside the standard finish path - still have to go through ``_free_request`` (or equivalent) for - block / connector / coordinator cleanup. This sweep only - reclaims the ``self.running`` slot reference. + A resumable ``FINISHED_STOPPED`` request may legitimately remain + in ``self.running`` between realtime segments. Preserve such a + request unless it belongs to this finish call. Other finished or + untracked entries are stale and can be swept defensively. + + When ``target_request_ids`` is ``None`` or empty, every resumable + ``FINISHED_STOPPED`` entry is preserved. Production + ``finish_requests`` always passes its resolved finish set. In-place via ``self.running[:] = ...`` for minor consistency with idiomatic vLLM scheduler mutation; upstream @@ -594,4 +857,38 @@ class OmniSchedulerMixin: """ if not self.running: return - self.running[:] = [req for req in self.running if not req.is_finished() and req.request_id in self.requests] + target_request_ids = target_request_ids or set() + + def keep_running(req: Request) -> bool: + if req.request_id not in self.requests: + return False + if not req.is_finished(): + return True + resumable_segment_stop = bool( + getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED + ) + return resumable_segment_stop and req.request_id not in target_request_ids + + self.running[:] = [req for req in self.running if keep_running(req)] + + def _drop_aborted_queued_requests(self) -> None: + for queue in (self.waiting, self.skipped_waiting): + aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED] + if aborted: + queue.remove_requests(aborted) + self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED] + + def _resync_streaming_input_counter(self) -> None: + counter = getattr(self, "num_waiting_for_streaming_input", None) + if counter is None: + return + parked = sum( + 1 + for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None)) + if queue + for request in queue + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ + ) + if parked != counter: + logger.debug("[Omni] resynced streaming-input counter %s -> %s", counter, parked) + self.num_waiting_for_streaming_input = parked diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 75f234d8c..93f40daa0 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -22,75 +22,19 @@ from vllm_omni.core.sched.output import OmniChunkRecvHandle logger = init_logger(__name__) -# (arch, model_stage) pairs that route their full_payload stage input via -# the worker connector and therefore need the scheduler-side coordinator to -# park requests in WAITING_FOR_INPUT until the recv side delivers. This set -# must stay aligned with the arch scope of `init_omni_connectors` in -# gpu_ar_model_runner.py and gpu_generation_model_runner.py. Adding a stage -# here without also wiring its worker connector init produces a permanent -# Stage 1 hang (gate parks the request, no transport ever releases it). -# -_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset( - { - ("Qwen3OmniMoeForConditionalGeneration", "talker"), - ("Qwen3OmniMoeForConditionalGeneration", "code2wav"), - # qwen2_5_omni thinker->talker uses the real full-payload - # producer builder (text_hidden_states routed via - # pooler_output["hidden"] -> accumulator -> connector). Both - # stages of qwen2_5_omni are enabled. - ("Qwen2_5OmniForConditionalGeneration", "talker"), - ("Qwen2_5OmniForConditionalGeneration", "code2wav"), - # covo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1). - ("CovoAudioForConditionalGeneration", "code2wav"), - # mimo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1). - ("MiMoAudioModel", "code2wav"), - # qwen3_tts: Qwen3TTSTalkerForConditionalGeneration (Stage 0) - # -> Qwen3TTSCode2Wav (Stage 1). Stage 1 is the consumer. - ("Qwen3TTSCode2Wav", "code2wav"), - # minicpmo_4_5: Talker (Stage 1) -> Code2Wav (Stage 2). - ("MiniCPMO45Code2Wav", "code2wav"), - # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1). - ("CosyVoice3Model", "cosyvoice3_code2wav"), - # nemotron_voicechat: talker (Stage 1) -> code2wav (Stage 2). Stage 2 - # waits for the talker's full-payload code stacks; the thinker (Stage 0) - # -> talker (Stage 1) hop is token-path only and must NOT be listed. - ("NemotronVoiceChatCode2Wav", "code2wav"), - # audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1). - # The default deploy is async_chunk; this covers async_chunk: false. - ("AudexCode2Wav", "audex_code2wav"), - # audex TTA: tta thinker (Stage 0) -> XCodec1 (Stage 1, always sync - # full-payload — CNN codec decoded over the full sequence). - ("AudexXCodec1", "audex_xcodec"), - # indextts2 / indextts2_5: talker (Stage 0) -> s2mel decoder - # (Stage 1). Stage 1 consumes the complete mel/optional-latent payload. - ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"), - ("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"), - # dynin: token2text (Stage 0) -> token2image (Stage 1) -> - # token2audio (Stage 2). Producer wires via - # custom_process_next_stage_input_func: *_full_payload in deploy yaml. - ("DyninOmniForConditionalGeneration", "token2image"), - ("DyninOmniForConditionalGeneration", "token2audio"), - } -) - - def uses_full_payload_input_coordinator(model_config: Any) -> bool: """Returns True if this stage parks pending requests in WAITING_FOR_INPUT awaiting a full_payload delivery on the worker connector. - Gated by (model_arch, model_stage) — see _FULL_PAYLOAD_INPUT_STAGES for the - rationale on why this is a whitelist instead of a marker-driven structural - gate. + Gated by the topology-declared ``requires_full_payload_input`` capability on + downstream (stage_id > 0) stages, and only on the non-async-chunk path + (async-chunk stages are fed through the streamed connector instead). """ if getattr(model_config, "stage_id", 0) <= 0: return False if getattr(model_config, "async_chunk", False): return False - key = ( - getattr(model_config, "model_arch", None), - getattr(model_config, "model_stage", None), - ) - return key in _FULL_PAYLOAD_INPUT_STAGES + return bool(getattr(model_config, "requires_full_payload_input", False)) class OmniSchedulingCoordinator: diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index db2cf15ea..2ce1eed70 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Structured payload types for inter-stage communication. Categories under ``OmniPayload``: @@ -12,12 +12,16 @@ Categories under ``OmniPayload``: from __future__ import annotations +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, TypedDict import msgspec import numpy as np import torch +REQUEST_ARTIFACT_DIRS_KEY = "_omni_request_artifact_dirs" +TRANSFORM_OWNED_META_KEYS = frozenset({"minimax_h3_prepared_reference_videos"}) + if TYPE_CHECKING: from vllm_omni.engine import AdditionalInformationEntry, AdditionalInformationPayload @@ -55,6 +59,7 @@ class Ids(TypedDict, total=False): output: list[int] speech_token: list[int] prior_image: list[int] + streaming_prompt_previous_codes: list[int] class OmniPayloadMeta(TypedDict, total=False): @@ -74,7 +79,11 @@ class OmniPayloadMeta(TypedDict, total=False): override_keys: list[tuple[str, str]] num_processed_tokens: int next_stage_prompt_len: int + next_stage_generation_tokens: int replace_streaming_prompt: bool + streaming_prompt_recompute: bool + streaming_condition_seq: int + replace_runtime_additional_information: bool ar_width: int eol_token_id: int visual_token_start_id: int @@ -91,6 +100,13 @@ class OmniPayloadMeta(TypedDict, total=False): ref_context_included: bool talker_prefill_offset: int omni_final_stage_id: int + # Per-request audio seed. Stages that draw their own noise (flow-matching + # / diffusion decoders) need the producing stage's seed to stay + # reproducible, and the producing stage's SamplingParams do not travel + # with the payload. + audio_seed: int + token_role_ids: torch.Tensor + minimax_h3_prepared_reference_videos: str class OmniPayload(TypedDict, total=False): @@ -111,7 +127,9 @@ class OmniPayload(TypedDict, total=False): # ── msgspec.Struct mirror of the TypedDicts (runtime-validated) ── -class _StructBase(msgspec.Struct, omit_defaults=True, kw_only=True, forbid_unknown_fields=True): +class _StructBase( # type: ignore[call-arg] + msgspec.Struct, omit_defaults=True, kw_only=True, forbid_unknown_fields=True +): pass @@ -153,6 +171,7 @@ class IdsStruct(_StructBase): output: list[int] | None = None speech_token: list[int] | None = None prior_image: list[int] | None = None + streaming_prompt_previous_codes: list[int] | None = None class MetaStruct(_StructBase): @@ -169,7 +188,11 @@ class MetaStruct(_StructBase): override_keys: list[tuple[str, str]] | None = None num_processed_tokens: int | None = None next_stage_prompt_len: int | None = None + next_stage_generation_tokens: int | None = None replace_streaming_prompt: bool | None = None + streaming_prompt_recompute: bool | None = None + streaming_condition_seq: int | None = None + replace_runtime_additional_information: bool | None = None ar_width: int | None = None eol_token_id: int | None = None visual_token_start_id: int | None = None @@ -189,6 +212,9 @@ class MetaStruct(_StructBase): codec_left_context_frames: int | None = None code_flat_numel: int | None = None omni_final_stage_id: int | None = None + audio_seed: int | None = None + token_role_ids: torch.Tensor | None = None + minimax_h3_prepared_reference_videos: str | None = None class OmniPayloadStruct(_StructBase): @@ -302,7 +328,7 @@ def _dtype_to_name(dtype: torch.dtype) -> str: _NESTED_KEYS = frozenset({"hidden_states", "embed", "ids", "codes", "meta"}) -def flatten_payload(payload: dict[str, Any]) -> dict[str, Any]: +def flatten_payload(payload: Mapping[str, Any]) -> dict[str, Any]: """Flatten a nested ``OmniPayload`` to dotted keys. Nested sub-dicts under ``_NESTED_KEYS`` are expanded: @@ -353,17 +379,25 @@ def _serialize_tensor(t: torch.Tensor) -> AdditionalInformationEntry: t_cpu = t.detach().to("cpu").contiguous() return AdditionalInformationEntry( - tensor_data=t_cpu.numpy().tobytes(), + tensor_data=t_cpu.flatten().view(torch.uint8).numpy().tobytes(), tensor_shape=list(t_cpu.shape), tensor_dtype=_dtype_to_name(t_cpu.dtype), ) def _deserialize_tensor(entry: AdditionalInformationEntry) -> torch.Tensor: - dt = np.dtype(entry.tensor_dtype or "float32") - arr = np.frombuffer(entry.tensor_data, dtype=dt) # type: ignore[arg-type] - arr = arr.reshape(entry.tensor_shape) - return torch.from_numpy(arr.copy()) + dtype_name = entry.tensor_dtype or "float32" + dtype = getattr(torch, dtype_name, None) + if not isinstance(dtype, torch.dtype): + raise ValueError(f"Unsupported tensor dtype: {dtype_name}") + + if entry.tensor_shape is None: + raise ValueError("Tensor shape is required") + if not entry.tensor_data: + return torch.empty(0, dtype=dtype).reshape(entry.tensor_shape) + + data = torch.frombuffer(bytearray(entry.tensor_data), dtype=torch.uint8) + return data.view(dtype).reshape(entry.tensor_shape) def serialize_payload( diff --git a/vllm_omni/deploy/README_higgs_audio_v3.md b/vllm_omni/deploy/README_higgs_audio_v3.md index ea1363069..44f7f75c7 100644 --- a/vllm_omni/deploy/README_higgs_audio_v3.md +++ b/vllm_omni/deploy/README_higgs_audio_v3.md @@ -44,6 +44,65 @@ vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml ``` +## Why Stage 0 Pins Native FlashInfer + +The three Stage 0 profiles set `attention_backend: FLASHINFER` and also set: + +```yaml +attention_config: + use_trtllm_attention: false +``` + +`FLASHINFER` names the overall vLLM attention backend. Starting with vLLM +0.25, that backend can internally route Hopper decode to the TRT-LLM XQA +kernel. The XQA eligibility check answers whether the kernel can run, not +whether it is the fastest kernel for the current workload. In particular, the +decode builder can select XQA statically when all of the following hold: + +- the explicit `use_trtllm_attention` override is not `false`; +- the device is Hopper/SM90 and the XQA cubin is available; and +- the query-head count is divisible by the KV-head count. + +Higgs has 32 query heads and 8 KV heads, so it passes the head-layout check. +The static gate does not include a model dtype, KV-cache dtype, context-length, +batch-shape, or measured-performance comparison. This matters because kernel +support is not a performance guarantee. + +Higgs also has an unusual output shape that must not be confused with its +attention shape. One autoregressive audio step advances one token per request. +The resulting hidden row is then projected into eight codebook distributions +with shape `[1, 8, codebook_size]`. The eight codebooks therefore do not turn +the attention query into `q_len=8`; at concurrency one, XQA still receives a +single BF16 query row. At this small per-step shape, any XQA advantage for +other KV dtypes, batch sizes, or context lengths does not necessarily apply. +The difference is paid in every transformer layer and every generated audio +step, so even a small per-call disadvantage accumulates in TTFP and RTF. + +The upstream XQA integration also showed that the result is dtype-sensitive: +its FP8 benchmark favored XQA, while its BF16 benchmark reported 21.16 ms mean +TPOT for XQA versus 19.47 ms for FA3. That is not a direct native-FlashInfer +comparison, but it demonstrates why SM90 compatibility alone is insufficient +to choose the fastest BF16 kernel. See the +[vLLM XQA benchmark](https://github.com/vllm-project/vllm/pull/43232#issuecomment-4501176335). + +The end-to-end H800 A/B in +[issue #5584](https://github.com/vllm-project/vllm-omni/issues/5584#issuecomment-5188700218) +isolated the routing decision on the same vLLM 0.25 stack: + +| Decode path | TTFP (ms) | RTF | +|---|---:|---:| +| XQA (automatic default) | 175.76 | 0.3405 | +| Native FlashInfer (`use_trtllm_attention: false`) | 166.39 | 0.3233 | +| vLLM 0.24 reference | 166.68 | 0.3209 | + +This proves that XQA selection causes the observed local regression and that +native FlashInfer recovers the pre-v0.25 range. It does not by itself identify +one defective XQA instruction or memory transaction. Establishing that deeper +kernel-level cause requires an attention-only microbenchmark and GPU profiler +trace for the exact Higgs batch/context distribution. Until vLLM has a +workload-aware selector or XQA is faster for this BF16 shape, the model-level +pin is the deterministic choice. + ## Notes - Stage 1 remains `enforce_eager: true` in both profiles. diff --git a/vllm_omni/deploy/bagel_single_stage.yaml b/vllm_omni/deploy/bagel_single_stage.yaml index bcfbad253..fb0872af8 100644 --- a/vllm_omni/deploy/bagel_single_stage.yaml +++ b/vllm_omni/deploy/bagel_single_stage.yaml @@ -9,7 +9,7 @@ # --deploy-config vllm_omni/deploy/bagel_single_stage.yaml # # Or programmatically: -# Omni(model="...", deploy_config_path="vllm_omni/deploy/bagel_single_stage.yaml") +# Omni(model="...", deploy_config="vllm_omni/deploy/bagel_single_stage.yaml") pipeline: bagel_single_stage async_chunk: false diff --git a/vllm_omni/deploy/dots_tts.yaml b/vllm_omni/deploy/dots_tts.yaml new file mode 100644 index 000000000..acd28de52 --- /dev/null +++ b/vllm_omni/deploy/dots_tts.yaml @@ -0,0 +1,28 @@ +# dots.tts deployment config (single-stage AR talker). +# +# Conservative first-integration settings: eager execution, no CUDA +# Graphs, no batched side-path optimizations yet — see voxcpm2.yaml for +# the runtime knobs a mature deployment of this model family grows. +# +# Verified locally on RTX 5080 16GB. + +async_chunk: false +dtype: bfloat16 + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.8 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 4096 + repetition_penalty: 1.0 diff --git a/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml index 360b72974..b6e5bc7fe 100644 --- a/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml +++ b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml @@ -11,7 +11,7 @@ # # Or programmatically: # Omni(model="fishaudio/s2-pro", -# deploy_config_path="vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml") +# deploy_config="vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml") async_chunk: true enable_chunked_prefill: false diff --git a/vllm_omni/deploy/gepard.yaml b/vllm_omni/deploy/gepard.yaml new file mode 100644 index 000000000..26000d616 --- /dev/null +++ b/vllm_omni/deploy/gepard.yaml @@ -0,0 +1,40 @@ +# Gepard-1.0: single-stage native-AR FSQ/NanoCodec TTS (~0.5B). +# Target 1x L4 24GB. Zero-shot only; enforce_eager until the perf PR. +# +# The pipeline pin is required, not cosmetic: the checkpoint's config.json +# self-identifies as qwen3_5_text, so model_type inference cannot reach the +# gepard registry entry, and the architectures fallback would match +# Qwen3_5ForCausalLM against the diffusion pipeline. +pipeline: gepard +async_chunk: false +dtype: bfloat16 +trust_remote_code: true + +stages: + - stage_id: 0 + max_num_seqs: 4 + gpu_memory_utilization: 0.3 # AR LM + NanoCodec fit in ~2 GiB + enforce_eager: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + # Override architectures so vLLM resolves our talker rather than the native + # Qwen3.5 text model. Do NOT override model_type: the backbone's own + # model_type check needs it to stay qwen3_5_text. + hf_overrides: + architectures: ["GepardTalkerForConditionalGeneration"] + default_sampling_params: + # All 32 heads are sampled in-model, so vLLM's sampler is pinned + # deterministic and just argmaxes what forward chose. + temperature: 0.0 + top_p: 1.0 + top_k: -1 + # One output token is one frame, so this is the frame budget: 1000 frames + # is ~46 s at 21.5 fps. A backstop, not a target -- normal utterances are + # 40-420 frames -- but a load-bearing one, because this checkpoint's stop + # head does not fire on every batched trajectory. + max_tokens: 1000 + seed: 42 + detokenize: false diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml index da299415b..41e1a9677 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml @@ -38,6 +38,11 @@ stages: gpu_memory_utilization: 0.6 enforce_eager: true attention_backend: FLASHINFER + # Eight codebooks share one attention row, so decode remains BF16 q_len=1 + # per request. vLLM's SM90 XQA capability gate does not benchmark this + # shape; keep the faster native path (README_higgs_audio_v3.md, #5584). + attention_config: + use_trtllm_attention: false trust_remote_code: true # Stage-0 prefix caching is on by default. The talker's HS opt-out and # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) @@ -78,3 +83,10 @@ stages: max_tokens: 65536 seed: 42 repetition_penalty: 1.0 + +platforms: + rocm: + stages: + - stage_id: 0 + # FlashInfer is not supported on ROCm. + attention_backend: TRITON_ATTN diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml index 8a751483c..bb37bdb4c 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml @@ -37,6 +37,11 @@ stages: gpu_memory_utilization: 0.6 enforce_eager: true attention_backend: FLASHINFER + # Eight codebooks share one attention row, so decode remains BF16 q_len=1 + # per request. vLLM's SM90 XQA capability gate does not benchmark this + # shape; keep the faster native path (README_higgs_audio_v3.md, #5584). + attention_config: + use_trtllm_attention: false trust_remote_code: true # Stage-0 prefix caching is on by default. The talker's HS opt-out and # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml index 0d6ae8218..13c6961e6 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml @@ -43,6 +43,11 @@ stages: gpu_memory_utilization: 0.6 enforce_eager: false attention_backend: FLASHINFER + # Eight codebooks share one attention row, so decode remains BF16 q_len=1 + # per request. vLLM's SM90 XQA capability gate does not benchmark this + # shape; keep the faster native path (README_higgs_audio_v3.md, #5584). + attention_config: + use_trtllm_attention: false trust_remote_code: true enable_prefix_caching: true async_scheduling: false diff --git a/vllm_omni/deploy/mimo_audio_5090d.yaml b/vllm_omni/deploy/mimo_audio_5090d.yaml new file mode 100644 index 000000000..5e6b7b035 --- /dev/null +++ b/vllm_omni/deploy/mimo_audio_5090d.yaml @@ -0,0 +1,70 @@ +# MiMo-Audio deploy for 1x RTX 5090 / 5090D 32GB (Blackwell, SM 12.x). +# +# Verified offline TTS (`tts_sft`) with async-chunk streaming on a single GPU. +# Stage 0 uses TRITON_ATTN to avoid FlashAttention-2 PTX/toolchain issues on +# Blackwell. Stage 1 keeps a small VRAM budget; pair with +# MIMO_AUDIO_TOKENIZER_DEVICE=cpu in the recipe when tokenizer VRAM is tight. +# +# Default async-chunk layout matches vllm_omni/deploy/mimo_audio.yaml (both +# stages on device 0). For the legacy 2-GPU sync pipeline, see the header +# comment in mimo_audio.yaml. +async_chunk: true +dtype: bfloat16 + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.001 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 30 + codec_left_context_frames: 40 + +stages: + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.78 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + attention_backend: TRITON_ATTN + enable_flashinfer_autotune: false + skip_mm_profiling: true + mm_processor_cache_gb: 0 + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.6 + top_p: 0.95 + top_k: 50 + max_tokens: 18192 + seed: 42 + repetition_penalty: 1.1 + + - stage_id: 1 + max_num_seqs: 1 + gpu_memory_utilization: 0.12 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + attention_backend: TRITON_ATTN + enable_flashinfer_autotune: false + skip_mm_profiling: true + mm_processor_cache_gb: 0 + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 18192 + seed: 42 diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index bc34bfe9f..023101223 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -1,6 +1,15 @@ # MiniCPM-o 4.5 three-stage deploy on one large-memory GPU. pipeline: minicpmo_4_5 async_chunk: true +# Native full-duplex shares this profile with /v1/chat/completions. Every other +# duplex_session knob keeps its runtime default. +session_mode: duplex +# Keep the active window aligned with the session and stage capacities below. +# A smaller window serializes first-packet generation for concurrent chat turns. +active_stream_window: 4 +duplex_session: + # Matches Stage 0/1 max_num_seqs below. + max_sessions: 4 trust_remote_code: true enable_prefix_caching: false @@ -32,7 +41,7 @@ stages: max_num_batched_tokens: 16384 devices: "0" # Required for MiniCPM interleaved image/audio Daily-Omni packs (~78% recipe). - # Keep these in YAML so they survive both --deploy-config and --stage-configs-path. + # Keep these in YAML so they are applied through --deploy-config. interleave_mm_strings: true media_io_kwargs: video: @@ -55,6 +64,13 @@ stages: - stage_id: 1 max_num_seqs: 4 + # Native duplex uses MiniCPMTTS.generate_with_buffer's supported + # sliding_recompute policy: retain one completed condition/codec chunk + # when the next condition arrives. Other deploy profiles continue to use + # the checkpoint's full_attention policy. + hf_overrides: + tts_config: + attention_type: sliding_recompute # Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace. gpu_memory_utilization: 0.15 enforce_eager: false @@ -64,15 +80,20 @@ stages: devices: "0" output_connectors: to_stage_2: connector_of_shared_memory - # Codec sampling comes from the checkpoint's tts_config. These parameters - # govern only vLLM's binary continue/stop token. + # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler + # (head_code logits). EOS is tts_config.eos_token_id (num_audio_tokens-1). + # Values track upstream utils.TTSSamplingParams (min_p/win_size/tau_r are + # declared there but never read, so they stay out). repetition_penalty is + # consumed by the Talker, which rescores it over upstream's 16-frame window + # instead of the Sampler's whole-stream presence penalty. + # min_tokens is upstream's min_new_token. default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.85 + top_k: 25 + repetition_penalty: 1.05 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false @@ -111,3 +132,9 @@ platforms: max_num_batched_tokens: 8192 compilation_config: cudagraph_mode: PIECEWISE + - stage_id: 2 + # Stage 2 stays outer-eager and captures only exact-shape CFM DiT + # estimator kernels through the NPU platform patch. + additional_config: + code2wav_enable_npu_graph: true + code2wav_max_npu_graphs: 32 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 2a4abed42..9c3818813 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -4,6 +4,15 @@ # isolated from codec/vocoder contention while preserving stage-2 batching. pipeline: minicpmo_4_5 async_chunk: true +# Native full-duplex shares this profile with /v1/chat/completions. Every other +# duplex_session knob keeps its runtime default. +session_mode: duplex +# Keep the active window aligned with the session and stage capacities below. +# A smaller window serializes first-packet generation for concurrent chat turns. +active_stream_window: 4 +duplex_session: + # Matches Stage 0/1 max_num_seqs below. + max_sessions: 4 trust_remote_code: true enable_prefix_caching: false @@ -52,16 +61,14 @@ stages: devices: "1" output_connectors: to_stage_2: connector_of_shared_memory - # Codec sampling is model-owned by the checkpoint's tts_config (defaults: - # T=0.8, top-k=25, top-p=0.85, repetition penalty=1.05, seed=42). - # These parameters only govern vLLM's binary continue/stop token. default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.85 + top_k: 25 + # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window. + repetition_penalty: 1.05 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false @@ -93,3 +100,9 @@ platforms: max_num_batched_tokens: 8192 compilation_config: cudagraph_mode: PIECEWISE + - stage_id: 2 + # Stage 2 stays outer-eager and captures only exact-shape CFM DiT + # estimator kernels through the NPU platform patch. + additional_config: + code2wav_enable_npu_graph: true + code2wav_max_npu_graphs: 32 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index b34673910..5eddbb6f4 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -4,6 +4,15 @@ # Code2Wav run independently on GPUs 0, 1, and 2. pipeline: minicpmo_4_5 async_chunk: true +# Native full-duplex shares this profile with /v1/chat/completions. Every other +# duplex_session knob keeps its runtime default. +session_mode: duplex +# Keep the active window aligned with the session and stage capacities below. +# A smaller window serializes first-packet generation for concurrent chat turns. +active_stream_window: 4 +duplex_session: + # Matches Stage 0/1 max_num_seqs below. + max_sessions: 4 trust_remote_code: true enable_prefix_caching: false @@ -57,15 +66,14 @@ stages: devices: "1" output_connectors: to_stage_2: connector_of_shared_memory - # vLLM samples the binary continue/stop row; codec sampling comes from - # the checkpoint's tts_config and currently uses a deterministic seed. default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.85 + top_k: 25 + # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window. + repetition_penalty: 1.05 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false @@ -86,4 +94,20 @@ stages: max_tokens: 65536 detokenize: true -platforms: {} +platforms: + npu: + stages: + - stage_id: 0 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE + - stage_id: 1 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE + - stage_id: 2 + # Stage 2 stays outer-eager and captures only exact-shape CFM DiT + # estimator kernels through the NPU platform patch. + additional_config: + code2wav_enable_npu_graph: true + code2wav_max_npu_graphs: 32 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml index 7a0513167..806f3330e 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml @@ -5,7 +5,6 @@ # Stage 1: replicas on GPU 1 and 2 base_config: minicpmo_4_5.yaml pipeline: minicpmo_4_5 -session_mode: duplex stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml index da28456fe..16d9f79da 100644 --- a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml @@ -5,7 +5,6 @@ # Stage 1: replicas on GPU 1, 2, and 3 base_config: minicpmo_4_5.yaml pipeline: minicpmo_4_5 -session_mode: duplex stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 813a2e2eb..eacc069d1 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -15,6 +15,10 @@ # so 4096 is the empirically-verified safe ceiling. pipeline: minicpmo_4_5 async_chunk: true +# Native full-duplex shares this profile with /v1/chat/completions. The default +# duplex_session caps sessions at 1, matching Stage 0/1 max_num_seqs below. +session_mode: duplex +active_stream_window: 1 trust_remote_code: true enable_prefix_caching: false @@ -69,15 +73,14 @@ stages: devices: "4" output_connectors: to_stage_2: connector_of_shared_memory - # vLLM samples the binary continue/stop row; codec sampling comes from - # the checkpoint's tts_config and currently uses a deterministic seed. default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.85 + top_k: 25 + # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window. + repetition_penalty: 1.05 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml index f4ad2c6b4..f377ac040 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml @@ -5,7 +5,6 @@ # Stage 1: replicas on GPU 4, 5, 6, and 7 base_config: minicpmo_4_5_8x4090.yaml pipeline: minicpmo_4_5 -session_mode: duplex stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml deleted file mode 100644 index df0e40be9..000000000 --- a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# MiniCPM-o 4.5 native duplex deployment. -# -# The registered minicpmo_4_5 pipeline owns the fixed -# Thinker -> Talker -> Code2Wav topology. This overlay changes session -# lifetime and bounds all three stages to two live duplex sessions. -base_config: minicpmo_4_5.yaml -pipeline: minicpmo_4_5 -session_mode: duplex -active_stream_window: 1 -duplex_session: - idle_ttl_s: 300 - disconnect_grace_s: 30 - reaper_interval_s: 5 - resume_replay_ttl_s: 60 - resume_replay_max_bytes_per_session: 8388608 - max_pending_input_bytes_per_session: 16777216 - max_pending_turns_per_session: 4 - # Matches the configured Stage0/Stage1 sequence capacity below. - max_sessions: 2 - completed_append_cache_size: 256 - -stages: - - stage_id: 0 - max_num_seqs: 2 - async_scheduling: false - - - stage_id: 1 - max_num_seqs: 2 - async_scheduling: false - # Talker context is 4096 (tts_config.max_position_embeddings). - max_model_len: 4096 - # The split Talker exposes a binary continue/stop vocabulary; native - # segment and turn boundaries travel as explicit output metadata. - default_sampling_params: - max_tokens: 2048 - min_tokens: 0 - stop_token_ids: [1] - - - stage_id: 2 - max_num_seqs: 2 - async_scheduling: false diff --git a/vllm_omni/deploy/minimax_h3_disaggregated.yaml b/vllm_omni/deploy/minimax_h3_disaggregated.yaml new file mode 100644 index 000000000..1c1ee6c40 --- /dev/null +++ b/vllm_omni/deploy/minimax_h3_disaggregated.yaml @@ -0,0 +1,28 @@ +# MiniMax H3: vLLM Qwen3-VL text encoder + diffusion video/audio generation. +pipeline: minimax_h3_disaggregated +async_chunk: false +trust_remote_code: true +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0,1" + max_num_seqs: 1 + tensor_parallel_size: 2 + default_sampling_params: + max_tokens: 1 + temperature: 0.0 + detokenize: false + + - stage_id: 1 + devices: "2,3,4,5" + max_num_seqs: 1 + model_loaded: + text_encoder: false + parallel_config: + tensor_parallel_size: 1 + ulysses_degree: 4 + vae_patch_parallel_size: 4 + default_sampling_params: + seed: 42 + num_inference_steps: 50 diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml new file mode 100644 index 000000000..d7e84db42 --- /dev/null +++ b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml @@ -0,0 +1,31 @@ +# MiniMax H3 Turbo: five-sigma defaults for the disaggregated topology. +pipeline: minimax_h3_disaggregated +async_chunk: false +trust_remote_code: true +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0,1" + max_num_seqs: 1 + tensor_parallel_size: 2 + default_sampling_params: + max_tokens: 1 + temperature: 0.0 + detokenize: false + + - stage_id: 1 + devices: "2,3,4,5" + max_num_seqs: 1 + model_loaded: + text_encoder: false + parallel_config: + tensor_parallel_size: 1 + ulysses_degree: 4 + vae_patch_parallel_size: 4 + default_sampling_params: + seed: 42 + num_inference_steps: 5 + extra_args: + flow_shift: 6.0 + audio_flow_shift: 3.0 diff --git a/vllm_omni/deploy/minimax_music3.yaml b/vllm_omni/deploy/minimax_music3.yaml new file mode 100644 index 000000000..6e8a13d9d --- /dev/null +++ b/vllm_omni/deploy/minimax_music3.yaml @@ -0,0 +1,130 @@ +# MiniMax-Music3 deploy: Stage 0 (AR talker) -> Stage 1 (acoustic decoder), +# both on one 140 GB card, streaming 200-frame windows over shared memory. +# +# Auto-discovered default for model_type=minimax_music3. +# +# Classifier-free guidance is mandatory for this checkpoint: every request +# decodes as a cond/uncond pair in stage 0, so stage-0 KV, max_num_seqs and +# prefill budget are all sized for TWO rows per user request. The stage-0 +# extra_args block below is what turns on has_sampling_extra_args (so per-row +# sampling extra args reach the model's forward) and is also how the CFG +# scheduler patches learn this stage is paired. +pipeline: minimax_music3 +async_chunk: true + +# The repo-root config.json reports model_type=minimax_music3, which upstream +# transformers does not ship a config class for. +trust_remote_code: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + # Stream stereo PCM back to the client as stage 1 finishes each window. + codec_streaming: true + # Window geometry is fixed by the checkpoint (200-frame windows on a + # 100-frame hop, models/minimax_music3/constants.py), so there are no + # codec_chunk_frames knobs here: changing the window changes the audio. + shm_threshold_bytes: 65536 + connector_get_sleep_s: 0.01 + # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer + # for a chunk than a per-step codec stage would. First chunk additionally + # covers stage-0 prefill of both CFG rows. + connector_get_max_wait_first_chunk: 30000 + connector_get_max_wait: 10000 + +stages: + - stage_id: 0 + # 16 rows = 8 concurrent guided requests, since CFG doubles every request. + max_num_seqs: 16 + # ~16 GB of bf16 Qwen3-8B weights plus KV for 16 full-length rows + # (36 layers x 8 KV heads x 128 head_dim ~= 144 KB/token, so ~1.4 GB per + # 10240-token row) fits comfortably inside 0.6 of a 140 GB card. + gpu_memory_utilization: 0.6 + dtype: bfloat16 + # The checkpoint's max_position_embeddings is 10240 and the whole song is + # decoded in-context. Note this cannot hold a max-length prompt AND a + # max-length song at the same time: the prompt cap is 5000 tokens and the + # frame budget is 9000, which overruns 10240. That is a property of the + # checkpoint, not of this file; long captions shorten the maximum song. + max_model_len: 10240 + # One scheduler step must be able to prefill both rows of a pair + # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked. + max_num_batched_tokens: 10240 + # Two rows at 8B is a launch-bound decode step, so the backbone is + # captured. FULL_DECODE_ONLY only ever dispatches a graph for a uniform + # decode batch; prefill and mixed steps stay eager, which is what the + # talker's frame-feedback fallback is for. The talker's own sampler and + # depth decoder run after the graph, not inside it. + # + # What makes the forward capturable: the frame feedback writes a + # fixed-size slab in prepare_runner_inputs (outside the captured region) + # and the forward selects over it with torch.where, instead of an + # index_copy against an index tensor whose address and length changed + # every step. The row-to-request binding the sampler reads is likewise + # rebuilt in prepare_runner_inputs, because a replayed step never runs + # the forward's Python body. + enforce_eager: false + compilation_config: + # No compilation. The talker walks the backbone's decoder layers itself, + # so the compiled Qwen3Model.forward is never entered and inductor has + # nothing to do; what the mode still decides is which kernels vLLM + # defaults to. Left at its non-eager default it picks the elementwise + # PyTorch norms over the fused vllm_c ones, which is both slower and + # numerically different: measured here, that alone moved the sampled + # codes off the eager trajectory by frame 4. Graph capture is + # independent of this and still happens. + mode: 0 + # 16 rows is the stage ceiling, i.e. 8 guided requests. + cudagraph_capture_sizes: [1, 2, 4, 8, 16] + cudagraph_mode: FULL_DECODE_ONLY + cudagraph_num_of_warmups: 1 + # Belt and braces on the same point: the non-eager default is "none", + # which hands the norm and activation ops to elementwise PyTorch. + custom_ops: ["all"] + # A CFG pair must be sampled inside one engine step for the logits blend + # and the cond->uncond token copy to line up. + async_scheduling: false + # Off during bring-up: the uncond row is a long run of identical + # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than + # its cond partner, which is exactly the num_computed_tokens skew the CFG + # pair patches then have to undo. + enable_prefix_caching: false + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + # The frame budget is enforced by the model. This ceiling only has to + # leave room for the extra step that delivers the final window. + max_tokens: 9002 + extra_args: + cfg_role: cond + cfg_scale: 1.5 + + - stage_id: 1 + # One row per user request: the uncond companion is stage-0-final and + # never reaches the acoustic stage. + max_num_seqs: 8 + # No KV cache on this stage; the budget only covers the DiT and DAV + # weights plus per-window latents and activations. + gpu_memory_utilization: 0.2 + # The flow-matching solver degrades measurably in bfloat16 + # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py). + dtype: float32 + enforce_eager: true + # Async scheduling builds per-request payloads against the NEXT step's + # input batch, which breaks when chunk arrival staggers requests. + async_scheduling: false + enable_prefix_caching: false + # One placeholder prompt token per arriving window, so a six-minute song + # (90 windows) needs a tiny context. + max_model_len: 4096 + max_num_batched_tokens: 4096 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + max_tokens: 4096 + seed: 0 diff --git a/vllm_omni/deploy/minimax_music3_2gpu.yaml b/vllm_omni/deploy/minimax_music3_2gpu.yaml new file mode 100644 index 000000000..8c463e0e7 --- /dev/null +++ b/vllm_omni/deploy/minimax_music3_2gpu.yaml @@ -0,0 +1,97 @@ +# MiniMax-Music3 two-GPU layout: AR talker on device 0, acoustic decoder on +# device 1. +# +# Splitting the stages removes the memory-budget compromise of the single-card +# file: the talker's CFG-doubled KV no longer competes with the DiT and DAV +# weights, and the flow-matching solver no longer competes with AR decode for +# SMs. Everything else (window geometry, CFG contract, sampling defaults) +# matches minimax_music3.yaml, so the two files stay behaviorally identical. +pipeline: minimax_music3 +async_chunk: true + +# The repo-root config.json reports model_type=minimax_music3, which upstream +# transformers does not ship a config class for. +trust_remote_code: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + # Stream stereo PCM back to the client as stage 1 finishes each window. + codec_streaming: true + # Window geometry is fixed by the checkpoint (200-frame windows on a + # 100-frame hop, models/minimax_music3/constants.py), so there are no + # codec_chunk_frames knobs here: changing the window changes the audio. + shm_threshold_bytes: 65536 + connector_get_sleep_s: 0.01 + # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer + # for a chunk than a per-step codec stage would. First chunk additionally + # covers stage-0 prefill of both CFG rows. + connector_get_max_wait_first_chunk: 30000 + connector_get_max_wait: 10000 + +stages: + - stage_id: 0 + # 32 rows = 16 concurrent guided requests, since CFG doubles every request. + # Twice the single-card figure: this stage now owns its whole device. + max_num_seqs: 32 + gpu_memory_utilization: 0.85 + dtype: bfloat16 + # The checkpoint's max_position_embeddings is 10240 and the whole song is + # decoded in-context. Note this cannot hold a max-length prompt AND a + # max-length song at the same time: the prompt cap is 5000 tokens and the + # frame budget is 9000, which overruns 10240. That is a property of the + # checkpoint, not of this file; long captions shorten the maximum song. + max_model_len: 10240 + # One scheduler step must be able to prefill both rows of a pair + # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked. + max_num_batched_tokens: 10240 + # The talker owns its sampler and drives a depth decoder per AR step; + # leave graph capture off until that path is verified graph-safe. + enforce_eager: true + # A CFG pair must be sampled inside one engine step for the logits blend + # and the cond->uncond token copy to line up. + async_scheduling: false + # Off during bring-up: the uncond row is a long run of identical + # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than + # its cond partner, which is exactly the num_computed_tokens skew the CFG + # pair patches then have to undo. + enable_prefix_caching: false + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + # The frame budget is enforced by the model. This ceiling only has to + # leave room for the extra step that delivers the final window. + max_tokens: 9002 + extra_args: + cfg_role: cond + cfg_scale: 1.5 + + - stage_id: 1 + # One row per user request: the uncond companion is stage-0-final and + # never reaches the acoustic stage. + max_num_seqs: 16 + # No KV cache on this stage, but it now has a device to itself, so the + # solver can hold more windows in flight than the shared-card profile. + gpu_memory_utilization: 0.6 + # The flow-matching solver degrades measurably in bfloat16 + # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py). + dtype: float32 + enforce_eager: true + # Async scheduling builds per-request payloads against the NEXT step's + # input batch, which breaks when chunk arrival staggers requests. + async_scheduling: false + enable_prefix_caching: false + # One placeholder prompt token per arriving window, so a six-minute song + # (90 windows) needs a tiny context. + max_model_len: 4096 + max_num_batched_tokens: 4096 + devices: "1" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + max_tokens: 4096 + seed: 0 diff --git a/vllm_omni/deploy/moss_tts_nano.yaml b/vllm_omni/deploy/moss_tts_nano.yaml index 2c8fc54c0..e9eab184b 100644 --- a/vllm_omni/deploy/moss_tts_nano.yaml +++ b/vllm_omni/deploy/moss_tts_nano.yaml @@ -6,8 +6,8 @@ # * enforce_eager=true — the model uses upstream inference_stream() with # trust_remote_code; CUDA graph capture is not wired up. # * gpu_memory_utilization=0.3 — the AR LM + codec fit in ~2 GiB. -# * max_num_seqs=4 — mirrors the legacy stage_configs contract; per-request -# streaming generators are stored in self._stream_gens. +# * max_num_seqs=4 — supports per-request streaming generators stored in +# self._stream_gens. # # MOSS-TTS-Nano has no multimodal input preprocessing, so skip_mm_profiling # bypasses the dummy-MM pass at profile time. diff --git a/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml b/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml new file mode 100644 index 000000000..1a69b24e6 --- /dev/null +++ b/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml @@ -0,0 +1,47 @@ +# Nemotron VoiceChat native Full-Duplex deployment. +# +# One scheduler segment consumes one 1280-sample (80 ms) microphone frame. +# Stage 0 remains one resumable request, so its NemotronH KV/Mamba state is +# preserved between frame appends. Stage 1 and the codec overlap per frame. +base_config: nemotron_labs_voicechat_streaming.yaml +pipeline: nemotron_voicechat +session_mode: duplex +active_stream_window: 1 + +duplex_session: + idle_ttl_s: 300 + disconnect_grace_s: 30 + reaper_interval_s: 5 + resume_replay_ttl_s: 60 + resume_replay_max_bytes_per_session: 8388608 + max_pending_input_bytes_per_session: 16777216 + max_pending_turns_per_session: 4 + max_sessions: 1 + completed_append_cache_size: 256 + +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_chunk_frames: 1 + +stages: + - stage_id: 0 + max_num_seqs: 1 + async_scheduling: false + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + ignore_eos: true + seed: 0 + + - stage_id: 1 + max_num_seqs: 1 + async_scheduling: false + + - stage_id: 2 + max_num_seqs: 1 + async_scheduling: false diff --git a/vllm_omni/deploy/pi0.yaml b/vllm_omni/deploy/pi0.yaml new file mode 100644 index 000000000..8d93855d8 --- /dev/null +++ b/vllm_omni/deploy/pi0.yaml @@ -0,0 +1,52 @@ +# π0 (Pi-Zero) VLA deploy: single diffusion stage. +# +# Topology (PI0_PIPELINE) is declared in vllm_omni/diffusion/models/pi0/pipeline_pi0.py. +# π0 takes multi-camera images + a language instruction + robot state and emits +# a continuous action chunk [chunk_size, max_action_dim] via flow matching. +# Actions are exposed through the OpenPI realtime serving API +# (/v1/realtime/robot/openpi); they ride in multimodal_output["actions"]. +# +# Default: one GPU, TP=1, eager (flow matching needs no CUDA graph), bsz 1 +# (OpenPI policy serving is one observation -> one action reply). float32 to +# match the LeRobot pi0_base reference numerics. + +pipeline: pi0 +async_chunk: false +distributed_executor_backend: mp +dtype: float32 + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + model_class_name: Pi0Pipeline + model_config: + paligemma_variant: gemma_2b + action_expert_variant: gemma_300m + chunk_size: 50 + max_action_dim: 32 + max_state_dim: 32 + num_inference_steps: 10 + image_resolution: [224, 224] + tokenizer_max_length: 48 + max_cameras: 3 + # Camera identities (the order the model attends to). Serving maps raw obs + # keys onto these via image_key_map; leave image_key_map empty to match + # incoming keys verbatim. + image_feature_keys: + - observation.images.base_0_rgb + - observation.images.left_wrist_0_rgb + - observation.images.right_wrist_0_rgb + image_key_map: {} + # OpenPI handshake metadata advertised to the robot client. + policy_server_config: + image_resolution: [224, 224] + n_external_cameras: 1 + needs_wrist_camera: true + needs_stereo_camera: false + needs_session_id: false + action_space: joint_position + action_horizon: 50 + action_dim: 32 + max_cameras: 3 diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 76878804d..59378d06b 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -69,6 +69,15 @@ stages: repetition_penalty: 1.1 platforms: + # vLLM 0.27's compiled native MRoPE path can vertically fuse the thinker's + # Q/K RMSNorm reductions into the rotary operations and materially change + # Qwen3-Omni audio output (#6090). Keep MRoPE on its CUDA path instead. + cuda: + stages: + - stage_id: 0 + compilation_config: + custom_ops: ["+rotary_embedding"] + # The supported Qwen3-Omni ModelOpt checkpoint quantizes the thinker only. Clear the # checkpoint-level metadata for the BF16 audio stages on the verified MUSA # profile before vLLM performs its normal quantization auto-detection. diff --git a/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml b/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml index 7ccc64750..d0b3d5ac7 100644 --- a/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml @@ -1,6 +1,8 @@ -# Deploy config for Qwen3-Omni-MoE thinker-only checkpoints. +# Deploy config for Qwen3-Omni-MoE thinker-only (text output, no talker/code2wav). +# Works for Captioner/Thinking checkpoints and for Instruct weights when this +# YAML is passed via --deploy-config (pipeline key forces the single-stage topology). # Verified on 2x H100-80G GPUs. -pipeline: qwen3_omni_moe +pipeline: qwen3_omni_moe_thinker_only distributed_executor_backend: mp enable_prefix_caching: true diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 08531869a..747b68852 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -33,6 +33,17 @@ connectors: # initial_codec_chunk_frames (including the per-request API field). # Example: # codec_chunk_ramp: [4, 4, 8, 16, 25] + # Adaptive chunk ramp (opt-in, takes precedence over fixed codec_chunk_ramp): + # Chunk 0 uses dynamic IC; chunk 1+ sizes from live buffer-depth feedback. + # Self-tunes across hardware (NPU/NVIDIA) and load conditions. + # codec_chunk_adaptive: true + # codec_chunk_min_frames: 2 + # Code2Wav decode + transfer margin + # codec_chunk_safety_margin_ms: 50 + # EWMA->delta divisor (delta=max(delta_min, int(ewma/divisor))); higher=gentler ramp + # codec_chunk_ramp_divisor: 15 + # min chunk-size growth per emit + # codec_chunk_ramp_delta_min: 2 # Decoder CUDA Graph sizing: # * async_chunk=true derives all stateful ICL/xvec shapes from # codec_chunk_ramp when enabled, otherwise from @@ -73,6 +84,10 @@ stages: # audio (HTTP 400). See #4962. min_tokens: 2 repetition_penalty: 1.05 + # Suppress codebook-0 silence tokens for the first N decode frames of a + # request (#4966). Applies to every request on this stage, so intended + # for x-vector-only (Base voice-clone) deploys. 0 = off. + silence_ban_frames: 0 subtalker_sampling_params: do_sample: true temperature: 0.9 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index ff086beb5..7a73dd50a 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -38,6 +38,17 @@ connectors: # initial_codec_chunk_frames (including the per-request API field). # Example: # codec_chunk_ramp: [4, 4, 8, 16, 25] + # Adaptive chunk ramp (opt-in, takes precedence over fixed codec_chunk_ramp): + # Chunk 0 uses dynamic IC; chunk 1+ sizes from live buffer-depth feedback. + # Self-tunes across hardware (NPU/NVIDIA) and load conditions. + # codec_chunk_adaptive: true + # codec_chunk_min_frames: 2 + # Code2Wav decode + transfer margin + # codec_chunk_safety_margin_ms: 50 + # EWMA->delta divisor (delta=max(delta_min, int(ewma/divisor))); higher=gentler ramp + # codec_chunk_ramp_divisor: 15 + # min chunk-size growth per emit + # codec_chunk_ramp_delta_min: 2 # async_chunk derives stateful ICL/xvec graph shapes from the ramp when # enabled, otherwise from the fixed chunk settings above, and ignores # decode_cudagraph_capture_sizes. Keep B>1 diff --git a/vllm_omni/deploy/soulxsinger_svc.yaml b/vllm_omni/deploy/soulxsinger_svc.yaml deleted file mode 100644 index a9d815806..000000000 --- a/vllm_omni/deploy/soulxsinger_svc.yaml +++ /dev/null @@ -1,19 +0,0 @@ -# SoulX-Singer SVC: single-stage DiT (preprocess inline in pre_process_func). -# -# Use: vllm serve --omni -pipeline: soulxsinger_svc -async_chunk: false -dtype: float16 -trust_remote_code: true -distributed_executor_backend: "mp" - -stages: - - stage_id: 0 - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - enforce_eager: false - devices: "0" - default_sampling_params: - num_inference_steps: 32 - guidance_scale: 3.0 - seed: 42 diff --git a/vllm_omni/deploy/soulxsinger_svs.yaml b/vllm_omni/deploy/soulxsinger_svs.yaml deleted file mode 100644 index 3f021cb53..000000000 --- a/vllm_omni/deploy/soulxsinger_svs.yaml +++ /dev/null @@ -1,19 +0,0 @@ -# SoulX-Singer SVS: single-stage DiT (preprocess inline in pre_process_func). -# -# Use: vllm serve --omni -pipeline: soulxsinger_svs -async_chunk: false -dtype: float16 -trust_remote_code: true -distributed_executor_backend: "mp" - -stages: - - stage_id: 0 - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - enforce_eager: false - devices: "0" - default_sampling_params: - num_inference_steps: 32 - guidance_scale: 3.0 - seed: 42 diff --git a/vllm_omni/deploy/wan2_2_ti2v.yaml b/vllm_omni/deploy/wan2_2_ti2v.yaml index bdb5cfecb..0d0a9124c 100644 --- a/vllm_omni/deploy/wan2_2_ti2v.yaml +++ b/vllm_omni/deploy/wan2_2_ti2v.yaml @@ -8,7 +8,7 @@ distributed_executor_backend: mp stages: - stage_id: 0 devices: "0" - max_num_seqs: 1 + max_num_seqs: 4 # in-flight request batch width (CLI: --max-num-seqs) tensor_parallel_size: 1 enforce_eager: true # Keep VAE tiling on by default for higher resolutions / longer clips. diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index f3a5f99ef..c0b9a738e 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import Any, Generic, TypeVar +from typing import Any, Generic, Literal, TypeVar import torch @@ -15,6 +15,11 @@ class AttentionBackend(ABC): accept_output_buffer: bool = False supports_piecewise_spans: bool = False + # A backend that supports this capability can consume the opaque paged-KV + # context prepared by the diffusion Worker data plane. Keeping the + # capability on the backend class prevents a paged request from silently + # falling back to dense attention on an incompatible implementation. + supports_paged_kv: bool = False # The backend can represent a contiguous valid K/V prefix by slicing the # tensors instead of materializing a padding mask. Models may use this to # avoid a slower masked-attention plan when tail padding is not semantic. @@ -22,12 +27,30 @@ class AttentionBackend(ABC): @classmethod def supports_packed_mask_free(cls) -> bool: - """Whether packed attention never reads attn_mask on this platform. + """Whether [real, pad] packed layouts can run without attn_mask. When True, models that pack a [real, pad] two-document layout and - carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip - constructing the padding mask entirely. Backends whose mask-free - behavior is platform-dependent must check current_omni_platform. + provide ``AttentionMetadata.packed_padding`` alongside the packed + cu_seqlens/max_seqlen metadata may skip constructing the padding mask + entirely. Backends whose mask-free behavior is platform-dependent must + check current_omni_platform. + """ + return False + + @classmethod + def supports_multi_doc_packed_varlen(cls) -> bool: + """Whether this backend keeps N-document packed boundaries isolated. + + When True, the backend consumes ``AttentionMetadata.extra`` cu_seqlens + as a genuine block-diagonal attention plan (a dedicated varlen kernel, + not a padding-mask rebuild), so a caller may pack multiple real + requests into one forward without attention crossing document + boundaries. When False, callers packing more than one real document + must run one forward per document; otherwise a backend that only + supports a ``[real, pad]`` two-document contract, or that ignores + cu_seqlens outright, will silently attend across request boundaries. + Backends whose kernel selection is platform-dependent must consult + ``current_omni_platform``. """ return False @@ -81,6 +104,19 @@ class AttentionBackend(ABC): supported_head_sizes = cls.get_supported_head_sizes() return (not supported_head_sizes) or head_size in supported_head_sizes + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + """Whether this backend reads K/V pages by the runtime block stride. + + Returning ``True`` means the physical cache layout has ``num_blocks`` + as its outer stride, so native vLLM may safely use page-size padding + when it unifies cache layouts across layers. Dense diffusion backends + conservatively keep the default ``False``; a paged backend should + override this only when its kernel actually follows that layout. + """ + + return False + @dataclass(frozen=True, slots=True) class QueryRange: @@ -89,9 +125,22 @@ class QueryRange: global_start: int +@dataclass(frozen=True, slots=True) +class VideoTokenSpan: + """One physical video-grid slice in packed document 0.""" + + start: int + latent_grid: tuple[int, int, int] + role: Literal["reference", "target"] + + @property + def length(self) -> int: + return self.latent_grid[0] * self.latent_grid[1] * self.latent_grid[2] + + @dataclass(frozen=True, slots=True) class VideoTokenLayout: - """Where the video segment sits inside a packed multimodal sequence. + """Video-grid slices in a packed multimodal sequence. A model that packs its sequence as ``[prefix | t*h*w video rows | padding]`` publishes this so backends can recover spatiotemporal locality; the prefix @@ -99,11 +148,32 @@ class VideoTokenLayout: Publishing it also asserts that any ``attn_mask`` masks only the trailing padding, so ``prefix_len + t*h*w`` is the used length of the sequence. + ``prefix_len``/``latent_grid`` retain the original one-tail contract. A + Ref2VA layout instead publishes ``used_len`` and every physical video + span, so audio and image rows between videos remain dense. + Plain ints, so reading it never forces a device-to-host sync. """ - prefix_len: int - latent_grid: tuple[int, int, int] + prefix_len: int | None = None + latent_grid: tuple[int, int, int] | None = None + used_len: int | None = None + video_spans: tuple[VideoTokenSpan, ...] = () + + +@dataclass(frozen=True, slots=True) +class PackedPaddingMetadata: + """Producer-validated mask-free view of padding in a [real, pad] packing. + + The cumulative-length tensors are canonical two-element ``[0, length]`` + views. Consumers may use them without reading device scalars because the + producer owns the packing. + """ + + q_length: int + kv_length: int + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor @dataclass @@ -153,6 +223,10 @@ class AttentionMetadata: # locality (block-sparse selection, tiled masks). Dense backends ignore it. video_layout: VideoTokenLayout | None = None + # Canonical mask-free view of structural suffix padding. Backends that do + # not advertise supports_packed_mask_free ignore it. + packed_padding: PackedPaddingMetadata | None = None + T = TypeVar("T", bound=AttentionMetadata) @@ -208,6 +282,18 @@ class AttentionImpl(ABC, Generic[T]): else: raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}") + def forward_paged(self, paged_kv_context: Any) -> torch.Tensor: + """Execute one Worker-prepared paged-KV attention call. + + The context is intentionally opaque to the common attention layer. + Backends opt in by setting ``supports_paged_kv`` on their backend + class and implementing this method. Dense callers continue to use + ``forward`` unchanged. + """ + + del paged_kv_context + raise NotImplementedError(f"{type(self).__name__} does not support paged KV attention") + def forward_cuda( self, query: torch.Tensor, diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index ffcc2fb44..ab814f713 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py @@ -85,6 +85,24 @@ class CuDNNAttentionImpl(AttentionImpl): enable_gqa = query.shape[2] != key.shape[2] query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + # cuDNN has no kernel for degenerate attention shapes when either + # sequence has one token. Route those degenerate shapes before entering the + # cuDNN-only context: under torch.compile, FakeTensor evaluation raises + # before the eager RuntimeError fallback below can run. + if query.shape[-2] == 1 or key.shape[-2] == 1: + with sdpa_kernel([SDPBackend.MATH]): + output = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=self.causal, + scale=self.softmax_scale, + enable_gqa=enable_gqa, + ) + return output.permute(0, 2, 1, 3) + # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask, # cuDNN gets runtime-disabled in the same call and the dispatcher falls @@ -93,7 +111,7 @@ class CuDNNAttentionImpl(AttentionImpl): # # Fall back to the default SDPA dispatcher if cuDNN rejects the shape, # e.g. under torch.compile where Dynamo sees a symbolic head_dim and - # cuDNN's kernel selection fails (observed in LTX-2 audio attention). + # cuDNN's kernel selection fails for some symbolic attention shapes. # The unpinned dispatcher then picks EFFICIENT/MATH instead of raising. try: with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]): diff --git a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py new file mode 100644 index 000000000..fcd592dd8 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py @@ -0,0 +1,431 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import functools +import math +from collections.abc import Mapping +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, +) +from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + +logger = init_logger(__name__) + + +# Keep the external FastVideo pybind/CUDA kernel opaque to torch.compile. +# This mirrors the SageAttention3 backend pattern: tracing the raw extension +# through Dynamo can reach Inductor scheduling with unstable internal op names +# (e.g. KeyError: "op12"). The custom op gives Dynamo a single Tensor->Tensor +# boundary and lets Inductor schedule the surrounding Wan block normally. +if not hasattr(torch.ops.vllm_omni, "fastvideo_vsa_bshd"): + + @torch.library.custom_op("vllm_omni::fastvideo_vsa_bshd", mutates_args=()) + def _fastvideo_vsa_bshd_op( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + variable_block_sizes: torch.Tensor, + q_variable_block_sizes: torch.Tensor, + compress_attn_weight: torch.Tensor, + topk: int, + block_t: int, + block_h: int, + block_w: int, + ) -> torch.Tensor: + from fastvideo_kernel import video_sparse_attn_bshd + + return video_sparse_attn_bshd( + query, + key, + value, + variable_block_sizes=variable_block_sizes, + q_variable_block_sizes=q_variable_block_sizes, + topk=topk, + block_size=(block_t, block_h, block_w), + compress_attn_weight=compress_attn_weight if compress_attn_weight.numel() else None, + ) + + @_fastvideo_vsa_bshd_op.register_fake + def _( + query, + key, + value, + variable_block_sizes, + q_variable_block_sizes, + compress_attn_weight, + topk, + block_t, + block_h, + block_w, + ): + del ( + key, + value, + variable_block_sizes, + q_variable_block_sizes, + compress_attn_weight, + topk, + block_t, + block_h, + block_w, + ) + return torch.empty_like(query) + + +_fastvideo_vsa_bshd_op = torch.ops.vllm_omni.fastvideo_vsa_bshd + + +@functools.lru_cache(maxsize=32) +def _get_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.Tensor: + t_size, h_size, w_size = dit_seq_shape + tile_t, tile_h, tile_w = tile_size + indices = torch.arange(t_size * h_size * w_size, device=device, dtype=torch.long).reshape(t_size, h_size, w_size) + tiles = [] + for tile_t_idx in range(math.ceil(t_size / tile_t)): + for tile_h_idx in range(math.ceil(h_size / tile_h)): + for tile_w_idx in range(math.ceil(w_size / tile_w)): + tiles.append( + indices[ + tile_t_idx * tile_t : min((tile_t_idx + 1) * tile_t, t_size), + tile_h_idx * tile_h : min((tile_h_idx + 1) * tile_h, h_size), + tile_w_idx * tile_w : min((tile_w_idx + 1) * tile_w, w_size), + ].flatten() + ) + return torch.cat(tiles, dim=0) + + +@functools.lru_cache(maxsize=32) +def _construct_variable_block_sizes( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.Tensor: + num_tiles = tuple(math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, tile_size)) + + def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.Tensor: + sizes = torch.full((n_tiles,), tile, dtype=torch.int32, device=device) + remainder = dim_len - (n_tiles - 1) * tile + sizes[-1] = remainder if remainder > 0 else tile + return sizes + + t_sizes = _sizes(dit_seq_shape[0], tile_size[0], num_tiles[0]) + h_sizes = _sizes(dit_seq_shape[1], tile_size[1], num_tiles[1]) + w_sizes = _sizes(dit_seq_shape[2], tile_size[2], num_tiles[2]) + return (t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :]).reshape(-1) + + +@functools.lru_cache(maxsize=32) +def _get_non_pad_index(variable_block_sizes: torch.Tensor, max_block_size: int) -> torch.Tensor: + num_blocks = variable_block_sizes.shape[0] + device = variable_block_sizes.device + starts = torch.arange(num_blocks, device=device) * max_block_size + padded_index = starts[:, None] + torch.arange(max_block_size, device=device)[None, :] + valid = torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None] + return padded_index[valid] + + +@torch.compiler.disable +def _get_tile_metadata( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + block_elements: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, tile_size, device) + variable_block_sizes = _construct_variable_block_sizes(dit_seq_shape, tile_size, device) + non_pad_index = _get_non_pad_index(variable_block_sizes, block_elements) + untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)] + return tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index + + +def _get_vsa_dit_seq_shape(attn_metadata: AttentionMetadata | None) -> tuple[int, int, int] | None: + if attn_metadata is None: + return None + value = attn_metadata.extra.get("vsa_dit_seq_shape") + if value is None: + return None + if not isinstance(value, (list, tuple)) or len(value) != 3: + return None + return (int(value[0]), int(value[1]), int(value[2])) + + +def _get_gate_compress(attn_metadata: AttentionMetadata | None) -> torch.Tensor | None: + if attn_metadata is None: + return None + value = attn_metadata.extra.get("gate_compress") + return value if isinstance(value, torch.Tensor) else None + + +def _preserve_vsa_all_blocks(attn_metadata: AttentionMetadata | None) -> bool: + if attn_metadata is None: + return False + return attn_metadata.extra.get("preserve_vsa_all_blocks") is True + + +class FastVideoVSABackend(AttentionBackend): + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + # FastVideo VSA is intended for video DiT head sizes such as 64/128. + # Keep this permissive and let the runtime fallback handle unsupported + # cases from the installed fastvideo-kernel build. + return [] + + @staticmethod + def get_name() -> str: + return "FASTVIDEO_VSA" + + @staticmethod + def get_impl_cls() -> type[FastVideoVSAImpl]: + return FastVideoVSAImpl + + +class FastVideoVSAImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + qkv_layout: str | None = None, + backend_kwargs: Mapping[str, Any] | None = None, + **extra_impl_args, + ) -> None: + backend_kwargs = backend_kwargs or {} + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.head_size = head_size + self.softmax_scale = softmax_scale + self.causal = causal + self.qkv_layout = qkv_layout + + self.topk = int(backend_kwargs.get("topk", 64)) + self.block_size = self._parse_block_size(backend_kwargs.get("block_size", (4, 8, 8))) + self.block_elements = self.block_size[0] * self.block_size[1] * self.block_size[2] + self.min_seq_len = int(backend_kwargs.get("min_seq_len", self.block_elements * 2)) + self.fallback_on_error = bool(backend_kwargs.get("fallback_on_error", True)) + self.disable_when_sp_active = bool(backend_kwargs.get("disable_when_sp_active", True)) + + self.sdpa_fallback = SDPAImpl( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + qkv_layout=qkv_layout, + ) + + if self.block_elements != 256: + logger.warning( + "FASTVIDEO_VSA currently uses fastvideo_kernel.video_sparse_attn_bshd, " + "which supports only 256-token blocks. Configured block_size=%s " + "(product=%d) will fall back to SDPA.", + self.block_size, + self.block_elements, + ) + + @staticmethod + def _parse_block_size(value: Any) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if isinstance(value, (list, tuple)) and len(value) == 3: + return (int(value[0]), int(value[1]), int(value[2])) + raise ValueError(f"FASTVIDEO_VSA block_size must be an int or length-3 tuple/list, got {value!r}") + + def _fallback( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + reason: str, + ) -> torch.Tensor: + logger.warning_once("FASTVIDEO_VSA falling back to SDPA: %s", reason) + return self.sdpa_fallback.forward(query, key, value, attn_metadata) + + def _fallback_reason( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ) -> str | None: + if self.causal: + return "causal attention is not supported" + if self.block_elements != 256: + return f"block_elements must be 256, got {self.block_elements}" + if self.topk <= 0: + return f"topk must be positive, got {self.topk}" + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + return f"expected [B, S, H, D] tensors, got {query.shape}, {key.shape}, {value.shape}" + if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]: + return "batch dimensions must match" + if query.shape[2:] != key.shape[2:] or query.shape[2:] != value.shape[2:]: + return "head/head_dim dimensions must match" + if query.shape[1] != key.shape[1] or query.shape[1] != value.shape[1]: + return "initial VSA backend supports self-attention with Sq == Skv only" + if query.shape[1] < self.min_seq_len: + return f"sequence length {query.shape[1]} is below min_seq_len {self.min_seq_len}" + dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata) + if dit_seq_shape is None: + return "vsa_dit_seq_shape metadata is required" + if math.prod(dit_seq_shape) != query.shape[1]: + return f"vsa_dit_seq_shape product {math.prod(dit_seq_shape)} != sequence length {query.shape[1]}" + num_blocks = math.prod( + math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size) + ) + if self.topk > num_blocks: + return f"topk {self.topk} > num_blocks {num_blocks}" + if query.dtype not in (torch.float16, torch.bfloat16): + return f"dtype {query.dtype} is not supported" + if key.dtype != query.dtype or value.dtype != query.dtype: + return "q/k/v dtypes must match" + if query.device.type != "cuda" or key.device.type != "cuda" or value.device.type != "cuda": + return "q/k/v must be CUDA tensors" + expected_scale = self.head_size**-0.5 + if abs(float(self.softmax_scale) - float(expected_scale)) > 1e-6: + return f"softmax_scale {self.softmax_scale} differs from FastVideo VSA scale {expected_scale}" + if attn_metadata is not None and attn_metadata.attn_mask is not None: + return "attention masks are not supported" + if attn_metadata is not None and attn_metadata.full_attn_spans is not None: + return "piecewise/full attention spans are not supported" + if self.num_heads != self.num_kv_heads: + return "GQA/MQA is not supported" + if self.disable_when_sp_active and is_forward_context_available(): + ctx = get_forward_context() + if getattr(ctx, "sp_active", False): + return "sequence parallel context is active" + return None + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + reason = self._fallback_reason(query, key, value, attn_metadata) + if reason is not None: + return self._fallback(query, key, value, attn_metadata, reason) + + seq_len = query.shape[1] + dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata) + assert dit_seq_shape is not None + num_blocks = math.prod( + math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size) + ) + preserve_all_blocks = _preserve_vsa_all_blocks(attn_metadata) + use_native_sdpa = self.topk == num_blocks and not preserve_all_blocks + route = "SDPA" if use_native_sdpa else "VSA_ALL_BLOCKS" if self.topk == num_blocks else "VSA" + checkpoint_mode = "fastvideo_dmd" if preserve_all_blocks else "native" + logger.info_once( + "FASTVIDEO_VSA routing: seq_len=%d, dit_seq_shape=%s, block_size=%s, num_blocks=%d, " + "topk=%d, keep_ratio=%.1f%%, checkpoint_mode=%s, route=%s", + seq_len, + dit_seq_shape, + self.block_size, + num_blocks, + self.topk, + 100.0 * self.topk / num_blocks, + checkpoint_mode, + route, + ) + if use_native_sdpa: + return self._fallback( + query, + key, + value, + attn_metadata, + f"topk {self.topk} selects all blocks for a native checkpoint", + ) + + try: + tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index = _get_tile_metadata( + dit_seq_shape, + self.block_size, + self.block_elements, + query.device, + ) + + padded_len = variable_block_sizes.numel() * self.block_elements + target_shape = (query.shape[0], padded_len, query.shape[2], query.shape[3]) + query_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype) + key_tiled = torch.zeros_like(query_tiled) + value_tiled = torch.zeros_like(query_tiled) + query_tiled[:, non_pad_index] = query[:, tile_partition_indices] + key_tiled[:, non_pad_index] = key[:, tile_partition_indices] + value_tiled[:, non_pad_index] = value[:, tile_partition_indices] + # Gate behavior is checkpoint-driven, not user-configured. + # Wan VSA layers always provide a gate projection. Its zero + # initialization makes checkpoints without gate weights sparse-only; + # checkpoints containing to_gate_compress weights use the learned gate. + gate_compress = _get_gate_compress(attn_metadata) + if gate_compress is None: + gate_compress = torch.zeros_like(query) + elif gate_compress.shape != query.shape: + raise ValueError(f"gate_compress shape {gate_compress.shape} must match query shape {query.shape}") + gate_tiled = torch.zeros_like(query_tiled) + gate_tiled[:, non_pad_index] = gate_compress[:, tile_partition_indices] + compress_attn_weight = gate_tiled + + output = _fastvideo_vsa_bshd_op( + query_tiled.contiguous(), + key_tiled.contiguous(), + value_tiled.contiguous(), + variable_block_sizes, + variable_block_sizes, + compress_attn_weight, + self.topk, + self.block_size[0], + self.block_size[1], + self.block_size[2], + ) + return output[:, untile_combined_index].contiguous() + except Exception as exc: + if not self.fallback_on_error: + raise + return self._fallback(query, key, value, attn_metadata, f"VSA kernel failed: {exc}") + + def forward_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + return self.sdpa_fallback.forward_npu(query, key, value, attn_metadata) + + def forward_xpu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + return self.sdpa_fallback.forward_xpu(query, key, value, attn_metadata) + + def forward_musa( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + return self.sdpa_fallback.forward_musa(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index bdef1dd65..a376216fa 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import os from functools import partial @@ -7,14 +7,11 @@ from functools import partial import torch from vllm.logger import init_logger -from vllm_omni.diffusion.attention.backends.abstract import ( - AttentionBackend, - AttentionImpl, - AttentionMetadata, -) +from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( piecewise_attn, + run_paged_piecewise_plan, ) from vllm_omni.diffusion.config import get_current_diffusion_config_or_none from vllm_omni.platforms import current_omni_platform @@ -25,6 +22,7 @@ logger = init_logger(__name__) class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True supports_piecewise_spans: bool = True + supports_paged_kv: bool = True @classmethod def supports_packed_mask_free(cls) -> bool: @@ -35,6 +33,20 @@ class FlashAttentionBackend(AttentionBackend): # constructing it there. return current_omni_platform.is_cuda() or current_omni_platform.is_npu() + @classmethod + def supports_multi_doc_packed_varlen(cls) -> bool: + # CUDA / ROCm / MUSA all route through ``forward_cuda``, which + # dispatches ``_forward_varlen_packed`` -> ``flash_attn_varlen_func`` + # over the caller's cu_seqlens without a mask, so an arbitrary + # N-document packing keeps its boundaries. NPU's forward path + # (``_resolve_packed_seq_npu``) only accepts a ``[real, pad]`` + # two-document layout (cu_seqlens shape <= 3) and otherwise falls + # back to a padding-mask rebuild that spans the whole packed row, + # so N>=2 real documents would silently attend across request + # boundaries. XPU never consumes cu_seqlens from ``extra``. Both + # therefore stay False here. + return current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa() + @classmethod def supports_attention_mask(cls) -> bool: return True @@ -77,12 +89,14 @@ class FlashAttentionImpl(AttentionImpl): prefix: str = "", qkv_layout: str | None = None, backend_kwargs: dict | None = None, + role: str = "self", **extra_impl_args, ) -> None: self.num_heads = num_heads self.causal = causal self.softmax_scale = softmax_scale self.qkv_layout = qkv_layout + self.is_cross_attn = role == "cross" cfg = get_current_diffusion_config_or_none() self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False if backend_kwargs: @@ -136,6 +150,7 @@ class FlashAttentionImpl(AttentionImpl): attention_mask: torch.Tensor, ) -> torch.Tensor: from vllm_omni.diffusion.attention.backends.utils.fa import ( + _index_first_axis, _pad_input, _unpad_input, _upad_input, @@ -143,10 +158,21 @@ class FlashAttentionImpl(AttentionImpl): ) assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)" - query_length = query.size(1) - q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input( - query, key, value, attention_mask, query_length, _unpad_input - ) + batch_size, query_length = query.shape[:2] + if not self.is_cross_attn and query_length == key.size(1): + q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input( + query, key, value, attention_mask, query_length, _unpad_input + ) + else: + # Cross-attention: the mask covers keys only, so keep every query row. + k, indices_k, cu_seq_lens_k, max_length_k, _ = _unpad_input(key, attention_mask) + v = _index_first_axis(value, indices_k) + q = query.flatten(0, 1) + cu_seq_lens_q = torch.arange( + 0, (batch_size + 1) * query_length, query_length, dtype=torch.int32, device=query.device + ) + max_length_q = query_length + indices_q = None out_unpad = flash_attn_varlen_func( q, @@ -162,7 +188,9 @@ class FlashAttentionImpl(AttentionImpl): }, ) out_unpad = self._unwrap_flash_output(out_unpad) - return _pad_input(out_unpad, indices_q, query.size(0), query_length) + if indices_q is None: + return out_unpad.reshape(batch_size, query_length, *out_unpad.shape[1:]) + return _pad_input(out_unpad, indices_q, batch_size, query_length) def _forward_varlen_packed( self, @@ -246,6 +274,105 @@ class FlashAttentionImpl(AttentionImpl): # (b s) h d -> b s h d return out.reshape(batch_size, q_len, *out.shape[1:]) + def forward_paged(self, paged_kv_context) -> torch.Tensor: + """Run platform-native paged attention through the Omni backend. + + ``DiffusionPagedAttentionAdapter`` owns BlockTables and prepares the + rank-local native cache context. The adapter no longer performs the + attention call itself; this method is the backend-owned execution + boundary. Its native layer wrapper keeps vLLM version-specific cache + and kernel details out of Omni's common ``Attention`` layer. CUDA uses + the native vLLM FlashAttention writer/kernel contract. Ascend writes + the complete K/V span once before its piecewise FIA calls. + """ + + layer = getattr(paged_kv_context, "layer", None) + if layer is None: + raise TypeError("paged_kv_context must expose a native diffusion attention layer") + kv_cache = getattr(layer, "kv_cache", None) + if kv_cache is None: + raise RuntimeError(f"Native KV cache is not bound for diffusion layer {layer.layer_name!r}") + native_impl = getattr(layer, "impl", None) + if native_impl is None: + raise RuntimeError(f"Native attention implementation is not bound for diffusion layer {layer.layer_name!r}") + prewrite_kv = current_omni_platform.requires_diffusion_paged_kv_prewrite() + read_kv_from_cache = prewrite_kv and layer.attn_backend.forward_includes_kv_cache_update + # The GPU/default path preserves native cache-update ownership. Ascend + # prewrites through vLLM-Ascend's normal-layout writer so piecewise FIA + # segments do not repeatedly scatter the same layer K/V. + if not layer.attn_backend.forward_includes_kv_cache_update or prewrite_kv: + cache_update = getattr(native_impl, "do_kv_cache_update", None) + if not callable(cache_update): + raise RuntimeError( + f"Native attention implementation for {layer.layer_name!r} cannot update the KV cache" + ) + cache_update( + layer, + paged_kv_context.key_write, + paged_kv_context.value_write, + kv_cache, + paged_kv_context.slot_mapping, + ) + + def run_native_attention( + query: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, + native_metadata, + output: torch.Tensor | None, + ) -> torch.Tensor: + if output is None: + output = torch.empty( + (query.shape[0], layer.num_heads, layer.head_size_v), + dtype=query.dtype, + device=query.device, + ) + return native_impl.forward( + layer, + query, + None if read_kv_from_cache else key, + None if read_kv_from_cache else value, + kv_cache, + native_metadata, + output, + ) + + if paged_kv_context.piecewise_plan is not None: + # Ascend FIA can execute identical CFG rows as one batched call + # per piece. Keep that batch layout through the piecewise runner + # so it can concatenate row-major results instead of emitting an + # indexed ScatterUpdate for every segment. The CUDA path keeps + # its existing output-buffer contract (including graph capture). + use_homogeneous_batch = ( + current_omni_platform.is_npu() and paged_kv_context.piecewise_plan.homogeneous_batch_shape is not None + ) + output = None + if not use_homogeneous_batch: + output = torch.empty( + (paged_kv_context.query.shape[0], layer.num_heads, layer.head_size_v), + dtype=paged_kv_context.query.dtype, + device=paged_kv_context.query.device, + ) + output = run_paged_piecewise_plan( + paged_kv_context.query, + None if read_kv_from_cache else paged_kv_context.key_write, + None if read_kv_from_cache else paged_kv_context.value_write, + paged_kv_context.piecewise_plan, + paged_kv_context.piecewise_native_metadata, + run_native_attention, + output_buffer=output, + use_homogeneous_batch=use_homogeneous_batch, + ) + else: + output = run_native_attention( + paged_kv_context.query, + paged_kv_context.key_write, + paged_kv_context.value_write, + paged_kv_context.native_metadata, + None, + ) + return paged_kv_context.restore_output(output) + def forward_cuda( self, query: torch.Tensor, @@ -417,15 +544,8 @@ class FlashAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: - try: - from mindiesd import attention_forward - except ImportError: - raise ImportError( - "FlashAttentionBackend NPU implementation requires MindIE-SD. " - "Please install MindIE-SD to enable NPU attention support. " - "For installation details, see https://gitcode.com/Ascend/MindIE-SD" - "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" - ) + from mindiesd import attention_forward + # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the # model marks extra["npu_attn_varlen"] and carries packed metadata, so # the padding document is excluded without reading or materializing @@ -534,15 +654,8 @@ class FlashAttentionImpl(AttentionImpl): return None seq_q, seq_k = resolved - try: - from mindiesd import attention_forward_varlen - except ImportError: - raise ImportError( - "FlashAttentionBackend NPU implementation requires MindIE-SD. " - "Please install MindIE-SD to enable NPU attention support. " - "For installation details, see https://gitcode.com/Ascend/MindIE-SD" - "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" - ) + from mindiesd import attention_forward_varlen + q = query.squeeze(0) # [T, N, D] == TND k = key.squeeze(0) v = value.squeeze(0) @@ -590,15 +703,8 @@ class FlashAttentionImpl(AttentionImpl): _, seq_k = resolved used_k = seq_k[0] # real document length (first cumulative end) - try: - from mindiesd import attention_forward - except ImportError: - raise ImportError( - "FlashAttentionBackend NPU implementation requires MindIE-SD. " - "Please install MindIE-SD to enable NPU attention support. " - "For installation details, see https://gitcode.com/Ascend/MindIE-SD" - "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" - ) + from mindiesd import attention_forward + # mindiesd always takes BSND input regardless of `layout`; the arg # selects the op-internal layout and laser only supports BNSD, so do # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here. diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py index 28b032712..9bbde02fa 100644 --- a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py +++ b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +import functools +import inspect import math from dataclasses import dataclass from typing import Any @@ -50,6 +52,27 @@ _MISSING_MINDIESD = ( "Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN" ) +_INCOMPATIBLE_MINDIESD = ( + "RAINFUSION_ATTN requires a MindIE-SD build whose sparse_attention supports the video_spans " + "argument. Please upgrade MindIE-SD or select FLASH_ATTN." +) + + +# Whether the installed mindiesd ``sparse_attention`` accepts ``precision=``. +# Releases without it accept the kwarg through ``**kwargs`` but silently ignore +# it, so a requested mix/fp8 mode would silently run the BF16 path. Cached +# because forwards run per layer per denoise step. +@functools.cache +def _mindiesd_supports_precision() -> bool: + try: + from inspect import signature + + from mindiesd import sparse_attention + + return "precision" in signature(sparse_attention).parameters + except Exception: + return False + def _try_extract_layer_index(prefix: str) -> int | None: if not prefix: @@ -60,6 +83,13 @@ def _try_extract_layer_index(prefix: str) -> int | None: return None +def _supports_video_spans(sparse_attention: Any) -> bool: + try: + return "video_spans" in inspect.signature(sparse_attention).parameters + except (TypeError, ValueError): + return False + + @dataclass(frozen=True) class RainFusionConfig: """Resolved RainFusion controls for one attention layer. @@ -72,6 +102,8 @@ class RainFusionConfig: sparsity: float = 0.0 start_step: int = 0 + end_step: int = 0 + precision: str = "bf16" skip_layers: frozenset[int] = frozenset() @classmethod @@ -80,6 +112,8 @@ class RainFusionConfig: return cls( sparsity=float(bk.get("sparsity", 0.0)), start_step=int(bk.get("start_step", 0)), + end_step=int(bk.get("end_step", 0)), + precision=str(bk.get("precision", "bf16")), skip_layers=frozenset(bk.get("skip_layers") or ()), ) @@ -92,9 +126,10 @@ class RainFusionConfig: class RainFusionPlan: """Per-forward geometry handed to the rf_v2 kernel.""" - prefix_len: int used_len: int - latent_shape: list[int] + prefix_len: int | None = None + latent_shape: list[int] | None = None + video_spans: list[dict[str, object]] | None = None class RainFusionAttentionBackend(AttentionBackend): @@ -107,6 +142,12 @@ class RainFusionAttentionBackend(AttentionBackend): if find_spec("mindiesd") is None: raise ValueError(_MISSING_MINDIESD) + try: + from mindiesd import sparse_attention + except ImportError as exc: + raise ValueError(_MISSING_MINDIESD) from exc + if not _supports_video_spans(sparse_attention): + raise ValueError(_INCOMPATIBLE_MINDIESD) @staticmethod def get_supported_head_sizes() -> list[int]: @@ -232,8 +273,17 @@ class RainFusionAttentionImpl(AttentionImpl): return None if is_forward_context_available(): step_idx = get_forward_context().denoise_step_idx + total_steps = get_forward_context().total_denoise_steps if step_idx is not None and step_idx < rf.start_step: return None + # Tail fallback: keep the last ``end_step`` denoise steps dense. + if ( + rf.end_step > 0 + and step_idx is not None + and total_steps is not None + and step_idx >= total_steps - rf.end_step + ): + return None if self.qkv_layout is None: # The sparse path reads the sequence off dim 1, which the tensors alone # do not establish, and the dense fallback resolves an absent layout its @@ -266,6 +316,14 @@ class RainFusionAttentionImpl(AttentionImpl): ) return None + if layout.video_spans: + return self._resolve_multi_span_plan(layout, max_seqlen_q) + + if layout.prefix_len is None or layout.latent_grid is None: + logger.warning_once( + "RAINFUSION_ATTN staying dense: video layout has neither a legacy video tail nor multi-video spans." + ) + return None prefix_len = int(layout.prefix_len) latent_shape = [int(dim) for dim in layout.latent_grid] # rf_v2 splits the sequence as [prefix | t*h*w video rows]. Document 0 of @@ -304,11 +362,97 @@ class RainFusionAttentionImpl(AttentionImpl): video_len, ) return RainFusionPlan( - prefix_len=prefix_len, used_len=used_len, + prefix_len=prefix_len, latent_shape=latent_shape, ) + def _resolve_multi_span_plan(self, layout, max_seqlen_q: int) -> RainFusionPlan | None: + """Validate Ref2VA's non-contiguous video grids before sparse dispatch.""" + used_len = layout.used_len + if used_len is None or int(used_len) != int(max_seqlen_q): + logger.warning_once( + "RAINFUSION_ATTN staying dense: multi-video layout used_len=%r does not match packed document 0 (%d).", + used_len, + int(max_seqlen_q), + ) + return None + + spans: list[dict[str, object]] = [] + span_summaries: list[str] = [] + previous_end = 0 + target_count = 0 + video_seqlen = 0 + boundary_dense_seqlen = 0 + previous_video_length: int | None = None + for span in sorted(layout.video_spans, key=lambda item: item.start): + grid = tuple(int(dim) for dim in span.latent_grid) + length = math.prod(grid) + start = int(span.start) + if ( + len(grid) != 3 + or any(dim <= 0 for dim in grid) + or start < previous_end + or start + length > int(used_len) + ): + logger.warning_once( + "RAINFUSION_ATTN staying dense: invalid multi-video span start=%d grid=%s used_len=%d.", + start, + grid, + int(used_len), + ) + return None + role = span.role + if role not in ("reference", "target"): + logger.warning_once("RAINFUSION_ATTN staying dense: unsupported multi-video span role %r.", role) + return None + if role == "target": + target_count += 1 + if previous_video_length is not None: + # rf_v2 works on fixed 128-token blocks. The preceding clip + # needs these dense rows to complete its tail block before + # this clip begins, otherwise one sparse block would cross a + # clip boundary. + boundary_dense_seqlen += (-previous_video_length) % _BLOCK_SIZE + spans.append({"start": start, "latent_shape": list(grid)}) + span_summaries.append(f"role={role}, start={start}, seqlen={length}, latent_shape={grid}") + previous_end = start + length + video_seqlen += length + previous_video_length = length + + if target_count != 1: + logger.warning_once( + "RAINFUSION_ATTN staying dense: Ref2VA layout must contain exactly one target video span, got %d.", + target_count, + ) + return None + if video_seqlen < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE: + logger.warning_once( + "RAINFUSION_ATTN staying dense: multi-video seqlen=%d is under the sparse threshold seqlen=%d.", + video_seqlen, + _MIN_VIDEO_BLOCKS * _BLOCK_SIZE, + ) + return None + dense_context_seqlen = int(used_len) - video_seqlen + if boundary_dense_seqlen > dense_context_seqlen: + logger.warning_once( + "RAINFUSION_ATTN staying dense: multi-video spans need dense_context_seqlen=%d " + "to isolate clip block boundaries, but this layout has only %d.", + boundary_dense_seqlen, + dense_context_seqlen, + ) + return None + logger.info_once( + "RAINFUSION_ATTN multi-video active: sparsity=%.2f, spans=[%s], " + "video_seqlen=%d, dense_context_seqlen=%d, valid_packed_seqlen=%d.", + self.rainfusion.sparsity, + "; ".join(span_summaries), + video_seqlen, + dense_context_seqlen, + int(used_len), + ) + return RainFusionPlan(used_len=int(used_len), video_spans=spans) + def _forward_sparse_npu( self, query: torch.Tensor, @@ -320,26 +464,45 @@ class RainFusionAttentionImpl(AttentionImpl): from mindiesd import sparse_attention except ImportError: raise ImportError(_MISSING_MINDIESD) + if self.rainfusion.precision != "bf16" and not _mindiesd_supports_precision(): + raise RuntimeError( + f"block_sparse.precision={self.rainfusion.precision!r} requires MindIE-SD " + "with sparse_attention(precision=...) support; the installed mindiesd " + "silently ignores it and would run the BF16 path. Install a compatible " + "MindIE-SD release or use precision='bf16'." + ) used = plan.used_len q, k, v = (tensor[:, :used] for tensor in (query, key, value)) # Ulysses has already gathered the full sequence onto this rank and split # the heads, so read the head count off the tensor rather than num_heads. - out = sparse_attention( - q, - k, - v, - scale=self.softmax_scale, - head_num=query.shape[-2], - input_layout=_INPUT_LAYOUT, - inner_precise=_INNER_PRECISE, - sparse_type="rf_v2", - txt_len=plan.prefix_len, - block_size=_BLOCK_SIZE, - latent_shape_q=plan.latent_shape, - latent_shape_k=plan.latent_shape, - sparsity=self.rainfusion.sparsity, - ) + common_kwargs: dict[str, object] = { + "scale": self.softmax_scale, + "head_num": query.shape[-2], + "input_layout": _INPUT_LAYOUT, + "inner_precise": _INNER_PRECISE, + "block_size": _BLOCK_SIZE, + "sparsity": self.rainfusion.sparsity, + "precision": self.rainfusion.precision, + } + if plan.video_spans is not None: + out = sparse_attention( + q, + k, + v, + sparse_type="rf_v2", + video_spans=plan.video_spans, + **common_kwargs, + ) + else: + assert plan.prefix_len is not None and plan.latent_shape is not None + common_kwargs.update( + sparse_type="rf_v2", + txt_len=plan.prefix_len, + latent_shape_q=plan.latent_shape, + latent_shape_k=plan.latent_shape, + ) + out = sparse_attention(q, k, v, **common_kwargs) if used == query.shape[1]: return out padded = torch.zeros_like(query) diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index fdfaf5f9c..1af079696 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/vllm_omni/diffusion/attention/backends/registry.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Diffusion attention backend registry. This module provides an enum-based registry for diffusion attention backends, @@ -23,7 +23,7 @@ logger = init_logger(__name__) class _DiffusionBackendEnumMeta(EnumMeta): """Metaclass for DiffusionAttentionBackendEnum to provide better error messages.""" - def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum": + def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum": # type: ignore[override] """Get backend by name with helpful error messages.""" try: return super().__getitem__(name) # type: ignore[return-value] @@ -62,6 +62,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta): SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend" CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend" FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend" + FASTVIDEO_VSA = "vllm_omni.diffusion.attention.backends.fastvideo_vsa.FastVideoVSABackend" FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend" FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend" TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend" diff --git a/vllm_omni/diffusion/attention/backends/trtllm_attn.py b/vllm_omni/diffusion/attention/backends/trtllm_attn.py index c63280fc8..967e81c6e 100644 --- a/vllm_omni/diffusion/attention/backends/trtllm_attn.py +++ b/vllm_omni/diffusion/attention/backends/trtllm_attn.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import functools import inspect import math from dataclasses import dataclass, replace +from typing import NamedTuple, cast import torch from vllm.logger import init_logger @@ -13,6 +14,7 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionBackend, AttentionImpl, AttentionMetadata, + PackedPaddingMetadata, ) logger = init_logger(__name__) @@ -42,8 +44,9 @@ class SkipSoftmaxConfig: return cls( threshold=_validate_control(bk.get("skip_softmax_threshold"), "skip_softmax_threshold", 0.0, None), target_sparsity=_validate_control(bk.get("target_sparsity"), "target_sparsity", 0.0, 1.0), - disabled_until_timestep=_validate_control( - bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0 + disabled_until_timestep=cast( + float, + _validate_control(bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0), ), ) @@ -131,7 +134,7 @@ class QuantConfig: return self.dtype_qk is not None def quantize(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, quantize_fn): - qk_quant_dtype = _QK_QUANT_DTYPES[self.dtype_qk] + qk_quant_dtype = _QK_QUANT_DTYPES[cast(str, self.dtype_qk)] q_q, k_q, v_q, q_sfs, k_sfs, v_sfs = quantize_fn( q, k, @@ -145,6 +148,18 @@ class QuantConfig: return q_q, k_q, v_q, sage_attn_sfs, num_elts_per_sage_attn_blk +class _PackedLayout(NamedTuple): + """Active token ranges and kernel metadata for a flattened ragged batch.""" + + q_tokens: int + kv_tokens: int + cu_seqlens_q: torch.Tensor + cu_seqlens_kv: torch.Tensor + batch_size: int + seq_lens: torch.Tensor + known_min_kv_length: int | None = None + + def _workspace_bytes() -> int: import vllm.envs as envs @@ -154,6 +169,14 @@ def _workspace_bytes() -> int: class TrtllmAttentionBackend(AttentionBackend): accept_output_buffer: bool = True + @classmethod + def supports_packed_mask_free(cls) -> bool: + return True + + @classmethod + def supports_multi_doc_packed_varlen(cls) -> bool: + return True + @staticmethod def get_supported_head_sizes() -> list[int]: return [128] @@ -250,6 +273,97 @@ class TrtllmAttentionImpl(AttentionImpl): cls._workspace = ws return ws + @staticmethod + def _prepare_packed_padding_layout( + physical_batch: int, + q: torch.Tensor, + k: torch.Tensor, + packed_padding: PackedPaddingMetadata, + extra: dict, + ) -> _PackedLayout: + """Prepare a [real, pad] layout from an explicit host-side boundary.""" + valid_q_tokens = packed_padding.q_length + valid_kv_tokens = packed_padding.kv_length + max_q_len = extra["max_seqlen_q"] + max_kv_len = extra["max_seqlen_k"] + lengths = (valid_q_tokens, valid_kv_tokens, max_q_len, max_kv_len) + if any(isinstance(value, bool) or not isinstance(value, int) for value in lengths): + raise ValueError("Mask-free packed TRTLLM attention lengths must be Python integers") + if physical_batch != 1: + raise ValueError("Mask-free packed TRTLLM attention requires a single packed batch") + + total_q_tokens = q.shape[0] + total_kv_tokens = k.shape[0] + if not 0 < valid_q_tokens <= total_q_tokens: + raise ValueError( + "PackedPaddingMetadata.q_length must be within the packed Q sequence, " + f"got {valid_q_tokens} for length {total_q_tokens}" + ) + if not 0 < valid_kv_tokens <= total_kv_tokens: + raise ValueError( + "PackedPaddingMetadata.kv_length must be within the packed K/V sequence, " + f"got {valid_kv_tokens} for length {total_kv_tokens}" + ) + if max_q_len != valid_q_tokens or max_kv_len != valid_kv_tokens: + raise ValueError("Mask-free packed TRTLLM attention lengths must match max_seqlen_q/k") + published_kv_length = extra.get("valid_kv_length") + if published_kv_length is not None: + if isinstance(published_kv_length, bool) or not isinstance(published_kv_length, int): + raise ValueError("valid_kv_length must be a Python integer") + if published_kv_length != valid_kv_tokens: + raise ValueError("PackedPaddingMetadata.kv_length must match valid_kv_length") + + cu_seq_lens_q = packed_padding.cu_seqlens_q + cu_seq_lens_kv = packed_padding.cu_seqlens_k + if cu_seq_lens_q.dtype != torch.int32 or cu_seq_lens_kv.dtype != torch.int32: + raise ValueError("Mask-free packed TRTLLM attention requires int32 cu_seqlens") + if cu_seq_lens_q.device != q.device or cu_seq_lens_kv.device != k.device: + raise ValueError("Mask-free packed TRTLLM attention metadata must be on the Q/K device") + if cu_seq_lens_q.shape != (2,) or cu_seq_lens_kv.shape != (2,): + raise ValueError("Mask-free packed TRTLLM attention requires canonical two-element cu_seqlens") + return _PackedLayout( + q_tokens=valid_q_tokens, + kv_tokens=valid_kv_tokens, + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_kv=cu_seq_lens_kv, + batch_size=1, + # Canonical packed-padding metadata is [0, valid_kv_tokens], so this slice + # is the one-element sequence-length view expected by the kernel. + seq_lens=cu_seq_lens_kv[1:], + known_min_kv_length=valid_kv_tokens, + ) + + @staticmethod + def _prepare_generic_packed_layout( + q: torch.Tensor, + k: torch.Tensor, + cu_seq_lens_q: torch.Tensor, + cu_seq_lens_kv: torch.Tensor, + ) -> _PackedLayout: + """Validate and normalize a general ragged batch.""" + if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]: + raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens") + + # Packed producers may retain a zero-length trailing segment when the + # input is already aligned. It is not a real sequence and is removed + # before deriving the batch size and per-sequence lengths. + while ( + cu_seq_lens_q.numel() > 2 + and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item()) + and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item()) + ): + cu_seq_lens_q = cu_seq_lens_q[:-1] + cu_seq_lens_kv = cu_seq_lens_kv[:-1] + + return _PackedLayout( + q_tokens=q.shape[0], + kv_tokens=k.shape[0], + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_kv=cu_seq_lens_kv, + batch_size=cu_seq_lens_q.numel() - 1, + seq_lens=(cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous(), + ) + def forward_cuda( self, query: torch.Tensor, @@ -266,12 +380,16 @@ class TrtllmAttentionImpl(AttentionImpl): has_packed_metadata = len(present_packed_keys) == len(packed_keys) attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None - if attn_mask is not None and not has_packed_metadata: + packed_padding = getattr(attn_metadata, "packed_padding", None) if attn_metadata is not None else None + if packed_padding is not None and not isinstance(packed_padding, PackedPaddingMetadata): + raise ValueError("packed_padding must be PackedPaddingMetadata") + if packed_padding is not None and not has_packed_metadata: + raise ValueError("PackedPaddingMetadata requires complete packed TRTLLM attention metadata") + if attn_mask is not None: raise ValueError( - "TRTLLM_ATTN does not support attn_mask (Wan sets one only under SP with " - "mask_sp_padding=True and a non-divisible seq len). Either set " - "parallel_config.mask_sp_padding=False, or select a mask-capable backend " - "(e.g. CUDNN_ATTN / TORCH_SDPA)." + "TRTLLM_ATTN does not support attn_mask. Represent structural suffix padding " + "with packed-padding metadata, or select a mask-capable backend such as " + "CUDNN_ATTN or TORCH_SDPA." ) if not HAS_FLASHINFER: @@ -282,7 +400,6 @@ class TrtllmAttentionImpl(AttentionImpl): physical_batch, q_len, num_q_heads, head_dim = query.shape kv_len, num_kv_heads = key.shape[1], key.shape[2] - device = query.device q = query.reshape(physical_batch * q_len, num_q_heads, head_dim).contiguous() @@ -290,6 +407,7 @@ class TrtllmAttentionImpl(AttentionImpl): v = value.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous() output_tokens = q.shape[0] + known_min_kv_length: int | None = kv_len if has_packed_metadata: cu_seq_lens_q = extra["cu_seqlens_q"] cu_seq_lens_kv = extra["cu_seqlens_k"] @@ -297,48 +415,36 @@ class TrtllmAttentionImpl(AttentionImpl): raise ValueError("Packed TRTLLM attention cu_seqlens tensors must be one-dimensional") if cu_seq_lens_q.numel() != cu_seq_lens_kv.numel() or cu_seq_lens_q.numel() < 2: raise ValueError("Packed TRTLLM attention requires matching non-empty Q and KV sequence batches") - if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]: - raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens") - # Drop trailing empty sequences retained by packed producers when - # alignment adds no padding tokens. - while ( - cu_seq_lens_q.numel() > 2 - and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item()) - and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item()) - ): - cu_seq_lens_q = cu_seq_lens_q[:-1] - cu_seq_lens_kv = cu_seq_lens_kv[:-1] + # PackedPaddingMetadata is an optional producer-published shortcut + # for structural suffix padding. Inputs without it retain the + # general cu_seqlens path for arbitrary ragged Q/KV batches. + if packed_padding is not None: + packed_layout = self._prepare_packed_padding_layout( + physical_batch, + q, + k, + packed_padding, + extra, + ) + else: + # Generic metadata can describe any number of real sequences, + # not only a single [real, pad] pair. + packed_layout = self._prepare_generic_packed_layout( + q, + k, + cu_seq_lens_q, + cu_seq_lens_kv, + ) - # A packed prefix mask denotes structural suffix padding, not another - # document. Exclude it before quantization and attention; in particular, - # SAGE block quantization must not span the valid/padding boundary. - if attn_mask is not None: - if q.shape[0] != k.shape[0] or attn_mask.numel() != q.shape[0]: - raise ValueError("Packed TRTLLM prefix masks require equal Q/K lengths") - flat_mask = attn_mask.reshape(-1).to(dtype=torch.bool) - valid_tokens = int(flat_mask.sum().item()) - expected_mask = torch.arange(flat_mask.numel(), device=device) < valid_tokens - if not torch.equal(flat_mask, expected_mask): - raise ValueError("TRTLLM_ATTN only supports prefix-valid packed attention masks") - - q_boundary = (cu_seq_lens_q == valid_tokens).nonzero(as_tuple=False) - kv_boundary = (cu_seq_lens_kv == valid_tokens).nonzero(as_tuple=False) - if q_boundary.numel() != 1 or kv_boundary.numel() != 1: - raise ValueError("Packed TRTLLM prefix length must be a Q/K sequence boundary") - q_end = int(q_boundary.item()) + 1 - kv_end = int(kv_boundary.item()) + 1 - if q_end != kv_end: - raise ValueError("Packed TRTLLM prefix masks require matching Q/K sequence batches") - - q = q[:valid_tokens] - k = k[:valid_tokens] - v = v[:valid_tokens] - cu_seq_lens_q = cu_seq_lens_q[:q_end].contiguous() - cu_seq_lens_kv = cu_seq_lens_kv[:kv_end].contiguous() - - batch = cu_seq_lens_q.numel() - 1 - seq_lens = (cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous() + q = q[: packed_layout.q_tokens] + k = k[: packed_layout.kv_tokens] + v = v[: packed_layout.kv_tokens] + cu_seq_lens_q = packed_layout.cu_seqlens_q + cu_seq_lens_kv = packed_layout.cu_seqlens_kv + batch = packed_layout.batch_size + seq_lens = packed_layout.seq_lens + known_min_kv_length = packed_layout.known_min_kv_length max_q_len = int(extra["max_seqlen_q"]) max_kv_len = int(extra["max_seqlen_k"]) else: @@ -361,7 +467,13 @@ class TrtllmAttentionImpl(AttentionImpl): sage_kwargs: dict = {} # The SAGE kernel requires every KV sequence to contain at least one full # quantization block. Small auxiliary attention sites use the dense kernel. - use_sage = self.quant.enabled and bool(torch.all(seq_lens >= self.quant.k_block_size).item()) + use_sage = False + if self.quant.enabled: + if known_min_kv_length is not None: + sage_lengths_supported = known_min_kv_length >= self.quant.k_block_size + else: + sage_lengths_supported = bool(torch.all(seq_lens >= self.quant.k_block_size).item()) + use_sage = sage_lengths_supported if self.quant.enabled and not use_sage: message = ( f"TRTLLM_ATTN SAGE quantization is configured for attention role {self.role!r}, but at least one " diff --git a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py index 5d83d8dbf..a8676c34a 100644 --- a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py +++ b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Piecewise attention for mixed causal / full (bidirectional) masks. @@ -15,6 +15,8 @@ Per segment: from __future__ import annotations +from collections.abc import Callable, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Literal, NamedTuple import torch @@ -90,6 +92,7 @@ def piecewise_attn( ): _check_homogeneous(full_attn_spans) spans = full_attn_spans[0] + ranges: tuple[tuple[int, int, int], ...] if query_ranges is None: query_len = query.shape[1] ranges = ((0, query_len, key.shape[1] - query_len),) @@ -123,3 +126,282 @@ def piecewise_attn( if len(outputs) == 1: return outputs[0] return torch.cat(outputs, dim=1) + + +@dataclass(frozen=True, slots=True) +class PagedPiecewiseSegment: + """One aligned segment across rows in a packed paged batch.""" + + row_segments: tuple[Segment, ...] + query_indices: torch.Tensor + query_range: tuple[int, int] | None + # Local [start, end) range shared by every row when the batch can retain + # its original [B, S] layout. The packed query_indices remain available + # for metadata construction and the heterogeneous fallback. + local_query_range: tuple[int, int] | None = None + + +@dataclass(frozen=True, slots=True) +class PagedPiecewisePlan: + """Piecewise segment batches for one packed paged-attention batch.""" + + spans: tuple[tuple[tuple[int, int], ...], ...] + segments: tuple[PagedPiecewiseSegment, ...] + num_query_tokens: int + segments_cover_query_contiguously: bool + # (batch_size, query_len) when all rows have the same local segment + # layout. Native FIA still receives flattened row-major tokens, but the + # runner can restore outputs with one batch concat instead of one indexed + # scatter per segment. + homogeneous_batch_shape: tuple[int, int] | None = None + + +PagedPiecewiseRunner = Callable[ + [torch.Tensor, torch.Tensor | None, torch.Tensor | None, object, torch.Tensor | None], + torch.Tensor, +] + + +def build_paged_piecewise_plan( + full_attn_spans: Sequence[Sequence[tuple[int, int]]], + query_offsets: Sequence[int], + query_lens: Sequence[int], + seq_lens: Sequence[int], + *, + device: torch.device | str | None = None, +) -> PagedPiecewisePlan: + """Build packed indices for corresponding piecewise segments in each row.""" + + row_count = len(full_attn_spans) + if row_count == 0 or not (row_count == len(query_offsets) == len(query_lens) == len(seq_lens)): + raise ValueError("Paged piecewise inputs must have one entry per row") + + spans: tuple[tuple[tuple[int, int], ...], ...] = tuple( + tuple((span_start, span_end) for span_start, span_end in row_spans) for row_spans in full_attn_spans + ) + + packed_offsets = [0] + for row_spans, query_offset, query_len, seq_len in zip(spans, query_offsets, query_lens, seq_lens, strict=True): + previous_end = 0 + for start, end in row_spans: + if start < previous_end or not 0 <= start < end <= seq_len: + raise ValueError("Paged piecewise spans must be sorted, non-overlapping, and within the sequence") + previous_end = end + if query_offset < 0 or query_len <= 0 or query_offset + query_len > seq_len: + raise ValueError("Paged piecewise query range must be within the sequence") + packed_offsets.append(packed_offsets[-1] + query_len) + + segments_by_row = tuple( + tuple(build_segments(row_spans, query_offset, query_len)) + for row_spans, query_offset, query_len in zip(spans, query_offsets, query_lens, strict=True) + ) + if len({len(row_segments) for row_segments in segments_by_row}) != 1: + raise ValueError("Paged piecewise rows must produce aligned segments") + + packed_segments = [] + for row_segments in zip(*segments_by_row, strict=True): + if len({segment.mode for segment in row_segments}) != 1: + raise ValueError("Paged piecewise rows must use the same attention mode per segment") + query_indices: list[int] = [] + local_ranges: list[tuple[int, int]] = [] + for row_index, (segment, query_offset) in enumerate(zip(row_segments, query_offsets, strict=True)): + local_start = segment.q_start - query_offset + local_end = segment.q_end - query_offset + local_ranges.append((local_start, local_end)) + query_indices.extend(range(packed_offsets[row_index] + local_start, packed_offsets[row_index] + local_end)) + query_range = None + if query_indices and query_indices == list(range(query_indices[0], query_indices[-1] + 1)): + query_range = (query_indices[0], query_indices[-1] + 1) + local_query_range = local_ranges[0] if len(set(local_ranges)) == 1 else None + packed_segments.append( + PagedPiecewiseSegment( + row_segments=tuple(row_segments), + query_indices=torch.tensor(query_indices, dtype=torch.long, device=device), + query_range=query_range, + local_query_range=local_query_range, + ) + ) + + covered = 0 + for segment in packed_segments: + if segment.query_range is None or segment.query_range[0] != covered: + break + covered = segment.query_range[1] + segments_cover_query_contiguously = covered == packed_offsets[-1] + homogeneous_batch_shape: tuple[int, int] | None = None + if ( + row_count > 1 + and len(set(query_lens)) == 1 + and all(segment.local_query_range is not None for segment in packed_segments) + ): + # A common local range is sufficient even when rows have different + # global offsets or KV lengths; those values remain row-specific in + # native metadata. + homogeneous_batch_shape = (row_count, int(query_lens[0])) + return PagedPiecewisePlan( + spans=spans, + segments=tuple(packed_segments), + num_query_tokens=packed_offsets[-1], + segments_cover_query_contiguously=segments_cover_query_contiguously, + homogeneous_batch_shape=homogeneous_batch_shape, + ) + + +def _run_homogeneous_paged_piecewise_plan( + query: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, + plan: PagedPiecewisePlan, + segment_metadata: Sequence[object], + segment_runner: PagedPiecewiseRunner, + output_buffer: torch.Tensor | None, +) -> torch.Tensor: + """Run a homogeneous batch while retaining the old [B, S] layout. + + The native FIA contract is flattened TND, so each segment is flattened + only for the kernel call. Results are reshaped back to [B, L] and + concatenated along the sequence dimension, which avoids the per-segment + ``index_copy_`` used by heterogeneous packed rows. + """ + + batch_shape = plan.homogeneous_batch_shape + assert batch_shape is not None + batch_size, query_len = batch_shape + if query.shape[0] != batch_size * query_len: + raise ValueError( + f"Homogeneous paged piecewise plan expects {batch_size * query_len} query tokens, got {query.shape[0]}" + ) + query_rows = query.reshape(batch_size, query_len, *query.shape[1:]) + key_rows = None if key is None else key.reshape(batch_size, query_len, *key.shape[1:]) + value_rows = None if value is None else value.reshape(batch_size, query_len, *value.shape[1:]) + + segment_outputs: list[torch.Tensor] = [] + for segment, metadata in zip(plan.segments, segment_metadata, strict=True): + local_range = segment.local_query_range + if local_range is None: + raise ValueError("Homogeneous paged piecewise execution requires a shared local segment range") + start, end = local_range + segment_query = query_rows[:, start:end].contiguous().reshape(-1, *query.shape[1:]) + segment_key = None if key_rows is None else key_rows[:, start:end].contiguous().reshape(-1, *key.shape[1:]) + segment_value = ( + None if value_rows is None else value_rows[:, start:end].contiguous().reshape(-1, *value.shape[1:]) + ) + segment_output = segment_runner(segment_query, segment_key, segment_value, metadata, None) + expected_tokens = batch_size * (end - start) + if segment_output.shape[0] != expected_tokens: + raise ValueError( + f"Paged piecewise runner returned {segment_output.shape[0]} tokens " + f"for a {expected_tokens}-token homogeneous segment" + ) + segment_outputs.append(segment_output.reshape(batch_size, end - start, *segment_output.shape[1:])) + + if not segment_outputs: + result = torch.empty_like(query) + elif len(segment_outputs) == 1: + result = segment_outputs[0].reshape(plan.num_query_tokens, *segment_outputs[0].shape[2:]) + else: + result_rows = torch.cat(segment_outputs, dim=1) + result = result_rows.reshape(plan.num_query_tokens, *result_rows.shape[2:]) + if output_buffer is not None: + if output_buffer.shape != result.shape: + raise ValueError( + f"Paged piecewise output buffer shape {tuple(output_buffer.shape)} does not match " + f"homogeneous result shape {tuple(result.shape)}" + ) + output_buffer.copy_(result) + return output_buffer + return result + + +def run_paged_piecewise_plan( + query: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, + plan: PagedPiecewisePlan, + segment_metadata: Sequence[object], + segment_runner: PagedPiecewiseRunner, + output_buffer: torch.Tensor | None = None, + *, + use_homogeneous_batch: bool = False, +) -> torch.Tensor: + """Run piecewise attention with view fast paths for contiguous segments. + + ``use_homogeneous_batch`` is an opt-in for native backends that can keep + identical rows in the legacy batch layout. The default remains the + indexed packed path so existing GPU/heterogeneous contracts are unchanged. + """ + + if query.shape[0] != plan.num_query_tokens: + raise ValueError(f"Paged piecewise plan has {plan.num_query_tokens} query tokens, got {query.shape[0]}") + if output_buffer is not None and output_buffer.shape[0] != plan.num_query_tokens: + raise ValueError( + f"Paged piecewise output buffer has {output_buffer.shape[0]} tokens; expected {plan.num_query_tokens}" + ) + + if use_homogeneous_batch and plan.homogeneous_batch_shape is not None: + return _run_homogeneous_paged_piecewise_plan( + query, + key, + value, + plan, + segment_metadata, + segment_runner, + output_buffer, + ) + + output = output_buffer + contiguous_outputs = [] + for segment, metadata in zip(plan.segments, segment_metadata, strict=True): + indices = segment.query_indices + segment_output_buffer = None + if segment.query_range is None: + segment_query = query.index_select(0, indices) + segment_key = None if key is None else key.index_select(0, indices) + segment_value = None if value is None else value.index_select(0, indices) + else: + start, end = segment.query_range + segment_query = query[start:end] + segment_key = None if key is None else key[start:end] + segment_value = None if value is None else value[start:end] + if output is not None: + segment_output_buffer = output[start:end] + segment_output = segment_runner( + segment_query, + segment_key, + segment_value, + metadata, + segment_output_buffer, + ) + if segment_output.shape[0] != indices.shape[0]: + raise ValueError( + f"Paged piecewise runner returned {segment_output.shape[0]} tokens " + f"for a {indices.shape[0]}-token segment" + ) + if segment_output_buffer is not None and segment_output is not segment_output_buffer: + if segment_output.shape != segment_output_buffer.shape: + raise ValueError( + f"Paged piecewise runner returned shape {tuple(segment_output.shape)} " + f"for output buffer shape {tuple(segment_output_buffer.shape)}" + ) + segment_output_buffer.copy_(segment_output) + segment_output = segment_output_buffer + if plan.segments_cover_query_contiguously: + if output_buffer is None: + contiguous_outputs.append(segment_output) + continue + if output is None: + output = segment_output.new_empty((query.shape[0], *segment_output.shape[1:])) + if segment.query_range is None: + output.index_copy_(0, indices, segment_output) + else: + start, end = segment.query_range + output[start:end].copy_(segment_output) + + if output_buffer is not None and plan.segments_cover_query_contiguously: + return output_buffer + if contiguous_outputs: + if len(contiguous_outputs) == 1: + return contiguous_outputs[0] + return torch.cat(contiguous_outputs, dim=0) + assert output is not None + return output diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 7eb74a941..7dca3c1e6 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright (c) Microsoft Corporation and Jiarui Fang # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team & Jiarui Fang @@ -11,8 +11,10 @@ from dataclasses import replace import torch import torch.nn as nn +from vllm.config import VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.models.utils import extract_layer_index +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend @@ -21,8 +23,13 @@ from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role from vllm_omni.diffusion.config import get_current_diffusion_config_or_none +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.distributed.parallel_state import get_sp_group -from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.forward_context import ( + get_forward_context, + get_ulysses_mode, + is_forward_context_available, +) from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -38,6 +45,8 @@ def _try_extract_layer_index(prefix: str) -> int | None: class Attention(nn.Module): + _scheduler_paged_kv = False + def __init__( self, num_heads: int, @@ -61,12 +70,28 @@ class Attention(nn.Module): # perf for this layer (e.g. Wan2.2 cross-attn has short sequences and # block-FP8 quant offers no win). Default False = follow global config. disable_kv_quant: bool = False, + # Opt-in marker for Scheduler-managed paged KV. Unmarked diffusion + # attention remains dense and contributes no native KVCacheSpec. + paged_kv_cache_role: str | None = None, + paged_kv_cache_dtype: torch.dtype | None = None, ): super().__init__() self.role = role self.role_category = role_category self.qkv_layout = qkv_layout + # ``prefix`` is also the stable layer identity used by vLLM's native + # KV-cache metadata. Keep it on the Omni layer so the active paged + # adapter can dispatch the already-resharded Q/K/V to the matching + # native cache tensor without replacing this Omni execution path. + self.prefix = prefix + if paged_kv_cache_role == "": + raise ValueError("paged_kv_cache_role must be non-empty when provided") + self.paged_kv_cache_role = paged_kv_cache_role + self.paged_kv_cache_dtype = paged_kv_cache_dtype + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.head_size = head_size # Resolve backend via role-aware config. # The global diffusion config is set during model init via @@ -89,6 +114,44 @@ class Attention(nn.Module): role_category=role_category, allow_trtllm_default=allow_trtllm_default, ) + scheduler_paged_kv = ( + config is not None + and getattr(config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ) + self._scheduler_paged_kv = scheduler_paged_kv + if ( + scheduler_paged_kv + and paged_kv_cache_role is not None + and spec is None + and not attn_backend_cls.supports_paged_kv + ): + # FLASH_ATTN is an Omni selector, not a device-specific kernel. + # vLLM resolves it to CUDA FlashAttention on GPU and Ascend native + # attention on NPU. A platform may legitimately default dense + # attention to SDPA (for example when an optional diffusion FA + # package is absent), but a paged layer must advertise the + # capability so the Worker can register its native cache view. + # Resolve the selector only for this marked layer; unmarked dense + # attention keeps the platform default unchanged. + from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec + + dense_backend_name = attn_backend_cls.get_name() + paged_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")) + attn_backend_cls, spec = get_attn_backend_for_role( + role=role, + head_size=head_size, + attention_config=paged_attention_config, + role_category=role_category, + allow_trtllm_default=allow_trtllm_default, + ) + logger.info( + "Resolved marked paged diffusion attention role=%r to %r because the platform default %r " + "does not support Scheduler-owned KV", + role, + attn_backend_cls.get_name(), + dense_backend_name, + ) parallel_config = getattr(config, "parallel_config", None) allgather_degree = getattr(parallel_config, "allgather_degree", 1) # TODO: Move AllGather-KV compatibility into an AttentionBackend capability @@ -171,6 +234,25 @@ class Attention(nn.Module): self._disable_kv_quant: bool = disable_kv_quant self._init_kv_cache_quantization(config) + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + """Return native rank-local geometry for an opted-in paged cache.""" + + if self.paged_kv_cache_role is None: + return None + dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype + # Keep backend layout discovery under the same config context used by + # upstream vLLM's attention-spec collector. + with set_current_vllm_config(vllm_config): + indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride() + return FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + dtype=dtype, + indexes_kv_by_block_stride=indexes_kv_by_block_stride, + non_causal=not self.causal, + ) + def _get_active_parallel_strategy(self): """Get the parallel strategy based on current SP active state. @@ -279,19 +361,94 @@ class Attention(nn.Module): ) -> torch.Tensor: # Get the appropriate parallel strategy based on SP active state strategy = self._get_active_parallel_strategy() + paged_adapter = self._active_paged_kv_adapter() + in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile + if ( + self._scheduler_paged_kv + and self.paged_kv_cache_role is not None + and paged_adapter is None + and not in_kv_memory_profile + ): + raise RuntimeError( + "Scheduler-paged diffusion attention reached model forward without an active Worker adapter. " + "Only the startup KV memory profile may execute before paged KV initialization." + ) + use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None + if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False): + raise NotImplementedError( + f"Diffusion paged KV requires an Omni backend with paged support; " + f"selected {self.attn_backend.get_name()}" + ) + if use_paged_attention and strategy is not self._no_parallel_strategy: + strategy_name = strategy.name + if self.use_ring or strategy_name == "ring": + raise NotImplementedError( + "paged Scheduler KV is not supported with Ring attention; use strict Ulysses or no SP" + ) + if strategy_name == "allgather_kv": + raise NotImplementedError("paged Scheduler KV is not supported with AllGather-KV sequence parallelism") + if strategy_name == "ulysses" and get_ulysses_mode(default="strict") != "strict": + raise NotImplementedError("paged Scheduler KV currently supports only strict Ulysses") # 1. Prepare inputs (Communication / Resharding) # For Ulysses: AllToAll Q/K/V; Slicing joint_q/k/v # For Ring: Concat joint_q query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata) - attn_metadata = self._with_kv_cache_dtype(attn_metadata) + # Scheduler rows describe the logical sequence, while strict Ulysses + # may append synthetic tokens solely to make the image shard divisible. + # Remove those tokens after the all-to-all and put zero placeholders + # back before the reverse all-to-all. + paged_sp_padding = 0 + paged_sp_padding_offset = query.shape[1] + if use_paged_attention and strategy is not self._no_parallel_strategy and strategy.name == "ulysses": + forward_ctx = get_forward_context() if is_forward_context_available() else None + paged_sp_padding = int(getattr(forward_ctx, "sp_padding_size", 0)) + if paged_sp_padding: + joint_len = int(getattr(ctx, "joint_len", 0)) + joint_strategy = str(getattr(ctx, "joint_strategy", "front")) + paged_sp_padding_offset = query.shape[1] - (joint_len if joint_strategy == "rear" else 0) + padding_start = paged_sp_padding_offset - paged_sp_padding + if padding_start < 0: + raise ValueError( + "Paged Ulysses padding exceeds the post-all-to-all sequence: " + f"padding={paged_sp_padding}, sequence={query.shape[1]}" + ) - # 2. Kernel Execution (Computation) - if self.use_ring and strategy is not self._no_parallel_strategy: - out = self._run_ring_attention(query, key, value, attn_metadata) + def _remove_paged_sp_padding(tensor: torch.Tensor) -> torch.Tensor: + return torch.cat( + (tensor[:, :padding_start], tensor[:, paged_sp_padding_offset:]), + dim=1, + ).contiguous() + + query = _remove_paged_sp_padding(query) + key = _remove_paged_sp_padding(key) + value = _remove_paged_sp_padding(value) + + # 2. This is the shared GPU/NPU boundary. The Worker adapter prepares + # the native page-table context after SP has produced rank-local Q/K/V; + # backend resolution below selects CUDA or Ascend execution. + if use_paged_attention: + assert paged_adapter is not None + paged_kv_context = paged_adapter.prepare_layer_context( + self.prefix, + query, + key, + value, + omni_attn_metadata=attn_metadata, + ) + out = self.attention.forward_paged(paged_kv_context) else: - out = self._run_local_attention(query, key, value, attn_metadata) + attn_metadata = self._with_kv_cache_dtype(attn_metadata) + if self.use_ring and strategy is not self._no_parallel_strategy: + out = self._run_ring_attention(query, key, value, attn_metadata) + else: + out = self._run_local_attention(query, key, value, attn_metadata) + + if paged_sp_padding: + padding_start = paged_sp_padding_offset - paged_sp_padding + output_padding = out.new_zeros((out.shape[0], paged_sp_padding, *out.shape[2:])) + out = torch.cat((out[:, :padding_start], output_padding, out[:, padding_start:]), dim=1) # 3. Post-processing (Reverse Communication) # For Ulysses: AllToAll Output, and AllGather Joint Output @@ -299,6 +456,19 @@ class Attention(nn.Module): return out + @staticmethod + def _active_paged_kv_adapter(): + """Return the Worker adapter selected by Runner-owned metadata.""" + + if not is_forward_context_available(): + return None + return getattr(get_forward_context(), "paged_kv_adapter", None) + + def is_paged_kv_active(self) -> bool: + """Return whether this layer will use Scheduler-managed paged KV.""" + + return self.paged_kv_cache_role is not None and self._active_paged_kv_adapter() is not None + def _run_local_attention(self, query, key, value, attn_metadata): self._assert_piecewise_compatible(attn_metadata) @@ -309,7 +479,25 @@ class Attention(nn.Module): ) return self.sdpa_fallback.forward(query, key, value, attn_metadata) - # Fallback to standard attention + in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile + # NPU only: the startup KV-capacity profile needs tensor shapes, not a + # paged attention result. If dense MindIE-SD is absent, SDPA provides + # that profile forward. Formal GPU/NPU paged requests never use this + # branch because their Worker adapter is active. + if ( + self._scheduler_paged_kv + and self.paged_kv_cache_role is not None + and in_kv_memory_profile + and current_omni_platform.is_npu() + and self.attn_backend.get_name() == "FLASH_ATTN" + and not current_omni_platform.supports_diffusion_dense_flash_attention() + ): + logger.warning_once( + "The startup KV memory profile is using SDPA because MindIE-SD is unavailable. " + "Formal paged requests still use the platform-native paged attention backend." + ) + return self.sdpa_fallback.forward(query, key, value, attn_metadata) + return self.attention.forward(query, key, value, attn_metadata) def _assert_piecewise_compatible(self, attn_metadata: AttentionMetadata | None) -> None: diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py new file mode 100644 index 000000000..08e9e0e4a --- /dev/null +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused MiniMax H3 modulation with FP32 accumulation.""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _indexed_scale_shift_kernel( + output_ptr, + x_ptr, + shift_ptr, + scale_ptr, + indices_ptr, + hidden_size, + stride_x_row, + stride_shift_row, + stride_scale_row, + stride_indices, + block_n: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.arange(0, block_n) + mask = columns < hidden_size + index = tl.load(indices_ptr + row * stride_indices) + + x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32) + scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32) + + tl.store( + output_ptr + row * stride_x_row + columns, + x * (1.0 + scale) + shift, + mask=mask, + ) + + +@triton.jit +def _indexed_gate_kernel( + output_ptr, + x_ptr, + gate_ptr, + other_ptr, + indices_ptr, + hidden_size, + stride_output_row, + stride_x_row, + stride_gate_row, + stride_other_row, + stride_indices, + block_n: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.arange(0, block_n) + mask = columns < hidden_size + index = tl.load(indices_ptr + row * stride_indices) + + x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32) + other = tl.load(other_ptr + row * stride_other_row + columns, mask=mask, other=0.0).to(tl.float32) + + tl.store( + output_ptr + row * stride_output_row + columns, + x + gate * other, + mask=mask, + ) + + +@triton.jit +def _rms_norm_indexed_scale_shift_kernel( + output_ptr, + x_ptr, + weight_ptr, + shift_ptr, + scale_ptr, + indices_ptr, + hidden_size: tl.constexpr, + eps: tl.constexpr, + stride_x_row, + stride_shift_row, + stride_scale_row, + stride_indices, + block_n: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.arange(0, block_n) + mask = columns < hidden_size + index = tl.load(indices_ptr + row * stride_indices) + + x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32) + weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(x * x, axis=0) / hidden_size + normalized = x * tl.rsqrt(variance + eps) * weight + + shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32) + scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32) + tl.store( + output_ptr + row * hidden_size + columns, + normalized * (1.0 + scale) + shift, + mask=mask, + ) + + +@triton.jit +def _indexed_gate_rms_norm_scale_shift_kernel( + residual_out_ptr, + modulated_out_ptr, + residual_ptr, + gate_ptr, + branch_ptr, + weight_ptr, + shift_ptr, + scale_ptr, + indices_ptr, + hidden_size: tl.constexpr, + eps: tl.constexpr, + stride_residual_row, + stride_gate_row, + stride_branch_row, + stride_shift_row, + stride_scale_row, + stride_indices, + block_n: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.arange(0, block_n) + mask = columns < hidden_size + index = tl.load(indices_ptr + row * stride_indices) + + residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32) + branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0).to(tl.float32) + updated = residual + gate * branch + tl.store(residual_out_ptr + row * hidden_size + columns, updated, mask=mask) + + weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(updated * updated, axis=0) / hidden_size + normalized = updated * tl.rsqrt(variance + eps) * weight + shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32) + scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32) + tl.store( + modulated_out_ptr + row * hidden_size + columns, + normalized * (1.0 + scale) + shift, + mask=mask, + ) + + +def indexed_scale_shift_( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, +) -> torch.Tensor: + """Apply indexed scale/shift in-place to a disposable contiguous input.""" + if x.is_cpu: + x.copy_((x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(x.dtype)) + return x + rows, hidden_size = x.shape + if rows == 0: + return x + _indexed_scale_shift_kernel[(rows,)]( + x, + x, + shift, + scale, + indices, + hidden_size, + x.stride(0), + shift.stride(0), + scale.stride(0), + indices.stride(0), + block_n=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return x + + +def indexed_gate( + x: torch.Tensor, + gate: torch.Tensor, + other: torch.Tensor, + indices: torch.Tensor, +) -> torch.Tensor: + """Return ``x + gate[indices] * other`` without indexed temporaries.""" + if x.is_cpu: + return (x + gate.index_select(0, indices) * other).to(x.dtype) + output = torch.empty_like(x) + rows, hidden_size = x.shape + if rows == 0: + return output + _indexed_gate_kernel[(rows,)]( + output, + x, + gate, + other, + indices, + hidden_size, + output.stride(0), + x.stride(0), + gate.stride(0), + other.stride(0), + indices.stride(0), + block_n=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return output + + +def rms_norm_indexed_scale_shift( + x: torch.Tensor, + weight: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Fuse H3 RMSNorm with its indexed AdaLN affine transform.""" + if x.is_cpu: + input_dtype = x.dtype + normalized = x.float() + variance = normalized.pow(2).mean(-1, keepdim=True) + normalized = normalized * torch.rsqrt(variance + eps) + normalized = (weight.float() * normalized).to(input_dtype) + return (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(input_dtype) + output = torch.empty_like(x) + rows, hidden_size = x.shape + if rows: + _rms_norm_indexed_scale_shift_kernel[(rows,)]( + output, + x, + weight, + shift, + scale, + indices, + hidden_size, + eps, + x.stride(0), + shift.stride(0), + scale.stride(0), + indices.stride(0), + block_n=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return output + + +def indexed_gate_rms_norm_scale_shift( + residual: torch.Tensor, + gate: torch.Tensor, + branch: torch.Tensor, + weight: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse gated residual, the following RMSNorm, and AdaLN affine.""" + if residual.is_cpu: + input_dtype = residual.dtype + residual_out = (residual + gate.index_select(0, indices) * branch).to(input_dtype) + normalized = residual_out.float() + variance = normalized.pow(2).mean(-1, keepdim=True) + normalized = normalized * torch.rsqrt(variance + eps) + normalized = (weight.float() * normalized).to(input_dtype) + modulated_out = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to( + input_dtype + ) + return residual_out, modulated_out + residual_out = torch.empty_like(residual) + modulated_out = torch.empty_like(residual) + rows, hidden_size = residual.shape + if rows: + _indexed_gate_rms_norm_scale_shift_kernel[(rows,)]( + residual_out, + modulated_out, + residual, + gate, + branch, + weight, + shift, + scale, + indices, + hidden_size, + eps, + residual.stride(0), + gate.stride(0), + branch.stride(0), + shift.stride(0), + scale.stride(0), + indices.stride(0), + block_n=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return residual_out, modulated_out + + +__all__ = [ + "indexed_gate", + "indexed_gate_rms_norm_scale_shift", + "indexed_scale_shift_", + "rms_norm_indexed_scale_shift", +] diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py index 76f7a78da..5b644ca45 100644 --- a/vllm_omni/diffusion/attention/parallel/factory.py +++ b/vllm_omni/diffusion/attention/parallel/factory.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -47,6 +47,7 @@ def build_parallel_attention_strategy( ulysses_degree = getattr(p, "ulysses_degree", 1) ring_degree = getattr(p, "ring_degree", 1) allgather_degree = getattr(p, "allgather_degree", 1) + ulysses_a2a_permute = getattr(p, "ulysses_a2a_permute", False) try: sp_group = get_sp_group() @@ -83,6 +84,7 @@ def build_parallel_attention_strategy( scatter_idx=scatter_idx, gather_idx=gather_idx, use_sync=use_sync, + ulysses_a2a_permute=ulysses_a2a_permute, ) # Pure Ring Attention diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index fc903d050..69f6a810e 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -16,6 +16,11 @@ from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGr from vllm_omni.diffusion.forward_context import get_ulysses_mode +def _a2a_permute_enabled(enabled: bool, scatter_idx: int, gather_idx: int, world_size: int) -> bool: + """Whether the fused permute-free all-to-all applies (strict layout only).""" + return enabled and world_size > 1 and scatter_idx == 2 and gather_idx == 1 + + def _ceil_div(n: int, d: int) -> int: return (n + d - 1) // d @@ -181,12 +186,23 @@ class UlyssesParallelAttention: scatter_idx: int, gather_idx: int, use_sync: bool, + ulysses_a2a_permute: bool = False, ) -> None: self._sp_group = sp_group self._ulysses_pg = sp_group.ulysses_group self._scatter_idx = scatter_idx self._gather_idx = gather_idx self._use_sync = use_sync + self._ulysses_a2a_permute = ulysses_a2a_permute + if _a2a_permute_enabled( + ulysses_a2a_permute, + scatter_idx, + gather_idx, + sp_group.ulysses_world_size, + ): + from vllm_omni.diffusion.distributed.a2a_permute import ensure_a2a_permute_available + + ensure_a2a_permute_available() @property def enabled(self) -> bool: @@ -347,9 +363,26 @@ class UlyssesParallelAttention: ) # (bs, seq_len/P, head_cnt, head_size) -> (bs, seq_len, head_cnt/P, head_size) - query = SeqAllToAll4D.apply(self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync) - key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync) - value = SeqAllToAll4D.apply(self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync) + if _a2a_permute_enabled( + self._ulysses_a2a_permute, + self._scatter_idx, + self._gather_idx, + ulysses_world_size, + ): + from vllm_omni.diffusion.distributed.a2a_permute import ulysses_qkv_fwd + + group_name = self._ulysses_pg.group_name + query = ulysses_qkv_fwd(query, group_name, ulysses_world_size) + key = ulysses_qkv_fwd(key, group_name, ulysses_world_size) + value = ulysses_qkv_fwd(value, group_name, ulysses_world_size) + else: + query = SeqAllToAll4D.apply( + self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync + ) + key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync) + value = SeqAllToAll4D.apply( + self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync + ) seq_lens = [] local_seq_len = 0 orig_head_cnt = 0 @@ -456,6 +489,19 @@ class UlyssesParallelAttention: orig_head_cnt=ctx.orig_head_cnt, use_sync=ctx.use_sync, ) + elif _a2a_permute_enabled( + self._ulysses_a2a_permute, + ctx.scatter_idx, + ctx.gather_idx, + dist.get_world_size(ctx.ulysses_pg), + ): + from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev + + output_img = ulysses_o_rev( + output_img, + ctx.ulysses_pg.group_name, + dist.get_world_size(ctx.ulysses_pg), + ) else: output_img = SeqAllToAll4D.apply( ctx.ulysses_pg, output_img, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync @@ -488,4 +534,17 @@ class UlyssesParallelAttention: orig_head_cnt=ctx.orig_head_cnt, use_sync=ctx.use_sync, ) + if _a2a_permute_enabled( + self._ulysses_a2a_permute, + ctx.scatter_idx, + ctx.gather_idx, + dist.get_world_size(ctx.ulysses_pg), + ): + from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev + + return ulysses_o_rev( + attn_output, + ctx.ulysses_pg.group_name, + dist.get_world_size(ctx.ulysses_pg), + ) return SeqAllToAll4D.apply(ctx.ulysses_pg, attn_output, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync) diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 5f1ca1d22..7f0242d57 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Model-specific extractors for TeaCache. @@ -13,6 +13,7 @@ all model-specific information needed for generic caching, including preprocessi transformer execution, and postprocessing logic. """ +import math from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -723,11 +724,11 @@ def extract_flux2_klein_context( def extract_longcat_context( module: nn.Module, # LongCatImageTransformer2DModel hidden_states, - timestep, - guidance, - encoder_hidden_states, - txt_ids, - img_ids, + encoder_hidden_states=None, + timestep=None, + img_ids=None, + txt_ids=None, + guidance=None, **kwargs, ) -> CacheContext: """Extract the cache context for LongCat Image. @@ -1259,11 +1260,10 @@ def extract_minimax_h3_context( preprocessing via ``_embed``, cacheable ``blocks`` loop, and ``final_layer`` postprocessing with row selection and update masks. """ + from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import indexed_scale_shift_ from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( - _BF16_DTYPE, _FORWARD_SUPPORTED_KWARGS, MINIMAX_H3_ADALN_MODALITY_NUM, - _modulate_scale_shift, _required_kwarg, ) @@ -1304,6 +1304,9 @@ def extract_minimax_h3_context( psp = _required_kwarg(kwargs, "packed_seq_params") cu_seqlens = module._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32) max_seqlen = int(module._psp_field(psp, "packed_seq_params", "max_seqlen_q")) + # Mirror forward()'s host-side scalar so the refiner sees the number of + # packed requests without reading cu_seqlens off-device. + num_requests = int(module._psp_optional(psp, "num_requests", 1)) refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") refiner_cu = module._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32) refiner_max = int(module._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")) @@ -1317,8 +1320,20 @@ def extract_minimax_h3_context( if inverse_indices.shape[0] != seq_len: raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}") device = x.device - - rope_freqs = module.rope(img_position_ids).to(device) + local_span = module._rope_local_span(seq_len) + local_start, local_len = local_span + rope_table = kwargs.get("rope_table") + if rope_table is None: + rope_table = module.prepare_rope_table( + img_position_ids, + seq_len=seq_len, + ) + else: + module._validate_prepared_rope_table( + rope_table, + local_len=local_len, + device=device, + ) decoder_input, t_emb = module._embed( x=x, @@ -1330,50 +1345,103 @@ def extract_minimax_h3_context( text_pos=text_pos.to(device), refiner_cu_seqlens=refiner_cu.to(device), refiner_max_seqlen=refiner_max, + num_requests=num_requests, seq_len=seq_len, device=device, + local_span=local_span, ) combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device) inverse_indices = inverse_indices.to(device) cu_seqlens = cu_seqlens.to(device) + # In strict SP, _embed already returns this rank's packed rows while + # combined_indices still describes the full packed sequence. TeaCache uses + # the first block's modulation as part of its rank-local cache state, so + # the modulation indices must be sliced to the same local layout. Keep the + # full indices below: local_sp_prepare owns the global-to-local conversion + # for transformer-block inputs. + state_combined_indices = combined_indices.narrow(0, local_start, local_len) + if local_len == seq_len: + state_combined_indices = combined_indices + shift_msa, scale_msa, *_ = module.blocks[0].adaln_proj(t_emb) - modulated_hidden = module.blocks[0].norm1(decoder_input) - modulated_input = _modulate_scale_shift( - modulated_hidden, + modulated_input = module.blocks[0].norm1(decoder_input) + modulated_input = indexed_scale_shift_( + modulated_input, shift_msa, scale_msa, - combined_indices, - dtype=_BF16_DTYPE, + state_combined_indices, ) - def run_transformer_blocks() -> tuple[torch.Tensor, ...]: - hidden, block_rope, block_combined = module.sp_prepare( - decoder_input, - rope_freqs, - combined_indices, + def synchronize_cache_decision(local_should_compute: bool) -> bool: + if local_len == seq_len: + return local_should_compute + + from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + decision = torch.tensor( + int(local_should_compute), + dtype=torch.int32, + device=device, ) + get_sp_group().all_reduce(decision, op=torch.distributed.ReduceOp.MAX) + return bool(decision.item()) + + def run_transformer_blocks() -> tuple[torch.Tensor, ...]: + if local_len == seq_len: + hidden, block_rope, block_combined = module.sp_prepare( + decoder_input, + rope_table, + combined_indices, + ) + else: + hidden, block_rope, block_combined = module.local_sp_prepare( + decoder_input, + rope_table, + combined_indices, + ) for block in module.blocks: hidden = block( hidden, t_emb=t_emb, combined_indices=block_combined, - rope_freqs=block_rope, + rope_table=block_rope, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=seq_len, video_layout=video_layout, ) - hidden = module.sp_gather(hidden) + # TeaCache stores residuals as block_output - decoder_input. Strict SP + # must keep both tensors rank-local until postprocess projects local + # rows to compact logits and gathers them exactly once. + if local_len == seq_len: + hidden = module.sp_gather(hidden) return (hidden,) def postprocess(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - video_logits, audio_logits = module.final_layer( - hidden, - t_emb=t_emb, - inverse_indices=inverse_indices, - ) + if local_len == seq_len: + video_logits, audio_logits = module.final_layer( + hidden, + t_emb=t_emb, + inverse_indices=inverse_indices, + ) + else: + local_inverse_indices = inverse_indices.narrow( + 0, + local_start, + local_len, + ) + video_logits, audio_logits = module.final_layer( + hidden, + t_emb=t_emb, + inverse_indices=local_inverse_indices, + ) + compact_logits = torch.cat((video_logits, audio_logits), dim=-1) + compact_logits = module.sp_gather(compact_logits) + video_width = module.arch.latents_dim * math.prod(module.arch.patch_size) + video_logits = compact_logits[..., :video_width] + audio_logits = compact_logits[..., video_width:] video_logits = video_logits.index_select(0, infer_out_pos.to(device)) audio_logits = audio_logits.index_select(0, audio_pos.to(device)) if not skip_mask_out_condition: @@ -1392,6 +1460,9 @@ def extract_minimax_h3_context( temb=t_emb, run_transformer_blocks=run_transformer_blocks, postprocess=postprocess, + extra_states={ + "synchronize_cache_decision": synchronize_cache_decision, + }, ) diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py index 85a42adbb..b60e17836 100644 --- a/vllm_omni/diffusion/cache/teacache/hook.py +++ b/vllm_omni/diffusion/cache/teacache/hook.py @@ -138,7 +138,16 @@ class TeaCacheHook(ModelHook): state = self.state_manager.get_state() # Decide whether to compute or cache based on modulated input similarity - should_compute = self._should_compute_full_transformer(state, ctx.modulated_input) + local_should_compute = self._should_compute_full_transformer(state, ctx.modulated_input) + sync_cache_decision = (ctx.extra_states or {}).get("synchronize_cache_decision") + if sync_cache_decision is not None: + should_compute = sync_cache_decision(local_should_compute) + # A rank that locally chose the cache path must reset its counter + # when another SP rank requires a full collective block execution. + if should_compute and not local_should_compute: + state.accumulated_rel_l1_distance = 0.0 + else: + should_compute = local_should_compute if not should_compute and state.previous_residual is not None: # ============================================================================ diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 4a050ab38..11fc2c113 100644 --- a/vllm_omni/diffusion/compile.py +++ b/vllm_omni/diffusion/compile.py @@ -56,26 +56,50 @@ def regionally_compile( # Build all compiled callables before mutating the model. This keeps setup # failures atomic: callers can safely continue with the uncompiled model if # torch.compile raises synchronously for any repeated block. - compiled_forwards: list[tuple[nn.Module, Any]] = [] + # Keep the pre-hook callable separately. DLO and the ordinary layerwise + # offloader install a HookRegistry wrapper in ``module.forward``. That + # wrapper performs stream/event synchronization and storage rebinding, so + # compiling it pulls offload control flow into the graph. Compile the + # original block compute instead and leave the wrapper outside the graph. + compiled_forwards: list[tuple[nn.Module, Any | None, Any]] = [] for name, submod in model.named_modules(): if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs): # Compile the block compute while keeping nn.Module.__call__ hooks - # outside the compiled graph. NOTE: anything that wraps this - # callable must stay signature-transparent — cache-dit's - # BlockAdapter matches blocks by inspecting forward's parameter - # names and return annotation, which torch.compile preserves. + # outside the compiled graph. If a HookRegistry is already + # installed, ``submod.forward`` is the hook dispatcher and the + # original callable is stored in ``_omni_original_forward``. + # NOTE: anything that wraps this callable must stay + # signature-transparent — cache-dit's BlockAdapter matches blocks + # by inspecting forward's parameter names and return annotation, + # which torch.compile preserves. + original_forward = getattr(submod, "_omni_original_forward", None) + forward_to_compile = original_forward if original_forward is not None else submod.forward compiled_forwards.append( ( submod, - torch.compile(submod.forward, *compile_args, **compile_kwargs), + original_forward, + torch.compile(forward_to_compile, *compile_args, **compile_kwargs), ) ) if not compiled_forwards: logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.") else: - for submod, compiled_forward in compiled_forwards: - submod.forward = compiled_forward + for submod, original_forward, compiled_forward in compiled_forwards: + if original_forward is None: + submod.forward = compiled_forward + continue + + # Keep HookRegistry's dispatcher as ``forward`` and replace only + # the callable it dispatches to. Hooks such as MagCache retain a + # reference to the original callable as well, so update those + # references when they still point at the pre-compile function. + submod._omni_original_forward = compiled_forward + registry = getattr(submod, "_hook_registry", None) + for hook in getattr(registry, "_hooks", {}).values(): + fn_ref = getattr(hook, "fn_ref", None) + if fn_ref is not None and fn_ref.original_forward is original_forward: + fn_ref.original_forward = compiled_forward logger.info( "Regional compilation applied to %d module(s) for repeated blocks %s.", len(compiled_forwards), diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 137fb9860..4c9c37c35 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1,13 +1,13 @@ # adapted from sglang and fastvideo # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import copy import math import os import random from collections.abc import Callable, Mapping from dataclasses import dataclass, field, fields -from enum import Enum +from enum import Enum, StrEnum from typing import TYPE_CHECKING, Any import diffusers @@ -77,11 +77,13 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any] # Handle "diffusion_attention_backend" shorthand: merge into # diffusion_attention_config before field filtering. diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None) - if diffusion_attn_backend is not None: + fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None) + if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None: existing = normalized.get("diffusion_attention_config") normalized["diffusion_attention_config"] = parse_attention_config( existing, attention_backend=diffusion_attn_backend, + fastvideo_vsa_topk=fastvideo_vsa_topk, ) # Check environment variable as fallback for cache_backend. @@ -98,12 +100,42 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any] return normalized +def validate_host_weight_runtime_options(*, mode: object, root: object) -> None: + """Validate HWR policy without touching the configured storage domain. + + Filesystem locality and store construction belong to the eligible loader + path. Keeping this check purely structural is what lets disabled and + AllGather DLO configurations retain zero HWR interaction. + """ + if mode not in {"disabled", "preferred", "required"}: + raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required") + if mode != "disabled" and (not isinstance(root, str) or not root.strip()): + raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root") + + +def validate_dlo_host_registration_options( + *, + limit_gib: object, + enable_dlo: bool, + use_allgather: bool, + hwr_mode: object, +) -> float: + """Validate the optional transport budget without probing CUDA or HWR.""" + value = float(limit_gib) + if not math.isfinite(value) or value < 0: + raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0") + if value and (not enable_dlo or use_allgather or hwr_mode == "disabled"): + raise ValueError("dlo_host_registration_limit_gib requires enabled no-AllGather DLO and Host Weight Runtime") + return value + + def parse_kv_cache_skip_selector( selector: str | list[int] | tuple[int, ...] | set[int] | None, ) -> set[int] | None: """Parse a non-negative index selector such as "0-9,20,25-30".""" if selector is None: return None + values: set[int] if isinstance(selector, set): values = selector elif isinstance(selector, (list, tuple)): @@ -112,7 +144,7 @@ def parse_kv_cache_skip_selector( text = selector.strip() if not text: return None - values: set[int] = set() + values = set() for chunk in text.split(","): token = chunk.strip() if not token: @@ -154,8 +186,8 @@ class DiffusionParallelConfig: pipeline_parallel_size: int = 1 """Number of pipeline parallel stages.""" - data_parallel_size: int = 1 - """Number of data parallel groups.""" + data_parallel_size: int | None = None + """Number of data parallel groups. Inferred from WORLD when omitted.""" tensor_parallel_size: int = 1 """Number of tensor parallel groups.""" @@ -191,6 +223,9 @@ class DiffusionParallelConfig: sequence shapes across the ring group. """ + ulysses_a2a_permute: bool = False + """Use fused permute-free all-to-all for eligible strict Ulysses exchanges.""" + cfg_parallel_size: int = 1 """Number of ranks used to execute guidance passes in parallel.""" @@ -243,8 +278,9 @@ class DiffusionParallelConfig: def _validate_parallel_config(self) -> Self: """Validates the config relationships among the parallel strategies.""" assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0" - assert self.data_parallel_size > 0, "Data parallel size must be > 0" + assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0" assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0" + assert self.sequence_parallel_size is not None assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0" assert self.ulysses_degree > 0, "Ulysses degree must be > 0" assert self.ring_degree > 0, "Ring degree must be > 0" @@ -283,10 +319,12 @@ class DiffusionParallelConfig: self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree ) - # Calculate world_size from other parallelism dimensions + # Until the runtime WORLD size is known, an omitted DP dimension means + # one replica. OmniDiffusionConfig resolves it against num_gpus below. + data_parallel_size = self.data_parallel_size or 1 other_parallel_world_size = ( self.pipeline_parallel_size - * self.data_parallel_size + * data_parallel_size * self.tensor_parallel_size * self.sequence_parallel_size * self.cfg_parallel_size @@ -297,13 +335,17 @@ class DiffusionParallelConfig: # 1. Standalone: when other parallelism is all 1, HSDP determines world_size # 2. Combined: HSDP overlays on top of other parallelism if self.use_hsdp: - if self.tensor_parallel_size > 1 or self.data_parallel_size > 1: - raise ValueError( - "HSDP (use_hsdp=True) cannot be used with TP or DP " - f"(tensor_parallel_size={self.tensor_parallel_size}, " - f"data_parallel_size={self.data_parallel_size}). " - "Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP." - ) + incompatible = [] + if self.tensor_parallel_size > 1: + incompatible.append("TP") + if data_parallel_size > 1: + incompatible.append("DP") + if self.pipeline_parallel_size > 1: + incompatible.append("PP") + if self.enable_expert_parallel: + incompatible.append("EP") + if incompatible: + raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible)) if self.hsdp_shard_size == -1: # Auto-calculate: use other_parallel_world_size as shard_size if self.hsdp_replicate_size <= 0: @@ -338,6 +380,43 @@ class DiffusionParallelConfig: else: self.world_size = other_parallel_world_size + def resolve_data_parallel_size(self, world_size: int) -> int: + """Resolve or validate diffusion DP against the actual WORLD size.""" + if world_size <= 0: + raise ValueError(f"WORLD size must be > 0, but got {world_size}") + + if self.use_hsdp: + if self.data_parallel_size not in (None, 1): + raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1") + expected_world_size = self.hsdp_replicate_size * self.hsdp_shard_size + if world_size != expected_world_size: + raise ValueError( + f"WORLD size ({world_size}) must equal HSDP size " + f"({self.hsdp_replicate_size} x {self.hsdp_shard_size} = {expected_world_size})" + ) + self.data_parallel_size = 1 + self.world_size = world_size + return 1 + + assert self.sequence_parallel_size is not None + non_dp_size = ( + self.pipeline_parallel_size + * self.tensor_parallel_size + * self.sequence_parallel_size + * self.cfg_parallel_size + ) + if world_size % non_dp_size != 0: + raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})") + inferred_data_parallel_size = world_size // non_dp_size + if self.data_parallel_size is not None and self.data_parallel_size != inferred_data_parallel_size: + raise ValueError( + f"data_parallel_size ({self.data_parallel_size}) does not match WORLD-derived value " + f"({inferred_data_parallel_size})" + ) + self.data_parallel_size = inferred_data_parallel_size + self.world_size = world_size + return inferred_data_parallel_size + @classmethod def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig": """ @@ -538,7 +617,11 @@ class DiffusionCacheConfig: raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'") -def resolve_model_class_name(model: str | None, diffusion_load_format: str = "default") -> str | None: +def resolve_model_class_name( + model: str | None, + diffusion_load_format: str = "default", + revision: str | None = None, +) -> str | None: """Resolve the diffusion pipeline class name from the model config. Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call @@ -550,11 +633,14 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de if not model: return None - is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"} - # Diffusers models: read _class_name from the pipeline index. - model_index = get_diffusion_model_index(model) + # Diffusers models: read _class_name from the pipeline index. Missing + # local paths can otherwise be interpreted as invalid Hub repo IDs. + try: + model_index = get_diffusion_model_index(model, revision=revision) + except Exception: + model_index = None if model_index is not None: return model_index.get("_class_name") if diffusion_load_format == "diffusers": @@ -562,12 +648,16 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de # Other models: map model_type / architecture from config.json. try: - cfg = get_hf_file_to_dict("config.json", model) or {} + cfg = get_hf_file_to_dict("config.json", model, revision=revision) or {} except Exception: cfg = {} model_type = cfg.get("model_type") architectures = cfg.get("architectures") or [] + from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1 + + if _looks_like_hidream_o1(model, cfg): + return "HiDreamO1ImagePipeline" if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: return "BagelPipeline" if ( @@ -598,6 +688,19 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de return None +def uses_diffusers_adapter(od_config: object) -> bool: + """Return whether execution uses the Diffusers adapter backend. + + A custom pipeline takes precedence over ``diffusion_load_format`` in the + worker, so adapter-specific behavior must only apply when no custom + pipeline override is configured. + """ + return ( + getattr(od_config, "custom_pipeline_args", None) is None + and getattr(od_config, "diffusion_load_format", "default") == "diffusers" + ) + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) @@ -651,8 +754,10 @@ class OmniDiffusionConfig: # in both directions. See docs/features/session_state_manager.md. enable_session_state_manager: bool = False - # Distributed executor backend - distributed_executor_backend: str = "mp" + # Distributed executor backend. ``None`` means auto: ``"uni"`` at + # ``num_gpus == 1``, ``"mp"`` otherwise. Explicit ``"mp"`` keeps a + # worker subprocess (process isolation and RPC timeouts). + distributed_executor_backend: str | None = None nccl_port: int | None = None # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class`` @@ -665,6 +770,9 @@ class OmniDiffusionConfig: # Local Diffusion KV ownership and cache-layout mode. diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY + # Maximum number of native BlockTable rows one public request can own + # (sequences plus independent contexts). The model adapter defines it. + diffusion_kv_max_rows_per_request: int | None = None # Optional override for the diffusion model runner class (import path). # Precedence in the worker: this override > the runner declared by the @@ -707,6 +815,14 @@ class OmniDiffusionConfig: dlo_use_allgather: bool = True # Leading main-DiT blocks kept resident by distributed layerwise offload. dlo_resident_layers: int = 0 + # Final-layout Host Weight Runtime policy. The loader only activates this + # for eligible no-AllGather DLO; all other configurations preserve their + # existing loader/storage path. + host_weight_runtime_mode: str = "disabled" + host_weight_runtime_root: str | None = None + # Optional per-worker ceiling for registering final-layout HWR mappings. + # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried. + dlo_host_registration_limit_gib: float = 0.0 pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading @@ -785,6 +901,7 @@ class OmniDiffusionConfig: default_factory=lambda: { "transformer": True, "vae": True, + "text_encoder": True, } ) override_transformer_cls_name: str | None = None @@ -848,6 +965,14 @@ class OmniDiffusionConfig: # Maximum number of sequences to generate in a batch max_num_seqs: int = 1 + # Native vLLM KV-cache sizing inputs used only by paged_scheduler mode. + # These mirror the structured stage cache/scheduler config until all + # diffusion runtime construction consumes VllmOmniConfig directly. + kv_cache_memory_bytes: int | None = None + gpu_memory_utilization: float = 0.9 + max_num_batched_tokens: int | None = None + max_model_len: int | None = None + # Request-mode batch admission: wait briefly for compatible requests to # accumulate in the scheduler waiting queue before the first schedule() of # a wave. Improves fused forward batch sizes under bursty HTTP ingress. @@ -933,6 +1058,23 @@ class OmniDiffusionConfig: if not isinstance(self.diffusion_compile_dynamic, bool): raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}") self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) + if self.diffusion_kv_max_rows_per_request is not None and ( + type(self.diffusion_kv_max_rows_per_request) is not int or self.diffusion_kv_max_rows_per_request <= 0 + ): + raise ValueError("diffusion_kv_max_rows_per_request must be a positive integer when set") + if ( + self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + and self.diffusion_kv_max_rows_per_request is None + ): + raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set") + if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0: + raise ValueError("kv_cache_memory_bytes must be non-negative") + if not 0.0 < self.gpu_memory_utilization <= 1.0: + raise ValueError("gpu_memory_utilization must be in (0, 1]") + if self.max_num_batched_tokens is not None and self.max_num_batched_tokens <= 0: + raise ValueError("max_num_batched_tokens must be positive") + if self.max_model_len is not None and self.max_model_len != -1 and self.max_model_len <= 0: + raise ValueError("max_model_len must be positive or -1") if self.omni_kv_config is None: self.omni_kv_config = {} @@ -944,10 +1086,9 @@ class OmniDiffusionConfig: "need_recv_cache", False ): raise ValueError( - "paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; " - "disable need_recv_cache until connector-aware admission is implemented" + "paged_scheduler Diffusion KV does not support imported AR KV; " + "disable need_recv_cache until connector-aware import is implemented" ) - self.master_port = self._resolve_master_port() self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0) if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0: @@ -980,13 +1121,11 @@ class OmniDiffusionConfig: else: self.num_gpus = 1 - if self.num_gpus < self.parallel_config.world_size: - raise ValueError( - f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})" - ) + self.parallel_config.resolve_data_parallel_size(self.num_gpus) if self.diffusion_compile_granularity == "full": incompatible_features = [] + assert self.parallel_config.sequence_parallel_size is not None if self.parallel_config.use_hsdp: incompatible_features.append("HSDP") if self.parallel_config.sequence_parallel_size > 1: @@ -1060,6 +1199,17 @@ class OmniDiffusionConfig: elif self.max_cpu_loras < 1: raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + validate_host_weight_runtime_options( + mode=self.host_weight_runtime_mode, + root=self.host_weight_runtime_root, + ) + self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options( + limit_gib=self.dlo_host_registration_limit_gib, + enable_dlo=self.enable_distributed_layerwise_offload, + use_allgather=self.dlo_use_allgather, + hwr_mode=self.host_weight_runtime_mode, + ) + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): raise ValueError( "diffusers_load_kwargs and diffusers_call_kwargs are only " @@ -1182,10 +1332,7 @@ class OmniDiffusionConfig: from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index - # Default model_class_name for diffusers adapter - if self.model_class_name is None and self.diffusion_load_format == "diffusers": - self.model_class_name = "DiffusersAdapterPipeline" - + assert self.model is not None try: config_dict = get_diffusion_model_index( self.model, @@ -1194,6 +1341,8 @@ class OmniDiffusionConfig: if config_dict is not None: if self.model_class_name is None: self.model_class_name = config_dict.get("_class_name", None) + if self.model_class_name is None and self.diffusion_load_format == "diffusers": + self.model_class_name = "DiffusersAdapterPipeline" self.update_multimodal_support() # Skip transformer config loading for diffusers adapter @@ -1210,9 +1359,20 @@ class OmniDiffusionConfig: exc, ) else: - tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) + from vllm_omni.model_extras import get_transformer_config_subfolder + + transformer_subfolder = get_transformer_config_subfolder( + self.model_class_name, + model=self.model, + revision=self.revision, + ) + tf_config_dict = get_hf_file_to_dict( + f"{transformer_subfolder}/config.json", + self.model, + revision=self.revision, + ) if tf_config_dict is None: - tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model) + tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision) if tf_config_dict is not None: self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) else: @@ -1223,6 +1383,8 @@ class OmniDiffusionConfig: # Skip transformer config loading for diffusers adapter # (non-DiT models don't have a separate transformer folder/config) if self.diffusion_load_format == "diffusers": + if self.model_class_name is None: + self.model_class_name = "DiffusersAdapterPipeline" self.set_tf_model_config(TransformerConfig()) logger.warning( "Could not find a valid pipeline index per Diffusers format. " @@ -1232,7 +1394,7 @@ class OmniDiffusionConfig: "that require additional inputs." ) else: - cfg = get_hf_file_to_dict("config.json", self.model) + cfg = get_hf_file_to_dict("config.json", self.model, revision=self.revision) if cfg is None: # Lance ships its top-level config.json one directory above # the per-checkpoint subfolders (``Lance_3B/`` or @@ -1249,7 +1411,18 @@ class OmniDiffusionConfig: model_type = cfg.get("model_type") architectures = cfg.get("architectures") or [] - if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1 + + assert self.model is not None + is_hidream_o1 = _looks_like_hidream_o1(self.model, cfg) + if self.model_class_name == "HiDreamO1ImagePipeline" and not is_hidream_o1: + raise ValueError(f"Checkpoint {self.model} does not have the HiDream-O1 signature") + + if is_hidream_o1: + self.model_class_name = "HiDreamO1ImagePipeline" + self.set_tf_model_config(TransformerConfig()) + self.update_multimodal_support() + elif model_type == "bagel" or "BagelForConditionalGeneration" in architectures: self.model_class_name = "BagelPipeline" self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() @@ -1298,12 +1471,19 @@ class OmniDiffusionConfig: elif model_type == "vla": from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero + assert self.model is not None if _looks_like_dreamzero(self.model): self.model_class_name = "DreamZeroPipeline" self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() else: raise + elif cfg.get("type") == "pi0": + # π0 (Pi-Zero) VLA — the LeRobot config.json uses ``type: "pi0"``. + if self.model_class_name is None: + self.model_class_name = "Pi0Pipeline" + self.set_tf_model_config(TransformerConfig()) + self.update_multimodal_support() elif architectures and len(architectures) == 1: architecture = architectures[0] from vllm_omni.diffusion.registry import DiffusionModelRegistry @@ -1368,6 +1548,9 @@ class DiffusionOutput: # memory usage info peak_memory_mb: float = 0.0 + # KV-recv wall-clock (ms) from the runner; feeds vllm_omni:diffusion_kv_load_s. + kv_recv_ms: float = 0.0 + # When True, move tensor fields to CPU at construction time. Useful when # the output is shipped across process boundaries (e.g. step-execution # mode) and the receiving side must not initialise a stray CUDA context. @@ -1490,10 +1673,11 @@ class AttnQuantSpec: for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)): if v is not None and v not in self._VALID_DTYPES: raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.") - for name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)): - if v not in self._VALID_BLOCK_SIZES: + for block_name, block_value in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)): + if block_value not in self._VALID_BLOCK_SIZES: raise ValueError( - f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}." + f"quant.{block_name}={block_value!r} unsupported; " + f"kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}." ) @property @@ -1506,6 +1690,20 @@ class AttnQuantSpec: BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"}) +class RainFusionPrecision(StrEnum): + """Execution precision for block-sparse RainFusion (rf_v3) attention. + + ``bf16``: no quantization, pure BF16 sparse attention (official baseline). + ``fp8``: BSA FP8 path - Hadamard rotation then full FP8 block quantization + of Q/K/V before the BSA kernel. + ``mix``: EagleQBSA mixed precision - Q/K per-block INT8 + V per-channel FP8. + """ + + BF16 = "bf16" + FP8 = "fp8" + MIX = "mix" + + @dataclass class BlockSparseSpec: """User-facing controls shared by block-sparse attention backends. @@ -1514,10 +1712,13 @@ class BlockSparseSpec: ``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an index selector such as "0-3,38") exempts individual DiT blocks. Those two are the accuracy knobs to trade back quality at a fixed ``sparsity``. + ``precision`` selects the kernel precision mode (see ``RainFusionPrecision``). """ sparsity: float = 0.8 start_step: int = 0 + end_step: int = 0 + precision: str = RainFusionPrecision.BF16.value skip_layers: str | list[int] | None = None skip_layer_indices: set[int] | None = field(default=None, repr=False) @@ -1525,6 +1726,13 @@ class BlockSparseSpec: self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0 if self.start_step < 0: raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.") + if self.end_step < 0: + raise ValueError(f"block_sparse.end_step must be >= 0; got {self.end_step!r}.") + if self.precision not in {p.value for p in RainFusionPrecision}: + raise ValueError( + f"block_sparse.precision must be one of {sorted(p.value for p in RainFusionPrecision)}; " + f"got {self.precision!r}." + ) self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers) @@ -1535,6 +1743,7 @@ class AttentionSpec: backend: str skip_softmax: SkipSoftmaxSpec | None = None quant: AttnQuantSpec | None = None + fastvideo_vsa_topk: int | None = None block_sparse: BlockSparseSpec | None = None skip_calibration: dict | None = field(default=None, repr=False) @@ -1554,6 +1763,11 @@ class AttentionSpec: f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but " f"backend={self.backend!r}. Remove quant or set a supported backend." ) + if self.fastvideo_vsa_topk is not None: + if self.backend.upper() != "FASTVIDEO_VSA": + raise ValueError("fastvideo_vsa_topk is only supported by the FASTVIDEO_VSA backend.") + if self.fastvideo_vsa_topk <= 0: + raise ValueError("fastvideo_vsa_topk must be positive.") if self.backend.upper() in BLOCK_SPARSE_BACKENDS: # Selecting the backend is the opt-in; without an explicit block the # defaults apply rather than silently running dense. @@ -1595,12 +1809,17 @@ class AttentionSpec: if q.flashinfer_backend is not None: quant_kw["flashinfer_backend"] = q.flashinfer_backend kw["quant"] = quant_kw + if self.fastvideo_vsa_topk is not None: + kw["topk"] = self.fastvideo_vsa_topk if self.block_sparse is not None: bs = self.block_sparse kw["sparsity"] = bs.sparsity kw["start_step"] = bs.start_step + kw["end_step"] = bs.end_step + kw["precision"] = bs.precision if bs.skip_layer_indices: kw["skip_layers"] = sorted(bs.skip_layer_indices) + return kw or None @@ -1670,7 +1889,7 @@ class AttentionConfig: normalized[role] = node return - spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"} + spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse"} node_dict = dict(node) node_keys = set(node_dict) if node_keys & spec_keys: @@ -1709,6 +1928,7 @@ def parse_attention_config( attention_config: AttentionConfig | Mapping[str, Any] | None = None, *, attention_backend: str | None = None, + fastvideo_vsa_topk: int | None = None, ) -> AttentionConfig: """Pure type-conversion: coerce *attention_config* to an AttentionConfig. @@ -1736,6 +1956,14 @@ def parse_attention_config( if attention_backend.lower() != "auto": normalized.default = AttentionSpec(backend=attention_backend) + if fastvideo_vsa_topk is not None: + if normalized.default is None: + raise ValueError("--fastvideo-vsa-topk requires --diffusion-attention-backend FASTVIDEO_VSA.") + if normalized.default.backend.upper() != "FASTVIDEO_VSA": + raise ValueError("--fastvideo-vsa-topk is only valid with the FASTVIDEO_VSA backend.") + normalized.default.fastvideo_vsa_topk = fastvideo_vsa_topk + normalized.default.__post_init__() + return normalized diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 08218122e..0929e9419 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import asyncio import concurrent.futures +import copy import inspect +import os import queue import threading import time @@ -17,18 +19,24 @@ from typing import TYPE_CHECKING, Any import numpy as np import PIL.Image import torch +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.kv_cache_interface import KVCacheConfig from vllm_omni.diffusion.data import ( DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig, + uses_diffusers_adapter, ) +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode +from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.io_support import ( get_dummy_run_num_frames, + get_dummy_run_num_image_inputs, image_color_format, supports_audio_output, supports_multimodal_input, @@ -49,13 +57,61 @@ from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput from vllm_omni.errors import client_error_from_metadata, is_client_error_status from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.metrics.utils import ( + diffusion_scheduler_waiting_metrics, + extract_diffusion_denoise_ms, + extract_diffusion_vae_decode_ms, +) if TYPE_CHECKING: from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -_ASYNC_OUTPUT_TIMEOUT = 30.0 # seconds +_ASYNC_OUTPUT_TIMEOUT_ENV = "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT" +_ASYNC_OUTPUT_TIMEOUT_DEFAULT = 600.0 # seconds + + +def _async_output_timeout() -> float: + """Seconds to wait for one step's background D2H/SHM copy. + + The copy itself finishes in milliseconds, but it is queued behind the GPU + work for that step, so the wall-clock wait tracks step time — a single-GPU + box legitimately runs tens of seconds per step on large shapes. A tight + bound therefore does not catch a hung engine (worker death and a dead + result pump are surfaced by the worker monitor and ``check_health``); it + only aborts renders that are still making progress, throwing away the + denoise that already completed. The default matches + ``_DLO_DP_WAVE_TIMEOUT_S`` in the same subsystem. + + Resolved here rather than at import so a malformed value degrades to the + default instead of raising on the request path: this runs inside + ``step_streaming``/``add_req_and_wait_for_response``, where a typo in the + environment must not start failing generations. + """ + raw = os.environ.get(_ASYNC_OUTPUT_TIMEOUT_ENV) + if raw is None: + return _ASYNC_OUTPUT_TIMEOUT_DEFAULT + try: + timeout = float(raw) + except ValueError: + logger.warning_once( + "Ignoring %s=%r: not a number. Using the default %.1fs.", + _ASYNC_OUTPUT_TIMEOUT_ENV, + raw, + _ASYNC_OUTPUT_TIMEOUT_DEFAULT, + ) + return _ASYNC_OUTPUT_TIMEOUT_DEFAULT + if timeout <= 0: + logger.warning_once( + "Ignoring %s=%r: must be positive. Using the default %.1fs.", + _ASYNC_OUTPUT_TIMEOUT_ENV, + raw, + _ASYNC_OUTPUT_TIMEOUT_DEFAULT, + ) + return _ASYNC_OUTPUT_TIMEOUT_DEFAULT + return timeout + __all__ = [ "DiffusionEngine", @@ -102,8 +158,14 @@ def _resolve_custom_pipeline_cls(custom_pipeline_args: dict[str, Any] | None) -> def supports_request_batch(od_config: OmniDiffusionConfig) -> bool: model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None)) - if model_cls is None: - model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None)) + if model_cls is not None: + return bool(getattr(model_cls, "supports_request_batch", False)) + + if uses_diffusers_adapter(od_config): + model_cls = DiffusionModelRegistry._try_load_model_cls("DiffusersAdapterPipeline") + return bool(getattr(model_cls, "supports_request_batch", False)) + + model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None)) if model_cls is None: return False return bool(getattr(model_cls, "supports_request_batch", False)) @@ -184,12 +246,39 @@ class DiffusionEngine: from the resolved execution mode. """ self.od_config = od_config + # Set after the paged-KV profile request has gone through model-owned + # preprocessing. Real requests are admitted only within this measured + # activation envelope: (max execution sequences, max seq_len, + # max target_len). + self._diffusion_kv_profile_limits: tuple[int, int, int] | None = None self._init_process_hooks(od_config) self.execution_mode = self._resolve_execution_mode(od_config) self._init_executor(od_config) try: - self._init_scheduler(od_config, scheduler) + profile_requests = self._prepare_diffusion_kv_profile_requests() + kv_control_plane = initialize_diffusion_kv_control_plane( + self.executor, + od_config, + profile_requests=profile_requests, + ) + if kv_control_plane is None: + self._init_scheduler(od_config, scheduler) + else: + ( + kv_cache_config, + scheduler_block_size, + hash_block_size, + kv_vllm_config, + ) = kv_control_plane + self._init_scheduler( + od_config, + scheduler, + kv_cache_config, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + kv_vllm_config=kv_vllm_config, + ) self._init_runtime_state() self._init_execute_fn() self._log_execution_mode(od_config) @@ -242,6 +331,11 @@ class DiffusionEngine: self, od_config: OmniDiffusionConfig, scheduler: BaseScheduler | None = None, + kv_cache_config: KVCacheConfig | None = None, + *, + scheduler_block_size: int | None = None, + hash_block_size: int | None = None, + kv_vllm_config: VllmConfig | None = None, ) -> None: if scheduler is not None: self.scheduler = scheduler @@ -249,7 +343,16 @@ class DiffusionEngine: self.scheduler = StepScheduler() else: self.scheduler = RequestScheduler() - self.scheduler.initialize(od_config) + if kv_cache_config is None: + self.scheduler.initialize(od_config) + else: + self.scheduler.initialize( + od_config, + kv_cache_config=kv_cache_config, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + kv_vllm_config=kv_vllm_config, + ) def _init_runtime_state(self) -> None: # DP multi-concurrency: allow batching dp_size requests so each @@ -284,6 +387,9 @@ class DiffusionEngine: self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() + # Copied onto the existing output metrics payload so queue monitoring + # reuses the normal diffusion result path without additional IPC. + self._scheduler_num_waiting_reqs = 0 def _init_execute_fn(self) -> None: if self.execution_mode == DiffusionExecutionMode.STEP_BATCH: @@ -339,9 +445,28 @@ class DiffusionEngine: # Async mode: wait for background D2H/SHM to complete. if output.async_output_id: fut = self.executor.wait_output_ready(output.async_output_id) - output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=_ASYNC_OUTPUT_TIMEOUT) + timeout = _async_output_timeout() + try: + output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout) + except (TimeoutError, asyncio.TimeoutError): + describe = getattr(self.executor, "describe_pending_state", None) + logger.error( + "Timed out after %.1fs waiting for async output; set %s to a larger value " + "to allow slower steps. Executor state: %s", + timeout, + _ASYNC_OUTPUT_TIMEOUT_ENV, + describe(output.async_output_id) if describe else "unavailable", + ) + raise postprocess_start_time = time.perf_counter() - formatted_outputs = self.postprocess_output(request, output) + scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0)) + try: + formatted_outputs = self.postprocess_output(request, output) + except Exception as exc: + # Preserve the latest scheduler snapshot across terminal + # abort/error paths, which do not produce formatted outputs. + setattr(exc, "diffusion_metrics", scheduler_metrics) + raise postprocess_time = time.perf_counter() - postprocess_start_time step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000 logger.debug( @@ -353,14 +478,24 @@ class DiffusionEngine: step_total_ms, ) for request_output in formatted_outputs: - request_output.metrics.update( - { - "preprocess_time_ms": preprocess_time * 1000, - "diffusion_engine_exec_time_ms": exec_total_time * 1000, - "diffusion_engine_total_time_ms": step_total_ms, - "postprocess_time_ms": postprocess_time * 1000, - } - ) + metrics_update = { + "preprocess_time_ms": preprocess_time * 1000, + "diffusion_engine_exec_time_ms": exec_total_time * 1000, + "postprocess_time_ms": postprocess_time * 1000, + **scheduler_metrics, + } + if request.scheduler_queue_wait_ms is not None: + metrics_update["scheduler_queue_wait_ms"] = request.scheduler_queue_wait_ms + vae_decode_ms = extract_diffusion_vae_decode_ms(output) + if vae_decode_ms is not None: + metrics_update["vae_decode_time_ms"] = vae_decode_ms + forward_ms = extract_diffusion_denoise_ms(output) + if forward_ms is not None: + metrics_update["forward_time_ms"] = forward_ms + kv_recv_ms = getattr(output, "kv_recv_ms", 0.0) + if kv_recv_ms > 0: + metrics_update["kv_recv_time_ms"] = kv_recv_ms + request_output.metrics.update(metrics_update) yield formatted_outputs async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: @@ -451,6 +586,7 @@ class DiffusionEngine: self._wait_for_admission_if_needed_locked() sched_output = self.scheduler.schedule() + self._scheduler_num_waiting_reqs = max(int(sched_output.num_waiting_reqs), 0) if sched_output.is_empty: self._emit_finished_outputs(sched_output.finished_req_ids, None) @@ -578,12 +714,25 @@ class DiffusionEngine: if not task.future.done(): task.future.set_exception(exc) + def _remove_diffusion_kv_requests(self, request_ids: Iterable[str]) -> None: + """Clear terminal Worker rows while Scheduler owns the allocations.""" + + od_config = getattr(self, "od_config", None) + if od_config is None or not is_scheduler_paged_kv_mode( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ): + return + unique_request_ids = list(dict.fromkeys(request_ids)) + if unique_request_ids: + self.executor.remove_diffusion_kv_requests(unique_request_ids) + def _emit_finished_outputs( self, finished_ids: set[str], runner_output: BaseRunnerOutput | None = None, missing_result_error: str = "Diffusion execution finished without a final output", ) -> None: + self._remove_diffusion_kv_requests(finished_ids) for rid in finished_ids: if runner_output is not None: _output = runner_output.get_request_output(rid) @@ -603,6 +752,8 @@ class DiffusionEngine: self._emit_finished_outputs(finished_ids, runner_output) return + self._remove_diffusion_kv_requests(finished_ids) + delivered_finished_req_ids: set[str] = set() # finished_ids may have some requests that are not scheduler in this round. @@ -694,9 +845,36 @@ class DiffusionEngine: """Run model-owned preprocessing once, before entering Engine locks.""" pre_process_func = getattr(self, "pre_process_func", None) - if pre_process_func is None: - return request - return pre_process_func(request) + if pre_process_func is not None: + request = pre_process_func(request) + self._validate_diffusion_kv_profile_limits(request) + return request + + def _validate_diffusion_kv_profile_limits(self, request: OmniDiffusionRequest) -> None: + """Keep admitted paged-KV requests within the profiled activation shape.""" + + profile_limits = getattr(self, "_diffusion_kv_profile_limits", None) + if profile_limits is None: + return + kv_requests = request.diffusion_kv_requests + if not kv_requests: + return + + request_limits = ( + len(kv_requests), + max(kv_request.seq_len for kv_request in kv_requests), + max(kv_request.target_len for kv_request in kv_requests), + ) + if any(actual > profiled for actual, profiled in zip(request_limits, profile_limits, strict=True)): + request_sequences, request_seq_len, request_target_len = request_limits + profile_sequences, profile_seq_len, profile_target_len = profile_limits + raise ValueError( + f"Diffusion KV request {request.request_id!r} exceeds the startup memory-profile envelope: " + f"sequences={request_sequences} (profiled={profile_sequences}), " + f"max_seq_len={request_seq_len} (profiled={profile_seq_len}), " + f"max_target_len={request_target_len} (profiled={profile_target_len}). " + "Reduce the request shape or extend the model's paged-KV profile recipe." + ) def _add_prepared_request(self, request: OmniDiffusionRequest) -> str: """Admit a request whose model-owned preprocessing is complete.""" @@ -769,6 +947,7 @@ class DiffusionEngine: sched_output = self.scheduler.schedule() if sched_output.is_empty: if target_request_id in sched_output.finished_req_ids: + self._remove_diffusion_kv_requests([target_request_id]) return self._finalize_finished_request(target_request_id) if not self.scheduler.has_requests(): raise RuntimeError("Diffusion scheduler has no runnable requests.") @@ -798,6 +977,7 @@ class DiffusionEngine: if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1: raise ValueError("Sync func should receive one result at one time") if target_request_id in finished_req_ids: + self._remove_diffusion_kv_requests([target_request_id]) req_output = runner_output.get_request_output(target_request_id) output = self._finalize_finished_request( target_request_id, @@ -806,7 +986,7 @@ class DiffusionEngine: ) if output.async_output_id: fut = self.executor.wait_output_ready(output.async_output_id) - output = fut.result(timeout=_ASYNC_OUTPUT_TIMEOUT) + output = fut.result(timeout=_async_output_timeout()) return output def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: @@ -860,47 +1040,111 @@ class DiffusionEngine: self.close() raise e - def _dummy_run(self): - """A dummy run to warm up the model.""" - num_inference_steps = 1 - height = 512 - width = 512 - prompt: OmniTextPrompt = {"prompt": "dummy run"} + def _make_dummy_request( + self, + *, + height: int, + width: int, + guidance_scale: float, + num_image_inputs: int = 1, + ) -> OmniDiffusionRequest | None: + """Build a one-step model request for startup profiling or warmup.""" + prompt: OmniTextPrompt = {"prompt": "dummy run"} supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config) if supports_image_input: - # Provide a dummy image input if the model supports it color_format = image_color_format(self.od_config.model_class_name) - dummy_image = PIL.Image.new(color_format, (width, height)) - prompt.setdefault("multi_modal_data", {})["image"] = dummy_image + images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)] + prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images if supports_audio_input: audio_sr = 16000 - dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32) - prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio + prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32) num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) if num_frames <= 0: - logger.info("Skipping dummy warmup run (num_frames=0)") - return - req = OmniDiffusionRequest( + return None + return OmniDiffusionRequest( prompt=prompt, request_id=DUMMY_DIFFUSION_REQUEST_ID, sampling_params=OmniDiffusionSamplingParams( height=height, width=width, - num_inference_steps=num_inference_steps, + num_inference_steps=1, num_frames=num_frames, - # Keep warmup path minimal and robust across text encoders. - # Some models may fail when warmup implicitly triggers - # classifier-free guidance with an empty negative prompt. - guidance_scale=0.0, + guidance_scale=guidance_scale, num_outputs_per_prompt=1, - # Disable CFG for warmup to avoid triggering CFG parallel - # validation when cfg_parallel_size > 1. extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0}, ), ) + + def _prepare_diffusion_kv_profile_requests(self) -> list[OmniDiffusionRequest] | None: + """Prepare the per-rank request batch used to profile paged-KV headroom. + + The profile executes directly on each Worker before the Scheduler and + its KV manager exist. It uses the maximum number of requests that one + rank can execute together. DLO+DP is the exception because each rank + executes one request from the collective wave. + + Hunyuan is currently the only model integrated with + ``paged_scheduler``; 1024x1024, enabled CFG, and the maximum advertised + reference-image count exercise its first-step activation peak. + Admission compares each preprocessed request's CFG count and tokenized + sequence/target shape with the resulting per-request profile envelope. + Future paged model integrations must extend this recipe for their + serving limits (for example, video frame count) rather than reusing it + silently. + """ + + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is not DiffusionKVCacheMode.PAGED_SCHEDULER + ): + return None + request = self._make_dummy_request( + height=1024, + width=1024, + guidance_scale=5.0, + num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name), + ) + if request is None: + raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request") + request = self._prepare_request_for_admission(request) + kv_requests = request.diffusion_kv_requests + if not kv_requests: + raise RuntimeError("paged_scheduler profile preprocessing must produce Diffusion KV requests") + self._diffusion_kv_profile_limits = ( + len(kv_requests), + max(kv_request.seq_len for kv_request in kv_requests), + max(kv_request.target_len for kv_request in kv_requests), + ) + dlo_dp_request_mode = self.execution_mode is DiffusionExecutionMode.REQUEST_BATCH and _uses_dlo_dp_concurrency( + self.od_config + ) + profile_batch_size = 1 if dlo_dp_request_mode else _max_num_seqs(self.od_config) + profile_requests: list[OmniDiffusionRequest] = [] + for index in range(profile_batch_size): + profile_request = copy.copy(request) + profile_request.request_id = f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-{index}" + profile_request.sampling_params = copy.deepcopy(request.sampling_params) + # Preprocessing owns request geometry and prepared model inputs, + # but native KV requests remain Scheduler-only mutable state. The + # profile bypasses Scheduler admission and must not send them to + # the Worker. + profile_request.diffusion_kv_requests = None + profile_requests.append(profile_request) + return profile_requests + + def _dummy_run(self): + """A dummy run to warm up the model.""" + req = self._make_dummy_request( + height=512, + width=512, + guidance_scale=0.0, + ) + if req is None: + logger.info("Skipping dummy warmup run (num_frames=0)") + return logger.info("dummy run to warm up the model") output = self.add_req_and_wait_for_response(req) if output.error: @@ -1085,8 +1329,11 @@ class DiffusionEngine: def _abort_requests(self, request_ids: str | Iterable[str]) -> None: request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids) + request_ids = list(dict.fromkeys(request_ids)) - for request_id in dict.fromkeys(request_ids): + self._remove_diffusion_kv_requests(request_ids) + + for request_id in request_ids: if self.scheduler.get_request_state(request_id) is not None: self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED) @@ -1118,4 +1365,7 @@ class DiffusionEngine: if runner_output is not None and runner_output.async_output_id is not None: return DiffusionOutput(async_output_id=runner_output.async_output_id) + if state.status == DiffusionRequestStatus.FINISHED_ERROR and state.error: + return DiffusionOutput(error=state.error) + return DiffusionOutput(error=missing_result_error) diff --git a/vllm_omni/diffusion/diffusion_kv/config.py b/vllm_omni/diffusion/diffusion_kv/config.py index cd69d8988..0a7d6070b 100644 --- a/vllm_omni/diffusion/diffusion_kv/config.py +++ b/vllm_omni/diffusion/diffusion_kv/config.py @@ -37,3 +37,8 @@ def parse_diffusion_kv_cache_mode(value: object) -> DiffusionKVCacheMode: "use 'dense_legacy' or 'paged_scheduler'" ) return mode + + +def is_scheduler_paged_kv_mode(mode: DiffusionKVCacheMode) -> bool: + """Return whether a parsed cache mode uses Scheduler-owned paging.""" + return mode is DiffusionKVCacheMode.PAGED_SCHEDULER diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py new file mode 100644 index 000000000..31aaaa46e --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/initialization.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from vllm.config import VllmConfig +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, + get_kv_cache_configs, + resolve_kv_cache_block_sizes, +) +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config +from vllm_omni.platforms import current_omni_platform + +if TYPE_CHECKING: + from vllm_omni.diffusion.executor.abstract import DiffusionExecutor + from vllm_omni.diffusion.request import OmniDiffusionRequest + + +def build_native_kv_cache_configs( + vllm_config: VllmConfig, + worker_specs: list[dict[str, KVCacheSpec]], + available_memory: list[int], +) -> tuple[list[KVCacheConfig], KVCacheConfig]: + """Build rank-local and Scheduler-native configs using vLLM utilities.""" + + from vllm.v1.core import single_type_kv_cache_manager + + single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config) + worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory) + return worker_configs, generate_scheduler_kv_cache_config(worker_configs) + + +def initialize_diffusion_kv_control_plane( + executor: DiffusionExecutor, + od_config: OmniDiffusionConfig, + *, + profile_requests: list[OmniDiffusionRequest] | None = None, + device: torch.device | None = None, +) -> tuple[KVCacheConfig, int, int, VllmConfig] | None: + """Run the native Worker-spec to Scheduler-config initialization chain.""" + + if ( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is not DiffusionKVCacheMode.PAGED_SCHEDULER + ): + return None + if not profile_requests: + raise ValueError("paged_scheduler Diffusion KV initialization requires prepared memory profile requests") + if device is None: + device = current_omni_platform.get_torch_device(0) + vllm_config = create_diffusion_vllm_config(device, od_config) + worker_specs = executor.get_kv_cache_specs() + available_memory = executor.determine_available_kv_memory(profile_requests) + if len(worker_specs) != od_config.num_gpus or len(available_memory) != od_config.num_gpus: + raise ValueError( + "Diffusion KV initialization rank count mismatch: " + f"expected={od_config.num_gpus}, specs={len(worker_specs)}, memory={len(available_memory)}" + ) + worker_configs, scheduler_kv_cache_config = build_native_kv_cache_configs( + vllm_config, + worker_specs, + available_memory, + ) + scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( + scheduler_kv_cache_config, + vllm_config, + ) + # Native sizing may auto-fit ``max_model_len=-1`` against the profiled + # memory budget. Forward that final value so every Worker sizes its + # physical BlockTables from the same bound as the Scheduler. + resolved_max_model_len = vllm_config.model_config.max_model_len + if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0: + raise ValueError("native Diffusion KV sizing must resolve a positive max_model_len") + executor.set_kv_cache_configs(worker_configs, resolved_max_model_len) + return ( + scheduler_kv_cache_config, + scheduler_block_size, + hash_block_size, + vllm_config, + ) diff --git a/vllm_omni/diffusion/diffusion_kv/manager.py b/vllm_omni/diffusion/diffusion_kv/manager.py new file mode 100644 index 000000000..eee51b864 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/manager.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Sequence + +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.kv_cache_interface import KVCacheConfig + +from vllm_omni.diffusion.diffusion_kv.metadata import ( + DiffusionKVMetadata, + DiffusionKVSequenceMetadata, +) +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest + + +class DiffusionKVAdmissionError(ValueError): + """A request-specific KV admission failure safe to return to the caller.""" + + +class DiffusionKVCacheManager: + """Atomic public-request facade over native vLLM KVCacheManager.""" + + def __init__( + self, + kv_cache_config: KVCacheConfig, + *, + max_model_len: int, + scheduler_block_size: int, + hash_block_size: int, + max_in_flight_tokens: int | None = None, + ) -> None: + if max_model_len <= 0: + raise ValueError(f"max_model_len must be positive, got {max_model_len}") + if kv_cache_config.num_blocks <= 0 or not kv_cache_config.kv_cache_groups: + raise ValueError("Diffusion KVCacheConfig must contain a positive block pool and at least one group") + if scheduler_block_size <= 0 or hash_block_size <= 0: + raise ValueError("scheduler_block_size and hash_block_size must be positive") + if max_in_flight_tokens is not None and max_in_flight_tokens <= 0: + raise ValueError(f"max_in_flight_tokens must be positive when provided, got {max_in_flight_tokens}") + self.native_manager = KVCacheManager( + kv_cache_config=kv_cache_config, + max_model_len=max_model_len, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + max_in_flight_tokens=max_in_flight_tokens, + enable_caching=False, + ) + self.max_model_len = max_model_len + # Native vLLM may reserve a null block, so an idle BlockPool does not + # necessarily report ``kv_cache_config.num_blocks`` free blocks. + self._empty_pool_num_free_blocks = self.native_manager.block_pool.get_num_free_blocks() + self._requests: dict[str, tuple[DiffusionKVRequest, ...]] = {} + self._metadata: dict[str, DiffusionKVMetadata] = {} + self._internal_request_ids: set[str] = set() + self._next_allocation_generation = 1 + + def _get_empty_pool_required_blocks(self, requests: Sequence[DiffusionKVRequest]) -> int: + """Return the native full-sequence reservation for one public request. + + vLLM performs this calculation inside ``allocate_slots`` when + ``full_sequence_must_fit`` is enabled. Diffusion CFG adds one semantic + requirement that native requests do not have: every execution sequence + must fit atomically. Sum the native per-sequence result so a request + that can never fit is rejected independently of current pool usage. + """ + + coordinator = self.native_manager.coordinator + empty_blocks = self.native_manager.empty_kv_cache_blocks.blocks + return sum( + coordinator.get_num_blocks_to_allocate( + request_id=request.request_id, + num_tokens=request.num_tokens, + new_computed_blocks=empty_blocks, + num_encoder_tokens=0, + total_computed_tokens=0, + num_local_computed_tokens=0, + num_tokens_main_model=request.num_tokens, + apply_admission_cap=True, + ) + for request in requests + ) + + def has_request(self, public_request_id: str) -> bool: + return public_request_id in self._requests + + def reserve_request( + self, + public_request_id: str, + kv_requests: Sequence[DiffusionKVRequest], + ) -> DiffusionKVMetadata | None: + """Reserve all CFG sequences or roll back the complete public request.""" + + if not public_request_id: + raise ValueError("public_request_id must be non-empty") + if public_request_id in self._requests: + raise ValueError(f"Diffusion KV request {public_request_id!r} is already allocated") + requests = tuple(kv_requests) + if not requests: + raise ValueError("Diffusion KV allocation requires at least one sequence") + + sequence_ids = [request.sequence_id for request in requests] + internal_ids = [request.request_id for request in requests] + expected_sequence_ids = list(range(len(requests))) + if sequence_ids != expected_sequence_ids: + raise ValueError( + "Diffusion KV requests must be ordered by contiguous sequence_id values: " + f"expected={expected_sequence_ids}, got={sequence_ids}" + ) + if len(internal_ids) != len(set(internal_ids)): + raise ValueError(f"Diffusion KV internal request IDs must be unique, got {internal_ids}") + conflicts = self._internal_request_ids.intersection(internal_ids) + if conflicts: + raise ValueError(f"Diffusion KV internal request IDs are already allocated: {sorted(conflicts)}") + for request in requests: + if request.kv_contexts: + raise DiffusionKVAdmissionError("independent DiffusionKVContext allocation is not implemented yet") + if request.seq_len > self.max_model_len: + raise DiffusionKVAdmissionError( + f"Diffusion KV sequence {request.request_id!r} exceeds max_model_len: " + f"seq_len={request.seq_len}, max_model_len={self.max_model_len}" + ) + + required_blocks = self._get_empty_pool_required_blocks(requests) + if required_blocks > self._empty_pool_num_free_blocks: + raise DiffusionKVAdmissionError( + f"Diffusion KV request {public_request_id!r} cannot fit even when the block pool is empty: " + f"required_blocks={required_blocks}, available_blocks={self._empty_pool_num_free_blocks}; " + "increase KV cache capacity or reduce the request sequence count/length" + ) + + allocated: list[DiffusionKVRequest] = [] + sequence_metadata: list[DiffusionKVSequenceMetadata] = [] + try: + for request in requests: + blocks = self.native_manager.allocate_slots( + request, + num_new_tokens=request.seq_len, + delay_cache_blocks=True, + full_sequence_must_fit=True, + ) + if blocks is None: + self._rollback(allocated) + allocated.clear() + return None + allocated.append(request) + sequence_metadata.append( + DiffusionKVSequenceMetadata( + sequence_id=request.sequence_id, + prefix_len=request.prefix_len, + target_len=request.target_len, + seq_len=request.seq_len, + block_ids=blocks.get_block_ids(), + ) + ) + except Exception: + self._rollback(allocated) + raise + + metadata = DiffusionKVMetadata( + request_id=public_request_id, + allocation_generation=self._next_allocation_generation, + sequences=tuple(sequence_metadata), + ) + self._next_allocation_generation += 1 + self._requests[public_request_id] = requests + self._metadata[public_request_id] = metadata + self._internal_request_ids.update(internal_ids) + return metadata + + def get_metadata(self, public_request_id: str) -> DiffusionKVMetadata: + try: + return self._metadata[public_request_id] + except KeyError as exc: + raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc + + def free_request(self, public_request_id: str) -> None: + requests = self._requests.pop(public_request_id, ()) + self._metadata.pop(public_request_id, None) + for request in reversed(requests): + self.native_manager.free(request) + self._internal_request_ids.discard(request.request_id) + + def close(self) -> None: + for public_request_id in tuple(self._requests): + self.free_request(public_request_id) + + def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None: + for request in reversed(requests): + self.native_manager.free(request) diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py new file mode 100644 index 000000000..c5ab3fb16 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import copy +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed import get_dcp_group +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec +from vllm.v1.worker import block_table as native_block_table +from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache +from vllm.v1.worker.gpu.block_table import BlockTables + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode +from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata +from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( + DiffusionPagedAttentionAdapter, + DiffusionPagedAttentionLayerAdapter, + DiffusionPagedAttentionMetadata, + DiffusionPagedAttentionRowBinding, + DiffusionPagedAttentionRuntime, +) +from vllm_omni.platforms import current_omni_platform + +DiffusionKVIdentity = tuple[str, int | None, str | None] +DiffusionKVSnapshot = tuple[object, ...] + + +@dataclass(frozen=True) +class _DiffusionKVRowInstall: + identity: DiffusionKVIdentity + token_len: int + block_ids: tuple[tuple[int, ...], ...] + + +@dataclass(frozen=True) +class _DiffusionKVRequestState: + generation: int + snapshot: DiffusionKVSnapshot + row_installs: tuple[_DiffusionKVRowInstall, ...] + + +class DiffusionKVModelRunnerBackend: + """Native paged-KV state and BlockTable operations for a model runner.""" + + def __init__( + self, + *, + vllm_config: VllmConfig, + od_config: OmniDiffusionConfig, + device: torch.device, + ) -> None: + self.vllm_config = vllm_config + self.od_config = od_config + self.device = device + self.kv_cache_config: KVCacheConfig | None = None + self.kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + self.block_tables: BlockTables | None = None + self.paged_attention_adapter: DiffusionPagedAttentionAdapter | None = None + self._kv_cache_layer_adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {} + self._diffusion_kv_identity_to_row: dict[DiffusionKVIdentity, int] = {} + self._diffusion_kv_free_rows: list[int] = [] + self._diffusion_kv_request_states: dict[str, _DiffusionKVRequestState] = {} + + def register_kv_cache_layers( + self, + layers: Mapping[str, tuple[Any, KVCacheSpec]], + ) -> dict[str, KVCacheSpec]: + """Register native adapters and return their canonical cache specs.""" + + if self.kv_cache_config is not None: + raise RuntimeError("Diffusion KV cache layers cannot be changed after physical initialization") + + if layers: + current_omni_platform.get_diffusion_kv_block_tables_cls() + if not callable(current_omni_platform.build_diffusion_kv_attn_metadata): + raise NotImplementedError( + "Diffusion paged KV requires platform-native BlockTables and attention metadata hooks" + ) + + forward_context = self.vllm_config.compilation_config.static_forward_context + previous_adapters = self._kv_cache_layer_adapters + for layer_name, (layer, spec) in layers.items(): + if not isinstance(spec, AttentionSpec): + raise TypeError( + f"Diffusion KV layer {layer_name!r} produced unsupported spec {type(spec).__name__}; " + "only native attention specs are supported" + ) + if not layer.attn_backend.supports_paged_kv: + raise NotImplementedError( + f"Diffusion paged KV layer {layer_name!r} requires an Omni backend with paged support; " + f"selected {layer.attn_backend.get_name()}" + ) + existing = forward_context.get(layer_name) + if existing is not None and existing is not layer and existing is not previous_adapters.get(layer_name): + raise RuntimeError(f"Diffusion KV layer name {layer_name!r} collides with the native forward context") + + previous_use_non_causal = self.vllm_config.attention_config.use_non_causal + self.vllm_config.attention_config.use_non_causal = any( + bool(getattr(spec, "non_causal", False)) for _, spec in layers.values() + ) + adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {} + try: + with set_current_vllm_config(self.vllm_config): + for layer_name, (layer, spec) in layers.items(): + adapters[layer_name] = DiffusionPagedAttentionLayerAdapter( + layer_name=layer_name, + layer=layer, + spec=spec, + vllm_config=self.vllm_config, + device=self.device, + ulysses_degree=self._get_paged_attention_ulysses_degree(layer), + ) + except Exception: + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise + + geometries = {(adapter.num_heads, adapter.num_kv_heads, adapter.head_size) for adapter in adapters.values()} + if len(geometries) > 1: + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise ValueError( + "Native paged attention metadata builders require one rank-local attention geometry; " + f"got {sorted(geometries)!r}" + ) + attention_geometry = next(iter(geometries)) if geometries else None + set_attention_geometry = ( + getattr(self.vllm_config.model_config, "set_attention_geometry", None) + if attention_geometry is not None + else None + ) + if attention_geometry is not None and not callable(set_attention_geometry): + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise RuntimeError( + "Diffusion native paged attention requires a model config that exposes set_attention_geometry" + ) + + updated_forward_context = dict(forward_context) + for layer_name in set(layers) | set(previous_adapters): + entry = updated_forward_context.get(layer_name) + layer_entry = layers.get(layer_name) + if entry is previous_adapters.get(layer_name) or (layer_entry is not None and entry is layer_entry[0]): + updated_forward_context.pop(layer_name, None) + updated_forward_context.update(adapters) + + if attention_geometry is not None: + num_heads, num_kv_heads, head_size = attention_geometry + try: + set_attention_geometry( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_size, + ) + except Exception: + self.vllm_config.attention_config.use_non_causal = previous_use_non_causal + raise + + forward_context.clear() + forward_context.update(updated_forward_context) + self._kv_cache_layer_adapters = adapters + return {layer_name: adapter.spec for layer_name, adapter in adapters.items()} + + def _get_paged_attention_ulysses_degree(self, layer: Any) -> int: + if bool(getattr(layer, "skip_sequence_parallel", False)): + return 1 + + parallel_config = getattr(self.od_config, "parallel_config", None) + ulysses_degree = int(getattr(parallel_config, "ulysses_degree", 1)) + ring_degree = int(getattr(parallel_config, "ring_degree", 1)) + allgather_degree = int(getattr(parallel_config, "allgather_degree", 1)) + ulysses_mode = str(getattr(parallel_config, "ulysses_mode", "strict")) + if ring_degree > 1: + raise NotImplementedError( + "Diffusion paged attention does not support Ring or hybrid Ulysses+Ring attention" + ) + if allgather_degree > 1: + raise NotImplementedError("Diffusion paged attention does not support AllGather-KV sequence parallelism") + if ulysses_degree > 1 and ulysses_mode != "strict": + raise NotImplementedError( + f"Diffusion paged attention currently supports only strict Ulysses; got ulysses_mode={ulysses_mode!r}" + ) + return ulysses_degree + + def _get_max_rows_per_request(self) -> int: + max_rows = getattr(self.od_config, "diffusion_kv_max_rows_per_request", None) + if type(max_rows) is not int or max_rows <= 0: + raise ValueError( + "paged_scheduler requires a positive integer diffusion_kv_max_rows_per_request adapter limit" + ) + return max_rows + + def _get_max_model_len(self) -> int: + max_len = getattr(self.vllm_config.model_config, "max_model_len", None) + if type(max_len) is not int or max_len <= 0: + raise ValueError("paged_scheduler requires a positive model_config.max_model_len") + return max_len + + def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: + """Allocate and bind native attention KV tensors for this rank.""" + if self.kv_cache_config is not None: + raise RuntimeError("Diffusion native KV cache is already initialized") + if not self._kv_cache_layer_adapters: + raise RuntimeError("Diffusion KV cache layers must be registered before physical initialization") + + configured_layers = { + layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names + } + registered_layers = set(self._kv_cache_layer_adapters) + if configured_layers != registered_layers: + raise ValueError( + "Rank-local Diffusion KVCacheConfig layer mismatch: " + f"expected={sorted(registered_layers)}, configured={sorted(configured_layers)}" + ) + + kv_cache_config = copy.deepcopy(kv_cache_config) + max_len = self._get_max_model_len() + max_rows_per_request = self._get_max_rows_per_request() + scheduler_config = self.vllm_config.scheduler_config + max_num_seqs = getattr(scheduler_config, "max_num_seqs", None) + if type(max_num_seqs) is not int or max_num_seqs <= 0: + raise ValueError("paged_scheduler requires a positive scheduler_config.max_num_seqs") + + kv_cache_groups = kv_cache_config.kv_cache_groups + if not kv_cache_groups: + raise ValueError("paged_scheduler requires at least one native KV cache group") + block_sizes = [group.kv_cache_spec.block_size for group in kv_cache_groups] + if any(type(block_size) is not int or block_size <= 0 for block_size in block_sizes): + raise ValueError("native diffusion KV cache groups must have positive integer block sizes") + unaligned_max_num_blocks_per_group = [ + group.kv_cache_spec.max_num_blocks_per_req(self.vllm_config, max_len) for group in kv_cache_groups + ] + if any(type(capacity) is not int or capacity <= 0 for capacity in unaligned_max_num_blocks_per_group): + raise ValueError("native diffusion KV cache groups must have positive row capacities") + max_num_blocks_per_group = [ + native_block_table.get_block_table_width(capacity, block_size) + for capacity, block_size in zip(unaligned_max_num_blocks_per_group, block_sizes, strict=True) + ] + + max_num_reqs = max_num_seqs * max_rows_per_request + max_num_batched_tokens = getattr(scheduler_config, "max_num_batched_tokens", None) + if type(max_num_batched_tokens) is not int or max_num_batched_tokens <= 0: + raise ValueError("scheduler_config.max_num_batched_tokens must be a positive integer") + parallel_config = self.vllm_config.parallel_config + cp_size = parallel_config.decode_context_parallel_size + cp_rank = get_dcp_group().rank_in_group if cp_size > 1 else 0 + cp_interleave = parallel_config.cp_kv_cache_interleave_size + + # Native metadata builders size per-row buffers from max_num_seqs. + # A diffusion public request can occupy several sequence/context rows. + scheduler_config.max_num_seqs = max_num_reqs + kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + previous_adapter_caches = { + layer_name: adapter.kv_cache for layer_name, adapter in self._kv_cache_layer_adapters.items() + } + try: + with set_current_vllm_config(self.vllm_config): + attn_groups, _, kernel_block_sizes = init_attn_backend( + kv_cache_config, + self.vllm_config, + self.device, + ) + block_tables_cls = current_omni_platform.get_diffusion_kv_block_tables_cls() + block_tables = block_tables_cls( + block_sizes=block_sizes, + max_num_reqs=max_num_reqs, + max_num_batched_tokens=max_num_batched_tokens, + max_num_blocks_per_group=max_num_blocks_per_group, + device=self.device, + kernel_block_sizes=kernel_block_sizes, + cp_size=cp_size, + cp_rank=cp_rank, + cp_interleave=cp_interleave, + ) + paged_attention_adapter = DiffusionPagedAttentionAdapter( + vllm_config=self.vllm_config, + device=self.device, + kv_cache_config=kv_cache_config, + block_tables=block_tables, + attn_groups=attn_groups, + layers=self._kv_cache_layer_adapters, + resolve_row=self._resolve_paged_attention_row, + ) + init_kv_cache( + kv_caches, + self.vllm_config.compilation_config.static_forward_context, + kv_cache_config, + attn_groups, + self.device, + self.vllm_config.cache_config.cache_dtype, + kernel_block_sizes, + self.vllm_config, + ) + except Exception: + scheduler_config.max_num_seqs = max_num_seqs + for layer_name, previous_cache in previous_adapter_caches.items(): + self._kv_cache_layer_adapters[layer_name].kv_cache = previous_cache + raise + + self.kv_caches = kv_caches + self.block_tables = block_tables + self.paged_attention_adapter = paged_attention_adapter + self._diffusion_kv_identity_to_row.clear() + self._diffusion_kv_free_rows = list(range(max_num_reqs - 1, -1, -1)) + self._diffusion_kv_request_states.clear() + self.kv_cache_config = kv_cache_config + + def _validate_row( + self, + *, + identity: DiffusionKVIdentity, + token_len: int, + block_ids: tuple[list[int], ...], + seen_identities: set[DiffusionKVIdentity], + ) -> _DiffusionKVRowInstall: + if identity in seen_identities: + raise ValueError(f"Duplicate diffusion KV row identity: {identity!r}") + seen_identities.add(identity) + if type(token_len) is not int or token_len < 0: + raise ValueError(f"Diffusion KV row {identity!r} has invalid token length {token_len!r}") + + assert self.kv_cache_config is not None + groups = self.kv_cache_config.kv_cache_groups + if not isinstance(block_ids, tuple) or len(block_ids) != len(groups): + actual_groups = len(block_ids) if isinstance(block_ids, (tuple, list)) else "invalid" + raise ValueError(f"Diffusion KV row {identity!r} has {actual_groups} block groups; expected {len(groups)}") + + max_len = self._get_max_model_len() + normalized_groups: list[tuple[int, ...]] = [] + for group_index, (group_block_ids, kv_cache_group) in enumerate(zip(block_ids, groups, strict=True)): + if not isinstance(group_block_ids, list): + raise ValueError(f"Diffusion KV row {identity!r} group {group_index} block IDs must be a list") + spec = kv_cache_group.kv_cache_spec + expected_count = spec.max_num_blocks_per_req(self.vllm_config, token_len) + if len(group_block_ids) != expected_count: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} has {len(group_block_ids)} blocks; " + f"expected {expected_count} for {token_len} tokens" + ) + row_capacity = spec.max_num_blocks_per_req(self.vllm_config, max_len) + if len(group_block_ids) > row_capacity: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} exceeds row capacity {row_capacity}" + ) + normalized_ids: list[int] = [] + for block_id in group_block_ids: + if type(block_id) is not int: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} has non-integer block ID {block_id!r}" + ) + if not 0 <= block_id < self.kv_cache_config.num_blocks: + raise ValueError( + f"Diffusion KV row {identity!r} group {group_index} block ID {block_id} is outside " + f"[0, {self.kv_cache_config.num_blocks})" + ) + normalized_ids.append(block_id) + normalized_groups.append(tuple(normalized_ids)) + + return _DiffusionKVRowInstall( + identity=identity, + token_len=token_len, + block_ids=tuple(normalized_groups), + ) + + def _prepare_install( + self, + metadata: DiffusionKVMetadata, + ) -> tuple[list[_DiffusionKVRowInstall], DiffusionKVSnapshot]: + if self.block_tables is None or self.kv_cache_config is None: + raise RuntimeError("paged_scheduler native KV cache is not initialized") + if type(metadata.request_id) is not str or not metadata.request_id: + raise ValueError("Diffusion KV metadata request_id must be a non-empty string") + if type(metadata.allocation_generation) is not int or metadata.allocation_generation < 0: + raise ValueError("Diffusion KV allocation_generation must be a non-negative integer") + if type(self.kv_cache_config.num_blocks) is not int or self.kv_cache_config.num_blocks <= 0: + raise ValueError("native diffusion KV cache config must provide a positive num_blocks") + if not metadata.sequences: + raise ValueError(f"Diffusion KV request {metadata.request_id!r} must contain at least one sequence") + + seen_identities: set[DiffusionKVIdentity] = set() + installs: list[_DiffusionKVRowInstall] = [] + sequence_snapshots: list[object] = [] + for sequence in metadata.sequences: + if type(sequence.sequence_id) is not int or sequence.sequence_id < 0: + raise ValueError(f"Diffusion KV sequence_id must be a non-negative integer: {sequence.sequence_id!r}") + sequence_identity = (metadata.request_id, sequence.sequence_id, None) + sequence_install = self._validate_row( + identity=sequence_identity, + token_len=sequence.seq_len, + block_ids=sequence.block_ids, + seen_identities=seen_identities, + ) + installs.append(sequence_install) + sequence_snapshots.append( + ( + sequence.sequence_id, + sequence.prefix_len, + sequence.target_len, + sequence.seq_len, + sequence_install.block_ids, + sequence.context_ids, + ) + ) + + context_snapshots: list[object] = [] + for context in metadata.contexts: + if type(context.context_id) is not str or not context.context_id: + raise ValueError("Diffusion KV context_id must be a non-empty string") + context_identity = (metadata.request_id, None, context.context_id) + context_install = self._validate_row( + identity=context_identity, + token_len=context.num_tokens, + block_ids=context.block_ids, + seen_identities=seen_identities, + ) + installs.append(context_install) + context_snapshots.append( + ( + context.context_id, + context.cache_role, + context.num_tokens, + context_install.block_ids, + ) + ) + + max_rows_per_request = self._get_max_rows_per_request() + if len(installs) > max_rows_per_request: + raise ValueError( + f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows; " + f"adapter limit is {max_rows_per_request}" + ) + snapshot: DiffusionKVSnapshot = ( + metadata.request_id, + tuple(sequence_snapshots), + tuple(context_snapshots), + ) + return installs, snapshot + + def _rollback_staged_writes(self, rows: list[int], previous_num_blocks: object | None) -> None: + block_tables = self.block_tables + if block_tables is None: + return + for group_table in getattr(block_tables, "block_tables", ()): + clear_staged_writes = getattr(group_table, "clear_staged_writes", None) + if clear_staged_writes is not None: + clear_staged_writes() + if previous_num_blocks is None: + return + native_num_blocks = getattr(block_tables, "num_blocks", None) + if native_num_blocks is None or not hasattr(native_num_blocks, "np"): + return + native_num_blocks.np[:, rows] = previous_num_blocks + copy_to_uva = getattr(native_num_blocks, "copy_to_uva", None) + if copy_to_uva is not None: + copy_to_uva() + + def _apply_rows(self, rows: list[int], installs: list[_DiffusionKVRowInstall]) -> None: + assert self.block_tables is not None + if self.paged_attention_adapter is not None: + self.paged_attention_adapter.invalidate_prepared_batches() + native_num_blocks = getattr(self.block_tables, "num_blocks", None) + previous_num_blocks = None + if native_num_blocks is not None and hasattr(native_num_blocks, "np"): + previous_num_blocks = native_num_blocks.np[:, rows].copy() + try: + for row, install in zip(rows, installs, strict=True): + self.block_tables.append_block_ids( + row, + tuple(list(group_ids) for group_ids in install.block_ids), + overwrite=True, + ) + self.block_tables.apply_staged_writes() + except Exception: + self._rollback_staged_writes(rows, previous_num_blocks) + raise + + def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool: + """Install one Scheduler allocation snapshot into native Worker rows.""" + if not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ): + raise ValueError("Dense diffusion execution must not install Diffusion KV metadata") + + installs, snapshot = self._prepare_install(metadata) + current_state = self._diffusion_kv_request_states.get(metadata.request_id) + if current_state is not None: + if metadata.allocation_generation < current_state.generation: + raise ValueError( + f"Stale Diffusion KV allocation generation {metadata.allocation_generation} for request " + f"{metadata.request_id!r}; installed generation is {current_state.generation}" + ) + if snapshot != current_state.snapshot: + raise ValueError( + f"Conflicting Diffusion KV allocation snapshot for active request {metadata.request_id!r}; " + "remove the request before replacing its rows" + ) + if metadata.allocation_generation == current_state.generation: + return False + self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState( + generation=metadata.allocation_generation, + snapshot=snapshot, + row_installs=current_state.row_installs, + ) + return False + + existing_identities = [ + install.identity for install in installs if install.identity in self._diffusion_kv_identity_to_row + ] + if existing_identities: + raise RuntimeError(f"Diffusion KV row registry contains orphaned identities: {existing_identities!r}") + if len(installs) > len(self._diffusion_kv_free_rows): + raise ValueError( + f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows, but only " + f"{len(self._diffusion_kv_free_rows)} rows are free" + ) + + rows = list(reversed(self._diffusion_kv_free_rows[-len(installs) :])) if installs else [] + if rows: + self._apply_rows(rows, installs) + + if installs: + del self._diffusion_kv_free_rows[-len(installs) :] + self._diffusion_kv_identity_to_row.update( + (install.identity, row) for install, row in zip(installs, rows, strict=True) + ) + self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState( + generation=metadata.allocation_generation, + snapshot=snapshot, + row_installs=tuple(installs), + ) + return True + + def get_diffusion_kv_row( + self, + request_id: str, + sequence_id: int | None, + context_id: str | None = None, + ) -> int: + """Resolve a Scheduler allocation identity to its native table row.""" + if ( + not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ) + or self.block_tables is None + ): + raise RuntimeError("paged_scheduler native BlockTables are not initialized") + identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None) + try: + return self._diffusion_kv_identity_to_row[identity] + except KeyError as exc: + raise KeyError(f"No native diffusion KV row is installed for {identity!r}") from exc + + def get_paged_attention_adapter(self) -> DiffusionPagedAttentionAdapter: + if self.paged_attention_adapter is None: + raise RuntimeError("paged_scheduler native attention adapter is not initialized") + return self.paged_attention_adapter + + def activate_paged_attention_metadata(self, metadata: DiffusionPagedAttentionMetadata): + """Activate Runner-owned request metadata for an internal denoise loop.""" + + runtime = DiffusionPagedAttentionRuntime(self.get_paged_attention_adapter(), metadata) + return runtime, runtime.activate() + + def _resolve_paged_attention_row( + self, + request_id: str, + sequence_id: int | None, + context_id: str | None, + ) -> DiffusionPagedAttentionRowBinding: + row_index = self.get_diffusion_kv_row(request_id, sequence_id, context_id) + identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None) + request_state = self._diffusion_kv_request_states.get(request_id) + if request_state is not None: + for install in request_state.row_installs: + if install.identity == identity: + return DiffusionPagedAttentionRowBinding( + row_index=row_index, + max_seq_len=install.token_len, + block_ids=install.block_ids, + ) + raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}") + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int: + """Retire Worker rows without logically freeing Scheduler-owned blocks.""" + if ( + not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ) + or self.block_tables is None + ): + return 0 + request_id_set = set(request_ids) + identities_and_rows = [ + (identity, row) + for identity, row in self._diffusion_kv_identity_to_row.items() + if identity[0] in request_id_set + ] + rows = [row for _, row in identities_and_rows] + if rows: + assert self.kv_cache_config is not None + num_groups = len(self.kv_cache_config.kv_cache_groups) + installs = [ + _DiffusionKVRowInstall( + identity=identity, + token_len=0, + block_ids=tuple(tuple() for _ in range(num_groups)), + ) + for identity, _ in identities_and_rows + ] + self._apply_rows(rows, installs) + + for identity, _ in identities_and_rows: + del self._diffusion_kv_identity_to_row[identity] + for request_id in request_id_set: + self._diffusion_kv_request_states.pop(request_id, None) + self._diffusion_kv_free_rows.extend(rows) + self._diffusion_kv_free_rows.sort(reverse=True) + return len(rows) + + def refresh_block_table_layout(self) -> None: + """Refresh native pointer tensors after a CuMem KV-cache wake-up.""" + if ( + not is_scheduler_paged_kv_mode( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + ) + or self.block_tables is None + ): + return + if self.paged_attention_adapter is not None: + self.paged_attention_adapter.invalidate_prepared_batches() + self.block_tables.init_block_table_layout_tensors() diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py new file mode 100644 index 000000000..6a7862669 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py @@ -0,0 +1,946 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass, field, replace +from typing import Any + +import torch +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens +from vllm.v1.attention.selector import get_attn_backend +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.block_table import BlockTables + +from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( + PagedPiecewisePlan, + build_paged_piecewise_plan, +) +from vllm_omni.diffusion.forward_context import override_paged_kv_adapter +from vllm_omni.platforms import current_omni_platform + + +@dataclass(frozen=True) +class DiffusionPagedAttentionRowBinding: + """Worker row and logical length installed for one allocation identity.""" + + row_index: int + max_seq_len: int + block_ids: tuple[tuple[int, ...], ...] = () + + +DiffusionKVRowResolver = Callable[ + [str, int | None, str | None], + DiffusionPagedAttentionRowBinding, +] + + +class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase): + """Register a diffusion layer with vLLM's native cache machinery. + + This object deliberately does *not* subclass ``vllm.Attention``. The + latter owns a second execution path and would bypass Omni's sequence + parallel pre/post hooks. The wrapper only supplies the small + ``AttentionLayerBase`` contract needed by ``init_attn_backend`` and keeps + the platform-native attention implementation/cache view available to the + diffusion adapter. + """ + + def __init__( + self, + *, + layer_name: str, + layer: Any, + spec: AttentionSpec, + vllm_config: VllmConfig, + device: torch.device, + ulysses_degree: int = 1, + ) -> None: + if type(ulysses_degree) is not int or ulysses_degree <= 0: + raise ValueError(f"ulysses_degree must be a positive integer, got {ulysses_degree!r}") + num_heads = int(layer.num_heads) + num_kv_heads = int(spec.num_kv_heads) + if num_heads <= 0 or num_kv_heads <= 0 or num_heads % num_kv_heads != 0: + raise ValueError( + "Paged attention requires positive Q/KV heads with num_heads divisible by num_kv_heads: " + f"num_heads={num_heads}, num_kv_heads={num_kv_heads}" + ) + if num_heads % ulysses_degree != 0 or num_kv_heads % ulysses_degree != 0: + raise ValueError( + "Paged attention requires Q/KV heads divisible by ulysses_degree: " + f"num_heads={num_heads}, num_kv_heads={num_kv_heads}, ulysses_degree={ulysses_degree}" + ) + num_heads //= ulysses_degree + num_kv_heads //= ulysses_degree + + attention_config = vllm_config.attention_config + previous_backend = attention_config.backend + previous_backend_per_kind = attention_config.backend_per_kind + try: + # This is a portable vLLM backend request. The active platform + # resolves it to its native implementation, such as FlashAttention + # on CUDA or AscendAttentionBackend on NPU. + attention_config.backend = AttentionBackendEnum.FLASH_ATTN + attention_config.backend_per_kind = {} + with set_current_vllm_config(vllm_config): + attn_backend = get_attn_backend( + head_size=spec.head_size, + dtype=vllm_config.model_config.dtype, + kv_cache_dtype=vllm_config.cache_config.cache_dtype, + attn_type=AttentionType.DECODER, + num_heads=num_heads, + ) + finally: + attention_config.backend = previous_backend + attention_config.backend_per_kind = previous_backend_per_kind + attn_backend = current_omni_platform.get_diffusion_paged_kv_attn_backend( + attn_backend, + ulysses_degree=ulysses_degree, + ) + canonical_spec = replace( + spec, + num_kv_heads=num_kv_heads, + indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), + ) + self.layer_name = layer_name + self.spec = canonical_spec + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_size = int(canonical_spec.head_size) + self.head_size_v = int(getattr(canonical_spec, "head_size_v", canonical_spec.head_size)) + if self.head_size_v != self.head_size: + raise ValueError( + "Diffusion native paged attention requires equal key/value head sizes; " + f"got head_size={self.head_size}, head_size_v={self.head_size_v}" + ) + self.softmax_scale = float(layer.softmax_scale) + self.attn_backend = attn_backend + self.kv_cache: torch.Tensor | None = None + # Native attention backends consume either device tensors or host + # scalar copies of these scales. The paged diffusion path is currently + # unquantized, so initialize both contracts to identity. + self._q_scale = torch.ones(1, device=device, dtype=torch.float32) + self._k_scale = torch.ones(1, device=device, dtype=torch.float32) + self._v_scale = torch.ones(1, device=device, dtype=torch.float32) + self._q_scale_float = 1.0 + self._k_scale_float = 1.0 + self._v_scale_float = 1.0 + with set_current_vllm_config(vllm_config): + self.impl = self._create_native_impl(vllm_config) + + def get_attn_backend(self): + return self.attn_backend + + def _create_native_impl(self, vllm_config: VllmConfig): + impl_cls = self.attn_backend.get_impl_cls() + impl = impl_cls( + num_heads=self.num_heads, + head_size=self.head_size, + scale=self.softmax_scale, + num_kv_heads=self.num_kv_heads, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype=vllm_config.cache_config.cache_dtype, + logits_soft_cap=None, + attn_type=AttentionType.DECODER, + kv_sharing_target_layer_name=None, + ) + if not callable(getattr(impl, "forward", None)): + raise RuntimeError( + f"Native paged-attention backend {self.attn_backend.get_name()!r} has no forward implementation" + ) + needs_explicit_cache_update = ( + not self.attn_backend.forward_includes_kv_cache_update + or current_omni_platform.requires_diffusion_paged_kv_prewrite() + ) + if needs_explicit_cache_update and not callable(getattr(impl, "do_kv_cache_update", None)): + raise RuntimeError( + f"Native paged-attention backend {self.attn_backend.get_name()!r} cannot update the KV cache" + ) + if hasattr(impl, "vllm_flash_attn_version") and impl.vllm_flash_attn_version is None: + raise RuntimeError(f"Native paged-attention kernel is unavailable for diffusion layer {self.layer_name!r}") + return impl + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> AttentionSpec: + del vllm_config + return self.spec + + +@dataclass(frozen=True) +class DiffusionPagedAttentionRow: + """One logical BlockTable row participating in a paged attention call.""" + + request_id: str + query_len: int + seq_len: int + kv_start_pos: int = 0 + sequence_id: int | None = None + context_id: str | None = None + + def __post_init__(self) -> None: + if type(self.request_id) is not str or not self.request_id: + raise ValueError("Paged attention request_id must be a non-empty string") + if (self.sequence_id is None) == (self.context_id is None): + raise ValueError("Paged attention row requires exactly one of sequence_id or context_id") + if self.sequence_id is not None and (type(self.sequence_id) is not int or self.sequence_id < 0): + raise ValueError("Paged attention sequence_id must be a non-negative integer") + if self.context_id is not None and (type(self.context_id) is not str or not self.context_id): + raise ValueError("Paged attention context_id must be a non-empty string") + if type(self.query_len) is not int or self.query_len <= 0: + raise ValueError("Paged attention query_len must be a positive integer") + if type(self.seq_len) is not int or self.seq_len <= 0: + raise ValueError("Paged attention seq_len must be a positive integer") + if type(self.kv_start_pos) is not int or self.kv_start_pos < 0: + raise ValueError("Paged attention kv_start_pos must be a non-negative integer") + if self.kv_start_pos + self.query_len > self.seq_len: + raise ValueError( + "Paged attention write span exceeds seq_len: " + f"start={self.kv_start_pos}, query_len={self.query_len}, seq_len={self.seq_len}" + ) + + @property + def identity(self) -> tuple[str, int | None, str | None]: + return (self.request_id, self.sequence_id, self.context_id) + + +@dataclass(frozen=True) +class DiffusionPagedAttentionMetadata: + """Runner-owned row layouts for one request-level denoise loop.""" + + prefill_rows: tuple[DiffusionPagedAttentionRow, ...] + denoise_rows: tuple[DiffusionPagedAttentionRow, ...] + + def __post_init__(self) -> None: + if not self.prefill_rows or not self.denoise_rows: + raise ValueError("Paged attention metadata requires both prefill and denoise rows") + prefill_identities = tuple(row.identity for row in self.prefill_rows) + denoise_identities = tuple(row.identity for row in self.denoise_rows) + if prefill_identities != denoise_identities: + raise ValueError( + "Paged attention prefill and denoise rows must use the same ordered identities: " + f"prefill={prefill_identities!r}, denoise={denoise_identities!r}" + ) + + def rows_for_step(self, step_idx: int | None) -> tuple[DiffusionPagedAttentionRow, ...]: + if type(step_idx) is not int or step_idx < 0: + raise ValueError(f"Paged attention requires a non-negative denoise step index, got {step_idx!r}") + return self.prefill_rows if step_idx == 0 else self.denoise_rows + + +@dataclass(frozen=True) +class PreparedDiffusionPagedAttentionBatch: + """Native metadata shared by all paged attention layers in one forward.""" + + rows: tuple[DiffusionPagedAttentionRow, ...] + row_indices: torch.Tensor + query_start_loc: torch.Tensor + seq_lens: torch.Tensor + positions: torch.Tensor + block_tables: tuple[torch.Tensor, ...] + slot_mappings: torch.Tensor + attn_metadata: dict[str, Any] + slot_mappings_by_layer: dict[str, torch.Tensor] + num_tokens: int + _owner: object = field(repr=False, compare=False) + _generation: int = field(repr=False, compare=False) + + +@dataclass(frozen=True, slots=True) +class DiffusionPagedAttentionContext: + """One layer's native inputs for an Omni paged-backend invocation.""" + + layer: DiffusionPagedAttentionLayerAdapter + query: torch.Tensor + key_write: torch.Tensor + value_write: torch.Tensor + slot_mapping: torch.Tensor + native_metadata: Any + piecewise_plan: PagedPiecewisePlan | None + piecewise_native_metadata: tuple[Any, ...] + query_token_shape: tuple[int, ...] + query_has_head_dims: bool + output_scatter_indices: torch.Tensor | None = None + + def restore_output(self, output: torch.Tensor) -> torch.Tensor: + if self.output_scatter_indices is not None: + flat_output = output.reshape(-1, self.layer.num_heads, self.layer.head_size_v) + restored = output.new_zeros( + (*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v) + ).reshape(-1, self.layer.num_heads, self.layer.head_size_v) + restored.index_copy_(0, self.output_scatter_indices, flat_output) + if self.query_has_head_dims: + return restored.reshape( + *self.query_token_shape, + self.layer.num_heads, + self.layer.head_size_v, + ) + return restored.reshape( + *self.query_token_shape, + self.layer.num_heads * self.layer.head_size_v, + ) + if self.query_has_head_dims: + return output.reshape(*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v) + return output.reshape(*self.query_token_shape, self.layer.num_heads * self.layer.head_size_v) + + +class DiffusionPagedAttentionRuntime: + """Activate Runner-prepared rows lazily as the denoise loop advances.""" + + def __init__( + self, + adapter: DiffusionPagedAttentionAdapter, + metadata: DiffusionPagedAttentionMetadata, + ) -> None: + self.adapter = adapter + self.metadata = metadata + self._active_phase: bool | None = None + self._activation: Any | None = None + self._closed = False + + def ensure_active(self, step_idx: int | None) -> DiffusionPagedAttentionAdapter: + if self._closed: + raise RuntimeError("Paged attention runtime is already closed") + is_prefill = step_idx == 0 + rows = self.metadata.rows_for_step(step_idx) + if self._active_phase == is_prefill: + return self.adapter + + self._deactivate() + batch = self.adapter.prepare_batch(rows) + activation = self.adapter.activate(batch) + activation.__enter__() + self._activation = activation + self._active_phase = is_prefill + return self.adapter + + def _deactivate(self) -> None: + activation = self._activation + self._activation = None + self._active_phase = None + if activation is not None: + activation.__exit__(None, None, None) + + def close(self) -> None: + if self._closed: + return + self._deactivate() + self._closed = True + + @contextmanager + def activate(self) -> Iterator[DiffusionPagedAttentionRuntime]: + try: + yield self + finally: + self.close() + + +class DiffusionPagedAttentionAdapter: + """Translate diffusion rows into vLLM metadata for an Omni paged backend.""" + + def __init__( + self, + *, + vllm_config: VllmConfig, + device: torch.device, + kv_cache_config: KVCacheConfig, + block_tables: BlockTables, + attn_groups: list[list[Any]], + layers: Mapping[str, DiffusionPagedAttentionLayerAdapter], + resolve_row: DiffusionKVRowResolver, + ) -> None: + if not layers: + raise ValueError("Paged attention requires at least one native attention layer") + if len(attn_groups) != len(kv_cache_config.kv_cache_groups): + raise ValueError( + "Paged attention group mismatch: " + f"builders={len(attn_groups)}, cache_groups={len(kv_cache_config.kv_cache_groups)}" + ) + self.vllm_config = vllm_config + self.device = torch.device(device) + self.kv_cache_config = kv_cache_config + self.block_tables = block_tables + self.attn_groups = attn_groups + self.layers = dict(layers) + self.resolve_row = resolve_row + self._owner = object() + self._prepare_generation = 0 + self._active_batch: PreparedDiffusionPagedAttentionBatch | None = None + self._active_piecewise_plan: PagedPiecewisePlan | None = None + self._active_piecewise_native_metadata: tuple[dict[str, Any], ...] | None = None + self._causal_by_group = self._resolve_group_causality() + self._reorder_batch_threshold = self._resolve_reorder_batch_threshold() + + def _resolve_group_causality(self) -> dict[int, bool]: + causal_by_group: dict[int, bool] = {} + for group_index, cache_group in enumerate(self.kv_cache_config.kv_cache_groups): + layer_causality = { + not bool(getattr(self.layers[layer_name].spec, "non_causal", False)) + for layer_name in cache_group.layer_names + } + if len(layer_causality) != 1: + raise ValueError(f"Paged attention cache group {group_index} mixes causal and non-causal layers") + causal_by_group[group_index] = layer_causality.pop() + return causal_by_group + + def _resolve_reorder_batch_threshold(self) -> int | None: + thresholds = [ + group.get_metadata_builder(0).reorder_batch_threshold + for group_index, groups in enumerate(self.attn_groups) + if self._causal_by_group[group_index] + for group in groups + ] + concrete_thresholds = [threshold for threshold in thresholds if threshold is not None] + return min(concrete_thresholds, default=None) + + def _validate_native_batch_order(self, rows: tuple[DiffusionPagedAttentionRow, ...]) -> None: + threshold = self._reorder_batch_threshold + if threshold is None or not any(self._causal_by_group.values()): + return + + found_long_query = False + for row in rows: + if row.query_len > threshold: + found_long_query = True + elif found_long_query: + raise ValueError( + "Causal paged attention rows must place native decode/short-query rows before " + f"prefill/long-query rows (decode threshold={threshold})" + ) + + def _validate_row_capacity( + self, + rows: tuple[DiffusionPagedAttentionRow, ...], + row_indices: list[int], + ) -> None: + native_num_blocks = self.block_tables.num_blocks.np + blocks_per_kv_block = self.block_tables.blocks_per_kv_block + for row, row_index in zip(rows, row_indices, strict=True): + if type(row_index) is not int or not 0 <= row_index < self.block_tables.max_num_reqs: + raise ValueError(f"Paged attention row {row.identity!r} resolved to invalid Worker row {row_index!r}") + for group_index, (block_size, block_multiplier) in enumerate( + zip( + self.block_tables.block_sizes, + blocks_per_kv_block, + strict=True, + ) + ): + required_blocks = cdiv( + row.seq_len, + block_size * self.block_tables.cp_size, + ) + required_kernel_blocks = required_blocks * block_multiplier + installed_kernel_blocks = int(native_num_blocks[group_index, row_index]) + if required_kernel_blocks > installed_kernel_blocks: + raise ValueError( + f"Paged attention row {row.identity!r} requires {required_blocks} blocks in " + f"cache group {group_index} for seq_len={row.seq_len}, but its installed Worker " + f"row contains only {installed_kernel_blocks // block_multiplier} blocks" + ) + + def _build_native_metadata( + self, + *, + query_lens: Sequence[int], + query_start_loc: torch.Tensor, + query_start_loc_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + positions: torch.Tensor, + block_tables: Sequence[torch.Tensor], + slot_mappings: torch.Tensor, + causal: bool | Mapping[int, bool], + ) -> dict[str, Any]: + dcp_local_seq_lens = None + if self.block_tables.cp_size > 1: + dcp_local_seq_lens = get_dcp_local_seq_lens( + seq_lens, + dcp_size=self.block_tables.cp_size, + dcp_rank=self.block_tables.cp_rank, + cp_kv_cache_interleave_size=self.block_tables.cp_interleave, + ) + with set_current_vllm_config(self.vllm_config): + return current_omni_platform.build_diffusion_kv_attn_metadata( + attn_groups=self.attn_groups, + num_reqs=len(query_lens), + num_tokens=sum(query_lens), + query_start_loc_gpu=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=max(query_lens), + seq_lens=seq_lens, + max_seq_len=max(int(seq_len) for seq_len in seq_lens_cpu), + block_tables=block_tables, + slot_mappings=slot_mappings, + kv_cache_config=self.kv_cache_config, + seq_lens_cpu=seq_lens_cpu, + seq_lens_cpu_upper_bound=seq_lens_cpu, + dcp_local_seq_lens=dcp_local_seq_lens, + positions=positions, + causal=causal, + ) + + def prepare_batch( + self, + rows: Sequence[DiffusionPagedAttentionRow], + ) -> PreparedDiffusionPagedAttentionBatch: + rows = tuple(rows) + if not rows: + raise ValueError("Paged attention batch must contain at least one row") + if self._active_batch is not None: + raise RuntimeError("Cannot prepare paged attention metadata during an active forward") + if len(rows) > self.block_tables.max_num_reqs: + raise ValueError( + f"Paged attention batch has {len(rows)} rows; Worker capacity is {self.block_tables.max_num_reqs}" + ) + + identities = [row.identity for row in rows] + if len(set(identities)) != len(identities): + raise ValueError("Paged attention batch contains duplicate row identities") + row_bindings = [self.resolve_row(row.request_id, row.sequence_id, row.context_id) for row in rows] + row_indices_list = [binding.row_index for binding in row_bindings] + if len(set(row_indices_list)) != len(row_indices_list): + raise ValueError("Paged attention batch resolves multiple inputs to the same Worker row") + for row, binding in zip(rows, row_bindings, strict=True): + if type(binding.max_seq_len) is not int or binding.max_seq_len < 0: + raise ValueError( + f"Paged attention row {row.identity!r} resolved to invalid logical capacity {binding.max_seq_len!r}" + ) + if row.seq_len > binding.max_seq_len: + raise ValueError( + f"Paged attention row {row.identity!r} uses seq_len={row.seq_len}, but its installed " + f"allocation has logical length {binding.max_seq_len}" + ) + self._validate_native_batch_order(rows) + self._validate_row_capacity(rows, row_indices_list) + + query_lens = [row.query_len for row in rows] + num_tokens = sum(query_lens) + if num_tokens > self.block_tables.max_num_batched_tokens: + raise ValueError( + f"Paged attention batch has {num_tokens} tokens; Worker capacity is " + f"{self.block_tables.max_num_batched_tokens}" + ) + query_offsets = [0] + for query_len in query_lens: + query_offsets.append(query_offsets[-1] + query_len) + + # Native BlockTables reuses persistent gather/slot buffers. Starting a + # new preparation invalidates every batch prepared before it, even if a + # later metadata builder raises. + self._prepare_generation += 1 + generation = self._prepare_generation + row_indices = torch.tensor(row_indices_list, dtype=torch.int32, device=self.device) + query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32) + query_start_loc = query_start_loc_cpu.to(self.device) + seq_lens_cpu = torch.tensor([row.seq_len for row in rows], dtype=torch.int32) + seq_lens = seq_lens_cpu.to(self.device) + positions = torch.cat( + [ + torch.arange( + row.kv_start_pos, + row.kv_start_pos + row.query_len, + dtype=torch.int64, + device=self.device, + ) + for row in rows + ] + ) + + block_tables = self.block_tables.gather_block_tables( + row_indices, + num_reqs_padded=len(rows), + ) + slot_mappings = self.block_tables.compute_slot_mappings( + row_indices, + query_start_loc, + positions, + num_tokens_padded=num_tokens, + ) + causal: bool | Mapping[int, bool] + causal_values = set(self._causal_by_group.values()) + causal = causal_values.pop() if len(causal_values) == 1 else self._causal_by_group + attn_metadata = self._build_native_metadata( + query_lens=query_lens, + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + positions=positions, + block_tables=block_tables, + slot_mappings=slot_mappings, + causal=causal, + ) + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, + self.kv_cache_config, + ) + return PreparedDiffusionPagedAttentionBatch( + rows=rows, + row_indices=row_indices, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + positions=positions, + block_tables=tuple(block_tables), + slot_mappings=slot_mappings, + attn_metadata=attn_metadata, + slot_mappings_by_layer=slot_mappings_by_layer, + num_tokens=num_tokens, + _owner=self._owner, + _generation=generation, + ) + + def invalidate_prepared_batches(self) -> None: + """Invalidate native buffer views after BlockTable state changes.""" + + if self._active_batch is not None: + raise RuntimeError("Cannot change paged attention BlockTables during an active forward") + self._prepare_generation += 1 + + @contextmanager + def activate( + self, + batch: PreparedDiffusionPagedAttentionBatch, + ) -> Iterator[DiffusionPagedAttentionAdapter]: + if batch._owner is not self._owner: + raise ValueError("Prepared paged attention batch belongs to a different adapter") + if batch._generation != self._prepare_generation: + raise ValueError("Prepared paged attention batch is stale after a newer batch preparation") + if self._active_batch is not None: + raise RuntimeError("Paged attention adapter already has an active forward") + self._active_batch = batch + self._active_piecewise_plan = None + self._active_piecewise_native_metadata = None + try: + with override_paged_kv_adapter(self), set_current_vllm_config(self.vllm_config): + yield self + finally: + self._active_piecewise_plan = None + self._active_piecewise_native_metadata = None + self._active_batch = None + + @staticmethod + def _flatten_tensor( + tensor: torch.Tensor, + *, + num_heads: int, + head_size: int, + name: str, + ) -> tuple[torch.Tensor, tuple[int, ...], bool]: + if tensor.ndim >= 3 and tensor.shape[-2:] == (num_heads, head_size): + return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-2]), True + hidden_size = num_heads * head_size + if tensor.ndim >= 1 and tensor.shape[-1] == hidden_size: + return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-1]), False + raise ValueError( + f"Paged attention {name} must end in ({num_heads}, {head_size}) or ({hidden_size},); " + f"got shape={tuple(tensor.shape)}" + ) + + @staticmethod + def _token_packing_indices( + token_shape: tuple[int, ...], + batch: PreparedDiffusionPagedAttentionBatch, + *, + name: str, + device: torch.device, + ) -> torch.Tensor | None: + if len(token_shape) == 1: + if token_shape[0] == batch.num_tokens: + return None + raise ValueError( + f"Paged attention {name} token count must match the prepared write batch: " + f"tokens={token_shape[0]}, prepared={batch.num_tokens}" + ) + if len(token_shape) == 2: + batch_size, tokens_per_row = token_shape + query_lens = tuple(row.query_len for row in batch.rows) + if batch_size != len(batch.rows) or any(query_len > tokens_per_row for query_len in query_lens): + raise ValueError( + f"Paged attention batched {name} layout must match prepared rows for the current write: " + f"shape={token_shape}, row_query_lens={query_lens}" + ) + if all(query_len == tokens_per_row for query_len in query_lens): + return None + return torch.cat( + [ + torch.arange( + row_index * tokens_per_row, + row_index * tokens_per_row + query_len, + dtype=torch.long, + device=device, + ) + for row_index, query_len in enumerate(query_lens) + ] + ) + raise ValueError( + f"Paged attention {name} supports packed [T, ...] or uniform batched [B, T, ...] token layouts; " + f"got token shape={token_shape}" + ) + + def _get_piecewise_plan( + self, + full_attn_spans: list[list[tuple[int, int]]], + ) -> PagedPiecewisePlan: + batch = self._active_batch + if batch is None: + raise RuntimeError("Piecewise paged attention requires an active prepared batch") + if any(row.kv_start_pos + row.query_len != row.seq_len for row in batch.rows): + invalid_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len] + raise ValueError( + "Paged piecewise attention requires each query/write span to end at seq_len; " + f"invalid rows={invalid_rows!r}" + ) + + cached_plan = self._active_piecewise_plan + if cached_plan is not None: + frozen_spans = tuple(tuple(tuple(span) for span in row_spans) for row_spans in full_attn_spans) + if cached_plan.spans != frozen_spans: + raise ValueError("Paged piecewise attention metadata changed between layers in one active batch") + return cached_plan + + plan = build_paged_piecewise_plan( + full_attn_spans, + query_offsets=[row.kv_start_pos for row in batch.rows], + query_lens=[row.query_len for row in batch.rows], + seq_lens=[row.seq_len for row in batch.rows], + device=self.device, + ) + self._active_piecewise_plan = plan + return plan + + def _get_piecewise_native_metadata( + self, + plan: PagedPiecewisePlan, + ) -> tuple[dict[str, Any], ...]: + """Build one native metadata object per segment and cache it by batch.""" + + batch = self._active_batch + if batch is None: + raise RuntimeError("Piecewise paged attention requires an active prepared batch") + cached_metadata = self._active_piecewise_native_metadata + if cached_metadata is not None: + return cached_metadata + + metadata_by_segment: list[dict[str, Any]] = [] + for packed_segment in plan.segments: + row_segments = packed_segment.row_segments + query_lens = [segment.q_end - segment.q_start for segment in row_segments] + seq_lens_values = [segment.kv_end for segment in row_segments] + query_offsets = [0] + for query_len in query_lens: + query_offsets.append(query_offsets[-1] + query_len) + query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32) + seq_lens_cpu = torch.tensor(seq_lens_values, dtype=torch.int32) + if packed_segment.query_range is None: + positions = batch.positions.index_select(0, packed_segment.query_indices) + slot_mappings = batch.slot_mappings.index_select(-1, packed_segment.query_indices) + else: + start, end = packed_segment.query_range + positions = batch.positions[start:end] + slot_mappings = batch.slot_mappings[..., start:end] + segment_metadata = self._build_native_metadata( + query_lens=query_lens, + query_start_loc=query_start_loc_cpu.to(self.device), + query_start_loc_cpu=query_start_loc_cpu, + seq_lens=seq_lens_cpu.to(self.device), + seq_lens_cpu=seq_lens_cpu, + positions=positions, + block_tables=batch.block_tables, + slot_mappings=slot_mappings, + causal=(row_segments[0].mode == "causal"), + ) + # FA3's full-CUDA-graph metadata builder reuses one persistent + # scheduler buffer across builds. Piecewise attention prepares all + # segments before executing any of them, so each segment needs its + # own snapshot rather than a view overwritten by the next build. + snapshotted = set() + for native_metadata in segment_metadata.values(): + metadata_id = id(native_metadata) + if metadata_id in snapshotted: + continue + snapshotted.add(metadata_id) + for field_name in ("scheduler_metadata", "prefix_scheduler_metadata"): + scheduler_metadata = getattr(native_metadata, field_name, None) + if isinstance(scheduler_metadata, torch.Tensor): + setattr(native_metadata, field_name, scheduler_metadata.clone()) + metadata_by_segment.append(segment_metadata) + self._active_piecewise_native_metadata = tuple(metadata_by_segment) + return self._active_piecewise_native_metadata + + def prepare_layer_context( + self, + layer_name: str, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + omni_attn_metadata: Any | None = None, + ) -> DiffusionPagedAttentionContext: + if self._active_batch is None: + raise RuntimeError("Paged attention forward must run inside adapter.activate(batch)") + try: + layer = self.layers[layer_name] + except KeyError as exc: + raise KeyError(f"Unknown diffusion paged attention layer {layer_name!r}") from exc + batch = self._active_batch + full_attn_spans = self._validate_omni_attn_metadata(omni_attn_metadata) + + query_flat, query_token_shape, query_has_head_dims = self._flatten_tensor( + query, + num_heads=layer.num_heads, + head_size=layer.head_size, + name="query", + ) + key_flat, key_token_shape, _ = self._flatten_tensor( + key, + num_heads=layer.num_kv_heads, + head_size=layer.head_size, + name="key", + ) + value_flat, value_token_shape, _ = self._flatten_tensor( + value, + num_heads=layer.num_kv_heads, + head_size=layer.head_size_v, + name="value", + ) + if key_token_shape != query_token_shape: + if len(key_token_shape) == 1: + raise ValueError( + "Paged attention key token count must match the query layout: " + f"query={query_token_shape}, key={key_token_shape}" + ) + raise ValueError( + "Paged attention batched key layout must match prepared rows and the query layout: " + f"query={query_token_shape}, key={key_token_shape}" + ) + if value_token_shape != query_token_shape: + if len(value_token_shape) == 1: + raise ValueError( + "Paged attention value token count must match the query layout: " + f"query={query_token_shape}, value={value_token_shape}" + ) + raise ValueError( + "Paged attention batched value layout must match prepared rows and the query layout: " + f"query={query_token_shape}, value={value_token_shape}" + ) + output_scatter_indices = self._token_packing_indices( + query_token_shape, + batch, + name="Q/K/V", + device=query.device, + ) + if output_scatter_indices is not None: + query_flat = query_flat.index_select(0, output_scatter_indices) + key_flat = key_flat.index_select(0, output_scatter_indices) + value_flat = value_flat.index_select(0, output_scatter_indices) + if query.device != self.device or key.device != self.device or value.device != self.device: + raise ValueError( + f"Paged attention Q/K/V must be on {self.device}; " + f"got query={query.device}, key={key.device}, value={value.device}" + ) + if query.dtype != key.dtype or query.dtype != value.dtype: + raise ValueError(f"Paged attention Q/K/V dtypes must match; got {query.dtype}, {key.dtype}, {value.dtype}") + expected_dtype = self.vllm_config.model_config.dtype + if query.dtype != expected_dtype: + raise ValueError( + f"Paged attention Q/K/V dtype must match model activation dtype {expected_dtype}; got {query.dtype}" + ) + if not bool(getattr(layer.spec, "non_causal", False)): + non_suffix_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len] + if non_suffix_rows: + raise ValueError( + "Causal paged attention requires each query/write span to end at seq_len; " + f"invalid rows={non_suffix_rows!r}" + ) + + slot_mapping = batch.slot_mappings_by_layer.get(layer_name) + if slot_mapping is None: + raise KeyError(f"No native slot mapping was built for diffusion layer {layer_name!r}") + slot_mapping = slot_mapping.reshape(-1)[: batch.num_tokens] + if slot_mapping.numel() != batch.num_tokens: + raise ValueError( + f"Paged attention slot mapping for {layer_name!r} has {slot_mapping.numel()} entries; " + f"expected {batch.num_tokens}" + ) + + # vLLM's CUDA writer consumes the source token strides directly, so + # preserve strided QKV-projection views instead of materializing an + # extra K/V copy here. Platform writers with stricter contracts (for + # example, Ascend PA scatter) make their inputs contiguous locally. + # The native CUDA writer still assumes heads and head dimensions are + # packed within each token; materialize unusual inner layouts only. + if key_flat.stride(-1) != 1 or (key_flat.shape[-2] > 1 and key_flat.stride(-2) != key_flat.shape[-1]): + key_flat = key_flat.contiguous() + if value_flat.stride(-1) != 1 or (value_flat.shape[-2] > 1 and value_flat.stride(-2) != value_flat.shape[-1]): + value_flat = value_flat.contiguous() + # The writer also consumes slot_mapping directly, so a span that starts + # inside a physical block naturally writes only the requested suffix. + try: + native_metadata = batch.attn_metadata[layer_name] + except KeyError as exc: + raise KeyError(f"No native attention metadata was built for diffusion layer {layer_name!r}") from exc + piecewise_plan: PagedPiecewisePlan | None = None + piecewise_native_metadata: tuple[Any, ...] = () + if full_attn_spans is not None: + piecewise_plan = self._get_piecewise_plan(full_attn_spans) + native_metadata_by_segment = self._get_piecewise_native_metadata(piecewise_plan) + try: + piecewise_native_metadata = tuple( + segment_metadata[layer_name] for segment_metadata in native_metadata_by_segment + ) + except KeyError as exc: + raise KeyError( + f"No piecewise native attention metadata was built for diffusion layer {layer_name!r}" + ) from exc + return DiffusionPagedAttentionContext( + layer=layer, + query=query_flat.contiguous(), + key_write=key_flat, + value_write=value_flat, + slot_mapping=slot_mapping, + native_metadata=native_metadata, + piecewise_plan=piecewise_plan, + piecewise_native_metadata=piecewise_native_metadata, + query_token_shape=query_token_shape, + query_has_head_dims=query_has_head_dims, + output_scatter_indices=output_scatter_indices, + ) + + @staticmethod + def _validate_omni_attn_metadata( + metadata: Any | None, + ) -> list[list[tuple[int, int]]] | None: + if metadata is None: + return None + full_attn_spans = getattr(metadata, "full_attn_spans", None) + attn_mask = getattr(metadata, "attn_mask", None) + unsupported_fields = [ + field_name + for field_name in ( + "joint_attn_mask", + "query_ranges", + "video_layout", + ) + if getattr(metadata, field_name, None) is not None + ] + if attn_mask is not None: + if full_attn_spans is None: + unsupported_fields.append("attn_mask") + elif not isinstance(attn_mask, torch.Tensor) or attn_mask.ndim != 4: + unsupported_fields.append("attn_mask (piecewise paging requires a 4D tensor)") + extra = getattr(metadata, "extra", None) + if extra: + unsupported_fields.append(f"extra={sorted(extra)}") + if unsupported_fields: + raise NotImplementedError( + "Diffusion paged attention cannot translate Omni attention metadata fields " + f"{unsupported_fields!r} to native FlashAttention metadata" + ) + return full_attn_spans diff --git a/vllm_omni/diffusion/diffusion_kv/request.py b/vllm_omni/diffusion/diffusion_kv/request.py index 4db9701a1..fd4740099 100644 --- a/vllm_omni/diffusion/diffusion_kv/request.py +++ b/vllm_omni/diffusion/diffusion_kv/request.py @@ -105,7 +105,7 @@ class DiffusionKVRequest: self.status = RequestStatus.WAITING self.num_preemptions = 0 self.num_in_flight_tokens = 0 - # vLLM 0.26 uses this optional boundary when publishing cached blocks; + # vLLM 0.26+ uses this optional boundary when publishing cached blocks; # zero means that no sparse-retention boundary is active. self.shared_prefix_boundary = 0 diff --git a/vllm_omni/diffusion/distributed/__init__.py b/vllm_omni/diffusion/distributed/__init__.py index 668ace06d..2a1159026 100644 --- a/vllm_omni/diffusion/distributed/__init__.py +++ b/vllm_omni/diffusion/distributed/__init__.py @@ -6,11 +6,6 @@ from vllm_omni.diffusion.distributed.hsdp import ( HSDPInferenceConfig, apply_hsdp_to_model, ) -from vllm_omni.diffusion.distributed.parallel_state import ( - get_fs_group, - get_fully_shard_rank, - get_fully_shard_world_size, -) from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelConfig, SequenceParallelInput, @@ -46,8 +41,4 @@ __all__ = [ "get_sharding_validator", "HSDPInferenceConfig", "apply_hsdp_to_model", - # FS (Fully Shard) utilities - "get_fs_group", - "get_fully_shard_rank", - "get_fully_shard_world_size", ] diff --git a/vllm_omni/diffusion/distributed/a2a_permute.py b/vllm_omni/diffusion/distributed/a2a_permute.py new file mode 100644 index 000000000..d6ea05218 --- /dev/null +++ b/vllm_omni/diffusion/distributed/a2a_permute.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused permute-free Ulysses all-to-all over NCCL symmetric memory. + +JIT-compiles the CUDA kernel from pytorch/pytorch#178230 (``all_to_all_permute``) +and exposes it as two *functional* custom ops: + + ulysses_qkv_fwd(x, group_name, world_size) # (B, S/p, H, D) -> (B, S, H/p, D) + ulysses_o_rev(y, group_name, world_size) # (B, S, H/p, D) -> (B, S/p, H, D) + +These replace the synchronous ``all_to_all_4D`` (permute + NCCL all_to_all_single) +used by Ulysses SP. All symmetric-memory bookkeeping (a shape-keyed, rendezvoused +buffer cache + the copy-in) lives *inside* the ops, so from torch.compile's point +of view each op is an opaque ``input -> output`` function: no graph break, no +visible buffer mutation, no Python control flow in the traced graph. A +``register_fake`` provides output metadata so Dynamo keeps the op in-graph. +""" + +from __future__ import annotations + +import glob +import math +import os +import sysconfig +import threading +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from torch.distributed.distributed_c10d import _resolve_process_group +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_BUILD_LOCK = threading.Lock() +_BUILT = False +_WORKSPACE_LOCK = threading.Lock() + + +@dataclass(slots=True) +class _SymmWorkspace: + allocation: torch.Tensor + handle: Any + capacity_bytes: int + stream_id: int + + +# A single grow-only byte workspace per device/process-group. Shape changes +# reuse a typed view into the workspace instead of retaining one allocation per +# historical request shape. +_SYMM_WORKSPACES: dict[tuple[torch.device, str], _SymmWorkspace] = {} + + +def _ensure_built() -> None: + """JIT-compile + load the kernel once (process-wide). Cached by torch.""" + global _BUILT + if _BUILT: + return + with _BUILD_LOCK: + if _BUILT: + return + from torch.utils.cpp_extension import load + + here = os.path.dirname(__file__) + src = os.path.join(here, "csrc", "a2a_permute.cu") + site_packages = sysconfig.get_paths()["purelib"] + nvidia_root = os.path.join(site_packages, "nvidia") + include_paths = glob.glob(os.path.join(nvidia_root, "*", "include")) + nccl_libs = glob.glob(os.path.join(nvidia_root, "nccl", "lib", "libnccl.so*")) + if not nccl_libs: + raise RuntimeError("a2a_permute: could not locate nvidia-nccl libnccl.so") + if not any(os.path.isfile(os.path.join(path, "nccl.h")) for path in include_paths): + raise RuntimeError("a2a_permute: could not locate nvidia-nccl nccl.h") + load( + name="vllm_omni_a2a_permute", + sources=[src], + # PyTorch wheel-only installations keep headers such as cusparse.h + # under nvidia/{cu13,cusparse}/include rather than CUDA_HOME. + extra_include_paths=list(include_paths), + extra_cflags=["-DUSE_NCCL", "-DUSE_C10D_NCCL", "-O3"], + extra_cuda_cflags=["-DUSE_NCCL", "-DUSE_C10D_NCCL", "-O3", "--expt-relaxed-constexpr"], + extra_ldflags=[nccl_libs[0]], + is_python_module=False, + verbose=False, + ) + symm_mem.set_backend("NCCL") + logger.info("[a2a_permute] JIT kernel built and loaded; symm-mem backend=NCCL") + _BUILT = True + + +def ensure_a2a_permute_available() -> None: + """Build the extension during worker/model initialization, not a request.""" + _ensure_built() + + +def _required_nbytes(shape: tuple[int, ...], dtype: torch.dtype) -> int: + return math.prod(shape) * torch.empty((), dtype=dtype).element_size() + + +def _get_symm_buffer( + symm_shape: tuple[int, ...], dtype: torch.dtype, device: torch.device, group_name: str +) -> torch.Tensor: + """Get a typed view from a grow-only symmetric-memory workspace. + + Allocation/growth and rendezvous are collective. Every rank reaches them in + lockstep because sequence parallel execution is symmetric. Steady-state + requests at or below the peak capacity take no synchronization path. + """ + device = torch.device(device) + key = (device, group_name) + required_bytes = _required_nbytes(symm_shape, dtype) + stream_id = torch.cuda.current_stream(device).cuda_stream + with _WORKSPACE_LOCK: + workspace = _SYMM_WORKSPACES.get(key) + if workspace is not None and workspace.stream_id != stream_id: + raise RuntimeError("a2a_permute workspace must be used from a single CUDA stream") + if workspace is None or workspace.capacity_bytes < required_bytes: + if torch.cuda.is_current_stream_capturing(): + capacity = 0 if workspace is None else workspace.capacity_bytes + raise RuntimeError( + "a2a_permute: cannot grow the symmetric-memory workspace " + f"from {capacity} to {required_bytes} bytes during CUDA graph capture; " + "warm up the maximum request shape before capture" + ) + if workspace is None: + # NCCL symmetric-memory rendezvous requires the communicator to + # exist before the first allocation. + pg = _resolve_process_group(group_name) + warm = torch.ones(1, device=device) + dist.all_reduce(warm, group=pg) + # Growth is rare and must not release the previous allocation while + # an earlier kernel is still consuming it. + torch.accelerator.synchronize(device) + allocation = symm_mem.empty(required_bytes, dtype=torch.uint8, device=device) + handle = symm_mem.rendezvous(allocation, group_name) + workspace = _SymmWorkspace( + allocation=allocation, + handle=handle, + capacity_bytes=required_bytes, + stream_id=stream_id, + ) + _SYMM_WORKSPACES[key] = workspace + return workspace.allocation[:required_bytes].view(dtype).view(symm_shape) + + +def clear_a2a_permute_workspaces() -> None: + """Release cached symmetric-memory workspaces during worker shutdown.""" + with _WORKSPACE_LOCK: + if not _SYMM_WORKSPACES: + return + devices = {workspace.allocation.device for workspace in _SYMM_WORKSPACES.values()} + for device in devices: + torch.accelerator.synchronize(device) + count = len(_SYMM_WORKSPACES) + _SYMM_WORKSPACES.clear() + logger.info("[a2a_permute] Released %d symmetric-memory workspace(s)", count) + + +# --------------------------------------------------------------------------- +# Forward: (B, S_local, H, D) -> (B, S_global, H_local, D) +# scatter heads (dim 2), gather sequence (dim 1). Matches all_to_all_4D(x,2,1). +# --------------------------------------------------------------------------- +@torch.library.custom_op("vllm_omni_a2a::ulysses_qkv_fwd", mutates_args=(), device_types="cuda") +def ulysses_qkv_fwd(x: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor: + _ensure_built() + p = world_size + B, s_local, H, D = x.shape + Hl = H // p + lc = Hl * D + rows = B * s_local + symm_in = _get_symm_buffer((rows, p, lc), x.dtype, x.device, group_name) + # (B, S_local, H, D) -> (rows, p, lc); H is row-major so column block r = heads [r*Hl:(r+1)*Hl] + # Fused QKV projections produce row-strided views. Stage complete rows with + # the copy engine instead of materializing the view with TensorIterator. + torch.ops.a2ap.copy_rows(x.view(rows, p, lc), symm_in) + out = torch.empty(p, rows, lc, device=x.device, dtype=x.dtype) + torch.ops.a2ap.all_to_all_permute(symm_in, out, 1, 0, group_name) + # (p, rows, lc) -> (B, S_global, H_local, D), sequence ordered rank-major + return out.reshape(p, B, s_local, Hl, D).permute(1, 0, 2, 3, 4).reshape(B, p * s_local, Hl, D).contiguous() + + +@ulysses_qkv_fwd.register_fake +def _(x: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor: + B, s_local, H, D = x.shape + return x.new_empty(B, s_local * world_size, H // world_size, D) + + +# --------------------------------------------------------------------------- +# Reverse: (B, S_global, H_local, D) -> (B, S_local, H, D) +# scatter sequence (dim 1), gather heads (dim 2). Matches all_to_all_4D(y,1,2). +# --------------------------------------------------------------------------- +@torch.library.custom_op("vllm_omni_a2a::ulysses_o_rev", mutates_args=(), device_types="cuda") +def ulysses_o_rev(y: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor: + _ensure_built() + p = world_size + B, s_global, Hl, D = y.shape + s_local = s_global // p + H = Hl * p + cols = Hl * D + rows = B * s_local + symm_in = _get_symm_buffer((p, rows, cols), y.dtype, y.device, group_name) + # (B, p*S_local, Hl, D) -> (p, B*S_local, cols): block r = sequence shard destined to rank r + symm_in.copy_(y.reshape(B, p, s_local, Hl, D).permute(1, 0, 2, 3, 4).reshape(p, rows, cols)) + out = torch.empty(rows, p, cols, device=y.device, dtype=y.dtype) + torch.ops.a2ap.all_to_all_permute(symm_in, out, 0, 1, group_name) + # (rows, p, cols) -> (B, S_local, H, D), heads ordered rank-major + return out.reshape(B, s_local, p, Hl, D).reshape(B, s_local, H, D).contiguous() + + +@ulysses_o_rev.register_fake +def _(y: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor: + B, s_global, Hl, D = y.shape + return y.new_empty(B, s_global // world_size, Hl * world_size, D) diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py index d804740f9..e8ae6bf32 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -9,7 +9,7 @@ import torch import torch.distributed as dist from vllm.logger import init_logger -from vllm_omni.diffusion.distributed.parallel_state import get_dit_group +from vllm_omni.diffusion.distributed.parallel_state import get_world_group logger = init_logger(__name__) @@ -45,7 +45,8 @@ class DistributedVaeExecutor: """ def __init__(self): - self.group = get_dit_group() + # Use a dedicated process group spanning the complete worker WORLD. + self.group = get_world_group().device_group self.world_size = dist.get_world_size(self.group) self.rank = dist.get_rank(self.group) self.parallel_size = 1 @@ -189,7 +190,7 @@ class DistributedVaeMixin: if self.distributed_executor.parallel_size > pp_size: logger.warning( f"vae_patch_parallel_size={self.distributed_executor.parallel_size} " - f"is greater than dit_group={world_size};" - f" using dit_group size={world_size}" + f"is greater than WORLD={world_size};" + f" using WORLD size={world_size}" ) return True diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 55130e0cf..4c20f1587 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/vllm_omni/diffusion/distributed/cfg_parallel.py @@ -235,8 +235,8 @@ class CFGParallelMixin(metaclass=ABCMeta): Predict noise with N-branch CFG dispatch across M GPUs. This is the multi-branch counterpart of predict_noise_maybe_with_cfg(). - Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel, - DreamID). Existing 2-branch models should continue using + Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel). + Existing 2-branch models should continue using predict_noise_maybe_with_cfg(). Args: diff --git a/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu b/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu new file mode 100644 index 000000000..1c218a1cf --- /dev/null +++ b/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu @@ -0,0 +1,268 @@ +// JIT build of pytorch/pytorch#178230 all_to_all_permute (Ulysses-style fused +// permute-free all-to-all over NCCL symmetric memory). +// +// Kernel is verbatim from the PR; the host entry point is adapted to the +// NCCLDevCommManager API shipped in torch 2.11 (we obtain the host ncclComm_t +// via the manager and create our own ncclDevComm with enough LSA barriers, +// cached per group — the PR's keyed-devcomm manager API is not in 2.11). + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +using namespace c10d::symmetric_memory; + +#ifndef NCCL_HAS_SYMMEM_DEVICE_SUPPORT +#error "Build requires NCCL >= 2.28 with symmetric-memory device support" +#endif + +void copy_rows(const at::Tensor& input, at::Tensor& out) { + TORCH_CHECK(input.is_cuda() && out.is_cuda(), "a2a_permute: copy_rows requires CUDA tensors"); + TORCH_CHECK(input.device() == out.device(), "a2a_permute: copy_rows device mismatch"); + TORCH_CHECK(input.scalar_type() == out.scalar_type(), "a2a_permute: copy_rows dtype mismatch"); + TORCH_CHECK(input.dim() == 3, "a2a_permute: copy_rows input must be 3-D"); + TORCH_CHECK(out.sizes() == input.sizes(), "a2a_permute: copy_rows shape mismatch"); + TORCH_CHECK(out.is_contiguous(), "a2a_permute: copy_rows output must be contiguous"); + TORCH_CHECK( + input.stride(2) == 1 && input.stride(1) == input.size(2), + "a2a_permute: copy_rows requires contiguous rows"); + + const int64_t row_elements = input.size(1) * input.size(2); + TORCH_CHECK(input.stride(0) >= row_elements, "a2a_permute: copy_rows input rows overlap"); + const size_t rows = static_cast(input.size(0)); + const size_t row_bytes = static_cast(row_elements) * input.element_size(); + const size_t input_pitch = + static_cast(input.stride(0) * input.element_size()); + + c10::cuda::CUDAGuard guard(input.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + C10_CUDA_CHECK(cudaMemcpy2DAsync( + out.data_ptr(), + row_bytes, + input.data_ptr(), + input_pitch, + row_bytes, + rows, + cudaMemcpyDeviceToDevice, + stream)); +} + +// ---- kernel (verbatim from PR) ---------------------------------------------- +__device__ inline void copy_bytes_vec16_aligned( + const char* src, char* dst, size_t nbytes, size_t tid, size_t stride) { + const size_t n_vec = nbytes / 16; + constexpr int kUnroll = 4; + size_t vec_idx = tid; + for (; vec_idx + static_cast(kUnroll - 1) * stride < n_vec; + vec_idx += static_cast(kUnroll) * stride) { + at::native::memory::Vec<16> chunk[kUnroll]; +#pragma unroll 4 + for (int u = 0; u < kUnroll; ++u) { + const size_t i = vec_idx + static_cast(u) * stride; + chunk[u] = at::native::memory::ld_vec<16>(src + i * 16); + } +#pragma unroll 4 + for (int u = 0; u < kUnroll; ++u) { + const size_t i = vec_idx + static_cast(u) * stride; + at::native::memory::st_vec<16>(dst + i * 16, chunk[u]); + } + } + for (; vec_idx < n_vec; vec_idx += stride) { + auto v = at::native::memory::ld_vec<16>(src + vec_idx * 16); + at::native::memory::st_vec<16>(dst + vec_idx * 16, v); + } +} + +constexpr int A2A_MAX_SLOTS = 64; +constexpr int A2A_MAX_CTAS_PER_SLOT = 128; +constexpr int A2A_THREADS_PER_CTA = 128; +constexpr int A2A_MAX_CTA_COUNT = A2A_MAX_SLOTS * A2A_MAX_CTAS_PER_SLOT; + +__global__ void all_to_all_lsa_kernel( + ncclWindow_t window, + size_t base_src_byte_offset, + unsigned char* out, + int num_rows, + size_t src_row_stride_bytes, + size_t copy_row_bytes, + size_t peer_stride_bytes, + size_t dst_row_stride_bytes, + ncclDevComm devComm) { + const int peer_idx = blockIdx.x; + const int local_block = blockIdx.y; + const ncclCoopCta coop{}; + + ncclLsaBarrierSession bar{ + coop, devComm, ncclTeamLsa(devComm), devComm.lsaBarrier, + blockIdx.x * gridDim.y + blockIdx.y}; + bar.sync(coop, cuda::memory_order_acquire); + + unsigned char* dst_peer_base = + out + static_cast(peer_idx) * peer_stride_bytes; + CUDA_KERNEL_ASSERT((base_src_byte_offset & 15) == 0); + CUDA_KERNEL_ASSERT((reinterpret_cast(dst_peer_base) & 15) == 0); + + for (int i = local_block; i < num_rows; i += gridDim.y) { + const size_t row_byte_offset = + base_src_byte_offset + static_cast(i) * src_row_stride_bytes; + const void* src_row = ncclGetLsaPointer(window, row_byte_offset, peer_idx); + unsigned char* dst_row = + dst_peer_base + static_cast(i) * dst_row_stride_bytes; + copy_bytes_vec16_aligned( + reinterpret_cast(src_row), + reinterpret_cast(dst_row), + copy_row_bytes, coop.thread_rank(), coop.size()); + } + bar.sync(coop, cuda::memory_order_release); +} + +// ---- own devcomm cache (replaces PR's keyed manager API) -------------------- +ncclDevComm get_or_create_devcomm(ncclComm_t comm, const std::string& group_name) { + static std::mutex mu; + static std::map cache; + std::lock_guard lk(mu); + auto it = cache.find(group_name); + if (it != cache.end()) return it->second; +#ifdef NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER + ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER; +#else + ncclDevCommRequirements reqs; + memset(&reqs, 0, sizeof(ncclDevCommRequirements)); +#endif + reqs.lsaBarrierCount = A2A_MAX_CTA_COUNT; + ncclDevComm dc; + ncclResult_t res = ncclDevCommCreate(comm, &reqs, &dc); + TORCH_CHECK(res == ncclSuccess, "ncclDevCommCreate (a2a_permute) failed: ", ncclGetErrorString(res)); + cache.emplace(group_name, dc); + return dc; +} + +// ---- host entry ------------------------------------------------------------- +void all_to_all_permute( + const at::Tensor& input, + at::Tensor& out, + int64_t scatter_dim, + int64_t gather_dim, + std::string group_name) { + TORCH_CHECK(input.stride(-1) == 1, "a2a_permute: innermost dim must be contiguous"); + const bool col_scatter = (scatter_dim == 1 && gather_dim == 0); + const bool row_scatter = (scatter_dim == 0 && gather_dim == 1); + TORCH_CHECK(col_scatter || row_scatter, + "a2a_permute: unsupported (scatter_dim, gather_dim)=(", scatter_dim, ",", gather_dim, ")"); + + auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name); + TORCH_CHECK(symm_mem != nullptr, "a2a_permute: input must be NCCL symmetric memory"); + auto* nccl_hdl = dynamic_cast(symm_mem.get()); + TORCH_CHECK(nccl_hdl != nullptr, "a2a_permute: requires NCCL symmetric memory backend"); + + c10::cuda::CUDAGuard guard(input.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto device = input.device(); + + auto& manager = NCCLDevCommManager::get(device); + ncclComm_t comm = manager.get_comm(group_name); + ncclDevComm devcomm = get_or_create_devcomm(comm, group_name); + + const int my_rank = nccl_hdl->get_rank(); + const int p = nccl_hdl->get_world_size(); + TORCH_CHECK(p <= A2A_MAX_SLOTS, "a2a_permute: group size ", p, " exceeds ", A2A_MAX_SLOTS); + TORCH_CHECK(out.is_contiguous(), "a2a_permute: out must be contiguous"); + TORCH_CHECK(out.scalar_type() == input.scalar_type(), "a2a_permute: dtype mismatch"); + + auto window = nccl_hdl->get_window(); + TORCH_CHECK(window != nullptr, "a2a_permute: NCCL window is null"); + + const size_t window_base_offset = nccl_hdl->get_offset(); + const int64_t esize = input.element_size(); + const size_t tensor_leading_offset = + window_base_offset + static_cast(input.storage_offset()) * static_cast(esize); + TORCH_CHECK(tensor_leading_offset % 16 == 0, "a2a_permute: tensor byte offset must be 16B aligned"); + TORCH_CHECK(reinterpret_cast(input.data_ptr()) % 16 == 0, "a2a_permute: input ptr must be 16B aligned"); + TORCH_CHECK(reinterpret_cast(out.data_ptr()) % 16 == 0, "a2a_permute: out ptr must be 16B aligned"); + + const int unroll = 4 * 16 / static_cast(esize); + const int elems_per_cta = A2A_THREADS_PER_CTA * unroll; + int ctas_per_slot = 1; + + if (col_scatter) { + const int rows = static_cast(input.size(0)); + int64_t total_cols = 0; int local_cols = 0; + if (input.dim() == 2) { + total_cols = input.size(1); + TORCH_CHECK(total_cols % p == 0, "a2a_permute: cols must divide group size"); + local_cols = static_cast(total_cols / p); + } else { + TORCH_CHECK(input.dim() == 3, "a2a_permute: input must be 2-D or 3-D"); + TORCH_CHECK(input.size(1) == p, "a2a_permute: 3-D input dim1 must equal group size"); + const int64_t lc = input.size(2); + total_cols = static_cast(p) * lc; + TORCH_CHECK(input.stride(1) == lc && input.stride(0) == total_cols, + "a2a_permute: 3-D input must be row-major contiguous"); + local_cols = static_cast(lc); + } + const bool ok3 = out.dim() == 3 && out.size(0) == p && out.size(1) == rows && out.size(2) == local_cols; + const bool ok2 = out.dim() == 2 && out.size(0) == (int64_t)p * rows && out.size(1) == local_cols; + TORCH_CHECK(ok3 || ok2, "a2a_permute: bad out shape for (1,0)"); + const size_t row_bytes = (size_t)local_cols * (size_t)esize; + TORCH_CHECK(row_bytes % 16 == 0, "a2a_permute: local_cols*esize must be 16B-divisible"); + ctas_per_slot = std::max(1, std::min((rows * local_cols + elems_per_cta - 1) / elems_per_cta, A2A_MAX_CTAS_PER_SLOT)); + const size_t esz = (size_t)esize; + const size_t base_src = tensor_leading_offset + (size_t)my_rank * (size_t)local_cols * esz; + all_to_all_lsa_kernel<<>>( + window, base_src, reinterpret_cast(out.data_ptr()), rows, + (size_t)total_cols * esz, row_bytes, (size_t)rows * local_cols * esz, (size_t)local_cols * esz, devcomm); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + int local_rows = 0; int cols = 0; + if (input.dim() == 2) { + const int64_t total_rows = input.size(0); + cols = static_cast(input.size(1)); + TORCH_CHECK(total_rows % p == 0, "a2a_permute: rows must divide group size"); + local_rows = static_cast(total_rows / p); + } else { + TORCH_CHECK(input.dim() == 3, "a2a_permute: input must be 2-D or 3-D"); + TORCH_CHECK(input.size(0) == p, "a2a_permute: 3-D input dim0 must equal group size"); + local_rows = static_cast(input.size(1)); + const int64_t c = input.size(2); cols = static_cast(c); + const int64_t s01 = (int64_t)local_rows * c; + TORCH_CHECK(input.stride(1) == c && input.stride(0) == s01, "a2a_permute: 3-D input must be row-major contiguous"); + } + const bool ok3 = out.dim() == 3 && out.size(0) == local_rows && out.size(1) == p && out.size(2) == cols; + const bool ok2 = out.dim() == 2 && out.size(0) == local_rows && out.size(1) == (int64_t)p * cols; + TORCH_CHECK(ok3 || ok2, "a2a_permute: bad out shape for (0,1)"); + const size_t row_bytes = (size_t)cols * (size_t)esize; + TORCH_CHECK(row_bytes % 16 == 0, "a2a_permute: cols*esize must be 16B-divisible"); + ctas_per_slot = std::max(1, std::min((local_rows * cols + elems_per_cta - 1) / elems_per_cta, A2A_MAX_CTAS_PER_SLOT)); + const size_t esz = (size_t)esize; const size_t cu = (size_t)cols; + const size_t base_src = tensor_leading_offset + (size_t)(my_rank * local_rows) * cu * esz; + all_to_all_lsa_kernel<<>>( + window, base_src, reinterpret_cast(out.data_ptr()), local_rows, + cu * esz, row_bytes, cu * esz, (size_t)p * cu * esz, devcomm); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } +} + +} // namespace + +TORCH_LIBRARY_FRAGMENT(a2ap, m) { + m.def("copy_rows(Tensor input, Tensor(a!) out) -> ()"); + m.def("all_to_all_permute(Tensor input, Tensor(a!) out, int scatter_dim, int gather_dim, str group_name) -> ()"); +} +TORCH_LIBRARY_IMPL(a2ap, CUDA, m) { + m.impl("copy_rows", TORCH_FN(copy_rows)); + m.impl("all_to_all_permute", TORCH_FN(all_to_all_permute)); +} diff --git a/vllm_omni/diffusion/distributed/csrc/nccl_extension.hpp b/vllm_omni/diffusion/distributed/csrc/nccl_extension.hpp new file mode 100644 index 000000000..e9fe0e19c --- /dev/null +++ b/vllm_omni/diffusion/distributed/csrc/nccl_extension.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include +#include + +namespace c10d::nccl_extension { + +TORCH_API bool is_nccl_symmem_available(); + +TORCH_API void nccl_put(at::Tensor& tensor, const int64_t peer); + +TORCH_API void nccl_get(at::Tensor& tensor, const int64_t peer); + +TORCH_API void nccl_wait_for_signal(at::Tensor& sigpad, int64_t signal); + +TORCH_API void nccl_put_with_signal( + at::Tensor& tensor, + int64_t signal, + int64_t peer); + +// Simultaneously reduce N blocks of a 2-D input tensor from a shared symmetric +// memory buffer, routing each to a specific destination rank. Blocks are +// described by inclusive-prefix-sum offsets along `dim` (0 or 1); all blocks +// must have equal size. +TORCH_API void nccl_reduce_scatter_offset( + const at::Tensor& input, + at::TensorList out, + const std::string& group_name, + int64_t dim, + std::optional offsets, + std::optional dst_ranks, + const std::string& red_op); + +// Permute-free all-to-all: scatter shards along `scatter_dim`, gather peer +// chunks along `gather_dim`. Supported: (scatter_dim=1, gather_dim=0) -> [rows, +// p*lc] to [p, rows, lc]; (scatter_dim=0, gather_dim=1) -> [p*lr, cols] to [lr, +// p, cols]. +TORCH_API void nccl_all_to_all_permute( + const at::Tensor& input, + at::Tensor& out, + int64_t scatter_dim, + int64_t gather_dim, + const std::string& group_name); +} // namespace c10d::nccl_extension diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 14847007f..ac2059f56 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -991,6 +991,18 @@ class PipelineGroupCoordinator(GroupCoordinator): def _pipeline_isend_skip(self, tensor: torch.tensor): return torch.distributed.isend(tensor, dst=self.skip_rank, group=self.skip_device_group) + def destroy(self): + groups = [*self.device_groups, *self.cpu_groups, self.skip_device_group] + seen = {id(self.device_group), id(self.cpu_group)} + for group in groups: + if group is not None and id(group) not in seen: + torch.distributed.destroy_process_group(group) + seen.add(id(group)) + self.device_groups = [] + self.cpu_groups = [] + self.skip_device_group = None + super().destroy() + class SequenceParallelGroupCoordinator(GroupCoordinator): def __init__( @@ -1031,3 +1043,15 @@ class SequenceParallelGroupCoordinator(GroupCoordinator): self.ring_rank = torch.distributed.get_rank(self.ring_group) self.allgather_world_size = torch.distributed.get_world_size(self.allgather_group) self.allgather_rank = torch.distributed.get_rank(self.allgather_group) + + def destroy(self): + groups = [self.ulysses_group, self.ring_group, self.allgather_group] + seen = {id(self.device_group), id(self.cpu_group)} + for group in groups: + if group is not None and id(group) not in seen: + torch.distributed.destroy_process_group(group) + seen.add(id(group)) + self.ulysses_group = None + self.ring_group = None + self.allgather_group = None + super().destroy() diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index cae4356f4..c496804bc 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/vllm_omni/diffusion/distributed/hsdp.py @@ -15,12 +15,7 @@ from torch.distributed.fsdp import ( ) from vllm.logger import init_logger -from vllm_omni.diffusion.distributed.parallel_state import ( - get_fs_group, - get_fully_shard_rank, - get_fully_shard_world_size, - get_world_group, -) +from vllm_omni.diffusion.distributed.parallel_state import get_world_group from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -54,36 +49,36 @@ class HSDPInferenceConfig: def _create_hsdp_mesh( device_type: str, replicate_size: int, - shard_pg: torch.distributed.ProcessGroup, + shard_size: int, ) -> DeviceMesh: - """Create a 2D DeviceMesh for HSDP using an existing ProcessGroup for the shard dimension. + """Create the FSDP2 DeviceMesh; it owns the HSDP process groups. Args: device_type: The device type (e.g., "cuda", "npu") replicate_size: Number of replica groups - shard_pg: The ProcessGroup for the shard dimension (from FS GroupCoordinator) + shard_size: Number of ranks in each FSDP shard group Returns: A 2D DeviceMesh with dimensions ("replicate", "shard") """ - shard_size = torch.distributed.get_world_size(shard_pg) world_size = replicate_size * shard_size + actual_world_size = torch.distributed.get_world_size() + if world_size != actual_world_size: + raise ValueError( + f"HSDP mesh dimensions ({replicate_size} x {shard_size} = {world_size}) " + f"must equal WORLD size ({actual_world_size})" + ) # Build 2D mesh tensor: shape (replicate_size, shard_size) # Ranks are arranged so that each row is a shard group mesh_tensor = torch.arange(world_size).reshape(replicate_size, shard_size) - # Create DeviceMesh with the shard ProcessGroup - # For the shard dimension, we reuse the existing FS ProcessGroup device_mesh = init_device_mesh( device_type, mesh_shape=(replicate_size, shard_size), mesh_dim_names=("replicate", "shard"), ) - # Note: init_device_mesh creates new ProcessGroups internally. - # For consistency, we verify the mesh structure matches our FS group. - # In a future optimization, we could pass the existing ProcessGroups directly. logger.debug( "Created HSDP mesh: replicate_size=%d, shard_size=%d, mesh=%s", replicate_size, @@ -119,34 +114,24 @@ def apply_hsdp_to_model( if not hsdp_config.enabled: raise ValueError("HSDP is not enabled in config") - # Use GroupCoordinator for distributed info world_group = get_world_group() - fs_group = get_fs_group() - world_size = world_group.world_size rank = world_group.rank_in_group - fs_world_size = get_fully_shard_world_size() - fs_rank = get_fully_shard_rank() hsdp_replicate_size = hsdp_config.hsdp_replicate_size hsdp_shard_size = hsdp_config.hsdp_shard_size - # Validate that the FS group matches the HSDP shard size - if fs_world_size != hsdp_shard_size: + if hsdp_replicate_size * hsdp_shard_size != world_size: raise ValueError( - f"FS group world_size ({fs_world_size}) does not match " - f"HSDP shard_size ({hsdp_shard_size}). " - "Ensure fully_shard_degree is set correctly in initialize_model_parallel." + f"HSDP mesh dimensions ({hsdp_replicate_size} x {hsdp_shard_size}) must equal WORLD size ({world_size})" ) logger.info( - "HSDP Inference: replicate_size=%d, shard_size=%d, world_size=%d, rank=%d, fs_world_size=%d, fs_rank=%d", + "HSDP Inference (FSDP2): replicate_size=%d, shard_size=%d, world_size=%d, rank=%d", hsdp_replicate_size, hsdp_shard_size, world_size, rank, - fs_world_size, - fs_rank, ) # When the model contains FP8 parameters (online quantization), let FSDP @@ -163,14 +148,14 @@ def apply_hsdp_to_model( device_type = current_omni_platform.device_type - # Create 2D DeviceMesh for HSDP using the FS group's ProcessGroup for shard dimension + # DeviceMesh is the single source of truth for FSDP2 shard/replicate groups # The mesh shape is (replicate, shard) where: # - replicate: groups of ranks that hold the same shard (for gradient all-reduce in training) # - shard: groups of ranks that each hold different shards (for parameter all-gather) device_mesh = _create_hsdp_mesh( device_type=device_type, replicate_size=hsdp_replicate_size, - shard_pg=fs_group.device_group, + shard_size=hsdp_shard_size, ) hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 55324ccf4..85ba9cd6c 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py # Copyright 2023 The vLLM team. # Adapted from @@ -29,11 +30,11 @@ If you only need to use the distributed environment without model parallelism, """ import inspect +from math import prod import torch import torch.distributed import vllm.distributed.parallel_state as vllm_parallel_state -from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm_omni.diffusion import envs @@ -60,7 +61,7 @@ _PP: PipelineGroupCoordinator | None = None _CFG: GroupCoordinator | None = None _DP: GroupCoordinator | None = None _FS: GroupCoordinator | None = None # Fully Sharded (HSDP shard dimension) -_DIT: GroupCoordinator | None = None +_HSDP_REPLICATE: GroupCoordinator | None = None # HSDP replica dimension # Rank-layout metadata for expert parallelism. This is not a process group; # it is reused by platform-specific runtimes that must build companion groups @@ -184,7 +185,6 @@ class RankGenerator: pp: int, cfg: int, dp: int, - fs: int = 1, order: str = "tp-sp-pp-cfg-dp", rank_offset: int = 0, ) -> None: @@ -193,9 +193,7 @@ class RankGenerator: self.pp = pp self.cfg = cfg self.dp = dp - self.fs = fs self.rank_offset = rank_offset - self.world_size = tp * sp * pp * cfg * dp self.name_to_size = { "tp": self.tp, @@ -203,15 +201,10 @@ class RankGenerator: "pp": self.pp, "cfg": self.cfg, "dp": self.dp, - "fs": self.fs, } order = order.lower() for name in self.name_to_size.keys(): - # Skip 'fs' validation - it's handled separately with independent_ranks=True - # and doesn't participate in the main orthogonal rank generation - if name == "fs": - continue if name not in order and self.name_to_size[name] != 1: raise RuntimeError( f"The size of ({name}) is ({self.name_to_size[name]}), " @@ -234,7 +227,7 @@ class RankGenerator: mask[ordered_token.index(t)] = True return mask - def get_ranks(self, token, independent_ranks: bool = False): + def get_ranks(self, token): """Get rank group by input token. Arguments: @@ -243,27 +236,11 @@ class RankGenerator: to obtain multiple parallel types, we can use a hyphen '-' to separate them. For example, if we want to obtain the TP_DP group, the token should be 'tp-dp'. - independent_ranks (bool): - If True, generate independent rank groups that divide the world - into groups of the specified size. Used for FS (fully shard) groups - which operate independently from the main parallelism hierarchy. """ - if independent_ranks and token == "fs": - # FS groups divide world into groups of size fs - # e.g., world_size=8, fs=4 -> [[0,1,2,3], [4,5,6,7]] - ranks = [] - num_groups = self.world_size // self.fs - for i in range(num_groups): - group = list(range(i * self.fs + self.rank_offset, (i + 1) * self.fs + self.rank_offset)) - ranks.append(group) - return ranks - mask = self.get_mask(self.order, token) - ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask) - if self.rank_offset > 0: - for rank_group in ranks: - for i in range(len(rank_group)): - rank_group[i] += self.rank_offset + ranks = generate_masked_orthogonal_rank_groups(prod(self.ordered_size), self.ordered_size, mask) + if self.rank_offset: + return [[rank + self.rank_offset for rank in rank_group] for rank_group in ranks] return ranks @@ -382,14 +359,9 @@ def get_fs_group() -> GroupCoordinator: return _FS -def get_fully_shard_world_size(): - """Return world size for the fully shard group.""" - return get_fs_group().world_size - - -def get_fully_shard_rank(): - """Return my rank for the fully shard group.""" - return get_fs_group().rank_in_group +def get_hsdp_replicate_group() -> GroupCoordinator: + assert _HSDP_REPLICATE is not None, "HSDP replicate group is not initialized" + return _HSDP_REPLICATE def is_dp_last_group(): @@ -401,17 +373,6 @@ def is_dp_last_group(): ) -def get_dit_world_size(): - """Return world size for the DiT model.""" - return ( - get_data_parallel_world_size() - * get_classifier_free_guidance_world_size() - * get_sequence_parallel_world_size() - * get_pipeline_parallel_world_size() - * get_tensor_model_parallel_world_size() - ) - - # * SET @@ -546,19 +507,6 @@ def init_vllm_model_parallel_group( ) -def init_dit_group( - dit_parallel_size: int, - backend: str, -): - global _DIT - _DIT = torch.distributed.new_group(ranks=list(range(dit_parallel_size)), backend=backend) - - -def get_dit_group(): - assert _DIT is not None, "DIT group is not initialized" - return _DIT - - # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py def set_seq_parallel_pg( sp_ulysses_degree: int, @@ -737,8 +685,8 @@ def set_seq_parallel_pg( return ulysses_pg, ring_pg, allgather_pg -def initialize_model_parallel( - data_parallel_size: int = 1, +def _initialize_model_parallel( + data_parallel_size: int | None = None, cfg_parallel_size: int = 1, sequence_parallel_size: int | None = None, ulysses_degree: int = 1, @@ -747,10 +695,12 @@ def initialize_model_parallel( tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, fully_shard_degree: int = 1, - hsdp_replicate_size: int = 1, enable_expert_parallel: bool = False, + use_hsdp: bool = False, backend: str | None = None, ) -> None: + global _FS, _HSDP_REPLICATE + if backend is None: backend = current_omni_platform.dist_backend """ @@ -768,7 +718,7 @@ def initialize_model_parallel( (causal=False only). Mutually exclusive with ulysses/ring in v1. tensor_parallel_size: number of GPUs used for tensor parallelism. pipeline_parallel_size: number of GPUs used for pipeline parallelism. - fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension). + fully_shard_degree: number of GPUs used for the HSDP shard dimension. backend: distributed backend of pytorch collective comm. Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we @@ -798,16 +748,20 @@ def initialize_model_parallel( ranks 8 to 15 belong to the second box. """ # Get world size and rank. Ensure some consistencies. - assert torch.distributed.is_initialized() + if not torch.distributed.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before model parallel initialization") + if model_parallel_is_initialized(): + raise RuntimeError("model parallel groups are already initialized") world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) if allgather_degree > 1: - assert ulysses_degree == 1 and ring_degree == 1, ( - "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " - f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " - f"allgather_degree={allgather_degree}." - ) + if ulysses_degree != 1 or ring_degree != 1: + raise ValueError( + "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " + f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " + f"allgather_degree={allgather_degree}." + ) expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree if sequence_parallel_size is None: @@ -820,42 +774,44 @@ def initialize_model_parallel( f"but got {sequence_parallel_size}" ) - dit_parallel_size = ( - data_parallel_size * cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size - ) - - # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1 - is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1 - - # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size) as dit_parallel_size - # This ensures orthogonal rank generation works correctly for all HSDP workers - if is_standalone_hsdp: - dit_parallel_size = fully_shard_degree * hsdp_replicate_size - - if world_size < dit_parallel_size: - raise RuntimeError( - f"world_size ({world_size}) is less than " - f"tensor_parallel_size ({tensor_parallel_size}) x " - f"pipeline_parallel_size ({pipeline_parallel_size}) x" - f"sequence_parallel_size ({sequence_parallel_size}) x" - f"cfg_parallel_size " - f"({cfg_parallel_size}) x" - f"data_parallel_size ({data_parallel_size})" - ) - - # For standalone HSDP, use (fully_shard_degree * hsdp_replicate_size) as data_parallel_size - # so that RankGenerator.world_size matches the actual number of workers - effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size + non_dp_size = cfg_parallel_size * sequence_parallel_size * pipeline_parallel_size * tensor_parallel_size + if world_size % non_dp_size != 0: + raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})") + if use_hsdp: + if data_parallel_size not in (None, 1): + raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1") + if non_dp_size not in (1, world_size): + raise ValueError(f"HSDP non-DP parallel size must be 1 or WORLD size ({world_size}), but got {non_dp_size}") + if fully_shard_degree <= 0: + raise ValueError(f"fully_shard_degree must be positive, got {fully_shard_degree}") + if world_size % fully_shard_degree != 0: + raise ValueError( + f"WORLD size ({world_size}) must be divisible by fully_shard_degree ({fully_shard_degree})" + ) + data_parallel_size = 1 + else: + inferred_data_parallel_size = world_size // non_dp_size + if data_parallel_size is not None and data_parallel_size != inferred_data_parallel_size: + raise ValueError( + f"data_parallel_size ({data_parallel_size}) does not match WORLD-derived value " + f"({inferred_data_parallel_size})" + ) + data_parallel_size = inferred_data_parallel_size rank_generator: RankGenerator = RankGenerator( tensor_parallel_size, sequence_parallel_size, pipeline_parallel_size, cfg_parallel_size, - effective_dp_size, - fs=fully_shard_degree, + data_parallel_size, order="tp-sp-pp-cfg-dp", ) + + def get_rank_groups(token: str) -> list[list[int]]: + if use_hsdp and non_dp_size == 1: + return [[rank] for rank in range(world_size)] + return rank_generator.get_ranks(token) + use_moe_parallel_mapping = False if enable_expert_parallel: od_config = get_forward_context().omni_diffusion_config @@ -863,11 +819,11 @@ def initialize_model_parallel( if not use_moe_parallel_mapping: raise RuntimeError("Expert parallelism enabled for a non-MoE model") - sp_group_ranks = rank_generator.get_ranks("sp") + sp_group_ranks = get_rank_groups("sp") global _DP assert _DP is None, "data parallel group is already initialized" _DP = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("dp"), + group_ranks=get_rank_groups("dp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="data", @@ -877,7 +833,7 @@ def initialize_model_parallel( global _CFG assert _CFG is None, "classifier_free_guidance group is already initialized" _CFG = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("cfg"), + group_ranks=get_rank_groups("cfg"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="classifier_free_guidance", @@ -885,7 +841,7 @@ def initialize_model_parallel( global _PP assert _PP is None, "pipeline model parallel group is already initialized" _PP = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("pp"), + group_ranks=get_rank_groups("pp"), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="pipeline", @@ -899,7 +855,7 @@ def initialize_model_parallel( sp_ring_degree=ring_degree, sp_allgather_degree=allgather_degree, rank=get_world_group().rank_in_group, - world_size=dit_parallel_size, + world_size=world_size, sp_group_ranks=sp_group_ranks, ) _SP = init_model_parallel_group( @@ -925,7 +881,7 @@ def initialize_model_parallel( ) assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized" - tp_group_ranks = rank_generator.get_ranks("tp") + tp_group_ranks = get_rank_groups("tp") if use_moe_parallel_mapping: vllm_parallel_state._TP = init_vllm_model_parallel_group( group_ranks=tp_group_ranks, @@ -944,22 +900,40 @@ def initialize_model_parallel( # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP # only when constructing the vLLM EP layout for expert-parallel paths. vllm_parallel_state._DP = init_vllm_model_parallel_group( - group_ranks=rank_generator.get_ranks("cfg-dp"), + group_ranks=get_rank_groups("cfg-dp"), local_rank=get_world_group().local_rank, backend=backend, group_name="dp", ) - global _FS, _EXPERT_PARALLEL_GROUP_RANKS - assert _FS is None, "fully shard group is already initialized" - _FS = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("fs", independent_ranks=True), - local_rank=get_world_group().local_rank, - backend=backend, - parallel_mode="fully_shard", - ) + if use_hsdp: + assert _FS is None, "fully shard group is already initialized" + # HSDP builds its mesh from arange(world_size).reshape(replicate, shard), + # so each consecutive rank run is one fully-sharded group. + fs_group_ranks = [ + list(range(start, start + fully_shard_degree)) for start in range(0, world_size, fully_shard_degree) + ] + _FS = init_model_parallel_group( + group_ranks=fs_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="fully_shard", + ) + if world_size > fully_shard_degree: + # The HSDP mesh's columns are replica groups: each column contains + # the same shard position across all replica rows. + hsdp_replicate_group_ranks = [ + list(range(offset, world_size, fully_shard_degree)) for offset in range(fully_shard_degree) + ] + _HSDP_REPLICATE = init_model_parallel_group( + group_ranks=hsdp_replicate_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="fully_shard", + ) - _EXPERT_PARALLEL_GROUP_RANKS = rank_generator.get_ranks("tp-sp-cfg-dp") + global _EXPERT_PARALLEL_GROUP_RANKS + _EXPERT_PARALLEL_GROUP_RANKS = get_rank_groups("tp-sp-cfg-dp") if use_moe_parallel_mapping: vllm_parallel_state._EP = init_vllm_model_parallel_group( group_ranks=_EXPERT_PARALLEL_GROUP_RANKS, @@ -969,12 +943,67 @@ def initialize_model_parallel( use_all2all=True, ) - init_dit_group(dit_parallel_size, backend) + +def initialize_model_parallel( + data_parallel_size: int | None = None, + cfg_parallel_size: int = 1, + sequence_parallel_size: int | None = None, + ulysses_degree: int = 1, + ring_degree: int = 1, + allgather_degree: int = 1, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + fully_shard_degree: int = 1, + enable_expert_parallel: bool = False, + use_hsdp: bool = False, + backend: str | None = None, +) -> None: + """Atomically initialize diffusion parallel groups. + + Configuration is validated before group creation. Any failure tears down + groups created by this attempt, including vLLM-bound coordinators. + """ + existing_groups = { + "dp": _DP, + "cfg": _CFG, + "sp": _SP, + "pp": _PP, + "fs": _FS, + "hsdp_replicate": _HSDP_REPLICATE, + "vllm_tp": vllm_parallel_state._TP, + "vllm_dp": vllm_parallel_state._DP, + "vllm_pp": vllm_parallel_state._PP, + "vllm_ep": vllm_parallel_state._EP, + "vllm_pcp": vllm_parallel_state._PCP, + } + initialized_names = [name for name, group in existing_groups.items() if group is not None] + if initialized_names: + raise RuntimeError( + "model parallel state must be empty before initialization; found " + ", ".join(initialized_names) + ) + try: + _initialize_model_parallel( + data_parallel_size=data_parallel_size, + cfg_parallel_size=cfg_parallel_size, + sequence_parallel_size=sequence_parallel_size, + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + allgather_degree=allgather_degree, + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + fully_shard_degree=fully_shard_degree, + enable_expert_parallel=enable_expert_parallel, + use_hsdp=use_hsdp, + backend=backend, + ) + except BaseException: + destroy_model_parallel() + raise def destroy_model_parallel(): """Set the groups to none and destroy them.""" - global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS + global _DP, _CFG, _SP, _PP, _FS, _HSDP_REPLICATE, _EXPERT_PARALLEL_GROUP_RANKS if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP: vllm_parallel_state._DP.destroy() @@ -984,6 +1013,14 @@ def destroy_model_parallel(): _DP.destroy() _DP = None + if _FS: + _FS.destroy() + _FS = None + + if _HSDP_REPLICATE: + _HSDP_REPLICATE.destroy() + _HSDP_REPLICATE = None + if _CFG: _CFG.destroy() _CFG = None @@ -1005,14 +1042,14 @@ def destroy_model_parallel(): vllm_parallel_state._EP = None _EXPERT_PARALLEL_GROUP_RANKS = None + if vllm_parallel_state._PP and vllm_parallel_state._PP is not _PP: + vllm_parallel_state._PP.destroy() + vllm_parallel_state._PP = None + if _PP: _PP.destroy() _PP = None - if _FS: - _FS.destroy() - _FS = None - def destroy_distributed_environment(): global _WORLD diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index 8ed8fbefd..32e006f7e 100644 --- a/vllm_omni/diffusion/envs.py +++ b/vllm_omni/diffusion/envs.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2024 xDiT team. # Adapted from # https://github.com/xdit-project/xDiT/blob/main/xfuser/envs.py diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index ec6b3486a..f15db04bd 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -5,10 +5,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm_omni.diffusion.data import OmniDiffusionConfig if TYPE_CHECKING: + from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput from vllm_omni.diffusion.worker.utils import BaseRunnerOutput @@ -22,10 +24,15 @@ class DiffusionExecutor(ABC): def get_class(od_config: OmniDiffusionConfig) -> type[DiffusionExecutor]: executor_class: type[DiffusionExecutor] distributed_executor_backend = od_config.distributed_executor_backend - # Keep backward-compatible behavior for callers/configs that omit this - # field and rely on the historical diffusion default backend. + # Mirror vLLM's `world_size == 1 -> "uni"` default + # (vllm/config/parallel.py). A single-GPU diffusion deployment has + # nothing to distribute: spawning a worker only adds MessageQueues, + # /dev/shm output segments, and a second model load. Explicit "mp" + # keeps process isolation and RPC timeouts. Multi-GPU still defaults + # to "mp". if distributed_executor_backend is None: - distributed_executor_backend = "mp" + num_gpus = od_config.num_gpus or 1 + distributed_executor_backend = "uni" if num_gpus == 1 else "mp" if isinstance(distributed_executor_backend, type): if not issubclass(distributed_executor_backend, DiffusionExecutor): @@ -40,6 +47,10 @@ class DiffusionExecutor(ABC): from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor executor_class = MultiprocDiffusionExecutor + elif distributed_executor_backend == "uni": + from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor + + executor_class = UniProcDiffusionExecutor elif distributed_executor_backend == "external_launcher": raise NotImplementedError("external_launcher backend is not yet supported.") elif isinstance(distributed_executor_backend, str): @@ -115,6 +126,50 @@ class DiffusionExecutor(ABC): """ return None + def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]: + """Collect rank-local native specs after every Worker loads its model.""" + + result = self.collective_rpc( + "get_kv_cache_specs", + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, list): + raise TypeError(f"get_kv_cache_specs returned {type(result).__name__}, expected list") + return result + + def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]: + """Profile and collect the KV memory budget on every Worker rank.""" + + result = self.collective_rpc( + "determine_available_kv_memory", + args=(profile_requests,), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, list) or not all(isinstance(value, int) for value in result): + raise TypeError("determine_available_kv_memory must return list[int]") + return result + + def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig], resolved_max_model_len: int) -> None: + """Send rank-local configs and the resolved model length to all Workers.""" + + # The default control-plane RPC mode executes on every rank and has + # rank 0 return the gathered rank statuses, so failures on nonzero + # ranks are not silently dropped. + self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs, resolved_max_model_len)) + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> None: + """Clear request rows on every Worker after Scheduler retirement.""" + + unique_request_ids = list(dict.fromkeys(request_ids)) + if not unique_request_ids: + return + self.collective_rpc( + "remove_diffusion_kv_requests", + args=(unique_request_ids,), + ) + @abstractmethod def shutdown(self) -> None: """Shutdown the executor and release resources.""" diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 7de46b5cf..ff9253962 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -18,11 +18,13 @@ import zmq from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError +from vllm.v1.executor.multiproc_executor import set_multiprocessing_worker_envs from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput from vllm_omni.diffusion.executor.abstract import DiffusionExecutor from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key +from vllm_omni.diffusion.utils.future_utils import try_set_exception, try_set_result from vllm_omni.diffusion.worker import WorkerProc if TYPE_CHECKING: @@ -294,6 +296,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): logger.info("Starting server...") num_gpus = cast(int, od_config.num_gpus) + # Without this, every worker inherits one Torch thread per core, so an + # N-GPU run oversubscribes the host by N x core_count. Honours a + # user-provided OMP_NUM_THREADS. + set_multiprocessing_worker_envs() mp.set_start_method("spawn", force=True) processes = [] @@ -459,16 +465,14 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): for new_req in new_reqs: validate_new_request_data_identity(new_req) - has_diffusion_kv_metadata = any(new_req.diffusion_kv_metadata is not None for new_req in new_reqs) - # DP multi-concurrency: when DLO+AllGather is active and multiple - # requests are scheduled, send ALL requests in one broadcast RPC. - # Each rank picks req[rank % len(reqs)] and computes independently. + # requests are scheduled, send every complete NewRequestData envelope + # in one broadcast RPC. Each rank picks one envelope, keeping its + # request and Diffusion KV metadata inseparable. # All ranks reply (unique_reply_rank=None) so we collect dp_size # responses and match by dp_rank. if ( len(new_reqs) > 1 - and not has_diffusion_kv_metadata and getattr(self.od_config, "enable_distributed_layerwise_offload", False) and getattr(self.od_config, "dlo_use_allgather", True) ): @@ -501,12 +505,11 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): f"empty prompt request IDs: {empty_prompt_ids}." ) - reqs_list = [nr.req for nr in new_reqs] try: results = self.collective_rpc( "execute_model", timeout=_DLO_DP_WAVE_TIMEOUT_S, - args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job), + args=(new_reqs, self.od_config, scheduler_output.kv_prefetch_job), unique_reply_rank=None, exec_all_ranks=True, ) @@ -636,8 +639,23 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): async_output_id=per_req_id, ) ) + # The worker's background D2H/SHM thread can publish OUTPUT_READY + # before this call returns. In that case the pump found no split + # map and cached the whole batch output under batch_id, so adopt it + # here instead of registering a map nothing will ever consume. with self._futures_lock: - self._batch_split_map[batch_id] = per_req_map + early = self._completed_outputs.pop(batch_id, None) + if early is None: + self._batch_split_map[batch_id] = per_req_map + if early is not None: + logger.debug("Batch %s output arrived before split map; splitting now", batch_id) + try: + batch_output = early.result() + error = None + except Exception as exc: + batch_output = None + error = str(exc) + self._deliver_batch_split(per_req_map, batch_output, error) return BatchRunnerOutput.from_list(runner_outputs) if not isinstance(result, BatchRunnerOutput): raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}") @@ -850,9 +868,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): fut = self._rpc_futures.pop(msg.rpc_id, None) if msg.rpc_id else None if fut is not None and not fut.done(): if msg.error: - fut.set_exception(RuntimeError(msg.error)) + try_set_exception(fut, RuntimeError(msg.error)) else: - fut.set_result(msg) + try_set_result(fut, msg) elif msg.kind == AsyncOutputKind.OUTPUT_READY: batch_id = msg.async_output_id with self._futures_lock: @@ -863,52 +881,76 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): unpack_diffusion_output_shm(msg.output) except Exception: logger.exception("SHM unpack failed for batch %s", batch_id) - batch_output = msg.output - for per_req_id, req_id in per_req_map.items(): - req_output = batch_output.get_request_output(req_id) - per_req_result: DiffusionOutput - if req_output is not None and req_output.result is not None: - per_req_result = req_output.result - elif msg.error: - per_req_result = DiffusionOutput(error=msg.error) - else: - per_req_result = DiffusionOutput(error="No output result for batch request") - fut: concurrent.futures.Future = concurrent.futures.Future() - fut.set_result(per_req_result) - with self._futures_lock: - pending = self._output_futures.pop(per_req_id, None) - if pending is not None and not pending.done(): - pending.set_result(per_req_result) - else: - self._completed_outputs[per_req_id] = fut + self._deliver_batch_split(per_req_map, msg.output, msg.error) else: - with self._futures_lock: - fut = self._output_futures.pop(batch_id, None) if batch_id else None - if fut is not None and not fut.done(): - if msg.error: - fut.set_exception(RuntimeError(msg.error)) - else: - try: - unpack_diffusion_output_shm(msg.output) - except Exception as e: - logger.exception("SHM unpack failed in result pump") - fut.set_exception(e) - continue - fut.set_result(msg.output) - elif batch_id: - fut = concurrent.futures.Future() - if msg.error: - fut.set_exception(RuntimeError(msg.error)) - else: - try: - unpack_diffusion_output_shm(msg.output) - except Exception as e: - logger.exception("SHM unpack failed in result pump (cached)") - fut.set_exception(e) - else: - fut.set_result(msg.output) + # Single-request result: unpack SHM first, then resolve or cache atomically. + output_result: DiffusionOutput | None = None + exc: Exception | None = None + if msg.error: + exc = RuntimeError(msg.error) + else: + try: + unpack_diffusion_output_shm(msg.output) + output_result = msg.output + except Exception as e: + logger.exception("SHM unpack failed in result pump") + exc = e + + if batch_id: with self._futures_lock: - self._completed_outputs[batch_id] = fut + pending = self._output_futures.pop(batch_id, None) + if pending is not None and not pending.done(): + if exc is not None: + try_set_exception(pending, exc) + else: + try_set_result(pending, output_result) + else: + fut = concurrent.futures.Future() + if exc is not None: + fut.set_exception(exc) + else: + fut.set_result(output_result) + self._completed_outputs[batch_id] = fut + + def _deliver_batch_split( + self, + per_req_map: dict[str, str], + batch_output: Any, + error: str | None = None, + ) -> None: + """Resolve per-request futures from one batch-level output.""" + for per_req_id, req_id in per_req_map.items(): + req_output = batch_output.get_request_output(req_id) if batch_output is not None else None + per_req_result: DiffusionOutput + if req_output is not None and req_output.result is not None: + per_req_result = req_output.result + elif error: + per_req_result = DiffusionOutput(error=error) + else: + per_req_result = DiffusionOutput(error="No output result for batch request") + with self._futures_lock: + pending = self._output_futures.pop(per_req_id, None) + if pending is not None and not pending.done(): + try_set_result(pending, per_req_result) + else: + fut: concurrent.futures.Future = concurrent.futures.Future() + fut.set_result(per_req_result) + self._completed_outputs[per_req_id] = fut + + def describe_pending_state(self, async_output_id: str | None = None) -> str: + """Summarize async-output bookkeeping for diagnosing stuck waits.""" + with self._futures_lock: + waiting = list(self._output_futures) + cached = list(self._completed_outputs) + batches = list(self._batch_split_map) + rpcs = list(self._rpc_futures) + return ( + f"async_output_id={async_output_id} " + f"waiting={len(waiting)}{waiting[:5]} " + f"cached={len(cached)}{cached[:5]} " + f"unsplit_batches={len(batches)}{batches[:5]} " + f"pending_rpcs={len(rpcs)}{rpcs[:5]}" + ) def _next_rpc_id(self) -> str: with self._rpc_id_lock: @@ -955,10 +997,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): with self._futures_lock: for fut in self._rpc_futures.values(): if not fut.done(): - fut.set_exception(RuntimeError("Executor shut down")) + try_set_exception(fut, RuntimeError("Executor shut down")) for fut in self._output_futures.values(): if not fut.done(): - fut.set_exception(RuntimeError("Executor shut down")) + try_set_exception(fut, RuntimeError("Executor shut down")) self._rpc_futures.clear() self._output_futures.clear() self._batch_split_map.clear() diff --git a/vllm_omni/diffusion/executor/uniproc_executor.py b/vllm_omni/diffusion/executor/uniproc_executor.py new file mode 100644 index 000000000..8eb098323 --- /dev/null +++ b/vllm_omni/diffusion/executor/uniproc_executor.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""In-process diffusion executor for single-GPU deployments. + +Counterpart to :class:`MultiprocDiffusionExecutor` for ``num_gpus == 1``. +The worker is constructed and driven in the engine process, so none of the +multiproc machinery is created: + +* no spawned worker subprocess (and no second model load / CUDA context) +* no ``MessageQueue`` shared-memory ring buffers or zmq ``ipc://`` sockets +* no POSIX ``/dev/shm`` segments for output tensors (``ipc.py`` pack/unpack) + +Mirrors vLLM's ``UniProcExecutor`` (``vllm/v1/executor/uniproc_executor.py``). + +``timeout`` on :meth:`collective_rpc` is accepted for interface parity but is +not enforced: the call runs on the calling thread, so a hung worker blocks it. +The multiproc executor enforces deadlines via its result-queue dequeue. +""" + +from __future__ import annotations + +import gc +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import torch +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.engine.exceptions import EngineDeadError + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.executor.abstract import DiffusionExecutor +from vllm_omni.platforms import current_omni_platform + +if TYPE_CHECKING: + from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput + from vllm_omni.diffusion.worker.utils import BaseRunnerOutput + +logger = init_logger(__name__) + + +class UniProcDiffusionExecutor(DiffusionExecutor): + """Runs a single diffusion worker inline, with no IPC of any kind.""" + + def _init_executor(self) -> None: + from vllm_omni.diffusion.worker.diffusion_worker import WorkerWrapperBase + + self._closed = False + self._is_failed = False + self.driver_worker = None + self._failure_callbacks: list[Callable[[], None]] = [] + self._warned_about_timeout = False + + num_gpus = self.od_config.num_gpus or 1 + if num_gpus != 1: + raise ValueError( + f"UniProcDiffusionExecutor supports a single GPU only, got num_gpus={num_gpus}. " + 'Use distributed_executor_backend="mp" for multi-GPU deployments.' + ) + + # The worker still initializes a one-rank distributed environment so + # that group lookups (`get_dp_group()` and friends) resolve. At world + # size 1, vLLM's GroupCoordinator skips the device communicator and + # the shared-memory MessageQueue, and every collective short-circuits. + worker_cls_path = current_omni_platform.get_diffusion_worker_cls() + self.driver_worker = WorkerWrapperBase( + gpu_id=0, + od_config=self.od_config, + worker_extension_cls=self.od_config.worker_extension_cls, + custom_pipeline_args=getattr(self.od_config, "custom_pipeline_args", None), + base_worker_class=resolve_obj_by_qualname(worker_cls_path), + ) + logger.info("Diffusion worker initialized in-process (uniproc executor)") + + @property + def is_dead(self) -> bool: + return self._closed or self._is_failed + + def _ensure_open(self) -> None: + if self._closed: + raise RuntimeError("DiffusionExecutor is closed.") + if self.driver_worker is None: + raise RuntimeError("Diffusion worker is not initialized.") + + def collective_rpc( + self, + method: str, + timeout: float | None = None, + args: tuple = (), + kwargs: dict | None = None, + unique_reply_rank: int | None = None, + exec_all_ranks: bool = False, + ) -> Any: + """Invoke ``method`` on the single in-process worker. + + ``unique_reply_rank`` / ``exec_all_ranks`` only select which ranks + participate, which is degenerate at world size 1. The return shape + still matches :class:`MultiprocDiffusionExecutor` (a bare result when a + reply rank is named, a one-element list otherwise). + """ + self._ensure_open() + if timeout is not None and not self._warned_about_timeout: + self._warned_about_timeout = True + logger.info( + "Diffusion RPC deadlines are not enforced by the in-process executor: " + "the worker runs on the calling thread, so a hung call blocks it." + ) + try: + result = self.driver_worker.execute_method(method, *(args or ()), **(kwargs or {})) + except Exception: + if not self._device_is_usable(): + self._mark_failed() + raise + return result if unique_reply_rank is not None else [result] + + def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + from vllm_omni.diffusion.sched.interface import validate_new_request_data_identity + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + self._ensure_open() + runner_outputs: list[RunnerOutput] = [] + for new_req in scheduler_output.scheduled_new_reqs: + validate_new_request_data_identity(new_req) + try: + args: tuple = (new_req.req, self.od_config, scheduler_output.kv_prefetch_job) + if new_req.diffusion_kv_metadata is not None: + args += (new_req.diffusion_kv_metadata,) + result = self.collective_rpc( + "execute_model", + args=args, + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, DiffusionOutput): + raise RuntimeError(f"Unexpected response type: {type(result)!r}") + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=result, + ) + ) + except Exception as exc: + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=DiffusionOutput(error=str(exc)), + ) + ) + return BatchRunnerOutput.from_list(runner_outputs) + + def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput + + self._ensure_open() + if len(scheduler_output.scheduled_new_reqs) <= 1: + return self.execute_request(scheduler_output) + + result = self.collective_rpc( + "execute_model_batch", + args=(scheduler_output, self.od_config), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, BatchRunnerOutput): + raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}") + return result + + def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + from vllm_omni.diffusion.worker.utils import BaseRunnerOutput + + self._ensure_open() + result = self.collective_rpc( + "execute_stepwise", + args=(scheduler_output,), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if isinstance(result, BaseRunnerOutput): + return result + raise RuntimeError(f"Unexpected response type for execute_step: {type(result)!r}") + + def _device_is_usable(self) -> bool: + """Whether the accelerator context survived the failure we just caught. + + Running the worker inline means a worker fault no longer kills a + process, so nothing else notices it. Recoverable faults (OOM, an + invalid request) leave the context intact and must not take the + deployment down. A sticky fault (illegal memory access, ECC error) + poisons the context: every later launch fails, so the engine needs + to be restarted. ``synchronize()`` re-raises a pending sticky error + on the registered accelerator (CUDA, ROCm, NPU). + """ + if not torch.accelerator.is_available(): + return True + accelerator = torch.accelerator.current_accelerator() + device_mod = getattr(torch, accelerator.type, None) if accelerator is not None else None + is_initialized = getattr(device_mod, "is_initialized", None) + if callable(is_initialized) and not is_initialized(): + return True + try: + torch.accelerator.synchronize() + except Exception as exc: + logger.error("Device context is unusable after a worker failure: %s", exc) + return False + return True + + def _mark_failed(self) -> None: + if self._is_failed: + return + self._is_failed = True + logger.error("Diffusion worker failed unrecoverably in-process; marking the executor dead.") + for callback in self._failure_callbacks: + try: + callback() + except Exception: + logger.exception("Diffusion executor failure callback raised") + + def register_failure_callback(self, callback: Callable[[], None]) -> None: + """Register a callback invoked when the inline worker fails fatally. + + The multiproc executor fires these from its process monitor. There is + no process to monitor here, so they are fired from ``collective_rpc`` + instead — without them a poisoned device context would leave + ``check_health`` reporting healthy while every request fails. + """ + self._failure_callbacks.append(callback) + + def check_health(self) -> None: + if self._is_failed: + raise EngineDeadError() + self._ensure_open() + + def shutdown(self) -> None: + if self._closed: + return + self._closed = True + worker = self.driver_worker + self.driver_worker = None + self._failure_callbacks.clear() + if worker is not None: + try: + worker.shutdown() + except Exception as exc: + logger.warning("Diffusion worker shutdown encountered an error: %s", exc) + # Unlike the multiprocess executor, tearing down an inline worker does + # not exit its process. Drop the final model reference and return cached + # allocations to the driver so a subsequent engine can use the device. + del worker + gc.collect() + try: + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + except Exception as exc: + logger.warning("Failed to release accelerator cache during diffusion worker shutdown: %s", exc) diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index bb589d7bb..8924e175e 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -1,8 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from contextlib import contextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch import vllm.ir @@ -26,9 +29,23 @@ class ForwardContext: vllm_config: VllmConfig | None = None omni_diffusion_config: OmniDiffusionConfig | None = None attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None + # Runner-owned paged execution metadata/runtime. Attention resolves the + # active Worker adapter from it; model code must not construct BlockTable + # rows or activate the runtime directly. + paged_kv_runtime: object | None = None + # Active Worker-side paged KV adapter. The adapter is installed only for + # the duration of a paged forward; dense forwards leave this as ``None``. + # Keep the field opaque here to avoid coupling the common context module to + # the diffusion_kv implementation. + paged_kv_adapter: Any | None = None + # Startup-only memory profiling runs before Scheduler-owned pages exist. + # Attention layers use this explicit scope to distinguish that probe from + # a malformed paged request whose Worker adapter was not activated. + in_diffusion_kv_memory_profile: bool = False split_text_embed_in_sp: bool = False denoise_step_idx: int | None = None denoise_timestep: float | None = None + total_denoise_steps: int | None = None # Per-request reference latent for img2img DiT models (e.g. Ming) ref_latent: torch.Tensor | None = None # whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel @@ -145,6 +162,8 @@ def create_forward_context( vllm_config: VllmConfig | None = None, omni_diffusion_config: OmniDiffusionConfig | None = None, attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None, + paged_kv_runtime: object | None = None, + in_diffusion_kv_memory_profile: bool = False, split_text_embed_in_sp: bool = False, denoise_step_idx: int | None = None, ): @@ -152,6 +171,8 @@ def create_forward_context( vllm_config=vllm_config, omni_diffusion_config=omni_diffusion_config, attn_metadata=attn_metadata, + paged_kv_runtime=paged_kv_runtime, + in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, split_text_embed_in_sp=split_text_embed_in_sp, denoise_step_idx=denoise_step_idx, ) @@ -177,6 +198,8 @@ def set_forward_context( vllm_config: VllmConfig | None = None, omni_diffusion_config: OmniDiffusionConfig | None = None, attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None, + paged_kv_runtime: object | None = None, + in_diffusion_kv_memory_profile: bool = False, split_text_embed_in_sp: bool = False, denoise_step_idx: int | None = None, ): @@ -188,6 +211,8 @@ def set_forward_context( vllm_config=vllm_config, omni_diffusion_config=omni_diffusion_config, attn_metadata=attn_metadata, + paged_kv_runtime=paged_kv_runtime, + in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, split_text_embed_in_sp=split_text_embed_in_sp, denoise_step_idx=denoise_step_idx, ) @@ -209,10 +234,60 @@ def set_forward_context( yield +@contextmanager +def override_paged_kv_adapter(adapter: Any | None): + """Temporarily expose a Worker paged-KV adapter to Omni Attention. + + This is deliberately a small context override instead of a second global + forward context. The model runner owns the outer context and the adapter + only replaces one opaque field while its prepared native metadata is live. + """ + + if _forward_context is None: + # Unit-level adapter users can prepare/activate metadata without an + # Omni model forward context. In that case there is nothing to + # override and the adapter's explicit ``forward`` API remains usable. + yield + return + + previous = _forward_context.paged_kv_adapter + _forward_context.paged_kv_adapter = adapter + try: + yield + finally: + _forward_context.paged_kv_adapter = previous + + def set_forward_context_denoise_step_idx(step_idx: int | None) -> None: """Set the current diffusion denoise step on the active ForwardContext.""" if _forward_context is not None: _forward_context.denoise_step_idx = step_idx + if step_idx is not None: + paged_kv_runtime = getattr(_forward_context, "paged_kv_runtime", None) + ensure_active = getattr(paged_kv_runtime, "ensure_active", None) + if callable(ensure_active): + ensure_active(step_idx) + + +def set_forward_context_denoise_timestep(timestep: float | None) -> None: + """Set the normalized (descending, 1 -> 0) denoise timestep. + + Timestep-gated attention features read this; pipelines that drive their own + denoise loop can publish it directly instead of going through + :meth:`DenoiseProgressMixin.record_denoise_step`. + """ + if _forward_context is not None: + _forward_context.denoise_timestep = None if timestep is None else float(timestep) + + +def set_forward_context_denoise_total_steps(total_steps: int | None) -> None: + """Set the total denoise step count on the active ForwardContext. + + Denoise loops publish it so tail-fallback gates (e.g. ``end_step`` in + RAINFUSION_ATTN) know when the final denoise steps begin. + """ + if _forward_context is not None: + _forward_context.total_denoise_steps = total_steps class DenoiseProgressMixin: @@ -222,8 +297,11 @@ class DenoiseProgressMixin: timestep=None, scheduler=None, normalized_timestep: float | None = None, + total_steps: int | None = None, ) -> None: set_forward_context_denoise_step_idx(step_idx) + if _forward_context is not None: + _forward_context.total_denoise_steps = total_steps if _forward_context is None: return if normalized_timestep is not None: diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 5082796a2..3d3bc4683 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime. Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process @@ -23,6 +26,7 @@ from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt +from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID, diffusion_exception_metrics from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -44,7 +48,6 @@ class InlineStageDiffusionClient(StageClientBase): model: str, od_config: OmniDiffusionConfig, metadata: StageMetadata, - batch_size: int = 1, ) -> None: self.model = model self.od_config = od_config @@ -56,8 +59,9 @@ class InlineStageDiffusionClient(StageClientBase): self.default_sampling_params = metadata.default_sampling_params self.requires_multimodal_data = metadata.requires_multimodal_data self.custom_process_input_func = metadata.custom_process_input_func + self.prompt_transform_func = None + self.prompt_expand_func = None self.engine_input_source = metadata.engine_input_source - self.batch_size = batch_size self._enrich_config() self._engine = DiffusionEngine.make_engine(self.od_config) @@ -73,10 +77,10 @@ class InlineStageDiffusionClient(StageClientBase): self._engine.executor.register_failure_callback(self._mark_engine_dead) logger.info( - "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)", + "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (max_num_seqs=%d)", self.stage_id, self.replica_id, - self.batch_size, + self.od_config.max_num_seqs, ) def _enrich_config(self) -> None: @@ -158,6 +162,15 @@ class InlineStageDiffusionClient(StageClientBase): self._output_queue.put_nowait(result) except DiffusionRequestAbortedError as e: logger.info("request_id: %s aborted: %s", request_id, str(e)) + metrics = diffusion_exception_metrics(e) + if metrics: + self._output_queue.put_nowait( + OmniRequestOutput( + request_id=DIFFUSION_METRICS_ONLY_REQUEST_ID, + finished=True, + metrics=metrics, + ) + ) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) status_code, error_type = client_error_metadata(e) @@ -167,6 +180,7 @@ class InlineStageDiffusionClient(StageClientBase): status_code=status_code, error_type=error_type, ) + error_output.metrics.update(diffusion_exception_metrics(e)) self._output_queue.put_nowait(error_output) finally: self._tasks.pop(request_id, None) @@ -181,9 +195,8 @@ class InlineStageDiffusionClient(StageClientBase): async def abort_requests_async(self, request_ids: list[str]) -> None: for rid in request_ids: - task = self._tasks.pop(rid, None) - if task: - task.cancel() + # Keep the consumer alive until the engine emits its terminal + # abort output. That path carries the refreshed scheduler gauge. self._engine.abort(rid) async def submit_interaction_async( diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index 8d1405a6c..f4cf5bb36 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/vllm_omni/diffusion/io_support.py @@ -6,11 +6,16 @@ from __future__ import annotations import inspect from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import resolve_diffusers_pipeline_class from vllm_omni.diffusion.registry import DiffusionModelRegistry def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]: - if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None: + if ( + od_config.diffusion_load_format == "diffusers" + and (pipe_cls := resolve_diffusers_pipeline_class(od_config)) is not None + ): signature = inspect.signature(pipe_cls.__call__) support_image_input = "image" in signature.parameters support_audio_input = ( @@ -40,6 +45,16 @@ def supports_audio_output(model_class_name: str) -> bool: return bool(getattr(model_cls, "support_audio_output", False)) +def get_diffusion_output_type(model_class_name: str | None) -> str: + """Return the declared final modality for a diffusion model.""" + declared_output_type = get_diffusion_model_metadata(model_class_name).final_output_type + if declared_output_type is not None: + return declared_output_type + if model_class_name and supports_audio_output(model_class_name): + return "audio" + return "image" + + def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int: """Get num_frames for the dummy warmup run. Returns 0 to skip warmup.""" @@ -47,3 +62,9 @@ def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"): return int(getattr(model_cls, "dummy_run_num_frames")) return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1 + + +def get_dummy_run_num_image_inputs(model_class_name: str) -> int: + """Return the maximum advertised image-input count for profiling.""" + + return get_diffusion_model_metadata(model_class_name).max_multimodal_image_inputs or 1 diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 061f89838..6eb24c540 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -121,8 +121,19 @@ def _pack_tensor_if_large( val: torch.Tensor, d2h_stream: torch.Stream | None = None, ) -> torch.Tensor | dict: - """Replace a tensor with an SHM handle if it exceeds the threshold.""" - if val.nelement() * val.element_size() > _SHM_TENSOR_THRESHOLD: + """Replace a tensor with an SHM handle if it exceeds the threshold. + + Batch outputs are split into per-request views that share one storage. + Pickle serializes a view's whole storage, so a batch of small views costs + the full batch tensor once per request. Size the decision on the storage + to keep those views off the wire; packing copies just the view. + """ + view_bytes = val.nelement() * val.element_size() + try: + storage_bytes = val.untyped_storage().nbytes() + except Exception: + storage_bytes = view_bytes + if max(view_bytes, storage_bytes) > _SHM_TENSOR_THRESHOLD: return _tensor_to_shm(val, d2h_stream=d2h_stream) return val @@ -147,7 +158,7 @@ def _pack_value_if_large( Walks the container shapes pipelines return as ``DiffusionOutput.output``: bare tensors, dicts (e.g. Cosmos3 ``{"image"/"video": ...}``), and - tuples/lists (e.g. LTX2 and DreamID ``(video, audio)``). Other values pass + tuples/lists (e.g. LTX2 ``(video, audio)``). Other values pass through unchanged. ``_unpack_if_shm_handle`` must mirror these shapes — keep the two in sync. """ diff --git a/vllm_omni/diffusion/layers/activation.py b/vllm_omni/diffusion/layers/activation.py new file mode 100644 index 000000000..97a5a60fc --- /dev/null +++ b/vllm_omni/diffusion/layers/activation.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Custom activation functions for diffusion models.""" + +import torch +import torch.nn.functional as F + +from vllm_omni.diffusion.layers.custom_op import CustomOp +from vllm_omni.platforms import current_omni_platform + + +class SiluAndMul(CustomOp): + """SwiGLU activation dispatched by the current vLLM-Omni platform.""" + + def __init__(self) -> None: + super().__init__() + if current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa(): + self.op = torch.ops._C.silu_and_mul + + @staticmethod + def forward_native(x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + if x.device.type == "cpu": + return self.forward_native(x) + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + self.op(out, x) + return out + + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: + import torch_npu + + return torch_npu.npu_swiglu(x, dim=-1) + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py new file mode 100644 index 000000000..f7e27675c --- /dev/null +++ b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Q/K RMSNorm followed by packed non-interleaved RoPE. + +The public contract is shared by diffusion attention implementations: + +* ``q`` and ``k`` are ``[tokens, heads, head_dim]``; +* norm weights are one-dimensional ``[head_dim]`` tensors; +* ``rope_table`` is ``[tokens, rotary_dim]`` and stores + ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. + +The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K +or rotary-product intermediates. Ascend composes its RMSNorm and rotary fused +primitives; unsupported inputs use the eager reference. +""" + +from __future__ import annotations + +from importlib.util import find_spec + +import torch +import torch.nn.functional as F +from torch.library import Library +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + +from vllm_omni.platforms import current_omni_platform + +_FUSED_HEAD_DIM = 128 +_FUSED_ROTARY_DIM = 96 +_HEADS_PER_PROGRAM = 8 + + +def _apply_rope_table( + x: torch.Tensor, + rope_table: torch.Tensor, + rotary_dim: int, +) -> torch.Tensor: + half = rotary_dim // 2 + cos = rope_table[..., :half].to(x.dtype).unsqueeze(1) + sin = rope_table[..., half:].to(x.dtype).unsqueeze(1) + first = x[..., :half] + second = x[..., half:rotary_dim] + return torch.cat( + ( + first * cos - second * sin, + second * cos + first * sin, + x[..., rotary_dim:], + ), + dim=-1, + ) + + +if HAS_TRITON: + + @triton.jit + def _rms_norm_rope_kernel( + x_ptr, + weight_ptr, + rope_table_ptr, + out_ptr, + x_stride_t, + x_stride_h, + x_stride_d, + rope_stride_t, + out_stride_t, + out_stride_h, + out_stride_d, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + rotary_half: tl.constexpr, + eps: tl.constexpr, + heads_per_program: tl.constexpr, + ): + token = tl.program_id(0) + head_group = tl.program_id(1) + heads = head_group * heads_per_program + tl.arange(0, heads_per_program) + dims = tl.arange(0, head_dim) + mask = heads[:, None] < num_heads + offsets = token * x_stride_t + heads[:, None] * x_stride_h + dims[None, :] * x_stride_d + + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + weight = tl.load(weight_ptr + dims).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) + normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16) + + rotary_dim = rotary_half * 2 + pair_dims = tl.where( + dims < rotary_half, + dims + rotary_half, + tl.where(dims < rotary_dim, dims - rotary_half, dims), + ) + pair_offsets = token * x_stride_t + heads[:, None] * x_stride_h + pair_dims[None, :] * x_stride_d + pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + pair_weight = tl.load(weight_ptr + pair_dims).to(tl.float32) + pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16) + + freq_dims = tl.where( + dims < rotary_half, + dims, + tl.where(dims < rotary_dim, dims - rotary_half, 0), + ) + table_offsets = token * rope_stride_t + freq_dims + cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32) + sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32) + first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin + second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin + output = tl.where( + dims < rotary_dim, + tl.where(dims < rotary_half, first, second), + normalized.to(tl.float32), + ) + + out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d + tl.store(out_ptr + out_offsets, output, mask=mask) + + +def _eager_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + q_norm = F.rms_norm(q, (head_dim,), q_weight, eps) + k_norm = F.rms_norm(k, (head_dim,), k_weight, eps) + return ( + _apply_rope_table(q_norm, rope_table, rotary_dim), + _apply_rope_table(k_norm, rope_table, rotary_dim), + ) + + +def _npu_apply_rope_table( + x: torch.Tensor, + rope_table: torch.Tensor, + rotary_dim: int, +) -> torch.Tensor: + """Apply H3's rotary dimensions through Ascend's fused rotary kernel. + + ``mindiesd.rotary_position_embedding`` takes a 4-D BSND tensor, whereas + MiniMax-H3 keeps packed activations as ``[tokens, heads, head_dim]``. + The shared MindIE-SD wrapper normalizes this 3-D layout to BSND and + restores it afterwards. It receives the half-width cos/sin values from + the packed table; the wrapper expands them to H3's non-interleaved 96-D + rotary layout. The 32 non-rotary head dimensions bypass the kernel. + + Some CANN environments package ``torch_npu`` without MindIE-SD. Keep + those deployments on the same Ascend fused rotary primitive rather than + silently falling back to eager elementwise RoPE. + """ + import torch_npu + + half = rotary_dim // 2 + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + cos = rope_table[..., :half] + sin = rope_table[..., half:] + + if find_spec("mindiesd") is not None: + from vllm_omni.diffusion.layers.rope import apply_rotary_emb_mindiesd + + x_rot = apply_rotary_emb_mindiesd( + x_rot, + cos, + sin, + interleaved=False, + half_head_dim=True, + ) + else: + # npu_rotary_mul uses BSND and full rotary-width cos/sin. H3 uses + # NeoX/rotated-half ordering, so duplicate rather than interleave the + # half-width table along the final dimension. + cos = cos.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) + sin = sin.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) + x_rot = torch_npu.npu_rotary_mul(x_rot.unsqueeze(0), cos, sin).squeeze(0) + + return torch.cat((x_rot, x_pass), dim=-1) + + +def _npu_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Use Ascend RMSNorm and RoPE fused primitives for MiniMax-H3 DiT.""" + import torch_npu + + q_norm = torch_npu.npu_rms_norm(q, q_weight, epsilon=eps)[0] + k_norm = torch_npu.npu_rms_norm(k, k_weight, epsilon=eps)[0] + return ( + _npu_apply_rope_table(q_norm, rope_table, rotary_dim), + _npu_apply_rope_table(k_norm, rope_table, rotary_dim), + ) + + +def _fused_cuda_supported( + q: torch.Tensor, + k: torch.Tensor, + head_dim: int, + rotary_dim: int, +) -> bool: + return ( + HAS_TRITON + and current_platform.is_cuda() + and q.is_cuda + and k.is_cuda + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and head_dim == _FUSED_HEAD_DIM + and rotary_dim == _FUSED_ROTARY_DIM + ) + + +def _fused_npu_supported( + q: torch.Tensor, + k: torch.Tensor, + head_dim: int, + rotary_dim: int, +) -> bool: + """Return whether the MiniMax-H3 Ascend fused-op contract is satisfied.""" + return ( + current_omni_platform.is_npu() + and q.device.type == "npu" + and k.device.type == "npu" + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and head_dim == _FUSED_HEAD_DIM + and rotary_dim == _FUSED_ROTARY_DIM + ) + + +def _launch_fused_rms_norm_rope( + x: torch.Tensor, + weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, +) -> torch.Tensor: + tokens, heads, head_dim = x.shape + rotary_half = rope_table.shape[-1] // 2 + out = torch.empty(x.shape, dtype=x.dtype, device=x.device) + if tokens == 0: + return out + grid = (tokens, triton.cdiv(heads, _HEADS_PER_PROGRAM)) + _rms_norm_rope_kernel[grid]( + x, + weight, + rope_table, + out, + x.stride(0), + x.stride(1), + x.stride(2), + rope_table.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + num_heads=heads, + head_dim=head_dim, + rotary_half=rotary_half, + eps=eps, + heads_per_program=_HEADS_PER_PROGRAM, + num_warps=8, + ) + return out + + +def _fused_qk_norm_rope_impl( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if _fused_npu_supported(q, k, head_dim, rotary_dim): + return _npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + rotary_dim, + ) + if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + return _eager_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + head_dim, + rotary_dim, + ) + return ( + _launch_fused_rms_norm_rope(q, q_weight, rope_table, eps), + _launch_fused_rms_norm_rope(k, k_weight, rope_table, eps), + ) + + +def _fused_qk_norm_rope_fake( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + head_dim: int, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim + return torch.empty_like(q), torch.empty_like(k) + + +_OMNI_OP_LIB = Library("vllm_omni", "FRAGMENT") +if not hasattr(torch.ops.vllm_omni, "fused_qk_norm_rope"): + direct_register_custom_op( + op_name="fused_qk_norm_rope", + op_func=_fused_qk_norm_rope_impl, + fake_impl=_fused_qk_norm_rope_fake, + mutates_args=[], + target_lib=_OMNI_OP_LIB, + ) + + +def fused_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + *, + head_dim: int | None = None, + rotary_dim: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply Q/K RMSNorm and packed non-interleaved RoPE.""" + if q.ndim != 3 or k.ndim != 3: + raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}") + if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]: + raise ValueError(f"q and k shapes are incompatible: {q.shape} and {k.shape}") + if q.dtype != k.dtype or q.device != k.device: + raise ValueError("q and k must have the same dtype and device") + if q.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise TypeError(f"Fused QK RMSNorm/RoPE requires floating inputs, got {q.dtype}") + + head_dim = q.shape[-1] if head_dim is None else head_dim + rotary_dim = rope_table.shape[-1] if rotary_dim is None else rotary_dim + if q.shape[-1] != head_dim: + raise ValueError(f"Expected q/k head_dim={head_dim}, got {q.shape[-1]}") + if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2: + raise ValueError(f"rotary_dim must be even and in [2, {head_dim}], got {rotary_dim}") + if q_weight.shape != (head_dim,) or k_weight.shape != (head_dim,): + raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}") + if q_weight.device != q.device or k_weight.device != q.device: + raise ValueError("Q/K norm weights must be on the activation device") + if rope_table.device != q.device or rope_table.dtype != q.dtype: + raise ValueError("rope_table must have the same dtype and device as q/k") + if rope_table.shape != (q.shape[0], rotary_dim): + raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}") + + q_weight = q_weight.contiguous() + k_weight = k_weight.contiguous() + rope_table = rope_table.contiguous() + if _fused_npu_supported(q, k, head_dim, rotary_dim): + return _npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + rotary_dim, + ) + if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + return _fused_qk_norm_rope_impl( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + head_dim, + rotary_dim, + ) + return torch.ops.vllm_omni.fused_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + head_dim, + rotary_dim, + ) + + +__all__ = ["fused_qk_norm_rope"] diff --git a/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py index ce83d8842..235fc3342 100644 --- a/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py +++ b/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import torch @@ -91,6 +94,21 @@ class MoTQKVParallelLinear(QKVParallelLinear): # boilerplate Module subclass just for the VAE pathway. self.gen_exp = torch.nn.Module() + # ``gen_exp`` participates in the same tensor-parallel layout as this + # layer. vLLM's online quantizers inspect these attributes while + # processing weights (for example, FP8 uses them to share scales across + # TP ranks), so the lightweight container must expose the parallel + # linear metadata rather than only holding parameters. + for attr in ( + "tp_rank", + "tp_size", + "input_size", + "input_size_per_partition", + "output_size", + "output_size_per_partition", + ): + setattr(self.gen_exp, attr, getattr(self, attr)) + # Use the same weight_loader as text vae_weight_loader = ( self.weight_loader_v2 diff --git a/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py index 2811b110f..6e846af47 100644 --- a/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py +++ b/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import torch @@ -84,6 +87,21 @@ class MoTRowParallelLinear(RowParallelLinear): # boilerplate Module subclass just for the VAE pathway. self.gen_exp = torch.nn.Module() + # ``gen_exp`` participates in the same tensor-parallel layout as this + # layer. vLLM's online quantizers inspect these attributes while + # processing weights (for example, FP8 uses them to share scales across + # TP ranks), so the lightweight container must expose the parallel + # linear metadata rather than only holding parameters. + for attr in ( + "tp_rank", + "tp_size", + "input_size", + "input_size_per_partition", + "output_size", + "output_size_per_partition", + ): + setattr(self.gen_exp, attr, getattr(self, attr)) + # Select weight_loader consistent with text vae_weight_loader = ( self.weight_loader_v2 diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 8d91df8f0..3618388e6 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -85,7 +85,10 @@ class RMSNorm(CustomOp): def forward_cuda( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + return self.forward_native(x, residual) # During torch.compile tracing, fused_rms_norm writes to `out` in-place # (returns None) and accesses self.weight.data, which is a DTensor under # HSDP. Both patterns confuse inductor's compute_ancestors scheduler. @@ -101,7 +104,10 @@ class RMSNorm(CustomOp): def forward_hip( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + return self.forward_native(x, residual) if torch.compiler.is_compiling(): return self.forward_native(x) try: @@ -112,7 +118,10 @@ class RMSNorm(CustomOp): def forward_musa( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + return self.forward_native(x, residual) # Preserve the aten::rms_norm graph so dynamic Inductor can fuse the # H3 Q/K norm with its inline RoPE path on MUSA. return F.rms_norm( @@ -125,9 +134,19 @@ class RMSNorm(CustomOp): def forward_npu( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: import torch_npu + if residual is not None: + output, _, updated_residual = torch_npu.npu_add_rms_norm( + x, + residual, + self.weight, + self.variance_epsilon, + ) + return output, updated_residual + output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0] return output @@ -135,19 +154,27 @@ class RMSNorm(CustomOp): def forward_xpu( self, x: torch.Tensor, - ) -> torch.Tensor: - return self.forward_native(x) + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is None: + return self.forward_native(x) + return self.forward_native(x, residual) def forward_native( self, x: torch.Tensor, - ) -> torch.Tensor: + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is not None: + residual = residual + x + x = residual input_dtype = x.dtype x = x.to(torch.float32) variance = x.pow(2).mean(-1, keepdim=True) out = x * torch.rsqrt(variance + self.variance_epsilon) out = self.weight.to(torch.float32) * out - return out.to(input_dtype) + out = out.to(input_dtype) + return (out, residual) if residual is not None else out class RMSNormVAE(CustomOp): diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index fae1a3414..e78574b04 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -423,16 +423,16 @@ class WanS2VRotaryPosEmbed(torch.nn.Module): h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long() w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long() - freqs_0 = freqs_split[0][f_sam] if f_o >= 0 else freqs_split[0][f_sam].conj() + freqs_0 = torch.index_select(freqs_split[0] if f_o >= 0 else freqs_split[0].conj(), 0, f_sam) freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1) freqs_i = torch.cat( [ freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs_split[1][h_sam] + torch.index_select(freqs_split[1], 0, h_sam) .view(1, seq_h_int, 1, -1) .expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs_split[2][w_sam] + torch.index_select(freqs_split[2], 0, w_sam) .view(1, 1, seq_w_int, -1) .expand(seq_f_int, seq_h_int, seq_w_int, -1), ], @@ -523,14 +523,18 @@ class RotaryEmbeddingS2VGrid(torch.nn.Module): h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long() w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long() - freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj() + freqs_0 = torch.index_select(freqs[0] if f_o >= 0 else freqs[0].conj(), 0, f_sam) freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1) freqs_i = torch.cat( [ freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs[1][h_sam].view(1, seq_h_int, 1, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1), - freqs[2][w_sam].view(1, 1, seq_w_int, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1), + torch.index_select(freqs[1], 0, h_sam) + .view(1, seq_h_int, 1, -1) + .expand(seq_f_int, seq_h_int, seq_w_int, -1), + torch.index_select(freqs[2], 0, w_sam) + .view(1, 1, seq_w_int, -1) + .expand(seq_f_int, seq_h_int, seq_w_int, -1), ], dim=-1, ).reshape(seg_len, 1, -1) diff --git a/vllm_omni/diffusion/lora/layers/base_linear.py b/vllm_omni/diffusion/lora/layers/base_linear.py index fe32868d0..fae385c36 100644 --- a/vllm_omni/diffusion/lora/layers/base_linear.py +++ b/vllm_omni/diffusion/lora/layers/base_linear.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -22,6 +22,19 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): is inherited from vLLM's BaseLinearLayerWithLoRA. """ + lora_a_stacked: tuple[torch.Tensor, ...] + lora_b_stacked: tuple[torch.Tensor, ...] + + def _move_lora_buffers(self, device: torch.device) -> None: + self.lora_a_stacked = tuple(tensor.to(device=device) for tensor in self.lora_a_stacked) + self.lora_b_stacked = tuple(tensor.to(device=device) for tensor in self.lora_b_stacked) + + def _set_diffusion_lora_buffer_device(self, device: torch.device) -> None: + """Keep dynamic LoRA sidecars resident on the compute device.""" + + self._diffusion_lora_buffer_device = device + self._move_lora_buffers(device) + def create_lora_weights( self, max_loras: int, @@ -29,6 +42,9 @@ class DiffusionBaseLinearLayerWithLoRA(BaseLinearLayerWithLoRA): model_config=None, ) -> None: super().create_lora_weights(max_loras, lora_config, model_config) + buffer_device = getattr(self, "_diffusion_lora_buffer_device", None) + if buffer_device is not None: + self._move_lora_buffers(buffer_device) # Keep a direct reference for attribute forwarding: `base_layer` is a # registered submodule (stored under `_modules`), so direct access via # `object.__getattribute__` will not find it. We stash a ref in diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index d11b1b9a2..1b63692fc 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math import time @@ -9,7 +9,7 @@ from typing import get_args import torch import torch.nn as nn -from vllm.config.lora import MaxLoRARanks +from vllm.config.lora import LoRAConfig, MaxLoRARanks from vllm.logger import init_logger from vllm.lora.layers import BaseLayerWithLoRA from vllm.lora.lora_model import LoRAModel @@ -23,7 +23,6 @@ from vllm.lora.utils import ( ) from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear -from vllm_omni.config.lora import LoRAConfig from vllm_omni.diffusion.lora.utils import ( _expand_expected_modules_for_packed_layers, _match_target_modules, @@ -69,6 +68,12 @@ class DiffusionLoRAManager: self.pipeline = pipeline self.device = device self.dtype = dtype + od_config = getattr(pipeline, "od_config", None) + # DLO owns the base-weight lifecycle. Keep request-switchable LoRA + # sidecars resident instead of rebuilding DLO host shards per request. + self._resident_lora_device = ( + device if getattr(od_config, "enable_distributed_layerwise_offload", False) is True else None + ) # Cache supported/expected module suffixes once, before any layer # replacement happens. After LoRA layers are injected, the original @@ -290,11 +295,35 @@ class DiffusionLoRAManager: lora_path = get_adapter_absolute_path(lora_request.lora_path) logger.debug("Resolved LoRA path: %s", lora_path) - peft_helper = PEFTHelper.from_local_dir( - lora_path, - max_position_embeddings=None, # no need in diffusion - tensorizer_config_dict=lora_request.tensorizer_config_dict, - ) + model_loader = getattr(self.pipeline, "_load_diffusion_lora_adapter", None) + loaded = None + if callable(model_loader): + loaded = model_loader( + lora_request=lora_request, + lora_path=lora_path, + dtype=self.dtype, + ) + + if loaded is None: + peft_helper = PEFTHelper.from_local_dir( + lora_path, + max_position_embeddings=None, # no need in diffusion + tensorizer_config_dict=lora_request.tensorizer_config_dict, + ) + + lora_model = LoRAModel.from_local_checkpoint( + lora_path, + expected_lora_modules=self._expected_lora_modules, + peft_helper=peft_helper, + lora_model_id=lora_request.lora_int_id, + device="cpu", # consistent w/ vllm's behavior + dtype=self.dtype, + model_vocab_size=None, + tensorizer_config_dict=lora_request.tensorizer_config_dict, + weights_mapper=None, + ) + else: + lora_model, peft_helper = loaded logger.info( "Loaded PEFT config: r=%d, lora_alpha=%d, target_modules=%s", @@ -303,18 +332,6 @@ class DiffusionLoRAManager: peft_helper.target_modules, ) - lora_model = LoRAModel.from_local_checkpoint( - lora_path, - expected_lora_modules=self._expected_lora_modules, - peft_helper=peft_helper, - lora_model_id=lora_request.lora_int_id, - device="cpu", # consistent w/ vllm's behavior - dtype=self.dtype, - model_vocab_size=None, - tensorizer_config_dict=lora_request.tensorizer_config_dict, - weights_mapper=None, - ) - logger.info( "Loaded LoRA model: id=%d, num_modules=%d, modules=%s", lora_model.id, @@ -439,6 +456,13 @@ class DiffusionLoRAManager: ) if lora_layer is not module and isinstance(lora_layer, BaseLayerWithLoRA): + if self._resident_lora_device is not None: + set_buffer_device = getattr(lora_layer, "_set_diffusion_lora_buffer_device", None) + if not callable(set_buffer_device): + raise RuntimeError( + f"{type(lora_layer).__name__} cannot keep dynamic LoRA buffers resident for DLO" + ) + set_buffer_device(self._resident_lora_device) replace_submodule(component, module_name, lora_layer) self._lora_modules[full_module_name] = lora_layer logger.debug("Replaced layer: %s -> %s", full_module_name, type(lora_layer).__name__) @@ -523,13 +547,17 @@ class DiffusionLoRAManager: matches_scale = self._adapter_scales.get(adapter_id) == rounded_scale return is_active and matches_scale - def _activate_adapter(self, adapter_id: int, scale: float) -> None: - if self._is_active_at_scale(adapter_id, scale): - logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale) - return + def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None: + binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None) + lora_names_by_id = ( + {id(weights): name for name, weights in lora_model.loras.items()} if callable(binding_validator) else {} + ) + bound_lora_names: set[str] = set() - logger.info("Activating adapter: id=%d", adapter_id) - lora_model = self._registered_adapters[adapter_id] + def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None: + name = lora_names_by_id.get(id(weights)) + if name is not None: + bound_lora_names.add(name) # activate weights in each LoRA layer for full_module_name, lora_layer in self._lora_modules.items(): @@ -571,6 +599,9 @@ class DiffusionLoRAManager: lora_b_list.append(sub_lora.lora_b * scale) lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list) + for sub_lora in sub_loras: + if sub_lora is not None: + _record_bound(sub_lora) logger.debug( "Activated packed LoRA for %s via submodules=%s (scale=%.2f)", full_module_name, @@ -589,6 +620,7 @@ class DiffusionLoRAManager: for b in lora_weights.lora_b ] lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list) + _record_bound(lora_weights) logger.debug( "Activated packed LoRA for %s (scale=%.2f)", full_module_name, @@ -619,6 +651,7 @@ class DiffusionLoRAManager: lora_a_list = [lora_weights.lora_a] * n_slices lora_b_list = [b * scale for b in b_splits] lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list) + _record_bound(lora_weights) logger.debug( "Activated fused LoRA for packed layer %s (scale=%.2f)", full_module_name, @@ -628,6 +661,7 @@ class DiffusionLoRAManager: scaled_lora_b = lora_weights.lora_b * scale lora_layer.set_lora(index=0, lora_a=lora_weights.lora_a, lora_b=scaled_lora_b) + _record_bound(lora_weights) logger.debug( "Activated LoRA for %s: lora_a shape=%s, lora_b shape=%s, scale=%.2f", full_module_name, @@ -636,6 +670,33 @@ class DiffusionLoRAManager: scale, ) + if callable(binding_validator): + binding_validator( + lora_model=lora_model, + bound_lora_names=frozenset(bound_lora_names), + ) + + def _reset_lora_layers(self) -> None: + for lora_layer in self._lora_modules.values(): + lora_layer.reset_lora(0) + + def _activate_adapter(self, adapter_id: int, scale: float) -> None: + if self._is_active_at_scale(adapter_id, scale): + logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale) + return + + logger.info("Activating adapter: id=%d", adapter_id) + lora_model = self._registered_adapters[adapter_id] + # Binding overwrites slot 0 incrementally. Invalidate the fast-path + # state before the first mutation and leave every wrapper inactive if + # any set_lora() call or model validator fails. + self._active_adapter_id = None + try: + self._bind_adapter_weights(lora_model, scale) + except Exception: + self._reset_lora_layers() + raise + self._active_adapter_id = adapter_id self._update_adapter_scale(adapter_id, scale) @@ -644,8 +705,7 @@ class DiffusionLoRAManager: logger.debug("All adapters already inactive") return logger.info("Deactivating all adapters: %d layers", len(self._lora_modules)) - for lora_layer in self._lora_modules.values(): - lora_layer.reset_lora(0) + self._reset_lora_layers() self._active_adapter_id = None logger.debug("All adapters deactivated") diff --git a/vllm_omni/diffusion/lora/utils.py b/vllm_omni/diffusion/lora/utils.py index 465edc99c..a8fdc2776 100644 --- a/vllm_omni/diffusion/lora/utils.py +++ b/vllm_omni/diffusion/lora/utils.py @@ -5,8 +5,8 @@ from __future__ import annotations import torch.nn as nn from transformers import PretrainedConfig +from vllm.config.lora import LoRAConfig -from vllm_omni.config.lora import LoRAConfig from vllm_omni.diffusion.lora.layers import ( DiffusionColumnParallelLinearWithLoRA, DiffusionMergedColumnParallelLinearWithLoRA, diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py index 584562193..804699bc7 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py @@ -9,6 +9,13 @@ from .modelopt import ( ) +def get_direct_mmap_adapter(model: nn.Module): + """Resolve direct-mmap adapters lazily to avoid model-loader cycles.""" + from .direct_mmap import get_direct_mmap_adapter as _get_direct_mmap_adapter + + return _get_direct_mmap_adapter(model) + + def get_checkpoint_adapter( model: nn.Module, source: object, @@ -29,4 +36,5 @@ __all__ = [ "ModelOptMixedPrecisionCheckpointAdapter", "ModelOptNvFp4CheckpointAdapter", "get_checkpoint_adapter", + "get_direct_mmap_adapter", ] diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py new file mode 100644 index 000000000..86b23429b --- /dev/null +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Loader-side adapters for proven direct-checkpoint mmap layouts.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from functools import partial +from typing import Protocol + +import torch +from torch import nn + + +@dataclass(frozen=True) +class DirectMmapTensorPolicy: + """Loader-owned proof for one tensor's direct mmap behavior.""" + + allow_custom_loader: bool = False + transform: Callable[[torch.Tensor], torch.Tensor] | None = None + + +class DirectMmapAdapter(Protocol): + def policy_for( + self, + runtime_name: str, + target: torch.Tensor, + ) -> DirectMmapTensorPolicy | None: ... + + +class _MiniMaxH3DirectMmapAdapter: + """TP1 runtime-layout contract already enforced by MiniMax-H3's loader.""" + + def __init__(self, pipeline: nn.Module) -> None: + self.pipeline = pipeline + + def policy_for( + self, + runtime_name: str, + target: torch.Tensor, + ) -> DirectMmapTensorPolicy: + del target + transform = None + suffix = ".qkv_proj.weight" + if runtime_name.endswith(suffix): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + _reorder_grouped_qkv_to_qkv, + ) + + attention_path = runtime_name[: -len(suffix)] + attention = self.pipeline.get_submodule(attention_path) + if not isinstance(attention, MiniMaxH3Attention): + raise ValueError( + f"MiniMax-H3 direct-mmap adapter expected attention at {attention_path!r}, " + f"got {type(attention).__name__}" + ) + transform = partial( + _reorder_grouped_qkv_to_qkv, + num_query_groups=attention.total_num_heads, + heads_per_group=1, + head_dim=attention.head_dim, + ) + return DirectMmapTensorPolicy( + allow_custom_loader=True, + transform=transform, + ) + + +class _Cosmos3DirectMmapAdapter: + """Preserve Cosmos3's existing TP1 direct-mmap load contract.""" + + def policy_for( + self, + runtime_name: str, + target: torch.Tensor, + ) -> DirectMmapTensorPolicy: + del runtime_name, target + # Cosmos3 has no packed module mapping or TP1 semantic transform after + # its checkpoint-key remap. Its custom vLLM loaders reduce to exact + # copies at TP1; topology and metadata are checked by plan preflight. + return DirectMmapTensorPolicy(allow_custom_loader=True) + + +def get_direct_mmap_adapter(pipeline: nn.Module) -> DirectMmapAdapter | None: + """Return a loader-owned adapter without requiring a pipeline flag.""" + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTModel, + ) + + if any(isinstance(module, MiniMaxH3DiTModel) for module in pipeline.modules()): + return _MiniMaxH3DirectMmapAdapter(pipeline) + + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + Cosmos3VFMTransformer, + ) + + if any(isinstance(module, Cosmos3VFMTransformer) for module in pipeline.modules()): + return _Cosmos3DirectMmapAdapter() + return None + + +__all__ = [ + "DirectMmapAdapter", + "DirectMmapTensorPolicy", + "get_direct_mmap_adapter", +] diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index d1a1ed15e..7069e5012 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -1,30 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import contextlib import dataclasses import glob +import json import os import re import time -from collections.abc import Generator, Iterable +from collections.abc import Generator, Iterable, Sequence from pathlib import Path from typing import cast +import huggingface_hub import torch from torch import nn from vllm.config.load import LoadConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import ( - download_safetensors_index_file_from_hf, download_weights_from_hf, - filter_duplicate_safetensors_files, filter_files_not_needed_for_inference, maybe_download_from_modelscope, multi_thread_safetensors_weights_iterator, safetensors_weights_iterator, ) -from vllm.transformers_utils.repo_utils import file_exists +from vllm.transformers_utils.repo_utils import hf_api from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.torch_utils import set_default_torch_dtype @@ -34,9 +34,16 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) +from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + build_checkpoint_mmap_plan, + has_online_quantization, +) from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.registry import initialize_model +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific # download_gguf was removed from upstream vLLM (commit 6635279d8). @@ -83,6 +90,23 @@ def _natural_sort_key(filepath: str) -> list: DIFFUSION_MODEL_WEIGHTS_INDEX = "diffusion_pytorch_model.safetensors.index.json" TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json" INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX] +SHARDED_SAFETENSORS_PATTERN = re.compile(r"^(?P.+)-\d+-of-(?P\d+)\.safetensors$") + + +def _validate_unindexed_safetensors_layout(weight_files: Sequence[str]) -> None: + """Reject ambiguous shard families when no index is available.""" + shard_counts: dict[str, set[int]] = {} + for weight_file in weight_files: + match = SHARDED_SAFETENSORS_PATTERN.fullmatch(os.path.basename(weight_file)) + if match is not None: + shard_counts.setdefault(match.group("family"), set()).add(int(match.group("count"))) + + conflicts = {family: sorted(counts) for family, counts in shard_counts.items() if len(counts) > 1} + if conflicts: + raise ValueError( + "Ambiguous unindexed safetensors checkpoint with conflicting shard totals: " + f"{conflicts}. Refusing to load potentially stale checkpoint shards." + ) def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> type: @@ -102,7 +126,7 @@ def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> typ ) -class DiffusersPipelineLoader: +class DiffusersPipelineLoader(HWRLoaderMixin): """Model loader that can load diffusers pipeline components from disk.""" @dataclasses.dataclass @@ -135,6 +159,70 @@ class DiffusersPipelineLoader: self.od_config = od_config self.quant_config = od_config.quantization_config self.parallel_config = od_config.parallel_config + self.host_weight_plan: HostWeightPlan | None = None + self._hwr_state: dict[str, object] | None = None + self._last_load_request: dict[str, object] | None = None + self._force_canonical_load = False + + def take_host_weight_plan(self) -> HostWeightPlan | None: + """Transfer the loader-produced plan to the offload backend.""" + plan = self.host_weight_plan + self.host_weight_plan = None + return plan + + @staticmethod + def _repo_relative_path(subfolder: str | None, filename: str) -> str: + if subfolder is None: + return filename + prefix = f"{subfolder.rstrip('/')}/" + return filename if filename.startswith(prefix) else f"{prefix}{filename.lstrip('/')}" + + def _resolve_weight_index( + self, + model_name_or_path: Path | str, + subfolder: str | None, + revision: str | None, + ) -> list[str] | None: + """Resolve an index and return its authoritative shard manifest.""" + is_local = os.path.isdir(model_name_or_path) + index_paths: list[tuple[str, Path]] = [] + for index_file in INDEX_FILES: + repo_index_path = self._repo_relative_path(subfolder, index_file) + if is_local: + index_path = Path(model_name_or_path) / repo_index_path + if index_path.is_file(): + index_paths.append((index_file, index_path)) + continue + + try: + index_path = hf_api().hf_hub_download( + repo_id=str(model_name_or_path), + filename=repo_index_path, + cache_dir=self.load_config.download_dir, + revision=revision, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + except huggingface_hub.errors.EntryNotFoundError: + continue + index_paths.append((index_file, Path(index_path))) + + if len(index_paths) > 1: + raise ValueError( + f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: " + f"{[index_file for index_file, _ in index_paths]}" + ) + if not index_paths: + return None + + index_file, index_path = index_paths[0] + with open(index_path) as f: + index = json.load(f) + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Weight index {index_file} must contain a non-empty `weight_map`") + if not all(isinstance(filename, str) and filename for filename in weight_map.values()): + raise ValueError(f"Weight index {index_file} contains an invalid shard filename") + return sorted(set(weight_map.values())) def _prepare_weights( self, @@ -152,17 +240,11 @@ class DiffusersPipelineLoader: is_local = os.path.isdir(model_name_or_path) load_format = self.load_config.load_format use_safetensors = False - possible_index_files = [ - f"{subfolder}/{index_file}" if subfolder is not None else index_file for index_file in INDEX_FILES - ] - available_index_file = [ - f for f in possible_index_files if file_exists(model_name_or_path, f, revision=revision) - ] - if len(available_index_file) > 1: - raise ValueError( - f"Multiple index files found in {model_name_or_path} with subfolder {subfolder}: {available_index_file}" - ) - index_file = available_index_file[0] if available_index_file else "" + indexed_weight_files = ( + self._resolve_weight_index(model_name_or_path, subfolder, revision) + if allow_patterns_overrides is None + else None + ) # only hf is supported currently if load_format == "auto": @@ -180,7 +262,16 @@ class DiffusersPipelineLoader: if allow_patterns_overrides is not None: allow_patterns = allow_patterns_overrides - if not is_local: + if not is_local and indexed_weight_files is not None: + hf_folder: Path | str = download_weights_from_hf_specific( + model_name_or_path=str(model_name_or_path), + cache_dir=self.load_config.download_dir, + allow_patterns=[self._repo_relative_path(subfolder, filename) for filename in indexed_weight_files], + revision=revision, + ignore_patterns=self.load_config.ignore_patterns, + require_all=True, + ) + elif not is_local: hf_folder = download_weights_from_hf( model_name_or_path, self.load_config.download_dir, @@ -195,31 +286,24 @@ class DiffusersPipelineLoader: if subfolder is not None: hf_folder = os.path.join(hf_folder, subfolder) - hf_weights_files: list[str] = [] - for pattern in allow_patterns: - hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) - if len(hf_weights_files) > 0: - # Decide by actual files rather than pattern name (patterns may include subfolders). - use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files) - break - - if use_safetensors: - # For models like Mistral-7B-Instruct-v0.3 - # there are both sharded safetensors files and a consolidated - # safetensors file. Using both breaks. - # Here, we download the `model.safetensors.index.json` and filter - # any files not found in the index. - if not is_local: - download_safetensors_index_file_from_hf( - model_name_or_path, - index_file, - cache_dir=self.load_config.download_dir, - subfolder=subfolder, - revision=revision, - ) - hf_weights_files = filter_duplicate_safetensors_files(hf_weights_files, hf_folder, index_file) + if indexed_weight_files is not None: + hf_weights_files = [os.path.join(hf_folder, filename) for filename in indexed_weight_files] + missing_files = [filename for filename in hf_weights_files if not os.path.isfile(filename)] + if missing_files: + raise FileNotFoundError(f"Weight files referenced in index but missing: {missing_files}") + use_safetensors = any(filename.endswith(".safetensors") for filename in hf_weights_files) else: - hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) + hf_weights_files = [] + for pattern in allow_patterns: + hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) + if hf_weights_files: + # Decide by actual files rather than pattern name (patterns may include subfolders). + use_safetensors = any(f.endswith(".safetensors") for f in hf_weights_files) + break + if use_safetensors: + _validate_unindexed_safetensors_layout(hf_weights_files) + else: + hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) if len(hf_weights_files) == 0: raise RuntimeError(f"Cannot find any model weights with `{model_name_or_path}`") @@ -273,8 +357,9 @@ class DiffusersPipelineLoader: def _get_source_quant_config(self, source: "ComponentSource") -> object | None: quant_config = self.quant_config - if hasattr(quant_config, "resolve"): - return quant_config.resolve(source.prefix.rstrip(".")) + resolve = getattr(quant_config, "resolve", None) + if resolve is not None: + return resolve(source.prefix.rstrip(".")) return quant_config def _get_checkpoint_adapter( @@ -293,11 +378,67 @@ class DiffusersPipelineLoader: def get_all_weights( self, model: nn.Module, + sources: Sequence["ComponentSource"] | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: - sources = self._get_weight_sources(model) + if sources is None: + sources = self._get_weight_sources(model) for source in sources: yield from self._get_weights_iterator(source, model=model) + @staticmethod + def _stream_online_quant_weights_to_cpu( + model: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Generator[tuple[str, torch.Tensor], None, None]: + """Offload each online-quantized layer as soon as it is complete. + + Upstream vLLM's online layerwise loader materializes and quantizes a + layer synchronously while consuming ``weights``. Generator execution + resumes after that weight has been consumed, which gives us a safe + point to move completed layers to CPU before the next layer is loaded. + This bounds accelerator residency during CPU-offloaded model startup + instead of retaining the entire quantized model until loading ends. + """ + from vllm.model_executor.model_loader.reload.layerwise import ( + get_layerwise_info, + ) + + pending = { + module + for module in model.modules() + if getattr(getattr(module, "quant_method", None), "uses_meta_device", False) + and get_layerwise_info(module).can_load() + } + offloaded = 0 + + def offload_completed() -> None: + nonlocal offloaded + for module in tuple(pending): + if get_layerwise_info(module).can_load(): + continue + module.to("cpu") + pending.remove(module) + offloaded += 1 + + for weight in weights: + # This runs after the consumer has handled the previous yield. + offload_completed() + yield weight + offload_completed() + + # Quantization workspaces and the old accelerator-side parameter + # storages are now reusable cache blocks. Release them before the + # remaining (unquantized) model tensors are copied to CPU; otherwise + # the cached quantization footprint and final offload overlap in the + # process-level startup peak. + if offloaded: + torch.accelerator.empty_cache() + + logger.info( + "Stream-offloaded %d online-quantized layers to CPU during weight loading", + offloaded, + ) + def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]: return tuple( cast( @@ -332,6 +473,14 @@ class DiffusersPipelineLoader: device: torch.device | None = None, ) -> nn.Module: """Load a model with the given configurations.""" + self.host_weight_plan = None + self._hwr_state = None + self._last_load_request = { + "load_device": load_device, + "load_format": load_format, + "custom_pipeline_name": custom_pipeline_name, + "device": device, + } if load_format is None: load_format = "default" # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), @@ -363,34 +512,103 @@ class DiffusersPipelineLoader: else: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) - # Skip load_weights only for DLO+AllGather when the model - # supports mmap loading and no online quantization is active. - # This condition MUST match the gate in - # DistributedLayerwiseOffloadBackend.enable() so that the - # loader skips ⟺ enable() uses mmap. - from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading - _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) _use_ag = getattr(self.od_config, "dlo_use_allgather", True) _has_online_quant = self._has_online_quant(model) - _supports_mmap = supports_mmap_loading(model) + _tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1)) + _use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False)) + _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1)) + _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1)) + _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size - _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant - - if _skip_load: - logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())") - else: - if _dist_offload and _use_ag and _has_online_quant: - raise ValueError( - "Online quantization is incompatible with DLO+AllGather: " - "the sharding + AllGather mechanism flattens weights by " - "dtype, which breaks quantized weight/scale layouts. " - "Please use --dlo-no-use-allgather or disable online " - "quantization." + plan_result = None + weight_sources = self._get_weight_sources(model) + hwr_state = None + if not self._force_canonical_load: + try: + hwr_state = self._resolve_hwr( + model, + ModuleDiscovery.discover(model), + dist_offload=_dist_offload, + use_allgather=_use_ag, + load_format=load_format, + sources=weight_sources, ) - if _dist_offload and _use_ag and not _supports_mmap: + except _HWRCommitError: + from vllm_omni.host_weight_runtime import RuntimeMode + + mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled")) + if mode is not RuntimeMode.PREFERRED: + raise + logger.warning( + "HWR restore commit failed; discarding the model and retrying a fresh canonical load", + exc_info=True, + ) + del model + return self.load_fresh_canonical_model() + self._hwr_state = hwr_state + hwr_active = hwr_state is not None + if hwr_active and hwr_state is not None: + self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan")) + if _dist_offload and not hwr_active and not self._force_canonical_load: + modules = ModuleDiscovery.discover(model) + plan_result = build_checkpoint_mmap_plan( + model, + dit_modules=tuple(zip(modules.dit_names, modules.dits)), + sources=weight_sources, + model_path=str(getattr(self.od_config, "model", "")) or None, + tensor_parallel_size=_tp_size, + use_hsdp=_use_hsdp, + online_quantization=_has_online_quant, + ) + self.host_weight_plan = plan_result.plan + + host_weight_plan = self.host_weight_plan + + if host_weight_plan is not None: + logger.info( + "DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s", + "AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local", + host_weight_plan.backing_kind, + sorted(host_weight_plan.planned_source_prefixes) or "legacy DiT sources", + ) + ordinary_sources = tuple( + source + for source in weight_sources + if source.prefix not in host_weight_plan.planned_source_prefixes + ) + if ordinary_sources: logger.info( - "DLO+AllGather active but model does not support mmap: loading weights via regular loader." + "Loading %d component weight source(s) outside the DLO host-weight plan", + len(ordinary_sources), + ) + self.load_weights( + model, + sources=ordinary_sources, + planned_weights=host_weight_plan.bindings, + ) + else: + if _dist_offload and _use_ag and _dlo_group_size > 1 and _has_online_quant: + # An effective DLO group size of one performs no weight + # collective, so the AllGather layout allowlist does not + # apply there even with dlo_use_allgather=True. + unsupported_methods = self._unsupported_dlo_allgather_online_quant_methods(model) + if unsupported_methods: + raise ValueError( + "DLO+AllGather supports online quantization only for " + "per-tensor FP8, INT8, and MXFP8 linears; unsupported online " + f"methods: {', '.join(unsupported_methods)}. Please use " + "--dlo-no-use-allgather or disable online quantization." + ) + logger.info( + "Validated online methods (per-tensor FP8, INT8, MXFP8) with " + "DLO+AllGather: using the ordinary loader before sharding " + "finalized weights and scales" + ) + if _dist_offload and plan_result is not None: + logger.info( + "DLO direct checkpoint mmap unavailable; using ordinary loader: %s", + plan_result.fallback_reason, ) logger.debug("Loading weights on %s ...", load_device) if offload_after_quant: @@ -403,15 +621,49 @@ class DiffusersPipelineLoader: if load_format == "diffusers": cast(DiffusersAdapterPipeline, model).load_weights() else: - self.load_weights(model) + if offload_after_quant: + self.load_weights(model, stream_online_quant_to_cpu=True) + else: + self.load_weights(model) self._process_weights_after_loading(model, target_device) + # A warm final-layout hit has already completed all + # byte-changing work through the restorer. Shared runtime + # finalization happens once at the end for both cold and warm + # paths; the warm path never re-enters the ordinary + # materialization/finalization pipeline. + if offload_after_quant: model.to("cpu") logger.info("Quantization complete, offloaded model back to CPU") - self._apply_skip_softmax_calibration(model) - return model.eval() + try: + self._apply_skip_softmax_calibration(model) + model = model.eval() + if self._hwr_state is not None: + warm_snapshot = self._hwr_state.get("warm_snapshot") + if warm_snapshot is not None: + self._assert_final_layout_tensors_unchanged(model, cast(dict[str, tuple[int, str]], warm_snapshot)) + self._publish_hwr_after_load(model, ModuleDiscovery.discover(model), self._hwr_state) + except Exception: + hwr_plan = self._hwr_state.get("plan") if self._hwr_state is not None else None + if isinstance(hwr_plan, HostWeightPlan): + carrier = hwr_plan.lease_carrier + if carrier is not None: + carrier.close() + from vllm_omni.host_weight_runtime import RuntimeMode + + mode = RuntimeMode(getattr(self.od_config, "host_weight_runtime_mode", "disabled")) + if mode is RuntimeMode.PREFERRED: + logger.warning( + "HWR warm finalization failed; discarding the model and retrying a fresh canonical load", + exc_info=True, + ) + del model + return self.load_fresh_canonical_model() + raise + self._attach_offload_startup_state(model) + return model @staticmethod def _request_offload_after_quant(model: nn.Module) -> int: @@ -431,9 +683,10 @@ class DiffusersPipelineLoader: marked = 0 for module in model.modules(): quant_method = getattr(module, "quant_method", None) - if getattr(quant_method, "supports_offload_after_quant", False): - quant_method.enable_offload_after_quant() - marked += 1 + if quant_method is None or not getattr(quant_method, "supports_offload_after_quant", False): + continue + quant_method.enable_offload_after_quant() + marked += 1 return marked @staticmethod @@ -441,11 +694,70 @@ class DiffusersPipelineLoader: """Whether any layer uses an online-quant method that defers weight materialization onto the ``meta`` device (upstream vLLM ``uses_meta_device=True``, e.g. online FP8).""" + return has_online_quantization(model) + + @staticmethod + def _unsupported_dlo_allgather_online_quant_methods(model: nn.Module) -> tuple[str, ...]: + """Return unsupported online-quant methods for DLO AllGather. + + Per-tensor online FP8, online INT8, and online MXFP8 are safe after + the ordinary loader has finalized their weight and scale parameters. + DLO shards those runtime tensors by dtype and reconstructs their + recorded shapes and strides before the kernel consumes them. They all + keep plain transportable 1-byte dtypes over ordinary strided views: + + - online INT8: int8 weight plus fp32 scale, either contiguous (NPU, + pre-transposed (K, N)) or a transposed view (CUDA, stride (1, K)); + - online MXFP8: fp8 weight plus e8m0 block scale, either contiguous + (NPU: (K, N) weight with (K_groups/2, N, 2) scale) or with the + scale stored as a transposed view (vLLM kernel, .t() over a + contiguous (K/32, N) buffer). + + Both shape families are already covered by the physical-order packing + that online FP8 requires. Other online methods may create different + scale, packing, or aliasing layouts (e.g. dual-scale fp4 pairs, + swizzled or NZ hardware formats) and remain fail-closed until + validated. + """ + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + from vllm_omni.quantization.int8_config import ( + Int8OnlineLinearMethod, + NPUInt8OnlineLinearMethod, + ) + + try: + from vllm_omni.quantization.mxfp8_config import ( + NPUMxfp8OnlineLinearMethod, + VllmMxfp8OnlineLinearMethod, + ) + + mxfp8_online_methods: tuple[type, ...] = ( + NPUMxfp8OnlineLinearMethod, + VllmMxfp8OnlineLinearMethod, + ) + except ImportError: + # MXFP8 requires a vLLM build with MXFP8 kernel support; treat it + # as absent when the module cannot be imported. + mxfp8_online_methods = () + + allowed_online_methods: tuple[type, ...] = ( + Fp8PerTensorOnlineLinearMethod, + Int8OnlineLinearMethod, + NPUInt8OnlineLinearMethod, + *mxfp8_online_methods, + ) + + unsupported: set[str] = set() for module in model.modules(): quant_method = getattr(module, "quant_method", None) - if getattr(quant_method, "uses_meta_device", False): - return True - return False + if not getattr(quant_method, "uses_meta_device", False): + continue + if not isinstance(quant_method, allowed_online_methods): + unsupported.add(type(quant_method).__name__) + return tuple(sorted(unsupported)) def _apply_skip_softmax_calibration(self, model: nn.Module) -> None: from vllm_omni.diffusion.attention.backends.trtllm_calibration import ( @@ -498,6 +810,14 @@ class DiffusersPipelineLoader: continue if has_online_quant: + # finalize_layerwise_processing() and the synchronous online + # loader already processed these layers. Avoid moving their + # quantized CPU weights back to the accelerator merely to call + # an idempotent no-op; doing so rebuilds a large CUDA allocator + # cache and defeats streaming CPU offload's startup-memory bound. + if getattr(module, "_already_called_process_weights_after_loading", False): + continue + # Online quant may leave straggler params on the ``meta`` device. # Move only real (non-meta) params onto the target device for # processing and restore them afterward, mirroring upstream vLLM's @@ -535,9 +855,21 @@ class DiffusersPipelineLoader: if needs_device_move: module.to(module_device) - def load_weights(self, model: nn.Module) -> None: + def load_weights( + self, + model: nn.Module, + *, + stream_online_quant_to_cpu: bool = False, + sources: Sequence["ComponentSource"] | None = None, + planned_weights: Iterable[str] = (), + ) -> None: weights_to_load = self._get_expected_parameter_names(model) - loaded_weights = model.load_weights(self.get_all_weights(model)) + weights = self.get_all_weights(model) if sources is None else self.get_all_weights(model, sources=sources) + if stream_online_quant_to_cpu: + weights = self._stream_online_quant_weights_to_cpu(model, weights) + loaded_weights = model.load_weights(weights) + if loaded_weights is not None: + loaded_weights = set(loaded_weights).union(planned_weights) self.counter_after_loading_weights = time.perf_counter() logger.info_once( diff --git a/vllm_omni/diffusion/model_loader/host_weight_loader.py b/vllm_omni/diffusion/model_loader/host_weight_loader.py new file mode 100644 index 000000000..4b9853555 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weight_loader.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Host Weight Runtime integration for diffusion model loaders. + +This mixin keeps HWR identity, restore, publication, and startup-recovery +policy out of the ordinary diffusers loader. The host loader supplies the +canonical source-preparation and weight-loading hooks; this module owns only +the optional final-layout transaction. +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import os +from collections.abc import Iterable, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, cast + +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan, TensorBinding + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext + from vllm_omni.diffusion.model_loader.host_weights.source_identity import ( + NodeSourceDigestCache, + PreparedWeightSource, + ) + + +class _HWRCommitError(RuntimeError): + """A committed warm restore made the current model disposable.""" + + +class HWRLoaderMixin: + """Optional final-layout HWR behavior shared by diffusion loaders.""" + + @staticmethod + def _identity_value(value: object) -> object: + """Convert config objects into deterministic identity metadata.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, torch.dtype): + return str(value) + if isinstance(value, dict): + return {str(key): HWRLoaderMixin._identity_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [HWRLoaderMixin._identity_value(item) for item in value] + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + try: + return HWRLoaderMixin._identity_value(to_dict()) + except TypeError: + pass + if dataclasses.is_dataclass(value): + return HWRLoaderMixin._identity_value(dataclasses.asdict(value)) + return f"{type(value).__module__}.{type(value).__qualname__}:{value!r}" + + @staticmethod + def _identity_fingerprint(value: object) -> str: + from vllm_omni.host_weight_runtime import CanonicalJson + + return hashlib.sha256(CanonicalJson.from_value(value).encoded).hexdigest() + + @staticmethod + def _snapshot_final_layout_tensors(model: nn.Module, names: Iterable[str]) -> dict[str, tuple[int, str]]: + snapshot: dict[str, tuple[int, str]] = {} + for name in names: + parent_path, _, leaf_name = name.rpartition(".") + parent = model.get_submodule(parent_path) + tensor = parent._parameters.get(leaf_name) + if tensor is None: + tensor = parent._buffers.get(leaf_name) + if tensor is None or tensor.is_meta: + raise RuntimeError(f"cannot snapshot missing or meta final-layout tensor {name!r}") + value = tensor.detach() + if value.device.type != "cpu": + value = value.cpu() + value = value.contiguous() + digest = hashlib.sha256(memoryview(value.view(torch.uint8).numpy())).hexdigest() + snapshot[name] = (tensor.untyped_storage().data_ptr(), digest) + return snapshot + + @classmethod + def _assert_final_layout_tensors_unchanged( + cls, + model: nn.Module, + snapshot: dict[str, tuple[int, str]], + ) -> None: + current = cls._snapshot_final_layout_tensors(model, snapshot) + changed = [name for name in snapshot if current[name] != snapshot[name]] + if changed: + raise RuntimeError( + "shared warm finalization changed restored final-layout tensor bytes or backing pointers: " + f"{changed[:5]}" + ) + + def _hwr_eligibility_mode( + self, + model: nn.Module, + modules: object, + *, + dist_offload: bool, + use_allgather: bool, + load_format: str, + ) -> object | None: + """Return the enabled HWR mode only after all zero-interaction gates.""" + from vllm_omni.host_weight_runtime import RuntimeMode + + raw_mode = getattr(self.od_config, "host_weight_runtime_mode", "disabled") + try: + mode = RuntimeMode(raw_mode) + except ValueError as exc: + raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required") from exc + + # These gates intentionally precede HWR imports, source preparation, + # identity construction, and store creation. + if mode is RuntimeMode.DISABLED or not dist_offload or use_allgather: + return None + + parallel = self.parallel_config + reasons: list[str] = [] + if load_format != "default": + reasons.append("load_format must be 'default'") + if bool(getattr(parallel, "use_hsdp", False)): + reasons.append("HSDP layouts are not supported by the final-layout BF16 consumer") + if self.quant_config is not None: + reasons.append("quantized layouts require a representation-specific HWR producer") + if getattr(self.od_config, "lora_path", None): + reasons.append("adapted weights are not reusable base-model artifacts") + + dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ()))) + if not dit_modules: + reasons.append("no DiT modules were discovered") + elif any( + getattr(dit, "host_weight_restore_contract", None) is None + or not callable(getattr(dit, "validate_restored_host_weights", None)) + for _, dit in dit_modules + ): + reasons.append("every owned DiT must declare the final-layout restore contract") + + if reasons: + message = "; ".join(reasons) + if mode is RuntimeMode.REQUIRED: + raise ValueError(f"required Host Weight Runtime path is ineligible: {message}") + logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message) + return None + return mode + + def _prepare_hwr_sources( + self, + model: nn.Module, + modules: object, + sources: Sequence[object], + ) -> tuple[PreparedWeightSource, ...]: + from vllm_omni.diffusion.model_loader.host_weights import ( + ImplementationIdentity, + PreparedWeightSource, + WeightSourceKind, + ) + + dit_prefixes = tuple(f"{name}." for name in getattr(modules, "dit_names", ())) + selected_sources = tuple( + source + for source in sources + if not source.prefix or any(prefix.startswith(source.prefix) for prefix in dit_prefixes) + ) + if not selected_sources: + raise ValueError("final-layout HWR requires canonical weight sources covering every DiT") + + prepared: list[PreparedWeightSource] = [] + for source in selected_sources: + resolved_root, weight_files, use_safetensors = self._prepare_weights( + source.model_or_path, + source.subfolder, + source.revision, + source.fall_back_to_pt, + source.allow_patterns_overrides, + ) + adapter = self._get_checkpoint_adapter(model, source, use_safetensors) + adapter_identity = None + if adapter is not None: + adapter_name = f"{type(adapter).__module__}.{type(adapter).__qualname__}" + adapter_identity = ImplementationIdentity( + implementation_id=adapter_name, + version="1", + fingerprint=self._identity_fingerprint({"adapter": adapter_name}), + ) + source_kind = ( + WeightSourceKind.LOCAL_PATH + if os.path.isdir(os.fspath(source.model_or_path)) + else WeightSourceKind.HUGGING_FACE_HUB + ) + prepared.append( + PreparedWeightSource( + model_or_path=os.fspath(source.model_or_path), + subfolder=source.subfolder, + requested_revision=source.revision, + prefix=source.prefix, + resolved_root=Path(os.fspath(resolved_root)), + weight_files=tuple(Path(os.fspath(path)) for path in weight_files), + use_safetensors=use_safetensors, + checkpoint_adapter=adapter_identity, + source_kind=source_kind, + ) + ) + return tuple(prepared) + + def _build_hwr_context( + self, + model: nn.Module, + modules: object, + *, + load_format: str, + sources: Sequence[object], + source_digest_cache: NodeSourceDigestCache | None = None, + ) -> FinalLayoutIdentityContext: + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + + from vllm_omni.diffusion.model_loader.host_weights import ( + FINAL_LAYOUT_BF16_POLICY, + FinalLayoutLoaderIdentity, + FinalLayoutParallelIdentity, + FinalLayoutRequest, + ImplementationIdentity, + build_final_layout_identity, + ) + + parallel = self.parallel_config + tp_size = int(getattr(parallel, "tensor_parallel_size", 1)) + try: + tp_rank = int(get_tensor_model_parallel_rank()) if tp_size > 1 else 0 + except Exception: + tp_rank = int(getattr(parallel, "tensor_parallel_rank", 0)) + sp_size = int(getattr(parallel, "sequence_parallel_size", 1) or 1) + loader_config = { + "dtype": str(self.od_config.dtype), + "model_class": f"{type(model).__module__}.{type(model).__qualname__}", + "model_config": self._identity_value(getattr(self.od_config, "tf_model_config", None)), + "load_format": load_format, + "quantization": self._identity_value(self.quant_config), + } + contracts = [ + self._identity_value(getattr(dit, "host_weight_restore_contract")) for dit in getattr(modules, "dits", ()) + ] + loader_identity = FinalLayoutLoaderIdentity( + implementation=ImplementationIdentity( + implementation_id="vllm-omni.diffusion.diffusers-loader", + version="final-layout-v1", + fingerprint=self._identity_fingerprint( + { + "loader": "diffusers-loader-final-layout-v1", + "pipeline": f"{type(model).__module__}.{type(model).__qualname__}", + } + ), + ), + model_config_fingerprint=self._identity_fingerprint(loader_config), + weight_transform_fingerprint=self._identity_fingerprint( + { + "contracts": contracts, + "transform": "diffusion-final-layout-loader-transforms-v1", + } + ), + ) + semantic_parallel = FinalLayoutParallelIdentity( + tensor_parallel_size=tp_size, + tensor_parallel_rank=tp_rank, + sequence_parallel_size=sp_size, + ulysses_degree=int(getattr(parallel, "ulysses_degree", 1)), + ring_degree=int(getattr(parallel, "ring_degree", 1)), + allgather_degree=int(getattr(parallel, "allgather_degree", 1)), + ulysses_mode=str(getattr(parallel, "ulysses_mode", "strict")), + pipeline_parallel_size=int(getattr(parallel, "pipeline_parallel_size", 1)), + cfg_parallel_size=int(getattr(parallel, "cfg_parallel_size", 1)), + use_hsdp=bool(getattr(parallel, "use_hsdp", False)), + enable_expert_parallel=bool(getattr(parallel, "enable_expert_parallel", False)), + ) + model_id = str(getattr(self.od_config, "model", "") or "") + if not model_id: + raise ValueError("final-layout HWR requires a canonical model identifier") + request = FinalLayoutRequest( + model_id=model_id, + loader=loader_identity, + parallel=semantic_parallel, + load_format=load_format, + ) + prepared_sources = self._prepare_hwr_sources(model, modules, sources) + dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ()))) + return build_final_layout_identity( + model, + dit_modules=dit_modules, + prepared_sources=prepared_sources, + request=request, + policy=FINAL_LAYOUT_BF16_POLICY, + source_digest_cache=source_digest_cache, + ) + + def _resolve_hwr( + self, + model: nn.Module, + modules: object, + *, + dist_offload: bool, + use_allgather: bool, + load_format: str, + sources: Sequence[object], + ) -> dict[str, object] | None: + """Resolve an eligible no-AllGather final-layout HWR transaction.""" + from vllm_omni.diffusion.model_loader.host_weights import ( + FinalLayoutTensorRestorer, + NodeSourceDigestCache, + ) + from vllm_omni.host_weight_runtime import ( + HostWeightLeaseCarrier, + HostWeightRuntime, + HostWeightRuntimeConfig, + ProductionPolicy, + ResolutionOutcome, + RuntimeMode, + StorageDomainPolicy, + ) + from vllm_omni.host_weight_runtime.filesystem import detect_storage_class + + mode = self._hwr_eligibility_mode( + model, + modules, + dist_offload=dist_offload, + use_allgather=use_allgather, + load_format=load_format, + ) + if mode is None: + return None + assert isinstance(mode, RuntimeMode) + expected_prefixes = frozenset(f"{name}." for name in getattr(modules, "dit_names", ())) + available_prefixes = frozenset(getattr(source, "prefix", "") for source in sources) + if not expected_prefixes <= available_prefixes: + message = "final-layout HWR requires one dedicated source prefix per owned DiT" + if mode is RuntimeMode.REQUIRED: + raise ValueError(f"required Host Weight Runtime path is ineligible: {message}") + logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message) + return None + overlapping_prefixes = frozenset( + prefix + for prefix in available_prefixes + if any(dit_prefix.startswith(prefix) for dit_prefix in expected_prefixes) + ) + if not overlapping_prefixes <= expected_prefixes: + message = "final-layout HWR requires dedicated DiT sources and rejects mixed component sources" + if mode is RuntimeMode.REQUIRED: + raise ValueError(f"required Host Weight Runtime path is ineligible: {message}") + logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message) + return None + root_value = getattr(self.od_config, "host_weight_runtime_root", None) + if not isinstance(root_value, str) or not root_value.strip(): + raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root") + root = Path(root_value).expanduser() + runtime = HostWeightRuntime.from_config( + HostWeightRuntimeConfig( + mode=mode, + domain=StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)), + production=ProductionPolicy( + allow_local_build=False, + allow_post_load_publish=True, + ), + ) + ) + source_digest_cache = NodeSourceDigestCache( + root, + timeout_seconds=runtime.config.wait.coordination_timeout_seconds, + ) + context = self._build_hwr_context( + model, + modules, + load_format=load_format, + sources=sources, + source_digest_cache=source_digest_cache, + ) + resolution = runtime.resolve(context.identity) + state: dict[str, object] = { + "mode": mode, + "context": context, + "runtime": runtime, + "outcome": resolution.report.outcome, + } + if resolution.report.outcome is ResolutionOutcome.FAILED: + failure = next( + (attempt.failure for attempt in reversed(resolution.report.attempts) if attempt.failure is not None), + None, + ) + detail = failure.message if failure is not None else "resolution failed without a typed detail" + raise RuntimeError(f"Host Weight Runtime resolution failed: {detail}") + if resolution.report.outcome is not ResolutionOutcome.LOCAL_HIT: + return state + + lease = resolution.lease + if lease is None: + raise RuntimeError("Host Weight Runtime returned LOCAL_HIT without a lease") + restorer = FinalLayoutTensorRestorer(context) + try: + restore_plan = restorer.plan_restore(model, lease) + except Exception: + lease.close() + if mode is RuntimeMode.REQUIRED: + raise + logger.warning("HWR warm restore planning failed; falling back to canonical loading", exc_info=True) + state["outcome"] = ResolutionOutcome.CANONICAL_FALLBACK + return state + + try: + restore_plan.commit() + except Exception as exc: + lease.close() + raise _HWRCommitError( + "Host Weight Runtime restore commit failed; the partially restored model must be discarded" + ) from exc + + try: + carrier = HostWeightLeaseCarrier(lease) + warm_snapshot = self._snapshot_final_layout_tensors(model, context.tensor_names) + except Exception as exc: + lease.close() + raise _HWRCommitError( + "Host Weight Runtime committed restore could not establish its startup ownership boundary" + ) from exc + source_metadata = context.identity.source.metadata.to_value() + target_bindings = source_metadata.get("target_bindings") if isinstance(source_metadata, dict) else None + planned_prefixes = frozenset( + binding["source_prefix"] + for binding in (target_bindings or ()) + if isinstance(binding, dict) and isinstance(binding.get("source_prefix"), str) + ) + if not planned_prefixes: + lease.close() + raise _HWRCommitError("Host Weight Runtime identity did not record owned canonical source prefixes") + state["plan"] = HostWeightPlan( + backing_kind="host_weight_runtime", + bindings={name: TensorBinding(name, "") for name in context.tensor_names}, + planned_source_prefixes=planned_prefixes, + lease_carrier=carrier, + ) + state["warm_snapshot"] = warm_snapshot + return state + + def load_fresh_canonical_model(self) -> nn.Module: + """Reload a disposable startup model without HWR or checkpoint mmap.""" + if self._last_load_request is None: + raise RuntimeError("cannot construct a fresh canonical model before the initial load") + request = dict(self._last_load_request) + self._force_canonical_load = True + try: + return self.load_model(**cast(dict[str, object], request)) + finally: + self._force_canonical_load = False + + def _publish_hwr_after_load(self, model: nn.Module, modules: object, state: dict[str, object] | None) -> None: + from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutBF16Producer + from vllm_omni.host_weight_runtime import PostLoadPublicationOutcome, ResolutionOutcome + + if state is None or state.get("outcome") is not ResolutionOutcome.CANONICAL_FALLBACK: + return + context = state["context"] + runtime = state["runtime"] + dit_modules = tuple(zip(getattr(modules, "dit_names", ()), getattr(modules, "dits", ()))) + try: + producer = FinalLayoutBF16Producer(context, model, dit_modules) + report = runtime.publish_after_load(context.identity, producer=producer) + except Exception: + logger.warning("Host Weight Runtime post-load publication failed", exc_info=True) + return + if report.outcome is PostLoadPublicationOutcome.FAILED: + logger.warning("Host Weight Runtime post-load publication failed: %s", report.failure) + + def _attach_offload_startup_state(self, model: nn.Module) -> None: + """Publish generic offload startup state on the loaded pipeline.""" + if self.host_weight_plan is None: + return + from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state + + hwr_state = self._hwr_state + allow_retry = False + if hwr_state is not None and isinstance(hwr_state.get("plan"), HostWeightPlan): + from vllm_omni.host_weight_runtime import RuntimeMode + + allow_retry = hwr_state.get("mode") is RuntimeMode.PREFERRED + attach_offload_startup_state( + model, + OffloadStartupState( + host_weight_plan=self.host_weight_plan, + fresh_model_loader=self.load_fresh_canonical_model if allow_retry else None, + allow_fresh_retry=allow_retry, + ), + ) diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py new file mode 100644 index 000000000..ca755b082 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Loader-owned host-weight plans shared with diffusion offload backends.""" + +from __future__ import annotations + +import glob +import json +import os +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from safetensors import safe_open +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( + get_direct_mmap_adapter, +) + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm_omni.host_weight_runtime import HostWeightLeaseCarrier + +TensorTransform = Callable[[torch.Tensor], torch.Tensor] + + +@dataclass(frozen=True) +class TensorBinding: + """One runtime tensor backed by one safetensors entry.""" + + checkpoint_key: str + file_path: str + transform: TensorTransform | None = None + + +@dataclass(frozen=True) +class HostWeightPlan: + """Complete, prevalidated host backing consumed by an offload backend.""" + + backing_kind: str + bindings: dict[str, TensorBinding] + planned_source_prefixes: frozenset[str] = frozenset() + lease_carrier: HostWeightLeaseCarrier | None = None + + +@dataclass(frozen=True) +class HostWeightPlanResult: + """A complete plan or a fail-closed reason to use the ordinary loader.""" + + plan: HostWeightPlan | None + fallback_reason: str | None = None + + +class _PlanIncompatibleError(RuntimeError): + pass + + +_SAFETENSORS_DTYPES: dict[str, torch.dtype] = { + "BOOL": torch.bool, + "U8": torch.uint8, + "I8": torch.int8, + "I16": torch.int16, + "I32": torch.int32, + "I64": torch.int64, + "F16": torch.float16, + "BF16": torch.bfloat16, + "F32": torch.float32, + "F64": torch.float64, +} +for _safe_name, _torch_name in ( + ("F8_E4M3", "float8_e4m3fn"), + ("F8_E5M2", "float8_e5m2"), +): + if (_dtype := getattr(torch, _torch_name, None)) is not None: + _SAFETENSORS_DTYPES[_safe_name] = _dtype + + +def has_online_quantization(model: nn.Module) -> bool: + """Whether any module defers online-quantized weights on ``meta``.""" + return any(getattr(getattr(module, "quant_method", None), "uses_meta_device", False) for module in model.modules()) + + +def _resolve_source_root(source: object) -> str: + source_root = os.fspath(getattr(source, "model_or_path")) + subfolder = getattr(source, "subfolder", None) + if not os.path.isdir(source_root): + from vllm.model_executor.model_loader.weight_utils import ( + download_weights_from_hf, + ) + + source_root = download_weights_from_hf( + model_name_or_path=source_root, + cache_dir=None, + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + revision=getattr(source, "revision", None), + subfolder=subfolder, + ) + return os.path.join(source_root, subfolder) if subfolder is not None else source_root + + +def _record_binding( + model_to_ckpt: dict[str, tuple[str, str]], + model_name: str, + checkpoint_key: str, + file_path: str, +) -> None: + existing = model_to_ckpt.get(model_name) + candidate = (checkpoint_key, file_path) + if existing is not None and existing != candidate: + raise _PlanIncompatibleError( + f"multiple checkpoint entries map to runtime tensor {model_name!r}: {existing!r} and {candidate!r}" + ) + model_to_ckpt[model_name] = candidate + + +def _add_safetensors_source( + checkpoint_dir: str, + prefix: str, + remap_fn: Callable[[str], str | None] | None, + model_to_ckpt: dict[str, tuple[str, str]], +) -> int: + indexed_keys = 0 + index_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors.index.json"))) + if index_files: + for index_file in index_files: + with open(index_file, encoding="utf-8") as handle: + index = json.load(handle) + for checkpoint_key, filename in index.get("weight_map", {}).items(): + source_name = prefix + checkpoint_key + model_name = remap_fn(source_name) if remap_fn is not None else source_name + if model_name is not None: + _record_binding( + model_to_ckpt, + model_name, + checkpoint_key, + os.path.join(checkpoint_dir, filename), + ) + indexed_keys += 1 + return indexed_keys + + for filename in sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))): + with safe_open(filename, framework="pt", device="cpu") as handle: + for checkpoint_key in handle.keys(): + source_name = prefix + checkpoint_key + model_name = remap_fn(source_name) if remap_fn is not None else source_name + if model_name is not None: + _record_binding(model_to_ckpt, model_name, checkpoint_key, filename) + indexed_keys += 1 + return indexed_keys + + +def _build_source_map( + *, + sources: Sequence[object], + model_path: str | None, + remap_fn: Callable[[str], str | None] | None, +) -> dict[str, tuple[str, str]]: + model_to_ckpt: dict[str, tuple[str, str]] = {} + indexed_keys = 0 + + for source in sources: + indexed_keys += _add_safetensors_source( + _resolve_source_root(source), + getattr(source, "prefix", ""), + remap_fn, + model_to_ckpt, + ) + + if not sources: + if remap_fn is None: + raise _PlanIncompatibleError( + "the loader has no component weight sources and the pipeline has no checkpoint-key remapper" + ) + if not model_path: + raise _PlanIncompatibleError("the loader has no model path for legacy checkpoint remapping") + + resolved_model_path = model_path + if not os.path.isdir(resolved_model_path): + from vllm.model_executor.model_loader.weight_utils import ( + download_weights_from_hf, + ) + + resolved_model_path = download_weights_from_hf( + model_name_or_path=resolved_model_path, + cache_dir=None, + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + ) + checkpoint_dirs = { + os.path.dirname(path) + for pattern in ("*.safetensors.index.json", "*.safetensors") + for path in glob.glob(os.path.join(resolved_model_path, "**", pattern), recursive=True) + } + for checkpoint_dir in sorted(checkpoint_dirs): + indexed_keys += _add_safetensors_source( + checkpoint_dir, + "", + remap_fn, + model_to_ckpt, + ) + + if not model_to_ckpt: + raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources") + logger.info( + "Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight", + len(model_to_ckpt), + indexed_keys, + ) + return model_to_ckpt + + +def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool: + parts = local_name.split(".") + owner = module.get_submodule(".".join(parts[:-1])) if len(parts) > 1 else module + return parts[-1] not in owner._non_persistent_buffers_set + + +def _collect_required_dit_tensors( + dit_modules: Sequence[tuple[str, nn.Module]], +) -> dict[str, torch.Tensor]: + required: dict[str, torch.Tensor] = {} + for dit_name, dit_module in dit_modules: + for local_name, tensor in dit_module.named_parameters(): + required[f"{dit_name}.{local_name}"] = tensor + for local_name, tensor in dit_module.named_buffers(): + if _is_persistent_buffer(dit_module, local_name): + required[f"{dit_name}.{local_name}"] = tensor + return required + + +def _planned_source_prefixes( + sources: Sequence[object], + dit_modules: Sequence[tuple[str, nn.Module]], +) -> frozenset[str]: + """Identify component sources that the DiT-only plan owns completely.""" + if not sources: + return frozenset() + + expected = frozenset(f"{dit_name}." for dit_name, _ in dit_modules) + available = {getattr(source, "prefix", "") for source in sources} + missing = sorted(expected - available) + if missing: + raise _PlanIncompatibleError( + f"direct mmap requires a dedicated component weight source for each DiT; missing prefixes: {missing}" + ) + return expected + + +def _validate_source_metadata( + required: dict[str, torch.Tensor], + bindings: dict[str, TensorBinding], +) -> None: + by_file: dict[str, list[tuple[str, torch.Tensor, TensorBinding]]] = {} + for runtime_name, target in required.items(): + by_file.setdefault(bindings[runtime_name].file_path, []).append((runtime_name, target, bindings[runtime_name])) + + for file_path, entries in by_file.items(): + if not os.path.isfile(file_path): + raise _PlanIncompatibleError(f"checkpoint file does not exist: {file_path}") + with safe_open(file_path, framework="pt", device="cpu") as handle: + available = set(handle.keys()) + for runtime_name, target, binding in entries: + if binding.checkpoint_key not in available: + raise _PlanIncompatibleError( + f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}" + ) + source = handle.get_slice(binding.checkpoint_key) + source_shape = tuple(source.get_shape()) + source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype()) + if source_shape != tuple(target.shape): + raise _PlanIncompatibleError( + f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, runtime={tuple(target.shape)}" + ) + if source_dtype is None or source_dtype != target.dtype: + raise _PlanIncompatibleError( + f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, runtime={target.dtype}" + ) + + +def build_checkpoint_mmap_plan( + pipeline: nn.Module, + *, + dit_modules: Sequence[tuple[str, nn.Module]], + sources: Sequence[object], + model_path: str | None, + tensor_parallel_size: int, + use_hsdp: bool, + online_quantization: bool, +) -> HostWeightPlanResult: + """Build a complete direct-checkpoint plan or return a fallback reason.""" + if tensor_parallel_size != 1: + return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader") + if use_hsdp: + return HostWeightPlanResult(None, "HSDP requires the ordinary loader") + if online_quantization: + return HostWeightPlanResult(None, "online quantization requires the ordinary loader") + + remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) + if not callable(remap_fn): + remap_fn = None + adapter = get_direct_mmap_adapter(pipeline) + + try: + planned_source_prefixes = _planned_source_prefixes(sources, dit_modules) + planned_sources = tuple( + source for source in sources if getattr(source, "prefix", "") in planned_source_prefixes + ) + model_to_ckpt = _build_source_map( + sources=planned_sources, + model_path=model_path, + remap_fn=remap_fn, + ) + required = _collect_required_dit_tensors(dit_modules) + if not required: + raise _PlanIncompatibleError("no DiT parameters or persistent buffers were discovered") + + missing = sorted(set(required) - set(model_to_ckpt)) + if missing: + raise _PlanIncompatibleError( + f"{len(missing)} required DiT tensors have no checkpoint binding (first 5: {missing[:5]})" + ) + + bindings: dict[str, TensorBinding] = {} + for runtime_name, target in required.items(): + policy = adapter.policy_for(runtime_name, target) if adapter is not None else None + has_custom_loader = callable(getattr(target, "weight_loader", None)) + if has_custom_loader and not (policy and policy.allow_custom_loader): + raise _PlanIncompatibleError( + f"{runtime_name!r} requires a custom weight_loader with no direct-mmap adapter" + ) + checkpoint_key, file_path = model_to_ckpt[runtime_name] + bindings[runtime_name] = TensorBinding( + checkpoint_key=checkpoint_key, + file_path=file_path, + transform=policy.transform if policy is not None else None, + ) + + _validate_source_metadata(required, bindings) + except (OSError, ValueError, _PlanIncompatibleError) as exc: + return HostWeightPlanResult(None, str(exc)) + + return HostWeightPlanResult( + HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings=bindings, + planned_source_prefixes=planned_source_prefixes, + ) + ) + + +__all__ = [ + "HostWeightPlan", + "HostWeightPlanResult", + "TensorBinding", + "build_checkpoint_mmap_plan", + "has_online_quantization", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/__init__.py b/vllm_omni/diffusion/model_loader/host_weights/__init__.py new file mode 100644 index 000000000..1a1927fdb --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/__init__.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed diffusion adapters for runtime-ready final-layout artifacts.""" + +from .contracts import ( + FinalLayoutArtifactSpec, + FinalLayoutLoaderIdentity, + FinalLayoutParallelIdentity, + FinalLayoutRequest, + ImplementationIdentity, +) +from .identity_adapter import ( + FinalLayoutIdentityContext, + build_final_layout_identity, +) +from .producers import ( + FINAL_LAYOUT_BF16_POLICY, + FINAL_LAYOUT_BF16_SPEC, + FinalLayoutBF16Producer, +) +from .restorer import FinalLayoutTensorRestorer +from .source_identity import NodeSourceDigestCache, PreparedWeightSource, WeightSourceKind + +__all__ = [ + "FINAL_LAYOUT_BF16_POLICY", + "FINAL_LAYOUT_BF16_SPEC", + "FinalLayoutArtifactSpec", + "FinalLayoutBF16Producer", + "FinalLayoutIdentityContext", + "FinalLayoutLoaderIdentity", + "FinalLayoutParallelIdentity", + "FinalLayoutRequest", + "FinalLayoutTensorRestorer", + "ImplementationIdentity", + "NodeSourceDigestCache", + "PreparedWeightSource", + "WeightSourceKind", + "build_final_layout_identity", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/contracts.py b/vllm_omni/diffusion/model_loader/host_weights/contracts.py new file mode 100644 index 000000000..f62bb7bb0 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/contracts.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed semantic contracts shared by final-layout artifact representations.""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING, Protocol + +import torch + +from vllm_omni.diffusion.models.host_weight_contract import ( + FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, +) +from vllm_omni.host_weight_runtime import ( + AdaptationIdentity, + CanonicalJson, + FailureCode, + HostWeightError, + HostWeightFailure, + ProducerIdentity, + ResolutionStage, + TensorKind, + WeightRepresentation, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from .tensor_layout import RuntimeTensorTarget + +FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA = "diffusion-final-layout-tensor-restorer-v1" + + +def _require_text(name: str, value: object) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{name} must not be empty") + return value + + +def _require_positive_int(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{name} must be a positive integer") + return value + + +class FinalLayoutContractCode(str, Enum): + SOURCE_CHANGED = "source_changed" + SOURCE_COVERAGE_INVALID = "source_coverage_invalid" + MODEL_CONTRACT_UNSUPPORTED = "model_contract_unsupported" + OWNERSHIP_AMBIGUOUS = "ownership_ambiguous" + TENSOR_UNSUPPORTED = "tensor_unsupported" + DTYPE_UNSUPPORTED = "dtype_unsupported" + IDENTITY_INCOMPATIBLE = "identity_incompatible" + TENSOR_CONTRACT_CHANGED = "tensor_contract_changed" + + +class FinalLayoutContractError(ValueError): + """Typed fail-closed error at a final-layout model integration boundary.""" + + def __init__(self, code: FinalLayoutContractCode, message: str) -> None: + if not isinstance(code, FinalLayoutContractCode): + raise ValueError("final-layout contract errors require FinalLayoutContractCode") + super().__init__(message) + self.code = code + + +def final_layout_producer_error(error: FinalLayoutContractError) -> HostWeightError: + """Translate representation failures into the runtime's stable taxonomy.""" + if error.code is FinalLayoutContractCode.SOURCE_CHANGED: + stage = ResolutionStage.CANONICAL_LOADING + code = FailureCode.CANONICAL_SOURCE_FAILED + elif error.code in { + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + }: + stage = ResolutionStage.PRODUCTION + code = FailureCode.PRODUCER_UNSUPPORTED + else: + stage = ResolutionStage.PRODUCTION + code = FailureCode.PRODUCER_FAILED + return HostWeightError( + HostWeightFailure( + stage=stage, + code=code, + retryable=False, + message=str(error), + details=CanonicalJson.from_value({"final_layout_code": error.code.value}), + ) + ) + + +@dataclass(frozen=True) +class ImplementationIdentity: + """Explicit stable identity for one semantic implementation boundary.""" + + implementation_id: str + version: str + fingerprint: str + + def __post_init__(self) -> None: + _require_text("implementation_id", self.implementation_id) + _require_text("implementation version", self.version) + _require_text("implementation fingerprint", self.fingerprint) + + def to_dict(self) -> dict[str, str]: + return { + "implementation_id": self.implementation_id, + "version": self.version, + "fingerprint": self.fingerprint, + } + + +@dataclass(frozen=True) +class FinalLayoutLoaderIdentity: + """Loader ABI and exact configuration facts that may change final bytes.""" + + implementation: ImplementationIdentity + model_config_fingerprint: str + weight_transform_fingerprint: str + + def __post_init__(self) -> None: + if not isinstance(self.implementation, ImplementationIdentity): + raise ValueError("loader implementation must use ImplementationIdentity") + _require_text("model config fingerprint", self.model_config_fingerprint) + _require_text("weight transform fingerprint", self.weight_transform_fingerprint) + + def to_dict(self) -> dict[str, object]: + return { + "implementation": self.implementation.to_dict(), + "model_config_fingerprint": self.model_config_fingerprint, + "weight_transform_fingerprint": self.weight_transform_fingerprint, + } + + +@dataclass(frozen=True) +class FinalLayoutParallelIdentity: + """Semantic parallel layout with DP and SP rank intentionally absent.""" + + tensor_parallel_size: int = 1 + tensor_parallel_rank: int = 0 + sequence_parallel_size: int = 1 + ulysses_degree: int = 1 + ring_degree: int = 1 + allgather_degree: int = 1 + ulysses_mode: str = "strict" + pipeline_parallel_size: int = 1 + cfg_parallel_size: int = 1 + use_hsdp: bool = False + enable_expert_parallel: bool = False + + def __post_init__(self) -> None: + tp_size = _require_positive_int("tensor_parallel_size", self.tensor_parallel_size) + if isinstance(self.tensor_parallel_rank, bool) or not isinstance(self.tensor_parallel_rank, int): + raise ValueError("tensor_parallel_rank must be an integer") + if not 0 <= self.tensor_parallel_rank < tp_size: + raise ValueError("tensor_parallel_rank must be within tensor_parallel_size") + for name, value in ( + ("sequence_parallel_size", self.sequence_parallel_size), + ("ulysses_degree", self.ulysses_degree), + ("ring_degree", self.ring_degree), + ("allgather_degree", self.allgather_degree), + ("pipeline_parallel_size", self.pipeline_parallel_size), + ("cfg_parallel_size", self.cfg_parallel_size), + ): + _require_positive_int(name, value) + _require_text("ulysses_mode", self.ulysses_mode) + if not isinstance(self.use_hsdp, bool) or not isinstance(self.enable_expert_parallel, bool): + raise ValueError("HSDP and expert-parallel flags must be booleans") + + @property + def sequence_parallel_backend(self) -> str: + if self.sequence_parallel_size == 1: + return "none" + active = [ + name + for name, degree in ( + ("ulysses", self.ulysses_degree), + ("ring", self.ring_degree), + ("allgather", self.allgather_degree), + ) + if degree > 1 + ] + return "+".join(active) if active else "replicated-sp" + + def to_dict(self) -> dict[str, object]: + return { + "tensor_parallel_size": self.tensor_parallel_size, + "tensor_parallel_rank": self.tensor_parallel_rank, + "sequence_parallel_size": self.sequence_parallel_size, + "sequence_parallel_backend": self.sequence_parallel_backend, + "ulysses_degree": self.ulysses_degree, + "ring_degree": self.ring_degree, + "allgather_degree": self.allgather_degree, + "ulysses_mode": self.ulysses_mode, + "pipeline_parallel_size": self.pipeline_parallel_size, + "cfg_parallel_size": self.cfg_parallel_size, + "use_hsdp": self.use_hsdp, + "enable_expert_parallel": self.enable_expert_parallel, + } + + +@dataclass(frozen=True) +class FinalLayoutRequest: + """Representation-independent semantic request assembled by the loader.""" + + model_id: str + loader: FinalLayoutLoaderIdentity + parallel: FinalLayoutParallelIdentity = field(default_factory=FinalLayoutParallelIdentity) + load_format: str = "default" + adaptation: AdaptationIdentity = field(default_factory=AdaptationIdentity) + + def __post_init__(self) -> None: + _require_text("canonical model ID or path", self.model_id) + if not isinstance(self.loader, FinalLayoutLoaderIdentity): + raise ValueError("loader must use FinalLayoutLoaderIdentity") + if not isinstance(self.parallel, FinalLayoutParallelIdentity): + raise ValueError("parallel must use FinalLayoutParallelIdentity") + _require_text("load_format", self.load_format) + if not isinstance(self.adaptation, AdaptationIdentity): + raise ValueError("adaptation must use AdaptationIdentity") + + +def implementation_abi_fingerprint(abi: CanonicalJson) -> str: + if not isinstance(abi, CanonicalJson): + raise ValueError("implementation ABI must use CanonicalJson") + return hashlib.sha256(abi.encoded).hexdigest() + + +@dataclass(frozen=True) +class FinalLayoutArtifactSpec: + """Concrete representation, layout, producer, and restoration ABI.""" + + representation: WeightRepresentation + producer: ProducerIdentity + implementation_abi: CanonicalJson + layout_name: str = "diffusion-final-module-layout-v1" + component_name: str = "diffusion-dit" + component_ownership: str = "complete-final-layout-tensors" + model_contract_schema: str = FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA + + def __post_init__(self) -> None: + if not isinstance(self.representation, WeightRepresentation): + raise ValueError("artifact representation must use WeightRepresentation") + if not isinstance(self.producer, ProducerIdentity): + raise ValueError("artifact producer must use ProducerIdentity") + if not isinstance(self.implementation_abi, CanonicalJson): + raise ValueError("artifact implementation ABI must use CanonicalJson") + for name, value in ( + ("layout_name", self.layout_name), + ("component_name", self.component_name), + ("component_ownership", self.component_ownership), + ("model_contract_schema", self.model_contract_schema), + ): + _require_text(name, value) + abi = self.implementation_abi.to_value() + required_abi_fields = { + "artifact_identity", + "model_contract", + "producer", + "representation_policy", + "restorer", + "source_identity", + "tensor_contract", + } + if not isinstance(abi, dict) or not required_abi_fields <= set(abi): + raise ValueError( + "artifact implementation ABI must version identity, source, tensor, policy, producer, " + "restorer, and model contracts" + ) + for name in required_abi_fields: + _require_text(f"implementation ABI {name}", abi[name]) + if self.producer.implementation_fingerprint != implementation_abi_fingerprint(self.implementation_abi): + raise ValueError("producer implementation fingerprint differs from its explicit ABI descriptor") + + +class FinalLayoutTensorPolicy(Protocol): + """Representation-specific validation over shared final-layout mechanics.""" + + @property + def spec(self) -> FinalLayoutArtifactSpec: ... + + def validate_request(self, request: FinalLayoutRequest) -> None: ... + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: ... + + def validate_target(self, target: RuntimeTensorTarget) -> None: ... + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: ... + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: ... + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: ... + + +__all__ = [ + "FinalLayoutArtifactSpec", + "FinalLayoutContractCode", + "FinalLayoutContractError", + "FinalLayoutLoaderIdentity", + "FinalLayoutParallelIdentity", + "FinalLayoutRequest", + "FinalLayoutTensorPolicy", + "FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA", + "ImplementationIdentity", + "final_layout_producer_error", + "implementation_abi_fingerprint", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py b/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py new file mode 100644 index 000000000..30c9baac1 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact pre-load identity assembly for policy-defined final-layout artifacts.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +from torch import nn + +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract +from vllm_omni.host_weight_runtime import ( + CanonicalJson, + ComponentIdentity, + RuntimeWeightLayout, + WeightArtifactIdentity, + WeightRepresentation, +) + +from .contracts import ( + FinalLayoutArtifactSpec, + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutRequest, + FinalLayoutTensorPolicy, +) +from .source_identity import ( + FinalLayoutSourceContext, + NodeSourceDigestCache, + PreparedWeightSource, + resolve_final_layout_source_identity, +) +from .tensor_layout import ( + RuntimeTensorTarget, + collect_final_layout_targets, + tensor_contract_sha256, + validate_final_layout_model_contract, +) + + +@dataclass(frozen=True) +class FinalLayoutIdentityContext: + """Exact semantic identity plus process-local transaction guards.""" + + identity: WeightArtifactIdentity + spec: FinalLayoutArtifactSpec + policy: FinalLayoutTensorPolicy + request: FinalLayoutRequest + tensor_names: frozenset[str] + dit_names: tuple[str, ...] + model_contracts: tuple[FinalLayoutModelContract, ...] + source: FinalLayoutSourceContext + pipeline_type: type[nn.Module] + dit_types: tuple[type[nn.Module], ...] + + def sources_unchanged(self) -> bool: + return self.source.sources_unchanged() + + def ensure_sources_unchanged(self) -> None: + self.source.ensure_sources_unchanged() + + +def _require_metadata_object(name: str, metadata: CanonicalJson) -> Mapping[str, object]: + value = metadata.to_value() + if not isinstance(value, dict): + raise ValueError(f"{name} must encode one JSON object") + return value + + +def _model_contract_values(contracts: Sequence[FinalLayoutModelContract]) -> list[dict[str, str]]: + return [contract.to_dict() for contract in contracts] + + +def build_final_layout_identity( + pipeline: nn.Module, + *, + dit_modules: Sequence[tuple[str, nn.Module]], + prepared_sources: Sequence[PreparedWeightSource], + request: FinalLayoutRequest, + policy: FinalLayoutTensorPolicy, + source_digest_cache: NodeSourceDigestCache | None = None, +) -> FinalLayoutIdentityContext: + """Build exact source, representation, layout, and implementation identity.""" + if not isinstance(request, FinalLayoutRequest): + raise TypeError("final-layout identity requires FinalLayoutRequest") + policy.validate_request(request) + spec = policy.spec + bindings = validate_final_layout_model_contract( + dit_modules, + expected_schema=spec.model_contract_schema, + ) + structural_targets = collect_final_layout_targets( + pipeline, + dit_modules, + policy=policy, + require_materialized=False, + ) + target_names = frozenset(target.name for target in structural_targets) + source = resolve_final_layout_source_identity( + prepared_sources, + model_id=request.model_id, + target_names=target_names, + digest_cache=source_digest_cache, + ) + dit_names = tuple(name for name, _ in dit_modules) + model_contracts = tuple(binding.contract for binding in bindings) + + representation = WeightRepresentation( + name=spec.representation.name, + dtype=spec.representation.dtype, + metadata=CanonicalJson.from_value( + { + "artifact_policy": _require_metadata_object( + "artifact representation metadata", + spec.representation.metadata, + ), + "load_format": request.load_format, + "loader": request.loader.to_dict(), + } + ), + ) + identity = WeightArtifactIdentity( + schema_version=1, + source=source.identity, + component=ComponentIdentity( + name=spec.component_name, + ownership=spec.component_ownership, + metadata=CanonicalJson.from_value( + { + "component_names": list(dit_names), + "model_contracts": _model_contract_values(model_contracts), + } + ), + ), + representation=representation, + layout=RuntimeWeightLayout( + name=spec.layout_name, + tensor_parallel_size=request.parallel.tensor_parallel_size, + tensor_parallel_rank=request.parallel.tensor_parallel_rank, + sequence_parallel_size=request.parallel.sequence_parallel_size, + sequence_parallel_backend=request.parallel.sequence_parallel_backend, + metadata=CanonicalJson.from_value( + { + "model_contract_schema": spec.model_contract_schema, + "parallel": request.parallel.to_dict(), + "tensor_contract_sha256": tensor_contract_sha256(structural_targets), + } + ), + ), + adaptation=request.adaptation, + producer=spec.producer, + ) + return FinalLayoutIdentityContext( + identity=identity, + spec=spec, + policy=policy, + request=request, + tensor_names=target_names, + dit_names=dit_names, + model_contracts=model_contracts, + source=source, + pipeline_type=type(pipeline), + dit_types=tuple(type(module) for _, module in dit_modules), + ) + + +def validate_final_layout_identity( + context: FinalLayoutIdentityContext, + targets: Sequence[RuntimeTensorTarget], +) -> str: + """Validate a context against current targets and return their digest.""" + identity = context.identity + spec = context.spec + request = context.request + digest = tensor_contract_sha256(targets) + if identity.source != context.source.identity or identity.adaptation != request.adaptation: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity source or adaptation differs from the exact context", + ) + if identity.producer != spec.producer: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity producer differs from the artifact specification", + ) + if identity.representation.name != spec.representation.name or identity.representation.dtype != ( + spec.representation.dtype + ): + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity representation differs from the artifact specification", + ) + if identity.component.name != spec.component_name or identity.component.ownership != spec.component_ownership: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity component ownership differs from the artifact specification", + ) + if identity.layout.name != spec.layout_name: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity runtime layout differs from the artifact specification", + ) + + component = _require_metadata_object("component metadata", identity.component.metadata) + representation = _require_metadata_object("representation metadata", identity.representation.metadata) + layout = _require_metadata_object("layout metadata", identity.layout.metadata) + if representation != { + "artifact_policy": _require_metadata_object( + "artifact representation metadata", + spec.representation.metadata, + ), + "load_format": request.load_format, + "loader": request.loader.to_dict(), + }: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity loader or representation policy differs from the exact context", + ) + if component.get("component_names") != list(context.dit_names) or component.get("model_contracts") != ( + _model_contract_values(context.model_contracts) + ): + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity model ownership differs from the exact context", + ) + if layout.get("model_contract_schema") != spec.model_contract_schema: + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity model restore schema differs from the artifact specification", + ) + if ( + identity.layout.tensor_parallel_size != request.parallel.tensor_parallel_size + or identity.layout.tensor_parallel_rank != request.parallel.tensor_parallel_rank + or identity.layout.sequence_parallel_size != request.parallel.sequence_parallel_size + or identity.layout.sequence_parallel_backend != request.parallel.sequence_parallel_backend + or layout.get("parallel") != request.parallel.to_dict() + ): + raise FinalLayoutContractError( + FinalLayoutContractCode.IDENTITY_INCOMPATIBLE, + "identity parallel semantics differ from the exact context", + ) + if layout.get("tensor_contract_sha256") != digest: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_CONTRACT_CHANGED, + "identity tensor contract differs from the current model", + ) + if frozenset(target.name for target in targets) != context.tensor_names: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_CONTRACT_CHANGED, + "current tensor ownership differs from the exact identity context", + ) + return digest + + +__all__ = [ + "FinalLayoutIdentityContext", + "build_final_layout_identity", + "validate_final_layout_identity", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/producers/__init__.py b/vllm_omni/diffusion/model_loader/host_weights/producers/__init__.py new file mode 100644 index 000000000..268e26981 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/producers/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Diffusion Host Weight Runtime producers.""" + +from .final_layout_bf16 import ( + FINAL_LAYOUT_BF16_MANIFEST_SCHEMA, + FINAL_LAYOUT_BF16_POLICY, + FINAL_LAYOUT_BF16_PRODUCER_ID, + FINAL_LAYOUT_BF16_REPRESENTATION, + FINAL_LAYOUT_BF16_SPEC, + FINAL_LAYOUT_BF16_VERSION, + FinalLayoutBF16Policy, + FinalLayoutBF16Producer, +) + +__all__ = [ + "FINAL_LAYOUT_BF16_MANIFEST_SCHEMA", + "FINAL_LAYOUT_BF16_POLICY", + "FINAL_LAYOUT_BF16_PRODUCER_ID", + "FINAL_LAYOUT_BF16_REPRESENTATION", + "FINAL_LAYOUT_BF16_SPEC", + "FINAL_LAYOUT_BF16_VERSION", + "FinalLayoutBF16Policy", + "FinalLayoutBF16Producer", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py b/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py new file mode 100644 index 000000000..04a48356e --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Concrete BF16-with-preserved-FP32 final-layout artifact policy and producer.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +from torch import nn + +from vllm_omni.diffusion.models.host_weight_contract import ( + FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, +) +from vllm_omni.host_weight_runtime import ( + ArtifactWriter, + CanonicalJson, + CoordinationScope, + LookupPhase, + ProducerIdentity, + ProductionMetadata, + ProductionSourceMode, + TensorKind, + TensorWriteSpec, + WeightProductionSpec, + WeightRepresentation, +) + +from ..contracts import ( + FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + FinalLayoutArtifactSpec, + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutRequest, + final_layout_producer_error, + implementation_abi_fingerprint, +) +from ..identity_adapter import ( + FinalLayoutIdentityContext, + validate_final_layout_identity, +) +from ..tensor_layout import ( + RuntimeTensorTarget, + collect_final_layout_targets, + validate_final_layout_model_contract, +) + +FINAL_LAYOUT_BF16_PRODUCER_ID = "vllm-omni.diffusion.final-layout-bf16" +FINAL_LAYOUT_BF16_VERSION = "1" +FINAL_LAYOUT_BF16_REPRESENTATION = "diffusion-final-layout-bf16" +FINAL_LAYOUT_BF16_MANIFEST_SCHEMA = "diffusion-final-layout-bf16-manifest-v1" +DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3 + +_BF16_PARAMETER_DTYPES = {torch.bfloat16, torch.float32} +_BF16_BUFFER_DTYPES = { + torch.bool, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, +} +for _dtype_name in ( + "float8_e4m3fn", + "float8_e5m2", + "float8_e4m3fnuz", + "float8_e5m2fnuz", +): + if (_dtype := getattr(torch, _dtype_name, None)) is not None: + _BF16_BUFFER_DTYPES.add(_dtype) + +_BF16_IMPLEMENTATION_ABI = CanonicalJson.from_value( + { + "artifact_identity": "diffusion-final-layout-identity-v1", + "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA, + "producer": "finalized-bf16-tensor-writer-v1", + "representation_policy": "bf16-with-preserved-fp32-v1", + "restorer": "exact-final-layout-tensor-rebind-v1", + "source_identity": "prepared-diffusion-weight-source-v2", + "tensor_contract": "complete-strided-tensor-ownership-v1", + } +) + +FINAL_LAYOUT_BF16_SPEC = FinalLayoutArtifactSpec( + representation=WeightRepresentation( + name=FINAL_LAYOUT_BF16_REPRESENTATION, + dtype=str(torch.bfloat16), + metadata=CanonicalJson.from_value( + { + "format": FINAL_LAYOUT_BF16_REPRESENTATION, + "mixed_precision_policy": "bf16-with-preserved-fp32", + } + ), + ), + producer=ProducerIdentity( + producer_id=FINAL_LAYOUT_BF16_PRODUCER_ID, + version=FINAL_LAYOUT_BF16_VERSION, + implementation_fingerprint=implementation_abi_fingerprint(_BF16_IMPLEMENTATION_ABI), + manifest_schema=FINAL_LAYOUT_BF16_MANIFEST_SCHEMA, + restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + ), + implementation_abi=_BF16_IMPLEMENTATION_ABI, +) + + +class FinalLayoutBF16Policy: + """Concrete compatibility and metadata policy for finalized BF16 tensors.""" + + @property + def spec(self) -> FinalLayoutArtifactSpec: + return FINAL_LAYOUT_BF16_SPEC + + def validate_request(self, request: FinalLayoutRequest) -> None: + if request.load_format != "default": + raise ValueError(f"final-layout BF16 artifacts require load_format='default', got {request.load_format!r}") + if request.adaptation.kind != "base" or request.adaptation.fingerprint is not None: + raise ValueError("merged or static LoRA requires a representation-specific producer") + parallel = request.parallel + if parallel.pipeline_parallel_size != 1: + raise ValueError("pipeline-parallel ownership is not supported by the final-layout BF16 producer") + if parallel.cfg_parallel_size != 1: + raise ValueError("CFG-parallel ownership is not supported by the final-layout BF16 producer") + if parallel.use_hsdp: + raise ValueError("HSDP/DTensor layouts are not supported by the final-layout BF16 producer") + if parallel.enable_expert_parallel: + raise ValueError("expert-parallel ownership is not supported by the final-layout BF16 producer") + + def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str: + del name, tensor + if kind is TensorKind.PARAMETER: + return "weight" + return "persistent_buffer" + + def validate_target(self, target: RuntimeTensorTarget) -> None: + if target.kind is TensorKind.PARAMETER and target.tensor.dtype not in _BF16_PARAMETER_DTYPES: + raise FinalLayoutContractError( + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + f"{target.name!r} must be BF16 or an explicitly preserved FP32 parameter, got {target.tensor.dtype}", + ) + if target.kind is TensorKind.BUFFER and target.tensor.dtype not in _BF16_BUFFER_DTYPES: + raise FinalLayoutContractError( + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + f"{target.name!r} uses unsupported buffer dtype {target.tensor.dtype}", + ) + + def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None: + if not any(target.kind is TensorKind.PARAMETER and target.tensor.dtype is torch.bfloat16 for target in targets): + raise FinalLayoutContractError( + FinalLayoutContractCode.DTYPE_UNSUPPORTED, + "final-layout BF16 representation requires at least one BF16 parameter", + ) + + def build_format_metadata( + self, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> CanonicalJson: + return CanonicalJson.from_value( + { + "component_names": list(component_names), + "format": FINAL_LAYOUT_BF16_REPRESENTATION, + "mixed_precision_policy": "bf16-with-preserved-fp32", + "tensor_contract_sha256": tensor_contract_digest, + "tensor_count": tensor_count, + "tensor_layout": "contiguous", + } + ) + + def validate_format_metadata( + self, + metadata: CanonicalJson, + *, + component_names: tuple[str, ...], + tensor_contract_digest: str, + tensor_count: int, + ) -> None: + expected = self.build_format_metadata( + component_names=component_names, + tensor_contract_digest=tensor_contract_digest, + tensor_count=tensor_count, + ) + if metadata != expected: + raise ValueError("lease format metadata differs from the BF16 artifact policy") + + +FINAL_LAYOUT_BF16_POLICY = FinalLayoutBF16Policy() + + +def _split_shards( + records: Sequence[RuntimeTensorTarget], + max_shard_bytes: int, +) -> tuple[tuple[RuntimeTensorTarget, ...], ...]: + if max_shard_bytes <= 0: + raise ValueError("final-layout BF16 shard size must be positive") + shards: list[tuple[RuntimeTensorTarget, ...]] = [] + current: list[RuntimeTensorTarget] = [] + current_bytes = 0 + for record in records: + if current and current_bytes + record.nbytes > max_shard_bytes: + shards.append(tuple(current)) + current = [] + current_bytes = 0 + current.append(record) + current_bytes += record.nbytes + if current: + shards.append(tuple(current)) + return tuple(shards) + + +class FinalLayoutBF16Producer: + """Publish finalized BF16-policy tensors through a store-scoped writer.""" + + def __init__( + self, + context: FinalLayoutIdentityContext, + pipeline: nn.Module, + dit_modules: Sequence[tuple[str, nn.Module]], + *, + max_shard_bytes: int = DEFAULT_SHARD_SIZE_BYTES, + ) -> None: + if max_shard_bytes <= 0: + raise ValueError("final-layout BF16 shard size must be positive") + if context.spec != FINAL_LAYOUT_BF16_SPEC: + raise ValueError("producer requires the exact BF16 artifact specification") + self._context = context + self._pipeline = pipeline + self._dit_modules = tuple(dit_modules) + self._max_shard_bytes = max_shard_bytes + if tuple(name for name, _ in self._dit_modules) != context.dit_names: + raise ValueError("producer DiT components differ from the exact identity context") + if type(pipeline) is not context.pipeline_type or tuple(type(module) for _, module in self._dit_modules) != ( + context.dit_types + ): + raise ValueError("producer model implementation differs from the exact identity context") + self._spec = WeightProductionSpec( + producer_id=FINAL_LAYOUT_BF16_PRODUCER_ID, + outputs=(context.identity,), + source_mode=ProductionSourceMode.FINALIZED_MODEL, + coordination_scope=CoordinationScope.SINGLE_PROCESS, + lookup_phase=LookupPhase.POST_LOAD_ONLY, + ) + + @property + def spec(self) -> WeightProductionSpec: + return self._spec + + def produce(self, writer: ArtifactWriter) -> ProductionMetadata: + try: + return self._produce(writer) + except FinalLayoutContractError as exc: + raise final_layout_producer_error(exc) from exc + + def _produce(self, writer: ArtifactWriter) -> ProductionMetadata: + self._context.ensure_sources_unchanged() + bindings = validate_final_layout_model_contract( + self._dit_modules, + expected_schema=self._context.spec.model_contract_schema, + ) + records = collect_final_layout_targets( + self._pipeline, + self._dit_modules, + policy=FINAL_LAYOUT_BF16_POLICY, + require_materialized=True, + ) + contract_digest = validate_final_layout_identity(self._context, records) + for binding in bindings: + binding.validator() + + shards = _split_shards(records, self._max_shard_bytes) + shard_count = len(shards) + for index, shard in enumerate(shards, start=1): + specs = tuple( + TensorWriteSpec( + name=record.name, + shape=tuple(record.tensor.shape), + dtype=record.tensor.dtype, + kind=record.kind, + role=record.role, + ) + for record in shard + ) + file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors" + with writer.open_tensor_file(file_name, specs, ordered=True) as output: + for record in shard: + output.write_tensor(record.name, record.tensor.detach()) + + self._context.ensure_sources_unchanged() + return ProductionMetadata( + producer_schema=FINAL_LAYOUT_BF16_MANIFEST_SCHEMA, + restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, + format_metadata=FINAL_LAYOUT_BF16_POLICY.build_format_metadata( + component_names=self._context.dit_names, + tensor_contract_digest=contract_digest, + tensor_count=len(records), + ), + ) + + +__all__ = [ + "DEFAULT_SHARD_SIZE_BYTES", + "FINAL_LAYOUT_BF16_MANIFEST_SCHEMA", + "FINAL_LAYOUT_BF16_POLICY", + "FINAL_LAYOUT_BF16_PRODUCER_ID", + "FINAL_LAYOUT_BF16_REPRESENTATION", + "FINAL_LAYOUT_BF16_SPEC", + "FINAL_LAYOUT_BF16_VERSION", + "FinalLayoutBF16Policy", + "FinalLayoutBF16Producer", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/restorer.py b/vllm_omni/diffusion/model_loader/host_weights/restorer.py new file mode 100644 index 000000000..04581172f --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/restorer.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generic validation-first restorer for exact final-layout tensor artifacts.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import nn + +from vllm_omni.host_weight_runtime import HostWeightLease, TensorKind, WeightRestorePlan + +from .identity_adapter import ( + FinalLayoutIdentityContext, + validate_final_layout_identity, +) +from .tensor_layout import collect_final_layout_targets, validate_final_layout_model_contract + + +@dataclass(frozen=True) +class _Replacement: + name: str + parent: nn.Module + leaf_name: str + target: torch.Tensor + source: torch.Tensor + kind: TensorKind + + def current_target(self) -> torch.Tensor | None: + if self.kind is TensorKind.PARAMETER: + return self.parent._parameters.get(self.leaf_name) + return self.parent._buffers.get(self.leaf_name) + + +class FinalLayoutTensorRestorePlan(WeightRestorePlan): + """One-shot exact tensor rebinding validated before the first mutation.""" + + def __init__( + self, + lease: HostWeightLease, + replacements: tuple[_Replacement, ...], + validators: tuple[Callable[[], None], ...], + source_guard: Callable[[], None], + ) -> None: + self._lease = lease + self._replacements = replacements + self._validators = validators + self._source_guard = source_guard + self._committed = False + + def commit(self) -> None: + if self._committed: + raise RuntimeError("final-layout tensor restore plan was already committed") + if self._lease.closed: + raise RuntimeError("cannot commit a final-layout tensor restore from a closed lease") + self._source_guard() + + for replacement in self._replacements: + current = replacement.current_target() + if current is not replacement.target: + raise RuntimeError(f"restore target {replacement.name!r} changed after preflight") + if ( + tuple(current.shape) != tuple(replacement.source.shape) + or current.dtype != replacement.source.dtype + or tuple(current.stride()) != tuple(replacement.source.stride()) + or not replacement.source.is_contiguous() + ): + raise RuntimeError(f"restore target {replacement.name!r} changed layout after preflight") + + # Mark first: assignment or model-validator failure makes the plan + # permanently non-retryable and the partially restored model disposable. + self._committed = True + for replacement in self._replacements: + if replacement.target.is_meta: + if replacement.kind is TensorKind.PARAMETER: + replacement.parent._parameters[replacement.leaf_name] = nn.Parameter( + replacement.source, + requires_grad=replacement.target.requires_grad, + ) + else: + replacement.parent._buffers[replacement.leaf_name] = replacement.source + else: + replacement.target.data = replacement.source + for validator in self._validators: + validator() + + +class FinalLayoutTensorRestorer: + """Validate and plan restoration for one exact policy-defined lease.""" + + def __init__(self, context: FinalLayoutIdentityContext) -> None: + if not context.dit_names or any(not name for name in context.dit_names): + raise ValueError("final-layout restorer requires DiT component names") + self._context = context + + @property + def schema(self) -> str: + return self._context.spec.producer.restorer_schema + + def plan_restore(self, model: object, lease: HostWeightLease) -> FinalLayoutTensorRestorePlan: + if not isinstance(model, nn.Module): + raise TypeError("final-layout restoration requires an nn.Module pipeline") + if type(model) is not self._context.pipeline_type: + raise ValueError("restore model implementation differs from the exact identity context") + if lease.closed: + raise ValueError("cannot restore from a closed HostWeightLease") + if lease.identity.canonical_bytes != self._context.identity.canonical_bytes: + raise ValueError("lease semantic identity differs from the exact restore request") + if lease.manifest.identity.canonical_bytes != self._context.identity.canonical_bytes: + raise ValueError("lease manifest identity differs from the exact restore request") + if lease.identity.producer.restorer_schema != self.schema or lease.manifest.restorer_schema != self.schema: + raise ValueError("lease restorer schema is incompatible with the exact restore request") + self._context.ensure_sources_unchanged() + + try: + dit_modules = tuple((name, model.get_submodule(name)) for name in self._context.dit_names) + except AttributeError as exc: + raise ValueError("one or more lease-owned DiT modules do not exist in the pipeline") from exc + if tuple(type(module) for _, module in dit_modules) != self._context.dit_types: + raise ValueError("restore DiT implementation differs from the exact identity context") + bindings = validate_final_layout_model_contract( + dit_modules, + expected_schema=self._context.spec.model_contract_schema, + ) + targets = collect_final_layout_targets( + model, + dit_modules, + policy=self._context.policy, + require_materialized=False, + ) + contract_digest = validate_final_layout_identity(self._context, targets) + + manifest_entries = {entry.name: entry for entry in lease.manifest.tensors} + target_names = {target.name for target in targets} + if set(manifest_entries) != target_names or target_names != self._context.tensor_names: + missing = sorted(target_names - set(manifest_entries)) + unexpected = sorted(set(manifest_entries) - target_names) + raise ValueError( + "lease tensor coverage differs from the structurally initialized DiT: " + f"missing={missing[:5]}, unexpected={unexpected[:5]}" + ) + self._context.policy.validate_format_metadata( + lease.manifest.format_metadata, + component_names=self._context.dit_names, + tensor_contract_digest=contract_digest, + tensor_count=len(manifest_entries), + ) + + replacements: list[_Replacement] = [] + source_storages: dict[tuple[int, int], str] = {} + for target in targets: + entry = manifest_entries[target.name] + source = lease.tensors[target.name] + if entry.kind is not target.kind or entry.role != target.role: + raise ValueError(f"lease tensor kind or role differs for {target.name!r}") + if ( + tuple(source.shape) != tuple(target.tensor.shape) + or source.dtype != target.tensor.dtype + or tuple(source.stride()) != tuple(target.tensor.stride()) + or not source.is_contiguous() + or source.device.type != "cpu" + ): + raise ValueError(f"lease tensor layout differs for {target.name!r}") + if source.numel(): + storage = source.untyped_storage() + storage_id = (storage.data_ptr(), storage.nbytes()) + owner = source_storages.setdefault(storage_id, target.name) + if owner != target.name: + raise ValueError(f"lease tensor {target.name!r} aliases {owner!r}") + parent_path, _, leaf_name = target.name.rpartition(".") + replacements.append( + _Replacement( + name=target.name, + parent=model.get_submodule(parent_path), + leaf_name=leaf_name, + target=target.tensor, + source=source, + kind=target.kind, + ) + ) + + self._context.ensure_sources_unchanged() + return FinalLayoutTensorRestorePlan( + lease, + tuple(replacements), + tuple(binding.validator for binding in bindings), + self._context.ensure_sources_unchanged, + ) + + +__all__ = [ + "FinalLayoutTensorRestorePlan", + "FinalLayoutTensorRestorer", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/source_identity.py b/vllm_omni/diffusion/model_loader/host_weights/source_identity.py new file mode 100644 index 000000000..5591928c8 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/source_identity.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Canonical diffusion weight-source identity and one-load replacement guard.""" + +from __future__ import annotations + +import hashlib +import json +import os +import stat +import time +import uuid +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from enum import Enum +from pathlib import Path + +import regex as re + +from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity +from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, FileLockTimeoutError +from vllm_omni.host_weight_runtime.identity import canonical_json + +from .contracts import ( + FinalLayoutContractCode, + FinalLayoutContractError, + ImplementationIdentity, +) + +_HASH_CHUNK_BYTES = 8 * 1024**2 +_IMMUTABLE_REVISION_RE = re.compile(r"(?:[0-9a-fA-F]{40}|[0-9a-fA-F]{64})\Z") + + +class WeightSourceKind(str, Enum): + LOCAL_PATH = "local_path" + HUGGING_FACE_HUB = "hugging_face_hub" + + +@dataclass(frozen=True) +class PreparedWeightSource: + """Canonical source files and adapter ABI selected by the ordinary loader.""" + + model_or_path: str + subfolder: str | None + requested_revision: str | None + prefix: str + resolved_root: Path + weight_files: tuple[Path, ...] + use_safetensors: bool + checkpoint_adapter: ImplementationIdentity | None = None + source_kind: WeightSourceKind = WeightSourceKind.LOCAL_PATH + + def __post_init__(self) -> None: + if not self.model_or_path or not isinstance(self.prefix, str): + raise ValueError("prepared source model and prefix must be valid strings") + if not isinstance(self.resolved_root, Path): + raise ValueError("prepared source root must use pathlib.Path") + if not isinstance(self.weight_files, tuple) or not self.weight_files: + raise ValueError("prepared source must contain an immutable tuple of weight files") + if any(not isinstance(path, Path) for path in self.weight_files): + raise ValueError("prepared source weight files must use pathlib.Path") + if not isinstance(self.use_safetensors, bool): + raise ValueError("prepared source safetensors mode must be a boolean") + if self.checkpoint_adapter is not None and not isinstance( + self.checkpoint_adapter, + ImplementationIdentity, + ): + raise ValueError("checkpoint adapter must use ImplementationIdentity") + if not isinstance(self.source_kind, WeightSourceKind): + raise ValueError("prepared source kind must use WeightSourceKind") + + +@dataclass(frozen=True) +class _FileSnapshot: + path: Path + relative_name: str + size: int + device: int + inode: int + mtime_ns: int + ctime_ns: int + symlink_target: str | None + content_id: str + + def semantic_dict(self) -> dict[str, object]: + return { + "relative_name": self.relative_name, + "size": self.size, + "content_id": self.content_id, + } + + def unchanged(self) -> bool: + try: + current = self.path.stat() + symlink_target = os.readlink(self.path) if self.path.is_symlink() else None + except OSError: + return False + return ( + current.st_size == self.size + and current.st_dev == self.device + and current.st_ino == self.inode + and current.st_mtime_ns == self.mtime_ns + and current.st_ctime_ns == self.ctime_ns + and symlink_target == self.symlink_target + ) + + +@dataclass(frozen=True) +class _FileState: + size: int + device: int + inode: int + mtime_ns: int + ctime_ns: int + symlink_target: str | None + + def to_dict(self) -> dict[str, object]: + return { + "size": self.size, + "device": self.device, + "inode": self.inode, + "mtime_ns": self.mtime_ns, + "ctime_ns": self.ctime_ns, + "symlink_target": self.symlink_target, + } + + +def _observe(path: Path) -> _FileState: + try: + current = path.stat() + except OSError as exc: + raise ValueError(f"cannot stat canonical weight file {path}") from exc + if not stat.S_ISREG(current.st_mode): + raise ValueError(f"canonical weight source is not a regular file: {path}") + return _FileState( + size=current.st_size, + device=current.st_dev, + inode=current.st_ino, + mtime_ns=current.st_mtime_ns, + ctime_ns=current.st_ctime_ns, + symlink_target=os.readlink(path) if path.is_symlink() else None, + ) + + +def _sha256_file(path: Path, expected: _FileState) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + before = os.fstat(handle.fileno()) + while chunk := handle.read(_HASH_CHUNK_BYTES): + digest.update(chunk) + after = os.fstat(handle.fileno()) + observed = (before.st_size, before.st_dev, before.st_ino, before.st_mtime_ns, before.st_ctime_ns) + completed = (after.st_size, after.st_dev, after.st_ino, after.st_mtime_ns, after.st_ctime_ns) + expected_stat = (expected.size, expected.device, expected.inode, expected.mtime_ns, expected.ctime_ns) + if observed != expected_stat or completed != expected_stat or _observe(path) != expected: + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_CHANGED, + f"canonical weight file changed while hashing: {path}", + ) + return digest.hexdigest() + + +class NodeSourceDigestCache: + """Share content-backed local-source fingerprints across node workers.""" + + def __init__(self, root: Path, *, timeout_seconds: float) -> None: + self.entries = root / "source-digests-v1" / "entries" + self.locks = root / "source-digests-v1" / "locks" + self.timeout_seconds = timeout_seconds + self._available = True + try: + self.entries.mkdir(parents=True, exist_ok=True, mode=0o700) + self.locks.mkdir(parents=True, exist_ok=True, mode=0o700) + except OSError: + self._available = False + + @staticmethod + def _cached_content_id(entry: Path, path: Path, state: _FileState) -> str | None: + try: + cached = json.loads(entry.read_bytes()) + if not isinstance(cached, dict): + return None + cached_state = cached["state"] + content_id = cached["content_id"] + record_sha256 = cached["record_sha256"] + except (FileNotFoundError, OSError, json.JSONDecodeError, KeyError, TypeError): + return None + record = { + "content_id": content_id, + "schema_version": cached.get("schema_version"), + "state": cached_state, + } + if ( + record["schema_version"] != 1 + or cached_state != state.to_dict() + or not isinstance(content_id, str) + or re.fullmatch(r"sha256:[0-9a-f]{64}", content_id) is None + or not isinstance(record_sha256, str) + or hashlib.sha256(canonical_json(record)).hexdigest() != record_sha256 + or _observe(path) != state + ): + return None + return content_id + + def content_id(self, path: Path, state: _FileState) -> str: + if not self._available: + return f"sha256:{_sha256_file(path, state)}" + key = hashlib.sha256(os.fsencode(path.absolute())).hexdigest() + entry = self.entries / f"{key}.json" + content_id: str | None = None + try: + with FileLock( + self.locks / f"{key}.lock", + exclusive=True, + deadline=time.monotonic() + self.timeout_seconds, + ): + if (cached := self._cached_content_id(entry, path, state)) is not None: + return cached + content_id = f"sha256:{_sha256_file(path, state)}" + record: dict[str, object] = { + "content_id": content_id, + "schema_version": 1, + "state": state.to_dict(), + } + record["record_sha256"] = hashlib.sha256(canonical_json(record)).hexdigest() + temporary = entry.with_name(f".{entry.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp") + try: + temporary.write_bytes(canonical_json(record)) + os.chmod(temporary, 0o600) + os.replace(temporary, entry) + finally: + temporary.unlink(missing_ok=True) + return content_id + except (FileLockTimeoutError, OSError): + # Digest sharing is an optimization. Coordination or cache I/O + # failures must not make a readable canonical source unavailable. + if content_id is not None and _observe(path) == state: + return content_id + return f"sha256:{_sha256_file(path, state)}" + + +@dataclass(frozen=True) +class _SourceSnapshot: + semantic: CanonicalJson + revision: str + files: tuple[_FileSnapshot, ...] + + def unchanged(self) -> bool: + return all(file.unchanged() for file in self.files) + + +@dataclass(frozen=True) +class FinalLayoutSourceContext: + """Resolved source identity plus process-local replacement observations.""" + + identity: WeightSourceIdentity + snapshots: tuple[_SourceSnapshot, ...] + + def sources_unchanged(self) -> bool: + return all(snapshot.unchanged() for snapshot in self.snapshots) + + def ensure_sources_unchanged(self) -> None: + if not self.sources_unchanged(): + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_CHANGED, + "canonical source changed after final-layout identity resolution", + ) + + +def _hf_blob_content_id( + path: Path, + *, + source: PreparedWeightSource, + source_root: Path, +) -> str | None: + if source.source_kind is not WeightSourceKind.HUGGING_FACE_HUB or not path.is_symlink(): + return None + if Path(source.model_or_path).expanduser().exists(): + return None + + parts = source_root.parts + snapshot_indexes = [index for index, part in enumerate(parts[:-1]) if part == "snapshots"] + if not snapshot_indexes: + return None + snapshot_index = snapshot_indexes[-1] + revision = parts[snapshot_index + 1] + if _IMMUTABLE_REVISION_RE.fullmatch(revision) is None: + return None + repo_root = Path(*parts[:snapshot_index]) + expected_repo_name = f"models--{source.model_or_path.replace('/', '--')}" + if repo_root.name != expected_repo_name: + return None + snapshot_root = repo_root / "snapshots" / revision + try: + path.relative_to(snapshot_root) + except ValueError: + return None + blobs_root = repo_root / "blobs" + try: + resolved_target = path.resolve(strict=True) + resolved_blobs_root = blobs_root.resolve(strict=True) + except OSError: + return None + if resolved_target.parent != resolved_blobs_root: + return None + blob_name = resolved_target.name + if _IMMUTABLE_REVISION_RE.fullmatch(blob_name) is None: + return None + return f"immutable-blob:{blob_name.lower()}" + + +def _snapshot_revision(path: Path) -> str | None: + parts = path.absolute().parts + for index, part in enumerate(parts[:-1]): + if part != "snapshots": + continue + candidate = parts[index + 1] + if _IMMUTABLE_REVISION_RE.fullmatch(candidate) is not None: + return candidate.lower() + return None + + +def _logical_model_id(value: str) -> str: + candidate = Path(value).expanduser() + try: + if candidate.exists(): + return str(candidate.resolve()) + except OSError: + pass + return value + + +def _snapshot_source( + source: PreparedWeightSource, + digest_cache: NodeSourceDigestCache | None, +) -> _SourceSnapshot: + root = source.resolved_root.expanduser().resolve() + if not root.is_dir(): + raise ValueError(f"canonical source root is not a directory: {root}") + immutable_revision = _snapshot_revision(root) + if immutable_revision is None and source.requested_revision is not None: + requested = source.requested_revision.strip() + if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None: + immutable_revision = requested.lower() + + def snapshot_file(path: Path) -> _FileSnapshot: + candidate = path.expanduser() + if not candidate.is_absolute(): + candidate = root / candidate + candidate = candidate.absolute() + try: + relative_name = candidate.relative_to(root).as_posix() + except ValueError as exc: + raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc + current = _observe(candidate) + content_id = _hf_blob_content_id( + candidate, + source=source, + source_root=root, + ) + if content_id is None: + content_id = ( + digest_cache.content_id(candidate, current) + if digest_cache is not None + else f"sha256:{_sha256_file(candidate, current)}" + ) + return _FileSnapshot( + path=candidate, + relative_name=relative_name, + size=current.size, + device=current.device, + inode=current.inode, + mtime_ns=current.mtime_ns, + ctime_ns=current.ctime_ns, + symlink_target=current.symlink_target, + content_id=content_id, + ) + + paths = sorted(source.weight_files, key=lambda item: str(item)) + with ThreadPoolExecutor(max_workers=min(8, len(paths))) as executor: + files = list(executor.map(snapshot_file, paths)) + + file_semantics = [file.semantic_dict() for file in files] + file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest() + revision = immutable_revision or f"content-{file_fingerprint}" + semantic = CanonicalJson.from_value( + { + "checkpoint_adapter": ( + source.checkpoint_adapter.to_dict() if source.checkpoint_adapter is not None else None + ), + "files": file_semantics, + "model_or_path": _logical_model_id(source.model_or_path), + "prefix": source.prefix, + "resolved_revision": revision, + "source_kind": source.source_kind.value, + "subfolder": source.subfolder, + "use_safetensors": source.use_safetensors, + } + ) + return _SourceSnapshot(semantic=semantic, revision=revision, files=tuple(files)) + + +def _resolve_target_sources( + prepared_sources: Sequence[PreparedWeightSource], + target_names: frozenset[str], +) -> dict[str, PreparedWeightSource]: + """Bind every target to one source using deterministic longest-prefix wins.""" + if not target_names: + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_COVERAGE_INVALID, + "final-layout source identity requires at least one owned tensor", + ) + + bindings: dict[str, PreparedWeightSource] = {} + for target_name in sorted(target_names): + matches = [source for source in prepared_sources if target_name.startswith(source.prefix)] + if not matches: + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_COVERAGE_INVALID, + f"no canonical weight source covers final-layout tensor {target_name!r}", + ) + longest_prefix = max(len(source.prefix) for source in matches) + winners = [source for source in matches if len(source.prefix) == longest_prefix] + if len(winners) != 1: + candidates = sorted( + f"{source.model_or_path}:{source.subfolder or ''}:{source.prefix}" for source in winners + ) + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_COVERAGE_INVALID, + f"multiple equally specific canonical sources cover {target_name!r}: {candidates}", + ) + bindings[target_name] = winners[0] + return bindings + + +def resolve_final_layout_source_identity( + prepared_sources: Sequence[PreparedWeightSource], + *, + model_id: str, + target_names: frozenset[str], + digest_cache: NodeSourceDigestCache | None = None, +) -> FinalLayoutSourceContext: + """Resolve exact source identity for files covering the owned tensors.""" + bindings = _resolve_target_sources(prepared_sources, target_names) + selected_sources = frozenset(bindings.values()) + snapshot_by_source = {source: _snapshot_source(source, digest_cache) for source in selected_sources} + snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded)) + source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots] + target_bindings = [ + { + "source_fingerprint": hashlib.sha256(snapshot_by_source[source].semantic.encoded).hexdigest(), + "source_prefix": source.prefix, + "target_name": target_name, + } + for target_name, source in sorted(bindings.items()) + ] + source_document = { + "sources": source_semantics, + "target_bindings": target_bindings, + } + source_fingerprint = hashlib.sha256(canonical_json(source_document)).hexdigest() + unique_revisions = tuple(sorted({snapshot.revision for snapshot in snapshots})) + aggregate_revision = ( + unique_revisions[0] + if len(unique_revisions) == 1 + else f"aggregate-{hashlib.sha256(canonical_json(unique_revisions)).hexdigest()}" + ) + return FinalLayoutSourceContext( + identity=WeightSourceIdentity( + model_id=_logical_model_id(model_id), + revision=aggregate_revision, + fingerprint=source_fingerprint, + metadata=CanonicalJson.from_value(source_document), + ), + snapshots=snapshots, + ) + + +__all__ = [ + "FinalLayoutSourceContext", + "NodeSourceDigestCache", + "PreparedWeightSource", + "WeightSourceKind", + "resolve_final_layout_source_identity", +] diff --git a/vllm_omni/diffusion/model_loader/host_weights/tensor_layout.py b/vllm_omni/diffusion/model_loader/host_weights/tensor_layout.py new file mode 100644 index 000000000..d8759417c --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weights/tensor_layout.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Representation-neutral ownership discovery for final-layout tensors.""" + +from __future__ import annotations + +import hashlib +from collections.abc import Callable, Iterator, Sequence +from dataclasses import dataclass +from itertools import chain + +import torch +from torch import nn + +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract +from vllm_omni.host_weight_runtime import TensorKind +from vllm_omni.host_weight_runtime.identity import canonical_json + +from .contracts import ( + FinalLayoutContractCode, + FinalLayoutContractError, + FinalLayoutTensorPolicy, +) + + +@dataclass(frozen=True) +class ModelRestoreBinding: + """Model-declared restore ABI and its post-commit validator.""" + + name: str + contract: FinalLayoutModelContract + validator: Callable[[], None] + + +@dataclass(frozen=True) +class RuntimeTensorTarget: + """One complete final-layout tensor owned by a DiT component.""" + + name: str + tensor: torch.Tensor + kind: TensorKind + role: str + + @property + def nbytes(self) -> int: + return self.tensor.numel() * self.tensor.element_size() + + +def validate_final_layout_model_contract( + dit_modules: Sequence[tuple[str, nn.Module]], + *, + expected_schema: str, +) -> tuple[ModelRestoreBinding, ...]: + """Require an explicit tensor-complete restore ABI from every DiT.""" + if not dit_modules: + raise FinalLayoutContractError( + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + "final-layout artifacts require at least one DiT component", + ) + + bindings: list[ModelRestoreBinding] = [] + seen_names: set[str] = set() + for name, module in dit_modules: + if not name or name in seen_names: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"DiT component name {name!r} is empty or duplicated", + ) + seen_names.add(name) + contract = getattr(module, "host_weight_restore_contract", None) + if not isinstance(contract, FinalLayoutModelContract) or contract.schema != expected_schema: + raise FinalLayoutContractError( + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + f"DiT component {name!r} does not declare restore schema {expected_schema!r}", + ) + validator = getattr(module, "validate_restored_host_weights", None) + if not callable(validator): + raise FinalLayoutContractError( + FinalLayoutContractCode.MODEL_CONTRACT_UNSUPPORTED, + f"DiT component {name!r} does not implement validate_restored_host_weights()", + ) + bindings.append(ModelRestoreBinding(name, contract, validator)) + return tuple(bindings) + + +def _named_parameters_with_duplicates(module: nn.Module) -> Iterator[tuple[str, nn.Parameter]]: + try: + return module.named_parameters(remove_duplicate=False) + except TypeError: # pragma: no cover - compatibility with old torch + return module.named_parameters() + + +def _named_buffers_with_duplicates(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]: + try: + return module.named_buffers(remove_duplicate=False) + except TypeError: # pragma: no cover - compatibility with old torch + return module.named_buffers() + + +def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool: + parent_path, _, leaf_name = local_name.rpartition(".") + owner = module.get_submodule(parent_path) + return leaf_name not in owner._non_persistent_buffers_set + + +def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None: + parent_path, _, leaf_name = runtime_name.rpartition(".") + parent = pipeline.get_submodule(parent_path) + tensor = parent._parameters.get(leaf_name) + if tensor is None: + tensor = parent._buffers.get(leaf_name) + return tensor + + +def collect_final_layout_targets( + pipeline: nn.Module, + dit_modules: Sequence[tuple[str, nn.Module]], + *, + policy: FinalLayoutTensorPolicy, + require_materialized: bool, +) -> tuple[RuntimeTensorTarget, ...]: + """Return complete, alias-free targets validated by one representation policy.""" + validate_final_layout_model_contract( + dit_modules, + expected_schema=policy.spec.model_contract_schema, + ) + records: dict[str, RuntimeTensorTarget] = {} + for dit_name, dit_module in dit_modules: + for local_name, tensor in _named_parameters_with_duplicates(dit_module): + runtime_name = f"{dit_name}.{local_name}" + candidate = RuntimeTensorTarget( + runtime_name, + tensor, + TensorKind.PARAMETER, + policy.tensor_role(runtime_name, tensor, TensorKind.PARAMETER), + ) + existing = records.get(runtime_name) + if existing is not None and existing.tensor is not tensor: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"multiple DiT tensors resolve to {runtime_name!r}", + ) + records[runtime_name] = candidate + for local_name, tensor in _named_buffers_with_duplicates(dit_module): + if not _is_persistent_buffer(dit_module, local_name): + continue + runtime_name = f"{dit_name}.{local_name}" + candidate = RuntimeTensorTarget( + runtime_name, + tensor, + TensorKind.BUFFER, + policy.tensor_role(runtime_name, tensor, TensorKind.BUFFER), + ) + existing = records.get(runtime_name) + if existing is not None and existing.tensor is not tensor: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"multiple DiT tensors resolve to {runtime_name!r}", + ) + records[runtime_name] = candidate + + if not records: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + "no DiT parameters or persistent buffers were discovered", + ) + + object_owners: dict[int, str] = {} + storage_owners: dict[tuple[int, int], str] = {} + for record in records.values(): + tensor = record.tensor + if not record.role: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} has an empty semantic role", + ) + try: + pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name) + except AttributeError as exc: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"DiT tensor {record.name!r} is not owned by the pipeline", + ) from exc + if pipeline_tensor is not tensor: + raise FinalLayoutContractError( + FinalLayoutContractCode.OWNERSHIP_AMBIGUOUS, + f"DiT tensor {record.name!r} does not resolve to the discovered object", + ) + + object_owner = object_owners.setdefault(id(tensor), record.name) + if object_owner != record.name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} aliases tensor object {object_owner!r}", + ) + if hasattr(tensor, "to_local"): + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} is a distributed tensor", + ) + if tensor.layout != torch.strided: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} uses unsupported layout {tensor.layout}", + ) + if tensor.device.type not in {"cpu", "meta"}: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} must be a CPU or meta tensor, got {tensor.device}", + ) + if not tensor.is_contiguous(): + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} is non-contiguous with stride {tensor.stride()}", + ) + policy.validate_target(record) + + if not require_materialized: + continue + if tensor.device.type != "cpu" or tensor.is_meta: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}", + ) + storage = tensor.untyped_storage() + if tensor.storage_offset() != 0 or storage.nbytes() != record.nbytes: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} is a view into a larger storage", + ) + if record.nbytes: + storage_id = (storage.data_ptr(), storage.nbytes()) + storage_owner = storage_owners.setdefault(storage_id, record.name) + if storage_owner != record.name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"{record.name!r} shares storage with {storage_owner!r}", + ) + + for pipeline_name, tensor in chain( + _named_parameters_with_duplicates(pipeline), + _named_buffers_with_duplicates(pipeline), + ): + owner = object_owners.get(id(tensor)) + if owner is not None and owner != pipeline_name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"cached tensor {owner!r} aliases pipeline tensor {pipeline_name!r}", + ) + if not require_materialized or tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0: + continue + storage = tensor.untyped_storage() + pipeline_storage_owner = storage_owners.get((storage.data_ptr(), storage.nbytes())) + if pipeline_storage_owner is not None and pipeline_storage_owner != pipeline_name: + raise FinalLayoutContractError( + FinalLayoutContractCode.TENSOR_UNSUPPORTED, + f"cached tensor {pipeline_storage_owner!r} shares storage with pipeline tensor {pipeline_name!r}", + ) + + ordered = tuple(records[name] for name in sorted(records)) + policy.validate_collection(ordered) + return ordered + + +def tensor_contract_sha256(records: Sequence[RuntimeTensorTarget]) -> str: + """Hash exact structural ownership and representation-specific tensor roles.""" + contract = [ + { + "dtype": str(record.tensor.dtype), + "kind": record.kind.value, + "name": record.name, + "role": record.role, + "shape": list(record.tensor.shape), + "stride": list(record.tensor.stride()), + } + for record in records + ] + return hashlib.sha256(canonical_json(contract)).hexdigest() + + +__all__ = [ + "ModelRestoreBinding", + "RuntimeTensorTarget", + "collect_final_layout_targets", + "tensor_contract_sha256", + "validate_final_layout_model_contract", +] diff --git a/vllm_omni/diffusion/model_loader/hub_prefetch.py b/vllm_omni/diffusion/model_loader/hub_prefetch.py index 9a9295ab9..3594aa4e1 100644 --- a/vllm_omni/diffusion/model_loader/hub_prefetch.py +++ b/vllm_omni/diffusion/model_loader/hub_prefetch.py @@ -265,6 +265,7 @@ def prefetch_subfolders( subfolders: Iterable[str], *, local_files_only: bool | None = None, + revision: str | None = None, include_root_metadata: bool = True, ) -> None: """Materialise ``model``'s ``subfolders`` in the HF cache before loading. @@ -278,6 +279,7 @@ def prefetch_subfolders( local_files_only: When ``True``, skip the prefetch entirely. When ``None`` (default), auto-detect: skip if *model* is a local directory, run otherwise. + revision: Optional Hub revision shared with component loaders. include_root_metadata: When True, also pull ``*.json`` at the repo root so ``model_index.json`` / ``config.json`` resolution during ``from_pretrained`` also hits a warm cache. @@ -334,6 +336,7 @@ def prefetch_subfolders( with _repo_prefetch_lock(model): snapshot_download( repo_id=model, + revision=revision, allow_patterns=allow_patterns, ) logger.info("Prefetch complete for %s", model) @@ -450,6 +453,7 @@ def from_pretrained_with_prefetch( on the first failure exactly as before. """ prefetch_list = list(prefetch_list) + revision = from_pretrained_kwargs.get("revision") can_heal = not local_files_only and bool(model) and not os.path.isdir(model) last_exc: BaseException | None = None @@ -481,7 +485,7 @@ def from_pretrained_with_prefetch( # Force a fresh, verified snapshot of every component this pipeline # needs - not just ``subfolder`` - so a sibling component that was # also half-written gets repaired in the same pass. - prefetch_subfolders(model, prefetch_list, local_files_only=False) + prefetch_subfolders(model, prefetch_list, local_files_only=False, revision=revision) assert last_exc is not None # loop only exits via return or a caught exc raise last_exc diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index ecc53cdc6..1665d6a16 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import dataclass @@ -12,6 +12,7 @@ class DiffusionModelMetadata: max_multimodal_image_inputs: int | None = None supports_mixed_reference_inputs: bool = False attention_mask_free: bool = False + final_output_type: str | None = None QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4 @@ -41,6 +42,7 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { supports_multimodal_inputs=True, max_multimodal_image_inputs=9, supports_mixed_reference_inputs=True, + final_output_type="video", # H3 represents alignment padding as a second packed sequence. The # packed TRTLLM backend consumes cu_seqlens and isolates that padding. attention_mask_free=True, @@ -51,11 +53,52 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { supports_multimodal_inputs=True, max_multimodal_image_inputs=9, supports_mixed_reference_inputs=True, + final_output_type="video", ), - "WanPipeline": DiffusionModelMetadata(attention_mask_free=True), - "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True), - "WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True), - "WanS2VPipeline": DiffusionModelMetadata(attention_mask_free=True), + "WanPipeline": DiffusionModelMetadata( + attention_mask_free=True, + final_output_type="video", + ), + "WanImageToVideoPipeline": DiffusionModelMetadata( + attention_mask_free=True, + final_output_type="video", + ), + "WanVACEPipeline": DiffusionModelMetadata( + attention_mask_free=True, + final_output_type="video", + ), + "WanS2VPipeline": DiffusionModelMetadata( + attention_mask_free=True, + final_output_type="video", + ), + "WanT2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"), + "WanI2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"), + "LTX2Pipeline": DiffusionModelMetadata(final_output_type="video"), + "LTX2DistilledPipeline": DiffusionModelMetadata(final_output_type="video"), + "LTX2T2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"), + "LTX2I2VDMD2Pipeline": DiffusionModelMetadata(final_output_type="video"), + "HeliosPipeline": DiffusionModelMetadata(final_output_type="video"), + "HeliosPyramidPipeline": DiffusionModelMetadata(final_output_type="video"), + "HunyuanVideo15Pipeline": DiffusionModelMetadata(final_output_type="video"), + "HunyuanVideo15ImageToVideoPipeline": DiffusionModelMetadata(final_output_type="video"), + "LingBotVideoPipeline": DiffusionModelMetadata(final_output_type="video"), + "LongCatVideoAvatarPipeline": DiffusionModelMetadata(final_output_type="video"), + "MagiHumanPipeline": DiffusionModelMetadata(final_output_type="video"), + "DreamIDOmniPipeline": DiffusionModelMetadata(final_output_type="video"), + "Cosmos3OmniDiffusersPipeline": DiffusionModelMetadata(final_output_type="video"), + "Cosmos3OmniPipeline": DiffusionModelMetadata(final_output_type="video"), + "SanaVideoPipeline": DiffusionModelMetadata(final_output_type="video"), + "SanaWmPipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=1, + ), +} + +_DIFFUSION_MODEL_METADATA_ALIASES = { + "WanDMDPipeline": "WanPipeline", + "LTX2TwoStagePipeline": "LTX2Pipeline", + "LTX2DistilledOneStagePipeline": "LTX2DistilledPipeline", + "LingBotWorldCausalDMDPipeline": "LingBotVideoPipeline", } @@ -67,6 +110,9 @@ def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModel metadata = _DIFFUSION_MODEL_METADATA.get(model_class_name) if metadata is not None: return metadata + canonical_name = _DIFFUSION_MODEL_METADATA_ALIASES.get(model_class_name) + if canonical_name is not None: + return _DIFFUSION_MODEL_METADATA[canonical_name] # Some checkpoints report the HF architecture name diff from internal pipeline class name # (e.g. HunyuanImage3ForCausalMM, WanVACEPipeline, OmniVoice ...). from vllm_omni.diffusion.registry import _DIFFUSION_MODELS @@ -80,5 +126,6 @@ def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModel # key spaces. Aliases whose pipeline class has no metadata entry (e.g. # Wan22VACEPipeline, OmniVoicePipeline) still fall back to the defaults; # that is not a regression, it just means no capability override. - return _DIFFUSION_MODEL_METADATA.get(pipeline_cls_name, DiffusionModelMetadata()) + canonical_name = _DIFFUSION_MODEL_METADATA_ALIASES.get(pipeline_cls_name, pipeline_cls_name) + return _DIFFUSION_MODEL_METADATA.get(canonical_name, DiffusionModelMetadata()) return DiffusionModelMetadata() diff --git a/vllm_omni/diffusion/models/audiox/__init__.py b/vllm_omni/diffusion/models/audiox/__init__.py deleted file mode 100644 index cec1470f2..000000000 --- a/vllm_omni/diffusion/models/audiox/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from .audiox_transformer import MMDiffusionTransformer -from .pipeline_audiox import AudioXPipeline, get_audiox_post_process_func - -__all__ = [ - "AudioXPipeline", - "MMDiffusionTransformer", - "get_audiox_post_process_func", -] diff --git a/vllm_omni/diffusion/models/audiox/audiox_transformer.py b/vllm_omni/diffusion/models/audiox/audiox_transformer.py deleted file mode 100644 index 9377f2b03..000000000 --- a/vllm_omni/diffusion/models/audiox/audiox_transformer.py +++ /dev/null @@ -1,487 +0,0 @@ -from __future__ import annotations - -import logging -from collections.abc import Iterable -from typing import Any - -import torch -import torch.nn as nn -from einops import rearrange -from torch.nn import functional as F -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_gather, - tensor_model_parallel_all_reduce, -) -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - MergedColumnParallelLinear, - QKVParallelLinear, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader, sharded_weight_loader - -from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.layers.fourier import GaussianFourierProjection -from vllm_omni.diffusion.layers.rope import RotaryEmbedding - - -class AudioXCrossAttention(nn.Module): - def __init__(self, dim: int, nheads: int, prefix: str = ""): - super().__init__() - self.dim = dim - self.nheads = nheads - head_dim = dim // nheads - self.head_dim = head_dim - - # to_kv bundle weights arrive in (head, dim, VK-index) interleaved order; pipeline's - # load_weights restacks them to [V|K] before MergedColumnParallelLinear consumes them. - self.to_q = ColumnParallelLinear(dim, dim, bias=False, gather_output=False, prefix=f"{prefix}.to_q") - self.to_kv = MergedColumnParallelLinear( - input_size=dim, output_sizes=[dim, dim], bias=False, gather_output=False, prefix=f"{prefix}.to_kv" - ) - self.q_norm = AudioXRMSNorm(head_dim) - self.k_norm = AudioXRMSNorm(head_dim) - local_nheads = nheads // get_tensor_model_parallel_world_size() - self.attn = Attention( - num_heads=local_nheads, - head_size=head_dim, - softmax_scale=head_dim**-0.5, - causal=False, - ) - - def forward(self, x: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor: - tp = get_tensor_model_parallel_world_size() - local_h = self.nheads // tp - d = self.head_dim - - # ROCm's ``wvSplitK`` GEMM bypasses autocast and rejects mismatched dtypes, - # so pre-match the input to the linear's weight dtype. - weight_dtype = self.to_q.weight.dtype - q_flat, _ = self.to_q(x.to(weight_dtype)) - q = rearrange(q_flat, "b n (h d) -> b n h d", h=local_h, d=d) - kv_flat, _ = self.to_kv(context.to(weight_dtype)) - v_flat, k_flat = kv_flat.chunk(2, dim=-1) - # Upstream `CrossAttention.forward` unpacks `pre_attention()` as `q, v, k = ...` (K/V - # swapped), AND only normalizes the first chunk of to_kv via k_norm. Trained weights - # depend on this quirk: first chunk normalized -> V; second chunk unnormalized -> K. - v = rearrange(v_flat, "b n (h d) -> b n h d", h=local_h, d=d) - k = rearrange(k_flat, "b n (h d) -> b n h d", h=local_h, d=d) - q = self.q_norm(q) - v = self.k_norm(v) - - out = self.attn(q.contiguous(), k.contiguous(), v.contiguous(), attn_metadata=None) - out = rearrange(out, "b n h d -> b n (h d)").contiguous() - if tp > 1: - out = tensor_model_parallel_all_gather(out, dim=-1) - return out - - -logger = logging.getLogger(__name__) - -__all__ = [ - "AudioXMMChannelLastConv1d", - "AudioXMMConvFeedForward", - "AudioXMMDiTSelfAttention", - "AudioXMMDiTBlock", - "ContinuousMMDiTTransformer", - "MMDiffusionTransformer", -] - - -class AudioXMMChannelLastConv1d(nn.Conv1d): - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = rearrange(x, "b n c -> b c n") - x = super().forward(x) - x = rearrange(x, "b c n -> b n c") - return x - - -class _ColumnParallelChannelLastConv1d(AudioXMMChannelLastConv1d): - def __init__(self, in_channels: int, out_channels_total: int, **kwargs: Any): - tp_size = get_tensor_model_parallel_world_size() - assert out_channels_total % tp_size == 0, (out_channels_total, tp_size) - super().__init__(in_channels, out_channels_total // tp_size, **kwargs) - self.weight.weight_loader = sharded_weight_loader(0) - - -class _RowParallelChannelLastConv1d(AudioXMMChannelLastConv1d): - def __init__(self, in_channels_total: int, out_channels: int, **kwargs: Any): - tp_size = get_tensor_model_parallel_world_size() - assert in_channels_total % tp_size == 0, (in_channels_total, tp_size) - super().__init__(in_channels_total // tp_size, out_channels, **kwargs) - self._tp_size = tp_size - self.weight.weight_loader = sharded_weight_loader(1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - y = super().forward(x) - if self._tp_size > 1: - y = tensor_model_parallel_all_reduce(y) - return y - - -class AudioXMMConvFeedForward(nn.Module): - def __init__( - self, - dim: int, - hidden_dim: int, - multiple_of: int = 256, - kernel_size: int = 3, - padding: int = 1, - ): - super().__init__() - hidden_dim = int(2 * hidden_dim / 3) - hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) - - self.w1 = _ColumnParallelChannelLastConv1d( - dim, hidden_dim, bias=False, kernel_size=kernel_size, padding=padding - ) - self.w2 = _RowParallelChannelLastConv1d(hidden_dim, dim, bias=False, kernel_size=kernel_size, padding=padding) - self.w3 = _ColumnParallelChannelLastConv1d( - dim, hidden_dim, bias=False, kernel_size=kernel_size, padding=padding - ) - - def forward(self, x): - return self.w2(F.silu(self.w1(x)) * self.w3(x)) - - -class AudioXRMSNorm(nn.Module): - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - super().__init__() - self.eps = eps - - def forward(self, x: torch.Tensor) -> torch.Tensor: - mean_sq = torch.mean(x**2, dim=-1, keepdim=True) - scale = torch.rsqrt(mean_sq + self.eps) - return x * scale - - -class AudioXMMDiTSelfAttention(nn.Module): - def __init__(self, dim: int, nheads: int, prefix: str = ""): - super().__init__() - self.dim = dim - self.nheads = nheads - - head_dim = dim // nheads - self.head_dim = head_dim - # Bundle weights arrive in interleaved (h, d, qkv) layout; pipeline's load_weights - # restacks to [Q|K|V] before QKVParallelLinear's weight_loader consumes them. - self.qkv = QKVParallelLinear( - hidden_size=dim, - head_size=head_dim, - total_num_heads=nheads, - bias=True, - prefix=f"{prefix}.qkv", - ) - self.q_norm = AudioXRMSNorm(head_dim) - self.k_norm = AudioXRMSNorm(head_dim) - - self.rope = RotaryEmbedding(is_neox_style=False) - self.attn = Attention( - num_heads=self.qkv.num_heads, - head_size=head_dim, - softmax_scale=head_dim**-0.5, - causal=False, - ) - - def apply_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - out = self.attn(q.contiguous(), k.contiguous(), v.contiguous(), attn_metadata=None) - out = rearrange(out, "b n h d -> b n (h d)").contiguous() - # Downstream linear1/ffn are replicated and expect full hidden dim. - if get_tensor_model_parallel_world_size() > 1: - out = tensor_model_parallel_all_gather(out, dim=-1) - return out - - def pre_attention(self, x: torch.Tensor, rot: tuple[torch.Tensor, torch.Tensor] | None = None): - # ROCm GEMM dtype cast — see AudioXCrossAttention.forward. - qkv, _ = self.qkv(x.to(self.qkv.weight.dtype)) - local_h = self.qkv.num_heads - d = self.head_dim - q_size = local_h * d - q, k, v = qkv.split([q_size, q_size, q_size], dim=-1) - q = rearrange(q, "b n (h d) -> b n h d", h=local_h, d=d) - k = rearrange(k, "b n (h d) -> b n h d", h=local_h, d=d) - v = rearrange(v, "b n (h d) -> b n h d", h=local_h, d=d) - q = self.q_norm(q) - k = self.k_norm(k) - - if rot is not None: - cos, sin = rot - cos = cos.to(dtype=q.dtype) - sin = sin.to(dtype=q.dtype) - q = self.rope(q, cos, sin) - k = self.rope(k, cos, sin) - - return q, k, v - - def forward( - self, - x: torch.Tensor, - rot: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> torch.Tensor: - q, k, v = self.pre_attention(x, rot=rot) - return self.apply_attention(q, k, v) - - -class AudioXMMDiTBlock(nn.Module): - def __init__( - self, - dim: int, - nhead: int, - mlp_ratio: float = 4.0, - prefix: str = "", - ): - super().__init__() - self.norm1 = nn.LayerNorm(dim, elementwise_affine=False) - self.attn = AudioXMMDiTSelfAttention(dim, nhead, prefix=f"{prefix}.attn") - self.cross_attn = AudioXCrossAttention(dim, nhead, prefix=f"{prefix}.cross_attn") - self.linear1 = AudioXMMChannelLastConv1d(dim, dim, kernel_size=3, padding=1) - self.norm2 = nn.LayerNorm(dim, elementwise_affine=False) - self.ffn = AudioXMMConvFeedForward(dim, int(dim * mlp_ratio), kernel_size=3, padding=1) - self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True)) - - def pre_attention(self, x: torch.Tensor, c: torch.Tensor, rot: tuple[torch.Tensor, torch.Tensor] | None): - modulation = self.adaLN_modulation(c) - (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = modulation.chunk(6, dim=-1) - x = self.norm1(x) * (1 + scale_msa) + shift_msa - q, k, v = self.attn.pre_attention(x, rot) - return (q, k, v), (gate_msa, shift_mlp, scale_mlp, gate_mlp) - - def post_attention(self, x: torch.Tensor, attn_out: torch.Tensor, c: tuple[torch.Tensor, ...], context=None): - (gate_msa, shift_mlp, scale_mlp, gate_mlp) = c - x = x + self.linear1(attn_out) * gate_msa - - x = x + self.cross_attn(x, context=context) - - r = self.norm2(x) * (1 + scale_mlp) + shift_mlp - x = x + self.ffn(r) * gate_mlp - return x - - def forward( - self, - x: torch.Tensor, - cond: torch.Tensor, - rot: tuple[torch.Tensor, torch.Tensor] | None, - context: torch.Tensor = None, - ) -> torch.Tensor: - x_qkv, x_conditions = self.pre_attention(x, cond, rot) - attn_out = self.attn.apply_attention(*x_qkv) - x = self.post_attention(x, attn_out, x_conditions, context=context) - return x - - -class MMDiffusionTransformer(nn.Module): - """AudioX MMDiT, specialized for the published bundle (`zhangj1an/AudioX`). - - The bundle fixes patch_size=1, transformer_type="continuous_transformer", - cond_token_dim=768 (>0, project_cond_tokens=False), and never sets - prepend_cond_dim or input_concat_dim, so those code paths are removed. - """ - - def __init__( - self, - io_channels: int, - embed_dim: int, - cond_token_dim: int, - global_cond_dim: int, - depth: int, - num_heads: int, - project_cond_tokens: bool = False, - project_global_cond: bool = True, - **kwargs, - ): - super().__init__() - if kwargs: - logger.debug("MMDiffusionTransformer ignoring unused config keys: %s", sorted(kwargs.keys())) - if project_cond_tokens: - raise ValueError("AudioX bundle requires project_cond_tokens=False to match official checkpoints.") - - self.cond_token_dim = cond_token_dim - - timestep_features_dim = 256 - self.timestep_features = GaussianFourierProjection( - in_features=1, - embedding_size=timestep_features_dim // 2, - scale=1.0, - trainable=False, - ) - self.to_timestep_embed = nn.Sequential( - nn.Linear(timestep_features_dim, embed_dim, bias=True), - nn.SiLU(), - nn.Linear(embed_dim, embed_dim, bias=True), - ) - - cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim - self.to_cond_embed = nn.Sequential( - nn.Linear(cond_token_dim, cond_embed_dim, bias=False), - nn.SiLU(), - nn.Linear(cond_embed_dim, cond_embed_dim, bias=False), - ) - - # ``to_global_embed`` weights live in the bundle but global conditioning is always None - # at inference; kept so AutoWeightsLoader has a slot to load them into. - global_embed_dim = global_cond_dim if not project_global_cond else embed_dim - self.to_global_embed = nn.Sequential( - nn.Linear(global_cond_dim, global_embed_dim, bias=False), - nn.SiLU(), - nn.Linear(global_embed_dim, global_embed_dim, bias=False), - ) - - self.transformer = ContinuousMMDiTTransformer( - dim=embed_dim, - depth=depth, - dim_heads=embed_dim // num_heads, - dim_in=io_channels, - dim_out=io_channels, - ) - - self.preprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) - nn.init.zeros_(self.preprocess_conv.weight) - self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False) - nn.init.zeros_(self.postprocess_conv.weight) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - def _forward(self, x, t, cross_attn_cond): - timestep_embed = self.to_timestep_embed(self.timestep_features(t[:, None])) - prepend_inputs = timestep_embed.unsqueeze(1) - prepend_length = prepend_inputs.shape[1] - - x = self.preprocess_conv(x) + x - x = rearrange(x, "b c n -> b n c") - output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond) - output = rearrange(output, "b n c -> b c n")[:, :, prepend_length:] - return self.postprocess_conv(output) + output - - def forward( - self, - x, - t, - cross_attn_cond, - negative_cross_attn_cond=None, - negative_cross_attn_mask=None, - cfg_scale: float = 1.0, - scale_phi: float = 0.0, - **kwargs, - ): - if cfg_scale == 1.0: - return self._forward(x, t, cross_attn_cond) - - # Classifier-free guidance: batch the conditional + unconditional pass. - null_embed = torch.zeros_like(cross_attn_cond) - if negative_cross_attn_cond is not None and negative_cross_attn_mask is not None: - mask = negative_cross_attn_mask.to(torch.bool).unsqueeze(2) - negative_cross_attn_cond = torch.where(mask, negative_cross_attn_cond, null_embed) - uncond = negative_cross_attn_cond if negative_cross_attn_cond is not None else null_embed - - batch_output = self._forward( - torch.cat([x, x], dim=0), - torch.cat([t, t], dim=0), - torch.cat([cross_attn_cond, uncond], dim=0), - ) - cond_output, uncond_output = torch.chunk(batch_output, 2, dim=0) - cfg_output = uncond_output + (cond_output - uncond_output) * cfg_scale - - if scale_phi == 0.0: - return cfg_output - cond_std = cond_output.std(dim=1, keepdim=True) - cfg_std = cfg_output.std(dim=1, keepdim=True) - return scale_phi * (cfg_output * (cond_std / cfg_std)) + (1 - scale_phi) * cfg_output - - -class ContinuousMMDiTTransformer(nn.Module): - def __init__( - self, - dim, - depth, - *, - dim_in=None, - dim_out=None, - dim_heads=64, - _latent_seq_len=237, - ): - super().__init__() - - self.dim = dim - self.depth = depth - - self.project_in = nn.Linear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity() - self.project_out = nn.Linear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity() - - hidden_dim = dim - num_heads = dim_heads - mlp_ratio = 4.0 - self._latent_seq_len = _latent_seq_len - - self.layers = nn.ModuleList( - [ - AudioXMMDiTBlock( - hidden_dim, - num_heads, - mlp_ratio=mlp_ratio, - prefix=f"layers.{i}", - ) - for i in range(depth) - ] - ) - self.proj_mm_tokens = nn.Linear(768, hidden_dim) if dim != 768 else nn.Identity() - self.proj_mm_seq_len = nn.Linear(384, self._latent_seq_len) if self._latent_seq_len != 384 else nn.Identity() - - # AudioX RoPE: interleaved (GPT-J) pair layout, theta=10000. - head_dim = hidden_dim // num_heads - pos = torch.arange(self._latent_seq_len, dtype=torch.float32, device=self.device) - inv_freq = 1.0 / (10000 ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=self.device) / head_dim)) - ang = torch.outer(pos, inv_freq) - self.register_buffer("latent_rope_cos", torch.cos(ang), persistent=False) - self.register_buffer("latent_rope_sin", torch.sin(ang), persistent=False) - - @property - def device(self): - return next(self.parameters()).device - - def forward( - self, - x, - prepend_embeds=None, - context=None, - ): - x = self.project_in(x) - - if prepend_embeds is not None: - prepend_dim = prepend_embeds.shape[-1] - assert prepend_dim == x.shape[-1], "prepend dimension must match sequence dimension" - x = torch.cat((prepend_embeds, x), dim=-2) - - time_cond = prepend_embeds.squeeze(1) - mm_tokens = context - - mm_tokens = self.proj_mm_tokens(mm_tokens) - mm_tokens = rearrange(mm_tokens, "b s d -> b d s") - mm_tokens = self.proj_mm_seq_len(mm_tokens) - mm_tokens = rearrange(mm_tokens, "b d s -> b s d") - - time_cond = time_cond.unsqueeze(1) - rot = ( - self.latent_rope_cos.to(device=x.device, dtype=x.dtype), - self.latent_rope_sin.to(device=x.device, dtype=x.dtype), - ) - for block in self.layers: - x = block(x, mm_tokens, rot, context=time_cond) - - x = self.project_out(x) - return x - - -def __getattr__(name: str): - if name in __all__: - return globals()[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py deleted file mode 100644 index 395e06dc8..000000000 --- a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py +++ /dev/null @@ -1,933 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import json -import math -import os -from collections.abc import Iterable -from typing import Any, ClassVar - -import torch -import torch.nn.functional as F -import torchsde -from diffusers import AutoencoderOobleck -from einops import rearrange -from torch import einsum, nn -from torchvision import transforms -from transformers import AutoConfig, CLIPVisionModelWithProjection, T5EncoderModel, T5TokenizerFast -from vllm.logger import init_logger -from vllm.model_executor.models.utils import AutoWeightsLoader - -from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.audiox.audiox_transformer import MMDiffusionTransformer -from vllm_omni.diffusion.models.interface import SupportAudioOutput -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from vllm_omni.transformers_utils.processors import audiox as _audiox_transforms - -_VIDEO_ONLY_TASKS = _audiox_transforms.VIDEO_ONLY_TASKS -_TEXT_VIDEO_TASKS = _audiox_transforms.TEXT_VIDEO_TASKS -_VIDEO_CONDITIONED_TASKS = _audiox_transforms.VIDEO_CONDITIONED_TASKS -_normalize_prompts = _audiox_transforms.normalize_prompts -prepare_audio_reference = _audiox_transforms.prepare_audio_reference -prepare_video_reference = _audiox_transforms.prepare_video_reference - -# Polyexponential sigma schedule defaults; mirror upstream AudioX gradio interface -# (``audiox/interface/gradio.py`` ``generate_cond`` defaults: ``sigma_min=0.03``, ``sigma_max=1000``). -_DEFAULT_UPSTREAM_SIGMA_MIN = 0.03 -_DEFAULT_UPSTREAM_SIGMA_MAX = 1000.0 - -logger = init_logger(__name__) - - -def _load_audiox_bundle_config(model_root: str) -> dict[str, Any]: - with open(os.path.join(os.path.abspath(model_root), "config.json"), encoding="utf-8") as f: - return json.load(f) - - -def _audio_conditioning_input_samples(model_config: dict[str, Any]) -> int: - configs = model_config["model"]["conditioning"]["configs"] - cfg = next(c["config"] for c in configs if c["id"] == "audio_prompt") - return int(cfg["latent_seq_len"]) * int(cfg["pretransform_config"]["config"]["downsampling_ratio"]) - - -def get_audiox_post_process_func(od_config: OmniDiffusionConfig): - """Convert the pipeline's float audio tensor to a CPU numpy array for serving.""" - - def post_process_func(audio: torch.Tensor) -> Any: - if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): - payload = dict(audio["payload"]) - audio_payload = payload.get("audio") - if isinstance(audio_payload, torch.Tensor): - payload["audio"] = audio_payload.detach().cpu().float().numpy() - return { - "payload": payload, - "metadata": audio.get("metadata") or {}, - } - if isinstance(audio, torch.Tensor): - return audio.detach().cpu().float().numpy() - return audio - - return post_process_func - - -class SA_PreNorm(nn.Module): - def __init__(self, dim, fn): - super().__init__() - self.norm = nn.LayerNorm(dim) - self.fn = fn - - def forward(self, x, **kwargs): - return self.fn(self.norm(x), **kwargs) - - -class SA_FeedForward(nn.Module): - def __init__(self, dim, hidden_dim): - super().__init__() - # Dropout p=0 preserves upstream ``net.{2,4}`` state-dict keys so the upstream weights - # load into the right slots; inference is identical to no dropout. - self.net = nn.Sequential( - nn.Linear(dim, hidden_dim), - nn.GELU(), - nn.Dropout(0.0), - nn.Linear(hidden_dim, dim), - nn.Dropout(0.0), - ) - - def forward(self, x): - return self.net(x) - - -# Manual einsum+softmax only. SDPA/diffusion Attention here degrades conditioning vs upstream. -class SA_Attention(nn.Module): - def __init__(self, dim, heads=8, dim_head=64): - super().__init__() - inner_dim = dim_head * heads - project_out = not (heads == 1 and dim_head == dim) - - self.heads = heads - self.scale = dim_head**-0.5 - - self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False) - - self.to_out = ( - nn.Sequential( - nn.Linear(inner_dim, dim), - nn.Dropout(0.0), - ) - if project_out - else nn.Identity() - ) - - def forward(self, x): - h = self.heads - qkv = self.to_qkv(x).chunk(3, dim=-1) - q, k, v = (rearrange(t, "b n (h d) -> b h n d", h=h) for t in qkv) - - dots = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale - - attn = dots.softmax(dim=-1) - - out = einsum("b h i j, b h j d -> b h i d", attn, v) - out = rearrange(out, "b h n d -> b n (h d)") - out = self.to_out(out) - return out - - -class SA_Transformer(nn.Module): - def __init__(self, dim, depth, heads, dim_head, mlp_dim): - super().__init__() - self.layers = nn.ModuleList([]) - self.norm = nn.LayerNorm(dim) - for _ in range(depth): - self.layers.append( - nn.ModuleList( - [ - SA_PreNorm(dim, SA_Attention(dim, heads=heads, dim_head=dim_head)), - SA_PreNorm(dim, SA_FeedForward(dim, mlp_dim)), - ] - ) - ) - - def forward(self, x): - for attn, ff in self.layers: - x = attn(x) + x - x = ff(x) + x - return self.norm(x) - - -_AUDIOX_OOBLECK_CONFIG = { - "audio_channels": 2, - "channel_multiples": [1, 2, 4, 8, 16], - "decoder_channels": 128, - "decoder_input_channels": 64, - "downsampling_ratios": [2, 4, 4, 8, 8], - "encoder_hidden_size": 128, - "sampling_rate": 44100, -} - - -def _build_audiox_oobleck(scaling_factor: float = 1.0) -> AutoencoderOobleck: - vae = AutoencoderOobleck(**_AUDIOX_OOBLECK_CONFIG) - vae.audiox_scaling_factor = float(scaling_factor) # type: ignore[attr-defined] - return vae.eval().requires_grad_(False) - - -class AudioVaePromptAdapter(nn.Module): - def __init__(self, *, cond_dim: int, latent_seq_len: int = 215): - super().__init__() - self.pretransform = _build_audiox_oobleck() - in_ch = int(self.pretransform.config.decoder_input_channels) - self.proj_features_128 = nn.Linear(latent_seq_len, 128) - self.proj_out = nn.Linear(in_ch, cond_dim) if in_ch != cond_dim else nn.Identity() - - def forward(self, audio: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - z = self.pretransform.encode(audio, return_dict=True).latent_dist.sample() - latents = z / float(self.pretransform.audiox_scaling_factor) - latents = rearrange(self.proj_features_128(latents), "b c s -> b s c") - latents = self.proj_out(latents) - ones = torch.ones(latents.shape[0], latents.shape[2], device=latents.device) - return latents, ones - - -class _MAFCrossAttentionBlock(nn.Module): - def __init__(self, dim: int, num_heads: int): - super().__init__() - assert dim % num_heads == 0, "dim must be divisible by num_heads" - self._num_heads = num_heads - self._head_dim = dim // num_heads - self._scale = self._head_dim**-0.5 - self.to_q = nn.Linear(dim, dim, bias=True) - self.to_kv = nn.Linear(dim, dim * 2, bias=True) - self.to_out = nn.Linear(dim, dim, bias=True) - nn.init.zeros_(self.to_out.weight) - if self.to_out.bias is not None: - nn.init.zeros_(self.to_out.bias) - - def forward(self, experts: torch.Tensor, full_context: torch.Tensor) -> torch.Tensor: - nh, hd = self._num_heads, self._head_dim - q = rearrange(self.to_q(experts), "b n (h d) -> b h n d", h=nh, d=hd) - k, v = (rearrange(t, "b n (h d) -> b h n d", h=nh, d=hd) for t in self.to_kv(full_context).chunk(2, dim=-1)) - out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False, scale=self._scale) - return self.to_out(rearrange(out, "b h n d -> b n (h d)")) - - -class _MAFFusionBlock(nn.Module): - def __init__(self, dim: int, num_heads: int, mlp_ratio: float): - super().__init__() - assert dim % num_heads == 0, "dim must be divisible by num_heads" - head_dim = dim // num_heads - self._num_heads = num_heads - hidden = int(dim * mlp_ratio) - self.norm1 = nn.LayerNorm(dim) - self.to_qkv = nn.Linear(dim, dim * 3, bias=True) - self.self_attn = Attention( - num_heads=num_heads, - head_size=head_dim, - softmax_scale=head_dim**-0.5, - causal=False, - ) - self.out_proj = nn.Linear(dim, dim, bias=True) - self.norm2 = nn.LayerNorm(dim) - self.ff = nn.Sequential( - nn.Linear(dim, hidden), - nn.GELU(), - nn.Linear(hidden, dim), - ) - nn.init.zeros_(self.out_proj.weight) - if self.out_proj.bias is not None: - nn.init.zeros_(self.out_proj.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - h = self.norm1(x) - q, k, v = self.to_qkv(h).chunk(3, dim=-1) - nh = self._num_heads - q_bsn = rearrange(q, "b n (h d) -> b n h d", h=nh).contiguous() - k_bsn = rearrange(k, "b n (h d) -> b n h d", h=nh).contiguous() - v_bsn = rearrange(v, "b n (h d) -> b n h d", h=nh).contiguous() - out = self.self_attn(q_bsn, k_bsn, v_bsn, attn_metadata=None) - out = rearrange(out, "b n h d -> b n (h d)") - out = self.out_proj(out) - x = x + out - x = x + self.ff(self.norm2(x)) - return x - - -class MAF_Block(nn.Module): - DIM = 768 - MLP_RATIO = 4.0 - - def __init__( - self, - *, - dim: int = 768, - num_experts_per_modality: int = 64, - num_heads: int = 24, - num_fusion_layers: int = 8, - mlp_ratio: float = 4.0, - ): - super().__init__() - total_experts = num_experts_per_modality * 3 - - self.gating_network = nn.Sequential( - nn.Linear(dim * 3, dim), - nn.GELU(), - nn.Linear(dim, 3), - nn.Sigmoid(), - ) - - self.unified_experts = nn.Parameter(torch.randn(total_experts, dim)) - - self.cross_block = _MAFCrossAttentionBlock(dim, num_heads) - self.norm1 = nn.LayerNorm(dim) - self.fusion_blocks = nn.ModuleList( - [_MAFFusionBlock(dim, num_heads, mlp_ratio) for _ in range(num_fusion_layers)] - ) - - self.norm_v2 = nn.LayerNorm(dim) - self.norm_t2 = nn.LayerNorm(dim) - self.norm_a2 = nn.LayerNorm(dim) - self.bypass_gate_v = nn.Parameter(torch.tensor(-10.0)) - self.bypass_gate_t = nn.Parameter(torch.tensor(-10.0)) - self.bypass_gate_a = nn.Parameter(torch.tensor(-10.0)) - - def forward( - self, - video_tokens: torch.Tensor, - text_tokens: torch.Tensor, - audio_tokens: torch.Tensor, - ) -> dict[str, torch.Tensor]: - batch_size = video_tokens.shape[0] - - v_global = video_tokens.mean(dim=1) - t_global = text_tokens.mean(dim=1) - a_global = audio_tokens.mean(dim=1) - - all_global = torch.cat([v_global, t_global, a_global], dim=1) - gates = self.gating_network(all_global) - w_v, w_t, w_a = gates.chunk(3, dim=-1) - - gated_v = video_tokens * w_v.unsqueeze(-1) - gated_t = text_tokens * w_t.unsqueeze(-1) - gated_a = audio_tokens * w_a.unsqueeze(-1) - - full_context = torch.cat([gated_v, gated_t, gated_a], dim=1) - - experts = self.unified_experts.unsqueeze(0).expand(batch_size, -1, -1) - cross_out = self.cross_block(experts, full_context) - updated_experts = self.norm1(experts + cross_out) - - for blk in self.fusion_blocks: - updated_experts = blk(updated_experts) - - fused_v_experts, fused_t_experts, fused_a_experts = updated_experts.chunk(3, dim=1) - - refinement_v = fused_v_experts.mean(dim=1) - refinement_t = fused_t_experts.mean(dim=1) - refinement_a = fused_a_experts.mean(dim=1) - - alpha_v = torch.sigmoid(self.bypass_gate_v) - alpha_t = torch.sigmoid(self.bypass_gate_t) - alpha_a = torch.sigmoid(self.bypass_gate_a) - - final_v = video_tokens + alpha_v * self.norm_v2(refinement_v).unsqueeze(1) - final_t = text_tokens + alpha_t * self.norm_t2(refinement_t).unsqueeze(1) - final_a = audio_tokens + alpha_a * self.norm_a2(refinement_a).unsqueeze(1) - - return { - "video": final_v, - "text": final_t, - "audio": final_a, - } - - -class _BrownianTreeNoiseSampler: - """Brownian-tree noise sampler for DPM-Solver++ SDE, ported from k-diffusion. - - Returns a scaled Brownian increment between two sigma levels; the tree is indexed by - transformed ``sigma`` (here linear, same as k-diffusion's default), so re-querying - identical (sigma, sigma_next) pairs returns the same noise. Entropy must come from - the user's seed — otherwise the unseeded global RNG would make sampling non-deterministic. - """ - - def __init__(self, x: torch.Tensor, sigma_min: torch.Tensor, sigma_max: torch.Tensor, entropy: int): - t0 = torch.as_tensor(sigma_min) - t1 = torch.as_tensor(sigma_max) - self._tree = torchsde.BrownianTree(t0, torch.zeros_like(x), t1, entropy=entropy) - - def __call__(self, sigma: torch.Tensor, sigma_next: torch.Tensor) -> torch.Tensor: - # AudioX denoises from sigma_max → 0, so sigma > sigma_next throughout the loop. - t0 = torch.as_tensor(sigma_next) - t1 = torch.as_tensor(sigma) - return -self._tree(t0, t1) / (t1 - t0).sqrt() - - -class AudioXPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMixin): - supports_request_batch = False - support_audio_output: ClassVar[bool] = True - audio_sample_rate: ClassVar[int] = 44100 - audio_channels: ClassVar[int] = 2 - _PROFILER_TARGETS: ClassVar[list[str]] = ["diffuse"] - _CLIP_SYNC_DURATION_SEC: ClassVar[float] = 10.0 - _VIDEO_SYNC_FRAME_COUNT: ClassVar[int] = 240 - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - self.device = get_local_device() - if od_config.model is None: - raise ValueError( - "AudioXPipeline requires od_config.model (directory with unified safetensors; " - "see https://huggingface.co/zhangj1an/AudioX)." - ) - - if os.path.exists(od_config.model): - self._model_root = os.path.abspath(od_config.model) - else: - from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific - - self._model_root = download_weights_from_hf_specific(od_config.model, None, ["*"]) - self._model_config = _load_audiox_bundle_config(self._model_root) - - model_config = self._model_config["model"] - diffusion_config = model_config["diffusion"] - - self.model = MMDiffusionTransformer(**dict(diffusion_config["config"])) - - cond_configs = {c["id"]: c["config"] for c in model_config["conditioning"]["configs"]} - self.audio_vae_adapter = AudioVaePromptAdapter( - cond_dim=int(model_config["conditioning"]["cond_dim"]), - latent_seq_len=int(cond_configs["audio_prompt"]["latent_seq_len"]), - ) - - t5_name = cond_configs["text_prompt"]["t5_model_name"] - self._t5_max_length = int(cond_configs["text_prompt"]["max_length"]) - self.tokenizer = T5TokenizerFast.from_pretrained(t5_name) - t5_config = AutoConfig.from_pretrained(t5_name) - self.text_encoder = T5EncoderModel(t5_config).train(False).requires_grad_(False).to(torch.float16) - - clip_name = cond_configs["video_prompt"]["clip_model_name"] - clip_config = AutoConfig.from_pretrained(clip_name) - self.clip_encoder = CLIPVisionModelWithProjection(clip_config.vision_config) - _CLIP_PATCH_TOKENS, _VIDEO_FPS, _DURATION_SEC, _DIM = 50, 5, 10, 768 - _in_features = _CLIP_PATCH_TOKENS * _VIDEO_FPS * _DURATION_SEC - self._clip_in_features = _in_features - self._clip_out_features = 128 - self.clip_proj = nn.Linear(_in_features, self._clip_out_features) - self.clip_proj_sync = nn.Linear(240, self._clip_out_features) - self.clip_sync_weight = nn.Parameter(torch.tensor(0.0)) - self.clip_temp_transformer = SA_Transformer(_DIM, depth=4, heads=16, dim_head=64, mlp_dim=_DIM * 4) - self.clip_temp_pos_embedding = nn.Parameter(torch.randn(1, _VIDEO_FPS * _DURATION_SEC, _DIM)) - self.clip_empty_visual_feat = nn.Parameter(torch.zeros(1, self._clip_out_features, _DIM), requires_grad=False) - _CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) - _CLIP_STD = (0.26862954, 0.26130258, 0.27577711) - self._clip_normalize = transforms.Compose([transforms.Normalize(mean=list(_CLIP_MEAN), std=list(_CLIP_STD))]) - - self.pretransform = _build_audiox_oobleck( - scaling_factor=float(model_config["pretransform"].get("scale", 1.0)), - ) - - self.io_channels = model_config["io_channels"] - self.diffusion_objective = "v" - - gate_type_config = diffusion_config["gate_type_config"] - self.maf_block = MAF_Block( - dim=768, - num_experts_per_modality=int(gate_type_config["num_experts_per_modality"]), - num_heads=int(gate_type_config["num_heads"]), - num_fusion_layers=int(gate_type_config["num_fusion_layers"]), - ) - - logger.debug("AudioX model built from %s", self._model_root) - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=self._model_root, - subfolder="transformer", - revision=getattr(od_config, "revision", None), - prefix="", - ), - ] - sample_rate = int(self._model_config.get("sample_rate", 48000)) - self._sample_rate = sample_rate - self._sample_size = int(self._model_config.get("sample_size", sample_rate * 10)) - self._target_fps = int(self._model_config.get("video_fps", 5)) - self._audio_conditioning_samples = _audio_conditioning_input_samples(self._model_config) - - self.setup_diffusion_pipeline_profiler( - profiler_targets=list(self._PROFILER_TARGETS), - enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - _legacy_prefix = "conditioner.conditioners.audio_prompt." - - # DiT self-attn QKV: bundle stores weights with last dim (h, d, qkv) interleaved - # (i.e. Q/K/V rows for the same head sit next to each other), whereas - # QKVParallelLinear expects stacked [Q|K|V] blocks along the output axis. - # Reshape once at load time so the parallel loader can consume them. - nheads = int(self._model_config["model"]["diffusion"]["config"]["num_heads"]) - embed_dim = int(self._model_config["model"]["diffusion"]["config"]["embed_dim"]) - total_heads = embed_dim // nheads - head_dim = embed_dim // total_heads - qkv_mid = ".attn.qkv." - to_kv_mid = ".cross_attn.to_kv." - - def _restack_interleaved(tensor: torch.Tensor, n_slots: int) -> torch.Tensor: - """Turn bundle's (h, d, slot) interleaved-last-dim layout into stacked [slot|slot|...].""" - if tensor.dim() == 2: # weight: (out, in) - out, inp = tensor.shape - assert out == n_slots * total_heads * head_dim, (out, n_slots, total_heads, head_dim) - return ( - tensor.view(total_heads, head_dim, n_slots, inp).permute(2, 0, 1, 3).reshape(out, inp).contiguous() - ) - out = tensor.shape[0] - assert out == n_slots * total_heads * head_dim - return tensor.view(total_heads, head_dim, n_slots).permute(2, 0, 1).reshape(out).contiguous() - - def _remap(items): - for name, tensor in items: - if name.startswith(_legacy_prefix): - name = "audio_vae_adapter." + name[len(_legacy_prefix) :] - if qkv_mid in name and (name.endswith(".weight") or name.endswith(".bias")): - tensor = _restack_interleaved(tensor, 3) - elif to_kv_mid in name and (name.endswith(".weight") or name.endswith(".bias")): - tensor = _restack_interleaved(tensor, 2) - yield name, tensor - - loaded = AutoWeightsLoader(self).load_weights(_remap(weights)) - - self.to(torch.float32) - self.eval().requires_grad_(False) - - return loaded - - def _conditioning_dtype(self) -> torch.dtype: - p = next(self.model.parameters()) - return p.dtype if p.dtype.is_floating_point else torch.float32 - - @staticmethod - def _normalize_task(task: str | None) -> str | None: - if task is None: - return None - t = str(task).strip().lower() - return t or None - - @staticmethod - def _text_for_task(task_norm: str | None, prompt: str) -> str: - if task_norm in _VIDEO_ONLY_TASKS: - return "" - return prompt - - @staticmethod - def _ensure_text_video_prompts(task_norm: str | None, prompts: list[str]) -> None: - if task_norm not in _TEXT_VIDEO_TASKS: - return - for i, p in enumerate(prompts): - if not str(p).strip(): - raise ValueError( - f"audiox_task={task_norm!r} requires a non-empty text prompt for item {i}; " - "use v2a/v2m for video-only generation." - ) - - def _audio_prompt_tensors( - self, - *, - raw_prompts: list[Any], - extra: dict[str, Any], - device: torch.device, - cond_dtype: torch.dtype = torch.float32, - ) -> list[torch.Tensor]: - target_len = self._audio_conditioning_samples - sample_rate = self._sample_rate - seconds_start = float(extra.get("seconds_start", 0.0)) - seconds_total = float(target_len) / float(sample_rate) - out: list[torch.Tensor] = [] - for raw in raw_prompts: - src = extra.get("audio_path") - if src is None: - out.append(torch.zeros(2, target_len, device=device, dtype=cond_dtype)) - continue - wav = prepare_audio_reference( - src, - model_sample_rate=sample_rate, - seconds_start=seconds_start, - seconds_total=seconds_total, - device=device, - ) - out.append(wav.to(dtype=cond_dtype)) - return out - - def _video_feature_tensors( - self, - *, - task_norm: str | None, - raw_prompts: list[Any], - extra: dict[str, Any], - seconds_start: float, - target_fps: int, - device: torch.device, - cond_dtype: torch.dtype = torch.float32, - ) -> list[torch.Tensor]: - clip_frames = int(round(self._CLIP_SYNC_DURATION_SEC * target_fps)) - if task_norm not in _VIDEO_CONDITIONED_TASKS: - empty = torch.zeros(clip_frames, 3, 224, 224, device=device, dtype=cond_dtype) - return [empty for _ in raw_prompts] - - tensors: list[torch.Tensor] = [] - for _ in raw_prompts: - src = extra.get("video_path") - if src is None: - raise ValueError(f"audiox_task={task_norm!r} requires video input: set extra_args['video_path'].") - vt = prepare_video_reference( - src, - duration=float(self._CLIP_SYNC_DURATION_SEC), - target_fps=target_fps, - seek_time=seconds_start, - ) - tensors.append(vt.to(device=device, dtype=cond_dtype)) - return tensors - - def get_conditioning_inputs(self, conditioning_tensors: dict[str, Any], negative: bool = False) -> dict[str, Any]: - video_feature, video_mask = conditioning_tensors["video_prompt"] - text_feature, text_mask = conditioning_tensors["text_prompt"] - audio_feature, audio_mask = conditioning_tensors["audio_prompt"] - - refined = self.maf_block(text_feature, video_feature, audio_feature) - fused = torch.cat(list(refined.values()), dim=1) - masks = torch.cat([video_mask, text_mask, audio_mask], dim=1) - - if negative: - return {"negative_cross_attn_cond": fused, "negative_cross_attn_mask": masks} - return {"cross_attn_cond": fused} - - def diffuse( - self, - *, - steps: int, - guidance_scale: float, - conditioning_tensors: dict[str, Any], - negative_conditioning_tensors: dict[str, Any] | None, - batch_size: int, - sigma_min: float, - sigma_max: float, - generator: torch.Generator, - cfg_rescale: float, - ) -> torch.Tensor: - device = self.device - model_dtype = next(self.model.parameters()).dtype - - # Match upstream AudioX: disable TF32 matmul + fp16 reduced precision + cudnn benchmark - # for numerical parity with audiox/inference/generation.py:152-156. - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cudnn.allow_tf32 = False - torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False - torch.backends.cudnn.benchmark = False - - latent_len = self._sample_size // int(self.pretransform.hop_length) - noise = torch.randn( - [batch_size, self.io_channels, latent_len], device=device, generator=generator, dtype=model_dtype - ) - - def _cast(d: dict[str, Any]) -> dict[str, Any]: - return {k: (v.type(model_dtype) if isinstance(v, torch.Tensor) else v) for k, v in d.items()} - - cond = _cast(self.get_conditioning_inputs(conditioning_tensors)) - neg = ( - _cast(self.get_conditioning_inputs(negative_conditioning_tensors, negative=True)) - if negative_conditioning_tensors is not None - else {} - ) - - # Inlined k-diffusion VDenoiser + sample_dpmpp_3m_sde, matching upstream AudioX exactly. - # diffusers' EDMDPMSolverMultistepScheduler uses different v-prediction preconditioning - # and a different stochastic update rule, which here produces a fixed ~861 Hz resonance - # in the decoded audio regardless of conditioning. - def denoise(x_in: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: - s2 = sigma * sigma - c_skip = 1.0 / (s2 + 1.0) - c_out = -sigma / (s2 + 1.0).sqrt() - c_in = 1.0 / (s2 + 1.0).sqrt() - t_cond = sigma.atan() * (2.0 / math.pi) - v = self.model( - x_in * c_in, - t_cond, - cross_attn_cond=cond["cross_attn_cond"], - negative_cross_attn_cond=neg.get("negative_cross_attn_cond"), - negative_cross_attn_mask=neg.get("negative_cross_attn_mask"), - cfg_scale=guidance_scale, - scale_phi=cfg_rescale, - ) - return v * c_out + x_in * c_skip - - ramp = torch.linspace(1.0, 0.0, steps, device=device) - sigmas = torch.cat( - [ - torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min)), - torch.zeros(1, device=device), - ] - ) - x = noise * sigmas[0] - - # Match upstream AudioX: sampler runs under fp16 autocast (see audiox/inference/sampling.py:184). - with torch.cuda.amp.autocast(): - if steps <= 1: - s_in = x.new_ones([x.shape[0]]) - sampled = denoise(x, sigmas[0] * s_in) - else: - # DPM-Solver++(3M) SDE loop (k-diffusion sample_dpmpp_3m_sde), eta=1.0, s_noise=1.0. - noise_sampler = _BrownianTreeNoiseSampler( - x, - sigmas[sigmas > 0].min(), - sigmas.max(), - entropy=generator.initial_seed(), - ) - s_in = x.new_ones([x.shape[0]]) - denoised_1 = denoised_2 = None - h_1 = h_2 = None - eta = 1.0 - for i in range(len(sigmas) - 1): - denoised = denoise(x, sigmas[i] * s_in) - if sigmas[i + 1] == 0: - x = denoised - else: - t_, s_ = -sigmas[i].log(), -sigmas[i + 1].log() - h = s_ - t_ - h_eta = h * (eta + 1) - x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised - if h_2 is not None: - r0 = h_1 / h - r1 = h_2 / h - d1_0 = (denoised - denoised_1) / r0 - d1_1 = (denoised_1 - denoised_2) / r1 - d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1) - d2 = (d1_0 - d1_1) / (r0 + r1) - phi_2 = h_eta.neg().expm1() / h_eta + 1 - phi_3 = phi_2 / h_eta - 0.5 - x = x + phi_2 * d1 - phi_3 * d2 - elif h_1 is not None: - r = h_1 / h - d = (denoised - denoised_1) / r - phi_2 = h_eta.neg().expm1() / h_eta + 1 - x = x + phi_2 * d - x = ( - x - + noise_sampler(sigmas[i], sigmas[i + 1]) - * sigmas[i + 1] - * (-2 * h * eta).expm1().neg().sqrt() - ) - denoised_1, denoised_2 = denoised, denoised_1 - h_1, h_2 = h, h_1 - sampled = x - - vae = self.pretransform.to(device=sampled.device, dtype=torch.float32).eval() - return vae.decode(sampled.to(torch.float32) * float(vae.audiox_scaling_factor), return_dict=True).sample - - def _encode_text(self, texts: list[str], device: torch.device) -> list[torch.Tensor]: - self.text_encoder.to(device) - encoded = self.tokenizer( - texts, - truncation=True, - max_length=self._t5_max_length, - padding="max_length", - return_tensors="pt", - ) - input_ids = encoded["input_ids"].to(device) - attention_mask = encoded["attention_mask"].to(device).to(torch.bool) - - self.text_encoder.eval() - with torch.no_grad(): - embeddings = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask)["last_hidden_state"] - - embeddings = embeddings.float() * attention_mask.unsqueeze(-1).float() - return [embeddings, attention_mask] - - def _encode_video(self, video_list: list[dict], device: torch.device) -> list[torch.Tensor]: - self.clip_encoder.to(device).eval() - - video_tensors = [item["video_tensors"] for item in video_list] - video_sync_frames = torch.cat([item["video_sync_frames"] for item in video_list], dim=0).to(device) - - original_videos = torch.cat(video_tensors, dim=0).to(device) - batch_size, time_length, _, _, _ = original_videos.size() - is_zero = torch.all(original_videos == 0, dim=(1, 2, 3, 4)) - - frames = original_videos.flatten(0, 1) - pixel_values = self._clip_normalize(frames).to(device) - - with torch.no_grad(): - outputs = self.clip_encoder(pixel_values=pixel_values) - hidden = outputs.last_hidden_state - hidden = rearrange(hidden, "(b t) p d -> (b p) t d", b=batch_size, t=time_length) - hidden = hidden + self.clip_temp_pos_embedding - hidden = self.clip_temp_transformer(hidden) - hidden = rearrange(hidden, "(b p) t d -> b (t p) d", b=batch_size) - hidden = self.clip_proj(hidden.view(-1, self._clip_in_features)) - hidden = hidden.view(batch_size, self._clip_out_features, -1) - - sync = self.clip_proj_sync(video_sync_frames.view(-1, 240)) - sync = sync.view(batch_size, self._clip_out_features, -1) - hidden = hidden + self.clip_sync_weight * sync - - empty = self.clip_empty_visual_feat.expand(batch_size, -1, -1) - hidden = torch.where(is_zero.view(batch_size, 1, 1), empty, hidden) - return [hidden, torch.ones(batch_size, 1, device=device)] - - def _encode_conditioning_tensors(self, batch_metadata: list[dict[str, Any]]) -> dict[str, Any]: - device = self.device - audio = torch.cat([item["audio_prompt"] for item in batch_metadata], dim=0).to(device) - return { - "audio_prompt": list(self.audio_vae_adapter(audio)), - "text_prompt": self._encode_text([item["text_prompt"] for item in batch_metadata], device), - "video_prompt": self._encode_video([item["video_prompt"] for item in batch_metadata], device), - } - - def _build_conditioning_batch( - self, - *, - texts: list[str], - video_tensors_list: list[torch.Tensor], - audio_prompt_list: list[torch.Tensor], - sync_features: torch.Tensor, - seconds_start: float, - seconds_model: float, - num_outputs_per_prompt: int, - task_norm: str | None, - ) -> list[dict[str, Any]]: - batch: list[dict[str, Any]] = [] - for i, text in enumerate(texts): - for _ in range(num_outputs_per_prompt): - batch.append( - { - "video_prompt": { - "video_tensors": video_tensors_list[i].unsqueeze(0), - "video_sync_frames": sync_features, - }, - "text_prompt": self._text_for_task(task_norm, text), - "audio_prompt": audio_prompt_list[i].unsqueeze(0), - "seconds_start": seconds_start, - "seconds_total": seconds_model, - } - ) - return batch - - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - if req.prompts is None or len(req.prompts) == 0: - raise ValueError("AudioXPipeline requires at least one prompt.") - normalized_prompts = _normalize_prompts(list(req.prompts)) - prompts = [p["prompt"] for p in normalized_prompts] - - sampling_params = req.sampling_params - if sampling_params.num_inference_steps is None: - raise ValueError("AudioXPipeline requires sampling_params.num_inference_steps.") - num_inference_steps = int(sampling_params.num_inference_steps) - extra_args = sampling_params.extra_args or {} - task_norm = self._normalize_task(extra_args.get("audiox_task")) - self._ensure_text_video_prompts(task_norm, prompts) - - neg: list[str] | None = None - if not all(p.get("negative_prompt") is None for p in normalized_prompts): - neg = [str(p.get("negative_prompt") or "") for p in normalized_prompts] - - guidance_scale = float(sampling_params.guidance_scale) - if sampling_params.num_outputs_per_prompt <= 0: - raise ValueError("AudioXPipeline requires sampling_params.num_outputs_per_prompt > 0.") - num_outputs_per_prompt = int(sampling_params.num_outputs_per_prompt) - batch_size = len(prompts) * num_outputs_per_prompt - - seconds_start = float(extra_args.get("seconds_start", 0.0)) - seconds_model = self._sample_size / self._sample_rate - seconds_total = float(extra_args.get("seconds_total", seconds_model)) - sigma_min = float(extra_args.get("sigma_min", _DEFAULT_UPSTREAM_SIGMA_MIN)) - sigma_max = float(extra_args.get("sigma_max", _DEFAULT_UPSTREAM_SIGMA_MAX)) - cfg_rescale = float(extra_args.get("cfg_rescale", 0.0)) - device = self.device - generator = sampling_params.generator - if generator is None: - raise ValueError("AudioXPipeline requires sampling_params.generator.") - target_fps = self._target_fps - cond_dtype = self._conditioning_dtype() - - sync_features = torch.zeros(1, self._VIDEO_SYNC_FRAME_COUNT, 768, device=device, dtype=cond_dtype) - - audio_prompt_list = self._audio_prompt_tensors( - raw_prompts=normalized_prompts, - extra=extra_args, - device=device, - cond_dtype=cond_dtype, - ) - - video_tensors_list = self._video_feature_tensors( - task_norm=task_norm, - raw_prompts=normalized_prompts, - extra=extra_args, - seconds_start=seconds_start, - target_fps=target_fps, - device=device, - cond_dtype=cond_dtype, - ) - - conditioning_batch = self._build_conditioning_batch( - texts=prompts, - video_tensors_list=video_tensors_list, - audio_prompt_list=audio_prompt_list, - sync_features=sync_features, - seconds_start=seconds_start, - seconds_model=seconds_model, - num_outputs_per_prompt=num_outputs_per_prompt, - task_norm=task_norm, - ) - - negative_conditioning_batch: list[dict[str, Any]] | None = None - if neg is not None and guidance_scale > 1.0: - negative_conditioning_batch = self._build_conditioning_batch( - texts=neg, - video_tensors_list=video_tensors_list, - audio_prompt_list=audio_prompt_list, - sync_features=sync_features, - seconds_start=seconds_start, - seconds_model=seconds_model, - num_outputs_per_prompt=num_outputs_per_prompt, - task_norm=task_norm, - ) - - conditioning_tensors = self._encode_conditioning_tensors(conditioning_batch) - negative_conditioning_tensors: dict[str, Any] | None = None - if negative_conditioning_batch is not None: - negative_conditioning_tensors = self._encode_conditioning_tensors(negative_conditioning_batch) - - audio = self.diffuse( - steps=num_inference_steps, - guidance_scale=guidance_scale, - conditioning_tensors=conditioning_tensors, - negative_conditioning_tensors=negative_conditioning_tensors, - batch_size=batch_size, - sigma_min=sigma_min, - sigma_max=sigma_max, - generator=generator, - cfg_rescale=cfg_rescale, - ) - - # Trim decoded audio to the requested duration (matches upstream AudioX sample script). - if 0.0 < seconds_total < seconds_model: - target_samples = int(seconds_total * self._sample_rate) - audio = audio[..., :target_samples] - - stage_durations = self.stage_durations if getattr(self, "enable_diffusion_pipeline_profiler", False) else None - return DiffusionOutput( - output={ - "payload": {"audio": audio}, - "metadata": {"audio": {"audiox_task": task_norm}}, - }, - stage_durations=stage_durations, - ) diff --git a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py index 7ef6877c8..a3b56051d 100644 --- a/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py +++ b/vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Native vLLM-Omni pipeline for Boogu-Image-0.1. @@ -220,7 +220,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery DiffusersPipelineLoader.ComponentSource( model_or_path=od_config.model, subfolder="transformer", - revision=None, + revision=od_config.revision, prefix="transformer.", fall_back_to_pt=True, ) @@ -233,10 +233,18 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder # race; prefetch the whole component set before any from_pretrained. boogu_subfolders = ["scheduler", "vae", "mllm", "processor"] - prefetch_subfolders(model, boogu_subfolders, local_files_only=local_files_only) + prefetch_subfolders( + model, + boogu_subfolders, + local_files_only=local_files_only, + revision=od_config.revision, + ) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - model, subfolder="scheduler", local_files_only=local_files_only + model, + subfolder="scheduler", + local_files_only=local_files_only, + revision=od_config.revision, ) mllm = from_pretrained_with_prefetch( @@ -246,6 +254,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery prefetch_list=boogu_subfolders, local_files_only=local_files_only, torch_dtype=od_config.dtype, + revision=od_config.revision, ) # Upstream reuses the full VLM as an optional instruction rewriter and # encodes with its inner model (no ``lm_head``); the rewriter is not @@ -255,7 +264,10 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery self.mllm = mllm.to(self._execution_device) self.processor = Qwen3VLProcessor.from_pretrained( - model, subfolder="processor", local_files_only=local_files_only + model, + subfolder="processor", + local_files_only=local_files_only, + revision=od_config.revision, ) self.vae = from_pretrained_with_prefetch( @@ -264,6 +276,7 @@ class BooguImagePipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery subfolder="vae", prefetch_list=boogu_subfolders, local_files_only=local_files_only, + revision=od_config.revision, ).to(self._execution_device) self.transformer = BooguImageTransformer2DModel(od_config=od_config) diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index c5fe41a91..dc755b39c 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -15,6 +15,8 @@ from typing import Any import numpy as np import torch +from .utils import VIDEO_RES_SIZE_INFO + ACTION_MODE_POLICY = "policy" ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" @@ -46,38 +48,6 @@ EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { } -VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { - "256": { - "1,1": (256, 256), - "4,3": (320, 256), - "3,4": (256, 320), - "16,9": (320, 192), - "9,16": (192, 320), - }, - "480": { - "1,1": (640, 640), - "4,3": (736, 544), - "3,4": (544, 736), - "16,9": (832, 480), - "9,16": (480, 832), - }, - "704": { - "1,1": (960, 960), - "4,3": (1088, 832), - "3,4": (832, 1088), - "16,9": (1280, 704), - "9,16": (704, 1280), - }, - "720": { - "1,1": (960, 960), - "4,3": (1104, 832), - "3,4": (832, 1104), - "16,9": (1280, 720), - "9,16": (720, 1280), - }, -} - - def normalize_action_mode(mode: Any) -> str | None: if mode is None: return None diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 97cf92e26..7a924dd52 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -126,6 +126,7 @@ from .utils import ( ROBOLAB_DEFAULT_RAW_ACTION_DIM, ROBOLAB_DEFAULT_RESOLUTION, ROBOLAB_MIDTRAIN_RAW_ACTION_DIM, + VIDEO_RES_SIZE_INFO, RoboLabActionPostprocessInputs, RoboLabPolicyInputs, build_abs_pose_from_components, @@ -163,6 +164,14 @@ COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{widt # NOTE: Intentional typo in "give" instead of "given" to match training setup. COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a give prompt." COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a give prompt." +COSMOS3_TRANSFER_SYSTEM_PROMPT = ( + "You are a helpful assistant that generates images or videos following the user's instructions" + " and control signals (edge maps, blur, depth, or segmentation)." +) +COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE = ( + "Follow the {hint_names} control video precisely: shape, contour, silhouette, position, and motion of every " + "visible structure must align with the {hint_names} signal at every frame." +) COSMOS3_T2V_DEFAULT_HEIGHT = 720 COSMOS3_T2V_DEFAULT_WIDTH = 1280 @@ -191,6 +200,7 @@ COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS = "FlowMatchEulerDiscreteScheduler" # are tokenized to their natural length (no padding); this only bounds the # UND pathway / GEN cross-attention cost for pathologically long prompts. COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096 +COSMOS3_TRANSFER_REQUESTED_SIZE_KEY = "_cosmos3_transfer_requested_size" def _ceil_video_num_frames(num_frames: int, temporal_compression_factor: int) -> int: @@ -235,6 +245,50 @@ def _format_json_object_prompt( return json.dumps(prompt_obj) +def _json_object_aspect_ratio(prompt: str) -> Any | None: + """Return an aspect-ratio field from a JSON-object prompt, if present.""" + try: + prompt_obj = json.loads(prompt) + except (json.JSONDecodeError, TypeError, ValueError): + return None + if not isinstance(prompt_obj, dict): + return None + return prompt_obj.get("aspect_ratio") + + +def _normalize_aspect_ratio(value: Any) -> str | None: + """Normalize supported ratio spellings for comparison and prompt metadata.""" + if value is None: + return None + parts = str(value).strip().replace(":", ",").split(",") + if len(parts) != 2: + return None + try: + width, height = (int(part.strip()) for part in parts) + except ValueError: + return None + if width <= 0 or height <= 0: + return None + divisor = math.gcd(width, height) + return f"{width // divisor},{height // divisor}" + + +def _aspect_ratio_for_dimensions(width: int, height: int) -> str: + """Return the nearest supported Cosmos aspect-ratio label.""" + if width <= 0 or height <= 0: + raise ValueError(f"Cosmos3 canvas dimensions must be positive, got {width}x{height}.") + canvas_ratio = width / height + + def ratio_value(label: str) -> float: + ratio_width, ratio_height = (int(part) for part in label.split(",")) + return ratio_width / ratio_height + + return min( + (label for sizes in VIDEO_RES_SIZE_INFO.values() for label in sizes), + key=lambda label: abs(canvas_ratio - ratio_value(label)), + ) + + def resolve_cosmos3_transformer_cls(model_config: Any) -> type[Cosmos3VFMTransformer]: """Select the Cosmos3 transformer implementation from transformer/config.json.""" backbone_type = _tf_config_get(model_config, "backbone_type", None) @@ -492,6 +546,17 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): # Resolve missing H/W. if transfer_requested: + requested_width = request.sampling_params.width + requested_height = request.sampling_params.height + if ( + requested_width is not None + and requested_height is not None + and COSMOS3_TRANSFER_REQUESTED_SIZE_KEY not in extra + ): + extra[COSMOS3_TRANSFER_REQUESTED_SIZE_KEY] = { + "width": int(requested_width), + "height": int(requested_height), + } _set_transfer_size_from_image(request, image) elif request.sampling_params.height is None or request.sampling_params.width is None: if action_mode is not None: @@ -2068,6 +2133,12 @@ class Cosmos3OmniDiffusersPipeline( sp: OmniDiffusionSamplingParams, use_system_prompt: bool = False, is_t2i: bool = False, + system_prompt: str | None = None, + prompt_suffix: str | None = None, + use_duration_template: bool | None = None, + use_resolution_template: bool | None = None, + negative_metadata_mode: str | None = None, + aspect_ratio_override: str | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Format prompts with metadata templates and tokenize. @@ -2075,24 +2146,67 @@ class Cosmos3OmniDiffusersPipeline( For T2I (``is_t2i=True``) the duration template is suppressed (no FPS or duration concept for a single image) and the image-flavored - resolution template is used. + resolution template is used. ``prompt_suffix`` is appended only to the + positive prompt, after metadata formatting. ``negative_metadata_mode`` + controls whether the negative branch receives no metadata, the same + metadata as the positive branch, or inverse metadata. """ # Route cosmos3-specific controls through ``_get_sp_param`` so they # are picked up from ``extra_args`` (OpenAI endpoint path) as well # as from direct attributes. - use_duration_template = bool(self._get_sp_param(sp, "use_duration_template", False)) and not is_t2i + if use_duration_template is None: + use_duration_template = bool(self._get_sp_param(sp, "use_duration_template", False)) + if use_resolution_template is None: + use_resolution_template = bool(self._get_sp_param(sp, "use_resolution_template", False)) + if negative_metadata_mode is None: + negative_metadata_mode = str(self._get_sp_param(sp, "negative_metadata_mode", "inverse")) + negative_metadata_mode = negative_metadata_mode.strip().lower() + if negative_metadata_mode not in {"none", "same", "inverse"}: + raise ValueError( + "Cosmos3 negative_metadata_mode must be one of 'none', 'same', or 'inverse'; " + f"got {negative_metadata_mode!r}." + ) + + use_duration_template = use_duration_template and not is_t2i dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None - if bool(self._get_sp_param(sp, "use_resolution_template", False)): + if use_resolution_template: res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE else: res_tmpl = None + requested_aspect_ratio = self._get_sp_param(sp, "aspect_ratio", None) + prompt_aspect_ratio = _json_object_aspect_ratio(prompt) + metadata_aspect_ratio = requested_aspect_ratio if requested_aspect_ratio is not None else prompt_aspect_ratio + canvas_aspect_ratio = aspect_ratio_override or _aspect_ratio_for_dimensions(width, height) + normalized_metadata_ratio = _normalize_aspect_ratio(metadata_aspect_ratio) + has_aspect_ratio_conflict = ( + aspect_ratio_override is None + and metadata_aspect_ratio is not None + and normalized_metadata_ratio != canvas_aspect_ratio + ) + if has_aspect_ratio_conflict: + aspect_ratio_source = ( + "requested aspect_ratio" if requested_aspect_ratio is not None else "JSON prompt aspect_ratio" + ) + if _is_rank_zero(): + logger.warning( + "Cosmos3 %s=%r conflicts with the generated %dx%d canvas (WxH); using %s for prompt metadata.", + aspect_ratio_source, + metadata_aspect_ratio, + width, + height, + canvas_aspect_ratio, + ) + metadata_aspect_ratio = canvas_aspect_ratio + elif aspect_ratio_override is not None: + metadata_aspect_ratio = aspect_ratio_override + json_prompt = _format_json_object_prompt( prompt, num_frames=num_frames, frame_rate=frame_rate, height=height, width=width, - aspect_ratio=self._get_sp_param(sp, "aspect_ratio", None), + aspect_ratio=metadata_aspect_ratio, ) if json_prompt is not None: prompt = json_prompt @@ -2106,37 +2220,44 @@ class Cosmos3OmniDiffusersPipeline( duration_template=dur_tmpl, resolution_template=res_tmpl, ) + if prompt_suffix: + prompt = f"{prompt.rstrip()} {prompt_suffix.lstrip()}".strip() if _is_rank_zero(): logger.info("Final prompt: '%s'", prompt) - # Negative prompt: inverse templates ("not {duration}...", "not {height}x{width}..."). - # Applied whenever the matching positive template is enabled; an empty - # negative_prompt yields output that starts with the template, not a dot. - inv_dur = COSMOS3_INVERSE_DURATION_TEMPLATE if dur_tmpl else None - if res_tmpl is None: - inv_res = None - elif is_t2i: - inv_res = COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE + if negative_metadata_mode == "none": + negative_dur_tmpl = None + negative_res_tmpl = None + elif negative_metadata_mode == "same": + negative_dur_tmpl = dur_tmpl + negative_res_tmpl = res_tmpl else: - inv_res = COSMOS3_INVERSE_RESOLUTION_TEMPLATE + negative_dur_tmpl = COSMOS3_INVERSE_DURATION_TEMPLATE if dur_tmpl else None + if res_tmpl is None: + negative_res_tmpl = None + elif is_t2i: + negative_res_tmpl = COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE + else: + negative_res_tmpl = COSMOS3_INVERSE_RESOLUTION_TEMPLATE negative_prompt = self._apply_metadata_templates( negative_prompt, num_frames, frame_rate, height, width, - duration_template=inv_dur, - resolution_template=inv_res, - force_duration_template=True, + duration_template=negative_dur_tmpl, + resolution_template=negative_res_tmpl, + force_duration_template=negative_metadata_mode == "inverse", ) - default_sys_prompt = COSMOS3_T2I_SYSTEM_PROMPT if is_t2i else COSMOS3_SYSTEM_PROMPT - sys_prompt = self._get_sp_param(sp, "system_prompt", default_sys_prompt) or default_sys_prompt + if system_prompt is None: + default_sys_prompt = COSMOS3_T2I_SYSTEM_PROMPT if is_t2i else COSMOS3_SYSTEM_PROMPT + system_prompt = self._get_sp_param(sp, "system_prompt", default_sys_prompt) or default_sys_prompt cond_ids, cond_mask = self._tokenize_prompt( - prompt, max_sequence_length, use_system_prompt, system_prompt=sys_prompt + prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt ) uncond_ids, uncond_mask = self._tokenize_prompt( - negative_prompt, max_sequence_length, use_system_prompt, system_prompt=sys_prompt + negative_prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt ) return cond_ids, cond_mask, uncond_ids, uncond_mask @@ -2857,11 +2978,76 @@ class Cosmos3OmniDiffusersPipeline( self, sp: OmniDiffusionSamplingParams, source_hw: tuple[int, int] | None, - ) -> tuple[int, int]: + ) -> tuple[int, int, str]: resolution = self._get_sp_param(sp, "resolution", self._get_sp_param(sp, "image_size", 720)) source_h, source_w = source_hw or (COSMOS3_T2V_DEFAULT_HEIGHT, COSMOS3_T2V_DEFAULT_WIDTH) target_w, target_h = find_closest_target_size(int(source_h), int(source_w), resolution) - return int(target_h), int(target_w) + aspect_ratio = next( + ratio for ratio, size in VIDEO_RES_SIZE_INFO[str(resolution)].items() if size == (target_w, target_h) + ) + return int(target_h), int(target_w), aspect_ratio + + def _warn_transfer_bucket_conflicts( + self, + sp: OmniDiffusionSamplingParams, + prompt: str, + *, + source_hw: tuple[int, int] | None, + height: int, + width: int, + aspect_ratio: str, + ) -> None: + """Explain transfer's bucket selection when it overrides request metadata.""" + if not _is_rank_zero(): + return + + requested_size_data = self._get_sp_param(sp, COSMOS3_TRANSFER_REQUESTED_SIZE_KEY, None) + if isinstance(requested_size_data, Mapping): + requested_width = requested_size_data.get("width") + requested_height = requested_size_data.get("height") + else: + requested_width = getattr(sp, "width", None) + requested_height = getattr(sp, "height", None) + if requested_width is not None and requested_height is not None: + try: + requested_size = (int(requested_width), int(requested_height)) + except (TypeError, ValueError): + requested_size = None + if requested_size is not None and requested_size != (width, height): + source_description = ( + f"control input {source_hw[1]}x{source_hw[0]} (WxH)" + if source_hw is not None + else "the default transfer geometry" + ) + resolution = self._get_sp_param(sp, "resolution", self._get_sp_param(sp, "image_size", 720)) + logger.warning( + "Cosmos3 transfer ignores requested size=%dx%d (WxH); %s selected the %s resolution=%s " + "bucket, %dx%d (WxH). Transfer supports canonical Cosmos buckets, not arbitrary output sizes.", + *requested_size, + source_description, + aspect_ratio, + resolution, + width, + height, + ) + + request_aspect_ratio = self._get_sp_param(sp, "aspect_ratio", None) + aspect_ratio_source = "requested aspect_ratio" + if request_aspect_ratio is None: + request_aspect_ratio = _json_object_aspect_ratio(prompt) + aspect_ratio_source = "JSON prompt aspect_ratio" + normalized_requested_ratio = _normalize_aspect_ratio(request_aspect_ratio) + if normalized_requested_ratio is not None and normalized_requested_ratio != aspect_ratio: + logger.warning( + "Cosmos3 transfer %s=%r conflicts with the control-selected %s bucket, %dx%d (WxH); " + "using %s for prompt metadata.", + aspect_ratio_source, + request_aspect_ratio, + aspect_ratio, + width, + height, + aspect_ratio, + ) @staticmethod def _first_transfer_control_hw(transfer_config: Cosmos3TransferConfig) -> tuple[int, int] | None: @@ -2890,6 +3076,7 @@ class Cosmos3OmniDiffusersPipeline( control_latents: list[torch.Tensor], shared_kwargs: dict[str, Any], *, + control_weights: list[float] | None = None, velocity_mask: torch.Tensor, condition_latents: torch.Tensor, guidance_interval: tuple[float, float] | None = None, @@ -2925,6 +3112,7 @@ class Cosmos3OmniDiffusersPipeline( text_ids=cond_ids, text_mask=cond_mask, control_latents=control_latents, + control_weights=control_weights, _cosmos3_cache_key="transfer_cond_full", **shared_kwargs, ) @@ -2946,6 +3134,7 @@ class Cosmos3OmniDiffusersPipeline( text_ids=uncond_ids, text_mask=uncond_mask, control_latents=control_latents, + control_weights=control_weights, _cosmos3_cache_key="transfer_uncond_full", **shared_kwargs, ), @@ -2993,6 +3182,7 @@ class Cosmos3OmniDiffusersPipeline( text_ids=uncond_ids, text_mask=uncond_mask, control_latents=control_latents, + control_weights=control_weights, _cosmos3_cache_key="transfer_uncond_full", **shared_kwargs, ), @@ -3043,7 +3233,15 @@ class Cosmos3OmniDiffusersPipeline( source_hw = (int(input_frames.shape[-2]), int(input_frames.shape[-1])) else: source_hw = self._first_transfer_control_hw(transfer_config) - height, width = self._transfer_bucket_size(sp, source_hw) + height, width, transfer_aspect_ratio = self._transfer_bucket_size(sp, source_hw) + self._warn_transfer_bucket_conflicts( + sp, + prompt, + source_hw=source_hw, + height=height, + width=width, + aspect_ratio=transfer_aspect_ratio, + ) if input_frames is not None: if tuple(input_frames.shape[-2:]) != (height, width): @@ -3113,7 +3311,10 @@ class Cosmos3OmniDiffusersPipeline( self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH) or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH ) - use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False)) + transfer_prompt_suffix = None + if transfer_config.emphasize_control_in_prompt: + hint_names = ", ".join(hint.key for hint in transfer_config.ordered_hints) + transfer_prompt_suffix = COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names=hint_names) self._guidance_scale = guidance_scale self._num_timesteps = num_inference_steps @@ -3124,6 +3325,13 @@ class Cosmos3OmniDiffusersPipeline( if generator is None: generator = torch.Generator(device=self.device).manual_seed(seed) + for ignored_key in ("system_prompt", "use_system_prompt"): + if self._get_sp_param(sp, ignored_key, None) is not None: + logger.warning_once( + "Cosmos3 transfer ignores the '%s' request parameter and always " + "tokenizes both CFG branches with the transfer-specific system prompt.", + ignored_key, + ) cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts( prompt, negative_prompt, @@ -3133,8 +3341,14 @@ class Cosmos3OmniDiffusersPipeline( width, max_sequence_length, sp, - use_system_prompt, + use_system_prompt=True, is_t2i=False, + system_prompt=COSMOS3_TRANSFER_SYSTEM_PROMPT, + prompt_suffix=transfer_prompt_suffix, + use_duration_template=bool(self._get_sp_param(sp, "use_duration_template", True)), + use_resolution_template=bool(self._get_sp_param(sp, "use_resolution_template", True)), + negative_metadata_mode=str(self._get_sp_param(sp, "negative_metadata_mode", "same")), + aspect_ratio_override=transfer_aspect_ratio, ) output_chunks: list[torch.Tensor] = [] @@ -3227,6 +3441,7 @@ class Cosmos3OmniDiffusersPipeline( control_guidance=transfer_config.control_guidance, control_guidance_interval=transfer_config.control_guidance_interval, control_latents=control_latents, + control_weights=transfer_config.normalized_control_weights, shared_kwargs=shared_kwargs, velocity_mask=velocity_mask, condition_latents=condition_latents, diff --git a/vllm_omni/diffusion/models/cosmos3/transfer.py b/vllm_omni/diffusion/models/cosmos3/transfer.py index 42cbea559..9c8a52f3f 100644 --- a/vllm_omni/diffusion/models/cosmos3/transfer.py +++ b/vllm_omni/diffusion/models/cosmos3/transfer.py @@ -14,6 +14,7 @@ generation, or action generation. from __future__ import annotations import importlib +import math from collections.abc import Mapping from dataclasses import dataclass, field from pathlib import Path @@ -25,6 +26,14 @@ import torch import torch.nn.functional as F TRANSFER_HINT_KEYS: tuple[str, ...] = ("edge", "blur", "depth", "seg", "wsm") +_TRANSFER_HINT_COMMON_FIELDS = frozenset({"control_path", "control", "control_weight"}) +_TRANSFER_HINT_FIELDS: dict[str, frozenset[str]] = { + "edge": _TRANSFER_HINT_COMMON_FIELDS | {"preset_edge_threshold"}, + "blur": _TRANSFER_HINT_COMMON_FIELDS | {"preset_blur_strength"}, + "depth": _TRANSFER_HINT_COMMON_FIELDS, + "seg": _TRANSFER_HINT_COMMON_FIELDS, + "wsm": _TRANSFER_HINT_COMMON_FIELDS, +} TRANSFER_SAMPLE_DEFAULTS: dict[str, Any] = { "num_video_frames_per_chunk": 93, "num_conditional_frames": 1, @@ -33,6 +42,7 @@ TRANSFER_SAMPLE_DEFAULTS: dict[str, Any] = { "show_input": False, "num_first_chunk_conditional_frames": 0, "share_vision_temporal_positions": True, + "emphasize_control_in_prompt": True, } TRANSFER_DEFAULTS: dict[str, dict[str, Any]] = { "edge": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, @@ -85,6 +95,7 @@ class Cosmos3TransferHint: key: str control_path: str | None = None control: Any | None = None + control_weight: float = 1.0 preset_edge_threshold: str = "medium" preset_blur_strength: str = "medium" @@ -103,6 +114,7 @@ class Cosmos3TransferConfig: show_input: bool = False num_first_chunk_conditional_frames: int = 0 share_vision_temporal_positions: bool = True + emphasize_control_in_prompt: bool = True num_frames: int | None = None fps: float | None = None @@ -110,6 +122,16 @@ class Cosmos3TransferConfig: def ordered_hints(self) -> list[Cosmos3TransferHint]: return [self.hints[key] for key in TRANSFER_HINT_KEYS if key in self.hints] + @property + def normalized_control_weights(self) -> list[float]: + weights = [hint.control_weight for hint in self.ordered_hints] + if any(not math.isfinite(weight) or weight < 0.0 for weight in weights): + raise ValueError(f"Cosmos3 transfer control_weight values must be finite and non-negative, got {weights}.") + total = sum(weights) + if total <= 0.0: + raise ValueError("Cosmos3 transfer control_weight values must have a positive sum.") + return [weight / total for weight in weights] + def _extra_args(sp: Any) -> Mapping[str, Any]: extra = getattr(sp, "extra_args", None) @@ -184,7 +206,7 @@ def _as_interval(value: Any) -> tuple[float, float] | None: return None if isinstance(value, str): value = [item.strip() for item in value.split(",") if item.strip()] - if not isinstance(value, (list, tuple)) or len(value) != 2: + if not isinstance(value, list | tuple) or len(value) != 2: raise ValueError("Cosmos3 transfer control_guidance_interval must contain exactly two values.") lo, hi = float(value[0]), float(value[1]) if lo > hi: @@ -220,10 +242,27 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer raise TypeError( f"Cosmos3 transfer hint '{key}' must be an object, path string, or true; got {type(raw)!r}." ) + unknown_fields = set(raw) - _TRANSFER_HINT_FIELDS[key] + if unknown_fields: + raise ValueError( + f"Unsupported Cosmos3 transfer hint '{key}' fields: {sorted(unknown_fields)}. " + f"Supported fields are: {sorted(_TRANSFER_HINT_FIELDS[key])}." + ) + try: + control_weight = float(raw.get("control_weight", 1.0)) + except (TypeError, ValueError) as exc: + raise TypeError( + f"Cosmos3 transfer hint '{key}' control_weight must be a finite non-negative number." + ) from exc + if not math.isfinite(control_weight) or control_weight < 0.0: + raise ValueError( + f"Cosmos3 transfer hint '{key}' control_weight must be finite and non-negative, got {control_weight!r}." + ) hints[key] = Cosmos3TransferHint( key=key, control_path=str(raw["control_path"]) if raw.get("control_path") is not None else None, control=raw.get("control"), + control_weight=control_weight, preset_edge_threshold=str(raw.get("preset_edge_threshold") or "medium").lower(), preset_blur_strength=str(raw.get("preset_blur_strength") or "medium").lower(), ) @@ -239,6 +278,7 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer "show_input", "num_first_chunk_conditional_frames", "share_vision_temporal_positions", + "emphasize_control_in_prompt", ) if any(_is_user_field(extra, sp, prompt_data, key) for key in transfer_only): raise ValueError("Cosmos3 transfer options were provided, but no transfer hint was selected.") @@ -293,6 +333,10 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer _param(extra, sp, prompt_data, "share_vision_temporal_positions", None), bool(TRANSFER_SAMPLE_DEFAULTS["share_vision_temporal_positions"]), ), + emphasize_control_in_prompt=_as_bool( + _param(extra, sp, prompt_data, "emphasize_control_in_prompt", None), + bool(TRANSFER_SAMPLE_DEFAULTS["emphasize_control_in_prompt"]), + ), num_frames=( int(_param(extra, sp, prompt_data, "num_frames")) if _param(extra, sp, prompt_data, "num_frames", None) is not None @@ -328,6 +372,7 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer raise ValueError(f"Unsupported Cosmos3 edge preset: {hint.preset_edge_threshold!r}.") if hint.key == "blur" and hint.preset_blur_strength not in BLUR_DOWNUP_PRESETS: raise ValueError(f"Unsupported Cosmos3 blur preset: {hint.preset_blur_strength!r}.") + _ = config.normalized_control_weights return config diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index cea5a32ff..c946eb4ca 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -651,6 +651,19 @@ class Cosmos3CrossAttention(nn.Module): softmax_scale=1.0 / (self.head_dim**0.5), num_kv_heads=self.num_kv_heads, ) + # Multi-control attention operates on one full [control_i, target] + # sequence at a time. Keep those sequences replicated when Ulysses is + # active; sharding the concatenated [control_1, ..., control_N, target] + # layout would split the per-control ranges across ranks. + self.multi_control_attn = FrameworkAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=1.0 / (self.head_dim**0.5), + num_kv_heads=self.num_kv_heads, + prefix=f"{prefix}.multi_control_attn", + skip_sequence_parallel=True, + ) # TODO(follow-up): collapse _forward_local and _forward_sp into a single # joint-based path when NoParallelAttention can process joint_key/value. @@ -701,6 +714,61 @@ class Cosmos3CrossAttention(nn.Module): out = self.attn(q, k, v, attn_metadata) return out.reshape(B, S_gen, -1) + def _forward_multi_control( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_und: torch.Tensor, + v_und: torch.Tensor, + control_token_sizes: tuple[int, ...], + control_weights: tuple[float, ...], + ) -> torch.Tensor: + """Run independent attention per control and weight target outputs. + + For control ``i``, both its queries and the target queries attend to + ``[text, control_i, target]``. Each control keeps the output from its + own pass, while target outputs are combined using normalized weights. + """ + if len(control_token_sizes) != len(control_weights): + raise ValueError( + "Cosmos3 control token sizes and weights must have the same length: " + f"sizes={len(control_token_sizes)}, weights={len(control_weights)}." + ) + control_tokens = sum(control_token_sizes) + if control_tokens <= 0 or control_tokens >= q.shape[1]: + raise ValueError( + "Cosmos3 multi-control attention requires non-empty control and target token ranges, " + f"got control_tokens={control_tokens}, total_tokens={q.shape[1]}." + ) + + q_target = q[:, control_tokens:] + k_target = k[:, control_tokens:] + v_target = v[:, control_tokens:] + control_outputs: list[torch.Tensor] = [] + target_output: torch.Tensor | None = None + start = 0 + for size, weight in zip(control_token_sizes, control_weights, strict=True): + if size <= 0: + raise ValueError(f"Cosmos3 control token sizes must be positive, got {control_token_sizes}.") + end = start + size + q_control = q[:, start:end] + k_control = k[:, start:end] + v_control = v[:, start:end] + q_pair = torch.cat([q_control, q_target], dim=1) + k_pair = torch.cat([k_und, k_control, k_target], dim=1) + v_pair = torch.cat([v_und, v_control, v_target], dim=1) + pair_output = self.multi_control_attn(q_pair, k_pair, v_pair) + control_outputs.append(pair_output[:, :size]) + weighted_target = pair_output[:, size:] * weight + target_output = weighted_target if target_output is None else target_output + weighted_target + start = end + + if start != control_tokens or target_output is None: + raise RuntimeError("Cosmos3 multi-control attention failed to cover all control token ranges.") + output = torch.cat([*control_outputs, target_output], dim=1) + return output.reshape(q.shape[0], q.shape[1], -1) + # -- Public forward: routes to the appropriate path ---------------------- def forward( @@ -710,6 +778,8 @@ class Cosmos3CrossAttention(nn.Module): v_und: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, + control_token_sizes: tuple[int, ...] | None = None, + control_weights: tuple[float, ...] | None = None, ) -> torch.Tensor: """ Args: @@ -733,7 +803,19 @@ class Cosmos3CrossAttention(nn.Module): # Qwen3-style RoPE q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) - if _is_sp_active(): + if control_token_sizes is not None or control_weights is not None: + if control_token_sizes is None or control_weights is None: + raise ValueError("Cosmos3 multi-control attention requires both control token sizes and weights.") + out = self._forward_multi_control( + q, + k, + v, + k_und, + v_und, + control_token_sizes, + control_weights, + ) + elif _is_sp_active(): out = self._forward_sp(q, k, v, k_und, v_und) else: out = self._forward_local(q, k, v, k_und, v_und) @@ -847,6 +929,8 @@ class Cosmos3GenDecoderLayer(nn.Module): freqs_sin: torch.Tensor | None = None, cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None, freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None, + control_token_sizes: tuple[int, ...] | None = None, + control_weights: tuple[float, ...] | None = None, ) -> torch.Tensor: if cached_kv is not None: if self.layer_idx is None: @@ -861,7 +945,13 @@ class Cosmos3GenDecoderLayer(nn.Module): hidden_states = self.input_layernorm(hidden_states) hidden_states = self.cross_attention( - hidden_states, k_und=k_und, v_und=v_und, freqs_cos=freqs_cos, freqs_sin=freqs_sin + hidden_states, + k_und=k_und, + v_und=v_und, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + control_token_sizes=control_token_sizes, + control_weights=control_weights, ) hidden_states = residual + hidden_states @@ -1579,6 +1669,7 @@ class Cosmos3VFMTransformer(nn.Module): sound_latents: torch.Tensor | None = None, noisy_frame_mask: torch.Tensor | None = None, control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None, transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, ...]: @@ -1601,6 +1692,8 @@ class Cosmos3VFMTransformer(nn.Module): frames are noisy, as in T2I/T2V. control_latents: Optional transfer-control latents. Controls are clean vision context and are packed before the noisy target. + control_weights: Optional non-negative relative weights for + transfer controls. Values are normalized to sum to one. Returns: [B, C, t, h, w] velocity prediction, or @@ -1629,6 +1722,31 @@ class Cosmos3VFMTransformer(nn.Module): else: control_latent_list = list(control_latents) has_control = len(control_latent_list) > 0 + if control_weights is None: + normalized_control_weights = ( + [1.0 / len(control_latent_list)] * len(control_latent_list) if has_control else [] + ) + else: + if isinstance(control_weights, torch.Tensor): + raw_control_weights = control_weights.detach().cpu().flatten().tolist() + else: + raw_control_weights = list(control_weights) + if len(raw_control_weights) != len(control_latent_list): + raise ValueError( + "Cosmos3 transfer control_weights length must match control_latents: " + f"weights={len(raw_control_weights)}, controls={len(control_latent_list)}." + ) + normalized_control_weights = [float(weight) for weight in raw_control_weights] + if any(not math.isfinite(weight) or weight < 0.0 for weight in normalized_control_weights): + raise ValueError( + "Cosmos3 transfer control_weights must be finite and non-negative, " + f"got {normalized_control_weights}." + ) + control_weight_sum = sum(normalized_control_weights) + if control_weight_sum <= 0.0: + raise ValueError("Cosmos3 transfer control_weights must have a positive sum.") + normalized_control_weights = [weight / control_weight_sum for weight in normalized_control_weights] + use_multi_control_attention = len(control_latent_list) > 1 if has_control and (has_action or has_sound): raise ValueError("Cosmos3 transfer control latents cannot be combined with action or sound latents.") if has_action and not self.action_gen: @@ -1646,6 +1764,13 @@ class Cosmos3VFMTransformer(nn.Module): # Query Ulysses state at runtime ulysses_size, _, _ = _get_ulysses_state() + if use_multi_control_attention and ulysses_size > 1: + logger.warning_once( + "Cosmos3 multi-control attention requires full [control_i | target] " + "sequences, so the GEN pathway runs replicated on every Ulysses rank. " + "Sequence parallelism will not reduce per-rank memory or latency for " + "multi-control transfer requests." + ) # Pack action/sound tokens (no learned weights) up front so the UND # cache sizing knows their token lengths. The modality projections are @@ -1722,6 +1847,9 @@ class Cosmos3VFMTransformer(nn.Module): ) hidden_controls.append(hidden_control) s_control += hidden_control.shape[1] + control_token_sizes = tuple(control.shape[1] for control in hidden_controls) + multi_control_weights = tuple(normalized_control_weights) if use_multi_control_attention else None + multi_control_token_sizes = control_token_sizes if use_multi_control_attention else None hidden_action = None hidden_sound = None if action_tokens is not None: @@ -1788,10 +1916,11 @@ class Cosmos3VFMTransformer(nn.Module): has_action=has_action, has_sound=has_sound, has_control=has_control, - ulysses_size=ulysses_size, + ulysses_size=1 if use_multi_control_attention else ulysses_size, ) freqs_cos, freqs_sin = self.cached_freqs_gen - hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) + if not use_multi_control_attention: + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) freqs_gen = (freqs_cos, freqs_sin) if len(self.gen_layers) == len(self.cached_kv): @@ -1802,6 +1931,8 @@ class Cosmos3VFMTransformer(nn.Module): v_und=v_und, freqs_cos=freqs_cos, freqs_sin=freqs_sin, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, ) # Cache-dit's block wrapper may return a tuple; unwrap it. if isinstance(hidden_gen, tuple): @@ -1813,11 +1944,14 @@ class Cosmos3VFMTransformer(nn.Module): hidden_gen, cached_kv=self.cached_kv, freqs_gen=freqs_gen, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, ) if isinstance(hidden_gen, tuple): hidden_gen = hidden_gen[0] - hidden_gen = self.gen_sp_gather(hidden_gen) + if not use_multi_control_attention: + hidden_gen = self.gen_sp_gather(hidden_gen) # Final norm and project back to latent space hidden_gen = self.norm_moe_gen(hidden_gen) diff --git a/vllm_omni/diffusion/models/cosmos3/utils.py b/vllm_omni/diffusion/models/cosmos3/utils.py index bfa3fd4db..76df64b3b 100644 --- a/vllm_omni/diffusion/models/cosmos3/utils.py +++ b/vllm_omni/diffusion/models/cosmos3/utils.py @@ -27,6 +27,37 @@ COSMOS3_DEFAULT_CONDITION_VIDEO_KEEP = "first" # pixel-frame budgets without instantiating the VAE. COSMOS3_VAE_TEMPORAL_COMPRESSION = 4 +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + ROBOLAB_DEFAULT_CONDITIONING_FPS = 15.0 ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE = 32 ROBOLAB_DEFAULT_IMAGE_HEIGHT = 540 diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index a4a9bc0dd..006af60e1 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -23,7 +23,11 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline from torch import nn from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import ( + BasePipelineUtils, + get_pipeline_utils_for_config, + resolve_diffusers_pipeline_class, +) from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import ( apply_diffusers_quantization_config, convert_diffusers_quantization_config, @@ -117,9 +121,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): } convert_diffusers_quantization_config(load_kwargs) - pipeline_class = self.od_config.diffusers_pipeline_cls - pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None - self._pipeline_utils = get_pipeline_utils(pipeline_class_name) + self._pipeline_utils = get_pipeline_utils_for_config(self.od_config) self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs) component_names = ( self._load_diffusers_component_names(model_id, load_kwargs) @@ -129,7 +131,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names) logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") - self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs) + self._pipeline = self._load_pipeline_from_pretrained(model_id, load_kwargs) self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config) self._pipeline.to(self.device) @@ -162,6 +164,23 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): # Attention backend self._set_attention_backend() + def _load_pipeline_from_pretrained( + self, + model_id: str, + load_kwargs: dict[str, Any], + ) -> DiffusionPipeline: + pipeline_class = resolve_diffusers_pipeline_class(self.od_config) + configured_pipeline_class = self.od_config.diffusers_pipeline_cls + # Preserve DiffusionPipeline's normal dynamic class resolution unless + # a model-specific adapter hook explicitly selected another class. + loader_class = ( + pipeline_class + if pipeline_class is not None and pipeline_class is not configured_pipeline_class + else DiffusionPipeline + ) + logger.info("Loading Diffusers pipeline class %s", loader_class.__name__) + return loader_class.from_pretrained(model_id, **load_kwargs) + # ------------------------------------------------------------------ # Step-wise execution — explicitly rejected # ------------------------------------------------------------------ @@ -376,6 +395,13 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): if self._accept_call_kwargs is None or key in self._accept_call_kwargs: kwargs[key] = value + self._pipeline_utils.update_call_kwargs( + req, + sampling, + self._accept_call_kwargs, + kwargs, + ) + # Special format fields in sampling params if output_type := sampling.output_type or self.od_config.output_type: kwargs["output_type"] = output_type diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py index 4c810bb68..2d00881b8 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py @@ -1,17 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from typing import Any from diffusers.pipelines.pipeline_utils import DiffusionPipeline from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.request import resolve_video_num_frames +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams class BasePipelineUtils: """No-op hooks for pipeline-specific diffusers adapter behavior.""" + def resolve_pipeline_class( + self, + od_config: OmniDiffusionConfig, + configured_pipeline_class: type[DiffusionPipeline] | None, + ) -> type[DiffusionPipeline] | None: + """Resolve the concrete Diffusers class used to load a checkpoint.""" + return configured_pipeline_class + def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None: pass @@ -19,9 +29,66 @@ class BasePipelineUtils: pass def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None: + if sampling.enable_frame_interpolation: + raise ValueError( + "Frame interpolation is not supported with the Diffusers backend because " + "Diffusers returns already postprocessed frames." + ) + + def update_call_kwargs( + self, + req: DiffusionRequestBatch, + sampling: OmniDiffusionSamplingParams, + accepted_call_kwargs: set[str] | None, + call_kwargs: dict[str, Any], + ) -> None: pass +class SanaVideoPipelineUtils(BasePipelineUtils): + def resolve_pipeline_class( + self, + od_config: OmniDiffusionConfig, + configured_pipeline_class: type[DiffusionPipeline] | None, + ) -> type[DiffusionPipeline] | None: + # Both released SANA-Video checkpoints declare SanaVideoPipeline in + # model_index.json. Honour an explicit Omni I2V class selection by + # loading the same components into Diffusers' I2V pipeline instead. + if od_config.model_class_name == "SanaImageToVideoPipeline": + from diffusers import SanaImageToVideoPipeline + + return SanaImageToVideoPipeline + return configured_pipeline_class + + def update_call_kwargs( + self, + req: DiffusionRequestBatch, + sampling: OmniDiffusionSamplingParams, + accepted_call_kwargs: set[str] | None, + call_kwargs: dict[str, Any], + ) -> None: + # Diffusers SANA-Video calls the shared request's `num_frames` + # parameter `frames`. + if ( + accepted_call_kwargs is not None + and "frames" in accepted_call_kwargs + and "num_frames" not in accepted_call_kwargs + ): + configured_default = call_kwargs.get("frames", 81) + if configured_default is None: + configured_default = 81 + call_kwargs["frames"] = resolve_video_num_frames( + sampling.num_frames, + default_num_frames=configured_default, + is_dummy_run=req.is_dummy_run(), + ) + + # SANA's startup dummy dimensions can map to a resolution bucket that + # conflicts with the upstream pipeline's divisibility validation. + if req.is_dummy_run() and accepted_call_kwargs is not None and "use_resolution_binning" in accepted_call_kwargs: + call_kwargs["use_resolution_binning"] = False + + class WanPipelineUtils(BasePipelineUtils): def update_load_kwargs(self, od_config: OmniDiffusionConfig, load_kwargs: dict[str, Any]) -> None: if od_config.boundary_ratio is not None: @@ -36,6 +103,7 @@ class WanPipelineUtils(BasePipelineUtils): ) def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None: + super().validate_runtime_sampling_params(sampling) if sampling.boundary_ratio is not None: raise ValueError( "Boundary ratio is not supported at runtime with the diffusers backend for Wan models. Please set " @@ -49,6 +117,8 @@ class WanPipelineUtils(BasePipelineUtils): PIPELINE_UTILS_REGISTRY: dict[str, type[BasePipelineUtils]] = { + "SanaVideoPipeline": SanaVideoPipelineUtils, + "SanaImageToVideoPipeline": SanaVideoPipelineUtils, "WanPipeline": WanPipelineUtils, "WanImageToVideoPipeline": WanPipelineUtils, "WanVACEPipeline": WanPipelineUtils, @@ -62,3 +132,20 @@ def get_pipeline_utils(pipeline_class_name: str | None) -> BasePipelineUtils: return BasePipelineUtils() pipeline_utils_cls = PIPELINE_UTILS_REGISTRY.get(pipeline_class_name, BasePipelineUtils) return pipeline_utils_cls() + + +def get_pipeline_utils_for_config(od_config: OmniDiffusionConfig) -> BasePipelineUtils: + """Return model-specific adapter hooks for the requested or configured class.""" + requested_class_name = od_config.model_class_name + if requested_class_name in PIPELINE_UTILS_REGISTRY: + return get_pipeline_utils(requested_class_name) + + configured_pipeline_class = od_config.diffusers_pipeline_cls + configured_class_name = configured_pipeline_class.__name__ if configured_pipeline_class is not None else None + return get_pipeline_utils(configured_class_name) + + +def resolve_diffusers_pipeline_class(od_config: OmniDiffusionConfig) -> type[DiffusionPipeline] | None: + """Resolve the concrete Diffusers pipeline used for loading and capability checks.""" + pipeline_utils = get_pipeline_utils_for_config(od_config) + return pipeline_utils.resolve_pipeline_class(od_config, od_config.diffusers_pipeline_cls) diff --git a/vllm_omni/diffusion/models/dreamid_omni/__init__.py b/vllm_omni/diffusion/models/dreamid_omni/__init__.py deleted file mode 100644 index 8c7f05bf8..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -__all__ = [ - "DreamIDOmniPipeline", -] - - -def __getattr__(name: str): - if name == "DreamIDOmniPipeline": - from .pipeline_dreamid_omni import DreamIDOmniPipeline - - return DreamIDOmniPipeline - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py deleted file mode 100644 index 8bcbbbb64..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ /dev/null @@ -1,460 +0,0 @@ -import re -from typing import TYPE_CHECKING - -import torch -import torch.nn as nn -from cache_dit import ForwardPattern -from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.logger import init_logger -from vllm.model_executor.layers.linear import ColumnParallelLinear - -from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig - -if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.base_config import QuantizationConfig - -try: - from dreamid_omni.modules.model import WanLayerNorm -except ImportError: - raise ImportError("Failed to import from dependency 'dreamid_omni'.") - -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import ( - DistributedRMSNorm, - WanModel, - rope_apply, -) - -logger = init_logger(__name__) - - -class FusedBlock(nn.Module): - """Wrapper pairing a video block and audio block for layerwise offloading. - - Registers both blocks as submodules so their parameters are visible to the offload hooks. - """ - - def __init__( - self, - vid_block: nn.Module, - audio_block: nn.Module, - device: torch.device, - ): - super().__init__() - self.vid_block = vid_block - self.audio_block = audio_block - self.device = device - - def _cross_attention_forward( - self, - attn: Attention, - cross_attn_block, - src_seq, - src_grid_sizes, - src_freqs, - target_seq, - target_seq_lens, - target_grid_sizes, - target_freqs, - context, - context_lens, - src_ref_lengths=None, - src_freqs_scaling=None, - target_ref_lengths=None, - target_freqs_scaling=None, - ): - b, n, d = src_seq.size(0), cross_attn_block.num_heads, cross_attn_block.head_dim - if hasattr(cross_attn_block, "k_img"): - q, k, v, k_img, v_img = cross_attn_block.qkv_fn(src_seq, context) - else: - q, k, v = cross_attn_block.qkv_fn(src_seq, context) - k_img = v_img = None - - x = attn(q, k, v) - - if k_img is not None: - img_x = attn(q, k_img, v_img) - x = x + img_x - - target_seq = cross_attn_block.pre_attn_norm_fusion(target_seq) - k_target = cross_attn_block.norm_k_fusion(cross_attn_block.k_fusion(target_seq)).view(b, -1, n, d) - v_target = cross_attn_block.v_fusion(target_seq).view(b, -1, n, d) - - q = rope_apply(q, src_grid_sizes, src_freqs, ref_lengths=src_ref_lengths, freqs_scaling=src_freqs_scaling) - k_target = rope_apply( - k_target, - target_grid_sizes, - target_freqs, - ref_lengths=target_ref_lengths, - freqs_scaling=target_freqs_scaling, - ) - - target_x = attn(q, k_target, v_target) - - x = x + target_x - x = x.flatten(2) - x = cross_attn_block.o(x) - return x - - def _cross_attention_ffn_forward( - self, - attn: Attention, - attn_block, - src_seq, - src_grid_sizes, - src_freqs, - target_seq, - target_seq_lens, - target_grid_sizes, - target_freqs, - context, - context_lens, - src_e, - src_ref_lengths=None, - src_freqs_scaling=None, - target_ref_lengths=None, - target_freqs_scaling=None, - ): - src_seq = src_seq + self._cross_attention_forward( - attn, - attn_block.cross_attn, - attn_block.norm3(src_seq), - src_grid_sizes=src_grid_sizes, - src_freqs=src_freqs, - target_seq=target_seq, - target_seq_lens=target_seq_lens, - target_grid_sizes=target_grid_sizes, - target_freqs=target_freqs, - context=context, - context_lens=context_lens, - src_ref_lengths=src_ref_lengths, - src_freqs_scaling=src_freqs_scaling, - target_ref_lengths=target_ref_lengths, - target_freqs_scaling=target_freqs_scaling, - ) - y = attn_block.ffn(attn_block.norm2(src_seq).bfloat16() * (1 + src_e[4].squeeze(2)) + src_e[3].squeeze(2)) - with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True): - src_seq = src_seq + y * src_e[5].squeeze(2) - return src_seq - - def forward( - self, - hidden_states, - encoder_hidden_states, - attn: Attention, - vid_e, - vid_seq_lens, - vid_grid_sizes, - vid_freqs, - vid_context, - vid_context_lens, - vid_ref_lengths, - vid_freqs_scaling, - audio_e, - audio_seq_lens, - audio_grid_sizes, - audio_freqs, - audio_context, - audio_context_lens, - audio_ref_lengths, - audio_freqs_scaling, - ): - vid = hidden_states - audio = encoder_hidden_states - vid_block = self.vid_block - audio_block = self.audio_block - - ## audio modulation - assert audio_e.dtype == torch.bfloat16 - assert len(audio_e.shape) == 4 and audio_e.size(2) == 6 and audio_e.shape[1] == audio.shape[1], ( - f"{audio_e.shape}, {audio.shape}" - ) - with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True): - audio_e = audio_block.modulation(audio_e).chunk(6, dim=2) - assert audio_e[0].dtype == torch.bfloat16 - - # audio self-attention - audio_y = audio_block.self_attn( - audio_block.norm1(audio).bfloat16() * (1 + audio_e[1].squeeze(2)) + audio_e[0].squeeze(2), - audio_seq_lens, - audio_grid_sizes, - audio_freqs, - ref_lengths=audio_ref_lengths, - freqs_scaling=audio_freqs_scaling, - ) - with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True): - audio = audio + audio_y * audio_e[2].squeeze(2) - - ## video modulation - assert len(vid_e.shape) == 4 and vid_e.size(2) == 6 and vid_e.shape[1] == vid.shape[1], ( - f"{vid_e.shape}, {vid.shape}" - ) - with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True): - vid_e = vid_block.modulation(vid_e).chunk(6, dim=2) - - # video self-attention - vid_y = vid_block.self_attn( - vid_block.norm1(vid).bfloat16() * (1 + vid_e[1].squeeze(2)) + vid_e[0].squeeze(2), - vid_seq_lens, - vid_grid_sizes, - vid_freqs, - ref_lengths=vid_ref_lengths, - ) - - with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True): - vid = vid + vid_y * vid_e[2].squeeze(2) - - og_audio = audio - - # audio cross-attention - audio = self._cross_attention_ffn_forward( - attn, - audio_block, - audio, - audio_grid_sizes, - audio_freqs, - vid, - vid_seq_lens, - vid_grid_sizes, - vid_freqs, - audio_context, - audio_context_lens, - audio_e, - src_ref_lengths=audio_ref_lengths, - src_freqs_scaling=audio_freqs_scaling, - target_ref_lengths=vid_ref_lengths, - target_freqs_scaling=None, - ) - - assert not torch.equal(og_audio, audio), "Audio should be changed after cross-attention!" - - # video cross-attention - vid = self._cross_attention_ffn_forward( - attn, - vid_block, - vid, - vid_grid_sizes, - vid_freqs, - og_audio, - audio_seq_lens, - audio_grid_sizes, - audio_freqs, - vid_context, - vid_context_lens, - vid_e, - src_ref_lengths=vid_ref_lengths, - src_freqs_scaling=None, - target_ref_lengths=audio_ref_lengths, - target_freqs_scaling=audio_freqs_scaling, - ) - - return vid, audio - - -class FusionModel(nn.Module): - _cache_dit_adapter_config = CacheDiTAdapterConfig( - block_forward_patterns={ - "fused_blocks": ForwardPattern.Pattern_0, - }, - has_separate_cfg=True, - ) - - _layerwise_offload_blocks_attrs = ["fused_blocks"] - - packed_modules_mapping = { - "to_qkv": ["q", "k", "v"], - } - - @staticmethod - def _is_fused_block(name: str, module) -> bool: - return "fused_blocks" in name and name.split(".")[-1].isdigit() - - _hsdp_shard_conditions = [_is_fused_block] - - def __init__( - self, - video_config=None, - audio_config=None, - *, - quant_config: "QuantizationConfig | None" = None, - ): - super().__init__() - has_video = True - has_audio = True - self.device = get_local_device() - if video_config is not None: - self.video_model = WanModel(quant_config=quant_config, prefix="video_model", **video_config) - else: - has_video = False - self.video_model = None - logger.warning("No video model is provided!") - - if audio_config is not None: - self.audio_model = WanModel(quant_config=quant_config, prefix="audio_model", **audio_config) - else: - has_audio = False - self.audio_model = None - logger.warning("No audio model is provided!") - - if has_video and has_audio: - assert len(self.video_model.blocks) == len(self.audio_model.blocks) - self.num_blocks = len(self.video_model.blocks) - - self.inject_cross_attention_kv_projections(quant_config=quant_config) - - tp_size = get_tensor_model_parallel_world_size() - self.full_num_heads = self.video_model.num_heads - self.num_heads = self.full_num_heads // tp_size - self.head_dim = self.video_model.dim // self.video_model.num_heads - self.attn = Attention( - num_heads=self.num_heads, - head_size=self.head_dim, - num_kv_heads=self.num_heads, - softmax_scale=1.0 / (self.head_dim**0.5), - causal=False, - ) - - if has_video and has_audio: - self.fused_blocks = nn.ModuleList( - [ - FusedBlock( - self.video_model.blocks[i], - self.audio_model.blocks[i], - self.device, - ) - for i in range(self.num_blocks) - ] - ) - - def load_state_dict(self, state_dict, strict=True, assign=False): - """Remap checkpoints where blocks are stored under - `video_model.blocks.N.*` / `audio_model.blocks.N.*` to the current - `fused_blocks.N.vid_block.*` / `fused_blocks.N.audio_block.*`. - """ - needs_remap = any(re.match(r"^(video_model|audio_model)\.blocks\.\d+\.", k) for k in state_dict) - if needs_remap: - remapped = {} - for k, v in state_dict.items(): - new_k = re.sub(r"^video_model\.blocks\.(\d+)\.", r"fused_blocks.\1.vid_block.", k) - new_k = re.sub(r"^audio_model\.blocks\.(\d+)\.", r"fused_blocks.\1.audio_block.", new_k) - remapped[new_k] = v - state_dict = remapped - - self._detach_blocks_from_backbones() - - return super().load_state_dict(state_dict, strict=strict, assign=assign) - - def inject_cross_attention_kv_projections( - self, - quant_config: "QuantizationConfig | None" = None, - ): - def _make_kv_proj(dim: int, prefix: str) -> ColumnParallelLinear: - # K/V projections for the cross-modal (a2v / v2a) attention path. - # Inputs are full latents, so we quantize matched to self_attn. - return ColumnParallelLinear( - dim, - dim, - bias=True, - gather_output=False, - return_bias=False, - quant_config=quant_config, - prefix=prefix, - ) - - tp_size = get_tensor_model_parallel_world_size() - for i, vid_block in enumerate(self.video_model.blocks): - cross = vid_block.cross_attn - block_prefix = f"video_model.blocks.{i}.cross_attn" - cross.k_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.k_fusion") - cross.v_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.v_fusion") - cross.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True) - cross.norm_k_fusion = ( - DistributedRMSNorm(vid_block.dim // tp_size, eps=1e-6) if vid_block.qk_norm else nn.Identity() - ) - - for i, audio_block in enumerate(self.audio_model.blocks): - cross = audio_block.cross_attn - block_prefix = f"audio_model.blocks.{i}.cross_attn" - cross.k_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.k_fusion") - cross.v_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.v_fusion") - cross.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True) - cross.norm_k_fusion = ( - DistributedRMSNorm(audio_block.dim // tp_size, eps=1e-6) if audio_block.qk_norm else nn.Identity() - ) - - def _detach_blocks_from_backbones(self) -> None: - """Keep offloadable blocks owned only by a single place. - - NOTE: This is a special workaround to support layerwise offloading. - The model registers the same Wan blocks under both the video/audio - backbones and `fused_blocks` which is a wrapper for unified blocks - walking through. However, layerwise offloading will only consider - `fused_blocks` as offloadable components and will materialize all - other modules onto device, including the same blocks owned by both - `fused_blocks` and `video_model` and `audio_model`. - """ - video_blocks = list(self.video_model.blocks) - audio_blocks = list(self.audio_model.blocks) - self.video_model._modules.pop("blocks", None) - self.audio_model._modules.pop("blocks", None) - self.video_model.blocks = tuple(video_blocks) - self.audio_model.blocks = tuple(audio_blocks) - - def merge_kwargs(self, vid_kwargs, audio_kwargs): - """ - keys in each kwarg: - e - seq_lens - grid_sizes - freqs - context - context_lens - """ - merged_kwargs = {} - for key in vid_kwargs: - merged_kwargs[f"vid_{key}"] = vid_kwargs[key] - for key in audio_kwargs: - merged_kwargs[f"audio_{key}"] = audio_kwargs[key] - return merged_kwargs - - def forward( - self, - vid, - audio, - t, - vid_context, - audio_context, - vid_seq_len, - audio_seq_len, - ref_ip_lengths=None, - ref_audio_lengths=None, - slg_layer=False, - freqs_scaling=None, - ): - vid, vid_e, vid_kwargs = self.video_model.prepare_transformer_block_kwargs( - x=vid, t=t, context=vid_context, seq_len=vid_seq_len, ref_lengths=ref_ip_lengths - ) - - audio, audio_e, audio_kwargs = self.audio_model.prepare_transformer_block_kwargs( - x=audio, - t=t, - context=audio_context, - seq_len=audio_seq_len, - ref_lengths=ref_audio_lengths, - freqs_scaling=freqs_scaling, - ) - - kwargs = self.merge_kwargs(vid_kwargs, audio_kwargs) - - for fused_block in self.fused_blocks: - vid, audio = fused_block(vid, audio, self.attn, **kwargs) - - vid = self.video_model.post_transformer_block_out(vid, vid_kwargs["grid_sizes"], vid_e) - audio = self.audio_model.post_transformer_block_out(audio, audio_kwargs["grid_sizes"], audio_e) - - return vid, audio - - def set_rope_params(self): - self.video_model.set_rope_params() - self.audio_model.set_rope_params() diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py deleted file mode 100644 index 5555e3708..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ /dev/null @@ -1,608 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import logging -import math -import os -from collections.abc import Iterable -from typing import ClassVar - -import torch -import torch.distributed -from diffusers import FlowMatchEulerDiscreteScheduler -from PIL import Image, ImageOps -from torch import nn -from torchvision.transforms import Compose, Normalize -from tqdm import tqdm -from vllm.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput, SupportsComponentDiscovery -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -try: - from dreamid_omni.utils.divisible_crop import DivisibleCrop - from dreamid_omni.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps - from dreamid_omni.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler - from dreamid_omni.utils.model_loading_utils import ( - init_mmaudio_vae, - init_text_model, - init_wan_vae_2_2, - ) - from dreamid_omni.utils.rearrange import Rearrange - from dreamid_omni.utils.resize import NaResize -except ImportError: - raise ImportError("Failed to import from dependency 'dreamid_omni'.") -from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel - -logger = logging.getLogger(__name__) - - -def get_dreamid_omni_post_process_func(*args, **kwargs): - def post_process(output): - if isinstance(output, tuple) and len(output) == 2: - video, audio = output - return { - "video": video, - "audio": audio, - "audio_sample_rate": 16000, - "fps": 24, - } - return output - - return post_process - - -AUDIO_CONFIG = { - "patch_size": [1], - "model_type": "t2a", - "dim": 3072, - "ffn_dim": 14336, - "freq_dim": 256, - "num_heads": 24, - "num_layers": 30, - "in_dim": 20, - "out_dim": 20, - "text_len": 512, - "window_size": [-1, -1], - "qk_norm": True, - "cross_attn_norm": True, - "eps": 1e-6, - "temporal_rope_scaling_factor": 0.19676, -} - -VIDEO_CONFIG = { - "patch_size": [1, 2, 2], - "model_type": "ti2v", - "dim": 3072, - "ffn_dim": 14336, - "freq_dim": 256, - "num_heads": 24, - "num_layers": 30, - "in_dim": 48, - "out_dim": 48, - "text_len": 512, - "window_size": [-1, -1], - "qk_norm": True, - "cross_attn_norm": True, - "eps": 1e-6, -} - - -class DreamIDOmniPipeline( - nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput, SupportsComponentDiscovery -): - """DreamID-Omni pipeline for vLLM-Omni.""" - - _dit_modules: ClassVar[list[str]] = ["model"] - _encoder_modules: ClassVar[list[str]] = ["text_model"] - _vae_modules: ClassVar[list[str]] = ["vae_model_video", "vae_model_audio"] - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - self.parallel_config = od_config.parallel_config - - self.device = get_local_device() - model = od_config.model - self.target_dtype = od_config.dtype - - # Init Models - ## Load VAEs - vae_model_video = init_wan_vae_2_2(model, rank=self.device) - vae_model_video.model.requires_grad_(False).eval() - vae_model_video.model = vae_model_video.model.bfloat16() - self.vae_model_video = vae_model_video - self.vae = self.vae_model_video - - vae_model_audio = init_mmaudio_vae(model, rank=self.device) - vae_model_audio.requires_grad_(False).eval() - self.vae_model_audio = vae_model_audio.bfloat16() - - # Load T5 text model - self.text_model = init_text_model(model, rank=self.device) - self.text_encoder = self.text_model.model - - # Fusion model — weights are loaded later via load_weights() - quant_config = getattr(self.od_config, "quantization_config", None) - self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG, quant_config=quant_config) - self.transformer = self.model - - fusion_path = self.od_config.tf_model_config.get("fusion", None) - assert fusion_path is not None, "fusion checkpoint path is None in transformer config" - fusion_subfolder = os.path.dirname(fusion_path) or None - fusion_filename = os.path.basename(fusion_path) - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=model, - subfolder=fusion_subfolder, - revision=None, - prefix="model.", - allow_patterns_overrides=[fusion_filename], - ) - ] - - # Fixed attributes, non-configurable - self.audio_latent_channel = AUDIO_CONFIG.get("in_dim") - self.video_latent_channel = VIDEO_CONFIG.get("in_dim") - self.video_latent_length = 31 - self.audio_latent_length = 157 - self.target_area = 960 * 960 - - # Fixed attributes, CFG scales - self.video_cfg_scale = 3.0 - self.video_ref_cfg_scale = 1.5 - self.audio_cfg_scale = 4.0 - self.audio_ref_cfg_scale = 2.0 - - # Schedulers will be set in forward - self.scheduler_video = None - self.scheduler_audio = None - - def load_image_latent_ref_ip_video( - self, - images, - audios, - video_frame_height_width, - ): - # Load size. - patch_size = self.model.video_model.patch_size - vae_stride = [4, 16, 16] - - # Load image and video. - ref_vae_latents = { - "image": [], - "audio": [], - } - video_h = video_frame_height_width[0] - video_w = video_frame_height_width[1] - ref_audio_lengths = [] - for img in images: - # Calculate the required size to keep aspect ratio and fill the rest with padding. - img_ratio = img.width / img.height - target_ratio = video_w / video_h - - if img_ratio > target_ratio: # Image is wider than target - new_width = video_w - new_height = int(new_width / img_ratio) - else: # Image is taller than target - new_height = video_h - new_width = int(new_height * img_ratio) - - img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) - - # Create a new image with the target size and place the resized image in the center - delta_w = video_w - img.size[0] - delta_h = video_h - img.size[1] - padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2)) - new_img = ImageOps.expand(img, padding, fill=(255, 255, 255)) - - # Transform to tensor and normalize. - image_transform = Compose( - [ - NaResize( - resolution=math.sqrt( - video_frame_height_width[0] * video_frame_height_width[1] - ), # 256*448, 480*832 - mode="area", - downsample_only=True, - ), - DivisibleCrop((vae_stride[1] * patch_size[1], vae_stride[2] * patch_size[2])), - Normalize(0.5, 0.5), - Rearrange("t c h w -> c t h w"), - ] - ) - new_img = image_transform([new_img]) - new_img = new_img.transpose(0, 1) - new_img = new_img.to(self.device) - new_img = new_img.to(self.target_dtype) - - img_vae_latent = self.vae_model_video.wrapped_encode(new_img[:, :, None]).to(self.target_dtype).squeeze(0) - ref_vae_latents["image"].append(img_vae_latent) - - for audio_array in audios: - audio_tensor = torch.from_numpy(audio_array).float().unsqueeze(0) - # print(audio_tensor.shape)#torch.Size([1, 81280]) - audio_vae_latent = self.vae_model_audio.wrapped_encode(audio_tensor) - # print(audio_vae_latent.shape)#torch.Size([1, 20, length]) - audio_length = audio_vae_latent.shape[2] - ref_audio_lengths.append(audio_length) - audio_vae_latent = audio_vae_latent.squeeze(0).transpose(0, 1) - ref_vae_latents["audio"].append(audio_vae_latent) - - ref_vae_latents["image"] = torch.cat(ref_vae_latents["image"], dim=1) - - ref_vae_latents["audio"] = torch.cat(ref_vae_latents["audio"], dim=0) - - return ref_vae_latents, ref_audio_lengths - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - pipeline_prefix = "model." - - qkv_stacked_mapping = [ - (".self_attn.to_qkv", ".self_attn.q", "q"), - (".self_attn.to_qkv", ".self_attn.k", "k"), - (".self_attn.to_qkv", ".self_attn.v", "v"), - ] - - params_dict = dict(self.model.named_parameters()) - loaded_local: set[str] = set() - - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0 - - def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor: - """Slice a checkpoint tensor to this rank's TP shard when the - param is a TP-local shape but the checkpoint stores the full one. - """ - if tp_size <= 1 or weight.shape == param.shape: - return weight - if weight.ndim == 1 and weight.numel() == param.numel() * tp_size: - return weight.chunk(tp_size, dim=0)[tp_rank] - if weight.ndim == 2: - if weight.shape[0] == param.shape[0] * tp_size: - return weight.chunk(tp_size, dim=0)[tp_rank] - if weight.shape[1] == param.shape[1] * tp_size: - return weight.chunk(tp_size, dim=1)[tp_rank] - return weight - - for full_name, loaded_weight in weights: - if not full_name.startswith(pipeline_prefix): - continue - name = full_name[len(pipeline_prefix) :] - - # Fused-QKV stacking: route checkpoint's separate - # - ...self_attn.{q,k,v}: weights into the corresponding - # - ...self_attn.to_qkv : via its sharded weight_loader. - stacked_hit = False - for fused_name, raw_name, shard_id in qkv_stacked_mapping: - if raw_name not in name: - continue - stacked = name.replace(raw_name, fused_name) - if stacked not in params_dict: - continue - param = params_dict[stacked] - param.weight_loader(param, loaded_weight, shard_id) - loaded_local.add(stacked) - stacked_hit = True - break - if stacked_hit: - continue - - if name not in params_dict: - logger.warning(f"Skipping DreamID-Omni weight {full_name} -- not found in model parameters") - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", None) - if weight_loader is not None: - # Parallel Linears handle TP sharding internally - weight_loader(param, loaded_weight) - else: - default_weight_loader(param, _maybe_shard_weight(loaded_weight, param)) - loaded_local.add(name) - - # Detach must happen before `DiffusersPipelineLoader._process_weights_after_loading` - # and before any forward pass that triggers layerwise-offload hooks. - self.model._detach_blocks_from_backbones() - - return {pipeline_prefix + k for k in loaded_local} - - def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0): - torch.manual_seed(4) - - if solver_name == "unipc": - sample_scheduler = FlowUniPCMultistepScheduler( - num_train_timesteps=1000, shift=1, use_dynamic_shifting=False - ) - sample_scheduler.set_timesteps(sampling_steps, device=device, shift=shift) - timesteps = sample_scheduler.timesteps - - elif solver_name == "dpm++": - sample_scheduler = FlowDPMSolverMultistepScheduler( - num_train_timesteps=1000, shift=1, use_dynamic_shifting=False - ) - sampling_sigmas = get_sampling_sigmas(sampling_steps, shift=shift) - timesteps, _ = retrieve_timesteps(sample_scheduler, device=device, sigmas=sampling_sigmas) - - elif solver_name == "euler": - sample_scheduler = FlowMatchEulerDiscreteScheduler(shift=shift) - timesteps, sampling_steps = retrieve_timesteps( - sample_scheduler, - sampling_steps, - device=device, - ) - - else: - raise NotImplementedError("Unsupported solver.") - - return sample_scheduler, timesteps - - def predict_noise(self, **kwargs): - pred_vid, pred_audio = self.model(**kwargs) - return (pred_vid[0], pred_audio[0]) - - def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False): - vid_pos, audio_pos = predictions[0] - vid_neg, audio_neg = predictions[1] - vid_ip_neg, _ = predictions[2] - _, refaudio_neg = predictions[3] - - pred_video = ( - vid_neg - + true_cfg_scale["video_cfg_scale"] * (vid_pos - vid_neg) - + true_cfg_scale["video_ref_cfg_scale"] * (vid_pos - vid_ip_neg) - ) - pred_audio = ( - audio_neg - + true_cfg_scale["audio_cfg_scale"] * (audio_pos - audio_neg) - + true_cfg_scale["audio_ref_cfg_scale"] * (audio_pos - refaudio_neg) - ) - return (pred_video, pred_audio) - - def diffuse( - self, - video_noise: torch.Tensor, - audio_noise: torch.Tensor, - latents_ref_image: torch.Tensor, - latents_ref_audio: torch.Tensor, - timesteps_video: torch.Tensor, - timesteps_audio: torch.Tensor, - text_embeddings_video_pos: torch.Tensor, - text_embeddings_video_neg: torch.Tensor, - text_embeddings_audio_pos: torch.Tensor, - text_embeddings_audio_neg: torch.Tensor, - max_seq_len_video: int, - max_seq_len_audio: int, - freqs_scaling_tensor: torch.Tensor, - ref_ip_num: int, - ref_audio_length: int, - ref_audio_lengths: list, - scheduler_video, - scheduler_audio, - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Diffusion loop with CFG parallel support for DreamID-Omni. - """ - for i, (t_v, t_a) in tqdm(enumerate(zip(timesteps_video, timesteps_audio))): - timestep_input = torch.full((1,), t_v, device=self.device) - - model_input_video = torch.cat([video_noise[:, :-ref_ip_num], latents_ref_image], dim=1) - model_input_video_neg = torch.cat( - [video_noise[:, :-ref_ip_num], torch.zeros_like(latents_ref_image)], dim=1 - ) - - model_input_audio = torch.cat([audio_noise[:-ref_audio_length, :], latents_ref_audio], dim=0) - - model_input_audio_neg = torch.cat( - [audio_noise[:-ref_audio_length, :], torch.zeros_like(latents_ref_audio)], dim=0 - ) - ref_ip_lengths = [ref_ip_num] - - common_args = { - "vid_seq_len": max_seq_len_video, - "audio_seq_len": max_seq_len_audio, - "freqs_scaling": freqs_scaling_tensor, - "ref_ip_lengths": [ref_ip_lengths], - "ref_audio_lengths": [ref_audio_lengths], - } - pos_args = { - **common_args, - "audio_context": [text_embeddings_audio_pos], - "vid_context": [text_embeddings_video_pos], - } - neg_args = { - **common_args, - "audio_context": [text_embeddings_audio_neg], - "vid_context": [text_embeddings_video_neg], - } - - branches_kwargs = [ - {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **pos_args}, - {"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **neg_args}, - {"vid": [model_input_video_neg], "audio": [model_input_audio], "t": timestep_input, **pos_args}, - {"vid": [model_input_video], "audio": [model_input_audio_neg], "t": timestep_input, **pos_args}, - ] - - pred_video_guided, pred_audio_guided = self.predict_noise_with_multi_branch_cfg( - do_true_cfg=True, - true_cfg_scale={ - "video_cfg_scale": self.video_cfg_scale, - "video_ref_cfg_scale": self.video_ref_cfg_scale, - "audio_cfg_scale": self.audio_cfg_scale, - "audio_ref_cfg_scale": self.audio_ref_cfg_scale, - }, - branches_kwargs=branches_kwargs, - ) - video_noise = scheduler_video.step( - pred_video_guided.unsqueeze(0), t_v, video_noise.unsqueeze(0), return_dict=False - )[0].squeeze(0) - audio_noise = scheduler_audio.step( - pred_audio_guided.unsqueeze(0), t_a, audio_noise.unsqueeze(0), return_dict=False - )[0].squeeze(0) - return video_noise, audio_noise - - def encode_prompt( - self, - prompt: str, - video_negative_prompt: str = "", - audio_negative_prompt: str = "", - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Encode the positive and negative prompts via ``self.text_model``. - - This is the single text-encoder entrypoint so that the runner-level - prompt-embedding cache (see - ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently - memoize results when the same prompts are submitted repeatedly (e.g. - GRPO rollouts that sample the same prompt with different seeds). - - Returns: - (audio_pos, video_pos, video_neg, audio_neg) embeddings cast to - ``self.target_dtype``. ``audio_pos`` and ``video_pos`` are the - same tensor, matching the original inline behavior. - """ - text_embeddings = self.text_model( - [prompt, video_negative_prompt, audio_negative_prompt], - device=self.device, - ) - text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings] - # Index 0 is the positive embedding, shared between video and audio branches. - return text_embeddings[0], text_embeddings[0], text_embeddings[1], text_embeddings[2] - - def forward( - self, - request: DiffusionRequestBatch, - **kwargs, - ) -> DiffusionOutput: - """Main forward pass for DreamID-Omni pipeline for R2AV task.""" - # Extract parameters from request - r_prompts = request.prompts[0] - if isinstance(r_prompts, str): - prompt = r_prompts - video_negative_prompt = "" - audio_negative_prompt = "" - else: - prompt = r_prompts.get("prompt", "") - video_negative_prompt = r_prompts.get("video_negative_prompt", "") - audio_negative_prompt = r_prompts.get("audio_negative_prompt", "") - - multi_modal_data = r_prompts.get("multi_modal_data", {}) if not isinstance(r_prompts, str) else None - raw_images = multi_modal_data.get("image", None) if multi_modal_data is not None else None - raw_audios = multi_modal_data.get("audio", None) if multi_modal_data is not None else None - if raw_images is None or raw_audios is None: - raise ValueError("This model requires image and audio to run.") - if not isinstance(raw_images, list): - raw_images = [raw_images] - if not isinstance(raw_audios, list): - raw_audios = [raw_audios] - if len(raw_images) != len(raw_audios): - raise ValueError("'image' and 'audio' lists must have the same length.") - - height = request.sampling_params.height - width = request.sampling_params.width - num_inference_steps = request.sampling_params.num_inference_steps - shift = request.sampling_params.extra_args.get("shift", 5.0) - solver_name = request.sampling_params.extra_args.get("solver_name", "unipc") - seed = request.sampling_params.seed if request.sampling_params.seed is not None else 42 - - # 1. Prepare reference latents - ref_vae_latents, ref_audio_lengths = self.load_image_latent_ref_ip_video( - images=raw_images, - audios=raw_audios, - video_frame_height_width=(height, width), - ) - - latents_ref_image = ref_vae_latents["image"] - latents_ref_audio = ref_vae_latents["audio"] - ref_ip_num = latents_ref_image.shape[1] - ref_audio_length = latents_ref_audio.shape[0] - - # 2. scheduler - scheduler_video, timesteps_video = self.get_scheduler_time_steps( - sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift - ) - scheduler_audio, timesteps_audio = self.get_scheduler_time_steps( - sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift - ) - - # 3. text embedding - text_embeddings_audio_pos, text_embeddings_video_pos, text_embeddings_video_neg, text_embeddings_audio_neg = ( - self.encode_prompt( - prompt=prompt, - video_negative_prompt=video_negative_prompt, - audio_negative_prompt=audio_negative_prompt, - ) - ) - - video_latent_h, video_latent_w = height // 16, width // 16 - - video_noise_len = self.video_latent_length + ref_ip_num - audio_noise_len = self.audio_latent_length + ref_audio_length - freqs_scaling_tensor = torch.tensor( - self.video_latent_length / self.audio_latent_length, device=self.device, dtype=self.target_dtype - ) - - video_noise = torch.randn( - (self.video_latent_channel, video_noise_len, video_latent_h, video_latent_w), - device=self.device, - dtype=self.target_dtype, - generator=torch.Generator(device=self.device).manual_seed(seed), - ) - audio_noise = torch.randn( - (audio_noise_len, self.audio_latent_channel), - device=self.device, - dtype=self.target_dtype, - generator=torch.Generator(device=self.device).manual_seed(seed), - ) - - _patch_size_h, _patch_size_w = self.model.video_model.patch_size[1], self.model.video_model.patch_size[2] - - max_seq_len_video = ( - video_noise.shape[1] * video_noise.shape[2] * video_noise.shape[3] // (_patch_size_h * _patch_size_w) - ) - - max_seq_len_audio = audio_noise_len - - with torch.amp.autocast("cuda", enabled=self.target_dtype != torch.float32, dtype=self.target_dtype): - video_noise, audio_noise = self.diffuse( - video_noise, - audio_noise, - latents_ref_image, - latents_ref_audio, - timesteps_video, - timesteps_audio, - text_embeddings_video_pos, - text_embeddings_video_neg, - text_embeddings_audio_pos, - text_embeddings_audio_neg, - max_seq_len_video, - max_seq_len_audio, - freqs_scaling_tensor, - ref_ip_num, - ref_audio_length, - ref_audio_lengths, - scheduler_video, - scheduler_audio, - ) - - video_noise_for_decode = video_noise[:, :-ref_ip_num] - audio_noise_for_decode = audio_noise[:-ref_audio_length, :] - - audio_latents_for_vae = audio_noise_for_decode.unsqueeze(0).transpose(1, 2) - generated_audio = self.vae_model_audio.wrapped_decode(audio_latents_for_vae).squeeze().cpu().float().numpy() - - video_latents_for_vae = video_noise_for_decode.unsqueeze(0) - generated_video = self.vae_model_video.wrapped_decode(video_latents_for_vae).squeeze(0).cpu().float().numpy() - - return DiffusionOutput(output=(generated_video, generated_audio)) diff --git a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py b/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py deleted file mode 100644 index c7e464cc8..000000000 --- a/vllm_omni/diffusion/models/dreamid_omni/wan2_2.py +++ /dev/null @@ -1,694 +0,0 @@ -# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. -import math - -import torch -import torch.amp as amp -import torch.nn as nn -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.modeling_utils import ModelMixin -from vllm.distributed import ( - get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, -) -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.quantization.base_config import QuantizationConfig - -from vllm_omni.diffusion.attention.layer import Attention - -try: - from dreamid_omni.modules.model import ( - ChannelLastConv1d, - ConvMLP, - Head, - MLPProj, - ModulationAdd, - WanLayerNorm, - rope_apply, # diff from wan2.2 - rope_params, - sinusoidal_embedding_1d, - ) -except ImportError: - raise ImportError("Failed to import from dependency 'dreamid_omni'.") - - -class DistributedRMSNorm(nn.Module): - """RMSNorm that computes global RMS across tensor parallel ranks. - - Mirrors `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py::DistributedRMSNorm` - """ - - def __init__(self, hidden_size: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(hidden_size)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - tp_size = get_tensor_model_parallel_world_size() - x_dtype = x.dtype - x_float = x.float() - local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True) - local_count = x.shape[-1] - - if tp_size > 1: - local_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq) - local_count = local_count * tp_size - - rms = torch.sqrt(local_sum_sq / local_count + self.eps) - return (x_float / rms * self.weight.float()).to(x_dtype) - - -class WanSelfAttention(nn.Module): - """Optimized self-attention module using vLLM layers.""" - - def __init__( - self, - dim, - num_heads, - window_size=(-1, -1), - qk_norm=True, - eps=1e-6, - *, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - assert dim % num_heads == 0 - super().__init__() - self.dim = dim - self.head_dim = dim // num_heads - self.window_size = window_size - self.qk_norm = qk_norm - self.eps = eps - - self.to_qkv = QKVParallelLinear( - hidden_size=dim, - head_size=self.head_dim, - total_num_heads=num_heads, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.to_qkv" if prefix else "to_qkv", - ) - self.num_heads = self.to_qkv.num_heads - self.num_kv_heads = self.to_qkv.num_kv_heads - self.tp_inner_dim = self.num_heads * self.head_dim - self.tp_kv_dim = self.num_kv_heads * self.head_dim - - self.o = RowParallelLinear( - dim, - dim, - bias=True, - input_is_parallel=True, - return_bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o" if prefix else "o", - ) - self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = DistributedRMSNorm(self.tp_kv_dim, eps=eps) if qk_norm else nn.Identity() - # Unified attention layer - self.attn = Attention( - num_heads=self.num_heads, - head_size=self.head_dim, - num_kv_heads=self.num_kv_heads, - softmax_scale=1.0 / (self.head_dim**0.5), - causal=False, - prefix=prefix, - ) - - # query, key, value function - def qkv_fn(self, x): - b, s = x.shape[:2] - qkv, _ = self.to_qkv(x) - q_size = self.num_heads * self.head_dim - kv_size = self.num_kv_heads * self.head_dim - q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) - - q = self.norm_q(q).view(b, s, self.num_heads, self.head_dim) - k = self.norm_k(k).view(b, s, self.num_kv_heads, self.head_dim) - v = v.view(b, s, self.num_kv_heads, self.head_dim) - return q, k, v - - def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None): - r""" - Args: - x(Tensor): Shape [B, L, C] - seq_lens(Tensor): Shape [B] - grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) - freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] - ref_lengths(Tensor): Shape [B] - """ - q, k, v = self.qkv_fn(x) - - x = self.attn( - rope_apply(q, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling), - rope_apply(k, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling), - v, - ) - - # output - x = x.flatten(2) - x = self.o(x) - return x - - -class WanT2VCrossAttention(nn.Module): - """Text -> latent cross-attention.""" - - def __init__( - self, - dim, - num_heads, - window_size=(-1, -1), - qk_norm=True, - eps=1e-6, - *, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - assert dim % num_heads == 0 - self.dim = dim - self.head_dim = dim // num_heads - self.window_size = window_size - self.qk_norm = qk_norm - self.eps = eps - - # Separate Q/K/V — Q's input differs from K/V's input (text context) - self.q = ColumnParallelLinear( - dim, - dim, - bias=True, - gather_output=False, - return_bias=False, - quant_config=quant_config, - prefix=f"{prefix}.q" if prefix else "q", - ) - self.k = ColumnParallelLinear( - dim, - dim, - bias=True, - gather_output=False, - return_bias=False, - quant_config=quant_config, - prefix=f"{prefix}.k" if prefix else "k", - ) - self.v = ColumnParallelLinear( - dim, - dim, - bias=True, - gather_output=False, - return_bias=False, - quant_config=quant_config, - prefix=f"{prefix}.v" if prefix else "v", - ) - - tp_size = get_tensor_model_parallel_world_size() - self.num_heads = num_heads // tp_size - self.num_kv_heads = self.num_heads - self.tp_inner_dim = self.num_heads * self.head_dim - - self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() - - self.o = RowParallelLinear( - dim, - dim, - bias=True, - input_is_parallel=True, - return_bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o" if prefix else "o", - ) - - # Kept for interface parity with WanSelfAttention; - # the active cross-attention path in DreamID-Omni uses the shared FusionModel.attn. - # Disable kv quant to prevent affecting quality. - self.attn = Attention( - num_heads=self.num_heads, - head_size=self.head_dim, - num_kv_heads=self.num_kv_heads, - softmax_scale=1.0 / (self.head_dim**0.5), - causal=False, - prefix=prefix, - disable_kv_quant=True, - ) - - def qkv_fn(self, x, context): - b = x.size(0) - - q = self.q(x) - k = self.k(context) - v = self.v(context) - q = self.norm_q(q).view(b, -1, self.num_heads, self.head_dim) - k = self.norm_k(k).view(b, -1, self.num_kv_heads, self.head_dim) - v = v.view(b, -1, self.num_kv_heads, self.head_dim) - - return q, k, v - - def forward(self, x, context, context_lens): - r""" - Args: - x(Tensor): Shape [B, L1, C] - context(Tensor): Shape [B, L2, C] - context_lens(Tensor): Shape [B] - """ - q, k, v = self.qkv_fn(x, context) - - # compute attention - x = self.attn(q, k, v) - - # output - x = x.flatten(2) - x = self.o(x) - return x - - -class WanI2VCrossAttention(nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - - def forward(self, *args, **kwargs): - raise RuntimeError("WanT2VCrossAttention is currently disabled and should not be used.") - - -WAN_CROSSATTENTION_CLASSES = { - "t2v_cross_attn": WanT2VCrossAttention, - "i2v_cross_attn": WanI2VCrossAttention, -} - - -class WanAttentionBlock(nn.Module): - def __init__( - self, - cross_attn_type, - dim, - ffn_dim, - num_heads, - window_size=(-1, -1), - qk_norm=True, - cross_attn_norm=False, - eps=1e-6, - additional_emb_length=None, - *, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.dim = dim - self.ffn_dim = ffn_dim - self.num_heads = num_heads - self.window_size = window_size - self.qk_norm = qk_norm - self.cross_attn_norm = cross_attn_norm - self.eps = eps - - # layers - self.norm1 = WanLayerNorm(dim, eps) - self.self_attn = WanSelfAttention( - dim, - num_heads, - window_size, - qk_norm, - eps, - quant_config=quant_config, - prefix=f"{prefix}.self_attn" if prefix else "self_attn", - ) - self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() - if cross_attn_type == "i2v_cross_attn": - assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn" - self.cross_attn = WanI2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, additional_emb_length) - else: - assert additional_emb_length is None, "additional_emb_length should be None for t2v_cross_attn" - self.cross_attn = WanT2VCrossAttention( - dim, - num_heads, - (-1, -1), - qk_norm, - eps, - quant_config=quant_config, - prefix=f"{prefix}.cross_attn" if prefix else "cross_attn", - ) - self.norm2 = WanLayerNorm(dim, eps) - # FFN layout: Stay compatible with upstream DreamID-Omni checkpoints - self.ffn = nn.Sequential( - ColumnParallelLinear( - dim, - ffn_dim, - bias=True, - gather_output=False, - return_bias=False, - quant_config=quant_config, - prefix=f"{prefix}.ffn.0" if prefix else "ffn.0", - ), - nn.GELU(approximate="tanh"), - RowParallelLinear( - ffn_dim, - dim, - bias=True, - input_is_parallel=True, - return_bias=False, - quant_config=quant_config, - prefix=f"{prefix}.ffn.2" if prefix else "ffn.2", - ), - ) - - # modulation - self.modulation = ModulationAdd(dim, 6) - - def forward( - self, - x, - e, - seq_lens, - grid_sizes, - freqs, - context, - context_lens, - ): - r""" - Args: - x(Tensor): Shape [B, L, C] - e(Tensor): Shape [B, L1, 6, C] - seq_lens(Tensor): Shape [B], length of each sequence in batch - grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) - freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] - """ - assert e.dtype == torch.bfloat16 - assert len(e.shape) == 4 and e.size(2) == 6 and e.shape[1] == x.shape[1], f"{e.shape}, {x.shape}" - with amp.autocast("cuda", dtype=torch.bfloat16): - e = self.modulation(e).chunk(6, dim=2) - assert e[0].dtype == torch.bfloat16 - - # self-attention - y = self.self_attn( - self.norm1(x).bfloat16() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), seq_lens, grid_sizes, freqs - ) - with amp.autocast("cuda", dtype=torch.bfloat16): - x = x + y * e[2].squeeze(2) - - # cross-attention & ffn function - def cross_attn_ffn(x, context, context_lens, e): - x = x + self.cross_attn(self.norm3(x), context, context_lens) - y = self.ffn(self.norm2(x).bfloat16() * (1 + e[4].squeeze(2)) + e[3].squeeze(2)) - with amp.autocast("cuda", dtype=torch.bfloat16): - x = x + y * e[5].squeeze(2) - return x - - x = cross_attn_ffn(x, context, context_lens, e) - return x - - -class WanModel(ModelMixin, ConfigMixin): - r""" - Wan diffusion backbone supporting both text-to-video and image-to-video, text-to-audio. - """ - - @register_to_config - def __init__( - self, - model_type="t2v", - patch_size=(1, 2, 2), - text_len=512, - in_dim=16, - dim=2048, - ffn_dim=8192, - freq_dim=256, - text_dim=4096, - additional_emb_dim=None, - additional_emb_length=None, - out_dim=16, - num_heads=16, - num_layers=32, - window_size=(-1, -1), - qk_norm=True, - cross_attn_norm=True, - temporal_rope_scaling_factor=1.0, - eps=1e-6, - *, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - r""" - Initialize the diffusion model backbone. - - Args: - model_type (`str`, *optional*, defaults to 't2v'): - Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) - patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): - 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) - text_len (`int`, *optional*, defaults to 512): - Fixed length for text embeddings - in_dim (`int`, *optional*, defaults to 16): - Input video channels (C_in) - dim (`int`, *optional*, defaults to 2048): - Hidden dimension of the transformer - ffn_dim (`int`, *optional*, defaults to 8192): - Intermediate dimension in feed-forward network - freq_dim (`int`, *optional*, defaults to 256): - Dimension for sinusoidal time embeddings - text_dim (`int`, *optional*, defaults to 4096): - Input dimension for text embeddings - out_dim (`int`, *optional*, defaults to 16): - Output video channels (C_out) - num_heads (`int`, *optional*, defaults to 16): - Number of attention heads - num_layers (`int`, *optional*, defaults to 32): - Number of transformer blocks - window_size (`tuple`, *optional*, defaults to (-1, -1)): - Window size for local attention (-1 indicates global attention) - qk_norm (`bool`, *optional*, defaults to True): - Enable query/key normalization - cross_attn_norm (`bool`, *optional*, defaults to False): - Enable cross-attention normalization - eps (`float`, *optional*, defaults to 1e-6): - Epsilon value for normalization layers - """ - - super().__init__() - - assert model_type in [ - "t2v", - "i2v", - "t2a", - "tt2a", - "ti2v", - ] ## tt2a means text transcript + text description to audio (to support both TTS and T2A - self.model_type = model_type - is_audio_type = "a" in self.model_type - is_video_type = "v" in self.model_type - assert is_audio_type ^ is_video_type, "Either audio or video model should be specified" - if is_audio_type: - ## audio model - assert len(patch_size) == 1 and patch_size[0] == 1, ( - "Audio model should only accept 1 dimensional input, and we dont do patchify" - ) - - self.patch_size = patch_size - self.text_len = text_len - self.in_dim = in_dim - self.dim = dim - self.ffn_dim = ffn_dim - self.freq_dim = freq_dim - self.text_dim = text_dim - self.out_dim = out_dim - self.num_heads = num_heads - self.num_layers = num_layers - self.window_size = window_size - self.qk_norm = qk_norm - self.cross_attn_norm = cross_attn_norm - self.eps = eps - self.temporal_rope_scaling_factor = temporal_rope_scaling_factor - self.is_audio_type = is_audio_type - self.is_video_type = is_video_type - # embeddings - if is_audio_type: - ## hardcoded to MMAudio - self.patch_embedding = nn.Sequential( - ChannelLastConv1d(in_dim, dim, kernel_size=7, padding=3), - nn.SiLU(), - ConvMLP(dim, dim * 4, kernel_size=7, padding=3), - ) - else: - self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) - - self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim)) - - self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) - self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) - - # blocks - ## so i2v and tt2a share the same cross attention while t2v and t2a share the same cross attention - cross_attn_type = "t2v_cross_attn" if model_type in ["t2v", "t2a", "ti2v"] else "i2v_cross_attn" - - if cross_attn_type == "t2v_cross_attn": - assert additional_emb_dim is None and additional_emb_length is None, ( - "additional_emb_length should be None for t2v and t2a model" - ) - else: - assert additional_emb_dim is not None and additional_emb_length is not None, ( - "additional_emb_length should be specified for i2v and tt2a model" - ) - - self.blocks = nn.ModuleList( - [ - WanAttentionBlock( - cross_attn_type, - dim, - ffn_dim, - num_heads, - window_size, - qk_norm, - cross_attn_norm, - eps, - additional_emb_length, - quant_config=quant_config, - prefix=f"{prefix}.blocks.{layer_idx}" if prefix else f"blocks.{layer_idx}", - ) - for layer_idx in range(num_layers) - ] - ) - - # head - self.head = Head(dim, out_dim, patch_size, eps) - - self.set_rope_params() - - if model_type in ["i2v", "tt2a"]: - self.img_emb = MLPProj(additional_emb_dim, dim) - - def set_rope_params(self): - # buffers (don't use register_buffer otherwise dtype will be changed in to()) - dim = self.dim - num_heads = self.num_heads - assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 - d = dim // num_heads - - if self.is_audio_type: - ## to be determined - # self.freqs = rope_params(1024, d, freqs_scaling=temporal_rope_scaling_factor) - self.freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=self.temporal_rope_scaling_factor) - else: - self.freqs = torch.cat( - [rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))], - dim=1, - ) - - def prepare_transformer_block_kwargs( - self, - x, - t, - context, - seq_len, - ref_lengths=None, - freqs_scaling=None, - ): - # params - ## need to change! - device = next(self.patch_embedding.parameters()).device - - if self.is_audio_type and freqs_scaling is not None: - if isinstance(freqs_scaling, torch.Tensor): - scale_val = freqs_scaling.item() - else: - scale_val = freqs_scaling - d = self.dim // self.num_heads - - current_freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=scale_val).to(device) - else: - current_freqs = self.freqs - if current_freqs.device != device: - current_freqs = current_freqs.to(device) - - # embeddings - x = [self.patch_embedding(u.unsqueeze(0)) for u in x] ## x is list of [B L D] or [B C F H W] - if self.is_audio_type: - # [B, 1] - grid_sizes = torch.stack([torch.tensor(u.shape[1:2], dtype=torch.long) for u in x]) - else: - # [B, 3] - grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) - x = [u.flatten(2).transpose(1, 2) for u in x] # [B C F H W] -> [B (F H W) C] -> [B L C] - - seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) - assert seq_lens.max() <= seq_len, f"Sequence length {seq_lens.max()} exceeds maximum {seq_len}." - x = torch.cat( - [torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x] - ) # single [B, L, C] - - # time embeddings - if t.dim() == 1: - t = t.unsqueeze(1).expand(t.size(0), seq_len) - with amp.autocast("cuda", dtype=torch.bfloat16): - bt = t.size(0) - t = t.flatten() - e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, seq_len)).bfloat16()) - e0 = self.time_projection(e).unflatten(2, (6, self.dim)) # [1, 26784, 6, 3072] - B, seq_len, 6, dim - assert e.dtype == torch.bfloat16 and e0.dtype == torch.bfloat16 - - # context - context_lens = None - context = self.text_embedding( - torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context]) - ) - - # arguments - kwargs = dict( - e=e0, - seq_lens=seq_lens, - grid_sizes=grid_sizes, - freqs=current_freqs, - context=context, - context_lens=context_lens, - ref_lengths=ref_lengths, - freqs_scaling=freqs_scaling, - ) - - return x, e, kwargs - - def post_transformer_block_out(self, x, grid_sizes, e): - # head - x = self.head(x, e) - # unpatchify - if self.is_audio_type: - ## grid_sizes is [B 1] where 1 is L, - # converting grid_sizes from [B 1] -> [B] - grid_sizes = [gs[0] for gs in grid_sizes] - assert len(x) == len(grid_sizes) - x = [u[:gs] for u, gs in zip(x, grid_sizes)] - else: - ## grid_sizes is [B 3] where 3 is F H w - x = self.unpatchify(x, grid_sizes) - - return [u.bfloat16() for u in x] - - def forward(self, *args, **kwargs): - raise NotImplementedError("WanModel model does not support forward pass.") - - def unpatchify(self, x, grid_sizes) -> list[torch.Tensor]: - r""" - Reconstruct video tensors from patch embeddings. - - Args: - x (List[Tensor]): - List of patchified features, each with shape [L, C_out * prod(patch_size)] - grid_sizes (Tensor): - Original spatial-temporal grid dimensions before patching, - shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) - - Returns: - List[Tensor]: - Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] - """ - - c = self.out_dim - out = [] - for u, v in zip(x, grid_sizes.tolist()): - # v is [F H w] F * H * 80, 100, it was right padded by 20. - u = u[: math.prod(v)].view(*v, *self.patch_size, c) - u = torch.einsum("fhwpqrc->cfphqwr", u) - u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) - out.append(u) - # out is list of [C F H W] - return out diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index b2668b0a5..693766332 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -70,6 +70,9 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 +# Shipped DreamZero geometry retains 24 Wan VAE causal-convolution cache +# entries. This is the measured persistent CUDA upper bound per live session. +DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION = 603 * 1024 * 1024 # The pipeline's per-session state is a bespoke ``DreamZeroState`` by default, or # a ``DreamZeroStateAdapter`` view when the opt-in session manager is enabled. @@ -149,6 +152,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): session_capacity=MAX_DREAMZERO_SESSIONS, cross_attention=tuple(cross_attention), max_scratch_tokens_per_branch=int(transformer.num_action_per_block + transformer.num_state_per_block), + model_owned_state_bytes_per_session=DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION, ) @contextmanager diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 0dbfe0f8c..139b5275a 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -39,16 +39,11 @@ from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import resolve_component_quant_config as _resolve_component_quant_config logger = logging.getLogger(__name__) -def _resolve_component_quant_config(quant_config, component: str): - if hasattr(quant_config, "resolve"): - return quant_config.resolve(component) - return quant_config - - class Flux2ImageProcessor(VaeImageProcessor): """Image processor to preprocess the reference image for Flux2.""" @@ -443,6 +438,12 @@ class Flux2Pipeline( self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) + + if hasattr(self.text_encoder.config, "text_encoder_out_layers"): + self.text_encoder_out_layers = tuple(self.text_encoder.config.text_encoder_out_layers) + else: + self.text_encoder_out_layers = (10, 20, 30) + self._current_timestep = None self._interrupt = False @@ -682,7 +683,6 @@ class Flux2Pipeline( num_images_per_prompt: int = 1, prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 512, - text_encoder_out_layers: tuple[int, ...] = (10, 20, 30), ): device = device or self._execution_device @@ -699,7 +699,7 @@ class Flux2Pipeline( device=device, max_sequence_length=max_sequence_length, system_message=self.system_message, - hidden_states_layers=text_encoder_out_layers, + hidden_states_layers=self.text_encoder_out_layers, ) batch_size, seq_len, _ = prompt_embeds.shape @@ -891,7 +891,6 @@ class Flux2Pipeline( callback_on_step_end: Callable[[int, int, dict], None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, - text_encoder_out_layers: tuple[int, ...] = (10, 20, 30), caption_upsample_temperature: float = None, ) -> DiffusionOutput: if len(req.prompts) > 1: @@ -927,7 +926,6 @@ class Flux2Pipeline( else num_images_per_prompt ) max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers) caption_upsample_temperature = req.sampling_params.extra_args.get( "caption_upsample_temperature", caption_upsample_temperature ) @@ -982,7 +980,6 @@ class Flux2Pipeline( device=device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, - text_encoder_out_layers=text_encoder_out_layers, ) has_neg_prompt = negative_prompt_embeds is not None or any(req_negative_prompt) @@ -998,7 +995,6 @@ class Flux2Pipeline( device=device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, - text_encoder_out_layers=text_encoder_out_layers, ) # 4. process images diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index 72bed54cb..b10b63586 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -49,12 +49,23 @@ from vllm_omni.diffusion.distributed.sp_plan import ( ) from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract logger = init_logger(__name__) if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +_FLUX2_STACKED_PARAMS_MAPPING = ( + (".to_qkv.", ".to_q.", "q"), + (".to_qkv.", ".to_k.", "k"), + (".to_qkv.", ".to_v.", "v"), + (".add_kv_proj", ".add_q_proj", "q"), + (".add_kv_proj", ".add_k_proj", "k"), + (".add_kv_proj", ".add_v_proj", "v"), +) + + def _join_prefix(prefix: str, name: str) -> str: return f"{prefix}.{name}" if prefix else name @@ -770,6 +781,14 @@ class Flux2Transformer2DModel(nn.Module): Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. """ + # Constructor state plus complete final-layout parameters and persistent + # buffers is sufficient to reconstruct a ready base-model transformer. + # Changes to the packed mapping or restore invariants require a version bump. + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="flux2-klein-dit", + version="1", + ) + _cache_dit_adapter_config = CacheDiTAdapterConfig( block_forward_patterns={ "transformer_blocks": ForwardPattern.Pattern_1, @@ -817,6 +836,9 @@ class Flux2Transformer2DModel(nn.Module): quant_config: "QuantizationConfig | None" = None, ): super().__init__() + # Warm HWR restoration skips load_weights(), so loader consumers such + # as LoRA discovery must not depend on that method to create this state. + self.stacked_params_mapping = list(_FLUX2_STACKED_PARAMS_MAPPING) self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.config = SimpleNamespace( @@ -1004,18 +1026,78 @@ class Flux2Transformer2DModel(nn.Module): return (output,) return Transformer2DModelOutput(sample=output) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - (".to_qkv.", ".to_q.", "q"), - (".to_qkv.", ".to_k.", "k"), - (".to_qkv.", ".to_v.", "v"), - (".add_kv_proj", ".add_q_proj", "q"), - (".add_kv_proj", ".add_k_proj", "k"), - (".add_kv_proj", ".add_v_proj", "v"), - ] - # Expose packed shard mappings for LoRA handling of fused projections. - self.stacked_params_mapping = stacked_params_mapping + def validate_restored_host_weights(self) -> None: + """Validate the tensor-complete BF16 base-model restore contract.""" + if not isinstance(self.stacked_params_mapping, (list, tuple)) or ( + tuple(self.stacked_params_mapping) != _FLUX2_STACKED_PARAMS_MAPPING + ): + raise ValueError("FLUX.2-klein packed QKV mapping is not constructor-stable") + stack_specs = ( + ("transformer_blocks", self.config.num_layers, Flux2TransformerBlock), + ("single_transformer_blocks", self.config.num_single_layers, Flux2SingleTransformerBlock), + ) + for name, expected_count, block_type in stack_specs: + blocks = getattr(self, name, None) + if not isinstance(blocks, nn.ModuleList) or len(blocks) != expected_count or expected_count <= 0: + raise ValueError( + f"FLUX.2-klein {name} is incomplete: expected {expected_count} blocks, " + f"got {len(blocks) if isinstance(blocks, nn.ModuleList) else 'missing'}" + ) + if any(not isinstance(block, block_type) for block in blocks): + raise ValueError(f"FLUX.2-klein {name} contains an unexpected block implementation") + + required_modules = ( + "pos_embed", + "rope_prepare", + "time_guidance_embed", + "double_stream_modulation_img", + "double_stream_modulation_txt", + "single_stream_modulation", + "x_embedder", + "context_embedder", + "norm_out", + "proj_out", + ) + for name in required_modules: + if not isinstance(getattr(self, name, None), nn.Module): + raise ValueError(f"FLUX.2-klein required module {name!r} is missing") + + parameter_count = 0 + stack_parameter_coverage = {name: False for name, _, _ in stack_specs} + for name, parameter in self.named_parameters(): + parameter_count += 1 + for stack_name in stack_parameter_coverage: + if name.startswith(f"{stack_name}."): + stack_parameter_coverage[stack_name] = True + if parameter.is_meta or parameter.device.type != "cpu": + raise ValueError(f"FLUX.2-klein parameter {name!r} is not materialized on CPU") + if parameter.dtype is not torch.bfloat16: + raise ValueError(f"FLUX.2-klein parameter {name!r} must stay bf16, got {parameter.dtype}") + if parameter.layout is not torch.strided or not parameter.is_contiguous(): + raise ValueError(f"FLUX.2-klein parameter {name!r} must use contiguous strided layout") + if parameter.numel() == 0: + raise ValueError(f"FLUX.2-klein parameter {name!r} is empty") + if parameter_count == 0 or not all(stack_parameter_coverage.values()): + raise ValueError("FLUX.2-klein restore did not materialize both transformer block stacks") + + for name, buffer in self.named_buffers(): + parent_path, _, leaf_name = name.rpartition(".") + owner = self.get_submodule(parent_path) + persistent = leaf_name not in owner._non_persistent_buffers_set + loader_buffer = name.endswith((".beta", ".eps")) + if loader_buffer and not persistent: + raise ValueError(f"FLUX.2-klein loader buffer {name!r} must be persistent") + if not persistent: + continue + if buffer.is_meta or buffer.device.type != "cpu": + raise ValueError(f"FLUX.2-klein persistent buffer {name!r} is not materialized on CPU") + if buffer.layout is not torch.strided or not buffer.is_contiguous(): + raise ValueError(f"FLUX.2-klein persistent buffer {name!r} must use contiguous strided layout") + if loader_buffer and (buffer.dtype is not torch.bfloat16 or buffer.numel() != 1): + raise ValueError(f"FLUX.2-klein loader buffer {name!r} must be a scalar bf16 tensor") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) for name, buffer in self.named_buffers(): @@ -1026,7 +1108,7 @@ class Flux2Transformer2DModel(nn.Module): for name, loaded_weight in weights: original_name = name mapped = False - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, shard_id in self.stacked_params_mapping: if weight_name not in original_name: continue name = original_name.replace(weight_name, param_name) diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index a4350373b..beb0b3787 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -1020,10 +1020,8 @@ class Flux2KleinPipeline( if reference_image_latents is not None: latent_model_input = torch.cat([latents, reference_image_latents], dim=1) - latent_image_ids = latent_ids elif image_latents is not None: latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype) - latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1) positive_kwargs = { "hidden_states": latent_model_input, diff --git a/vllm_omni/diffusion/models/hidream_o1_image/__init__.py b/vllm_omni/diffusion/models/hidream_o1_image/__init__.py new file mode 100644 index 000000000..fe2723c76 --- /dev/null +++ b/vllm_omni/diffusion/models/hidream_o1_image/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import ( + HiDreamO1ImageTransformer, +) +from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import ( + HiDreamO1ImagePipeline, + get_hidream_o1_image_post_process_func, +) + +__all__ = [ + "HiDreamO1ImageTransformer", + "HiDreamO1ImagePipeline", + "get_hidream_o1_image_post_process_func", +] diff --git a/vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py b/vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py new file mode 100644 index 000000000..8a0b1955c --- /dev/null +++ b/vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. +# +# 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. + +from __future__ import annotations + +import math +from collections.abc import Iterable +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from cache_dit import ForwardPattern +from torch import nn +from transformers import Qwen3VLConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3 import Qwen3MLP + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + + +def _prefix(parent: str, child: str) -> str: + return f"{parent}.{child}" if parent else child + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rotary_pos_emb( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return ( + query * cos + _rotate_half(query) * sin, + key * cos + _rotate_half(key) * sin, + ) + + +class HiDreamO1RotaryEmbedding(nn.Module): + def __init__(self, config) -> None: + super().__init__() + rope_parameters = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {} + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + rope_theta = rope_parameters.get("rope_theta", getattr(config, "rope_theta", 10000.0)) + inv_freq = 1.0 / (float(rope_theta) ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.mrope_section = rope_parameters.get("mrope_section", [24, 20, 20]) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if position_ids.ndim == 2: + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + inv_freq = self.inv_freq[None, None, :, None].float().to(hidden_states.device) + inv_freq = inv_freq.expand(3, position_ids.shape[1], -1, 1) + positions = position_ids[:, :, None, :].float() + device_type = hidden_states.device.type if hidden_states.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq @ positions).transpose(2, 3) + freqs_t = freqs[0] + for dim, offset in enumerate((1, 2), start=1): + length = self.mrope_section[dim] * 3 + freqs_t[..., offset:length:3] = freqs[dim, ..., offset:length:3] + emb = torch.cat((freqs_t, freqs_t), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(hidden_states.dtype), sin.to(hidden_states.dtype) + + +class HiDreamO1Attention(nn.Module): + def __init__( + self, + config, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + tp_size = get_tensor_model_parallel_world_size() + if config.num_attention_heads % tp_size != 0: + raise ValueError( + f"num_attention_heads={config.num_attention_heads} must be divisible by tensor_parallel_size={tp_size}" + ) + if config.num_key_value_heads >= tp_size: + if config.num_key_value_heads % tp_size != 0: + raise ValueError( + f"num_key_value_heads={config.num_key_value_heads} must be divisible by " + f"tensor_parallel_size={tp_size}" + ) + elif tp_size % config.num_key_value_heads != 0: + raise ValueError( + f"tensor_parallel_size={tp_size} must be divisible by num_key_value_heads={config.num_key_value_heads}" + ) + + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + config.num_attention_heads, + config.num_key_value_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=_prefix(prefix, "qkv_proj"), + ) + self.o_proj = RowParallelLinear( + config.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=_prefix(prefix, "o_proj"), + ) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.scaling = self.head_dim**-0.5 + self.attn = Attention( + num_heads=self.qkv_proj.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.scaling, + num_kv_heads=self.qkv_proj.num_kv_heads, + prefix=_prefix(prefix, "attn"), + role="self", + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + full_attn_spans: list[list[tuple[int, int]]] | None, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + qkv, _ = self.qkv_proj(hidden_states) + q_size = self.qkv_proj.num_heads * self.head_dim + kv_size = self.qkv_proj.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + + query = query.view(batch_size, seq_len, self.qkv_proj.num_heads, self.head_dim).transpose(1, 2) + key = key.view(batch_size, seq_len, self.qkv_proj.num_kv_heads, self.head_dim).transpose(1, 2) + value = value.view(batch_size, seq_len, self.qkv_proj.num_kv_heads, self.head_dim).transpose(1, 2) + query = self.q_norm(query) + key = self.k_norm(key) + query, key = _apply_rotary_pos_emb(query, key, *position_embeddings) + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + attn_output = self.attn( + query, + key, + value, + AttentionMetadata( + attn_mask=attention_mask, + full_attn_spans=full_attn_spans, + ), + ) + attn_output = attn_output.reshape(batch_size, seq_len, q_size) + attn_output, _ = self.o_proj(attn_output) + return attn_output + + +class HiDreamO1DecoderLayer(nn.Module): + def __init__( + self, + config, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = HiDreamO1Attention( + config, + quant_config=quant_config, + prefix=_prefix(prefix, "self_attn"), + ) + self.mlp = Qwen3MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=_prefix(prefix, "mlp"), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + full_attn_spans: list[list[tuple[int, int]]] | None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states, + position_embeddings, + attention_mask, + full_attn_spans, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class HiDreamO1TextModel(nn.Module): + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={"layers": ForwardPattern.Pattern_3}, + has_separate_cfg=True, + ) + + def __init__( + self, + config, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=_prefix(prefix, "embed_tokens"), + ) + self.layers = nn.ModuleList( + [ + HiDreamO1DecoderLayer( + config, + quant_config=quant_config, + prefix=_prefix(prefix, f"layers.{layer_idx}"), + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = HiDreamO1RotaryEmbedding(config) + + def forward( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor | None, + full_attn_spans: list[list[tuple[int, int]]] | None, + ) -> torch.Tensor: + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer( + hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + full_attn_spans=full_attn_spans, + ) + return self.norm(hidden_states) + + +class TimestepEmbedder(nn.Module): + def __init__( + self, + hidden_size: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + frequency_embedding_size: int = 256, + ) -> None: + super().__init__() + self.mlp = nn.Sequential( + ColumnParallelLinear( + frequency_embedding_size, + hidden_size, + bias=True, + gather_output=False, + quant_config=quant_config, + prefix=_prefix(prefix, "mlp.0"), + return_bias=False, + ), + nn.SiLU(), + RowParallelLinear( + hidden_size, + hidden_size, + bias=True, + input_is_parallel=True, + quant_config=quant_config, + prefix=_prefix(prefix, "mlp.2"), + return_bias=False, + ), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(timestep: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor: + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(0, half, dtype=torch.float32, device=timestep.device) / half + ) + args = timestep[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + embedding = self.timestep_embedding(timestep * 1000, self.frequency_embedding_size) + return self.mlp(embedding.to(self.mlp[0].weight.dtype)) + + +class BottleneckPatchEmbed(nn.Module): + def __init__( + self, + hidden_size: int, + *, + patch_size: int, + in_channels: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + bottleneck_dim = hidden_size // 4 + self.proj1 = ColumnParallelLinear( + patch_size * patch_size * in_channels, + bottleneck_dim, + bias=False, + gather_output=False, + quant_config=quant_config, + prefix=_prefix(prefix, "proj1"), + return_bias=False, + ) + self.proj2 = RowParallelLinear( + bottleneck_dim, + hidden_size, + bias=True, + input_is_parallel=True, + quant_config=quant_config, + prefix=_prefix(prefix, "proj2"), + return_bias=False, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.proj2(self.proj1(hidden_states)) + + +class FinalLayer(nn.Module): + def __init__( + self, + hidden_size: int, + *, + patch_size: int, + out_channels: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.linear = ColumnParallelLinear( + hidden_size, + patch_size * patch_size * out_channels, + bias=True, + gather_output=True, + quant_config=quant_config, + prefix=_prefix(prefix, "linear"), + return_bias=False, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.linear(hidden_states) + + +@dataclass +class HiDreamO1ImageOutput: + x_pred: torch.Tensor + + +class HiDreamO1ImageModel(nn.Module): + def __init__( + self, + config: Qwen3VLConfig, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.language_model = HiDreamO1TextModel( + config.text_config, + quant_config=quant_config, + prefix=_prefix(prefix, "language_model"), + ) + self.patch_size = 32 + self.in_channels = 3 + hidden_size = config.text_config.hidden_size + self.t_embedder1 = TimestepEmbedder( + hidden_size, + quant_config=quant_config, + prefix=_prefix(prefix, "t_embedder1"), + ) + self.x_embedder = BottleneckPatchEmbed( + hidden_size, + patch_size=self.patch_size, + in_channels=self.in_channels, + quant_config=quant_config, + prefix=_prefix(prefix, "x_embedder"), + ) + self.final_layer2 = FinalLayer( + hidden_size, + patch_size=self.patch_size, + out_channels=self.in_channels, + quant_config=quant_config, + prefix=_prefix(prefix, "final_layer2"), + ) + self.tms_token_id = 151673 + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + vinputs: torch.Tensor, + timestep: torch.Tensor, + attention_mask: torch.Tensor | None, + full_attn_spans: list[list[tuple[int, int]]] | None, + ) -> HiDreamO1ImageOutput: + inputs_embeds = self.language_model.embed_tokens(input_ids) + timestep_embeds = self.t_embedder1(timestep.to(inputs_embeds.device)) + timestep_mask = input_ids == self.tms_token_id + inputs_embeds = torch.where( + timestep_mask.unsqueeze(-1), + timestep_embeds.unsqueeze(1).expand_as(inputs_embeds), + inputs_embeds, + ) + + image_embeds = self.x_embedder(vinputs.to(inputs_embeds.device)).to(inputs_embeds.dtype) + inputs_embeds = torch.cat([inputs_embeds, image_embeds], dim=1) + batch_size, total_seq_len, _ = inputs_embeds.shape + + if attention_mask is not None: + if attention_mask.shape[0] == 1 and batch_size > 1: + attention_mask = attention_mask.expand(batch_size, -1, -1, -1) + expected_mask_shape = (batch_size, 1, total_seq_len, total_seq_len) + if attention_mask.shape != expected_mask_shape: + raise ValueError( + f"attention_mask must have shape {expected_mask_shape}, got {tuple(attention_mask.shape)}" + ) + attention_mask = attention_mask.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + + if full_attn_spans is None or len(full_attn_spans) != batch_size: + raise ValueError( + f"full_attn_spans must contain one entry per batch item; got {full_attn_spans!r} " + f"for batch_size={batch_size}" + ) + + hidden_states = self.language_model( + inputs_embeds, + position_ids, + attention_mask, + full_attn_spans, + ) + return HiDreamO1ImageOutput(x_pred=self.final_layer2(hidden_states)) + + +class HiDreamO1ImageTransformer(nn.Module): + def __init__( + self, + config: Qwen3VLConfig, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = HiDreamO1ImageModel( + config, + quant_config=quant_config, + prefix=_prefix(prefix, "model"), + ) + + @property + def supports_piecewise_attention(self) -> bool: + return all( + layer.self_attn.attn.attn_backend.supports_piecewise_spans for layer in self.model.language_model.layers + ) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + vinputs: torch.Tensor, + timestep: torch.Tensor, + attention_mask: torch.Tensor | None, + full_attn_spans: list[list[tuple[int, int]]] | None = None, + **_: object, + ) -> HiDreamO1ImageOutput: + return self.model( + input_ids=input_ids, + position_ids=position_ids, + vinputs=vinputs, + timestep=timestep, + attention_mask=attention_mask, + full_attn_spans=full_attn_spans, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + params = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + if name.startswith("model.visual.") or name == "lm_head.weight": + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params[name] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + if name not in params: + continue + param = params[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params + + +__all__ = ["HiDreamO1ImageTransformer"] diff --git a/vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py b/vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py new file mode 100644 index 000000000..de16df9de --- /dev/null +++ b/vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py @@ -0,0 +1,732 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Iterable +from typing import ClassVar, TypedDict + +import numpy as np +import torch +from einops import rearrange +from PIL import Image +from torch import nn +from transformers import AutoProcessor, PreTrainedTokenizerBase, Qwen3VLConfig +from transformers.processing_utils import ProcessorMixin +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import ( + HiDreamO1ImageTransformer, +) +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( + DiffusionPipelineProfilerMixin, +) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +logger = init_logger(__name__) + + +# upstream: HiDream-O1-Image models/pipeline.py L14-18 @21bcd30471ac +PATCH_SIZE = 32 +TIMESTEP_TOKEN_NUM = 1 +NOISE_SCALE = 8.0 +T_EPS = 0.001 + +# upstream: HiDream-O1-Image models/utils.py L7-19 @21bcd30471ac +PREDEFINED_RESOLUTIONS: tuple[tuple[int, int], ...] = ( + (2048, 2048), + (2304, 1728), + (1728, 2304), + (2560, 1440), + (1440, 2560), + (2496, 1664), + (1664, 2496), + (3104, 1312), + (1312, 3104), + (2304, 1792), + (1792, 2304), +) + + +class HiDreamO1TextSample(TypedDict): + input_ids: torch.Tensor + position_ids: torch.Tensor + attention_mask: torch.Tensor | None + full_attn_spans: list[list[tuple[int, int]]] + vinput_mask: torch.Tensor + + +def find_closest_resolution(width: int, height: int) -> tuple[int, int]: + img_ratio = width / height + best_res: tuple[int, int] | None = None + min_diff = float("inf") + for w, h in PREDEFINED_RESOLUTIONS: + diff = abs(w / h - img_ratio) + if diff < min_diff: + min_diff = diff + best_res = (w, h) + assert best_res is not None + return best_res + + +def get_rope_index_fix_point( + *, + spatial_merge_size: int, + image_token_id: int, + video_token_id: int, + vision_start_token_id: int, + input_ids: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + skip_vision_start_token: list[int], + fix_point: int = 4096, +) -> tuple[torch.Tensor, torch.Tensor]: + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) + video_grid_thw[:, 0] = 1 + + mrope_position_deltas = [] + if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + total_input_ids = input_ids + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + position_ids = torch.ones( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + for i, input_ids in enumerate(total_input_ids): + input_ids = input_ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_tokens = input_ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + input_tokens = input_ids.tolist() + llm_pos_ids_list: list[torch.Tensor] = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_index += 1 + remain_videos -= 1 + ed = ed_video + llm_grid_t, llm_grid_h, llm_grid_w = ( + t.item(), + h.item() // spatial_merge_size, + w.item() // spatial_merge_size, + ) + text_len = ed - st + + text_len -= skip_vision_start_token[image_index - 1] + text_len = max(0, text_len) + + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() + + if skip_vision_start_token[image_index - 1]: + if fix_point > 0: + fix_point = fix_point - st_idx + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + fix_point + st_idx) + fix_point = 0 + else: + llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + else: + if attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) + max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] + mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + else: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(3, input_ids.shape[0], -1) + ) + mrope_position_deltas = torch.zeros( + [input_ids.shape[0], 1], + device=input_ids.device, + dtype=input_ids.dtype, + ) + return position_ids, mrope_position_deltas + + +def build_t2i_text_sample( + *, + prompt: str, + height: int, + width: int, + tokenizer: PreTrainedTokenizerBase, + processor: ProcessorMixin, + model_config: Qwen3VLConfig, +) -> dict[str, torch.Tensor]: + image_token_id = model_config.image_token_id + video_token_id = model_config.video_token_id + vision_start_token_id = model_config.vision_start_token_id + image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE) + + boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>") + tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>") + + messages = [{"role": "user", "content": prompt}] + template_caption = ( + processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + boi_token + + tms_token * TIMESTEP_TOKEN_NUM + ) + input_ids = tokenizer.encode(template_caption, return_tensors="pt", add_special_tokens=False) + + image_grid_thw = torch.tensor([1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64).unsqueeze(0) + + vision_tokens = torch.zeros((1, image_len), dtype=input_ids.dtype) + image_token_id + vision_tokens[0, 0] = vision_start_token_id + input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1) + + position_ids, _ = get_rope_index_fix_point( + spatial_merge_size=1, + image_token_id=image_token_id, + video_token_id=video_token_id, + vision_start_token_id=vision_start_token_id, + input_ids=input_ids_pad, + image_grid_thw=image_grid_thw, + video_grid_thw=None, + attention_mask=None, + skip_vision_start_token=[1], + ) + + txt_seq_len = input_ids.shape[-1] + all_seq_len = position_ids.shape[-1] + + token_types = torch.zeros((1, all_seq_len), dtype=input_ids.dtype) + bgn = txt_seq_len - TIMESTEP_TOKEN_NUM + token_types[0, bgn : bgn + image_len + TIMESTEP_TOKEN_NUM] = 1 + token_types[0, txt_seq_len - TIMESTEP_TOKEN_NUM : txt_seq_len] = 3 + + vinput_mask = token_types == 1 + token_types_bin = (token_types > 0).to(token_types.dtype) + + return { + "input_ids": input_ids, + "position_ids": position_ids, + "token_types": token_types_bin, + "vinput_mask": vinput_mask, + } + + +def build_hidream_o1_attention_mask( + token_types: torch.Tensor, + *, + dtype: torch.dtype, +) -> torch.Tensor: + """Build the mixed causal/full mask once while preparing a request. + + Text rows stay causal; image and timestep rows get full attention. + """ + if token_types.ndim == 1: + token_types = token_types.unsqueeze(0) + elif token_types.ndim == 2 and token_types.shape[1] == 1 and token_types.shape[0] != 1: + token_types = token_types.transpose(0, 1) + if token_types.ndim != 2: + raise ValueError(f"token_types must have one or two dimensions, got shape {tuple(token_types.shape)}") + + batch_size, seq_len = token_types.shape + min_value = torch.finfo(dtype).min + attention_mask = torch.triu( + torch.full( + (batch_size, seq_len, seq_len), + min_value, + dtype=dtype, + device=token_types.device, + ), + diagonal=1, + ) + attention_mask[token_types.bool()] = 0 + return attention_mask.unsqueeze(1) + + +def build_hidream_o1_full_attn_spans( + token_types: torch.Tensor, +) -> list[list[tuple[int, int]]]: + """Encode the mixed causal/full rows without allocating an ``S x S`` mask.""" + if token_types.ndim == 1: + token_types = token_types.unsqueeze(0) + elif token_types.ndim == 2 and token_types.shape[1] == 1 and token_types.shape[0] != 1: + token_types = token_types.transpose(0, 1) + if token_types.ndim != 2: + raise ValueError(f"token_types must have one or two dimensions, got shape {tuple(token_types.shape)}") + + spans: list[list[tuple[int, int]]] = [] + for sample_types in token_types: + full_rows = sample_types.bool().nonzero(as_tuple=False).flatten() + if full_rows.numel() == 0: + spans.append([]) + continue + + start = int(full_rows[0].item()) + end = int(full_rows[-1].item()) + 1 + expected = torch.arange(start, end, device=full_rows.device) + if not torch.equal(full_rows, expected): + raise ValueError("HiDream-O1 full-attention rows must form one contiguous span") + spans.append([(start, end)]) + return spans + + +def build_hidream_o1_scheduler( + *, + scheduler_name: str = "default", + num_inference_steps: int = 50, + shift: float = 3.0, + device: torch.device | str | None = None, + timesteps_list: list[int] | None = None, +) -> FlowUniPCMultistepScheduler: + if scheduler_name != "default": + raise NotImplementedError(f"unsupported scheduler_name={scheduler_name!r}") + + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, + shift=shift, + prediction_type="flow_prediction", + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(num_inference_steps, device=device) + + if timesteps_list is not None: + scheduler.timesteps = torch.tensor(timesteps_list, device=device, dtype=torch.long) + sigmas = [t.item() / 1000.0 for t in scheduler.timesteps] + sigmas.append(0.0) + scheduler.sigmas = torch.tensor(sigmas, device=device) + + return scheduler + + +def get_hidream_o1_image_post_process_func(od_config: OmniDiffusionConfig): + del od_config + + def post_process_func( + output_data: tuple[torch.Tensor, int, int], + ) -> Image.Image: + """Convert HiDream-O1 patch-space output to a PIL RGB image.""" + if not isinstance(output_data, tuple) or len(output_data) != 3: + raise TypeError( + f"HiDreamO1 postprocess expects (patch_tensor, height, width), got {type(output_data).__name__}" + ) + + z, height, width = output_data + + if not torch.is_tensor(z): + raise TypeError( + f"HiDreamO1 postprocess expects the first tuple element to be a torch.Tensor, got {type(z).__name__}" + ) + + height = int(height) + width = int(width) + + if height <= 0 or width <= 0: + raise ValueError( + f"HiDreamO1 postprocess requires positive height/width, got height={height}, width={width}" + ) + + if height % PATCH_SIZE != 0 or width % PATCH_SIZE != 0: + raise ValueError( + f"HiDreamO1 output resolution must be divisible by PATCH_SIZE={PATCH_SIZE}, got {width}x{height}" + ) + + h_patches = height // PATCH_SIZE + w_patches = width // PATCH_SIZE + expected_image_len = h_patches * w_patches + expected_patch_dim = 3 * PATCH_SIZE * PATCH_SIZE + + if z.ndim != 3: + raise ValueError( + f"HiDreamO1 patch tensor must have shape (B, image_len, patch_dim), got shape={tuple(z.shape)}" + ) + + if z.shape[0] != 1: + raise ValueError(f"HiDreamO1 postprocess currently expects batch size 1, got batch_size={z.shape[0]}") + + if z.shape[1] != expected_image_len: + raise ValueError( + "HiDreamO1 patch count mismatch: " + f"got {z.shape[1]}, expected {expected_image_len} " + f"for resolution {width}x{height}" + ) + + if z.shape[2] != expected_patch_dim: + raise ValueError(f"HiDreamO1 patch dimension mismatch: got {z.shape[2]}, expected {expected_patch_dim}") + + # Preserve upstream decode order exactly: + # 1. map patch values from [-1, 1] to [0, 1] + # 2. restore BCHW image layout + # 3. convert to float32 on CPU + img = (z + 1) / 2 + + img = rearrange( + img.cpu().float(), + "B (H W) (C p1 p2) -> B C (H p1) (W p2)", + H=h_patches, + W=w_patches, + p1=PATCH_SIZE, + p2=PATCH_SIZE, + ) + + arr = np.round( + np.clip( + img[0].numpy().transpose(1, 2, 0) * 255, + 0, + 255, + ) + ).astype(np.uint8) + + return Image.fromarray(arr).convert("RGB") + + return post_process_func + + +class HiDreamO1ImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin): + supports_request_batch = False + dummy_run_num_frames: ClassVar[int] = 0 + _dit_modules: ClassVar[list[str]] = ["model.model.language_model"] + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.prefix = prefix + self.model_dir = od_config.model + self.dtype = od_config.dtype if od_config.dtype is not None else torch.bfloat16 + self.device = get_local_device() + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder=None, + revision=od_config.revision, + prefix="model.", + fall_back_to_pt=True, + ) + ] + + self._validate_static_config() + + self.setup_diffusion_pipeline_profiler( + profiler_targets=["forward"], + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler, + ) + + self._init_processor_and_model() + self._validate_tms_token_id() + + def _validate_static_config(self) -> None: + if self.dtype != torch.bfloat16: + raise NotImplementedError(f"only bfloat16 is supported; got {self.dtype!r}") + cfg_parallel_size = int(self.od_config.parallel_config.cfg_parallel_size) + if cfg_parallel_size > 1: + raise NotImplementedError(f"cfg_parallel_size={cfg_parallel_size} not supported") + + def _init_processor_and_model(self) -> None: + logger.info( + "HiDreamO1ImagePipeline: loading processor and model config from %s", + self.model_dir, + ) + self.processor = AutoProcessor.from_pretrained( + self.model_dir, + revision=self.od_config.revision, + ) + model_config = Qwen3VLConfig.from_pretrained( + self.model_dir, + revision=self.od_config.revision, + ) + self.model = HiDreamO1ImageTransformer( + model_config, + quant_config=self.od_config.quantization_config, + prefix=self.prefix, + ) + + self._add_special_tokens(self.processor) + self.tokenizer = ( + self.processor if isinstance(self.processor, PreTrainedTokenizerBase) else self.processor.tokenizer + ) + + def _validate_tms_token_id(self) -> None: + actual_tms_ids = self.tokenizer.encode(self.tokenizer.tms_token, add_special_tokens=False) + expected_tms_id = int(self.model.model.tms_token_id) + if actual_tms_ids != [expected_tms_id]: + raise RuntimeError(f"tms_token maps to {actual_tms_ids}, model expects [{expected_tms_id}]") + + @staticmethod + def _add_special_tokens(processor: ProcessorMixin | PreTrainedTokenizerBase) -> None: + tok = processor if isinstance(processor, PreTrainedTokenizerBase) else processor.tokenizer + tok.boi_token = "<|boi_token|>" + tok.bor_token = "<|bor_token|>" + tok.eor_token = "<|eor_token|>" + tok.bot_token = "<|bot_token|>" + tok.tms_token = "<|tms_token|>" + + def _resolve_generation_params( + self, + req: DiffusionRequestBatch, + ) -> tuple[str, int, int, int, int, float]: + sp = req.sampling_params + + if sp.timesteps is not None: + raise NotImplementedError("sp.timesteps not supported") + if sp.latents is not None: + raise NotImplementedError("sp.latents not supported") + if sp.num_outputs_per_prompt != 1: + raise NotImplementedError(f"num_outputs_per_prompt={sp.num_outputs_per_prompt} not supported") + + if not req.prompts: + raise ValueError("req.prompts is empty") + if len(req.prompts) != 1: + raise NotImplementedError(f"only one prompt per request; got {len(req.prompts)}") + first_prompt = req.prompts[0] + if isinstance(first_prompt, str): + prompt = first_prompt + elif isinstance(first_prompt, dict): + raw_prompt = first_prompt.get("prompt") + if not isinstance(raw_prompt, str): + raise TypeError(f"dict prompt must contain a str 'prompt' field; got {type(raw_prompt).__name__}") + prompt = raw_prompt + else: + raise TypeError(f"prompt must be str or dict; got {type(first_prompt).__name__}") + + raw_h = 2048 if sp.height is None else int(sp.height) + raw_w = 2048 if sp.width is None else int(sp.width) + if raw_h <= 0 or raw_w <= 0: + raise ValueError(f"h/w must be positive, got h={raw_h} w={raw_w}") + snapped_w, snapped_h = find_closest_resolution(raw_w, raw_h) + if (snapped_w, snapped_h) != (raw_w, raw_h): + logger.info( + "HiDreamO1: resolution snapped from %dx%d to %dx%d", + raw_w, + raw_h, + snapped_w, + snapped_h, + ) + + steps = 50 if sp.num_inference_steps is None else int(sp.num_inference_steps) + if steps <= 0: + raise ValueError(f"num_inference_steps must be positive, got {steps}") + + generator = sp.generator + if isinstance(generator, list): + if len(generator) != 1: + raise NotImplementedError("generator list only supported for single-output requests with one generator") + generator = generator[0] + + generator_seed: int | None = None + if generator is not None: + if not hasattr(generator, "initial_seed"): + raise TypeError(f"sp.generator must expose initial_seed(); got {type(generator).__name__}") + generator_seed = int(generator.initial_seed()) + + if sp.seed is not None and generator_seed is not None and int(sp.seed) != generator_seed: + raise ValueError(f"seed/generator mismatch: seed={int(sp.seed)} generator.initial_seed()={generator_seed}") + + if sp.seed is None and generator_seed is None: + raise RuntimeError("expected request initialization to resolve a concrete seed") + seed = int(sp.seed) if sp.seed is not None else generator_seed + assert seed is not None + guidance_scale = float(sp.guidance_scale) + return prompt, snapped_h, snapped_w, steps, seed, guidance_scale + + def _prepare_noise_and_patchify( + self, + height: int, + width: int, + seed: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + gen = torch.Generator("cpu").manual_seed(seed + 1) + noise = NOISE_SCALE * torch.randn((1, 3, height, width), generator=gen) + noise = noise.to(device, dtype) + return rearrange( + noise, + "B C (H p1) (W p2) -> B (H W) (C p1 p2)", + p1=PATCH_SIZE, + p2=PATCH_SIZE, + ) + + def _forward_once( + self, + sample: HiDreamO1TextSample, + z_in: torch.Tensor, + t_pixeldit: torch.Tensor, + ) -> torch.Tensor: + with torch.autocast(self.device.type, dtype=self.dtype, cache_enabled=False): + outputs = self.model( + input_ids=sample["input_ids"], + position_ids=sample["position_ids"], + vinputs=z_in, + timestep=t_pixeldit.reshape(-1).to(self.device), + attention_mask=sample["attention_mask"], + full_attn_spans=sample["full_attn_spans"], + ) + + if outputs.x_pred is None: + raise RuntimeError("model returned x_pred=None") + if outputs.x_pred.shape[-1] != z_in.shape[-1]: + raise RuntimeError(f"x_pred patch_dim {outputs.x_pred.shape[-1]} != z_in {z_in.shape[-1]}") + + mask = sample["vinput_mask"][0] + return outputs.x_pred[0, mask].unsqueeze(0) + + def _build_and_validate_sample( + self, + prompt: str, + height: int, + width: int, + expected_image_len: int, + ) -> HiDreamO1TextSample: + sample = build_t2i_text_sample( + prompt=prompt, + height=height, + width=width, + tokenizer=self.tokenizer, + processor=self.processor, + model_config=self.model.config, + ) + actual_image_len = int(sample["vinput_mask"].sum().item()) + if actual_image_len != expected_image_len: + raise RuntimeError( + f"vinput_mask has {actual_image_len} slots, expected {expected_image_len} for {height}x{width}" + ) + full_attn_spans = build_hidream_o1_full_attn_spans(sample["token_types"]) + inputs = {key: value.to(self.device) for key, value in sample.items()} + attention_mask = None + if not self.model.supports_piecewise_attention: + attention_mask = build_hidream_o1_attention_mask( + inputs["token_types"], + dtype=self.dtype, + ) + return { + "input_ids": inputs["input_ids"], + "position_ids": inputs["position_ids"], + "attention_mask": attention_mask, + "full_attn_spans": full_attn_spans, + "vinput_mask": inputs["vinput_mask"], + } + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + prompt, height, width, num_inference_steps, seed, guidance_scale = self._resolve_generation_params(req) + + expected_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE) + cond_sample = self._build_and_validate_sample( + prompt, + height, + width, + expected_image_len, + ) + samples: list[HiDreamO1TextSample] = [cond_sample] + if guidance_scale > 1.0: + uncond_sample = self._build_and_validate_sample( + " ", + height, + width, + expected_image_len, + ) + samples.append(uncond_sample) + + z = self._prepare_noise_and_patchify( + height, + width, + seed, + self.dtype, + self.device, + ) + + sched = build_hidream_o1_scheduler( + scheduler_name="default", + num_inference_steps=num_inference_steps, + shift=3.0, + device=self.device, + ) + + with self.progress_bar(total=len(sched.timesteps)) as pbar: + for step_t in sched.timesteps: + step_t_f32 = step_t.float() + t_pixeldit = 1.0 - step_t_f32 / 1000.0 + sigma = (step_t_f32 / 1000.0).to(dtype=torch.float32).clamp_min(T_EPS) + + x_pred_cond = self._forward_once(samples[0], z, t_pixeldit) + v_cond = (x_pred_cond.to(torch.float32) - z.to(torch.float32)) / sigma + + if len(samples) > 1: + x_pred_uncond = self._forward_once( + samples[1], + z, + t_pixeldit, + ) + v_uncond = (x_pred_uncond.to(torch.float32) - z.to(torch.float32)) / sigma + v_guided = v_uncond + guidance_scale * (v_cond - v_uncond) + else: + v_guided = v_cond + + model_output = -v_guided + + z = sched.step( + model_output, + step_t_f32, + z.float(), + return_dict=False, + )[0].to(self.dtype) + pbar.update() + + return DiffusionOutput( + output=(z, height, width), + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/host_weight_contract.py b/vllm_omni/diffusion/models/host_weight_contract.py new file mode 100644 index 000000000..939c9af2e --- /dev/null +++ b/vllm_omni/diffusion/models/host_weight_contract.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Model-owned declaration for complete final-layout tensor restoration.""" + +from __future__ import annotations + +from dataclasses import dataclass + +FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA = "vllm-omni.diffusion.final-layout-tensors-v1" + + +@dataclass(frozen=True) +class FinalLayoutModelContract: + """Versioned model ABI for reconstruction from exact final-layout tensors.""" + + implementation_id: str + version: str + schema: str = FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA + + def __post_init__(self) -> None: + for name, value in ( + ("schema", self.schema), + ("implementation_id", self.implementation_id), + ("version", self.version), + ): + if not isinstance(value, str) or not value: + raise ValueError(f"final-layout model contract {name} must not be empty") + + def to_dict(self) -> dict[str, str]: + return { + "schema": self.schema, + "implementation_id": self.implementation_id, + "version": self.version, + } + + +__all__ = [ + "FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA", + "FinalLayoutModelContract", +] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/__init__.py index 625895d99..4c93e2397 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/__init__.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/__init__.py @@ -10,4 +10,8 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import ( HunyuanImage3Pipeline, ) -__all__ = ["HunyuanImage3Pipeline", "HunyuanImage3Model", "HunyuanImage3Text2ImagePipeline"] +__all__ = [ + "HunyuanImage3Pipeline", + "HunyuanImage3Model", + "HunyuanImage3Text2ImagePipeline", +] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py b/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py index efba2f274..1ccbafa65 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import math from dataclasses import dataclass import numpy as np @@ -15,6 +14,10 @@ from diffusers.utils.torch_utils import randn_tensor from einops import rearrange from torch import Tensor, nn +# ResnetBlock is platform-dispatched in the layers package __init__: CUDA gets a +# fused GroupNorm+SiLU kernel, every other backend gets the plain PyTorch block. +from vllm_omni.diffusion.models.hunyuan_image3.layers import ResnetBlock + class DiagonalGaussianDistribution: def __init__(self, parameters: torch.Tensor, deterministic: bool = False): @@ -53,10 +56,6 @@ class DecoderOutput(BaseOutput): posterior: DiagonalGaussianDistribution | None = None -def swish(x: Tensor) -> Tensor: - return x * torch.sigmoid(x) - - def forward_with_checkpointing(module, *inputs, use_checkpointing=False): def create_custom_forward(module): def custom_forward(*inputs): @@ -70,46 +69,6 @@ def forward_with_checkpointing(module, *inputs, use_checkpointing=False): return module(*inputs) -class Conv3d(nn.Conv3d): - """ - Perform Conv3d on patches with numerical differences from nn.Conv3d within 1e-5. - Only symmetric padding is supported. - """ - - def forward(self, input): - B, C, T, H, W = input.shape - memory_count = (C * T * H * W) * 2 / 1024**3 - if memory_count > 2: - n_split = math.ceil(memory_count / 2) - assert n_split >= 2 - chunks = torch.chunk(input, chunks=n_split, dim=-3) - padded_chunks = [] - for i in range(len(chunks)): - if self.padding[0] > 0: - padded_chunk = F.pad( - chunks[i], - (0, 0, 0, 0, self.padding[0], self.padding[0]), - mode="constant" if self.padding_mode == "zeros" else self.padding_mode, - value=0, - ) - if i > 0: - padded_chunk[:, :, : self.padding[0]] = chunks[i - 1][:, :, -self.padding[0] :] - if i < len(chunks) - 1: - padded_chunk[:, :, -self.padding[0] :] = chunks[i + 1][:, :, : self.padding[0]] - else: - padded_chunk = chunks[i] - padded_chunks.append(padded_chunk) - padding_bak = self.padding - self.padding = (0, self.padding[1], self.padding[2]) - outputs = [] - for i in range(len(padded_chunks)): - outputs.append(super().forward(padded_chunks[i])) - self.padding = padding_bak - return torch.cat(outputs, dim=-3) - else: - return super().forward(input) - - class AttnBlock(nn.Module): """Attention with torch sdpa implementation.""" @@ -119,10 +78,12 @@ class AttnBlock(nn.Module): self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) - self.q = Conv3d(in_channels, in_channels, kernel_size=1) - self.k = Conv3d(in_channels, in_channels, kernel_size=1) - self.v = Conv3d(in_channels, in_channels, kernel_size=1) - self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1) + # The memory concern is already handled via VAE tiling, + # so there's no need to chunk here — calling nn.Conv3d directly is fine. + self.q = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv3d(in_channels, in_channels, kernel_size=1) def attention(self, h_: Tensor) -> Tensor: h_ = self.norm(h_) @@ -142,41 +103,12 @@ class AttnBlock(nn.Module): return x + self.proj_out(self.attention(x)) -class ResnetBlock(nn.Module): - def __init__(self, in_channels: int, out_channels: int): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - - self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) - self.conv1 = Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) - self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) - self.conv2 = Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) - if self.in_channels != self.out_channels: - self.nin_shortcut = Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) - - def forward(self, x): - h = x - h = self.norm1(h) - h = swish(h) - h = self.conv1(h) - - h = self.norm2(h) - h = swish(h) - h = self.conv2(h) - - if self.in_channels != self.out_channels: - x = self.nin_shortcut(x) - return x + h - - class DownsampleDCAE(nn.Module): def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 assert out_channels % factor == 0 - self.conv = Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + self.conv = nn.Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) self.add_temporal_downsample = add_temporal_downsample self.group_size = factor * in_channels // out_channels @@ -196,7 +128,7 @@ class UpsampleDCAE(nn.Module): def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): super().__init__() factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 - self.conv = Conv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) + self.conv = nn.Conv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) self.add_temporal_upsample = add_temporal_upsample self.repeats = factor * out_channels // in_channels @@ -233,7 +165,7 @@ class Encoder(nn.Module): self.num_res_blocks = num_res_blocks # downsampling - self.conv_in = Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + self.conv_in = nn.Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) self.down = nn.ModuleList() block_in = block_out_channels[0] @@ -265,7 +197,7 @@ class Encoder(nn.Module): # end self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) - self.conv_out = Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + self.conv_out = nn.Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) self.gradient_checkpointing = False @@ -291,7 +223,7 @@ class Encoder(nn.Module): group_size = self.block_out_channels[-1] // (2 * self.z_channels) shortcut = rearrange(h, "b (c r) f h w -> b c r f h w", r=group_size).mean(dim=2) h = self.norm_out(h) - h = swish(h) + h = F.silu(h) h = self.conv_out(h) h += shortcut return h @@ -321,7 +253,7 @@ class Decoder(nn.Module): # z to block_in block_in = block_out_channels[0] - self.conv_in = Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + self.conv_in = nn.Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) # middle self.mid = nn.Module() @@ -351,7 +283,7 @@ class Decoder(nn.Module): # end self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) - self.conv_out = Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) + self.conv_out = nn.Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) self.gradient_checkpointing = False @@ -376,7 +308,7 @@ class Decoder(nn.Module): # end h = self.norm_out(h) - h = swish(h) + h = F.silu(h) h = self.conv_out(h) return h diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 877fd1c3f..a1ad2c986 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import inspect import logging @@ -79,6 +79,13 @@ from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step from vllm_omni.diffusion.layers.fused_moe import FusedMoE from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +# ResBlock is platform-dispatched in the layers package __init__: CUDA gets +# fused GroupNorm+SiLU and AdaGN kernels, every other backend gets the plain +# PyTorch block. UNetDown and UNetUp below instantiate whichever one this +# resolves to. +from vllm_omni.diffusion.models.hunyuan_image3.layers import ResBlock +from vllm_omni.diffusion.models.hunyuan_image3.layers.common import conv_nd, normalization from vllm_omni.diffusion.utils.kv_utils import repeat_kv from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding from vllm_omni.platforms import current_omni_platform @@ -171,44 +178,6 @@ def real_batched_index_select(t, dim, idx): return torch.stack([torch.index_select(t[i], dim - 1, idx[i]) for i in range(len(t))]) -def conv_nd(dims, *args, **kwargs): # noqa: N802 - """ - Create a 1D, 2D, or 3D convolution module. - """ - if dims == 1: - return nn.Conv1d(*args, **kwargs) - elif dims == 2: - return nn.Conv2d(*args, **kwargs) - elif dims == 3: - return nn.Conv3d(*args, **kwargs) - raise ValueError(f"unsupported dimensions: {dims}") - - -def normalization(channels, **kwargs): - """ - Make a standard normalization layer. - :param channels: number of input channels. - :return: a nn.Module for normalization. - """ - return nn.GroupNorm(32, channels, **kwargs) - - -def linear(*args, **kwargs): - """ - Create a linear module. - """ - return nn.Linear(*args, **kwargs) - - -def zero_module(module): - """ - Zero out the parameters of a module and return it. - """ - for p in module.parameters(): - p.detach().zero_() - return module - - def _to_tuple(x, dim=2): if isinstance(x, int): return (x,) * dim @@ -913,7 +882,7 @@ class HunYuanRotary2DEmbedder: return q, k -class ImageKVCacheManager: +class ImageKVCacheManager(nn.Module): """ Manages specialized caching and updating of KV-Cache for image tokens in multimodal models. """ @@ -932,6 +901,7 @@ class ImageKVCacheManager: image_token_len: Number of tokens per image (including special placeholders), default 4097 (timestamp + 4096 image tokens). """ + super().__init__() # attention related self.num_heads = num_heads self.num_kv_heads = num_kv_heads @@ -953,6 +923,10 @@ class ImageKVCacheManager: softmax_scale=self.scaling, num_kv_heads=self.num_kv_heads, prefix=f"{prefix}.attn" if prefix else "", + paged_kv_cache_role="primary", + # Hunyuan image RoPE produces bfloat16 K and denoising runs under + # bfloat16 autocast, independent of the weight-loading dtype. + paged_kv_cache_dtype=torch.bfloat16, ) @staticmethod @@ -1095,18 +1069,98 @@ class ImageKVCacheManager: ) return key, value - def __call__( + def clear_legacy_prompt_kv_cache(self) -> None: + """Drop model-owned dense KV state before Scheduler-paged execution.""" + + self.image_kv_cache_map = None + self.image_kv_cache_lens = None + + def _forward_paged( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attention_mask: torch.Tensor | None = None, + *, + first_step: bool, + query_lens: list[int], + seq_lens: list[int], + shard_image_size: int | None, + full_attn_spans: list[list[tuple[int, int]]], + uncond_cfg_prefill: bool, + ) -> torch.Tensor: + """Run Hunyuan attention against Worker-managed Scheduler pages. + + This model boundary is shared by GPU and NPU: it describes Hunyuan's + Q/K/V layout and mixed-attention spans, while the common Attention + layer resolves the actual CUDA or Ascend paged kernel. + """ + + if uncond_cfg_prefill: + raise RuntimeError("Hunyuan negative-CFG prefill must run before paged row activation") + if self._injected_ar_kv is not None: + raise NotImplementedError("Hunyuan Scheduler-paged KV does not support imported AR KV") + if not query_lens or len(seq_lens) != len(query_lens): + raise ValueError("Hunyuan Scheduler-paged KV requires matching query and sequence lengths") + if len(full_attn_spans) != len(query_lens): + raise ValueError("Hunyuan Scheduler-paged KV requires one full-attention span row per sequence") + + self.clear_legacy_prompt_kv_cache() + bs = len(query_lens) + q_len = query_lens[0] + seq_len = seq_lens[0] + assert query.shape[0] == bs * q_len, f"{query.shape[0]} != {bs * q_len}" + + head_num_per_rank = query.shape[1] + kv_head_num_per_rank = key.shape[1] + head_dim = query.shape[2] + + query = query.reshape(bs, q_len, head_num_per_rank, head_dim) + key = key.reshape(bs, q_len, kv_head_num_per_rank, head_dim) + value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim) + + joint_text_query = joint_text_key = joint_text_value = None + if self.sp_size > 1 and first_step: + if shard_image_size is None or shard_image_size <= 0: + raise ValueError("Hunyuan paged Ulysses requires a positive local image shard size") + local_prompt_len = seq_len - shard_image_size + joint_query_len = query.shape[1] - shard_image_size + if local_prompt_len != joint_query_len: + raise ValueError( + "Hunyuan paged Ulysses prompt layout mismatch: " + f"key_prompt={local_prompt_len}, query_prompt={joint_query_len}" + ) + joint_text_query = query[:, :joint_query_len] + joint_text_key = key[:, :local_prompt_len] + joint_text_value = value[:, :local_prompt_len] + query = query[:, joint_query_len:] + key = key[:, local_prompt_len:] + value = value[:, local_prompt_len:] + + if joint_text_query is None: + attn_metadata = AttentionMetadata(full_attn_spans=full_attn_spans) + else: + attn_metadata = AttentionMetadata( + joint_query=joint_text_query, + joint_key=joint_text_key, + joint_value=joint_text_value, + joint_strategy="front", + full_attn_spans=full_attn_spans, + ) + attn_output = self.attn(query, key, value, attn_metadata) + return attn_output.reshape(bs * q_len, head_num_per_rank, head_dim) + + def _forward_dense_legacy( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, **kwargs, ) -> torch.Tensor: - self.image_token_len = kwargs.get("num_image_tokens") + """Run the device-neutral, model-owned dense prompt-KV compatibility path.""" + first_step = kwargs.get("first_step") uncond_cfg_prefill = kwargs.get("uncond_cfg_prefill", False) - query_lens = kwargs.get("query_lens") seq_lens = kwargs.get("seq_lens") shard_image_size = kwargs.get("shard_image_size") if self.sp_size > 1 else None @@ -1135,8 +1189,7 @@ class ImageKVCacheManager: key = key[:, :0, :, :] value = value[:, :0, :, :] elif first_step: - self.image_kv_cache_map = None # reset first - self.image_kv_cache_lens = None + self.clear_legacy_prompt_kv_cache() key, value = self._cache_prompt_kv( key, value, @@ -1189,6 +1242,34 @@ class ImageKVCacheManager: attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim) return attn_output + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + self.image_token_len = kwargs.get("num_image_tokens") + if self.attn.is_paged_kv_active(): + full_attn_spans = kwargs.get("full_attn_spans") + if full_attn_spans is None: + raise ValueError("Hunyuan Scheduler-paged KV requires full_attn_spans metadata") + return self._forward_paged( + query, + key, + value, + first_step=bool(kwargs.get("first_step")), + query_lens=kwargs.get("query_lens"), + seq_lens=kwargs.get("seq_lens"), + shard_image_size=kwargs.get("shard_image_size") if self.sp_size > 1 else None, + full_attn_spans=full_attn_spans, + uncond_cfg_prefill=kwargs.get("uncond_cfg_prefill", False), + ) + if attention_mask is None: + raise ValueError("Hunyuan dense attention requires an attention mask") + return self._forward_dense_legacy(query, key, value, attention_mask, **kwargs) + @dataclass class CausalMMOutputWithPast(CausalLMOutputWithPast): @@ -2732,7 +2813,6 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): # List[List[...]] per-sample metadata indexed along the CFG batch dim if isinstance(model_kwargs.get("full_attn_spans"), list): model_kwargs["full_attn_spans"] = model_kwargs["full_attn_spans"][s.start : s.stop] - # custom_pos_emb: tuple of (cos, sin) if "custom_pos_emb" in model_kwargs and model_kwargs["custom_pos_emb"] is not None: cos, sin = model_kwargs["custom_pos_emb"] @@ -3277,99 +3357,6 @@ class TimestepEmbedder(nn.Module): return t_emb -class ResBlock(nn.Module): - r""" - A residual block that can optionally change the number of channels. - - Args: - in_channels (`int`): - The number of input channels. - emb_channels (`int`): - The number of timestep embedding channels. - dropout (`float`): - The rate of dropout. - out_channels (`int`, *optional*): - If specified, the number of output channels. - use_conv (`bool`, *optional*): - If True and out_channels is specified, use a spatial convolution instead of a - smaller 1x1 convolution to change the channels in the skip connection. - dims (`int`, *optional*): - Determines if the signal is 1D, 2D, or 3D. - up (`bool`, *optional*): - If True, use this block for upsampling. - down (`bool`, *optional*): - If True, use this block for downsampling. - - """ - - def __init__( - self, - in_channels, - emb_channels, - out_channels=None, - dropout=0.0, - use_conv=False, - dims=2, - up=False, - down=False, - device=None, - dtype=None, - ) -> None: - factory_kwargs = {"dtype": dtype, "device": device} - super().__init__() - self.in_channels = in_channels - self.dropout = dropout - self.out_channels = out_channels or self.in_channels - self.use_conv = use_conv - - self.in_layers = nn.Sequential( - normalization(self.in_channels, **factory_kwargs), - nn.SiLU(), - conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 - ) - - self.updown = up or down - self.h_upd = self.x_upd = nn.Identity() - - self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) - - self.out_layers = nn.Sequential( - normalization(self.out_channels, **factory_kwargs), - nn.SiLU(), - nn.Dropout(p=dropout), - zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 - ) - - if self.out_channels == self.in_channels: - self.skip_connection = nn.Identity() - elif use_conv: - self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 - else: - self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 - - def forward(self, x, emb) -> torch.Tensor: - if self.updown: - in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] - h = in_rest(x) - h = self.h_upd(h) - x = self.x_upd(x) - h = in_conv(h) - else: - h = self.in_layers(x) - - emb_out = self.emb_layers(emb) - while len(emb_out.shape) < len(h.shape): - emb_out = emb_out[..., None] - - # Adaptive Group Normalization - out_norm, out_rest = self.out_layers[0], self.out_layers[1:] - scale, shift = torch.chunk(emb_out, 2, dim=1) - h = out_norm(h) * (1.0 + scale) + shift - h = out_rest(h) - - return self.skip_connection(x) + h - - class UNetDown(nn.Module): def __init__( self, patch_size, in_channels, emb_channels, hidden_channels, out_channels, dropout=0.0, device=None, dtype=None diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/__init__.py new file mode 100644 index 000000000..39081df17 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Platform-dispatched layer implementations for HunyuanImage3. + +The two blocks are selected by platform: + +Layout:: + native/autoencoder_blocks.py default ResnetBlock + native/transformer_blocks.py default ResBlock + nvidia/autoencoder_blocks.py CUDA ResnetBlock (fused) + nvidia/transformer_blocks.py CUDA ResBlock (fused) +""" + +from vllm_omni.platforms import current_omni_platform + +if current_omni_platform.is_cuda(): + from vllm_omni.diffusion.models.hunyuan_image3.layers.nvidia import ResBlock, ResnetBlock +else: + from vllm_omni.diffusion.models.hunyuan_image3.layers.native import ResBlock, ResnetBlock + +__all__ = ["ResBlock", "ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/common/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/__init__.py new file mode 100644 index 000000000..5272e9d51 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .utils import ( + conv_nd, + linear, + normalization, + zero_module, +) + +__all__ = ["conv_nd", "linear", "normalization", "zero_module"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/common/utils.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/utils.py new file mode 100644 index 000000000..01dd1cc4e --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/common/utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Small platform-agnostic helpers for UNet-style blocks (adapted from diffusers). + +``conv_nd`` / ``normalization`` / ``linear`` / ``zero_module`` are used by the +``ResBlock`` / ``UNetDown`` / ``UNetUp`` implementations across HunyuanImage3's +platform-specific layer packages. They are plain ``nn.ConvNd`` / ``nn.GroupNorm`` +/ ``nn.Linear`` wrappers with no hardware-specific behaviour, so they live here +instead of being duplicated per backend. +""" + +from torch import nn + + +def conv_nd(dims, *args, **kwargs): # noqa: N802 + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def normalization(channels, **kwargs): + """ + Make a standard normalization layer. + :param channels: number of input channels. + :return: a nn.Module for normalization. + """ + return nn.GroupNorm(32, channels, **kwargs) + + +def linear(*args, **kwargs): + """ + Create a linear module. + """ + return nn.Linear(*args, **kwargs) + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/native/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/__init__.py new file mode 100644 index 000000000..55a695af7 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Native (plain PyTorch) layer implementations.""" + +from vllm_omni.diffusion.models.hunyuan_image3.layers.native.autoencoder_blocks import ResnetBlock +from vllm_omni.diffusion.models.hunyuan_image3.layers.native.transformer_blocks import ResBlock + +__all__ = ["ResBlock", "ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/native/autoencoder_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/autoencoder_blocks.py new file mode 100644 index 000000000..6e85a52f6 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/autoencoder_blocks.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""ResnetBlock for the HunyuanImage3 autoencoder — default fallback implementation. + +Uses native PyTorch ``nn.GroupNorm`` + ``F.silu``. Hardware-specific optimized +versions live in subdirectories (e.g., nvidia/) and are selected via __init__.py +dispatch. +""" + +import torch.nn.functional as F +from torch import nn + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = F.silu(h) + h = self.conv1(h) + + h = self.norm2(h) + h = F.silu(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +__all__ = ["ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/native/transformer_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/transformer_blocks.py new file mode 100644 index 000000000..2322fa6e3 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/native/transformer_blocks.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""DiT ``ResBlock`` for HunyuanImage3 -- default implementation. + +Plain PyTorch: ``GroupNorm -> SiLU`` runs as two ops and the adaptive +GroupNorm as ``norm(h) * (1 + scale) + shift``. Hardware-specific variants +live in sibling packages (e.g. ``nvidia/``) and are picked by the dispatch in +``hunyuan_image3/__init__.py``. + +""" + +import torch +from torch import nn + +from vllm_omni.diffusion.models.hunyuan_image3.layers.common import conv_nd, linear, normalization, zero_module + + +class ResBlock(nn.Module): + r""" + A residual block that can optionally change the number of channels. + Args: + in_channels (`int`): + The number of input channels. + emb_channels (`int`): + The number of timestep embedding channels. + dropout (`float`): + The rate of dropout. + out_channels (`int`, *optional*): + If specified, the number of output channels. + use_conv (`bool`, *optional*): + If True and out_channels is specified, use a spatial convolution instead of a + smaller 1x1 convolution to change the channels in the skip connection. + dims (`int`, *optional*): + Determines if the signal is 1D, 2D, or 3D. + up (`bool`, *optional*): + If True, use this block for upsampling. + down (`bool`, *optional*): + If True, use this block for downsampling. + """ + + def __init__( + self, + in_channels, + emb_channels, + out_channels=None, + dropout=0.0, + use_conv=False, + dims=2, + up=False, + down=False, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.in_channels = in_channels + self.dropout = dropout + self.out_channels = out_channels or self.in_channels + self.use_conv = use_conv + + self.in_layers = nn.Sequential( + normalization(self.in_channels, **factory_kwargs), + nn.SiLU(), + conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 + ) + + self.updown = up or down + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) + + self.out_layers = nn.Sequential( + normalization(self.out_channels, **factory_kwargs), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 + ) + + if self.out_channels == self.in_channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 + else: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 + + def forward(self, x, emb) -> torch.Tensor: + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + # Adaptive Group Normalization + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1.0 + scale) + shift + h = out_rest(h) + + return self.skip_connection(x) + h + + +__all__ = ["ResBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/__init__.py new file mode 100644 index 000000000..78d3304a5 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""NVIDIA (CUDA) implementations.""" + +from .autoencoder_blocks import ResnetBlock # noqa: E402 +from .transformer_blocks import ResBlock # noqa: E402 + +__all__ = ["ResnetBlock", "ResBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/_cudnn.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/_cudnn.py new file mode 100644 index 000000000..4c541591d --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/_cudnn.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import contextmanager + +import torch + + +@contextmanager +def cudnn_settings(benchmark: bool, deterministic: bool): + """Temporarily set cuDNN backend flags, restoring the old values on exit. + + Used to scope ``torch.backends.cudnn.benchmark``/``deterministic`` to the + few blocks whose convolutions actually benefit from autotune, instead of + setting them globally at import time (which would silently change every + other model's conv behaviour in the same process). + """ + old_benchmark = torch.backends.cudnn.benchmark + old_deterministic = torch.backends.cudnn.deterministic + + torch.backends.cudnn.benchmark = benchmark + torch.backends.cudnn.deterministic = deterministic + try: + yield + finally: + torch.backends.cudnn.benchmark = old_benchmark + torch.backends.cudnn.deterministic = old_deterministic diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/autoencoder_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/autoencoder_blocks.py new file mode 100644 index 000000000..c6c68957e --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/autoencoder_blocks.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""ResnetBlock for the HunyuanImage3 autoencoder — NVIDIA CUDA + Triton implementation. + +Split out from autoencoder.py because its ``GroupNorm -> SiLU`` pairs are served +by :func:`fused_group_norm_silu`, a single Triton kernel that falls back to +native ``F.silu(F.group_norm(...))`` when Triton is unavailable. +""" + +from torch import nn + +from vllm_omni.diffusion.models.hunyuan_image3.layers.nvidia._cudnn import cudnn_settings +from vllm_omni.model_executor.models.common.ops import fused_group_norm_silu + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + # cuDNN autotune is scoped to this block rather than set at import time. + # These ResnetBlocks carry the bulk of the VAE's conv work: ~29 instances + # x 2 convs across every resolution (up to the 1024x1024 decode), versus + # the scattered convs left in autoencoder.py (AttnBlock's 1x1 at the + # lowest-res mid, down/up-samplers, conv_in/out) whose activations are + # negligible. A process-wide flag would instead silently change every + # other model's conv behaviour in the same process. + with cudnn_settings(benchmark=True, deterministic=True): + h = x + h = fused_group_norm_silu( + h, self.norm1.weight, self.norm1.bias, num_groups=self.norm1.num_groups, eps=self.norm1.eps + ) + h = self.conv1(h) + + h = fused_group_norm_silu( + h, self.norm2.weight, self.norm2.bias, num_groups=self.norm2.num_groups, eps=self.norm2.eps + ) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +__all__ = ["ResnetBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/transformer_blocks.py b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/transformer_blocks.py new file mode 100644 index 000000000..411dc6351 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/layers/nvidia/transformer_blocks.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""``ResBlock`` for HunyuanImage3 -- NVIDIA CUDA + Triton implementation. + +Two fusions over the default block: + +* ``in_layers``: ``GroupNorm -> SiLU`` becomes one :func:`fused_group_norm_silu` + kernel instead of two ops. +* ``out_layers``: ``GroupNorm -> adaptive modulation -> SiLU`` (i.e. + ``SiLU(norm(h) * (1 + scale) + shift)``) becomes one + :func:`fused_adaptive_group_norm_silu` kernel instead of four. +""" + +import torch +from torch import nn + +from vllm_omni.diffusion.models.hunyuan_image3.layers.common import conv_nd, normalization, zero_module +from vllm_omni.diffusion.models.hunyuan_image3.layers.nvidia._cudnn import cudnn_settings +from vllm_omni.model_executor.models.common.ops import ( + fused_adaptive_group_norm_silu, + fused_group_norm_silu, +) + + +class ResBlock(nn.Module): + r""" + A residual block that can optionally change the number of channels. + Args: + in_channels (`int`): + The number of input channels. + emb_channels (`int`): + The number of timestep embedding channels. + dropout (`float`): + The rate of dropout. + out_channels (`int`, *optional*): + If specified, the number of output channels. + use_conv (`bool`, *optional*): + If True and out_channels is specified, use a spatial convolution instead of a + smaller 1x1 convolution to change the channels in the skip connection. + dims (`int`, *optional*): + Determines if the signal is 1D, 2D, or 3D. + up (`bool`, *optional*): + If True, use this block for upsampling. + down (`bool`, *optional*): + If True, use this block for downsampling. + """ + + def __init__( + self, + in_channels, + emb_channels, + out_channels=None, + dropout=0.0, + use_conv=False, + dims=2, + up=False, + down=False, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.in_channels = in_channels + self.dropout = dropout + self.out_channels = out_channels or self.in_channels + self.use_conv = use_conv + + self.in_layers = nn.Sequential( + normalization(self.in_channels, **factory_kwargs), + nn.SiLU(), + conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 + ) + + self.updown = up or down + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential(nn.SiLU(), nn.Linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) + + self.out_layers = nn.Sequential( + normalization(self.out_channels, **factory_kwargs), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 + ) + + if self.out_channels == self.in_channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 + else: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 + + def forward(self, x, emb) -> torch.Tensor: + # cuDNN autotune is scoped to this block rather than set at import time. + # The benchmark=True win is concentrated in the convolutions right here: + # this ResBlock sits inside UNetDown/UNetUp (patch_embed / final_layer), + # the heaviest conv users in the DiT. A process-wide flag would instead + # silently change every other model's conv behaviour in the same process. + with cudnn_settings(benchmark=True, deterministic=True): + # ``in_layers``/``out_layers`` stay nn.Sequential and are indexed into + # rather than unpacked into named submodules -- that is what keeps the + # state_dict keys (``in_layers.0/.2``, ``out_layers.0/.3``) identical to + # the unfused block, so checkpoints load unchanged. + in_norm, in_conv = self.in_layers[0], self.in_layers[-1] + + # GroupNorm -> SiLU, one kernel instead of two. + h = fused_group_norm_silu(x, in_norm.weight, in_norm.bias, num_groups=in_norm.num_groups, eps=in_norm.eps) + if self.updown: + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + # Adaptive GroupNorm -> SiLU: GroupNorm -> mul -> add -> silu, one + # kernel instead of four. The ``nn.SiLU`` at ``out_layers[1]`` is + # skipped here but stays in the Sequential, so state_dict keys are + # unchanged. + out_norm, out_rest = self.out_layers[0], self.out_layers[2:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = fused_adaptive_group_norm_silu( + h, + out_norm.weight, + out_norm.bias, + scale, + shift, + num_groups=out_norm.num_groups, + eps=out_norm.eps, + ) + h = out_rest(h) + + return self.skip_connection(x) + h + + +__all__ = ["ResBlock"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 4b1da8911..ce2d1d7e2 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import copy import logging @@ -72,7 +72,6 @@ _STEP_OUTPUT_SIZE = "hunyuan_output_size" _STEP_COT_TEXT_LIST = "hunyuan_cot_text_list" _STEP_AR_KV = "hunyuan_ar_kv" _STEP_PROMPT_KV = "hunyuan_prompt_kv" - _HUNYUAN_DEFAULT_OUTPUT_TYPE = "pil" @@ -466,6 +465,11 @@ class HunyuanImage3Pipeline( return self._pipeline def _validate_step_request(self, state: "StepRequestState") -> None: + if self._uses_scheduler_paged_kv(): + raise ValueError( + "HunyuanImage3 paged_scheduler currently supports request-level execution only; " + "disable step execution or use dense_legacy." + ) prompt = state.prompt sampling = state.sampling if prompt is None: @@ -663,6 +667,12 @@ class HunyuanImage3Pipeline( prefix_lens.append(int(prompt_kv[0]["lens"][branch].item())) return prefix_lens + def _uses_scheduler_paged_kv(self) -> bool: + return ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ) + def _merge_step_model_inputs( self, states: list["StepRequestState"], @@ -693,6 +703,7 @@ class HunyuanImage3Pipeline( values, row_branches, prefix_lens or [None] * len(row_branches), + strict=True, ) ] continue @@ -1871,8 +1882,7 @@ class HunyuanImage3Pipeline( row_branches = [branch for branch in range(cfg_factor) for _ in states] return row_state_indexes, row_branches - @staticmethod - def _validate_step_group_states(states: list["StepRequestState"]) -> tuple[bool, int]: + def _validate_step_group_states(self, states: list["StepRequestState"]) -> tuple[bool, int]: if not states: raise ValueError("HunyuanImage3 denoise_step received an empty group.") @@ -2035,6 +2045,11 @@ class HunyuanImage3Pipeline( states = list(input_batch.states) if not states: raise ValueError("HunyuanImage3 denoise_step received an empty batch.") + if self._uses_scheduler_paged_kv(): + raise ValueError( + "HunyuanImage3 paged_scheduler currently supports request-level execution only; " + "disable step execution or use dense_legacy." + ) self._ensure_grouped_attention_backend_supported(len(states)) outputs: dict[str, torch.Tensor] = {} for group in self._split_step_groups(states): diff --git a/vllm_omni/diffusion/models/lingbot_world/__init__.py b/vllm_omni/diffusion/models/lingbot_world/__init__.py new file mode 100644 index 000000000..a7fe05807 --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_world/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .actions import LingBotCameraControlReducer +from .pipeline import ( + LingBotWorldCausalDMDPipeline, + get_lingbot_world_post_process_func, + get_lingbot_world_pre_process_func, +) +from .transformer import CausalLingBotWorldTransformer3DModel + +__all__ = [ + "CausalLingBotWorldTransformer3DModel", + "LingBotCameraControlReducer", + "LingBotWorldCausalDMDPipeline", + "get_lingbot_world_post_process_func", + "get_lingbot_world_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/lingbot_world/actions.py b/vllm_omni/diffusion/models/lingbot_world/actions.py new file mode 100644 index 000000000..e725c698a --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_world/actions.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Realtime keyboard controls for LingBot World 2.0.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from vllm_omni.diffusion.models.lingbot_world.camera import ( + CameraTrajectory, +) +from vllm_omni.experimental.ar_diffusion.session import ( + ARDiffusionPreparedControls, + ARDiffusionSessionEvent, +) +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionControlInput, +) + +LINGBOT_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1" +LINGBOT_CAMERA_TRAJECTORY_SCHEMA = "lingbot.camera_trajectory.v1" +_ACTION_ORDER = ("w", "a", "s", "d", "i", "j", "k", "l") +_VALID_ACTIONS = frozenset(_ACTION_ORDER) +_REFERENCE_HEIGHT = 480 +_REFERENCE_WIDTH = 832 + + +def _normalize_actions(value: object, *, field: str) -> tuple[str, ...]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"{field} must be a sequence of LingBot action keys.") + actions: set[str] = set() + for item in value: + if not isinstance(item, str) or item.lower() not in _VALID_ACTIONS: + raise ValueError(f"{field} supports only W/A/S/D/I/J/K/L action keys.") + actions.add(item.lower()) + return tuple(action for action in _ACTION_ORDER if action in actions) + + +def _normalize_frames(value: object, *, field: str) -> tuple[tuple[str, ...], ...]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"{field} must be a sequence of per-frame action lists.") + return tuple(_normalize_actions(actions, field=f"{field}[{index}]") for index, actions in enumerate(value)) + + +def parse_lingbot_camera_action_frames( + data: Mapping[str, Any], + *, + expected_frames: int, +) -> tuple[tuple[str, ...], ...]: + """Validate the chunk-sized payload emitted by the session reducer.""" + + if data.get("mode") != "frames": + raise ValueError("lingbot.camera_actions.v1 model ticks require mode='frames'.") + frames = _normalize_frames(data.get("frames"), field="camera action frames") + if len(frames) != expected_frames: + raise ValueError( + "lingbot.camera_actions.v1 must contain exactly one action list per " + f"latent frame; expected {expected_frames}, got {len(frames)}." + ) + return frames + + +@dataclass(frozen=True) +class _LingBotCameraReducerState: + mode: str | None = None + held_actions: tuple[str, ...] = () + script_frames: tuple[tuple[str, ...], ...] = () + + +class LingBotCameraControlReducer: + """Sample live state/script controls into one LingBot AR block. + + State mode preserves held keys across chunks and guarantees a one-frame + pulse when a press and release both arrive before the next chunk. Script + mode is a finite FIFO padded with neutral frames after exhaustion. + """ + + def __init__(self, *, frames_per_block: int = 3) -> None: + if isinstance(frames_per_block, bool) or not isinstance(frames_per_block, int) or frames_per_block <= 0: + raise ValueError("frames_per_block must be a positive integer.") + self._frames_per_block = frames_per_block + self._state = _LingBotCameraReducerState() + + @staticmethod + def _state_transitions(data: Mapping[str, Any]) -> tuple[tuple[str, ...], ...]: + transitions = data.get("transitions") + if not isinstance(transitions, Sequence) or isinstance(transitions, (str, bytes)) or not transitions: + raise ValueError("LingBot state-mode camera actions require non-empty transitions.") + normalized: list[tuple[str, ...]] = [] + previous_timestamp: int | None = None + for index, transition in enumerate(transitions): + if not isinstance(transition, Mapping): + raise ValueError("Each LingBot camera state transition must be a mapping.") + timestamp = transition.get("client_ts_ms") + if timestamp is not None: + if isinstance(timestamp, bool) or not isinstance(timestamp, int) or timestamp < 0: + raise ValueError("client_ts_ms must be a non-negative integer when provided.") + if previous_timestamp is not None and timestamp < previous_timestamp: + raise ValueError("client_ts_ms values must be monotonic within one event.") + previous_timestamp = timestamp + normalized.append( + _normalize_actions( + transition.get("actions"), + field=f"camera state transitions[{index}].actions", + ) + ) + return tuple(normalized) + + @staticmethod + def _script(data: Mapping[str, Any]) -> tuple[tuple[str, ...], ...]: + return _normalize_frames(data.get("frames"), field="camera action script frames") + + def prepare( + self, + *, + current_controls: Mapping[str, ARDiffusionControlInput], + events: Sequence[ARDiffusionSessionEvent], + chunk_index: int, + ) -> ARDiffusionPreparedControls: + del chunk_index + controls = { + track: control + for track, control in current_controls.items() + if not (track == "camera" and control.schema == LINGBOT_CAMERA_ACTION_SCHEMA) + } + state = self._state + pending_transitions: tuple[tuple[str, ...], ...] = () + + for event in events: + for control in event.controls: + if control.track != "camera" or control.schema != LINGBOT_CAMERA_ACTION_SCHEMA: + controls[control.track] = control + if control.track == "camera": + state = _LingBotCameraReducerState() + pending_transitions = () + continue + + mode = control.data.get("mode") + controls.pop("camera", None) + if mode == "script": + state = _LingBotCameraReducerState( + mode="script", + script_frames=self._script(control.data), + ) + pending_transitions = () + elif mode == "state": + transitions = self._state_transitions(control.data) + if state.mode != "state": + state = _LingBotCameraReducerState(mode="state") + else: + state = _LingBotCameraReducerState( + mode="state", + held_actions=state.held_actions, + ) + pending_transitions += transitions + else: + raise ValueError("lingbot.camera_actions.v1 events require mode='state' or mode='script'.") + + if state.mode == "script": + frames = state.script_frames[: self._frames_per_block] + frames += ((),) * (self._frames_per_block - len(frames)) + state = _LingBotCameraReducerState( + mode="script", + script_frames=state.script_frames[self._frames_per_block :], + ) + elif state.mode == "state": + final_actions = pending_transitions[-1] if pending_transitions else state.held_actions + pulse = next( + (actions for actions in reversed(pending_transitions) if actions), + None, + ) + if pulse is not None and pulse != final_actions: + frames = (pulse,) + (final_actions,) * (self._frames_per_block - 1) + else: + frames = (final_actions,) * self._frames_per_block + state = _LingBotCameraReducerState( + mode="state", + held_actions=final_actions, + ) + else: + frames = () + + if frames: + controls["camera"] = ARDiffusionControlInput( + track="camera", + schema=LINGBOT_CAMERA_ACTION_SCHEMA, + data={ + "mode": "frames", + "frames": [list(actions) for actions in frames], + }, + ) + return ARDiffusionPreparedControls( + controls=tuple(controls[track] for track in sorted(controls)), + private_state=state, + ) + + def commit(self, prepared: ARDiffusionPreparedControls) -> None: + if not isinstance(prepared.private_state, _LingBotCameraReducerState): + raise TypeError("LingBot reducer prepared state has an unexpected type.") + self._state = prepared.private_state + + def reset(self) -> None: + self._state = _LingBotCameraReducerState() + + +def _rotation_matrix(axis: str, angle: float) -> np.ndarray: + cosine = math.cos(angle) + sine = math.sin(angle) + if axis == "x": + return np.array( + ((1.0, 0.0, 0.0), (0.0, cosine, -sine), (0.0, sine, cosine)), + dtype=np.float64, + ) + return np.array( + ((cosine, 0.0, sine), (0.0, 1.0, 0.0), (-sine, 0.0, cosine)), + dtype=np.float64, + ) + + +def integrate_lingbot_camera_actions( + frames: Sequence[Sequence[str]], + *, + width: int, + height: int, + initial_pose: torch.Tensor | None = None, + initial_pitch: float = 0.0, +) -> tuple[CameraTrajectory, float]: + """Convert latent-frame WASD/IJKL controls to cumulative C2W poses. + + Calibration and motion constants intentionally match SGLang's LingBot + adapter: movement 0.05, pitch 4 degrees, yaw 6 degrees, pitch clamp 85 + degrees, and target-resolution focal lengths of 500 pixels. + """ + + normalized_frames = tuple( + _normalize_actions(actions, field=f"camera action frames[{index}]") for index, actions in enumerate(frames) + ) + if not normalized_frames: + raise ValueError("camera action frames must not be empty.") + if width <= 0 or height <= 0: + raise ValueError("camera action resolution must be positive.") + + if initial_pose is None: + current_pose = np.eye(4, dtype=np.float64) + else: + if initial_pose.shape != (4, 4) or not torch.isfinite(initial_pose).all(): + raise ValueError("initial_pose must be one finite 4x4 camera-to-world matrix.") + current_pose = initial_pose.detach().cpu().double().numpy().copy() + current_pitch = float(initial_pitch) + pitch_limit = math.radians(85.0) + poses: list[np.ndarray] = [] + + for actions in normalized_frames: + rotation = current_pose[:3, :3] + translation = current_pose[:3, 3] + pitch_delta = math.radians(4.0) * (("i" in actions) - ("k" in actions)) + if not -pitch_limit <= current_pitch + pitch_delta <= pitch_limit: + pitch_delta = 0.0 + else: + current_pitch += pitch_delta + yaw_delta = math.radians(6.0) * (("l" in actions) - ("j" in actions)) + new_rotation = _rotation_matrix("y", yaw_delta) @ rotation @ _rotation_matrix("x", pitch_delta) + + forward = np.array((new_rotation[0, 2], 0.0, new_rotation[2, 2])) + right = np.array((new_rotation[0, 0], 0.0, new_rotation[2, 0])) + forward_norm = np.linalg.norm(forward) + right_norm = np.linalg.norm(right) + if forward_norm > 0: + forward /= forward_norm + 1e-6 + if right_norm > 0: + right /= right_norm + 1e-6 + + movement = np.zeros(3, dtype=np.float64) + movement += forward * 0.05 * (("w" in actions) - ("s" in actions)) + movement += right * 0.05 * (("d" in actions) - ("a" in actions)) + current_pose = np.eye(4, dtype=np.float64) + current_pose[:3, :3] = new_rotation + current_pose[:3, 3] = translation + movement + poses.append(current_pose) + + # build_plucker_embedding expects intrinsics in the 832x480 reference + # coordinate system. These values become SGLang's [500, 500, W/2, H/2] + # after that function scales them to the requested resolution. + reference_intrinsics = ( + 500.0 * _REFERENCE_WIDTH / width, + 500.0 * _REFERENCE_HEIGHT / height, + _REFERENCE_WIDTH / 2, + _REFERENCE_HEIGHT / 2, + ) + trajectory = CameraTrajectory( + # Keep the cumulative controller pose in FP64 between chunks, matching + # SGLang's NumPy integration and avoiding long-session drift. + poses=torch.from_numpy(np.stack(poses)), + intrinsics=torch.tensor(reference_intrinsics, dtype=torch.float32).repeat(len(poses), 1), + ) + return trajectory, current_pitch diff --git a/vllm_omni/diffusion/models/lingbot_world/camera.py b/vllm_omni/diffusion/models/lingbot_world/camera.py new file mode 100644 index 000000000..0e78e53e6 --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_world/camera.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Camera trajectory loading and geometry for LingBot World v2 conditioning.""" + +from __future__ import annotations + +import io +import math +import os +import stat +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO + +import numpy as np +import torch + +_REFERENCE_HEIGHT = 480 +_REFERENCE_WIDTH = 832 +_MAX_ACTION_FRAMES = 117 +_MAX_SOURCE_ACTION_FRAMES = 4096 +_MAX_NPY_HEADER_BYTES = 10_000 +_DIRECTORY_OPEN_FLAGS = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW +_FILE_OPEN_FLAGS = os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK + + +@dataclass(frozen=True) +class CameraTrajectory: + """Camera-to-world poses and ``(fx, fy, cx, cy)`` intrinsics by frame.""" + + poses: torch.Tensor + intrinsics: torch.Tensor + + +@dataclass(frozen=True) +class TrustedActionDirectory: + """Canonical action location tied to the trusted root's filesystem identity.""" + + root: Path + relative: Path + root_device: int + root_inode: int + + +def resolve_trusted_action_directory( + action_path: str | os.PathLike[str], + trusted_root: str | os.PathLike[str], +) -> TrustedActionDirectory: + """Resolve an action directory beneath a trusted root without exposing paths.""" + + try: + root = Path(trusted_root).expanduser().resolve(strict=True) + except (OSError, RuntimeError): + raise ValueError("The configured LingBot trusted action root is unavailable.") from None + if not root.is_dir(): + raise ValueError("The configured LingBot trusted action root must be a directory.") + + try: + candidate = Path(action_path).expanduser() + if not candidate.is_absolute(): + candidate = root / candidate + candidate = candidate.resolve(strict=True) + relative = candidate.relative_to(root) + except (OSError, RuntimeError, ValueError): + raise ValueError( + "sampling_params.extra_args.action_path must be contained by the trusted action root." + ) from None + if not candidate.is_dir(): + raise ValueError("sampling_params.extra_args.action_path must identify a directory in the trusted action root.") + + try: + root_stat = root.stat() + except OSError: + raise ValueError("The configured LingBot trusted action root is unavailable.") from None + return TrustedActionDirectory( + root=root, + relative=relative, + root_device=root_stat.st_dev, + root_inode=root_stat.st_ino, + ) + + +def _validate_trajectory( + trajectory: CameraTrajectory, + *, + file_suffix: str = "", +) -> None: + poses = trajectory.poses + intrinsics = trajectory.intrinsics + if poses.ndim != 3 or poses.shape[1:] != (4, 4): + raise ValueError(f"poses{file_suffix} must have shape [frames, 4, 4], got {tuple(poses.shape)}") + if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,): + raise ValueError( + f"intrinsics{file_suffix} must have shape [frames, 4] ordered as fx, fy, cx, cy, " + f"got {tuple(intrinsics.shape)}" + ) + if poses.shape[0] == 0: + raise ValueError("camera trajectory must contain at least one frame") + if poses.shape[0] != intrinsics.shape[0]: + raise ValueError("poses and intrinsics must contain the same number of frames") + if not torch.isfinite(poses).all() or not torch.isfinite(intrinsics).all(): + raise ValueError("camera trajectory values must all be finite") + + +@contextmanager +def _open_action_directory(action_directory: TrustedActionDirectory) -> Iterator[int]: + root_fd: int | None = None + action_fd: int | None = None + try: + root_fd = os.open(action_directory.root, _DIRECTORY_OPEN_FLAGS) + root_stat = os.fstat(root_fd) + if (root_stat.st_dev, root_stat.st_ino) != ( + action_directory.root_device, + action_directory.root_inode, + ): + raise ValueError("The configured LingBot trusted action root changed before camera files were opened.") + + action_fd = os.dup(root_fd) + for component in action_directory.relative.parts: + next_fd = os.open(component, _DIRECTORY_OPEN_FLAGS, dir_fd=action_fd) + os.close(action_fd) + action_fd = next_fd + except OSError: + if action_fd is not None: + os.close(action_fd) + if root_fd is not None: + os.close(root_fd) + raise ValueError("Unable to securely open the LingBot action directory within the trusted root.") from None + except ValueError: + if action_fd is not None: + os.close(action_fd) + if root_fd is not None: + os.close(root_fd) + raise + assert action_fd is not None + assert root_fd is not None + try: + yield action_fd + finally: + os.close(action_fd) + os.close(root_fd) + + +@contextmanager +def _open_camera_file(action_fd: int, filename: str) -> Iterator[tuple[BinaryIO, os.stat_result]]: + file_fd = -1 + try: + file_fd = os.open(filename, _FILE_OPEN_FLAGS, dir_fd=action_fd) + file_stat = os.fstat(file_fd) + if not stat.S_ISREG(file_stat.st_mode): + raise ValueError(f"{filename} must be a regular NPY file.") + with os.fdopen(file_fd, "rb", closefd=True) as file_handle: + file_fd = -1 + yield file_handle, file_stat + except FileNotFoundError: + raise FileNotFoundError(f"camera trajectory file not found: {filename}") from None + except OSError: + raise ValueError(f"{filename} could not be opened securely from the trusted action directory.") from None + finally: + if file_fd >= 0: + os.close(file_fd) + + +def _load_bounded_npy( + action_fd: int, + filename: str, + *, + trailing_shape: tuple[int, ...], +) -> np.ndarray: + with _open_camera_file(action_fd, filename) as (file_handle, file_stat): + max_payload_bytes = _MAX_SOURCE_ACTION_FRAMES * math.prod(trailing_shape) * 8 + max_file_bytes = _MAX_NPY_HEADER_BYTES + max_payload_bytes + 32 + if file_stat.st_size > max_file_bytes: + raise ValueError(f"{filename} byte size exceeds the bounded camera-array limit.") + snapshot = file_handle.read(max_file_bytes + 1) + if len(snapshot) != file_stat.st_size: + raise ValueError(f"{filename} changed while its bounded contents were read.") + snapshot_file = io.BytesIO(snapshot) + try: + version = np.lib.format.read_magic(snapshot_file) + if version == (1, 0): + shape, _fortran_order, dtype = np.lib.format.read_array_header_1_0(snapshot_file) + elif version == (2, 0): + shape, _fortran_order, dtype = np.lib.format.read_array_header_2_0(snapshot_file) + else: + raise ValueError(f"{filename} must use NPY format version 1.0 or 2.0.") + except (EOFError, TypeError, ValueError): + raise ValueError(f"{filename} has an invalid or oversized NPY header.") from None + + if len(shape) != len(trailing_shape) + 1 or tuple(shape[1:]) != trailing_shape: + expected = ", ".join(str(value) for value in trailing_shape) + raise ValueError(f"{filename} must have shape [frames, {expected}], got {shape}.") + frames = shape[0] + if isinstance(frames, bool) or not isinstance(frames, int) or not 1 <= frames <= _MAX_SOURCE_ACTION_FRAMES: + raise ValueError(f"{filename} must contain between 1 and {_MAX_SOURCE_ACTION_FRAMES} source frames.") + if dtype.kind not in "fiu" or dtype.hasobject or dtype.itemsize > 8: + raise ValueError(f"{filename} must contain real numeric values with at most 8 bytes per element.") + + data_offset = snapshot_file.tell() + expected_bytes = math.prod(shape) * dtype.itemsize + if len(snapshot) != data_offset + expected_bytes: + raise ValueError(f"{filename} byte size does not match its NPY header.") + snapshot_file.seek(0) + try: + array = np.load(snapshot_file, allow_pickle=False) + except (EOFError, TypeError, ValueError): + raise ValueError(f"{filename} does not contain a valid numeric NPY array.") from None + # Official LingBot trajectories may be longer than one request. Only + # materialize the bounded prefix that any supported request can consume. + result: np.ndarray = np.asarray(array[:_MAX_ACTION_FRAMES], dtype=np.float32) + return result + + +def load_camera_trajectory(action_directory: TrustedActionDirectory) -> CameraTrajectory: + """Load bounded numeric camera arrays through a trusted directory handle.""" + + if not isinstance(action_directory, TrustedActionDirectory): + raise TypeError("action_directory must be resolved against a trusted action root before loading.") + with _open_action_directory(action_directory) as action_fd: + poses_array = _load_bounded_npy(action_fd, "poses.npy", trailing_shape=(4, 4)) + intrinsics_array = _load_bounded_npy(action_fd, "intrinsics.npy", trailing_shape=(4,)) + + trajectory = CameraTrajectory( + poses=torch.from_numpy(poses_array), + intrinsics=torch.from_numpy(intrinsics_array), + ) + _validate_trajectory(trajectory, file_suffix=".npy") + return trajectory + + +def _linear_interpolate(values: torch.Tensor, target_times: np.ndarray) -> torch.Tensor: + source_times = np.linspace(0.0, 1.0, values.shape[0]) + values_array = values.detach().cpu().double().numpy() + columns = [ + np.interp(target_times, source_times, values_array[:, column]) for column in range(values_array.shape[1]) + ] + interpolated = torch.from_numpy(np.stack(columns, axis=1)) + return interpolated.to(device=values.device, dtype=values.dtype) + + +def _rotation_matrices_to_quaternions(matrices: np.ndarray) -> np.ndarray: + """Convert rotation matrices to normalized ``(w, x, y, z)`` quaternions.""" + + quaternions: np.ndarray = np.empty((matrices.shape[0], 4), dtype=np.float64) + for index, matrix in enumerate(matrices): + trace = float(np.trace(matrix)) + if trace > 0.0: + scale = math.sqrt(trace + 1.0) * 2.0 + quaternion = np.array( + [ + 0.25 * scale, + (matrix[2, 1] - matrix[1, 2]) / scale, + (matrix[0, 2] - matrix[2, 0]) / scale, + (matrix[1, 0] - matrix[0, 1]) / scale, + ] + ) + else: + axis = int(np.argmax(np.diag(matrix))) + if axis == 0: + scale = math.sqrt(max(0.0, 1.0 + matrix[0, 0] - matrix[1, 1] - matrix[2, 2])) * 2.0 + quaternion = np.array( + [ + (matrix[2, 1] - matrix[1, 2]) / scale, + 0.25 * scale, + (matrix[0, 1] + matrix[1, 0]) / scale, + (matrix[0, 2] + matrix[2, 0]) / scale, + ] + ) + elif axis == 1: + scale = math.sqrt(max(0.0, 1.0 + matrix[1, 1] - matrix[0, 0] - matrix[2, 2])) * 2.0 + quaternion = np.array( + [ + (matrix[0, 2] - matrix[2, 0]) / scale, + (matrix[0, 1] + matrix[1, 0]) / scale, + 0.25 * scale, + (matrix[1, 2] + matrix[2, 1]) / scale, + ] + ) + else: + scale = math.sqrt(max(0.0, 1.0 + matrix[2, 2] - matrix[0, 0] - matrix[1, 1])) * 2.0 + quaternion = np.array( + [ + (matrix[1, 0] - matrix[0, 1]) / scale, + (matrix[0, 2] + matrix[2, 0]) / scale, + (matrix[1, 2] + matrix[2, 1]) / scale, + 0.25 * scale, + ] + ) + norm = np.linalg.norm(quaternion) + if not np.isfinite(norm) or norm <= np.finfo(np.float64).eps: + raise ValueError("camera rotation matrices must describe valid rotations") + quaternions[index] = quaternion / norm + return quaternions + + +def _quaternions_to_rotation_matrices(quaternions: np.ndarray) -> np.ndarray: + quaternions = quaternions / np.linalg.norm(quaternions, axis=1, keepdims=True) + w, x, y, z = quaternions.T + matrices: np.ndarray = np.stack( + ( + 1 - 2 * (y * y + z * z), + 2 * (x * y - z * w), + 2 * (x * z + y * w), + 2 * (x * y + z * w), + 1 - 2 * (x * x + z * z), + 2 * (y * z - x * w), + 2 * (x * z - y * w), + 2 * (y * z + x * w), + 1 - 2 * (x * x + y * y), + ), + axis=1, + ).reshape(-1, 3, 3) + return matrices + + +def _slerp_rotations(matrices: np.ndarray, target_times: np.ndarray) -> np.ndarray: + quaternions = _rotation_matrices_to_quaternions(matrices) + interpolated = np.empty((target_times.shape[0], 4), dtype=np.float64) + for target_index, target_time in enumerate(target_times): + scaled_time = float(target_time) * (matrices.shape[0] - 1) + left = min(max(math.floor(scaled_time), 0), matrices.shape[0] - 2) + fraction = scaled_time - left + first = quaternions[left] + second = quaternions[left + 1] + cosine = float(np.dot(first, second)) + if cosine < 0.0: + second = -second + cosine = -cosine + cosine = min(cosine, 1.0) + if cosine > 0.9995: + quaternion = first + fraction * (second - first) + interpolated[target_index] = quaternion / np.linalg.norm(quaternion) + continue + angle = math.acos(cosine) + denominator = math.sin(angle) + interpolated[target_index] = ( + math.sin((1.0 - fraction) * angle) / denominator * first + math.sin(fraction * angle) / denominator * second + ) + return _quaternions_to_rotation_matrices(interpolated) + + +def interpolate_camera_trajectory( + trajectory: CameraTrajectory, + num_frames: int, +) -> CameraTrajectory: + """Resample translations/intrinsics linearly and rotations with quaternion SLERP.""" + + _validate_trajectory(trajectory) + if num_frames <= 0: + raise ValueError(f"num_frames must be positive, got {num_frames}") + + source_frames = trajectory.poses.shape[0] + if num_frames == source_frames: + return trajectory + if source_frames == 1: + return CameraTrajectory( + poses=trajectory.poses.repeat(num_frames, 1, 1), + intrinsics=trajectory.intrinsics.repeat(num_frames, 1), + ) + + # Rotations use quaternion SLERP; translations and intrinsics are linear. + target_times = np.linspace(0.0, 1.0, num_frames) + rotation_matrices = trajectory.poses[:, :3, :3].detach().cpu().double().numpy() + interpolated_rotations = _slerp_rotations(rotation_matrices, target_times) + translations = _linear_interpolate(trajectory.poses[:, :3, 3], target_times) + intrinsics = _linear_interpolate(trajectory.intrinsics, target_times) + + poses = torch.eye( + 4, + device=trajectory.poses.device, + dtype=trajectory.poses.dtype, + ).repeat(num_frames, 1, 1) + poses[:, :3, :3] = torch.from_numpy(interpolated_rotations).to( + device=trajectory.poses.device, + dtype=trajectory.poses.dtype, + ) + poses[:, :3, 3] = translations + return CameraTrajectory(poses=poses, intrinsics=intrinsics) + + +def _invert_camera_poses(poses: torch.Tensor) -> torch.Tensor: + """Invert rigid transforms using R^-1 = R^T and t^-1 = -R^T t.""" + + rotations = poses[:, :3, :3] + translations = poses[:, :3, 3:] + inverse_rotations = rotations.transpose(-1, -2) + inverse = torch.eye(4, device=poses.device, dtype=poses.dtype).repeat(poses.shape[0], 1, 1) + inverse[:, :3, :3] = inverse_rotations + inverse[:, :3, 3:] = -torch.bmm(inverse_rotations, translations) + return inverse + + +def _prepare_framewise_poses(raw_poses: torch.Tensor) -> torch.Tensor: + """Convert raw C2W poses to normalized first-relative framewise deltas.""" + + # Remove global placement, then convert the sequence to framewise deltas. + first_world_to_camera = _invert_camera_poses(raw_poses[:1]) + relative_poses = torch.matmul(first_world_to_camera, raw_poses).clone() + relative_poses[0] = torch.eye(4, device=raw_poses.device, dtype=raw_poses.dtype) + if relative_poses.shape[0] > 1: + relative_poses[1:] = torch.bmm( + _invert_camera_poses(relative_poses[:-1]), + relative_poses[1:], + ) + + # Only relative direction/magnitude matters to the checkpoint; normalizing + # by the largest step removes the trajectory's arbitrary translation scale. + translations = relative_poses[:, :3, 3] + max_translation_norm = torch.linalg.vector_norm(translations, dim=-1).max() + if max_translation_norm > 0: + relative_poses[:, :3, 3] = translations / max_translation_norm + return relative_poses + + +def build_plucker_embedding( + trajectory: CameraTrajectory, + *, + height: int, + width: int, + target_height: int, + target_width: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Build checkpoint-ordered ``(ray origin, ray direction)`` channels.""" + + _validate_trajectory(trajectory) + dimensions = { + "height": height, + "width": width, + "target_height": target_height, + "target_width": target_width, + } + for name, value in dimensions.items(): + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + if dtype not in (torch.float16, torch.bfloat16, torch.float32, torch.float64): + raise ValueError(f"dtype must be floating point, got {dtype}") + + # Compute geometry in at least FP32, then cast at the model boundary. + compute_dtype = torch.float32 if dtype in (torch.float16, torch.bfloat16) else dtype + poses = _prepare_framewise_poses(trajectory.poses.to(device=device, dtype=compute_dtype)) + intrinsics = trajectory.intrinsics.to(device=device, dtype=compute_dtype).clone() + intrinsics[:, (0, 2)] *= width / _REFERENCE_WIDTH + intrinsics[:, (1, 3)] *= height / _REFERENCE_HEIGHT + if torch.any(intrinsics[:, :2] == 0): + raise ValueError("camera focal lengths fx and fy must be non-zero") + + # Sample requested-pixel coordinates directly at the conditioning resolution. + x_coordinates = (torch.arange(target_width, device=device, dtype=compute_dtype) + 0.5) * (width / target_width) + y_coordinates = (torch.arange(target_height, device=device, dtype=compute_dtype) + 0.5) * (height / target_height) + grid_y, grid_x = torch.meshgrid(y_coordinates, x_coordinates, indexing="ij") + grid_x = grid_x.unsqueeze(0) + grid_y = grid_y.unsqueeze(0) + + # Pinhole back-projection: pixel (x, y) becomes the camera-space direction + # ((x-cx)/fx, (y-cy)/fy, 1), which is normalized before rotation. + fx, fy, cx, cy = intrinsics.unbind(dim=1) + camera_x = (grid_x - cx[:, None, None]) / fx[:, None, None] + camera_y = (grid_y - cy[:, None, None]) / fy[:, None, None] + camera_directions = torch.stack( + (camera_x, camera_y, torch.ones_like(camera_x)), + dim=-1, + ) + camera_directions = camera_directions / torch.linalg.vector_norm(camera_directions, dim=-1, keepdim=True) + + # A ray is represented by its camera center (origin) and unit direction in + # the prepared relative coordinate frame. Channel order is checkpoint + # data, not an implementation preference: [origin_xyz, direction_xyz]. + world_directions = torch.einsum( + "fij,fhwj->fhwi", + poses[:, :3, :3], + camera_directions, + ) + world_directions = world_directions / torch.linalg.vector_norm(world_directions, dim=-1, keepdim=True) + ray_origins = poses[:, None, None, :3, 3].expand_as(world_directions) + + embedding = torch.cat((ray_origins, world_directions), dim=-1) + return embedding.permute(0, 3, 1, 2).contiguous().to(dtype=dtype) diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py new file mode 100644 index 000000000..5f71979e8 --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -0,0 +1,1292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request-scoped LingBot-World v2 causal DMD pipeline.""" + +from __future__ import annotations + +import math +import os +from collections.abc import Callable, Iterable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, ClassVar, cast + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import AutoTokenizer, UMT5EncoderModel +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery +from vllm_omni.diffusion.models.lingbot_world.actions import ( + LINGBOT_CAMERA_ACTION_SCHEMA, + LINGBOT_CAMERA_TRAJECTORY_SCHEMA, + integrate_lingbot_camera_actions, + parse_lingbot_camera_action_frames, +) +from vllm_omni.diffusion.models.lingbot_world.camera import ( + CameraTrajectory, + build_plucker_embedding, + interpolate_camera_trajectory, + load_camera_trajectory, + resolve_trusted_action_directory, +) +from vllm_omni.diffusion.models.lingbot_world.transformer import ( + CausalLingBotWorldTransformer3DModel, + LingBotAttentionCache, + LingBotTransformerCache, +) +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler +from vllm_omni.diffusion.models.utils import _load_json +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import load_transformer_config, retrieve_latents +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.experimental.ar_diffusion.capability import ( + ARDiffusionCrossAttentionKVSpec, + ARDiffusionKVBranchSpec, + ARDiffusionKVCacheSpec, +) +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionChunkMetadata, + ARDiffusionTickRequest, +) + +if TYPE_CHECKING: + from tqdm.std import tqdm as TqdmProgressBar + + from vllm_omni.experimental.ar_diffusion.kv_cache.state import ( + ARDiffusionKVState, + ) + +LINGBOT_DMD_TIMESTEPS = (1000, 750, 500, 250) +_CAMERA_SPATIAL_FOLD = 8 +_MAX_PIXEL_AREA = 480 * 832 +_MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096 +_MAX_RAW_FRAMES = 117 +_MAX_SEQUENCE_LENGTH = 512 +_ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT" +_PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory" +_PREPROCESSED_CAMERA_ACTIONS_KEY = "_lingbot_camera_actions" +_SOURCE_IMAGE_ERROR = ( + "Unable to load multi_modal_data.image; expected a decodable image within 4096 * 4096 source pixels." +) + + +@dataclass(frozen=True) +class _LingBotRequestInputs: + """Validated model-specific request boundary.""" + + prompt: str + image: PIL.Image.Image | torch.Tensor + camera_trajectory: CameraTrajectory | None + camera_actions: tuple[tuple[str, ...], ...] | None + height: int + width: int + num_frames: int + num_latent_frames: int + output_type: str + max_sequence_length: int + flow_shift: float + generator: torch.Generator + + +@dataclass +class _LingBotARSessionState: + """Small model-owned state; attention tensors remain runner-owned.""" + + next_chunk_index: int = 0 + prompt: str | None = None + generator_state: torch.Tensor | None = None + image_condition: torch.Tensor | None = None + camera_tail: CameraTrajectory | None = None + camera_pitch: float = 0.0 + + +def _positive_finite_flow_shift(value: Any) -> float: + if isinstance(value, bool): + raise ValueError("flow_shift must be a positive finite number.") + try: + flow_shift = float(value) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError("flow_shift must be a positive finite number.") from exc + if not math.isfinite(flow_shift) or flow_shift <= 0: + raise ValueError("flow_shift must be a positive finite number.") + return flow_shift + + +def _build_shifted_flow_schedule( + *, + flow_shift: float, +) -> tuple[tuple[float, float], ...]: + """Map checkpoint DMD labels to request-local warped timestep/sigma pairs.""" + + # The checkpoint labels index the unshifted [1, ..., 1 / N] training + # lattice. Apply the request's flow shift before using the value both as + # the transformer timestep and the sampling sigma. + base_sigmas = torch.tensor(LINGBOT_DMD_TIMESTEPS, dtype=torch.float64) / 1000 + shifted_numerators = flow_shift * base_sigmas + shifted_sigmas = shifted_numerators / ((1.0 - base_sigmas) + shifted_numerators) + warped_timesteps = shifted_sigmas * 1000 + return tuple( + (float(timestep), float(sigma)) + for timestep, sigma in zip(warped_timesteps.tolist(), shifted_sigmas.tolist(), strict=True) + ) + + +def _validate_scheduler_config(config: dict[str, Any]) -> None: + contract = { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "prediction_type": "flow_prediction", + "predict_x0": True, + "use_flow_sigmas": True, + "use_dynamic_shifting": False, + "use_beta_sigmas": False, + "use_exponential_sigmas": False, + "use_karras_sigmas": False, + "final_sigmas_type": "zero", + "timestep_spacing": "linspace", + } + for name, expected in contract.items(): + if config.get(name) != expected: + raise ValueError(f"LingBot scheduler config {name} must be {expected!r}, got {config.get(name)!r}.") + + +def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None: + if getattr(od_config, "quantization_config", None) is not None: + raise NotImplementedError("LingBot World v1 does not support quantization.") + parallel_config = getattr(od_config, "parallel_config", None) + if parallel_config is None: + return + unsupported_sizes = { + "pipeline_parallel_size": "pipeline parallelism", + "sequence_parallel_size": "sequence parallelism", + "cfg_parallel_size": "CFG parallelism", + "vae_patch_parallel_size": "VAE parallelism", + } + for field, feature in unsupported_sizes.items(): + size = getattr(parallel_config, field, 1) or 1 + if size > 1: + raise NotImplementedError(f"LingBot World v1 does not support {feature} ({field}={size}).") + if getattr(parallel_config, "use_hsdp", False): + raise NotImplementedError("LingBot World v1 does not support HSDP.") + if getattr(parallel_config, "enable_expert_parallel", False): + raise NotImplementedError("LingBot World v1 does not support expert parallelism.") + + +def _validate_source_image_size(image: PIL.Image.Image) -> None: + width, height = image.size + if ( + isinstance(width, bool) + or not isinstance(width, int) + or width <= 0 + or isinstance(height, bool) + or not isinstance(height, int) + or height <= 0 + ): + raise ValueError("source image width and height must be positive integers.") + if width * height > _MAX_SOURCE_IMAGE_PIXELS: + raise ValueError("source image pixel count must not exceed 4096 * 4096.") + + +def _decode_source_image(image: PIL.Image.Image) -> PIL.Image.Image: + _validate_source_image_size(image) + try: + return image.convert("RGB") + except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError): + raise ValueError(_SOURCE_IMAGE_ERROR) from None + + +def _load_source_image(path: str | os.PathLike[str]) -> PIL.Image.Image: + try: + source_image = PIL.Image.open(path) + except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError): + raise ValueError(_SOURCE_IMAGE_ERROR) from None + try: + return _decode_source_image(source_image) + finally: + source_image.close() + + +def get_lingbot_world_pre_process_func( + od_config: OmniDiffusionConfig, +) -> Callable[[OmniDiffusionRequest], OmniDiffusionRequest]: + """Materialize request-local files once before dispatching GPU workers.""" + + model_config = getattr(od_config, "model_config", None) or {} + configured_action_root = model_config.get("lingbot_action_root") or os.environ.get(_ACTION_ROOT_ENV) + + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + prompt = request.prompt + if not isinstance(prompt, dict): + raise ValueError("LingBot World requires a prompt mapping containing multi_modal_data.image.") + multi_modal_data = prompt.get("multi_modal_data") or {} + additional_information = prompt.get("additional_information") or {} + if not isinstance(multi_modal_data, dict): + raise ValueError("prompt.multi_modal_data must be a mapping containing image.") + if not isinstance(additional_information, dict): + raise ValueError("prompt.additional_information must be a mapping.") + if "action_path" in additional_information: + raise ValueError( + "prompt.additional_information.action_path is not supported; use " + "sampling_params.extra_args.action_path." + ) + + image = multi_modal_data.get("image") + if isinstance(image, list): + raise ValueError("LingBot World requires one image and does not accept an image list.") + if image is None: + raise ValueError("LingBot World requires exactly one image in multi_modal_data.image.") + if isinstance(image, (str, os.PathLike)): + image = _load_source_image(image) + elif isinstance(image, PIL.Image.Image): + image = _decode_source_image(image) + elif not isinstance(image, torch.Tensor): + raise ValueError("multi_modal_data.image must be a PIL image, tensor, or file path.") + + extra_args = getattr(request.sampling_params, "extra_args", None) or {} + if not isinstance(extra_args, dict): + raise ValueError("sampling_params.extra_args must be a mapping.") + tick = ARDiffusionTickRequest.from_extra_args(extra_args) + if tick is not None: + camera_controls = [control for control in tick.controls if control.track == "camera"] + if len(camera_controls) != 1: + raise ValueError("LingBot realtime ticks require exactly one camera control.") + camera_control = camera_controls[0] + if camera_control.schema == LINGBOT_CAMERA_TRAJECTORY_SCHEMA: + camera_data = camera_control.data + try: + poses = torch.as_tensor(camera_data["poses"], dtype=torch.float32) + intrinsics = torch.as_tensor( + camera_data["intrinsics"], + dtype=torch.float32, + ) + except (KeyError, TypeError, ValueError): + raise ValueError( + "lingbot.camera_trajectory.v1 requires numeric poses and intrinsics fields." + ) from None + if poses.ndim != 3 or poses.shape[1:] != (4, 4): + raise ValueError("camera control poses must have shape [frames, 4, 4].") + if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,): + raise ValueError("camera control intrinsics must have shape [frames, 4].") + if poses.shape[0] != intrinsics.shape[0] or poses.shape[0] == 0: + raise ValueError( + "camera control poses and intrinsics must contain the same positive number of frames." + ) + if not torch.isfinite(poses).all() or not torch.isfinite(intrinsics).all(): + raise ValueError("camera control values must all be finite.") + trajectory = CameraTrajectory( + poses=poses, + intrinsics=intrinsics, + ) + camera_actions = None + elif camera_control.schema == LINGBOT_CAMERA_ACTION_SCHEMA: + trajectory = None + camera_actions = parse_lingbot_camera_action_frames( + camera_control.data, + expected_frames=3, + ) + else: + raise ValueError( + "LingBot realtime camera controls require schema " + "'lingbot.camera_trajectory.v1' or 'lingbot.camera_actions.v1'." + ) + else: + action_path = extra_args.get("action_path") + if not isinstance(action_path, (str, os.PathLike)) or not str(action_path): + raise ValueError("action_path is required in sampling_params.extra_args.action_path.") + if not configured_action_root: + raise ValueError( + "sampling_params.extra_args.action_path requires a trusted action root configured by " + f"model_config.lingbot_action_root or {_ACTION_ROOT_ENV}." + ) + action_directory = resolve_trusted_action_directory( + action_path, + configured_action_root, + ) + try: + trajectory = load_camera_trajectory(action_directory) + except OSError: + raise ValueError( + "Unable to load camera trajectory from action_path; expected poses.npy and intrinsics.npy." + ) from None + camera_actions = None + + updated_prompt = dict(prompt) + updated_multi_modal_data = dict(multi_modal_data) + updated_multi_modal_data["image"] = image + updated_prompt["multi_modal_data"] = updated_multi_modal_data + request.prompt = updated_prompt + request.sampling_params.extra_args = { + **extra_args, + _PREPROCESSED_CAMERA_KEY: trajectory, + _PREPROCESSED_CAMERA_ACTIONS_KEY: camera_actions, + } + return request + + return pre_process_func + + +def _fold_camera_embedding( + camera_embedding: torch.Tensor, +) -> torch.Tensor: + """Pixel-unshuffle ``[frames, 6, H, W]`` onto the Wan latent grid.""" + + folded = F.pixel_unshuffle(camera_embedding, _CAMERA_SPATIAL_FOLD) + return folded.permute(1, 0, 2, 3).unsqueeze(0).contiguous() + + +def get_lingbot_world_post_process_func(od_config: OmniDiffusionConfig) -> Callable[..., Any]: + del od_config + from diffusers.video_processor import VideoProcessor + + video_processor = VideoProcessor(vae_scale_factor=8) + + def post_process_func( + video: torch.Tensor, + output_type: str = "np", + sampling_params: Any | None = None, + ) -> Any: + if isinstance(video, dict) and isinstance(video.get("payload"), dict): + return video + if sampling_params is not None: + output_type = getattr(sampling_params, "output_type", None) or output_type + if output_type == "latent": + return video + return { + "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, + "metadata": {}, + } + + return post_process_func + + +class LingBotWorldCausalDMDPipeline( + nn.Module, + SupportImageInput, + SupportsComponentDiscovery, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, +): + """LingBot-World v2 I2V generation with a request-local causal cache.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + # Generic warmup cannot synthesize the required camera action directory. + dummy_run_num_frames: ClassVar[int] = 0 + _AR_BRANCH = "main" + _AR_TEXT_CACHE = "text" + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ) -> None: + super().__init__() + del prefix + _validate_parallel_config(od_config) + self.od_config = od_config + self.device = get_local_device() + dtype = getattr(od_config, "dtype", torch.bfloat16) + model = od_config.model + local_files_only = os.path.exists(model) + managed_component_placement = bool( + getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False) + ) + + # Standard components use from_pretrained; the custom transformer uses the loader. + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + subfolders = ["tokenizer", "text_encoder", "vae"] + prefetch_subfolders(model, subfolders, local_files_only=local_files_only) + self.tokenizer = from_pretrained_with_prefetch( + AutoTokenizer.from_pretrained, + model, + subfolder="tokenizer", + prefetch_list=subfolders, + local_files_only=local_files_only, + ) + self.text_encoder = from_pretrained_with_prefetch( + UMT5EncoderModel.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ) + if not managed_component_placement: + self.text_encoder = self.text_encoder.to(self.device) + self.vae = from_pretrained_with_prefetch( + DistributedAutoencoderKLWan.from_pretrained, + model, + subfolder="vae", + prefetch_list=subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ) + if not managed_component_placement: + self.vae = self.vae.to(self.device) + + transformer_config = load_transformer_config(model, "transformer", local_files_only) + self.transformer = CausalLingBotWorldTransformer3DModel.from_config( + transformer_config, + quant_config=getattr(od_config, "quantization_config", None), + prefix="transformer", + ) + + scheduler_config = _load_json(model, "scheduler/scheduler_config.json", local_files_only) + _validate_scheduler_config(scheduler_config) + configured_shift = getattr(od_config, "flow_shift", None) + checkpoint_shift = scheduler_config.get("flow_shift", scheduler_config.get("shift", 5.0)) + scheduler_shift = _positive_finite_flow_shift( + checkpoint_shift if configured_shift is None else configured_shift + ) + scheduler_keys = { + "num_train_timesteps", + "solver_order", + "prediction_type", + "use_dynamic_shifting", + "thresholding", + "dynamic_thresholding_ratio", + "sample_max_value", + "predict_x0", + "solver_type", + "lower_order_final", + "disable_corrector", + "timestep_spacing", + "steps_offset", + "final_sigmas_type", + } + scheduler_kwargs = {name: value for name, value in scheduler_config.items() if name in scheduler_keys} + scheduler_kwargs["shift"] = scheduler_shift + self.scheduler = FlowUniPCMultistepScheduler(**scheduler_kwargs) + + self.vae_scale_factor_temporal = int(getattr(self.vae.config, "scale_factor_temporal", 4)) + self.vae_scale_factor_spatial = int(getattr(self.vae.config, "scale_factor_spatial", 8)) + model_config = getattr(od_config, "model_config", None) or {} + self._ar_height = int(model_config.get("ar_diffusion_height", 480)) + self._ar_width = int(model_config.get("ar_diffusion_width", 832)) + self._ar_diffusion_kv_state: ARDiffusionKVState | None = None + self._ar_sessions: dict[str, _LingBotARSessionState] = {} + self.setup_diffusion_pipeline_profiler( + profiler_targets=[ + "vae.encode", + "vae.decode", + "_generate_block", + "text_encoder.forward", + "tokenizer.forward", + ], + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, + ) + + def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec: + """Describe the fixed worker-local cache geometry for realtime ticks.""" + patch_frames, patch_height, patch_width = self.transformer.config.patch_size + if patch_frames != 1: + raise ValueError("LingBot AR-Diffusion requires temporal patch size 1.") + spatial = self.vae_scale_factor_spatial + if self._ar_height % (spatial * patch_height) or self._ar_width % (spatial * patch_width): + raise ValueError("ar_diffusion_height/width must align with VAE and DiT patch sizes.") + latent_height = self._ar_height // spatial + latent_width = self._ar_width // spatial + tokens_per_frame = (latent_height // patch_height) * (latent_width // patch_width) + tp_size = get_tensor_model_parallel_world_size() + num_local_heads = int(self.transformer.config.num_attention_heads) // tp_size + total_window_frames = ( + int(self.transformer.config.local_attn_size) + if int(self.transformer.config.local_attn_size) != -1 + else int(self.transformer.config.sliding_window_num_frames) + ) + sink_frames = int(self.transformer.config.sink_size) + recent_window_frames = total_window_frames - sink_frames + if recent_window_frames <= 0: + raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.") + horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 + condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels) + condition_bytes_per_session = ( + condition_channels + * horizon_latent_frames + * latent_height + * latent_width + * torch.empty((), dtype=self.transformer.dtype).element_size() + ) + return ARDiffusionKVCacheSpec( + num_layers=int(self.transformer.config.num_layers), + num_kv_heads=num_local_heads, + head_size=int(self.transformer.config.attention_head_dim), + tokens_per_frame=tokens_per_frame, + frames_per_block=int(self.transformer.config.num_frames_per_block), + window_frames=recent_window_frames, + sink_frames=sink_frames, + kv_branches=(ARDiffusionKVBranchSpec(self._AR_BRANCH, 0),), + session_capacity=2, + cross_attention=( + ARDiffusionCrossAttentionKVSpec( + self._AR_TEXT_CACHE, + _MAX_SEQUENCE_LENGTH, + ), + ), + model_owned_state_bytes_per_session=condition_bytes_per_session, + ) + + @contextmanager + def bind_ar_diffusion_state( + self, + session_id: str, + state: ARDiffusionKVState, + ) -> Iterator[None]: + if self._ar_diffusion_kv_state is not None: + raise RuntimeError("LingBot AR-Diffusion state is already bound.") + if state.session_id != session_id: + raise ValueError(f"LingBot bound session mismatch: {state.session_id!r} != {session_id!r}.") + self._ar_diffusion_kv_state = state + try: + yield + finally: + self._ar_diffusion_kv_state = None + + def reset_ar_diffusion_session(self, session_id: str) -> None: + self._ar_sessions.pop(session_id, None) + + def close_ar_diffusion_session(self, session_id: str) -> None: + self._ar_sessions.pop(session_id, None) + + def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs: + if req.num_reqs != 1 or len(req.prompts) != 1: + raise ValueError("LingBot World supports a single prompt request, not request batching.") + sampling = req.sampling_params + if int(sampling.num_outputs_per_prompt or 1) != 1: + raise ValueError("LingBot World requires num_outputs_per_prompt=1.") + generator = getattr(sampling, "generator", None) + if isinstance(generator, list): + raise ValueError("LingBot World accepts one torch.Generator, not a generator list.") + if not isinstance(generator, torch.Generator): + raise ValueError("LingBot World requires the runner-provided torch.Generator.") + if getattr(sampling, "latents", None) is not None: + raise ValueError("LingBot World does not support caller-provided latents.") + + prompt_value = req.prompts[0] + multi_modal_data: dict[str, Any] + if isinstance(prompt_value, str): + prompt = prompt_value + multi_modal_data = {} + elif isinstance(prompt_value, dict): + prompt = prompt_value.get("prompt") or "" + multi_modal_data = prompt_value.get("multi_modal_data") or {} + else: + raise ValueError("prompt must be a string or prompt mapping.") + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError("prompt must contain non-empty text.") + if not isinstance(multi_modal_data, dict): + raise ValueError("prompt.multi_modal_data must be a mapping containing image.") + + image = multi_modal_data.get("image") + if isinstance(image, list): + raise ValueError("LingBot World requires one image and does not accept an image list.") + if image is None: + raise ValueError("LingBot World requires exactly one image in multi_modal_data.image.") + if isinstance(image, (str, os.PathLike)): + raise ValueError("file-path images must be materialized by the LingBot pre-process function.") + if not isinstance(image, (PIL.Image.Image, torch.Tensor)): + raise ValueError("multi_modal_data.image must be a PIL image, tensor, or file path.") + + extra_args = getattr(sampling, "extra_args", None) or {} + if not isinstance(extra_args, dict): + raise ValueError("sampling_params.extra_args must be a mapping.") + camera_trajectory = extra_args.get(_PREPROCESSED_CAMERA_KEY) + camera_actions = extra_args.get(_PREPROCESSED_CAMERA_ACTIONS_KEY) + if camera_trajectory is not None and not isinstance(camera_trajectory, CameraTrajectory): + raise ValueError("LingBot camera trajectory must be materialized by the pre-process function.") + if camera_actions is not None: + camera_actions = tuple(tuple(actions) for actions in camera_actions) + if (camera_trajectory is None) == (camera_actions is None): + raise ValueError("LingBot pre-processing must materialize exactly one camera input.") + + request_flow_shift = ( + extra_args["flow_shift"] if "flow_shift" in extra_args else getattr(self.scheduler.config, "shift", 5.0) + ) + flow_shift = _positive_finite_flow_shift(request_flow_shift) + + height = getattr(sampling, "height", None) + width = getattr(sampling, "width", None) + if isinstance(image, PIL.Image.Image): + image_width, image_height = image.size + elif image.ndim == 3 and image.shape[0] == 3: + image_height, image_width = image.shape[-2:] + elif image.ndim == 4 and image.shape[:2] == (1, 3): + image_height, image_width = image.shape[-2:] + else: + raise ValueError("tensor image must have shape [3, height, width] or [1, 3, height, width].") + if image_height <= 0 or image_width <= 0: + raise ValueError("source image width and height must be positive integers.") + if image_height * image_width > _MAX_SOURCE_IMAGE_PIXELS: + raise ValueError("source image pixel count must not exceed 4096 * 4096.") + if (height is None) != (width is None): + raise ValueError("height and width must either both be provided or both be omitted.") + patch_size = tuple(self.transformer.config.patch_size) + if len(patch_size) != 3 or patch_size[0] != 1: + raise RuntimeError( + "transformer.config.patch_size must be a three-dimensional tuple with temporal patch size 1." + ) + height_divisor = self.vae_scale_factor_spatial * patch_size[1] + width_divisor = self.vae_scale_factor_spatial * patch_size[2] + if height is None or width is None: + aspect_ratio = image_height / image_width + latent_height = round( + math.sqrt(_MAX_PIXEL_AREA * aspect_ratio) + // self.vae_scale_factor_spatial + // patch_size[1] + * patch_size[1] + ) + latent_width = round( + math.sqrt(_MAX_PIXEL_AREA / aspect_ratio) + // self.vae_scale_factor_spatial + // patch_size[2] + * patch_size[2] + ) + height = latent_height * self.vae_scale_factor_spatial + width = latent_width * self.vae_scale_factor_spatial + if isinstance(height, bool) or not isinstance(height, int) or height <= 0: + raise ValueError(f"height must be a positive integer, got {height!r}.") + if isinstance(width, bool) or not isinstance(width, int) or width <= 0: + raise ValueError(f"width must be a positive integer, got {width!r}.") + if height % height_divisor: + raise ValueError(f"height must be divisible by {height_divisor}, got {height}.") + if width % width_divisor: + raise ValueError(f"width must be divisible by {width_divisor}, got {width}.") + if height * width > _MAX_PIXEL_AREA: + raise ValueError("height * width pixel area must not exceed 480 * 832.") + + num_inference_steps = getattr(sampling, "num_inference_steps", None) + if num_inference_steps is None: + num_inference_steps = len(LINGBOT_DMD_TIMESTEPS) + if num_inference_steps != len(LINGBOT_DMD_TIMESTEPS): + raise ValueError( + "num_inference_steps must be 4 for LingBot World causal DMD " + f"timesteps {list(LINGBOT_DMD_TIMESTEPS)}, got {num_inference_steps}." + ) + + num_frames = getattr(sampling, "num_frames", None) + if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: + raise ValueError(f"num_frames must be a positive integer, got {num_frames!r}.") + if num_frames > _MAX_RAW_FRAMES: + raise ValueError(f"num_frames must not exceed {_MAX_RAW_FRAMES}.") + temporal_factor = self.vae_scale_factor_temporal + if (num_frames - 1) % temporal_factor: + raise ValueError( + "num_frames must satisfy the causal Wan VAE geometry " + f"(num_frames - 1) divisible by {temporal_factor}, got {num_frames}." + ) + num_latent_frames = (num_frames - 1) // temporal_factor + 1 + block_frames = int(self.transformer.config.num_frames_per_block) + if num_latent_frames % block_frames: + raise ValueError( + "num_frames must map to a whole number of configured three-frame latent blocks; " + f"got num_frames={num_frames}, latent_frames={num_latent_frames}, block_frames={block_frames}." + ) + + max_sequence_length = getattr(sampling, "max_sequence_length", None) + if max_sequence_length is None: + max_sequence_length = _MAX_SEQUENCE_LENGTH + if isinstance(max_sequence_length, bool) or not isinstance(max_sequence_length, int): + raise ValueError(f"max_sequence_length must be exactly {_MAX_SEQUENCE_LENGTH}.") + if max_sequence_length != _MAX_SEQUENCE_LENGTH: + raise ValueError(f"max_sequence_length must be exactly {_MAX_SEQUENCE_LENGTH}.") + + return _LingBotRequestInputs( + prompt=prompt.strip(), + image=image, + camera_trajectory=camera_trajectory, + camera_actions=camera_actions, + height=height, + width=width, + num_frames=num_frames, + num_latent_frames=num_latent_frames, + output_type=getattr(sampling, "output_type", None) or "np", + max_sequence_length=max_sequence_length, + flow_shift=flow_shift, + generator=generator, + ) + + def _prepare_image_tensor(self, image: PIL.Image.Image | torch.Tensor, *, height: int, width: int) -> torch.Tensor: + if isinstance(image, PIL.Image.Image): + array = np.asarray(image, dtype=np.float32).copy() + image_tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) / 255.0 + else: + image_tensor = image.detach() + if image_tensor.ndim == 3: + image_tensor = image_tensor.unsqueeze(0) + image_tensor = image_tensor.to(dtype=torch.float32) + if not torch.isfinite(image_tensor).all(): + raise ValueError("tensor image values must all be finite.") + minimum = image_tensor.min().item() + maximum = image_tensor.max().item() + if minimum >= 0.0 and maximum > 255.0: + raise ValueError("tensor image values must be in [0, 1], [0, 255], or [-1, 1].") + if minimum < -1.0 or (minimum < 0.0 and maximum > 1.0): + raise ValueError("tensor image values must be in [0, 1], [0, 255], or [-1, 1].") + if maximum > 1.0: + image_tensor = image_tensor / 255.0 + if image_tensor.min().item() >= 0.0: + image_tensor = image_tensor * 2.0 - 1.0 + if image_tensor.shape[-2:] != (height, width): + image_tensor = F.interpolate( + image_tensor, + size=(height, width), + mode="bicubic", + align_corners=False, + ) + return image_tensor.to(device=self.device, dtype=torch.float32) + + def _vae_latent_stats(self, reference: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + shape = (1, -1, 1, 1, 1) + latent_mean = torch.as_tensor( + self.vae.config.latents_mean, + device=reference.device, + dtype=reference.dtype, + ).view(*shape) + latent_std = torch.as_tensor( + self.vae.config.latents_std, + device=reference.device, + dtype=reference.dtype, + ).view(*shape) + return latent_mean, latent_std + + def _prepare_condition(self, inputs: _LingBotRequestInputs, *, dtype: torch.dtype) -> torch.Tensor: + """Encode the first frame as ``[mask4, image_latent16]``.""" + + image = self._prepare_image_tensor(inputs.image, height=inputs.height, width=inputs.width) + video_condition = image.new_zeros(1, 3, inputs.num_frames, inputs.height, inputs.width) + video_condition[:, :, 0] = image + latent_condition = retrieve_latents( + self.vae.encode(video_condition.to(dtype=self.vae.dtype)), + sample_mode="argmax", + ) + if latent_condition.shape != ( + 1, + self.transformer.config.out_channels, + inputs.num_latent_frames, + inputs.height // self.vae_scale_factor_spatial, + inputs.width // self.vae_scale_factor_spatial, + ): + raise RuntimeError( + f"vae.encode returned an incompatible image latent shape: got {tuple(latent_condition.shape)}." + ) + latent_mean, latent_std = self._vae_latent_stats(latent_condition) + latent_condition = (latent_condition - latent_mean) / latent_std + temporal_mask = latent_condition.new_zeros( + 1, + self.vae_scale_factor_temporal, + inputs.num_latent_frames, + latent_condition.shape[-2], + latent_condition.shape[-1], + ) + temporal_mask[:, :, 0] = 1 + condition = torch.cat((temporal_mask, latent_condition), dim=1).to(dtype=dtype) + if condition.shape[1] != 20: + raise RuntimeError( + "LingBot image condition must contain 4 temporal-mask then 16 image-latent channels, " + f"got {condition.shape[1]}." + ) + return condition + + def _prepare_camera( + self, + inputs: _LingBotRequestInputs, + *, + dtype: torch.dtype, + previous: CameraTrajectory | None = None, + ) -> tuple[torch.Tensor, CameraTrajectory]: + """Convert raw camera frames to a latent-aligned ray tensor.""" + + trajectory = inputs.camera_trajectory + if trajectory is None: + raise RuntimeError("LingBot camera trajectory was not prepared before camera embedding.") + available_frames = int(trajectory.poses.shape[0]) + if inputs.camera_actions is not None: + if available_frames != inputs.num_latent_frames: + raise ValueError( + "camera actions must produce exactly one pose per latent frame; " + f"got camera_frames={available_frames}, latent_frames={inputs.num_latent_frames}." + ) + else: + if available_frames < inputs.num_frames: + raise ValueError( + "camera trajectory frames must be at least num_frames; " + f"got camera_frames={available_frames}, num_frames={inputs.num_frames}." + ) + trajectory = CameraTrajectory( + poses=trajectory.poses[: inputs.num_frames], + intrinsics=trajectory.intrinsics[: inputs.num_frames], + ) + trajectory = interpolate_camera_trajectory(trajectory, inputs.num_latent_frames) + embedding_trajectory = trajectory + drop_anchor = False + if previous is not None: + embedding_trajectory = CameraTrajectory( + poses=torch.cat((previous.poses, trajectory.poses), dim=0), + intrinsics=torch.cat( + (previous.intrinsics, trajectory.intrinsics), + dim=0, + ), + ) + drop_anchor = True + elif inputs.camera_actions is not None: + # The action integrator returns post-action poses. Keep the first + # action visible to framewise-delta conditioning by prepending its + # known pre-action state (identity for a new realtime session). + identity = torch.eye( + 4, + device=trajectory.poses.device, + dtype=trajectory.poses.dtype, + ).unsqueeze(0) + embedding_trajectory = CameraTrajectory( + poses=torch.cat((identity, trajectory.poses), dim=0), + intrinsics=torch.cat( + (trajectory.intrinsics[:1], trajectory.intrinsics), + dim=0, + ), + ) + drop_anchor = True + camera_embedding = build_plucker_embedding( + embedding_trajectory, + height=inputs.height, + width=inputs.width, + target_height=inputs.height, + target_width=inputs.width, + device=self.device, + dtype=dtype, + ) + if drop_anchor: + camera_embedding = camera_embedding[1:] + tail = CameraTrajectory( + poses=trajectory.poses[-1:].clone(), + intrinsics=trajectory.intrinsics[-1:].clone(), + ) + return _fold_camera_embedding(camera_embedding), tail + + def _ar_text_caches( + self, + prompt_embeds: torch.Tensor, + *, + invalidate: bool, + ) -> list[LingBotAttentionCache]: + state = self._ar_diffusion_kv_state + if state is None: + raise RuntimeError("LingBot AR text cache requested without a bound state.") + if invalidate: + state.clear_cross_attention() + if not state.is_cross_attention_populated( + self._AR_BRANCH, + self._AR_TEXT_CACHE, + ): + projected_text = self.transformer.text_embedding(prompt_embeds) + + def layer_kv() -> Iterator[tuple[torch.Tensor, torch.Tensor]]: + for block in self.transformer.blocks: + cross_attention = block.cross_attn + key = cross_attention.norm_k(cross_attention.k(projected_text)).unflatten( + 2, + ( + cross_attention.num_local_heads, + cross_attention.head_dim, + ), + ) + value = cross_attention.v(projected_text).unflatten( + 2, + ( + cross_attention.num_local_heads, + cross_attention.head_dim, + ), + ) + yield key, value + + state.populate_cross_attention( + self._AR_BRANCH, + self._AR_TEXT_CACHE, + layer_kv(), + ) + return [ + LingBotAttentionCache( + key=layer["k"], + value=layer["v"], + end=prompt_embeds.shape[1], + absolute_end=prompt_embeds.shape[1], + last_start=0, + ) + for layer in state.get_cross_attention_kv( + self._AR_BRANCH, + self._AR_TEXT_CACHE, + ) + ] + + def _ar_transformer_cache( + self, + *, + condition: torch.Tensor, + cross_attention: list[LingBotAttentionCache], + commit_current: bool, + ) -> LingBotTransformerCache: + state = self._ar_diffusion_kv_state + if state is None: + raise RuntimeError("LingBot AR cache requested without a bound state.") + patch_frames, patch_height, patch_width = self.transformer.config.patch_size + frames = condition.shape[2] // patch_frames + height = condition.shape[3] // patch_height + width = condition.shape[4] // patch_width + seq_len = frames * height * width + return LingBotTransformerCache( + self_attention=state.get_kv_caches( + self._AR_BRANCH, + seq_len=seq_len, + commit_current=commit_current, + ), + cross_attention=cross_attention, + ) + + def _generate_block( + self, + *, + condition: torch.Tensor, + camera: torch.Tensor, + prompt_embeds: torch.Tensor, + cache: LingBotTransformerCache | None, + ar_cross_attention: list[LingBotAttentionCache] | None, + start_frame: int, + schedule: tuple[tuple[float, float], ...], + generator: torch.Generator, + progress_bar: TqdmProgressBar[Any], + ) -> torch.Tensor: + block_shape = ( + 1, + self.transformer.config.out_channels, + condition.shape[2], + condition.shape[3], + condition.shape[4], + ) + current_latents = randn_tensor( + block_shape, + generator=generator, + device=self.device, + dtype=torch.float32, + ) + for step_index, (timestep_value, sigma) in enumerate(schedule): + if not self.od_config.enforce_eager: + torch.compiler.cudagraph_mark_step_begin() + set_forward_context_denoise_step_idx(step_index) + timestep = torch.full((1,), float(timestep_value), device=self.device, dtype=torch.float32) + # Checkpoint channel contract: + # [noise/x_t(16), temporal_mask(4), image_latent(16)] -> 36. + model_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1) + current_cache = ( + self._ar_transformer_cache( + condition=condition, + cross_attention=cast( + list[LingBotAttentionCache], + ar_cross_attention, + ), + commit_current=False, + ) + if ar_cross_attention is not None + else cast(LingBotTransformerCache, cache) + ) + flow_prediction = self.transformer( + hidden_states=model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + camera_hidden_states=camera, + cache=current_cache, + start_frame=start_frame, + update_cache=False, + ) + if flow_prediction.shape != current_latents.shape: + raise RuntimeError( + "transformer flow prediction shape must match the 16-channel noise latent, " + f"got {tuple(flow_prediction.shape)} and {tuple(current_latents.shape)}." + ) + # The checkpoint's flow parameterization is inverted by + # x0 = x_t - sigma * flow. Intermediate steps re-noise that x0 + # estimate at the next sigma; the final step keeps x0 as this + # block's generated latent. + x0 = current_latents - sigma * flow_prediction.float() + if step_index + 1 < len(schedule): + next_sigma = schedule[step_index + 1][1] + noise = randn_tensor( + current_latents.shape, + generator=generator, + device=self.device, + dtype=torch.float32, + ) + current_latents = (1.0 - next_sigma) * x0 + next_sigma * noise + else: + current_latents = x0 + progress_bar.update() + + # Commit K/V only for the final clean block, never for noisy probes. + cache_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1) + commit_cache = ( + self._ar_transformer_cache( + condition=condition, + cross_attention=cast( + list[LingBotAttentionCache], + ar_cross_attention, + ), + commit_current=True, + ) + if ar_cross_attention is not None + else cast(LingBotTransformerCache, cache) + ) + if not self.od_config.enforce_eager: + torch.compiler.cudagraph_mark_step_begin() + self.transformer( + hidden_states=cache_input, + timestep=torch.zeros(1, device=self.device, dtype=torch.float32), + encoder_hidden_states=prompt_embeds, + camera_hidden_states=camera, + cache=commit_cache, + start_frame=start_frame, + update_cache=True, + ) + if ar_cross_attention is not None: + state = self._ar_diffusion_kv_state + if state is None: + raise RuntimeError("LingBot AR state disappeared before KV commit.") + state.commit_paged_context(self._AR_BRANCH) + return current_latents + + def encode_prompt( + self, + prompt: str, + *, + max_sequence_length: int, + dtype: torch.dtype, + ) -> torch.Tensor: + text_inputs = self.tokenizer( + [" ".join(prompt.strip().split())], + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + input_ids = text_inputs.input_ids.to(self.device) + attention_mask = text_inputs.attention_mask.to(self.device) + prompt_embeds = self.text_encoder(input_ids, attention_mask).last_hidden_state + prompt_embeds = prompt_embeds.to(device=self.device, dtype=dtype) + return prompt_embeds * attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype) + + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + inputs = self._parse_request(req) + tick = ARDiffusionTickRequest.from_extra_args(req.sampling_params.extra_args) + if tick is not None and self._ar_diffusion_kv_state is None: + raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.") + if tick is None and self._ar_diffusion_kv_state is not None: + raise ValueError("LingBot ARDiffusionEngine requests must carry ar_diffusion_tick.") + session_state: _LingBotARSessionState | None = None + if tick is not None: + if tick.prompt is not None and " ".join(tick.prompt.split()) != " ".join(inputs.prompt.split()): + raise ValueError("ar_diffusion_tick.prompt must match the standard request prompt.") + block_frames = int(self.transformer.config.num_frames_per_block) + if inputs.num_latent_frames != block_frames: + raise ValueError("A LingBot AR-Diffusion request must generate exactly one three-latent-frame block.") + if (inputs.height, inputs.width) != ( + self._ar_height, + self._ar_width, + ): + raise ValueError( + "LingBot AR-Diffusion request resolution must match the " + "fixed cache geometry " + f"{self._ar_height}x{self._ar_width}." + ) + horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 + max_realtime_ticks = horizon_latent_frames // block_frames + if tick.chunk_index >= max_realtime_ticks: + raise ValueError( + "LingBot realtime generation currently supports at most " + f"{max_realtime_ticks} ticks per generation epoch " + f"(chunk_index 0 through {max_realtime_ticks - 1}) because " + f"the image-condition horizon is {_MAX_RAW_FRAMES} pixel " + "frames; reset or create a session to start a new world." + ) + session_state = self._ar_sessions.setdefault( + tick.session_id, + _LingBotARSessionState(), + ) + if tick.chunk_index != session_state.next_chunk_index: + raise ValueError( + "LingBot chunk_index must be contiguous within a session: " + f"expected {session_state.next_chunk_index}, got " + f"{tick.chunk_index}." + ) + schedule = _build_shifted_flow_schedule( + flow_shift=inputs.flow_shift, + ) + dtype = self.transformer.dtype + # Phase 1: turn all three user inputs into DiT-ready conditions. + prompt_embeds = self.encode_prompt( + inputs.prompt, + max_sequence_length=inputs.max_sequence_length, + dtype=dtype, + ) + if tick is None: + condition = self._prepare_condition(inputs, dtype=dtype) + else: + assert session_state is not None + if session_state.image_condition is None: + horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 + session_state.image_condition = self._prepare_condition( + replace( + inputs, + num_frames=_MAX_RAW_FRAMES, + num_latent_frames=horizon_latent_frames, + ), + dtype=dtype, + ) + block_frames = int(self.transformer.config.num_frames_per_block) + condition_start = tick.chunk_index * block_frames + condition_stop = condition_start + block_frames + if condition_stop > session_state.image_condition.shape[2]: + raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") + condition = session_state.image_condition[ + :, + :, + condition_start:condition_stop, + ] + camera_pitch: float | None = None + if inputs.camera_actions is not None: + assert session_state is not None + action_trajectory, camera_pitch = integrate_lingbot_camera_actions( + inputs.camera_actions, + width=inputs.width, + height=inputs.height, + initial_pose=(session_state.camera_tail.poses[-1] if session_state.camera_tail is not None else None), + initial_pitch=session_state.camera_pitch, + ) + inputs = replace(inputs, camera_trajectory=action_trajectory) + camera, camera_tail = self._prepare_camera( + inputs, + dtype=dtype, + previous=session_state.camera_tail if session_state is not None else None, + ) + if camera.shape[2:] != condition.shape[2:]: + raise RuntimeError( + "folded camera and image condition must share latent frame/height/width geometry; " + f"got camera={tuple(camera.shape)}, condition={tuple(condition.shape)}." + ) + + block_frames = int(self.transformer.config.num_frames_per_block) + cache: LingBotTransformerCache | None = None + ar_cross_attention: list[LingBotAttentionCache] | None = None + if tick is None: + # Direct/offline oracle keeps request-local contiguous K/V. + cache = self.transformer.allocate_cache( + batch_size=1, + latent_height=condition.shape[-2], + latent_width=condition.shape[-1], + device=self.device, + dtype=dtype, + ) + else: + assert session_state is not None + prompt_changed = session_state.prompt is not None and session_state.prompt != inputs.prompt + if session_state.generator_state is not None: + inputs.generator.set_state(session_state.generator_state) + ar_cross_attention = self._ar_text_caches( + prompt_embeds, + invalidate=prompt_changed, + ) + generated_blocks: list[torch.Tensor] = [] + total_steps = (inputs.num_latent_frames // block_frames) * len(LINGBOT_DMD_TIMESTEPS) + first_start_frame = tick.chunk_index * block_frames if tick is not None else 0 + with self.progress_bar(total=total_steps) as progress_bar: + # Phase 3: generate left-to-right in latent-frame blocks. The + # monotonically increasing ``start_frame`` is also the temporal + # position and cache offset used by the Transformer. + for local_start_frame in range( + 0, + inputs.num_latent_frames, + block_frames, + ): + start_frame = first_start_frame + local_start_frame + stop_frame = local_start_frame + block_frames + generated_blocks.append( + self._generate_block( + condition=condition[:, :, local_start_frame:stop_frame], + camera=camera[:, :, local_start_frame:stop_frame], + prompt_embeds=prompt_embeds, + cache=cache, + ar_cross_attention=ar_cross_attention, + start_frame=start_frame, + schedule=schedule, + generator=inputs.generator, + progress_bar=progress_bar, + ) + ) + generated_latents = torch.cat(generated_blocks, dim=2) + # The direct/offline cache is request-local and no longer needed once + # all latent blocks have been generated. Release it before VAE decode + # so the two large allocations do not overlap. + cache = None + if tick is not None: + assert session_state is not None + session_state.prompt = inputs.prompt + session_state.generator_state = inputs.generator.get_state() + session_state.camera_tail = camera_tail + if camera_pitch is not None: + session_state.camera_pitch = camera_pitch + session_state.next_chunk_index += 1 + + # Phase 4: either expose model-space latents or invert the checkpoint's + # latent normalization and decode to pixel-space video. + if tick is not None: + if inputs.output_type != "latent": + raise ValueError( + "LingBot realtime ticks currently require output_type='latent'; " + "stateful streaming VAE decode is a separate integration step." + ) + output = { + "payload": {"latents": generated_latents}, + "metadata": {"ar_diffusion": ARDiffusionChunkMetadata.from_tick(tick).to_dict()}, + } + elif inputs.output_type == "latent": + output = generated_latents + else: + latent_mean, latent_std = self._vae_latent_stats(generated_latents) + vae_latents = (generated_latents * latent_std + latent_mean).to(dtype=self.vae.dtype) + output = self.vae.decode(vae_latents, return_dict=False)[0] + if output.shape[2] != inputs.num_frames: + raise RuntimeError( + "vae.decode returned an incompatible temporal geometry: " + f"expected {inputs.num_frames} frames, got {output.shape[2]}." + ) + return DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return cast(set[str], loader.load_weights(weights)) diff --git a/vllm_omni/diffusion/models/lingbot_world/transformer.py b/vllm_omni/diffusion/models/lingbot_world/transformer.py new file mode 100644 index 000000000..35ec18801 --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_world/transformer.py @@ -0,0 +1,1201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Checkpoint-compatible causal DiT for the LingBot World v2 model package.""" + +from __future__ import annotations + +import math +from collections.abc import Iterable +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Self, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.conv import Conv3dLayer +from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.utils import set_weight_attrs + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.layers.norm import LayerNorm +from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan +from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( + ARDiffusionPagedLayerContext, + ARDiffusionPagedLayerInputs, + ar_diffusion_paged_attention, + paged_write_attn, +) + + +@dataclass +class LingBotAttentionCache: + """K/V storage plus its logical cursor metadata. + + ``end`` is occupied storage, ``absolute_end`` is the next global token + position, and ``sink_end`` separates permanently retained prefix tokens + from the sliding local window. ``last_start`` rejects overlapping or + out-of-order causal chunks. + """ + + key: torch.Tensor + value: torch.Tensor + end: int = 0 + absolute_end: int = 0 + last_start: int | None = None + sink_end: int = 0 + + +@dataclass +class LingBotTransformerCache: + """One request's per-layer video K/V and reusable text K/V.""" + + self_attention: list[LingBotAttentionCache | ARDiffusionPagedLayerContext | ARDiffusionPagedLayerInputs] + cross_attention: list[LingBotAttentionCache | None] + + +def allocate_lingbot_cache( + *, + batch_size: int, + num_layers: int, + max_tokens: int, + num_local_heads: int, + head_dim: int, + device: torch.device, + dtype: torch.dtype, +) -> LingBotTransformerCache: + # Cross-attention starts empty because its token count is known after text encoding. + shape = (batch_size, max_tokens, num_local_heads, head_dim) + self_attention = [ + LingBotAttentionCache( + key=torch.zeros(shape, device=device, dtype=dtype), + value=torch.zeros(shape, device=device, dtype=dtype), + ) + for _ in range(num_layers) + ] + return LingBotTransformerCache( + self_attention=self_attention, + cross_attention=[None for _ in range(num_layers)], + ) + + +class _LingBotRMSNorm(nn.Module): + """RMSNorm over a tensor-parallel sharded projection.""" + + def __init__(self, hidden_size: int, eps: float) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + if param.shape == loaded_weight.shape: + param.data.copy_(loaded_weight) + return + + tp_size = get_tensor_model_parallel_world_size() + if loaded_weight.shape[0] % tp_size != 0: + raise ValueError( + f"Cannot shard RMSNorm weight of shape {tuple(loaded_weight.shape)} across tp_size={tp_size}." + ) + shard_size = loaded_weight.shape[0] // tp_size + shard_start = get_tensor_model_parallel_rank() * shard_size + shard = loaded_weight.narrow(0, shard_start, shard_size) + if param.shape != shard.shape: + raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.") + param.data.copy_(shard) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + value_float = value.float() + sum_of_squares = value_float.pow(2).sum(dim=-1, keepdim=True) + element_count = value.shape[-1] + if tp_size > 1: + sum_of_squares = tensor_model_parallel_all_reduce(sum_of_squares) + element_count *= tp_size + rms = torch.sqrt(sum_of_squares / element_count + self.eps) + return (value_float / rms * self.weight.float()).to(value.dtype) + + +def _projection_prefix(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + +class LingBotSelfAttention(nn.Module): + """Block-causal self-attention over retained history and one full chunk.""" + + def __init__( + self, + dim: int, + num_heads: int, + *, + eps: float = 1e-6, + prefix: str = "", + ) -> None: + super().__init__() + if dim % num_heads != 0: + raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.") + + tp_size = get_tensor_model_parallel_world_size() + if num_heads % tp_size != 0: + raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.") + + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_heads, + bias=True, + prefix=_projection_prefix(prefix, "qkv"), + ) + self.num_local_heads = self.qkv.num_heads + self.tp_inner_dim = self.num_local_heads * self.head_dim + self.o = RowParallelLinear( + dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + prefix=_projection_prefix(prefix, "o"), + ) + self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps) + self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps) + self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) + self.attn = Attention( + num_heads=self.num_local_heads, + head_size=self.head_dim, + num_kv_heads=self.num_local_heads, + softmax_scale=self.head_dim**-0.5, + causal=False, + role="self", + qkv_layout="BSND", + prefix=prefix, + skip_sequence_parallel=True, + ) + + def _update_cache( + self, + cache: LingBotAttentionCache, + key: torch.Tensor, + value: torch.Tensor, + current_start: int, + *, + sink_tokens: int, + update_cache: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + chunk_tokens = key.shape[1] + capacity = cache.key.shape[1] + next_sink_end = cache.sink_end + + if cache.last_start is None: + if chunk_tokens > capacity: + raise ValueError( + f"Current chunk has {chunk_tokens} tokens but the cache can hold only {capacity}; " + "the full current chunk must remain visible." + ) + next_key = key + next_value = value + next_sink_end = max(0, min(chunk_tokens, sink_tokens - current_start)) + elif current_start == cache.last_start: + if current_start + chunk_tokens != cache.absolute_end: + raise ValueError("A repeated current_start must overwrite the same-size current chunk.") + if chunk_tokens > cache.end: + raise ValueError("The current chunk is no longer fully retained in the cache.") + prefix_end = cache.end - chunk_tokens + next_key = torch.cat((cache.key[:, :prefix_end], key), dim=1) + next_value = torch.cat((cache.value[:, :prefix_end], value), dim=1) + else: + if current_start < cache.last_start: + raise ValueError(f"current_start={current_start} precedes the latest chunk start {cache.last_start}.") + if current_start < cache.absolute_end: + raise ValueError( + f"current_start={current_start} overlaps cached tokens ending at {cache.absolute_end}." + ) + if current_start > cache.absolute_end: + raise ValueError( + f"New chunks must be contiguous: current_start={current_start}, expected {cache.absolute_end}." + ) + + incoming_sink_tokens = max(0, min(chunk_tokens, sink_tokens - current_start)) + old_sink_key = cache.key[:, : cache.sink_end] + old_sink_value = cache.value[:, : cache.sink_end] + new_sink_key = key[:, :incoming_sink_tokens] + new_sink_value = value[:, :incoming_sink_tokens] + next_sink_end = cache.sink_end + incoming_sink_tokens + + old_local_key = cache.key[:, cache.sink_end : cache.end] + old_local_value = cache.value[:, cache.sink_end : cache.end] + new_local_key = key[:, incoming_sink_tokens:] + new_local_value = value[:, incoming_sink_tokens:] + local_capacity = capacity - next_sink_end + if new_local_key.shape[1] > local_capacity: + raise ValueError("The configured cache cannot retain all sink tokens and the full current chunk.") + retained_local_tokens = min( + old_local_key.shape[1], + local_capacity - new_local_key.shape[1], + ) + if retained_local_tokens: + old_local_key = old_local_key[:, -retained_local_tokens:] + old_local_value = old_local_value[:, -retained_local_tokens:] + else: + old_local_key = old_local_key[:, :0] + old_local_value = old_local_value[:, :0] + + next_key = torch.cat((old_sink_key, new_sink_key, old_local_key, new_local_key), dim=1) + next_value = torch.cat((old_sink_value, new_sink_value, old_local_value, new_local_value), dim=1) + + next_end = next_key.shape[1] + if update_cache: + cache.key[:, :next_end].copy_(next_key) + cache.value[:, :next_end].copy_(next_value) + cache.end = next_end + cache.absolute_end = current_start + chunk_tokens + cache.last_start = current_start + cache.sink_end = next_sink_end + return next_key, next_value + + def forward( + self, + hidden_states: torch.Tensor, + *, + cache: LingBotAttentionCache | ARDiffusionPagedLayerInputs, + current_start: int, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + sink_tokens: int, + update_cache: bool = True, + ) -> torch.Tensor: + # Project logical [B, S, D] tokens into TP-local + # [B, S, num_local_heads, head_dim]. + qkv, _ = self.qkv(hidden_states) + query, key, value = qkv.split((self.tp_inner_dim, self.tp_inner_dim, self.tp_inner_dim), dim=-1) + query = self.norm_q(query) + key = self.norm_k(key) + + query = query.unflatten(2, (self.num_local_heads, self.head_dim)) + key = key.unflatten(2, (self.num_local_heads, self.head_dim)) + value = value.unflatten(2, (self.num_local_heads, self.head_dim)) + if rotary_emb is not None: + cos, sin = rotary_emb + query = self.rotary_embedding(query, cos, sin) + key = self.rotary_embedding(key, cos, sin) + + if isinstance(cache, ARDiffusionPagedLayerInputs): + if query.shape[0] != 1: + raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.") + output = paged_write_attn( + cache, + query[0], + key[0], + value[0], + None, + None, + self.head_dim**-0.5, + ).unsqueeze(0) + else: + # The direct path sees history + current K/V even when the caller + # asks not to mutate persistent cache state. + visible_key, visible_value = self._update_cache( + cache, + key, + value, + current_start, + sink_tokens=sink_tokens, + update_cache=update_cache, + ) + if query.is_cuda and query.shape[0] == 1: + # Use the same block-table FlashAttention entry point as the + # realtime path so direct replay is a numerical oracle for + # paged execution, rather than a comparison between two + # different attention kernels. + block_size = key.shape[1] + key_cache = visible_key[0].unflatten(0, (-1, block_size)) + value_cache = visible_value[0].unflatten(0, (-1, block_size)) + block_count = key_cache.shape[0] + block_table = torch.arange( + block_count, + dtype=torch.int32, + device=query.device, + ).unsqueeze(0) + query_start_loc = torch.tensor( + [0, query.shape[1]], + dtype=torch.int32, + device=query.device, + ) + seq_lens = torch.tensor( + [visible_key.shape[1]], + dtype=torch.int32, + device=query.device, + ) + output = ar_diffusion_paged_attention( + query, + key_cache, + value_cache, + block_table=block_table, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + max_query_len=query.shape[1], + max_seq_len=visible_key.shape[1], + softmax_scale=self.head_dim**-0.5, + ) + else: + output = self.attn(query, visible_key, visible_value) + return self.o(output.flatten(2, 3)) + + +class LingBotCrossAttention(nn.Module): + """Cross-attention with caller-owned request-local encoder K/V reuse.""" + + def __init__( + self, + dim: int, + num_heads: int, + *, + eps: float = 1e-6, + prefix: str = "", + ) -> None: + super().__init__() + if dim % num_heads != 0: + raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.") + + tp_size = get_tensor_model_parallel_world_size() + if num_heads % tp_size != 0: + raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.") + + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.num_local_heads = num_heads // tp_size + self.tp_inner_dim = self.num_local_heads * self.head_dim + + self.q = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + prefix=_projection_prefix(prefix, "q"), + ) + self.k = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + prefix=_projection_prefix(prefix, "k"), + ) + self.v = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + prefix=_projection_prefix(prefix, "v"), + ) + self.o = RowParallelLinear( + dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + prefix=_projection_prefix(prefix, "o"), + ) + self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps) + self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps) + self.attn = Attention( + num_heads=self.num_local_heads, + head_size=self.head_dim, + num_kv_heads=self.num_local_heads, + softmax_scale=self.head_dim**-0.5, + causal=False, + role="cross", + qkv_layout="BSND", + prefix=prefix, + skip_sequence_parallel=True, + disable_kv_quant=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + *, + cache: LingBotAttentionCache | None, + ) -> tuple[torch.Tensor, LingBotAttentionCache]: + query = self.norm_q(self.q(hidden_states)) + query = query.unflatten(2, (self.num_local_heads, self.head_dim)) + + # Text K/V is constant within a request and is projected once per layer. + if cache is None: + if encoder_hidden_states is None: + raise ValueError("encoder_hidden_states are required when the cross-attention cache is empty.") + key = self.norm_k(self.k(encoder_hidden_states)) + value = self.v(encoder_hidden_states) + key = key.unflatten(2, (self.num_local_heads, self.head_dim)) + value = value.unflatten(2, (self.num_local_heads, self.head_dim)) + cache = LingBotAttentionCache( + key=key, + value=value, + end=key.shape[1], + absolute_end=key.shape[1], + last_start=0, + ) + else: + key = cache.key[:, : cache.end] + value = cache.value[:, : cache.end] + + output = self.attn(query, key, value) + return self.o(output.flatten(2, 3)), cache + + +class LingBotAttentionBlock(nn.Module): + """Checkpoint-compatible LingBot block with causal video attention.""" + + def __init__( + self, + dim: int, + num_heads: int, + *, + ffn_dim: int | None = None, + cross_attn_norm: bool = True, + eps: float = 1e-6, + prefix: str = "", + ) -> None: + super().__init__() + ffn_dim = ffn_dim or dim * 4 + self.dim = dim + self.norm1 = LayerNorm(dim, eps=eps, elementwise_affine=False) + self.self_attn = LingBotSelfAttention( + dim, + num_heads, + eps=eps, + prefix=_projection_prefix(prefix, "self_attn"), + ) + self.cross_attn = LingBotCrossAttention( + dim, + num_heads, + eps=eps, + prefix=_projection_prefix(prefix, "cross_attn"), + ) + self.norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = LayerNorm(dim, eps=eps) if cross_attn_norm else nn.Identity() + self.ffn = nn.Sequential( + ColumnParallelLinear( + dim, + ffn_dim, + bias=True, + gather_output=False, + return_bias=False, + prefix=_projection_prefix(prefix, "ffn.0"), + ), + nn.GELU(approximate="tanh"), + RowParallelLinear( + ffn_dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + prefix=_projection_prefix(prefix, "ffn.2"), + ), + ) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / math.sqrt(dim)) + self.cam_injector_layer1 = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + prefix=_projection_prefix(prefix, "cam_injector_layer1"), + ) + self.cam_injector_layer2 = RowParallelLinear( + dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + prefix=_projection_prefix(prefix, "cam_injector_layer2"), + ) + self.cam_scale_layer = nn.Linear(dim, dim) + self.cam_shift_layer = nn.Linear(dim, dim) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + timestep_projection: torch.Tensor, + camera_hidden_states: torch.Tensor, + *, + self_cache: LingBotAttentionCache | ARDiffusionPagedLayerInputs, + cross_cache: LingBotAttentionCache | None, + current_start: int, + sink_tokens: int, + update_cache: bool, + rotary_emb: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, LingBotAttentionCache]: + batch_size, token_count, dim = hidden_states.shape + num_frames = timestep_projection.shape[1] + tokens_per_frame = token_count // num_frames + # Timestep, camera, and text remain separate conditioning paths. + modulation = self.modulation.unsqueeze(1) + timestep_projection.float() + shift_msa, scale_msa, gate_msa, shift_ffn, scale_ffn, gate_ffn = modulation.chunk(6, dim=2) + + hidden_grid = hidden_states.unflatten(1, (num_frames, tokens_per_frame)) + normalized = self.norm1(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame)) + normalized = (normalized * (1 + scale_msa) + shift_msa).flatten(1, 2).to(hidden_states.dtype) + attention_output = self.self_attn( + normalized, + cache=self_cache, + current_start=current_start, + rotary_emb=rotary_emb, + sink_tokens=sink_tokens, + update_cache=update_cache, + ) + hidden_grid = hidden_grid + attention_output.unflatten(1, (num_frames, tokens_per_frame)) * gate_msa + hidden_states = hidden_grid.flatten(1, 2).to(hidden_states.dtype) + + camera_features = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(camera_hidden_states))) + camera_features = camera_features + camera_hidden_states + camera_scale = self.cam_scale_layer(camera_features) + camera_shift = self.cam_shift_layer(camera_features) + hidden_states = ((1 + camera_scale) * hidden_states + camera_shift).to(hidden_states.dtype) + + attention_output, cross_cache = self.cross_attn( + self.norm3(hidden_states), + encoder_hidden_states, + cache=cross_cache, + ) + hidden_states = hidden_states + attention_output + + hidden_grid = hidden_states.unflatten(1, (num_frames, tokens_per_frame)) + normalized = self.norm2(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame)) + normalized = (normalized * (1 + scale_ffn) + shift_ffn).flatten(1, 2).to(hidden_states.dtype) + ffn_output = self.ffn(normalized).unflatten(1, (num_frames, tokens_per_frame)) + hidden_states = (hidden_grid + ffn_output * gate_ffn).flatten(1, 2).to(hidden_states.dtype) + return hidden_states, cast(LingBotAttentionCache, cross_cache) + + +class _LingBotCameraPatchEmbedding(nn.Module): + """Linear camera patchifier matching the checkpoint's direct parameter names.""" + + def __init__( + self, + in_channels: int, + dim: int, + patch_size: tuple[int, int, int], + ) -> None: + super().__init__() + self.patch_size = patch_size + self.in_features = in_channels * math.prod(patch_size) + self.weight = nn.Parameter(torch.empty(dim, self.in_features)) + self.bias = nn.Parameter(torch.empty(dim)) + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + bound = 1 / math.sqrt(self.in_features) + nn.init.uniform_(self.bias, -bound, bound) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, channels, frames, height, width = hidden_states.shape + patch_frames, patch_height, patch_width = self.patch_size + # Preserve Conv3d patch order while retaining checkpoint weight/bias names. + hidden_states = hidden_states.reshape( + batch_size, + channels, + frames // patch_frames, + patch_frames, + height // patch_height, + patch_height, + width // patch_width, + patch_width, + ) + hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7) + hidden_states = hidden_states.reshape(batch_size, -1, self.in_features) + return F.linear(hidden_states, self.weight, self.bias) + + +class _LingBotHead(nn.Module): + def __init__( + self, + dim: int, + out_channels: int, + patch_size: tuple[int, int, int], + eps: float, + ) -> None: + super().__init__() + self.norm = LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_channels * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / math.sqrt(dim)) + + def forward( + self, + hidden_states: torch.Tensor, + timestep_embedding: torch.Tensor, + ) -> torch.Tensor: + num_frames = timestep_embedding.shape[1] + tokens_per_frame = hidden_states.shape[1] // num_frames + modulation = self.modulation.unsqueeze(1) + timestep_embedding.unsqueeze(2).float() + shift, scale = modulation.chunk(2, dim=2) + normalized = self.norm(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame)) + normalized = (normalized * (1 + scale) + shift).flatten(1, 2).to(hidden_states.dtype) + return self.head(normalized) + + +def _sinusoidal_embedding(dim: int, timestep: torch.Tensor) -> torch.Tensor: + if dim % 2: + raise ValueError(f"freq_dim must be even, got {dim}.") + half_dim = dim // 2 + timestep = timestep.to(torch.float64) + frequencies = torch.pow( + 10000, + -torch.arange(half_dim, device=timestep.device, dtype=torch.float64) / half_dim, + ) + phase = torch.outer(timestep, frequencies) + return torch.cat((phase.cos(), phase.sin()), dim=1) + + +def _rope_axis(max_seq_len: int, dim: int) -> tuple[torch.Tensor, torch.Tensor]: + if dim == 0: + empty = torch.empty(max_seq_len, 0, dtype=torch.float32) + return empty, empty.clone() + if dim % 2: + raise ValueError(f"RoPE axis dimension must be even, got {dim}.") + frequencies = 1.0 / torch.pow( + 10000, + torch.arange(0, dim, 2, dtype=torch.float64) / dim, + ) + phase = torch.outer(torch.arange(max_seq_len, dtype=torch.float64), frequencies) + return phase.cos().float(), phase.sin().float() + + +class CausalLingBotWorldTransformer3DModel(nn.Module): + """Checkpoint-compatible causal LingBot World video transformer.""" + + _repeated_blocks = ["LingBotAttentionBlock"] + packed_modules_mapping = {"qkv": ["q", "k", "v"]} + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__( + self, + patch_size: tuple[int, int, int] = (1, 2, 2), + num_attention_heads: int = 40, + attention_head_dim: int = 128, + in_channels: int = 36, + out_channels: int = 16, + text_dim: int = 4096, + freq_dim: int = 256, + ffn_dim: int = 13824, + num_layers: int = 40, + cross_attn_norm: bool = True, + eps: float = 1e-6, + image_dim: int | None = None, + added_kv_proj_dim: int | None = None, + rope_max_seq_len: int = 1024, + pos_embed_seq_len: int | None = None, + qk_norm: str = "rms_norm_across_heads", + sink_size: int = 9, + num_frames_per_block: int = 3, + sliding_window_num_frames: int = 18, + local_attn_size: int = -1, + quant_config: object | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if len(patch_size) != 3 or any(size <= 0 for size in patch_size): + raise ValueError("patch_size must contain three positive integers.") + if num_attention_heads <= 0 or attention_head_dim <= 0: + raise ValueError("Attention head counts and dimensions must be positive.") + if attention_head_dim % 2: + raise ValueError("attention_head_dim must be even for RoPE.") + if num_layers <= 0 or ffn_dim <= 0 or freq_dim <= 0: + raise ValueError("num_layers, ffn_dim, and freq_dim must be positive.") + if sink_size < 0 or num_frames_per_block <= 0 or sliding_window_num_frames <= 0: + raise ValueError("Causal frame-window sizes must be positive, except sink_size which may be zero.") + if local_attn_size != -1 and local_attn_size <= 0: + raise ValueError("local_attn_size must be -1 or a positive frame count.") + cache_window_frames = local_attn_size if local_attn_size != -1 else sliding_window_num_frames + if cache_window_frames < sink_size + num_frames_per_block: + raise ValueError("The causal cache window must fit all sink frames and one complete current block.") + if qk_norm != "rms_norm_across_heads": + raise ValueError( + f"qk_norm must be 'rms_norm_across_heads' for the LingBot World v2 checkpoint, got {qk_norm!r}." + ) + for field_name, field_value in ( + ("image_dim", image_dim), + ("added_kv_proj_dim", added_kv_proj_dim), + ("pos_embed_seq_len", pos_embed_seq_len), + ): + if field_value is not None: + raise ValueError(f"{field_name} must be None because LingBot World v2 has no image embedding path.") + if quant_config is not None: + raise RuntimeError( + "quant_config is not supported by the LingBot World transformer; construct the unquantized model." + ) + + dim = num_attention_heads * attention_head_dim + self.dim = dim + self.config = SimpleNamespace( + patch_size=patch_size, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=in_channels, + out_channels=out_channels, + text_dim=text_dim, + freq_dim=freq_dim, + ffn_dim=ffn_dim, + num_layers=num_layers, + cross_attn_norm=cross_attn_norm, + eps=eps, + image_dim=image_dim, + added_kv_proj_dim=added_kv_proj_dim, + rope_max_seq_len=rope_max_seq_len, + pos_embed_seq_len=pos_embed_seq_len, + qk_norm=qk_norm, + sink_size=sink_size, + num_frames_per_block=num_frames_per_block, + sliding_window_num_frames=sliding_window_num_frames, + local_attn_size=local_attn_size, + ) + + self.patch_embedding = Conv3dLayer( + in_channels=in_channels, + out_channels=dim, + kernel_size=patch_size, + stride=patch_size, + ) + self.patch_embedding_wancamctrl = _LingBotCameraPatchEmbedding(6 * 8 * 8, dim, patch_size) + self.c2ws_hidden_states_layer1 = ColumnParallelLinear( + dim, + dim, + bias=True, + gather_output=False, + return_bias=False, + prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer1"), + ) + self.c2ws_hidden_states_layer2 = RowParallelLinear( + dim, + dim, + bias=True, + input_is_parallel=True, + return_bias=False, + prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer2"), + ) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate="tanh"), + nn.Linear(dim, dim), + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim), + ) + self.time_projection = nn.Sequential( + nn.SiLU(), + nn.Linear(dim, dim * 6), + ) + self.blocks = nn.ModuleList( + [ + LingBotAttentionBlock( + dim, + num_attention_heads, + ffn_dim=ffn_dim, + cross_attn_norm=cross_attn_norm, + eps=eps, + prefix=_projection_prefix(prefix, f"blocks.{index}"), + ) + for index in range(num_layers) + ] + ) + self.head = _LingBotHead(dim, out_channels, patch_size, eps) + + temporal_dim = attention_head_dim - 4 * (attention_head_dim // 6) + height_dim = width_dim = 2 * (attention_head_dim // 6) + for axis, axis_dim in (("temporal", temporal_dim), ("height", height_dim), ("width", width_dim)): + cosine, sine = _rope_axis(rope_max_seq_len, axis_dim) + self.register_buffer(f"_rope_{axis}_cosine", cosine, persistent=False) + self.register_buffer(f"_rope_{axis}_sine", sine, persistent=False) + + @property + def dtype(self) -> torch.dtype: + """Return the dtype used by the transformer parameters.""" + + return next(self.parameters()).dtype + + @classmethod + def from_config( + cls, + config: dict[str, Any], + *, + quant_config: Any | None = None, + prefix: str = "", + ) -> Self: + checkpoint_contract = { + "patch_size": (1, 2, 2), + "num_attention_heads": 40, + "attention_head_dim": 128, + "in_channels": 36, + "out_channels": 16, + "text_dim": 4096, + "freq_dim": 256, + "ffn_dim": 13824, + "num_layers": 40, + "cross_attn_norm": True, + "eps": 1e-6, + "image_dim": None, + "added_kv_proj_dim": None, + "rope_max_seq_len": 1024, + "pos_embed_seq_len": None, + "qk_norm": "rms_norm_across_heads", + "sink_size": 9, + "num_frames_per_block": 3, + "sliding_window_num_frames": 18, + "local_attn_size": -1, + } + class_name = config.get("_class_name") + if class_name is not None and class_name != "CausalLingBotWorldTransformer3DModel": + raise ValueError( + "_class_name must be 'CausalLingBotWorldTransformer3DModel' for the LingBot World v2 checkpoint." + ) + kwargs = {name: value for name, value in config.items() if name in checkpoint_contract} + if "patch_size" in kwargs: + kwargs["patch_size"] = tuple(kwargs["patch_size"]) + for name, expected in checkpoint_contract.items(): + actual = kwargs.get(name, expected) + if actual != expected: + raise ValueError(f"LingBot checkpoint config {name} must be {expected!r}, got {actual!r}.") + return cls(**kwargs, quant_config=quant_config, prefix=prefix) + + def allocate_cache( + self, + *, + batch_size: int, + latent_height: int, + latent_width: int, + device: torch.device, + dtype: torch.dtype, + ) -> LingBotTransformerCache: + """Allocate one caller-owned cache using this transformer's geometry.""" + + if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError("batch_size must be a positive integer for LingBot cache allocation.") + if ( + isinstance(latent_height, bool) + or not isinstance(latent_height, int) + or latent_height <= 0 + or isinstance(latent_width, bool) + or not isinstance(latent_width, int) + or latent_width <= 0 + ): + raise ValueError("latent_height and latent_width must be positive integers.") + + # Cache capacity is measured in post-patch tokens, not raw/latent + # pixels: window_frames * (latent_height/patch_h) * (latent_width/patch_w). + patch_frames, patch_height, patch_width = self.config.patch_size + if patch_frames != 1 or latent_height % patch_height or latent_width % patch_width: + raise ValueError("latent height/width must align with the configured LingBot patch size.") + post_patch_height = latent_height // patch_height + post_patch_width = latent_width // patch_width + window_frames = ( + self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames + ) + max_tokens = int(window_frames * post_patch_height * post_patch_width) + tp_size = get_tensor_model_parallel_world_size() + num_local_heads = self.config.num_attention_heads // tp_size + return allocate_lingbot_cache( + batch_size=batch_size, + num_layers=self.config.num_layers, + max_tokens=max_tokens, + num_local_heads=num_local_heads, + head_dim=self.config.attention_head_dim, + device=device, + dtype=dtype, + ) + + def _rotary_embedding( + self, + *, + frames: int, + height: int, + width: int, + start_frame: int, + dtype: torch.dtype, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + if start_frame + frames > self.config.rope_max_seq_len: + raise ValueError("Temporal RoPE positions exceed rope_max_seq_len.") + if height > self.config.rope_max_seq_len or width > self.config.rope_max_seq_len: + raise ValueError("Spatial RoPE positions exceed rope_max_seq_len.") + + def expand_axis(table: torch.Tensor, axis: str) -> torch.Tensor: + if axis == "temporal": + return ( + table[start_frame : start_frame + frames].view(frames, 1, 1, -1).expand(frames, height, width, -1) + ) + if axis == "height": + return table[:height].view(1, height, 1, -1).expand(frames, height, width, -1) + return table[:width].view(1, 1, width, -1).expand(frames, height, width, -1) + + cosine = torch.cat( + ( + expand_axis(self._rope_temporal_cosine, "temporal"), + expand_axis(self._rope_height_cosine, "height"), + expand_axis(self._rope_width_cosine, "width"), + ), + dim=-1, + ) + sine = torch.cat( + ( + expand_axis(self._rope_temporal_sine, "temporal"), + expand_axis(self._rope_height_sine, "height"), + expand_axis(self._rope_width_sine, "width"), + ), + dim=-1, + ) + return ( + cosine.reshape(frames * height * width, -1).to(device=device, dtype=dtype), + sine.reshape(frames * height * width, -1).to(device=device, dtype=dtype), + ) + + def _validate_inputs( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + camera_hidden_states: torch.Tensor, + cache: LingBotTransformerCache, + start_frame: int, + ) -> tuple[int, int, int, int, int]: + if hidden_states.ndim != 5: + raise ValueError("hidden_states must be rank 5 [batch, channels, frames, height, width].") + if hidden_states.shape[1] != self.config.in_channels: + raise ValueError(f"hidden_states must have {self.config.in_channels} channels.") + if not isinstance(start_frame, int) or isinstance(start_frame, bool) or start_frame < 0: + raise ValueError("start_frame must be a non-negative integer latent-frame offset.") + patch_frames, patch_height, patch_width = self.config.patch_size + batch_size, _, frames, height, width = hidden_states.shape + if frames % patch_frames or height % patch_height or width % patch_width: + raise ValueError("hidden_states dimensions must be divisible by patch_size.") + post_patch_frames = frames // patch_frames + if post_patch_frames != self.config.num_frames_per_block: + raise ValueError( + f"Each LingBot forward must contain exactly {self.config.num_frames_per_block} post-patch frames, " + f"got {post_patch_frames}." + ) + if start_frame % patch_frames: + raise ValueError("start_frame must align to the temporal patch size.") + if encoder_hidden_states.ndim != 3: + raise ValueError("encoder_hidden_states must be rank 3 [batch, tokens, text_dim].") + if encoder_hidden_states.shape[0] != batch_size or encoder_hidden_states.shape[2] != self.config.text_dim: + raise ValueError("encoder_hidden_states batch/width do not match the LingBot config.") + if encoder_hidden_states.shape[1] == 0: + raise ValueError("encoder_hidden_states must contain at least one token.") + expected_camera_shape = (batch_size, 6 * 8 * 8, frames, height, width) + if camera_hidden_states.shape != expected_camera_shape: + raise ValueError( + "camera_hidden_states must use folded [batch, 6*8*8, frames, height, width] layout " + f"matching the video; expected {expected_camera_shape}, got {tuple(camera_hidden_states.shape)}." + ) + if hidden_states.device != encoder_hidden_states.device or hidden_states.device != camera_hidden_states.device: + raise ValueError("Video, text, and camera tensors must be on the same device.") + if hidden_states.dtype != encoder_hidden_states.dtype or hidden_states.dtype != camera_hidden_states.dtype: + raise ValueError("Video, text, and camera tensors must use the same dtype.") + if timestep.device != hidden_states.device: + raise ValueError("timestep must be on the same device as hidden_states.") + if len(cache.self_attention) != self.config.num_layers or len(cache.cross_attention) != self.config.num_layers: + raise ValueError("LingBotTransformerCache layer counts must match num_layers.") + return batch_size, frames, height, width, patch_frames + + def _timestep_embeddings( + self, + timestep: torch.Tensor, + *, + batch_size: int, + frames: int, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + if timestep.ndim == 0: + if batch_size != 1: + raise ValueError("A scalar timestep is only valid for batch_size=1.") + timestep = timestep.reshape(1) + if timestep.ndim == 1: + if timestep.shape[0] != batch_size: + raise ValueError("A rank-1 timestep must contain one value per batch item.") + timestep = timestep.unsqueeze(1).expand(batch_size, frames) + elif timestep.ndim == 2: + if timestep.shape != (batch_size, frames): + raise ValueError(f"A rank-2 timestep must have shape {(batch_size, frames)}.") + else: + raise ValueError("timestep must be scalar, rank 1 [batch], or rank 2 [batch, frames].") + + frequency_embedding = _sinusoidal_embedding(self.config.freq_dim, timestep.reshape(-1)).to(dtype=dtype) + timestep_embedding = self.time_embedding(frequency_embedding).unflatten(0, (batch_size, frames)) + timestep_projection = self.time_projection(timestep_embedding).unflatten(2, (6, self.dim)) + return timestep_embedding, timestep_projection + + def _unpatchify( + self, + hidden_states: torch.Tensor, + *, + batch_size: int, + frames: int, + height: int, + width: int, + ) -> torch.Tensor: + patch_frames, patch_height, patch_width = self.config.patch_size + hidden_states = hidden_states.reshape( + batch_size, + frames, + height, + width, + patch_frames, + patch_height, + patch_width, + self.config.out_channels, + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + camera_hidden_states: torch.Tensor, + *, + cache: LingBotTransformerCache, + start_frame: int, + update_cache: bool, + ) -> torch.Tensor: + if not isinstance(update_cache, bool): + raise ValueError("update_cache must be a boolean.") + batch_size, frames, height, width, patch_frames = self._validate_inputs( + hidden_states, + timestep, + encoder_hidden_states, + camera_hidden_states, + cache, + start_frame, + ) + patch_height, patch_width = self.config.patch_size[1:] + patched_frames = frames // patch_frames + patched_height = height // patch_height + patched_width = width // patch_width + tokens_per_frame = patched_height * patched_width + patched_start_frame = start_frame // patch_frames + current_start = patched_start_frame * tokens_per_frame + sink_tokens = self.config.sink_size * tokens_per_frame + # Phase 1: create absolute 3D positions for this causal block. The + # spatial coordinates restart for every frame; temporal coordinates + # begin at ``start_frame`` so cached blocks retain their true ordering. + rotary_emb = self._rotary_embedding( + frames=patched_frames, + height=patched_height, + width=patched_width, + start_frame=patched_start_frame, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + # Phase 2: independently patchify video and camera grids to identical + # token layouts, then project timestep and text conditions. + hidden_states = self.patch_embedding(hidden_states).flatten(2).transpose(1, 2) + camera_hidden_states = self.patch_embedding_wancamctrl(camera_hidden_states) + camera_hidden_states = camera_hidden_states + self.c2ws_hidden_states_layer2( + F.silu(self.c2ws_hidden_states_layer1(camera_hidden_states)) + ) + if hidden_states.shape != camera_hidden_states.shape: + raise ValueError("Patched video and camera token shapes must match.") + + timestep_embedding, timestep_projection = self._timestep_embeddings( + timestep, + batch_size=batch_size, + frames=patched_frames, + dtype=hidden_states.dtype, + ) + projected_text = ( + self.text_embedding(encoder_hidden_states) + if any(value is None for value in cache.cross_attention) + else None + ) + if cache.self_attention and isinstance(cache.self_attention[0], ARDiffusionPagedLayerContext): + if batch_size != 1: + raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.") + forward_context = cache.self_attention[0].forward_ctx + expected_seq_len = patched_frames * tokens_per_frame + if forward_context.seq_len != expected_seq_len: + raise RuntimeError( + "LingBot AR-Diffusion paged context token count does not " + f"match this block: {forward_context.seq_len} != " + f"{expected_seq_len}." + ) + forward_context.prepare( + device=hidden_states.device, + action_len=0, + query_len=hidden_states.shape[1], + ) + cache.self_attention = [layer_context.to_layer_inputs() for layer_context in cache.self_attention] + # Phase 3: each layer receives its own cache entry. Text K/V is passed + # only when absent; the returned cache is stored for subsequent DMD + # steps and causal blocks in this request. + for index, block in enumerate(self.blocks): + hidden_states, cross_cache = block( + hidden_states, + projected_text if cache.cross_attention[index] is None else None, + timestep_projection, + camera_hidden_states, + self_cache=cache.self_attention[index], + cross_cache=cache.cross_attention[index], + current_start=current_start, + sink_tokens=sink_tokens, + update_cache=update_cache, + rotary_emb=rotary_emb, + ) + cache.cross_attention[index] = cross_cache + + # Phase 4: map tokens to per-patch 16-channel flow values and restore + # [B, C, F, H, W] for the Pipeline's sampler update. + hidden_states = self.head(hidden_states, timestep_embedding) + return self._unpatchify( + hidden_states, + batch_size=batch_size, + frames=patched_frames, + height=patched_height, + width=patched_width, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load exact checkpoint names, delegating TP sharding to parameters.""" + + params = dict(self.named_parameters()) + loaded: set[str] = set() + qkv_mapping = (("q", "q"), ("k", "k"), ("v", "v")) + for checkpoint_name, loaded_weight in weights: + name = checkpoint_name + shard_id = None + for projection_name, projection_shard in qkv_mapping: + marker = f".self_attn.{projection_name}." + if marker in name: + name = name.replace(marker, ".self_attn.qkv.") + shard_id = projection_shard + break + if name not in params: + raise KeyError(f"Unexpected LingBot model weight name: {checkpoint_name}") + param = params[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + weight_loader(param, loaded_weight) + else: + if weight_loader is default_weight_loader: + raise RuntimeError(f"Fused LingBot QKV parameter {name} has no stacked weight loader.") + weight_loader(param, loaded_weight, shard_id) + # DiffusersPipelineLoader compares this return value with the + # model parameter namespace. Packed Q/K/V checkpoint tensors + # therefore need to report both their source names and the fused + # model parameter they populated, matching Wan's loader contract. + loaded.add(checkpoint_name) + loaded.add(name) + return loaded diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py index b95f59228..112426002 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py @@ -233,7 +233,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder # race; prefetch the whole component set before any from_pretrained. - longcat_subfolders = ["scheduler", "text_encoder", "tokenizer", "vae"] + longcat_subfolders = ["scheduler", "text_encoder", "text_processor", "tokenizer", "vae"] prefetch_subfolders(model, longcat_subfolders, local_files_only=local_files_only) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( @@ -248,7 +248,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile local_files_only=local_files_only, ) self.text_processor = Qwen2VLProcessor.from_pretrained( - model, subfolder="tokenizer", local_files_only=local_files_only + model, subfolder="text_processor", local_files_only=local_files_only ) self.vae = from_pretrained_with_prefetch( AutoencoderKL.from_pretrained, @@ -627,6 +627,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile negative_kwargs = { "hidden_states": latents, "timestep": timestep / 1000, + "guidance": guidance, "encoder_hidden_states": negative_prompt_embeds, "txt_ids": negative_text_ids, "img_ids": latent_image_ids, diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py index 4821425b0..ac08b8def 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py @@ -37,6 +37,7 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import ( from vllm_omni.diffusion.models.longcat_image.pipeline_longcat_image import calculate_shift from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.size_utils import normalize_min_aligned_size from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( @@ -93,17 +94,16 @@ def get_longcat_image_edit_pre_process_func( height = request.sampling_params.height or calculated_height width = request.sampling_params.width or calculated_width - # Store calculated dimensions in request - prompt["additional_information"]["calculated_height"] = calculated_height - prompt["additional_information"]["calculated_width"] = calculated_width + # Ensure dimensions are multiples of vae_scale_factor * 2 + height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) request.sampling_params.height = height request.sampling_params.width = width # Preprocess image if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels): - image = image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image_processor.resize(image, calculated_height // 2, calculated_width // 2) - image = image_processor.preprocess(image, calculated_height, calculated_width) + image = image_processor.resize(image, height, width) + prompt_image = image_processor.resize(image, height // 2, width // 2) + image = image_processor.preprocess(image, height, width) # Store preprocessed image and prompt image in request prompt["additional_information"]["preprocessed_image"] = image @@ -575,15 +575,13 @@ class LongCatImageEditPipeline( ): prompt_image = additional_information.get("prompt_image") image = additional_information.get("preprocessed_image") - calculated_height = additional_information.get("calculated_height", height) - calculated_width = additional_information.get("calculated_width", width) else: raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.") self.check_inputs( prompt, - calculated_height, - calculated_width, + height, + width, negative_prompt=negative_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, @@ -610,8 +608,8 @@ class LongCatImageEditPipeline( image, batch_size * num_images_per_prompt, num_channels_latents, - calculated_height, - calculated_width, + height, + width, prompt_embeds.dtype, prompt_embeds.shape[1], device, @@ -667,6 +665,7 @@ class LongCatImageEditPipeline( negative_kwargs = { "hidden_states": latent_model_input, "timestep": timestep / 1000, + "guidance": guidance, "encoder_hidden_states": negative_prompt_embeds, "txt_ids": negative_text_ids, "img_ids": latent_image_ids, @@ -691,7 +690,7 @@ class LongCatImageEditPipeline( if output_type == "latent": image = latents else: - latents = self._unpack_latents(latents, calculated_height, calculated_width, self.vae_scale_factor) + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor if latents.dtype != self.vae.dtype: diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index 300295bd8..a493f1c0b 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -8,17 +8,25 @@ from vllm_omni.diffusion.models.ltx2.ltx2_components import ( ) from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2DistilledOneStagePipeline, LTX2I2VDMD2Pipeline, LTX2Pipeline, LTX2T2VDMD2Pipeline, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( + LTX2DistilledPipeline, + LTX2DistilledTwoStagePipeline, + LTX2TwoStagePipeline, +) __all__ = [ "LTX2Pipeline", + "LTX2DistilledOneStagePipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", + "LTX2TwoStagePipeline", "LTX2DistilledPipeline", + "LTX2DistilledTwoStagePipeline", "get_ltx2_post_process_func", "load_transformer_config", "create_transformer_from_config", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py b/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py new file mode 100644 index 000000000..d8228370e --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Parse official LTX distilled adapters without materializing their tensors. + +The official LTX adapters are raw safetensors files rather than PEFT +checkpoints. This module is deliberately LTX-specific: it translates the +official names into the locally constructed Diffusers transformer names and +describes the resulting adapter targets. Applying those targets is the job +of :mod:`ltx2_phase_adapter`. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.nn as nn +from safetensors import safe_open + +_IGNORED_OFFICIAL_LORA_MODULES = {"text_embedding_projection.aggregate_embed"} +_OFFICIAL_TO_DIFFUSERS_PREFIXES = ( + ("audio_prompt_adaln_single.", "audio_prompt_adaln."), + ("prompt_adaln_single.", "prompt_adaln."), + ("audio_adaln_single.", "audio_time_embed."), + ("adaln_single.", "time_embed."), + ("audio_patchify_proj", "audio_proj_in"), + ("patchify_proj", "proj_in"), + ("av_ca_video_scale_shift_adaln_single.", "av_cross_attn_video_scale_shift."), + ("av_ca_audio_scale_shift_adaln_single.", "av_cross_attn_audio_scale_shift."), + ("av_ca_a2v_gate_adaln_single.", "av_cross_attn_video_a2v_gate."), + ("av_ca_v2a_gate_adaln_single.", "av_cross_attn_audio_v2a_gate."), +) + + +@dataclass(frozen=True) +class AdapterTarget: + """One A/B pair mapped onto a local transformer linear layer.""" + + # Keep the original Diffusers module name for diagnostics after LTX packs + # Q/K/V into the runtime module path. + source_module: str + module: str + slice: str | None + a_key: str + b_key: str + a_shape: tuple[int, ...] + b_shape: tuple[int, ...] + + +@dataclass(frozen=True) +class AdapterManifest: + """Tensor-free description of the fixed refinement adapter.""" + + path: str + targets: tuple[AdapterTarget, ...] + + +def _to_diffusers_module_name(name: str) -> str: + for official_prefix, diffusers_prefix in _OFFICIAL_TO_DIFFUSERS_PREFIXES: + if name.startswith(official_prefix): + return diffusers_prefix + name[len(official_prefix) :] + return name + + +def _resolve_lora_target(transformer: nn.Module, module_name: str) -> tuple[str, str | int | None] | None: + modules = dict(transformer.named_modules()) + if module_name in modules: + return module_name, None + for packed_suffix, sub_suffix, shard_id in getattr(transformer, "stacked_params_mapping", ()): + if sub_suffix in module_name: + packed_name = module_name.replace(sub_suffix, packed_suffix) + if packed_name in modules: + return packed_name, shard_id + return None + + +def parse_ltx_adapter(transformer: nn.Module, path: str) -> AdapterManifest: + """Translate one official raw LTX adapter into local Transformer targets.""" + path = str(Path(path)) + targets: list[AdapterTarget] = [] + unresolved: list[str] = [] + seen: set[tuple[str, str | int | None]] = set() + + with safe_open(path, framework="pt", device="cpu") as handle: + tensor_names = set(handle.keys()) + official_name_set = { + tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0] + for tensor_name in tensor_names + if tensor_name.endswith(".lora_A.weight") + } + official_b_names = { + tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0] + for tensor_name in tensor_names + if tensor_name.endswith(".lora_B.weight") + } + missing_a = sorted(official_b_names - official_name_set - _IGNORED_OFFICIAL_LORA_MODULES) + if missing_a: + raise ValueError(f"Missing paired LoRA A tensors for {missing_a} in {path}.") + + official_names = sorted(official_name_set) + for official_name in official_names: + if official_name in _IGNORED_OFFICIAL_LORA_MODULES: + continue + key_prefix = f"diffusion_model.{official_name}" + key_a = f"{key_prefix}.lora_A.weight" + key_b = f"{key_prefix}.lora_B.weight" + if key_b not in tensor_names: + raise ValueError(f"Missing paired LoRA tensor {key_b!r} in {path}.") + + a_shape = tuple(handle.get_slice(key_a).get_shape()) + b_shape = tuple(handle.get_slice(key_b).get_shape()) + if len(a_shape) != 2 or len(b_shape) != 2 or b_shape[1] != a_shape[0]: + raise ValueError(f"Invalid LoRA pair for {official_name}: A={a_shape}, B={b_shape}.") + + source_module = _to_diffusers_module_name(official_name) + resolved = _resolve_lora_target(transformer, source_module) + if resolved is None: + unresolved.append(official_name) + continue + module, shard_id = resolved + target_id = (module, shard_id) + if target_id in seen: + raise ValueError(f"Official LTX LoRA contains duplicate target {module!r}, slice {shard_id!r}.") + seen.add(target_id) + targets.append( + AdapterTarget( + source_module=source_module, + module=module, + slice=shard_id, + a_key=key_a, + b_key=key_b, + a_shape=a_shape, + b_shape=b_shape, + ) + ) + + if unresolved: + raise ValueError(f"Official LTX LoRA contains unmapped modules: {unresolved}.") + if not targets: + raise ValueError(f"No Transformer LoRA tensors were found in {path}.") + return AdapterManifest(path=path, targets=tuple(targets)) + + +def iter_adapter_tensors(manifest: AdapterManifest) -> Iterator[tuple[AdapterTarget, torch.Tensor, torch.Tensor]]: + """Stream adapter pairs in manifest order without a full checkpoint copy.""" + + with safe_open(manifest.path, framework="pt", device="cpu") as handle: + for target in manifest.targets: + yield ( + target, + handle.get_tensor(target.a_key), + handle.get_tensor(target.b_key), + ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 54c866eb7..6325df870 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared component construction helpers for the LTX model family.""" @@ -9,17 +9,21 @@ import inspect import json import logging import os -from dataclasses import dataclass +from dataclasses import dataclass, replace +from pathlib import Path from typing import TYPE_CHECKING, Any import torch from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler +from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from diffusers.video_processor import VideoProcessor from huggingface_hub import hf_hub_download -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration +from safetensors import safe_open +from safetensors.torch import load_file +from transformers import AutoModelForImageTextToText, AutoTokenizer, Gemma3ForConditionalGeneration from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as OmniAttention @@ -32,6 +36,13 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from .ltx2_diffusion_decoder import ( + LTX25_NATIVE_ARTIFACT_REVISION, + LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, +) +from .ltx2_diffusion_decoder_distributed import DistributedLTX2VideoDiffusionDecoderModel +from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint from .ltx2_transformer import ( LTX2VideoTransformer3DModel, apply_interleaved_rotary_emb, @@ -44,6 +55,14 @@ try: except ImportError: LTX2VocoderWithBWE = None +try: + from transformers import Gemma4UnifiedForConditionalGeneration as _Gemma4UnifiedForConditionalGeneration +except ImportError: + _Gemma4UnifiedForConditionalGeneration = None + + +_LTX25_TEXT_ENCODER_CLS = AutoModelForImageTextToText if _Gemma4UnifiedForConditionalGeneration is not None else None + _LTX_COMPONENT_SUBFOLDERS = ( "tokenizer", @@ -57,6 +76,18 @@ _LTX_COMPONENT_SUBFOLDERS = ( ) logger = logging.getLogger(__name__) +_LTX2_CONV_VAE_EXTRA = "ltx2_use_conv_vae" +_LTX2_DIFFUSION_DECODER_SUBFOLDER = "diffusion_decoder" + + +def _ltx2_use_diffusion_decoder(od_config: Any, model_version: str) -> bool: + """Select DiffVAE by default for LTX-2.5, with an explicit ConvVAE opt-in.""" + extras = getattr(od_config, "extras", {}) or {} + use_conv_vae = extras.get(_LTX2_CONV_VAE_EXTRA, False) + if not isinstance(use_conv_vae, bool): + raise TypeError(f"{_LTX2_CONV_VAE_EXTRA} must be a bool, got {type(use_conv_vae)!r}") + return model_version == "2.5" and not use_conv_vae + @dataclass(frozen=True) class LTXComponentProfile: @@ -69,7 +100,16 @@ class LTXComponentProfile: resident_modules: tuple[str, ...] = () video_vae_cls: type = AutoencoderKLLTX2Video vocoder_cls: type = LTX2Vocoder + text_encoder_cls: type | None = Gemma3ForConditionalGeneration vocoder_fallback_cls: type | None = None + artifact_repo_id: str | None = None + artifact_revision: str | None = None + latent_upsampler_filename: str | None = None + distilled_lora_filename: str | None = None + transformer_subfolder: str = "transformer" + scheduler_use_dynamic_shifting: bool = False + scheduler_shift_terminal: float | None = None + preserve_connector_attention_mask: bool = False LTX2_COMPONENT_PROFILE = LTXComponentProfile( @@ -92,6 +132,23 @@ LTX23_COMPONENT_PROFILE = LTXComponentProfile( vocoder_fallback_cls=LTX2Vocoder, ) +LTX25_FULL_COMPONENT_PROFILE = LTXComponentProfile( + name="ltx2_5_full", + dit_modules=("transformer",), + encoder_modules=("text_encoder", "connectors"), + vae_modules=("vae", "audio_vae"), + resident_modules=("vocoder",), + video_vae_cls=DistributedAutoencoderKLLTX2Video, + vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder, + vocoder_fallback_cls=LTX2Vocoder, + text_encoder_cls=_LTX25_TEXT_ENCODER_CLS, + transformer_subfolder="transformer_full", + scheduler_use_dynamic_shifting=True, + scheduler_shift_terminal=0.1, + preserve_connector_attention_mask=True, +) + + LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( name="ltx2_distilled", dit_modules=("transformer",), @@ -101,16 +158,183 @@ LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( video_vae_cls=DistributedAutoencoderKLLTX2Video, ) +LTX25_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile( + name="ltx2_5_distilled", + dit_modules=("transformer",), + encoder_modules=("text_encoder", "connectors"), + vae_modules=("vae", "audio_vae"), + resident_modules=("vocoder", "latent_upsampler"), + video_vae_cls=DistributedAutoencoderKLLTX2Video, + vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder, + vocoder_fallback_cls=LTX2Vocoder, + text_encoder_cls=_LTX25_TEXT_ENCODER_CLS, + preserve_connector_attention_mask=True, +) + +LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace( + LTX2_DISTILLED_COMPONENT_PROFILE, + name="ltx2_distilled_one_stage", + resident_modules=LTX2_COMPONENT_PROFILE.resident_modules, +) + +LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace( + LTX25_DISTILLED_COMPONENT_PROFILE, + name="ltx2_5_distilled_one_stage", + resident_modules=LTX25_FULL_COMPONENT_PROFILE.resident_modules, +) + + +LTX23_DISTILLED_COMPONENT_PROFILE = replace( + LTX23_COMPONENT_PROFILE, + name="ltx2_3_distilled", + resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2.3", + latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", +) + +LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace( + LTX23_DISTILLED_COMPONENT_PROFILE, + name="ltx2_3_distilled_one_stage", + resident_modules=LTX23_COMPONENT_PROFILE.resident_modules, +) + +LTX2_TWO_STAGE_COMPONENT_PROFILE = replace( + LTX2_COMPONENT_PROFILE, + name="ltx2_two_stage", + resident_modules=(*LTX2_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2", + latent_upsampler_filename="ltx-2-spatial-upscaler-x2-1.0.safetensors", + distilled_lora_filename="ltx-2-19b-distilled-lora-384.safetensors", +) + +LTX23_TWO_STAGE_COMPONENT_PROFILE = replace( + LTX23_COMPONENT_PROFILE, + name="ltx2_3_two_stage", + resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2.3", + latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors", + distilled_lora_filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors", +) + + +LTX25_TWO_STAGE_COMPONENT_PROFILE = replace( + LTX25_FULL_COMPONENT_PROFILE, + name="ltx2_5_two_stage", + resident_modules=(*LTX25_FULL_COMPONENT_PROFILE.resident_modules, "latent_upsampler"), + artifact_repo_id="Lightricks/LTX-2.5", + artifact_revision=LTX25_NATIVE_ARTIFACT_REVISION, + latent_upsampler_filename=("latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors"), + distilled_lora_filename="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors", +) _COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = { ("one_stage", "2"): LTX2_COMPONENT_PROFILE, ("one_stage", "2.3"): LTX23_COMPONENT_PROFILE, + ("one_stage", "2.5"): LTX25_FULL_COMPONENT_PROFILE, + ("two_stage", "2"): LTX2_TWO_STAGE_COMPONENT_PROFILE, + ("two_stage", "2.3"): LTX23_TWO_STAGE_COMPONENT_PROFILE, + ("two_stage", "2.5"): LTX25_TWO_STAGE_COMPONENT_PROFILE, + ("distilled_one_stage", "2"): LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, + ("distilled_one_stage", "2.3"): LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, + ("distilled_one_stage", "2.5"): LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE, ("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE, + ("distilled_two_stage", "2.3"): LTX23_DISTILLED_COMPONENT_PROFILE, + ("distilled_two_stage", "2.5"): LTX25_DISTILLED_COMPONENT_PROFILE, ("dmd2", "2"): LTX2_COMPONENT_PROFILE, ("dmd2", "2.3"): LTX23_COMPONENT_PROFILE, } +def resolve_ltx_checkpoint_kind(pipeline_kind: str) -> LTXCheckpointKind | None: + """Derive checkpoint requirements from the execution contract.""" + if pipeline_kind in {"one_stage", "two_stage"}: + return "regular" + if pipeline_kind in {"distilled_one_stage", "distilled_two_stage"}: + return "distilled" + if pipeline_kind == "dmd2": + return None + raise ValueError(f"Unsupported LTX pipeline kind: {pipeline_kind!r}.") + + +def resolve_ltx_artifact( + model: str, + repo_id: str, + filename: str, + *, + model_revision: str | None, + artifact_revision: str | None, +) -> str: + """Resolve an official LTX sidecar without crossing repository revisions. + + A local model path selects where the primary components are loaded from; + it does not imply that independently hosted sidecars must be offline. Hub + offline behavior remains controlled by huggingface_hub itself. + """ + candidate = Path(model) / filename + if candidate.is_file(): + return str(candidate) + + # Hub revisions are repository-scoped. Reuse the model revision only when + # the model and artifact are in the same repository; otherwise use the + # independently pinned artifact revision. + revision = model_revision if model == repo_id else artifact_revision + try: + return hf_hub_download( + repo_id=repo_id, + filename=filename, + revision=revision, + ) + except Exception as exc: + raise FileNotFoundError( + f"Unable to resolve LTX artifact {filename!r}. Searched {candidate}; " + f"place the file in the model root or make {repo_id} available." + ) from exc + + +def _create_ltx25_natten_processor() -> LTX2VideoVaeNeighborhoodNattenProcessor: + """Create the required production attention backend with an actionable error.""" + try: + return LTX2VideoVaeNeighborhoodNattenProcessor() + except (ImportError, OSError, RuntimeError, ValueError) as exc: + raise RuntimeError( + "LTX-2.5 DiffVAE requires the shi-labs/natten Hub kernel. " + "Install kernels==0.14.1, use a supported GPU, leave " + "DIFFUSERS_DISABLE_REMOTE_CODE unset, and allow Hub access during kernel initialization." + ) from exc + + +def _load_ltx_latent_upsampler_single_file(path: str, dtype: torch.dtype) -> LTX2LatentUpsamplerModel: + """Load an official single-file upsampler into the Diffusers module.""" + with safe_open(path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + raw_config = json.loads(metadata.get("config", "{}")) + config = { + "in_channels": raw_config.get("in_channels", 128), + "mid_channels": raw_config.get("mid_channels", 1024), + "num_blocks_per_stage": raw_config.get("num_blocks_per_stage", 4), + "dims": raw_config.get("dims", 3), + "spatial_upsample": raw_config.get("spatial_upsample", True), + "temporal_upsample": raw_config.get("temporal_upsample", False), + "rational_spatial_scale": raw_config.get("spatial_scale", 2.0), + "use_rational_resampler": raw_config.get("rational_resampler", False), + } + with torch.device("cpu"): + upsampler = LTX2LatentUpsamplerModel(**config).to(dtype=dtype) + + state_dict = load_file(path, device="cpu") + if "upsampler.0.weight" in state_dict and hasattr(upsampler.upsampler, "conv"): + state_dict["upsampler.conv.weight"] = state_dict.pop("upsampler.0.weight") + state_dict["upsampler.conv.bias"] = state_dict.pop("upsampler.0.bias") + missing, unexpected = upsampler.load_state_dict(state_dict, strict=False) + unresolved_missing = set(missing) - {"upsampler.blur_down.kernel"} + if unresolved_missing or unexpected: + raise ValueError( + f"Invalid LTX latent upsampler {path}: missing={sorted(unresolved_missing)}, " + f"unexpected={sorted(unexpected)}." + ) + return upsampler + + def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile: """Resolve component construction independently from execution recipes.""" try: @@ -119,7 +343,7 @@ def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTX raise ValueError(f"Unsupported LTX component kind/version: {pipeline_kind!r}/{model_version!r}.") from exc -def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]: +def _load_ltx_metadata_json(model: str, filename: str, revision: str | None = None) -> dict[str, Any]: """Load small checkpoint metadata without relying on repository names.""" if os.path.isdir(model): path = os.path.join(model, filename) @@ -127,7 +351,7 @@ def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]: return {} else: try: - path = hf_hub_download(repo_id=model, filename=filename) + path = hf_hub_download(repo_id=model, filename=filename, revision=revision) except Exception: return {} try: @@ -138,17 +362,34 @@ def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]: return value if isinstance(value, dict) else {} -def detect_ltx_model_version(model: str) -> str: - """Detect LTX-2 versus LTX-2.3 from checkpoint component metadata. +def detect_ltx_model_version(model: str, revision: str | None = None) -> str: + """Detect the LTX model version from checkpoint component metadata.""" + model_index = _load_ltx_metadata_json(model, "model_index.json", revision) + model_version = str(model_index.get("model_version", "")) + if model_version.startswith("2.5"): + return "2.5" - Official checkpoints use ``model_version`` metadata. Diffusers repositories - expose the equivalent distinction through the BWE vocoder introduced by - LTX-2.3. Unknown conversions retain the official LTX-2 fallback. - """ - model_index = _load_ltx_metadata_json(model, "model_index.json") - if str(model_index.get("model_version", "")).startswith("2.3"): + text_encoder_entry = model_index.get("text_encoder") + if isinstance(text_encoder_entry, (list, tuple)) and text_encoder_entry: + text_encoder_class = str(text_encoder_entry[-1]) + elif isinstance(text_encoder_entry, dict): + text_encoder_class = str(text_encoder_entry.get("_class_name", "")) + else: + text_encoder_class = "" + if text_encoder_class == "Gemma4UnifiedForConditionalGeneration": + return "2.5" + + text_encoder_config = _load_ltx_metadata_json(model, "text_encoder/config.json", revision) + if text_encoder_config.get("model_type") in ("gemma4_unified", "gemma4"): + return "2.5" + + if model_version.startswith("2.3"): return "2.3" + transformer_config = _load_ltx_metadata_json(model, "transformer/config.json", revision) + if transformer_config.get("ff_bias") is False: + return "2.5" + vocoder_entry = model_index.get("vocoder") if isinstance(vocoder_entry, (list, tuple)) and vocoder_entry: vocoder_class = str(vocoder_entry[-1]) @@ -159,21 +400,32 @@ def detect_ltx_model_version(model: str) -> str: if vocoder_class == "LTX2VocoderWithBWE": return "2.3" - vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json") + vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json", revision) if str(vocoder_config.get("model_version", "")).startswith("2.3"): return "2.3" if vocoder_config.get("_class_name") == "LTX2VocoderWithBWE": return "2.3" - - logger.info("Using LTX-2 defaults for checkpoint %s", model) + if "bwe" in vocoder_config or "bwe_config" in vocoder_config: + return "2.3" return "2" +def preserves_reference_image_size(*, model: str | None, revision: str | None = None) -> bool: + """Preserve source geometry only for the LTX-2.5 CRF-18 I2V path.""" + return model is not None and detect_ltx_model_version(model, revision=revision) == "2.5" + + class _LTXConnectorAttnProcessor: """Preserve official connector math around Omni attention dispatch.""" - def __init__(self, *, has_learned_registers: bool) -> None: + def __init__( + self, + *, + has_learned_registers: bool = False, + preserve_learned_register_mask: bool = False, + ) -> None: self.has_learned_registers = has_learned_registers + self.preserve_learned_register_mask = preserve_learned_register_mask def __call__( self, @@ -187,12 +439,20 @@ class _LTXConnectorAttnProcessor: encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None - query = attn.norm_q(attn.to_q(hidden_states)) - key = attn.norm_k(attn.to_k(encoder_hidden_states)) + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) + # Offload hooks may execute affine-free Q/K norms in FP32. Restore the + # projection dtype before attention, matching the fully resident path. + query = attn.norm_q(query).to(dtype=value.dtype) + key = attn.norm_k(key).to(dtype=value.dtype) if query_rotary_emb is not None: + # Diffusers builds connector RoPE in FP32, while the official + # connector materializes it in the hidden-state dtype. + query_rotary_emb = tuple(component.to(value.dtype) for component in query_rotary_emb) key_rotary_emb = key_rotary_emb if key_rotary_emb is not None else query_rotary_emb + key_rotary_emb = tuple(component.to(value.dtype) for component in key_rotary_emb) if attn.rope_type == "interleaved": query = apply_interleaved_rotary_emb(query, query_rotary_emb) key = apply_interleaved_rotary_emb(key, key_rotary_emb) @@ -202,6 +462,10 @@ class _LTXConnectorAttnProcessor: else: raise ValueError(f"Unsupported LTX connector RoPE type: {attn.rope_type}") + # Keep Q/K in the projection dtype expected by the attention backend. + query = query.to(dtype=value.dtype) + key = key.to(dtype=value.dtype) + batch_size, _, inner_dim = query.shape head_dim = inner_dim // attn.heads kv_heads = attn.inner_kv_dim // attn.head_dim @@ -209,10 +473,10 @@ class _LTXConnectorAttnProcessor: key = key.view(batch_size, -1, kv_heads, head_dim) value = value.view(batch_size, -1, kv_heads, head_dim) - # The connector replaces padding tokens with learned registers before - # entering its blocks, so every key is valid and the old padding mask - # becomes an all-keep no-op. - if self.has_learned_registers: + # LTX-2/2.3 replace padding tokens with learned registers, so their + # old padding mask becomes a no-op. LTX-2.5 preserves that all-zero + # additive mask to match the official masked-SDPA dispatch. + if self.has_learned_registers and not self.preserve_learned_register_mask: attention_mask = None elif attention_mask is not None and attn.omni_attention.attn_backend.get_name().upper() == "FLASH_ATTN": attention_mask = to_ltx_padding_mask(attention_mask) @@ -229,7 +493,11 @@ class _LTXConnectorAttnProcessor: return attn.to_out[1](hidden_states) -def _install_connector_attention(connectors: LTX2TextConnectors) -> None: +def _install_connector_attention( + connectors: LTX2TextConnectors, + *, + preserve_learned_register_mask: bool = False, +) -> None: for connector_name in ("video_connector", "audio_connector"): connector = getattr(connectors, connector_name, None) has_learned_registers = getattr(connector, "learnable_registers", None) is not None @@ -248,15 +516,20 @@ def _install_connector_attention(connectors: LTX2TextConnectors) -> None: skip_sequence_parallel=True, disable_kv_quant=True, ) - attention.set_processor(_LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)) + attention.set_processor( + _LTXConnectorAttnProcessor( + has_learned_registers=has_learned_registers, + preserve_learned_register_mask=preserve_learned_register_mask, + ) + ) -def _detect_vocoder_output_sample_rate(model: str) -> int | None: +def _detect_vocoder_output_sample_rate(model: str, revision: str | None = None) -> int | None: """Read the generated waveform sample rate from the vocoder config.""" vocoder_config_path = os.path.join(model, "vocoder", "config.json") if not os.path.exists(vocoder_config_path): try: - vocoder_config_path = hf_hub_download(model, "vocoder/config.json") + vocoder_config_path = hf_hub_download(model, "vocoder/config.json", revision=revision) except Exception: return None try: @@ -268,7 +541,10 @@ def _detect_vocoder_output_sample_rate(model: str) -> int | None: def get_ltx2_post_process_func(od_config: Any): """Build the common LTX engine-output adapter.""" - output_sample_rate = _detect_vocoder_output_sample_rate(od_config.model) + output_sample_rate = _detect_vocoder_output_sample_rate( + od_config.model, + revision=getattr(od_config, "revision", None), + ) def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): if not (isinstance(output, tuple) and len(output) == 2): @@ -291,17 +567,50 @@ def _load_component( *, local_files_only: bool, dtype: torch.dtype, + revision: str | None, + prefetch_list: tuple[str, ...] = _LTX_COMPONENT_SUBFOLDERS, ) -> Any: return from_pretrained_with_prefetch( component_cls.from_pretrained, model, subfolder=subfolder, - prefetch_list=_LTX_COMPONENT_SUBFOLDERS, + prefetch_list=prefetch_list, local_files_only=local_files_only, + revision=revision, torch_dtype=dtype, ) +def _load_ltx25_native_diffusion_decoder( + model: str, + *, + local_files_only: bool, + dtype: torch.dtype, + revision: str | None, +) -> DistributedLTX2VideoDiffusionDecoderModel: + """Load canonical Native weights into the local Diffusers-compatible class.""" + config = DistributedLTX2VideoDiffusionDecoderModel.load_config( + model, + subfolder=_LTX2_DIFFUSION_DECODER_SUBFOLDER, + local_files_only=local_files_only, + revision=revision, + ) + checkpoint_path = resolve_ltx_artifact( + model, + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, + LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, + model_revision=revision, + artifact_revision=LTX25_NATIVE_ARTIFACT_REVISION, + ) + decoder = DistributedLTX2VideoDiffusionDecoderModel.from_ltx25_native_checkpoint( + checkpoint_path, + config, + dtype, + ) + decoder.init_distributed() + return decoder + + def _place_aux_components(pipeline: Any) -> None: parallel_config = getattr(pipeline.od_config, "parallel_config", None) use_managed_placement = bool( @@ -324,31 +633,37 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: pipeline.device = get_local_device() dtype = getattr(od_config, "dtype", torch.bfloat16) model = od_config.model + revision = getattr(od_config, "revision", None) local_files_only = os.path.exists(model) + use_diffusion_decoder = pipeline.use_diffusion_decoder pipeline.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=model, - subfolder="transformer", - revision=None, + subfolder=profile.transformer_subfolder, + revision=revision, prefix="transformer.", fall_back_to_pt=True, ), ] - prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only) + prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only, revision=revision) pipeline.tokenizer = AutoTokenizer.from_pretrained( model, subfolder="tokenizer", local_files_only=local_files_only, + revision=revision, ) + if profile.text_encoder_cls is None: + raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15.") with torch.device("cpu"): pipeline.text_encoder = _load_component( - Gemma3ForConditionalGeneration, + profile.text_encoder_cls, model, "text_encoder", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) pipeline.connectors = _load_component( LTX2TextConnectors, @@ -356,21 +671,48 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "connectors", local_files_only=local_files_only, dtype=dtype, + revision=revision, + ) + _install_connector_attention( + pipeline.connectors, + preserve_learned_register_mask=profile.preserve_connector_attention_mask, ) - _install_connector_attention(pipeline.connectors) pipeline.vae = _load_component( profile.video_vae_cls, model, "vae", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) + if use_diffusion_decoder: + pipeline.diffusion_decoder = _load_ltx25_native_diffusion_decoder( + model, + local_files_only=local_files_only, + dtype=dtype, + revision=revision, + ) + # The canonical decoder is trained and validated with NATTEN. Its + # portable FlexAttention fallback materializes an impractically large + # block mask at production video sizes, so fail early if the matching + # Hub kernel cannot be loaded instead of failing later during decode. + pipeline.diffusion_decoder.set_attn_processor(_create_ltx25_natten_processor()) + vae_patch_parallel_size = int( + getattr(getattr(od_config, "parallel_config", None), "vae_patch_parallel_size", 1) + ) + if vae_patch_parallel_size > 1 or getattr(od_config, "vae_use_tiling", False): + pipeline.diffusion_decoder.enable_tiling() + pipeline.diffusion_decoder.set_parallel_size( + vae_patch_parallel_size, + mode=getattr(getattr(od_config, "parallel_config", None), "vae_parallel_mode", "tile"), + ) pipeline.audio_vae = _load_component( AutoencoderKLLTX2Audio, model, "audio_vae", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) try: pipeline.vocoder = _load_component( @@ -379,6 +721,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "vocoder", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) except (TypeError, OSError, ValueError): if profile.vocoder_fallback_cls is None or profile.vocoder_fallback_cls is profile.vocoder_cls: @@ -389,21 +732,51 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: "vocoder", local_files_only=local_files_only, dtype=dtype, + revision=revision, ) if "latent_upsampler" in profile.resident_modules: - # BlurDownsample constructs an integer kernel that must be initialized - # on CPU; component placement is handled uniformly after construction. - with torch.device("cpu"): - pipeline.latent_upsampler = _load_component( - LTX2LatentUpsamplerModel, + upsampler_config = os.path.join(model, "latent_upsampler", "config.json") + if os.path.isfile(upsampler_config) or not local_files_only: + try: + # Diffusers builds the rational-resampler kernel with Long + # tensors; constructing it under a CUDA default-device context + # attempts an unsupported Long addmm before weights are loaded. + with torch.device("cpu"): + pipeline.latent_upsampler = _load_component( + LTX2LatentUpsamplerModel, + model, + "latent_upsampler", + local_files_only=local_files_only, + dtype=dtype, + revision=revision, + ) + except (OSError, ValueError): + if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None: + raise + upsampler_path = resolve_ltx_artifact( + model, + profile.artifact_repo_id, + profile.latent_upsampler_filename, + model_revision=revision, + artifact_revision=profile.artifact_revision, + ) + pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype) + else: + if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None: + raise FileNotFoundError(f"LTX latent upsampler component not found under {model}.") + upsampler_path = resolve_ltx_artifact( model, - "latent_upsampler", - local_files_only=local_files_only, - dtype=dtype, + profile.artifact_repo_id, + profile.latent_upsampler_filename, + model_revision=revision, + artifact_revision=profile.artifact_revision, ) + pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype) - transformer_config = load_transformer_config(model, "transformer", local_files_only) + transformer_config = load_transformer_config( + model, profile.transformer_subfolder, local_files_only, revision=revision + ) quant_config = getattr(od_config, "quantization_config", None) pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) _place_aux_components(pipeline) @@ -411,6 +784,18 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: model, subfolder="scheduler", local_files_only=local_files_only, + revision=revision, + ) + if profile.scheduler_use_dynamic_shifting: + pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config( + pipeline.scheduler.config, + use_dynamic_shifting=True, + shift_terminal=profile.scheduler_shift_terminal, + ) + validate_ltx_checkpoint( + pipeline.scheduler.config, + expected_kind=resolve_ltx_checkpoint_kind(pipeline.pipeline_kind), + pipeline_name=type(pipeline).__name__, ) pipeline.vae_spatial_compression_ratio = pipeline.vae.spatial_compression_ratio @@ -438,6 +823,8 @@ def load_transformer_config( model_path: str, subfolder: str = "transformer", local_files_only: bool = True, + *, + revision: str | None = None, ) -> dict: """Load an LTX transformer config from a local model or the HF Hub.""" if local_files_only: @@ -448,6 +835,7 @@ def load_transformer_config( config_path = hf_hub_download( repo_id=model_path, filename=f"{subfolder}/config.json", + revision=revision, ) with open(config_path) as config_file: return json.load(config_file) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py index f4d2f26f6..8b4491e71 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Task-specific conditioning shared by LTX model versions.""" @@ -7,6 +7,7 @@ from __future__ import annotations import math from dataclasses import dataclass +from io import BytesIO from typing import TYPE_CHECKING, Any import numpy as np @@ -52,19 +53,18 @@ def _preprocess_i2v_pil_images( *, height: int, width: int, + crf: int = 0, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> torch.Tensor: - """Match official LTX aspect-preserving resize and center crop.""" + """Match the version-specific official LTX PIL preprocessing.""" image_list = images if isinstance(images, list) else [images] processed = [] for image in image_list: - pixels = ( - torch.from_numpy(np.array(image.convert("RGB"), copy=True)) - .permute(2, 0, 1) - .unsqueeze(0) - .to(device=device, dtype=torch.float32) - ) + image_array = np.array(image.convert("RGB"), copy=True) + if crf: + image_array = _apply_image_conditioning_crf(image_array, crf) + pixels = torch.from_numpy(image_array).permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=torch.float32) source_height, source_width = pixels.shape[-2:] scale = max(height / source_height, width / source_width) resized_height = math.ceil(source_height * scale) @@ -82,6 +82,47 @@ def _preprocess_i2v_pil_images( return torch.cat(processed, dim=0) +def _apply_image_conditioning_crf(image: np.ndarray, crf: int) -> np.ndarray: + """Round-trip one RGB frame through H.264 to match LTX-2.5 training.""" + if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3: + raise ValueError("LTX image-conditioning CRF expects a uint8 RGB image.") + if min(image.shape[:2]) < 2: + return image + try: + import av + except ImportError as exc: + raise ImportError("PyAV with a libx264 encoder is required for LTX-2.5 image conditioning (CRF 18).") from exc + + with BytesIO() as output_file: + container = av.open(output_file, "w", format="mp4") + try: + stream = container.add_stream("libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"}) + height = image.shape[0] // 2 * 2 + width = image.shape[1] // 2 * 2 + stream.height = height + stream.width = width + frame = av.VideoFrame.from_ndarray(image[:height, :width], format="rgb24").reformat(format="yuv420p") + container.mux(stream.encode(frame)) + container.mux(stream.encode()) + except Exception as exc: + raise RuntimeError( + "The installed PyAV/FFmpeg build cannot encode libx264 video required " + "for LTX-2.5 image conditioning (CRF 18)." + ) from exc + finally: + container.close() + video_bytes = output_file.getvalue() + + with BytesIO(video_bytes) as video_file: + container = av.open(video_file) + try: + stream = next(stream for stream in container.streams if stream.type == "video") + frame = next(container.decode(stream)) + finally: + container.close() + return frame.to_ndarray(format="rgb24") + + class LTXTextConditioningMixin: """Shared Gemma encoding and text connector orchestration.""" @@ -106,13 +147,30 @@ class LTXTextConditioningMixin: if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token - text_inputs = self.tokenizer( - [text.strip() for text in prompt], + bos_token_id = self.tokenizer.bos_token_id + if bos_token_id is None: + raise ValueError("Tokenizer is missing bos_token_id; LTX prompt encoding requires a leading BOS.") + + input_ids = [] + for text in prompt: + encoded = self.tokenizer( + text.strip(), + padding=False, + truncation=True, + max_length=max_sequence_length, + return_tensors=None, + ) + ids = encoded.input_ids + if not ids or ids[0] != bos_token_id: + ids = [bos_token_id, *ids][:max_sequence_length] + input_ids.append(ids) + + text_inputs = self.tokenizer.pad( + {"input_ids": input_ids}, padding="max_length", max_length=max_sequence_length, - truncation=True, - add_special_tokens=True, return_tensors="pt", + return_attention_mask=True, ) text_input_ids = text_inputs.input_ids.to(device) prompt_attention_mask = text_inputs.attention_mask.to(device) @@ -121,7 +179,6 @@ class LTXTextConditioningMixin: attention_mask=prompt_attention_mask, output_hidden_states=True, ).hidden_states - prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype) prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) @@ -445,10 +502,12 @@ class LTXI2VConditioningMixin: init_latents = torch.cat(init_latents, dim=0).to(dtype=dtype) if num_videos_per_prompt > 1: init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0) + latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self) return latent_ops.normalize_latents( init_latents, - self.vae.latents_mean, - self.vae.latents_std, + latents_mean, + latents_std, + scaling_factor, ) @staticmethod @@ -501,6 +560,7 @@ class LTXI2VConditioningMixin: vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, ) mask_shape = (batch_size, 1, num_frames, height, width) + is_ltx25 = getattr(self, "model_version", None) == "2.5" if latents is not None: unpacked_latents = latents.ndim == 5 @@ -512,12 +572,14 @@ class LTXI2VConditioningMixin: ) conditioning_mask = latents.new_zeros(mask_shape) conditioning_mask[:, :, 0] = 1.0 + latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self) latents = latent_ops.normalize_latents( latents, - self.vae.latents_mean, - self.vae.latents_std, - self.vae.config.scaling_factor, + latents_mean, + latents_std, + scaling_factor, ) + clean_latents = latents if image is not None: image_latents = self._encode_i2v_image_latents( image, @@ -525,14 +587,19 @@ class LTXI2VConditioningMixin: generator=generator, dtype=dtype, ) - # Official distilled Stage 2 starts from the upsampled Stage 1 - # state, then reapplies the source image at final resolution. - latents = self._replace_first_latent_frame(latents, image_latents) + # Official Stage 2 starts from the upsampled Stage 1 state, + # then reapplies the source image at final resolution. + clean_latents = self._replace_first_latent_frame(clean_latents, image_latents) latents = latent_ops.pack_latents( latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ) + clean_latents = latent_ops.pack_latents( + clean_latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) else: conditioning_mask = latents.new_zeros(mask_shape) conditioning_mask[:, :, 0] = 1.0 @@ -548,14 +615,14 @@ class LTXI2VConditioningMixin: f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}." ) if unpacked_latents: - # Official LTX patchifies the latent state before drawing noise. - # Sampling the packed shape preserves its seed-to-output contract. - latents = latent_ops.create_noised_state( + latents = latent_ops.create_conditioned_noised_state( latents, - noise_scale * (1 - conditioning_mask).unsqueeze(-1), + clean_latents, + (1 - conditioning_mask.float()).unsqueeze(-1), + noise_scale, generator, ) - elif image is not None: + elif not unpacked_latents and image is not None: image_latents = self._encode_i2v_image_latents( image, batch_size=batch_size, @@ -595,13 +662,27 @@ class LTXI2VConditioningMixin: self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ) - conditioning_mask = latent_ops.pack_latents( + packed_conditioning_mask = latent_ops.pack_latents( conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ).squeeze(-1) - noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype) - latents = init_latents * conditioning_mask.unsqueeze(-1) + noise * (1 - conditioning_mask).unsqueeze(-1) + if is_ltx25: + clean_latents = init_latents + init_latents = torch.zeros_like(init_latents) + latents = latent_ops.create_conditioned_noised_state( + init_latents, + clean_latents, + (1 - packed_conditioning_mask.float()).unsqueeze(-1), + noise_scale, + generator=generator, + ) + else: + noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype) + latents = init_latents * packed_conditioning_mask.unsqueeze(-1) + noise * ( + 1 - packed_conditioning_mask + ).unsqueeze(-1) + conditioning_mask = packed_conditioning_mask return latents, conditioning_mask def _step_video_latents_i2v( @@ -665,6 +746,17 @@ class LTXI2VConditioningMixin: ) if image is not None: + default_crf = 18 if getattr(self, "model_version", "2") == "2.5" else 0 + image_crf = getattr(request_inputs, "image_crf", None) + image_crf = default_crf if image_crf is None else int(image_crf) + is_pil_input = isinstance(image, PIL.Image.Image) or ( + isinstance(image, list) and image and all(isinstance(item, PIL.Image.Image) for item in image) + ) + if image_crf != 0 and not is_pil_input: + raise ValueError( + "`image_crf` re-compression requires PIL images. Pass PIL images, or set `image_crf=0` " + "to skip re-compression." + ) if isinstance(image, torch.Tensor): if image.ndim == 3: image = image.unsqueeze(0) @@ -677,6 +769,7 @@ class LTXI2VConditioningMixin: image, height=request_inputs.height, width=request_inputs.width, + crf=image_crf, device=device, dtype=prompt_context.positive_connector_prompt_embeds.dtype, ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index 4a149be85..7ebd94143 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -42,6 +42,7 @@ class LTXForwardContext: original_audio_num_frames: int padded_audio_num_frames: int timesteps: torch.Tensor + sampler: str audio_scheduler: Any video_audio_step_adapter: Any @@ -62,13 +63,14 @@ class LTXDenoiseContext: audio_latents: torch.Tensor video_coords: torch.Tensor audio_coords: torch.Tensor + audio_attention_mask: torch.Tensor | None = None conditioning_mask: torch.Tensor | None = None conditioning_mask_for_model: torch.Tensor | None = None @dataclass class LTXPhaseResult: - """Denoised, unpacked AV latents and the context used to produce them.""" + """Denoised AV latents and the context used to produce them.""" forward_context: LTXForwardContext video: torch.Tensor @@ -87,6 +89,11 @@ class LTXDenoisePipeline(Protocol): LTXDenoiseStep = Callable[[int, torch.Tensor, LTXAVState], LTXAVState] +LTX_ANCESTRAL_ETA = 1.0 +LTX_ANCESTRAL_S_NOISE = 1.0 +LTX_ANCESTRAL_NOISE_SEED_OFFSET = 10000 +LTX_OFFICIAL_DEFAULT_SEED = 10 + class LTXDenoiseExecutor: """Run the one shared LTX denoise loop. @@ -137,6 +144,9 @@ class LTXVideoAudioStepAdapter: latent_width: int, *, image_conditioned: bool, + sampler: str = "euler", + generator: torch.Generator | list[torch.Generator] | None = None, + conditioning_mask: torch.Tensor | None = None, ) -> None: self._pipeline = pipeline self._audio_scheduler = audio_scheduler @@ -144,11 +154,27 @@ class LTXVideoAudioStepAdapter: self._latent_height = latent_height self._latent_width = latent_width self._image_conditioned = image_conditioned + self._sampler = sampler + self._conditioning_mask = conditioning_mask + self._ancestral_generators = ( + _make_ancestral_generators(generator, pipeline.device) if sampler == "euler_ancestral" else None + ) self._step_index = 0 def step(self, noise_pred, t, latents, return_dict=False, generator=None): del t, return_dict, generator - if self._image_conditioned: + if self._sampler == "euler_ancestral": + video_out = _ancestral_euler_step_from_velocity( + latents[0], + noise_pred[0], + self._pipeline.scheduler.sigmas, + self._step_index, + self._ancestral_generators, + ) + if self._conditioning_mask is not None: + mask = self._conditioning_mask.unsqueeze(-1).to(video_out.dtype) + video_out = torch.lerp(video_out, latents[0], mask) + elif self._image_conditioned: video_out = self._pipeline._step_video_latents_i2v( noise_pred[0], latents[0], @@ -164,17 +190,96 @@ class LTXVideoAudioStepAdapter: self._pipeline.scheduler.sigmas, self._step_index, ) - audio_out = euler_step_from_velocity( - latents[1], - noise_pred[1], - self._audio_scheduler.sigmas, - self._step_index, - ) + if self._sampler == "euler_ancestral": + audio_out = _ancestral_euler_step_from_velocity( + latents[1], + noise_pred[1], + self._audio_scheduler.sigmas, + self._step_index, + self._ancestral_generators, + ) + else: + audio_out = euler_step_from_velocity( + latents[1], + noise_pred[1], + self._audio_scheduler.sigmas, + self._step_index, + ) self._step_index += 1 return ((video_out, audio_out),) -def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> torch.Tensor: +def _make_ancestral_generators( + generator: torch.Generator | list[torch.Generator] | None, + device: torch.device, +) -> torch.Generator | list[torch.Generator]: + """Create the official independent ``seed + 10000`` noise stream.""" + + def _offset(source: torch.Generator | None) -> torch.Generator: + seed = LTX_OFFICIAL_DEFAULT_SEED if source is None else source.initial_seed() + return torch.Generator(device=device).manual_seed(seed + LTX_ANCESTRAL_NOISE_SEED_OFFSET) + + if isinstance(generator, list): + return [_offset(item) for item in generator] + return _offset(generator) + + +def _randn_like_with_generators( + sample: torch.Tensor, + generators: torch.Generator | list[torch.Generator], +) -> torch.Tensor: + if not isinstance(generators, list): + return torch.randn(sample.shape, generator=generators, dtype=sample.dtype, device=sample.device) + if len(generators) != sample.shape[0]: + raise ValueError( + f"LTX ancestral sampling received {len(generators)} generators for batch size {sample.shape[0]}." + ) + return torch.cat( + [ + torch.randn((1, *sample.shape[1:]), generator=item, dtype=sample.dtype, device=sample.device) + for item in generators + ] + ) + + +def _ancestral_euler_step_from_velocity( + sample: torch.Tensor, + velocity: torch.Tensor, + sigmas: torch.Tensor, + step_index: int, + generators: torch.Generator | list[torch.Generator] | None, +) -> torch.Tensor: + """Apply the official LTX-2.5 rectified-flow ancestral Euler step.""" + sigma = sigmas[step_index].to(torch.float32) + sigma_next = sigmas[step_index + 1].to(torch.float32) + # Official X0Model materializes the denoised prediction in the model + # dtype before the ancestral step promotes it back to float32. + denoised_model_dtype = (sample.float() - velocity.float() * sigma).to(sample.dtype) + denoised = denoised_model_dtype.float() + if sigma_next == 0: + return denoised.to(sample.dtype) + if generators is None: + raise ValueError("LTX ancestral Euler sampling requires a noise generator.") + + downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * LTX_ANCESTRAL_ETA + sigma_down = sigma_next * downstep_ratio + sigma_down_ratio = sigma_down / sigma + x_next = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised + + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp(min=0).sqrt() + noise = _randn_like_with_generators(sample, generators) + x_next = (alpha_next / alpha_down) * x_next + noise.float() * LTX_ANCESTRAL_S_NOISE * renoise_coeff + result = x_next.to(sample.dtype) + return result + + +def _official_ltx_sigmas( + scheduler: Any, + steps: int, + device: torch.device, +) -> torch.Tensor: """Build the official LTX one-stage sigma schedule in torch fp32.""" config = scheduler.config base_anchor = config.get("base_image_seq_len", 1024) @@ -184,7 +289,8 @@ def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> to sigmas = torch.linspace(1.0, 0.0, steps + 1) slope = (max_shift - base_shift) / (max_anchor - base_anchor) - # Official LTX2 one-stage intentionally uses the max sequence anchor, so the shift stays at max_shift. + # Official LTX one-stage pipelines omit the latent when constructing + # this schedule, so the shift stays at the max sequence anchor. sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor) exp_shift = math.exp(sigma_shift) sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0) @@ -223,6 +329,9 @@ def prepare_scheduler_stage( latent_width: int, use_official_sigma_schedule: bool, image_conditioned: bool = False, + sampler: str = "euler", + generator: torch.Generator | list[torch.Generator] | None = None, + conditioning_mask: torch.Tensor | None = None, ) -> tuple[Any, Any, torch.Tensor]: if sigmas is not None and timesteps is not None: raise ValueError("Only one of `sigmas` or `timesteps` may be provided.") @@ -235,6 +344,9 @@ def prepare_scheduler_stage( latent_height, latent_width, image_conditioned=image_conditioned, + sampler=sampler, + generator=generator, + conditioning_mask=conditioning_mask, ) if sigmas is not None: scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device) @@ -284,7 +396,8 @@ def prepare_rope_coords_stage( latents: torch.Tensor, audio_latents: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - video_coords = pipeline.transformer.rope.prepare_video_coords( + transformer = pipeline.transformer + video_coords = transformer.rope.prepare_video_coords( latents.shape[0], forward_ctx.latent_num_frames, forward_ctx.latent_height, @@ -292,7 +405,7 @@ def prepare_rope_coords_stage( latents.device, fps=forward_ctx.request_inputs.frame_rate, ) - audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords( + audio_coords = transformer.audio_rope.prepare_audio_coords( audio_latents.shape[0], forward_ctx.padded_audio_num_frames, audio_latents.device, @@ -300,6 +413,20 @@ def prepare_rope_coords_stage( return video_coords, audio_coords +def _first_frame_keyframes_mask(reference: torch.Tensor, latent_num_frames: int) -> torch.Tensor: + """Mark the first causal latent frame, matching official LTX-2.5.""" + if latent_num_frames <= 0: + raise ValueError(f"LTX latent frame count must be positive, got {latent_num_frames}.") + tokens_per_frame, remainder = divmod(reference.shape[1], latent_num_frames) + if remainder: + raise ValueError( + f"LTX video token count {reference.shape[1]} is not divisible by {latent_num_frames} latent frames." + ) + mask = reference.new_zeros((reference.shape[0], reference.shape[1], 1)) + mask[:, :tokens_per_frame] = 1 + return mask + + def build_transformer_kwargs( pipeline: Any, forward_ctx: LTXForwardContext, @@ -315,9 +442,13 @@ def build_transformer_kwargs( attention_kwargs: dict[str, Any] | None = None, ) -> dict[str, Any]: del encoder_attention_mask, audio_encoder_attention_mask + keyframes_mask = None + if getattr(pipeline.transformer.config, "use_keyframes_abs_pos_embedding", False): + keyframes_mask = _first_frame_keyframes_mask(hidden_states, forward_ctx.latent_num_frames) return { "hidden_states": hidden_states, "audio_hidden_states": audio_hidden_states, + "keyframes_mask": keyframes_mask, "encoder_hidden_states": encoder_hidden_states, "audio_encoder_hidden_states": audio_encoder_hidden_states, **pipeline._denoise_timestep_kwargs( @@ -331,6 +462,7 @@ def build_transformer_kwargs( # with a learned register, making all output context tokens valid. "encoder_attention_mask": None, "audio_encoder_attention_mask": None, + "audio_attention_mask": denoise_ctx.audio_attention_mask, "num_frames": forward_ctx.latent_num_frames, "height": forward_ctx.latent_height, "width": forward_ctx.latent_width, @@ -423,6 +555,9 @@ class LTXPhaseExecutor: latent_width=latent_width, use_official_sigma_schedule=phase_recipe.use_official_sigma_schedule, image_conditioned=conditioning_mask is not None, + sampler=phase_recipe.sampler, + generator=request_inputs.generator, + conditioning_mask=conditioning_mask, ) forward_ctx = LTXForwardContext( req=req, @@ -438,15 +573,30 @@ class LTXPhaseExecutor: original_audio_num_frames=original_audio_num_frames, padded_audio_num_frames=padded_audio_num_frames, timesteps=timesteps_tensor, + sampler=phase_recipe.sampler, audio_scheduler=audio_scheduler, video_audio_step_adapter=video_audio_step_adapter, ) video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents) + ring_degree = getattr(pipeline.od_config.parallel_config, "ring_degree", 1) or 1 + if padded_audio_num_frames > original_audio_num_frames and ring_degree > 1: + raise ValueError( + "LTX audio padding requires an attention mask, which Ring sequence parallelism does not support. " + "Use Ulysses-only SP or choose a request whose audio latent length is divisible by the SP size." + ) denoise_ctx = LTXDenoiseContext( latents=latents, audio_latents=audio_latents, video_coords=video_coords, audio_coords=audio_coords, + audio_attention_mask=( + torch.arange(padded_audio_num_frames, device=audio_latents.device) + .lt(original_audio_num_frames) + .unsqueeze(0) + .expand(audio_latents.shape[0], -1) + if padded_audio_num_frames > original_audio_num_frames + else None + ), conditioning_mask=conditioning_mask, ) denoise_ctx = pipeline._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py new file mode 100644 index 000000000..73db2ea6b --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved. +# +# Native conversion rules and decoder behavior copied and modified from +# Diffusers commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0). + +"""LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder.""" + +from collections.abc import Mapping +from typing import Any + +import torch +from diffusers.models.autoencoders import ( + LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel, +) +from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( + LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d, +) +from diffusers.models.autoencoders.vae import DecoderOutput +from safetensors import safe_open + +LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5" +LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1" +LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors" + +_NATIVE_DECODER_PREFIXES = ("vae.decoder.", "decoder.") +_NATIVE_STATISTICS_KEYS = { + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", +} +_NATIVE_KEY_REPLACEMENTS = ( + ("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1."), + ("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2."), + (".attn.proj.", ".attn.to_out.0."), + (".attn.q_norm.", ".attn.norm_q."), + (".attn.k_norm.", ".attn.norm_k."), +) +_GATE_FOLD_TARGETS = { + ".attn.to_out.0.weight": ".gate_msa", + ".attn.to_out.0.bias": ".gate_msa", + ".mlp.w_down.weight": ".gate_mlp", + ".context_proj.weight": ".gate_ctx", + ".context_proj.bias": ".gate_ctx", +} +_GATE_SUFFIXES = tuple(_GATE_FOLD_TARGETS.values()) + + +def _strip_native_decoder_prefix(key: str) -> str | None: + for prefix in _NATIVE_DECODER_PREFIXES: + if key.startswith(prefix): + return key.removeprefix(prefix) + return None + + +def _native_statistics_target(key: str) -> str | None: + return _NATIVE_STATISTICS_KEYS.get(key.removeprefix("vae.")) + + +def _rename_native_decoder_key(key: str) -> str: + for source, target in _NATIVE_KEY_REPLACEMENTS: + key = key.replace(source, target) + return key + + +def _fold_native_gate(key: str, value: torch.Tensor, gates: Mapping[str, torch.Tensor]) -> torch.Tensor: + target_suffix = next((suffix for suffix in _GATE_FOLD_TARGETS if key.endswith(suffix)), None) + if target_suffix is None: + return value + gate_key = key[: -len(target_suffix)] + _GATE_FOLD_TARGETS[target_suffix] + gate = gates.get(gate_key) + if gate is None: + return value + gate = gate.to(device=value.device, dtype=torch.float32) + value_float = value.to(dtype=torch.float32) + folded = gate.unsqueeze(1) * value_float if value.ndim == 2 else gate * value_float + return folded.to(dtype=value.dtype) + + +def convert_ltx25_native_diffusion_decoder_state_dict( + native_state_dict: Mapping[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Convert the canonical Native LTX-2.5 DiffVAE state dict to Diffusers.""" + decoder_state_dict: dict[str, torch.Tensor] = {} + converted: dict[str, torch.Tensor] = {} + for key, value in native_state_dict.items(): + statistics_target = _native_statistics_target(key) + if statistics_target is not None: + converted[statistics_target] = value + continue + decoder_key = _strip_native_decoder_prefix(key) + if decoder_key is not None: + decoder_state_dict[decoder_key] = value + + gates = {key: value for key, value in decoder_state_dict.items() if key.endswith(_GATE_SUFFIXES)} + for key, value in decoder_state_dict.items(): + # ``type_emb`` is an unused shipping-checkpoint residue. Preview heads + # and legacy static gates are not parameters of the runtime decoder. + if key == "type_emb" or key.startswith("coarse_") or ".coarse_" in key or key in gates: + continue + + converted_key = _rename_native_decoder_key(key) + value = _fold_native_gate(converted_key, value, gates) + if converted_key.endswith((".qkv.weight", ".qkv.bias")): + leaf = "weight" if converted_key.endswith(".weight") else "bias" + prefix = converted_key[: -len(f"qkv.{leaf}")] + if value.shape[0] % 3 != 0: + raise ValueError( + f"Fused LTX-2.5 DiffVAE parameter {key!r} has leading dimension " + f"{value.shape[0]}, which is not divisible by 3." + ) + chunk = value.shape[0] // 3 + converted[f"decoder.{prefix}to_q.{leaf}"] = value[:chunk].clone() + converted[f"decoder.{prefix}to_k.{leaf}"] = value[chunk : 2 * chunk].clone() + converted[f"decoder.{prefix}to_v.{leaf}"] = value[2 * chunk :].clone() + continue + + converted[f"decoder.{converted_key}"] = value + + return converted + + +def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch.Tensor]: + """Read and convert only DiffVAE tensors from the canonical Native VAE file.""" + with safe_open(path, framework="pt", device="cpu") as handle: + native_state_dict = { + key: handle.get_tensor(key) + for key in handle.keys() + if _strip_native_decoder_prefix(key) is not None or _native_statistics_target(key) is not None + } + return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict) + + +class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d): + """Diffusers decoder core with the short-clip NATTEN context fix.""" + + def forward_stage_4( + self, + hidden_states: torch.Tensor, + drop_leading_frame: bool = True, + crop_trailing_ghost: bool = True, + ) -> torch.Tensor: + blocks = self.det_stages[-1] + block_mask = blocks[0].attn.build_block_mask(hidden_states) + for block in blocks: + hidden_states = block(hidden_states, block_mask) + hidden_states = self.upsamples[-1](hidden_states, drop_leading_frame=drop_leading_frame) + + num_pad = self.trailing_pad_latent_frames + if crop_trailing_ghost and num_pad > 0: + content_frames = max(hidden_states.shape[1] - num_pad * self.temporal_compression_ratio, 1) + keep_frames = min(hidden_states.shape[1], max(content_frames, self.stage5_kernel[0])) + hidden_states = hidden_states[:, :keep_frames] + return hidden_states + + +class LTX2VideoDiffusionDecoderModel(DiffusersLTX2VideoDiffusionDecoderModel): + """Diffusers 0.40 decoder plus canonical Native loading and short-clip support.""" + + def __init__( + self, + out_channels: int = 3, + latent_channels: int = 128, + patch_size: int = 4, + scaling_factor: float = 1.0, + decoder_head_dim: int = 64, + decoder_stage_channels: tuple[int, ...] = (2048, 1024, 512, 512, 256), + decoder_stage_depths: tuple[int, ...] = (4, 6, 4, 2, 8), + decoder_stage_kernels: tuple[tuple[int, int, int], ...] = ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)), + decoder_upsample_strides: tuple[tuple[int, int, int], ...] = ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)), + decoder_upsample_channel_reductions: tuple[int, ...] = (2, 2, 1, 2), + decoder_stage5_kernel: tuple[int, int, int] = (11, 11, 11), + decoder_t_emb_dim: int = 384, + decoder_timestep_scale_multiplier: float = 1000.0, + decoder_model_output_type: str = "x0", + decoder_num_inference_steps: int = 1, + spatial_compression_ratio: int = 32, + temporal_compression_ratio: int = 8, + ) -> None: + super().__init__( + out_channels=out_channels, + latent_channels=latent_channels, + patch_size=patch_size, + scaling_factor=scaling_factor, + decoder_head_dim=decoder_head_dim, + decoder_stage_channels=decoder_stage_channels, + decoder_stage_depths=decoder_stage_depths, + decoder_stage_kernels=decoder_stage_kernels, + decoder_upsample_strides=decoder_upsample_strides, + decoder_upsample_channel_reductions=decoder_upsample_channel_reductions, + decoder_stage5_kernel=decoder_stage5_kernel, + decoder_t_emb_dim=decoder_t_emb_dim, + decoder_timestep_scale_multiplier=decoder_timestep_scale_multiplier, + decoder_model_output_type=decoder_model_output_type, + decoder_num_inference_steps=decoder_num_inference_steps, + spatial_compression_ratio=spatial_compression_ratio, + temporal_compression_ratio=temporal_compression_ratio, + ) + # Diffusers constructs its private decoder core internally. Switching + # to this behavior-only subclass preserves parameters and state-dict keys. + self.decoder.__class__ = LTX2VideoDiffusionDecoder3d + self.decoder.stage5_kernel = tuple(decoder_stage5_kernel) + + @classmethod + def from_ltx25_native_checkpoint( + cls, + checkpoint_path: str, + config: Mapping[str, Any], + dtype: torch.dtype, + ) -> "LTX2VideoDiffusionDecoderModel": + """Construct this decoder class directly from the canonical Native checkpoint.""" + with torch.device("meta"): + model = cls.from_config(dict(config)) + state_dict = load_ltx25_native_diffusion_decoder_state_dict(checkpoint_path) + try: + model.load_state_dict(state_dict, strict=True, assign=True) + except RuntimeError as exc: + raise ValueError(f"Invalid LTX-2.5 Native DiffVAE checkpoint {checkpoint_path!r}.") from exc + model.to(device="cpu", dtype=dtype) + return model + + def decode( + self, + z: torch.Tensor, + generator: torch.Generator | None = None, + num_inference_steps: int | None = None, + return_dict: bool = True, + ) -> DecoderOutput | tuple[torch.Tensor]: + decoded = super().decode( + z, + generator=generator, + num_inference_steps=num_inference_steps, + return_dict=False, + )[0] + + # Short clips retain replicated trailing context through stage 5. Do + # not expose those context pixels to the caller. + target_num_frames = (z.shape[2] - 1) * self.temporal_compression_ratio + 1 + target_height = z.shape[3] * self.spatial_compression_ratio + target_width = z.shape[4] * self.spatial_compression_ratio + decoded = decoded[:, :, :target_num_frames, :target_height, :target_width] + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py new file mode 100644 index 000000000..7f486dbdf --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py @@ -0,0 +1,429 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved. +# +# Distributed tiling is copied and modified from Diffusers' serial LTX-2.5 +# tiling at commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0). + +"""LTX-2.5-specific distributed execution for the diffusion VAE decoder.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +from diffusers.utils import logging +from diffusers.utils.torch_utils import randn_tensor + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _tile_intervals(length: int, tile_size: int, stride: int, min_size: int) -> list[tuple[int, int]]: + """Build overlapping intervals, merging a too-small trailing remnant.""" + if length <= tile_size: + return [(0, length)] + starts = list(range(0, length, stride)) + while len(starts) > 1 and length - starts[-1] < min_size: + starts.pop() + return [(start, min(start + tile_size, length)) for start in starts[:-1]] + [(starts[-1], length)] + + +@dataclass(frozen=True) +class LTX2VideoDiffusionTilePlan: + """Geometry shared by all ranks for one distributed DiffVAE decode.""" + + temporal_tiles: tuple[tuple[int, int], ...] + height_tiles: tuple[tuple[int, int], ...] + width_tiles: tuple[tuple[int, int], ...] + num_frames: int + height: int + width: int + scale_t: int + scale_h: int + scale_w: int + stride_t: int + stride_h: int + stride_w: int + blend_frames: int + blend_height: int + blend_width: int + single_step_x0: bool + + +@dataclass +class LTX2VideoDiffusionTileTask(TileTask): + """One stage-4 + stage-5 tile and its rank-independent noise.""" + + drop_leading_frame: bool = False + crop_trailing_ghost: bool = False + noise_generator: torch.Generator | list[torch.Generator] | None = None + x_t: torch.Tensor | None = None + + +class DistributedLTX2VideoDiffusionDecoderModel(LTX2VideoDiffusionDecoderModel, DistributedVaeMixin): + """LTX-2.5 diffusion decoder with model-specific tile parallelism. + + Stages 1-3 run over the complete low-resolution feature volume on every + participating rank. Stage 4 and the diffusion stage run as independent + overlapping tile tasks, and rank 0 performs the reference blend/merge. + """ + + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None: + if mode != "tile": + raise ValueError(f"LTX-2.5 DiffVAE only supports vae_parallel_mode='tile', got {mode!r}.") + super().set_parallel_size(parallel_size, mode=mode) + + def _build_tiled_decode_plan( + self, + features: torch.Tensor, + num_inference_steps: int, + ) -> LTX2VideoDiffusionTilePlan: + """Build the rank-independent tile schedule around Diffusers' decoder.""" + decoder = self.decoder + patch_size = decoder.patch_size + upsample_stride = decoder.upsamples[-1].stride + scale_t, scale_h, scale_w = ( + upsample_stride[0], + upsample_stride[1] * patch_size, + upsample_stride[2] * patch_size, + ) + tile_t = self.tile_sample_min_num_frames // scale_t + tile_h = self.tile_sample_min_height // scale_h + tile_w = self.tile_sample_min_width // scale_w + stride_t = self.tile_sample_stride_num_frames // scale_t + stride_h = self.tile_sample_stride_height // scale_h + stride_w = self.tile_sample_stride_width // scale_w + min_sizes = [ + max(kernel_4, -(-kernel_5 // stride)) + for kernel_4, kernel_5, stride in zip( + self.config.decoder_stage_kernels[-1], + self.config.decoder_stage5_kernel, + upsample_stride, + ) + ] + + ghost_frames = decoder.trailing_pad_latent_frames * math.prod( + upsample.stride[0] for upsample in decoder.upsamples[:-1] + ) + num_frames = features.shape[1] - ghost_frames + height, width = features.shape[2:4] + return LTX2VideoDiffusionTilePlan( + temporal_tiles=tuple(_tile_intervals(num_frames, tile_t, stride_t, min_sizes[0])), + height_tiles=tuple(_tile_intervals(height, tile_h, stride_h, min_sizes[1])), + width_tiles=tuple(_tile_intervals(width, tile_w, stride_w, min_sizes[2])), + num_frames=num_frames, + height=height, + width=width, + scale_t=scale_t, + scale_h=scale_h, + scale_w=scale_w, + stride_t=stride_t, + stride_h=stride_h, + stride_w=stride_w, + blend_frames=(tile_t - stride_t) * scale_t, + blend_height=(tile_h - stride_h) * scale_h, + blend_width=(tile_w - stride_w) * scale_w, + single_step_x0=num_inference_steps == 1 and decoder.model_output_type == "x0", + ) + + @staticmethod + def _iter_tiled_decode_coords(plan: LTX2VideoDiffusionTilePlan): + for t_idx in range(len(plan.temporal_tiles)): + for h_idx in range(len(plan.height_tiles)): + for w_idx in range(len(plan.width_tiles)): + yield t_idx, h_idx, w_idx + + @staticmethod + def _get_tiled_feature_slice( + features: torch.Tensor, + coord: tuple[int, int, int], + plan: LTX2VideoDiffusionTilePlan, + ) -> tuple[torch.Tensor, bool, bool]: + t_idx, h_idx, w_idx = coord + t0, t1 = plan.temporal_tiles[t_idx] + h0, h1 = plan.height_tiles[h_idx] + w0, w1 = plan.width_tiles[w_idx] + is_origin = t0 == 0 + is_trailing = t1 == plan.num_frames + feature_t1 = features.shape[1] if is_trailing else t1 + return features[:, t0:feature_t1, h0:h1, w0:w1], is_origin, is_trailing + + def _tiled_pixel_shape_from_features( + self, + feature_tile: torch.Tensor, + *, + drop_leading_frame: bool, + crop_trailing_ghost: bool, + ) -> tuple[int, int, int, int, int]: + """Predict one output tile's shape without materializing stage 4.""" + decoder = self.decoder + stride_t, stride_h, stride_w = decoder.upsamples[-1].stride + num_frames = feature_tile.shape[1] * stride_t + if stride_t == 2 and drop_leading_frame: + num_frames -= 1 + if crop_trailing_ghost and decoder.trailing_pad_latent_frames > 0: + content_frames = max( + num_frames - decoder.trailing_pad_latent_frames * decoder.temporal_compression_ratio, + 1, + ) + num_frames = min(num_frames, max(content_frames, decoder.stage5_kernel[0])) + return ( + feature_tile.shape[0], + decoder.out_channels, + num_frames, + feature_tile.shape[2] * stride_h * decoder.patch_size, + feature_tile.shape[3] * stride_w * decoder.patch_size, + ) + + @staticmethod + def _slice_tiled_noise( + x_t_full: torch.Tensor, + coord: tuple[int, int, int], + tile_pixel_shape: tuple[int, int, int, int, int], + plan: LTX2VideoDiffusionTilePlan, + ) -> torch.Tensor: + t_idx, h_idx, w_idx = coord + t0, _ = plan.temporal_tiles[t_idx] + h0, _ = plan.height_tiles[h_idx] + w0, _ = plan.width_tiles[w_idx] + pixel_t0 = t0 * plan.scale_t - (1 if t0 != 0 and plan.scale_t == 2 else 0) + return x_t_full[ + :, + :, + pixel_t0 : pixel_t0 + tile_pixel_shape[2], + h0 * plan.scale_h : h0 * plan.scale_h + tile_pixel_shape[3], + w0 * plan.scale_w : w0 * plan.scale_w + tile_pixel_shape[4], + ] + + def _merge_tiled_decode( + self, + tiles: dict[tuple[int, int, int], torch.Tensor], + plan: LTX2VideoDiffusionTilePlan, + ) -> torch.Tensor: + """Blend distributed RGB tiles in Diffusers' serial order.""" + frame_groups = [] + for t_idx in range(len(plan.temporal_tiles)): + rows = [ + [tiles[(t_idx, h_idx, w_idx)] for w_idx in range(len(plan.width_tiles))] + for h_idx in range(len(plan.height_tiles)) + ] + result_rows = [] + for h_idx, row in enumerate(rows): + result_row = [] + for w_idx, tile in enumerate(row): + if h_idx > 0: + tile = self.blend_v(rows[h_idx - 1][w_idx], tile, plan.blend_height) + if w_idx > 0: + tile = self.blend_h(row[w_idx - 1], tile, plan.blend_width) + keep_height = plan.stride_h * plan.scale_h if h_idx < len(rows) - 1 else tile.shape[3] + keep_width = plan.stride_w * plan.scale_w if w_idx < len(row) - 1 else tile.shape[4] + result_row.append(tile[:, :, :, :keep_height, :keep_width]) + result_rows.append(torch.cat(result_row, dim=4)) + frame_groups.append(torch.cat(result_rows, dim=3)) + + result = [] + for t_idx, group in enumerate(frame_groups): + if t_idx > 0: + group = self.blend_t(frame_groups[t_idx - 1], group, plan.blend_frames) + if t_idx < len(frame_groups) - 1: + keep_frames = plan.stride_t * plan.scale_t - (1 if t_idx == 0 and plan.scale_t == 2 else 0) + group = group[:, :, :keep_frames] + result.append(group) + return torch.cat(result, dim=2) + + def _default_generator(self, device: torch.device) -> torch.Generator: + if device.type == "cpu": + return torch.default_generator + device_module = getattr(torch, device.type, None) + default_generators = getattr(device_module, "default_generators", None) + if default_generators is None: + raise ValueError( + f"Distributed LTX-2.5 diffusion decode on {device.type!r} requires an explicit torch.Generator." + ) + device_index = device.index + if device_index is None: + current_device = getattr(device_module, "current_device", None) + device_index = current_device() if current_device is not None else 0 + return default_generators[device_index] + + def _sync_generators( + self, + generator: torch.Generator | list[torch.Generator] | None, + device: torch.device, + ) -> torch.Generator | list[torch.Generator]: + """Make every rank start tile-noise generation from rank 0's state.""" + generators = self._default_generator(device) if generator is None else generator + generator_list = generators if isinstance(generators, list) else [generators] + for item in generator_list: + state_on_device = item.get_state().to(device=device) + dist.broadcast(state_on_device, src=0, group=self.distributed_executor.group) + item.set_state(state_on_device.cpu()) + return generators + + @staticmethod + def _clone_generator(generator: torch.Generator) -> torch.Generator: + cloned = torch.Generator(device=generator.device) + cloned.set_state(generator.get_state()) + return cloned + + def _clone_generators( + self, + generators: torch.Generator | list[torch.Generator], + ) -> torch.Generator | list[torch.Generator]: + if isinstance(generators, list): + return [self._clone_generator(item) for item in generators] + return self._clone_generator(generators) + + def _distributed_tile_split( + self, + z: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None, + num_inference_steps: int, + ) -> tuple[list[LTX2VideoDiffusionTileTask], GridSpec]: + decoder = self.decoder + features = decoder.forward_stages_1_to_3(z) + plan = self._build_tiled_decode_plan(features, num_inference_steps) + generators = self._sync_generators(generator, z.device) + + x_t_full = None + if not plan.single_step_x0: + pixel_frames = plan.num_frames * plan.scale_t - (1 if plan.scale_t == 2 else 0) + x_t_full = randn_tensor( + ( + z.shape[0], + decoder.out_channels, + pixel_frames, + plan.height * plan.scale_h, + plan.width * plan.scale_w, + ), + generator=generators, + device=z.device, + dtype=z.dtype, + ) + + tasks = [] + for tile_id, coord in enumerate(self._iter_tiled_decode_coords(plan)): + feature_tile, is_origin, is_trailing = self._get_tiled_feature_slice(features, coord, plan) + tile_pixel_shape = self._tiled_pixel_shape_from_features( + feature_tile, + drop_leading_frame=is_origin, + crop_trailing_ghost=is_trailing, + ) + tile_generator = None + x_t = None + if plan.single_step_x0: + # Save the canonical serial state for this tile, then advance + # the request generator exactly as the single-rank tiled path + # does. Task placement therefore cannot change tile noise. + tile_generator = self._clone_generators(generators) + discarded_noise = randn_tensor( + tile_pixel_shape, + generator=generators, + device=z.device, + dtype=z.dtype, + ) + del discarded_noise + else: + x_t = self._slice_tiled_noise(x_t_full, coord, tile_pixel_shape, plan) + + tasks.append( + LTX2VideoDiffusionTileTask( + tile_id=tile_id, + grid_coord=coord, + tensor=feature_tile, + workload=math.prod(tile_pixel_shape), + drop_leading_frame=is_origin, + crop_trailing_ghost=is_trailing, + noise_generator=tile_generator, + x_t=x_t, + ) + ) + + return tasks, GridSpec( + split_dims=(2, 3, 4), + grid_shape=(len(plan.temporal_tiles), len(plan.height_tiles), len(plan.width_tiles)), + tile_spec={"plan": plan}, + output_dtype=z.dtype, + ) + + def _distributed_tile_exec( + self, + task: LTX2VideoDiffusionTileTask, + num_inference_steps: int, + ) -> torch.Tensor: + context = self.decoder.forward_stage_4( + task.tensor, + drop_leading_frame=task.drop_leading_frame, + crop_trailing_ghost=task.crop_trailing_ghost, + ) + tile_pixel_shape = ( + context.shape[0], + self.decoder.out_channels, + context.shape[1], + context.shape[2] * self.decoder.patch_size, + context.shape[3] * self.decoder.patch_size, + ) + if task.x_t is None: + x_t = randn_tensor( + tile_pixel_shape, + generator=task.noise_generator, + device=context.device, + dtype=context.dtype, + ) + else: + x_t = task.x_t + return self.decoder.denoise(context, x_t, num_inference_steps) + + def _distributed_tile_merge( + self, + tiles: dict[tuple[int, int, int], torch.Tensor], + grid_spec: GridSpec, + ) -> torch.Tensor: + plan = grid_spec.tile_spec["plan"] + if not isinstance(plan, LTX2VideoDiffusionTilePlan): + raise TypeError(f"Expected an LTX2VideoDiffusionTilePlan, got {type(plan)!r}.") + return self._merge_tiled_decode(tiles, plan) + + def tiled_decode( + self, + z: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None = None, + num_inference_steps: int | None = None, + ) -> torch.Tensor: + if not self.is_distributed_enabled(): + return super().tiled_decode(z, generator=generator, num_inference_steps=num_inference_steps) + + logger.debug("LTX-2.5 diffusion decoder running with distributed stage-4/stage-5 tiles") + num_inference_steps = num_inference_steps or self.decoder.default_num_inference_steps + result = self.distributed_executor.execute( + z, + DistributedOperator( + split=lambda tensor: self._distributed_tile_split(tensor, generator, num_inference_steps), + exec=lambda task: self._distributed_tile_exec(task, num_inference_steps), + merge=self._distributed_tile_merge, + ), + broadcast_result=False, + ) + if result.numel() == 0: + # The base decode method crops five dimensions before the LTX + # runtime discards non-output-rank results. + return torch.empty((0, self.decoder.out_channels, 0, 0, 0), device=z.device, dtype=z.dtype) + return result diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py index f35f5b998..9bb8ded75 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -126,7 +126,10 @@ def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch. def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor: """Convert x0 back to velocity using the official fp32 arithmetic order.""" - sigma = sigma.to(torch.float32) + # Official LTX materializes the scalar schedule value on the host before + # dividing. CUDA tensor-scalar division can cross bfloat16 rounding bounds + # even when the scalar has the same float32 value. + sigma = sigma.to(torch.float32).item() return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype) @@ -148,6 +151,7 @@ def combine_guided_x0( uncond_perturbed: torch.Tensor | float, uncond_modality: torch.Tensor | float, guidance: LTXModalityGuidance, + rescale_token_count: int | None = None, ) -> torch.Tensor: dtype = cond.dtype cond = cond.float() @@ -163,9 +167,18 @@ def combine_guided_x0( # uses its full-tensor standard deviation. In Omni, dimension 0 may # instead contain independently scheduled requests. Reduce each item # separately so its result cannot depend on unrelated co-batched work. - reduce_dims = tuple(range(1, pred.ndim)) - cond_std = cond.std(dim=reduce_dims, keepdim=True) - pred_std = pred.std(dim=reduce_dims, keepdim=True) + stats_cond = cond + stats_pred = pred + if rescale_token_count is not None: + if cond.ndim < 2 or not 0 < rescale_token_count <= cond.shape[1]: + raise ValueError( + f"Guidance rescale token count must be in [1, {cond.shape[1]}], got {rescale_token_count}." + ) + stats_cond = cond[:, :rescale_token_count] + stats_pred = pred[:, :rescale_token_count] + reduce_dims = tuple(range(1, stats_pred.ndim)) + cond_std = stats_cond.std(dim=reduce_dims, keepdim=True) + pred_std = stats_pred.std(dim=reduce_dims, keepdim=True) eps = torch.finfo(pred.dtype).eps factor = cond_std / pred_std.clamp_min(eps) factor = torch.where(pred_std > eps, factor, torch.ones_like(factor)) @@ -175,6 +188,8 @@ def combine_guided_x0( def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor: + if repeats == 1: + return tensor return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1)) @@ -278,6 +293,11 @@ class LTXGuidanceExecutor: if model_pass_count > 1: denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, model_pass_count) denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, model_pass_count) + if denoise_ctx.audio_attention_mask is not None: + denoise_ctx.audio_attention_mask = _repeat_batch( + denoise_ctx.audio_attention_mask, + model_pass_count, + ) return denoise_ctx @staticmethod @@ -285,10 +305,17 @@ class LTXGuidanceExecutor: ts: torch.Tensor, video_token_count: int, audio_token_count: int, + *, + expand_for_sequence_parallel: bool = False, ) -> dict[str, torch.Tensor]: + timestep = ts + audio_timestep = ts + if expand_for_sequence_parallel: + timestep = ts.reshape(-1, 1).expand(-1, video_token_count) + audio_timestep = ts.reshape(-1, 1).expand(-1, audio_token_count) return { - "timestep": ts.reshape(-1, 1).expand(-1, video_token_count), - "audio_timestep": ts.reshape(-1, 1).expand(-1, audio_token_count), + "timestep": timestep, + "audio_timestep": audio_timestep, "sigma": ts, "audio_sigma": ts, } @@ -322,7 +349,11 @@ class LTXGuidanceExecutor: guidance: LTXModalityGuidance, *, model_sigma: torch.Tensor | None = None, + preserve_positive_velocity: bool = False, + rescale_token_count: int | None = None, ) -> torch.Tensor: + if preserve_positive_velocity and tuple(splits) == ("cond",): + return splits["cond"] model_sigma = sigma if model_sigma is None else model_sigma # Official guidance reduces contiguous BSC tensors. The fused # transformer output may be channel-major after splitting, which @@ -334,6 +365,7 @@ class LTXGuidanceExecutor: uncond_perturbed=x0.get("ptb", 0.0), uncond_modality=x0.get("mod", 0.0), guidance=guidance, + rescale_token_count=rescale_token_count, ) return velocity_from_x0(sample, guided, sigma) @@ -346,6 +378,7 @@ class LTXGuidanceExecutor: state: LTXAVState, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, + preserve_positive_velocity: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: guidance_world_size = get_guidance_parallel_world_size() self.validate_guidance_world_size(plan, guidance_world_size) @@ -428,12 +461,15 @@ class LTXGuidanceExecutor: video_sigma, plan.spec.video, model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx), + preserve_positive_velocity=preserve_positive_velocity, ), self._guide_modality( state.audio, audio_splits, forward_ctx.audio_scheduler.sigmas[index], plan.spec.audio, + preserve_positive_velocity=preserve_positive_velocity, + rescale_token_count=forward_ctx.original_audio_num_frames, ), ) @@ -446,9 +482,19 @@ class LTXGuidanceExecutor: state: LTXAVState, forward_ctx: LTXForwardContext, denoise_ctx: LTXDenoiseContext, + preserve_positive_velocity: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if forward_ctx.guidance_parallel_ready: - return self.predict_parallel_guidance(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx) + return self.predict_parallel_guidance( + pipeline, + plan, + index, + timestep, + state, + forward_ctx, + denoise_ctx, + preserve_positive_velocity=preserve_positive_velocity, + ) prompt = forward_ctx.prompt_context video_contexts: list[torch.Tensor] = [] @@ -503,12 +549,15 @@ class LTXGuidanceExecutor: pipeline.scheduler.sigmas[index], denoise_ctx, ), + preserve_positive_velocity=preserve_positive_velocity, ), self._guide_modality( state.audio, audio_splits, forward_ctx.audio_scheduler.sigmas[index], plan.spec.audio, + preserve_positive_velocity=preserve_positive_velocity, + rescale_token_count=forward_ctx.original_audio_num_frames, ), ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py index bb0155e79..82de85504 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared latent layout and normalization primitives for LTX pipelines.""" @@ -20,6 +20,11 @@ class LTXAVState: audio: torch.Tensor +def official_video_token_layout(latents: torch.Tensor) -> torch.Tensor: + """Match the token-major view produced by the official LTX patchifier.""" + return latents.transpose(1, 2).contiguous().transpose(1, 2) + + def pack_latents( latents: torch.Tensor, patch_size: int = 1, @@ -64,6 +69,13 @@ def unpack_latents( return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) +def resolve_video_latent_statistics(pipeline: Any) -> tuple[torch.Tensor, torch.Tensor, float]: + """Return statistics owned by the active video decoder.""" + decoder = pipeline.diffusion_decoder if getattr(pipeline, "use_diffusion_decoder", False) else pipeline.vae + scaling_factor = getattr(getattr(decoder, "config", None), "scaling_factor", 1.0) + return decoder.latents_mean, decoder.latents_std, scaling_factor + + def normalize_latents( latents: torch.Tensor, latents_mean: torch.Tensor, @@ -117,7 +129,23 @@ def create_noised_state( device=latents.device, dtype=latents.dtype, ) - return noise_scale * noise + (1 - noise_scale) * latents + # Official GaussianNoisier blends in fp32, then materializes the next + # latent state in the model dtype. Batched request weights may arrive in + # the model dtype, so promote tensor weights alongside both endpoints. + lerp_weight = noise_scale.float() if isinstance(noise_scale, torch.Tensor) else noise_scale + return torch.lerp(latents.float(), noise.float(), lerp_weight).to(latents.dtype) + + +def create_conditioned_noised_state( + latents: torch.Tensor, + clean_latents: torch.Tensor, + denoise_mask: torch.Tensor, + noise_scale: float | torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None = None, +) -> torch.Tensor: + """Match the official conditioned Gaussian noisier operation order.""" + noised_latents = create_noised_state(latents, noise_scale, generator) + return torch.lerp(clean_latents.float(), noised_latents.float(), denoise_mask.float()).to(latents.dtype) def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor: @@ -135,6 +163,16 @@ def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: return latents[:, :num_frames] +def clear_audio_padding(latents: torch.Tensor, num_frames: int) -> torch.Tensor: + """Keep SP-only audio padding outside the logical sampler state.""" + if not 0 < num_frames <= latents.shape[1]: + raise ValueError(f"Audio frame count must be in [1, {latents.shape[1]}], got {num_frames}.") + if num_frames == latents.shape[1]: + return latents + padding = latents.new_zeros(latents.shape[0], latents.shape[1] - num_frames, latents.shape[2]) + return torch.cat([latents[:, :num_frames], padding], dim=1) + + def get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int: if sp_size > 1: audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size @@ -171,11 +209,12 @@ def prepare_video_latents( ) -> torch.Tensor: if latents is not None: if latents.ndim == 5: + latents_mean, latents_std, scaling_factor = resolve_video_latent_statistics(pipeline) latents = normalize_latents( latents, - pipeline.vae.latents_mean, - pipeline.vae.latents_std, - pipeline.vae.config.scaling_factor, + latents_mean, + latents_std, + scaling_factor, ) latents = pack_latents( latents, @@ -184,7 +223,8 @@ def prepare_video_latents( ) if latents.ndim != 3: raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") - return create_noised_state(latents, noise_scale, generator).to(device=device, dtype=dtype) + latents = create_noised_state(latents, noise_scale, generator).to(device=device, dtype=dtype) + return official_video_token_layout(latents) num_frames, height, width = resolve_video_latent_shape( height, @@ -205,7 +245,8 @@ def prepare_video_latents( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) - return randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return official_video_token_layout(latents) def prepare_audio_latents( @@ -226,6 +267,16 @@ def prepare_audio_latents( sp_size = getattr(pipeline.od_config.parallel_config, "sequence_parallel_size", 1) or 1 padded_latent_length = get_sp_padded_audio_latent_length(original_latent_length, int(sp_size)) + def pad_logical_latents(logical_latents: torch.Tensor) -> torch.Tensor: + if padded_latent_length == original_latent_length: + return logical_latents + padding = logical_latents.new_zeros( + logical_latents.shape[0], + padded_latent_length - original_latent_length, + logical_latents.shape[2], + ) + return torch.cat([logical_latents, padding], dim=1) + if latents is not None: if latents.ndim == 4: latents = pack_audio_latents(latents) @@ -237,29 +288,24 @@ def prepare_audio_latents( pipeline.audio_vae.latents_mean, pipeline.audio_vae.latents_std, ) - latents = create_noised_state(latents, noise_scale, generator) if latents.shape[1] not in {original_latent_length, padded_latent_length}: raise ValueError( "Provided `audio_latents` has incompatible audio frame count " f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}." ) - if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length: - padding = torch.zeros( - latents.shape[0], - padded_latent_length - original_latent_length, - latents.shape[2], - dtype=latents.dtype, - device=latents.device, - ) - latents = torch.cat([latents, padding], dim=1) + # Padding is an Omni implementation detail, not part of the official + # audio state. Draw request RNG only for logical tokens so later phases + # remain seed-invariant across SP degrees, then append deterministic 0s. + latents = create_noised_state(latents[:, :original_latent_length], noise_scale, generator) + latents = pad_logical_latents(latents) return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length - shape = (batch_size, padded_latent_length, num_channels_latents * latent_mel_bins) + shape = (batch_size, original_latent_length, num_channels_latents * latent_mel_bins) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError( f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" f" size of {batch_size}. Make sure the batch size matches the length of the generators." ) - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = pad_logical_latents(randn_tensor(shape, generator=generator, device=device, dtype=dtype)) return latents, original_latent_length, padded_latent_length diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py new file mode 100644 index 000000000..6b0433bbe --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py @@ -0,0 +1,437 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Construction and execution of a fixed LTX refinement-phase adapter. + +Checkpoint key parsing remains in :mod:`ltx2_adapter_parser`; this module +resolves the declared artifact and makes the existing Transformer linears +adapter-capable while retaining their quantization and tensor-parallel paths. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +import vllm.envs as envs +from vllm.distributed import ( + split_tensor_along_last_dim, + tensor_model_parallel_all_reduce, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.batch_invariant import linear_batch_invariant +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.utils import dispatch_unquantized_gemm +from vllm.platforms import current_platform + +from .ltx2_adapter_parser import AdapterManifest, AdapterTarget, iter_adapter_tensors, parse_ltx_adapter +from .ltx2_components import resolve_ltx_artifact +from .ltx2_recipes import LTX_DISTILLED_ADAPTER_SLOT + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class _AdapterLayout: + """Rank-local storage and execution layout for one LoRA pair.""" + + kind: str + input_size: int + output_size: int + a_input_start: int + a_input_size: int + b_output_start: int + b_output_size: int + output_start: int + reduce_results: bool = False + + +def _module_device(module: nn.Module) -> torch.device: + for tensor in (*module.parameters(), *module.buffers()): + if tensor.device.type != "meta": + return tensor.device + raise RuntimeError(f"Cannot materialize an LTX phase adapter for meta module {module._get_name()}.") + + +def _build_layout(layer: nn.Module, target: AdapterTarget) -> _AdapterLayout: + if isinstance(layer, QKVParallelLinear): + if target.slice not in {"q", "k", "v"}: + raise ValueError( + f"LTX packed QKV target {target.module!r} requires a q/k/v adapter slice, got {target.slice!r}." + ) + q_local = layer.num_heads * layer.head_size + k_local = layer.num_kv_heads * layer.head_size + v_local = layer.num_kv_heads * layer.v_head_size + if target.slice == "q": + output_size = layer.total_num_heads * layer.head_size + b_output_start = layer.tp_rank * q_local + b_output_size = q_local + output_start = 0 + elif target.slice == "k": + output_size = layer.total_num_kv_heads * layer.head_size + shard_rank = layer.tp_rank // layer.num_kv_head_replicas + b_output_start = shard_rank * k_local + b_output_size = k_local + output_start = q_local + else: + output_size = layer.total_num_kv_heads * layer.v_head_size + shard_rank = layer.tp_rank // layer.num_kv_head_replicas + b_output_start = shard_rank * v_local + b_output_size = v_local + output_start = q_local + k_local + return _AdapterLayout( + kind="qkv", + input_size=layer.input_size, + output_size=output_size, + a_input_start=0, + a_input_size=layer.input_size, + b_output_start=b_output_start, + b_output_size=b_output_size, + output_start=output_start, + ) + + if isinstance(layer, ColumnParallelLinear): + if layer.gather_output: + raise ValueError( + f"LTX phase adapter does not support gathered ColumnParallelLinear target {target.module!r}." + ) + return _AdapterLayout( + kind="column", + input_size=layer.input_size, + output_size=layer.output_size, + a_input_start=0, + a_input_size=layer.input_size, + b_output_start=layer.tp_rank * layer.output_size_per_partition, + b_output_size=layer.output_size_per_partition, + output_start=0, + ) + + if isinstance(layer, RowParallelLinear): + return _AdapterLayout( + kind="row", + input_size=layer.input_size, + output_size=layer.output_size, + a_input_start=layer.tp_rank * layer.input_size_per_partition, + a_input_size=layer.input_size_per_partition, + b_output_start=0, + b_output_size=layer.output_size, + output_start=0, + reduce_results=layer.reduce_results, + ) + + if isinstance(layer, nn.Linear): + return _AdapterLayout( + kind="replicated", + input_size=layer.in_features, + output_size=layer.out_features, + a_input_start=0, + a_input_size=layer.in_features, + b_output_start=0, + b_output_size=layer.out_features, + output_start=0, + ) + + raise TypeError(f"LTX phase adapter target {target.module!r} has unsupported layer type {type(layer).__name__}.") + + +class _AdapterPiece(nn.Module): + """One rank-local A/B pair stored as non-persistent buffers.""" + + def __init__(self, target: AdapterTarget, layout: _AdapterLayout) -> None: + super().__init__() + self.target = target + self.layout = layout + # The fixed graph is installed before base checkpoint loading. The + # empty buffers are replaced only after the base layer's quant/TP layout + # has settled and before offload/compile snapshots are taken. + self.register_buffer("lora_a", torch.empty(0), persistent=False) + self.register_buffer("lora_b", torch.empty(0), persistent=False) + + def load(self, lora_a: torch.Tensor, lora_b: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> None: + if tuple(lora_a.shape) != self.target.a_shape or tuple(lora_b.shape) != self.target.b_shape: + raise ValueError( + f"LTX adapter source shape changed for {self.target.source_module!r}: " + f"A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}." + ) + if lora_a.ndim != 2 or lora_b.ndim != 2: + raise ValueError(f"LTX adapter target {self.target.source_module!r} must use rank-2 A/B tensors.") + rank = self.target.a_shape[0] + if lora_a.shape[1] != self.layout.input_size or lora_b.shape[0] != self.layout.output_size: + raise ValueError( + f"LTX adapter shape does not match {self.target.module!r}: " + f"expected A=({rank}, {self.layout.input_size}), " + f"B=({self.layout.output_size}, {rank}); " + f"got A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}." + ) + + local_a = lora_a.narrow(1, self.layout.a_input_start, self.layout.a_input_size) + local_b = lora_b.narrow(0, self.layout.b_output_start, self.layout.b_output_size) + self.lora_a = local_a.to(device=device, dtype=dtype) + self.lora_b = local_b.to(device=device, dtype=dtype) + + def delta(self, input_: torch.Tensor) -> torch.Tensor: + if self.lora_a.numel() == 0 or self.lora_b.numel() == 0: + raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.") + if self.layout.kind == "row" and input_.shape[-1] != self.layout.a_input_size: + input_ = input_.narrow(-1, self.layout.a_input_start, self.layout.a_input_size) + input_ = input_.to(dtype=self.lora_a.dtype) + delta = F.linear(F.linear(input_, self.lora_a), self.lora_b) + if self.layout.kind == "row" and self.layout.reduce_results: + delta = tensor_model_parallel_all_reduce(delta) + return delta + + def weight_delta(self) -> torch.Tensor: + """Materialize the rank-local BF16 weight delta used by official fusion.""" + if self.lora_a.numel() == 0 or self.lora_b.numel() == 0: + raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.") + return torch.matmul(self.lora_b, self.lora_a) + + +def _apply_unquantized_gemm_with_weight( + layer: nn.Module, + input_: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + # GEMM dispatch is selected from the host platform, not the tensor device. + # In CPU tests on a ROCm host that selects rocm_unquantized_gemm, which has + # no CPU kernel. Keep that combination on PyTorch's portable implementation. + if current_platform.is_rocm() and input_.device.type == "cpu": + return F.linear(input_, weight, bias) + if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike(): + return linear_batch_invariant(input_, weight, bias) + return dispatch_unquantized_gemm()(layer, input_, weight, bias) + + +def _forward_with_weight( + layer: nn.Module, + input_: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]: + """Run an unquantized base layer with an ephemeral fused weight.""" + if isinstance(layer, nn.Linear): + return F.linear(input_, weight, layer.bias) + + quant_method = getattr(layer, "quant_method", None) + if not isinstance(quant_method, UnquantizedLinearMethod): + raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.") + + if isinstance(layer, RowParallelLinear): + if layer.input_is_parallel: + input_parallel = input_ + else: + input_parallel = split_tensor_along_last_dim(input_, num_partitions=layer.tp_size)[ + layer.tp_rank + ].contiguous() + bias = None if layer.tp_rank > 0 or layer.skip_bias_add else layer.bias + output_parallel = _apply_unquantized_gemm_with_weight(layer, input_parallel, weight, bias) + if layer.reduce_results and layer.tp_size > 1: + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + elif isinstance(layer, ColumnParallelLinear): + bias = layer.bias if not layer.skip_bias_add else None + output = _apply_unquantized_gemm_with_weight(layer, input_, weight, bias) + else: + raise TypeError(f"Unsupported LTX layer-fused linear type {type(layer).__name__}.") + + if not layer.return_bias: + return output + return output, layer.bias if layer.skip_bias_add else None + + +class _PhaseAdapterLinear(nn.Module): + """Delegate the base linear unchanged and add an enabled adapter slot.""" + + def __init__( + self, + base_layer: nn.Module, + targets: Iterable[AdapterTarget], + *, + adapter_dtype: torch.dtype, + layer_fused: bool = False, + ) -> None: + super().__init__() + self.base_layer = base_layer + self.adapter_dtype = adapter_dtype + if layer_fused and not ( + isinstance(base_layer, nn.Linear) + or isinstance(getattr(base_layer, "quant_method", None), UnquantizedLinearMethod) + ): + raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.") + self.layer_fused = layer_fused + pieces = [_AdapterPiece(target, _build_layout(base_layer, target)) for target in targets] + self.pieces = nn.ModuleList(pieces) + self._pieces_by_target = {piece.target: piece for piece in pieces} + self._enabled = False + + def load_target(self, target: AdapterTarget, lora_a: torch.Tensor, lora_b: torch.Tensor) -> None: + try: + piece = self._pieces_by_target[target] + except KeyError as exc: + raise ValueError(f"Adapter target {target.module!r} is not installed on this layer.") from exc + piece.load(lora_a, lora_b, device=_module_device(self.base_layer), dtype=self.adapter_dtype) + + def set_enabled(self, enabled: bool) -> None: + self._enabled = enabled + + def _add_adapter_delta(self, input_: torch.Tensor, output: torch.Tensor) -> torch.Tensor: + if not self._enabled: + return output + if output.requires_grad: + output = output.clone() + for piece in self.pieces: + delta = piece.delta(input_).to(dtype=output.dtype) + start = piece.layout.output_start + output[..., start : start + piece.layout.b_output_size].add_(delta) + return output + + @torch.no_grad() + def _fused_weight(self) -> torch.Tensor: + base_weight = self.base_layer.weight + fused_weight = base_weight.clone() + for piece in self.pieces: + start = piece.layout.output_start + base_slice = base_weight.narrow(0, start, piece.layout.b_output_size) + delta = piece.weight_delta() + if delta.shape != base_slice.shape: + raise ValueError( + f"LTX layer-fused delta shape {tuple(delta.shape)} does not match base weight slice " + f"{tuple(base_slice.shape)} for {piece.target.module!r}." + ) + # Match official fusion exactly: materialize B@A in the adapter + # dtype, then add the base tensor in-place before the main GEMM. + delta.add_(base_slice) + fused_weight.narrow(0, start, piece.layout.b_output_size).copy_(delta) + return fused_weight + + def forward(self, input_: torch.Tensor): + if not self._enabled: + return self.base_layer(input_) + if self.layer_fused: + return _forward_with_weight(self.base_layer, input_, self._fused_weight()) + output = self.base_layer(input_) + if isinstance(output, tuple): + return (self._add_adapter_delta(input_, output[0]), *output[1:]) + return self._add_adapter_delta(input_, output) + + +class LTXPhaseAdapterRuntime: + """Install and operate one fixed phase adapter on an existing transformer.""" + + def __init__( + self, + transformer: nn.Module, + manifest: AdapterManifest, + *, + dtype: torch.dtype, + layer_fused: bool = False, + ) -> None: + self.transformer = transformer + self.manifest = manifest + self._wrappers: dict[str, _PhaseAdapterLinear] = {} + self._target_names: tuple[str, ...] = () + self._materialized = False + grouped: dict[str, list[AdapterTarget]] = defaultdict(list) + for target in manifest.targets: + grouped[target.module].append(target) + + original_layers = {name: self.transformer.get_submodule(name) for name in grouped} + for name, targets in grouped.items(): + wrapper = _PhaseAdapterLinear( + original_layers[name], + targets, + adapter_dtype=dtype, + layer_fused=layer_fused, + ) + parent_name, _, child_name = name.rpartition(".") + parent = self.transformer.get_submodule(parent_name) if parent_name else self.transformer + if child_name not in parent._modules: + raise ValueError(f"LTX phase adapter target {name!r} is no longer a direct module child.") + parent._modules[child_name] = wrapper + self._wrappers[name] = wrapper + + self._target_names = tuple(sorted(self._wrappers, key=len, reverse=True)) + # LTX's custom loader receives the original checkpoint names. Keep the + # adaptation local to that loader rather than changing AutoWeightsLoader. + self.transformer._phase_adapter_parameter_name = self.base_parameter_name + logger.info( + "Installed %d LTX %s phase-adapter linears", + len(self._wrappers), + "layer-fused" if layer_fused else "dynamic", + ) + + def base_parameter_name(self, name: str) -> str: + for target_name in self._target_names: + prefix = f"{target_name}." + if name.startswith(prefix): + return f"{target_name}.base_layer.{name[len(prefix) :]}" + return name + + def finalize(self) -> None: + """Stream rank-local A/B buffers after base weight processing completes.""" + if self._materialized: + return + for target, lora_a, lora_b in iter_adapter_tensors(self.manifest): + self._wrappers[target.module].load_target(target, lora_a, lora_b) + self._materialized = True + logger.info("Materialized %d LTX phase-adapter tensor pairs", len(self.manifest.targets)) + + def activate(self, adapter_slot: str | None) -> None: + if adapter_slot is not None: + if adapter_slot != LTX_DISTILLED_ADAPTER_SLOT: + raise ValueError(f"Unknown LTX phase adapter slot {adapter_slot!r}.") + if not self._materialized: + raise RuntimeError("LTX phase adapter data must be materialized before activation.") + for wrapper in self._wrappers.values(): + wrapper.set_enabled(adapter_slot is not None) + + +def build_ltx_phase_adapter(pipeline: Any) -> LTXPhaseAdapterRuntime | None: + """Build the fixed adapter required by a multi-phase LTX recipe.""" + adapter_slots = {phase.adapter_slot for phase in pipeline.pipeline_recipe.phases if phase.adapter_slot is not None} + if not adapter_slots: + return None + if adapter_slots != {LTX_DISTILLED_ADAPTER_SLOT}: + raise ValueError(f"Unsupported LTX phase adapter slots: {sorted(adapter_slots)}.") + + profile = pipeline.component_profile + if profile.distilled_lora_filename is None or profile.artifact_repo_id is None: + raise ValueError(f"{profile.name} does not declare the required distilled adapter artifact.") + if getattr(pipeline.od_config, "lora_path", None) is not None: + raise ValueError( + f"{pipeline.__class__.__name__} reserves LoRA execution for its phase adapter; " + "request or static LoRA composition is not supported yet." + ) + + quantization_config = getattr(pipeline.od_config, "quantization_config", None) + layer_fused = quantization_config is None + adapter_path = resolve_ltx_artifact( + pipeline.od_config.model, + profile.artifact_repo_id, + profile.distilled_lora_filename, + model_revision=getattr(pipeline.od_config, "revision", None), + artifact_revision=profile.artifact_revision, + ) + + if layer_fused and pipeline.od_config.dtype is not torch.bfloat16: + raise ValueError("LTX layer-fused phase LoRA requires a bfloat16 pipeline dtype.") + + manifest = parse_ltx_adapter(pipeline.transformer, adapter_path) + runtime = LTXPhaseAdapterRuntime( + pipeline.transformer, + manifest, + dtype=pipeline.od_config.dtype, + layer_fused=layer_fused, + ) + return runtime diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py index a70f349ac..0caf93f64 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py @@ -3,7 +3,7 @@ """Declarative execution recipes for the LTX model family.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Literal from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance @@ -22,8 +22,14 @@ LTX_DEFAULT_NEGATIVE_PROMPT = ( "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." ) +LTX25_DEFAULT_NEGATIVE_PROMPT = ( + "has_subtitles, has_blurbox, transition from black, transition to black, speech_ending_short, " + + LTX_DEFAULT_NEGATIVE_PROMPT +) + LTX_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0) LTX_STAGE_2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0) +LTX_DISTILLED_ADAPTER_SLOT = "ltx_distilled" @dataclass(frozen=True) @@ -36,6 +42,8 @@ class LTXPhaseRecipe: sigmas: tuple[float, ...] | None = None noise_scale: float = 0.0 input_transform: Literal["initial", "spatial_upsample"] = "initial" + adapter_slot: str | None = None + sampler: Literal["euler", "euler_ancestral"] = "euler" allow_guidance_override: bool = True use_official_sigma_schedule: bool = True @@ -44,6 +52,8 @@ class LTXPhaseRecipe: raise ValueError("LTX phase spatial_downscale must be positive.") if self.sigmas is not None and len(self.sigmas) < 2: raise ValueError("An explicit LTX sigma schedule must contain at least two denoise sigmas.") + if self.sampler not in ("euler", "euler_ancestral"): + raise ValueError(f"Unsupported LTX sampler: {self.sampler!r}.") @property def num_inference_steps(self) -> int | None: @@ -64,9 +74,11 @@ class LTXPipelineRecipe: video_output_phase: int = -1 audio_output_phase: int = -1 allow_request_sigmas: bool = True + allow_request_phase_sigmas: bool = False allow_request_latents: bool = True allow_negative_prompt: bool = True fixed_num_inference_steps: bool = False + supports_cache_dit: bool = False def __post_init__(self) -> None: if not self.phases: @@ -107,13 +119,29 @@ def _official_guidance(stg_block: int) -> LTXGuidanceSpec: LTX2_ONE_STAGE_RECIPE = LTXPipelineRecipe( + supports_cache_dit=True, phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(29)),), ) LTX23_ONE_STAGE_RECIPE = LTXPipelineRecipe( + supports_cache_dit=True, num_inference_steps=30, phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(28)),), ) +LTX25_FULL_RECIPE = LTXPipelineRecipe( + height=544, + width=960, + num_inference_steps=30, + negative_prompt=LTX25_DEFAULT_NEGATIVE_PROMPT, + phases=( + LTXPhaseRecipe( + name="generate", + guidance=_official_guidance(28), + noise_scale=1.0, + ), + ), +) LTX_POSITIVE_ONLY_RECIPE = LTXPipelineRecipe( + supports_cache_dit=True, phases=( LTXPhaseRecipe( name="generate", @@ -157,11 +185,133 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe( fixed_num_inference_steps=True, ) +LTX25_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe( + height=1088, + width=1920, + num_inference_steps=len(LTX_DISTILLED_SIGMAS) - 1, + negative_prompt="", + phases=( + LTXPhaseRecipe( + name="generate_lowres", + guidance=LTXGuidanceSpec.positive_only(), + spatial_downscale=2, + sigmas=LTX_DISTILLED_SIGMAS, + noise_scale=1.0, + allow_guidance_override=False, + use_official_sigma_schedule=False, + sampler="euler_ancestral", + ), + LTXPhaseRecipe( + name="refine", + guidance=LTXGuidanceSpec.positive_only(), + sigmas=LTX_STAGE_2_DISTILLED_SIGMAS, + noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0], + input_transform="spatial_upsample", + allow_guidance_override=False, + use_official_sigma_schedule=False, + ), + ), + video_output_phase=1, + audio_output_phase=1, + allow_request_sigmas=False, + allow_request_phase_sigmas=True, + allow_request_latents=False, + allow_negative_prompt=False, + fixed_num_inference_steps=True, +) + +# LTX-2.3 full-distilled uses the same official fixed 8 + 3 sigma schedules +# and request contract. Its distinct component profile selects the 22B merged +# distilled Transformer, BWE vocoder, and matching spatial upsampler; no +# distilled LoRA is loaded for either phase. +LTX23_DISTILLED_TWO_STAGE_RECIPE = LTX2_DISTILLED_TWO_STAGE_RECIPE + + +def _distilled_one_stage_recipe( + two_stage_recipe: LTXPipelineRecipe, + *, + supports_cache_dit: bool = True, +) -> LTXPipelineRecipe: + """Run the merged-distilled checkpoint at its native Stage-1 resolution.""" + generate_phase = replace( + two_stage_recipe.phases[0], + name="generate", + spatial_downscale=1, + ) + return replace( + two_stage_recipe, + height=two_stage_recipe.height // 2, + width=two_stage_recipe.width // 2, + phases=(generate_phase,), + video_output_phase=0, + audio_output_phase=0, + allow_request_phase_sigmas=False, + supports_cache_dit=supports_cache_dit, + ) + + +LTX2_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe(LTX2_DISTILLED_TWO_STAGE_RECIPE) +LTX23_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe(LTX23_DISTILLED_TWO_STAGE_RECIPE) +LTX25_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe( + LTX25_DISTILLED_TWO_STAGE_RECIPE, + supports_cache_dit=False, +) + + +def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipelineRecipe: + return LTXPipelineRecipe( + height=one_stage_recipe.height * 2, + width=one_stage_recipe.width * 2, + num_frames=one_stage_recipe.num_frames, + frame_rate=one_stage_recipe.frame_rate, + num_inference_steps=one_stage_recipe.num_inference_steps, + negative_prompt=one_stage_recipe.negative_prompt, + phases=( + LTXPhaseRecipe( + name="generate_lowres", + guidance=one_stage_recipe.request_guidance, + spatial_downscale=2, + noise_scale=1.0, + ), + LTXPhaseRecipe( + name="refine", + guidance=LTXGuidanceSpec.positive_only(), + sigmas=LTX_STAGE_2_DISTILLED_SIGMAS, + noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0], + input_transform="spatial_upsample", + adapter_slot=LTX_DISTILLED_ADAPTER_SLOT, + allow_guidance_override=False, + use_official_sigma_schedule=False, + ), + ), + video_output_phase=1, + # The official second stage refines video only and deliberately + # discards its audio result. Decode the full-context Stage-1 audio. + audio_output_phase=0, + allow_request_sigmas=False, + allow_request_phase_sigmas=True, + allow_request_latents=False, + ) + + +LTX2_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX2_ONE_STAGE_RECIPE) +LTX23_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX23_ONE_STAGE_RECIPE) +LTX25_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX25_FULL_RECIPE) + _PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = { ("one_stage", "2"): LTX2_ONE_STAGE_RECIPE, ("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE, + ("one_stage", "2.5"): LTX25_FULL_RECIPE, + ("two_stage", "2"): LTX2_TWO_STAGE_RECIPE, + ("two_stage", "2.3"): LTX23_TWO_STAGE_RECIPE, + ("two_stage", "2.5"): LTX25_TWO_STAGE_RECIPE, + ("distilled_one_stage", "2"): LTX2_DISTILLED_ONE_STAGE_RECIPE, + ("distilled_one_stage", "2.3"): LTX23_DISTILLED_ONE_STAGE_RECIPE, + ("distilled_one_stage", "2.5"): LTX25_DISTILLED_ONE_STAGE_RECIPE, ("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE, + ("distilled_two_stage", "2.3"): LTX23_DISTILLED_TWO_STAGE_RECIPE, + ("distilled_two_stage", "2.5"): LTX25_DISTILLED_TWO_STAGE_RECIPE, ("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE, ("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE, } diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index e2eef19df..22230e486 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Request normalization and validation shared by LTX pipeline variants.""" from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass, replace -from typing import Any +from typing import Any, Literal import torch @@ -42,6 +43,7 @@ class LTXRequestInputs: output_type: str max_sequence_length: int audio_latents_normalized: bool = False + image_crf: int | None = None @property def guidance_scale(self) -> float: @@ -49,6 +51,34 @@ class LTXRequestInputs: return self.guidance.video.cfg_scale +LTXCheckpointKind = Literal["distilled", "regular"] + + +def validate_ltx_checkpoint( + scheduler_config: Mapping[str, Any], + *, + expected_kind: LTXCheckpointKind | None, + pipeline_name: str, +) -> None: + """Ensure scheduler metadata matches the checkpoint required by a profile.""" + if expected_kind is None: + return + + if expected_kind == "distilled": + if scheduler_config.get("use_dynamic_shifting", True) or scheduler_config.get("shift_terminal") is not None: + raise ValueError( + f"{pipeline_name} requires a merged distilled LTX2 checkpoint; " + "the loaded scheduler metadata describes a non-distilled checkpoint." + ) + return + + if not scheduler_config.get("use_dynamic_shifting", False) or scheduler_config.get("shift_terminal") is None: + raise ValueError( + f"{pipeline_name} requires a regular non-distilled LTX checkpoint; " + "the loaded scheduler metadata describes a distilled checkpoint." + ) + + def validate_pipeline_request( request_inputs: LTXRequestInputs, *, @@ -56,7 +86,9 @@ def validate_pipeline_request( vae_spatial_compression_ratio: int, vae_temporal_compression_ratio: int, pipeline_name: str, + min_video_latent_spatial_size: tuple[int, int] | None = None, request_sigmas: list[float] | None = None, + request_phase_sigmas: tuple[list[float] | None, ...] | None = None, ) -> None: """Validate the request capabilities declared by an execution recipe.""" alignment = vae_spatial_compression_ratio * pipeline_recipe.max_spatial_downscale @@ -66,6 +98,18 @@ def validate_pipeline_request( f"{request_inputs.width}x{request_inputs.height}." ) + if min_video_latent_spatial_size is not None: + min_latent_height, min_latent_width = min_video_latent_spatial_size + latent_height = request_inputs.height // vae_spatial_compression_ratio + latent_width = request_inputs.width // vae_spatial_compression_ratio + if latent_height < min_latent_height or latent_width < min_latent_width: + raise ValueError( + f"{pipeline_name} DiffVAE requires video latent spatial dimensions of at least " + f"{min_latent_width}x{min_latent_height} for neighborhood attention, got " + f"{latent_width}x{latent_height} from request resolution " + f"{request_inputs.width}x{request_inputs.height}." + ) + if request_inputs.num_frames < 1 or (request_inputs.num_frames - 1) % vae_temporal_compression_ratio != 0: raise ValueError( f"{pipeline_name} `num_frames` must be {vae_temporal_compression_ratio} * k + 1, " @@ -79,9 +123,28 @@ def validate_pipeline_request( if not pipeline_recipe.allow_request_sigmas and request_sigmas is not None: raise ValueError(f"{pipeline_name} uses fixed phase sigma schedules.") + if request_sigmas is not None and len(request_sigmas) < 2: + raise ValueError("A custom LTX sigma schedule must contain at least two values.") - if pipeline_recipe.fixed_num_inference_steps and ( - request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps + has_phase_sigmas = request_phase_sigmas is not None and any(item is not None for item in request_phase_sigmas) + if request_sigmas is not None and has_phase_sigmas: + raise ValueError("Use either `sigmas` or per-stage sigma schedules, not both.") + if has_phase_sigmas and not pipeline_recipe.allow_request_phase_sigmas: + raise ValueError(f"{pipeline_name} does not accept per-stage sigma schedules.") + if request_phase_sigmas is not None and len(request_phase_sigmas) != len(pipeline_recipe.phases): + raise ValueError( + f"{pipeline_name} has {len(pipeline_recipe.phases)} phases but received " + f"{len(request_phase_sigmas)} per-stage sigma schedules." + ) + for index, phase_sigmas in enumerate(request_phase_sigmas or ()): + if phase_sigmas is not None and len(phase_sigmas) < 2: + raise ValueError(f"`stage_{index + 1}_sigmas` must contain at least two values.") + + if ( + request_sigmas is None + and (request_phase_sigmas is None or request_phase_sigmas[0] is None) + and pipeline_recipe.fixed_num_inference_steps + and (request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps) ): raise ValueError( f"{pipeline_name} uses {pipeline_recipe.num_inference_steps} fixed Stage 1 denoise steps; " @@ -136,6 +199,16 @@ def _get_extra_arg(sampling: Any, *names: str) -> Any: return None +def _normalize_image_crf(value: Any) -> int | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError("`image_crf` must be an integer from 0 through 51.") + if not 0 <= value <= 51: + raise ValueError("`image_crf` must be an integer from 0 through 51.") + return value + + def _normalize_stg_blocks(value: Any, default: tuple[int, ...]) -> tuple[int, ...]: if value is None: return default @@ -245,6 +318,7 @@ class LTXRequestMixin: req: DiffusionRequestBatch, *, image: Any | None = None, + image_crf: int | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, height: int | None = None, @@ -253,6 +327,8 @@ class LTXRequestMixin: frame_rate: float | None = None, num_inference_steps: int | None = None, sigmas: list[float] | None = None, + stage_1_sigmas: list[float] | None = None, + stage_2_sigmas: list[float] | None = None, timesteps: list[int] | None = None, guidance_scale: float | None = None, guidance_rescale: float | None = None, @@ -286,6 +362,7 @@ class LTXRequestMixin: return self._forward_request( req, image=image, + image_crf=image_crf, prompt=prompt, negative_prompt=negative_prompt, height=height, @@ -294,6 +371,8 @@ class LTXRequestMixin: frame_rate=frame_rate, num_inference_steps=num_inference_steps, sigmas=sigmas, + stage_1_sigmas=stage_1_sigmas, + stage_2_sigmas=stage_2_sigmas, guidance_scale=guidance_scale, num_videos_per_prompt=num_videos_per_prompt, generator=generator, @@ -314,12 +393,30 @@ class LTXRequestMixin: req: DiffusionRequestBatch, fallback: list[float] | None, ) -> list[float] | None: - request_sigmas = [sampling.sigmas for sampling in req.sampling_params_list] + request_sigmas = [ + sampling.sigmas if sampling.sigmas is not None else _get_extra_arg(sampling, "sigmas") + for sampling in req.sampling_params_list + ] first = request_sigmas[0] if request_sigmas else None if any(item != first for item in request_sigmas[1:]): raise ValueError("Batched LTX requests must use identical custom sigmas.") return first if first is not None else fallback + @staticmethod + def _resolve_request_phase_sigmas( + req: DiffusionRequestBatch, + stage_1_fallback: list[float] | None, + stage_2_fallback: list[float] | None, + ) -> tuple[list[float] | None, list[float] | None] | None: + resolved: list[list[float] | None] = [] + for name, fallback in (("stage_1_sigmas", stage_1_fallback), ("stage_2_sigmas", stage_2_fallback)): + values = [_get_extra_arg(sampling, name) for sampling in req.sampling_params_list] + first = values[0] if values else None + if any(item != first for item in values[1:]): + raise ValueError(f"Batched LTX requests must use identical `{name}` values.") + resolved.append(first if first is not None else fallback) + return tuple(resolved) if any(item is not None for item in resolved) else None + def check_inputs( self, prompt: str | list[str] | None, @@ -379,6 +476,7 @@ class LTXRequestMixin: decode_noise_scale: float | list[float] | None, output_type: str, max_sequence_length: int | None, + image_crf: int | None = None, ) -> LTXRequestInputs: sampling_params_list = req.sampling_params_list for sampling in sampling_params_list: @@ -392,6 +490,12 @@ class LTXRequestMixin: if _get_extra_arg(sampling, "flow_shift") is not None: raise ValueError("LTX does not support `flow_shift`; use `sigmas` for a custom schedule.") + request_image_crfs = [_get_extra_arg(item, "image_crf") for item in sampling_params_list] + first_image_crf = request_image_crfs[0] if request_image_crfs else None + if any(item != first_image_crf for item in request_image_crfs[1:]): + raise ValueError("Batched LTX requests must use identical `image_crf` values.") + image_crf = _normalize_image_crf(first_image_crf if first_image_crf is not None else image_crf) + sampling = sampling_params_list[0] is_dummy_run = req.is_dummy_run() prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt @@ -522,4 +626,5 @@ class LTXRequestMixin: decode_noise_scale=decode_noise_scale, output_type=output_type, max_sequence_length=int(max_sequence_length), + image_crf=image_crf, ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 4b74ebcb8..fc589fcb8 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared recipe-driven runtime for LTX pipeline variants.""" @@ -29,6 +29,7 @@ from vllm_omni.platforms import current_omni_platform from . import ltx2_latents as latent_ops from .ltx2_components import ( LTXComponentProfile, + _ltx2_use_diffusion_decoder, detect_ltx_model_version, initialize_pipeline_components, resolve_ltx_component_profile, @@ -47,6 +48,7 @@ from .ltx2_guidance import ( LTXGuidanceExecutor, LTXGuidancePlan, ) +from .ltx2_phase_adapter import LTXPhaseAdapterRuntime, build_ltx_phase_adapter from .ltx2_recipes import ( LTXPhaseRecipe, LTXPipelineRecipe, @@ -59,6 +61,29 @@ from .ltx2_request import ( ) +def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.Tensor: + """Run the BWE vocoder in FP32, matching the official LTX pipeline.""" + if not hasattr(vocoder, "bwe_generator"): + return vocoder(generated_mel) + + input_dtype = generated_mel.dtype + device_type = generated_mel.device.type + module_dtype = next(vocoder.parameters()).dtype + if device_type == "mps": + if module_dtype != torch.float32: + vocoder.float() + try: + return vocoder(generated_mel.float()).to(input_dtype) + finally: + if module_dtype != torch.float32: + vocoder.to(module_dtype) + + # BF16 errors compound through the BWE model's long convolution stack. + # FP32 autocast upcasts each op without materializing a second FP32 model. + with torch.autocast(device_type=device_type, dtype=torch.float32): + return vocoder(generated_mel.float()).to(input_dtype) + + def _expand_per_prompt_decode_value( value: float | list[float], *, @@ -137,19 +162,34 @@ class LTXRuntime( def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: del prefix - self.model_version = detect_ltx_model_version(od_config.model) + parallel_config = getattr(od_config, "parallel_config", None) + if getattr(parallel_config, "ulysses_mode", "strict") == "advanced_uaa": + raise ValueError( + f"{self.__class__.__name__} does not support ulysses_mode='advanced_uaa'. " + "Use the default ulysses_mode='strict' for LTX sequence parallelism." + ) + self.model_version = detect_ltx_model_version(od_config.model, revision=getattr(od_config, "revision", None)) + self.use_diffusion_decoder = _ltx2_use_diffusion_decoder(od_config, self.model_version) self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version) self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version) + if getattr(od_config, "cache_backend", "none") == "cache_dit" and not self.pipeline_recipe.supports_cache_dit: + raise ValueError( + f"{self.__class__.__name__} does not support cache_backend='cache_dit'. " + "Cache-DiT is not qualified for this LTX recipe." + ) self._dit_modules = list(self.component_profile.dit_modules) self._encoder_modules = list(self.component_profile.encoder_modules) self._vae_modules = list(self.component_profile.vae_modules) + if self.use_diffusion_decoder: + self._vae_modules.append("diffusion_decoder") self._resident_modules = list(self.component_profile.resident_modules) - if self.model_version == "2.3": + if self.model_version in ("2.3", "2.5"): self.preserve_sp_padded_audio_duration = True self.reports_stage_durations = True super().__init__() self._guidance_plan = LTXGuidancePlan.build(self.pipeline_recipe.request_guidance) initialize_pipeline_components(self, od_config) + self._phase_adapter: LTXPhaseAdapterRuntime | None = build_ltx_phase_adapter(self) self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) @@ -159,6 +199,7 @@ class LTXRuntime( req: DiffusionRequestBatch, *, image: Any | None = None, + image_crf: int | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, height: int | None = None, @@ -167,6 +208,8 @@ class LTXRuntime( frame_rate: float | None = None, num_inference_steps: int | None = None, sigmas: list[float] | None = None, + stage_1_sigmas: list[float] | None = None, + stage_2_sigmas: list[float] | None = None, guidance_scale: float | None = None, num_videos_per_prompt: int | None = 1, generator: torch.Generator | list[torch.Generator] | None = None, @@ -203,18 +246,41 @@ class LTXRuntime( decode_noise_scale=decode_noise_scale, output_type=output_type, max_sequence_length=max_sequence_length, + image_crf=image_crf, ) image = self._resolve_request_image(req, image, request_inputs) + if image is not None and not self.support_image_input: + raise ValueError(f"{self.__class__.__name__} does not support `image` input.") request_sigmas = self._resolve_request_sigmas(req, sigmas) + request_phase_sigmas = self._resolve_request_phase_sigmas(req, stage_1_sigmas, stage_2_sigmas) + min_video_latent_spatial_size = None + if getattr(self, "use_diffusion_decoder", False): + first_stage_kernel = self.diffusion_decoder.config.decoder_stage_kernels[0] + min_video_latent_spatial_size = (int(first_stage_kernel[-2]), int(first_stage_kernel[-1])) validate_pipeline_request( request_inputs, pipeline_recipe=self.pipeline_recipe, vae_spatial_compression_ratio=self.vae_spatial_compression_ratio, vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, pipeline_name=self.__class__.__name__, + min_video_latent_spatial_size=min_video_latent_spatial_size, request_sigmas=request_sigmas, + request_phase_sigmas=request_phase_sigmas, + ) + phase_adapter = getattr(self, "_phase_adapter", None) + if phase_adapter is not None and any( + getattr(sampling, "lora_request", None) is not None for sampling in req.sampling_params_list + ): + raise ValueError( + f"{self.__class__.__name__} cannot compose a request LoRA with its internal phase adapter." + ) + return self._run_recipe( + req, + request_inputs, + request_sigmas=request_sigmas, + request_phase_sigmas=request_phase_sigmas, + image=image, ) - return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image) def _run_recipe( self, @@ -222,24 +288,39 @@ class LTXRuntime( request_inputs: LTXRequestInputs, *, request_sigmas: list[float] | None, + request_phase_sigmas: tuple[list[float] | None, ...] | None = None, image: Any | None = None, ) -> DiffusionOutput | list[DiffusionOutput]: """Execute one- and multi-phase recipes through the same control flow.""" phase_results: list[LTXPhaseResult] = [] prompt_context = None - - for phase_recipe in self.pipeline_recipe.phases: + for phase_index, phase_recipe in enumerate(self.pipeline_recipe.phases): + override_sigmas = None if request_phase_sigmas is None else request_phase_sigmas[phase_index] + phase_sigmas = ( + override_sigmas + if override_sigmas is not None + else ( + request_sigmas + if request_sigmas is not None + else (list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else None) + ) + ) self._enter_phase(phase_recipe) phase_inputs = self._build_phase_inputs( request_inputs, phase_recipe, phase_results[-1] if phase_results else None, ) + if phase_sigmas is not None: + phase_inputs = replace(phase_inputs, num_inference_steps=len(phase_sigmas) - 1) + noise_scale = phase_recipe.noise_scale + if override_sigmas is not None and phase_recipe.input_transform == "spatial_upsample": + noise_scale = float(override_sigmas[0]) phase_result = self.run_phase( req, phase_inputs, - noise_scale=phase_recipe.noise_scale, - sigmas=list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else request_sigmas, + noise_scale=noise_scale, + sigmas=phase_sigmas, timesteps=None, attention_kwargs=None, phase_recipe=phase_recipe, @@ -258,8 +339,21 @@ class LTXRuntime( return self.decode_phase(output_phase) def _enter_phase(self, phase: LTXPhaseRecipe) -> None: - """Hook for a future phase-weight strategy.""" self._active_phase_name = phase.name + self._guidance_plan = LTXGuidancePlan.build(phase.guidance) + phase_adapter = getattr(self, "_phase_adapter", None) + if phase_adapter is None: + if phase.adapter_slot is not None: + raise RuntimeError(f"LTX phase {phase.name!r} requires adapter slot {phase.adapter_slot!r}.") + return + phase_adapter.activate(phase.adapter_slot) + + def eval(self): + result = super().eval() + phase_adapter = getattr(self, "_phase_adapter", None) + if phase_adapter is not None: + phase_adapter.finalize() + return result def prepare_latents( self, @@ -457,7 +551,10 @@ class LTXRuntime( return self._make_output((latents, audio_latents)) latents = latents.to(connector_prompt_embeds.dtype) - if not self.vae.config.timestep_conditioning: + use_diffusion_decoder = getattr(self, "use_diffusion_decoder", False) + if use_diffusion_decoder: + timestep_decode = None + elif not self.vae.config.timestep_conditioning: timestep_decode = None else: noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) @@ -474,14 +571,22 @@ class LTXRuntime( dist_initialized = torch.distributed.is_initialized() is_output_rank = not dist_initialized or torch.distributed.get_rank() == 0 vae_decode_needs_all_ranks = False - is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None) + video_decoder = self.diffusion_decoder if use_diffusion_decoder else self.vae + is_distributed_vae_enabled = getattr(video_decoder, "is_distributed_enabled", None) if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled): # Distributed tiled decode is collective, so every rank must enter it. vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled()) should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks if should_decode_video: - video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0] + if use_diffusion_decoder: + video = self.diffusion_decoder.decode( + latents.to(self.diffusion_decoder.dtype), + generator=generator, + return_dict=False, + )[0] + else: + video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0] else: video = torch.empty(0, device=latents.device, dtype=latents.dtype) @@ -496,7 +601,7 @@ class LTXRuntime( if video.numel() > 0: video = self.video_processor.postprocess_video(video, output_type=output_type) generated_mel = self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0] - audio = self.vocoder(generated_mel) + audio = _run_ltx_vocoder(self.vocoder, generated_mel) return self._make_output((video, audio)) def _prepare_video_latents_stage( @@ -610,7 +715,12 @@ class LTXRuntime( audio_token_count: int, ) -> dict[str, torch.Tensor]: del forward_ctx, denoise_ctx - return self.guidance_executor.timestep_kwargs(ts, video_token_count, audio_token_count) + return self.guidance_executor.timestep_kwargs( + ts, + video_token_count, + audio_token_count, + expand_for_sequence_parallel=True, + ) def _video_guidance_model_sigma( self, @@ -665,6 +775,7 @@ class LTXRuntime( state, forward_ctx, denoise_ctx, + preserve_positive_velocity=forward_ctx.sampler == "euler_ancestral", ) def _denoise_step( @@ -692,6 +803,7 @@ class LTXRuntime( noise_pred_audio, timestep, ) + audio = latent_ops.clear_audio_padding(audio, forward_ctx.original_audio_num_frames) return latent_ops.LTXAVState(video=video, audio=audio) def _unpack_and_denormalize_stage( @@ -708,11 +820,12 @@ class LTXRuntime( self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ) + latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self) latents = latent_ops.denormalize_latents( latents, - self.vae.latents_mean, - self.vae.latents_std, - self.vae.config.scaling_factor, + latents_mean, + latents_std, + scaling_factor, ) audio_latents = latent_ops.unpad_audio_latents(audio_latents, forward_ctx.original_audio_num_frames) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 95724bd19..0b0e7fc08 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -55,6 +55,18 @@ logger = init_logger(__name__) _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters +def apply_keyframes_absolute_embedding( + hidden_states: torch.Tensor, + keyframes_mask: torch.Tensor | None, + embedding: torch.Tensor | None, +) -> torch.Tensor: + """Add the official learned marker to first-frame/keyframe tokens.""" + if keyframes_mask is None or embedding is None: + return hidden_states + mask = (keyframes_mask > 0).to(dtype=hidden_states.dtype) + return hidden_states + mask * embedding.to(device=hidden_states.device, dtype=hidden_states.dtype) + + def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> nn.Module: """Bridge diffusers' RMSNorm API onto vLLM's `has_weight` variant.""" if not elementwise_affine: @@ -82,15 +94,10 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: cos, sin = freqs - if cos.dtype != x.dtype: - cos = cos.to(dtype=x.dtype) - if sin.dtype != x.dtype: - sin = sin.to(dtype=x.dtype) # Concrete pair count instead of -1 keeps SDPA shape static under torch.compile. x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1) # [B, S, C // 2] x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2) - out = x * cos + x_rotated * sin - return out + return x * cos + x_rotated * sin def apply_split_rotary_emb( @@ -101,12 +108,6 @@ def apply_split_rotary_emb( ) -> torch.Tensor: # `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile. cos, sin = freqs - if cos.dtype != x.dtype: - cos = cos.to(dtype=x.dtype) - if sin.dtype != x.dtype: - sin = sin.to(dtype=x.dtype) - - x_dtype = x.dtype needs_reshape = False if x.ndim != 4 and cos.ndim == 4: # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head) @@ -142,7 +143,6 @@ def apply_split_rotary_emb( if needs_reshape: out = out.swapaxes(1, 2).reshape(b, t, h * head_dim) - out = out.to(dtype=x_dtype) return out @@ -269,6 +269,7 @@ class LTX2FeedForward(nn.Module): RowParallelLinear( inner_dim, dim_out, + bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config, @@ -355,6 +356,19 @@ def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor: return attention_mask +class _LTX2ParallelAttention(Attention): + """Preserve LTX SP collectives while applying model-specific padding masks.""" + + def _run_local_attention(self, query, key, value, attn_metadata): + attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None + if attention_mask is not None: + # LTX masks describe K/V padding. Keep this semantic local to LTX: + # Flash varlen applies a 2D mask to both Q and K, while SP uses a + # broadcast 4D key mask that only SDPA consumes correctly. + return self.sdpa_fallback.forward(query, key, value, attn_metadata) + return super()._run_local_attention(query, key, value, attn_metadata) + + class LTX2AudioVideoAttnProcessor: r""" Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model. @@ -396,11 +410,10 @@ class LTX2AudioVideoAttnProcessor: return None if self._is_sp_enabled(): - # In SP, Ulysses expects a 2D padding mask that matches query length. - # For cross-attention, encoder sequence length != query length, so drop the mask. - if encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]: - return None - return to_ltx_padding_mask(attention_mask) + # Keep the full key mask replicated while strict Ulysses reshards + # Q/K/V. The broadcast Q dimension also works when cross-attention + # has different Q and K lengths, without changing generic Ulysses. + return to_ltx_padding_mask(attention_mask)[:, None, None, :] attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) @@ -505,8 +518,8 @@ class LTX2AudioVideoAttnProcessor: is_self_attention=is_self_attention, ) - query = attn.norm_q(query) - key = attn.norm_k(key) + query = attn.norm_q(query).to(dtype=value.dtype) + key = attn.norm_k(key).to(dtype=value.dtype) if query_rotary_emb is not None: query_rotary_emb = self._slice_rope_for_tp(query_rotary_emb, attn) @@ -525,6 +538,10 @@ class LTX2AudioVideoAttnProcessor: head_dim=attn.head_dim, ) + # Keep FP32 RoPE computation while satisfying attention dtype contracts. + query = query.to(dtype=value.dtype) + key = key.to(dtype=value.dtype) + query = query.unflatten(2, (attn.heads, attn.head_dim)) key = key.unflatten(2, (attn.heads, attn.head_dim)) value_heads = value.unflatten(2, (attn.heads, attn.head_dim)) @@ -696,7 +713,7 @@ class LTX2Attention(torch.nn.Module): torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(), ] ) - self.attn = Attention( + self.attn = _LTX2ParallelAttention( num_heads=self.query_num_heads, head_size=dim_head, num_kv_heads=self.kv_num_heads, @@ -709,7 +726,15 @@ class LTX2Attention(torch.nn.Module): # LTX-2.3: per-head gated attention # leave unquantized for this linear if apply_gated_attention: - self.to_gate_logits = nn.Linear(query_dim, self.query_num_heads, bias=True) + self.to_gate_logits = ColumnParallelLinear( + query_dim, + heads, + bias=True, + gather_output=False, + return_bias=False, + quant_config=None, + prefix=f"{prefix}.to_gate_logits" if prefix else "to_gate_logits", + ) else: self.to_gate_logits = None @@ -826,6 +851,8 @@ class LTX2VideoTransformerBlock(nn.Module): elementwise_affine: bool = False, rope_type: str = "interleaved", perturbed_attn: bool = False, + ff_bias: bool = True, + audio_ff_bias: bool = True, quant_config: "QuantizationConfig | None" = None, prefix: str = "", ): @@ -945,6 +972,7 @@ class LTX2VideoTransformerBlock(nn.Module): self.ff = LTX2FeedForward( dim, activation_fn=activation_fn, + bias=ff_bias, quant_config=quant_config, prefix=f"{prefix}.ff" if prefix else "ff", ) @@ -953,6 +981,7 @@ class LTX2VideoTransformerBlock(nn.Module): self.audio_ff = LTX2FeedForward( audio_dim, activation_fn=activation_fn, + bias=audio_ff_bias, quant_config=quant_config, prefix=f"{prefix}.audio_ff" if prefix else "audio_ff", ) @@ -1058,21 +1087,34 @@ class LTX2VideoTransformerBlock(nn.Module): audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa # 2. Video and Audio Cross-Attention with text embeddings (Q: Video/Audio; K,V: Text) - # LTX-2.3: compute prompt modulation params for K/V - if self.cross_attn_adaln and temb_prompt is not None: - video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size) - shift_text_kv, scale_text_kv = video_prompt_ada_params + if self.cross_attn_adaln: + if temb_prompt is not None: + shift_text_kv, scale_text_kv = self.get_mod_params( + self.prompt_scale_shift_table, temb_prompt, batch_size + ) + else: + shift_text_kv, scale_text_kv = ( + self.prompt_scale_shift_table[None, None] + .to(device=hidden_states.device, dtype=hidden_states.dtype) + .unbind(dim=2) + ) - audio_prompt_ada_params = self.get_mod_params( - self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size - ) - audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params + if temb_prompt_audio is not None: + audio_shift_text_kv, audio_scale_text_kv = self.get_mod_params( + self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size + ) + else: + audio_shift_text_kv, audio_scale_text_kv = ( + self.audio_prompt_scale_shift_table[None, None] + .to(device=audio_hidden_states.device, dtype=audio_hidden_states.dtype) + .unbind(dim=2) + ) # 2.1. Video-Text Cross-Attention norm_hidden_states = self.norm2(hidden_states) if self.video_cross_attn_adaln: norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q - if self.cross_attn_adaln and temb_prompt is not None: + if self.cross_attn_adaln: encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv attn_hidden_states = self.attn2( @@ -1089,7 +1131,7 @@ class LTX2VideoTransformerBlock(nn.Module): norm_audio_hidden_states = self.audio_norm2(audio_hidden_states) if self.audio_cross_attn_adaln: norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_text_q) + audio_shift_text_q - if self.cross_attn_adaln and temb_prompt is not None: + if self.cross_attn_adaln: audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv attn_audio_hidden_states = self.audio_attn2( @@ -1488,7 +1530,10 @@ class LTX2VideoTransformer3DModel(nn.Module): block_forward_patterns={ "transformer_blocks": ForwardPattern.Pattern_0, }, - has_separate_cfg=True, + # LTX fuses all guidance passes into the batch dimension and invokes + # the transformer once per denoise step. Cache-DiT must therefore + # advance its step counter for every transformer forward. + has_separate_cfg=False, check_forward_pattern=False, ) @@ -1498,6 +1543,14 @@ class LTX2VideoTransformer3DModel(nn.Module): _layerwise_offload_blocks_attrs = ["transformer_blocks"] _hsdp_shard_conditions = [is_transformer_block_module] _sp_plan: dict[str, Any] | None = None + stacked_params_mapping = ( + (".attn1.to_qkv", ".attn1.to_q", "q"), + (".attn1.to_qkv", ".attn1.to_k", "k"), + (".attn1.to_qkv", ".attn1.to_v", "v"), + (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"), + (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"), + (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"), + ) packed_modules_mapping = { "to_qkv": ["to_q", "to_k", "to_v"], } @@ -1518,6 +1571,7 @@ class LTX2VideoTransformer3DModel(nn.Module): # Shard video/audio latents across sequence "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), "audio_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), + "keyframes_mask": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), # Shard prompt embeds across sequence "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), "audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False), @@ -1592,6 +1646,10 @@ class LTX2VideoTransformer3DModel(nn.Module): cross_attn_mod: bool = False, audio_gated_attn: bool = False, audio_cross_attn_mod: bool = False, + ff_bias: bool = True, + audio_ff_bias: bool = True, + use_prompt_adaln_single: bool = True, + use_keyframes_abs_pos_embedding: bool = False, quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() @@ -1638,12 +1696,19 @@ class LTX2VideoTransformer3DModel(nn.Module): timestep_scale_multiplier=timestep_scale_multiplier, cross_attn_timestep_scale_multiplier=cross_attn_timestep_scale_multiplier, rope_type=rope_type, + ff_bias=ff_bias, + audio_ff_bias=audio_ff_bias, + use_prompt_adaln_single=use_prompt_adaln_single, + use_keyframes_abs_pos_embedding=use_keyframes_abs_pos_embedding, ) # 1. Patchification input projections self.proj_in = nn.Linear(in_channels, inner_dim) self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim) + if use_keyframes_abs_pos_embedding: + self.keyframes_abs_pos_embedding = nn.Parameter(torch.zeros(1, inner_dim)) + # 2. Prompt embeddings # LTX-2 (use_prompt_embeddings=True): caption projection in the transformer # LTX-2.3 (use_prompt_embeddings=False): caption projection in the connectors @@ -1667,7 +1732,7 @@ class LTX2VideoTransformer3DModel(nn.Module): # 3.3. LTX-2.3: Prompt modulation from sigma self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod - if self.prompt_modulation: + if self.prompt_modulation and use_prompt_adaln_single: self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False) self.audio_prompt_adaln = LTX2AdaLayerNormSingle( audio_inner_dim, num_mod_params=2, use_additional_conditions=False @@ -1787,6 +1852,8 @@ class LTX2VideoTransformer3DModel(nn.Module): elementwise_affine=norm_elementwise_affine, rope_type=rope_type, perturbed_attn=perturbed_attn, + ff_bias=ff_bias, + audio_ff_bias=audio_ff_bias, quant_config=quant_config, prefix=f"transformer_blocks.{layer_idx}", ) @@ -1824,10 +1891,12 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, audio_timestep: torch.LongTensor | None = None, + keyframes_mask: torch.Tensor | None = None, sigma: torch.Tensor | None = None, audio_sigma: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, + audio_attention_mask: torch.Tensor | None = None, num_frames: int | None = None, height: int | None = None, width: int | None = None, @@ -1861,6 +1930,8 @@ class LTX2VideoTransformer3DModel(nn.Module): Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`. audio_encoder_attention_mask (`torch.Tensor`, *optional*): Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling. + audio_attention_mask (`torch.Tensor`, *optional*): + Optional audio-token key padding mask shared by audio self-attention and audio-to-video attention. num_frames (`int`, *optional*): The number of latent video frames. Used if calculating the video coordinates for RoPE. height (`int`, *optional*): @@ -1921,9 +1992,7 @@ class LTX2VideoTransformer3DModel(nn.Module): video_rotary_emb = self.rope(video_coords, device=hidden_states.device, out_dtype=hidden_states.dtype) audio_rotary_emb = self.audio_rope( - audio_coords, - device=audio_hidden_states.device, - out_dtype=audio_hidden_states.dtype, + audio_coords, device=audio_hidden_states.device, out_dtype=audio_hidden_states.dtype ) video_cross_attn_rotary_emb = self.cross_attn_rope( @@ -1939,6 +2008,11 @@ class LTX2VideoTransformer3DModel(nn.Module): # 2. Patchify input projections hidden_states = self.proj_in(hidden_states) + hidden_states = apply_keyframes_absolute_embedding( + hidden_states, + keyframes_mask, + getattr(self, "keyframes_abs_pos_embedding", None), + ) audio_hidden_states = self.audio_proj_in(audio_hidden_states) # 3. Prepare timestep embeddings and modulation parameters @@ -1997,7 +2071,7 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1]) # 3.3. LTX-2.3: Compute prompt modulation from sigma - if self.prompt_modulation and sigma is not None: + if self.prompt_modulation and self.config.use_prompt_adaln_single and sigma is not None: temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype) temb_prompt_audio, _ = self.audio_prompt_adaln( audio_sigma.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype @@ -2044,6 +2118,8 @@ class LTX2VideoTransformer3DModel(nn.Module): "ca_audio_rotary_emb": audio_cross_attn_rotary_emb, "encoder_attention_mask": encoder_attention_mask, "audio_encoder_attention_mask": audio_encoder_attention_mask, + "audio_self_attention_mask": audio_attention_mask, + "a2v_cross_attention_mask": audio_attention_mask, "video_self_attention_perturbation_mask": perturbation_mask_for("video_self_attention", block_idx), "audio_self_attention_perturbation_mask": perturbation_mask_for("audio_self_attention", block_idx), "a2v_cross_attention_perturbation_mask": perturbation_mask_for("a2v_cross_attention", block_idx), @@ -2097,16 +2173,8 @@ class LTX2VideoTransformer3DModel(nn.Module): Returns: Set of parameter names that were successfully loaded. """ - stacked_params_mapping = [ - (".attn1.to_qkv", ".attn1.to_q", "q"), - (".attn1.to_qkv", ".attn1.to_k", "k"), - (".attn1.to_qkv", ".attn1.to_v", "v"), - (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"), - (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"), - (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"), - ] - params_dict = dict(self.named_parameters()) + phase_adapter_parameter_name = getattr(self, "_phase_adapter_parameter_name", None) tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0 loaded_params: set[str] = set() @@ -2127,10 +2195,12 @@ class LTX2VideoTransformer3DModel(nn.Module): return weight for name, loaded_weight in weights: - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, shard_id in self.stacked_params_mapping: if weight_name not in name: continue packed_name = name.replace(weight_name, param_name) + if callable(phase_adapter_parameter_name): + packed_name = phase_adapter_parameter_name(packed_name) if packed_name not in params_dict: continue name = packed_name @@ -2139,6 +2209,8 @@ class LTX2VideoTransformer3DModel(nn.Module): weight_loader(param, loaded_weight, shard_id) break else: + if callable(phase_adapter_parameter_name): + name = phase_adapter_parameter_name(name) if name not in params_dict: logger.warning( "Skipping transformer weight %s -- not found in model " diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index ff66af5e1..ff2ab0457 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -10,15 +10,12 @@ from typing import ClassVar from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin -from .ltx2_components import LTX2_COMPONENT_PROFILE +from .ltx2_components import LTX2_COMPONENT_PROFILE, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE from .ltx2_components import ( get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401 ) from .ltx2_conditioning import LTXI2VConditioningMixin -from .ltx2_recipes import ( - LTX2_ONE_STAGE_RECIPE, - LTX_POSITIVE_ONLY_RECIPE, -) +from .ltx2_recipes import LTX2_DISTILLED_ONE_STAGE_RECIPE, LTX2_ONE_STAGE_RECIPE, LTX_POSITIVE_ONLY_RECIPE from .ltx2_runtime import LTXRuntime @@ -36,6 +33,14 @@ class LTX2Pipeline(LTXI2VConditioningMixin, LTXRuntime): supports_request_batch = True +class LTX2DistilledOneStagePipeline(LTX2Pipeline): + """Merged-distilled checkpoint at its native one-stage resolution.""" + + pipeline_kind = "distilled_one_stage" + component_profile = LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE + pipeline_recipe = LTX2_DISTILLED_ONE_STAGE_RECIPE + + class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline): """LTX2 T2V entry for FastGen DMD2-distilled models.""" diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py index 0e5bbc558..5d50baf42 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py @@ -1,27 +1,52 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Distilled two-stage entry points for the LTX model family.""" +"""Two-stage entry points for the LTX model family.""" -from vllm_omni.diffusion.data import OmniDiffusionConfig +from typing import ClassVar -from .ltx2_components import LTX2_DISTILLED_COMPONENT_PROFILE +from .ltx2_components import ( + LTX2_DISTILLED_COMPONENT_PROFILE, + LTX2_TWO_STAGE_COMPONENT_PROFILE, +) from .ltx2_components import ( get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401 ) -from .ltx2_recipes import LTX2_DISTILLED_TWO_STAGE_RECIPE -from .pipeline_ltx2 import LTX2Pipeline +from .ltx2_conditioning import LTXI2VConditioningMixin +from .ltx2_recipes import ( + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX2_TWO_STAGE_RECIPE, +) +from .ltx2_runtime import LTXRuntime -class LTX2DistilledPipeline(LTX2Pipeline): - """Unified LTX-2 distilled two-stage T2V/I2V entry.""" +class _LTX2TwoStageBase(LTXI2VConditioningMixin, LTXRuntime): + """Static discovery contract shared by the two-stage public entries.""" - pipeline_kind = "distilled_two_stage" - component_profile = LTX2_DISTILLED_COMPONENT_PROFILE - pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE + _dit_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.dit_modules) + _encoder_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.encoder_modules) + _vae_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.vae_modules) + _resident_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.resident_modules) supports_request_batch = False support_image_input = True unified_text_image_entry = True - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__(od_config=od_config, prefix=prefix) + +class LTX2TwoStagePipeline(_LTX2TwoStageBase): + """Regular checkpoint with low-resolution generation and LoRA refinement.""" + + pipeline_kind = "two_stage" + component_profile = LTX2_TWO_STAGE_COMPONENT_PROFILE + pipeline_recipe = LTX2_TWO_STAGE_RECIPE + + +class LTX2DistilledTwoStagePipeline(_LTX2TwoStageBase): + """Merged-distilled checkpoint with latent upsampling and refinement.""" + + pipeline_kind = "distilled_two_stage" + component_profile = LTX2_DISTILLED_COMPONENT_PROFILE + pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE + + +# Backwards-compatible name retained for existing LTX-2/LTX-2.3 recipes. +LTX2DistilledPipeline = LTX2DistilledTwoStagePipeline diff --git a/vllm_omni/diffusion/models/magi_human/magi_human_dit.py b/vllm_omni/diffusion/models/magi_human/magi_human_dit.py deleted file mode 100644 index 491b1b3c4..000000000 --- a/vllm_omni/diffusion/models/magi_human/magi_human_dit.py +++ /dev/null @@ -1,1624 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 SandAI. All Rights Reserved. -# Ported from daVinci-MagiHuman inference/model/dit/dit_module.py -# Adaptations: removed Ulysses context-parallelism, inlined Modality/VarlenHandler. - -from __future__ import annotations - -import importlib -from collections.abc import Callable -from dataclasses import dataclass, field -from enum import Enum, IntEnum -from typing import TYPE_CHECKING, Any, Literal - -import torch -import torch.nn as nn -from einops import rearrange, repeat -from torch.nn import Parameter -from torch.nn import functional as F -from vllm.distributed import ( - get_tensor_model_parallel_world_size, -) -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.vllm_flash_attn import flash_attn_varlen_func as _vllm_fa_varlen - -try: - from magi_compiler.api import magi_register_custom_op - from magi_compiler.config import CompileConfig -except Exception: - - class CompileConfig: # type: ignore[no-redef] - pass - - def magi_register_custom_op(*args, **kwargs): # type: ignore[no-redef] - def decorator(func): - return func - - return decorator - - -def magi_compile(*args, **kwargs): - """No-op stub — vllm-omni handles execution; magi compilation is skipped.""" - - def decorator(cls_or_fn): - return cls_or_fn - - return decorator - - -# --------------------------------------------------------------------------- -# Inlined from inference/common/sequence_schema.py -# --------------------------------------------------------------------------- -class Modality(IntEnum): - VIDEO = 0 - AUDIO = 1 - TEXT = 2 - - -@dataclass -class VarlenHandler: - cu_seqlens_q: torch.Tensor - cu_seqlens_k: torch.Tensor - max_seqlen_q: int - max_seqlen_k: int - - -def _is_hopper_arch() -> bool: - if not torch.cuda.is_available(): - return False - return torch.cuda.get_device_capability()[0] == 9 - - -# --------------------------------------------------------------------------- -# FFA handler for local / flex attention -# --------------------------------------------------------------------------- -@dataclass -class FFAHandler: - q_ranges: torch.Tensor - k_ranges: torch.Tensor - max_seqlen_q: int - max_seqlen_k: int - attn_type_map: torch.Tensor - softmax_scale: float - - -# --------------------------------------------------------------------------- -# Activation helpers -# --------------------------------------------------------------------------- -class MLPActivationType(Enum): - SWIGLU7 = "swiglu7" - GELU7 = "gelu7" - - -def swiglu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None): - out_dtype = x.dtype if out_dtype is None else out_dtype - x = x.to(torch.float32) - x_glu, x_linear = x[..., ::2], x[..., 1::2] - x_glu = x_glu.clamp(min=None, max=limit) - x_linear = x_linear.clamp(min=-limit, max=limit) - out_glu = x_glu * torch.sigmoid(alpha * x_glu) - return (out_glu * (x_linear + 1)).to(out_dtype) - - -def gelu7(x, alpha: float = 1.702, limit: float = 7.0, out_dtype: torch.dtype | None = None): - out_dtype = x.dtype if out_dtype is None else out_dtype - x = x.to(torch.float32) - x_glu = x.clamp(min=None, max=limit) - out_glu = x_glu * torch.sigmoid(alpha * x_glu) - return out_glu.to(out_dtype) - - -def create_activation_func(activation_type: MLPActivationType) -> Callable: - match activation_type: - case MLPActivationType.SWIGLU7: - return swiglu7 - case MLPActivationType.GELU7: - return gelu7 - case _: - raise ValueError(f"Unknown activation type: {activation_type}") - - -# --------------------------------------------------------------------------- -# Modality dispatcher (permutation helper) -# --------------------------------------------------------------------------- -class ModalityDispatcher: - permuted_modality_mapping: torch.Tensor - group_size: torch.Tensor - group_size_cpu: list[int] - num_modalities: int - - def __init__(self, modality_mapping: torch.Tensor, num_modalities: int): - self.modality_mapping = modality_mapping - self.num_modalities = num_modalities - self.permuted_modality_mapping = self._precompute_permute_mapping(modality_mapping) - self.group_size = torch.bincount(self.permuted_modality_mapping, minlength=num_modalities).to(torch.int32) - self.group_size_cpu: list[int] = [int(x) for x in self.group_size.to("cpu").tolist()] - - def _precompute_permute_mapping(self, modality_mapping): - self.permute_mapping = torch.argsort(modality_mapping) - self.inv_permute_mapping = torch.argsort(self.permute_mapping) - return modality_mapping[self.permute_mapping] - - def dispatch(self, x: torch.Tensor) -> list[torch.Tensor]: - return list(torch.split(x, self.group_size_cpu, dim=0)) - - def undispatch(self, *processed_groups: list[torch.Tensor]) -> torch.Tensor: - return torch.cat(processed_groups, dim=0) - - @staticmethod - def permute(x: torch.Tensor, permute_mapping: torch.Tensor) -> torch.Tensor: - return x[permute_mapping] - - @staticmethod - def inv_permute(x: torch.Tensor, inv_permute_mapping: torch.Tensor) -> torch.Tensor: - return x[inv_permute_mapping] - - -# --------------------------------------------------------------------------- -# Positional / rotary embedding helpers -# --------------------------------------------------------------------------- -def freq_bands( - num_bands: int, temperature: float = 10000.0, step: int = 2, device: torch.device | None = None -) -> torch.Tensor: - exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands - return 1.0 / (temperature**exp) - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2) - - -def apply_rotary_emb_torch(x, cos, sin, interleaved=False): - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)") - return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], dim=-1) - - -# --------------------------------------------------------------------------- -# Fourier positional embedding -# --------------------------------------------------------------------------- -class ElementWiseFourierEmbed(nn.Module): - def __init__( - self, - dim: int, - max_res: int = 224, - temperature: float = 10000.0, - in_pixels: bool = True, - linear_bands: bool = False, - learnable: bool = False, - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, - ): - super().__init__() - self.dim = dim - self.in_pixels = in_pixels - self.learnable = learnable - self.temperature = temperature - self.max_res = max_res - self.linear_bands = linear_bands - self.device = device - self.dtype = dtype - bands = self.get_default_bands() - self.bands = nn.Parameter(bands, requires_grad=self.learnable) - - def forward(self, coords: torch.Tensor) -> torch.Tensor: - coords_xyz = coords[:, :3] - sizes = coords[:, 3:6] - refs = coords[:, 6:9] - - scales = (refs - 1) / (sizes - 1) - scales[(refs == 1) & (sizes == 1)] = 1 - assert not scales.isnan().any(), "scales has nan" - assert not scales.isinf().any(), "scales has inf" - - centers = (sizes - 1) / 2 - centers[:, 0] = 0 - coords_xyz = coords_xyz - centers - - bands = self.bands.to(coords.device, coords.dtype) - proj = coords_xyz.unsqueeze(-1) * scales.unsqueeze(-1) * bands - sin_proj = proj.sin() - cos_proj = proj.cos() - return torch.cat((sin_proj, cos_proj), dim=1).flatten(1) - - def reset_parameters(self): - self.bands.copy_(self.get_default_bands()) - - def get_default_bands(self): - if self.in_pixels: - raise NotImplementedError("in_pixels are not implemented yet") - return freq_bands(self.dim // 8, temperature=self.temperature, step=1, device=self.device).to(self.dtype) - - -# --------------------------------------------------------------------------- -# Multi-modality RMSNorm -# --------------------------------------------------------------------------- -class MultiModalityRMSNorm(nn.Module): - __constants__ = ["dim", "eps", "num_modality"] - - def __init__(self, dim: int, eps: float = 1e-6, device: torch.device | None = None, num_modality: int = 1): - super().__init__() - self.dim = dim - self.eps = eps - self.num_modality = num_modality - self.weight = nn.Parameter(torch.zeros(dim * num_modality, device=device, dtype=torch.float32)) - if num_modality > 1: - self.forward = self.forward_multi_experts - else: - self.forward = self.forward_single_expert - self.reset_parameters() - - def reset_parameters(self): - nn.init.zeros_(self.weight) - - def rms(self, x: torch.Tensor) -> torch.Tensor: - t = x.float() - return t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) - - def forward_multi_experts(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: - original_dtype = x.dtype - t = self.rms(x) - weight_chunked = self.weight.chunk(self.num_modality, dim=0) - t_list = modality_dispatcher.dispatch(t) - for i in range(self.num_modality): - t_list[i] = t_list[i] * (weight_chunked[i] + 1) - t = modality_dispatcher.undispatch(*t_list) - return t.to(original_dtype) - - def forward_single_expert( - self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher | None = None - ) -> torch.Tensor: - t, original_dtype = x.float(), x.dtype - t = t * torch.rsqrt(torch.mean(t**2, dim=-1, keepdim=True) + self.eps) - return (t * (self.weight + 1)).to(original_dtype) - - -# --------------------------------------------------------------------------- -# Linear layers with bf16 compute and MoE dispatch -# --------------------------------------------------------------------------- -class _BF16ComputeLinear(torch.autograd.Function): - @staticmethod - def forward( - ctx, - input: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - output_dtype: torch.dtype | None, - compute_dtype: torch.dtype = torch.bfloat16, - ): - input_cast = input.to(compute_dtype) - weight_cast = weight.to(compute_dtype) - output = torch.matmul(input_cast, weight_cast.t()) - if bias is not None: - output = output + bias.to(compute_dtype) - return output.to(output_dtype) - - -class BaseLinear(nn.Module): - __constants__ = ["in_features", "out_features", "num_layers", "num_experts"] - - def __init__( - self, in_features, out_features, num_layers_for_initialization, num_experts, bias=True, device=None, dtype=None - ): - super().__init__() - factory_kwargs = {"device": device, "dtype": torch.bfloat16} - self.in_features = in_features - self.out_features = out_features - self.num_layers_for_initialization = num_layers_for_initialization - self.num_experts = num_experts - self.use_bias = bias - self.weight = Parameter(torch.empty((out_features * num_experts, in_features), **factory_kwargs)) - if bias: - self.bias = Parameter(torch.empty(out_features * num_experts, **factory_kwargs)) - else: - self.register_parameter("bias", None) - - def forward( - self, - input: torch.Tensor, - output_dtype: torch.dtype | None = None, - modality_dispatcher: ModalityDispatcher | None = None, - ) -> torch.Tensor: - output_dtype = input.dtype if output_dtype is None else output_dtype - return _BF16ComputeLinear.apply(input, self.weight, self.bias, output_dtype, torch.bfloat16) - - -class NativeMoELinear(BaseLinear): - def forward( - self, - input: torch.Tensor, - output_dtype: torch.dtype | None = None, - modality_dispatcher: ModalityDispatcher | None = None, - ) -> torch.Tensor: - output_dtype = input.dtype if output_dtype is None else output_dtype - input_list = modality_dispatcher.dispatch(input) # type: ignore - weight_chunked = self.weight.chunk(self.num_experts, dim=0) - if self.bias is not None: - bias_chunked = self.bias.chunk(self.num_experts, dim=0) - for i in range(self.num_experts): - input_list[i] = _BF16ComputeLinear.apply( - input_list[i], - weight_chunked[i], - bias_chunked[i] if self.bias is not None else None, - output_dtype, - torch.bfloat16, - ) - return modality_dispatcher.undispatch(*input_list) # type: ignore - - -def create_linear( - in_features, out_features, num_layers=1, num_experts=1, bias=True, device=None, dtype=None -) -> BaseLinear | NativeMoELinear: - if num_experts == 1: - return BaseLinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) - else: - return NativeMoELinear(in_features, out_features, num_layers, num_experts, bias, device, dtype) - - -# --------------------------------------------------------------------------- -# MoE TP parallel linear wrappers: per-expert vLLM parallel layers -# --------------------------------------------------------------------------- -class MoEQKVParallelLinear(nn.Module): - """Per-expert QKVParallelLinear with modality dispatch. - - Wraps ``num_experts`` independent QKVParallelLinear instances. - Forward: dispatch tokens by modality → per-expert QKV matmul (TP-sharded) - → undispatch. - """ - - def __init__( - self, - hidden_size: int, - head_size: int, - total_num_heads: int, - total_num_kv_heads: int, - num_experts: int, - bias: bool = False, - ): - super().__init__() - self.num_experts = num_experts - self.experts = nn.ModuleList( - [ - QKVParallelLinear( - hidden_size=hidden_size, - head_size=head_size, - total_num_heads=total_num_heads, - total_num_kv_heads=total_num_kv_heads, - bias=bias, - return_bias=False, - ) - for _ in range(num_experts) - ] - ) - # Expose per-rank head info from the first expert (all are identical). - self.num_heads = self.experts[0].num_heads - self.num_kv_heads = self.experts[0].num_kv_heads - self.head_size = head_size - - def forward( - self, - x: torch.Tensor, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - x_list = modality_dispatcher.dispatch(x) - out_list: list[torch.Tensor] = [] - for i in range(self.num_experts): - out = self.experts[i](x_list[i]) - out_list.append(out) - return modality_dispatcher.undispatch(*out_list) - - -class MoEColumnParallelLinear(nn.Module): - """Per-expert ColumnParallelLinear with modality dispatch. - - Forward: dispatch → per-expert column-parallel matmul → undispatch. - Output stays TP-local (no gather). - """ - - def __init__( - self, - input_size: int, - output_size: int, - num_experts: int, - bias: bool = False, - ): - super().__init__() - self.num_experts = num_experts - self.experts = nn.ModuleList( - [ - ColumnParallelLinear( - input_size=input_size, - output_size=output_size, - bias=bias, - gather_output=False, - return_bias=False, - ) - for _ in range(num_experts) - ] - ) - - def forward( - self, - x: torch.Tensor, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - x_list = modality_dispatcher.dispatch(x) - out_list: list[torch.Tensor] = [] - for i in range(self.num_experts): - out = self.experts[i](x_list[i]) - out_list.append(out) - return modality_dispatcher.undispatch(*out_list) - - -class MoERowParallelLinear(nn.Module): - """Per-expert RowParallelLinear with modality dispatch. - - Forward: dispatch → per-expert row-parallel matmul (includes all-reduce) - → undispatch. - """ - - def __init__( - self, - input_size: int, - output_size: int, - num_experts: int, - bias: bool = False, - ): - super().__init__() - self.num_experts = num_experts - self.experts = nn.ModuleList( - [ - RowParallelLinear( - input_size=input_size, - output_size=output_size, - bias=bias, - input_is_parallel=True, - return_bias=False, - ) - for _ in range(num_experts) - ] - ) - - def forward( - self, - x: torch.Tensor, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - x_list = modality_dispatcher.dispatch(x) - out_list: list[torch.Tensor] = [] - for i in range(self.num_experts): - out = self.experts[i](x_list[i]) - out_list.append(out) - return modality_dispatcher.undispatch(*out_list) - - -def validate_magi_human_tp_constraints( - *, - hidden_size: int, - num_heads_q: int, - num_heads_kv: int, - tensor_parallel_size: int, -) -> None: - """Validate MagiHuman TP divisibility constraints. - - Both shared layers (num_modality == 1) and MoE layers (num_modality == 3) - support TP via vLLM's parallel linear layers (QKVParallelLinear / - ColumnParallelLinear / RowParallelLinear). MoE layers use per-expert - parallel layers with modality dispatch. - - Supported tp_sizes given default config (hidden=5120, heads_q=40, kv=8): 1, 2, 4. - """ - tp = tensor_parallel_size - if tp <= 1: - return - errors: list[str] = [] - if num_heads_q % tp != 0: - errors.append(f"num_heads_q ({num_heads_q}) must be divisible by tensor_parallel_size ({tp})") - if num_heads_kv % tp != 0: - errors.append(f"num_heads_kv ({num_heads_kv}) must be divisible by tensor_parallel_size ({tp})") - # SWIGLU layers use intermediate = int(hidden * 8/3) // 4 * 4 - intermediate_swiglu = int(hidden_size * 4 * 2 / 3) // 4 * 4 - if intermediate_swiglu % tp != 0: - errors.append( - f"swiglu intermediate_size ({intermediate_swiglu}) must be divisible by " - f"tensor_parallel_size ({tp}). Supported tp values: 1, 2, 4" - ) - # GELU7 MoE layers use intermediate = hidden * 4 - intermediate_gelu = hidden_size * 4 - if intermediate_gelu % tp != 0: - errors.append(f"gelu intermediate_size ({intermediate_gelu}) must be divisible by tensor_parallel_size ({tp})") - if errors: - raise ValueError("MagiHuman TP constraint violations:\n" + "\n".join(f" - {e}" for e in errors)) - - -# --------------------------------------------------------------------------- -# Flash attention (no context-parallelism) — uses vllm's flash attention -# --------------------------------------------------------------------------- - -HAS_MAGI_ATTENTION = importlib.util.find_spec("magi_attention") is not None - - -def _fa_varlen_simple( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, -) -> torch.Tensor: - had_batch = query.ndim == 4 - if had_batch: - query = query.squeeze(0) - key = key.squeeze(0) - value = value.squeeze(0) - seq_len = query.shape[0] - cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=query.device) - out = _vllm_fa_varlen( - q=query, - k=key, - v=value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=seq_len, - max_seqlen_k=seq_len, - ) - if had_batch: - out = out.unsqueeze(0) - return out - - -@magi_register_custom_op(name="infra::flash_attn_func", is_subgraph_boundary=True) -def flash_attn_func(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: - return _fa_varlen_simple(query, key, value) - - -def _split_q_range_with_no_overlap( - q_ranges: torch.Tensor, k_ranges: torch.Tensor -) -> tuple[list[list[int]], list[list[list[int]]]]: - range_boundary = torch.unique(q_ranges, sorted=True).tolist() - candidates = [[start, end, []] for start, end in zip(range_boundary[:-1], range_boundary[1:])] - q_ranges = q_ranges.tolist() - k_ranges = k_ranges.tolist() - for q_range, k_range in zip(q_ranges, k_ranges): - q_start, q_end = q_range - for q_range_cand in candidates: - if q_start <= q_range_cand[0] and q_range_cand[1] <= q_end: - q_range_cand[2].append(k_range) - q_ranges_out = [] - k_ranges_out = [] - for q_range_cand in candidates: - if len(q_range_cand[2]) > 0: - q_ranges_out.append(q_range_cand[0:2]) - k_ranges_out.append(q_range_cand[2]) - return q_ranges_out, k_ranges_out - - -def _flash_attn_with_correction( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - q_ranges: list[list[int]], - k_range_list: list[list[list[int]]], -): - output = torch.zeros_like(query) - output_lse = torch.zeros((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) - - for q_range, k_ranges in zip(q_ranges, k_range_list): - q_start, q_end = q_range - q_chunk = query[q_start:q_end] - q_len = q_chunk.shape[0] - - # Concatenate all k_ranges into a single key/value block, then run one - # flash-attention call. This avoids the need to merge per-chunk LSEs. - k_parts = [key[ks:ke] for ks, ke in k_ranges] - v_parts = [value[ks:ke] for ks, ke in k_ranges] - k_combined = torch.cat(k_parts, dim=0) if len(k_parts) > 1 else k_parts[0] - v_combined = torch.cat(v_parts, dim=0) if len(v_parts) > 1 else v_parts[0] - k_len = k_combined.shape[0] - - cu_q = torch.tensor([0, q_len], dtype=torch.int32, device=query.device) - cu_k = torch.tensor([0, k_len], dtype=torch.int32, device=query.device) - qo_out = _vllm_fa_varlen( - q=q_chunk, - k=k_combined, - v=v_combined, - cu_seqlens_q=cu_q, - cu_seqlens_k=cu_k, - max_seqlen_q=q_len, - max_seqlen_k=k_len, - ) - output[q_start:q_end] = qo_out - return output, output_lse - - -def _flex_flash_attn_func_infer_output_meta( - query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - output = torch.empty_like(query) - output_lse = torch.empty((query.shape[0], query.shape[1]), dtype=torch.float32, device=query.device) - return output, output_lse - - -@magi_register_custom_op( - name="infra::flex_flash_attn_func", - mutates_args=(), - infer_output_meta_fn=_flex_flash_attn_func_infer_output_meta, - is_subgraph_boundary=True, -) -def flex_flash_attn_func( - query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, q_ranges: torch.Tensor, k_ranges: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - if HAS_MAGI_ATTENTION and _is_hopper_arch(): - from magi_attention.api import flex_flash_attn_func as magi_flex_flash_attn_func - - return magi_flex_flash_attn_func(query, key, value, q_ranges, k_ranges) - else: - q_ranges_split, k_range_list = _split_q_range_with_no_overlap(q_ranges, k_ranges) - return _flash_attn_with_correction(query, key, value, q_ranges_split, k_range_list) - - -def flash_attn_no_cp(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: - q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16) - return flash_attn_func(q, k, v).squeeze(0) - - -def flex_flash_attn_no_cp( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - q_ranges: torch.Tensor, - k_ranges: torch.Tensor, -) -> torch.Tensor: - q, k, v = q.to(torch.bfloat16).squeeze(0), k.to(torch.bfloat16).squeeze(0), v.to(torch.bfloat16).squeeze(0) - out, _ = flex_flash_attn_func(q, k, v, q_ranges=q_ranges, k_ranges=k_ranges) - return out - - -# --------------------------------------------------------------------------- -# Attention module (no context-parallelism) -# --------------------------------------------------------------------------- -@dataclass -class AttentionConfig: - hidden_size: int - num_heads_q: int - num_heads_kv: int - head_dim: int - params_dtype: torch.dtype - checkpoint_qk_layernorm_rope: bool - num_modality: int - num_layers: int - use_local_attn: bool = False - enable_attn_gating: bool = False - - -class Attention(torch.nn.Module): - config: AttentionConfig - - def __init__(self, config: AttentionConfig): - super().__init__() - self.config = config - self.pre_norm = MultiModalityRMSNorm(config.hidden_size, eps=1e-6, num_modality=config.num_modality) - self.gating_size = config.num_heads_q if config.enable_attn_gating else 0 - - # Both shared blocks (num_modality == 1) and MoE blocks (num_modality > 1) - # use vLLM's parallel linear layers for TP support. - # MoE blocks wrap per-expert parallel layers with modality dispatch. - if config.num_modality == 1: - # QKVParallelLinear handles GQA head-sharding for any tp_size. - # The combined checkpoint weight [Q, K, V, G] is split during - # load_weights: Q+K+V → linear_qkv, G → linear_gating. - self.linear_qkv = QKVParallelLinear( - hidden_size=config.hidden_size, - head_size=config.head_dim, - total_num_heads=config.num_heads_q, - total_num_kv_heads=config.num_heads_kv, - bias=False, - return_bias=False, - ) - self.linear_proj = RowParallelLinear( - input_size=config.num_heads_q * config.head_dim, - output_size=config.hidden_size, - bias=False, - input_is_parallel=True, - return_bias=False, - ) - if config.enable_attn_gating: - self.linear_gating = ColumnParallelLinear( - input_size=config.hidden_size, - output_size=config.num_heads_q, - bias=False, - gather_output=False, - return_bias=False, - ) - else: - self.linear_gating = None - else: - # MoE blocks: per-expert TP-sharded parallel layers. - self.linear_qkv = MoEQKVParallelLinear( - hidden_size=config.hidden_size, - head_size=config.head_dim, - total_num_heads=config.num_heads_q, - total_num_kv_heads=config.num_heads_kv, - num_experts=config.num_modality, - bias=False, - ) - self.linear_proj = MoERowParallelLinear( - input_size=config.num_heads_q * config.head_dim, - output_size=config.hidden_size, - num_experts=config.num_modality, - bias=False, - ) - if config.enable_attn_gating: - self.linear_gating = MoEColumnParallelLinear( - input_size=config.hidden_size, - output_size=config.num_heads_q, - num_experts=config.num_modality, - bias=False, - ) - else: - self.linear_gating = None - - self.q_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) - self.k_norm = MultiModalityRMSNorm(config.head_dim, num_modality=config.num_modality) - - # q_size / kv_size reflect the per-rank head count when tp > 1. - # Both shared and MoE QKV layers expose .num_heads / .num_kv_heads. - if config.num_modality == 1: - self.q_size = self.linear_qkv.num_heads * config.head_dim - self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim - self._local_heads_q = self.linear_qkv.num_heads - self._local_heads_kv = self.linear_qkv.num_kv_heads - else: - self.q_size = self.linear_qkv.num_heads * config.head_dim - self.kv_size = self.linear_qkv.num_kv_heads * config.head_dim - self._local_heads_q = self.linear_qkv.num_heads - self._local_heads_kv = self.linear_qkv.num_kv_heads - - def forward( - self, - hidden_states: torch.Tensor, - rope: torch.Tensor, - permute_mapping: torch.Tensor, - inv_permute_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - hidden_states = self.pre_norm(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) - - if self.config.num_modality == 1: - # vLLM parallel layers with return_bias=False return a single tensor. - qkv = self.linear_qkv(hidden_states).to(torch.float32) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - if self.linear_gating is not None: - g = self.linear_gating(hidden_states).to(torch.float32) - else: - g = hidden_states.new_empty(hidden_states.shape[0], 0) - else: - # MoE TP path: per-expert QKV parallel layers. - qkv = self.linear_qkv(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - if self.linear_gating is not None: - g = self.linear_gating(hidden_states, modality_dispatcher=modality_dispatcher).to(torch.float32) - else: - g = hidden_states.new_empty(hidden_states.shape[0], 0) - - q = q.view(-1, self._local_heads_q, self.config.head_dim) - k = k.view(-1, self._local_heads_kv, self.config.head_dim) - v = v.view(-1, self._local_heads_kv, self.config.head_dim) - g = g.view(k.shape[0], self._local_heads_q, -1) - - q = self.q_norm(q, modality_dispatcher=modality_dispatcher) - k = self.k_norm(k, modality_dispatcher=modality_dispatcher) - - q = ModalityDispatcher.inv_permute(q, inv_permute_mapping).unsqueeze(0) - k = ModalityDispatcher.inv_permute(k, inv_permute_mapping).unsqueeze(0) - v = ModalityDispatcher.inv_permute(v, inv_permute_mapping).unsqueeze(0) - - sin_emb, cos_emb = rope.tensor_split(2, -1) - q = apply_rotary_emb_torch(q, cos_emb, sin_emb) - k = apply_rotary_emb_torch(k, cos_emb, sin_emb) - - if self.config.use_local_attn and local_attn_handler is not None: - self_attn_out = flex_flash_attn_no_cp(q, k, v, local_attn_handler.q_ranges, local_attn_handler.k_ranges) - else: - self_attn_out = flash_attn_no_cp(q, k, v) - self_attn_out = ModalityDispatcher.permute(self_attn_out, permute_mapping) - - if self.config.enable_attn_gating: - self_attn_out = self_attn_out * torch.sigmoid(g) - - self_attn_out = self_attn_out.view(-1, self._local_heads_q * self.config.head_dim).to(torch.bfloat16) - if self.config.num_modality == 1: - return self.linear_proj(self_attn_out) - return self.linear_proj(self_attn_out, modality_dispatcher=modality_dispatcher) - - -# --------------------------------------------------------------------------- -# MLP module -# --------------------------------------------------------------------------- -@dataclass -class MLPConfig: - hidden_size: int - intermediate_size: int - activation_type: MLPActivationType - params_dtype: torch.dtype - num_modality: int = 1 - num_layers: int = 1 - gated_act: bool = False - - -class MLP(torch.nn.Module): - config: MLPConfig - - def __init__(self, config: MLPConfig): - super().__init__() - num_experts = config.num_modality - self.pre_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=config.num_modality) - intermediate_size_up = config.intermediate_size * 2 if config.gated_act else config.intermediate_size - - # Both shared blocks (num_experts == 1) and MoE blocks (num_experts > 1) - # use vLLM's parallel linear layers for TP support. - if num_experts == 1: - # ColumnParallelLinear shards the output dim uniformly. For - # SWIGLU7 the interleaved [up0, gate0, up1, gate1, ...] format - # is preserved within each rank's contiguous slice, so swiglu7 - # (which uses x[..., ::2] / x[..., 1::2]) still works correctly. - self.up_gate_proj = ColumnParallelLinear( - input_size=config.hidden_size, - output_size=intermediate_size_up, - bias=False, - gather_output=False, - return_bias=False, - ) - self.down_proj = RowParallelLinear( - input_size=config.intermediate_size, - output_size=config.hidden_size, - bias=False, - input_is_parallel=True, - return_bias=False, - ) - else: - # MoE blocks: per-expert TP-sharded parallel layers. - self.up_gate_proj = MoEColumnParallelLinear( - input_size=config.hidden_size, - output_size=intermediate_size_up, - num_experts=num_experts, - bias=False, - ) - self.down_proj = MoERowParallelLinear( - input_size=config.intermediate_size, - output_size=config.hidden_size, - num_experts=num_experts, - bias=False, - ) - self.activation_func = create_activation_func(config.activation_type) - - def forward(self, x: torch.Tensor, modality_dispatcher: ModalityDispatcher) -> torch.Tensor: - x = self.pre_norm(x, modality_dispatcher=modality_dispatcher).to(torch.bfloat16) - if isinstance(self.up_gate_proj, ColumnParallelLinear): - x = self.up_gate_proj(x).to(torch.float32) - x = self.activation_func(x).to(torch.bfloat16) - return self.down_proj(x).to(torch.float32) - # MoE TP path: per-expert column/row parallel layers. - x = self.up_gate_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32) - x = self.activation_func(x).to(torch.bfloat16) - x = self.down_proj(x, modality_dispatcher=modality_dispatcher).to(torch.float32) - return x - - -# --------------------------------------------------------------------------- -# Adapter (per-modality embedders + RoPE) -# --------------------------------------------------------------------------- -@dataclass -class AdapterConfig: - hidden_size: int - num_attention_heads: int - text_in_channels: int - video_in_channels: int - audio_in_channels: int - params_dtype: torch.dtype - - -class Adapter(torch.nn.Module): - config: AdapterConfig - - def __init__(self, config: AdapterConfig): - super().__init__() - self.config = config - self.video_embedder = nn.Linear(config.video_in_channels, config.hidden_size, bias=True, dtype=torch.float32) - self.text_embedder = nn.Linear(config.text_in_channels, config.hidden_size, bias=True, dtype=torch.float32) - self.audio_embedder = nn.Linear(config.audio_in_channels, config.hidden_size, bias=True, dtype=torch.float32) - self.rope = ElementWiseFourierEmbed( - config.hidden_size // config.num_attention_heads, in_pixels=False, learnable=False - ) - - def forward(self, x, coords_mapping, video_mask, audio_mask, text_mask): - rope = self.rope(coords_mapping) - - text_input = x[text_mask, : self.config.text_in_channels] - audio_input = x[audio_mask, : self.config.audio_in_channels] - video_input = x[video_mask, : self.config.video_in_channels] - - text_out = self.text_embedder(text_input) - audio_out = self.audio_embedder(audio_input) - video_out = self.video_embedder(video_input) - - output_x = torch.zeros(x.shape[0], self.config.hidden_size, device=x.device, dtype=x.dtype) - output_x[text_mask] = text_out - output_x[audio_mask] = audio_out - output_x[video_mask] = video_out - return output_x, rope - - -# --------------------------------------------------------------------------- -# Transformer layer (no CP) -# --------------------------------------------------------------------------- -class TransFormerLayer(torch.nn.Module): - def __init__(self, config: Any, layer_idx: int): - super().__init__() - num_modality = 3 if layer_idx in config.mm_layers else 1 - use_local_attn = layer_idx in config.local_attn_layers - self.post_norm = layer_idx in config.post_norm_layers - attention_config = AttentionConfig( - hidden_size=config.hidden_size, - num_heads_q=config.num_heads_q, - num_heads_kv=config.num_heads_kv, - head_dim=config.head_dim, - params_dtype=config.params_dtype, - checkpoint_qk_layernorm_rope=config.checkpoint_qk_layernorm_rope, - num_modality=num_modality, - num_layers=config.num_layers, - use_local_attn=use_local_attn, - enable_attn_gating=config.enable_attn_gating, - ) - self.attention: Attention = Attention(attention_config) - - activation_type = MLPActivationType.GELU7 if layer_idx in config.gelu7_layers else MLPActivationType.SWIGLU7 - if activation_type == MLPActivationType.SWIGLU7: - gated_act = True - intermediate_size = int(config.hidden_size * 4 * 2 / 3) // 4 * 4 - else: - gated_act = False - intermediate_size = config.hidden_size * 4 - mlp_config = MLPConfig( - hidden_size=config.hidden_size, - intermediate_size=intermediate_size, - activation_type=activation_type, - params_dtype=config.params_dtype, - num_modality=num_modality, - num_layers=config.num_layers, - gated_act=gated_act, - ) - self.mlp: MLP = MLP(mlp_config) - if self.post_norm: - self.attn_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) - self.mlp_post_norm = MultiModalityRMSNorm(config.hidden_size, num_modality=num_modality) - - def forward( - self, - hidden_states: torch.Tensor, - rope: torch.Tensor, - permute_mapping: torch.Tensor, - inv_permute_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - attn_out = self.attention( - hidden_states, - rope, - permute_mapping, - inv_permute_mapping, - varlen_handler, - local_attn_handler, - modality_dispatcher, - ) - if self.post_norm: - attn_out = self.attn_post_norm(attn_out, modality_dispatcher=modality_dispatcher) - hidden_states = hidden_states + attn_out - - mlp_out = self.mlp(hidden_states, modality_dispatcher) - if self.post_norm: - mlp_out = self.mlp_post_norm(mlp_out, modality_dispatcher=modality_dispatcher) - hidden_states = hidden_states + mlp_out - return hidden_states - - -# --------------------------------------------------------------------------- -# TransformerBlock with magi_compile -# --------------------------------------------------------------------------- -is_base_model = True - - -def config_patch(compile_config: CompileConfig) -> CompileConfig: - global is_base_model - if is_base_model: - is_base_model = False - else: - compile_config.offload_config.gpu_resident_weight_ratio = 0.0 - return compile_config - - -@magi_compile( - config_patch=config_patch, dynamic_arg_dims={"x": 0, "rope": 0, "permute_mapping": 0, "inv_permute_mapping": 0} -) -class TransformerBlock(torch.nn.Module): - def __init__(self, model_config: Any): - super().__init__() - self.layers: list[TransFormerLayer] = nn.ModuleList() - for layer_idx in range(model_config.num_layers): - self.layers.append(TransFormerLayer(model_config, layer_idx)) - - def forward( - self, - x: torch.Tensor, - rope: torch.Tensor, - permute_mapping: torch.Tensor, - inv_permute_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - modality_dispatcher: ModalityDispatcher, - ) -> torch.Tensor: - for layer in self.layers: - x = layer( - x, rope, permute_mapping, inv_permute_mapping, varlen_handler, local_attn_handler, modality_dispatcher - ) - return x - - -# --------------------------------------------------------------------------- -# Internal config for TransformerBlock / DiTModel construction -# --------------------------------------------------------------------------- -@dataclass -class TransformerConfig: - hidden_size: int - video_in_channels: int - audio_in_channels: int - text_in_channels: int - params_dtype: torch.dtype - post_process_dtype: torch.dtype - - -# --------------------------------------------------------------------------- -# DiTModel (no context-parallelism) -# --------------------------------------------------------------------------- -class DiTModel(torch.nn.Module): - config: TransformerConfig - _layerwise_offload_blocks_attr = "blocks" - - @property - def blocks(self) -> nn.ModuleList: - return self.block.layers - - def __init__(self, model_config: Any): - super().__init__() - validate_magi_human_tp_constraints( - hidden_size=model_config.hidden_size, - num_heads_q=model_config.hidden_size // model_config.head_dim, - num_heads_kv=model_config.num_query_groups, - tensor_parallel_size=get_tensor_model_parallel_world_size(), - ) - self.config = TransformerConfig( - hidden_size=model_config.hidden_size, - video_in_channels=model_config.video_in_channels, - audio_in_channels=model_config.audio_in_channels, - text_in_channels=model_config.text_in_channels, - params_dtype=model_config.params_dtype, - post_process_dtype=torch.float32, - ) - adapter_config = AdapterConfig( - hidden_size=model_config.hidden_size, - num_attention_heads=model_config.num_heads_q, - text_in_channels=model_config.text_in_channels, - video_in_channels=model_config.video_in_channels, - audio_in_channels=model_config.audio_in_channels, - params_dtype=torch.float32, - ) - self.adapter: Adapter = Adapter(adapter_config) - self.block: TransformerBlock = TransformerBlock(model_config=model_config) - self.final_norm_video = MultiModalityRMSNorm(self.config.hidden_size) - self.final_norm_audio = MultiModalityRMSNorm(self.config.hidden_size) - self.final_linear_video = nn.Linear( - self.config.hidden_size, self.config.video_in_channels, bias=False, dtype=torch.float32 - ) - self.final_linear_audio = nn.Linear( - self.config.hidden_size, self.config.audio_in_channels, bias=False, dtype=torch.float32 - ) - - def forward( - self, - x: torch.Tensor, - coords_mapping: torch.Tensor, - modality_mapping: torch.Tensor, - varlen_handler: VarlenHandler, - local_attn_handler: FFAHandler | None, - ): - modality_dispatcher = ModalityDispatcher(modality_mapping, 3) - permute_mapping = modality_dispatcher.permute_mapping - inv_permute_mapping = modality_dispatcher.inv_permute_mapping - video_mask = modality_mapping == Modality.VIDEO - audio_mask = modality_mapping == Modality.AUDIO - text_mask = modality_mapping == Modality.TEXT - - x, rope = self.adapter(x, coords_mapping, video_mask, audio_mask, text_mask) - - x = x.to(self.config.params_dtype) - x = ModalityDispatcher.permute(x, permute_mapping) - - x = self.block( - x, - rope, - permute_mapping=permute_mapping, - inv_permute_mapping=inv_permute_mapping, - varlen_handler=varlen_handler, - local_attn_handler=local_attn_handler, - modality_dispatcher=modality_dispatcher, - ) - - x = ModalityDispatcher.inv_permute(x, inv_permute_mapping) - - x_video = x[video_mask].to(self.final_norm_video.weight.dtype) - x_video = self.final_norm_video(x_video) - x_video = self.final_linear_video(x_video) - - x_audio = x[audio_mask].to(self.final_norm_audio.weight.dtype) - x_audio = self.final_norm_audio(x_audio) - x_audio = self.final_linear_audio(x_audio) - - x_out = torch.zeros( - x.shape[0], - max(self.config.video_in_channels, self.config.audio_in_channels), - device=x.device, - dtype=x.dtype, - ) - x_out[video_mask, : self.config.video_in_channels] = x_video - x_out[audio_mask, : self.config.audio_in_channels] = x_audio - - return x_out - - -# --------------------------------------------------------------------------- -# Public config dataclass for building DiTModel from JSON -# --------------------------------------------------------------------------- -@dataclass -class MagiHumanDiTConfig: - num_layers: int = 40 - hidden_size: int = 5120 - head_dim: int = 128 - num_query_groups: int = 8 - video_in_channels: int = 48 * 4 - audio_in_channels: int = 64 - text_in_channels: int = 3584 - checkpoint_qk_layernorm_rope: bool = False - params_dtype: torch.dtype = torch.float32 - mm_layers: list = field(default_factory=lambda: [0, 1, 2, 3, 36, 37, 38, 39]) - local_attn_layers: list = field(default_factory=list) - enable_attn_gating: bool = True - gelu7_layers: list = field(default_factory=lambda: [0, 1, 2, 3]) - post_norm_layers: list = field(default_factory=list) - - def __post_init__(self): - self.num_heads_q = self.hidden_size // self.head_dim - self.num_heads_kv = self.num_query_groups - - -if TYPE_CHECKING: - from .pipeline_magi_human import EvalInput - - -# =========================================================================== -# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py) -# =========================================================================== -def _unfold_3d( - x: torch.Tensor, - kernel_size: tuple[int, int, int], - stride: tuple[int, int, int], -) -> torch.Tensor: - """Pure-PyTorch 3D unfold matching UnfoldAnd behavior. - - After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW). - UnfoldAnd permutes kernel dims next to channel before reshape so that the - col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics. - Without this permute, .view() interleaves spatial and kernel positions. - - Args: - x: (N, C, D, H, W) - kernel_size: (kD, kH, kW) - stride: (sD, sH, sW) - Returns: - (N, C*kD*kH*kW, L) where L = product of output spatial dims. - """ - ndim = len(kernel_size) - for d in range(ndim): - x = x.unfold(d + 2, kernel_size[d], stride[d]) - # x: (N, C, oD, oH, oW, kD, kH, kW) - # Permute to (N, C, kD, kH, kW, oD, oH, oW) so that view groups correctly - perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2)) - x = x.permute(*perm).contiguous() - - batch_size = x.shape[0] - col_dim = 1 - for i in range(1, ndim + 2): - col_dim *= x.shape[i] - spatial = 1 - for i in range(ndim + 2, 2 * ndim + 2): - spatial *= x.shape[i] - return x.view(batch_size, col_dim, spatial) - - -def calc_local_qk_range( - num_video_tokens, - num_audio_and_txt_tokens, - num_frames, - frame_receptive_field, -): - token_per_frame = num_video_tokens // num_frames - total_tokens = num_video_tokens + num_audio_and_txt_tokens - - q_range_list = [] - k_range_list = [] - for i in range(num_frames): - q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])) - k_range_list.append( - torch.tensor( - [ - (i - frame_receptive_field) * token_per_frame, - (i + frame_receptive_field + 1) * token_per_frame, - ] - ) - ) - local_q_range = torch.stack(q_range_list, dim=0) - local_k_range = torch.stack(k_range_list, dim=0) - - local_k_range[local_k_range < 0] = 0 - local_k_range[local_k_range > num_video_tokens] = num_video_tokens - - video_q_range = torch.tensor([[0, num_video_tokens]]) - video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]]) - - at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]]) - at_k_ranges = torch.tensor([[0, total_tokens]]) - - q_ranges = ( - torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - k_ranges = ( - torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - return q_ranges, k_ranges - - -def calc_local_attn_ffa_handler( - num_video_tokens, - num_audio_and_txt_tokens, - num_frames, - frame_receptive_field, -): - q_ranges, k_ranges = calc_local_qk_range( - num_video_tokens, - num_audio_and_txt_tokens, - num_frames, - frame_receptive_field, - ) - total = num_video_tokens + num_audio_and_txt_tokens - return FFAHandler( - q_ranges=q_ranges, - k_ranges=k_ranges, - max_seqlen_q=total, - max_seqlen_k=total, - attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32), - softmax_scale=None, - ) - - -def get_coords( - shape: list[int], - ref_feat_shape: list[int], - offset_thw: list[int] | None = None, - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, -): - if offset_thw is None: - offset_thw = [0, 0, 0] - ori_t, ori_h, ori_w = shape - ref_t, ref_h, ref_w = ref_feat_shape - - offset_t, offset_h, offset_w = offset_thw - time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t - height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h - width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w - - time_grid, height_grid, width_grid = torch.meshgrid( - time_rng, - height_rng, - width_rng, - indexing="ij", - ) - coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3) - - meta = torch.tensor( - [ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], - device=device, - dtype=dtype, - ) - meta_expanded = meta.expand(coords_flat.size(0), -1) - return torch.cat([coords_flat, meta_expanded], dim=-1) - - -@dataclass -class SingleData: - video_x_t: torch.Tensor - audio_x_t: torch.Tensor - audio_feat_len: int - txt_feat: torch.Tensor - txt_feat_len: int - t: int - h: int - w: int - patch_size: int - t_patch_size: int - spatial_rope_interpolation: Literal["inter", "extra"] - ref_audio_offset: int - text_offset: int - coords_style: Literal["v1", "v2"] = "v1" - - def __post_init__(self): - self.video_token_num = self.video_x_t.shape[0] - self.audio_x_t = self.audio_x_t[: self.audio_feat_len] - self.txt_feat = self.txt_feat[: self.txt_feat_len] - self.video_channel = self.video_x_t.shape[-1] - self.audio_channel = self.audio_x_t.shape[-1] - self.txt_channel = self.txt_feat.shape[-1] - - @property - def device(self): - return self.video_x_t.device - - @property - def default_dtype(self): - return self.video_x_t.dtype - - @property - def total_token_num(self): - return self.video_token_num + self.audio_feat_len + self.txt_feat_len - - @property - def token_sequence(self): - tensors = [self.video_x_t, self.audio_x_t, self.txt_feat] - max_channel = max(t.shape[-1] for t in tensors) - padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors] - return torch.cat(padded, dim=0) - - @property - def modality_mapping(self): - v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device) - a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device) - t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device) - return torch.cat([v_map, a_map, t_map], dim=0) - - def default_coords(self, shape, ref_feat_shape, offset_thw=None): - if offset_thw is None: - offset_thw = [0, 0, 0] - return get_coords( - shape=shape, - ref_feat_shape=ref_feat_shape, - offset_thw=offset_thw, - device=self.device, - dtype=self.default_dtype, - ) - - @property - def coords_mapping(self): - if self.spatial_rope_interpolation == "inter": - video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32) - else: - video_ref_feat_shape = ( - self.t // self.t_patch_size, - self.h // self.patch_size, - self.w // self.patch_size, - ) - - video_coords = self.default_coords( - shape=( - self.t // self.t_patch_size, - self.h // self.patch_size, - self.w // self.patch_size, - ), - ref_feat_shape=video_ref_feat_shape, - ) - - if self.coords_style == "v1": - audio_coords = self.default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(self.t // self.t_patch_size, 1, 1), - ) - text_coords = self.default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(2, 1, 1), - offset_thw=[self.text_offset, 0, 0], - ) - elif self.coords_style == "v2": - magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1 - audio_coords = self.default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1), - ) - text_coords = self.default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(1, 1, 1), - offset_thw=[-self.txt_feat_len, 0, 0], - ) - else: - raise ValueError(f"Unknown coords_style: {self.coords_style}") - - return torch.cat([video_coords, audio_coords, text_coords], dim=0) - - def depack_token_sequence(self, token_sequence): - video_x_t = token_sequence[: self.video_token_num, : self.video_channel] - video_x_t = rearrange( - video_x_t, - "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", - H=self.h // self.patch_size, - W=self.w // self.patch_size, - pT=self.t_patch_size, - pH=self.patch_size, - pW=self.patch_size, - ).contiguous() - audio_x_t = token_sequence[ - self.video_token_num : self.video_token_num + self.audio_feat_len, - : self.audio_channel, - ] - return video_x_t, audio_x_t - - -@dataclass -class SimplePackedData: - items: list[SingleData] - - @property - def token_sequence(self): - return torch.cat([item.token_sequence for item in self.items], dim=0) - - @property - def modality_mapping(self): - return torch.cat([item.modality_mapping for item in self.items], dim=0) - - @property - def coords_mapping(self): - return torch.cat([item.coords_mapping for item in self.items], dim=0) - - @property - def total_token_num(self): - return sum(item.total_token_num for item in self.items) - - def __getitem__(self, index): - return self.items[index] - - @property - def cu_seqlen(self): - cu = torch.cumsum( - torch.tensor([item.total_token_num for item in self.items]), - dim=0, - ) - return F.pad(cu, (1, 0)) - - @property - def max_seqlen(self): - return torch.tensor(max(item.total_token_num for item in self.items)) - - def depack_token_sequence(self, token_sequence): - video_list, audio_list = [], [] - parts = torch.split( - token_sequence, - [item.total_token_num for item in self.items], - dim=0, - ) - for item, part in zip(self.items, parts): - v, a = item.depack_token_sequence(part) - video_list.append(v) - audio_list.append(a) - return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0) - - -class MagiDataProxy: - def __init__( - self, - patch_size: int = 2, - t_patch_size: int = 1, - frame_receptive_field: int = 11, - spatial_rope_interpolation: str = "extra", - ref_audio_offset: int = 1000, - text_offset: int = 0, - coords_style: str = "v2", - ): - self.patch_size = patch_size - self.t_patch_size = t_patch_size - self.frame_receptive_field = frame_receptive_field - self.spatial_rope_interpolation = spatial_rope_interpolation - self.ref_audio_offset = ref_audio_offset - self.text_offset = text_offset - self.coords_style = coords_style - self._kernel = (t_patch_size, patch_size, patch_size) - self._stride = (t_patch_size, patch_size, patch_size) - self._saved_data: dict[str, Any] = {} - - def saved_for_output(self, **kwargs): - self._saved_data.update(kwargs) - - def get_saved_data(self, key: str): - return self._saved_data[key] - - def img2tokens(self, x_t: torch.Tensor): - x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride) - return rearrange( - x_t_unfolded, - "N col_dim num_tokens -> N num_tokens col_dim", - ).contiguous() - - def process_input(self, transported_data: EvalInput): - batch_size, _, t, h, w = transported_data.x_t.shape - x_t = self.img2tokens(transported_data.x_t) - audio_x_t = transported_data.audio_x_t.contiguous() - text_in = transported_data.txt_feat.contiguous() - - simple_packed_data = SimplePackedData(items=[]) - for i in range(batch_size): - single_data = SingleData( - video_x_t=x_t[i], - audio_x_t=audio_x_t[i], - audio_feat_len=transported_data.audio_feat_len[i], - txt_feat=text_in[i], - txt_feat_len=transported_data.txt_feat_len[i], - t=t, - h=h, - w=w, - patch_size=self.patch_size, - t_patch_size=self.t_patch_size, - spatial_rope_interpolation=self.spatial_rope_interpolation, - ref_audio_offset=self.ref_audio_offset, - text_offset=self.text_offset, - coords_style=self.coords_style, - ) - simple_packed_data.items.append(single_data) - - if self.frame_receptive_field != -1: - assert batch_size == 1, "local attention only supports batch size 1" - local_attn_handler = calc_local_attn_ffa_handler( - num_video_tokens=simple_packed_data[0].video_token_num, - num_audio_and_txt_tokens=(simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len), - num_frames=t, - frame_receptive_field=self.frame_receptive_field, - ) - if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor): - local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item() - if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor): - local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item() - else: - local_attn_handler = None - - varlen_handler = VarlenHandler( - cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - ) - - self.saved_for_output(simple_packed_data=simple_packed_data) - - x = simple_packed_data.token_sequence - coords_mapping = simple_packed_data.coords_mapping - modality_mapping = simple_packed_data.modality_mapping - return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler) - - def process_output(self, x: torch.Tensor): - simple_packed_data: SimplePackedData = self.get_saved_data("simple_packed_data") - return simple_packed_data.depack_token_sequence(x) diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py deleted file mode 100644 index 513d93335..000000000 --- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py +++ /dev/null @@ -1,2306 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 SandAI. All Rights Reserved. -# Ported from daVinci-MagiHuman inference/pipeline/video_generate.py -# Adapted for vllm-omni: single-GPU, diffusers VAE, configurable dit_subfolder. - -from __future__ import annotations - -import json -import logging -import math -import os -from collections.abc import Iterable -from dataclasses import dataclass -from pathlib import Path -from typing import Any, ClassVar, Literal - -import numpy as np -import torch -import torch.nn as nn -import whisper -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.schedulers.scheduling_utils import ( - KarrasDiffusionSchedulers, - SchedulerMixin, - SchedulerOutput, -) -from diffusers.utils import deprecate, load_image -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from einops import rearrange -from PIL import Image -from safetensors.torch import load_file -from torch.nn import functional as F -from torch.nn.utils import weight_norm -from transformers import AutoTokenizer -from transformers.models.t5gemma import T5GemmaEncoderModel -from vllm.distributed import ( - get_tensor_model_parallel_world_size, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( - DistributedAutoencoderKLWan, -) -from vllm_omni.diffusion.model_loader.diffusers_loader import ( - DiffusersPipelineLoader, -) -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( - DiffusionPipelineProfilerMixin, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -from .magi_human_dit import ( - DiTModel, - FFAHandler, - MagiHumanDiTConfig, - Modality, - VarlenHandler, -) - -logger = logging.getLogger(__name__) - - -# =========================================================================== -# Scheduler (ported from daVinci-MagiHuman inference/pipeline/scheduler_unipc.py) -# =========================================================================== -class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): - _compatibles = [e.name for e in KarrasDiffusionSchedulers] - order = 1 - - @register_to_config - def __init__( - self, - num_train_timesteps: int = 1000, - solver_order: int = 2, - prediction_type: str = "flow_prediction", - shift: float = 1.0, - use_dynamic_shifting=False, - thresholding: bool = False, - dynamic_thresholding_ratio: float = 0.995, - sample_max_value: float = 1.0, - predict_x0: bool = True, - solver_type: str = "bh2", - lower_order_final: bool = True, - disable_corrector: list[int] = [], - solver_p: SchedulerMixin = None, - timestep_spacing: str = "linspace", - steps_offset: int = 0, - final_sigmas_type: str | None = "zero", - ): - if solver_type not in ["bh1", "bh2"]: - if solver_type in ["midpoint", "heun", "logrho"]: - self.register_to_config(solver_type="bh2") - else: - raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}") - - self.predict_x0 = predict_x0 - self.num_inference_steps = None - alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps)[::-1].copy() - sigmas = 1.0 - alphas - sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32) - - if not use_dynamic_shifting: - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - - self.sigmas = sigmas - self.timesteps = sigmas * num_train_timesteps - - self.model_outputs = [None] * solver_order - self.timestep_list = [None] * solver_order - self.lower_order_nums = 0 - self.disable_corrector = disable_corrector - self.solver_p = solver_p - self.last_sample = None - self._step_index: int | None = None - self._begin_index: int | None = None - - self.sigmas = self.sigmas.to("cpu") - self.sigma_min = self.sigmas[-1].item() - self.sigma_max = self.sigmas[0].item() - - @property - def step_index(self): - return self._step_index - - @property - def begin_index(self): - return self._begin_index - - def set_begin_index(self, begin_index: int = 0): - self._begin_index = begin_index - - def set_timesteps( - self, - num_inference_steps: int | None = None, - device: str | torch.device = None, - sigmas: list[float] | None = None, - mu: float | None | None = None, - shift: float | None | None = None, - ): - if self.config.use_dynamic_shifting and mu is None: - raise ValueError(" you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") - - if sigmas is None: - sigmas = np.linspace(self.sigma_max, self.sigma_min, num_inference_steps + 1).copy()[:-1] - - if self.config.use_dynamic_shifting: - sigmas = self.time_shift(mu, 1.0, sigmas) - else: - if shift is None: - shift = self.config.shift - sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) - - if self.config.final_sigmas_type == "sigma_min": - sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 - elif self.config.final_sigmas_type == "zero": - sigma_last = 0 - else: - raise ValueError( - f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" - ) - - timesteps = sigmas * self.config.num_train_timesteps - sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) - - self.sigmas = torch.from_numpy(sigmas) - self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) - - self.num_inference_steps = len(timesteps) - - self.model_outputs = [None] * self.config.solver_order - self.lower_order_nums = 0 - self.last_sample = None - if self.solver_p: - self.solver_p.set_timesteps(self.num_inference_steps, device=device) - - self._step_index = None - self._begin_index = None - self.sigmas = self.sigmas.to("cpu") - - def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: - dtype = sample.dtype - batch_size, channels, *remaining_dims = sample.shape - - if dtype not in (torch.float32, torch.float64): - sample = sample.float() - - sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) - abs_sample = sample.abs() - s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) - s = torch.clamp(s, min=1, max=self.config.sample_max_value) - s = s.unsqueeze(1) - sample = torch.clamp(sample, -s, s) / s - sample = sample.reshape(batch_size, channels, *remaining_dims) - return sample.to(dtype) - - def _sigma_to_t(self, sigma): - return sigma * self.config.num_train_timesteps - - def _sigma_to_alpha_sigma_t(self, sigma): - return 1 - sigma, sigma - - def time_shift(self, mu: float, sigma: float, t: torch.Tensor): - return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) - - def convert_model_output( - self, model_output: torch.Tensor, *args, sample: torch.Tensor = None, **kwargs - ) -> torch.Tensor: - timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None) - if sample is None: - if len(args) > 1: - sample = args[1] - else: - raise ValueError("missing `sample` as a required keyword argument") - if timestep is not None: - deprecate( - "timesteps", - "1.0.0", - "Passing `timesteps` is deprecated and has no effect as model output " - "conversion is now handled via an internal counter `self.step_index`", - ) - - sigma = self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - - if self.predict_x0: - if self.config.prediction_type == "flow_prediction": - sigma_t = self.sigmas[self.step_index] - x0_pred = sample - sigma_t * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," - " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." - ) - if self.config.thresholding: - x0_pred = self._threshold_sample(x0_pred) - return x0_pred - else: - if self.config.prediction_type == "flow_prediction": - sigma_t = self.sigmas[self.step_index] - epsilon = sample - (1 - sigma_t) * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," - " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." - ) - if self.config.thresholding: - sigma_t = self.sigmas[self.step_index] - x0_pred = sample - sigma_t * model_output - x0_pred = self._threshold_sample(x0_pred) - epsilon = model_output + x0_pred - return epsilon - - def multistep_uni_p_bh_update( - self, - model_output: torch.Tensor, - *args, - sample: torch.Tensor | None = None, - order: int | None = None, - **kwargs, - ) -> torch.Tensor: - prev_timestep = args[0] if len(args) > 0 else kwargs.pop("prev_timestep", None) - if sample is None: - if len(args) > 1: - sample = args[1] - else: - raise ValueError(" missing `sample` as a required keyword argument") - if order is None: - if len(args) > 2: - order = args[2] - else: - raise ValueError(" missing `order` as a required keyword argument") - if prev_timestep is not None: - deprecate("prev_timestep", "1.0.0", "Passing `prev_timestep` is deprecated and has no effect.") - - model_output_list = self.model_outputs - s0 = self.timestep_list[-1] - m0 = model_output_list[-1] - x = sample - - if self.solver_p: - return self.solver_p.step(model_output, s0, x).prev_sample - - sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - h = lambda_t - lambda_s0 - device = sample.device - - rks = [] - D1s: list[Any] | None = [] - for i in range(1, order): - si = self.step_index - i - mi = model_output_list[-(i + 1)] - alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) - rk = (lambda_si - lambda_s0) / h - rks.append(rk) - D1s.append((mi - m0) / rk) - - rks.append(1.0) - rks = torch.tensor(rks, device=device) - - R = [] - b = [] - hh = -h if self.predict_x0 else h - h_phi_1 = torch.expm1(hh) - h_phi_k = h_phi_1 / hh - 1 - factorial_i = 1 - - if self.config.solver_type == "bh1": - B_h = hh - elif self.config.solver_type == "bh2": - B_h = torch.expm1(hh) - else: - raise NotImplementedError() - - for i in range(1, order + 1): - R.append(torch.pow(rks, i - 1)) - b.append(h_phi_k * factorial_i / B_h) - factorial_i *= i + 1 - h_phi_k = h_phi_k / hh - 1 / factorial_i - - R = torch.stack(R) - b = torch.tensor(b, device=device) - - if len(D1s) > 0: - D1s = torch.stack(D1s, dim=1) - if order == 2: - rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) - else: - rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) - else: - D1s = None - - if self.predict_x0: - x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 - pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 - x_t = x_t_ - alpha_t * B_h * pred_res - else: - x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 - pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 - x_t = x_t_ - sigma_t * B_h * pred_res - - return x_t.to(x.dtype) - - def multistep_uni_c_bh_update( - self, - this_model_output: torch.Tensor, - *args, - last_sample: torch.Tensor = None, - this_sample: torch.Tensor = None, - order: int | None = None, - **kwargs, - ) -> torch.Tensor: - this_timestep = args[0] if len(args) > 0 else kwargs.pop("this_timestep", None) - if last_sample is None: - if len(args) > 1: - last_sample = args[1] - else: - raise ValueError(" missing`last_sample` as a required keyword argument") - if this_sample is None: - if len(args) > 2: - this_sample = args[2] - else: - raise ValueError(" missing`this_sample` as a required keyword argument") - if order is None: - if len(args) > 3: - order = args[3] - else: - raise ValueError(" missing`order` as a required keyword argument") - if this_timestep is not None: - deprecate("this_timestep", "1.0.0", "Passing `this_timestep` is deprecated and has no effect.") - - model_output_list = self.model_outputs - m0 = model_output_list[-1] - x = last_sample - x_t = this_sample - model_t = this_model_output - - sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[self.step_index - 1] - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) - alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) - - lambda_t = torch.log(alpha_t) - torch.log(sigma_t) - lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) - h = lambda_t - lambda_s0 - device = this_sample.device - - rks = [] - D1s: list[Any] | None = [] - for i in range(1, order): - si = self.step_index - (i + 1) - mi = model_output_list[-(i + 1)] - alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) - lambda_si = torch.log(alpha_si) - torch.log(sigma_si) - rk = (lambda_si - lambda_s0) / h - rks.append(rk) - D1s.append((mi - m0) / rk) - - rks.append(1.0) - rks = torch.tensor(rks, device=device) - - R = [] - b = [] - hh = -h if self.predict_x0 else h - h_phi_1 = torch.expm1(hh) - h_phi_k = h_phi_1 / hh - 1 - factorial_i = 1 - - if self.config.solver_type == "bh1": - B_h = hh - elif self.config.solver_type == "bh2": - B_h = torch.expm1(hh) - else: - raise NotImplementedError() - - for i in range(1, order + 1): - R.append(torch.pow(rks, i - 1)) - b.append(h_phi_k * factorial_i / B_h) - factorial_i *= i + 1 - h_phi_k = h_phi_k / hh - 1 / factorial_i - - R = torch.stack(R) - b = torch.tensor(b, device=device) - - if len(D1s) > 0: - D1s = torch.stack(D1s, dim=1) - else: - D1s = None - - if order == 1: - rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) - else: - rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) - - if self.predict_x0: - x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 - corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 - D1_t = model_t - m0 - x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) - else: - x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 - corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) if D1s is not None else 0 - D1_t = model_t - m0 - x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) - return x_t.to(x.dtype) - - def index_for_timestep(self, timestep, schedule_timesteps=None): - if schedule_timesteps is None: - schedule_timesteps = self.timesteps - indices = (schedule_timesteps == timestep).nonzero() - pos = 1 if len(indices) > 1 else 0 - return indices[pos].item() - - def _init_step_index(self, timestep): - if self.begin_index is None: - if isinstance(timestep, torch.Tensor): - timestep = timestep.to(self.timesteps.device) - self._step_index = self.index_for_timestep(timestep) - else: - self._step_index = self._begin_index - - def step( - self, - model_output: torch.Tensor, - timestep: int | torch.Tensor, - sample: torch.Tensor, - return_dict: bool = True, - generator=None, - ) -> SchedulerOutput | tuple: - if self.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) - - if self.step_index is None: - self._init_step_index(timestep) - - use_corrector = ( - self.step_index > 0 and self.step_index - 1 not in self.disable_corrector and self.last_sample is not None - ) - - model_output_convert = self.convert_model_output(model_output, sample=sample) - if use_corrector: - sample = self.multistep_uni_c_bh_update( - this_model_output=model_output_convert, - last_sample=self.last_sample, - this_sample=sample, - order=self.this_order, - ) - - for i in range(self.config.solver_order - 1): - self.model_outputs[i] = self.model_outputs[i + 1] - self.timestep_list[i] = self.timestep_list[i + 1] - - self.model_outputs[-1] = model_output_convert - self.timestep_list[-1] = timestep - - if self.config.lower_order_final: - this_order = min(self.config.solver_order, len(self.timesteps) - self.step_index) - else: - this_order = self.config.solver_order - - self.this_order = min(this_order, self.lower_order_nums + 1) - assert self.this_order > 0 - - self.last_sample = sample - prev_sample = self.multistep_uni_p_bh_update(model_output=model_output, sample=sample, order=self.this_order) - - if self.lower_order_nums < self.config.solver_order: - self.lower_order_nums += 1 - - self._step_index += 1 - - if not return_dict: - return (prev_sample,) - return SchedulerOutput(prev_sample=prev_sample) - - def step_ddim( - self, - velocity: torch.FloatTensor, - t: int, - curr_state: torch.FloatTensor, - prev_state: torch.FloatTensor | None = None, - generator: torch.Generator | None = None, - ): - device = curr_state.device - curr_t = self.sigmas[t] - prev_t = self.sigmas[t + 1] - variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype) - cur_clean_ = curr_state - curr_t * velocity - return prev_t * variance_noise + (1 - prev_t) * cur_clean_ - - def step_sde( - self, - velocity: torch.FloatTensor, - t: int, - curr_state: torch.FloatTensor, - noise_theta: float = 1.0, - prev_state: torch.FloatTensor | None = None, - generator: torch.Generator | None = None, - ): - device = curr_state.device - curr_t = self.sigmas[t] - prev_t = self.sigmas[t + 1] - cos = torch.cos(torch.tensor(noise_theta) * torch.pi / 2).to(device) - sin = torch.sin(torch.tensor(noise_theta) * torch.pi / 2).to(device) - prev_sample_mean = (1 - prev_t + prev_t * cos) * (curr_state - curr_t * velocity) + prev_t * cos * velocity - std_dev_t = prev_t * sin - std_dev_t = torch.ones((1, 1)).to(curr_state) * std_dev_t - if prev_state is None: - variance_noise = randn_tensor(curr_state.shape, generator=generator, device=device, dtype=curr_state.dtype) - prev_state = prev_sample_mean + std_dev_t * variance_noise - else: - prev_state = prev_sample_mean + (prev_state - prev_sample_mean.detach()) - return prev_state - - def scale_model_input(self, sample: torch.Tensor, *args, **kwargs) -> torch.Tensor: - return sample - - def add_noise( - self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor - ) -> torch.Tensor: - sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) - if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): - schedule_timesteps = self.timesteps.to(original_samples.device, dtype=torch.float32) - timesteps = timesteps.to(original_samples.device, dtype=torch.float32) - else: - schedule_timesteps = self.timesteps.to(original_samples.device) - timesteps = timesteps.to(original_samples.device) - - if self.begin_index is None: - step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timesteps] - elif self.step_index is not None: - step_indices = [self.step_index] * timesteps.shape[0] - else: - step_indices = [self.begin_index] * timesteps.shape[0] - - sigma = sigmas[step_indices].flatten() - while len(sigma.shape) < len(original_samples.shape): - sigma = sigma.unsqueeze(-1) - - alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) - return alpha_t * original_samples + sigma_t * noise - - def __len__(self): - return self.config.num_train_timesteps - - -# =========================================================================== -# Audio VAE (ported from daVinci-MagiHuman inference/model/sa_audio/) -# =========================================================================== -def _snake_beta(x, alpha, beta): - return x + (1.0 / (beta + 1e-9)) * torch.pow(torch.sin(x * alpha), 2) - - -class _SnakeBeta(nn.Module): - def __init__(self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True, alpha_logscale: bool = True): - super().__init__() - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: - self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) - self.beta = nn.Parameter(torch.zeros(in_features) * alpha) - else: - self.alpha = nn.Parameter(torch.ones(in_features) * alpha) - self.beta = nn.Parameter(torch.ones(in_features) * alpha) - self.alpha.requires_grad = alpha_trainable - self.beta.requires_grad = alpha_trainable - - def forward(self, x): - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) - beta = self.beta.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - beta = torch.exp(beta) - return _snake_beta(x, alpha, beta) - - -def _vae_sample(mean, scale): - stdev = F.softplus(scale) + 1e-4 - var = stdev * stdev - logvar = torch.log(var) - latents = torch.randn_like(mean) * stdev + mean - kl = (mean * mean + var - logvar - 1).sum(1).mean() - return latents, kl - - -class _VAEBottleneck(nn.Module): - def encode(self, x, return_info=False, **kwargs): - info = {} - mean, scale = x.chunk(2, dim=1) - x, kl = _vae_sample(mean, scale) - info["kl"] = kl - return (x, info) if return_info else x - - def decode(self, x): - return x - - -def _WNConv1d(*args, **kwargs): - return weight_norm(nn.Conv1d(*args, **kwargs)) - - -def _WNConvTranspose1d(*args, **kwargs): - return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) - - -def _checkpoint(function, *args, **kwargs): - kwargs.setdefault("use_reentrant", False) - return torch.utils.checkpoint.checkpoint(function, *args, **kwargs) - - -def _get_activation(activation: Literal["elu", "snake", "none"], antialias: bool = False, channels=None) -> nn.Module: - if antialias: - raise NotImplementedError("antialias activation not supported") - if activation == "elu": - return nn.ELU() - if activation == "snake": - return _SnakeBeta(channels) - if activation == "none": - return nn.Identity() - raise ValueError(f"Unknown activation {activation}") - - -class _ResidualUnit(nn.Module): - def __init__(self, in_channels, out_channels, dilation, use_snake=False, antialias_activation=False): - super().__init__() - padding = (dilation * (7 - 1)) // 2 - self.layers = nn.Sequential( - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels), - _WNConv1d(in_channels, out_channels, kernel_size=7, dilation=dilation, padding=padding), - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels), - _WNConv1d(out_channels, out_channels, kernel_size=1), - ) - - def forward(self, x): - return (_checkpoint(self.layers, x) if self.training else self.layers(x)) + x - - -class _EncoderBlock(nn.Module): - def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False): - super().__init__() - self.layers = nn.Sequential( - _ResidualUnit(in_channels, in_channels, 1, use_snake=use_snake), - _ResidualUnit(in_channels, in_channels, 3, use_snake=use_snake), - _ResidualUnit(in_channels, in_channels, 9, use_snake=use_snake), - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels), - _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2)), - ) - - def forward(self, x): - return self.layers(x) - - -class _DecoderBlock(nn.Module): - def __init__( - self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False, use_nearest_upsample=False - ): - super().__init__() - if use_nearest_upsample: - upsample_layer = nn.Sequential( - nn.Upsample(scale_factor=stride, mode="nearest"), - _WNConv1d(in_channels, out_channels, kernel_size=2 * stride, stride=1, bias=False, padding="same"), - ) - else: - upsample_layer = _WNConvTranspose1d( - in_channels, out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2) - ) - self.layers = nn.Sequential( - _get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels), - upsample_layer, - _ResidualUnit(out_channels, out_channels, 1, use_snake=use_snake), - _ResidualUnit(out_channels, out_channels, 3, use_snake=use_snake), - _ResidualUnit(out_channels, out_channels, 9, use_snake=use_snake), - ) - - def forward(self, x): - return self.layers(x) - - -class _OobleckEncoder(nn.Module): - def __init__( - self, - in_channels=2, - channels=128, - latent_dim=32, - c_mults=[1, 2, 4, 8], - strides=[2, 4, 8, 8], - use_snake=False, - antialias_activation=False, - ): - super().__init__() - c_mults = [1] + c_mults - depth = len(c_mults) - layers = [_WNConv1d(in_channels, c_mults[0] * channels, kernel_size=7, padding=3)] - for i in range(depth - 1): - layers.append( - _EncoderBlock(c_mults[i] * channels, c_mults[i + 1] * channels, strides[i], use_snake=use_snake) - ) - layers.extend( - [ - _get_activation( - "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[-1] * channels - ), - _WNConv1d(c_mults[-1] * channels, latent_dim, kernel_size=3, padding=1), - ] - ) - self.layers = nn.Sequential(*layers) - - def forward(self, x): - return self.layers(x) - - -class _OobleckDecoder(nn.Module): - def __init__( - self, - out_channels=2, - channels=128, - latent_dim=32, - c_mults=[1, 2, 4, 8], - strides=[2, 4, 8, 8], - use_snake=False, - antialias_activation=False, - use_nearest_upsample=False, - final_tanh=True, - ): - super().__init__() - c_mults = [1] + c_mults - depth = len(c_mults) - layers = [_WNConv1d(latent_dim, c_mults[-1] * channels, kernel_size=7, padding=3)] - for i in range(depth - 1, 0, -1): - layers.append( - _DecoderBlock( - c_mults[i] * channels, - c_mults[i - 1] * channels, - strides[i - 1], - use_snake=use_snake, - antialias_activation=antialias_activation, - use_nearest_upsample=use_nearest_upsample, - ) - ) - layers.extend( - [ - _get_activation( - "snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[0] * channels - ), - _WNConv1d(c_mults[0] * channels, out_channels, kernel_size=7, padding=3, bias=False), - nn.Tanh() if final_tanh else nn.Identity(), - ] - ) - self.layers = nn.Sequential(*layers) - - def forward(self, x): - return self.layers(x) - - -class _AudioAutoencoder(nn.Module): - def __init__( - self, - encoder, - decoder, - latent_dim, - downsampling_ratio, - sample_rate, - io_channels=2, - bottleneck=None, - in_channels=None, - out_channels=None, - soft_clip=False, - ): - super().__init__() - self.downsampling_ratio = downsampling_ratio - self.sample_rate = sample_rate - self.latent_dim = latent_dim - self.io_channels = io_channels - self.in_channels = in_channels if in_channels is not None else io_channels - self.out_channels = out_channels if out_channels is not None else io_channels - self.bottleneck = bottleneck - self.encoder = encoder - self.decoder = decoder - self.soft_clip = soft_clip - - def encode(self, audio, skip_bottleneck=False, return_info=False, **kwargs): - info = {} - latents = self.encoder(audio) - info["pre_bottleneck_latents"] = latents - if self.bottleneck is not None and not skip_bottleneck: - latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs) - info.update(bottleneck_info) - return (latents, info) if return_info else latents - - def decode(self, latents, skip_bottleneck=False, **kwargs): - if self.bottleneck is not None and not skip_bottleneck: - latents = self.bottleneck.decode(latents) - decoded = self.decoder(latents, **kwargs) - if self.soft_clip: - decoded = torch.tanh(decoded) - return decoded - - -def _create_encoder_from_config(cfg: dict[str, Any]): - assert cfg.get("type") == "oobleck", f"Only 'oobleck' encoder supported, got: {cfg.get('type')}" - enc = _OobleckEncoder(**cfg["config"]) - if not cfg.get("requires_grad", True): - for p in enc.parameters(): - p.requires_grad = False - return enc - - -def _create_decoder_from_config(cfg: dict[str, Any]): - assert cfg.get("type") == "oobleck", f"Only 'oobleck' decoder supported, got: {cfg.get('type')}" - dec = _OobleckDecoder(**cfg["config"]) - if not cfg.get("requires_grad", True): - for p in dec.parameters(): - p.requires_grad = False - return dec - - -def _create_bottleneck_from_config(cfg: dict[str, Any]): - assert cfg.get("type") == "vae", f"Only 'vae' bottleneck supported, got: {cfg.get('type')}" - bn = _VAEBottleneck() - if not cfg.get("requires_grad", True): - for p in bn.parameters(): - p.requires_grad = False - return bn - - -def _create_autoencoder_from_config(config: dict[str, Any]): - ae_config = config["model"] - if ae_config.get("pretransform") is not None: - raise NotImplementedError("Nested pretransform not supported") - encoder = _create_encoder_from_config(ae_config["encoder"]) - decoder = _create_decoder_from_config(ae_config["decoder"]) - bottleneck_cfg = ae_config.get("bottleneck") - bottleneck = _create_bottleneck_from_config(bottleneck_cfg) if bottleneck_cfg else None - return _AudioAutoencoder( - encoder=encoder, - decoder=decoder, - latent_dim=ae_config["latent_dim"], - downsampling_ratio=ae_config["downsampling_ratio"], - sample_rate=config["sample_rate"], - io_channels=ae_config["io_channels"], - bottleneck=bottleneck, - in_channels=ae_config.get("in_channels"), - out_channels=ae_config.get("out_channels"), - soft_clip=ae_config["decoder"].get("soft_clip", False), - ) - - -class SAAudioFeatureExtractor: - def __init__(self, device, model_path): - self.device = device - self.vae_model, self.sample_rate = self._load_vae(model_path) - self.resampler = None - - def _load_vae(self, model_path): - if not (isinstance(model_path, str) and Path(model_path).is_dir()): - raise ValueError("model_path must be a local directory") - - model_config_path = os.path.join(model_path, "model_config.json") - with open(model_config_path) as f: - full_config = json.load(f) - - vae_config = full_config["model"]["pretransform"]["config"] - sample_rate = full_config["sample_rate"] - - autoencoder_config = { - "model_type": "autoencoder", - "sample_rate": sample_rate, - "model": vae_config, - } - vae_model = _create_autoencoder_from_config(autoencoder_config) - - weights_path = Path(model_path) / "model.safetensors" - if not weights_path.exists(): - raise FileNotFoundError(f"Weight file does not exist: {weights_path}") - - full_state_dict = load_file(weights_path, device=str(self.device)) - vae_state_dict = {} - for key, value in full_state_dict.items(): - if key.startswith("pretransform.model."): - vae_state_dict[key[len("pretransform.model.") :]] = value - - model_keys = set(vae_model.state_dict().keys()) - vae_keys = set(vae_state_dict.keys()) - missing = model_keys - vae_keys - extra = vae_keys - model_keys - if missing: - logger.warning("Audio VAE missing keys (%d): %s", len(missing), list(missing)[:5]) - if extra: - logger.warning("Audio VAE unexpected keys (%d): %s", len(extra), list(extra)[:5]) - - vae_model.load_state_dict(vae_state_dict) - vae_model.to(self.device) - return vae_model, sample_rate - - def decode(self, latents): - with torch.no_grad(): - return self.vae_model.decode(latents) - - def encode(self, waveform): - with torch.no_grad(): - return self.vae_model.encode(waveform) - - -# =========================================================================== -# Audio utilities (ported from daVinci-MagiHuman inference/pipeline/video_process.py) -# =========================================================================== -_SAMPLE_RATE = 51200 -_AUDIO_CHUNK_DURATION = 29 -_OVERLAP_RATIO = 0.5 - - -def _merge_overlapping_vae_features(audio_feats: list[torch.Tensor], overlap_ratio: float = 0.5) -> torch.Tensor | None: - if not audio_feats: - return None - if len(audio_feats) == 1: - return audio_feats[0] - - batch_size, total_frames, feature_dim = audio_feats[0].shape - overlap_frames = int(total_frames * overlap_ratio) - step_frames = total_frames - overlap_frames - final_length = (len(audio_feats) - 1) * step_frames + total_frames - output_feat = torch.zeros( - batch_size, final_length, feature_dim, device=audio_feats[0].device, dtype=audio_feats[0].dtype - ) - - for block_idx, current_feat in enumerate(audio_feats): - output_start = block_idx * step_frames - if block_idx == 0: - output_feat[:, output_start : output_start + total_frames, :] = current_feat - continue - - non_overlap_start = output_start + overlap_frames - non_overlap_end = output_start + total_frames - output_feat[:, non_overlap_start:non_overlap_end, :] = current_feat[:, overlap_frames:, :] - - for frame_idx in range(overlap_frames): - output_pos = output_start + frame_idx - prev_weight = (overlap_frames - frame_idx) / overlap_frames - curr_weight = frame_idx / overlap_frames - output_feat[:, output_pos, :] = ( - prev_weight * output_feat[:, output_pos, :] + curr_weight * current_feat[:, frame_idx, :] - ) - return output_feat - - -def load_audio_and_encode(audio_vae, audio_path: str, seconds: int | None = None) -> torch.Tensor: - """Load audio from file and encode to latent space using the Stable Audio VAE.""" - audio_full = whisper.load_audio(audio_path, sr=_SAMPLE_RATE) - if seconds is not None: - audio_full = audio_full[: min(int(seconds * _SAMPLE_RATE), audio_full.shape[0])] - total_samples = audio_full.shape[0] - - window_size = int(_AUDIO_CHUNK_DURATION * _SAMPLE_RATE) - step_size = int(window_size * (1 - _OVERLAP_RATIO)) - if total_samples <= window_size: - audio = torch.from_numpy(audio_full).cuda() - audio = audio.unsqueeze(0).expand(2, -1) - return audio_vae.vae_model.encode(audio) - - encoded_chunks = [] - latent_to_audio_ratio = None - for offset_start in range(0, total_samples, step_size): - offset_end = min(offset_start + window_size, total_samples) - chunk = whisper.pad_or_trim(audio_full[offset_start:offset_end], length=window_size) - chunk_tensor = torch.from_numpy(chunk).cuda().unsqueeze(0).expand(2, -1) - encoded_chunk = audio_vae.vae_model.encode(chunk_tensor) - - if latent_to_audio_ratio is None: - latent_to_audio_ratio = encoded_chunk.shape[-1] / window_size - - encoded_chunks.append(encoded_chunk.permute(0, 2, 1)) - if offset_end >= total_samples: - break - - final_feat = _merge_overlapping_vae_features(encoded_chunks, overlap_ratio=_OVERLAP_RATIO).permute(0, 2, 1) - final_target_len = math.ceil(total_samples * latent_to_audio_ratio) - return final_feat[:, :, :final_target_len] - - -# =========================================================================== -# Data proxy (ported from daVinci-MagiHuman inference/pipeline/data_proxy.py) -# =========================================================================== -def _unfold_3d(x: torch.Tensor, kernel_size: tuple[int, int, int], stride: tuple[int, int, int]) -> torch.Tensor: - """Pure-PyTorch 3D unfold matching UnfoldAnd behavior. - - After N unfold ops the shape is (batch, C, oD, oH, oW, kD, kH, kW). - UnfoldAnd permutes kernel dims next to channel before reshape so that the - col_dim axis is ordered as (C, kD, kH, kW) -- matching F.unfold semantics. - Without this permute, .view() interleaves spatial and kernel positions. - - Args: - x: (N, C, D, H, W) - kernel_size: (kD, kH, kW) - stride: (sD, sH, sW) - Returns: - (N, C*kD*kH*kW, L) where L = product of output spatial dims. - """ - ndim = len(kernel_size) - for d in range(ndim): - x = x.unfold(d + 2, kernel_size[d], stride[d]) - perm = [0, 1] + list(range(ndim + 2, 2 * ndim + 2)) + list(range(2, ndim + 2)) - x = x.permute(*perm).contiguous() - - batch_size = x.shape[0] - col_dim = 1 - for i in range(1, ndim + 2): - col_dim *= x.shape[i] - spatial = 1 - for i in range(ndim + 2, 2 * ndim + 2): - spatial *= x.shape[i] - return x.view(batch_size, col_dim, spatial) - - -def _calc_local_qk_range(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field): - token_per_frame = num_video_tokens // num_frames - total_tokens = num_video_tokens + num_audio_and_txt_tokens - - q_range_list = [] - k_range_list = [] - for i in range(num_frames): - q_range_list.append(torch.tensor([i * token_per_frame, (i + 1) * token_per_frame])) - k_range_list.append( - torch.tensor( - [ - (i - frame_receptive_field) * token_per_frame, - (i + frame_receptive_field + 1) * token_per_frame, - ] - ) - ) - local_q_range = torch.stack(q_range_list, dim=0) - local_k_range = torch.stack(k_range_list, dim=0) - - local_k_range[local_k_range < 0] = 0 - local_k_range[local_k_range > num_video_tokens] = num_video_tokens - - video_q_range = torch.tensor([[0, num_video_tokens]]) - video_k_range = torch.tensor([[num_video_tokens, num_video_tokens + num_audio_and_txt_tokens]]) - - at_q_ranges = torch.tensor([[num_video_tokens, total_tokens]]) - at_k_ranges = torch.tensor([[0, total_tokens]]) - - q_ranges = ( - torch.cat([local_q_range, video_q_range, at_q_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - k_ranges = ( - torch.cat([local_k_range, video_k_range, at_k_ranges], dim=0).to(torch.int32).to("cuda", non_blocking=True) - ) - return q_ranges, k_ranges - - -def _calc_local_attn_ffa_handler(num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field): - q_ranges, k_ranges = _calc_local_qk_range( - num_video_tokens, num_audio_and_txt_tokens, num_frames, frame_receptive_field - ) - total = num_video_tokens + num_audio_and_txt_tokens - return FFAHandler( - q_ranges=q_ranges, - k_ranges=k_ranges, - max_seqlen_q=total, - max_seqlen_k=total, - attn_type_map=torch.zeros([q_ranges.shape[0]], device="cuda", dtype=torch.int32), - softmax_scale=None, - ) - - -def _get_coords( - shape: list[int], - ref_feat_shape: list[int], - offset_thw: list[int] | None = None, - device: torch.device = torch.device("cpu"), - dtype: torch.dtype = torch.float32, -): - if offset_thw is None: - offset_thw = [0, 0, 0] - ori_t, ori_h, ori_w = shape - ref_t, ref_h, ref_w = ref_feat_shape - - offset_t, offset_h, offset_w = offset_thw - time_rng = torch.arange(ori_t, device=device, dtype=dtype) + offset_t - height_rng = torch.arange(ori_h, device=device, dtype=dtype) + offset_h - width_rng = torch.arange(ori_w, device=device, dtype=dtype) + offset_w - - time_grid, height_grid, width_grid = torch.meshgrid(time_rng, height_rng, width_rng, indexing="ij") - coords_flat = torch.stack([time_grid, height_grid, width_grid], dim=-1).reshape(-1, 3) - - meta = torch.tensor([ori_t, ori_h, ori_w, ref_t, ref_h, ref_w], device=device, dtype=dtype) - meta_expanded = meta.expand(coords_flat.size(0), -1) - return torch.cat([coords_flat, meta_expanded], dim=-1) - - -@dataclass -class _SingleData: - video_x_t: torch.Tensor - audio_x_t: torch.Tensor - audio_feat_len: int - txt_feat: torch.Tensor - txt_feat_len: int - t: int - h: int - w: int - patch_size: int - t_patch_size: int - spatial_rope_interpolation: Literal["inter", "extra"] - ref_audio_offset: int - text_offset: int - coords_style: Literal["v1", "v2"] = "v1" - - def __post_init__(self): - self.video_token_num = self.video_x_t.shape[0] - self.audio_x_t = self.audio_x_t[: self.audio_feat_len] - self.txt_feat = self.txt_feat[: self.txt_feat_len] - self.video_channel = self.video_x_t.shape[-1] - self.audio_channel = self.audio_x_t.shape[-1] - self.txt_channel = self.txt_feat.shape[-1] - - @property - def device(self): - return self.video_x_t.device - - @property - def default_dtype(self): - return self.video_x_t.dtype - - @property - def total_token_num(self): - return self.video_token_num + self.audio_feat_len + self.txt_feat_len - - @property - def token_sequence(self): - tensors = [self.video_x_t, self.audio_x_t, self.txt_feat] - max_channel = max(t.shape[-1] for t in tensors) - padded = [F.pad(t, (0, max_channel - t.shape[-1])) for t in tensors] - return torch.cat(padded, dim=0) - - @property - def modality_mapping(self): - v_map = torch.full((self.video_token_num,), Modality.VIDEO, dtype=torch.int64, device=self.device) - a_map = torch.full((self.audio_feat_len,), Modality.AUDIO, dtype=torch.int64, device=self.device) - t_map = torch.full((self.txt_feat_len,), Modality.TEXT, dtype=torch.int64, device=self.device) - return torch.cat([v_map, a_map, t_map], dim=0) - - def _default_coords(self, shape, ref_feat_shape, offset_thw=None): - if offset_thw is None: - offset_thw = [0, 0, 0] - return _get_coords( - shape=shape, - ref_feat_shape=ref_feat_shape, - offset_thw=offset_thw, - device=self.device, - dtype=self.default_dtype, - ) - - @property - def coords_mapping(self): - if self.spatial_rope_interpolation == "inter": - video_ref_feat_shape = (self.t // self.t_patch_size, 32, 32) - else: - video_ref_feat_shape = (self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size) - - video_coords = self._default_coords( - shape=(self.t // self.t_patch_size, self.h // self.patch_size, self.w // self.patch_size), - ref_feat_shape=video_ref_feat_shape, - ) - - if self.coords_style == "v1": - audio_coords = self._default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(self.t // self.t_patch_size, 1, 1), - ) - text_coords = self._default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(2, 1, 1), - offset_thw=[self.text_offset, 0, 0], - ) - elif self.coords_style == "v2": - magic_audio_ref_t = (self.audio_feat_len - 1) // 4 + 1 - audio_coords = self._default_coords( - shape=(self.audio_feat_len, 1, 1), - ref_feat_shape=(magic_audio_ref_t // self.t_patch_size, 1, 1), - ) - text_coords = self._default_coords( - shape=(self.txt_feat_len, 1, 1), - ref_feat_shape=(1, 1, 1), - offset_thw=[-self.txt_feat_len, 0, 0], - ) - else: - raise ValueError(f"Unknown coords_style: {self.coords_style}") - - return torch.cat([video_coords, audio_coords, text_coords], dim=0) - - def depack_token_sequence(self, token_sequence): - video_x_t = token_sequence[: self.video_token_num, : self.video_channel] - video_x_t = rearrange( - video_x_t, - "(T H W) (pT pH pW C) -> C (T pT) (H pH) (W pW)", - H=self.h // self.patch_size, - W=self.w // self.patch_size, - pT=self.t_patch_size, - pH=self.patch_size, - pW=self.patch_size, - ).contiguous() - audio_x_t = token_sequence[ - self.video_token_num : self.video_token_num + self.audio_feat_len, : self.audio_channel - ] - return video_x_t, audio_x_t - - -@dataclass -class _SimplePackedData: - items: list[_SingleData] - - @property - def token_sequence(self): - return torch.cat([item.token_sequence for item in self.items], dim=0) - - @property - def modality_mapping(self): - return torch.cat([item.modality_mapping for item in self.items], dim=0) - - @property - def coords_mapping(self): - return torch.cat([item.coords_mapping for item in self.items], dim=0) - - @property - def total_token_num(self): - return sum(item.total_token_num for item in self.items) - - def __getitem__(self, index): - return self.items[index] - - @property - def cu_seqlen(self): - cu = torch.cumsum(torch.tensor([item.total_token_num for item in self.items]), dim=0) - return F.pad(cu, (1, 0)) - - @property - def max_seqlen(self): - return torch.tensor(max(item.total_token_num for item in self.items)) - - def depack_token_sequence(self, token_sequence): - video_list, audio_list = [], [] - parts = torch.split(token_sequence, [item.total_token_num for item in self.items], dim=0) - for item, part in zip(self.items, parts): - v, a = item.depack_token_sequence(part) - video_list.append(v) - audio_list.append(a) - return torch.stack(video_list, dim=0), torch.stack(audio_list, dim=0) - - -class MagiDataProxy: - def __init__( - self, - patch_size: int = 2, - t_patch_size: int = 1, - frame_receptive_field: int = 11, - spatial_rope_interpolation: str = "extra", - ref_audio_offset: int = 1000, - text_offset: int = 0, - coords_style: str = "v2", - ): - self.patch_size = patch_size - self.t_patch_size = t_patch_size - self.frame_receptive_field = frame_receptive_field - self.spatial_rope_interpolation = spatial_rope_interpolation - self.ref_audio_offset = ref_audio_offset - self.text_offset = text_offset - self.coords_style = coords_style - self._kernel = (t_patch_size, patch_size, patch_size) - self._stride = (t_patch_size, patch_size, patch_size) - self._saved_data: dict[str, Any] = {} - - def saved_for_output(self, **kwargs): - self._saved_data.update(kwargs) - - def get_saved_data(self, key: str): - return self._saved_data[key] - - def img2tokens(self, x_t: torch.Tensor): - x_t_unfolded = _unfold_3d(x_t, self._kernel, self._stride) - return rearrange(x_t_unfolded, "N col_dim num_tokens -> N num_tokens col_dim").contiguous() - - def process_input(self, transported_data: EvalInput): - batch_size, _, t, h, w = transported_data.x_t.shape - x_t = self.img2tokens(transported_data.x_t) - audio_x_t = transported_data.audio_x_t.contiguous() - text_in = transported_data.txt_feat.contiguous() - - simple_packed_data = _SimplePackedData(items=[]) - for i in range(batch_size): - single_data = _SingleData( - video_x_t=x_t[i], - audio_x_t=audio_x_t[i], - audio_feat_len=transported_data.audio_feat_len[i], - txt_feat=text_in[i], - txt_feat_len=transported_data.txt_feat_len[i], - t=t, - h=h, - w=w, - patch_size=self.patch_size, - t_patch_size=self.t_patch_size, - spatial_rope_interpolation=self.spatial_rope_interpolation, - ref_audio_offset=self.ref_audio_offset, - text_offset=self.text_offset, - coords_style=self.coords_style, - ) - simple_packed_data.items.append(single_data) - - if self.frame_receptive_field != -1: - assert batch_size == 1, "local attention only supports batch size 1" - local_attn_handler = _calc_local_attn_ffa_handler( - num_video_tokens=simple_packed_data[0].video_token_num, - num_audio_and_txt_tokens=simple_packed_data[0].audio_feat_len + simple_packed_data[0].txt_feat_len, - num_frames=t, - frame_receptive_field=self.frame_receptive_field, - ) - if isinstance(local_attn_handler.max_seqlen_k, torch.Tensor): - local_attn_handler.max_seqlen_k = local_attn_handler.max_seqlen_k.item() - if isinstance(local_attn_handler.max_seqlen_q, torch.Tensor): - local_attn_handler.max_seqlen_q = local_attn_handler.max_seqlen_q.item() - else: - local_attn_handler = None - - varlen_handler = VarlenHandler( - cu_seqlens_q=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - cu_seqlens_k=simple_packed_data.cu_seqlen.to(torch.int32).cuda(), - max_seqlen_q=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - max_seqlen_k=simple_packed_data.max_seqlen.to(torch.int32).cuda(), - ) - - self.saved_for_output(simple_packed_data=simple_packed_data) - - x = simple_packed_data.token_sequence - coords_mapping = simple_packed_data.coords_mapping - modality_mapping = simple_packed_data.modality_mapping - return (x, coords_mapping, modality_mapping, varlen_handler, local_attn_handler) - - def process_output(self, x: torch.Tensor): - simple_packed_data: _SimplePackedData = self.get_saved_data("simple_packed_data") - return simple_packed_data.depack_token_sequence(x) - - -# =========================================================================== -# Pipeline helpers -# =========================================================================== -@dataclass -class EvalInput: - x_t: torch.Tensor - audio_x_t: torch.Tensor - audio_feat_len: torch.Tensor | list[int] - txt_feat: torch.Tensor - txt_feat_len: torch.Tensor | list[int] - - -class _T5GemmaEncoder: - def __init__(self, model_path: str, device: str, weight_dtype: torch.dtype, subfolder: str | None = None): - from vllm.distributed import get_tensor_model_parallel_world_size - - self.device = device - hf_kwargs: dict[str, Any] = {} - if subfolder is not None: - hf_kwargs["subfolder"] = subfolder - self.tokenizer = AutoTokenizer.from_pretrained(model_path, **hf_kwargs) - - tp_size = get_tensor_model_parallel_world_size() - if tp_size > 1: - from transformers.models.t5gemma.configuration_t5gemma import T5GemmaConfig - - config = T5GemmaConfig.from_pretrained(model_path, **hf_kwargs) - # The config we need is the encoder config - config_encoder = config.encoder - # Propagate some outer config values - config_encoder.vocab_size = config.vocab_size - config_encoder.rms_norm_eps = getattr(config, "rms_norm_eps", config_encoder.rms_norm_eps) - self.model = T5GemmaEncoderModelTP(config_encoder).to(device).to(weight_dtype) - self.is_tp = True - else: - self.model = T5GemmaEncoderModel.from_pretrained( - model_path, is_encoder_decoder=False, dtype=weight_dtype, **hf_kwargs - ).to(device) - self.is_tp = False - - @torch.inference_mode() - def encode(self, prompt: str) -> torch.Tensor: - inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device) - outputs = self.model(**inputs) - - if self.is_tp: - # T5GemmaEncoderModelTP just returns the hidden states tensor - return outputs.half() - else: - # HF model returns BaseModelOutput - return outputs["last_hidden_state"].half() - - -def _pad_or_trim(tensor: torch.Tensor, target_size: int, dim: int, pad_value: float = 0.0) -> tuple[torch.Tensor, int]: - current_size = tensor.size(dim) - if current_size < target_size: - padding_amount = target_size - current_size - padding_tuple = [0] * (2 * tensor.dim()) - padding_dim_index = tensor.dim() - 1 - dim - padding_tuple[2 * padding_dim_index + 1] = padding_amount - return F.pad(tensor, tuple(padding_tuple), "constant", pad_value), current_size - slicing = [slice(None)] * tensor.dim() - slicing[dim] = slice(0, target_size) - return tensor[tuple(slicing)], target_size - - -def _get_padded_t5_gemma_embedding( - prompt: str, - encoder: _T5GemmaEncoder, - target_length: int, -) -> tuple[torch.Tensor, int]: - txt_feat = encoder.encode(prompt) - txt_feat, original_len = _pad_or_trim(txt_feat, target_size=target_length, dim=1) - return txt_feat.to(torch.float32), original_len - - -def _resizecrop(img: Image.Image, target_height: int, target_width: int) -> Image.Image: - """Centre-crop resize keeping aspect ratio then letterbox to target.""" - pil_image = img.convert("RGB") - original_width, original_height = pil_image.size - scale_x = target_width / original_width - scale_y = target_height / original_height - scale = max(scale_x, scale_y) - new_width = int(round(original_width * scale)) - new_height = int(round(original_height * scale)) - resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS) - left = (new_width - target_width) // 2 - top = (new_height - target_height) // 2 - return resized_image.crop((left, top, left + target_width, top + target_height)) - - -class ZeroSNRDDPMDiscretization: - """ZeroSNR DDPM sigma schedule, ported from daVinci-MagiHuman. - Used to compute sigma values for SR noise injection. - """ - - def __init__( - self, - linear_start: float = 0.00085, - linear_end: float = 0.0120, - num_timesteps: int = 1000, - shift_scale: float = 1.0, - keep_start: bool = False, - post_shift: bool = False, - ): - from functools import partial - - if keep_start and not post_shift: - linear_start = linear_start / (shift_scale + (1 - shift_scale) * linear_start) - self.num_timesteps = num_timesteps - betas = torch.linspace(linear_start**0.5, linear_end**0.5, num_timesteps, dtype=torch.float64) ** 2 - alphas = 1.0 - betas.cpu().numpy() - self.alphas_cumprod = np.cumprod(alphas, axis=0) - self.to_torch = partial(torch.tensor, dtype=torch.float32) - if not post_shift: - self.alphas_cumprod = self.alphas_cumprod / (shift_scale + (1 - shift_scale) * self.alphas_cumprod) - self.post_shift = post_shift - self.shift_scale = shift_scale - - def __call__( - self, - n: int, - do_append_zero: bool = True, - device: str = "cpu", - flip: bool = False, - return_idx: bool = False, - ): - from functools import partial - - if n < self.num_timesteps: - timesteps = np.linspace(self.num_timesteps - 1, 0, n, endpoint=False).astype(int)[::-1] - alphas_cumprod = self.alphas_cumprod[timesteps] - elif n == self.num_timesteps: - alphas_cumprod = self.alphas_cumprod - else: - raise ValueError(f"n={n} > num_timesteps={self.num_timesteps}") - - to_torch = partial(torch.tensor, dtype=torch.float32, device=device) - alphas_cumprod = to_torch(alphas_cumprod) - alphas_cumprod_sqrt = alphas_cumprod.sqrt() - alphas_cumprod_sqrt_0 = alphas_cumprod_sqrt[0].clone() - alphas_cumprod_sqrt_T = alphas_cumprod_sqrt[-1].clone() - alphas_cumprod_sqrt -= alphas_cumprod_sqrt_T - alphas_cumprod_sqrt *= alphas_cumprod_sqrt_0 / (alphas_cumprod_sqrt_0 - alphas_cumprod_sqrt_T) - - if self.post_shift: - alphas_cumprod_sqrt = ( - alphas_cumprod_sqrt**2 / (self.shift_scale + (1 - self.shift_scale) * alphas_cumprod_sqrt**2) - ) ** 0.5 - - sigmas = torch.flip(alphas_cumprod_sqrt, (0,)) - sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) if do_append_zero else sigmas - if return_idx: - return sigmas if not flip else torch.flip(sigmas, (0,)), timesteps - return sigmas if not flip else torch.flip(sigmas, (0,)) - - -def _schedule_latent_step( - *, - video_scheduler: FlowUniPCMultistepScheduler, - audio_scheduler: FlowUniPCMultistepScheduler, - latent_video: torch.Tensor, - latent_audio: torch.Tensor, - t, - idx: int, - steps, - v_cfg_video: torch.Tensor, - v_cfg_audio: torch.Tensor, - is_a2v: bool, - cfg_number: int, - using_sde_flag: bool, - use_sr_model: bool = False, -): - # Fast DDIM path for cfg_number==1, only used during the BR stage - if cfg_number == 1 and not use_sr_model: - latent_video = video_scheduler.step_ddim(v_cfg_video, idx, latent_video) - latent_audio = audio_scheduler.step_ddim(v_cfg_audio, idx, latent_audio) - return latent_video, latent_audio - - if using_sde_flag: - if use_sr_model: - # SR stage with SDE: only update video, keep audio unchanged - latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0] - return latent_video, latent_audio - if idx < int(len(steps) * (3 / 4)): - noise_theta = 1.0 if (idx + 1) % 2 == 0 else 0.0 - else: - noise_theta = 1.0 if idx % 3 == 0 else 0.0 - latent_video = video_scheduler.step_sde(v_cfg_video, idx, latent_video, noise_theta=noise_theta) - if not is_a2v: - latent_audio = audio_scheduler.step_sde(v_cfg_audio, idx, latent_audio, noise_theta=noise_theta) - return latent_video, latent_audio - - latent_video = video_scheduler.step(v_cfg_video, t, latent_video, return_dict=False)[0] - # Do not update audio latent during the SR stage - if not is_a2v and not use_sr_model: - latent_audio = audio_scheduler.step(v_cfg_audio, t, latent_audio, return_dict=False)[0] - return latent_video, latent_audio - - -_NEGATIVE_PROMPT = ( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, " - "overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, " - "poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, " - "still picture, messy background, three legs, many people in the background, walking backwards" - ", low quality, worst quality, poor quality, noise, background noise, hiss, hum, buzz, crackle, static, " - "compression artifacts, MP3 artifacts, digital clipping, distortion, muffled, muddy, unclear, echo, " - "reverb, room echo, over-reverberated, hollow sound, distant, washed out, harsh, shrill, piercing, " - "grating, tinny, thin sound, boomy, bass-heavy, flat EQ, over-compressed, abrupt cut, jarring transition, " - "sudden silence, looping artifact, music, instrumental, sirens, alarms, crowd noise, unrelated sound " - "effects, chaotic, disorganized, messy, cheap sound" - ", emotionless, flat delivery, deadpan, lifeless, apathetic, robotic, mechanical, monotone, flat " - "intonation, undynamic, boring, reading from a script, AI voice, synthetic, text-to-speech, TTS, " - "insincere, fake emotion, exaggerated, overly dramatic, melodramatic, cheesy, cringey, hesitant, " - "unconfident, tired, weak voice, stuttering, stammering, mumbling, slurred speech, mispronounced, " - "bad articulation, lisp, vocal fry, creaky voice, mouth clicks, lip smacks, wet mouth sounds, heavy " - "breathing, audible inhales, plosives, p-pops, coughing, clearing throat, sneezing, speaking too fast, " - "rushed, speaking too slow, dragged out, unnatural pauses, awkward silence, choppy, disjointed, multiple " - "speakers, two voices, background talking, out of tune, off-key, autotune artifacts" -) - - -# =========================================================================== -# Pre/post process funcs (registered in registry) -# =========================================================================== -def get_magi_human_pre_process_func(*args, **kwargs): - def pre_process(request: OmniDiffusionRequest): - return request - - return pre_process - - -def get_magi_human_post_process_func(*args, **kwargs): - def post_process(output): - if isinstance(output, tuple) and len(output) == 2: - video, audio = output - return { - "video": video, - "audio": audio, - "audio_sample_rate": 44100, - "fps": 25, - } - return output - - return post_process - - -# =========================================================================== -# HF Hub / local path helpers -# =========================================================================== - - -def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict: - """Load a JSON config file from a local path or HuggingFace Hub repo.""" - if local_files_only: - path = os.path.join(model_path, *filename.split("/")) - with open(path) as f: - return json.load(f) - else: - from huggingface_hub import hf_hub_download - - cached = hf_hub_download(repo_id=model_path, filename=filename) - with open(cached) as f: - return json.load(f) - - -def _resolve_subdir( - model_path: str, - subfolder: str, - local_files_only: bool = True, - required_files: list[str] | None = None, -) -> str: - """Resolve a model subfolder to a local directory path. - - For HF Hub repos, downloads all ``required_files`` (default: ``["config.json"]``) - into the HF cache and returns the parent directory. - """ - if local_files_only: - return os.path.join(model_path, subfolder) - from huggingface_hub import hf_hub_download - - files = required_files or ["config.json"] - last_cached: str | None = None - for fname in files: - last_cached = hf_hub_download(repo_id=model_path, filename=f"{subfolder}/{fname}") - return os.path.dirname(last_cached) - - -# =========================================================================== -# Main Pipeline -# =========================================================================== -class MagiHumanPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin): - _dit_modules: ClassVar[list[str]] = ["dit", "sr_dit"] - _encoder_modules: ClassVar[list[str]] = ["text_encoder"] - _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] - - def __init__(self, od_config: OmniDiffusionConfig, **kwargs): - super().__init__() - model_path = od_config.model - local_files_only = os.path.exists(model_path) - device = f"cuda:{torch.accelerator.current_device_index()}" - self.device_str = device - self.dtype = od_config.dtype or torch.bfloat16 - - model_index = _load_json(model_path, "model_index.json", local_files_only) - eval_cfg = model_index - dp_cfg = model_index.get("data_proxy", {}) - - dit_subfolder = "transformer" - - dit_json = _load_json(model_path, f"{dit_subfolder}/config.json", local_files_only) - dit_model_config = MagiHumanDiTConfig(**dit_json) - - self.dit = DiTModel(dit_model_config) - self.dit.eval() - - self.vae = DistributedAutoencoderKLWan.from_pretrained(model_path, subfolder="vae") - self.vae.to(device) - self.vae.eval() - vae_cfg = _load_json(model_path, "vae/config.json", local_files_only) - self.vae_latent_mean = torch.tensor(vae_cfg["latents_mean"], dtype=torch.float32) - self.vae_latent_std = torch.tensor(vae_cfg["latents_std"], dtype=torch.float32) - - self.audio_vae = SAAudioFeatureExtractor( - device=device, - model_path=_resolve_subdir( - model_path, - "audio_vae", - local_files_only, - required_files=["config.json", "model_config.json", "model.safetensors"], - ), - ) - - logger.info("Loading T5Gemma text encoder from %s (subfolder=text_encoder)", model_path) - if local_files_only: - txt_enc_path = os.path.join(model_path, "text_encoder") - txt_enc_subfolder = None - else: - txt_enc_path = model_path - txt_enc_subfolder = "text_encoder" - self.text_encoder = _T5GemmaEncoder( - model_path=txt_enc_path, - device=device, - weight_dtype=self.dtype, - subfolder=txt_enc_subfolder, - ) - - self.data_proxy = MagiDataProxy( - patch_size=dp_cfg.get("patch_size", 2), - t_patch_size=dp_cfg.get("t_patch_size", 1), - frame_receptive_field=dp_cfg.get("frame_receptive_field", 11), - spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"), - ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000), - text_offset=dp_cfg.get("text_offset", 0), - coords_style=dp_cfg.get("coords_style", "v2"), - ) - # SR DataProxy forces v1 coordinate style (consistent with the original) - self.sr_data_proxy = MagiDataProxy( - patch_size=dp_cfg.get("patch_size", 2), - t_patch_size=dp_cfg.get("t_patch_size", 1), - frame_receptive_field=dp_cfg.get("frame_receptive_field", 11), - spatial_rope_interpolation=dp_cfg.get("spatial_rope_interpolation", "extra"), - ref_audio_offset=dp_cfg.get("ref_audio_offset", 1000), - text_offset=dp_cfg.get("text_offset", 0), - coords_style="v1", - ) - - self.fps = eval_cfg.get("fps", 25) - self.num_inference_steps_default = eval_cfg.get("num_inference_steps", 32) - self.video_txt_guidance_scale = eval_cfg.get("video_txt_guidance_scale", 5.0) - self.audio_txt_guidance_scale = eval_cfg.get("audio_txt_guidance_scale", 5.0) - self.shift = eval_cfg.get("shift", 5.0) - self.cfg_number = eval_cfg.get("cfg_number", 2) - self.use_cfg_trick = eval_cfg.get("use_cfg_trick", True) - self.cfg_trick_start_frame = eval_cfg.get("cfg_trick_start_frame", 13) - self.cfg_trick_value = eval_cfg.get("cfg_trick_value", 2.0) - self.using_sde_flag = eval_cfg.get("using_sde_flag", False) - self.t5_gemma_target_length = eval_cfg.get("t5_gemma_target_length", 640) - self.vae_stride = eval_cfg.get("vae_stride", [4, 16, 16]) - self.z_dim = eval_cfg.get("z_dim", 48) - self.patch_size = eval_cfg.get("patch_size", [1, 2, 2]) - # SR-specific hyperparameters - self.sr_num_inference_steps_default = eval_cfg.get("sr_num_inference_steps", 5) - self.sr_cfg_number = eval_cfg.get("sr_cfg_number", 2) - self.sr_video_txt_guidance_scale = eval_cfg.get("sr_video_txt_guidance_scale", 3.5) - self.noise_value = eval_cfg.get("noise_value", 220) - self.sr_audio_noise_scale = eval_cfg.get("sr_audio_noise_scale", 0.7) - # ZeroSNR sigma schedule for SR noise injection (flip=True, high to low) - self.zerosnr_sigmas = ZeroSNRDDPMDiscretization()(1000, do_append_zero=False, flip=True) - - self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding( - _NEGATIVE_PROMPT, - self.text_encoder, - self.t5_gemma_target_length, - ) - self.video_processor = VideoProcessor(vae_scale_factor=16) - - # SR DiT model (loaded from the sr/ subdirectory) - sr_dit_subfolder = "sr" - sr_dit_json = _load_json(model_path, f"{sr_dit_subfolder}/config.json", local_files_only) - sr_dit_model_config = MagiHumanDiTConfig(**sr_dit_json) - self.sr_dit = DiTModel(sr_dit_model_config) - self.sr_dit.eval() - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=model_path, - subfolder=dit_subfolder, - revision=None, - prefix="dit.", - fall_back_to_pt=True, - ), - DiffusersPipelineLoader.ComponentSource( - model_or_path=model_path, - subfolder=sr_dit_subfolder, - revision=None, - prefix="sr_dit.", - fall_back_to_pt=True, - ), - ] - if getattr(self.text_encoder, "is_tp", False): - self.weights_sources.append( - DiffusersPipelineLoader.ComponentSource( - model_or_path=model_path, - subfolder="text_encoder", - revision=None, - prefix="text_encoder.", - fall_back_to_pt=True, - ), - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # Weight loading for MagiHuman DiT with TP support. - # - # The checkpoint stores weights with these naming patterns: - # - attention.linear_qkv.weight: fused [Q, K, V, G] for shared layers, - # or stacked per-expert [expert0_Q|K|V|G, expert1_..., expert2_...] for MoE. - # - attention.linear_proj.weight: single for shared, stacked per-expert for MoE. - # - mlp.up_gate_proj.weight / mlp.down_proj.weight: similarly stacked for MoE. - # - # The model now uses per-expert vLLM parallel layers for MoE blocks: - # attention.linear_qkv.experts.{i}.weight (QKVParallelLinear per expert) - # attention.linear_gating.experts.{i}.weight (ColumnParallelLinear per expert) - # attention.linear_proj.experts.{i}.weight (RowParallelLinear per expert) - # mlp.up_gate_proj.experts.{i}.weight (ColumnParallelLinear per expert) - # mlp.down_proj.experts.{i}.weight (RowParallelLinear per expert) - # - # Shared layers keep the same naming (no .experts.). - params_dict = dict(self.named_parameters()) - modules_dict = dict(self.named_modules()) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - # ── Text Encoder weights ── - if name.startswith("text_encoder."): - if getattr(self.text_encoder, "is_tp", False): - # Strip "text_encoder." prefix for the T5Gemma TP model - # The T5GemmaEncoderModelTP load_weights handles the "encoder." prefix itself - sub_name = name[len("text_encoder.") :] - loaded_params.update( - f"text_encoder.{k}" for k in self.text_encoder.model.load_weights([(sub_name, loaded_weight)]) - ) - else: - loaded_params.add(name) - continue - - # ── Shared attention QKV + Gating split ── - # Checkpoint: attention.linear_qkv.weight = [Q, K, V, G] fused. - # Model: attention.linear_qkv.weight (QKVParallelLinear) + attention.linear_gating.weight. - if "attention.linear_qkv.weight" in name: - gating_name = name.replace("attention.linear_qkv.weight", "attention.linear_gating.weight") - # Check if this is a shared layer (direct param exists, no .experts.) - if name in params_dict and gating_name in params_dict: - qkv_param = params_dict[name] - gating_param = params_dict[gating_name] - - mod_path = name[: -len(".weight")] - qkv_mod = modules_dict.get(mod_path) - if qkv_mod is not None and hasattr(qkv_mod, "total_num_heads"): - total_heads_q = qkv_mod.total_num_heads - total_heads_kv = qkv_mod.total_num_kv_heads - head_dim = qkv_mod.head_size - else: - head_dim = 128 - tp_size = get_tensor_model_parallel_world_size() - total_heads_q = gating_param.data.shape[0] * tp_size - total_heads_kv = (loaded_weight.shape[0] - total_heads_q * head_dim - total_heads_q) // ( - 2 * head_dim - ) - - q_size = total_heads_q * head_dim - kv_size = total_heads_kv * head_dim - - q_w = loaded_weight[:q_size] - k_w = loaded_weight[q_size : q_size + kv_size] - v_w = loaded_weight[q_size + kv_size : q_size + 2 * kv_size] - g_w = loaded_weight[q_size + 2 * kv_size :] - - qkv_loader = getattr(qkv_param, "weight_loader", default_weight_loader) - qkv_loader(qkv_param, q_w, "q") - qkv_loader(qkv_param, k_w, "k") - qkv_loader(qkv_param, v_w, "v") - - gating_loader = getattr(gating_param, "weight_loader", default_weight_loader) - gating_loader(gating_param, g_w) - - loaded_params.add(name) - loaded_params.add(gating_name) - continue - - # ── MoE attention QKV + Gating split ── - # Checkpoint: attention.linear_qkv.weight = stacked [expert0_QKVG, expert1_QKVG, ...]. - # Model: attention.linear_qkv.experts.{i}.weight (QKVParallelLinear per expert) - # + attention.linear_gating.experts.{i}.weight (ColumnParallelLinear per expert). - expert0_name = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0.weight") - if expert0_name in params_dict: - # Determine num_experts by checking which expert indices exist. - moe_qkv_mod_path = name[: -len(".weight")] - moe_qkv_mod = modules_dict.get(moe_qkv_mod_path) - num_experts = moe_qkv_mod.num_experts if moe_qkv_mod is not None else 3 - - # Get head info from the first expert's QKVParallelLinear. - expert0_mod_path = name.replace("attention.linear_qkv.weight", "attention.linear_qkv.experts.0") - expert0_mod = modules_dict.get(expert0_mod_path) - if expert0_mod is not None and hasattr(expert0_mod, "total_num_heads"): - total_heads_q = expert0_mod.total_num_heads - total_heads_kv = expert0_mod.total_num_kv_heads - head_dim = expert0_mod.head_size - else: - head_dim = 128 - # Infer from checkpoint weight shape. - # We'll get exact sizes from model config below. - total_heads_q = 40 # fallback for default config - total_heads_kv = 8 - - q_size = total_heads_q * head_dim - kv_size = total_heads_kv * head_dim - # Check if gating is present. - gating_expert0_name = name.replace( - "attention.linear_qkv.weight", "attention.linear_gating.experts.0.weight" - ) - has_gating = gating_expert0_name in params_dict - - # Split stacked checkpoint weight into per-expert chunks. - expert_weights = loaded_weight.chunk(num_experts, dim=0) - - for i in range(num_experts): - expert_w = expert_weights[i] - # Each expert chunk: [Q, K, V, G (optional)]. - q_w = expert_w[:q_size] - k_w = expert_w[q_size : q_size + kv_size] - v_w = expert_w[q_size + kv_size : q_size + 2 * kv_size] - - expert_param_name = name.replace( - "attention.linear_qkv.weight", - f"attention.linear_qkv.experts.{i}.weight", - ) - expert_param = params_dict[expert_param_name] - expert_loader = getattr(expert_param, "weight_loader", default_weight_loader) - expert_loader(expert_param, q_w, "q") - expert_loader(expert_param, k_w, "k") - expert_loader(expert_param, v_w, "v") - loaded_params.add(expert_param_name) - - if has_gating: - g_w = expert_w[q_size + 2 * kv_size :] - gating_param_name = name.replace( - "attention.linear_qkv.weight", - f"attention.linear_gating.experts.{i}.weight", - ) - gating_param = params_dict[gating_param_name] - gating_loader = getattr(gating_param, "weight_loader", default_weight_loader) - gating_loader(gating_param, g_w) - loaded_params.add(gating_param_name) - continue - - # ── MoE stacked weight splitting for proj / MLP layers ── - # Checkpoint: x.y.weight (stacked [expert0, expert1, ...]). - # Model: x.y.experts.{i}.weight. - if name not in params_dict: - # Check if this is a stacked MoE weight by looking for .experts.0. - base, _, suffix = name.rpartition(".") - expert0_name = f"{base}.experts.0.{suffix}" if base else None - if expert0_name and expert0_name in params_dict: - # Determine num_experts. - moe_mod = modules_dict.get(base) - num_experts = getattr(moe_mod, "num_experts", 3) if moe_mod is not None else 3 - - # Split stacked weight into per-expert chunks. - expert_weights = loaded_weight.chunk(num_experts, dim=0) - for i in range(num_experts): - expert_param_name = f"{base}.experts.{i}.{suffix}" - if expert_param_name not in params_dict: - continue - expert_param = params_dict[expert_param_name] - expert_loader = getattr(expert_param, "weight_loader", default_weight_loader) - expert_loader(expert_param, expert_weights[i]) - loaded_params.add(expert_param_name) - continue - # Truly unknown weight — skip. - continue - - # ── Standard weight loading (shared layers + non-MoE params) ── - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - if getattr(self.text_encoder, "is_tp", False): - self.context_null, self.original_context_null_len = _get_padded_t5_gemma_embedding( - _NEGATIVE_PROMPT, - self.text_encoder, - self.t5_gemma_target_length, - ) - - return loaded_params - - def _dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]: - packed = self.data_proxy.process_input(eval_input) - noise_pred = self.dit(*packed) - return self.data_proxy.process_output(noise_pred) - - def _sr_dit_forward(self, eval_input: EvalInput) -> tuple[torch.Tensor, torch.Tensor]: - """SR stage uses sr_data_proxy (coords_style=v1) and sr_dit model.""" - packed = self.sr_data_proxy.process_input(eval_input) - noise_pred = self.sr_dit(*packed) - return self.sr_data_proxy.process_output(noise_pred) - - @torch.inference_mode() - def _evaluate_with_latent( - self, - context: torch.Tensor, - original_context_len: int, - latent_image: torch.Tensor | None, - latent_video: torch.Tensor, - latent_audio: torch.Tensor, - num_inference_steps: int, - is_a2v: bool = False, - use_sr_model: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Select cfg_number and guidance_scale based on BR/SR stage - cfg_number = self.sr_cfg_number if use_sr_model else self.cfg_number - video_guidance = self.sr_video_txt_guidance_scale if use_sr_model else self.video_txt_guidance_scale - forward_fn = self._sr_dit_forward if use_sr_model else self._dit_forward - - video_scheduler = FlowUniPCMultistepScheduler() - audio_scheduler = FlowUniPCMultistepScheduler() - video_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift) - audio_scheduler.set_timesteps(num_inference_steps, device=self.device_str, shift=self.shift) - timesteps = video_scheduler.timesteps - - latent_length = latent_video.shape[2] - cfg_trick_guidance = ( - torch.tensor(video_guidance, device=self.device_str).expand(1, 1, latent_length, 1, 1).clone() - ) - if self.use_cfg_trick: - cfg_trick_guidance[:, :, : self.cfg_trick_start_frame] = min(self.cfg_trick_value, video_guidance) - - with self.progress_bar(total=len(timesteps)) as pbar: - for idx, t in enumerate(timesteps): - if latent_image is not None: - latent_video[:, :, :1] = latent_image[:, :, :1] - - # Reduce guidance when t<=500 during BR stage (original behavior) - cur_video_guidance = video_guidance if (use_sr_model or t > 500) else 2.0 - - eval_input_cond = EvalInput( - x_t=latent_video, - audio_x_t=latent_audio, - audio_feat_len=[latent_audio.shape[1]], - txt_feat=context, - txt_feat_len=[original_context_len], - ) - - v_cond_video, v_cond_audio = forward_fn(eval_input_cond) - - if cfg_number == 1: - v_cfg_video = v_cond_video - v_cfg_audio = v_cond_audio - elif cfg_number == 2: - eval_input_uncond = EvalInput( - x_t=latent_video, - audio_x_t=latent_audio, - audio_feat_len=[latent_audio.shape[1]], - txt_feat=self.context_null, - txt_feat_len=[self.original_context_null_len], - ) - v_uncond_video, v_uncond_audio = forward_fn(eval_input_uncond) - v_cfg_video = v_uncond_video + cur_video_guidance * (v_cond_video - v_uncond_video) - v_cfg_audio = v_uncond_audio + self.audio_txt_guidance_scale * (v_cond_audio - v_uncond_audio) - else: - raise ValueError(f"Invalid cfg_number: {cfg_number}") - - latent_video, latent_audio = _schedule_latent_step( - video_scheduler=video_scheduler, - audio_scheduler=audio_scheduler, - latent_video=latent_video, - latent_audio=latent_audio, - t=t, - idx=idx, - steps=timesteps, - v_cfg_video=v_cfg_video, - v_cfg_audio=v_cfg_audio, - is_a2v=is_a2v, - cfg_number=cfg_number, - using_sde_flag=self.using_sde_flag, - use_sr_model=use_sr_model, - ) - - pbar.update() - - if latent_image is not None: - latent_video[:, :, :1] = latent_image[:, :, :1] - return latent_video, latent_audio - - def _encode_image(self, image: Image.Image, height: int, width: int) -> torch.Tensor: - image = load_image(image) - image = _resizecrop(image, height, width) - image = self.video_processor.preprocess(image, height=height, width=width) - image = image.to(device=self.device_str, dtype=self.dtype).unsqueeze(2) - vae_out = self.vae.encode(image) - if hasattr(vae_out, "latent_dist"): - return vae_out.latent_dist.mode().to(torch.float32) - return vae_out.to(torch.float32) - - def _decode_video(self, latent: torch.Tensor) -> list[np.ndarray]: - mean = self.vae_latent_mean.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) - std = self.vae_latent_std.to(latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1) - latent = latent * std + mean - - videos = self.vae.decode(latent.to(self.dtype)) - if hasattr(videos, "sample"): - videos = videos.sample - videos.mul_(0.5).add_(0.5).clamp_(0, 1) - videos = [v.float().cpu().permute(1, 2, 3, 0) * 255 for v in videos] - return [v.numpy().astype(np.uint8) for v in videos] - - def _decode_audio(self, latent_audio: torch.Tensor) -> np.ndarray: - latent_audio = latent_audio.squeeze(0).to(self.dtype) - audio_output = self.audio_vae.decode(latent_audio.T) - audio_np = audio_output.squeeze(0).T.float().cpu().numpy() - target_len = int(audio_np.shape[0] * 441 / 512) - from scipy.signal import resample - - return resample(audio_np, target_len) - - def encode_prompt( - self, - prompt: str, - target_length: int | None = None, - ) -> tuple[torch.Tensor, int]: - """Encode *prompt* with the T5-Gemma text encoder and pad to fixed length. - - This is the single text-encoder entrypoint so the runner-level - prompt-embedding cache (see - ``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently - memoize results when the same prompt is submitted repeatedly. - - Returns: - ``(context, original_context_len)`` matching - :func:`_get_padded_t5_gemma_embedding`. - """ - return _get_padded_t5_gemma_embedding( - prompt, - self.text_encoder, - target_length if target_length is not None else self.t5_gemma_target_length, - ) - - @torch.inference_mode() - def forward( - self, - req: DiffusionRequestBatch, - prompt: str | None = None, - height: int = 256, - width: int = 448, - num_inference_steps: int | None = None, - seconds: int = 10, - seed: int | None = None, - image_path: str | None = None, - audio_path: str | None = None, - **kwargs, - ) -> DiffusionOutput: - if len(req.prompts) >= 1: - p = req.prompts[0] - prompt = p if isinstance(p, str) else p.get("prompt", prompt) - if not isinstance(p, str): - image_path = p.get("image_path", image_path) - audio_path = p.get("audio_path", audio_path) - if prompt is None: - raise ValueError("prompt is required") - - height = req.sampling_params.height or height - width = req.sampling_params.width or width - seed = req.sampling_params.seed if req.sampling_params.seed is not None else seed - num_steps = req.sampling_params.num_inference_steps or num_inference_steps or self.num_inference_steps_default - sr_height: int | None = None - sr_width: int | None = None - sr_num_steps: int | None = None - if hasattr(req.sampling_params, "extra_args") and req.sampling_params.extra_args: - seconds = req.sampling_params.extra_args.get("seconds", seconds) - audio_path = req.sampling_params.extra_args.get("audio_path", audio_path) - image_path = req.sampling_params.extra_args.get("image_path", image_path) - sr_height = req.sampling_params.extra_args.get("sr_height", None) - sr_width = req.sampling_params.extra_args.get("sr_width", None) - sr_num_steps = req.sampling_params.extra_args.get("sr_num_inference_steps", None) - - device = self.device_str - - br_latent_height = height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1] - br_latent_width = width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2] - br_height = br_latent_height * self.vae_stride[1] - br_width = br_latent_width * self.vae_stride[2] - - if seed is not None: - torch.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - - if audio_path is not None: - latent_audio = load_audio_and_encode(self.audio_vae, audio_path, seconds) - latent_audio = latent_audio.permute(0, 2, 1) - num_frames = latent_audio.shape[1] - is_a2v = True - else: - num_frames = seconds * self.fps + 1 - latent_audio = torch.randn(1, num_frames, 64, dtype=torch.float32, device=device) - is_a2v = False - - latent_length = (num_frames - 1) // 4 + 1 - latent_video = torch.randn( - 1, - self.z_dim, - latent_length, - br_latent_height, - br_latent_width, - dtype=torch.float32, - device=device, - ) - - context, original_context_len = self.encode_prompt(prompt) - - if image_path is not None: - br_image = self._encode_image(load_image(image_path), br_height, br_width) - else: - br_image = None - - # ── BR stage ───────────────────────────────────────────────────────── - br_latent_video, br_latent_audio = self._evaluate_with_latent( - context, - original_context_len, - br_image, - latent_video.clone(), - latent_audio.clone(), - num_steps, - is_a2v, - use_sr_model=False, - ) - - # ── SR stage (optional, triggered when sr_height/sr_width are provided) ── - if sr_height is not None and sr_width is not None: - sr_latent_height = sr_height // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1] - sr_latent_width = sr_width // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2] - sr_height = sr_latent_height * self.vae_stride[1] - sr_width = sr_latent_width * self.vae_stride[2] - - # Image condition (at SR resolution) - if image_path is not None: - sr_image = self._encode_image(load_image(image_path), sr_height, sr_width) - else: - sr_image = None - - # Trilinear interpolation of BR latent to SR resolution - sr_latent_video = torch.nn.functional.interpolate( - br_latent_video, - size=(latent_length, sr_latent_height, sr_latent_width), - mode="trilinear", - align_corners=True, - ) - - # Noise injection: sigma-weighted blend (noise_value indexes the ZeroSNR sigma schedule) - if self.noise_value != 0: - noise = torch.randn_like(sr_latent_video) - sigma = self.zerosnr_sigmas.to(sr_latent_video.device)[self.noise_value] - sr_latent_video = sr_latent_video * sigma + noise * (1 - sigma**2) ** 0.5 - - # Audio: blend with noise (noised version used during SR inference; final audio keeps BR result) - sr_latent_audio = torch.randn_like(br_latent_audio) * self.sr_audio_noise_scale + br_latent_audio * ( - 1 - self.sr_audio_noise_scale - ) - - torch.accelerator.empty_cache() - sr_steps = sr_num_steps or self.sr_num_inference_steps_default - final_latent_video, _ = self._evaluate_with_latent( - context, - original_context_len, - sr_image, - sr_latent_video.clone(), - sr_latent_audio.clone(), - sr_steps, - is_a2v, - use_sr_model=True, - ) - # SR stage does not update audio; keep the BR result - final_latent_video = final_latent_video - final_latent_audio = br_latent_audio - else: - final_latent_video = br_latent_video - final_latent_audio = br_latent_audio - - torch.accelerator.empty_cache() - videos_np = self._decode_video(final_latent_video) - torch.accelerator.empty_cache() - audio_np = self._decode_audio(final_latent_audio) - - return DiffusionOutput(output=(videos_np, audio_np)) diff --git a/vllm_omni/diffusion/models/minimax_h3/__init__.py b/vllm_omni/diffusion/models/minimax_h3/__init__.py index 53052ad42..02937c058 100644 --- a/vllm_omni/diffusion/models/minimax_h3/__init__.py +++ b/vllm_omni/diffusion/models/minimax_h3/__init__.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """MiniMax H3 audio-video diffusion support.""" +from vllm_omni.model_executor.models.minimax_h3 import reference_video + from .pipeline_minimax_h3 import ( MiniMaxH3Pipeline, get_minimax_h3_post_process_func, @@ -9,4 +11,5 @@ from .pipeline_minimax_h3 import ( __all__ = [ "MiniMaxH3Pipeline", "get_minimax_h3_post_process_func", + "reference_video", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/batched_packing.py b/vllm_omni/diffusion/models/minimax_h3/batched_packing.py new file mode 100644 index 000000000..4d41bde61 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/batched_packing.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Batched multi-request packing for MiniMax H3 step-wise execution. + +Request mode forwards exactly one packed sequence per denoise step. Step mode +(continuous batching) may hold several requests at once, so this module +concatenates their packed layouts into a single sequence and rebuilds +``cu_seqlens`` with one document per request plus any nonempty alignment-padding +tail. Attention therefore never crosses a request boundary, while every request +keeps its own RoPE coordinates, token tags, and timesteps. +A batch of one is not rebuilt at all -- it delegates to the request-mode layout. + +Row order is the request order, so the concatenated video velocity returned by +the DiT can be split back per request by row count alone -- which is exactly +what the runner does with ``StepRequestState.latents``. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +from .denoise_loop import ( + MINIMAX_H3_AUDIO_ROW_WIDTH, + MINIMAX_H3_VIDEO_ROW_WIDTH, + MiniMaxH3DenoiseBranch, +) + + +def minimax_h3_batched_forward_kwargs( + *, + branches: Sequence[MiniMaxH3DenoiseBranch], + video_rows: Sequence[torch.Tensor], + audio_rows: Sequence[torch.Tensor], + t_video: Sequence[float], + t_audio: Sequence[float], + imgvid_cond_timesteps: Sequence[float], + audio_ref_cond_timesteps: Sequence[float], +) -> dict[str, Any]: + """Build one DiT forward kwargs dict covering every request in the batch. + + A single request delegates to :meth:`MiniMaxH3DenoiseBranch.forward_kwargs`, + so step mode and request mode are the same call rather than two layouts that + happen to agree -- which also keeps ``video_token_layout`` (the sparse- + attention hint) on the single-request step path. + """ + if len(branches) == 1: + return branches[0].forward_kwargs( + video_rows=video_rows[0], + audio_rows=audio_rows[0], + t_video=t_video[0], + t_audio=t_audio[0], + imgvid_cond_timestep=imgvid_cond_timesteps[0], + audio_ref_cond_timestep=audio_ref_cond_timesteps[0], + ) + + device = branches[0].device + total_seq = sum(branch.seq_len for branch in branches) + + img_pos_parts: list[torch.Tensor] = [] + audio_pos_parts: list[torch.Tensor] = [] + text_pos_parts: list[torch.Tensor] = [] + update_parts: list[torch.Tensor] = [] + tag_parts: list[torch.Tensor] = [] + position_id_parts: list[torch.Tensor] = [] + text_embed_parts: list[torch.Tensor] = [] + cu_bounds: list[int] = [0] + refiner_bounds: list[int] = [0] + max_seqlen = 0 + refiner_max_seqlen = 0 + + # Non-media rows (text and alignment padding) inherit their request's video + # timestep, matching the single-request packed-sequence semantics. + timesteps = torch.empty(total_seq, dtype=torch.float32, device=device) + + seq_offset = 0 + text_offset = 0 + for index, branch in enumerate(branches): + img_pos = branch.img_pos_dev + seq_offset + audio_pos = branch.audio_pos_dev + seq_offset + img_pos_parts.append(img_pos) + audio_pos_parts.append(audio_pos) + text_pos_parts.append(branch.text_pos_dev + seq_offset) + update_parts.append(branch.update_mask_dev) + tag_parts.append(branch.token_tags_dev) + position_id_parts.append(branch.img_position_ids_dev) + text_embed_parts.append(branch.text_embeddings_dev) + + timesteps[seq_offset : seq_offset + branch.seq_len] = float(t_video[index]) + timesteps[img_pos[branch.update_mask_dev]] = float(t_video[index]) + timesteps[img_pos[~branch.update_mask_dev]] = float(imgvid_cond_timesteps[index]) + timesteps[audio_pos[branch.audio_update_mask_dev]] = float(t_audio[index]) + timesteps[audio_pos[~branch.audio_update_mask_dev]] = float(audio_ref_cond_timesteps[index]) + + # Empty documents are omitted because not every varlen kernel accepts + # repeated interior boundaries. + real_end = seq_offset + branch.used_len + padded_end = seq_offset + branch.seq_len + cu_bounds.append(real_end) + if padded_end != real_end: + cu_bounds.append(padded_end) + max_seqlen = max(max_seqlen, branch.used_len, branch.seq_len - branch.used_len) + + text_offset += branch.text_len + refiner_bounds.append(text_offset) + refiner_max_seqlen = max(refiner_max_seqlen, branch.text_len) + seq_offset += branch.seq_len + + img_pos_all = torch.cat(img_pos_parts) + audio_pos_all = torch.cat(audio_pos_parts) + + x = torch.zeros(1, total_seq, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device) + x[0].index_copy_(0, img_pos_all, torch.cat([rows.to(device=device, dtype=torch.float32) for rows in video_rows])) + audio_x = torch.zeros(1, total_seq, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device) + audio_x[0].index_copy_( + 0, + audio_pos_all, + torch.cat([rows.to(device=device, dtype=torch.float32) for rows in audio_rows]), + ) + + unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True) + return { + "x": x, + "audio_x": audio_x, + "img_position_ids": torch.cat(position_id_parts)[None], + "update_mask": torch.cat(update_parts), + "token_tags": torch.cat(tag_parts), + "skip_mask_out_condition": False, + "prompt_embeds": torch.cat(text_embed_parts), + "img_pos_info": {"position_ids": img_pos_all}, + "audio_pos_info": {"position_ids": audio_pos_all}, + "text_pos_info": {"position_ids": torch.cat(text_pos_parts)}, + "img_pos_for_infer_output_info": {"position_ids": img_pos_all}, + "packed_seq_params": { + "cu_seqlens_q": torch.tensor(cu_bounds, dtype=torch.int32, device=device), + "max_seqlen_q": max_seqlen, + # Attention needs the request count on the host to know its valid + # rows are block-diagonal rather than a prefix. It describes the + # whole forward, so the DiT reads it once here and applies it to the + # token refiner too -- the refiner's text-only bounds cannot be told + # apart from a single padded request on their own. + "num_requests": len(branches), + }, + "refiner_packed_seq_params": { + "cu_seqlens_q": torch.tensor(refiner_bounds, dtype=torch.int32, device=device), + "max_seqlen_q": refiner_max_seqlen, + }, + "unique_timesteps": unique_timesteps, + "inverse_indices": inverse_indices, + } + + +__all__ = [ + "minimax_h3_batched_forward_kwargs", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index d22065c3a..653dfce71 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 cfg-distilled full denoise loop. Per step, the positive presentation is forwarded exactly once. Video and audio @@ -14,8 +15,13 @@ from typing import Any import torch -from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout -from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx +from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout, VideoTokenSpan +from vllm_omni.diffusion.forward_context import ( + set_forward_context_denoise_step_idx, + set_forward_context_denoise_timestep, + set_forward_context_denoise_total_steps, +) +from vllm_omni.platforms import current_omni_platform from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, @@ -31,6 +37,21 @@ MINIMAX_H3_VIDEO_ROW_WIDTH = 96 MINIMAX_H3_AUDIO_ROW_WIDTH = 32 +def minimax_h3_publish_denoise_progress( + step: int | None, sigma_video: float | None, total_steps: int | None = None +) -> None: + """Publish denoise progress for step-gated attention features. + + Both execution modes must publish the same trio: the step index drives the + dense warmup of RAINFUSION_ATTN, the normalized descending timestep + drives the TRTLLM_ATTN skip gate (which stays dense while it is unset), + and the total step count enables the ``end_step`` tail fallback. + """ + set_forward_context_denoise_step_idx(step) + set_forward_context_denoise_timestep(sigma_video) + set_forward_context_denoise_total_steps(total_steps) + + class MiniMaxH3DenoiseBranch: """Static per-branch state: packed layout + fixed forward kwargs. @@ -64,26 +85,39 @@ class MiniMaxH3DenoiseBranch: if int(token_tags.view(-1).shape[0]) != seq_len: raise ValueError(f"token_tags length {int(token_tags.view(-1).shape[0])} != seq_len {seq_len}") cu = packed["cu_seqlens"].to(torch.int32) + self.device = device + # ``used_len`` is the real (non-padding) document length; step-mode + # batching concatenates layouts and needs both bounds as Python ints so + # it can rebuild ``cu_seqlens`` without a device sync per step. + self.used_len = int(cu[1]) + self.text_len = text_len self.img_pos_dev = self.img_pos.to(device) self.audio_pos_dev = self.audio_pos.to(device) self.update_mask_dev = self.update_mask.to(device) self.audio_update_mask_dev = self.audio_update_mask.to(device) self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device) self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device) - text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device) + self.text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device) + self.token_tags_dev = token_tags.view(-1).to(torch.long).to(device) + self.img_position_ids_dev = packed["img_position_ids"].to(device) + self.text_embeddings_dev = text_embeddings.to(device) self.static_kwargs: dict[str, Any] = { - "img_position_ids": packed["img_position_ids"][None].to(device), + "img_position_ids": self.img_position_ids_dev[None], "update_mask": self.update_mask_dev, - "token_tags": token_tags.view(-1).to(torch.long).to(device), + "token_tags": self.token_tags_dev, "skip_mask_out_condition": False, - "prompt_embeds": text_embeddings.to(device), + "prompt_embeds": self.text_embeddings_dev, "img_pos_info": {"position_ids": self.img_pos_dev}, "audio_pos_info": {"position_ids": self.audio_pos_dev}, - "text_pos_info": {"position_ids": text_pos_dev}, + "text_pos_info": {"position_ids": self.text_pos_dev}, "img_pos_for_infer_output_info": {"position_ids": self.img_pos_dev}, "packed_seq_params": { "cu_seqlens_q": cu.to(device), - "max_seqlen_q": int(cu[1]), + "max_seqlen_q": self.used_len, + # One request: valid rows are a prefix, so attention may use a + # KV prefix length or a 1-D pad mask. See batched_packing for the + # co-batched layout, where neither can describe the valid rows. + "num_requests": 1, }, "refiner_packed_seq_params": { "cu_seqlens_q": torch.tensor([0, text_len, text_len], dtype=torch.int32, device=device), @@ -92,10 +126,38 @@ class MiniMaxH3DenoiseBranch: } # Where the video segment sits in the packed sequence. Resolved to plain # ints here so the attention layers never sync on it per step. - grid = packed["latent_grid"].tolist() - self.static_kwargs["video_token_layout"] = VideoTokenLayout( - prefix_len=int(packed["video_row_start"]), - latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])), + raw_spans = packed.get("video_spans") + if raw_spans: + spans = tuple( + VideoTokenSpan( + start=int(span["start"]), + latent_grid=tuple(int(dim) for dim in span["latent_grid"]), + role=str(span["role"]), + ) + for span in raw_spans + ) + self.static_kwargs["video_token_layout"] = VideoTokenLayout(used_len=int(cu[1]), video_spans=spans) + else: + grid = packed["latent_grid"].tolist() + self.static_kwargs["video_token_layout"] = VideoTokenLayout( + prefix_len=int(packed["video_row_start"]), + latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])), + ) + + def prepare_rope_table(self, model: Any) -> None: + """Materialize the branch-local DiT RoPE table once per denoise run. + + ``img_position_ids`` is immutable for this branch while latents and + timesteps change every scheduler step. Keeping the table in + ``static_kwargs`` makes every model call reuse the exact same BF16 + tensor without extending its lifetime beyond this request branch. + """ + prepare = getattr(model, "prepare_rope_table", None) + if not callable(prepare): + return + self.static_kwargs["rope_table"] = prepare( + self.static_kwargs["img_position_ids"], + seq_len=self.seq_len, ) def forward_kwargs( @@ -136,34 +198,21 @@ class MiniMaxH3DenoiseBranch: } -def minimax_h3_denoise_loop( +def minimax_h3_prepare_denoise_rows( *, - model: Any, positive: MiniMaxH3DenoiseBranch, initial_video_rows: torch.Tensor, initial_audio_rows: torch.Tensor, keyframe_cond_rows: torch.Tensor | None, - audio_ref_rows: torch.Tensor | None = None, - sigmas_video: list[float], - sigmas_audio: list[float], + audio_ref_rows: torch.Tensor | None, device: torch.device, - imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP, - audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, - on_step_end: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, - on_step_start: Callable[[int, float, float], None] | None = None, - step_profiler: Callable[[int], AbstractContextManager] | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Run the full denoise loop; returns final (video_rows, audio_rows). +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """Validate the initial rows against the layout and pin the condition rows. - ``initial_video_rows`` covers all image rows of the positive layout. For a - conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to - pin those rows across every step. The model's raw positive velocity is the - update signal; MiniMax H3 only supports cfg-distilled checkpoints. + Returns ``(video_rows, audio_rows, cond_anchor, audio_anchor)`` with the + anchors already written into their rows, which is the state both the + request-mode loop and step-mode ``prepare_encode()`` start from. """ - if len(sigmas_video) != len(sigmas_audio): - raise ValueError("video/audio sigma schedules must have equal length") - if len(sigmas_video) < 2: - raise ValueError("sigma schedules need at least 2 entries") n_cond = int((~positive.update_mask).sum()) if keyframe_cond_rows is None: if n_cond != 0: @@ -197,19 +246,63 @@ def minimax_h3_denoise_loop( video_rows[~update] = cond_anchor if audio_anchor is not None: audio_rows[~audio_update] = audio_anchor + return video_rows, audio_rows, cond_anchor, audio_anchor + + +def minimax_h3_denoise_loop( + *, + model: Any, + positive: MiniMaxH3DenoiseBranch, + initial_video_rows: torch.Tensor, + initial_audio_rows: torch.Tensor, + keyframe_cond_rows: torch.Tensor | None, + audio_ref_rows: torch.Tensor | None = None, + sigmas_video: list[float], + sigmas_audio: list[float], + device: torch.device, + imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP, + audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, + step_profiler: Callable[[int], AbstractContextManager] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the full denoise loop; returns final (video_rows, audio_rows). + + ``initial_video_rows`` covers all image rows of the positive layout. For a + conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to + pin those rows across every step. The model's raw positive velocity is the + update signal; MiniMax H3 only supports cfg-distilled checkpoints. + """ + if len(sigmas_video) != len(sigmas_audio): + raise ValueError("video/audio sigma schedules must have equal length") + if len(sigmas_video) < 2: + raise ValueError("sigma schedules need at least 2 entries") + # Keep CUDA/CPU on the reference path: its accuracy CI compares the + # generated video against the official artifact. The table reuse is an + # Ascend-specific performance optimization and must not alter that path. + if current_omni_platform.is_npu(): + positive.prepare_rope_table(model) + video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows( + positive=positive, + initial_video_rows=initial_video_rows, + initial_audio_rows=initial_audio_rows, + keyframe_cond_rows=keyframe_cond_rows, + audio_ref_rows=audio_ref_rows, + device=device, + ) + update = positive.update_mask_dev + audio_update = positive.audio_update_mask_dev num_steps = len(sigmas_video) - 1 for step in range(num_steps): step_cm = step_profiler(step) if step_profiler is not None else nullcontext() with step_cm: - # Publish the step index so step-gated attention features (e.g. the - # dense warmup of RAINFUSION_ATTN) can see where we are. - set_forward_context_denoise_step_idx(step) s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1] - if on_step_start is not None: - # Attention gates use the scheduler-style descending timestep. - on_step_start(step, s_v, s_a) + # Publish where we are so step-gated attention features (the dense + # warmup of RAINFUSION_ATTN, the timestep gate of TRTLLM_ATTN) can + # see it. Gates use the scheduler-style descending timestep, which + # for a rectified-flow schedule is the video sigma. + minimax_h3_publish_denoise_progress(step, s_v, num_steps) t_v, t_a = 1.0 - s_v, 1.0 - s_a imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference)) audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference)) @@ -250,10 +343,10 @@ def minimax_h3_denoise_loop( audio_rows[audio_update] = new_audio if audio_anchor is not None: audio_rows[~audio_update] = audio_anchor # per-step audio ref reset - if on_step_end is not None: - on_step_end(step, video_rows, audio_rows) + if on_step is not None: + on_step(step, video_rows, audio_rows) - set_forward_context_denoise_step_idx(None) + minimax_h3_publish_denoise_progress(None, None, None) return video_rows, audio_rows @@ -264,4 +357,6 @@ __all__ = [ "MINIMAX_H3_VIDEO_ROW_WIDTH", "MiniMaxH3DenoiseBranch", "minimax_h3_denoise_loop", + "minimax_h3_prepare_denoise_rows", + "minimax_h3_publish_denoise_progress", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 2e07ace96..a91cb180b 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 Qwen3-VL layer-50 text/vision encoder. The encoder is reimplemented on top of vLLM-style tensor-parallel building @@ -30,16 +31,22 @@ complete ``[seq, 5120]`` hidden state. from __future__ import annotations import itertools -import json import re +from collections.abc import Iterable from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.layers.norm import RMSNorm as DiffusionRMSNorm +from vllm_omni.diffusion.offloader.module_residency import ( + BoundedAllocatorCache, + PinnedModuleStager, +) MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50 MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 @@ -47,6 +54,19 @@ MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 logger = init_logger(__name__) +def _validate_encoder_quant_config(quant_config: QuantizationConfig | None) -> None: + if ( + quant_config is not None + and quant_config.get_name() == "fp8" + and bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False)) + ): + raise ValueError("MiniMax-H3 text encoder supports online FP8 only, not serialized FP8 checkpoints") + + +def _is_quantized_linear(module: LinearBase) -> bool: + return not isinstance(module.quant_method, UnquantizedLinearMethod) + + def _default_weight_loader( param: nn.Parameter, loaded_weight: torch.Tensor, @@ -114,7 +134,7 @@ class MiniMaxH3Qwen3VLVocabParallelEmbedding(nn.Module): param.data.copy_(loaded_weight.narrow(0, start_idx, shard_size)) -class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module): +class MiniMaxH3Qwen3VLMergedColumnParallelLinear(LinearBase): """Packed gate/up projection sharded along the output dimension.""" def __init__( @@ -123,23 +143,43 @@ class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module): input_size: int, intermediate_size: int, dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: - super().__init__() self.group = group - self.input_size = input_size self.intermediate_size = intermediate_size tp_rank, tp_size = _tp_range(group) assert intermediate_size % tp_size == 0, ( f"intermediate_size {intermediate_size} must be divisible by text_encoder_tp_size {tp_size}" ) self.intermediate_size_per_partition = intermediate_size // tp_size - self.weight = nn.Parameter(torch.empty(2 * self.intermediate_size_per_partition, input_size, dtype=dtype)) - self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + self.output_dtype = dtype + super().__init__( + input_size=input_size, + output_size=2 * intermediate_size, + bias=False, + params_dtype=dtype, + quant_config=quant_config, + prefix=prefix, + disable_tp=True, + ) + self.quant_method.create_weights( + self, + input_size_per_partition=input_size, + output_partition_sizes=[ + self.intermediate_size_per_partition, + self.intermediate_size_per_partition, + ], + input_size=input_size, + output_size=2 * intermediate_size, + params_dtype=dtype, + weight_loader=self.weight_loader, + ) self._tp_rank = tp_rank self._tp_size = tp_size def forward(self, input_: torch.Tensor) -> torch.Tensor: - return F.linear(input_, self.weight) + return self.quant_method.apply(self, input_) def weight_loader( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None @@ -152,7 +192,7 @@ class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module): param.data[0:shard_size].copy_(loaded_weight.narrow(0, start_idx, shard_size)) -class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): +class MiniMaxH3Qwen3VLQKVParallelLinear(LinearBase): """QKV projection with GQA head sharding along the output dimension.""" def __init__( @@ -163,8 +203,9 @@ class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): num_kv_heads: int, head_dim: int, dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: - super().__init__() self.group = group self.hidden_size = hidden_size self.num_heads = num_heads @@ -179,15 +220,33 @@ class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): ) self.local_num_heads = num_heads // tp_size self.local_num_kv_heads = num_kv_heads // tp_size + self.output_dtype = dtype q_local = self.local_num_heads * head_dim kv_local = self.local_num_kv_heads * head_dim - self.weight = nn.Parameter(torch.empty(q_local + 2 * kv_local, hidden_size, dtype=dtype)) - self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + output_size = (num_heads + 2 * num_kv_heads) * head_dim + super().__init__( + input_size=hidden_size, + output_size=output_size, + bias=False, + params_dtype=dtype, + quant_config=quant_config, + prefix=prefix, + disable_tp=True, + ) + self.quant_method.create_weights( + self, + input_size_per_partition=hidden_size, + output_partition_sizes=[q_local, kv_local, kv_local], + input_size=hidden_size, + output_size=output_size, + params_dtype=dtype, + weight_loader=self.weight_loader, + ) self._tp_rank = tp_rank self._tp_size = tp_size def forward(self, input_: torch.Tensor) -> torch.Tensor: - return F.linear(input_, self.weight) + return self.quant_method.apply(self, input_) def weight_loader( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None @@ -240,7 +299,7 @@ def _verify_fused_shards_complete( raise RuntimeError(f"MiniMax H3 text encoder fused weights are missing source shards: {details}") -class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): +class MiniMaxH3Qwen3VLRowParallelLinear(LinearBase): def __init__( self, group: Any, @@ -248,17 +307,33 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): output_size: int, dtype: torch.dtype, input_is_parallel: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: - super().__init__() self.group = group - self.input_size = input_size - self.output_size = output_size self.input_is_parallel = input_is_parallel + self.output_dtype = dtype tp_rank, tp_size = _tp_range(group) assert input_size % tp_size == 0, f"input_size {input_size} must be divisible by text_encoder_tp_size {tp_size}" self.input_size_per_partition = input_size // tp_size - self.weight = nn.Parameter(torch.empty(output_size, self.input_size_per_partition, dtype=dtype)) - self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + super().__init__( + input_size=input_size, + output_size=output_size, + bias=False, + params_dtype=dtype, + quant_config=quant_config, + prefix=prefix, + disable_tp=True, + ) + self.quant_method.create_weights( + self, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=[output_size], + input_size=input_size, + output_size=output_size, + params_dtype=dtype, + weight_loader=self.weight_loader, + ) self._tp_rank = tp_rank self._tp_size = tp_size @@ -268,7 +343,7 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): else: split_input = input_.split(self.input_size_per_partition, dim=-1) input_parallel = split_input[self._tp_rank].contiguous() - output_parallel = F.linear(input_parallel, self.weight) + output_parallel = self.quant_method.apply(self, input_parallel) if self._tp_size > 1: # Reduce in fp32: the reference path accumulates the full (K=8192) # dot product inside a single cuBLAS GEMM before rounding to bf16. @@ -276,7 +351,7 @@ class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): # and amplify error on this model's large-magnitude activations. output_parallel = output_parallel.float() self.group.all_reduce(output_parallel) - output_parallel = output_parallel.to(self.weight.dtype) + output_parallel = output_parallel.to(self.output_dtype) return output_parallel def weight_loader( @@ -327,6 +402,25 @@ def _apply_rotary_pos_emb( return q_embed, k_embed +def _scaled_dot_product_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> torch.Tensor: + """Run causal SDPA with the reference expanded-K/V GQA behavior.""" + num_key_value_groups = query.shape[1] // key.shape[1] + if num_key_value_groups != 1: + key = key.repeat_interleave(num_key_value_groups, dim=1) + value = value.repeat_interleave(num_key_value_groups, dim=1) + return F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.0, + is_causal=True, + ) + + def _apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor: """Reorganize chunked [TTT...HHH...WWW] into interleaved [THWTHW...] layout.""" freqs_t = freqs[0] @@ -665,7 +759,14 @@ class MiniMaxH3Qwen3VLVisionModel(nn.Module): class MiniMaxH3Qwen3VLTextAttention(nn.Module): - def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() self.group = group self.config = config @@ -683,6 +784,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): num_kv_heads=self.num_kv_heads, head_dim=self.head_dim, dtype=dtype, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", ) self.o_proj = MiniMaxH3Qwen3VLRowParallelLinear( group, @@ -690,6 +793,8 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): output_size=self.hidden_size, dtype=dtype, input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", ) self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps, dtype=dtype) @@ -703,38 +808,42 @@ class MiniMaxH3Qwen3VLTextAttention(nn.Module): hidden_shape = (*input_shape, -1, self.head_dim) q_local = self.qkv_proj.local_num_heads * self.head_dim kv_local = self.qkv_proj.local_num_kv_heads * self.head_dim - q_weight = self.qkv_proj.weight[0:q_local] - k_weight = self.qkv_proj.weight[q_local : q_local + kv_local] - v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local] - # Three separate GEMMs keep the numerics identical to the reference - # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding). - query_states = self.q_norm(F.linear(hidden_states, q_weight).view(hidden_shape)).transpose(1, 2) - key_states = self.k_norm(F.linear(hidden_states, k_weight).view(hidden_shape)).transpose(1, 2) - value_states = F.linear(hidden_states, v_weight).view(hidden_shape).transpose(1, 2) + if _is_quantized_linear(self.qkv_proj): + query, key, value = self.qkv_proj(hidden_states).split( + (q_local, kv_local, kv_local), + dim=-1, + ) + else: + q_weight = self.qkv_proj.weight[0:q_local] + k_weight = self.qkv_proj.weight[q_local : q_local + kv_local] + v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local] + # Three separate GEMMs keep the numerics identical to the reference + # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding). + query = F.linear(hidden_states, q_weight) + key = F.linear(hidden_states, k_weight) + value = F.linear(hidden_states, v_weight) + query_states = self.q_norm(query.view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(key.view(hidden_shape)).transpose(1, 2) + value_states = value.view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin) - # Expand GQA keys/values to the local query-head count so SDPA works - # uniformly across backends (mirrors the reference eager path). - num_key_value_groups = self.num_heads // self.num_kv_heads - key_states = key_states.repeat_interleave(num_key_value_groups, dim=1) - value_states = value_states.repeat_interleave(num_key_value_groups, dim=1) - - attn_output = F.scaled_dot_product_attention( - query_states, - key_states, - value_states, - dropout_p=0.0, - is_causal=True, - ) + attn_output = _scaled_dot_product_attention(query_states, key_states, value_states) attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output class MiniMaxH3Qwen3VLTextMLP(nn.Module): - def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size @@ -743,6 +852,8 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module): input_size=self.hidden_size, intermediate_size=self.intermediate_size, dtype=dtype, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", ) self.down_proj = MiniMaxH3Qwen3VLRowParallelLinear( group, @@ -750,22 +861,46 @@ class MiniMaxH3Qwen3VLTextMLP(nn.Module): output_size=self.hidden_size, dtype=dtype, input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", ) def forward(self, x: torch.Tensor) -> torch.Tensor: ip_local = self.gate_up_proj.intermediate_size_per_partition - gate = F.linear(x, self.gate_up_proj.weight[0:ip_local]) - up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local]) + if _is_quantized_linear(self.gate_up_proj): + gate, up = self.gate_up_proj(x).split(ip_local, dim=-1) + else: + gate = F.linear(x, self.gate_up_proj.weight[0:ip_local]) + up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local]) x = F.silu(gate) * up return self.down_proj(x) class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): - def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() self.hidden_size = config.hidden_size - self.self_attn = MiniMaxH3Qwen3VLTextAttention(group, config, dtype) - self.mlp = MiniMaxH3Qwen3VLTextMLP(group, config, dtype) + self.self_attn = MiniMaxH3Qwen3VLTextAttention( + group, + config, + dtype, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp = MiniMaxH3Qwen3VLTextMLP( + group, + config, + dtype, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype) self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm( config.hidden_size, eps=config.rms_norm_eps, dtype=dtype @@ -779,10 +914,7 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(hidden_states, position_embeddings=position_embeddings) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states @@ -791,7 +923,15 @@ class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): class MiniMaxH3Qwen3VLTextModel(nn.Module): """TP-sharded Qwen3-VL text model returning unnormalized layer-50 states.""" - def __init__(self, group: Any, config: Any, selected_layer: int, dtype: torch.dtype) -> None: + def __init__( + self, + group: Any, + config: Any, + selected_layer: int, + dtype: torch.dtype, + quant_config: QuantizationConfig | None = None, + prefix: str = "text_encoder.text_model", + ) -> None: super().__init__() self.group = group self.config = config @@ -804,7 +944,16 @@ class MiniMaxH3Qwen3VLTextModel(nn.Module): ) self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config) self.layers = nn.ModuleList( - [MiniMaxH3Qwen3VLTextDecoderLayer(group, config, dtype) for _ in range(self.num_layers)] + [ + MiniMaxH3Qwen3VLTextDecoderLayer( + group, + config, + dtype, + quant_config=quant_config, + prefix=f"{prefix}.layers.{layer_idx}", + ) + for layer_idx in range(self.num_layers) + ] ) def _deepstack_process( @@ -951,10 +1100,13 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): device: torch.device, load_model: bool, encoder_group: Any | None = None, + quant_config: QuantizationConfig | None = None, ) -> None: super().__init__() self.device_target = device self.encoder_group = encoder_group + self.quant_config = quant_config + _validate_encoder_quant_config(quant_config) self.image_token_id = 151655 self.video_token_id = 151656 self._tp_size = 1 @@ -975,9 +1127,9 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): config.text_config, MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, dtype, + quant_config=quant_config, ) self.text_model.to(dtype=dtype) - self._load_weights(model_path) logger.info( "MiniMax H3 Qwen3-VL encoder: %d retained decoder layers, text_encoder_tp_size=%d, vision replicated", MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, @@ -1032,59 +1184,54 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): return ("text_model." + rest, None) return None - def _load_weights(self, model_path: str) -> None: - import os - - import safetensors.torch - - index_path = os.path.join(model_path, "model.safetensors.index.json") - with open(index_path, encoding="utf-8") as handle: - index = json.load(handle) - weight_map: dict[str, str] = index["weight_map"] - files = sorted(set(weight_map.values())) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params = dict(self.named_parameters()) loaded: set[str] = set() - # Seeded from the modules that own a multi-shard loader, so a fused weight - # that receives no source shard at all is caught too, not only a partially - # filled one. + # Seeded from modules that own a multi-shard loader, so a fused weight + # receiving no source shard is caught as well as a partially filled one. expected_fused_shards: dict[str, tuple[tuple[str, str | int], ...]] = { f"{module_name}.weight": sources for module_name, module in self.named_modules() if (sources := _fused_source_shards(module)) is not None } loaded_fused_shards: dict[str, set[str | int]] = {name: set() for name in expected_fused_shards} - for shard in files: - tensors = safetensors.torch.load_file(os.path.join(model_path, shard), device="cpu") - for name, tensor in tensors.items(): - mapped = self._map_weight_name(name) - if mapped is None: - continue - param_name, shard_id = mapped - param = params.get(param_name) - if param is None: - logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name) - continue - weight_loader = getattr(param, "weight_loader", _default_weight_loader) - weight_loader(param, tensor, shard_id) - loaded.add(param_name) - if shard_id is not None and param_name in loaded_fused_shards: - loaded_fused_shards[param_name].add(shard_id) + for name, tensor in weights: + mapped = self._map_weight_name(name) + if mapped is None: + continue + param_name, shard_id = mapped + param = params.get(param_name) + if param is None: + logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name) + continue + weight_loader = getattr(param, "weight_loader", _default_weight_loader) + weight_loader(param, tensor, shard_id) + loaded.add(param_name) + if shard_id is not None and param_name in loaded_fused_shards: + loaded_fused_shards[param_name].add(shard_id) _verify_fused_shards_complete(expected_fused_shards, loaded_fused_shards) missing = sorted(set(params) - loaded) if missing: # Listed in full: this aborts startup, so the message is the only diagnostic. raise RuntimeError(f"MiniMax H3 text encoder weights not loaded: {len(missing)} params: {missing}") + return loaded def load_to_device(self) -> None: if not self.is_loaded: return if getattr(self, "_omni_layerwise_enabled", False): - for name, child in self.vision.named_children(): - if name != "blocks": - child.to(self.device_target) - for name, child in self.text_model.named_children(): - if name != "layers": - child.to(self.device_target) + stager = getattr(self, "_omni_non_block_stager", None) + if stager is None: + hooks = getattr(self, "_omni_layerwise_hooks", ()) + pin_memory = bool(getattr(hooks[0], "pin_memory", True)) if hooks else True + stager = PinnedModuleStager( + self._omni_non_block_modules(), + self.device_target, + pin_memory=pin_memory, + cache_retention=getattr(self, "_omni_component_cache", None), + ) + self._omni_non_block_stager = stager + stager.load() return self.vision.to(self.device_target) self.text_model.to(self.device_target) @@ -1095,18 +1242,39 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): if getattr(self, "_omni_layerwise_enabled", False): for hook in self._omni_layerwise_hooks: hook.offload_layer() - for name, child in self.vision.named_children(): - if name != "blocks": + stager = getattr(self, "_omni_non_block_stager", None) + if stager is not None: + stager.offload() + else: + # The initial DLO setup reaches this branch before the first + # load. These modules were constructed on CPU, so this is not + # a device-to-host rematerialization. + for child in self._omni_non_block_modules(): child.to("cpu") - for name, child in self.text_model.named_children(): - if name != "layers": - child.to("cpu") - torch.accelerator.empty_cache() + self._release_omni_component_cache() return self.vision.to("cpu") self.text_model.to("cpu") torch.accelerator.empty_cache() + def _omni_non_block_modules(self) -> list[nn.Module]: + modules = [child for name, child in self.vision.named_children() if name != "blocks"] + modules.extend(child for name, child in self.text_model.named_children() if name != "layers") + return modules + + def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: + self._omni_component_cache = cache + stager = getattr(self, "_omni_non_block_stager", None) + if stager is not None: + stager.set_cache_retention(cache) + + def _release_omni_component_cache(self) -> None: + cache = getattr(self, "_omni_component_cache", None) + if cache is None: + torch.accelerator.empty_cache() + else: + cache.release_if_needed() + def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None: """Stream the TP-local Qwen vision/text blocks for low-HBM serving. @@ -1296,15 +1464,19 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): if next(self.parameters()).device.type != self.device_target.type: raise RuntimeError("call load_to_device() before encode_ids()") + cudnn_sdp_enabled = torch.backends.cuda.cudnn_sdp_enabled() torch.backends.cuda.enable_cudnn_sdp(True) - hidden = self._encode( - input_ids, - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - pixel_values_videos=pixel_values_videos, - video_grid_thw=video_grid_thw, - ) - return hidden.cpu() + try: + hidden = self._encode( + input_ids, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + ) + return hidden.cpu() + finally: + torch.backends.cuda.enable_cudnn_sdp(cudnn_sdp_enabled) def forward( self, diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py new file mode 100644 index 000000000..21f77828a --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py @@ -0,0 +1,679 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""FastVideo FastH3: a four-step DMD2 student of MiniMax-H3. + +FastH3 replaces H3's 49 denoiser evaluations with four. It ships as an adapter +over the base checkpoint rather than as a full release, so it reuses H3's text +encoder, video VAE, audio VAE, tokenizers and schedulers unchanged. + +The artifact is *not* a PEFT LoRA, and it is not request-switchable. Its own +metadata states the reconstruction as:: + + W = W_base + lora_B @ lora_A; then .diff/.diff_b added and .set_weight assigned + +so besides rank-64 factors it carries full-rank ``.diff``/``.diff_b`` deltas for +RMSNorm weights, biases, patch projections and the final layer - none of which a +LoRA layer can express - and the VSA variants add ``.set_weight`` tensors for +compression gates that do not exist in the base transformer at all. The adapter +is therefore fused into the checkpoint stream at load time, before the weights +are sharded, which is also what the release's model card requires. + +The low-rank factors carry no alpha: the reconstruction adds ``lora_B @ lora_A`` +directly, i.e. a scale of exactly 1. + +Two checkpoint spellings meet here. The adapter is written in the diffusers +namespace (``transformer_blocks.0.attn.to_q``) while vLLM-Omni loads H3's native +one (``blocks.0.attn.qkv_proj``), whose attention and MLP projections are fused. +Every mapping and layout convention below was verified tensor by tensor against +the released full checkpoint (``FastVideo-FastH3-4-step-Preview-v1-Dense-DataFree``): +``W_base + delta`` reproduces it to bf16 rounding. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from safetensors import safe_open +from vllm.logger import init_logger + +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.errors import OmniClientError +from vllm_omni.platforms import current_omni_platform + +if TYPE_CHECKING: + from .minimax_h3_transformer import MiniMaxH3DiTModel + +logger = init_logger(__name__) + +# FastVideo's generic adapter container, not a FastH3 marker: their tools emit +# it for ordinary H3 adapters too. +FASTH3_FORMAT = "fastvideo-lora-v2" +FASTH3_MANIFEST = "adapter_manifest.json" +# The release identity: the distilled student the adapter came from, over the +# base it edits. The name is held to FastVideo's own namespace as well as the +# student's name, so an unrelated adapter that merely mentions FastH3 is not +# claimed. +FASTH3_BASE_MODEL = "MiniMaxAI/MiniMax-H3" +_FASTH3_IDENTITY_KEY = "finetuned_model" +_FASTH3_IDENTITY_NAMESPACE = "fastvideo/" +_FASTH3_IDENTITY_MARKER = "fasth3" + +# The rectified-flow positions the student was distilled at, and the ladder the +# server samples on. The release states them as `dmd_denoising_steps` +# [999, 749, 500, 250] (`sampling.base_timesteps` in the bundle manifest): +# timestep indices out of 1000, i.e. pre-shift positions, closed here with the +# terminal 0.0 every rectified-flow schedule ends on. The opening rung is 0.999 +# rather than 1.0 because training capped the noise level there +# (`max_timestep_ratio`). +# +# These are positions, not final sigmas: H3's schedulers apply their own +# per-modality shift on top (12 for video, 3 for audio), which is what +# reproduces the levels the student saw. Nothing here overrides those shifts. +FASTH3_BASE_SCHEDULE = DMD2SigmaSchedule.from_positions((0.999, 0.749, 0.5, 0.25, 0.0)) +# Its five points bound four transformer forwards, one per sigma interval. That +# count is the one a request states, the one Cache-DiT is refreshed with and the +# one step execution admits a request on - the same interval contract H3's +# pinned checkpoint schedules and native LoRAs already use. +FASTH3_DENOISE_STEPS = FASTH3_BASE_SCHEDULE.num_inference_steps +# Preview v1 distills the text-to-video-and-audio path only. +FASTH3_SUPPORTED_TASKS = frozenset({"t2va"}) + +_LORA_A = ".lora_A.weight" +_LORA_B = ".lora_B.weight" +_DIFF = ".diff" +_DIFF_B = ".diff_b" +_SET_WEIGHT = ".set_weight" + +# Adapter module prefix -> the native parameter it edits, minus the +# ``.weight``/``.bias`` suffix. +_MODEL_LEVEL_TARGETS = { + "proj_in": "video_patch_proj", + "proj_out": "final_layer.video_out", + "audio_proj_in": "audio_patch_proj", + "audio_proj_out": "final_layer.audio_out", + "context_embedder": "condition_proj", + "time_embedder.linear_1": "time_embedder.proj_in", + "time_embedder.linear_2": "time_embedder.proj_out", + "norm_out.linear": "final_layer.adaln_proj.linear", + "norm_out.norm": "final_layer.norm", +} + +# Per-block adapter suffix -> (native suffix, how a delta enters the native +# parameter). H3 stores attention as one grouped QKV matrix and the MLP as one +# fused gate/up matrix, so those deltas need placing rather than adding. +_PLAIN, _QKV_Q, _QKV_K, _QKV_V, _SWAP_HALVES = "plain", "q", "k", "v", "swap_halves" +_QKV_SLOTS = (_QKV_Q, _QKV_K, _QKV_V) +_BLOCK_TARGETS = { + "attn.to_q": ("attn.qkv_proj", _QKV_Q), + "attn.to_k": ("attn.qkv_proj", _QKV_K), + "attn.to_v": ("attn.qkv_proj", _QKV_V), + "attn.to_out.0": ("attn.out_proj", _PLAIN), + "ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES), + "ff.net.2": ("mlp.fc2", _PLAIN), + "adaln_proj.linear": ("adaln_proj.linear", _PLAIN), + "norm1": ("norm1", _PLAIN), + "norm2": ("norm2", _PLAIN), +} + +# Adapter block prefix -> native block prefix. +_BLOCK_PREFIXES = ( + ("token_refiner.refiner_blocks.", "token_refiner.blocks."), + ("transformer_blocks.", "blocks."), +) + + +class FastH3AdapterError(ValueError): + """The artifact is a FastH3 adapter, but it cannot be applied as one.""" + + +@dataclass +class _ParamPatch: + """Everything the adapter contributes to one native parameter.""" + + # layout -> (lora_A, lora_B). A grouped QKV parameter collects three. + low_rank: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = field(default_factory=dict) + diff: torch.Tensor | None = None + layout: str = _PLAIN + + +def _swap_halves(tensor: torch.Tensor) -> torch.Tensor: + """Exchange the two halves of a fused gate/up matrix. + + The diffusers export stores the feed-forward projection value-first while + H3's native ``mlp.fc1`` is gate-first, so a delta computed in the diffusers + layout has to be swapped before it can be added to the native parameter. + """ + if tensor.shape[0] % 2: + raise FastH3AdapterError(f"fused gate/up delta must split evenly, got {tuple(tensor.shape)}") + first, second = tensor.chunk(2, dim=0) + return torch.cat((second, first), dim=0) + + +def _place_in_grouped_qkv(deltas: Mapping[str, torch.Tensor], *, head_dim: int) -> torch.Tensor: + """Interleave per-projection deltas into H3's grouped QKV layout. + + The checkpoint stores one head group at a time as ``[q, k, v]``, which is + what :func:`_reorder_grouped_qkv_to_qkv` unpacks on the way in. A delta + built from the separate diffusers projections has to be folded back into + that order. + """ + missing = sorted(set(_QKV_SLOTS) - set(deltas)) + if missing: + raise FastH3AdapterError(f"grouped QKV delta is missing its {missing} projections") + parts = [] + for slot in _QKV_SLOTS: + delta = deltas[slot] + if delta.shape[0] % head_dim: + raise FastH3AdapterError( + f"QKV {slot} delta rows {delta.shape[0]} are not a multiple of head_dim {head_dim}" + ) + parts.append(delta.reshape(delta.shape[0] // head_dim, head_dim, *delta.shape[1:])) + groups = parts[0].shape[0] + if any(part.shape[0] != groups for part in parts): + raise FastH3AdapterError("QKV projections disagree on the number of head groups") + return torch.cat(parts, dim=1).reshape(groups * 3 * head_dim, *parts[0].shape[2:]) + + +def _resolve_native_target(module: str) -> tuple[str, str, tuple[str, int] | None] | None: + """Map an adapter module path to ``(native path, layout, block)``. + + ``block`` is the ``(native block prefix, index)`` the module sits in, or + ``None`` for a model-level one. Coverage is checked per block, so the index + has to survive the mapping rather than being folded into the path. + """ + native = _MODEL_LEVEL_TARGETS.get(module) + if native is not None: + return native, _PLAIN, None + for adapter_prefix, native_prefix in _BLOCK_PREFIXES: + if not module.startswith(adapter_prefix): + continue + remainder = module[len(adapter_prefix) :] + index, _, suffix = remainder.partition(".") + if not index.isdigit(): + return None + target = _BLOCK_TARGETS.get(suffix) + if target is None: + return None + native_suffix, layout = target + return f"{native_prefix}{index}.{native_suffix}", layout, (native_prefix, int(index)) + return None + + +def _split_adapter_key(name: str) -> tuple[str, str] | None: + """Split an adapter tensor name into ``(module path, role)``.""" + for marker, role in ( + (_LORA_A, "lora_a"), + (_LORA_B, "lora_b"), + (_DIFF_B, "diff_b"), + (_DIFF, "diff"), + (_SET_WEIGHT, "set_weight"), + ): + if name.endswith(marker): + return name[: -len(marker)], role + return None + + +def _is_fasth3_release(metadata: Mapping[str, str]) -> bool: + """Whether this is a FastH3 release rather than merely FastVideo's format. + + Claiming on the container alone would fuse an ordinary H3 adapter into the + checkpoint and put every request on a four-step schedule it was never + distilled for, and so would a bare ``fasth3`` substring: an adapter someone + else named after the student is not the student. + """ + if metadata.get("format") != FASTH3_FORMAT: + return False + identity = metadata.get(_FASTH3_IDENTITY_KEY, "").lower() + if not identity.startswith(_FASTH3_IDENTITY_NAMESPACE) or _FASTH3_IDENTITY_MARKER not in identity: + return False + # Only held against a file that states it. + base_model = metadata.get("base_model") + return base_model is None or base_model.casefold() == FASTH3_BASE_MODEL.casefold() + + +def _check_declared_counts( + metadata: Mapping[str, str], + counted: Mapping[str, int], + *, + weights_path: Path, +) -> None: + """Hold the artifact to the tensor counts its own metadata declares. + + The writer of the ``fastvideo-lora-v2`` format records how many tensors of + each kind it emitted, which is the one statement in the file about its own + completeness. A truncated or partially re-exported artifact is otherwise + indistinguishable from a small one, so a claimed file has to carry all of + them rather than opt out by omission. + """ + for key, seen in counted.items(): + declared = metadata.get(key) + if declared is None: + raise FastH3AdapterError( + f"{weights_path} declares the FastH3 identity but omits {key}; " + "a claimed adapter has to state what it emitted" + ) + try: + expected = int(declared) + except (TypeError, ValueError) as exc: + raise FastH3AdapterError(f"{weights_path} declares a non-numeric {key}={declared!r}") from exc + if expected != seen: + raise FastH3AdapterError( + f"{weights_path} declares {key}={expected} but carries {seen}; the adapter is incomplete " + "and would leave most of the transformer on base H3 weights" + ) + + +def _check_block_coverage( + seen: Mapping[str, set[int]], + *, + expected: Mapping[str, int], + weights_path: Path, +) -> None: + """Every block of the model this adapter is loaded against must be edited. + + The release drops tensors training left unchanged, so per-parameter + coverage is legitimately sparse - but a distilled student touches every + block, so a block with no edits at all means the artifact does not match + this model. + """ + for prefix, count in expected.items(): + indices = seen.get(prefix, set()) + wanted = set(range(count)) + if indices == wanted: + continue + missing = sorted(wanted - indices) + extra = sorted(indices - wanted) + raise FastH3AdapterError( + f"{weights_path} edits {len(indices)} of the model's {count} {prefix}* blocks " + f"(missing={missing[:5]}, unknown={extra[:5]}); it is not an adapter for this checkpoint" + ) + + +class FastH3WeightFusion: + """Fuse a FastH3 adapter into the H3 checkpoint stream as it is loaded.""" + + def __init__( + self, + *, + source: Path, + patches: Mapping[str, _ParamPatch], + head_dim: int, + requires_vsa: bool, + ) -> None: + self._source = source + self._patches = dict(patches) + self._head_dim = head_dim + self.requires_vsa = requires_vsa + self._applied: set[str] = set() + self._device: torch.device | None = None + + @property + def source(self) -> Path: + return self._source + + @property + def base_schedule(self) -> tuple[float, ...]: + """The rectified-flow positions this student samples on. + + The fused checkpoint is a four-step student, so the ladder comes from + the release rather than from the many-step teacher's metadata or from + the uniform one ``num_inference_steps`` would otherwise derive. + """ + return FASTH3_BASE_SCHEDULE.base_schedule + + @classmethod + def from_path( + cls, + path: str | Path, + *, + head_dim: int, + num_blocks: int, + num_refiner_blocks: int, + ) -> FastH3WeightFusion | None: + """Build a fusion from an adapter file or directory, else ``None``. + + Returning ``None`` keeps every other ``--lora-path`` artifact on the + dynamic LoRA route; only a file carrying the FastH3 release identity is + claimed here. + + The block counts are the model's, and the artifact has to cover them: + claiming a partial adapter would switch the server onto the four-step + contract while most of the transformer still held base H3 weights. + """ + weights_path = _resolve_adapter_file(path) + if weights_path is None: + return None + + patches: dict[str, _ParamPatch] = {} + gate_tensors: list[str] = [] + unmapped: list[str] = [] + blocks_seen: dict[str, set[int]] = {} + counted = {"low_rank_tensors": 0, "diff_tensors": 0} + with safe_open(weights_path, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if not _is_fasth3_release(metadata): + return None + for name in checkpoint.keys(): + split = _split_adapter_key(name) + if split is None: + unmapped.append(name) + continue + module, role = split + if role == "set_weight": + # A VSA compression gate: a module the base transformer does + # not have, so there is nothing to fuse it into. + gate_tensors.append(name) + continue + target = _resolve_native_target(module) + if target is None: + unmapped.append(name) + continue + native_module, layout, block = target + if block is not None: + blocks_seen.setdefault(block[0], set()).add(block[1]) + counted["diff_tensors" if role in ("diff", "diff_b") else "low_rank_tensors"] += 1 + native_param = f"{native_module}.{'bias' if role == 'diff_b' else 'weight'}" + patch = patches.setdefault(native_param, _ParamPatch(layout=layout)) + tensor = checkpoint.get_tensor(name) + if role in ("diff", "diff_b"): + if patch.diff is not None: + raise FastH3AdapterError(f"duplicate {role} for {native_param}") + patch.diff = tensor + else: + a, b = patch.low_rank.get(layout, (None, None)) + patch.low_rank[layout] = (tensor, b) if role == "lora_a" else (a, tensor) + + if unmapped: + raise FastH3AdapterError( + f"FastH3 adapter at {weights_path} has {len(unmapped)} tensors that name no known " + f"H3 parameter: {sorted(unmapped)[:5]}" + ) + _check_declared_counts( + metadata, + {**counted, "set_weight_tensors": len(gate_tensors)}, + weights_path=weights_path, + ) + _check_block_coverage( + blocks_seen, + expected={"blocks.": num_blocks, "token_refiner.blocks.": num_refiner_blocks}, + weights_path=weights_path, + ) + for native_param, patch in patches.items(): + for slot, (a, b) in patch.low_rank.items(): + if a is None or b is None: + raise FastH3AdapterError(f"FastH3 adapter has an unpaired factor for {native_param} slot {slot!r}") + # Only the low-rank factors are placed into H3's fused QKV and + # gate/up layouts; a full-rank delta is added as it comes, so one + # aimed at a fused parameter would silently land transposed. + if patch.diff is not None and patch.layout != _PLAIN: + raise FastH3AdapterError( + f"FastH3 adapter carries a full-rank delta for {native_param}, which H3 stores in the " + f"{patch.layout!r} fused layout; this loader can only place low-rank factors there" + ) + + fusion = cls( + source=weights_path, + patches=patches, + head_dim=head_dim, + requires_vsa=bool(gate_tensors), + ) + logger.info( + "FastH3 adapter %s: rank=%s, parameters patched=%d, low-rank=%s, diff=%s, set_weight=%d", + weights_path, + metadata.get("rank", "?"), + len(patches), + metadata.get("low_rank_tensors", "?"), + metadata.get("diff_tensors", "?"), + len(gate_tensors), + ) + return fusion + + def _compute_device(self, weight: torch.Tensor) -> torch.device: + """Where to reconstruct a delta. + + H3's per-block modulation projection is 96768x2688, so rebuilding all + 343 patched parameters is a few TFLOP of rank-64 products. On CPU that + adds minutes to a load that already has a startup deadline, so the + accelerator does the arithmetic whenever there is one. + """ + if weight.device.type != "cpu": + return weight.device + if self._device is None: + # Ask the platform rather than PyTorch's global accelerator + # registry, so an out-of-tree backend controls its own placement. + self._device = current_omni_platform.get_torch_device() + return self._device + + @staticmethod + def _widen(tensor: torch.Tensor, device: torch.device) -> torch.Tensor: + """Move to ``device``, then widen to float32. + + Asking ``Tensor.to`` for a device and a dtype at once converts on the + host and ships twice the bytes; splitting it moves bfloat16 and widens + on the accelerator. + """ + return tensor.to(device, non_blocking=True).to(torch.float32) + + def fuse(self, name: str, weight: torch.Tensor) -> torch.Tensor: + """Return ``weight`` with this adapter's contribution added.""" + patch = self._patches.get(name) + if patch is None: + return weight + self._applied.add(name) + + device = self._compute_device(weight) + + delta: torch.Tensor | None = None + if patch.low_rank: + if patch.layout in _QKV_SLOTS: + per_slot = { + slot: self._widen(b, device) @ self._widen(a, device) for slot, (a, b) in patch.low_rank.items() + } + delta = _place_in_grouped_qkv(per_slot, head_dim=self._head_dim) + else: + a, b = patch.low_rank[patch.layout] + if patch.layout == _SWAP_HALVES: + # Permuting the rows of B permutes the rows of the product, + # so swap the rank-64 factor instead of the full delta. + b = _swap_halves(b) + delta = self._widen(b, device) @ self._widen(a, device) + if patch.diff is not None: + diff = self._widen(patch.diff, device) + delta = diff if delta is None else delta + diff + if delta is None: + return weight + if delta.shape != weight.shape: + raise FastH3AdapterError( + f"FastH3 delta for {name} has shape {tuple(delta.shape)}, parameter is {tuple(weight.shape)}" + ) + # Leave the result on the compute device. These weights are bound for + # the accelerator anyway, so returning them to host memory would pay a + # device-to-host copy of the whole checkpoint only for the loader to + # send it straight back: measured at 152s against 15s for 60 GiB of + # patched projections, against 17s for the unavoidable upload alone. + # Fold the base weight into the freshly built delta in place. Promoting + # the weight to float32 on its own would allocate two more buffers the + # size of the parameter, and H3's largest patched projection is 0.5 GiB. + return delta.add_(weight.to(device, non_blocking=True)).to(weight.dtype) + + def apply(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]: + """Fuse every streamed checkpoint tensor on its way into the model.""" + if self._applied: + # ``validate_fully_applied`` released the deltas, so a second stream + # would fuse nothing and then pass its own completeness check: the + # server would serve base H3 weights on the student's ladder. + raise FastH3AdapterError(f"{self._source} has already been fused into this checkpoint") + for name, weight in weights: + yield name, self.fuse(name, weight) + + def validate_fully_applied(self) -> None: + """Close the fusion: every edit must have met its parameter. + + A silently unapplied delta is the failure mode that matters here: the + model would load and generate, just not as the distilled student. The + weights are loaded once, so the mapped payloads are dropped afterwards + rather than held for the life of the process. + """ + missing = sorted(set(self._patches) - self._applied) + if missing: + raise FastH3AdapterError( + f"FastH3 adapter edits {len(missing)} parameters the checkpoint never provided: {missing[:5]}" + ) + for patch in self._patches.values(): + patch.low_rank.clear() + patch.diff = None + + def check_serving_contract( + self, + *, + partition: str, + od_config: Any, + video_shift: float, + audio_shift: float, + ) -> None: + """Hold a starting server to the ladder this student was trained on.""" + if partition == "ref2va": + raise ValueError("FastH3 preview v1 distills T2VA only, so it cannot serve a Ref2VA partition") + offloads = [ + flag + for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload") + if getattr(od_config, flag, False) + ] + if offloads: + # A host-weight plan installs the transformer without going through + # load_weights(), which is where the fusion and its completeness + # check live. Serving base H3 weights under a four-step schedule + # would otherwise degrade output with nothing to signal it. + raise ValueError( + f"FastH3 is fused while the checkpoint streams in, so it cannot be combined with " + f"{sorted(offloads)}. Serve it without offload." + ) + if self.requires_vsa: + raise ValueError( + f"{self.source} is a Video Sparse Attention variant of FastH3, and its " + "compression gates have no counterpart in vLLM-Omni's dense H3 attention. Use the " + "release's Dense variant until VSA lands for H3's packed sequence." + ) + logger.info( + "FastH3 adapter active: sigma points %s for %d transformer forwards, " + "flow_shift=%g, audio_flow_shift=%g, tasks=%s", + list(self.base_schedule), + FASTH3_DENOISE_STEPS, + video_shift, + audio_shift, + sorted(FASTH3_SUPPORTED_TASKS), + ) + + def check_task(self, task: str) -> None: + """Refuse a task this preview never distilled.""" + if task not in FASTH3_SUPPORTED_TASKS: + raise OmniClientError( + f"FastH3 preview v1 distills {sorted(FASTH3_SUPPORTED_TASKS)} only, got task={task!r}" + ) + + def check_request(self, sampling: Any, *, video_shift: float, audio_shift: float) -> None: + """Refuse a request that would sample the student off its rungs.""" + if sampling.lora_request is not None: + # The adapter is already in the weights and the dynamic LoRA manager + # is skipped, so nothing would apply the requested one. Serving the + # request anyway would quietly ignore it. + raise OmniClientError( + f"this server fused {self.source} into the checkpoint at startup, so per-request " + "lora is unavailable; drop the lora field" + ) + if int(sampling.num_inference_steps or 0) != FASTH3_DENOISE_STEPS: + raise OmniClientError( + f"FastH3 is a four-step student and requires num_inference_steps={FASTH3_DENOISE_STEPS} " + "(one transformer forward per sigma interval)" + ) + # The checkpoint's per-modality shifts turn the release's positions into + # the noise levels the student was distilled at, so a request that moves + # them samples where it was never trained. + extra = sampling.extra_args or {} + for key, expected in (("flow_shift", video_shift), ("audio_flow_shift", audio_shift)): + try: + requested = float(extra.get(key, expected)) + except (TypeError, ValueError) as exc: + raise OmniClientError(f"FastH3 requires {key}={expected:g}") from exc + if not math.isclose(requested, expected): + raise OmniClientError(f"FastH3 requires {key}={expected:g}, got {requested:g}") + + +def resolve_fasth3_fusion(od_config: Any, transformer: MiniMaxH3DiTModel) -> FastH3WeightFusion | None: + """Claim ``--lora-path`` when it points at a FastH3 adapter. + + FastH3 rewrites RMSNorm weights and biases, so it cannot be expressed as a + request-switchable LoRA and is fused into the checkpoint instead. Any other + artifact returns None here and stays on the dynamic LoRA route. + """ + lora_path = getattr(od_config, "lora_path", None) + if isinstance(lora_path, (list, tuple)): + if len(lora_path) != 1: + return None + lora_path = lora_path[0] + if not lora_path: + return None + # Placing a delta into the fused QKV parameter needs the head size, and + # completeness is judged against the model's depth, so the architecture is + # only read once an adapter is actually configured. + arch = transformer.arch + return FastH3WeightFusion.from_path( + lora_path, + head_dim=arch.attention_head_dim, + num_blocks=arch.num_layers, + num_refiner_blocks=arch.token_refiner_num_layers, + ) + + +def _safetensors_metadata(path: Path) -> Mapping[str, str]: + """The header metadata of a safetensors file, or ``{}`` if unreadable.""" + try: + with safe_open(path, framework="pt", device="cpu") as checkpoint: + return checkpoint.metadata() or {} + except Exception: # noqa: BLE001 - a path that will not open is simply not ours + return {} + + +def _resolve_adapter_file(path: str | Path) -> Path | None: + """Find the single adapter file at ``path``, or ``None``.""" + candidate = Path(path) + if candidate.is_file(): + return candidate if candidate.suffix == ".safetensors" else None + if not candidate.is_dir(): + return None + named = candidate / "adapter_model.safetensors" + if named.is_file(): + return named + files = sorted(candidate.glob("*.safetensors")) + if len(files) == 1: + return files[0] + # The published repository bundles four variants under one root, marked by an + # adapter_manifest.json. Such a bundle is ambiguous rather than loadable - but + # only hold that against a directory that actually carries FastH3 artifacts; + # an unrelated multi-shard LoRA stays on the dynamic route via ``None``. + if (candidate / FASTH3_MANIFEST).is_file() or any( + _is_fasth3_release(_safetensors_metadata(file)) for file in files + ): + raise FastH3AdapterError( + f"{candidate} holds several FastH3 adapters; point --lora-path at one variant " + "(for example dense-datafree/adapter_model.safetensors)" + ) + return None + + +__all__ = [ + "FASTH3_BASE_MODEL", + "FASTH3_BASE_SCHEDULE", + "FASTH3_DENOISE_STEPS", + "FASTH3_FORMAT", + "FASTH3_SUPPORTED_TASKS", + "FastH3AdapterError", + "FastH3WeightFusion", + "resolve_fasth3_fusion", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/lora.py b/vllm_omni/diffusion/models/minimax_h3/lora.py new file mode 100644 index 000000000..026d9f6b3 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/lora.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Loader for the LightX2V Turbo MiniMax-H3 LoRA. + +Keyed by ``key_format=minimax-h3-diffusers``. The native FlashGen contract lives +in ``.npu.lora``; the two artifacts differ in rank, target naming and QKV +layout, so each owns its own parsing and packing. +""" + +from __future__ import annotations + +import math +from pathlib import Path + +import torch +from safetensors import safe_open +from vllm.lora.lora_model import LoRAModel +from vllm.lora.lora_weights import PackedLoRALayerWeights +from vllm.lora.peft_helper import PEFTHelper +from vllm.model_executor.models.utils import WeightsMapper + +from vllm_omni.lora.request import LoRARequest + +_TURBO_RANK = 128 +_TURBO_ALPHA = 128 +_TURBO_HIDDEN_SIZE = 5376 +_TURBO_ATTENTION_INNER_SIZE = 7168 +_TURBO_FFN_HIDDEN_SIZE = 14336 +_TURBO_FILENAME = "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors" +_LORA_A_SUFFIX = ".lora_A.default.weight" +_LORA_B_SUFFIX = ".lora_B.default.weight" +_TURBO_TARGETS = frozenset({"to_q", "to_k", "to_v", "out_proj", "fc1", "fc2"}) +_TURBO_RAW_TARGET_SUFFIXES = ( + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.to_out.0", + "ff.net.0.proj", + "ff.net.2", +) +_TURBO_TARGET_DIMS = { + "attn.to_q": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE), + "attn.to_k": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE), + "attn.to_v": (_TURBO_HIDDEN_SIZE, _TURBO_ATTENTION_INNER_SIZE), + "attn.to_out.0": (_TURBO_ATTENTION_INNER_SIZE, _TURBO_HIDDEN_SIZE), + "ff.net.0.proj": (_TURBO_HIDDEN_SIZE, 2 * _TURBO_FFN_HIDDEN_SIZE), + "ff.net.2": (_TURBO_FFN_HIDDEN_SIZE, _TURBO_HIDDEN_SIZE), +} +_TURBO_EXPECTED_RAW_TARGETS = frozenset( + f"{prefix}.{block_index}.{suffix}" + for prefix, block_count in ( + ("transformer_blocks", 50), + ("token_refiner.refiner_blocks", 2), + ) + for block_index in range(block_count) + for suffix in _TURBO_RAW_TARGET_SUFFIXES +) +_TURBO_TARGET_PATTERN = ( + r"^transformer\.(?:token_refiner\.blocks|blocks)\.\d+\." + r"(?:attn\.(?:to_q|to_k|to_v|out_proj)|mlp\.(?:fc1|fc2))$" +) + +_TURBO_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_substr={ + "token_refiner.refiner_blocks.": "token_refiner.blocks.", + "transformer_blocks.": "blocks.", + ".attn.to_out.0.": ".attn.out_proj.", + ".ff.net.0.proj.": ".mlp.fc1.", + ".ff.net.2.": ".mlp.fc2.", + ".lora_A.default.": ".lora_A.", + ".lora_B.default.": ".lora_B.", + } +) + + +def _select_turbo_file(artifact_path: str | Path) -> Path | None: + path = Path(artifact_path) + if path.is_file(): + return path if path.suffix == ".safetensors" else None + if not path.is_dir(): + return None + + candidate = path / _TURBO_FILENAME + return candidate if candidate.is_file() else None + + +def _validate_and_convert_tensors(checkpoint) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + pairs: dict[str, set[str]] = {} + raw_targets: set[str] = set() + for name in checkpoint.keys(): + if name.endswith(_LORA_A_SUFFIX): + raw_target = name[: -len(_LORA_A_SUFFIX)] + side = "a" + elif name.endswith(_LORA_B_SUFFIX): + raw_target = name[: -len(_LORA_B_SUFFIX)] + side = "b" + else: + raise ValueError(f"Unconsumed MiniMax-H3 Turbo tensor: {name!r}") + raw_targets.add(raw_target) + + mapped_name = _TURBO_WEIGHTS_MAPPER.apply_list([name])[0] + mapped_target = mapped_name.rsplit(".lora_", 1)[0] + if mapped_target.rsplit(".", 1)[-1] not in _TURBO_TARGETS: + raise ValueError(f"Unsupported MiniMax-H3 Turbo target: {raw_target!r}") + target_sides = pairs.setdefault(mapped_target, set()) + if side in target_sides: + raise ValueError(f"Duplicate MiniMax-H3 Turbo tensor for {mapped_target}.{side}") + target_sides.add(side) + + tensor = checkpoint.get_tensor(name) + if tensor.ndim != 2: + raise ValueError(f"MiniMax-H3 Turbo LoRA tensors must be matrices, got {name}={tuple(tensor.shape)}") + suffix = next((suffix for suffix in _TURBO_RAW_TARGET_SUFFIXES if raw_target.endswith(suffix)), None) + if suffix is None: + raise ValueError(f"MiniMax-H3 Turbo LoRA contains unsupported target: {raw_target}") + input_dim, output_dim = _TURBO_TARGET_DIMS[suffix] + expected_shape = (_TURBO_RANK, input_dim) if side == "a" else (output_dim, _TURBO_RANK) + if tuple(tensor.shape) != expected_shape: + raise ValueError( + f"MiniMax-H3 Turbo tensor has invalid global shape: {name}={tuple(tensor.shape)}, " + f"expected={expected_shape}" + ) + if side == "b" and ".ff.net.0.proj." in name: + value, gate = tensor.chunk(2, dim=0) + tensor = torch.cat((gate, value), dim=0).contiguous() + tensors[name] = tensor + + incomplete = sorted(target for target, sides in pairs.items() if sides != {"a", "b"}) + if incomplete: + raise ValueError(f"Incomplete MiniMax-H3 Turbo LoRA pairs: {incomplete}") + missing = sorted(_TURBO_EXPECTED_RAW_TARGETS - raw_targets) + unexpected = sorted(raw_targets - _TURBO_EXPECTED_RAW_TARGETS) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 Turbo target set does not match the supported v1.0 artifact: " + f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}" + ) + return tensors + + +def _pack_h3_turbo_fc1(lora_model: LoRAModel) -> None: + """Represent H3's fused gate/up projection without generic layout guesses.""" + + for module_name, weights in tuple(lora_model.loras.items()): + if not module_name.endswith(".mlp.fc1"): + continue + gate_b, up_b = weights.lora_b.chunk(2, dim=0) + lora_model.loras[module_name] = PackedLoRALayerWeights( + module_name=module_name, + rank=weights.rank, + lora_alphas=[weights.lora_alpha, weights.lora_alpha], + lora_a=[weights.lora_a, weights.lora_a], + lora_b=[gate_b.contiguous(), up_b.contiguous()], + scaling=[weights.scaling, weights.scaling], + ) + + +def load_minimax_h3_turbo_lora( + *, + partition: str, + lora_request: LoRARequest, + lora_path: str | Path, + dtype: torch.dtype, + unsupported_offload_mode: str | None = None, +) -> tuple[LoRAModel, PEFTHelper] | None: + """Load the published LightX2V Turbo v1.0 through the legacy manager.""" + + lora_file = _select_turbo_file(lora_path) + if lora_file is None: + return None + with safe_open(lora_file, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if metadata.get("key_format") != "minimax-h3-diffusers": + if lora_file.name == _TURBO_FILENAME: + raise ValueError( + "MiniMax-H3 Turbo v1.0 requires safetensors metadata key_format='minimax-h3-diffusers'" + ) + return None + if lora_file.name != _TURBO_FILENAME: + raise ValueError(f"MiniMax-H3 Turbo supports only {_TURBO_FILENAME!r}, got {lora_file.name!r}") + raw_alpha = metadata.get("alpha") + try: + alpha = float(raw_alpha) if raw_alpha is not None else math.nan + except ValueError as exc: + raise ValueError(f"MiniMax-H3 Turbo alpha must be numeric, got {raw_alpha!r}") from exc + if alpha != _TURBO_ALPHA: + raise ValueError(f"MiniMax-H3 Turbo v1.0 requires alpha={_TURBO_ALPHA}, got {raw_alpha!r}") + if partition == "ref2va": + raise ValueError("MiniMax-H3 Turbo LoRA supports FL2VA/T2VA only") + if unsupported_offload_mode is not None: + raise ValueError(f"MiniMax-H3 Turbo dynamic LoRA does not support {unsupported_offload_mode}") + tensors = _validate_and_convert_tensors(checkpoint) + + peft_helper = PEFTHelper.from_dict( + { + "r": _TURBO_RANK, + "lora_alpha": _TURBO_ALPHA, + "target_modules": _TURBO_TARGET_PATTERN, + } + ) + lora_model = LoRAModel.from_lora_tensors( + lora_model_id=lora_request.lora_int_id, + tensors=tensors, + peft_helper=peft_helper, + device="cpu", + dtype=dtype, + weights_mapper=_TURBO_WEIGHTS_MAPPER, + ) + _pack_h3_turbo_fc1(lora_model) + return lora_model, peft_helper diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 5e180aa23..5b5d9c213 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 packed-token audio/video DiT for vLLM-Omni. vLLM tensor parallel linears and the unified attention layer provide TP and @@ -26,15 +27,29 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, + PackedPaddingMetadata, + VideoTokenLayout, +) from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( + indexed_gate, + indexed_gate_rms_norm_scale_shift, + indexed_scale_shift_, + rms_norm_indexed_scale_shift, +) from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) +from vllm_omni.diffusion.layers.activation import SiluAndMul +from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract +from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -46,6 +61,26 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# Packed multi-request forwards require the attention backend to actually +# consume cu_seqlens as a block-diagonal plan (not a padding-mask rebuild that +# spans the full packed row). The pipeline gates on this capability before +# packing, and ``_run_packed_attention`` re-checks it per forward; a name-only +# gate would let FLASH_ATTN's NPU/XPU code paths through even though those +# variants would silently attend across request boundaries. +def _attention_isolates_packed_requests(attention_layer: Any) -> bool: + """True if this attention layer keeps N-document packed boundaries. + + Requires a backend advertising ``supports_multi_doc_packed_varlen`` *and* + that the layer is not running under ring sequence parallelism (the ring + kernel dispatches through its own attention that ignores the packed + cu_seqlens regardless of the configured backend). + """ + backend = getattr(attention_layer, "attn_backend", None) + if backend is None or not backend.supports_multi_doc_packed_varlen(): + return False + return not getattr(attention_layer, "use_ring", False) + + @dataclass class MiniMaxH3DiTArchConfig: num_layers: int = 50 @@ -106,6 +141,7 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"}) # video/text/audio tokens (padding is clamped to 0 before the embedding # lookup and masked out afterwards). MINIMAX_H3_ADALN_MODALITY_NUM = 3 +_LOCAL_SP_PREPARE_HOOK = "sp_input---local_sp_prepare" # Opt-in fp16-range protection for the NPU ascend_laser_attention kernel # (consumed only via the "laser_input_scale" extra key; other backends and @@ -144,6 +180,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset( "packed_seq_params", "refiner_packed_seq_params", "video_token_layout", + "rope_table", } ) @@ -187,28 +224,49 @@ def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> RMSNorm return RMSNorm(size, eps=eps, dtype=dtype) -def _modulate_scale_shift( - x: torch.Tensor, - shift: torch.Tensor, - scale: torch.Tensor, - indices: torch.Tensor, +def _sequence_parallel_local_span( + seq_len: int, *, - dtype: torch.dtype, -) -> torch.Tensor: - # Apply per-index affine modulation: x * (1 + scale[idx]) + shift[idx]. - return (x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(dtype) + hooks_applied: bool, +) -> tuple[int, int]: + """Return the packed-row span owned by this sequence-parallel rank.""" + from vllm_omni.diffusion.forward_context import ( + get_ulysses_mode, + is_forward_context_available, + ) + if not hooks_applied or not is_forward_context_available(): + return 0, seq_len + if get_ulysses_mode(default="strict") != "strict": + return 0, seq_len -def _modulate_gate( - x: torch.Tensor, - gate: torch.Tensor, - other: torch.Tensor, - indices: torch.Tensor, - *, - dtype: torch.dtype, -) -> torch.Tensor: - # Apply the per-index gated residual: x + gate[idx] * other. - return (x + gate.index_select(0, indices) * other).to(dtype) + try: + from vllm_omni.diffusion.distributed.parallel_state import ( + get_allgather_parallel_world_size, + get_ring_parallel_world_size, + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_ulysses_parallel_world_size, + ) + + world_size = int(get_sequence_parallel_world_size()) + rank = int(get_sequence_parallel_rank()) + ulysses_world_size = int(get_ulysses_parallel_world_size()) + ring_world_size = int(get_ring_parallel_world_size()) + allgather_world_size = int(get_allgather_parallel_world_size()) + except AssertionError: + return 0, seq_len + + if world_size <= 1 or ulysses_world_size != world_size: + return 0, seq_len + if ring_world_size != 1 or allgather_world_size != 1: + return 0, seq_len + if seq_len < world_size or seq_len % world_size: + return 0, seq_len + + chunk_size = seq_len // world_size + start = rank * chunk_size + return start, chunk_size class MiniMaxH3Rope(nn.Module): @@ -237,6 +295,15 @@ class MiniMaxH3Rope(nn.Module): return torch.cat((half, half), dim=-1) # [S, 96] +def _build_rope_table(freqs: torch.Tensor) -> torch.Tensor: + """Materialize H3's packed ``[cos(freqs[:48]), sin(freqs[:48])]`` table.""" + half = freqs.shape[-1] // 2 + return torch.cat( + (torch.cos(freqs[..., :half]), torch.sin(freqs[..., :half])), + dim=-1, + ).to(_BF16_DTYPE) + + class MiniMaxH3TimeEmbedder(nn.Module): def __init__( self, @@ -392,6 +459,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int, + num_requests: int = 1, video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """Run packed attention as a small eager island. @@ -401,29 +469,61 @@ class MiniMaxH3Attention(nn.Module): narrow lets regional compile fuse projections, norms, RoPE, and the surrounding DiT block without repeated graph breaks. """ - # max_seqlen is already the first (real) packed document length. Do - # not read the CUDA cu_seqlens scalars here: this function runs once - # per layer and .item() would serialize every attention launch. + # max_seqlen is already the longest packed document length. Do not read + # the CUDA cu_seqlens scalars here: this function runs once per layer + # and .item() would serialize every attention launch. ``num_requests`` + # is carried as a Python int for the same reason. if not 0 < max_seqlen <= packed_total: raise ValueError( f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}" ) - used = min(max_seqlen, packed_total) attn_mask = None - # Ring attention can dispatch to a different implementation from the - # configured backend, so the no-mask fast paths are local-only. - # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN). - # supports_packed_mask_free: backend consumes the packed metadata - # without ever reading attn_mask (CUDA packed varlen, NPU - # npu_attn_varlen opt-in with its own fallback rebuild). - no_mask = not getattr(self.attention, "use_ring", False) and ( - self.attention.attn_backend.supports_prefix_kv_slicing - or self.attention.attn_backend.supports_packed_mask_free() - ) - if used < packed_total and not no_mask: - attn_mask = torch.arange(packed_total, device=q.device)[None] < used + mask_free_packed_padding = False + if num_requests > 1: + # A step-mode batch packs one document per request, so its valid + # rows are block-diagonal rather than a prefix: neither a KV prefix + # length nor a 1-D key mask can describe them. Such a layout is + # only correct on a backend that actually attends by cu_seqlens as + # a block-diagonal plan. Check the capability (not the backend + # name): FLASH_ATTN's NPU/XPU variants would otherwise silently + # fall back to a padding-mask rebuild that spans the whole packed + # row and attend across request boundaries. + if not _attention_isolates_packed_requests(self.attention): + backend_name = self.attention.attn_backend.get_name() + raise ValueError( + f"MiniMax H3 packed a {num_requests}-request batch, but the resolved " + f"attention ({backend_name}, use_ring={getattr(self.attention, 'use_ring', False)}) " + "does not isolate multi-document packed cu_seqlens. Run one request " + "per forward on this backend." + ) + used = packed_total + else: + used = min(max_seqlen, packed_total) + # Ring attention can dispatch to a different implementation from the + # configured backend, so the no-mask fast paths are local-only. + # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN). + # supports_packed_mask_free: backend consumes the packed metadata + # without ever reading attn_mask (CUDA packed varlen, NPU + # npu_attn_varlen opt-in with its own fallback rebuild). + use_ring = getattr(self.attention, "use_ring", False) + mask_free_packed_padding = not use_ring and self.attention.attn_backend.supports_packed_mask_free() + no_mask = not use_ring and ( + self.attention.attn_backend.supports_prefix_kv_slicing or mask_free_packed_padding + ) + if used < packed_total and not no_mask: + attn_mask = torch.arange(packed_total, device=q.device)[None] < used metadata = AttentionMetadata( attn_mask=attn_mask, + packed_padding=( + PackedPaddingMetadata( + q_length=used, + kv_length=used, + cu_seqlens_q=cu_seqlens[:2], + cu_seqlens_k=cu_seqlens[:2], + ) + if mask_free_packed_padding + else None + ), extra={ "cu_seqlens_q": cu_seqlens, "cu_seqlens_k": cu_seqlens, @@ -453,10 +553,11 @@ class MiniMaxH3Attention(nn.Module): self, x: torch.Tensor, *, - rope_freqs: torch.Tensor | None, + rope_table: torch.Tensor | None, cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int | None = None, + num_requests: int = 1, sp_seq_lens: list[int] | None = None, video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: @@ -479,15 +580,22 @@ class MiniMaxH3Attention(nn.Module): q = q.view(total, self.num_heads, self.head_dim) k = k.view(total, self.num_kv_heads, self.head_dim) v = v.view(total, self.num_kv_heads, self.head_dim) - q = self.q_norm(q) - k = self.k_norm(k) - if rope_freqs is not None: - q = self._apply_rope(q, rope_freqs) - k = self._apply_rope(k, rope_freqs) + if rope_table is None: + q = self.q_norm(q) + k = self.k_norm(k) + else: + q, k = fused_qk_norm_rope( + q, + k, + self.q_norm.weight, + self.k_norm.weight, + rope_table, + self.q_norm.variance_epsilon, + ) - # The packed layout uses a second document for alignment padding. - # Local/Ulysses backends unpad it, while Ring keeps aligned rows for - # fixed-size P2P buffers. + # Each request contributes a document for its rows plus one for any + # nonempty alignment padding. Local/Ulysses backends unpad it, while + # Ring keeps aligned rows for fixed-size P2P buffers. out = self._run_packed_attention( q, k, @@ -498,6 +606,7 @@ class MiniMaxH3Attention(nn.Module): # backend receives the global sequence after all-to-all, so carry # its Python length explicitly instead of inferring it from q. packed_total=packed_total if packed_total is not None else q.shape[0], + num_requests=num_requests, video_layout=video_layout, ) out = out.reshape(total, self.num_heads * self.head_dim) @@ -523,6 +632,7 @@ class MiniMaxH3MLP(nn.Module): quant_config=quant_config, prefix=f"{prefix}.fc1", ) + self.act_fn = SiluAndMul() # Chunk the fused fc1 output as [gate, up], then compute # silu(gate) * up. self.fc2 = RowParallelLinear( @@ -537,8 +647,7 @@ class MiniMaxH3MLP(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: hidden, _ = self.fc1(x) - gate, up = hidden.chunk(2, dim=-1) - hidden = nn.functional.silu(gate) * up + hidden = self.act_fn(hidden) out, _ = self.fc2(hidden) return out @@ -625,12 +734,14 @@ class MiniMaxH3TokenRefinerBlock(nn.Module): *, cu_seqlens: torch.Tensor, max_seqlen: int, + num_requests: int = 1, ) -> torch.Tensor: x = x + self.attn( self.norm1(x), - rope_freqs=None, + rope_table=None, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + num_requests=num_requests, ) x = x + self.mlp(self.norm2(x)) return x @@ -663,9 +774,10 @@ class MiniMaxH3TokenRefiner(nn.Module): *, cu_seqlens: torch.Tensor, max_seqlen: int, + num_requests: int = 1, ) -> torch.Tensor: for block in self.blocks: - x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, num_requests=num_requests) return self.final_norm(x) @@ -707,10 +819,11 @@ class MiniMaxH3DiTBlock(nn.Module): *, t_emb: torch.Tensor, combined_indices: torch.Tensor, - rope_freqs: torch.Tensor, + rope_table: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int, + num_requests: int = 1, sp_seq_lens: list[int] | None = None, video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: @@ -731,24 +844,37 @@ class MiniMaxH3DiTBlock(nn.Module): ) = self.adaln_proj(t_emb) residual = x - h = self.norm1(x) - h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE) + h = rms_norm_indexed_scale_shift( + x, + self.norm1.weight, + shift_msa, + scale_msa, + combined_indices, + self.norm1.variance_epsilon, + ) h = self.attn( h, - rope_freqs=rope_freqs, + rope_table=rope_table, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=packed_total, + num_requests=num_requests, sp_seq_lens=sp_seq_lens, video_layout=video_layout, ) - x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE) - + x, h = indexed_gate_rms_norm_scale_shift( + residual, + gate_msa, + h, + self.norm2.weight, + shift_mlp, + scale_mlp, + combined_indices, + self.norm2.variance_epsilon, + ) residual = x - h = self.norm2(x) - h = _modulate_scale_shift(h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE) h = self.mlp(h) - return _modulate_gate(residual, gate_mlp, h, combined_indices, dtype=_BF16_DTYPE) + return indexed_gate(residual, gate_mlp, h, combined_indices) class MiniMaxH3FinalLayer(nn.Module): @@ -803,7 +929,7 @@ class MiniMaxH3FinalLayer(nn.Module): """ shift, scale = self.adaln_proj(t_emb) h = self.norm(x) - h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE) + h = indexed_scale_shift_(h, shift, scale, inverse_indices) # Preserve full precision through both final output projections. h = h.to(_FP32_DTYPE) video, _ = self.video_out(h) @@ -817,10 +943,10 @@ class MiniMaxH3SPPrepare(nn.Module): def forward( self, hidden_states: torch.Tensor, - rope_freqs: torch.Tensor, + rope_table: torch.Tensor, combined_indices: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return hidden_states, rope_freqs, combined_indices + return hidden_states, rope_table, combined_indices class MiniMaxH3SPGather(nn.Module): @@ -831,6 +957,15 @@ class MiniMaxH3SPGather(nn.Module): class MiniMaxH3DiTModel(nn.Module): + # Loading is tensor-complete: constructor state plus final-layout + # parameters and persistent buffers is sufficient to reconstruct a ready + # inference model. The model-specific validator below checks the preserved + # FP32 portion after a lease-backed restore commits. + host_weight_restore_contract = FinalLayoutModelContract( + implementation_id="minimax-h3-dit", + version="1", + ) + _cache_dit_adapter_config = CacheDiTAdapterConfig( block_forward_patterns={"blocks": ForwardPattern.Pattern_3}, # H3 is CFG-distilled and performs one transformer forward per step. @@ -871,6 +1006,13 @@ class MiniMaxH3DiTModel(nn.Module): split_output=True, ), }, + "local_sp_prepare": { + 2: SequenceParallelInput( + split_dim=0, + expected_dims=1, + split_output=True, + ), + }, "sp_gather": SequenceParallelOutput(gather_dim=0, expected_dims=2), } # The checkpoint already stores qkv and the MLP gate/up as single tensors @@ -878,6 +1020,14 @@ class MiniMaxH3DiTModel(nn.Module): # quantization or LoRA to map onto. Address the fused layers directly, e.g. # ignored_layers=["blocks.0.attn.qkv_proj"]. packed_modules_mapping = {} + # Turbo LoRA checkpoints publish separate Q/K/V adapters. This declaration + # lets the legacy diffusion LoRA manager bind them to the packed QKV layer; + # it does not change the fused base-checkpoint loading path above. + stacked_params_mapping = ( + (".attn.qkv_proj", ".attn.to_q", "q"), + (".attn.qkv_proj", ".attn.to_k", "k"), + (".attn.qkv_proj", ".attn.to_v", "v"), + ) def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None: if tp_size < 1: @@ -978,6 +1128,7 @@ class MiniMaxH3DiTModel(nn.Module): ] ) self.sp_prepare = MiniMaxH3SPPrepare() + self.local_sp_prepare = MiniMaxH3SPPrepare() self.sp_gather = MiniMaxH3SPGather() self.final_layer = MiniMaxH3FinalLayer( arch, @@ -990,6 +1141,54 @@ class MiniMaxH3DiTModel(nn.Module): for _, param in self.named_parameters(): param.missing_param_init = "error" + def _rope_local_span(self, seq_len: int) -> tuple[int, int]: + """Return the sequence-parallel rows owned by this DiT rank.""" + local_sp_registry = getattr(self.local_sp_prepare, "_hook_registry", None) + hooks_applied = local_sp_registry is not None + if local_sp_registry is not None: + local_sp_hook = local_sp_registry.get_hook(_LOCAL_SP_PREPARE_HOOK) + hooks_applied = local_sp_hook is not None + return _sequence_parallel_local_span( + seq_len, + hooks_applied=hooks_applied, + ) + + def prepare_rope_table( + self, + img_position_ids: torch.Tensor, + *, + seq_len: int, + ) -> torch.Tensor: + """Build the static local RoPE table once for one denoise branch. + + A MiniMax-H3 denoise branch reuses its packed position IDs at every + scheduler step. The returned table is local to the current sequence- + parallel rank, and therefore must be built by the model that will + consume it rather than cached globally across requests or ranks. + """ + local_start, local_len = self._rope_local_span(seq_len) + rope_position_ids = img_position_ids.narrow(1, local_start, local_len) + return _build_rope_table(self.rope(rope_position_ids).to(img_position_ids.device)) + + def _validate_prepared_rope_table( + self, + rope_table: torch.Tensor, + *, + local_len: int, + device: torch.device, + ) -> None: + expected_width = 6 * self.arch.rope_inv_freq_len + if rope_table.dim() != 2 or tuple(rope_table.shape) != (local_len, expected_width): + raise ValueError( + "rope_table must be [local_seq_len, rotary_dim] for the current " + f"sequence-parallel rank, got {list(rope_table.shape)}; expected " + f"[{local_len}, {expected_width}]." + ) + if rope_table.device != device: + raise ValueError(f"rope_table device {rope_table.device} must match x device {device}.") + if rope_table.dtype != _BF16_DTYPE: + raise ValueError(f"rope_table must be {_BF16_DTYPE}, got {rope_table.dtype}.") + def post_load_weights(self) -> None: for name, param in self.named_parameters(): if name in MINIMAX_H3_FP32_PARAM_NAMES and param.dtype != _FP32_DTYPE: @@ -998,6 +1197,10 @@ class MiniMaxH3DiTModel(nn.Module): if name in MINIMAX_H3_FP32_BUFFER_NAMES and buffer.dtype != _FP32_DTYPE: raise ValueError(f"{name} must stay fp32 after load, got {buffer.dtype}.") + def validate_restored_host_weights(self) -> None: + """Validate mixed-precision invariants after lease-backed restore.""" + self.post_load_weights() + def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], @@ -1056,6 +1259,11 @@ class MiniMaxH3DiTModel(nn.Module): raise ValueError(f"{key}.{field} is required") return value + @staticmethod + def _psp_optional(psp: Any, field: str, default: Any) -> Any: + value = psp.get(field) if isinstance(psp, dict) else getattr(psp, field, None) + return default if value is None else value + def _embed( self, *, @@ -1070,16 +1278,45 @@ class MiniMaxH3DiTModel(nn.Module): refiner_max_seqlen: int, seq_len: int, device: torch.device, + local_span: tuple[int, int], + num_requests: int = 1, ) -> tuple[torch.Tensor, torch.Tensor]: - """Build packed multimodal embeddings for the TP=1/SP=1 inference path. + """Build this rank's packed multimodal embedding rows. - Returns (decoder_input [S, H] bf16, t_emb [M, t_dim] fp32). + Returns (decoder_input [S_local, H] bf16, t_emb [M, t_dim] fp32). + + ``num_requests`` defaults to a single packed request so callers that + pre-date the continuous-batching change (e.g. TeaCache's extractor + contract) do not silently miss the kwarg. ``forward()`` reads the real + value from ``packed_seq_params["num_requests"]``. """ + local_start, local_len = local_span + local_end = local_start + local_len + local_only = local_len != seq_len + if local_only: + img_mask = (img_pos >= local_start) & (img_pos < local_end) + audio_mask = (audio_pos >= local_start) & (audio_pos < local_end) + text_mask = (text_pos >= local_start) & (text_pos < local_end) + img_global_pos = img_pos[img_mask] + audio_global_pos = audio_pos[audio_mask] + img_local_pos = img_global_pos - local_start + audio_local_pos = audio_global_pos - local_start + text_local_pos = text_pos[text_mask] - local_start + text_local_indices = torch.nonzero(text_mask, as_tuple=False).view(-1) + else: + img_global_pos = img_pos + audio_global_pos = audio_pos + img_local_pos = img_pos + audio_local_pos = audio_pos + text_local_pos = text_pos + text_local_indices = None + # Latent embedders stay fp32 in and out; their outputs are cast to the # bf16 sequence dtype only during indexed scattering. - x_rows = x.view(-1, x.shape[-1]).index_select(0, img_pos).to(_FP32_DTYPE) + x_rows = x.view(-1, x.shape[-1]).index_select(0, img_global_pos).to(_FP32_DTYPE) video_embed, _ = self.video_patch_proj(x_rows) - audio_rows = audio_x.view(-1, audio_x.shape[-1]).index_select(0, audio_pos).to(_FP32_DTYPE) + audio_rows = audio_x.view(-1, audio_x.shape[-1]) + audio_rows = audio_rows.index_select(0, audio_global_pos).to(_FP32_DTYPE) audio_embed, _ = self.audio_patch_proj(audio_rows) text_rows = text_embeddings_selected.to(device=device, dtype=_BF16_DTYPE) @@ -1088,12 +1325,31 @@ class MiniMaxH3DiTModel(nn.Module): text_embed, cu_seqlens=refiner_cu_seqlens, max_seqlen=refiner_max_seqlen, + num_requests=num_requests, ) + if text_local_indices is not None: + text_embed = text_embed.index_select(0, text_local_indices) - embeddings = torch.zeros((seq_len, self.hidden_size), device=device, dtype=_BF16_DTYPE) - embeddings.index_add_(0, text_pos, text_embed.to(_BF16_DTYPE)[: text_pos.shape[0]]) - embeddings.index_add_(0, img_pos, video_embed.to(_BF16_DTYPE)[: img_pos.shape[0]]) - embeddings.index_add_(0, audio_pos, audio_embed.to(_BF16_DTYPE)[: audio_pos.shape[0]]) + embeddings = torch.zeros( + (local_len, self.hidden_size), + device=device, + dtype=_BF16_DTYPE, + ) + embeddings.index_add_( + 0, + text_local_pos, + text_embed.to(_BF16_DTYPE)[: text_local_pos.shape[0]], + ) + embeddings.index_add_( + 0, + img_local_pos, + video_embed.to(_BF16_DTYPE)[: img_local_pos.shape[0]], + ) + embeddings.index_add_( + 0, + audio_local_pos, + audio_embed.to(_BF16_DTYPE)[: audio_local_pos.shape[0]], + ) t_emb = self.time_embedder(unique_timesteps) return embeddings, t_emb @@ -1123,7 +1379,6 @@ class MiniMaxH3DiTModel(nn.Module): update_mask = _required_kwarg(kwargs, "update_mask") token_tags = _required_kwarg(kwargs, "token_tags").view(-1).to(torch.long) skip_mask_out_condition = bool(kwargs.get("skip_mask_out_condition", False)) - text_selected = _required_kwarg(kwargs, "prompt_embeds") img_pos = self._pos_ids(_required_kwarg(kwargs, "img_pos_info"), "img_pos_info") @@ -1140,6 +1395,10 @@ class MiniMaxH3DiTModel(nn.Module): psp = _required_kwarg(kwargs, "packed_seq_params") cu_seqlens = self._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32) max_seqlen = int(self._psp_field(psp, "packed_seq_params", "max_seqlen_q")) + # How many requests share this packed sequence. Carried as a host int so + # attention never reads cu_seqlens scalars off the device; a producer + # that omits it is packing a single request. + num_requests = int(self._psp_optional(psp, "num_requests", 1)) refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32) refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")) @@ -1153,8 +1412,26 @@ class MiniMaxH3DiTModel(nn.Module): if inverse_indices.shape[0] != seq_len: raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}") device = x.device - # Compute RoPE frequencies over the full packed sequence. - rope_freqs = self.rope(img_position_ids).to(device) + local_span = self._rope_local_span(seq_len) + local_start, local_len = local_span + rope_table = kwargs.get("rope_table") + if rope_table is None: + if current_omni_platform.is_npu(): + rope_table = self.prepare_rope_table( + img_position_ids, + seq_len=seq_len, + ) + else: + # Keep CUDA/CPU numerically and structurally identical to the + # main-branch reference path used by the H100 accuracy suite. + rope_position_ids = img_position_ids.narrow(1, local_start, local_len) + rope_table = _build_rope_table(self.rope(rope_position_ids).to(device)) + else: + self._validate_prepared_rope_table( + rope_table, + local_len=local_len, + device=device, + ) decoder_input, t_emb = self._embed( x=x, @@ -1166,8 +1443,10 @@ class MiniMaxH3DiTModel(nn.Module): text_pos=text_pos.to(device), refiner_cu_seqlens=refiner_cu.to(device), refiner_max_seqlen=refiner_max, + num_requests=num_requests, seq_len=seq_len, device=device, + local_span=local_span, ) combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device) @@ -1175,32 +1454,56 @@ class MiniMaxH3DiTModel(nn.Module): hidden = decoder_input cu_seqlens = cu_seqlens.to(device) - block_rope = rope_freqs + block_rope = rope_table block_combined = combined_indices - hidden, block_rope, block_combined = self.sp_prepare( - hidden, - block_rope, - block_combined, - ) + if local_len == seq_len: + hidden, block_rope, block_combined = self.sp_prepare( + hidden, + block_rope, + block_combined, + ) + else: + hidden, block_rope, block_combined = self.local_sp_prepare( + hidden, + block_rope, + block_combined, + ) for block in self.blocks: hidden = block( hidden, t_emb=t_emb, combined_indices=block_combined, - rope_freqs=block_rope, + rope_table=block_rope, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=seq_len, + num_requests=num_requests, video_layout=video_layout, ) - hidden = self.sp_gather(hidden) - - video_logits, audio_logits = self.final_layer( - hidden, - t_emb=t_emb, - inverse_indices=inverse_indices, - ) + if local_len == seq_len: + hidden = self.sp_gather(hidden) + video_logits, audio_logits = self.final_layer( + hidden, + t_emb=t_emb, + inverse_indices=inverse_indices, + ) + else: + local_inverse_indices = inverse_indices.narrow( + 0, + local_start, + local_len, + ) + video_logits, audio_logits = self.final_layer( + hidden, + t_emb=t_emb, + inverse_indices=local_inverse_indices, + ) + compact_logits = torch.cat((video_logits, audio_logits), dim=-1) + compact_logits = self.sp_gather(compact_logits) + video_width = self.arch.latents_dim * math.prod(self.arch.patch_size) + video_logits = compact_logits[..., :video_width] + audio_logits = compact_logits[..., video_width:] # Select target and condition rows at inference-output positions, then # zero the condition rows. diff --git a/vllm_omni/diffusion/models/minimax_h3/npu/__init__.py b/vllm_omni/diffusion/models/minimax_h3/npu/__init__.py new file mode 100644 index 000000000..d7b986807 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/npu/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""MiniMax-H3 components for artifacts produced on Ascend NPU. + +The package name records where these artifacts come from, not where they can +run. Everything here is plain ``torch`` plus ``safetensors`` and carries no +``torch_npu`` dependency, so the same checkpoint loads and serves on CUDA and +CPU as well. See ``tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py`` +for the platform-agnostic and CUDA coverage that pins this. +""" + +from .lora import MINIMAX_H3_NATIVE_INFERENCE_STEPS, load_minimax_h3_native_lora + +__all__ = [ + "MINIMAX_H3_NATIVE_INFERENCE_STEPS", + "load_minimax_h3_native_lora", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/npu/lora.py b/vllm_omni/diffusion/models/minimax_h3/npu/lora.py new file mode 100644 index 000000000..2e2d3abd9 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/npu/lora.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Loader for the native-layout MiniMax-H3 distilled LoRA (FlashGen v1.0). + +This contract is keyed by ``key_format=minimax-h3-native`` and is distinct from +the LightX2V Turbo contract in the sibling ``..lora`` module: it keeps the fused +``qkv_proj``, adds ``adaln_proj.linear``, uses rank 64, and declares its own +rectified-flow schedule in safetensors metadata. + +The artifact is trained on Ascend NPU, but nothing below is NPU-specific. Only +``torch`` and ``safetensors`` are used, tensors are read on CPU, and the compute +device is chosen later by the pipeline, so the same file serves on CUDA too. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +import torch +from safetensors import safe_open +from vllm.lora.lora_model import LoRAModel +from vllm.lora.lora_weights import PackedLoRALayerWeights +from vllm.lora.peft_helper import PEFTHelper +from vllm.model_executor.models.utils import WeightsMapper + +from vllm_omni.diffusion.sched.sigma_schedule import BASE_SCHEDULE_KEY, DMD2SigmaSchedule +from vllm_omni.lora.request import LoRARequest + +from ..minimax_h3_transformer import _reorder_grouped_qkv_to_qkv + +_LORA_A_SUFFIX = ".lora_A.default.weight" +_LORA_B_SUFFIX = ".lora_B.default.weight" + +_NATIVE_RANK = 64 +_NATIVE_ALPHA = 64 +_NATIVE_HIDDEN_SIZE = 5376 +_NATIVE_ATTENTION_INNER_SIZE = 7168 +_NATIVE_FFN_HIDDEN_SIZE = 14336 +_NATIVE_TIME_EMBED_DIM = 2688 +_NATIVE_BLOCK_ADALN_OUT = 96768 +_NATIVE_FINAL_ADALN_OUT = 10752 +_NATIVE_NUM_QUERY_GROUPS = 56 +_NATIVE_HEADS_PER_GROUP = 1 +_NATIVE_HEAD_DIM = 128 +_NATIVE_QKV_SLICE = _NATIVE_ATTENTION_INNER_SIZE +_NATIVE_KEY_FORMAT = "minimax-h3-native" +_NATIVE_QKV_LAYOUT = "grouped" +# Public because request validation in the pipeline speaks the same contract. +MINIMAX_H3_NATIVE_INFERENCE_STEPS = 4 +_NATIVE_FILENAME = "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors" +_NATIVE_TARGET_SUFFIXES = ( + "attn.qkv_proj", + "attn.out_proj", + "mlp.fc1", + "mlp.fc2", + "adaln_proj.linear", +) +_NATIVE_TOKEN_REFINER_SUFFIXES = ( + "attn.qkv_proj", + "attn.out_proj", + "mlp.fc1", + "mlp.fc2", +) +_NATIVE_TARGET_DIMS = { + "attn.qkv_proj": (_NATIVE_HIDDEN_SIZE, 3 * _NATIVE_ATTENTION_INNER_SIZE), + "attn.out_proj": (_NATIVE_ATTENTION_INNER_SIZE, _NATIVE_HIDDEN_SIZE), + "mlp.fc1": (_NATIVE_HIDDEN_SIZE, 2 * _NATIVE_FFN_HIDDEN_SIZE), + "mlp.fc2": (_NATIVE_FFN_HIDDEN_SIZE, _NATIVE_HIDDEN_SIZE), + "adaln_proj.linear": (_NATIVE_TIME_EMBED_DIM, _NATIVE_BLOCK_ADALN_OUT), +} +_NATIVE_FINAL_ADALN_DIMS = (_NATIVE_TIME_EMBED_DIM, _NATIVE_FINAL_ADALN_OUT) +_NATIVE_EXPECTED_TARGETS = frozenset( + [*(f"blocks.{block_index}.{suffix}" for block_index in range(50) for suffix in _NATIVE_TARGET_SUFFIXES)] + + [ + *( + f"token_refiner.blocks.{block_index}.{suffix}" + for block_index in range(2) + for suffix in _NATIVE_TOKEN_REFINER_SUFFIXES + ) + ] + + ["final_layer.adaln_proj.linear"] +) +_NATIVE_TARGET_PATTERN = ( + r"^(?:transformer\.blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2)|adaln_proj\.linear)" + r"|transformer\.token_refiner\.blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))" + r"|transformer\.final_layer\.adaln_proj\.linear)$" +) +_NATIVE_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_substr={ + "transformer.": "", + ".lora_A.default.": ".lora_A.", + ".lora_B.default.": ".lora_B.", + } +) + + +def _select_native_file(artifact_path: str | Path) -> Path | None: + path = Path(artifact_path) + if path.is_file(): + return path if path.suffix == ".safetensors" else None + if not path.is_dir(): + return None + + candidate = path / _NATIVE_FILENAME + return candidate if candidate.is_file() else None + + +def _native_target_dims(target: str) -> tuple[int, int]: + if target == "final_layer.adaln_proj.linear": + return _NATIVE_FINAL_ADALN_DIMS + if target.endswith("adaln_proj.linear"): + return _NATIVE_TARGET_DIMS["adaln_proj.linear"] + for suffix in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"): + if target.endswith(suffix): + return _NATIVE_TARGET_DIMS[suffix] + raise ValueError(f"Unsupported MiniMax-H3 native target: {target!r}") + + +def _parse_native_base_schedule(metadata: Mapping[str, str]) -> DMD2SigmaSchedule | None: + """Read the adapter-declared schedule from safetensors string metadata. + + safetensors metadata is string-only, so the native artifact ships the sigma + positions comma separated. Absent means the adapter declares no schedule. + """ + raw = metadata.get(BASE_SCHEDULE_KEY) + if raw is None: + return None + positions = [part.strip() for part in str(raw).split(",") if part.strip()] + if not positions: + raise ValueError(f"MiniMax-H3 native LoRA {BASE_SCHEDULE_KEY} must not be empty") + try: + return DMD2SigmaSchedule.from_positions([float(position) for position in positions]) + except ValueError as exc: + raise ValueError(f"MiniMax-H3 native LoRA {BASE_SCHEDULE_KEY} is malformed: {exc}") from exc + + +def _validate_native_metadata(metadata: Mapping[str, str]) -> DMD2SigmaSchedule: + if metadata.get("key_format") != _NATIVE_KEY_FORMAT: + raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata key_format={_NATIVE_KEY_FORMAT!r}") + if metadata.get("qkv_layout") != _NATIVE_QKV_LAYOUT: + raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata qkv_layout={_NATIVE_QKV_LAYOUT!r}") + try: + rank = int(metadata.get("lora_rank", "")) + alpha = float(metadata.get("lora_alpha", "")) + except (TypeError, ValueError) as exc: + raise ValueError("MiniMax-H3 native LoRA metadata lora_rank/lora_alpha must be numeric") from exc + if rank != _NATIVE_RANK: + raise ValueError(f"MiniMax-H3 native LoRA requires lora_rank={_NATIVE_RANK}, got {rank!r}") + if alpha != _NATIVE_ALPHA: + raise ValueError(f"MiniMax-H3 native LoRA requires lora_alpha={_NATIVE_ALPHA}, got {alpha!r}") + tasks = {part.strip().lower() for part in str(metadata.get("tasks", "")).split(",") if part.strip()} + if tasks != {"t2va"}: + raise ValueError("MiniMax-H3 native LoRA v1.0 supports tasks=t2va only") + schedule = _parse_native_base_schedule(metadata) + if schedule is None: + raise ValueError(f"MiniMax-H3 native LoRA requires safetensors metadata {BASE_SCHEDULE_KEY}") + # Request validation speaks in interval counts, so a mislabeled schedule + # would otherwise silently change the step count instead of failing here. + if schedule.num_inference_steps != MINIMAX_H3_NATIVE_INFERENCE_STEPS: + raise ValueError( + f"MiniMax-H3 native LoRA v1.0 requires a {MINIMAX_H3_NATIVE_INFERENCE_STEPS}-interval base_schedule " + f"({MINIMAX_H3_NATIVE_INFERENCE_STEPS + 1} sigma positions), got {schedule.num_inference_steps}" + ) + return schedule + + +def _validate_native_tensors(checkpoint) -> dict[str, torch.Tensor]: + tensors: dict[str, torch.Tensor] = {} + pairs: dict[str, set[str]] = {} + raw_targets: set[str] = set() + for name in checkpoint.keys(): + if name.endswith(_LORA_A_SUFFIX): + raw_target = name[: -len(_LORA_A_SUFFIX)] + side = "a" + elif name.endswith(_LORA_B_SUFFIX): + raw_target = name[: -len(_LORA_B_SUFFIX)] + side = "b" + else: + raise ValueError(f"Unconsumed MiniMax-H3 native tensor: {name!r}") + if not raw_target.startswith("transformer."): + raise ValueError(f"MiniMax-H3 native LoRA tensors must use transformer.* keys, got {name!r}") + mapped_target = raw_target.removeprefix("transformer.") + raw_targets.add(mapped_target) + + target_sides = pairs.setdefault(mapped_target, set()) + if side in target_sides: + raise ValueError(f"Duplicate MiniMax-H3 native tensor for {mapped_target}.{side}") + target_sides.add(side) + + tensor = checkpoint.get_tensor(name) + if tensor.ndim != 2: + raise ValueError(f"MiniMax-H3 native LoRA tensors must be matrices, got {name}={tuple(tensor.shape)}") + input_dim, output_dim = _native_target_dims(mapped_target) + expected_shape = (_NATIVE_RANK, input_dim) if side == "a" else (output_dim, _NATIVE_RANK) + if tuple(tensor.shape) != expected_shape: + raise ValueError( + f"MiniMax-H3 native tensor has invalid global shape: {name}={tuple(tensor.shape)}, " + f"expected={expected_shape}" + ) + tensors[name] = tensor + + incomplete = sorted(target for target, sides in pairs.items() if sides != {"a", "b"}) + if incomplete: + raise ValueError(f"Incomplete MiniMax-H3 native LoRA pairs: {incomplete}") + missing = sorted(_NATIVE_EXPECTED_TARGETS - raw_targets) + unexpected = sorted(raw_targets - _NATIVE_EXPECTED_TARGETS) + if missing or unexpected: + raise ValueError( + "MiniMax-H3 native target set does not match the supported v1.0 artifact: " + f"missing={len(missing)} {missing[:5]}, unexpected={len(unexpected)} {unexpected[:5]}" + ) + return tensors + + +def _pack_h3_native_fc1(lora_model: LoRAModel) -> None: + for module_name, weights in tuple(lora_model.loras.items()): + if not module_name.endswith(".mlp.fc1"): + continue + gate_b, up_b = weights.lora_b.chunk(2, dim=0) + lora_model.loras[module_name] = PackedLoRALayerWeights( + module_name=module_name, + rank=weights.rank, + lora_alphas=[weights.lora_alpha, weights.lora_alpha], + lora_a=[weights.lora_a, weights.lora_a], + lora_b=[gate_b.contiguous(), up_b.contiguous()], + scaling=[weights.scaling, weights.scaling], + ) + + +def _pack_h3_native_qkv(lora_model: LoRAModel) -> None: + for module_name, weights in tuple(lora_model.loras.items()): + if not module_name.endswith(".attn.qkv_proj"): + continue + reordered_b = _reorder_grouped_qkv_to_qkv( + weights.lora_b, + num_query_groups=_NATIVE_NUM_QUERY_GROUPS, + heads_per_group=_NATIVE_HEADS_PER_GROUP, + head_dim=_NATIVE_HEAD_DIM, + ) + q_b, k_b, v_b = torch.split( + reordered_b, + [_NATIVE_QKV_SLICE, _NATIVE_QKV_SLICE, _NATIVE_QKV_SLICE], + dim=0, + ) + lora_model.loras[module_name] = PackedLoRALayerWeights( + module_name=module_name, + rank=weights.rank, + lora_alphas=[weights.lora_alpha, weights.lora_alpha, weights.lora_alpha], + lora_a=[weights.lora_a, weights.lora_a, weights.lora_a], + lora_b=[q_b.contiguous(), k_b.contiguous(), v_b.contiguous()], + scaling=[weights.scaling, weights.scaling, weights.scaling], + ) + + +def load_minimax_h3_native_lora( + *, + partition: str, + lora_request: LoRARequest, + lora_path: str | Path, + dtype: torch.dtype, + unsupported_offload_mode: str | None = None, +) -> tuple[LoRAModel, PEFTHelper, DMD2SigmaSchedule] | None: + """Load a native-layout MiniMax-H3 distilled LoRA through the legacy manager.""" + + lora_file = _select_native_file(lora_path) + if lora_file is None: + return None + with safe_open(lora_file, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if metadata.get("key_format") != _NATIVE_KEY_FORMAT: + if lora_file.name == _NATIVE_FILENAME: + raise ValueError( + f"MiniMax-H3 native LoRA requires safetensors metadata key_format={_NATIVE_KEY_FORMAT!r}" + ) + return None + if lora_file.name != _NATIVE_FILENAME: + raise ValueError(f"MiniMax-H3 native LoRA supports only {_NATIVE_FILENAME!r}, got {lora_file.name!r}") + sigma_schedule = _validate_native_metadata(metadata) + if partition == "ref2va": + raise ValueError("MiniMax-H3 native LoRA supports T2VA only") + if unsupported_offload_mode is not None: + raise ValueError(f"MiniMax-H3 native dynamic LoRA does not support {unsupported_offload_mode}") + tensors = _validate_native_tensors(checkpoint) + + peft_helper = PEFTHelper.from_dict( + { + "r": _NATIVE_RANK, + "lora_alpha": _NATIVE_ALPHA, + "target_modules": _NATIVE_TARGET_PATTERN, + } + ) + lora_model = LoRAModel.from_lora_tensors( + lora_model_id=lora_request.lora_int_id, + tensors=tensors, + peft_helper=peft_helper, + device="cpu", + dtype=dtype, + weights_mapper=_NATIVE_WEIGHTS_MAPPER, + ) + _pack_h3_native_qkv(lora_model) + _pack_h3_native_fc1(lora_model) + return lora_model, peft_helper, sigma_schedule diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/README.md b/vllm_omni/diffusion/models/minimax_h3/ops/README.md new file mode 100644 index 000000000..e882ee78e --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/README.md @@ -0,0 +1,50 @@ +# MiniMax-H3 optimized operators + +This package contains eager optimizations whose contracts are specific to the +MiniMax-H3 model implementation. Keeping them beside the model is intentional, +not an assumption that their underlying operations can never be shared. + +## Why model-owned + +The optimized path must preserve the H3 VAE's exact operation order, dtype and +rounding behavior, tensor layout, execution mode, and remote-code contract. A +kernel that implements a mathematically similar expression may still change the +decoded video, and a kernel that is beneficial on one platform may regress on +another. The model package therefore owns the capability checks, per-operation +input guards, reference fallbacks, numerical evidence, and performance evidence. + +This placement is one point in the design space being discussed in the +[diffusion operator-boundary RFC](https://github.com/vllm-project/vllm-omni/issues/6305). +If that RFC establishes a suitable shared tier for any of these operations, +migration can be handled separately without weakening the current contract. + +## Current boundary + +The VAE installer validates the official model structure before making any +change. It then binds the selected operator set, reuses Omni's existing +`SiluAndMul`, and materializes only the decoder-block Linear weights in the FP16 +dtype already selected by decode autocast. Unsupported model contracts, tensor +inputs, execution modes, and devices retain the original implementation. + +Hardware dispatch is an explicit allowlist. SM90, SM100, and SM103 are enabled; +other capabilities fall back to the reference path. Bit-exact full-decode and +operator evidence has been collected on SM90, with independent full-decode and +stress validation on SM103. Enabling a target and claiming it as validated are +kept separate so the evidence remains clear. + +## Extending platform support + +Validation and improvements for other platforms are welcome. A contribution +should keep the model-facing installation path unchanged and add a complete +operator-set entry with conservative fallback behavior. Please include: + +- the GPU product, compute capability, driver, CUDA, PyTorch, and Triton versions; +- representative production shapes plus edge-shape and unsupported-input tests; +- complete decoded-output comparison against the reference path; +- direct numerical checks for every optimized operation; +- warmed latency measurements for both complete decode and individual operations; +- tests proving that unsupported targets and contracts still use the reference path. + +Bitwise equality is the current numerical contract. If a platform cannot meet +it, please discuss the proposed quality contract and validation criteria before +changing the default path. diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py new file mode 100644 index 000000000..616478128 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 model-specific optimized operators.""" + +from .vae import install_h3_vae_optimizations + +__all__ = ["install_h3_vae_optimizations"] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py new file mode 100644 index 000000000..247e072c8 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Install bit-exact eager optimizations for the MiniMax H3 video VAE.""" + +from __future__ import annotations + +from types import MethodType +from typing import Any + +import torch +import torch.nn as nn + +from vllm_omni.diffusion.layers.activation import SiluAndMul + +from .dispatch import resolve_h3_vae_operators + + +def _is_boolean_flag(value: Any) -> bool: + """Accept bools and the JSON-derived 0/1 flags used by remote VAE code.""" + + return isinstance(value, bool) or type(value) is int and value in (0, 1) + + +def _uses_fp32_attention_norm(attention: nn.Module) -> bool: + """Probe the loaded remote module's Q/K normalization semantics.""" + + forward = getattr(type(attention), "forward", None) + namespace = getattr(forward, "__globals__", None) + norm_input = namespace.get("_vit_norm_input") if isinstance(namespace, dict) else None + if not callable(norm_input): + return False + + probe = torch.empty(1, dtype=torch.float16) + try: + normalized = norm_input(attention.norm_q, probe) + except Exception: + return False + return ( + isinstance(normalized, torch.Tensor) and normalized.shape == probe.shape and normalized.dtype == torch.float32 + ) + + +def _optimized_feed_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor: + if torch.compiler.is_compiling(): + return type(self).forward(self, hidden_states) + + hidden_states = self.w1(hidden_states) + if hidden_states.is_cuda and hidden_states.dtype == torch.float16: + hidden_states = self._omni_silu_and_mul(hidden_states) + else: + gate, hidden_states = hidden_states.chunk(2, dim=-1) + hidden_states = self.act_fn(gate) * hidden_states + return self.w2(hidden_states) + + +def _optimized_attention( + self: nn.Module, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + pack_info: dict[str, Any] | None = None, +) -> torch.Tensor: + pack_info = {} if pack_info is None else pack_info + if torch.compiler.is_compiling() or self.spatial_parallel or rotary_pos_emb is None: + return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + + batch_size, sequence, _ = hidden_states.shape + qkv = self.to_qkv(hidden_states).view( + batch_size, + sequence, + -1, + 3 * self.dim_head, + ) + query, key, value = qkv.chunk(3, dim=-1) + optimized_qk = self._omni_qk_norm_rope( + query, + key, + rotary_pos_emb, + float(self.norm_q.eps), + ) + if optimized_qk is None: + return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + + query, key = optimized_qk + hidden_states = self.perform_attention(query, key, value, pack_info) + hidden_states = hidden_states.reshape(batch_size, sequence, -1) + return self.to_out(hidden_states) + + +def _optimized_transformer_block( + self: nn.Module, + hidden_states: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + pack_info: dict[str, Any] | None = None, +) -> torch.Tensor: + pack_info = {} if pack_info is None else pack_info + if torch.compiler.is_compiling() or hidden_states.dtype != torch.float32: + return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + + normalized = self.norm1(hidden_states.float()).to(hidden_states.dtype) + attention_output = self.attn(normalized, rotary_pos_emb, pack_info) + updated = self._omni_scaled_residual( + hidden_states, + attention_output, + self.scale1, + ) + hidden_states = hidden_states + attention_output * self.scale1 if updated is None else updated + + normalized = self.norm2(hidden_states.float()).to(hidden_states.dtype) + feed_forward_output = self.ff(normalized) + updated = self._omni_scaled_residual( + hidden_states, + feed_forward_output, + self.scale2, + ) + return hidden_states + feed_forward_output * self.scale2 if updated is None else updated + + +def _decoder_block_linears(decoder: nn.Module) -> tuple[nn.Linear, ...] | None: + blocks = getattr(decoder, "transformer_blocks", None) + if not isinstance(blocks, nn.ModuleList) or not blocks: + return None + + linears: list[nn.Linear] = [] + for block in blocks: + attention = getattr(block, "attn", None) + feed_forward = getattr(block, "ff", None) + to_qkv = getattr(attention, "to_qkv", None) + to_out = getattr(attention, "to_out", None) + w1 = getattr(feed_forward, "w1", None) + w2 = getattr(feed_forward, "w2", None) + candidates = (to_qkv, to_out, w1, w2) + dim_head = getattr(attention, "dim_head", None) + scale1 = getattr(block, "scale1", None) + scale2 = getattr(block, "scale2", None) + if ( + not all(isinstance(linear, nn.Linear) for linear in candidates) + or not _is_boolean_flag(getattr(attention, "spatial_parallel", None)) + or not isinstance(dim_head, int) + or dim_head <= 0 + or not callable(getattr(attention, "perform_attention", None)) + or not isinstance(getattr(attention, "norm_q", None), nn.RMSNorm) + or not isinstance(getattr(attention, "norm_k", None), nn.RMSNorm) + or attention.norm_q.weight is not None + or attention.norm_k.weight is not None + or attention.norm_q.eps != attention.norm_k.eps + or not _uses_fp32_attention_norm(attention) + or not isinstance(getattr(block, "norm1", None), nn.RMSNorm) + or not isinstance(getattr(block, "norm2", None), nn.RMSNorm) + or block.norm1.weight is None + or block.norm2.weight is None + or not getattr(block, "use_scale", False) + or not isinstance(scale1, torch.Tensor) + or not isinstance(scale2, torch.Tensor) + or not getattr(feed_forward, "use_gated", False) + or not isinstance(getattr(feed_forward, "act_fn", None), nn.SiLU) + or not hasattr(feed_forward, "_compile_forward_enabled") + or getattr(feed_forward, "_compile_forward_enabled") + or not hasattr(feed_forward, "_compile_forward_fatal") + or not hasattr(feed_forward, "_compiled_forward") + ): + return None + + hidden_size = to_qkv.in_features + if ( + to_qkv.out_features != 3 * to_out.in_features + or to_out.in_features % dim_head != 0 + or to_out.out_features != hidden_size + or w1.in_features != hidden_size + or w1.out_features != 2 * w2.in_features + or w2.out_features != hidden_size + or attention.norm_q.normalized_shape != (dim_head,) + or attention.norm_k.normalized_shape != (dim_head,) + or block.norm1.normalized_shape != (hidden_size,) + or block.norm2.normalized_shape != (hidden_size,) + or scale1.shape != (hidden_size,) + or scale2.shape != (hidden_size,) + ): + return None + linears.extend(candidates) + return tuple(linears) + + +def install_h3_vae_optimizations( + decoder: nn.Module, + *, + device: torch.device, +) -> bool: + """Install the operators selected for ``device`` once.""" + + if getattr(decoder, "_omni_h3_vae_optimizations_installed", False): + return True + operators = resolve_h3_vae_operators(device) + if operators is None: + return False + + linears = _decoder_block_linears(decoder) + if linears is None or any(linear.weight.dtype != torch.float32 for linear in linears): + return False + + # The H3 decode path always uses FP16 CUDA autocast. Persisting these + # rounded decoder-block weights avoids rebuilding the same casts per tile. + for linear in linears: + linear.to(dtype=torch.float16) + + for block in decoder.transformer_blocks: + block.ff._omni_silu_and_mul = SiluAndMul() + block.ff.forward = MethodType(_optimized_feed_forward, block.ff) + block.attn._omni_qk_norm_rope = operators.qk_norm_rope + block.attn.forward = MethodType(_optimized_attention, block.attn) + block._omni_scaled_residual = operators.scaled_residual + block.forward = MethodType(_optimized_transformer_block, block) + + decoder._omni_h3_vae_optimizations_installed = True + return True + + +__all__ = ["install_h3_vae_optimizations"] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py new file mode 100644 index 000000000..d413ef850 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Hardware dispatch for MiniMax H3 VAE operators.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from vllm.triton_utils import HAS_TRITON + +from vllm_omni.platforms import current_omni_platform + +from .qk_norm_rope import try_qk_norm_rope_exact +from .scaled_residual import try_scaled_residual_exact + +QKNormRopeOp = Callable[ + [ + torch.Tensor, + torch.Tensor, + tuple[torch.Tensor, torch.Tensor], + float, + ], + tuple[torch.Tensor, torch.Tensor] | None, +] +ScaledResidualOp = Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor], + torch.Tensor | None, +] + + +@dataclass(frozen=True) +class H3VAEOperatorSet: + """One validated hardware implementation of the H3 VAE operations.""" + + supports: Callable[[torch.device], bool] + qk_norm_rope: QKNormRopeOp + scaled_residual: ScaledResidualOp + + +def _supports_cuda_capability(device: torch.device, expected: int) -> bool: + if ( + not HAS_TRITON + or device.type != "cuda" + or torch.version.hip is not None + or not current_omni_platform.is_cuda() + or not current_omni_platform.is_available() + ): + return False + device_index = 0 if device.index is None else int(device.index) + capability = current_omni_platform.get_device_capability(device_index) + return capability is not None and int(capability.to_int()) == expected + + +def _supports_cuda_sm90(device: torch.device) -> bool: + return _supports_cuda_capability(device, 90) + + +def _supports_cuda_sm100(device: torch.device) -> bool: + return _supports_cuda_capability(device, 100) + + +def _supports_cuda_sm103(device: torch.device) -> bool: + return _supports_cuda_capability(device, 103) + + +# Keep hardware selection flat: adding a backend means adding one operator set, +# without changing the installer or the model execution path. +H3_VAE_OPERATOR_TABLE: tuple[H3VAEOperatorSet, ...] = ( + H3VAEOperatorSet( + supports=_supports_cuda_sm90, + qk_norm_rope=try_qk_norm_rope_exact, + scaled_residual=try_scaled_residual_exact, + ), + H3VAEOperatorSet( + supports=_supports_cuda_sm100, + qk_norm_rope=try_qk_norm_rope_exact, + scaled_residual=try_scaled_residual_exact, + ), + H3VAEOperatorSet( + supports=_supports_cuda_sm103, + qk_norm_rope=try_qk_norm_rope_exact, + scaled_residual=try_scaled_residual_exact, + ), +) + + +def resolve_h3_vae_operators(device: torch.device) -> H3VAEOperatorSet | None: + for operators in H3_VAE_OPERATOR_TABLE: + if operators.supports(device): + return operators + return None + + +__all__ = [ + "H3VAEOperatorSet", + "H3_VAE_OPERATOR_TABLE", + "resolve_h3_vae_operators", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py new file mode 100644 index 000000000..40b7069f6 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bit-exact fused Q/K RMSNorm and RoPE for the MiniMax H3 video VAE.""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +if HAS_TRITON: + + @triton.jit + def _qk_rms_norm_rope_exact_kernel( + x_ptr, + cos_ptr, + sin_ptr, + output_ptr, + x_stride_token, + x_stride_head, + x_stride_dim, + rope_stride_token, + rope_stride_dim, + output_stride_token, + output_stride_head, + output_stride_dim, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + rotary_dim: tl.constexpr, + eps: tl.constexpr, + heads_per_program: tl.constexpr, + ): + token = tl.program_id(0) + head_group = tl.program_id(1) + heads = head_group * heads_per_program + tl.arange(0, heads_per_program) + dims = tl.arange(0, head_dim) + mask = heads[:, None] < num_heads + + offsets = token * x_stride_token + heads[:, None] * x_stride_head + dims[None, :] * x_stride_dim + x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) + normalized = (x * inv_rms[:, None]).to(tl.float16) + + rotary_half = rotary_dim // 2 + pair_dims = tl.where( + dims < rotary_half, + dims + rotary_half, + tl.where(dims < rotary_dim, dims - rotary_half, dims), + ) + pair_offsets = token * x_stride_token + heads[:, None] * x_stride_head + pair_dims[None, :] * x_stride_dim + pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + pair_normalized = (pair_x * inv_rms[:, None]).to(tl.float16) + + rope_mask = dims < rotary_dim + rope_offsets = token * rope_stride_token + dims * rope_stride_dim + cos = tl.load(cos_ptr + rope_offsets, mask=rope_mask, other=0.0).to(tl.float16) + sin = tl.load(sin_ptr + rope_offsets, mask=rope_mask, other=0.0).to(tl.float16) + signed_pair = tl.where( + dims[None, :] < rotary_half, + -pair_normalized, + pair_normalized, + ) + # Preserve multiply-round, multiply-round, add-round. Ordinary Triton + # arithmetic may contract this expression into an FMA. + rotated = tl.inline_asm_elementwise( + """ + { + .reg .b32 first; + .reg .b32 second; + mul.rn.f16x2 first, $1, $2; + mul.rn.f16x2 second, $3, $4; + add.rn.f16x2 $0, first, second; + } + """, + constraints="=r,r,r,r,r", + args=[normalized, cos[None, :], signed_pair, sin[None, :]], + dtype=tl.float16, + is_pure=True, + pack=2, + ) + output = tl.where(rope_mask[None, :], rotated, normalized) + + output_offsets = ( + token * output_stride_token + heads[:, None] * output_stride_head + dims[None, :] * output_stride_dim + ) + tl.store(output_ptr + output_offsets, output, mask=mask) + + +def _supported_inputs( + q: torch.Tensor, + k: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor], +) -> bool: + cos, sin = rotary_pos_emb + return ( + HAS_TRITON + and torch.version.hip is None + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + and q.is_cuda + and k.is_cuda + and q.device == k.device + and q.dtype == torch.float16 + and k.dtype == q.dtype + and q.ndim == 4 + and k.shape == q.shape + and q.numel() > 0 + and q.shape[-1] == 64 + and q.stride(-1) == 1 + and k.stride(-1) == 1 + and cos.shape == sin.shape + and cos.ndim == 4 + and cos.shape[:2] == q.shape[:2] + and cos.shape[2] == 1 + and cos.shape[-1] == 48 + and cos.device == q.device + and sin.device == q.device + and cos.dtype == q.dtype + and sin.dtype == q.dtype + and cos.is_contiguous() + and sin.is_contiguous() + and cos.stride() == sin.stride() + ) + + +def try_qk_norm_rope_exact( + q: torch.Tensor, + k: torch.Tensor, + rotary_pos_emb: tuple[torch.Tensor, torch.Tensor], + eps: float, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return the fused bit-exact result, or ``None`` outside its contract.""" + + if not _supported_inputs(q, k, rotary_pos_emb): + return None + + output_shape = q.shape + batch, sequence, heads, head_dim = output_shape + tokens = batch * sequence + cos, sin = rotary_pos_emb + q = q.reshape(tokens, heads, head_dim) + k = k.reshape(tokens, heads, head_dim) + cos = cos.reshape(tokens, cos.shape[-1]) + sin = sin.reshape(tokens, sin.shape[-1]) + q_output = torch.empty_like(q) + k_output = torch.empty_like(k) + # This launch layout is part of the exactness contract on the validated + # SM90 and SM103 targets: changing the warp-to-row mapping changes the + # FP32 RMS reduction order. + heads_per_program = 8 + grid = (tokens, triton.cdiv(heads, heads_per_program)) + launch_args = { + "num_heads": heads, + "head_dim": head_dim, + "rotary_dim": cos.shape[-1], + "eps": eps, + "heads_per_program": heads_per_program, + "num_warps": 4, + } + _qk_rms_norm_rope_exact_kernel[grid]( + q, + cos, + sin, + q_output, + q.stride(0), + q.stride(1), + q.stride(2), + cos.stride(0), + cos.stride(1), + q_output.stride(0), + q_output.stride(1), + q_output.stride(2), + **launch_args, + ) + _qk_rms_norm_rope_exact_kernel[grid]( + k, + cos, + sin, + k_output, + k.stride(0), + k.stride(1), + k.stride(2), + cos.stride(0), + cos.stride(1), + k_output.stride(0), + k_output.stride(1), + k_output.stride(2), + **launch_args, + ) + return q_output.reshape(output_shape), k_output.reshape(output_shape) + + +__all__ = ["try_qk_norm_rope_exact"] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/scaled_residual.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/scaled_residual.py new file mode 100644 index 000000000..290e10e5f --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/scaled_residual.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bit-exact scaled residual update for the MiniMax H3 video VAE.""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +if HAS_TRITON: + + @triton.jit + def _mul_rn_f32(x, y): + return tl.inline_asm_elementwise( + "mul.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + @triton.jit + def _scaled_residual_exact_kernel( + output_ptr, + residual_ptr, + branch_ptr, + scale_ptr, + residual_stride_row, + branch_stride_row, + hidden_size: tl.constexpr, + block_n: tl.constexpr, + ): + row = tl.program_id(0) + columns = tl.arange(0, block_n) + mask = columns < hidden_size + residual = tl.load( + residual_ptr + row * residual_stride_row + columns, + mask=mask, + other=0.0, + ).to(tl.float32) + branch = tl.load( + branch_ptr + row * branch_stride_row + columns, + mask=mask, + other=0.0, + ).to(tl.float32) + scale = tl.load(scale_ptr + columns, mask=mask, other=0.0).to(tl.float32) + # Prevent contraction into FMA and retain the eager rounding boundary. + updated = residual + _mul_rn_f32(branch, scale) + tl.store(output_ptr + row * hidden_size + columns, updated, mask=mask) + + +def try_scaled_residual_exact( + residual: torch.Tensor, + branch: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor | None: + """Return the optimized result, or ``None`` outside its exact contract.""" + + if not ( + HAS_TRITON + and torch.version.hip is None + and not torch.is_grad_enabled() + and not torch.compiler.is_compiling() + and residual.is_cuda + and branch.is_cuda + and residual.device == branch.device == scale.device + and residual.dtype == torch.float32 + and branch.dtype == torch.float16 + and scale.dtype == torch.float32 + and residual.shape == branch.shape + and residual.ndim >= 1 + and residual.shape[-1] == 2048 + and scale.shape == residual.shape[-1:] + and residual.numel() > 0 + and residual.is_contiguous() + and branch.is_contiguous() + and scale.is_contiguous() + ): + return None + + hidden_size = residual.shape[-1] + residual_2d = residual.reshape(-1, hidden_size) + branch_2d = branch.reshape(-1, hidden_size) + output = torch.empty_like(residual_2d) + rows = residual_2d.shape[0] + _scaled_residual_exact_kernel[(rows,)]( + output, + residual_2d, + branch_2d, + scale, + residual_2d.stride(0), + branch_2d.stride(0), + hidden_size=hidden_size, + block_n=triton.next_power_of_2(hidden_size), + num_warps=8, + ) + return output.reshape_as(residual) + + +__all__ = ["try_scaled_residual_exact"] diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py index a57eae700..0dcbd19b8 100644 --- a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py +++ b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py @@ -124,7 +124,7 @@ def minimax_h3_packed_sequence( include_keyframe_cond: bool, keyframe_frame_indices: list[int] | tuple[int, ...] | None = None, frame_count: int | None = None, -) -> dict[str, torch.Tensor]: +) -> dict[str, object]: """Build the packed-sequence structural fields for one CFG branch. The used length is padded up to a multiple of 64. @@ -239,6 +239,14 @@ def minimax_h3_packed_sequence( # Extent of the video segment within the packed sequence. "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64), "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64), + # The single target video uses the same span contract as Ref2VA. + "video_spans": ( + { + "start": video_sl.start, + "latent_grid": (latent_t, ph, pw), + "role": "target", + }, + ), } @@ -300,7 +308,7 @@ def minimax_h3_packed_sequence_ref2va_blocks( ref_blocks: Sequence[Mapping[str, object]], audio_channel: int = 2, seq_len: int | None = None, -) -> dict[str, torch.Tensor]: +) -> dict[str, object]: """General ref2va-family packed layout. ``ref_blocks`` are consumed in request/plan order: @@ -400,6 +408,33 @@ def minimax_h3_packed_sequence_ref2va_blocks( video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows) pad_sl = slice(video_sl.stop, seq_len) + # This is deliberately tensor-free metadata. It is constructed once by + # the denoise branch and consumed by every attention layer without a + # device-to-host synchronization. + video_spans: list[dict[str, object]] = [] + for item in block_slices: + if item["kind"] in ("video", "video_audio"): + visual_sl = item["visual_sl"] + assert isinstance(visual_sl, slice) + video_spans.append( + { + "start": visual_sl.start, + "latent_grid": ( + int(item["latent_t"]), + int(item["latent_h"]) // _PATCH_H, + int(item["latent_w"]) // _PATCH_W, + ), + "role": "reference", + } + ) + video_spans.append( + { + "start": video_sl.start, + "latent_grid": (latent_t, ph, pw), + "role": "target", + } + ) + input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64) input_ids[text_sl] = MINIMAX_H3_TEXT_ID image_mask = torch.zeros(seq_len, dtype=torch.bool) @@ -560,6 +595,7 @@ def minimax_h3_packed_sequence_ref2va_blocks( # Extent of the video segment within the packed sequence. "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64), "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64), + "video_spans": tuple(video_spans), } diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 3b96b8ab7..2a13baa36 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """vLLM-Omni pipeline for MiniMax H3 FL2VA and Ref2VA partitions.""" from __future__ import annotations @@ -8,10 +9,10 @@ import math import os import tempfile from collections.abc import Iterable, Mapping, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from itertools import groupby from pathlib import Path -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch @@ -20,6 +21,7 @@ import torch.nn as nn from PIL import Image from transformers import Qwen2TokenizerFast, Qwen3VLProcessor from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion import envs from vllm_omni.diffusion.cache.cachedit import ( @@ -28,7 +30,7 @@ from vllm_omni.diffusion.cache.cachedit import ( ) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.parallel_state import ( - get_dit_group, + get_world_group, init_world_group, ) from vllm_omni.diffusion.distributed.utils import get_local_device @@ -43,25 +45,89 @@ from vllm_omni.diffusion.models.interface import ( SupportsComponentDiscovery, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.offloader import OffloadPlan +from vllm_omni.diffusion.offloader import ( + BoundedAllocatorCache, + OffloadPlan, + apply_sequential_offload, + remove_sequential_offload, + sequential_offload_component, +) +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch -from vllm_omni.errors import OmniClientError +from vllm_omni.errors import OmniClientError, client_error_from_metadata from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) +from vllm_omni.model_executor.models.minimax_h3.checkpoint import ( + resolve_minimax_h3_partition, +) +from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MiniMaxH3TextConditioning, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + load_minimax_h3_images as _load_images, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + minimax_h3_multi_image_presentation, + minimax_h3_ref2va_presentation, + minimax_h3_ref2va_video_presentation, + minimax_h3_text_only_ids, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + resolve_minimax_h3_aspect_ratio as _resolve_minimax_h3_aspect_ratio, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + resolve_minimax_h3_output_canvas as _resolve_output_canvas, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + resolve_minimax_h3_reference_image_shape as _reference_image_shape, +) +from vllm_omni.model_executor.models.minimax_h3.reference_video import ( + MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, + deserialize_prepared_reference_videos, + load_audio_file, + load_video_audio, + load_video_frames, + prepare_reference_videos, + sample_reference_video_frames, + validate_reference_audio_files, + validate_reference_audio_waveforms, +) from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import ( + resolve_component_quant_config as _resolve_component_quant_config, +) +from vllm_omni.quantization.component_config import ( + resolve_encoder_quant_config as _resolve_encoder_quant_config, +) +from .batched_packing import minimax_h3_batched_forward_kwargs from .condition_noise import ( minimax_h3_audio_cond_noise_aug_rows, minimax_h3_imgvid_cond_noise_aug_rows, ) -from .denoise_loop import MiniMaxH3DenoiseBranch, minimax_h3_denoise_loop +from .denoise_loop import ( + MiniMaxH3DenoiseBranch, + minimax_h3_denoise_loop, + minimax_h3_prepare_denoise_rows, + minimax_h3_publish_denoise_progress, +) from .encoder import MiniMaxH3Qwen3VLEncoder -from .minimax_h3_transformer import MiniMaxH3DiTModel +from .fasth3 import FastH3WeightFusion, resolve_fasth3_fusion +from .lora import load_minimax_h3_turbo_lora +from .minimax_h3_transformer import ( + MiniMaxH3Attention, + MiniMaxH3DiTModel, + _attention_isolates_packed_requests, +) +from .npu.lora import ( + MINIMAX_H3_NATIVE_INFERENCE_STEPS, + load_minimax_h3_native_lora, +) from .packed_sequence import ( minimax_h3_packed_sequence, minimax_h3_packed_sequence_ref2va_blocks, @@ -71,22 +137,10 @@ from .packed_tokens import ( minimax_h3_unpack_audio_tokens, minimax_h3_unpatchify_video_tokens, ) -from .presentation import ( - minimax_h3_multi_image_presentation_ids, - minimax_h3_multi_image_presentation_token_tags, - minimax_h3_ref2va_presentation, - minimax_h3_ref2va_video_presentation, - minimax_h3_text_only_ids, -) from .quality_policy import MINIMAX_H3_GENERIC_CACHE_KEY, MiniMaxH3QualityPolicy -from .reference_video import ( - load_audio_file, - load_video_audio, - load_video_frames, - prepare_reference_videos, - sample_reference_video_frames, - validate_reference_audio_files, - validate_reference_audio_waveforms, +from .scheduling_minimax_h3_euler_ancestral import ( + minimax_h3_euler_eta0_step, + minimax_h3_rf_v_to_x0, ) from .time_request import ( MINIMAX_H3_SHAPE_PLANNER, @@ -95,8 +149,18 @@ from .time_request import ( ) from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE +if TYPE_CHECKING: + from vllm_omni.diffusion.worker.input_batch import InputBatch + from vllm_omni.diffusion.worker.utils import StepRequestState + logger = init_logger(__name__) +if TYPE_CHECKING: + from vllm.lora.lora_model import LoRAModel + from vllm.lora.peft_helper import PEFTHelper + + from vllm_omni.lora.request import LoRARequest + MINIMAX_H3_FPS = 24 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 @@ -117,6 +181,9 @@ MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024 MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) MINIMAX_H3_MIN_OUTPUT_SECONDS = 4.0 MINIMAX_H3_MAX_OUTPUT_SECONDS = 15.0 +MINIMAX_H3_TURBO_SIGMA_POINTS = 5 +MINIMAX_H3_TURBO_VIDEO_SHIFT = 6.0 +MINIMAX_H3_TURBO_AUDIO_SHIFT = 3.0 MINIMAX_H3_DOWNLOAD_PATTERNS = [ "FL2VA/**", "Ref2VA/model_index.json", @@ -126,39 +193,62 @@ MINIMAX_H3_TASK_DOWNLOAD_PATTERNS = { "fl2va": ["FL2VA/**"], "ref2va": ["Ref2VA/**"], } +MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS = { + "fl2va": [ + "FL2VA/model_index.json", + "FL2VA/transformer/**", + "FL2VA/video_vae/**", + "FL2VA/audio_vae/**", + ], + "ref2va": [ + "Ref2VA/model_index.json", + "Ref2VA/transformer/**", + "Ref2VA/video_vae/**", + "Ref2VA/audio_vae/**", + ], + "combined": [ + "FL2VA/model_index.json", + "FL2VA/transformer/**", + "FL2VA/video_vae/**", + "FL2VA/audio_vae/**", + "Ref2VA/model_index.json", + "Ref2VA/transformer/**", + ], +} + + +def _resolve_minimax_h3_text_encoder_quant_config( + quant_config: QuantizationConfig | None, +) -> QuantizationConfig | None: + resolved = _resolve_component_quant_config(quant_config, "text_encoder") + return _resolve_encoder_quant_config(resolved) def _minimax_h3_partition_for_task( task_type: str | None, model: str | None = None, ) -> str: - task = str(task_type or "auto").lower() - if task == "auto" and model is not None: - path = Path(model) - if path.is_dir() and path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file(): - return path.name.lower() - if task in {"auto", "combined"}: - return "combined" - if task in {"t2va", "fl2va"}: - return "fl2va" - if task == "ref2va": - return "ref2va" - raise ValueError(f"MiniMax-H3 task_type must be one of auto, t2va, fl2va, or ref2va; got {task_type!r}") + return resolve_minimax_h3_partition(model or "", task_type, auto_partition="combined") def _resolve_minimax_h3_model_root( model: str, revision: str | None, partition: str, + *, + load_text_encoder: bool, ) -> Path: path = Path(model) if path.is_dir(): - if path.name in {"FL2VA", "Ref2VA"} and (path / "model_index.json").is_file(): + if path.name in {"FL2VA", "Ref2VA"}: return path.parent return path - allow_patterns = ( - MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition] - ) + if load_text_encoder: + allow_patterns = ( + MINIMAX_H3_DOWNLOAD_PATTERNS if partition == "combined" else MINIMAX_H3_TASK_DOWNLOAD_PATTERNS[partition] + ) + else: + allow_patterns = MINIMAX_H3_DIFFUSION_DOWNLOAD_PATTERNS[partition] return Path( download_weights_from_hf_specific( model_name_or_path=model, @@ -170,15 +260,94 @@ def _resolve_minimax_h3_model_root( ) +# Keys of ``_prepare_request_inputs`` that feed ``diffuse`` / ``_build_denoise_inputs``. +_MINIMAX_H3_DENOISE_INPUT_KEYS = ( + "task", + "text_embeddings", + "text_tags", + "seed", + "latent_t", + "latent_h", + "latent_w", + "audio_t", + "num_frames", + "num_steps", + "video_shift", + "audio_shift", + "base_schedule", + "visual_condition", + "visual_condition_shape", + "audio_condition", + "ref_audio_t", + "ref_blocks", + "visual_condition_shapes", + "audio_condition_lengths", + "keyframe_frame_indices", +) + +# ``StepRequestState.extra`` keys owned by the step-execution path. +_STEP_BRANCH = "minimax_h3_branch" +_STEP_AUDIO_ROWS = "minimax_h3_audio_rows" +_STEP_AUDIO_NOISE_PRED = "minimax_h3_audio_noise_pred" +_STEP_SIGMAS_VIDEO = "minimax_h3_sigmas_video" +_STEP_SIGMAS_AUDIO = "minimax_h3_sigmas_audio" +_STEP_COND_ANCHOR = "minimax_h3_cond_anchor" +_STEP_AUDIO_ANCHOR = "minimax_h3_audio_anchor" +_STEP_SHAPE = "minimax_h3_shape" +_STEP_TRANSFORMER = "minimax_h3_transformer" + + +def _minimax_h3_step_schedule(state: StepRequestState) -> dict[str, float]: + """Return the sigma/timestep values this request needs for its current step. + + Mirrors the per-iteration arithmetic of ``minimax_h3_denoise_loop`` so step + mode and request mode advance identically. + """ + step = int(state.step_index) + sigmas_video = state.extra[_STEP_SIGMAS_VIDEO] + sigmas_audio = state.extra[_STEP_SIGMAS_AUDIO] + sigma_video = float(sigmas_video[step]) + sigma_audio = float(sigmas_audio[step]) + t_video = 1.0 - sigma_video + t_audio = 1.0 - sigma_audio + return { + "sigma_video": sigma_video, + "sigma_video_next": float(sigmas_video[step + 1]), + "sigma_audio": sigma_audio, + "sigma_audio_next": float(sigmas_audio[step + 1]), + "t_video": t_video, + "t_audio": t_audio, + "imgvid_cond_timestep": max(t_video, MINIMAX_H3_IMGVID_COND_TIMESTEP), + "audio_ref_cond_timestep": max(t_audio, MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + } + + def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None: """Read a partition's distilled schedule. An absent key means legacy uniform.""" return DMD2SigmaSchedule.from_metadata(release) -def _resolve_component_quant_config(quant_config, component: str): - if hasattr(quant_config, "resolve"): - return quant_config.resolve(component) - return quant_config +def resolve_minimax_h3_diffusion_model_path( + model: str, + revision: str | None, + task_type: str | None, +) -> str: + """Resolve a repository root or Hub ID to its startup partition.""" + partition = ( + "combined" + if str(task_type or "").lower() == "combined" + else resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va") + ) + model_root = _resolve_minimax_h3_model_root( + model, + revision, + partition, + load_text_encoder=False, + ) + if partition == "combined": + return str(model_root) + subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" + return str(model_root / subdir) def _minimax_h3_post_process(output, output_type: str = "np"): @@ -186,14 +355,24 @@ def _minimax_h3_post_process(output, output_type: str = "np"): The callable crosses the multiprocessing result queue, so it must remain a module-level function that the standard pickle module can resolve. + + ``_prepare_minimax_h3_video_output`` already quantises the video to uint8 + frames on the accelerator, so there is nothing left to scale or transpose + here. """ if not isinstance(output, tuple) or len(output) != 2: return output video, audio = output + if video.dtype != torch.uint8 or video.ndim != 5 or video.shape[-1] not in (3, 4): + # Float or channel-first frames would reach the muxer as a black or + # banded video rather than as an error. + raise ValueError( + f"MiniMax-H3 post-processing expects (B, T, H, W, C) uint8, got {tuple(video.shape)} {video.dtype}" + ) if output_type == "latent": return output if output_type == "np": - video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy() + video = video.detach().cpu().numpy() audio = audio.detach().float().cpu().numpy() video = [sample for sample in video] return { @@ -204,6 +383,22 @@ def _minimax_h3_post_process(output, output_type: str = "np"): } +def _prepare_minimax_h3_video_output(video: torch.Tensor) -> torch.Tensor: + """Quantize decoded frames in place before worker-to-engine transfer.""" + video = video.detach().float() + video.clamp_(0, 1).mul_(255).round_() + return video.permute(0, 2, 3, 4, 1).to( + dtype=torch.uint8, + memory_format=torch.contiguous_format, + ) + + +def _register_dlo_component_cache(cache: BoundedAllocatorCache, *components: Any) -> None: + for component in components: + if component is not None: + component.set_omni_component_cache(cache) + + def get_minimax_h3_post_process_func( od_config: OmniDiffusionConfig, ): @@ -222,39 +417,6 @@ def _load_image(value: Any) -> Image.Image: return images[0] -def _load_images(value: Any) -> list[Image.Image]: - if isinstance(value, (list, tuple)): - if not value: - raise OmniClientError("MiniMax H3 image input must not be empty") - return [_load_image(item) for item in value] - if isinstance(value, (str, os.PathLike)): - file_size = os.path.getsize(value) - if file_size > MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES: - raise OmniClientError("MiniMax H3 reference image exceeds the 30 MiB size limit") - with Image.open(value) as image: - image_format = str(image.format or "").lower() - if image_format and image_format not in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: - raise OmniClientError( - f"MiniMax H3 reference image must use JPG, JPEG, PNG, WEBP, HEIC, or HEIF, got {image.format}" - ) - return [image.convert("RGB")] - if isinstance(value, Image.Image): - return [value.convert("RGB")] - if isinstance(value, torch.Tensor): - tensor = value.detach().float().cpu() - if tensor.ndim == 4 and tensor.shape[0] == 1: - tensor = tensor[0] - if tensor.ndim != 3: - raise OmniClientError(f"image tensor must be [C,H,W], got {tuple(tensor.shape)}") - if tensor.shape[0] in (1, 3, 4): - tensor = tensor.permute(1, 2, 0) - array = tensor.numpy() - if array.max(initial=0) <= 1.0: - array = array * 255.0 - return [Image.fromarray(array.clip(0, 255).astype(np.uint8)).convert("RGB")] - raise OmniClientError(f"unsupported MiniMax H3 image input {type(value)!r}") - - def _load_audio(value: Any) -> tuple[torch.Tensor, int]: if isinstance(value, (list, tuple)) and not (len(value) == 2 and isinstance(value[1], (int, np.integer))): audios = _load_audios(value) @@ -263,7 +425,7 @@ def _load_audio(value: Any) -> tuple[torch.Tensor, int]: return audios[0] if isinstance(value, (str, os.PathLike)): return load_audio_file(str(value)) - if isinstance(value, tuple) and len(value) == 2: + if isinstance(value, (list, tuple)) and len(value) == 2: waveform, sample_rate = value waveform = torch.as_tensor(waveform).float() return waveform, int(sample_rate) @@ -317,6 +479,21 @@ def _resolve_fl2va_keyframe_indices(extra: Mapping[str, Any], image_count: int) return raw_indices +def _reuse_prepared_reference_videos( + prepared: list[dict[str, Any]] | None, + *, + expected_count: int, +) -> list[dict[str, Any]] | None: + if prepared is None: + return None + if len(prepared) != expected_count: + raise OmniClientError("MiniMax H3 prepared-reference-video count does not match the request") + for item in prepared: + if not os.path.isfile(item["prepared_path"]): + raise OmniClientError(f"MiniMax H3 prepared reference video is unavailable: {item['prepared_path']}") + return prepared + + def _validate_ref2va_reference_counts( image_count: int, video_count: int, @@ -337,55 +514,6 @@ def _validate_ref2va_reference_counts( raise OmniClientError("ref2va accepts at most 12 total references") -def _resolve_minimax_h3_aspect_ratio( - task: str, - value: Any, - image: Image.Image | None, -) -> float: - """Resolve H3's task-specific ratio policy. - - T2VA must name one of the official ratios. FL2VA always follows the - first input image, even when a generic client sends ``aspect_ratio``. - Ref2VA defaults to 16:9; ``adaptive``/``auto`` are retained as aliases - for that default for compatibility with existing clients. - """ - if task == "fl2va": - if image is None: - raise OmniClientError("fl2va requires an input image to resolve its aspect ratio") - return float(image.width) / float(image.height) - - if value is None: - if task == "t2va": - raise OmniClientError("t2va requires an explicit aspect_ratio") - return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"] - - if isinstance(value, str): - normalized = value.strip().lower() - if normalized in {"adaptive", "auto"}: - if task == "t2va": - raise OmniClientError("t2va requires an explicit named aspect_ratio, not adaptive") - return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"] - if normalized in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS: - return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS[normalized] - try: - numeric_value = float(normalized) - except (TypeError, ValueError) as exc: - supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS) - raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}") from exc - elif isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool): - numeric_value = float(value) - else: - raise OmniClientError(f"MiniMax H3 aspect_ratio must be a string ratio, got {value!r}") - - if not math.isfinite(numeric_value) or not any( - math.isclose(numeric_value, ratio, rel_tol=0.0, abs_tol=1e-6) - for ratio in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS.values() - ): - supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS) - raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}") - return numeric_value - - def _resolve_minimax_h3_num_outputs(value: Any) -> int: if value is None: return 1 @@ -415,10 +543,62 @@ def _validate_reference_image(image: Image.Image) -> None: def _dit_rank_world() -> tuple[Any, int, int]: if not dist.is_initialized(): return None, 0, 1 - group = get_dit_group() + group = get_world_group().device_group return group, dist.get_rank(group), dist.get_world_size(group) +def _broadcast_rank0_exception(exc: Exception | None) -> None: + """Synchronize a rank-0-only exception across every DiT rank. + + H3 reference-video preparation runs only on rank 0; the other DiT ranks + return ``None`` without touching disk. When rank 0 raises inside that + path it exits :meth:`prepare_encode` before reaching the downstream + ``dist.broadcast`` calls, and non-zero ranks then hang on those + collectives forever. Every rank calls this helper right after the + rank-0-only work, before any subsequent collective, so all ranks either + raise the same error together or all continue. + """ + group, rank, world_size = _dit_rank_world() + if world_size == 1: + if exc is not None: + raise exc + return + if rank == 0: + if exc is None: + payload: list[Any] = [None] + else: + payload = [ + { + "type": type(exc).__name__, + "message": str(exc), + "status_code": getattr(exc, "status_code", None), + "error_type": getattr(exc, "error_type", None), + } + ] + else: + payload = [None] + dist.broadcast_object_list(payload, src=0, group=group) + info = payload[0] + if info is None: + return + if rank == 0: + assert exc is not None + raise exc + # Rebuild a matching client-facing error on non-zero ranks so the runner's + # per-request try/except records the same 4xx status as rank 0. The exact + # subclass need not survive the wire; the message and status suffice. + status_code = info.get("status_code") + error_type = info.get("error_type") + message = f"[rank 0] {info['type']}: {info['message']}" + if status_code is not None: + raise client_error_from_metadata( + message, + status_code=int(status_code), + error_type=error_type, + ) + raise RuntimeError(message) + + def _broadcast_tensor( tensor: torch.Tensor | None, *, @@ -451,49 +631,6 @@ def _broadcast_tensor( return output -def _reference_image_shape(image: Image.Image) -> tuple[int, int]: - width, height = image.size - ratio = width / height - if not 0.4 <= ratio <= 2.5: - raise OmniClientError(f"reference image aspect ratio must be in [0.4, 2.5], got {width}x{height}") - if min(width, height) < 256 or max(width, height) > 5760: - raise OmniClientError(f"reference image dimensions must be in [256, 5760] pixels, got {width}x{height}") - scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height) - return ( - _align_multiple( - width * scale, - MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE, - ), - _align_multiple( - height * scale, - MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE, - ), - ) - - -def _resolve_output_canvas(aspect_ratio: float, short_edge: int) -> tuple[int, int]: - """Resolve the official H3 ratio/area policy to a 32-pixel canvas.""" - if not math.isfinite(float(aspect_ratio)) or float(aspect_ratio) <= 0: - raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be positive, got {aspect_ratio!r}") - if short_edge != MINIMAX_H3_OUTPUT_SHORT_EDGE: - raise OmniClientError(f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {short_edge}") - if aspect_ratio >= 1.0: - width = float(short_edge) * aspect_ratio - height = float(short_edge) - else: - width = float(short_edge) - height = float(short_edge) / aspect_ratio - area = width * height - if area > MINIMAX_H3_OUTPUT_MAX_PIXELS: - scale = (MINIMAX_H3_OUTPUT_MAX_PIXELS / area) ** 0.5 - width *= scale - height *= scale - return ( - _align_multiple(height, 32), - _align_multiple(width, 32), - ) - - class _SingleRankEncoderGroup: """Lightweight encoder group for ``text_encoder_tp_size == 1``. @@ -523,6 +660,8 @@ class MiniMaxH3Pipeline( ): """CFG-distilled joint video/audio generation for MiniMax H3.""" + supports_step_execution: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"] @@ -532,26 +671,182 @@ class MiniMaxH3Pipeline( encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}), ) - # H3's regular loader performs checkpoint-layout conversions (grouped QKV - # and fused MLP). Keep it on that path until mmap can run the same loader - # callbacks for every affected parameter. - _supports_mmap_loading: ClassVar[bool] = False - # TODO(offload): Re-enable after the generic rank-local mmap path can run - # this model's grouped-QKV reorder and fused-MLP packing before TP sharding. - # Do not bypass the regular loader until that equivalence is tested. _PROFILER_TARGETS: ClassVar[list[str]] = [ "_prepare_reference_videos", "encode_prompt", - "_encode_video_conditions", - "_encode_video_audio_conditions", - "_encode_audio_conditions", + "_encode_visual_conditions", + "_encode_reference_audio_conditions", "diffuse", "decode", + "prepare_encode", + "denoise_step", + "post_decode", ] dummy_run_num_frames: ClassVar[int] = 0 # Only distilled releases pin a schedule, so the default keeps the legacy # uniform path available to partially constructed pipelines. _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {} + # Set from --lora-path during construction; absent means no FastH3 adapter. + _fasth3: FastH3WeightFusion | None = None + + def _load_diffusion_lora_adapter( + self, + *, + lora_request: LoRARequest, + lora_path: str | Path, + dtype: torch.dtype, + ) -> tuple[LoRAModel, PEFTHelper] | None: + # A cache eviction may be followed by a different adapter reusing the + # same client-supplied ID. Every real load replaces the classification. + self._turbo_lora_adapter_ids.discard(lora_request.lora_int_id) + self._native_lora_adapter_ids.discard(lora_request.lora_int_id) + self._lora_sigma_schedules.pop(lora_request.lora_int_id, None) + od_config = getattr(self, "od_config", None) + offload_modes = [] + if getattr(od_config, "enable_cpu_offload", False): + offload_modes.append("model-level CPU offload (--enable-cpu-offload)") + if getattr(od_config, "enable_layerwise_offload", False): + offload_modes.append("layerwise offload (--enable-layerwise-offload)") + loaded = load_minimax_h3_turbo_lora( + partition=self.partition, + lora_request=lora_request, + lora_path=lora_path, + dtype=dtype, + unsupported_offload_mode=" or ".join(offload_modes) or None, + ) + if loaded is not None: + self._turbo_lora_adapter_ids.add(lora_request.lora_int_id) + return loaded + + # Selection is by the artifact's safetensors ``key_format``, not by the + # running platform: the native loader is checkpoint-format parsing with + # no ``torch_npu`` dependency, so it needs no ``current_omni_platform`` + # dispatch and binds the same adapter on NPU, CUDA and CPU. + native_loaded = load_minimax_h3_native_lora( + partition=self.partition, + lora_request=lora_request, + lora_path=lora_path, + dtype=dtype, + unsupported_offload_mode=" or ".join(offload_modes) or None, + ) + if native_loaded is not None: + lora_model, peft_helper, sigma_schedule = native_loaded + self._native_lora_adapter_ids.add(lora_request.lora_int_id) + self._lora_sigma_schedules[lora_request.lora_int_id] = sigma_schedule + return lora_model, peft_helper + return None + + def _validate_diffusion_lora_binding( + self, + *, + lora_model: LoRAModel, + bound_lora_names: frozenset[str], + ) -> None: + if lora_model.id in self._turbo_lora_adapter_ids: + missing = sorted(set(lora_model.loras) - bound_lora_names) + if missing: + raise ValueError( + "MiniMax-H3 Turbo LoRA binding is incomplete: " + f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}" + ) + return + if lora_model.id not in self._native_lora_adapter_ids: + return + missing = sorted(set(lora_model.loras) - bound_lora_names) + if missing: + raise ValueError( + "MiniMax-H3 native LoRA binding is incomplete: " + f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, missing={missing[:5]}" + ) + + def _has_active_turbo_lora(self, sampling: Any) -> bool: + lora_request = sampling.lora_request + return ( + lora_request is not None + and not math.isclose(0.0, float(sampling.lora_scale)) + and lora_request.lora_int_id in self._turbo_lora_adapter_ids + ) + + def _has_active_native_lora(self, sampling: Any) -> bool: + lora_request = sampling.lora_request + return ( + lora_request is not None + and not math.isclose(0.0, float(sampling.lora_scale)) + and lora_request.lora_int_id in self._native_lora_adapter_ids + ) + + def _validate_native_sampling(self, sampling: Any, *, task: str) -> None: + if task != "t2va": + raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only") + # Derive the expected count from the adapter's own schedule so the + # message can never disagree with the schedule the denoise loop runs. + schedule = self._lora_sigma_schedules.get(sampling.lora_request.lora_int_id) + expected_steps = MINIMAX_H3_NATIVE_INFERENCE_STEPS if schedule is None else schedule.num_inference_steps + # Only request mode can take the count from the adapter schedule: step + # mode admits the request in ``StepScheduler``, which reads + # ``num_inference_steps`` off it before any pipeline hook runs. Reject + # omission there rather than advertise a contract that would either fail + # admission or disagree with the denoise loop. + od_config = getattr(self, "od_config", None) + omission_allowed = not getattr(od_config, "step_execution", False) + or_omitted = " or omitted" if omission_allowed else "" + sigma_steps = sampling.num_inference_steps + if sigma_steps is None: + if omission_allowed: + return + raise OmniClientError( + f"MiniMax-H3 native LoRA requires an explicit num_inference_steps={expected_steps} " + "under step execution, because the step scheduler derives the total step count from " + "the request before the adapter schedule is known" + ) + if int(sigma_steps) == expected_steps + 1: + raise OmniClientError( + "MiniMax-H3 native LoRA uses the distilled interval-count contract; " + f"num_inference_steps must be {expected_steps}{or_omitted}, not {expected_steps + 1}" + ) + if int(sigma_steps) != expected_steps: + raise OmniClientError( + f"MiniMax-H3 native LoRA requires num_inference_steps={expected_steps} " + f"(one denoiser evaluation per sigma interval){or_omitted}" + ) + + def _sigma_schedule_for_request(self, sampling: Any, task: str) -> DMD2SigmaSchedule | None: + lora_request = sampling.lora_request + if ( + lora_request is not None + and not math.isclose(0.0, float(sampling.lora_scale)) + and lora_request.lora_int_id in self._lora_sigma_schedules + ): + adapter_schedule = self._lora_sigma_schedules[lora_request.lora_int_id] + checkpoint_schedule = self._base_schedule_for_task(task) + if checkpoint_schedule is not None: + raise OmniClientError( + "MiniMax-H3 native LoRA cannot be activated on a checkpoint that already pins base_schedule" + ) + return adapter_schedule + return self._base_schedule_for_task(task) + + def _validate_turbo_sampling(self, sampling: Any) -> None: + extra = sampling.extra_args or {} + sigma_points = sampling.num_inference_steps + if sigma_points != MINIMAX_H3_TURBO_SIGMA_POINTS: + raise OmniClientError( + "MiniMax-H3 Turbo requires num_inference_steps=5 (five sigma points produce four denoiser evaluations)" + ) + try: + video_shift = float(extra.get("flow_shift", self.default_video_shift)) + except (TypeError, ValueError) as exc: + raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}") from exc + if not math.isclose(video_shift, MINIMAX_H3_TURBO_VIDEO_SHIFT): + raise OmniClientError(f"MiniMax-H3 Turbo requires flow_shift={MINIMAX_H3_TURBO_VIDEO_SHIFT:g}") + try: + audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift)) + except (TypeError, ValueError) as exc: + raise OmniClientError( + f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}" + ) from exc + if not math.isclose(audio_shift, MINIMAX_H3_TURBO_AUDIO_SHIFT): + raise OmniClientError(f"MiniMax-H3 Turbo requires audio_flow_shift={MINIMAX_H3_TURBO_AUDIO_SHIFT:g}") def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None: """Adopt runner-installed generic Cache-DiT for request transitions.""" @@ -579,14 +874,19 @@ class MiniMaxH3Pipeline( if int(self.parallel_config.cfg_parallel_size) != 1: raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1") self.device = get_local_device() + self.load_text_encoder = od_config.model_loaded["text_encoder"] self.partition = _minimax_h3_partition_for_task( getattr(od_config, "task_type", None), str(od_config.model), ) + self._turbo_lora_adapter_ids: set[int] = set() + self._native_lora_adapter_ids: set[int] = set() + self._lora_sigma_schedules: dict[int, DMD2SigmaSchedule] = {} model_root = _resolve_minimax_h3_model_root( str(od_config.model), od_config.revision, self.partition, + load_text_encoder=self.load_text_encoder, ) model_path = model_root / ("Ref2VA" if self.partition == "ref2va" else "FL2VA") model_index = json.loads((model_path / "model_index.json").read_text(encoding="utf-8")) @@ -658,42 +958,71 @@ class MiniMaxH3Pipeline( quant_config=transformer_quant_config, ) - self.tokenizer = Qwen2TokenizerFast.from_pretrained( - str(model_path), - subfolder="tokenizer", - local_files_only=os.path.isdir(model_path), - ) - self.processor = Qwen3VLProcessor.from_pretrained( - str(model_path), - subfolder="processor", - local_files_only=os.path.isdir(model_path), - ) + self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer) + if self._fasth3 is not None: + self._fasth3.check_serving_contract( + partition=self.partition, + od_config=od_config, + video_shift=self.default_video_shift, + audio_shift=self.default_audio_shift, + ) + + if self.load_text_encoder: + self.tokenizer = Qwen2TokenizerFast.from_pretrained( + str(model_path), + subfolder="tokenizer", + local_files_only=os.path.isdir(model_path), + ) + self.processor = Qwen3VLProcessor.from_pretrained( + str(model_path), + subfolder="processor", + local_files_only=os.path.isdir(model_path), + ) + else: + self.tokenizer = None + self.processor = None _, rank, dit_world = _dit_rank_world() self._dit_rank = rank - text_encoder_tp_size = int(getattr(self.parallel_config, "text_encoder_tp_size", 1)) - if text_encoder_tp_size < 1: - raise ValueError(f"text_encoder_tp_size must be >= 1, got {text_encoder_tp_size}") - if text_encoder_tp_size > dit_world: - raise ValueError( - f"text_encoder_tp_size must not exceed the DiT group size ({dit_world}), got {text_encoder_tp_size}" + if self.load_text_encoder: + text_encoder_tp_size = int(getattr(self.parallel_config, "text_encoder_tp_size", 1)) + if text_encoder_tp_size < 1: + raise ValueError(f"text_encoder_tp_size must be >= 1, got {text_encoder_tp_size}") + if text_encoder_tp_size > dit_world: + raise ValueError( + f"text_encoder_tp_size must not exceed the DiT group size ({dit_world}), got {text_encoder_tp_size}" + ) + # The Qwen3-VL text model uses 64 attention heads / 8 KV heads. + if 64 % text_encoder_tp_size or 8 % text_encoder_tp_size: + raise ValueError( + "text_encoder_tp_size must divide both Qwen3-VL " + f"num_attention_heads (64) and num_key_value_heads (8), " + f"got {text_encoder_tp_size}" + ) + self.text_encoder_tp_size = text_encoder_tp_size + self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size) + self.text_encoder = MiniMaxH3Qwen3VLEncoder( + os.path.join(model_path, "text_encoder"), + device=self.device, + load_model=rank < text_encoder_tp_size, + encoder_group=self.text_encoder_group, + quant_config=_resolve_minimax_h3_text_encoder_quant_config(od_config.quantization_config), ) - # The Qwen3-VL text model uses 64 attention heads / 8 KV heads; the - # encoder shards them across the encoder TP ranks. - if 64 % text_encoder_tp_size or 8 % text_encoder_tp_size: - raise ValueError( - "text_encoder_tp_size must divide both Qwen3-VL " - f"num_attention_heads (64) and num_key_value_heads (8), " - f"got {text_encoder_tp_size}" - ) - self.text_encoder_tp_size = text_encoder_tp_size - self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size) - self.text_encoder = MiniMaxH3Qwen3VLEncoder( - os.path.join(model_path, "text_encoder"), - device=self.device, - load_model=rank < text_encoder_tp_size, - encoder_group=self.text_encoder_group, - ) + if rank < text_encoder_tp_size: + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=str(model_path), + subfolder="text_encoder", + revision=od_config.revision, + prefix="text_encoder.", + fall_back_to_pt=False, + ) + ) + else: + self.text_encoder_tp_size = 0 + self.text_encoder_group = None + self.text_encoder = None + self._encoder_modules = [] stage_components = bool( od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False) ) @@ -711,6 +1040,16 @@ class MiniMaxH3Pipeline( # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae + self._dlo_component_cache = None + if getattr(od_config, "enable_distributed_layerwise_offload", False): + self._dlo_component_cache = BoundedAllocatorCache(self.device) + _register_dlo_component_cache( + self._dlo_component_cache, + self.text_encoder, + self.video_vae, + self.audio_vae, + ) + self._quality_policy = MiniMaxH3QualityPolicy(od_config) self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self) @@ -725,7 +1064,7 @@ class MiniMaxH3Pipeline( def source_prefix(item: tuple[str, torch.Tensor]) -> str: name, _ = item prefix = name.partition(".")[0] + "." - if prefix in {"transformer.", "transformers_ref."}: + if prefix in {"transformer.", "transformers_ref.", "text_encoder."}: return prefix raise ValueError(f"unexpected MiniMax-H3 weight {name!r}") @@ -735,26 +1074,68 @@ class MiniMaxH3Pipeline( if prefix in loaded_prefixes: raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous") loaded_prefixes.add(prefix) - transformer = getattr(self, prefix.removesuffix(".")) - loaded = transformer.load_weights((name[len(prefix) :], tensor) for name, tensor in grouped_weights) - transformer.post_load_weights() + component = getattr(self, prefix.removesuffix(".")) + stream = ((name[len(prefix) :], tensor) for name, tensor in grouped_weights) + if prefix == "transformer." and self._fasth3 is not None: + # Fuse before the model shards anything, which is also the only + # point where the checkpoint's fused QKV/MLP layouts are intact. + stream = self._fasth3.apply(stream) + loaded = component.load_weights(stream) + if prefix != "text_encoder.": + component.post_load_weights() loaded_with_prefix.update(prefix + name for name in loaded) - # The text encoder and both VAEs load eagerly in ``__init__`` rather - # than through ``weights_sources``. Record them for the runner's strict - # missing-parameter check. For the text encoder that report is exact - # because its loader raises on any unloaded parameter or partially filled - # fused parameter; weaken that and gaps here become invisible. The two - # VAEs carry no equivalent guarantee. - for component_name in ("text_encoder", "video_vae", "audio_vae"): + # Both VAEs load eagerly in ``__init__`` rather than through + # ``weights_sources``. The text encoder uses the shared component + # loader so online quantization and offload processing follow the same + # path as the DiT. + for component_name in ("video_vae", "audio_vae"): component = getattr(self, component_name) + if component is None: + continue loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) + if self._fasth3 is not None: + self._fasth3.validate_fully_applied() return loaded_with_prefix + @property + def lora_is_fused(self) -> bool: + """True when --lora-path was consumed as a load-time weight fusion.""" + return self._fasth3 is not None + def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel: if task == "ref2va" and hasattr(self, "transformers_ref"): return self.transformers_ref return self.transformer + def _resolve_sigma_positions(self, task: str, sampling: Any) -> tuple[tuple[float, ...] | None, int]: + """Pick the rectified-flow positions this request denoises on. + + Returns them explicitly, or ``None`` to leave the uniform ladder to be + derived from the step count, together with the count the rest of the + request speaks in. + """ + if self._fasth3 is not None: + # A fused student carries its own positions; the checkpoint + # underneath it is the many-step teacher, whose schedule does not + # apply. Its five points bound four transformer forwards, and + # forwards is the unit ``check_request``, the pinned-schedule branch + # below and Cache-DiT all speak in. + positions = self._fasth3.base_schedule + return positions, len(positions) - 1 + sigma_schedule = self._sigma_schedule_for_request(sampling, task) + if sigma_schedule is None: + return None, int(sampling.num_inference_steps or 50) + # The schedule lists sigma boundaries; the denoise loop runs one step + # per interval, and that count is what requests and Cache-DiT speak in. + num_steps = sigma_schedule.num_inference_steps + requested_steps = sampling.num_inference_steps + if requested_steps is not None and int(requested_steps) != num_steps: + raise OmniClientError( + "this MiniMax H3 checkpoint pins a distilled sigma schedule; num_inference_steps " + f"must be {num_steps} or omitted, got {int(requested_steps)}" + ) + return sigma_schedule.base_schedule, num_steps + def _base_schedule_for_task(self, task: str) -> DMD2SigmaSchedule | None: """Return the distilled schedule of the partition that serves ``task``.""" partition = "ref2va" if task == "ref2va" else "fl2va" @@ -764,6 +1145,9 @@ class MiniMaxH3Pipeline( self, requested: str | None, multi_modal_data: dict[str, Any], + *, + has_turbo_lora: bool = False, + has_native_lora: bool = False, ) -> str: if requested is None: # A Ref2VA-only startup has no FL2VA transformer; preserve its @@ -781,6 +1165,12 @@ class MiniMaxH3Pipeline( raise OmniClientError( f"checkpoint partition {self.partition!r} supports {sorted(self.supported_tasks)}, got task={task!r}" ) + if task == "ref2va" and has_turbo_lora: + raise OmniClientError("MiniMax-H3 Turbo LoRA supports T2VA/FL2VA requests only") + if has_native_lora and task != "t2va": + raise OmniClientError("MiniMax-H3 native LoRA supports T2VA requests only") + if self._fasth3 is not None: + self._fasth3.check_task(task) return task def _resolve_shape( @@ -942,12 +1332,7 @@ class MiniMaxH3Pipeline( if task == "fl2va": if prepared_videos: raise OmniClientError("fl2va does not accept video conditions") - ids = minimax_h3_multi_image_presentation_ids( - self.tokenizer, - prompt=prompt, - image_token_counts=image_token_counts, - ) - tags = minimax_h3_multi_image_presentation_token_tags( + ids, tags = minimax_h3_multi_image_presentation( self.tokenizer, prompt=prompt, image_token_counts=image_token_counts, @@ -1117,9 +1502,7 @@ class MiniMaxH3Pipeline( input_ids: torch.Tensor, vision_kwargs: dict[str, torch.Tensor], ) -> torch.Tensor: - if self.od_config.enable_cpu_offload and not getattr( - self.od_config, "enable_distributed_layerwise_offload", False - ): + if getattr(self, "_model_cpu_offload_modules", None): # Invoke nn.Module.__call__ so the generic model-level offloader # swaps the resident DiT and encoder. return self.text_encoder(input_ids, **vision_kwargs) @@ -1129,11 +1512,8 @@ class MiniMaxH3Pipeline( ): # Layerwise DiT offload already provides the low-residency encoder # phase used by the checkpoint reference. - self.text_encoder.load_to_device() - try: + with self._component_on_device(self.text_encoder): return self.text_encoder.encode_ids(input_ids, **vision_kwargs) - finally: - self.text_encoder.offload_to_cpu() # Keep both Qwen and DiT resident across requests. Moving either model # here makes encoder latency include a tens-of-gigabytes PCIe transfer, @@ -1148,31 +1528,78 @@ class MiniMaxH3Pipeline( or getattr(od_config, "enable_distributed_layerwise_offload", False) ) + def enable_omni_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool, + use_hsdp: bool, + ) -> None: + if getattr(self, "_model_cpu_offload_modules", None): + return + + components = ModuleDiscovery.discover(self) + dits = components.dits + stages = [*components.encoders, *components.vaes] + modules = [*dits, *stages] + apply_sequential_offload( + dit_modules=dits, + encoder_modules=stages, + device=device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + offload_initial_dits=True, + ) + + self._model_cpu_offload_modules = modules + logger.info( + "MiniMax-H3 model-level CPU offload enabled for %d DiT(s), text encoder, video VAE, and audio VAE", + len(dits), + ) + + def disable_omni_model_cpu_offload(self) -> None: + modules = getattr(self, "_model_cpu_offload_modules", None) + if not modules: + return + remove_sequential_offload(modules) + self._model_cpu_offload_modules = [] + @contextmanager def _component_on_device(self, component: nn.Module): + if getattr(self, "_model_cpu_offload_modules", None): + with sequential_offload_component(component): + yield + return staged = self._uses_manual_component_offload() - if staged: - component.load_to_device() try: - yield - finally: if staged: - component.offload_to_cpu() - - def _encode_visual_condition( - self, - image: Image.Image, - ) -> torch.Tensor: - _, rank, _ = _dit_rank_world() - rows = None - if rank == 0: - with self._component_on_device(self.video_vae): - rows = self.video_vae.encode_image(image) - return _broadcast_tensor( - rows, - dtype=torch.float32, - device=self.device, - ) + component.load_to_device() + yield + except BaseException: + if staged: + try: + component.offload_to_cpu() + except BaseException: + logger.exception("Failed to release %s after component failure", component.__class__.__name__) + cache = getattr(self, "_dlo_component_cache", None) + if cache is not None: + try: + cache.release_if_needed(force=True) + except BaseException: + logger.exception("Failed to release retained allocator cache after component failure") + raise + else: + if staged: + try: + component.offload_to_cpu() + except BaseException: + cache = getattr(self, "_dlo_component_cache", None) + if cache is not None: + try: + cache.release_if_needed(force=True) + except BaseException: + logger.exception("Failed to release retained allocator cache after offload failure") + raise def _encode_visual_conditions( self, @@ -1183,44 +1610,34 @@ class MiniMaxH3Pipeline( ) -> tuple[torch.Tensor | None, list[tuple[int, int, int]]]: rows: list[torch.Tensor] = [] shapes: list[tuple[int, int, int]] = [] - for image in images: - rows.append(self._encode_visual_condition(image)) - shapes.append((1, image.height // 16, image.width // 16)) - if video_count: - video_rows, video_shapes = self._encode_video_conditions( - prepared_videos, - count=video_count, - ) - rows.append(video_rows) - shapes.extend(video_shapes) + _, rank, _ = _dit_rank_world() + # Keep image and video references in one residency window when both + # appear in a request; otherwise the video branch would reload the VAE. + needs_video_vae = video_count > 0 or (rank == 0 and bool(images)) + video_vae_context = self._component_on_device(self.video_vae) if needs_video_vae else nullcontext() + with video_vae_context: + if images: + image_rows = None + if rank == 0: + image_rows = torch.cat([self.video_vae.encode_image(image) for image in images]) + rows.append( + _broadcast_tensor( + image_rows, + dtype=torch.float32, + device=self.device, + ) + ) + shapes.extend((1, image.height // 16, image.width // 16) for image in images) + if video_count: + video_rows, video_shapes = self._encode_video_conditions_resident( + prepared_videos, + count=video_count, + ) + rows.append(video_rows) + shapes.extend(video_shapes) return (torch.cat(rows) if rows else None), shapes - def _encode_audio_condition( - self, - audio: tuple[torch.Tensor, int], - ) -> tuple[torch.Tensor, int]: - _, rank, _ = _dit_rank_world() - rows = None - audio_t = 0 - if rank == 0: - with self._component_on_device(self.audio_vae): - rows, audio_t = self.audio_vae.encode_waveform(*audio) - audio_t_tensor = torch.tensor( - [audio_t], - dtype=torch.long, - device=self.device, - ) - group, _, world_size = _dit_rank_world() - if world_size > 1: - dist.broadcast(audio_t_tensor, src=0, group=group) - rows = _broadcast_tensor( - rows, - dtype=torch.float32, - device=self.device, - ) - return rows, int(audio_t_tensor.item()) - - def _encode_audio_conditions( + def _encode_audio_conditions_resident( self, audios: list[tuple[torch.Tensor, int]], *, @@ -1257,15 +1674,6 @@ class MiniMaxH3Pipeline( [int(value) for value in lengths.tolist()], ) - def _encode_video_conditions( - self, - prepared_videos: list[dict[str, Any]] | None, - *, - count: int, - ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: - with self._component_on_device(self.video_vae): - return self._encode_video_conditions_resident(prepared_videos, count=count) - def _encode_video_conditions_resident( self, prepared_videos: list[dict[str, Any]] | None, @@ -1313,15 +1721,6 @@ class MiniMaxH3Pipeline( [tuple(int(value) for value in item) for item in shapes.tolist()], ) - def _encode_video_audio_conditions( - self, - prepared_videos: list[dict[str, Any]] | None, - *, - has_audio: list[bool], - ) -> tuple[torch.Tensor | None, list[int]]: - with self._component_on_device(self.audio_vae): - return self._encode_video_audio_conditions_resident(prepared_videos, has_audio=has_audio) - def _encode_video_audio_conditions_resident( self, prepared_videos: list[dict[str, Any]] | None, @@ -1365,6 +1764,34 @@ class MiniMaxH3Pipeline( [int(value) for value in lengths.tolist()], ) + def _encode_reference_audio_conditions( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + has_audio: list[bool], + standalone_audios: list[tuple[torch.Tensor, int]], + max_duration_seconds: float, + ) -> tuple[torch.Tensor | None, list[int], torch.Tensor | None, list[int]]: + # Embedded and standalone audio are consecutive direct Audio-VAE + # calls. Keep the component resident across both paths. + needs_audio_vae = any(has_audio) or bool(standalone_audios) + audio_vae_context = self._component_on_device(self.audio_vae) if needs_audio_vae else nullcontext() + with audio_vae_context: + embedded_condition, embedded_lengths = self._encode_video_audio_conditions_resident( + prepared_videos, + has_audio=has_audio, + ) + external_condition, external_lengths = self._encode_audio_conditions_resident( + standalone_audios, + max_duration_seconds=max_duration_seconds, + ) + return ( + embedded_condition, + embedded_lengths, + external_condition, + external_lengths, + ) + def _initial_noise( self, *, @@ -1408,7 +1835,7 @@ class MiniMaxH3Pipeline( if controller is not None and enabled: controller.offload_resident_layers() - def diffuse( + def _build_denoise_inputs( self, *, task: str, @@ -1432,7 +1859,12 @@ class MiniMaxH3Pipeline( visual_condition_shapes: list[tuple[int, int, int]] | None = None, audio_condition_lengths: list[int] | None = None, keyframe_frame_indices: list[int] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> dict[str, Any]: + """Build the packed layout, initial rows, anchors, and sigma schedules. + + Shared by request-mode :meth:`diffuse` and step-mode + :meth:`prepare_encode` so both paths start from identical state. + """ initial_video, initial_audio = self._initial_noise( seed=seed, latent_t=latent_t, @@ -1532,32 +1964,34 @@ class MiniMaxH3Pipeline( shift_scale=audio_shift, base_schedule=base_schedule, ) - transformer = self._transformer_for_task(task) - # The static DLO plan keeps leading blocks resident only for the - # primary ``transformer``. In combined mode ``transformers_ref`` is - # fully streamed, so a Ref2VA request must not stage the inactive - # FL2VA transformer resident blocks. - with self._resident_dit_layers_on_device(enabled=transformer is self.transformer): - with self.progress_bar(total=len(video_sigmas) - 1) as progress: - video_rows, audio_rows = minimax_h3_denoise_loop( - model=transformer, - positive=branch, - initial_video_rows=initial_video, - initial_audio_rows=initial_audio, - keyframe_cond_rows=visual_anchor, - audio_ref_rows=audio_anchor, - sigmas_video=video_sigmas, - sigmas_audio=audio_sigmas, - device=self.device, - imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), - audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), - on_step_start=lambda step, video_sigma, audio_sigma: self.record_denoise_step( - step, - normalized_timestep=video_sigma, - ), - on_step_end=lambda step, video, audio: progress.update(), - ) + return { + "branch": branch, + # The request-mode loop moves these onto the device itself; step mode + # keeps them resident across steps, so normalize once for both. + "video_rows": initial_video.to(device=self.device, dtype=torch.float32), + "audio_rows": initial_audio.to(device=self.device, dtype=torch.float32), + "cond_anchor": ( + None if visual_anchor is None else visual_anchor.to(device=self.device, dtype=torch.float32) + ), + "audio_anchor": ( + None if audio_anchor is None else audio_anchor.to(device=self.device, dtype=torch.float32) + ), + "sigmas_video": video_sigmas, + "sigmas_audio": audio_sigmas, + } + def _unpack_denoised_rows( + self, + branch: MiniMaxH3DenoiseBranch, + video_rows: torch.Tensor, + audio_rows: torch.Tensor, + *, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Select the target rows and unpack them back into VAE latents.""" target_video = video_rows[branch.update_mask_dev] video_latent = minimax_h3_unpatchify_video_tokens( target_video, @@ -1577,6 +2011,83 @@ class MiniMaxH3Pipeline( ) return video_latent, audio_latent + def diffuse( + self, + *, + task: str, + text_embeddings: torch.Tensor, + text_tags: torch.Tensor, + seed: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + num_frames: int, + num_steps: int, + video_shift: float, + audio_shift: float, + base_schedule: Sequence[float] | None, + visual_condition: torch.Tensor | None, + visual_condition_shape: tuple[int, int, int] | None, + audio_condition: torch.Tensor | None, + ref_audio_t: int | None, + ref_blocks: list[dict[str, Any]] | None = None, + visual_condition_shapes: list[tuple[int, int, int]] | None = None, + audio_condition_lengths: list[int] | None = None, + keyframe_frame_indices: list[int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + inputs = self._build_denoise_inputs( + task=task, + text_embeddings=text_embeddings, + text_tags=text_tags, + seed=seed, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + num_frames=num_frames, + num_steps=num_steps, + video_shift=video_shift, + audio_shift=audio_shift, + base_schedule=base_schedule, + visual_condition=visual_condition, + visual_condition_shape=visual_condition_shape, + audio_condition=audio_condition, + ref_audio_t=ref_audio_t, + ref_blocks=ref_blocks, + visual_condition_shapes=visual_condition_shapes, + audio_condition_lengths=audio_condition_lengths, + keyframe_frame_indices=keyframe_frame_indices, + ) + branch = inputs["branch"] + transformer = self._transformer_for_task(task) + with self._resident_dit_layers_on_device(enabled=transformer is self.transformer): + with self.progress_bar(total=len(inputs["sigmas_video"]) - 1) as progress: + video_rows, audio_rows = minimax_h3_denoise_loop( + model=transformer, + positive=branch, + initial_video_rows=inputs["video_rows"], + initial_audio_rows=inputs["audio_rows"], + keyframe_cond_rows=inputs["cond_anchor"], + audio_ref_rows=inputs["audio_anchor"], + sigmas_video=inputs["sigmas_video"], + sigmas_audio=inputs["sigmas_audio"], + device=self.device, + imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), + audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + on_step=lambda step, video, audio: progress.update(), + ) + + return self._unpack_denoised_rows( + branch, + video_rows, + audio_rows, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + ) + def decode( self, video_latent: torch.Tensor, @@ -1597,11 +2108,9 @@ class MiniMaxH3Pipeline( audio = self.audio_vae.decode_latent(audio_latent) return video, audio - @torch.no_grad() - def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: - if len(request.prompts) != 1: - raise OmniClientError("MiniMax H3 supports one request at a time") - raw_prompt = request.prompts[0] + @staticmethod + def _extract_prompt(raw_prompt: Any) -> tuple[str, dict[str, Any]]: + """Split a request prompt into its text and multimodal parts.""" if isinstance(raw_prompt, str): prompt = raw_prompt multi_modal_data: dict[str, Any] = {} @@ -1610,12 +2119,76 @@ class MiniMaxH3Pipeline( multi_modal_data = raw_prompt.get("multi_modal_data") or {} if not prompt: raise OmniClientError("MiniMax H3 requires a non-empty prompt") + return prompt, multi_modal_data - sampling = request.sampling_params + @staticmethod + def _extract_text_conditioning(raw_prompt: Any) -> MiniMaxH3TextConditioning | None: + if isinstance(raw_prompt, str): + return None + additional_information = raw_prompt.get("additional_information") or {} + text_encoder_output = additional_information.get("text_encoder_output") + if text_encoder_output is None: + return None + if not isinstance(text_encoder_output, Mapping): + raise OmniClientError("text_encoder_output must be a mapping") + try: + return MiniMaxH3TextConditioning.from_payload(text_encoder_output) + except ValueError as exc: + raise OmniClientError(str(exc)) from exc + + @staticmethod + def _extract_prepared_reference_videos(raw_prompt: Any) -> list[dict[str, Any]] | None: + if isinstance(raw_prompt, str): + return None + additional_information = raw_prompt.get("additional_information") or {} + meta = additional_information.get("meta") or {} + descriptor = meta.get(MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY) + if descriptor is None: + return None + if not isinstance(descriptor, str): + raise OmniClientError("MiniMax H3 prepared-reference-video descriptor must be a string") + try: + _, videos = deserialize_prepared_reference_videos(descriptor) + except ValueError as exc: + raise OmniClientError(str(exc)) from exc + return videos + + def _prepare_request_inputs( + self, + *, + prompt: str, + multi_modal_data: dict[str, Any], + sampling: Any, + text_conditioning: MiniMaxH3TextConditioning | None = None, + prepared_reference_videos: list[dict[str, Any]] | None = None, + ) -> dict[str, Any]: + """Resolve the task and output shape, then run every request-level encode. + + Shared by request-mode :meth:`forward` and step-mode + :meth:`prepare_encode`; the returned mapping feeds :meth:`diffuse` and + :meth:`_build_denoise_inputs` unchanged. + """ quality = sampling.quality logger.debug("MiniMax H3 request quality=%s", quality) extra = sampling.extra_args or {} - task = self._resolve_task(extra.get("task"), multi_modal_data) + has_turbo_lora = self._has_active_turbo_lora(sampling) + has_native_lora = self._has_active_native_lora(sampling) + task = self._resolve_task( + extra.get("task"), + multi_modal_data, + has_turbo_lora=has_turbo_lora, + has_native_lora=has_native_lora, + ) + if has_turbo_lora: + self._validate_turbo_sampling(sampling) + if has_native_lora: + self._validate_native_sampling(sampling, task=task) + if self._fasth3 is not None: + self._fasth3.check_request( + sampling, + video_shift=self.default_video_shift, + audio_shift=self.default_audio_shift, + ) raw_image = multi_modal_data.get("image") raw_videos = multi_modal_data.get("video") @@ -1681,12 +2254,32 @@ class MiniMaxH3Pipeline( video_count = 0 if raw_videos is not None: video_count = len(raw_videos) if isinstance(raw_videos, (list, tuple)) else 1 - prepared_videos = self._prepare_reference_videos( - raw_videos, - target_frame_count=num_frames, - workdir=workdir, - start_time_seconds=extra.get("start_time_seconds"), - ) + # File-based reference-video prep runs only on rank 0; other + # ranks return None without touching disk. If rank 0 raises + # (e.g. invalid file, unsupported codec) it must not exit + # ``prepare_encode`` before the downstream broadcasts below -- + # non-zero ranks would then deadlock on them forever. Capture + # the exception here and let every rank agree on the outcome + # before starting any subsequent collective. + prep_error: Exception | None = None + try: + _, rank, _ = _dit_rank_world() + if prepared_reference_videos is not None: + if rank == 0: + prepared_videos = _reuse_prepared_reference_videos( + prepared_reference_videos, + expected_count=video_count, + ) + else: + prepared_videos = self._prepare_reference_videos( + raw_videos, + target_frame_count=num_frames, + workdir=workdir, + start_time_seconds=extra.get("start_time_seconds"), + ) + except Exception as exc: + prep_error = exc + _broadcast_rank0_exception(prep_error) has_audio_tensor = torch.zeros( video_count, dtype=torch.long, @@ -1703,7 +2296,7 @@ class MiniMaxH3Pipeline( dist.broadcast( has_audio_tensor, src=0, - group=get_dit_group(), + group=get_world_group().device_group, ) has_audio = [bool(value) for value in has_audio_tensor.tolist()] @@ -1724,13 +2317,27 @@ class MiniMaxH3Pipeline( audio_index += 1 condition_labels.append(("audio", audio_index)) - text_embeddings, text_tags = self.encode_prompt( - task=task, - prompt=prompt, - images=prepared_images, - prepared_videos=prepared_videos, - condition_labels=condition_labels if task == "ref2va" else None, - ) + if text_conditioning is not None: + text_embeddings = text_conditioning.hidden_states.to( + device=self.device, + dtype=torch.bfloat16, + ) + text_tags = text_conditioning.token_tags.to( + device=self.device, + dtype=torch.long, + ) + elif getattr(self, "text_encoder", None) is not None: + text_embeddings, text_tags = self.encode_prompt( + task=task, + prompt=prompt, + images=prepared_images, + prepared_videos=prepared_videos, + condition_labels=condition_labels if task == "ref2va" else None, + ) + else: + raise OmniClientError( + "MiniMax H3 diffusion stage requires text_encoder_output when text_encoder is not loaded" + ) # ``prepared_videos`` is intentionally ``None`` on non-zero DiT # ranks; the distributed video encoder broadcasts the prepared @@ -1743,12 +2350,15 @@ class MiniMaxH3Pipeline( prepared_videos, video_count=video_count, ) - embedded_audio_condition, embedded_audio_lengths = self._encode_video_audio_conditions( + ( + embedded_audio_condition, + embedded_audio_lengths, + external_audio_condition, + external_audio_lengths, + ) = self._encode_reference_audio_conditions( prepared_videos, has_audio=has_audio, - ) - external_audio_condition, external_audio_lengths = self._encode_audio_conditions( - standalone_audios, + standalone_audios=standalone_audios, max_duration_seconds=float(num_frames) / float(sampling.fps or MINIMAX_H3_FPS), ) audio_parts = [ @@ -1795,22 +2405,7 @@ class MiniMaxH3Pipeline( ref_audio_t = audio_lengths[0] seed = int(sampling.seed if sampling.seed is not None else 42) - sigma_schedule = self._base_schedule_for_task(task) - if sigma_schedule is None: - base_schedule = None - num_steps = int(sampling.num_inference_steps or 50) - else: - # The schedule lists sigma boundaries; the denoise loop runs one - # step per interval, and that count is what requests and Cache-DiT - # speak in. - base_schedule = sigma_schedule.base_schedule - num_steps = sigma_schedule.num_inference_steps - requested_steps = sampling.num_inference_steps - if requested_steps is not None and int(requested_steps) != num_steps: - raise OmniClientError( - "this MiniMax H3 checkpoint pins a distilled sigma schedule; num_inference_steps " - f"must be {num_steps} or omitted, got {int(requested_steps)}" - ) + base_schedule, num_steps = self._resolve_sigma_positions(task, sampling) video_shift = float(extra.get("flow_shift", self.default_video_shift)) audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift)) quality_plan = self._quality_policy.resolve( @@ -1820,37 +2415,369 @@ class MiniMaxH3Pipeline( ) self._cache_dit_runtime.prepare(quality_plan.cache_dit) num_outputs = _resolve_minimax_h3_num_outputs(sampling.num_outputs_per_prompt) + return { + "task": task, + "height": height, + "width": width, + "num_frames": num_frames, + "latent_t": latent_t, + "latent_h": height // 16, + "latent_w": width // 16, + "audio_t": audio_t, + "text_embeddings": text_embeddings, + "text_tags": text_tags, + "visual_condition": visual_condition, + "visual_condition_shape": visual_shape, + "audio_condition": audio_condition, + "ref_audio_t": ref_audio_t, + "ref_blocks": ref_blocks, + "visual_condition_shapes": visual_shapes, + "audio_condition_lengths": audio_lengths, + "keyframe_frame_indices": keyframe_frame_indices, + "seed": seed, + "num_steps": num_steps, + "video_shift": video_shift, + "audio_shift": audio_shift, + "base_schedule": base_schedule, + "num_outputs": num_outputs, + } + + @staticmethod + def _denoise_kwargs(context: dict[str, Any]) -> dict[str, Any]: + """Select the denoise-input arguments from a prepared request context.""" + return {key: context[key] for key in _MINIMAX_H3_DENOISE_INPUT_KEYS} + + @torch.no_grad() + def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: + if len(request.prompts) != 1: + raise OmniClientError("MiniMax H3 supports one request at a time") + raw_prompt = request.prompts[0] + prompt, multi_modal_data = self._extract_prompt(raw_prompt) + context = self._prepare_request_inputs( + prompt=prompt, + multi_modal_data=multi_modal_data, + sampling=request.sampling_params, + text_conditioning=self._extract_text_conditioning(raw_prompt), + prepared_reference_videos=self._extract_prepared_reference_videos(raw_prompt), + ) + denoise_kwargs = self._denoise_kwargs(context) + num_outputs = context["num_outputs"] videos = [] audios = [] - for output_seed in _minimax_h3_output_seeds(seed, num_outputs): - video_latent, audio_latent = self.diffuse( - task=task, - text_embeddings=text_embeddings, - text_tags=text_tags, - seed=output_seed, - latent_t=latent_t, - latent_h=height // 16, - latent_w=width // 16, - audio_t=audio_t, - num_frames=num_frames, - num_steps=num_steps, - video_shift=video_shift, - audio_shift=audio_shift, - base_schedule=base_schedule, - visual_condition=visual_condition, - visual_condition_shape=visual_shape, - audio_condition=audio_condition, - ref_audio_t=ref_audio_t, - ref_blocks=ref_blocks, - visual_condition_shapes=visual_shapes, - audio_condition_lengths=audio_lengths, - keyframe_frame_indices=keyframe_frame_indices, + for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): + video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed}) + video, audio = self.decode( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], ) - video, audio = self.decode(video_latent, audio_latent, height=height, width=width) - videos.append(video) + videos.append(_prepare_minimax_h3_video_output(video)) audios.append(audio) - video = torch.cat(videos, dim=0) - audio = torch.cat(audios, dim=0) + video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) + audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) + return DiffusionOutput( + output=(video, audio), + post_process_func=get_minimax_h3_post_process_func(self.od_config), + stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}), + ) + + # ------------------------------------------------------------------ + # Step-wise execution (continuous batching) + # ------------------------------------------------------------------ + + @staticmethod + def _packed_batch_supported(transformer: MiniMaxH3DiTModel) -> bool: + """Whether every attention in this DiT honors multi-document cu_seqlens. + + A packed batch is only isolated if *all* of them do: the token refiner + runs under its own attention role and can resolve to a different backend + from the DiT blocks. Ring sequence parallelism dispatches through + ``RingParallelAttention``, whose kernels ignore the packed + ``cu_seqlens`` metadata regardless of the configured backend; packing + multiple requests under ring would let attention cross document + boundaries, so any layer running ring disqualifies the batch. + + The gate probes a per-backend capability rather than a fixed backend + name: FLASH_ATTN, for example, only isolates arbitrary N-document + packed cu_seqlens on CUDA/ROCm/MUSA. Its NPU path only accepts a + ``[real, pad]`` two-document layout and its XPU path ignores + cu_seqlens outright — either would silently attend across request + boundaries. + """ + attentions = [module for module in transformer.modules() if isinstance(module, MiniMaxH3Attention)] + if not attentions: + return False + return all(_attention_isolates_packed_requests(module.attention) for module in attentions) + + def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState: + """Run every request-level stage once and seed the per-request step state.""" + del kwargs + # Two request-mode features have no place in the shared step contract: + # a request state carries exactly one latent tensor, and distributed + # layerwise offload streams the DiT around one whole denoise loop rather + # than around a single scheduler-driven step. + num_outputs = _resolve_minimax_h3_num_outputs(state.sampling.num_outputs_per_prompt) + if num_outputs != 1: + raise OmniClientError( + f"MiniMax H3 step execution produces one output per request, got num_outputs_per_prompt={num_outputs}" + ) + if getattr(self, "_dlo_residency_controller", None) is not None: + raise ValueError( + "MiniMax H3 step execution is not compatible with distributed layerwise offload; " + "the resident-layer window spans a whole denoise loop, so per-step streaming would " + "reload the DiT every step. Drop --step-execution or --enable-distributed-layerwise-offload." + ) + # Request-scoped Cache-DiT (quality=high) mutates hook state on the + # shared transformer rather than on ``StepRequestState``. In step mode + # two requests can interleave denoise steps, or be co-batched into a + # single forward, and the second one would then re-enter the DiT with + # cache buffers shaped for the first. Reject the profile here rather + # than let it corrupt outputs at runtime; startup-configured Cache-DiT + # is already blocked in ``DiffusionModelRunner.execute_stepwise``. + if getattr(state.sampling, "quality", None) == "high": + raise OmniClientError( + "MiniMax H3 step execution does not support the high-quality Cache-DiT profile " + "(quality=high); its hooks live on the shared transformer, so interleaved or " + "co-batched requests would reuse incompatible cache state. Drop --step-execution " + "or omit quality=high." + ) + prompt, multi_modal_data = self._extract_prompt(state.prompt) + context = self._prepare_request_inputs( + prompt=prompt, + multi_modal_data=multi_modal_data, + sampling=state.sampling, + text_conditioning=self._extract_text_conditioning(state.prompt), + prepared_reference_videos=self._extract_prepared_reference_videos(state.prompt), + ) + inputs = self._build_denoise_inputs(**self._denoise_kwargs(context)) + + sigmas_video = inputs["sigmas_video"] + sigmas_audio = inputs["sigmas_audio"] + if len(sigmas_video) < 2: + raise OmniClientError( + f"MiniMax H3 step execution needs at least one denoise step, got num_inference_steps=" + f"{len(sigmas_video) - 1}" + ) + + branch = inputs["branch"] + video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows( + positive=branch, + initial_video_rows=inputs["video_rows"], + initial_audio_rows=inputs["audio_rows"], + keyframe_cond_rows=inputs["cond_anchor"], + audio_ref_rows=inputs["audio_anchor"], + device=self.device, + ) + + # The denoise loop consumes sigma pairs, so the schedule carries one more + # point than there are steps. ``timesteps`` holds the video branch because + # the shared contract gives a request exactly one timestep sequence; the + # audio schedule rides along in ``extra``. + state.timesteps = torch.tensor( + [1.0 - sigma for sigma in sigmas_video[:-1]], + dtype=torch.float32, + device=self.device, + ) + state.step_index = 0 + # Video rows are the batched tensor the runner slices per request; audio + # rows have a different width, so they stay request-private. + state.latents = video_rows + state.do_true_cfg = False # H3 checkpoints are CFG-distilled. + state.extra.update( + { + _STEP_BRANCH: branch, + _STEP_TRANSFORMER: self._transformer_for_task(context["task"]), + _STEP_AUDIO_ROWS: audio_rows, + _STEP_COND_ANCHOR: cond_anchor, + _STEP_AUDIO_ANCHOR: audio_anchor, + _STEP_SIGMAS_VIDEO: sigmas_video, + _STEP_SIGMAS_AUDIO: sigmas_audio, + _STEP_SHAPE: { + "height": context["height"], + "width": context["width"], + "latent_t": context["latent_t"], + "latent_h": context["latent_h"], + "latent_w": context["latent_w"], + "audio_t": context["audio_t"], + }, + } + ) + return state + + def denoise_step( + self, + input_batch: InputBatch, + *, + states: Sequence[StepRequestState] | None = None, + **kwargs: Any, + ) -> torch.Tensor | None: + """Run one denoise forward covering every request in the batch. + + Requests are concatenated into a single packed sequence that keeps one + attention document each, so the whole batch costs one DiT forward. + Backends that ignore ``cu_seqlens`` cannot express that isolation, so + they fall back to one forward per request. + """ + del kwargs + batch_states = list(states if states is not None else input_batch.states) + + branches = [state.extra[_STEP_BRANCH] for state in batch_states] + video_rows = [state.latents for state in batch_states] + audio_rows = [state.extra[_STEP_AUDIO_ROWS] for state in batch_states] + schedules = [_minimax_h3_step_schedule(state) for state in batch_states] + transformers = [state.extra[_STEP_TRANSFORMER] for state in batch_states] + mixed_transformers = len({id(transformer) for transformer in transformers}) > 1 + + # Both execution modes must publish denoise progress for step-gated + # attention features. Requests can differ in both step index and sigma + # schedule, so a batch that is not at one single point has nothing to + # publish and those gates stay dense -- which is their safe default. + progress = { + (state.step_index, schedule["sigma_video"], len(state.extra[_STEP_SIGMAS_VIDEO]) - 1) + for state, schedule in zip(batch_states, schedules) + } + minimax_h3_publish_denoise_progress(*(progress.pop() if len(progress) == 1 else (None, None, None))) + + if len(batch_states) > 1 and (mixed_transformers or not self._packed_batch_supported(transformers[0])): + if mixed_transformers: + logger.warning_once( + "MiniMax H3 step batch contains requests for different task-specific DiTs; " + "running %d requests one forward at a time.", + len(batch_states), + ) + elif any( + getattr(getattr(module, "attention", None), "use_ring", False) + for module in transformers[0].modules() + if isinstance(module, MiniMaxH3Attention) + ): + logger.warning_once( + "MiniMax H3 step batching is disabled when ring attention is active: " + "the ring kernels ignore packed cu_seqlens and would attend across request " + "boundaries. Running %d requests one forward at a time.", + len(batch_states), + ) + else: + logger.warning_once( + "MiniMax H3 step batching needs every attention on a backend that isolates " + "packed multi-document cu_seqlens (see AttentionBackend." + "supports_multi_doc_packed_varlen); running %d requests one forward at a time.", + len(batch_states), + ) + video_parts: list[torch.Tensor] = [] + audio_parts: list[torch.Tensor] = [] + for index, branch in enumerate(branches): + forward_kwargs = branch.forward_kwargs( + video_rows=video_rows[index], + audio_rows=audio_rows[index], + t_video=schedules[index]["t_video"], + t_audio=schedules[index]["t_audio"], + imgvid_cond_timestep=schedules[index]["imgvid_cond_timestep"], + audio_ref_cond_timestep=schedules[index]["audio_ref_cond_timestep"], + ) + request_video, request_audio = transformers[index](**forward_kwargs) + video_parts.append(request_video) + audio_parts.append(request_audio) + video_velocity = torch.cat(video_parts) + audio_velocity = torch.cat(audio_parts) + else: + forward_kwargs = minimax_h3_batched_forward_kwargs( + branches=branches, + video_rows=video_rows, + audio_rows=audio_rows, + t_video=[schedule["t_video"] for schedule in schedules], + t_audio=[schedule["t_audio"] for schedule in schedules], + imgvid_cond_timesteps=[schedule["imgvid_cond_timestep"] for schedule in schedules], + audio_ref_cond_timesteps=[schedule["audio_ref_cond_timestep"] for schedule in schedules], + ) + logger.debug( + "MiniMax H3 denoise step: %d request(s) packed into %d rows", + len(batch_states), + int(forward_kwargs["x"].shape[1]), + ) + video_velocity, audio_velocity = transformers[0](**forward_kwargs) + + # The shared contract carries one velocity tensor per step, and audio rows + # are a different width than video rows, so hand the audio branch to + # step_scheduler() through request-private state. + audio_parts_by_request = torch.split(audio_velocity, [int(branch.audio_pos.shape[0]) for branch in branches]) + for state, request_audio in zip(batch_states, audio_parts_by_request, strict=True): + state.extra[_STEP_AUDIO_NOISE_PRED] = request_audio + return video_velocity + + def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: + """Apply one Euler-eta0 update to this request's video and audio rows.""" + del kwargs + # denoise_step() stages the audio half of this step's velocity; popping + # it keeps a second step_scheduler() call from reusing a stale one. + audio_noise_pred = state.extra.pop(_STEP_AUDIO_NOISE_PRED) + + branch = state.extra[_STEP_BRANCH] + schedule = _minimax_h3_step_schedule(state) + update = branch.update_mask_dev + audio_update = branch.audio_update_mask_dev + video_rows = state.latents + audio_rows = state.extra[_STEP_AUDIO_ROWS] + cond_anchor = state.extra[_STEP_COND_ANCHOR] + audio_anchor = state.extra[_STEP_AUDIO_ANCHOR] + device = video_rows.device + + x0_video = minimax_h3_rf_v_to_x0( + video_rows[update], + noise_pred.float()[update], + torch.tensor(schedule["t_video"], dtype=torch.float32, device=device), + ) + new_video = minimax_h3_euler_eta0_step( + video_rows[update], + x0_video, + sigma_curr=schedule["sigma_video"], + sigma_next=schedule["sigma_video_next"], + ) + video_rows = video_rows.clone() + video_rows[update] = new_video + if cond_anchor is not None: + video_rows[~update] = cond_anchor # per-step imgvid cond reset + + x0_audio = minimax_h3_rf_v_to_x0( + audio_rows[audio_update], + audio_noise_pred.float()[audio_update], + torch.tensor(schedule["t_audio"], dtype=torch.float32, device=device), + ) + new_audio = minimax_h3_euler_eta0_step( + audio_rows[audio_update], + x0_audio, + sigma_curr=schedule["sigma_audio"], + sigma_next=schedule["sigma_audio_next"], + ) + audio_rows = audio_rows.clone() + audio_rows[audio_update] = new_audio + if audio_anchor is not None: + audio_rows[~audio_update] = audio_anchor # per-step audio ref reset + + state.latents = video_rows + state.extra[_STEP_AUDIO_ROWS] = audio_rows + state.step_index += 1 + + def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput: + """Unpack the denoised rows and run the joint video/audio VAE decode.""" + del kwargs + shape = state.extra[_STEP_SHAPE] + video_latent, audio_latent = self._unpack_denoised_rows( + state.extra[_STEP_BRANCH], + state.latents, + state.extra[_STEP_AUDIO_ROWS], + latent_t=shape["latent_t"], + latent_h=shape["latent_h"], + latent_w=shape["latent_w"], + audio_t=shape["audio_t"], + ) + video, audio = self.decode( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + ) + video = _prepare_minimax_h3_video_output(video) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index e28450c90..2e830fa04 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 remote-code VAE adapters and exact latent contracts.""" from __future__ import annotations import importlib import json -from contextlib import AbstractContextManager +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext from pathlib import Path from typing import Any @@ -14,13 +16,18 @@ import torch.distributed as dist import torch.nn as nn from PIL import Image from transformers.dynamic_module_utils import get_class_from_dynamic_module +from vllm.logger import init_logger from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedVaeMixin, ) -from vllm_omni.diffusion.distributed.parallel_state import get_dit_group -from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager +from vllm_omni.diffusion.distributed.parallel_state import get_world_group +from vllm_omni.diffusion.offloader.module_residency import ( + BoundedAllocatorCache, + PinnedModuleStager, +) +from .ops import install_h3_vae_optimizations from .packed_tokens import minimax_h3_patchify_video_latent MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 @@ -28,6 +35,9 @@ MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 MINIMAX_H3_AUDIO_CHANNELS = 2 +logger = init_logger(__name__) + + def _load_component_config(component_path: str) -> dict[str, Any]: config_path = Path(component_path) / "config.json" config = json.loads(config_path.read_text(encoding="utf-8")) @@ -119,11 +129,18 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): component_path, self.config_dict, ) - # Match the reference loader contract: video VAE weights stay FP32. - # Keyframe encoding is numerically sensitive to first casting the - # checkpoint through FP16; decode still runs under FP16 autocast. + # Match the reference loader contract before installing inference-only + # decoder fast paths. Keyframe encoding remains FP32; decoder Linear + # weights may be materialized in FP16 because reference decode casts + # those same tensors through CUDA autocast on every tile. initial_device = load_device or device self.remote.eval().to(device=initial_device, dtype=torch.float32) + decoder = getattr(self.remote.model, "decoder", None) + if decoder is not None: + install_h3_vae_optimizations( + decoder, + device=device, + ) self._stager = None if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): self._stager = PinnedModuleStager( @@ -143,12 +160,21 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): else: self.remote.to(self._device_target) + def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: + self._omni_component_cache = cache + if self._stager is not None: + self._stager.set_cache_retention(cache) + def offload_to_cpu(self) -> None: if self._stager is not None: self._stager.offload() else: self.remote.to("cpu") - torch.accelerator.empty_cache() + cache = getattr(self, "_omni_component_cache", None) + if cache is None: + torch.accelerator.empty_cache() + else: + cache.release_if_needed() def set_parallel_size( self, @@ -157,7 +183,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): ) -> None: if mode != "tile": raise ValueError(f"MiniMax H3 VAE supports its native tile parallel mode only, got {mode!r}") - group = get_dit_group() + group = get_world_group().device_group world_size = dist.get_world_size(group) rank = dist.get_rank(group) parallel_size = int(parallel_size) @@ -170,9 +196,7 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self.parallel_size = parallel_size enabled = parallel_size > 1 - package = self.remote.__class__.__module__.rsplit(".", 1)[0] - parallel_module = importlib.import_module(f"{package}.parallel") - state = parallel_module.get_parallel_state() + state = self._native_parallel_state() state.clear() state.update( group_size=parallel_size, @@ -187,6 +211,57 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): ) self.model.parallel_tiling = enabled + def _native_parallel_state(self) -> dict[str, Any]: + """Return the checkpoint's own mutable parallel-state dict.""" + + package = self.remote.__class__.__module__.rsplit(".", 1)[0] + parallel_module = importlib.import_module(f"{package}.parallel") + return parallel_module.get_parallel_state() + + def _decoder_tile_count(self, latent: torch.Tensor) -> int: + """Number of decoder tiles the checkpoint will split ``latent`` into. + + Mirrors the checkpoint's ``decode_tiled``: the grid is computed from the + pixel-space dimensions, so it is a pure function of the latent shape and + resolves identically on every rank. + """ + + ratio = int(self.model.vae_ratio) + rows, _, _ = self.model.split_tiles(int(latent.shape[-2]) * ratio, True) + cols, _, _ = self.model.split_tiles(int(latent.shape[-1]) * ratio, True) + return len(rows) * len(cols) + + @contextmanager + def _rank_local_tiling(self) -> Iterator[None]: + """Run one decode with tiling kept on this rank, then restore the group. + + Used only when there are fewer tiles than ranks. Every rank then decodes + every tile, which is slower than sharing the work but is correct and + involves no collective. + """ + + state = self._native_parallel_state() + saved_state = dict(state) + saved_tiling = self.model.parallel_tiling + state.update( + group_size=1, + group_rank=0, + local_process_group=None, + sp_size=1, + sp_rank=0, + sp_enabled=False, + sp_process_group=None, + tp_size=1, + tp_rank=0, + ) + self.model.parallel_tiling = False + try: + yield + finally: + state.clear() + state.update(saved_state) + self.model.parallel_tiling = saved_tiling + def is_distributed_enabled(self) -> bool: return self.parallel_size > 1 and dist.is_initialized() @@ -293,7 +368,27 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): device=latent.device, dtype=latent.dtype, ).view(1, channels, 1, 1, 1) - decoded = self.model.decode_base(latent * std + mean) + # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)`` + # and then rejects an empty share inside the gather. A rank with no + # tiles raises and leaves the collective while the others block in it + # forever, so too few tiles hangs the whole stage rather than failing + # it. Tile count depends only on the latent shape, so every rank takes + # this branch together. + num_tiles = self._decoder_tile_count(latent) + if self.parallel_size > 1 and num_tiles < self.parallel_size: + logger.warning_once( + "MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has " + "%d ranks; decoding rank-locally for this shape instead, which is " + "slower but avoids ranks without tiles hanging the collective.", + num_tiles, + self.parallel_size, + ) + tiling_context: AbstractContextManager = self._rank_local_tiling() + else: + tiling_context = nullcontext() + + with tiling_context: + decoded = self.model.decode_base(latent * std + mean) frames = self.model.processor.revert_tensor(decoded) if frames.ndim == 4: frames = frames.unsqueeze(0).transpose(1, 2) @@ -337,12 +432,21 @@ class MiniMaxH3AudioVAE(nn.Module): else: self.remote.to(self._device_target) + def set_omni_component_cache(self, cache: BoundedAllocatorCache | None) -> None: + self._omni_component_cache = cache + if self._stager is not None: + self._stager.set_cache_retention(cache) + def offload_to_cpu(self) -> None: if self._stager is not None: self._stager.offload() else: self.remote.to("cpu") - torch.accelerator.empty_cache() + cache = getattr(self, "_omni_component_cache", None) + if cache is None: + torch.accelerator.empty_cache() + else: + cache.release_if_needed() @torch.inference_mode() def encode_waveform( diff --git a/vllm_omni/diffusion/models/pi0/__init__.py b/vllm_omni/diffusion/models/pi0/__init__.py new file mode 100644 index 000000000..3862d1836 --- /dev/null +++ b/vllm_omni/diffusion/models/pi0/__init__.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""π0 (Pi-Zero) VLA model for vllm-omni. + +PaliGemma (SigLIP vision + Gemma-2B LM) + Gemma-300M action expert + a +flow-matching action head. Outputs a continuous action chunk +``[horizon, action_dim]`` rather than tokens. +""" + +from vllm_omni.diffusion.models.pi0.config import Pi0Config +from vllm_omni.diffusion.models.pi0.modeling_pi0 import ( + GemmaVariantConfig, + PaliGemmaWithActionExpert, + Pi0ForActionPrediction, + create_sinusoidal_pos_embedding, + get_gemma_config, + make_att_2d_masks, + prepare_attention_masks_4d, +) +from vllm_omni.diffusion.models.pi0.pipeline_pi0 import ( + Pi0Pipeline, + get_pi0_post_process_func, +) +from vllm_omni.diffusion.models.pi0.processor_pi0 import ( + Pi0ImageProcessor, + build_model_inputs, + pil_image_to_tensor, + resize_with_pad, + tokenize_prompt, +) + +__all__ = [ + "Pi0Config", + "Pi0ForActionPrediction", + "PaliGemmaWithActionExpert", + "GemmaVariantConfig", + "get_gemma_config", + "create_sinusoidal_pos_embedding", + "make_att_2d_masks", + "prepare_attention_masks_4d", + "Pi0ImageProcessor", + "build_model_inputs", + "pil_image_to_tensor", + "resize_with_pad", + "tokenize_prompt", + "Pi0Pipeline", + "get_pi0_post_process_func", +] diff --git a/vllm_omni/diffusion/models/pi0/config.py b/vllm_omni/diffusion/models/pi0/config.py new file mode 100644 index 000000000..ef0484f8e --- /dev/null +++ b/vllm_omni/diffusion/models/pi0/config.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Config surface for the π0 (Pi-Zero) VLA model in vllm-omni. + +Only the parameters that actually shape runtime behaviour live here; the +transformer dimensions (hidden size, head count, ...) are derived from +``paligemma_variant`` / ``action_expert_variant`` inside the model via +``get_gemma_config`` (see ``modeling_pi0``). + +The resolver reads the **raw LeRobot ``config.json``** (the field surface of +``lerobot.policies.pi0.PI0Config``): only the recognized dataclass fields are +used; any other keys in the file are ignored. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from dataclasses import fields as dataclass_fields +from pathlib import Path +from typing import Any + +# LeRobot / OpenPI observation key conventions. +ACTION = "action" +OBS_STR = "observation" +OBS_STATE = OBS_STR + ".state" +OBS_IMAGES = OBS_STR + ".images" + + +@dataclass +class Pi0Config: + """π0 VLA config (dataclass, not an HF ``PretrainedConfig``). + + Mirrors the runtime-relevant subset of LeRobot ``PI0Config`` plus a couple + of serving-side knobs (``max_cameras``, ``image_feature_keys``, + ``image_key_map``, ``norm_stats``). + """ + + # Backbone variants — mapped to Gemma dimensions by ``get_gemma_config``. + paligemma_variant: str = "gemma_2b" + action_expert_variant: str = "gemma_300m" + + # Action chunk shape. + chunk_size: int = 50 + max_action_dim: int = 32 + max_state_dim: int = 32 + + # Flow-matching denoising schedule. + num_inference_steps: int = 10 + # Sinusoidal timestep embedding periods (must match OpenPI/LeRobot). + min_period: float = 4e-3 + max_period: float = 4.0 + + # Image preprocessing. π0/SigLIP only support square inputs. + image_resolution: tuple[int, int] = (224, 224) + tokenizer_max_length: int = 48 + # Number of camera slots the model attends to (real + padded). + max_cameras: int = 3 + + # Weight dtype the checkpoint was saved in; parity passes this through to + # match LeRobot's ``PI0Config.dtype``. + dtype: str = "float32" + + # Per-dataset normalization stats, optional (schema matches LeRobot's + # ``NormalizerProcessorStep``; consumed by ``_build_norm_buffers``). ``None`` + # means identity / pass-through (the ``lerobot/pi0_base`` case). + norm_stats: dict | None = None + + # Ordered list of image feature keys (e.g. ``observation.images.base_0_rgb``) + # taken from the checkpoint's ``input_features``. This is the **camera order** + # that must be reproduced for LeRobot parity. + image_feature_keys: list[str] | None = None + + # Optional map from raw OpenPI obs keys → ``image_feature_keys`` entries, so + # serving can translate wire obs (e.g. ``observation/exterior_image_0_left``) + # into the checkpoint's camera identities. Empty → identity match on key. + image_key_map: dict[str, str] = field(default_factory=dict) + + # Stored for reference / camera-order derivation; not used at inference. + input_features: dict[str, Any] = field(default_factory=dict) + output_features: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + # Coerce list → tuple (JSON has no tuples) and validate squareness. + res = self.image_resolution + if not isinstance(res, (tuple, list)) or len(res) != 2 or res[0] != res[1]: + raise ValueError(f"π0 expects a square image_resolution (H == W); got {res!r}.") + self.image_resolution = (int(res[0]), int(res[1])) + + # Derive the camera order from input_features if not given explicitly. + if self.image_feature_keys is None and self.input_features: + self.image_feature_keys = [key for key in self.input_features if key.startswith(OBS_IMAGES + ".")] + + # ------------------------------------------------------------------ + # Constructors + # ------------------------------------------------------------------ + @classmethod + def from_pretrained(cls, checkpoint_dir: str | Path) -> Pi0Config: + """Build from a checkpoint directory's ``config.json``.""" + checkpoint_dir = Path(checkpoint_dir) + config_path = checkpoint_dir / "config.json" + if not config_path.exists(): + return cls() + with open(config_path, encoding="utf-8") as f: + raw = json.load(f) + return cls.from_model_config(raw) + + @classmethod + def from_model_config(cls, model_config: dict[str, Any] | None) -> Pi0Config: + """Build from a config dict (LeRobot ``config.json`` or deploy yaml). + + Keeps only the recognized dataclass fields (the LeRobot config also + carries many training-only keys) and coerces ``image_resolution`` to a + tuple. + """ + if not model_config: + return cls() + + raw = dict(model_config) + if "image_resolution" in raw: + raw["image_resolution"] = tuple(raw["image_resolution"]) + + allowed = {item.name for item in dataclass_fields(cls)} + filtered = {key: value for key, value in raw.items() if key in allowed} + return cls(**filtered) diff --git a/vllm_omni/diffusion/models/pi0/modeling_pi0.py b/vllm_omni/diffusion/models/pi0/modeling_pi0.py new file mode 100644 index 000000000..a45d931ef --- /dev/null +++ b/vllm_omni/diffusion/models/pi0/modeling_pi0.py @@ -0,0 +1,944 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only π0 (Pi-Zero) VLA math kernel for vllm-omni. + +A self-contained inference kernel: only the math that turns a robot observation +into an action chunk, with no serving/request glue — validated bit-for-bit +against the LeRobot ``PI0Policy`` reference (``max|Δ| < 1e-4``). + +π0 = PaliGemma (SigLIP vision + Gemma 2B LM) prefix + Gemma 300M action expert +suffix + flow-matching head. Inference: + +1. Embed prefix (images + language) → bidirectional prefix tokens. +2. Forward prefix through PaliGemma → a per-layer ``list[(k, v)]`` KV cache. +3. For each Euler step ``t = 1.0, 1-dt, ..., 0``: embed the suffix + (state + noisy actions + timestep), run the action expert with cross + attention over the cached prefix K/V, predict velocity ``v_t``, and + integrate ``x_t = x_t + dt * v_t``. +4. Return ``x_0`` as the action chunk ``(batch, action_horizon, action_dim)``. + +We walk Gemma decoder layers ourselves in ``_compute_layer_*`` so we never go +through ``GemmaModel.forward`` (which has version-dependent mask/KV-cache +behaviour). Targets modern transformers (PaliGemmaForConditionalGeneration with +an inner PaliGemmaModel + DynamicCache). + +Reference implementations: + - OpenPI: openpi/src/openpi/models_pytorch/pi0_pytorch.py, gemma_pytorch.py + - LeRobot: lerobot/src/lerobot/policies/pi0/modeling_pi0.py +Weight source: https://huggingface.co/lerobot/pi0_base +""" + +from __future__ import annotations + +import logging +import math +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.models.auto import CONFIG_MAPPING +from transformers.models.gemma.modeling_gemma import ( + GemmaForCausalLM, + apply_rotary_pos_emb, +) +from transformers.models.paligemma.modeling_paligemma import ( + PaliGemmaForConditionalGeneration, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +logger = logging.getLogger(__name__) + +# ────────────────────────────────────────────────────────────────────── +# Constants +# ────────────────────────────────────────────────────────────────────── +DEFAULT_ACTION_DIM = 32 +DEFAULT_ACTION_HORIZON = 50 +DEFAULT_MAX_TOKEN_LEN = 48 +DEFAULT_NUM_INFERENCE_STEPS = 10 +DEFAULT_IMAGE_RESOLUTION = (224, 224) + +# Large negative value to fill masked-out positions in a float attention mask. +# Matches OpenPI's constant exactly so that numerics line up during parity. +# Ref: openpi/src/openpi/models/gemma.py +OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38 + + +# ────────────────────────────────────────────────────────────────────── +# Gemma variant configs (matches openpi/models/gemma.py get_config) +# ────────────────────────────────────────────────────────────────────── + + +class GemmaVariantConfig: + def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim): + self.width = width + self.depth = depth + self.mlp_dim = mlp_dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + + +def get_gemma_config(variant: str) -> GemmaVariantConfig: + if variant == "gemma_2b": + return GemmaVariantConfig(2048, 18, 16384, 8, 1, 256) + elif variant == "gemma_300m": + return GemmaVariantConfig(1024, 18, 4096, 8, 1, 256) + else: + raise ValueError(f"Unknown variant: {variant}") + + +# ────────────────────────────────────────────────────────────────────── +# Utility functions (match openpi/models_pytorch/pi0_pytorch.py) +# ────────────────────────────────────────────────────────────────────── + + +def create_sinusoidal_pos_embedding( + time: torch.Tensor, + dimension: int, + min_period: float = 4e-3, + max_period: float = 4.0, + device: torch.device = None, +) -> torch.Tensor: + """Compute a sine/cosine positional embedding for scalar timesteps. + + Ref: openpi/models_pytorch/pi0_pytorch.py create_sinusoidal_pos_embedding + """ + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + if time.ndim != 1: + raise ValueError("time tensor must be 1-D (batch_size,)") + if device is None: + device = time.device + + # Use float64 for the log-linear sweep and for the inner products, to + # match the reference implementation's numerical behaviour exactly. + dtype = torch.float64 + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) + period = min_period * (max_period / min_period) ** fraction + scaling_factor = 1.0 / period * 2 * math.pi + sin_input = scaling_factor[None, :] * time[:, None].to(dtype) + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + + +def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor: + """Build a 2D attention mask from a padding mask and an autoregressive mask. + + Ref: openpi/models_pytorch/pi0_pytorch.py make_att_2d_masks + """ + if att_masks.ndim != 2: + raise ValueError(f"att_masks must be 2-D, got {att_masks.ndim}-D") + if pad_masks.ndim != 2: + raise ValueError(f"pad_masks must be 2-D, got {pad_masks.ndim}-D") + + cumsum = torch.cumsum(att_masks, dim=1) + att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d_masks & pad_2d_masks + + +def _build_norm_buffers(norm_stats: dict | None, key: str) -> dict[str, torch.Tensor] | None: + """Parse a ``norm_stats[key]`` entry into CPU tensors, or ``None``. + + Matches LeRobot's ``NormalizationMode``: + - ``mean_std`` : ``forward = (x - mean) / std``, ``inverse = x * std + mean`` + - ``min_max`` : forward maps into ``[-1, 1]``, inverse maps back. + """ + if not norm_stats or not isinstance(norm_stats, dict): + return None + entry = norm_stats.get(key) + if not entry: + return None + mode = str(entry.get("mode", "mean_std")).lower() + if mode == "mean_std": + mean = entry.get("mean") + std = entry.get("std") + if mean is None or std is None: + return None + return { + "mode": mode, + "mean": torch.as_tensor(mean, dtype=torch.float32), + "std": torch.as_tensor(std, dtype=torch.float32), + } + if mode == "min_max": + lo = entry.get("min") + hi = entry.get("max") + if lo is None or hi is None: + return None + return { + "mode": mode, + "min": torch.as_tensor(lo, dtype=torch.float32), + "max": torch.as_tensor(hi, dtype=torch.float32), + } + return None + + +def _apply_norm( + x: torch.Tensor, + stats: dict[str, torch.Tensor] | None, + inverse: bool, + eps: float = 1e-8, +) -> torch.Tensor: + """Apply (un)normalization using the given stats. No-op if ``stats`` is None. + + Broadcasts over leading dims; stats vectors align with the last dim. When + ``x`` has more last-dim entries than the stats (π0 pads state/action to + ``max_dim``), only the first ``len(stats)`` entries are transformed and + the padded tail is left untouched. + """ + if stats is None: + return x + mode = stats["mode"] + if mode == "mean_std": + mean = stats["mean"].to(device=x.device, dtype=x.dtype) + std = stats["std"].to(device=x.device, dtype=x.dtype) + valid = mean.shape[0] + head = x[..., :valid] + if inverse: + head = head * std + mean + else: + head = (head - mean) / (std + eps) + if valid == x.shape[-1]: + return head + return torch.cat([head, x[..., valid:]], dim=-1) + if mode == "min_max": + lo = stats["min"].to(device=x.device, dtype=x.dtype) + hi = stats["max"].to(device=x.device, dtype=x.dtype) + valid = lo.shape[0] + head = x[..., :valid] + denom = (hi - lo).clamp_min(eps) + if inverse: + head = (head + 1.0) * 0.5 * denom + lo + else: + head = 2.0 * (head - lo) / denom - 1.0 + if valid == x.shape[-1]: + return head + return torch.cat([head, x[..., valid:]], dim=-1) + return x + + +def prepare_attention_masks_4d(att_2d_masks: torch.Tensor) -> torch.Tensor: + """Convert ``(B, S, S)`` bool masks to ``(B, 1, S, S)`` float masks. + + ``True`` → 0.0 (attend), ``False`` → ``OPENPI_ATTENTION_MASK_VALUE``. + Ref: openpi PI0Pytorch._prepare_attention_masks_4d + """ + att_2d_masks_4d = att_2d_masks[:, None, :, :] + return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE) + + +# ────────────────────────────────────────────────────────────────────── +# Dual-backbone: PaliGemma + Action Expert +# Ref: openpi/models_pytorch/gemma_pytorch.py PaliGemmaWithExpertModel +# Ref: lerobot/policies/pi0/modeling_pi0.py PaliGemmaWithExpertModel +# ────────────────────────────────────────────────────────────────────── + + +def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Repeat KV heads for grouped-query attention. + + ``(B, num_kv_heads, S, D) → (B, num_kv_heads * n_rep, S, D)``. + Matches ``transformers.models.gemma.modeling_gemma.repeat_kv``. + """ + if n_rep == 1: + return hidden_states + b, nh, s, d = hidden_states.shape + hidden_states = hidden_states[:, :, None, :, :].expand(b, nh, n_rep, s, d) + return hidden_states.reshape(b, nh * n_rep, s, d) + + +def _attend(query_states, key_states, value_states, attention_mask, num_kv_groups, scaling): + """Manual eager attention: ``softmax(Q Kᵀ · scale + mask) · V``. + + Important: we slice ``attention_mask[..., :key_states.shape[-2]]`` so that + the same ``(B, 1, Q, prefix+suffix)`` mask produced by ``denoise_step`` + can be re-used both in the prefix pass (K length = prefix) and in the + suffix pass (K length = prefix + suffix after we concat the cached prefix). + Stock transformers' ``eager_attention_forward`` adds the mask without + slicing and therefore breaks when the two dims disagree. + """ + k = _repeat_kv(key_states, num_kv_groups) + v = _repeat_kv(value_states, num_kv_groups) + attn_weights = torch.matmul(query_states, k.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask[:, :, :, : k.shape[-2]] + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + return torch.matmul(attn_weights, v) + + +def _compute_layer_prefix_only(layer_idx, hidden_states, attention_mask, position_ids, paligemma): + """Run one PaliGemma LM layer on the prefix only, returning the layer + output **and** the post-RoPE ``(k, v)`` so the caller can cache them for + the suffix pass.""" + model = paligemma.model.language_model + layer = model.layers[layer_idx] + residual = hidden_states + x = layer.input_layernorm(hidden_states) + + hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim) + q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2) + k = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2) + v = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2) + + cos, sin = model.rotary_emb(v, position_ids) + q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1) + + att = _attend( + q, + k, + v, + attention_mask, + num_kv_groups=layer.self_attn.num_key_value_groups, + scaling=1.0 / math.sqrt(layer.self_attn.head_dim), + ) + att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim) + + out = layer.self_attn.o_proj(att) + residual + after_resid = out + out = layer.mlp(layer.post_attention_layernorm(out)) + after_resid + return out, (k, v) + + +def _compute_layer_suffix_only( + layer_idx, + hidden_states, + prefix_kv, + attention_mask, + position_ids, + paligemma, + gemma_expert, +): + """Run one action-expert layer on the suffix, attending to the cached + prefix K/V concatenated with freshly computed suffix K/V. + + ``prefix_kv`` is a ``(k_prefix, v_prefix)`` tuple, shape + ``(B, num_kv_heads, prefix_len, head_dim)``, post-RoPE, produced by the + corresponding ``_compute_layer_prefix_only`` call. Bypasses + ``GemmaModel.forward`` entirely — no HF mask rebuild, no cache format + gymnastics. + """ + layer = gemma_expert.model.layers[layer_idx] + residual = hidden_states + x = layer.input_layernorm(hidden_states) + + hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim) + q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2) + k_suf = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2) + v_suf = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2) + + # RoPE frequencies are shared between PaliGemma and the expert. + cos, sin = gemma_expert.model.rotary_emb(v_suf, position_ids) + q, k_suf = apply_rotary_pos_emb(q, k_suf, cos, sin, unsqueeze_dim=1) + + # Concatenate cached prefix K/V (possibly different dtype) with suffix K/V. + k_prefix, v_prefix = prefix_kv + k = torch.cat([k_prefix.to(k_suf.dtype), k_suf], dim=2) + v = torch.cat([v_prefix.to(v_suf.dtype), v_suf], dim=2) + + att = _attend( + q, + k, + v, + attention_mask, + num_kv_groups=layer.self_attn.num_key_value_groups, + scaling=1.0 / math.sqrt(layer.self_attn.head_dim), + ) + att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim) + + out = layer.self_attn.o_proj(att) + residual + after_resid = out + out = layer.mlp(layer.post_attention_layernorm(out)) + after_resid + return out + + +class PaliGemmaWithActionExpert(nn.Module): + """Dual-backbone transformer: PaliGemma (Gemma 2B) + Action Expert (Gemma 300M). + + ``forward`` has two inference modes: + - **prefix_only**: ``inputs_embeds=[prefix, None]`` + ``use_cache=True`` + → compute prefix hidden states + a layer-wise K/V cache. + - **suffix_only**: ``inputs_embeds=[None, suffix]`` + the cache from the + previous prefix pass → compute expert hidden states with cross-attention + over the concatenated (prefix, suffix) K/V. + + Both modes walk Gemma decoder layers manually through ``_compute_layer_*`` + instead of ``GemmaModel.forward`` — that way we own the attention mask + handling and the KV cache format (a plain ``list[(k, v)]`` one entry + per layer). + + Ref: lerobot/policies/pi0/modeling_pi0.py PaliGemmaWithExpertModel + """ + + def __init__(self, vlm_config, action_expert_config): + super().__init__() + + # Build HF PaliGemma config from the variant dims we expose. + # NOTE: we do NOT set use_adarms / adarms_cond_dim — those are + # OpenPI-only attributes used by π0.5, not π0. + vlm_config_hf = CONFIG_MAPPING["paligemma"]() + vlm_config_hf._vocab_size = 257152 + vlm_config_hf.image_token_index = 257152 + vlm_config_hf.text_config.hidden_size = vlm_config.width + vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim + vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads + vlm_config_hf.text_config.head_dim = vlm_config.head_dim + vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth + vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads + vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh" + # transformers ≥ 5 uses ``dtype``; older versions still accept + # ``torch_dtype``. We set the canonical name; HF's PretrainedConfig + # normalizes one to the other internally. + vlm_config_hf.text_config.dtype = "float32" + vlm_config_hf.text_config.vocab_size = 257152 + vlm_config_hf.vision_config.intermediate_size = 4304 + vlm_config_hf.vision_config.projection_dim = 2048 + vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast" + vlm_config_hf.vision_config.dtype = "float32" + + action_expert_config_hf = CONFIG_MAPPING["gemma"]( + head_dim=action_expert_config.head_dim, + hidden_size=action_expert_config.width, + intermediate_size=action_expert_config.mlp_dim, + num_attention_heads=action_expert_config.num_heads, + num_hidden_layers=action_expert_config.depth, + num_key_value_heads=action_expert_config.num_kv_heads, + vocab_size=257152, + hidden_activation="gelu_pytorch_tanh", + dtype="float32", + ) + + self.paligemma = PaliGemmaForConditionalGeneration(config=vlm_config_hf) + self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf) + # The action expert doesn't embed tokens — it only consumes the + # suffix state/action embeddings we feed in. + self.gemma_expert.model.embed_tokens = None + + def embed_image(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Encode images with SigLIP vision tower + PaliGemma projector. + + We run the two steps explicitly instead of calling + ``PaliGemmaModel.get_image_features`` because that convenience + method historically has returned either the projected tensor (old) + or the raw vision-tower output (new). Being explicit makes us + independent of which transformers release we're on. + """ + # Shapes: pixel_values (B, 3, 224, 224) → SigLIP (B, 256, 1152) + # → projector (B, 256, 2048) + vision_outputs = self.paligemma.model.vision_tower(pixel_values) + image_features = vision_outputs.last_hidden_state + return self.paligemma.model.multi_modal_projector(image_features) + + def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor: + """Embed language tokens with PaliGemma's embedding table, returning the + ``* sqrt(hidden)``-scaled embedding (PaliGemma/Gemma convention). + + The scaling location moved across transformers releases: + * ≤ 5.3: ``embed_tokens`` is a plain ``nn.Embedding`` and the + ``* sqrt(hidden)`` normalizer lives inside ``GemmaModel.forward`` + (which we bypass) — so we apply it explicitly here. + * ≥ 5.4: ``embed_tokens`` is a ``GemmaTextScaledWordEmbedding`` that + self-applies ``embed_scale = hidden_size ** 0.5`` — applying it again + would double-scale (≈45×). We detect this and skip the manual scale. + This makes ``embed_language_tokens`` return the canonical scaled embedding + on every transformers version. + """ + embed_tokens = self.paligemma.model.language_model.embed_tokens + lang_emb = embed_tokens(tokens) + # If the embedding already self-scales (transformers ≥ 5.4), don't repeat it. + if getattr(embed_tokens, "embed_scale", None) is None: + lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1]) + return lang_emb + + def forward( + self, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values=None, + inputs_embeds: list[torch.Tensor | None] | None = None, + use_cache: bool = False, + ): + """Dispatch to prefix_only / suffix_only and return + ``([prefix_out, suffix_out], past_key_values_or_None)``. + """ + num_layers = self.paligemma.config.text_config.num_hidden_layers + pali_lm = self.paligemma.model.language_model + expert_lm = self.gemma_expert.model + + if inputs_embeds[1] is None: + # Prefix-only: PaliGemma LM on (images + language) tokens; the + # per-layer post-RoPE K/V is collected into a list that the + # suffix pass will consume directly. + hidden_states = inputs_embeds[0] + kv_list: list[tuple[torch.Tensor, torch.Tensor]] = [] + for layer_idx in range(num_layers): + hidden_states, kv = _compute_layer_prefix_only( + layer_idx, + hidden_states, + attention_mask, + position_ids, + paligemma=self.paligemma, + ) + kv_list.append(kv) + hidden_states = pali_lm.norm(hidden_states) + return [hidden_states, None], (kv_list if use_cache else None) + + if inputs_embeds[0] is not None: + raise ValueError( + "PaliGemmaWithActionExpert.forward only supports prefix-only " + "or suffix-only dispatch; got both inputs_embeds populated." + ) + # Suffix-only: action expert on (state + noisy_actions + time), + # attending to the cached prefix K/V. + if not isinstance(past_key_values, list): + raise TypeError( + "suffix_only forward expects past_key_values to be the " + "list[(k, v)] produced by a previous prefix_only forward; " + f"got {type(past_key_values)}" + ) + hidden_states = inputs_embeds[1] + for layer_idx in range(num_layers): + hidden_states = _compute_layer_suffix_only( + layer_idx, + hidden_states, + past_key_values[layer_idx], + attention_mask, + position_ids, + paligemma=self.paligemma, + gemma_expert=self.gemma_expert, + ) + hidden_states = expert_lm.norm(hidden_states) + return [None, hidden_states], None + + +# ────────────────────────────────────────────────────────────────────── +# Main π0 Model +# ────────────────────────────────────────────────────────────────────── + + +class Pi0ForActionPrediction(nn.Module): + """π0 VLA model for robot action prediction via flow matching. + + Inference flow: + 1. Embed prefix (images + language) → prefix tokens. + 2. Forward prefix through PaliGemma → layer-wise KV cache. + 3. For each denoising step ``t = 1.0, 1-dt, ..., 0``: + a. Embed suffix (state + x_t + timestep) → suffix tokens. + b. Forward suffix through the action expert with the prefix cache. + c. ``x_t = x_t + dt * v_t`` (Euler integration). + 4. Return ``x_0`` as the predicted action chunk. + """ + + def __init__( + self, + config, + quant_config=None, + prefix: str = "", + ): + super().__init__() + # ``quant_config`` is accepted for interface compatibility but unused — + # π0 runs in full precision for flow-matching parity. + del quant_config + self.config = config + + self.action_dim = getattr(config, "max_action_dim", DEFAULT_ACTION_DIM) + # ``max_state_dim`` is independent from ``max_action_dim`` — e.g. Aloha + # has state_dim=14, action_dim=32, both padded to their own max. Fall + # back to action_dim for older configs that only expose a single size. + self.max_state_dim = getattr(config, "max_state_dim", self.action_dim) + self.action_horizon = getattr(config, "chunk_size", DEFAULT_ACTION_HORIZON) + self.num_inference_steps = getattr(config, "num_inference_steps", DEFAULT_NUM_INFERENCE_STEPS) + + paligemma_variant = getattr(config, "paligemma_variant", "gemma_2b") + action_expert_variant = getattr(config, "action_expert_variant", "gemma_300m") + vlm_config = get_gemma_config(paligemma_variant) + expert_config = get_gemma_config(action_expert_variant) + self.vlm_width = vlm_config.width + self.expert_width = expert_config.width + + # Dual backbone + self.paligemma_with_expert = PaliGemmaWithActionExpert(vlm_config, expert_config) + + # Action chunk projections (openpi pi0_pytorch.py). + self.action_in_proj = nn.Linear(self.action_dim, self.expert_width) + self.action_out_proj = nn.Linear(self.expert_width, self.action_dim) + + # State is projected to the expert dim once (π0 only, not π0.5). + # Ref: lerobot modeling_pi0.py ``nn.Linear(config.max_state_dim, ...)`` + self.state_proj = nn.Linear(self.max_state_dim, self.expert_width) + + # Timestep + action fusion MLP (2W → W → W). + self.action_time_mlp_in = nn.Linear(2 * self.expert_width, self.expert_width) + self.action_time_mlp_out = nn.Linear(self.expert_width, self.expert_width) + + # Optional mean/std normalization stats (LeRobot NormalizerProcessorStep + # semantics). Format on the config (all optional): + # norm_stats = { + # "state": {"mode": "mean_std", "mean": [...], "std": [...]}, + # "action": {"mode": "mean_std", "mean": [...], "std": [...]}, + # } + # ``mode`` may also be ``"min_max"``. + # + # Missing norm_stats is logged at INFO (not WARNING) because the most + # common setup is "client pre-normalizes state and post-normalizes + # actions using dataset stats" — that's a supported mode, not an error. + self._state_norm = _build_norm_buffers(getattr(config, "norm_stats", None), "state") + self._action_norm = _build_norm_buffers(getattr(config, "norm_stats", None), "action") + if self._state_norm is None: + logger.info( + "π0: no state normalization stats on config.norm_stats — state values will pass through unchanged." + ) + if self._action_norm is None: + logger.info( + "π0: no action normalization stats on config.norm_stats — " + "returned actions are in the model's normalized space." + ) + + # ── State / action normalization ───────────────────────────────── + + def _normalize_state(self, state: torch.Tensor) -> torch.Tensor: + return _apply_norm(state, self._state_norm, inverse=False) + + def _unnormalize_actions(self, actions: torch.Tensor) -> torch.Tensor: + return _apply_norm(actions, self._action_norm, inverse=True) + + # ── Prefix embedding ───────────────────────────────────────────── + + def embed_prefix( + self, + images: list[torch.Tensor], + image_masks: list[torch.Tensor], + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the prefix embeddings, per-token padding mask, and AR mask. + + Prefix tokens form a contiguous sequence + ``[img_cam_0..., img_cam_1..., ..., lang_tokens...]`` with + bidirectional attention; the returned ``att_masks`` are all zeros + because each token is free to attend to every other prefix token (the + suffix pass will put a causal boundary right before the state token). + + Ref: openpi PI0Pytorch.embed_prefix + """ + embs: list[torch.Tensor] = [] + pad_masks: list[torch.Tensor] = [] + att_masks: list[int] = [] + + for img, img_mask in zip(images, image_masks): + img_emb = self.paligemma_with_expert.embed_image(img) + bsize, num_img_embs = img_emb.shape[:2] + embs.append(img_emb) + pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs)) + att_masks += [0] * num_img_embs + + # ``embed_language_tokens`` already returns the ``* sqrt(hidden)``-scaled + # embedding (version-robust: self-scaled at transformers ≥ 5.4, manually + # scaled at ≤ 5.3). The image slice went through the projector and is + # NOT scaled. See embed_language_tokens for the version details. + lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens) + + embs.append(lang_emb) + pad_masks.append(lang_masks) + att_masks += [0] * lang_emb.shape[1] + + embs = torch.cat(embs, dim=1) + pad_masks = torch.cat(pad_masks, dim=1) + att_masks = torch.tensor(att_masks, dtype=torch.bool, device=embs.device) + att_masks = att_masks[None, :].expand(pad_masks.shape[0], -1) + + return embs, pad_masks, att_masks + + # ── Suffix embedding ───────────────────────────────────────────── + + def embed_suffix( + self, + state: torch.Tensor, + noisy_actions: torch.Tensor, + timestep: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the suffix embeddings + masks: ``[state_token, action_tokens×H]``. + + AR mask layout: ``[1, 1, 0, 0, ..., 0]`` — the state token is a causal + boundary (no later token attends backwards through it onto prefix + *by mistake*), the first action token starts a new causal block, and + the rest of the action tokens attend to each other bidirectionally. + + Ref: openpi PI0Pytorch.embed_suffix + """ + model_dtype = self.state_proj.weight.dtype + state = state.to(dtype=model_dtype) + noisy_actions = noisy_actions.to(dtype=model_dtype) + device = state.device + + # State → (B, 1, W) + state_emb = self.state_proj(state)[:, None, :] + bsize = state_emb.shape[0] + + # Sinusoidal timestep → (B, W) → expand across the horizon. + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.action_in_proj.out_features, + min_period=4e-3, + max_period=4.0, + device=device, + ).to(dtype=model_dtype) + + # Fuse action and time via (2W → W → W) MLP with SiLU. + action_emb = self.action_in_proj(noisy_actions) # (B, H, W) + action_time_emb = torch.cat([action_emb, time_emb[:, None, :].expand_as(action_emb)], dim=2) + action_time_emb = self.action_time_mlp_in(action_time_emb) + action_time_emb = F.silu(action_time_emb) + action_time_emb = self.action_time_mlp_out(action_time_emb) + + embs = torch.cat([state_emb, action_time_emb], dim=1) + pad_masks = torch.ones(bsize, embs.shape[1], dtype=torch.bool, device=device) + att_masks = torch.tensor( + [1] + [1] + [0] * (self.action_horizon - 1), + dtype=embs.dtype, + device=device, + )[None, :].expand(bsize, -1) + return embs, pad_masks, att_masks + + # ── Denoising step ─────────────────────────────────────────────── + + def denoise_step( + self, + state: torch.Tensor, + prefix_pad_masks: torch.Tensor, + past_key_values, + x_t: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + """Apply one flow-matching denoising step: predict ``v_t`` from ``x_t``. + + Uses the prefix KV cache from ``sample_actions`` and only runs the + action expert. Ref: openpi PI0Pytorch.denoise_step. + """ + suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, x_t, timestep) + + # Build the full (B, suffix_len, prefix_len + suffix_len) boolean mask: + # * suffix queries can see every *valid* prefix key (padded camera + # slots and padded language tokens are blocked). + # * within the suffix, the state token is a causal boundary; action + # tokens attend to each other bidirectionally. + batch_size = prefix_pad_masks.shape[0] + suffix_len = suffix_pad_masks.shape[1] + prefix_len = prefix_pad_masks.shape[1] + + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + + # Position IDs continue from where the prefix's last valid token left off. + prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] + position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 + + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) + + outputs_embeds, _ = self.paligemma_with_expert.forward( + attention_mask=full_att_2d_masks_4d, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=False, + ) + + # Drop the state token, keep only the action tokens, project down. + suffix_out = outputs_embeds[1][:, -self.action_horizon :] + suffix_out = suffix_out.to(dtype=self.action_out_proj.weight.dtype) + return self.action_out_proj(suffix_out) + + # ── Full action generation ─────────────────────────────────────── + + @torch.no_grad() + def sample_actions( + self, + images: list[torch.Tensor], + image_masks: list[torch.Tensor], + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + state: torch.Tensor, + noise: torch.Tensor | None = None, + num_steps: int | None = None, + ) -> torch.Tensor: + """Generate an action chunk via iterative flow-matching denoising. + + Convention: ``t=1`` is noise, ``t=0`` is the target — opposite of the + published π0 paper but matches both OpenPI and LeRobot. + Ref: openpi PI0Pytorch.sample_actions + """ + if num_steps is None: + num_steps = self.num_inference_steps + + bsize = state.shape[0] + device = state.device + if noise is None: + noise = torch.randn( + bsize, + self.action_horizon, + self.action_dim, + dtype=torch.float32, + device=device, + ) + + # 1. Prefix embeddings + mask building. + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( + images, image_masks, lang_tokens, lang_masks + ) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks) + + # 2. Forward prefix through PaliGemma LM, producing a list[(k, v)] cache. + _, past_key_values = self.paligemma_with_expert.forward( + attention_mask=prefix_att_2d_masks_4d, + position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=True, + ) + + # 3. Euler-integrated denoising from t=1 down to t=0. + dt = -1.0 / num_steps + x_t = noise + for step in range(num_steps): + t = 1.0 + step * dt + time_tensor = torch.full((bsize,), t, dtype=torch.float32, device=device) + v_t = self.denoise_step( + state=state, + prefix_pad_masks=prefix_pad_masks, + past_key_values=past_key_values, + x_t=x_t, + timestep=time_tensor, + ) + x_t = x_t + dt * v_t + return x_t + + # ── Weight loading ─────────────────────────────────────────────── + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + """Load weights from an ``lerobot/pi0_base`` safetensors checkpoint. + + Checkpoint keys (after stripping the leading ``model.`` prefix) look + like ``paligemma_with_expert.paligemma..*``. Modern transformers + nests those sub-modules one level deeper (``paligemma.model.``), + and PaliGemma ties ``lm_head.weight`` with ``embed_tokens.weight`` at + ``post_init`` time. Two remap rules make the load lossless: + + 1. ``paligemma.{vision_tower,multi_modal_projector,language_model}.*`` + → ``paligemma.model.{...}.*`` + 2. ``paligemma.lm_head.weight`` + → ``paligemma.model.language_model.embed_tokens.weight`` + (the checkpoint does not store ``embed_tokens.weight`` on its + own — only the tied ``lm_head.weight`` copy — and + ``PaliGemmaForConditionalGeneration`` does not register + ``lm_head.weight`` as a Parameter when tied, so without this + rule the language embedding would silently remain at its + random init.) + + On any remaining mismatch the loader logs a warning listing a sample + of the unmatched checkpoint keys and model params; downstream parity + tests surface these as a hard failure. + """ + params_dict = dict(self.named_parameters()) + buffers_dict = dict(self.named_buffers()) + + _PALIGEMMA_SUBMODULES = ( + "vision_tower", + "multi_modal_projector", + "language_model", + ) + + def _remap(name: str) -> str: + # Strip the leading "model." that LeRobot's PI0Policy wrapper adds. + if name.startswith("model."): + name = name[len("model.") :] + + # Historical MLP alias used by some early π0 checkpoints. + if name.startswith("time_mlp_in."): + name = "action_time_mlp_in." + name[len("time_mlp_in.") :] + elif name.startswith("time_mlp_out."): + name = "action_time_mlp_out." + name[len("time_mlp_out.") :] + + # Nested PaliGemma layout. + for sub in _PALIGEMMA_SUBMODULES: + flat = f"paligemma_with_expert.paligemma.{sub}." + nested = f"paligemma_with_expert.paligemma.model.{sub}." + if name.startswith(flat) and not name.startswith(nested): + return nested + name[len(flat) :] + + # Tied lm_head → embed_tokens redirect (see docstring). + if name == "paligemma_with_expert.paligemma.lm_head.weight": + return "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + + return name + + def _fix_vision_tower(name: str) -> str: + """Reconcile the SigLIP vision-tower nesting across transformers + versions. transformers ≤5.3 wraps the encoder in an extra + ``vision_tower.vision_model.*`` module; ≥5.4 flattens it to + ``vision_tower.*``. The ``lerobot/pi0_base`` checkpoint uses the + ``vision_model`` layout, so on newer transformers we must drop that + infix (else 437 vision params silently load nothing). Adaptive: + only rewrite when the rewritten key actually exists on the model. + """ + vt = "paligemma_with_expert.paligemma.model.vision_tower." + if not name.startswith(vt): + return name + rest = name[len(vt) :] + if name in params_dict or name in buffers_dict: + return name # model already expects this exact layout + if rest.startswith("vision_model."): + # Checkpoint has the infix; model (≥5.4) doesn't → strip it. + candidate = vt + rest[len("vision_model.") :] + else: + # Checkpoint lacks the infix; model (≤5.3) has it → insert it. + candidate = vt + "vision_model." + rest + return candidate if (candidate in params_dict or candidate in buffers_dict) else name + + loaded = 0 + skipped: list[str] = [] + filled_params: set = set() + for name, loaded_weight in weights: + mapped = _fix_vision_tower(_remap(name)) + if mapped in params_dict: + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded += 1 + filled_params.add(mapped) + elif mapped in buffers_dict: + buffers_dict[mapped].copy_(loaded_weight) + loaded += 1 + filled_params.add(mapped) + else: + skipped.append(mapped) + + # Reverse audit: any model param that got no checkpoint tensor at all + # would be running with random init. ``rotary_emb.inv_freq`` and + # friends are config-derived buffers, not stored in the checkpoint. + missing_params: list[str] = [] + for pname in params_dict: + if pname in filled_params: + continue + if "rotary_emb" in pname or pname.endswith(".inv_freq"): + continue + missing_params.append(pname) + + if missing_params or skipped: + parts: list[str] = [] + if skipped: + parts.append(f"{len(skipped)} checkpoint key(s) had no home in the model (first 5: {skipped[:5]})") + if missing_params: + parts.append(f"{len(missing_params)} model param(s) received NO weight (first 5: {missing_params[:5]})") + logger.warning( + "π0 load_weights: %d tensors loaded — %s.", + loaded, + "; ".join(parts), + ) + else: + logger.info( + "π0 load_weights: %d tensors loaded, 0 skipped, 0 missing.", + loaded, + ) + return filled_params diff --git a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py new file mode 100644 index 000000000..1eeca2969 --- /dev/null +++ b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""π0 (Pi-Zero) VLA pipeline for vllm-omni. + +Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the +DreamZero contract: the pipeline owns ALL preprocessing. It reads the raw robot +observation from ``req.sampling_params.extra_args["robot_obs"]`` (delivered by +the OpenPI realtime serving layer), builds model inputs, runs flow-matching +denoising, and returns ``DiffusionOutput(output={"actions": ndarray})`` — which +``diffusion_engine`` promotes to ``multimodal_output["actions"]`` for the client. + +π0 is stateless across calls (no KV reuse), so ``session_id`` / ``reset`` from +the OpenPI protocol are accepted but ignored. +""" + +from __future__ import annotations + +import os + +import numpy as np +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.pi0.config import Pi0Config +from vllm_omni.diffusion.models.pi0.modeling_pi0 import Pi0ForActionPrediction +from vllm_omni.diffusion.models.pi0.processor_pi0 import build_model_inputs +from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE as PI0_PIPELINE +from vllm_omni.diffusion.request import OmniDiffusionRequest + +logger = init_logger(__name__) + +# Default tokenizer for the PaliGemma prefix (matches LeRobot Pi0). +DEFAULT_PI0_TOKENIZER = "google/paligemma-3b-pt-224" + + +def _pi0_post_process(x): + """Module-level identity post-process (picklable across the orchestrator's + multiprocess boundary — a local closure is not).""" + return x + + +def get_pi0_post_process_func(od_config: OmniDiffusionConfig): + """π0 returns actions directly; post-processing is identity. The diffusion + engine resolves this via the registry (``_DIFFUSION_POST_PROCESS_FUNCS``) and + applies it engine-side, so the pipeline does NOT attach it to DiffusionOutput. + """ + del od_config + return _pi0_post_process + + +class Pi0Pipeline(nn.Module): + """π0 VLA pipeline: raw robot obs → continuous action chunk. + + Registered as ``"Pi0Pipeline"`` in the diffusion registry. Weights are + self-loaded in ``__init__`` from the checkpoint's ``model.safetensors`` via + the kernel's ``load_weights`` (which handles the LeRobot key remaps). + """ + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.prefix = prefix + # Resolve od_config.model to a LOCAL directory. A bare HF repo id (e.g. + # the documented ``lerobot/pi0_base``) must be snapshot-downloaded first; + # otherwise config/tokenizer/weights would silently fall back to random + # init (this pipeline self-loads from a local dir and has no + # ``weights_sources`` for the framework loader to download from). + self.model_dir = self._resolve_model_dir(od_config.model) + self.config = self._build_config(od_config) + + custom_args = od_config.custom_pipeline_args or {} + self.tokenizer_source = str(custom_args.get("tokenizer", self._resolve_tokenizer_source())) + + # Torch dtype/device from od_config. + self._torch_dtype = self._resolve_dtype(od_config) + self._device = self._resolve_device(od_config) + + self.tokenizer = self._load_tokenizer() + self.model = self._initialize_model() + + # ------------------------------------------------------------------ + # Construction helpers + # ------------------------------------------------------------------ + @staticmethod + def _resolve_model_dir(model: str | None) -> str | None: + """Return a local directory for ``model``; download an HF repo id if needed. + + Limits the snapshot to the files π0 actually reads (config, weights, + tokenizer) so a repo id doesn't pull unrelated artifacts. + """ + if not model: + return None + if os.path.isdir(model): + return model + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id=model, + allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"], + ) + + def _build_config(self, od_config: OmniDiffusionConfig) -> Pi0Config: + """Build Pi0Config from deploy-yaml model_config, falling back to the + checkpoint's config.json (raw LeRobot format).""" + if od_config.model_config: + config = Pi0Config.from_model_config(dict(od_config.model_config)) + # The deploy yaml is authoritative, but must not silently drop + # checkpoint-derived fields it omits. Backfill the camera order and + # any normalization stats the yaml didn't specify. + if self.model_dir and (not config.image_feature_keys or config.norm_stats is None): + ckpt = Pi0Config.from_pretrained(self.model_dir) + if not config.image_feature_keys: + config.image_feature_keys = ckpt.image_feature_keys + if not config.input_features: + config.input_features = ckpt.input_features + if config.norm_stats is None: + config.norm_stats = ckpt.norm_stats + return config + if self.model_dir: + return Pi0Config.from_pretrained(self.model_dir) + return Pi0Config() + + def _resolve_tokenizer_source(self) -> str: + """Prefer the checkpoint dir if it ships tokenizer files; else PaliGemma.""" + if self.model_dir and os.path.isdir(self.model_dir): + if os.path.exists(os.path.join(self.model_dir, "tokenizer_config.json")): + return self.model_dir + return DEFAULT_PI0_TOKENIZER + + @staticmethod + def _resolve_dtype(od_config: OmniDiffusionConfig) -> torch.dtype: + dt = od_config.dtype + if isinstance(dt, torch.dtype): + return dt + return getattr(torch, str(dt).split(".")[-1], torch.float32) + + @staticmethod + def _resolve_device(od_config: OmniDiffusionConfig) -> torch.device: + from vllm_omni.diffusion.distributed.utils import get_local_device + + try: + return get_local_device() + except Exception: # noqa: BLE001 + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def _load_tokenizer(self): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(self.tokenizer_source) + + def has_real_checkpoint(self) -> bool: + return bool(self.model_dir) and os.path.exists(os.path.join(self.model_dir, "model.safetensors")) + + def _initialize_model(self) -> Pi0ForActionPrediction: + model = Pi0ForActionPrediction(self.config) + if self.has_real_checkpoint(): + self._load_checkpoint(model) + else: + logger.info("Pi0Pipeline: no model.safetensors under %s; using random init.", self.model_dir) + model.to(device=self._device, dtype=self._torch_dtype) + model.eval() + return model + + def _load_checkpoint(self, model: Pi0ForActionPrediction) -> None: + import safetensors.torch + + path = os.path.join(self.model_dir, "model.safetensors") + logger.info("Pi0Pipeline: loading π0 weights from %s", path) + state = safetensors.torch.load_file(path) + model.load_weights(list(state.items())) + + # ------------------------------------------------------------------ + # Framework weight-loading hook + # ------------------------------------------------------------------ + def load_weights(self, weights=()): # noqa: D401 + """No-op for the diffusion loader: π0 self-loads its checkpoint in + ``__init__`` (the kernel's ``load_weights`` handles the LeRobot remaps). + We expose no ``weights_sources``, so the loader passes an empty iterator + here; returning ``None`` skips its strict unloaded-weights check. + """ + for _ in weights: + pass + return None + + # ------------------------------------------------------------------ + # Inference + # ------------------------------------------------------------------ + @torch.inference_mode() + def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: + extra_args = getattr(req.sampling_params, "extra_args", None) or {} + robot_obs = extra_args.get("robot_obs") + + if robot_obs is None: + # Dummy warmup path (no obs): return zeros so engine warmup/capture + # doesn't crash. Mirrors DreamZero's dummy-run handling. + first_prompt = req.prompts[0] if req.prompts else "" + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + num_steps = getattr(req.sampling_params, "num_inference_steps", None) + if prompt == "dummy run" or num_steps == 1: + logger.info("Pi0Pipeline: dummy warmup request without robot_obs — returning zeros.") + return DiffusionOutput( + output={ + "actions": np.zeros( + (self.config.chunk_size, self.config.max_action_dim), + dtype=np.float32, + ) + }, + ) + return DiffusionOutput( + error="Pi0Pipeline.forward requires sampling_params.extra_args['robot_obs'].", + ) + + images, image_masks, lang_tokens, lang_masks, state = build_model_inputs( + robot_obs, self.config, self.tokenizer, self._device + ) + + # State normalization (identity for pi0_base) inside the model. + state = self.model._normalize_state(state) + + noise = extra_args.get("noise") + if noise is not None and not isinstance(noise, torch.Tensor): + noise = torch.as_tensor(noise, dtype=torch.float32, device=self._device) + elif isinstance(noise, torch.Tensor): + noise = noise.to(device=self._device, dtype=torch.float32) + + num_steps = extra_args.get("num_inference_steps") + + actions = self.model.sample_actions( + images=images, + image_masks=image_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + state=state, + noise=noise, + num_steps=num_steps, + ) + actions = self.model._unnormalize_actions(actions) + + # (B=1, horizon, action_dim) → (horizon, action_dim) numpy for the wire. + actions_np = actions.squeeze(0).float().cpu().numpy() + + # Note: post-processing is applied engine-side via the registry + # (_DIFFUSION_POST_PROCESS_FUNCS), so we don't attach it here (a local + # closure wouldn't survive the orchestrator's multiprocess pickling). + return DiffusionOutput(output={"actions": actions_np}) diff --git a/vllm_omni/diffusion/models/pi0/processor_pi0.py b/vllm_omni/diffusion/models/pi0/processor_pi0.py new file mode 100644 index 000000000..0133e5729 --- /dev/null +++ b/vllm_omni/diffusion/models/pi0/processor_pi0.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Preprocessing for the π0 VLA model. + +Converts a raw robot observation (multi-camera images + language instruction + +proprioceptive state) into the tensors that ``Pi0ForActionPrediction.sample_actions`` +consumes: + + * images : list of ``(1, 3, 224, 224)`` float tensors in ``[-1, 1]``, + one per camera slot (real cameras first, then ``-1``-filled + padding) in the checkpoint's ``image_feature_keys`` order + * image_masks : list of ``(1,)`` bool tensors, ``True`` = real camera + * lang_tokens : ``(1, L)`` long, PaliGemma tokenizer, ``\\n``-terminated + * lang_masks : ``(1, L)`` bool + * state : ``(1, max_state_dim)`` float32, zero-padded + +Image/prompt preprocessing matches OpenPI and LeRobot bit-for-bit — see the +parity test (``test_pi0_parity.py``). A set of stateless helpers the diffusion +pipeline calls directly (the DreamZero contract: the pipeline owns its +preprocessing). + +Reference: + - OpenPI: openpi/src/openpi/shared/image_tools.py (resize_with_pad) + - OpenPI: openpi/src/openpi/models_pytorch/preprocessing_pytorch.py + - LeRobot: lerobot/src/lerobot/policies/pi0/modeling_pi0.py (Pi0NewLineProcessor) +""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +logger = logging.getLogger(__name__) + +# Defaults straight from the π0 reference configs. +PI0_IMAGE_SIZE = 224 # openpi/models/model.py IMAGE_RESOLUTION +PI0_NUM_IMAGE_TOKENS = 256 # SigLIP So400m/14 on 224×224 → (224/14)**2 +PI0_IMAGE_TOKEN_INDEX = 257152 # openpi/models_pytorch/gemma_pytorch.py +PI0_MAX_CAMERAS = 3 # openpi/models/model.py (3 camera slots) +PI0_MAX_TOKEN_LEN = 48 # openpi/models/pi0_config.py + + +def resize_with_pad( + images: torch.Tensor, + target_height: int, + target_width: int, + mode: str = "bilinear", +) -> torch.Tensor: + """Resize ``(B, C, H, W)`` images to the target shape, preserving aspect + ratio with -1 padding on the short side. + + Matches openpi ``image_tools.resize_with_pad_torch`` — the clamp to + [-1, 1] is what lets the padded region blend with SigLIP-normalized + pixels without adding signal at the boundary. + """ + if images.ndim != 4: + raise ValueError(f"Expected 4-D (B,C,H,W), got {images.ndim}-D") + _, _, cur_h, cur_w = images.shape + ratio = max(cur_w / target_width, cur_h / target_height) + rh, rw = int(cur_h / ratio), int(cur_w / ratio) + align_corners = False if mode == "bilinear" else None + resized = F.interpolate(images, size=(rh, rw), mode=mode, align_corners=align_corners) + resized = resized.clamp(-1.0, 1.0) + ph, rem_h = divmod(target_height - rh, 2) + pw, rem_w = divmod(target_width - rw, 2) + return F.pad(resized, (pw, pw + rem_w, ph, ph + rem_h), value=-1.0) + + +def pil_image_to_tensor(image: Image.Image) -> torch.Tensor: + """PIL → ``(1, C, H, W)`` float32 in ``[-1, 1]`` (SigLIP normalization).""" + if image.mode != "RGB": + image = image.convert("RGB") + arr = np.array(image, dtype=np.float32) / 255.0 * 2.0 - 1.0 + return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) + + +class Pi0ImageProcessor: + """Minimal image preprocessor: image → normalized + padded ``[-1,1]`` tensor.""" + + def __init__(self, image_size: int = PI0_IMAGE_SIZE): + self.image_size = image_size + + def preprocess_single(self, image: Any) -> torch.Tensor: + """Accept a PIL image, an HWC uint8/float ndarray, or a CHW tensor and + return a ``(1, 3, image_size, image_size)`` float tensor in ``[-1, 1]``.""" + t = self._to_tensor(image) + if t.shape[2] != self.image_size or t.shape[3] != self.image_size: + t = resize_with_pad(t, self.image_size, self.image_size) + return t + + def _to_tensor(self, image: Any) -> torch.Tensor: + if isinstance(image, Image.Image): + return pil_image_to_tensor(image) + if isinstance(image, np.ndarray): + arr = image + # HWC → normalize uint8/[0,255] to [-1,1]; assume already [-1,1] if float. + if arr.ndim == 3 and arr.shape[-1] in (1, 3): + if np.issubdtype(arr.dtype, np.integer) or arr.max() > 1.0: + arr = arr.astype(np.float32) / 255.0 * 2.0 - 1.0 + else: + arr = arr.astype(np.float32) + return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) + return torch.as_tensor(arr, dtype=torch.float32) + if isinstance(image, torch.Tensor): + t = image + if t.ndim == 3: # (C, H, W) → (1, C, H, W) + t = t.unsqueeze(0) + return t.to(dtype=torch.float32) + raise TypeError(f"Unsupported image type for π0 preprocessing: {type(image)}") + + def make_empty_image(self) -> torch.Tensor: + """Fill tensor for an unused camera slot — pure -1, matches OpenPI/LeRobot.""" + return torch.full((1, 3, self.image_size, self.image_size), -1.0) + + +def tokenize_prompt(tokenizer, text: str, max_token_len: int = PI0_MAX_TOKEN_LEN): + """Return ``(input_ids, attention_mask)`` lists, length ``max_token_len``. + + Matches LeRobot's ``Pi0NewLineProcessor``: append ``\\n`` if the caller + didn't, then run the PaliGemma tokenizer with right-padding. + """ + prompt = text or "" + if not prompt.endswith("\n"): + prompt = prompt + "\n" + enc = tokenizer( + prompt, + padding="max_length", + max_length=max_token_len, + truncation=True, + add_special_tokens=True, + return_tensors=None, + ) + return list(enc["input_ids"]), list(enc["attention_mask"]) + + +def _extract_images(robot_obs: dict, config) -> dict[str, Any]: + """Pull a ``{feature_key: image}`` map out of a raw robot obs. + + Accepts two shapes: + * ``robot_obs["images"]`` is a dict keyed by camera name; or + * flat ``observation/`` style keys on ``robot_obs`` itself. + Keys are translated through ``config.image_key_map`` (raw obs key → + checkpoint feature key) so serving wire names map onto ``input_features``. + """ + images = robot_obs.get("images") + if not isinstance(images, dict): + # Fall back to flat keys (anything that isn't a known scalar field). + images = {k: v for k, v in robot_obs.items() if k not in ("state", "prompt", "session_id", "reset", "images")} + key_map = getattr(config, "image_key_map", None) or {} + return {key_map.get(k, k): v for k, v in images.items()} + + +def build_model_inputs(robot_obs: dict, config, tokenizer, device: torch.device): + """Convert a raw robot observation into ``sample_actions`` inputs. + + Reproduces the LeRobot camera ordering exactly: iterate + ``config.image_feature_keys`` (the ordered camera identities from the + checkpoint's ``input_features``); for each, use the supplied image (mask + ``True``) or a ``-1``-filled empty image (mask ``False``). Then tokenize + the prompt and zero-pad/truncate the state to ``max_state_dim``. + + Returns ``(images, image_masks, lang_tokens, lang_masks, state)`` with a + leading batch dim of 1 on the tensors. + """ + image_size = int(config.image_resolution[0]) + img_proc = Pi0ImageProcessor(image_size=image_size) + + feature_keys = config.image_feature_keys or [] + if not feature_keys: + # No declared camera order — fall back to whatever the obs provides, + # capped at max_cameras (preserves insertion order). + provided = _extract_images(robot_obs, config) + feature_keys = list(provided.keys())[: config.max_cameras] + + obs_images = _extract_images(robot_obs, config) + + images: list[torch.Tensor] = [] + image_masks: list[torch.Tensor] = [] + for key in feature_keys[: config.max_cameras]: + img = obs_images.get(key) + if img is not None: + tensor = img_proc.preprocess_single(img).to(device=device) + mask = True + else: + tensor = img_proc.make_empty_image().to(device=device) + mask = False + images.append(tensor) + image_masks.append(torch.tensor([mask], dtype=torch.bool, device=device)) + + # Pad up to max_cameras with empty slots if the checkpoint declares fewer + # cameras than the model attends to. + while len(images) < config.max_cameras: + images.append(img_proc.make_empty_image().to(device=device)) + image_masks.append(torch.tensor([False], dtype=torch.bool, device=device)) + + prompt = robot_obs.get("prompt", "") or "" + ids, attn = tokenize_prompt(tokenizer, prompt, config.tokenizer_max_length) + lang_tokens = torch.tensor([ids], dtype=torch.long, device=device) + lang_masks = torch.tensor([attn], dtype=torch.bool, device=device) + + state = _build_state(robot_obs.get("state"), config.max_state_dim, device) + + return images, image_masks, lang_tokens, lang_masks, state + + +def _build_state(raw_state, max_state_dim: int, device: torch.device) -> torch.Tensor: + """Zero-pad / truncate the raw state to ``(1, max_state_dim)`` float32.""" + if raw_state is None: + return torch.zeros(1, max_state_dim, dtype=torch.float32, device=device) + state = torch.as_tensor(raw_state, dtype=torch.float32, device=device) + if state.dim() == 1: + state = state[None, :] + if state.shape[-1] < max_state_dim: + state = F.pad(state, (0, max_state_dim - state.shape[-1])) + elif state.shape[-1] > max_state_dim: + state = state[..., :max_state_dim] + return state diff --git a/vllm_omni/diffusion/models/pi0_pipeline_config.py b/vllm_omni/diffusion/models/pi0_pipeline_config.py new file mode 100644 index 000000000..d778ca81a --- /dev/null +++ b/vllm_omni/diffusion/models/pi0_pipeline_config.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lightweight pipeline topology for π0 (Pi-Zero).""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +# π0 is one diffusion stage (robot observation -> action chunk via flow +# matching). It is registered in ``OMNI_PIPELINES`` because π0 is online-served +# via ``vllm serve --deploy-config pi0.yaml``. Keep this topology outside the +# ``pi0`` package: importing a package submodule executes ``pi0.__init__``, +# which imports the runtime pipeline and ``diffusion.data``. +PI0_PIPELINE = PipelineConfig( + model_type="pi0", + model_arch="Pi0Pipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="action", + model_arch="Pi0Pipeline", + ), + ), +) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 70e0ef0f1..0aa00abd9 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import copy import inspect diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 67e2d0db7..c82ace9c4 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -642,10 +642,10 @@ class QwenImageCrossAttention(nn.Module): txt_query = self.norm_added_q(txt_query) txt_key = self.norm_added_k(txt_key) - img_cos = vid_freqs.real.to(img_query.dtype) - img_sin = vid_freqs.imag.to(img_query.dtype) - txt_cos = txt_freqs.real.to(txt_query.dtype) - txt_sin = txt_freqs.imag.to(txt_query.dtype) + img_cos = torch.real(vid_freqs).to(img_query.dtype) + img_sin = torch.imag(vid_freqs).to(img_query.dtype) + txt_cos = torch.real(txt_freqs).to(txt_query.dtype) + txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) img_query = self.rope(img_query, img_cos, img_sin) img_key = self.rope(img_key, img_cos, img_sin) diff --git a/vllm_omni/diffusion/models/sana_video/__init__.py b/vllm_omni/diffusion/models/sana_video/__init__.py new file mode 100644 index 000000000..cebefc4f8 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/__init__.py @@ -0,0 +1,32 @@ +# Copyright 2025 SANA-Video Authors and The HuggingFace Team. All rights reserved. +# +# 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. + +from .pipeline_output import SanaVideoPipelineOutput +from .pipeline_sana_video import SanaVideoPipeline, get_sana_video_post_process_func +from .pipeline_sana_video_i2v import ( + SanaImageToVideoPipeline, + get_sana_video_i2v_post_process_func, + get_sana_video_i2v_pre_process_func, +) +from .transformer_sana_video import SanaVideoTransformer3DModel + +__all__ = [ + "SanaImageToVideoPipeline", + "SanaVideoPipeline", + "SanaVideoPipelineOutput", + "SanaVideoTransformer3DModel", + "get_sana_video_i2v_post_process_func", + "get_sana_video_i2v_pre_process_func", + "get_sana_video_post_process_func", +] diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_output.py b/vllm_omni/diffusion/models/sana_video/pipeline_output.py new file mode 100644 index 000000000..6494b112c --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/pipeline_output.py @@ -0,0 +1,19 @@ +from dataclasses import dataclass + +import torch +from diffusers.utils import BaseOutput + + +@dataclass +class SanaVideoPipelineOutput(BaseOutput): + r""" + Output class for Sana-Video pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + """ + + frames: torch.Tensor diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py new file mode 100644 index 000000000..a1fa38d43 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video.py @@ -0,0 +1,1055 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# Copyright 2025 SANA-Video Authors and The HuggingFace Team. All rights reserved. +# +# 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. + +# ruff: noqa: E501 + +import inspect +import os +import warnings +from collections.abc import Iterable + +import numpy as np +import PIL.Image +import torch +from diffusers.schedulers import DPMSolverMultistepScheduler +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from torch import nn +from transformers import AutoModel, Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast +from vllm.transformers_utils.config import get_hf_file_to_dict + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, +) +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.utils import _load_json +from vllm_omni.diffusion.postprocess import interpolate_video_tensor +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import resolve_video_num_frames +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +from .pipeline_output import SanaVideoPipelineOutput +from .transformer_sana_video import SanaVideoTransformer3DModel + + +def _resolve_vae_class_and_dtype( + vae_class_name: str, + pipeline_dtype: torch.dtype, +) -> tuple[type[DistributedAutoencoderKLLTX2Video | DistributedAutoencoderKLWan], torch.dtype]: + vae_classes = { + "AutoencoderKLLTX2Video": DistributedAutoencoderKLLTX2Video, + "AutoencoderKLWan": DistributedAutoencoderKLWan, + } + if vae_class_name not in vae_classes: + raise ValueError(f"Unsupported SANA-Video VAE class: {vae_class_name!r}") + vae_dtype = torch.float32 if vae_class_name == "AutoencoderKLWan" else pipeline_dtype + return vae_classes[vae_class_name], vae_dtype + + +def _load_sana_tokenizer( + model: str, + component_subfolders: list[str], + local_files_only: bool, +) -> GemmaTokenizer: + # Transformers v5 AutoTokenizer first tries to resolve an AutoConfig from + # the requested subfolder. Released SANA-Video repositories correctly + # contain tokenizer metadata under tokenizer/, but their model config lives + # under text_encoder/, so that extra tokenizer/config.json lookup fails for + # Hub model IDs. Both released checkpoints use GemmaTokenizer; loading the + # concrete class also matches Diffusers and avoids the unrelated config + # lookup. + return from_pretrained_with_prefetch( + GemmaTokenizer.from_pretrained, + model, + subfolder="tokenizer", + prefetch_list=component_subfolders, + local_files_only=local_files_only, + ) + + +ASPECT_RATIO_480_BIN = { + "0.5": [448.0, 896.0], + "0.57": [480.0, 832.0], + "0.68": [528.0, 768.0], + "0.78": [560.0, 720.0], + "1.0": [624.0, 624.0], + "1.13": [672.0, 592.0], + "1.29": [720.0, 560.0], + "1.46": [768.0, 528.0], + "1.67": [816.0, 496.0], + "1.75": [832.0, 480.0], + "2.0": [896.0, 448.0], +} + + +ASPECT_RATIO_720_BIN = { + "0.5": [672.0, 1344.0], + "0.57": [704.0, 1280.0], + "0.68": [800.0, 1152.0], + "0.78": [832.0, 1088.0], + "1.0": [960.0, 960.0], + "1.13": [1024.0, 896.0], + "1.29": [1088.0, 832.0], + "1.46": [1152.0, 800.0], + "1.67": [1248.0, 736.0], + "1.75": [1280.0, 704.0], + "2.0": [1344.0, 672.0], +} + + +SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION = ( + "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for video generation. Evaluate the level of detail in the user prompt:", + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and temporal relationships to create vivid and dynamic scenes.", + "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.", + "Here are examples of how to transform or refine prompts:", + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat slowly settling into a curled position, peacefully falling asleep on a warm sunny windowsill, with gentle sunlight filtering through surrounding pots of blooming red flowers.", + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps gradually lighting up, a diverse crowd of people in colorful clothing walking past, and a double-decker bus smoothly passing by towering glass skyscrapers.", + "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:", + "User Prompt: ", +) + + +def get_sana_video_default_resolution(sample_size: int) -> tuple[int, int]: + if sample_size == 30: + return 480, 832 + if sample_size == 22: + return 704, 1280 + raise ValueError(f"Unsupported SANA-Video sample size: {sample_size}") + + +def resolve_sana_video_sample_size(od_config: OmniDiffusionConfig) -> int: + tf_model_config = getattr(od_config, "tf_model_config", None) + sample_size = tf_model_config.get("sample_size", None) if tf_model_config is not None else None + model = getattr(od_config, "model", None) + if sample_size is None and model: + try: + transformer_config = get_hf_file_to_dict("transformer/config.json", model) or {} + sample_size = transformer_config.get("sample_size") + except (OSError, ValueError): + pass + if sample_size is None: + sample_size = 22 if model and "720p" in str(model).lower() else 30 + return int(sample_size) + + +def get_sana_video_post_process_func(od_config: OmniDiffusionConfig): + video_processor = VideoProcessor() + is_diffusers_adapter = od_config.diffusion_load_format == "diffusers" + + def post_process_func(video: torch.Tensor, output_type: str = "np", sampling_params=None): + if output_type == "latent" or ( + sampling_params is not None and getattr(sampling_params, "output_type", None) == "latent" + ): + return video + + first_frame = video + while isinstance(first_frame, (list, tuple)) and first_frame: + first_frame = first_frame[0] + is_postprocessed_frame_sequence = isinstance(video, np.ndarray) or isinstance( + first_frame, + (PIL.Image.Image, np.ndarray), + ) + if is_diffusers_adapter and is_postprocessed_frame_sequence: + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + raise ValueError( + "Frame interpolation is not supported for SANA-Video with the Diffusers adapter because " + "Diffusers returns already postprocessed frames." + ) + return {"payload": {"video": video}, "metadata": {}} + + video_metadata = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + video_metadata["video_fps_multiplier"] = multiplier + return { + "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, + "metadata": {"video": video_metadata} if video_metadata else {}, + } + + return post_process_func + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +) -> tuple[torch.Tensor, int]: + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + scheduler_device = torch.device("cpu") + # vLLM initializes models under a CUDA default-device context. Diffusers' + # schedulers are not nn.Modules, so their tensors can be created on CUDA and + # are not moved by pipeline.to(). DPM solvers also perform NumPy operations + # while constructing their schedule. Move scheduler state and temporarily + # override the global default device so every intermediate remains on CPU. + for name, value in vars(scheduler).items(): + if isinstance(value, torch.Tensor) and value.device.type != "cpu": + setattr(scheduler, name, value.cpu()) + previous_default_device = torch.get_default_device() + try: + torch.set_default_device(scheduler_device) + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=scheduler_device, **kwargs) + num_inference_steps = len(scheduler.timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=scheduler_device, **kwargs) + num_inference_steps = len(scheduler.timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=scheduler_device, **kwargs) + finally: + torch.set_default_device(previous_default_device) + scheduler.timesteps = scheduler.timesteps.to(device) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class SanaVideoPipeline( + nn.Module, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, +): + r""" + vLLM-Omni pipeline for text-to-video generation using + [Sana](https://huggingface.co/papers/2509.24695). + + Args: + tokenizer ([`GemmaTokenizer`] or [`GemmaTokenizerFast`]): + The tokenizer used to tokenize the prompt. + text_encoder ([`Gemma2PreTrainedModel`]): + Text encoder model to encode the input prompts. + vae ([`DistributedAutoencoderKLWan`] or [`DistributedAutoencoderKLLTX2Video`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + transformer ([`SanaVideoTransformer3DModel`]): + Conditional Transformer to denoise the input latents. + scheduler ([`DPMSolverMultistepScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded video latents. + """ + + _dit_modules = ["transformer"] + _encoder_modules = ["text_encoder"] + _vae_modules = ["vae"] + supports_step_execution = False + + def __init__( + self, + tokenizer: GemmaTokenizer | GemmaTokenizerFast | None = None, + text_encoder: Gemma2PreTrainedModel | None = None, + vae: DistributedAutoencoderKLLTX2Video | DistributedAutoencoderKLWan | None = None, + transformer: SanaVideoTransformer3DModel | None = None, + scheduler: DPMSolverMultistepScheduler | None = None, + *, + od_config: OmniDiffusionConfig | None = None, + prefix: str = "", + ): + super().__init__() + + self.od_config = od_config + self.device = get_local_device() + self.weights_sources: list[DiffusersPipelineLoader.ComponentSource] = [] + if od_config is not None: + tokenizer, text_encoder, vae, transformer, scheduler = self._load_components(od_config, prefix) + + if tokenizer is None or text_encoder is None or vae is None or transformer is None or scheduler is None: + raise ValueError("SanaVideoPipeline requires tokenizer, text_encoder, vae, transformer, and scheduler.") + + self.tokenizer = tokenizer + self.text_encoder = text_encoder + self.vae = vae + self.transformer = transformer + self.scheduler = scheduler + if od_config is None: + first_parameter = next(self.transformer.parameters(), None) + if first_parameter is not None: + self.device = first_parameter.device + + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + self.vae_scale_factor_temporal = self.vae.config.temporal_compression_ratio + self.vae_scale_factor_spatial = self.vae.config.spatial_compression_ratio + elif isinstance(self.vae, DistributedAutoencoderKLWan): + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial + else: + raise TypeError(f"Unsupported SANA-Video VAE type: {type(self.vae).__name__}") + + self.vae_scale_factor = self.vae_scale_factor_spatial + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + if od_config is not None: + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, + profiler_targets=["forward"], + ) + + def _load_components( + self, + od_config: OmniDiffusionConfig, + prefix: str, + ) -> tuple[ + GemmaTokenizer | GemmaTokenizerFast, + Gemma2PreTrainedModel, + DistributedAutoencoderKLLTX2Video | DistributedAutoencoderKLWan, + SanaVideoTransformer3DModel, + DPMSolverMultistepScheduler, + ]: + model = od_config.model + local_files_only = os.path.exists(model) + dtype = getattr(od_config, "dtype", torch.bfloat16) + device = self.device + # Transformer weights are streamed by DiffusersPipelineLoader below. + # Prefetch only components loaded through from_pretrained here. + component_subfolders = ["tokenizer", "text_encoder", "vae", "scheduler"] + prefetch_subfolders(model, component_subfolders, local_files_only=local_files_only) + + tokenizer = _load_sana_tokenizer( + model, + component_subfolders, + local_files_only=local_files_only, + ) + text_encoder = from_pretrained_with_prefetch( + AutoModel.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=component_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ).to(device) + + model_index = _load_json(model, "model_index.json", local_files_only) + vae_class_name = model_index.get("vae", [None, "AutoencoderKLWan"])[1] + # The released 480p checkpoint uses Wan VAE weights that are decoded + # in FP32 in the upstream reference recipe. The 720p LTX2 VAE is + # validated in BF16. + vae_class, vae_dtype = _resolve_vae_class_and_dtype(vae_class_name, dtype) + vae = from_pretrained_with_prefetch( + vae_class.from_pretrained, + model, + subfolder="vae", + prefetch_list=component_subfolders, + local_files_only=local_files_only, + torch_dtype=vae_dtype, + ).to(device) + + transformer_config = _load_json(model, "transformer/config.json", local_files_only) + transformer = SanaVideoTransformer3DModel.from_config(transformer_config).to(dtype=dtype, device=device) + scheduler = DPMSolverMultistepScheduler.from_pretrained( + model, + subfolder="scheduler", + local_files_only=local_files_only, + ) + self.weights_sources[:] = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="transformer", + revision=None, + prefix=f"{prefix}transformer." if prefix else "transformer.", + fall_back_to_pt=True, + ) + ] + return tokenizer, text_encoder, vae, transformer, scheduler + + # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + device: torch.device, + dtype: torch.dtype, + clean_caption: bool = False, + max_sequence_length: int = 300, + complex_human_instruction: list[str] | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. + complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): + If `complex_human_instruction` is not empty, the function will use the complex Human instruction for + the prompt. + """ + prompt = [prompt] if isinstance(prompt, str) else prompt + + if getattr(self, "tokenizer", None) is not None: + self.tokenizer.padding_side = "right" + + prompt = self._text_preprocessing(prompt, clean_caption=clean_caption) + + # prepare complex human instruction + if not complex_human_instruction: + max_length_all = max_sequence_length + else: + chi_prompt = "\n".join(complex_human_instruction) + prompt = [chi_prompt + p for p in prompt] + num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt)) + max_length_all = num_chi_prompt_tokens + max_sequence_length - 2 + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_length_all, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + + prompt_attention_mask = text_inputs.attention_mask + prompt_attention_mask = prompt_attention_mask.to(device) + + prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask) + prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device) + + return prompt_embeds, prompt_attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + do_classifier_free_guidance: bool = True, + negative_prompt: str = "", + num_videos_per_prompt: int = 1, + device: torch.device | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + clean_caption: bool = False, + max_sequence_length: int = 300, + complex_human_instruction: list[str] | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + negative_prompt (`str` or `list[str]`, *optional*): + The prompt not to guide the video generation. If not defined, one has to pass `negative_prompt_embeds` + instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For + PixArt-Alpha, this should be "". + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): + whether to use classifier free guidance or not + num_videos_per_prompt (`int`, *optional*, defaults to 1): + number of videos that should be generated per prompt + device: (`torch.device`, *optional*): + torch device to place the resulting embeddings on + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. For Sana, it's should be the embeddings of the "" string. + clean_caption (`bool`, defaults to `False`): + If `True`, the function will preprocess and clean the provided caption before encoding. + max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt. + complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`): + If `complex_human_instruction` is not empty, the function will use the complex Human instruction for + the prompt. + """ + + if device is None: + device = self.device + + if self.text_encoder is not None: + dtype = self.text_encoder.dtype + else: + dtype = None + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if getattr(self, "tokenizer", None) is not None: + self.tokenizer.padding_side = "right" + + # See Section 3.1. of the paper. + max_length = max_sequence_length + select_index = [0] + list(range(-max_length + 1, 0)) + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + device=device, + dtype=dtype, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + ) + + prompt_embeds = prompt_embeds[:, select_index] + prompt_attention_mask = prompt_attention_mask[:, select_index] + + bs_embed, seq_len, _ = prompt_embeds.shape + # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(bs_embed * num_videos_per_prompt, seq_len, -1) + prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + # get unconditional embeddings for classifier free guidance + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = [negative_prompt] * batch_size if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + device=device, + dtype=dtype, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=False, + ) + + if do_classifier_free_guidance: + # duplicate unconditional embeddings for each generation per prompt, using mps friendly method + seq_len = negative_prompt_embeds.shape[1] + + negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) + + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1) + negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_videos_per_prompt, 1) + else: + negative_prompt_embeds = None + negative_prompt_attention_mask = None + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs + def prepare_extra_step_kwargs(self, generator, eta): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502 + # and should be between [0, 1] + + accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) + extra_step_kwargs = {} + if accepts_eta: + extra_step_kwargs["eta"] = eta + + # check if the scheduler accepts generator + accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys()) + if accepts_generator: + extra_step_kwargs["generator"] = generator + return extra_step_kwargs + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + patch_size = getattr(self.transformer.config, "patch_size", (1, 1, 1)) + spatial_alignment = self.vae_scale_factor_spatial * max(patch_size[-2:]) + if height % spatial_alignment != 0 or width % spatial_alignment != 0: + raise ValueError( + f"`height` and `width` have to be divisible by {spatial_alignment} but are {height} and {width}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + def _text_preprocessing(self, text, clean_caption=False): + if clean_caption: + raise ValueError( + "SANA-Video native pipeline does not support `clean_caption=True`; " + "pre-clean the prompt or use the Diffusers adapter." + ) + if not isinstance(text, (tuple, list)): + text = [text] + return [str(item).lower().strip() for item in text] + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int = 16, + height: int = 480, + width: int = 832, + num_frames: int = 81, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + int(height) // self.vae_scale_factor_spatial, + int(width) // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + transformer_weights = ( + (name.removeprefix("transformer."), tensor) for name, tensor in weights if name.startswith("transformer.") + ) + return {f"transformer.{name}" for name in self.transformer.load_weights(transformer_weights)} + + def __call__(self, req: DiffusionRequestBatch) -> DiffusionOutput: + return self.forward(req) + + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + if len(req.prompts) != 1: + raise ValueError("SANA-Video currently supports exactly one prompt per request.") + prompt_obj = req.prompts[0] + prompt = prompt_obj if isinstance(prompt_obj, str) else (prompt_obj.get("prompt") or "") + negative_prompt = "" if isinstance(prompt_obj, str) else (prompt_obj.get("negative_prompt") or "") + if not prompt: + raise ValueError("Prompt is required for SANA-Video generation.") + + sampling = req.sampling_params + default_height, default_width = get_sana_video_default_resolution(self.transformer.config.sample_size) + height = sampling.height or default_height + width = sampling.width or default_width + num_frames = resolve_video_num_frames( + sampling.num_frames, + default_num_frames=81, + is_dummy_run=req.is_dummy_run(), + ) + num_steps = sampling.num_inference_steps if sampling.num_inference_steps is not None else 50 + guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 6.0 + generator = sampling.generator + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=get_local_device()).manual_seed(sampling.seed) + + extra_args = sampling.extra_args or {} + motion_score = extra_args.get("motion_score") + if motion_score is not None: + prompt = f"{prompt} motion score: {int(motion_score)}." + + result = self._generate( + prompt=prompt, + negative_prompt=negative_prompt, + num_inference_steps=num_steps, + timesteps=sampling.timesteps, + sigmas=sampling.sigmas, + guidance_scale=guidance_scale, + num_videos_per_prompt=sampling.num_outputs_per_prompt or 1, + height=height, + width=width, + frames=num_frames, + eta=sampling.eta, + generator=generator, + latents=sampling.latents, + output_type="latent" if sampling.output_type == "latent" else "raw", + clean_caption=bool(extra_args.get("clean_caption", False)), + use_resolution_binning=bool(extra_args.get("use_resolution_binning", True)), + max_sequence_length=sampling.max_sequence_length or 300, + ) + return DiffusionOutput(output=result.frames) + + @torch.no_grad() + def _generate( + self, + prompt: str | list[str] = None, + negative_prompt: str = "", + num_inference_steps: int = 50, + timesteps: list[int] = None, + sigmas: list[float] = None, + guidance_scale: float = 6.0, + num_videos_per_prompt: int | None = 1, + height: int = 480, + width: int = 832, + frames: int = 81, + eta: float = 0.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + clean_caption: bool = False, + use_resolution_binning: bool = True, + max_sequence_length: int = 300, + complex_human_instruction: list[str] | None = None, + ) -> SanaVideoPipelineOutput | tuple: + """ + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. + instead. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the video generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality video at the + expense of slower inference. + timesteps (`list[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`list[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 4.5): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. + of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate videos that are closely linked to + the text `prompt`, usually at the expense of lower video quality. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + height (`int`, *optional*, defaults to 480): + The height in pixels of the generated video. + width (`int`, *optional*, defaults to 832): + The width in pixels of the generated video. + frames (`int`, *optional*, defaults to 81): + The number of frames in the generated video. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only + applies to [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will be generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for text embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not + provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + negative_prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated video. Choose between mp4 or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`SanaVideoPipelineOutput`] instead of a plain tuple. + clean_caption (`bool`, *optional*, defaults to `True`): + Advanced caption cleaning is not supported by the native pipeline. Passing `True` raises an error; + pre-clean the prompt or use the Diffusers adapter. + use_resolution_binning (`bool` defaults to `True`): + If set to `True`, the requested height and width are first mapped to the closest resolutions using + `ASPECT_RATIO_480_BIN` or `ASPECT_RATIO_720_BIN`. After the produced latents are decoded into videos, + they are resized back to the requested resolution. Useful for generating non-square videos. + max_sequence_length (`int` defaults to `300`): + Maximum sequence length to use with the `prompt`. + complex_human_instruction (`list[str]`, *optional*): + Instructions for complex human attention: + https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55. + + Examples: + + Returns: + [`~pipelines.sana_video.pipeline_output.SanaVideoPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.sana_video.pipeline_output.SanaVideoPipelineOutput`] is + returned, otherwise a `tuple` is returned where the first element is a list with the generated videos + """ + if complex_human_instruction is None: + complex_human_instruction = list(SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION) + + # 1. Check inputs. Raise error if not correct + if use_resolution_binning: + if self.transformer.config.sample_size == 30: + aspect_ratio_bin = ASPECT_RATIO_480_BIN + elif self.transformer.config.sample_size == 22: + aspect_ratio_bin = ASPECT_RATIO_720_BIN + else: + raise ValueError("Invalid sample size") + orig_height, orig_width = height, width + height, width = self.video_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin) + + self.check_inputs( + prompt=prompt, + height=height, + width=width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + + self._guidance_scale = guidance_scale + self._interrupt = False + + # 2. Default height and width to transformer + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self.device + + # 3. Encode input prompt + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt, + self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + device=device, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas + ) + + # 5. Prepare latents. + latent_channels = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + latent_channels, + height, + width, + frames, + torch.float32, + device, + generator, + latents, + ) + + # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + self._num_timesteps = len(timesteps) + + transformer_dtype = self.transformer.dtype + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timestep = t.expand(latent_model_input.shape[0]) + + # predict noise model_output + noise_pred = self.transformer( + latent_model_input.to(dtype=transformer_dtype), + encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype), + encoder_attention_mask=prompt_attention_mask, + timestep=timestep, + return_dict=False, + )[0] + noise_pred = noise_pred.float() + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # learned sigma + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + + # compute previous image: x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if output_type == "latent": + video = latents + else: + latents = latents.to(self.vae.dtype) + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + latents_mean = self.vae.latents_mean + latents_std = self.vae.latents_std + z_dim = self.vae.config.latent_channels + elif isinstance(self.vae, DistributedAutoencoderKLWan): + latents_mean = torch.tensor(self.vae.config.latents_mean) + latents_std = torch.tensor(self.vae.config.latents_std) + z_dim = self.vae.config.z_dim + else: + latents_mean = torch.zeros(latents.shape[1], device=latents.device, dtype=latents.dtype) + latents_std = torch.ones(latents.shape[1], device=latents.device, dtype=latents.dtype) + z_dim = latents.shape[1] + + latents_mean = latents_mean.view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = 1.0 / latents_std.view(1, z_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents / latents_std + latents_mean + try: + video = self.vae.decode(latents, return_dict=False)[0] + except torch.OutOfMemoryError as e: + warnings.warn( + f"{e}. \n" + f"Try to use VAE tiling for large images. For example: \n" + f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)" + ) + raise + + if use_resolution_binning: + video = self.video_processor.resize_and_crop_tensor(video, orig_width, orig_height) + + if output_type != "raw": + video = self.video_processor.postprocess_video(video, output_type=output_type) + + if not return_dict: + return (video,) + + return SanaVideoPipelineOutput(frames=video) diff --git a/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py new file mode 100644 index 000000000..8d0fe9919 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/pipeline_sana_video_i2v.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import os +import warnings +from math import sqrt + +import PIL.Image +import torch +from diffusers.utils.torch_utils import randn_tensor + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, +) +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan +from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import ( + ASPECT_RATIO_480_BIN, + ASPECT_RATIO_720_BIN, + SanaVideoPipeline, + get_sana_video_default_resolution, + get_sana_video_post_process_func, + resolve_sana_video_sample_size, + retrieve_timesteps, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest, resolve_video_num_frames +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniTextPrompt + +SANA_VIDEO_I2V_COMPLEX_HUMAN_INSTRUCTION = [ + ( + "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for " + "video generation. Evaluate the level of detail in the user prompt:" + ), + ( + "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, motion, and " + "temporal relationships to create vivid and dynamic scenes." + ), + ("- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating."), + "Here are examples of how to transform or refine prompts:", + ( + "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat slowly settling into a curled position, " + "peacefully falling asleep on a warm sunny windowsill, with gentle sunlight filtering through surrounding " + "pots of blooming red flowers." + ), + ( + "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street " + "lamps gradually lighting up, a diverse crowd of people in colorful clothing walking past, and a double-decker " + "bus smoothly passing by towering glass skyscrapers." + ), + ( + "Please generate only the enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:" + ), + "User Prompt: ", +] + + +def get_sana_video_i2v_post_process_func(od_config: OmniDiffusionConfig): + return get_sana_video_post_process_func(od_config) + + +def get_sana_video_i2v_pre_process_func(od_config: OmniDiffusionConfig): + default_height, default_width = get_sana_video_default_resolution(resolve_sana_video_sample_size(od_config)) + + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + prompt = request.prompt + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt, multi_modal_data={}) + + multi_modal_data = prompt.get("multi_modal_data") or {} + raw_image = multi_modal_data.get("image") + if raw_image is None: + raise ValueError('SANA-Video I2V requires `"multi_modal_data": {"image": }`.') + if isinstance(raw_image, str): + if not os.path.isfile(raw_image): + raise ValueError(f"SANA-Video I2V image path does not exist: {raw_image}") + image = PIL.Image.open(raw_image).convert("RGB") + elif isinstance(raw_image, PIL.Image.Image): + image = raw_image.convert("RGB") + else: + raise TypeError(f"Unsupported SANA-Video I2V image type: {type(raw_image).__name__}") + + if request.sampling_params.height is None or request.sampling_params.width is None: + max_area = default_height * default_width + aspect_ratio = image.height / image.width + inferred_height = round(sqrt(max_area * aspect_ratio)) // 32 * 32 + inferred_width = round(sqrt(max_area / aspect_ratio)) // 32 * 32 + if request.sampling_params.height is None: + request.sampling_params.height = inferred_height + if request.sampling_params.width is None: + request.sampling_params.width = inferred_width + + image = image.resize( + (request.sampling_params.width, request.sampling_params.height), + PIL.Image.Resampling.LANCZOS, + ) + multi_modal_data["image"] = image + prompt["multi_modal_data"] = multi_modal_data + request.prompt = prompt + return request + + return pre_process_func + + +def _retrieve_latents(encoder_output: torch.Tensor) -> torch.Tensor: + if hasattr(encoder_output, "latent_dist"): + return encoder_output.latent_dist.mode() + if hasattr(encoder_output, "latents"): + return encoder_output.latents + raise AttributeError("Could not access latents from the SANA-Video VAE encoder output.") + + +class SanaImageToVideoPipeline(SanaVideoPipeline, SupportImageInput): + """Native vLLM-Omni SANA-Video 2B image-to-video pipeline. + + The first latent frame is encoded from the input image and remains fixed; + denoising updates only the subsequent video frames. + """ + + def _prepare_i2v_latents( + self, + image: torch.Tensor, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + num_frames: int, + dtype: torch.dtype, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None, + ) -> torch.Tensor: + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + batch_size, + num_channels_latents, + num_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"Generator list length {len(generator)} does not match effective batch size {batch_size}." + ) + if latents is None: + latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype) + else: + latents = latents.to(device=self.device, dtype=dtype) + + image = image.unsqueeze(2).to(device=self.device, dtype=self.vae.dtype) + image_latents = _retrieve_latents(self.vae.encode(image)) + image_latents = image_latents.repeat(batch_size, 1, 1, 1, 1) + + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + latents_mean = self.vae.latents_mean + latents_std = self.vae.latents_std + elif isinstance(self.vae, DistributedAutoencoderKLWan): + latents_mean = torch.tensor(self.vae.config.latents_mean) + latents_std = torch.tensor(self.vae.config.latents_std) + else: + raise TypeError(f"Unsupported SANA-Video I2V VAE type: {type(self.vae).__name__}") + + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(image_latents.device, image_latents.dtype) + image_latents = (image_latents - latents_mean) / latents_std + latents[:, :, :1] = image_latents.to(dtype) + return latents + + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + if len(req.prompts) != 1: + raise ValueError("SANA-Video I2V currently supports exactly one prompt per request.") + prompt_obj = req.prompts[0] + if isinstance(prompt_obj, str): + raise ValueError("SANA-Video I2V requires an image in `multi_modal_data`.") + + prompt = prompt_obj.get("prompt") or "" + negative_prompt = prompt_obj.get("negative_prompt") or "" + image = (prompt_obj.get("multi_modal_data") or {}).get("image") + if not prompt: + raise ValueError("Prompt is required for SANA-Video I2V generation.") + if not isinstance(image, PIL.Image.Image): + raise TypeError("SANA-Video I2V expects a preprocessed PIL image.") + sampling = req.sampling_params + default_height, default_width = get_sana_video_default_resolution(self.transformer.config.sample_size) + height = sampling.height or default_height + width = sampling.width or default_width + frames = resolve_video_num_frames( + sampling.num_frames, + default_num_frames=81, + is_dummy_run=req.is_dummy_run(), + ) + num_steps = sampling.num_inference_steps if sampling.num_inference_steps is not None else 50 + guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 6.0 + generator = sampling.generator + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(sampling.seed) + + extra_args = sampling.extra_args or {} + motion_score = extra_args.get("motion_score") + if motion_score is not None: + prompt = f"{prompt} motion score: {int(motion_score)}." + + video = self._generate_i2v( + image=image, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + frames=frames, + num_inference_steps=num_steps, + timesteps=sampling.timesteps, + sigmas=sampling.sigmas, + guidance_scale=guidance_scale, + num_videos_per_prompt=sampling.num_outputs_per_prompt or 1, + eta=sampling.eta, + generator=generator, + latents=sampling.latents, + clean_caption=bool(extra_args.get("clean_caption", False)), + use_resolution_binning=bool(extra_args.get("use_resolution_binning", True)), + max_sequence_length=sampling.max_sequence_length or 300, + output_type="latent" if sampling.output_type == "latent" else "raw", + ) + return DiffusionOutput(output=video) + + @torch.no_grad() + def _generate_i2v( + self, + *, + image: PIL.Image.Image, + prompt: str, + negative_prompt: str, + height: int, + width: int, + frames: int, + num_inference_steps: int, + guidance_scale: float, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None, + use_resolution_binning: bool, + max_sequence_length: int, + timesteps: list[int] | torch.Tensor | None = None, + sigmas: list[float] | None = None, + num_videos_per_prompt: int = 1, + eta: float = 0.0, + clean_caption: bool = False, + output_type: str = "raw", + complex_human_instruction: list[str] = SANA_VIDEO_I2V_COMPLEX_HUMAN_INSTRUCTION, + ) -> torch.Tensor: + orig_height, orig_width = height, width + if use_resolution_binning: + if self.transformer.config.sample_size == 30: + aspect_ratio_bin = ASPECT_RATIO_480_BIN + elif self.transformer.config.sample_size == 22: + aspect_ratio_bin = ASPECT_RATIO_720_BIN + else: + raise ValueError(f"Unsupported SANA-Video sample size: {self.transformer.config.sample_size}") + height, width = self.video_processor.classify_height_width_bin( + height, + width, + ratios=aspect_ratio_bin, + ) + + self.check_inputs( + prompt=prompt, + height=height, + width=width, + negative_prompt=negative_prompt, + ) + self._guidance_scale = guidance_scale + self._interrupt = False + + prompt_embeds, prompt_mask, negative_embeds, negative_mask = self.encode_prompt( + prompt, + self.do_classifier_free_guidance, + negative_prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + device=self.device, + clean_caption=clean_caption, + max_sequence_length=max_sequence_length, + complex_human_instruction=complex_human_instruction, + ) + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_embeds, prompt_embeds]) + prompt_mask = torch.cat([negative_mask, prompt_mask]) + + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps=num_inference_steps, + device=self.device, + timesteps=timesteps, + sigmas=sigmas, + ) + image_tensor = self.video_processor.preprocess(image, height=height, width=width).to( + self.device, + dtype=torch.float32, + ) + latent_channels = self.transformer.config.in_channels + latents = self._prepare_i2v_latents( + image_tensor, + batch_size=num_videos_per_prompt, + num_channels_latents=latent_channels, + height=height, + width=width, + num_frames=frames, + dtype=torch.float32, + generator=generator, + latents=latents, + ) + + patch_t, patch_h, patch_w = self.transformer.config.patch_size + conditioning_mask = latents.new_zeros( + num_videos_per_prompt, + 1, + latents.shape[2] // patch_t, + latents.shape[3] // patch_h, + latents.shape[4] // patch_w, + ) + conditioning_mask[:, :, 0] = 1 + if self.do_classifier_free_guidance: + conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) + + self._num_timesteps = len(timesteps) + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + transformer_dtype = self.transformer.dtype + with self.progress_bar(total=num_inference_steps) as progress_bar: + for index, timestep_value in enumerate(timesteps): + latent_input = torch.cat([latents, latents]) if self.do_classifier_free_guidance else latents + timestep = timestep_value.expand(conditioning_mask.shape) * (1 - conditioning_mask) + noise_pred = self.transformer( + latent_input.to(transformer_dtype), + encoder_hidden_states=prompt_embeds.to(transformer_dtype), + encoder_attention_mask=prompt_mask, + timestep=timestep, + return_dict=False, + )[0].float() + if self.do_classifier_free_guidance: + uncond, cond = noise_pred.chunk(2) + noise_pred = uncond + guidance_scale * (cond - uncond) + if self.transformer.config.out_channels // 2 == latent_channels: + noise_pred = noise_pred.chunk(2, dim=1)[0] + + denoised = self.scheduler.step( + noise_pred[:, :, 1:], + timestep_value, + latents[:, :, 1:], + **extra_step_kwargs, + return_dict=False, + )[0] + latents = torch.cat([latents[:, :, :1], denoised], dim=2) + if index == len(timesteps) - 1 or (index + 1) % self.scheduler.order == 0: + progress_bar.update() + + if output_type == "latent": + return latents + + latents = latents.to(self.vae.dtype) + if isinstance(self.vae, DistributedAutoencoderKLLTX2Video): + latents_mean = self.vae.latents_mean + latents_std = self.vae.latents_std + latent_dim = self.vae.config.latent_channels + else: + latents_mean = torch.tensor(self.vae.config.latents_mean) + latents_std = torch.tensor(self.vae.config.latents_std) + latent_dim = self.vae.config.z_dim + latents_mean = latents_mean.view(1, latent_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, latent_dim, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std + latents_mean + + try: + video = self.vae.decode(latents, return_dict=False)[0] + except torch.OutOfMemoryError as error: + warnings.warn( + f"{error}. Enable VAE tiling for SANA-Video I2V high-resolution decoding.", + stacklevel=2, + ) + raise + if use_resolution_binning: + video = self.video_processor.resize_and_crop_tensor(video, orig_width, orig_height) + return video diff --git a/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py new file mode 100644 index 000000000..eb0f5843b --- /dev/null +++ b/vllm_omni/diffusion/models/sana_video/transformer_sana_video.py @@ -0,0 +1,1066 @@ +# Copyright 2025 The HuggingFace Team and SANA-Video Team. All rights reserved. +# +# 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. + +import math +import numbers +from collections.abc import Iterable +from dataclasses import dataclass, fields + +import torch +from torch import nn +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as OmniAttention + + +@dataclass +class SanaVideoTransformerConfig: + in_channels: int = 16 + out_channels: int | None = 16 + num_attention_heads: int = 20 + attention_head_dim: int = 112 + num_layers: int = 20 + num_cross_attention_heads: int | None = 20 + cross_attention_head_dim: int | None = 112 + cross_attention_dim: int | None = 2240 + caption_channels: int = 2304 + mlp_ratio: float = 2.5 + dropout: float = 0.0 + attention_bias: bool = False + sample_size: int = 30 + patch_size: tuple[int, int, int] = (1, 2, 2) + norm_elementwise_affine: bool = False + norm_eps: float = 1e-6 + interpolation_scale: int | None = None + guidance_embeds: bool = False + guidance_embeds_scale: float = 0.1 + qk_norm: str | None = "rms_norm_across_heads" + rope_max_seq_len: int = 1024 + + @classmethod + def from_dict(cls, config: dict) -> "SanaVideoTransformerConfig": + field_names = {field.name for field in fields(cls)} + unknown = {key for key in config if key not in field_names and not key.startswith("_")} + if unknown: + raise ValueError(f"Unsupported SANA-Video transformer config fields: {sorted(unknown)}") + values = {key: value for key, value in config.items() if key in field_names} + if "patch_size" in values: + values["patch_size"] = tuple(values["patch_size"]) + return cls(**values) + + +@dataclass +class SanaVideoTransformerOutput: + sample: torch.Tensor + + +class SanaRMSNorm(nn.Module): + """RMSNorm with the parameter and dtype semantics used by SANA checkpoints.""" + + def __init__( + self, + dim: int | Iterable[int], + eps: float, + elementwise_affine: bool = True, + bias: bool = False, + ) -> None: + super().__init__() + + self.eps = eps + self.elementwise_affine = elementwise_affine + + if isinstance(dim, numbers.Integral): + dim = (dim,) + self.dim = torch.Size(dim) + + self.weight: nn.Parameter | None = None + self.bias: nn.Parameter | None = None + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(self.dim)) + if bias: + self.bias = nn.Parameter(torch.zeros(self.dim)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + + if self.weight is not None: + if self.weight.dtype in (torch.float16, torch.bfloat16): + hidden_states = hidden_states.to(self.weight.dtype) + hidden_states = hidden_states * self.weight + if self.bias is not None: + hidden_states = hidden_states + self.bias + else: + hidden_states = hidden_states.to(input_dtype) + + return hidden_states + + +def _get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +) -> torch.Tensor: + """Create the sinusoidal timestep embeddings used by SANA checkpoints.""" + if timesteps.ndim != 1: + raise ValueError(f"Timesteps must be a 1D tensor, got shape {tuple(timesteps.shape)}") + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, + end=half_dim, + dtype=torch.float32, + device=timesteps.device, + ) + exponent = exponent / (half_dim - downscale_freq_shift) + embeddings = timesteps[:, None].float() * torch.exp(exponent)[None, :] + embeddings = scale * embeddings + embeddings = torch.cat([torch.sin(embeddings), torch.cos(embeddings)], dim=-1) + + if flip_sin_to_cos: + embeddings = torch.cat([embeddings[:, half_dim:], embeddings[:, :half_dim]], dim=-1) + if embedding_dim % 2 == 1: + embeddings = torch.nn.functional.pad(embeddings, (0, 1, 0, 0)) + + return embeddings + + +class SanaTimesteps(nn.Module): + """Parameter-free timestep projection compatible with Diffusers Timesteps.""" + + def __init__( + self, + num_channels: int, + flip_sin_to_cos: bool, + downscale_freq_shift: float, + scale: int = 1, + ) -> None: + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + return _get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + ) + + +class SanaTimestepEmbedding(nn.Module): + """Two-layer timestep MLP with Diffusers-compatible parameter names.""" + + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int | None = None, + post_act_fn: str | None = None, + cond_proj_dim: int | None = None, + sample_proj_bias: bool = True, + ) -> None: + super().__init__() + if act_fn != "silu": + raise ValueError(f"SANA timestep embeddings only support act_fn='silu', got {act_fn!r}") + if post_act_fn not in (None, "silu"): + raise ValueError(f"SANA timestep embeddings only support post_act_fn=None or 'silu', got {post_act_fn!r}") + + self.linear_1 = nn.Linear(in_channels, time_embed_dim, bias=sample_proj_bias) + self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) if cond_proj_dim is not None else None + self.act = nn.SiLU() + self.linear_2 = nn.Linear( + time_embed_dim, + out_dim if out_dim is not None else time_embed_dim, + bias=sample_proj_bias, + ) + self.post_act = nn.SiLU() if post_act_fn == "silu" else None + + def forward( + self, + sample: torch.Tensor, + condition: torch.Tensor | None = None, + ) -> torch.Tensor: + if condition is not None: + if self.cond_proj is None: + raise ValueError("condition was provided but cond_proj_dim was not configured") + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + sample = self.act(sample) + sample = self.linear_2(sample) + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class SanaCombinedTimestepSizeEmbeddings(nn.Module): + """PixArt timestep/size conditioning used by SANA's adaLN-single.""" + + def __init__( + self, + embedding_dim: int, + size_emb_dim: int, + use_additional_conditions: bool = False, + ) -> None: + super().__init__() + self.outdim = size_emb_dim + self.time_proj = SanaTimesteps( + num_channels=256, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.timestep_embedder = SanaTimestepEmbedding( + in_channels=256, + time_embed_dim=embedding_dim, + ) + + self.use_additional_conditions = use_additional_conditions + if use_additional_conditions: + self.additional_condition_proj = SanaTimesteps( + num_channels=256, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.resolution_embedder = SanaTimestepEmbedding( + in_channels=256, + time_embed_dim=size_emb_dim, + ) + self.aspect_ratio_embedder = SanaTimestepEmbedding( + in_channels=256, + time_embed_dim=size_emb_dim, + ) + + def forward( + self, + timestep: torch.Tensor, + resolution: torch.Tensor | None, + aspect_ratio: torch.Tensor | None, + batch_size: int | None, + hidden_dtype: torch.dtype | None, + ) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) + + if not self.use_additional_conditions: + return timesteps_emb + if resolution is None or aspect_ratio is None or batch_size is None: + raise ValueError( + "resolution, aspect_ratio, and batch_size are required when use_additional_conditions=True" + ) + + resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype) + resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1) + aspect_ratio_emb = self.additional_condition_proj(aspect_ratio.flatten()).to(hidden_dtype) + aspect_ratio_emb = self.aspect_ratio_embedder(aspect_ratio_emb).reshape(batch_size, -1) + return timesteps_emb + torch.cat([resolution_emb, aspect_ratio_emb], dim=1) + + +class SanaAdaLayerNormSingle(nn.Module): + """SANA's checkpoint-compatible adaLN-single conditioning module.""" + + def __init__( + self, + embedding_dim: int, + use_additional_conditions: bool = False, + ) -> None: + super().__init__() + self.emb = SanaCombinedTimestepSizeEmbeddings( + embedding_dim, + size_emb_dim=embedding_dim // 3, + use_additional_conditions=use_additional_conditions, + ) + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: dict[str, torch.Tensor] | None = None, + batch_size: int | None = None, + hidden_dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + added_cond_kwargs = added_cond_kwargs or { + "resolution": None, + "aspect_ratio": None, + } + embedded_timestep = self.emb( + timestep, + **added_cond_kwargs, + batch_size=batch_size, + hidden_dtype=hidden_dtype, + ) + return self.linear(self.silu(embedded_timestep)), embedded_timestep + + +class _SanaFP32SiLU(nn.Module): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.silu(inputs.float(), inplace=False).to(inputs.dtype) + + +class SanaPixArtAlphaTextProjection(nn.Module): + """Caption projection with Diffusers-compatible parameters and activations.""" + + def __init__( + self, + in_features: int, + hidden_size: int, + out_features: int | None = None, + act_fn: str = "gelu_tanh", + ) -> None: + super().__init__() + out_features = hidden_size if out_features is None else out_features + self.linear_1 = nn.Linear(in_features, hidden_size, bias=True) + if act_fn == "gelu_tanh": + self.act_1 = nn.GELU(approximate="tanh") + elif act_fn == "silu": + self.act_1 = nn.SiLU() + elif act_fn == "silu_fp32": + self.act_1 = _SanaFP32SiLU() + else: + raise ValueError(f"Unknown activation function: {act_fn}") + self.linear_2 = nn.Linear(hidden_size, out_features, bias=True) + + def forward(self, caption: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + return self.linear_2(hidden_states) + + +def _get_1d_rotary_pos_embed( + dim: int, + max_seq_len: int, + theta: float, + freqs_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Generate the real-valued 1D RoPE representation used by SANA-Video.""" + if dim % 2 != 0: + raise ValueError(f"RoPE dimension must be even, got {dim}") + + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype) / dim)) + positions = torch.arange(max_seq_len, dtype=freqs_dtype) + freqs = torch.outer(positions, freqs) + freqs_cos = freqs.cos().float().repeat_interleave(2, dim=-1) + freqs_sin = freqs.sin().float().repeat_interleave(2, dim=-1) + return freqs_cos, freqs_sin + + +class GLUMBTempConv(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + expand_ratio: float = 4, + norm_type: str | None = None, + residual_connection: bool = True, + ) -> None: + super().__init__() + + hidden_channels = int(expand_ratio * in_channels) + self.norm_type = norm_type + self.residual_connection = residual_connection + + self.nonlinearity = nn.SiLU() + self.conv_inverted = nn.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0) + self.conv_depth = nn.Conv2d(hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2) + self.conv_point = nn.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False) + + self.norm = None + if norm_type == "rms_norm": + self.norm = SanaRMSNorm(out_channels, eps=1e-5, elementwise_affine=True, bias=True) + + self.conv_temp = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 1), stride=1, padding=(1, 0), bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.residual_connection: + residual = hidden_states + batch_size, num_frames, height, width, num_channels = hidden_states.shape + hidden_states = hidden_states.view(batch_size * num_frames, height, width, num_channels).permute(0, 3, 1, 2) + + hidden_states = self.conv_inverted(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv_depth(hidden_states) + hidden_states, gate = torch.chunk(hidden_states, 2, dim=1) + hidden_states = hidden_states * self.nonlinearity(gate) + + hidden_states = self.conv_point(hidden_states) + + # Temporal aggregation + hidden_states_temporal = hidden_states.view(batch_size, num_frames, num_channels, height * width).permute( + 0, 2, 1, 3 + ) + hidden_states = hidden_states_temporal + self.conv_temp(hidden_states_temporal) + hidden_states = hidden_states.permute(0, 2, 3, 1).view(batch_size, num_frames, height, width, num_channels) + + if self.norm_type == "rms_norm": + # move channel to the last dimension so we apply RMSnorm across channel dimension + hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1) + + if self.residual_connection: + hidden_states = hidden_states + residual + + return hidden_states + + +class SanaLinearAttention(nn.Module): + """SANA's ReLU-kernel linear self-attention. + + This is intentionally not routed through the unified softmax attention + layer because doing so would change the model's attention algorithm. + """ + + def __init__( + self, + dim: int, + num_heads: int, + head_dim: int, + dropout: float, + bias: bool, + qk_norm: str | None, + ) -> None: + super().__init__() + inner_dim = num_heads * head_dim + self.heads = num_heads + self.to_q = nn.Linear(dim, inner_dim, bias=bias) + self.to_k = nn.Linear(dim, inner_dim, bias=bias) + self.to_v = nn.Linear(dim, inner_dim, bias=bias) + self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=True), nn.Dropout(dropout)]) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + if rotary_emb is None: + raise ValueError("SanaLinearAttention requires rotary_emb.") + + original_dtype = hidden_states.dtype + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.unflatten(2, (self.heads, -1)) + key = key.unflatten(2, (self.heads, -1)) + value = value.unflatten(2, (self.heads, -1)) + # B,N,H,C + + query = torch.relu(query) + key = torch.relu(key) + + def apply_rotary_emb( + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ): + x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + query_rotate = apply_rotary_emb(query, *rotary_emb) + key_rotate = apply_rotary_emb(key, *rotary_emb) + + # B,H,C,N + query = query.permute(0, 2, 3, 1) + key = key.permute(0, 2, 3, 1) + query_rotate = query_rotate.permute(0, 2, 3, 1) + key_rotate = key_rotate.permute(0, 2, 3, 1) + value = value.permute(0, 2, 3, 1) + + query_rotate, key_rotate, value = query_rotate.float(), key_rotate.float(), value.float() + + # Keep the unrotated denominator contraction in the input dtype. This + # intentionally matches Diffusers 0.38.0's SanaLinearAttnProcessor3_0; + # only the rotated numerator path is accumulated in FP32. + z = 1 / (key.sum(dim=-1, keepdim=True).transpose(-2, -1) @ query + 1e-15) + + scores = torch.matmul(value, key_rotate.transpose(-1, -2)) + hidden_states = torch.matmul(scores, query_rotate) + + hidden_states = hidden_states * z + # B,H,C,N + hidden_states = hidden_states.flatten(1, 2).transpose(1, 2) + hidden_states = hidden_states.to(original_dtype) + + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + + return hidden_states + + +class WanRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size: tuple[int, int, int], + max_seq_len: int, + theta: float = 10000.0, + ): + super().__init__() + + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + + self.t_dim = t_dim + self.h_dim = h_dim + self.w_dim = w_dim + + freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + + freqs_cos = [] + freqs_sin = [] + + for dim in [t_dim, h_dim, w_dim]: + freq_cos, freq_sin = _get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype) + freqs_cos.append(freq_cos) + freqs_sin.append(freq_sin) + + self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False) + self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w + + split_sizes = [self.t_dim, self.h_dim, self.w_dim] + + freqs_cos = self.freqs_cos.split(split_sizes, dim=1) + freqs_sin = self.freqs_sin.split(split_sizes, dim=1) + + freqs_cos_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_cos_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_sin_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + freqs_sin_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + + return freqs_cos, freqs_sin + + +class SanaModulatedNorm(nn.Module): + def __init__(self, dim: int, elementwise_affine: bool = False, eps: float = 1e-6): + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor, scale_shift_table: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(hidden_states) + shift, scale = (scale_shift_table[None, None] + temb[:, :, None].to(scale_shift_table.device)).unbind(dim=2) + hidden_states = hidden_states * (1 + scale) + shift + return hidden_states + + +class SanaCombinedTimestepGuidanceEmbeddings(nn.Module): + def __init__(self, embedding_dim): + super().__init__() + self.time_proj = SanaTimesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = SanaTimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.guidance_condition_proj = SanaTimesteps( + num_channels=256, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.guidance_embedder = SanaTimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + def forward(self, timestep: torch.Tensor, guidance: torch.Tensor = None, hidden_dtype: torch.dtype = None): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + + guidance_proj = self.guidance_condition_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(dtype=hidden_dtype)) + conditioning = timesteps_emb + guidance_emb + + return self.linear(self.silu(conditioning)), conditioning + + +class SanaCrossAttention(nn.Module): + """SANA cross-attention backed by vLLM-Omni's unified attention layer.""" + + def __init__( + self, + dim: int, + cross_attention_dim: int, + num_heads: int, + head_dim: int, + dropout: float, + bias: bool, + out_bias: bool, + qk_norm: str | None, + prefix: str = "", + ) -> None: + super().__init__() + inner_dim = num_heads * head_dim + self.heads = num_heads + self.head_dim = head_dim + self.to_q = nn.Linear(dim, inner_dim, bias=bias) + self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + self.norm_q = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.norm_k = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) if qk_norm is not None else None + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim, bias=out_bias), nn.Dropout(dropout)]) + self.attn = OmniAttention( + num_heads=num_heads, + head_size=head_dim, + num_kv_heads=num_heads, + softmax_scale=1.0 / (head_dim**0.5), + causal=False, + role="cross", + qkv_layout="BSND", + prefix=prefix, + skip_sequence_parallel=True, + disable_kv_quant=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + query = self.to_q(hidden_states) + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.unflatten(2, (self.heads, self.head_dim)) + key = key.unflatten(2, (self.heads, self.head_dim)) + value = value.unflatten(2, (self.heads, self.head_dim)) + + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + if attention_mask is not None: + # The Flash masked-varlen path applies one padding mask to Q/K/V, + # while SANA's text mask describes only cross-attention K/V. + # Use SDPA until Flash supports an independent key-padding mask. + hidden_states = self.attn.sdpa_fallback.forward(query, key, value, attn_metadata) + else: + hidden_states = self.attn(query, key, value, None) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.type_as(query) + + hidden_states = self.to_out[0](hidden_states) + return self.to_out[1](hidden_states) + + +class SanaVideoTransformerBlock(nn.Module): + r""" + Transformer block introduced in [Sana-Video](https://huggingface.co/papers/2509.24695). + """ + + def __init__( + self, + dim: int = 2240, + num_attention_heads: int = 20, + attention_head_dim: int = 112, + dropout: float = 0.0, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, + attention_bias: bool = True, + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + attention_out_bias: bool = True, + mlp_ratio: float = 3.0, + qk_norm: str | None = "rms_norm_across_heads", + rope_max_seq_len: int = 1024, + prefix: str = "", + ) -> None: + super().__init__() + + # 1. Self Attention + self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) + self.attn1 = SanaLinearAttention( + dim=dim, + num_heads=num_attention_heads, + head_dim=attention_head_dim, + dropout=dropout, + bias=attention_bias, + qk_norm=qk_norm, + ) + + # 2. Cross Attention + self.norm2 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.attn2 = None + if cross_attention_dim is not None: + self.attn2 = SanaCrossAttention( + dim=dim, + cross_attention_dim=cross_attention_dim, + num_heads=num_cross_attention_heads, + head_dim=cross_attention_head_dim, + dropout=dropout, + bias=True, + out_bias=attention_out_bias, + qk_norm=qk_norm, + prefix=f"{prefix}.attn2" if prefix else "attn2", + ) + + # 3. Feed-forward + self.ff = GLUMBTempConv(dim, dim, mlp_ratio, norm_type=None, residual_connection=False) + + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + timestep: torch.LongTensor | None = None, + frames: int = None, + height: int = None, + width: int = None, + rotary_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size = hidden_states.shape[0] + + # 1. Modulation + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None] + timestep.reshape(batch_size, timestep.shape[1], 6, -1) + ).unbind(dim=2) + + # 2. Self Attention + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.to(hidden_states.dtype) + + attn_output = self.attn1(norm_hidden_states, rotary_emb=rotary_emb) + hidden_states = hidden_states + gate_msa * attn_output + + # 3. Cross Attention + if self.attn2 is not None: + attn_output = self.attn2( + hidden_states, + encoder_hidden_states, + encoder_attention_mask, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + norm_hidden_states = norm_hidden_states.unflatten(1, (frames, height, width)) + ff_output = self.ff(norm_hidden_states) + ff_output = ff_output.flatten(1, 3) + hidden_states = hidden_states + gate_mlp * ff_output + + return hidden_states + + +class SanaVideoTransformer3DModel(nn.Module): + r""" + A 3D Transformer model introduced in [Sana-Video](https://huggingface.co/papers/2509.24695) family of models. + + Args: + in_channels (`int`, defaults to `16`): + The number of channels in the input. + out_channels (`int`, *optional*, defaults to `16`): + The number of channels in the output. + num_attention_heads (`int`, defaults to `20`): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, defaults to `112`): + The number of channels in each head. + num_layers (`int`, defaults to `20`): + The number of layers of Transformer blocks to use. + num_cross_attention_heads (`int`, *optional*, defaults to `20`): + The number of heads to use for cross-attention. + cross_attention_head_dim (`int`, *optional*, defaults to `112`): + The number of channels in each head for cross-attention. + cross_attention_dim (`int`, *optional*, defaults to `2240`): + The number of channels in the cross-attention output. + caption_channels (`int`, defaults to `2304`): + The number of channels in the caption embeddings. + mlp_ratio (`float`, defaults to `2.5`): + The expansion ratio to use in the GLUMBConv layer. + dropout (`float`, defaults to `0.0`): + The dropout probability. + attention_bias (`bool`, defaults to `False`): + Whether to use bias in the attention layer. + sample_size (`int`, defaults to `32`): + The base size of the input latent. + patch_size (`int`, defaults to `1`): + The size of the patches to use in the patch embedding layer. + norm_elementwise_affine (`bool`, defaults to `False`): + Whether to use elementwise affinity in the normalization layer. + norm_eps (`float`, defaults to `1e-6`): + The epsilon value for the normalization layer. + qk_norm (`str`, *optional*, defaults to `None`): + The normalization to use for the query and key. + """ + + _no_split_modules = ["SanaVideoTransformerBlock", "SanaModulatedNorm"] + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return AutoWeightsLoader(self).load_weights(weights) + + def __init__( + self, + in_channels: int = 16, + out_channels: int | None = 16, + num_attention_heads: int = 20, + attention_head_dim: int = 112, + num_layers: int = 20, + num_cross_attention_heads: int | None = 20, + cross_attention_head_dim: int | None = 112, + cross_attention_dim: int | None = 2240, + caption_channels: int = 2304, + mlp_ratio: float = 2.5, + dropout: float = 0.0, + attention_bias: bool = False, + sample_size: int = 30, + patch_size: tuple[int, int, int] = (1, 2, 2), + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-6, + interpolation_scale: int | None = None, + guidance_embeds: bool = False, + guidance_embeds_scale: float = 0.1, + qk_norm: str | None = "rms_norm_across_heads", + rope_max_seq_len: int = 1024, + ) -> None: + super().__init__() + + out_channels = out_channels or in_channels + self.config = SanaVideoTransformerConfig( + in_channels=in_channels, + out_channels=out_channels, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + num_layers=num_layers, + num_cross_attention_heads=num_cross_attention_heads, + cross_attention_head_dim=cross_attention_head_dim, + cross_attention_dim=cross_attention_dim, + caption_channels=caption_channels, + mlp_ratio=mlp_ratio, + dropout=dropout, + attention_bias=attention_bias, + sample_size=sample_size, + patch_size=tuple(patch_size), + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + interpolation_scale=interpolation_scale, + guidance_embeds=guidance_embeds, + guidance_embeds_scale=guidance_embeds_scale, + qk_norm=qk_norm, + rope_max_seq_len=rope_max_seq_len, + ) + inner_dim = num_attention_heads * attention_head_dim + + # 1. Patch & position embedding + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + + # 2. Additional condition embeddings + if guidance_embeds: + self.time_embed = SanaCombinedTimestepGuidanceEmbeddings(inner_dim) + else: + self.time_embed = SanaAdaLayerNormSingle(inner_dim) + + self.caption_projection = SanaPixArtAlphaTextProjection( + in_features=caption_channels, + hidden_size=inner_dim, + ) + self.caption_norm = SanaRMSNorm(inner_dim, eps=1e-5, elementwise_affine=True) + + # 3. Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + SanaVideoTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + num_cross_attention_heads=num_cross_attention_heads, + cross_attention_head_dim=cross_attention_head_dim, + cross_attention_dim=cross_attention_dim, + attention_bias=attention_bias, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + mlp_ratio=mlp_ratio, + qk_norm=qk_norm, + prefix=f"transformer_blocks.{layer_idx}", + ) + for layer_idx in range(num_layers) + ] + ) + + # 4. Output blocks + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + self.norm_out = SanaModulatedNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, math.prod(patch_size) * out_channels) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + @classmethod + def from_config(cls, config: dict | SanaVideoTransformerConfig) -> "SanaVideoTransformer3DModel": + if isinstance(config, dict): + config = SanaVideoTransformerConfig.from_dict(config) + return cls(**{field.name: getattr(config, field.name) for field in fields(config)}) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + guidance: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + controlnet_block_samples: tuple[torch.Tensor] | None = None, + return_dict: bool = True, + ) -> tuple[torch.Tensor, ...] | SanaVideoTransformerOutput: + """ + The [`SanaVideoTransformer3DModel`] forward method. + + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`): + Input `hidden_states`. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_len, embed_dims)`): + Conditional embeddings (embeddings computed from the input conditions such as prompts) to use. + timestep (`torch.LongTensor`): + Used to indicate denoising step. + guidance (`torch.Tensor`, *optional*): + Guidance scale embedding. + encoder_attention_mask (`torch.Tensor`, *optional*): + Cross-attention mask applied to `encoder_hidden_states`. + attention_mask (`torch.Tensor`, *optional*): + Retained for Diffusers API compatibility. SANA linear self-attention does not consume this mask. + controlnet_block_samples (`tuple` of `torch.Tensor`, *optional*): + A list of tensors that if specified are added to the residuals of transformer blocks. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return `SanaVideoTransformerOutput` instead of a plain tuple. + + Returns: + If `return_dict` is True, a [`SanaVideoTransformerOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + # SANA's linear self-attention does not consume ``attention_mask``; + # retain the public argument for Diffusers-compatible call sites. + if encoder_attention_mask is not None: + if encoder_attention_mask.ndim != 2: + raise ValueError( + "encoder_attention_mask must be a 2D padding mask with " + f"shape (batch_size, key_length); got {tuple(encoder_attention_mask.shape)}" + ) + if encoder_attention_mask.shape != encoder_hidden_states.shape[:2]: + raise ValueError( + "encoder_attention_mask shape must match encoder_hidden_states: " + f"expected {tuple(encoder_hidden_states.shape[:2])}, " + f"got {tuple(encoder_attention_mask.shape)}" + ) + if encoder_attention_mask.is_floating_point() or encoder_attention_mask.is_complex(): + raise TypeError( + "encoder_attention_mask must be a boolean or integer padding mask, " + "not a floating-point additive attention bias" + ) + encoder_attention_mask = encoder_attention_mask.to( + device=encoder_hidden_states.device, + dtype=torch.bool, + ) + + # 1. Input + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.config.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + rotary_emb = self.rope(hidden_states) + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if guidance is not None: + timestep, embedded_timestep = self.time_embed( + timestep.flatten(), guidance=guidance, hidden_dtype=hidden_states.dtype + ) + else: + timestep, embedded_timestep = self.time_embed( + timestep.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype + ) + + timestep = timestep.view(batch_size, -1, timestep.size(-1)) + embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) + + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) + + encoder_hidden_states = self.caption_norm(encoder_hidden_states) + + # 2. Transformer blocks + for index_block, block in enumerate(self.transformer_blocks): + hidden_states = block( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + post_patch_num_frames, + post_patch_height, + post_patch_width, + rotary_emb, + ) + if controlnet_block_samples is not None and 0 < index_block <= len(controlnet_block_samples): + hidden_states = hidden_states + controlnet_block_samples[index_block - 1] + + # 3. Normalization + hidden_states = self.norm_out(hidden_states, embedded_timestep, self.scale_shift_table) + + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + + return SanaVideoTransformerOutput(sample=output) diff --git a/vllm_omni/diffusion/models/sana_wm/__init__.py b/vllm_omni/diffusion/models/sana_wm/__init__.py new file mode 100644 index 000000000..f6f776194 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_wm/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SANA-WM diffusion model integration.""" + +from vllm_omni.diffusion.models.sana_wm.config import SanaWmConfig +from vllm_omni.diffusion.models.sana_wm.pipeline_sana_wm import ( + SANA_WM_MODEL_ID, + SANA_WM_OUTPUT_HEIGHT, + SANA_WM_OUTPUT_WIDTH, + SanaWmPipeline, + get_sana_wm_pre_process_func, +) +from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload +from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmTransformer3DModel + +__all__ = [ + "SANA_WM_MODEL_ID", + "SANA_WM_OUTPUT_HEIGHT", + "SANA_WM_OUTPUT_WIDTH", + "SanaWmConfig", + "SanaWmPipeline", + "SanaWmTransformer3DModel", + "get_sana_wm_pre_process_func", + "normalize_sana_wm_payload", +] diff --git a/vllm_omni/diffusion/models/sana_wm/camera_control.py b/vllm_omni/diffusion/models/sana_wm/camera_control.py new file mode 100644 index 000000000..914fbc0d3 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_wm/camera_control.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Camera-control helpers for SANA-WM. + +This module mirrors the public NVlabs/Sana SANA-WM camera preparation path: +WASD/IJKL action rollout or explicit camera-to-world poses are converted into +relative poses, per-latent-frame ray metadata, and per-VAE-chunk Plucker maps. +Model-local projection layers consume these tensors in the Stage-1 DiT. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +SANA_WM_DEFAULT_TRANSLATION_SPEED = 0.05 +SANA_WM_DEFAULT_ROTATION_SPEED_DEG = 1.2 +SANA_WM_DEFAULT_PITCH_LIMIT_DEG = 85.0 +SANA_WM_ALLOWED_ACTION_KEYS = frozenset("wasdijkl") +SANA_WM_DEFAULT_VAE_STRIDE = (8, 32, 32) + + +@dataclass(frozen=True) +class SanaWmCameraCondition: + poses: Any | None = None + intrinsics: Any | None = None + action: str | None = None + num_frames: int | None = None + height: int = 704 + width: int = 1280 + translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED + rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG + pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG + + +def _rot_x(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64) + + +def _rot_y(angle_rad: float) -> np.ndarray: + c, s = np.cos(angle_rad), np.sin(angle_rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def _parse_action_segments(action: str) -> list[tuple[list[str], int]]: + """Parse an action string into ``(keys, duration)`` pairs without expanding it. + + A single segment names an unbounded rollout (``"w-1000000000"`` is valid + syntax), so expanding here would allocate one list entry per frame before + anything checks the frame count. Callers total the durations and validate + against ``num_frames`` first; only :func:`_parse_action_string` expands. + """ + cleaned = "".join(action.replace(",", ",").split()) + if not cleaned: + raise ValueError("Sana-WM action string is empty.") + + segments: list[tuple[list[str], int]] = [] + for segment in cleaned.split(","): + if not segment or "-" not in segment: + raise ValueError(f"Invalid Sana-WM action segment {segment!r}: expected '-'.") + keys_part, duration = segment.rsplit("-", 1) + if not duration.isdigit() or int(duration) <= 0: + raise ValueError(f"Sana-WM action segment {segment!r} has a non-positive duration {duration!r}.") + keys_lower = keys_part.lower() + if keys_lower == "none": + keys: list[str] = [] + else: + bad = sorted({key for key in keys_lower if key not in SANA_WM_ALLOWED_ACTION_KEYS}) + if bad: + allowed = "".join(sorted(SANA_WM_ALLOWED_ACTION_KEYS)) + raise ValueError(f"Sana-WM action segment {segment!r} contains unknown keys {bad}; allowed: {allowed}.") + keys = sorted(set(keys_lower)) + segments.append((keys, int(duration))) + return segments + + +def action_rollout_num_frames(action: str) -> int: + """Frames ``action`` rolls out: the identity start pose plus one per step. + + Totals the durations arithmetically, so an oversized request is rejected + before anything allocates per-frame state. + """ + return sum(duration for _, duration in _parse_action_segments(action)) + 1 + + +def _parse_action_string(action: str) -> list[list[str]]: + per_frame: list[list[str]] = [] + for keys, duration in _parse_action_segments(action): + per_frame.extend([keys] * duration) + return per_frame + + +def action_string_to_c2w( + action: str, + *, + translation_speed: float = SANA_WM_DEFAULT_TRANSLATION_SPEED, + rotation_speed_deg: float = SANA_WM_DEFAULT_ROTATION_SPEED_DEG, + pitch_limit_deg: float = SANA_WM_DEFAULT_PITCH_LIMIT_DEG, +) -> np.ndarray: + """Roll out an OpenCV-convention ``(N + 1, 4, 4)`` c2w trajectory.""" + + per_frame = _parse_action_string(action) + rotate_rad = math.radians(rotation_speed_deg) + pitch_limit_rad = math.radians(pitch_limit_deg) + current = np.eye(4, dtype=np.float64) + current_pitch = 0.0 + poses = [current.copy()] + + for keys in per_frame: + held = set(keys) + rotation = current[:3, :3] + translation = current[:3, 3] + + pitch_delta = (rotate_rad if "i" in held else 0.0) - (rotate_rad if "k" in held else 0.0) + new_pitch = current_pitch + pitch_delta + if -pitch_limit_rad <= new_pitch <= pitch_limit_rad: + current_pitch = new_pitch + else: + pitch_delta = 0.0 + + yaw_delta = (rotate_rad if "l" in held else 0.0) - (rotate_rad if "j" in held else 0.0) + new_rotation = _rot_y(yaw_delta) @ rotation @ _rot_x(pitch_delta) + + forward = new_rotation[:, 2].copy() + forward[1] = 0.0 + right = new_rotation[:, 0].copy() + right[1] = 0.0 + forward_norm = float(np.linalg.norm(forward)) + right_norm = float(np.linalg.norm(right)) + if forward_norm > 0: + forward /= forward_norm + 1e-6 + if right_norm > 0: + right /= right_norm + 1e-6 + + move = np.zeros(3, dtype=np.float64) + if "w" in held: + move += forward * translation_speed + if "s" in held: + move -= forward * translation_speed + if "d" in held: + move += right * translation_speed + if "a" in held: + move -= right * translation_speed + + current = np.eye(4, dtype=np.float64) + current[:3, :3] = new_rotation + current[:3, 3] = translation + move + poses.append(current.copy()) + + return np.stack(poses, axis=0).astype(np.float32) + + +def get_pose_inverse(transform: torch.Tensor) -> torch.Tensor: + """Invert homogeneous rigid transforms with shape ``(..., 4, 4)``.""" + + rotation = transform[..., :3, :3] + translation = transform[..., :3, 3] + rotation_inv = rotation.transpose(-1, -2) + translation_inv = -torch.matmul(rotation_inv, translation.unsqueeze(-1)).squeeze(-1) + output = torch.eye(4, dtype=transform.dtype, device=transform.device).repeat(transform.shape[:-2] + (1, 1)) + output[..., :3, :3] = rotation_inv + output[..., :3, 3] = translation_inv + return output + + +def compute_raymap( + intrinsics: torch.Tensor, + poses: torch.Tensor, + height: int, + width: int, + *, + use_plucker: bool = True, +) -> torch.Tensor: + """Compute SANA-WM geometry ray maps. + + Args: + intrinsics: ``(T, 4)`` tensor in ``[fx, fy, cx, cy]`` order. + poses: ``(T, 4, 4)`` camera-to-world matrices. + height: raymap height in latent pixels. + width: raymap width in latent pixels. + use_plucker: return ``[direction, moment]`` if true, otherwise + ``[origin, direction]``. + """ + + if intrinsics.ndim != 2 or intrinsics.shape[-1] != 4: + raise ValueError(f"Sana-WM intrinsics must have shape (T, 4), got {tuple(intrinsics.shape)}.") + if poses.ndim != 3 or poses.shape[-2:] != (4, 4): + raise ValueError(f"Sana-WM poses must have shape (T, 4, 4), got {tuple(poses.shape)}.") + if intrinsics.shape[0] != poses.shape[0]: + raise ValueError("Sana-WM intrinsics and poses must have the same temporal length.") + + frames = intrinsics.shape[0] + device = intrinsics.device + dtype = intrinsics.dtype + y_grid, x_grid = torch.meshgrid( + torch.arange(height, device=device, dtype=dtype), + torch.arange(width, device=device, dtype=dtype), + indexing="ij", + ) + x_grid = x_grid[None].expand(frames, -1, -1) + y_grid = y_grid[None].expand(frames, -1, -1) + + fx = intrinsics[:, 0].view(frames, 1, 1) + fy = intrinsics[:, 1].view(frames, 1, 1) + cx = intrinsics[:, 2].view(frames, 1, 1) + cy = intrinsics[:, 3].view(frames, 1, 1) + dirs_cam = torch.stack([(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], dim=-1) + + rotation = poses[:, :3, :3] + translation = poses[:, :3, 3] + dirs_world = torch.einsum("tij,thwj->thwi", rotation, dirs_cam) + dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True).clamp_min(1e-12) + origins = translation.view(frames, 1, 1, 3).expand_as(dirs_world) + + if use_plucker: + moments = torch.cross(origins, dirs_world, dim=-1) + return torch.cat([dirs_world, moments], dim=-1) + return torch.cat([origins, dirs_world], dim=-1) + + +def intrinsics_to_vec4_array(intrinsics: Any, *, num_frames: int, height: int, width: int) -> np.ndarray: + """Normalize the intrinsics payload to an ``(F, 4)`` ``[fx, fy, cx, cy]`` array. + + Accepts either ``None`` (derive from the output resolution) or the + ``{fx, fy, cx, cy}`` mapping — the only form the request contract exposes. + """ + + if intrinsics is None: + focal = float(max(height, width)) + vec = np.array([focal, focal, width / 2.0, height / 2.0], dtype=np.float32) + return np.broadcast_to(vec, (num_frames, 4)).copy() + + if isinstance(intrinsics, dict): + vec = np.array( + [intrinsics["fx"], intrinsics["fy"], intrinsics["cx"], intrinsics["cy"]], + dtype=np.float32, + ) + return np.broadcast_to(vec, (num_frames, 4)).copy() + + raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.") + + +def _pack_camera_conditions( + poses: torch.Tensor, + intrinsics_latent: torch.Tensor, + *, + num_frames: int, + latent_height: int, + latent_width: int, + vae_time_stride: int, +) -> dict[str, torch.Tensor]: + # This yields exactly the latent frame count the VAE produces: + # ``len(arange(0, n, s)) == ceil(n / s) == (n - 1) // s + 1``. + time_indices = torch.arange(0, num_frames, vae_time_stride) + + raymap = torch.cat([poses[time_indices].reshape(len(time_indices), -1), intrinsics_latent[time_indices]], dim=-1) + + chunks: list[torch.Tensor] = [] + spatial_raymap_frames: list[torch.Tensor] = [] + chunk_starts = time_indices - (vae_time_stride - 1) + for start, t_idx in zip(chunk_starts, time_indices): + start_idx = max(0, int(start)) + end_idx = start_idx + vae_time_stride + chunk_poses = poses[start_idx:end_idx] + chunk_intrinsics = intrinsics_latent[start_idx:end_idx] + if chunk_poses.shape[0] < vae_time_stride: + pad = vae_time_stride - chunk_poses.shape[0] + chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)], dim=0) + chunk_intrinsics = torch.cat([chunk_intrinsics, chunk_intrinsics[-1:].repeat(pad, 1)], dim=0) + plucker = compute_raymap(chunk_intrinsics, chunk_poses, latent_height, latent_width, use_plucker=True) + chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_height, latent_width)) + + # Per-frame spatial ray-direction map for raymap_embedder: take the + # representative frame in each chunk (the anchor at t_idx) and use only + # the direction channels (last 3 of 6 in the origin+direction layout). + t = int(t_idx) + od = compute_raymap( + intrinsics_latent[t : t + 1], + poses[t : t + 1], + latent_height, + latent_width, + use_plucker=False, + ) # [1, H, W, 6] + spatial_raymap_frames.append(od[0, :, :, 3:]) # [H, W, 3] — direction only + + chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3) + # spatial_raymap: [3, F_latent, H, W] — consumed by raymap_embedder + spatial_raymap = torch.stack(spatial_raymap_frames).permute(3, 0, 1, 2) + return {"raymap": raymap, "chunk_plucker": chunk_plucker, "spatial_raymap": spatial_raymap} + + +def build_plucker_condition( + condition: SanaWmCameraCondition, + *, + vae_stride: tuple[int, int, int] | list[int] = SANA_WM_DEFAULT_VAE_STRIDE, +) -> dict[str, torch.Tensor]: + """Build native Stage-1 camera tensors from normalized request metadata.""" + + if bool(condition.action) == (condition.poses is not None): + raise ValueError("Sana-WM camera condition requires exactly one of action or poses.") + + if condition.action: + poses_c2w = action_string_to_c2w( + condition.action, + translation_speed=condition.translation_speed, + rotation_speed_deg=condition.rotation_speed_deg, + pitch_limit_deg=condition.pitch_limit_deg, + ) + else: + poses_c2w = np.asarray(condition.poses, dtype=np.float32) + + if poses_c2w.ndim != 3 or poses_c2w.shape[-2:] != (4, 4): + raise ValueError(f"Sana-WM c2w poses must have shape (F, 4, 4), got {poses_c2w.shape}.") + + num_frames = condition.num_frames or int(poses_c2w.shape[0]) + if poses_c2w.shape[0] < num_frames: + if condition.action: + num_frames = int(poses_c2w.shape[0]) + else: + raise ValueError(f"Sana-WM c2w poses length {poses_c2w.shape[0]} is shorter than num_frames {num_frames}.") + poses_c2w = poses_c2w[:num_frames] + + intrinsics = intrinsics_to_vec4_array( + condition.intrinsics, + num_frames=num_frames, + height=condition.height, + width=condition.width, + ) + + vae_time_stride = int(vae_stride[0]) + vae_spatial_stride = int(vae_stride[-1]) + latent_height = condition.height // vae_spatial_stride + latent_width = condition.width // vae_spatial_stride + + poses = torch.from_numpy(poses_c2w).float() + first_inv = get_pose_inverse(poses[0:1]).squeeze(0) + poses_rel = torch.matmul(first_inv, poses[1:]) + poses = torch.cat([torch.eye(4, dtype=poses.dtype).unsqueeze(0), poses_rel], dim=0) + + intrinsics_latent = torch.from_numpy(intrinsics).float() + intrinsics_latent[:, [0, 2]] *= latent_width / float(condition.width) + intrinsics_latent[:, [1, 3]] *= latent_height / float(condition.height) + + return _pack_camera_conditions( + poses, + intrinsics_latent, + num_frames=num_frames, + latent_height=latent_height, + latent_width=latent_width, + vae_time_stride=vae_time_stride, + ) diff --git a/vllm_omni/diffusion/models/sana_wm/config.py b/vllm_omni/diffusion/models/sana_wm/config.py new file mode 100644 index 000000000..bf4eb428c --- /dev/null +++ b/vllm_omni/diffusion/models/sana_wm/config.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Config for the SANA-WM Stage-1 transformer. + +Loaded from the standard diffusers ``transformer/config.json`` (a flat dict +keyed by the dataclass field names). The bespoke ``config.yaml`` release format +is handled once, offline, at checkpoint-conversion time — the runtime only ever +sees the converted, diffusers-standard config. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from pathlib import Path +from typing import Any + + +def _as_tuple3(value: Any, default: tuple[int, int, int]) -> tuple[int, int, int]: + if value is None: + return default + if isinstance(value, int): + return (value, value, value) + if isinstance(value, (list, tuple)): + if len(value) == 2: + return (1, int(value[0]), int(value[1])) + if len(value) == 3: + return (int(value[0]), int(value[1]), int(value[2])) + return default + + +# LTX-2 VAE compression ratios (SANA-WM ships ``AutoencoderKLLTX2Video``). +# They live here rather than in the pipeline so the request-normalisation module +# can validate latent geometry without importing the pipeline. +SANA_WM_VAE_SPATIAL_COMPRESSION = 32 +SANA_WM_VAE_TEMPORAL_COMPRESSION = 8 + + +@dataclass(frozen=True) +class SanaWmConfig: + """SANA-WM Stage-1 architecture and runtime config. + + Defaults mirror the first public HF release. ``from_json`` (reading + ``transformer/config.json``) is the normal construction path. + """ + + architecture_name: str | None = "SanaMSVideoCamCtrl_1600M_P1_D20" + num_blocks: int = 20 + hidden_size: int = 2240 + mlp_ratio: float = 3.0 + attn_type: str = "BidirectionalGDNTriton" + softmax_every_n: int = 4 + linear_head_dim: int = 112 + conv_kernel_size: int = 4 + t_kernel_size: int = 3 + k_conv_only: bool = True + ffn_type: str = "GLUMBConvTemp" + pos_embed_type: str = "wan_rope" + patch_size: tuple[int, int, int] = (1, 1, 1) + qk_norm: bool = True + cross_norm: bool = True + mixed_precision: str = "bf16" + fp32_attention: bool = True + image_size: int = 720 + cam_attn_compress: int = 1 + use_chunk_plucker_post_attn: bool = True + chunk_plucker_channels: int = 48 + chunk_plucker_post_attn_blocks: int = 20 + inference_flow_shift: float = 9.8 + scheduler_type: str = "flow_dpm-solver" + chi_prompt: list[str] = field(default_factory=list) + model_max_length: int = 300 + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> SanaWmConfig: + """Build from a flat mapping of field name -> value. + + Unknown keys (e.g. the diffusers ``_class_name`` marker or the + ``latent_channels`` / ``prompt_channels`` constructor kwargs) are + ignored. ``patch_size`` is coerced from a JSON list back to a tuple. + """ + field_names = {f.name for f in fields(cls)} + known = {k: v for k, v in data.items() if k in field_names} + if "patch_size" in known: + known["patch_size"] = _as_tuple3(known["patch_size"], cls.patch_size) + return cls(**known) + + @classmethod + def from_json(cls, path: str | Path) -> SanaWmConfig: + with open(path, encoding="utf-8") as f: + data = json.load(f) + return cls.from_dict(data if isinstance(data, Mapping) else {}) diff --git a/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py b/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py new file mode 100644 index 000000000..9c2843040 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_wm/pipeline_sana_wm.py @@ -0,0 +1,942 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sana-WM pipeline integration. + +This module wires the registry-visible surface, release-layout validation, and +the Stage-1 sampling loop. The DiT, Gated-DeltaNet and camera stack are +vLLM-Omni-native layers (``sana_wm_transformer.py`` / ``ucpe.py``); nothing here +calls into NVlabs code at runtime. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, ClassVar + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery +from vllm_omni.diffusion.models.ltx2.ltx2_latents import ( + denormalize_latents, + normalize_latents, + pack_latents, + resolve_video_latent_shape, + unpack_latents, +) +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.models.sana_wm.camera_control import ( + SanaWmCameraCondition, + build_plucker_condition, +) +from vllm_omni.diffusion.models.sana_wm.config import ( + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_TEMPORAL_COMPRESSION, + SanaWmConfig, +) +from vllm_omni.diffusion.models.sana_wm.request import normalize_sana_wm_payload +from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import ( + SANA_WM_STAGE1_PROMPT_CHANNELS, + SanaWmTransformer3DModel, +) +from vllm_omni.diffusion.models.schedulers import FlowMatchEulerDiscreteScheduler +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( + DiffusionPipelineProfilerMixin, +) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +logger = init_logger(__name__) + +# The loader expects the standard diffusers layout (transformer/config.json + +# transformer/diffusion_pytorch_model.safetensors + model_index.json), so this +# repo is an offline conversion of the NVlabs release into that layout. It holds +# Stage-1 only — transformer/ + vae/, with model_index.json naming +# SanaWmPipeline — because startup prefetches the whole model repo +# (allow_patterns=["*"]); a repo that also carried the refiner would pull tens of +# GB this path never loads. +SANA_WM_MODEL_ID = "BBBBruce/SANA-WM_bidirectional-stage1-diffusers" + +SANA_WM_STAGE1_DIT_FILE = "transformer/diffusion_pytorch_model.safetensors" +SANA_WM_CONFIG_FILE = "transformer/config.json" +SANA_WM_VAE_CONFIG_FILE = "vae/config.json" +SANA_WM_VAE_WEIGHT_FILE = "vae/diffusion_pytorch_model.safetensors" +SANA_WM_STAGE1_TEXT_ENCODER_ID = "google/gemma-2-2b-it" +SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID = "Efficient-Large-Model/gemma-2-2b-it" +SANA_WM_STAGE1_TEXT_ENCODER_ENV = "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER" +SANA_WM_OUTPUT_HEIGHT = 704 +SANA_WM_OUTPUT_WIDTH = 1280 +SANA_WM_NUM_TRAIN_TIMESTEPS = 1000 +# Maximum supported clip length (20s @16fps) — the native envelope, not a +# request default. Request-side defaults live in +# ``models/sana_wm/request.py``. +SANA_WM_NATIVE_NUM_FRAMES = 321 +# Fail-fast latent-token cap for the native Stage-1 path, sized to the native +# 704x1280 output at the maximum clip length so it only rejects genuinely +# oversized requests. Overridable per-request via ``sana_wm_native_max_tokens``. +SANA_WM_NATIVE_MAX_TOKENS = ( + ((SANA_WM_NATIVE_NUM_FRAMES - 1) // SANA_WM_VAE_TEMPORAL_COMPRESSION + 1) + * (SANA_WM_OUTPUT_HEIGHT // SANA_WM_VAE_SPATIAL_COMPRESSION) + * (SANA_WM_OUTPUT_WIDTH // SANA_WM_VAE_SPATIAL_COMPRESSION) +) +# The NVlabs production inference settings for this release. They belong in the +# model rather than in a deploy YAML: single-stage diffusion models resolve +# through ``create_default_diffusion``, which never reads a YAML's +# ``default_sampling_params``, so a caller that omits either field lands here. +SANA_WM_DEFAULT_NUM_INFERENCE_STEPS = 60 +SANA_WM_DEFAULT_GUIDANCE_SCALE = 5.0 + +SANA_WM_STAGE1_PATTERNS = ( + SANA_WM_CONFIG_FILE, + SANA_WM_STAGE1_DIT_FILE, + SANA_WM_VAE_CONFIG_FILE, + SANA_WM_VAE_WEIGHT_FILE, +) +SANA_WM_STAGE1_DIT_SUBFOLDER = "transformer" +SANA_WM_STAGE1_DIT_BASENAME = Path(SANA_WM_STAGE1_DIT_FILE).name + + +@dataclass(frozen=True) +class SanaWmLocalPaths: + """Resolved local file paths for a SANA-WM snapshot.""" + + root: Path + config: Path + stage1_dit: Path + vae_config: Path + vae_weights: Path + + +@dataclass(frozen=True) +class SanaWmNativeParams: + """Resolved generation settings for one Stage-1 request. + + Built by ``_native_params`` from the request payload and sampling params, + and consumed by ``_run_native_backend`` — this is the production path. + """ + + height: int + width: int + num_frames: int + num_inference_steps: int + seed: int + # Defaults to the NVlabs production scale, so a caller that never mentions + # guidance still gets the two-branch CFG forward the model was tuned for. + # An explicit ``guidance_scale <= 1.0`` still selects the single-branch path. + cfg_scale: float = SANA_WM_DEFAULT_GUIDANCE_SCALE + + +def build_sana_wm_download_patterns() -> tuple[str, ...]: + """Return the minimal HF allow-patterns needed for SANA-WM.""" + + return tuple(SANA_WM_STAGE1_PATTERNS) + + +def resolve_sana_wm_local_paths(snapshot_dir: str | Path) -> SanaWmLocalPaths: + root = Path(snapshot_dir) + return SanaWmLocalPaths( + root=root, + config=root / SANA_WM_CONFIG_FILE, + stage1_dit=root / SANA_WM_STAGE1_DIT_FILE, + vae_config=root / SANA_WM_VAE_CONFIG_FILE, + vae_weights=root / SANA_WM_VAE_WEIGHT_FILE, + ) + + +def validate_sana_wm_local_paths(paths: SanaWmLocalPaths) -> None: + required = [ + paths.config, + paths.stage1_dit, + paths.vae_config, + paths.vae_weights, + ] + missing = [path for path in required if not path.exists()] + if missing: + joined = ", ".join(str(path) for path in missing) + raise FileNotFoundError(f"Missing SANA-WM checkpoint files: {joined}") + + +def resolve_or_download_sana_wm_checkpoint( + model: str = SANA_WM_MODEL_ID, + *, + revision: str | None = None, + cache_dir: str | None = None, +) -> SanaWmLocalPaths: + """Resolve a local SANA-WM tree or download the required HF files.""" + + root = Path(model) + if root.is_dir(): + snapshot_dir = root + else: + snapshot_dir = Path( + download_weights_from_hf_specific( + model, + cache_dir, + list(build_sana_wm_download_patterns()), + revision=revision, + require_all=True, + ) + ) + paths = resolve_sana_wm_local_paths(snapshot_dir) + validate_sana_wm_local_paths(paths) + return paths + + +def build_sana_wm_output_envelope( + *, + output: Any, + output_type: str, + metadata: dict[str, Any], +) -> dict[str, Any]: + """Wrap a Stage-1 result in the canonical output envelope. + + ``normalize_diffusion_postprocess_output`` splits ``{"payload": ..., + "metadata": ...}`` into the API-facing payload and the metadata groups, so + the model-specific diagnostics ride along under a ``sana_wm`` group instead + of the removed ``DiffusionOutput.custom_output`` field. + """ + payload_key = "latents" if output_type == "latent" else "video" + return {"payload": {payload_key: output}, "metadata": {"sana_wm": metadata}} + + +def get_sana_wm_pre_process_func(od_config: OmniDiffusionConfig): + del od_config + + def pre_process_func(request: Any) -> Any: + from vllm_omni.inputs.data import OmniTextPrompt + + sampling_params = getattr(request, "sampling_params", None) + prompt = request.prompt + prompt_mapping = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt + prompt_mapping = dict(prompt_mapping) + + # The camera block travels as extra_params["sana_wm"], which serving + # merges into extra_args. Offline callers may also put it straight on + # the prompt. + if sampling_params is not None: + extra_args = getattr(sampling_params, "extra_args", None) + if isinstance(extra_args, Mapping): + block = extra_args.get("sana_wm") + if isinstance(block, Mapping): + merged = dict(block) + merged.update(prompt_mapping.get("sana_wm") or {}) + prompt_mapping["sana_wm"] = merged + # Geometry always comes from the standard sampling fields. + for field in ("num_frames", "height", "width"): + value = getattr(sampling_params, field, None) + if value is not None and not (field == "num_frames" and value == 1): + prompt_mapping[field] = value + + request.prompt = normalize_sana_wm_payload(prompt_mapping) + + # Preprocess is the only place the request is normalised, so publish the + # canonical geometry back onto sampling_params for the rest of the stack. + if sampling_params is not None: + payload = request.prompt["additional_information"]["sana_wm"] + sampling_params.num_frames = payload["num_frames"] + sampling_params.height = payload["height"] + sampling_params.width = payload["width"] + return request + + return pre_process_func + + +class SanaWmPipeline( + nn.Module, + CFGParallelMixin, + SupportImageInput, + SupportsComponentDiscovery, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, +): + """Stage-1 SANA-WM image-to-video pipeline.""" + + support_image_input: ClassVar[bool] = True + color_format: ClassVar[str] = "RGB" + # Skip the startup warmup: the generic dummy request carries a blank image + # with no camera poses and no action string, which SANA-WM's request + # validation always rejects. The first real request does the warming. + dummy_run_num_frames: ClassVar[int] = 0 + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + + def __init__(self, *, od_config: OmniDiffusionConfig | None = None, prefix: str = "") -> None: + super().__init__() + self.od_config = od_config + self.prefix = prefix + + self.sana_wm_config = SanaWmConfig() + self.quant_config = getattr(od_config, "quantization_config", None) if od_config is not None else None + + self.tokenizer: Any | None = None + self.text_encoder: nn.Module | None = None + self.vae: nn.Module | None = None + self.release_paths: SanaWmLocalPaths | None = None + self.weights_sources = [] + if od_config is not None and od_config.model is not None: + # Resolve the checkpoint once at construction (the standard + # startup-time flow, like every other diffusion pipeline) so the + # transformer below is built with the transformer/config.json shape + # rather than defaults. + self.resolve_checkpoint() + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder=SANA_WM_STAGE1_DIT_SUBFOLDER, + revision=od_config.revision, + # The subfolder file carries module-local keys (standard + # diffusers layout); prefix them into the pipeline + # namespace for the loader's strict-coverage check. + prefix="transformer.", + fall_back_to_pt=False, + allow_patterns_overrides=[SANA_WM_STAGE1_DIT_BASENAME], + ) + ] + # Built eagerly: the loader constructs the pipeline under the target + # device / default-dtype context, and the weight loader then streams + # checkpoint tensors directly into these modules. + self.transformer = SanaWmTransformer3DModel( + config=self.sana_wm_config, + quant_config=self.quant_config, + prefix=f"{prefix}.transformer" if prefix else "transformer", + ) + self.device = get_local_device() + if od_config is not None and od_config.model is not None: + self._build_aux_components() + + def resolve_checkpoint(self) -> SanaWmLocalPaths: + if self.release_paths is not None: + return self.release_paths + if self.od_config is None or self.od_config.model is None: + raise ValueError("Sana-WM checkpoint resolution requires od_config.model.") + self.release_paths = resolve_or_download_sana_wm_checkpoint( + self.od_config.model, + revision=self.od_config.revision, + ) + self.sana_wm_config = SanaWmConfig.from_json(self.release_paths.config) + return self.release_paths + + @staticmethod + def _extra_args(sampling_params: Any | None) -> dict[str, Any]: + extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None + return dict(extra_args or {}) + + def _native_params(self, payload: dict[str, Any], sampling_params: Any | None) -> SanaWmNativeParams: + extra_args = self._extra_args(sampling_params) + height = int(getattr(sampling_params, "height", None) or payload["height"]) + width = int(getattr(sampling_params, "width", None) or payload["width"]) + num_frames = int(getattr(sampling_params, "num_frames", None) or payload["num_frames"]) + steps = int( + getattr(sampling_params, "num_inference_steps", None) + or extra_args.get("num_inference_steps") + or SANA_WM_DEFAULT_NUM_INFERENCE_STEPS + ) + seed = int(getattr(sampling_params, "seed", None) or extra_args.get("seed", 0)) + if min(height, width, num_frames, steps) <= 0: + raise ValueError("Sana-WM native height, width, num_frames, and steps must be positive.") + cfg_scale = SANA_WM_DEFAULT_GUIDANCE_SCALE + if sampling_params is not None and getattr(sampling_params, "guidance_scale_provided", False): + cfg_scale = float(getattr(sampling_params, "guidance_scale", 1.0) or 1.0) + return SanaWmNativeParams( + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=steps, + seed=seed, + cfg_scale=cfg_scale, + ) + + def _native_dtype(self, device: torch.device) -> torch.dtype: + dtype = getattr(self.od_config, "dtype", None) if self.od_config is not None else None + if isinstance(dtype, torch.dtype): + return dtype + return torch.float32 if device.type == "cpu" else torch.bfloat16 + + def _runtime_device_dtype(self) -> tuple[torch.device, torch.dtype]: + return self.device, self._native_dtype(self.device) + + def _build_aux_components(self) -> None: + """Construct the text encoder and VAE at startup, on CPU. + + ``ModuleDiscovery`` skips attributes that are ``None``, so components + created lazily on the first request are invisible to sequential + offload, HSDP auxiliary placement, and the shared device/memory + lifecycle. Building them here and deferring placement to the common + logic mirrors LTX-2's ``initialize_pipeline_components``. + """ + dtype = self._native_dtype(self.device) + self._load_stage1_text_encoder(dtype=dtype) + self._load_vae(dtype=dtype) + self._place_aux_components() + + def _place_aux_components(self) -> None: + """Move the discovered auxiliary components onto the local device. + + Skipped when offload or HSDP owns placement — those paths expect the + components to still be on CPU when they take over. + """ + parallel_config = getattr(self.od_config, "parallel_config", None) + if ( + getattr(self.od_config, "enable_cpu_offload", False) + or getattr(self.od_config, "enable_layerwise_offload", False) + or getattr(parallel_config, "use_hsdp", False) + ): + return + modules = ModuleDiscovery.discover(self) + for module in (*modules.encoders, *modules.vaes, *modules.resident_modules): + module.to(self.device) + + @staticmethod + def _component_local_files_only(model_id: Any) -> bool: + """Decide download policy per component, from its own model id. + + The Stage-1 release directory is local once ``resolve_checkpoint`` has + run, but the Gemma text encoder lives in a separate Hub repo. Deriving + locality from the release paths forced ``local_files_only=True`` onto + Gemma as well, which forbids the download on a box that has not cached + it yet. + """ + return Path(str(model_id)).expanduser().exists() + + def _load_stage1_text_encoder(self, *, dtype: torch.dtype) -> None: + from transformers import AutoModel, AutoTokenizer + + model_ids = [os.environ.get(SANA_WM_STAGE1_TEXT_ENCODER_ENV, "").strip()] + model_ids.extend([SANA_WM_STAGE1_TEXT_ENCODER_ID, SANA_WM_STAGE1_TEXT_ENCODER_FALLBACK_ID]) + errors: list[str] = [] + for model_id in dict.fromkeys(model_id for model_id in model_ids if model_id): + try: + self.tokenizer = AutoTokenizer.from_pretrained( + model_id, + local_files_only=self._component_local_files_only(model_id), + ) + text_encoder_model_id = model_id + break + except OSError as exc: + errors.append(f"{model_id}: {exc}") + else: + raise OSError("Could not load SANA-WM Stage-1 text encoder. Tried: " + "; ".join(errors)) + if getattr(self.tokenizer, "pad_token", None) is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + self.tokenizer.padding_side = "right" + with torch.device("cpu"): + # AutoModel, not AutoModelForCausalLM: only the final hidden state is + # used, so the LM head would compute a [seq, 256k] logits tensor + # that is thrown away on every prompt encode. + self.text_encoder = AutoModel.from_pretrained( + text_encoder_model_id, + torch_dtype=dtype, + local_files_only=self._component_local_files_only(text_encoder_model_id), + ) + + def _ensure_stage1_text_encoder(self, *, device: torch.device, dtype: torch.dtype) -> None: + if self.text_encoder is None: + self._load_stage1_text_encoder(dtype=dtype) + self.text_encoder.to(device=device, dtype=dtype) + + def _load_vae(self, *, dtype: torch.dtype) -> None: + if self.release_paths is None and self.od_config is not None and self.od_config.model is not None: + self.resolve_checkpoint() + if self.release_paths is None: + raise ValueError("Sana-WM VAE loading requires a resolved checkpoint.") + + from diffusers import AutoencoderKLLTX2Video + + with torch.device("cpu"): + self.vae = AutoencoderKLLTX2Video.from_pretrained( + str(self.release_paths.root), + subfolder="vae", + torch_dtype=dtype, + local_files_only=self._component_local_files_only(self.release_paths.root), + ) + self.vae.use_framewise_encoding = True + self.vae.use_framewise_decoding = True + # NVlabs reads these from its YAML VAE config with 64/96 defaults; the + # diffusers VAE config.json only ships scaling_factor, so the defaults + # are load-bearing (the keys are genuinely absent here). + self.vae.tile_sample_stride_num_frames = int(getattr(self.vae.config, "tile_sample_stride_num_frames", 64)) + self.vae.tile_sample_min_num_frames = int(getattr(self.vae.config, "tile_sample_min_num_frames", 96)) + + def _ensure_vae(self, *, device: torch.device, dtype: torch.dtype) -> None: + if self.vae is None: + self._load_vae(dtype=dtype) + # Tiling is re-asserted here, not in _load_vae: registry.py assigns + # ``model.vae.use_tiling = od_config.vae_use_tiling`` after the pipeline + # is constructed, and that flag defaults to False. Matching NVlabs' + # inference_video_scripts/inference_sana_wm.py, the SANA VAE has to + # tile — a 161-frame decode costs about 9 GiB more without it, and a + # 321-frame request (41 latent frames) decodes as one 3D volume and + # OOMs. Overriding is deliberate, but say so rather than dropping the + # caller's setting silently. + if not self.vae.use_tiling: + logger.warning("Sana-WM forces VAE tiling on; ignoring vae_use_tiling=False (the VAE OOMs without it).") + self.vae.use_tiling = True + self.vae.to(device=device, dtype=dtype) + + @staticmethod + def _preprocess_first_frame( + image: Any, + *, + height: int, + width: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Return the first-frame image as ``[1, 3, 1, H, W]`` in ``[-1, 1]``. + + ``VideoProcessor.preprocess`` covers the PIL / ndarray / tensor cases. + ``resample="bicubic"`` keeps the filter PIL's ``Image.resize`` defaulted + to; the processor's own default is lanczos, which would move the + conditioning latent. + """ + import numpy as np + from diffusers.video_processor import VideoProcessor + + # The processor normalises to [-1, 1] assuming float inputs in [0, 1]. + if isinstance(image, np.ndarray) and image.dtype == np.uint8: + image = image.astype(np.float32) / 255.0 + + processor = VideoProcessor(vae_scale_factor=SANA_WM_VAE_SPATIAL_COMPRESSION, resample="bicubic") + frame = processor.preprocess(image, height=height, width=width) + return frame.unsqueeze(2).to(device=device, dtype=dtype) + + def _vae_normalize_latent(self, latent: torch.Tensor) -> torch.Tensor: + """LTX-2 VAE per-channel latent normalisation matching NVlabs. + + ``z_norm = (z_raw - latents_mean) * scaling_factor / latents_std`` + + The LTX-2 VAE (``AutoencoderKLLTX2Video`` in diffusers) ships with + per-channel ``latents_mean`` and ``latents_std`` tensors that must be + applied alongside ``scaling_factor`` for the round-trip to be + identity. Applying `* scaling_factor` alone gives decoded videos a + systematic G/B colour shift (pred RGB means + ~(110, 220, 208) vs ref ~(94, 114, 138)). + + Returns the normalised latent in the same dtype as the input. + """ + if self.vae is None: + raise RuntimeError("Sana-WM VAE did not initialize.") + return normalize_latents( + latent, + self.vae.latents_mean, + self.vae.latents_std, + float(self.vae.config.scaling_factor), + ) + + def _vae_denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: + """Inverse of :meth:`_vae_normalize_latent` — match NVlabs decode.""" + if self.vae is None: + raise RuntimeError("Sana-WM VAE did not initialize.") + return denormalize_latents( + latent, + self.vae.latents_mean, + self.vae.latents_std, + float(self.vae.config.scaling_factor), + ) + + def _vae_encode_first_frame( + self, + image: Any, + *, + height: int, + width: int, + latent_height: int, + latent_width: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Encode the first-frame image to latent space via the LTX-2 VAE. + + Returns shape ``[1, 128, 1, latent_height, latent_width]``. + """ + self._ensure_vae(device=device, dtype=dtype) + if self.vae is None: + raise RuntimeError("Sana-WM VAE did not initialize.") + frame = self._preprocess_first_frame(image, height=height, width=width, device=device, dtype=dtype) + with torch.inference_mode(): + encoded = self.vae.encode(frame) + first_latent = encoded.latent_dist.mean + first_latent = self._vae_normalize_latent(first_latent.to(dtype=dtype)) + # The encoder's own rounding can land a pixel off the latent grid the + # noise tensor was built on, and the two are concatenated below. + if first_latent.shape[-2:] != (latent_height, latent_width): + first_latent = F.interpolate( + first_latent.squeeze(2), + size=(latent_height, latent_width), + mode="bilinear", + align_corners=False, + ).unsqueeze(2) + return first_latent # [1, 128, 1, lh, lw] + + def _stage1_prompt_text(self, prompt: dict[str, Any]) -> str: + user_prompt = str(prompt.get("prompt") or "") + chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part) + if chi_prompt: + return chi_prompt + user_prompt + return user_prompt + + def _native_prompt_embeds( + self, + prompt: dict[str, Any], + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + prompt_text = self._stage1_prompt_text(prompt) + self._ensure_stage1_text_encoder(device=device, dtype=dtype) + tokenizer = self.tokenizer + if tokenizer is None or self.text_encoder is None: + raise RuntimeError("Sana-WM Stage-1 text encoder did not initialize.") + chi_prompt = "\n".join(part for part in self.sana_wm_config.chi_prompt if part) + model_max_length = int(self.sana_wm_config.model_max_length) + max_length_all = model_max_length + if chi_prompt: + # NVlabs encodes the chi-prompt-prefixed text with extra room for + # the chi tokens, then keeps BOS plus the final model window. + max_length_all = len(tokenizer.encode(chi_prompt)) + model_max_length - 2 + encoded = tokenizer( + [prompt_text], + padding="max_length", + max_length=max_length_all, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ).to(device) + # last_hidden_state is the same tensor the CausalLM wrapper exposed as + # hidden_states[-1] (both are post-final-norm), without retaining every + # intermediate layer. + hidden_states = self.text_encoder(**encoded).last_hidden_state + attention_mask = encoded.attention_mask + if chi_prompt: + select_index = [0] + list(range(-model_max_length + 1, 0)) + hidden_states = hidden_states[:, select_index] + attention_mask = attention_mask[:, select_index] + if hidden_states.shape[-1] != SANA_WM_STAGE1_PROMPT_CHANNELS: + raise ValueError( + "Sana-WM Stage-1 Gemma hidden size mismatch: expected " + f"{SANA_WM_STAGE1_PROMPT_CHANNELS}, got {hidden_states.shape[-1]}." + ) + # Return RAW Gemma hidden states: the transformer's ``attention_y_norm`` + # RMSNorm does the normalisation, so normalising here double-normalises. + prompt_attention_mask = attention_mask.to(device=device, dtype=torch.float32) + # An all-ones mask carries no information. Keep it ``None`` so + # cross-attention takes the mask-free attention path instead of handing + # every block, on every step, a mask the backend must reshape to no + # effect. Evaluated once per prompt encode, so the device sync is paid + # twice per request (cond + uncond), not once per denoising step. + if bool(prompt_attention_mask.all()): + prompt_attention_mask = None + self._last_prompt_attention_mask = prompt_attention_mask + return hidden_states.to(device=device, dtype=dtype) + + def predict_noise(self, **kwargs: Any) -> torch.Tensor: + """Single transformer forward. + + Overrides the mixin default, which assumes a tuple-returning transformer + and takes ``result[0]``; SANA-WM's returns the noise prediction directly. + """ + return self.transformer(**kwargs) + + def _decode_native_latents( + self, + latents: torch.Tensor, + *, + output_type: str, + device: torch.device, + dtype: torch.dtype, + ) -> Any: + if output_type == "latent": + return latents + self._ensure_vae(device=device, dtype=dtype) + if self.vae is None: + raise RuntimeError("Sana-WM VAE did not initialize.") + # NVlabs always passes ``temb=None`` to ``vae.decode`` for this VAE + # (``builder.py::vae_decode -> LTX2VAE_diffusers``); match that. + denorm = self._vae_denormalize_latent(latents.to(getattr(self.vae, "dtype", dtype))) + video = self.vae.decode(denorm, temb=None, return_dict=False)[0] + # Post-process to the requested output_type (np/pil/pt). Let failures + # surface: silently returning the raw decoder tensor would hand the + # caller a wrong-format/range video with no error. + from diffusers.video_processor import VideoProcessor + + processor = VideoProcessor(vae_scale_factor=getattr(self.vae, "spatial_compression_ratio", 32)) + return processor.postprocess_video(video, output_type=output_type) + + def _run_native_backend( + self, + *, + prompt: dict[str, Any], + payload: dict[str, Any], + sampling_params: Any | None, + ) -> DiffusionOutput: + params = self._native_params(payload, sampling_params) + latent_frames, latent_height, latent_width = resolve_video_latent_shape( + params.height, + params.width, + params.num_frames, + vae_spatial_compression_ratio=SANA_WM_VAE_SPATIAL_COMPRESSION, + vae_temporal_compression_ratio=SANA_WM_VAE_TEMPORAL_COMPRESSION, + ) + token_count = latent_frames * latent_height * latent_width + extra_args = self._extra_args(sampling_params) + max_tokens = int(extra_args.get("sana_wm_native_max_tokens", SANA_WM_NATIVE_MAX_TOKENS)) + if token_count > max_tokens: + raise ValueError( + "Sana-WM native latent token count exceeds the configured cap. " + f"Requested latent tokens={token_count}, max={max_tokens}. Request a smaller " + "`height`/`width`/`num_frames`, or raise the cap via " + "`sana_wm_native_max_tokens`." + ) + + device, dtype = self._runtime_device_dtype() + generator = torch.Generator(device=device) + generator.manual_seed(params.seed) + noise = torch.randn( + (1, 128, latent_frames, latent_height, latent_width), + device=device, + dtype=dtype, + generator=generator, + ) + + # Shared flow-matching Euler scheduler. NVlabs ``LTXFlowEuler`` drives + # diffusers' ``FlowMatchEulerDiscreteScheduler``; this is vLLM-Omni's + # vendored copy of it, including the ``per_token_timesteps`` step branch. + # The schedule applies the shift twice (once via the shifted ``sigma_min`` + # persisted by ``__init__``, once in ``set_timesteps``), which is what + # reproduces the NVlabs timestep table (e.g. shift=9.8, N=3 -> + # [1000, 909.0, 87.7]). + scheduler = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=SANA_WM_NUM_TRAIN_TIMESTEPS, + shift=self.sana_wm_config.inference_flow_shift, + ) + scheduler.set_timesteps(params.num_inference_steps, device=device) + timesteps = scheduler.timesteps + + # ``normalize_sana_wm_payload`` guarantees the key is present but says + # nothing about its type. Falling back to pure noise on an unusable + # image would silently drop the conditioning and produce a + # plausible-looking but wrong video, so this raises instead. + import numpy as _np + + first_frame_image = (prompt.get("multi_modal_data") or {}).get("image") + if not (hasattr(first_frame_image, "convert") or isinstance(first_frame_image, (_np.ndarray, torch.Tensor))): + raise TypeError( + "Sana-WM first-frame image must be a PIL Image, numpy ndarray, or " + f"torch.Tensor; got {type(first_frame_image).__name__}." + ) + # Per-frame timestep contract matching NVlabs ``LTXFlowEuler.sample``: + # frame 0 is the clean VAE-encoded conditioning latent and its timestep + # is held at 0 while the sampling loop drives the other frames from + # noise. + first_latent = self._vae_encode_first_frame( + first_frame_image, + height=params.height, + width=params.width, + latent_height=latent_height, + latent_width=latent_width, + device=device, + dtype=dtype, + ) + # Place the CLEAN VAE-encoded first frame at frame 0 and rely on + # per-frame timesteps + the `condition_mask` torch.where restore (below) + # to keep it invariant across denoising steps. + latents = torch.cat([first_latent, noise[:, :, 1:]], dim=2) + + prompt_embeds = self._native_prompt_embeds( + prompt, + device=device, + dtype=dtype, + ) + prompt_attention_mask = self._last_prompt_attention_mask + + # Under CFG, encode the negative prompt up front; each denoise step + # then runs a two-branch transformer forward. The negative prompt comes + # from the request (``prompt["negative_prompt"]``, set by + # ``serving_video``); absent, it defaults to the empty string matching + # NVlabs ``GenerationParams.negative_prompt = ""``. + do_cfg = params.cfg_scale > 1.0 + if do_cfg: + negative_prompt_obj = {"prompt": str(prompt.get("negative_prompt") or "")} + negative_prompt_embeds = self._native_prompt_embeds( + negative_prompt_obj, + device=device, + dtype=dtype, + ) + negative_prompt_attention_mask = self._last_prompt_attention_mask + # Restore the cond mask so later readers of + # `_last_prompt_attention_mask` see the positive branch. + self._last_prompt_attention_mask = prompt_attention_mask + else: + negative_prompt_embeds = None + negative_prompt_attention_mask = None + + camera = payload.get("camera") or {} + condition = SanaWmCameraCondition( + poses=camera.get("poses") if isinstance(camera, dict) else None, + intrinsics=payload.get("intrinsics"), + action=payload.get("action"), + num_frames=params.num_frames, + height=params.height, + width=params.width, + translation_speed=float(payload.get("translation_speed", 0.05)), + rotation_speed_deg=float(payload.get("rotation_speed_deg", 1.2)), + ) + camera_tensors = build_plucker_condition(condition) + plucker = camera_tensors["chunk_plucker"].to(device=device, dtype=dtype) + raymap = camera_tensors["raymap"].to(device=device, dtype=dtype) + spatial_raymap = camera_tensors.get("spatial_raymap") + if spatial_raymap is not None: + spatial_raymap = spatial_raymap.to(device=device, dtype=dtype) + + # Build a per-frame condition mask matching NVlabs ``LTXFlowEuler``: + # frame 0 is the conditioning frame at sigma=0 (preserved after + # each ``scheduler.step``), all other frames carry the current + # sampling sigma. We omit the legacy `add_noise_to_image_conditioning_latents` + # motion-continuity term because the public SANA-WM config sets + # ``condition_frame_info={0: 0.0}`` (image_cond_noise_scale=0). + cam_batch = latents.shape[0] + cam_frames = latents.shape[2] + condition_mask = torch.zeros(cam_batch, 1, cam_frames, 1, 1, device=latents.device, dtype=latents.dtype) + condition_mask[:, :, 0] = 1.0 + + for timestep in timesteps: + # (B, 1, F) per-frame timestep, frame 0 forced to 0. Stays fp32: + # casting through the bf16 latent dtype would quantise it before + # ``SanaWmTimestepEmbedder``'s sinusoidal embed. + model_timestep = timestep.float().expand(cam_batch, 1, cam_frames).clone() + model_timestep[:, :, 0] = 0.0 + positive_kwargs = { + "hidden_states": latents, + "timestep": model_timestep, + "encoder_hidden_states": prompt_embeds, + "encoder_attention_mask": prompt_attention_mask, + "plucker": plucker, + "raymap": raymap, + "spatial_raymap": spatial_raymap, + } + negative_kwargs = ( + { + **positive_kwargs, + "encoder_hidden_states": negative_prompt_embeds, + "encoder_attention_mask": negative_prompt_attention_mask, + } + if do_cfg + else None + ) + # ``cfg_normalize=False`` keeps the plain NVlabs combine, + # ``uncond + scale * (cond - uncond)``. Under cfg_parallel_size > 1 + # the mixin runs one branch per rank and all-gathers instead of + # evaluating both branches serially. + noise_pred = self.predict_noise_maybe_with_cfg( + do_cfg, + params.cfg_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize=False, + ) + + # The per-frame contract pairs with a per-token flow-matching Euler + # step. Both step branches end at ``prev = sample + dt * + # model_output``, but the scheduler defines dt as + # ``sigma - sigma_next`` per-token against ``sigma_next - sigma`` + # standard, so the per-token call takes the sign-flipped prediction + # (matching NVlabs ``LTXFlowEuler.sample``); passing the standard + # sign here would integrate towards more noise. + # + # Broadcast the (B, 1, F) timestep to (B, 1, F, H, W) and flatten + # F*H*W; frame-0 conditioning tokens are already 0. The flatten + # order (F, H, W) matches ``pack_latents``. + pt_t = ( + model_timestep.unsqueeze(-1) + .unsqueeze(-1) + .expand(cam_batch, 1, cam_frames, latents.shape[3], latents.shape[4]) + .reshape(cam_batch, -1) + ) + # The per-token branch consumes/returns the packed (B, N, C) layout + # and keeps its result in fp32; cast back so the latent dtype stays + # stable across steps. + stepped_packed = scheduler.step( + -pack_latents(noise_pred), + timestep, + pack_latents(latents), + per_token_timesteps=pt_t, + return_dict=False, + )[0] + stepped = unpack_latents( + stepped_packed.to(latents.dtype), + cam_frames, + latents.shape[3], + latents.shape[4], + ) + + # Match NVlabs LTXFlowEuler exactly. This is stricter than a plain + # condition-frame restore: with bf16 latents, the `t=1000` + # comparison rounds `1 - 1e-6` to `1`, so the first full-noise step + # is discarded for generated tokens as well. + tokens_to_denoise_mask = timestep / float(SANA_WM_NUM_TRAIN_TIMESTEPS) - 1e-6 < (1.0 - condition_mask) + latents = torch.where(tokens_to_denoise_mask, stepped, latents) + + output_type = getattr(sampling_params, "output_type", None) or "np" + output = self._decode_native_latents(latents, output_type=output_type, device=device, dtype=dtype) + return DiffusionOutput( + output=build_sana_wm_output_envelope( + output=output, + output_type=output_type, + metadata={ + "backend": "native_gdn", + "output_space": output_type, + "chi_prompt_applied": bool(self.sana_wm_config.chi_prompt), + "num_frames": params.num_frames, + "height": params.height, + "width": params.width, + "latent_tokens": token_count, + "sampling_steps": params.num_inference_steps, + }, + ), + ) + + def forward(self, req: DiffusionRequestBatch, *args: Any, **kwargs: Any) -> DiffusionOutput: + del args, kwargs + if len(req.prompts) != 1: + raise ValueError("Sana-WM native backend currently supports exactly one prompt per request.") + + prompt = req.prompts[0] + if isinstance(prompt, str): + raise ValueError("Sana-WM requires a mapping prompt with first-frame image and camera/action metadata.") + # get_sana_wm_pre_process_func is the single canonical entry point: it + # normalises the request and republishes the geometry onto + # sampling_params. normalize_sana_wm_payload is idempotent, so this only + # covers offline callers that bypass the preprocess hook. + prompt = normalize_sana_wm_payload(prompt) + payload = prompt["additional_information"]["sana_wm"] + + # The checkpoint is resolved once in __init__; this is a cached no-op + # kept as a guard for pipelines constructed without od_config.model. + if self.od_config is not None and self.od_config.model is not None: + self.resolve_checkpoint() + + return self._run_native_backend( + prompt=prompt, + payload=payload, + sampling_params=req.sampling_params, + ) + + def load_weights(self, weights: Iterable[tuple[str, Any]]) -> set[str]: + # Keys arrive pipeline-prefixed (``transformer.*``, per the component + # source). AutoWeightsLoader streams them to the submodule and returns + # the prefixed names the loader's strict-coverage check compares. + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/sana_wm/request.py b/vllm_omni/diffusion/models/sana_wm/request.py new file mode 100644 index 000000000..010675045 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_wm/request.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request normalization for Sana-WM image-to-video. + +The released SANA-WM checkpoint is first-frame image-to-video. This module keeps +request-side normalization lightweight: it validates the image/camera/action +contract and stores a canonical payload under +``additional_information["sana_wm"]``. Model-internal raymap / Plucker +projection lives with the transformer. + +This lives beside the model rather than under ``stage_input_processors`` +because nothing loads it as one: the pipeline's ``pre_process_func`` is +registered from ``pipeline_sana_wm``, the same as every other diffusion model, +and no deploy config wires a ``custom_process_input_func`` here. Keeping it in +the stage package forced the pipeline to import *upwards* — the only such +import in ``diffusion/models/`` — which in turn made the VAE compression +constants impossible to share and so duplicated. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from typing import Any + +from vllm_omni.diffusion.models.sana_wm.camera_control import action_rollout_num_frames +from vllm_omni.diffusion.models.sana_wm.config import ( + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_TEMPORAL_COMPRESSION, +) + +SANA_WM_CANONICAL_KEY = "sana_wm" +# Request fallbacks used whenever the field is absent, mirroring the deploy YAML +# (161 frames) and the model's native resolution. Both are genuinely reachable: +# a served request that omits num_frames never gets one injected, and callers +# may ask for a different resolution. +SANA_WM_DEFAULT_NUM_FRAMES = 161 +SANA_WM_DEFAULT_HEIGHT = 704 +SANA_WM_DEFAULT_WIDTH = 1280 +SANA_WM_DEFAULT_CAMERA_FORMAT = "c2w_4x4" +SANA_WM_DEFAULT_COORDINATE_SYSTEM = "official" + + +def _unwrap_single(value: Any) -> Any: + if isinstance(value, list): + return value[0] if value else None + return value + + +def _as_dict(value: Any, *, name: str) -> dict[str, Any]: + if value is None: + return {} + if not isinstance(value, Mapping): + raise ValueError(f"Sana-WM {name} must be a mapping, got {type(value).__name__}.") + return dict(value) + + +def _is_sequence(value: Any) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)) + + +def _as_finite_float(value: Any, *, name: str) -> float: + """Parse a float and reject NaN/inf. + + Non-finite geometry does not fail loudly downstream: it propagates through + the raymap into the Plucker maps and comes back as an all-NaN video with a + 200 response. + """ + try: + parsed = float(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Sana-WM {name} must be a number, got {value!r}.") from exc + if not math.isfinite(parsed): + raise ValueError(f"Sana-WM {name} must be finite, got {parsed!r}.") + return parsed + + +def _validate_numeric_matrix4x4(matrix: Any, *, name: str) -> None: + if not _is_sequence(matrix) or len(matrix) != 4: + raise ValueError(f"Sana-WM {name} must be a 4x4 matrix.") + for row_idx, row in enumerate(matrix): + if not _is_sequence(row) or len(row) != 4: + raise ValueError(f"Sana-WM {name}[{row_idx}] must contain 4 values.") + for col_idx, value in enumerate(row): + _as_finite_float(value, name=f"{name}[{row_idx}][{col_idx}]") + + +def _validate_camera_poses(poses: Any, *, num_frames: int | None = None) -> None: + if not _is_sequence(poses) or len(poses) == 0: + raise ValueError("Sana-WM camera poses must be a non-empty sequence of 4x4 matrices.") + if num_frames is not None and len(poses) != num_frames: + raise ValueError(f"Sana-WM camera poses length {len(poses)} must equal num_frames {num_frames}.") + for idx, pose in enumerate(poses): + _validate_numeric_matrix4x4(pose, name=f"camera.poses[{idx}]") + + +def _validate_intrinsics(intrinsics: Any) -> None: + # Only the ``{fx, fy, cx, cy}`` mapping is accepted (the one form the recipe + # and examples use). Camera intrinsics are optional; when omitted the model + # derives them from the output resolution. + if intrinsics is None: + return + if not isinstance(intrinsics, Mapping): + raise ValueError("Sana-WM intrinsics must be a {fx, fy, cx, cy} mapping.") + required = {"fx", "fy", "cx", "cy"} + missing = sorted(required - set(intrinsics)) + if missing: + raise ValueError(f"Sana-WM intrinsics mapping is missing keys: {missing}.") + for key in required: + value = _as_finite_float(intrinsics[key], name=f"intrinsics[{key!r}]") + # ``compute_raymap`` divides by fx/fy. cx/cy stay unconstrained beyond + # finiteness — a principal point outside the frame is unusual, not wrong. + if key in ("fx", "fy") and value <= 0: + raise ValueError(f"Sana-WM intrinsics[{key!r}] must be positive, got {value}.") + + +def _as_positive_int(value: Any, *, name: str) -> int: + try: + parsed = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Sana-WM {name} must be an integer, got {value!r}.") from exc + if parsed <= 0: + raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.") + return parsed + + +def _as_positive_float(value: Any, *, name: str) -> float: + # Finiteness first: every comparison against NaN is False, so a bare + # ``parsed <= 0`` would accept it. + parsed = _as_finite_float(value, name=name) + if parsed <= 0: + raise ValueError(f"Sana-WM {name} must be positive, got {parsed}.") + return parsed + + +def _extract_image(mm_data: Mapping[str, Any]) -> Any: + return _unwrap_single(mm_data.get("image")) + + +def _first_present(*values: Any) -> Any: + for value in values: + if value is not None: + return value + return None + + +def _validate_latent_geometry(*, num_frames: int, height: int, width: int) -> None: + """Reject geometry the VAE would silently floor. + + The SANA-WM VAE compresses 32x spatially and 8x temporally, so a request + that is not aligned loses pixels/frames instead of failing. + """ + if height % SANA_WM_VAE_SPATIAL_COMPRESSION or width % SANA_WM_VAE_SPATIAL_COMPRESSION: + raise ValueError( + f"Sana-WM height and width must be divisible by {SANA_WM_VAE_SPATIAL_COMPRESSION}, " + f"got {height} and {width}." + ) + if (num_frames - 1) % SANA_WM_VAE_TEMPORAL_COMPRESSION: + raise ValueError( + f"Sana-WM num_frames must satisfy (num_frames - 1) % {SANA_WM_VAE_TEMPORAL_COMPRESSION} == 0, " + f"got {num_frames}." + ) + + +def normalize_sana_wm_payload(prompt: Mapping[str, Any]) -> dict[str, Any]: + """Return a prompt copy with canonical Sana-WM request metadata. + + The ``sana_wm`` block is read from the top-level ``sana_wm`` key, falling + back to ``additional_information["sana_wm"]`` so the function is idempotent + when called again on its own output. The first-frame image is read from + ``multi_modal_data["image"]``. + """ + + if not isinstance(prompt, Mapping): + raise ValueError(f"Sana-WM prompt must be a mapping, got {type(prompt).__name__}.") + + result = dict(prompt) + mm_data = _as_dict(result.get("multi_modal_data"), name="multi_modal_data") + image = _extract_image(mm_data) + if image is None: + raise ValueError("Sana-WM requires a first-frame image in `multi_modal_data['image']`.") + + additional = _as_dict(result.get("additional_information"), name="additional_information") + raw = _first_present( + result.get(SANA_WM_CANONICAL_KEY), + additional.get(SANA_WM_CANONICAL_KEY), + {}, + ) + raw = _as_dict(raw, name=SANA_WM_CANONICAL_KEY) + + num_frames = _as_positive_int( + _first_present(raw.get("num_frames"), result.get("num_frames"), SANA_WM_DEFAULT_NUM_FRAMES), + name="num_frames", + ) + height = _as_positive_int( + _first_present(raw.get("height"), result.get("height"), SANA_WM_DEFAULT_HEIGHT), + name="height", + ) + width = _as_positive_int(_first_present(raw.get("width"), result.get("width"), SANA_WM_DEFAULT_WIDTH), name="width") + _validate_latent_geometry(num_frames=num_frames, height=height, width=width) + + camera = raw.get("camera") + camera = _as_dict(camera, name="camera") if camera is not None else None + action = raw.get("action") + if action is not None: + if not isinstance(action, str) or not action.strip(): + raise ValueError("Sana-WM action must be a non-empty string when provided.") + # Rejected here rather than in the model: a mismatch used to shorten the + # camera tensor silently and then fail deep in the transformer on a + # camera/latent frame mismatch, and an unbounded duration used to be + # expanded one entry per frame before any of that. + rollout_frames = action_rollout_num_frames(action) + if rollout_frames != num_frames: + raise ValueError( + f"Sana-WM action rolls out {rollout_frames} frames but num_frames is {num_frames}. " + "Action segment durations must sum to num_frames - 1." + ) + if camera is not None and action is not None: + raise ValueError("Sana-WM accepts exactly one of `camera` or `action`, not both.") + if camera is None and action is None: + raise ValueError("Sana-WM requires either explicit camera poses or an action DSL string.") + + canonical_camera = None + if camera is not None: + poses = camera.get("poses") + if poses is None: + raise ValueError("Sana-WM camera payload requires `poses`.") + _validate_camera_poses(poses, num_frames=num_frames) + # Only the defaults are accepted: nothing downstream reads these two + # fields, so honouring another value would mean interpreting w2c poses + # as c2w and silently generating the wrong trajectory. + camera_format = camera.get("format", SANA_WM_DEFAULT_CAMERA_FORMAT) + if camera_format != SANA_WM_DEFAULT_CAMERA_FORMAT: + raise ValueError( + f"Sana-WM only accepts camera format {SANA_WM_DEFAULT_CAMERA_FORMAT!r}; " + f"convert the poses before sending, got {camera_format!r}." + ) + coordinate_system = camera.get("coordinate_system", SANA_WM_DEFAULT_COORDINATE_SYSTEM) + if coordinate_system != SANA_WM_DEFAULT_COORDINATE_SYSTEM: + raise ValueError( + f"Sana-WM only accepts coordinate system {SANA_WM_DEFAULT_COORDINATE_SYSTEM!r}, " + f"got {coordinate_system!r}." + ) + canonical_camera = {"poses": poses} + + intrinsics = raw.get("intrinsics") + _validate_intrinsics(intrinsics) + translation_speed = raw.get("translation_speed") + rotation_speed_deg = raw.get("rotation_speed_deg") + + canonical = { + "num_frames": num_frames, + "height": height, + "width": width, + "camera": canonical_camera, + "action": action, + "intrinsics": intrinsics, + } + if translation_speed is not None: + canonical["translation_speed"] = _as_positive_float(translation_speed, name="translation_speed") + if rotation_speed_deg is not None: + canonical["rotation_speed_deg"] = _as_positive_float(rotation_speed_deg, name="rotation_speed_deg") + + mm_data["image"] = image + result["multi_modal_data"] = mm_data + additional[SANA_WM_CANONICAL_KEY] = canonical + result["additional_information"] = additional + result.pop(SANA_WM_CANONICAL_KEY, None) + return result diff --git a/vllm_omni/diffusion/models/sana_wm/sana_wm_transformer.py b/vllm_omni/diffusion/models/sana_wm/sana_wm_transformer.py new file mode 100644 index 000000000..07537ec09 --- /dev/null +++ b/vllm_omni/diffusion/models/sana_wm/sana_wm_transformer.py @@ -0,0 +1,2065 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SANA-WM Stage-1 transformer. + +Native vLLM-Omni port of the NVlabs SANA-WM DiT. Modules are built eagerly at +construction (under the loader's target-device context) and checkpoint tensors +are streamed in via ``load_weights``. The Bidirectional Gated DeltaNet +recurrence is implemented in pure PyTorch in this module. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable +from typing import Any, ClassVar + +import torch +import torch.nn.functional as F +from diffusers.models.embeddings import Timesteps +from torch import nn +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from vllm.distributed.parallel_state import get_tp_group +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm +from vllm_omni.diffusion.layers.norm import RMSNorm +from vllm_omni.diffusion.models.sana_wm.config import SanaWmConfig +from vllm_omni.diffusion.models.sana_wm.ucpe import ( + SanaWmCamGeometry, + cam_prep_func, + prepare_cam_geometry, +) + +SANA_WM_STAGE1_LATENT_CHANNELS = 128 +SANA_WM_STAGE1_PROMPT_CHANNELS = 2304 +SANA_WM_STAGE1_TIMESTEP_CHANNELS = 256 + + +def _shard_param_across_tp(param: torch.Tensor | None, dim: int = 0) -> None: + """Attach vLLM's TP shard loader to a plain (non-parallel-layer) parameter. + + vLLM parallel layers narrow full checkpoint tensors to the local shard via + their own ``weight_loader``. SANA-WM also has plain parameters that become + TP-local - the GDN vectors (``A_log``, ``dt_bias``), the depthwise temporal + convs and the q/k norms - so they need the same behavior at load time. + Attached only under TP > 1: at TP = 1 the checkpoint tensor already matches + the parameter and ``sharded_weight_loader`` would require an initialized TP + group that standalone/unit-test builds do not have. + """ + if param is None or get_tensor_model_parallel_world_size() <= 1: + return + param.weight_loader = sharded_weight_loader(dim) + + +def _is_sana_wm_transformer_block(name: str, module: Any) -> bool: + del module + parts = name.split(".") + return len(parts) == 2 and parts[0] == "blocks" and parts[1].isdigit() + + +def _to_3tuple(value: int | tuple[int, int] | tuple[int, int, int]) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if len(value) == 2: + return (1, int(value[0]), int(value[1])) + return (int(value[0]), int(value[1]), int(value[2])) + + +# --------------------------------------------------------------------------- +# Bidirectional Gated DeltaNet recurrence (pure PyTorch) +# +# Ported from the Apache-2.0 NVlabs/Sana reference. SANA-WM's bidirectional +# video-latent recurrence is a different contract than vLLM's autoregressive +# Qwen3-Next GDN cache path, so it stays model-local. +# --------------------------------------------------------------------------- + + +def _validate_gdn_inputs( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + spatial_tokens: int, +) -> tuple[int, int, int, int, int]: + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + raise ValueError("Sana-WM GDN query/key/value must be shaped [B, H, D, N].") + if query.shape != key.shape or query.shape != value.shape: + raise ValueError( + "Sana-WM GDN query/key/value shapes must match, got " + f"{tuple(query.shape)}, {tuple(key.shape)}, {tuple(value.shape)}." + ) + if beta.ndim != 4: + raise ValueError("Sana-WM GDN beta must be shaped [B, H, T, S].") + if decay.ndim != 3: + raise ValueError("Sana-WM GDN decay must be shaped [B, H, T].") + if spatial_tokens <= 0: + raise ValueError("Sana-WM GDN spatial_tokens must be positive.") + + batch_size, num_heads, head_dim, token_count = query.shape + if token_count % spatial_tokens != 0: + raise ValueError(f"Sana-WM GDN token count {token_count} is not divisible by spatial_tokens={spatial_tokens}.") + frames = token_count // spatial_tokens + if beta.shape != (batch_size, num_heads, frames, spatial_tokens): + raise ValueError( + "Sana-WM GDN beta shape mismatch: expected " + f"{(batch_size, num_heads, frames, spatial_tokens)}, got {tuple(beta.shape)}." + ) + if decay.shape != (batch_size, num_heads, frames): + raise ValueError( + f"Sana-WM GDN decay shape mismatch: expected {(batch_size, num_heads, frames)}, got {tuple(decay.shape)}." + ) + return batch_size, num_heads, head_dim, token_count, frames + + +def _delta_scan( + query_rot: torch.Tensor, + key_rot: torch.Tensor, + value: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + spatial_tokens: int, + query: torch.Tensor | None = None, + key: torch.Tensor | None = None, + skip_z: bool = False, + flip_output: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """One-directional gated delta-rule recurrence over frames. + + The KV stream consumes only the rotary-embedded Q/K; the Z (denominator) + stream consumes only the non-rotary Q/K. ``skip_z=True`` drops the Z stream + entirely and returns ``(numerator, None)`` — the numerator-only variant the + camera branch needs (NVlabs ``torch_recurrent_cam_single_path_delta_rule``), + so ``query``/``key`` may be omitted there. + + ``flip_output=True`` emits the frames in reverse of the order they were + computed, which is what a caller running on flipped inputs wants; it costs a + list reversal instead of a flip over the assembled result. + """ + if not skip_z and (query is None or key is None): + raise ValueError("Sana-WM delta scan needs non-rotary query/key unless skip_z is set.") + + batch_size, num_heads, head_dim, token_count = query_rot.shape + frames = beta.shape[2] + if token_count != frames * spatial_tokens: + raise ValueError(f"Sana-WM delta scan token_count={token_count} != frames*S={frames * spatial_tokens}.") + + def to_frames(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(batch_size, num_heads, head_dim, frames, spatial_tokens).permute(0, 1, 3, 2, 4) + + query_rot_f = to_frames(query_rot) + key_rot_f = to_frames(key_rot) + value_f = to_frames(value) + state_kv = torch.zeros(batch_size, num_heads, head_dim, head_dim, device=query_rot.device, dtype=query_rot.dtype) + numerators: list[torch.Tensor] = [] + + if skip_z: + query_f = key_f = None + state_z = None + denominators = None + else: + query_f = to_frames(query) + key_f = to_frames(key) + state_z = torch.zeros(batch_size, num_heads, head_dim, 1, device=query_rot.device, dtype=query_rot.dtype) + denominators = [] + + for frame_idx in range(frames): + query_rot_t = query_rot_f[:, :, frame_idx] + key_rot_t = key_rot_f[:, :, frame_idx] + value_t = value_f[:, :, frame_idx] + beta_t = beta[:, :, frame_idx].unsqueeze(2) + decay_t = decay[:, :, frame_idx].view(batch_size, num_heads, 1, 1) + + state_kv = state_kv * decay_t + value_pred = torch.matmul(state_kv, key_rot_t) + delta_value = (value_t - value_pred) * beta_t + state_kv = state_kv + torch.matmul(delta_value, key_rot_t.transpose(-1, -2)) + numerators.append(torch.matmul(state_kv, query_rot_t)) + + if skip_z: + continue + query_t = query_f[:, :, frame_idx] + key_t = key_f[:, :, frame_idx] + state_z = state_z * decay_t + z_pred = torch.matmul(state_z.transpose(-1, -2), key_t) + delta_z = (1.0 - z_pred) * beta_t + state_z = state_z + torch.matmul(key_t, delta_z.transpose(-1, -2)) + denominators.append(torch.matmul(state_z.transpose(-1, -2), query_t)) + + def restore(tensors: list[torch.Tensor], dim: int) -> torch.Tensor: + stacked = torch.stack(tensors[::-1] if flip_output else tensors, dim=2) + return stacked.permute(0, 1, 3, 2, 4).reshape(batch_size, num_heads, dim, token_count) + + return restore(numerators, head_dim), (None if skip_z else restore(denominators, 1)) + + +def _bidirectional_delta_scan( + query_rot: torch.Tensor, + key_rot: torch.Tensor, + value: torch.Tensor, + beta: torch.Tensor, + decay: torch.Tensor, + *, + spatial_tokens: int, + query: torch.Tensor | None = None, + key: torch.Tensor | None = None, + skip_z: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Forward + backward delta scan, summed. + + The backward direction shifts K/V/beta by one frame (zero pad) and the decay + by one frame (neutral 1.0 pad), matching the NVlabs ``flip_and_shift`` + convention. The backward scan also emits its frames in forward order, so the + two directions add directly instead of flipping the assembled result back. + """ + frames = beta.shape[2] + + def reverse(tensor: torch.Tensor, *, shift_value: float | None = None) -> torch.Tensor: + return _reverse_frames(tensor, frames=frames, spatial_tokens=spatial_tokens, shift_value=shift_value) + + num_fwd, den_fwd = _delta_scan( + query_rot, + key_rot, + value, + beta, + decay, + spatial_tokens=spatial_tokens, + query=query, + key=key, + skip_z=skip_z, + ) + + bwd_kwargs: dict[str, torch.Tensor] = {} + if not skip_z: + bwd_kwargs["query"] = reverse(query) + bwd_kwargs["key"] = reverse(key, shift_value=0.0) + + num_bwd, den_bwd = _delta_scan( + reverse(query_rot), + reverse(key_rot, shift_value=0.0), + reverse(value, shift_value=0.0), + _flip_and_shift(beta, dim=2, shift_value=0.0), + _flip_and_shift(decay, dim=2, shift_value=1.0), + spatial_tokens=spatial_tokens, + skip_z=skip_z, + flip_output=True, + **bwd_kwargs, + ) + + return num_fwd + num_bwd, (None if skip_z else den_fwd + den_bwd) + + +def _reverse_frames( + tensor: torch.Tensor, + *, + frames: int, + spatial_tokens: int, + shift_value: float | None = None, +) -> torch.Tensor: + """Reverse frame order on a ``[..., T*S]`` tensor, optionally shifting by one frame. + + The frame axis is the outer half of the flat token axis, so unfolding it in + place keeps the tensor's own layout: the flip writes a contiguous result and + the reshape back is a view. Routing this through a ``[..., T, D, S]`` + permutation instead would cost a second full materialization. + """ + lead = tensor.shape[:-1] + reversed_ = torch.flip(tensor.reshape(*lead, frames, spatial_tokens), dims=[-2]) + if shift_value is not None: + padding = torch.full((*lead, 1, spatial_tokens), shift_value, device=tensor.device, dtype=tensor.dtype) + reversed_ = torch.cat([padding, reversed_.narrow(-2, 0, frames - 1)], dim=-2) + return reversed_.reshape(*lead, frames * spatial_tokens) + + +def _flip_and_shift(tensor: torch.Tensor, *, dim: int, shift_value: float) -> torch.Tensor: + flipped = torch.flip(tensor, dims=[dim]) + shifted = flipped.narrow(dim, 0, tensor.shape[dim] - 1) + pad_shape = list(tensor.shape) + pad_shape[dim] = 1 + padding = torch.full(pad_shape, shift_value, device=tensor.device, dtype=tensor.dtype) + return torch.cat([padding, shifted], dim=dim) + + +def reference_bidirectional_gated_delta_net( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + beta: torch.Tensor, + decay: torch.Tensor, + spatial_tokens: int, + query_rot: torch.Tensor | None = None, + key_rot: torch.Tensor | None = None, + eps: float = 1e-8, +) -> torch.Tensor: + """Run the SANA-WM bidirectional gated delta recurrence in PyTorch. + + Inputs follow the layout used by the official Stage-1 operator: + ``query/key/value/query_rot/key_rot`` are ``[B, H, D, T*S]``, ``beta`` is + ``[B, H, T, S]``, and ``decay`` is ``[B, H, T]``. + """ + + _validate_gdn_inputs(query, key, value, beta, decay, spatial_tokens) + if query_rot is None: + query_rot = query + if key_rot is None: + key_rot = key + if query_rot.shape != query.shape or key_rot.shape != key.shape: + raise ValueError("Sana-WM GDN rotary query/key shapes must match query/key.") + + dtype_orig = query.dtype + numerator, denominator = _bidirectional_delta_scan( + query_rot.float(), + key_rot.float(), + value.float(), + beta.float(), + decay.float(), + spatial_tokens=spatial_tokens, + query=query.float(), + key=key.float(), + ) + return (numerator / (denominator + eps)).to(dtype_orig) + + +class TensorParallelRMSNorm(nn.Module): + """RMSNorm whose statistics span a tensor-parallel sharded last dimension. + + Models that normalize q/k across *all* heads (``qk_norm="rms_norm_across_heads"``) + cannot use a plain RMSNorm once the q/k projections are column-parallel: each + rank would then compute the RMS over its local head shard only. This + all-reduces the squared sum so the denominator matches the global width, + while the affine weight stays local. + + Model-local on purpose: LTX-2 and DreamZero carry their own near-identical + copies, and converging all three belongs in its own change rather than in + the SANA-WM port. + """ + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + elementwise_affine: bool = True, + tp_size: int = 1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.tp_size = max(int(tp_size), 1) + self.global_hidden_size = hidden_size * self.tp_size + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(hidden_size)) + else: + self.register_parameter("weight", None) + + def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + x_float = x.float() + return x_float, x_float.pow(2).sum(dim=-1, keepdim=True) + + def _scale(self, x_float: torch.Tensor, global_sum_sq: torch.Tensor, input_dtype: torch.dtype) -> torch.Tensor: + out = x_float * torch.rsqrt(global_sum_sq / self.global_hidden_size + self.eps) + if self.weight is not None: + out = out * self.weight.float() + return out.to(dtype=input_dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x_float, sum_sq = self._local_sum_sq(x) + if self.tp_size > 1: + sum_sq = tensor_model_parallel_all_reduce(sum_sq) + return self._scale(x_float, sum_sq, x.dtype) + + +def fused_qk_rms_norm( + norm_q: nn.Module, + norm_k: nn.Module, + q: torch.Tensor, + k: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply q/k :class:`TensorParallelRMSNorm` with a SINGLE fused all-reduce. + + Self-attention normalizes q and k every step, and under TP each norm issues + its own tiny all-reduce of the per-token sum-of-squares. These collectives + are latency-bound, so packing both sums into one tensor halves the count + (2 collectives -> 1) for free. + + Numerically identical to ``norm_q(q), norm_k(k)``: all-reduce is + elementwise, so packing along the last dim reduces each slice independently + with the same fp32 accumulation. Falls back to independent application when + either norm is not a :class:`TensorParallelRMSNorm` (e.g. ``nn.Identity`` + when qk_norm is off, or plain ``RMSNorm`` at TP=1). + """ + if not (isinstance(norm_q, TensorParallelRMSNorm) and isinstance(norm_k, TensorParallelRMSNorm)): + return norm_q(q), norm_k(k) + if norm_q.tp_size <= 1: + return norm_q(q), norm_k(k) + + q_float, q_sum_sq = norm_q._local_sum_sq(q) + k_float, k_sum_sq = norm_k._local_sum_sq(k) + packed = tensor_model_parallel_all_reduce(torch.cat([q_sum_sq, k_sum_sq], dim=-1)) + q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2] + return norm_q._scale(q_float, q_sum_sq, q.dtype), norm_k._scale(k_float, k_sum_sq, k.dtype) + + +def _make_sharded_qk_rms_norm(hidden_size: int, eps: float = 1e-6) -> nn.Module: + """Build the q/k norm, TP-aware. + + SANA-WM normalizes q/k across all heads, so once the projections are + column-parallel the RMS denominator has to span the TP ranks. + """ + tp_size = get_tensor_model_parallel_world_size() + if tp_size > 1: + return TensorParallelRMSNorm(hidden_size, eps=eps, tp_size=tp_size) + return RMSNorm(hidden_size, eps=eps) + + +class SanaWmTextProjection(nn.Module): + def __init__( + self, + prompt_channels: int, + hidden_size: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.fc1 = ColumnParallelLinear( + prompt_channels, + hidden_size, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fc1" if prefix else "fc1", + ) + self.fc2 = RowParallelLinear( + hidden_size, + hidden_size, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fc2" if prefix else "fc2", + ) + self.act = nn.SiLU() + + def forward(self, prompt_embeds: torch.Tensor) -> torch.Tensor: + return self.fc2(self.act(self.fc1(prompt_embeds))) + + +class SanaWmTextEmbedder(nn.Module): + def __init__( + self, + prompt_channels: int, + hidden_size: int, + max_length: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.y_embedding = nn.Parameter(torch.zeros(max_length, prompt_channels)) + self.y_proj = SanaWmTextProjection( + prompt_channels, + hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.y_proj" if prefix else "y_proj", + ) + + def forward( + self, + prompt_embeds: torch.Tensor | None, + *, + batch_size: int, + dtype: torch.dtype, + ) -> torch.Tensor: + if prompt_embeds is None: + prompt_embeds = self.y_embedding.unsqueeze(0).expand(batch_size, -1, -1) + return self.y_proj(prompt_embeds.to(dtype=dtype)) + + +class SanaWmTimestepEmbedder(nn.Module): + def __init__( + self, + in_features: int, + hidden_size: int, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + linear_1 = ColumnParallelLinear( + in_features, + hidden_size, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.mlp.0" if prefix else "mlp.0", + ) + linear_2 = RowParallelLinear( + hidden_size, + hidden_size, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.mlp.2" if prefix else "mlp.2", + ) + self.act = nn.SiLU() + self.mlp = nn.ModuleList([linear_1, self.act, linear_2]) + # diffusers' Timesteps reproduces this model's sinusoid exactly: + # downscale_freq_shift=0 makes the exponent denominator half_dim, and + # flip_sin_to_cos=True emits [cos, sin] in that order. The MLP stays + # model-local because its checkpoint keys are mlp.0/mlp.2 and it is + # tensor-parallel. + self.timesteps_proj = Timesteps(num_channels=in_features, flip_sin_to_cos=True, downscale_freq_shift=0) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + emb = self.timesteps_proj(timestep.flatten()) + hidden_states = self.mlp[0](emb.to(dtype=self.mlp[0].weight.dtype)) + hidden_states = self.act(hidden_states) + return self.mlp[2](hidden_states) + + +class SanaWmPatchEmbedMS3D(nn.Module): + """Official-style 3D patch embedder used by SANA-WM Stage-1.""" + + def __init__( + self, + patch_size: tuple[int, int, int], + in_channels: int, + hidden_size: int, + *, + kernel_size: tuple[int, int, int] | None = None, + bias: bool = True, + ) -> None: + super().__init__() + self.patch_size = _to_3tuple(patch_size) + self.kernel_size = _to_3tuple(kernel_size or patch_size) + self.proj = nn.Conv3d( + in_channels, + hidden_size, + kernel_size=self.kernel_size, + stride=self.patch_size, + bias=bias, + ) + self.norm = nn.Identity() + + def project_with_shape(self, latents: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int, int]]: + hidden_states = self.norm(self.proj(latents)) + _, _, frames, height, width = hidden_states.shape + return hidden_states.flatten(2).transpose(1, 2), (frames, height, width) + + def forward(self, latents: torch.Tensor) -> torch.Tensor: + return self.project_with_shape(latents)[0] + + +class SanaWmWanRotaryPosEmbed(nn.Module): + """Wan-style 3D RoPE table used by the official SANA-WM GDN blocks.""" + + def __init__( + self, + attention_head_dim: int, + *, + max_seq_len: int = 1024, + theta: float = 10000.0, + ) -> None: + super().__init__() + self.theta = theta + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + self._split_sizes = ( + t_dim // 2, + h_dim // 2, + w_dim // 2, + ) + self.freqs = self._build_freqs(max_seq_len) + + def _build_1d_freq(self, dim: int, positions: torch.Tensor) -> torch.Tensor: + if dim <= 0: + return torch.empty(positions.shape[0], 0, dtype=torch.complex128) + freqs = 1.0 / ( + self.theta ** (torch.arange(0, dim, 2, dtype=torch.float64, device=positions.device)[: dim // 2] / dim) + ) + phase = torch.outer(positions.to(torch.float64), freqs) + return torch.polar(torch.ones_like(phase), phase) + + def _build_freqs(self, max_seq_len: int) -> torch.Tensor: + positions = torch.arange(max_seq_len, dtype=torch.float64) + dims = ( + self._split_sizes[0] * 2, + self._split_sizes[1] * 2, + self._split_sizes[2] * 2, + ) + return torch.cat([self._build_1d_freq(dim, positions) for dim in dims], dim=1) + + def forward(self, spatial_shape: tuple[int, int, int], device: torch.device) -> torch.Tensor: + frames, height, width = spatial_shape + if max(spatial_shape) > self.freqs.shape[0]: + self.freqs = self._build_freqs(max(spatial_shape)).to(self.freqs.device) + freqs = self.freqs.to(device=device) + freqs_f, freqs_h, freqs_w = freqs.split(self._split_sizes, dim=1) + f_dim, h_dim, w_dim = self._split_sizes + parts = [ + freqs_f[:frames].view(frames, 1, 1, f_dim).expand(frames, height, width, f_dim), + freqs_h[:height].view(1, height, 1, h_dim).expand(frames, height, width, h_dim), + freqs_w[:width].view(1, 1, width, w_dim).expand(frames, height, width, w_dim), + ] + return torch.cat(parts, dim=-1).reshape(1, 1, frames * height * width, -1) + + +class SanaWmSelfAttention(nn.Module): + def __init__( + self, + config: SanaWmConfig, + *, + use_gdn: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + hidden_size = config.hidden_size + self.total_num_heads = max(hidden_size // max(config.linear_head_dim, 1), 1) + self.head_dim = hidden_size // self.total_num_heads + self.num_heads = self.total_num_heads + self.num_kv_heads = self.total_num_heads + self.eps = 1e-8 + self.use_gdn = use_gdn and "GDN" in config.attn_type + self.patch_size = _to_3tuple(config.patch_size) + # Total camera branch width from the checkpoint. TP ColumnParallel + # layers expose a per-rank local slice at runtime; keep both sizes so + # layer construction uses the global contract while local conv/norm + # tensors match q/k/v outputs. + self.total_cam_dim = hidden_size // max(config.cam_attn_compress, 1) + self.cam_dim = self.total_cam_dim + self.qkv = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + bias=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv" if prefix else "qkv", + ) + self.num_heads = self.qkv.num_heads + self.num_kv_heads = self.qkv.num_kv_heads + self.proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.proj" if prefix else "proj", + ) + self.beta_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.beta_proj" if prefix else "beta_proj", + ) + self.gate_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_proj" if prefix else "gate_proj", + ) + self.output_gate = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.output_gate" if prefix else "output_gate", + ) + self.q_proj_cam = ColumnParallelLinear( + hidden_size, + self.total_cam_dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj_cam" if prefix else "q_proj_cam", + ) + self.k_proj_cam = ColumnParallelLinear( + hidden_size, + self.total_cam_dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.k_proj_cam" if prefix else "k_proj_cam", + ) + self.v_proj_cam = ColumnParallelLinear( + hidden_size, + self.total_cam_dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.v_proj_cam" if prefix else "v_proj_cam", + ) + self.out_proj_cam = RowParallelLinear( + self.total_cam_dim, + hidden_size, + bias=True, + input_is_parallel=True, + # The camera contribution is summed into the TP-local GDN stream, so + # the cross-rank sum and the slice down to the local width collapse + # into a single reduce-scatter (see _reduce_scatter_cam_contrib). + # RowParallelLinear must therefore not all-reduce, and its bias is + # applied after the scatter rather than inside the matmul. + reduce_results=False, + skip_bias_add=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_proj_cam" if prefix else "out_proj_cam", + ) + local_inner_dim = self.num_heads * self.head_dim + norm_cls = _make_sharded_qk_rms_norm if config.qk_norm else (lambda *_args, **_kwargs: nn.Identity()) + self.q_norm = norm_cls(local_inner_dim) + self.k_norm = norm_cls(local_inner_dim) + + # These names mirror the official GDN checkpoint. + self.A_log = nn.Parameter(torch.zeros(self.num_heads)) + self.dt_bias = nn.Parameter(torch.zeros(self.num_heads)) + self.register_buffer("recall_gate", torch.zeros(1)) + self.conv_k = ( + nn.Conv1d( + local_inner_dim, + local_inner_dim, + kernel_size=config.conv_kernel_size, + groups=local_inner_dim, + bias=False, + ) + if self.use_gdn and config.conv_kernel_size > 0 + else None + ) + self.conv_q = None + self.conv_v = None + + self.cam_dim = int(getattr(self.q_proj_cam, "output_size_per_partition", self.total_cam_dim)) + + cam_compress = max(config.cam_attn_compress, 1) + self.cam_heads = max(self.num_heads // cam_compress, 1) + if self.cam_dim % self.cam_heads != 0: + raise ValueError(f"Sana-WM local cam_dim={self.cam_dim} must be divisible by cam_heads={self.cam_heads}.") + self.cam_head_dim = self.cam_dim // self.cam_heads + # Under TP, cam_head_dim must equal main head_dim so the WAN RoPE + # (built for main head_dim) can be correctly sliced for the cam branch. + # A mismatch indicates the cam ColumnParallelLinear layers were not + # partitioned consistently with the main QKVParallelLinear. + if self.cam_head_dim != self.head_dim: + raise ValueError( + f"Sana-WM cam_head_dim={self.cam_head_dim} != head_dim={self.head_dim}. " + f"Under TP the cam layers must be partitioned by the same TP degree as the " + f"main QKV (cam_dim={self.cam_dim}, cam_heads={self.cam_heads}, " + f"total_cam_dim={self.total_cam_dim})." + ) + self.q_norm_cam = norm_cls(self.cam_dim) + self.k_norm_cam = norm_cls(self.cam_dim) + self.conv_k_cam = ( + nn.Conv1d( + self.cam_dim, + self.cam_dim, + kernel_size=config.conv_kernel_size, + groups=self.cam_dim, + bias=False, + ) + if self.use_gdn and config.conv_kernel_size > 0 + else None + ) + self.conv_q_cam = None + self.conv_v_cam = None + # Softmax hybrid blocks (every ``softmax_every_n``-th) and the camera + # branch both run plain non-causal self-attention, so both go through + # the shared Attention for platform backend selection and the fp32 + # fallback. The camera branch gets its own module because + # ``cam_attn_compress`` may give it fewer heads. + self.softmax_attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + num_kv_heads=self.num_kv_heads, + softmax_scale=self.head_dim**-0.5, + causal=False, + role="self", + qkv_layout="BSND", + skip_sequence_parallel=True, + prefix=prefix, + ) + self.softmax_attn_cam = Attention( + num_heads=self.cam_heads, + head_size=self.cam_head_dim, + num_kv_heads=self.cam_heads, + softmax_scale=self.cam_head_dim**-0.5, + causal=False, + role="self", + qkv_layout="BSND", + skip_sequence_parallel=True, + prefix=prefix, + ) + self._init_short_convs() + self._mark_tp_sharded_params() + + def _mark_tp_sharded_params(self) -> None: + """Declare which plain parameters are TP-local (see ``_shard_param_across_tp``).""" + # Per-head GDN vectors: one entry per local head. + _shard_param_across_tp(self.A_log) + _shard_param_across_tp(self.dt_bias) + # Depthwise temporal convs and q/k norms: one entry per local channel. + for module in (self.conv_q, self.conv_k, self.conv_v, self.conv_q_cam, self.conv_k_cam, self.conv_v_cam): + if module is not None: + _shard_param_across_tp(module.weight) + for norm in (self.q_norm, self.k_norm, self.q_norm_cam, self.k_norm_cam): + _shard_param_across_tp(getattr(norm, "weight", None)) + + @staticmethod + def _init_short_conv(conv: nn.Conv1d | None) -> None: + if conv is None: + return + with torch.no_grad(): + conv.weight.zero_() + conv.weight[:, 0, -1] = 1.0 + + def _init_short_convs(self) -> None: + for conv in (self.conv_q, self.conv_k, self.conv_v, self.conv_q_cam, self.conv_k_cam, self.conv_v_cam): + self._init_short_conv(conv) + + @staticmethod + def _apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + rotated = torch.view_as_complex(hidden_states.permute(0, 1, 3, 2).to(torch.float64).unflatten(3, (-1, 2))) + output = torch.view_as_real(rotated * freqs).flatten(3, 4).permute(0, 1, 3, 2) + return output.type_as(hidden_states) + + @staticmethod + def _reshape_to_temporal( + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + ) -> tuple[torch.Tensor, int, int, int]: + batch_size, token_count, hidden_size = hidden_states.shape + frames, height, width = spatial_shape + spatial_tokens = height * width + if token_count != frames * spatial_tokens: + raise ValueError(f"Sana-WM temporal conv expects N=T*H*W, got N={token_count}, THW={spatial_shape}.") + hidden_states = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size) + hidden_states = hidden_states.permute(0, 2, 1, 3).reshape(batch_size * spatial_tokens, frames, hidden_size) + return hidden_states, batch_size, spatial_tokens, frames + + @staticmethod + def _reshape_from_temporal( + hidden_states: torch.Tensor, + batch_size: int, + spatial_tokens: int, + frames: int, + ) -> torch.Tensor: + hidden_size = hidden_states.shape[-1] + return ( + hidden_states.reshape(batch_size, spatial_tokens, frames, hidden_size) + .permute(0, 2, 1, 3) + .reshape(batch_size, frames * spatial_tokens, hidden_size) + ) + + @staticmethod + def _causal_conv_1d(hidden_states: torch.Tensor, conv: nn.Conv1d) -> torch.Tensor: + dtype = hidden_states.dtype + conv_input = hidden_states.transpose(1, 2).to(conv.weight.dtype) + conv_input = F.pad(conv_input, (conv.kernel_size[0] - 1, 0)) + output = conv(conv_input).transpose(1, 2) + return output.to(dtype) + + def _bidirectional_temporal_short_conv( + self, + hidden_states: torch.Tensor, + conv: nn.Conv1d, + spatial_shape: tuple[int, int, int], + ) -> torch.Tensor: + hidden_states, batch_size, spatial_tokens, frames = self._reshape_to_temporal(hidden_states, spatial_shape) + forward_states = self._causal_conv_1d(hidden_states, conv) + backward_states = self._causal_conv_1d(hidden_states.flip(1), conv).flip(1) + center_weight = conv.weight[:, 0, -1].to(hidden_states.dtype) + center_states = hidden_states * center_weight.view(1, 1, -1) + output = forward_states + backward_states - center_states + return self._reshape_from_temporal(output, batch_size, spatial_tokens, frames) + + def _compute_frame_gates( + self, + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, token_count, hidden_size = hidden_states.shape + frames, height, width = spatial_shape + spatial_tokens = height * width + if token_count != frames * spatial_tokens: + raise ValueError(f"Sana-WM GDN token layout mismatch: N={token_count}, expected {frames * spatial_tokens}.") + beta = torch.sigmoid(self.beta_proj(hidden_states)) + beta = beta.reshape(batch_size, frames, spatial_tokens, self.num_heads).permute(0, 3, 1, 2) + frame_states = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size).mean(dim=2) + gate = self.gate_proj(frame_states).float() + decay = torch.exp( + -self.A_log.float().exp().view(1, 1, -1) * F.softplus(gate + self.dt_bias.float().view(1, 1, -1)) + ) + return beta, decay.transpose(1, 2) + + def _forward_gdn_raw( + self, + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + rotary_emb: torch.Tensor | None, + *, + precomputed_gates: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Run main-branch bidirectional GDN and return the (B, N, q_size) + raw output BEFORE ``output_gate`` and ``proj`` are applied. + + Returns ``(raw_output, (beta, decay))`` so the caller can share the + beta/decay gates with the camera branch (matches NVlabs' shared + ``precomputed_gates`` plumbing). + """ + batch_size, token_count, _hidden_size = hidden_states.shape + frames, height, width = spatial_shape + spatial_tokens = height * width + if token_count != frames * spatial_tokens: + raise ValueError(f"Sana-WM GDN expects N=T*H*W, got N={token_count}, THW={spatial_shape}.") + + qkv = self.qkv(hidden_states) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + if self.conv_k is not None: + key = self._bidirectional_temporal_short_conv(key, self.conv_k, spatial_shape) + if precomputed_gates is None: + beta, decay = self._compute_frame_gates(hidden_states, spatial_shape) + else: + beta, decay = precomputed_gates + + query, key = fused_qk_rms_norm(self.q_norm, self.k_norm, query, key) + query = query.reshape(batch_size, token_count, self.num_heads, self.head_dim) + key = key.reshape(batch_size, token_count, self.num_heads, self.head_dim) + value = value.reshape(batch_size, token_count, self.num_heads, self.head_dim) + + query = F.relu(query).permute(0, 2, 3, 1) + key = F.relu(key).permute(0, 2, 3, 1) + value = value.permute(0, 2, 3, 1) + key = key * ((self.head_dim**-0.5) * (spatial_tokens**-0.5)) + if rotary_emb is not None: + query_rot = self._apply_rotary_emb(query, rotary_emb) + key_rot = self._apply_rotary_emb(key, rotary_emb) + else: + query_rot = query + key_rot = key + + output = reference_bidirectional_gated_delta_net( + query, + key, + value, + beta=beta, + decay=decay, + spatial_tokens=spatial_tokens, + query_rot=query_rot, + key_rot=key_rot, + eps=self.eps, + ) + output = output.permute(0, 3, 1, 2).reshape(batch_size, token_count, q_size) + return output, (beta, decay) + + def _apply_output_gate_and_proj( + self, + combined: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Apply SiLU output gate (driven by hidden_states) + final projection. + + The gate reads ``hidden_states``, not ``combined``, and NVlabs applies + it once to the main+camera sum rather than per branch. + """ + gate = F.silu(self.output_gate(hidden_states).float()) + gated = combined * gate + return self.proj(gated.to(self.proj.weight.dtype)) + + def _reduce_scatter_cam_contrib(self, contrib: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: + """Reduce the partial camera projection straight down to the local shard. + + ``out_proj_cam`` is row-parallel over the camera features, so each rank + produces a partial sum spanning the full hidden width. The GDN stream it + is added to is TP-local (``output_gate`` and ``proj`` both expect the + local width), so an all-reduce would hand every rank the full result + only for all but one slice of it to be discarded. Reduce-scatter does + both steps in one collective and moves 1/tp of the bytes. + + ``skip_bias_add`` keeps the bias out of the matmul so it can be applied + to the scattered shard; that slice is the only remaining use of the TP + rank here, and it is a 1-D view rather than a copy of the hidden tensor. + + ``tensor_model_parallel_reduce_scatter`` is not usable here: it calls + ``get_tp_group().reduce_scatter()``, and vLLM-Omni's diffusion + ``GroupCoordinator`` deliberately does not implement that method (see + the note in ``diffusion/distributed/parallel_state.py``). Drive + ``torch.distributed`` on the group's process group instead, and fall + back to all-reduce plus a slice if the group is not exposed. + """ + bias = self.out_proj_cam.bias + tp_size = get_tensor_model_parallel_world_size() + if tp_size == 1: + return contrib if bias is None else contrib + bias + + local_width = reference.shape[-1] + if local_width <= 0 or contrib.shape[-1] != local_width * tp_size: + raise ValueError( + "Sana-WM TP camera contribution width mismatch: " + f"contrib={contrib.shape[-1]} expected={local_width * tp_size} (local={local_width}, tp={tp_size})." + ) + + device_group = getattr(get_tp_group(), "device_group", None) + if device_group is None: + out = tensor_model_parallel_all_reduce(contrib.contiguous()) + start = get_tensor_model_parallel_rank() * local_width + out = out[..., start : start + local_width].contiguous() + else: + # reduce_scatter_tensor splits along dim 0, so move the rank axis to + # the front first. Rank r then receives the summed columns + # [r * local_width, (r + 1) * local_width), matching what the + # all-reduce-and-slice fallback produces. + batch, tokens, _ = contrib.shape + source = ( + contrib.reshape(batch, tokens, tp_size, local_width) + .permute(2, 0, 1, 3) + .contiguous() + .reshape(tp_size * batch, tokens, local_width) + ) + out = torch.empty((batch, tokens, local_width), device=contrib.device, dtype=contrib.dtype) + torch.distributed.reduce_scatter_tensor(out, source, group=device_group) + + if bias is not None: + start = get_tensor_model_parallel_rank() * local_width + out = out + bias[start : start + local_width] + return out + + def _forward_gdn( + self, + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + rotary_emb: torch.Tensor | None, + ) -> torch.Tensor: + """Main-branch only GDN forward with output_gate + proj applied.""" + raw, _ = self._forward_gdn_raw(hidden_states, spatial_shape, rotary_emb) + return self._apply_output_gate_and_proj(raw, hidden_states) + + @staticmethod + def _downscale_to_reference_rms( + ref: torch.Tensor, + transformed: torch.Tensor, + eps: float = 1e-6, + ) -> torch.Tensor: + """Downscale ``transformed`` so its per-token channel RMS does not + exceed ``ref``'s. Mirrors NVlabs ``_downscale_to_reference_rms``. + + Inputs are ``(B, H, N, D)`` (channel dim last). The RMS-scale is + clamped at 1.0 — the function only ever shrinks, never amplifies. + """ + ref_rms = ref.square().mean(dim=-1, keepdim=True).add(eps).sqrt() + tr_rms = transformed.square().mean(dim=-1, keepdim=True).add(eps).sqrt() + scale = (ref_rms / tr_rms.clamp_min(eps)).clamp(max=1.0) + return transformed * scale + + @staticmethod + def _ucpe_rotary_freqs(rotary_emb: torch.Tensor | None) -> torch.Tensor | None: + # SanaWmRope emits (1, 1, N, D//2) complex; the cam path wants (N, D//2). + if rotary_emb is None: + return None + return rotary_emb.squeeze(0).squeeze(0) + + def _forward_softmax_raw( + self, + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + rotary_emb: torch.Tensor | None, + ) -> torch.Tensor: + """Softmax self-attention raw output before output_gate/proj. + + Mirrors NVlabs ``_forward_softmax_attn(..., apply_output_gate=False)`` + for every-N-th hybrid blocks. GDN-specific gates are intentionally not + used here. + """ + batch_size, token_count, hidden_size = hidden_states.shape + frames, height, width = spatial_shape + if token_count != frames * height * width: + raise ValueError(f"Sana-WM softmax attention expects N=T*H*W, got N={token_count}, THW={spatial_shape}.") + + qkv = self.qkv(hidden_states) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) + + query, key = fused_qk_rms_norm(self.q_norm, self.k_norm, query, key) + query = query.reshape(batch_size, token_count, self.num_heads, self.head_dim) + key = key.reshape(batch_size, token_count, self.num_kv_heads, self.head_dim) + value = value.reshape(batch_size, token_count, self.num_kv_heads, self.head_dim) + + if rotary_emb is not None: + query = self._apply_rotary_emb(query.permute(0, 2, 3, 1), rotary_emb).permute(0, 3, 1, 2) + key = self._apply_rotary_emb(key.permute(0, 2, 3, 1), rotary_emb).permute(0, 3, 1, 2) + + # NVlabs runs the attention in bf16 even under ``fp32_attention=True``, + # since FlashAttention only supports bf16/fp16 and fp32 falls back to + # the math backend. Keep Q/K/V at the module dtype so an upstream step + # that promoted to fp32 does not silently change the kernel. See + # ``sana_gdn_camctrl_blocks.py::_forward_softmax_attn_sdpa``. + dtype_orig = hidden_states.dtype + attn = self.softmax_attn(query.to(dtype_orig), key.to(dtype_orig), value.to(dtype_orig)) + return attn.reshape(batch_size, token_count, q_size) + + def _forward_softmax_cam_branch( + self, + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + cam_geometry: SanaWmCamGeometry, + rotary_emb: torch.Tensor | None, + ) -> torch.Tensor: + """UCPE camera branch for softmax hybrid blocks. + + Matches NVlabs ``_forward_cam_branch_softmax``: camera Q/K/V are + projected from the same hidden stream, transformed by UCPE, run through + SDPA, inverse-transformed, then returned as raw ``cam_dim`` features. + """ + batch_size, token_count, _ = hidden_states.shape + frames, height, width = spatial_shape + if token_count != frames * height * width: + raise ValueError(f"Sana-WM softmax cam branch expects N=T*H*W, got N={token_count}, THW={spatial_shape}.") + + q_cam = self.q_proj_cam(hidden_states) + k_cam = self.k_proj_cam(hidden_states) + v_cam = self.v_proj_cam(hidden_states) + + if self.conv_q_cam is not None: + q_cam = self._bidirectional_temporal_short_conv(q_cam, self.conv_q_cam, spatial_shape) + if self.conv_k_cam is not None: + k_cam = self._bidirectional_temporal_short_conv(k_cam, self.conv_k_cam, spatial_shape) + if self.conv_v_cam is not None: + v_cam = self._bidirectional_temporal_short_conv(v_cam, self.conv_v_cam, spatial_shape) + + q_cam, k_cam = fused_qk_rms_norm(self.q_norm_cam, self.k_norm_cam, q_cam, k_cam) + q_cam = q_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim) + k_cam = k_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim) + v_cam = v_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim) + + q_cam = q_cam.transpose(1, 2).contiguous() # (B, H_cam, N, D) + k_cam = k_cam.transpose(1, 2).contiguous() + v_cam = v_cam.transpose(1, 2).contiguous() + + q_cam_trans = cam_geometry.apply_q(q_cam) + kv_cam_trans = cam_geometry.apply_kv(torch.cat([k_cam, v_cam], dim=1)) + k_cam_trans, v_cam_trans = torch.chunk(kv_cam_trans, chunks=2, dim=1) + + q_cam_trans = self._downscale_to_reference_rms(q_cam, q_cam_trans) + k_cam_trans = self._downscale_to_reference_rms(k_cam, k_cam_trans) + v_cam_trans = self._downscale_to_reference_rms(v_cam, v_cam_trans) + + # The shared Attention takes BSND and selects a backend that supports + # ``cam_head_dim`` directly, so odd head dims need no padding up to an + # SDPA-friendly size. Q/K/V stay at the module dtype for the same + # reason as the main softmax path above. + dtype_orig = hidden_states.dtype + out = self.softmax_attn_cam( + q_cam_trans.transpose(1, 2).to(dtype_orig), + k_cam_trans.transpose(1, 2).to(dtype_orig), + v_cam_trans.transpose(1, 2).to(dtype_orig), + ) + out = cam_geometry.apply_output(out.transpose(1, 2)) + return out.transpose(1, 2).reshape(batch_size, token_count, self.cam_dim) + + def _forward_cam_branch( + self, + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + cam_geometry: SanaWmCamGeometry, + rotary_emb: torch.Tensor | None, + *, + precomputed_gates: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """UCPE camera branch — matches the NVlabs production + ``BidirectionalGDNUCPESinglePathLiteLABothTriton`` variant declared + in the SANA-WM 1600M release config. + + Two things differ from the main GDN branch and are easy to get wrong: + Q/K/V all come from ``hidden_states`` (UCPE is not cross-attention — + camera information enters through the per-pixel projection matrices), + and the recurrence is single-path, i.e. numerator only, with no Z + denominator and no final divide. + + Returns ``(B, N, cam_dim)`` raw — caller applies ``out_proj_cam`` and + the shared ``output_gate`` + ``proj``. + """ + batch_size, token_count, _ = hidden_states.shape + frames, height, width = spatial_shape + spatial_tokens = height * width + if token_count != frames * spatial_tokens: + raise ValueError(f"Sana-WM cam branch expects N=T*H*W, got N={token_count}, THW={spatial_shape}.") + + q_cam = self.q_proj_cam(hidden_states) + k_cam = self.k_proj_cam(hidden_states) + v_cam = self.v_proj_cam(hidden_states) + + if self.conv_k_cam is not None: + k_cam = self._bidirectional_temporal_short_conv(k_cam, self.conv_k_cam, spatial_shape) + if self.conv_q_cam is not None: + q_cam = self._bidirectional_temporal_short_conv(q_cam, self.conv_q_cam, spatial_shape) + if self.conv_v_cam is not None: + v_cam = self._bidirectional_temporal_short_conv(v_cam, self.conv_v_cam, spatial_shape) + + # Normalise through the model's own norm modules: under TP these are + # TensorParallelRMSNorm, which all-reduces the squared sum so the + # denominator spans the global channel width. Doing the RMS inside + # cam_prep_func would normalise over this rank's head shard only. q and + # k share a shape here, so the two reductions fuse into one. + q_cam, k_cam = fused_qk_rms_norm(self.q_norm_cam, self.k_norm_cam, q_cam, k_cam) + q_normed = q_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim) + k_normed = k_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim) + + # cam_prep_func returns NVlabs' ``(B, H, D, N)`` scan layout directly, + # so nothing needs permuting before the recurrence below. + v_raw = v_cam.reshape(batch_size, token_count, self.cam_heads, self.cam_head_dim).contiguous() + k_scale = (self.cam_head_dim**-0.5) * (spatial_tokens**-0.5) + q_rot_bhdn, k_rot_bhdn, v_bhdn, inflation_sq = cam_prep_func( + q_normed.contiguous(), + k_normed.contiguous(), + v_raw, + proj_q=cam_geometry.proj_q, + proj_kv=cam_geometry.proj_kv, + rope_cos=cam_geometry.rope_cos, + rope_sin=cam_geometry.rope_sin, + k_scale=k_scale, + ) + + # Discount beta by the raw K inflation instead of shrinking Q/K/V first: + # the released NVlabs ``BothTriton`` cam path skips the Python + # ``PostUCPERenorm`` step and feeds the scan the raw UCPE-transformed + # tensors. + beta, decay = precomputed_gates + inflation_per_token = inflation_sq.reshape(batch_size, self.cam_heads, frames, spatial_tokens) + frame_inflation_sq = inflation_per_token.mean(dim=-1) # (B, H_cam, T) + # β/decay come from the main branch, so the cam branch can only reuse + # them head-for-head. cam_attn_compress > 1 would give the cam branch + # fewer heads and needs a grouping rule this port does not implement; + # __init__ already rejects it via the cam_head_dim check. + if beta.shape[1] != self.cam_heads: + raise ValueError( + "Sana-WM cam branch requires cam_heads == main heads (cam_attn_compress=1)," + f" got cam_heads={self.cam_heads} main_heads={beta.shape[1]}." + ) + beta_cam = beta / frame_inflation_sq.unsqueeze(-1).clamp_min(1.0) + decay_cam = decay + + cam_out_bhdn, _ = _bidirectional_delta_scan( + q_rot_bhdn.float(), + k_rot_bhdn.float(), + v_bhdn.float(), + beta_cam.float(), + decay_cam.float(), + spatial_tokens=spatial_tokens, + skip_z=True, + ) + cam_out_bhdn = cam_out_bhdn.to(q_rot_bhdn.dtype) + + # ``apply_output`` expects (B, H, N, D); the scan emits (B, H, D, N). + cam_out_bhnd = cam_out_bhdn.transpose(-1, -2).contiguous() + cam_out_bhnd = cam_geometry.apply_output(cam_out_bhnd) + cam_out_bhdn = cam_out_bhnd.transpose(-1, -2).contiguous() + + # (B, H_cam, D_cam, N) → (B, N, cam_dim) + return cam_out_bhdn.permute(0, 3, 1, 2).reshape(batch_size, token_count, self.cam_dim) + + # The GDN/UCPE branch runs a sequential recurrence whose trip count is the + # latent frame count, i.e. data-dependent at trace time. Run the whole + # attention eagerly (graph break) rather than letting torch.compile unroll + # it per shape; the rest of the block (MLP / norms / cross-attention) still + # compiles. + @torch.compiler.disable + def forward( + self, + hidden_states: torch.Tensor, + spatial_shape: tuple[int, int, int], + rotary_emb: torch.Tensor | None = None, + cam_geometry: SanaWmCamGeometry | None = None, + ) -> torch.Tensor: + """Forward with dual-branch GDN+UCPE when ``cam_geometry`` is + provided, otherwise main-branch-only GDN or softmax fallback. + + ``cam_geometry`` holds the per-pixel 4x4 transforms built once per + transformer forward from the raw ``(B, T_latent, 20)`` raymap. The cam + branch consumes ``hidden_states`` as its sole Q/K/V source and applies + those transforms to the Q/K/V channels via UCPE. + """ + if self.use_gdn: + if cam_geometry is None: + # No camera conditioning on this request — main branch only. + return self._forward_gdn(hidden_states, spatial_shape, rotary_emb) + # β/decay are computed here rather than inside each branch: NVlabs + # drives both recurrences from one set of gates, and the camera + # branch then discounts beta by its own K inflation. + beta, decay = self._compute_frame_gates(hidden_states, spatial_shape) + main_raw, _ = self._forward_gdn_raw( + hidden_states, + spatial_shape, + rotary_emb, + precomputed_gates=(beta, decay), + ) + cam_raw = self._forward_cam_branch( + hidden_states, + spatial_shape, + cam_geometry, + rotary_emb, + precomputed_gates=(beta, decay), + ) + cam_contrib = self.out_proj_cam(cam_raw) + cam_contrib = self._reduce_scatter_cam_contrib(cam_contrib, main_raw) + combined = main_raw + cam_contrib.to(main_raw.dtype) + attn_out = self._apply_output_gate_and_proj(combined, hidden_states) + return attn_out + + main_raw = self._forward_softmax_raw(hidden_states, spatial_shape, rotary_emb) + if cam_geometry is not None: + cam_raw = self._forward_softmax_cam_branch( + hidden_states, + spatial_shape, + cam_geometry, + rotary_emb, + ) + cam_contrib = self.out_proj_cam(cam_raw) + cam_contrib = self._reduce_scatter_cam_contrib(cam_contrib, main_raw) + main_raw = main_raw + cam_contrib.to(main_raw.dtype) + return self._apply_output_gate_and_proj(main_raw, hidden_states) + + +class SanaWmCrossAttention(nn.Module): + def __init__( + self, + config: SanaWmConfig, + *, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + hidden_size = config.hidden_size + self.total_num_heads = max(hidden_size // max(config.linear_head_dim, 1), 1) + self.head_dim = hidden_size // self.total_num_heads + self.num_heads = self.total_num_heads + inner = self.total_num_heads * self.head_dim + self.q_linear = ColumnParallelLinear( + hidden_size, + inner, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_linear" if prefix else "q_linear", + ) + # The checkpoint stores K and V fused as one (2 * inner, hidden) + # tensor. A plain ColumnParallelLinear would shard that by the + # concatenated output dim, handing all of K to one rank and all of V + # to another; MergedColumnParallelLinear knows the block structure + # and shards K and V by heads independently, so `chunk(2)` below + # yields this rank's K and V directly - no gather, no manual slice. + self.kv_linear = MergedColumnParallelLinear( + hidden_size, + [inner, inner], + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_linear" if prefix else "kv_linear", + ) + self.num_heads = self.q_linear.output_size_per_partition // self.head_dim + self.proj = RowParallelLinear( + inner, + hidden_size, + bias=True, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.proj" if prefix else "proj", + ) + local_inner = self.num_heads * self.head_dim + norm_cls = _make_sharded_qk_rms_norm if config.cross_norm else (lambda *_args, **_kwargs: nn.Identity()) + self.q_norm = norm_cls(local_inner) + self.k_norm = norm_cls(local_inner) + # q is column-parallel and k is sliced to the local width before k_norm, + # so both norms hold one weight entry per local channel. + _shard_param_across_tp(getattr(self.q_norm, "weight", None)) + _shard_param_across_tp(getattr(self.k_norm, "weight", None)) + self.softmax_attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + num_kv_heads=self.num_heads, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + role="cross", + qkv_layout="BSND", + skip_sequence_parallel=True, + disable_kv_quant=True, + prefix=prefix, + ) + + def _reshape_to_seq_heads(self, tensor: torch.Tensor) -> torch.Tensor: + batch, seq_len, hidden_size = tensor.shape + return tensor.reshape(batch, seq_len, self.num_heads, hidden_size // self.num_heads) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + query = self.q_norm(self.q_linear(hidden_states)) + # MergedColumnParallelLinear emits [K_local; V_local], so the chunk is + # already this rank's head shard. + key, value = self.kv_linear(encoder_hidden_states).chunk(2, dim=-1) + key = self.k_norm(key) + if attention_mask is not None and attention_mask.ndim != 2: + raise ValueError( + f"Sana-WM cross-attention mask must be shaped [B, text_len], got {tuple(attention_mask.shape)}." + ) + query = self._reshape_to_seq_heads(query) + key = self._reshape_to_seq_heads(key) + value = self._reshape_to_seq_heads(value) + # The prompt padding mask goes to the shared Attention through + # ``AttentionMetadata``: backends take the [B, text_len] form directly + # and convert/reshape it themselves (``_maybe_reshape_attn_mask``), so + # cross-attention keeps backend selection. ``None`` keeps the mask-free + # fast path - callers normalise an all-ones mask to ``None`` rather than + # paying for a no-op mask. + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + attn = self.softmax_attn(query, key, value, attn_metadata) + return self.proj(attn.flatten(2, 3)) + + +class _ConvWrapper(nn.Module): + def __init__(self, conv: nn.Module) -> None: + super().__init__() + self.conv = conv + + +class SanaWmMbConvFfn(nn.Module): + def __init__(self, config: SanaWmConfig) -> None: + super().__init__() + hidden_size = config.hidden_size + expanded = int(hidden_size * config.mlp_ratio) * 2 + t_padding = config.t_kernel_size // 2 + self.glu_act = nn.SiLU() + self.inverted_conv = _ConvWrapper(nn.Conv2d(hidden_size, expanded, kernel_size=1)) + self.depth_conv = _ConvWrapper(nn.Conv2d(expanded, expanded, kernel_size=3, padding=1, groups=expanded)) + self.point_conv = _ConvWrapper(nn.Conv2d(expanded // 2, hidden_size, kernel_size=1, bias=False)) + self.t_conv = nn.Conv2d( + hidden_size, + hidden_size, + kernel_size=(config.t_kernel_size, 1), + padding=(t_padding, 0), + bias=False, + ) + + def forward(self, hidden_states: torch.Tensor, spatial_shape: tuple[int, int, int]) -> torch.Tensor: + """Match NVlabs ``GLUMBConvTemp.forward`` exactly: + + 1. Reshape ``(B, N=F*H*W, C) → (B*F, H, W, C) → (B*F, C, H, W)`` so the + spatial MBConv runs PER FRAME with proper 2D (H, W) geometry — the + 3×3 ``depth_conv`` must not span an (F, H*W) plane. + 2. Apply expand → depthwise 3×3 → GLU → contract on per-frame + (B*F, C, H, W). + 3. Temporal aggregation: reshape back to ``(B, C, F, H*W)`` + and add ``t_conv`` (kernel ``(3, 1)``) along the F axis. + 4. Reshape to ``(B, N, C)``. + """ + batch, _, hidden_size = hidden_states.shape + frames, height, width = spatial_shape + spatial_tokens = height * width + # (B, F*H*W, C) → (B*F, H, W, C) → (B*F, C, H, W) + x = ( + hidden_states.reshape(batch * frames, spatial_tokens, hidden_size) + .reshape(batch * frames, height, width, hidden_size) + .permute(0, 3, 1, 2) + ) + # Keep this NHWC-derived NCHW view non-contiguous like NVlabs: on bf16 + # CUDA convs the stride selects a different kernel, and a contiguous + # copy breaks late-step parity. + # NVlabs ``ConvLayer`` is `Conv → norm → act` with + # `act=("silu", "silu", None)`, so inverted_conv activates and the other + # two do not. ``_ConvWrapper`` stores only the raw Conv2d (to match the + # `inverted_conv.conv.weight` checkpoint key), hence the explicit SiLU. + x_inv = self.glu_act(self.inverted_conv.conv(x)) + x_dep = self.depth_conv.conv(x_inv) + value, gate = x_dep.chunk(2, dim=1) + x_glu = value * self.glu_act(gate) + x_pt = self.point_conv.conv(x_glu) # (B*F, hidden, H, W) + x = x_pt + # Temporal aggregation: (B*F, C, H, W) → (B, F, C, H*W) → (B, C, F, H*W) + x_temporal = x.reshape(batch, frames, hidden_size, spatial_tokens).permute(0, 2, 1, 3) + t_conv_out = self.t_conv(x_temporal) + x_temporal = x_temporal + t_conv_out + # → (B, F, H*W, C) → (B, N, C) + final = x_temporal.permute(0, 2, 3, 1).reshape(batch, frames * spatial_tokens, hidden_size).contiguous() + return final + + +class SanaWmBlock(nn.Module): + def __init__( + self, + config: SanaWmConfig, + *, + block_idx: int = 0, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + hidden_size = config.hidden_size + use_gdn = config.softmax_every_n <= 0 or (block_idx + 1) % config.softmax_every_n != 0 + use_plucker_proj = config.use_chunk_plucker_post_attn and ( + config.chunk_plucker_post_attn_blocks < 0 or block_idx < config.chunk_plucker_post_attn_blocks + ) + self.norm1 = AdaLayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.attn = SanaWmSelfAttention( + config, + use_gdn=use_gdn, + quant_config=quant_config, + prefix=f"{prefix}.attn" if prefix else "attn", + ) + self.cross_attn = SanaWmCrossAttention( + config, + quant_config=quant_config, + prefix=f"{prefix}.cross_attn" if prefix else "cross_attn", + ) + self.norm2 = AdaLayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.mlp = SanaWmMbConvFfn(config) + if use_plucker_proj: + self.plucker_proj: nn.Module | None = ColumnParallelLinear( + hidden_size, + hidden_size, + bias=True, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.plucker_proj" if prefix else "plucker_proj", + ) + else: + self.plucker_proj = None + self.scale_shift_table = nn.Parameter(torch.zeros(6, hidden_size)) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep_modulation: torch.Tensor, + spatial_shape: tuple[int, int, int], + rotary_emb: torch.Tensor | None = None, + camera_hidden_states: torch.Tensor | None = None, + cam_geometry: SanaWmCamGeometry | None = None, + encoder_attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + # Per-frame timestep contract (see SanaWmTransformer3DModel.forward). + if timestep_modulation.ndim > 2: + return self._forward_frame_aware( + hidden_states, + encoder_hidden_states, + timestep_modulation, + spatial_shape, + rotary_emb, + camera_hidden_states, + cam_geometry, + encoder_attention_mask, + ) + batch_size = hidden_states.shape[0] + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep_modulation.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + attn_input = self.norm1(hidden_states, scale_msa, shift_msa) + attn_output = self.attn(attn_input, spatial_shape, rotary_emb, cam_geometry) + if camera_hidden_states is not None and self.plucker_proj is not None: + attn_output = attn_output + self.plucker_proj(camera_hidden_states) + hidden_states = hidden_states + gate_msa * attn_output + hidden_states = hidden_states + self.cross_attn( + hidden_states, + encoder_hidden_states, + encoder_attention_mask, + ) + mlp_input = self.norm2(hidden_states, scale_mlp, shift_mlp) + hidden_states = hidden_states + gate_mlp * self.mlp(mlp_input, spatial_shape) + return hidden_states + + def _forward_frame_aware( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep_modulation: torch.Tensor, + spatial_shape: tuple[int, int, int], + rotary_emb: torch.Tensor | None, + camera_hidden_states: torch.Tensor | None, + cam_geometry: SanaWmCamGeometry | None, + encoder_attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Per-frame timestep modulation path matching NVlabs + ``SanaVideoMSCamCtrlBlock.forward_frame_aware``. + + ``timestep_modulation`` is ``(B, 1, F, 6*D)``. We split into the + six per-frame ``(B, F, 1, D)`` chunks via ``scale_shift_table`` + and apply ``shift/scale/gate`` per frame, broadcasting over the + spatial tokens within each frame. + """ + batch_size, token_count, hidden_size = hidden_states.shape + frames, height, width = spatial_shape + spatial_tokens = height * width + if token_count != frames * spatial_tokens: + raise ValueError(f"Sana-WM frame-aware block expects N=T*H*W, got N={token_count}, THW={spatial_shape}.") + if timestep_modulation.shape[2] != frames: + raise ValueError( + "Sana-WM frame-aware block: timestep frame axis must match spatial frames, " + f"got modulation frames={timestep_modulation.shape[2]} vs spatial frames={frames}." + ) + + # (B, 1, F, 6*D) -> (B, F, 6, D), then add the (1, 1, 6, D) table. + t_per_frame = timestep_modulation.reshape(batch_size, frames, 6, -1) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None, :, :] + t_per_frame + ).chunk(6, dim=-2) # each: (B, F, 1, D) + + # Apply per-frame modulation by reshaping x to (B, F, S, D), broadcasting + # scale/shift over the spatial axis, then flattening back. The reshape + # comes first because the norm only reduces over the last axis, so it is + # unaffected, and the modulation then broadcasts as written. + x_4d = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size) + x_msa_in = self.norm1(x_4d, scale_msa, shift_msa).reshape(batch_size, token_count, hidden_size) + attn_output = self.attn(x_msa_in, spatial_shape, rotary_emb, cam_geometry) + if camera_hidden_states is not None and self.plucker_proj is not None: + attn_output = attn_output + self.plucker_proj(camera_hidden_states) + attn_output_4d = attn_output.reshape(batch_size, frames, spatial_tokens, hidden_size) + hidden_states = hidden_states + (gate_msa * attn_output_4d).reshape(batch_size, token_count, hidden_size) + + hidden_states = hidden_states + self.cross_attn( + hidden_states, + encoder_hidden_states, + encoder_attention_mask, + ) + + x_4d = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size) + x_mlp_in = self.norm2(x_4d, scale_mlp, shift_mlp).reshape(batch_size, token_count, hidden_size) + mlp_out = self.mlp(x_mlp_in, spatial_shape) + mlp_out_4d = mlp_out.reshape(batch_size, frames, spatial_tokens, hidden_size) + hidden_states = hidden_states + (gate_mlp * mlp_out_4d).reshape(batch_size, token_count, hidden_size) + return hidden_states + + +class SanaWmFinalLayer(nn.Module): + def __init__( + self, + hidden_size: int, + patch_size: tuple[int, int, int], + out_channels: int, + *, + quant_config: Any = None, + prefix: str = "", + ) -> None: + super().__init__() + self.norm_final = AdaLayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.scale_shift_table = nn.Parameter(torch.zeros(2, hidden_size)) + self.out_channels = out_channels + out_features = math.prod(patch_size) * out_channels + self.linear: nn.Module = ColumnParallelLinear( + hidden_size, + out_features, + bias=True, + gather_output=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.linear" if prefix else "linear", + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep_embed: torch.Tensor, + spatial_shape: tuple[int, int, int] | None = None, + ) -> torch.Tensor: + # Per-frame timestep contract (see SanaWmTransformer3DModel.forward). + if timestep_embed.ndim > 2: + return self._forward_frame_aware(hidden_states, timestep_embed, spatial_shape) + shift, scale = (self.scale_shift_table[None] + timestep_embed[:, None]).chunk(2, dim=1) + return self.linear(self.norm_final(hidden_states, scale, shift)) + + def _forward_frame_aware( + self, + hidden_states: torch.Tensor, + timestep_embed: torch.Tensor, + spatial_shape: tuple[int, int, int] | None, + ) -> torch.Tensor: + """Per-frame final-layer modulation matching NVlabs + ``T2IFinalLayer.forward_frame_aware``. + + ``timestep_embed`` is ``(B, 1, F, D)``. We transpose to + ``(B, F, 1, D)`` so it adds correctly to the ``(1, 1, 2, D)`` + ``scale_shift_table`` and produces per-frame ``(B, F, 1, D)`` + ``shift`` / ``scale`` that broadcast over the spatial tokens + within each frame. + """ + batch_size, token_count, hidden_size = hidden_states.shape + frames = timestep_embed.shape[2] + if spatial_shape is not None: + spatial_tokens = spatial_shape[1] * spatial_shape[2] + else: + spatial_tokens = token_count // frames + if frames * spatial_tokens != token_count: + raise ValueError( + "Sana-WM frame-aware final layer: token count mismatch " + f"(N={token_count}, F={frames}, S={spatial_tokens})." + ) + # (B, 1, F, D) -> (B, F, 1, D); add (1, 1, 2, D); chunk into shift/scale: each (B, F, 1, D). + t_per_frame = timestep_embed.transpose(1, 2) + shift, scale = (self.scale_shift_table[None, None, :, :] + t_per_frame).chunk(2, dim=-2) + x_4d = hidden_states.reshape(batch_size, frames, spatial_tokens, hidden_size) + x_mod = self.norm_final(x_4d, scale, shift).reshape(batch_size, token_count, hidden_size) + return self.linear(x_mod) + + +class SanaWmTransformer3DModel(nn.Module): + """SANA-WM Stage-1 DiT: bidirectional Gated DeltaNet blocks with a softmax + attention block every ``softmax_every_n``.""" + + _repeated_blocks: ClassVar[list[str]] = ["SanaWmBlock"] + _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["blocks"] + _hsdp_shard_conditions: ClassVar[list[Any]] = [_is_sana_wm_transformer_block] + + def __init__( + self, + config: SanaWmConfig | None = None, + *, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + latent_channels: int = SANA_WM_STAGE1_LATENT_CHANNELS, + prompt_channels: int = SANA_WM_STAGE1_PROMPT_CHANNELS, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config or SanaWmConfig() + self.quant_config = quant_config + self.prefix = prefix + self._latent_channels = latent_channels + self._prompt_channels = prompt_channels + self.patch_size = _to_3tuple(self.config.patch_size) + self.x_embedder = SanaWmPatchEmbedMS3D(self.patch_size, self._latent_channels, self.config.hidden_size) + self.y_embedder = SanaWmTextEmbedder( + self._prompt_channels, + self.config.hidden_size, + self.config.model_max_length, + quant_config=self.quant_config, + prefix=f"{self.prefix}.y_embedder" if self.prefix else "y_embedder", + ) + self.t_embedder = SanaWmTimestepEmbedder( + SANA_WM_STAGE1_TIMESTEP_CHANNELS, + self.config.hidden_size, + quant_config=self.quant_config, + prefix=f"{self.prefix}.t_embedder" if self.prefix else "t_embedder", + ) + # t_block: SiLU → Linear(hidden, 6*hidden). Weight key is t_block.1.{weight,bias} + # to match the Stage-1 checkpoint layout (nn.Sequential index 1). + _t_block_linear: nn.Module = ColumnParallelLinear( + self.config.hidden_size, + 6 * self.config.hidden_size, + bias=True, + gather_output=True, + return_bias=False, + quant_config=self.quant_config, + prefix=f"{self.prefix}.t_block.1" if self.prefix else "t_block.1", + ) + self.t_block = nn.Sequential(nn.SiLU(), _t_block_linear) + self.plucker_embedder = SanaWmPatchEmbedMS3D( + self.patch_size, + self.config.chunk_plucker_channels, + self.config.hidden_size, + ) + self.raymap_embedder = SanaWmPatchEmbedMS3D(self.patch_size, 3, self.config.hidden_size) + self.blocks = nn.ModuleList( + [ + SanaWmBlock( + self.config, + block_idx=i, + quant_config=self.quant_config, + prefix=f"{self.prefix}.blocks.{i}" if self.prefix else f"blocks.{i}", + ) + for i in range(self.config.num_blocks) + ] + ) + self.final_layer = SanaWmFinalLayer( + self.config.hidden_size, + self.patch_size, + self._latent_channels, + quant_config=self.quant_config, + prefix=f"{self.prefix}.final_layer" if self.prefix else "final_layer", + ) + self.pos_embed = nn.Parameter(torch.zeros(1, 484, self.config.hidden_size)) + self.rope = SanaWmWanRotaryPosEmbed(self.config.linear_head_dim) + self.attention_y_norm = RMSNorm(self.config.hidden_size) + if device is not None or dtype is not None: + self.to(device=device, dtype=dtype) + + def _positional_embedding(self, token_count: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + pos_embed = self.pos_embed.to(device=device, dtype=dtype) + if pos_embed.shape[1] == token_count: + return pos_embed + pos_embed = pos_embed.transpose(1, 2) + pos_embed = F.interpolate(pos_embed, size=token_count, mode="linear", align_corners=False) + return pos_embed.transpose(1, 2) + + @staticmethod + def _match_tokens(hidden_states: torch.Tensor, expected_tokens: int) -> torch.Tensor: + if hidden_states.shape[1] == expected_tokens: + return hidden_states + if hidden_states.shape[1] > expected_tokens: + return hidden_states[:, :expected_tokens] + pad = hidden_states[:, -1:].expand(-1, expected_tokens - hidden_states.shape[1], -1) + return torch.cat([hidden_states, pad], dim=1) + + def _camera_hidden_states_from_conditions( + self, + *, + plucker: torch.Tensor | None, + spatial_raymap: torch.Tensor | None, + spatial_shape: tuple[int, int, int], + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor | None: + if plucker is None: + return None + if plucker.ndim == 4: + plucker = plucker.unsqueeze(0) + plucker = plucker.to(device=device, dtype=dtype) + if plucker.shape[0] == 1 and batch_size > 1: + plucker = plucker.expand(batch_size, -1, -1, -1, -1) + camera_hidden_states = self.plucker_embedder(plucker) + expected_tokens = spatial_shape[0] * spatial_shape[1] * spatial_shape[2] + camera_hidden_states = self._match_tokens(camera_hidden_states, expected_tokens) + + # Fuse per-pixel ray-direction map via raymap_embedder only on the + # non-plucker path. NVlabs skips the absmap/raymap embedder whenever + # chunk-plucker post-attention injection is enabled. + if spatial_raymap is not None and not self.config.use_chunk_plucker_post_attn: + if spatial_raymap.ndim == 4: # [C, F, H, W] → [1, C, F, H, W] + spatial_raymap = spatial_raymap.unsqueeze(0) + spatial_raymap = spatial_raymap.to(device=device, dtype=dtype) + if spatial_raymap.shape[0] == 1 and batch_size > 1: + spatial_raymap = spatial_raymap.expand(batch_size, -1, -1, -1, -1) + ray_hidden = self._match_tokens(self.raymap_embedder(spatial_raymap), expected_tokens) + camera_hidden_states = camera_hidden_states + ray_hidden + + return camera_hidden_states + + def _unpatchify(self, hidden_states: torch.Tensor, spatial_shape: tuple[int, int, int]) -> torch.Tensor: + batch_size = hidden_states.shape[0] + frames, height, width = spatial_shape + patch_frames, patch_height, patch_width = self.patch_size + hidden_states = hidden_states.reshape( + batch_size, + frames, + height, + width, + patch_frames, + patch_height, + patch_width, + self._latent_channels, + ) + hidden_states = torch.einsum("nfhwopqc->ncfohpwq", hidden_states) + return hidden_states.reshape( + batch_size, + self._latent_channels, + frames * patch_frames, + height * patch_height, + width * patch_width, + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor | float | int, + *, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + camera_hidden_states: torch.Tensor | None = None, + plucker: torch.Tensor | None = None, + raymap: torch.Tensor | None = None, + spatial_raymap: torch.Tensor | None = None, + ) -> torch.Tensor: + if hidden_states.ndim != 5: + raise ValueError("Sana-WM transformer expects latent input shaped [B, C, F, H, W].") + batch_size = hidden_states.shape[0] + latent_shape = hidden_states.shape[2:] + hidden_states, spatial_shape = self.x_embedder.project_with_shape(hidden_states) + rotary_emb = None + if self.config.pos_embed_type == "wan_rope": + rotary_emb = self.rope(spatial_shape, hidden_states.device) + else: + hidden_states = hidden_states + self._positional_embedding( + hidden_states.shape[1], hidden_states.dtype, hidden_states.device + ) + + encoder_hidden_states = self.y_embedder( + encoder_hidden_states.to(device=hidden_states.device) if encoder_hidden_states is not None else None, + batch_size=batch_size, + dtype=hidden_states.dtype, + ) + encoder_hidden_states = self.attention_y_norm(encoder_hidden_states) + # No mask means "attend to every prompt token"; keep it ``None`` so + # cross-attention takes the mask-free attention path rather than + # building an all-ones mask the backend would have to reshape per block. + if encoder_attention_mask is not None: + encoder_attention_mask = encoder_attention_mask.to(device=hidden_states.device) + if encoder_attention_mask.shape != encoder_hidden_states.shape[:2]: + raise ValueError( + "Sana-WM encoder_attention_mask must match text tokens " + f"{tuple(encoder_hidden_states.shape[:2])}, got {tuple(encoder_attention_mask.shape)}." + ) + + if not torch.is_tensor(timestep): + timestep = torch.tensor([timestep], device=hidden_states.device, dtype=torch.float32) + timestep = timestep.to(device=hidden_states.device) + if timestep.ndim == 0: + timestep = timestep.expand(batch_size) + elif timestep.ndim == 1 and timestep.shape[0] == 1 and batch_size > 1: + timestep = timestep.expand(batch_size) + # Per-frame timestep contract: a rank > 1 ``timestep`` (``(B, 1, F)`` + # from the LTX flow-matching sampler) embeds on the flattened token axis + # and unflattens back, giving ``(B, 1, F, D)`` / ``(B, 1, F, 6*D)``. + # SanaWmBlock and SanaWmFinalLayer dispatch on ``ndim > 2`` to the + # frame-aware paths that broadcast modulation over spatial tokens. + timestep_shape = tuple(timestep.shape) + time_embed = self.t_embedder(timestep) # (numel, D) + _t_silu = self.t_block[0](time_embed) + timestep_modulation = self.t_block[1](_t_silu).to(hidden_states.dtype) + if len(timestep_shape) > 1: + time_embed = time_embed.unflatten(0, timestep_shape) + timestep_modulation = timestep_modulation.unflatten(0, timestep_shape) + + if camera_hidden_states is None: + camera_hidden_states = self._camera_hidden_states_from_conditions( + plucker=plucker, + spatial_raymap=spatial_raymap, + spatial_shape=spatial_shape, + batch_size=batch_size, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # Build (B, T_latent, 20) camera_conditions from the raw raymap, then + # derive the UCPE geometry ONCE for the whole forward. The ray grid, + # its SE(3) inverse and the camera RoPE tables are identical in every + # block, so building them per block cost ~2400 rebuilds per request + # (20 blocks x 60 steps x 2 CFG branches). When the caller passes + # pre-embedded camera_hidden_states without raymap, the cam branch + # becomes a no-op (main-branch only). + cam_geometry: SanaWmCamGeometry | None = None + if raymap is not None: + # ``_pack_camera_conditions`` emits an unbatched (T_latent, 20). + camera_conditions = raymap.unsqueeze(0) + if batch_size > 1: + camera_conditions = camera_conditions.expand(batch_size, -1, -1) + camera_conditions = camera_conditions.to(device=hidden_states.device, dtype=hidden_states.dtype) + # cam_head_dim comes from the attention module rather than being + # re-derived here: __init__ already asserts it equals the main + # head_dim and is consistent across blocks. + attn = self.blocks[0].attn + cam_geometry = prepare_cam_geometry( + camera_conditions=camera_conditions, + spatial_shape=spatial_shape, + patch_size=self.patch_size, + head_dim=attn.cam_head_dim, + rotary_emb=attn._ucpe_rotary_freqs(rotary_emb), + ) + + for block in self.blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_modulation, + spatial_shape, + rotary_emb, + camera_hidden_states, + cam_geometry, + encoder_attention_mask, + ) + + hidden_states = self.final_layer(hidden_states, time_embed.to(hidden_states.dtype), spatial_shape) + hidden_states = self._unpatchify(hidden_states, spatial_shape) + if hidden_states.shape[2:] != latent_shape: + hidden_states = F.interpolate(hidden_states, size=latent_shape, mode="trilinear", align_corners=False) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, Any]]) -> set[str]: + """Stream checkpoint tensors into the eagerly-built modules. + + Follows the wan2_2 idiom: one ``named_parameters`` lookup table plus the + per-tensor ``weight_loader``, which narrows full checkpoint tensors to + the TP-local shard at copy time for both vLLM parallel layers and the + plain parameters marked by ``_shard_param_across_tp``. The source tensor + is dropped each iteration — no copy of the checkpoint is retained. + """ + params_dict = dict(self.named_parameters()) + buffers_dict = dict(self.named_buffers()) + loaded: set[str] = set() + unmapped: list[str] = [] + duplicates: list[str] = [] + + for source_name, tensor in weights: + # Weights ship in the standard diffusers ``transformer/`` layout, so + # the checkpoint key is already the module-local parameter name. + if source_name in loaded: + duplicates.append(source_name) + continue + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"Sana-WM weight {source_name!r} must be a torch.Tensor, got {type(tensor).__name__}.") + target = params_dict.get(source_name) + if target is None: + target = buffers_dict.get(source_name) + if target is None: + unmapped.append(source_name) + continue + weight_loader = getattr(target, "weight_loader", None) + if callable(weight_loader): + # vLLM parallel layers and the plain params marked by + # `_shard_param_across_tp` narrow to the local shard here. + weight_loader(target, tensor) + else: + if tuple(target.shape) != tuple(tensor.shape): + raise ValueError( + f"Sana-WM weight shape mismatch for {source_name}: " + f"expected {tuple(target.shape)}, got {tuple(tensor.shape)}." + ) + with torch.no_grad(): + target.copy_(tensor.to(device=target.device, dtype=target.dtype)) + loaded.add(source_name) + + if unmapped or duplicates: + details = [] + if unmapped: + details.append(f"unmapped={unmapped[:10]}") + if duplicates: + details.append(f"duplicates={duplicates[:10]}") + raise ValueError("Invalid SANA-WM Stage-1 checkpoint keys: " + "; ".join(details)) + return loaded diff --git a/vllm_omni/diffusion/models/sana_wm/ucpe.py b/vllm_omni/diffusion/models/sana_wm/ucpe.py new file mode 100644 index 000000000..88a3a115d --- /dev/null +++ b/vllm_omni/diffusion/models/sana_wm/ucpe.py @@ -0,0 +1,536 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SANA-WM UCPE (Unified Camera Pose Embedding) per-block attention transforms. + +Ported from NVlabs/Sana ``sana_camctrl_blocks.py`` for the inference-only +SANA-WM bidirectional 1600M release. Scope is intentionally narrow: + +* Pinhole camera only (xi=0); the UCM ``xi`` parameter is fixed. +* Inference path only — no training-time camera-branch dropout. +* Online computation only — the precomputed ``cam_pos_embeds`` shortcut + used by NVlabs' fused kernels is omitted because vLLM-Omni builds the + matrices once per transformer forward and shares them across all blocks. +* ``apply_vo=True`` only — the SANA-WM checkpoint always uses the + inverse-output transform. + +Public surface: + ``prepare_cam_geometry(...)`` -> :class:`SanaWmCamGeometry`, built once per + transformer forward and threaded through every block. It carries both the + ``(apply_q, apply_kv, apply_output)`` closures used by the softmax camera + branch and the projection/RoPE tensors used by ``cam_prep_func(...)``, the + NVlabs fused camera-prep contract on the GDN camera branch. + +Each closure transforms a tensor of shape ``(B, num_heads, N, D)`` where +``N = T * H * W`` (matching the GDN token layout) and ``D == head_dim``, +returning a tensor of the same shape. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from functools import partial + +import torch +import torch.nn.functional as F + +# --------------------------------------------------------------------------- +# Pinhole ray-grid construction +# --------------------------------------------------------------------------- + + +def _make_pixel_grid( + height: int, + width: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Return ``(H, W, 3)`` grid ``(x, y, 1)`` of pixel coordinates.""" + xs = torch.arange(width, device=device, dtype=dtype) + ys = torch.arange(height, device=device, dtype=dtype) + y_grid, x_grid = torch.meshgrid(ys, xs, indexing="ij") + return torch.stack((x_grid, y_grid, torch.ones_like(x_grid)), dim=-1) + + +def _unproject_pinhole( + fx: torch.Tensor, + fy: torch.Tensor, + cx: torch.Tensor, + cy: torch.Tensor, + height: int, + width: int, +) -> torch.Tensor: + """Unproject a pixel grid into camera-frame direction vectors. + + Args: + fx, fy, cx, cy: shape ``(B, T)`` intrinsics in token-grid units — the + caller divides all four by patch size before calling. + height, width: token-grid size. + + Returns: + ``(B, T, H, W, 3)`` direction tensor (NOT normalised). + """ + B, T = fx.shape + device = fx.device + dtype = fx.dtype + grid = _make_pixel_grid(height, width, device, dtype) # (H, W, 3) + u = grid[..., 0] # (H, W) + v = grid[..., 1] + fx_e = fx.view(B, T, 1, 1) + fy_e = fy.view(B, T, 1, 1) + cx_e = cx.view(B, T, 1, 1) + cy_e = cy.view(B, T, 1, 1) + x = (u - cx_e) / fx_e + y = (v - cy_e) / fy_e + z = torch.ones_like(x) + # Match NVlabs UCM normalisation at xi=0: alpha = 1, gamma = 1/(1+r^2), + # X = gamma*x, Y = gamma*y, Z = gamma. The resulting vector has unit + # norm. Computed compactly via direct stack + normalise. + d_cam = torch.stack((x, y, z), dim=-1) # (B, T, H, W, 3) + d_cam = d_cam / d_cam.norm(dim=-1, keepdim=True).clamp_min(1e-12) + return d_cam + + +def _world_to_ray_mats( + d_cam: torch.Tensor, # (B, T, H, W, 3) unit rays in camera frame + c2w: torch.Tensor, # (B, T, 4, 4) camera-to-world SE(3) +) -> torch.Tensor: + """Build per-pixel ``ray <- world`` SE(3) transforms ``(B, T, H, W, 4, 4)``. + + The ray frame is anchored at the camera origin with the +Z axis along + the per-pixel viewing direction and a +Y axis defined by the camera-Y + cross product. This matches NVlabs ``world_to_ray_mats``. + """ + B, T, H, W, _ = d_cam.shape + device = d_cam.device + dtype = d_cam.dtype + rot_cam = c2w[..., :3, :3] + trans_cam = c2w[..., :3, 3] + d_world = torch.einsum("btij,bthwj->bthwi", rot_cam, d_cam) + cam_y = rot_cam[..., :, 1].unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3) + z_ray = F.normalize(d_world, dim=-1, eps=1e-6) + x_ray = F.normalize(torch.cross(cam_y, z_ray, dim=-1), dim=-1, eps=1e-6) + y_ray = F.normalize(torch.cross(z_ray, x_ray, dim=-1), dim=-1, eps=1e-6) + rot_local_to_world = torch.stack((x_ray, y_ray, z_ray), dim=-1) # (B,T,H,W,3,3) + rot_world_to_local = rot_local_to_world.transpose(-1, -2) + t_world = trans_cam.unsqueeze(2).unsqueeze(3).expand(B, T, H, W, 3) + t_w2l = -torch.einsum("bthwij,bthwj->bthwi", rot_world_to_local, t_world) + raymats = torch.zeros(B, T, H, W, 4, 4, device=device, dtype=dtype) + raymats[..., :3, :3] = rot_world_to_local + raymats[..., :3, 3] = t_w2l + raymats[..., 3, 3] = 1.0 + # Guard NaNs from degenerate pixels (e.g. nan direction at infinity). + nan_mask = torch.isnan(d_world).any(-1) + if nan_mask.any(): + eye = torch.eye(4, device=device, dtype=dtype) + raymats[nan_mask] = eye + return raymats + + +def _process_camera_conditions( + camera_conditions: torch.Tensor, # (B, T, 20) + hw: tuple[int, int, int], + patch_size: tuple[int, int, int] = (1, 1, 1), +) -> torch.Tensor: + """Convert ``(B, T, 20)`` SANA-WM camera conditions to ``(B, T, H, W, 4, 4)`` + ``ray <- world`` transforms. + + Layout of the last axis is ``[c2w_flat (16), fx, fy, cx, cy]``. + Intrinsics follow the NVlabs UCPE contract: the packed camera payload is + expressed on the pre-patch latent image grid, while ``HW`` is the + transformer token grid. ``patch_size`` maps those units to the token grid + before unprojection. The released SANA-WM P1 config uses ``(1, 1, 1)``. + + Args: + camera_conditions: ``(B, T, 20)`` from ``_pack_camera_conditions``. + hw: ``(T, H, W)`` token-grid shape. Equal to the latent grid only + because the released config patches by ``(1, 1, 1)``. + """ + B, T_cond, last = camera_conditions.shape + if last != 20: + raise ValueError(f"SANA-WM camera_conditions last dim must be 20, got {last}.") + T_latent, H_latent, W_latent = hw + if T_cond != T_latent: + raise ValueError(f"SANA-WM camera_conditions frames {T_cond} != latent frames {T_latent}.") + c2w = camera_conditions[..., :16].reshape(B, T_latent, 4, 4) + patch_t, patch_h, patch_w = patch_size + del patch_t + fx = camera_conditions[..., 16] / float(patch_w) + fy = camera_conditions[..., 17] / float(patch_h) + cx = camera_conditions[..., 18] / float(patch_w) + cy = camera_conditions[..., 19] / float(patch_h) + d_cam = _unproject_pinhole(fx, fy, cx, cy, H_latent, W_latent) + raymats = _world_to_ray_mats(d_cam, c2w) + return raymats + + +def _invert_se3(transforms: torch.Tensor) -> torch.Tensor: + """Closed-form inverse of a stack of 4x4 SE(3) matrices.""" + if transforms.shape[-2:] != (4, 4): + raise ValueError(f"SE3 inverse expects (..., 4, 4), got {tuple(transforms.shape)}.") + rot_inv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = rot_inv + out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv, transforms[..., :3, 3]) + out[..., 3, 3] = 1.0 + return out + + +# --------------------------------------------------------------------------- +# Per-token transform primitives +# --------------------------------------------------------------------------- + + +def _apply_ray_projmat( + feats: torch.Tensor, # (B, num_heads, N, D) + matrix: torch.Tensor, # (B, N, 4, 4) +) -> torch.Tensor: + """Apply a per-token 4x4 projection to features grouped by 4 channels. + + ``D`` must be a multiple of 4. Out-of-place; preserves ``feats.shape``. + """ + B, num_heads, N, D = feats.shape + if D % 4 != 0: + raise ValueError(f"UCPE projmat expects D divisible by 4, got D={D}.") + feats_grouped = feats.reshape(B, num_heads, N, D // 4, 4) + # matrix: (B, N, 4, 4); feats_grouped: (B, num_heads, N, K, 4) + out = torch.einsum("bnij,bhnkj->bhnki", matrix, feats_grouped) + return out.reshape(B, num_heads, N, D) + + +def _apply_complex_rope( + hidden_states: torch.Tensor, + freqs: torch.Tensor, + *, + inverse: bool = False, +) -> torch.Tensor: + """Apply complex RoPE: treat last-dim pairs as ``(real, imag)`` and multiply + by ``freqs``. Computation runs in fp64 to match the NVlabs reference. + + Args: + hidden_states: ``(..., D)`` with ``D % 2 == 0``. + freqs: complex tensor broadcastable to ``(..., D // 2)``. + inverse: when ``True`` apply ``freqs.conj()`` (used for ``apply_o``). + """ + x_real = hidden_states.to(torch.float64).contiguous() + x_complex = torch.view_as_complex(x_real.unflatten(-1, (-1, 2))) + if inverse: + freqs = freqs.conj() + x_out = torch.view_as_real(x_complex * freqs).flatten(-2, -1) + return x_out.to(hidden_states.dtype) + + +def _apply_block_diagonal( + feats: torch.Tensor, + func_size_pairs: list[tuple[Callable[[torch.Tensor], torch.Tensor], int]], +) -> torch.Tensor: + """Split ``feats`` along the last axis by ``block_sizes`` and apply one + callable per block, then concatenate. Preserves ``feats.shape``. + """ + funcs, block_sizes = zip(*func_size_pairs) + if feats.shape[-1] != sum(block_sizes): + raise ValueError(f"UCPE block-diagonal block sum {sum(block_sizes)} != feats D {feats.shape[-1]}.") + blocks = torch.split(feats, list(block_sizes), dim=-1) + return torch.cat([func(block) for func, block in zip(funcs, blocks)], dim=-1) + + +def _slice_rope_for_cam( + rotary_emb: torch.Tensor | None, + rope_dim: int, +) -> torch.Tensor | None: + """Re-slice WAN-style RoPE frequencies for a smaller ``rope_dim`` using the + same ``(T, H, W)`` split as the main branch. Mirrors NVlabs slicing logic. + + The T/H/W band layout is read off the input shape + (``rotary_emb.shape[-1] * 2``) rather than a separately supplied head_dim, + so the split stays correct without the caller passing it. ``cam_head_dim`` + still has to equal the main ``head_dim`` — ``SanaWmSelfAttention.__init__`` + rejects anything else. + """ + if rotary_emb is None: + return None + # Derive WAN RoPE T/H/W band offsets from the actual input tensor size, + # not from a head_dim argument. This is correct regardless of whether + # cam_head_dim == main_head_dim. + in_head_dim = rotary_emb.shape[-1] * 2 + orig_t = in_head_dim // 2 - 2 * (in_head_dim // 6) + orig_h = in_head_dim // 6 + new_t = rope_dim // 2 - 2 * (rope_dim // 6) + new_h = rope_dim // 6 + new_w = rope_dim // 6 + t_part = rotary_emb[..., :new_t] + h_part = rotary_emb[..., orig_t : orig_t + new_h] + w_part = rotary_emb[..., orig_t + orig_h : orig_t + orig_h + new_w] + return torch.cat([t_part, h_part, w_part], dim=-1) + + +# --------------------------------------------------------------------------- +# Apply-fn preparation +# --------------------------------------------------------------------------- + + +def _identity(x: torch.Tensor) -> torch.Tensor: + return x + + +def _prepare_ray_apply_fns( + head_dim: int, + proj: torch.Tensor, # (B, N, 4, 4) ray <- world + proj_t: torch.Tensor, # (B, N, 4, 4) ray <- world transpose + proj_inv: torch.Tensor, # (B, N, 4, 4) world <- ray + rotary_emb: torch.Tensor | None = None, +) -> tuple[Callable, Callable, Callable]: + """Build ``(apply_q, apply_kv, apply_o)`` block-diagonal callables. + + Each callable transforms a ``(B, num_heads, N, head_dim)`` tensor: + the first ``head_dim // 2`` channels go through the per-token 4x4 + projection, the remaining ``head_dim // 2`` channels go through complex + RoPE (or identity when ``rotary_emb`` is None). + """ + if rotary_emb is not None: + rope_fwd = partial(_apply_complex_rope, freqs=rotary_emb, inverse=False) + rope_inv = partial(_apply_complex_rope, freqs=rotary_emb, inverse=True) + else: + rope_fwd = _identity + rope_inv = _identity + + half = head_dim // 2 + + transforms_q = [ + (partial(_apply_ray_projmat, matrix=proj_t), half), + (rope_fwd, half), + ] + transforms_kv = [ + (partial(_apply_ray_projmat, matrix=proj_inv), half), + (rope_fwd, half), + ] + transforms_o = [ + (partial(_apply_ray_projmat, matrix=proj), half), + (rope_inv, half), + ] + + apply_q = partial(_apply_block_diagonal, func_size_pairs=transforms_q) + apply_kv = partial(_apply_block_diagonal, func_size_pairs=transforms_kv) + apply_o = partial(_apply_block_diagonal, func_size_pairs=transforms_o) + return apply_q, apply_kv, apply_o + + +@dataclass(frozen=True) +class SanaWmCamGeometry: + """UCPE geometry shared by every block of one transformer forward. + + The ray grid, its SE(3) inverse and the camera RoPE tables depend only on + the camera payload, the latent grid, the patch size and the head dim — all + fixed for the whole forward — yet each block's camera branch used to + rebuild them. At 20 blocks x 60 steps x 2 CFG branches that is ~2400 + rebuilds of an 18480-token ray grid and its 4x4 inverses per request. + + ``apply_q``/``apply_kv``/``apply_output`` serve the softmax camera branch; + ``proj_q``/``proj_kv``/``rope_cos``/``rope_sin`` serve the fused + ``cam_prep_func`` contract on the GDN branch. Both are derived from the + same projection matrices, so they are built together. + """ + + proj_q: torch.Tensor + proj_kv: torch.Tensor + rope_cos: torch.Tensor + rope_sin: torch.Tensor + apply_q: Callable[[torch.Tensor], torch.Tensor] + apply_kv: Callable[[torch.Tensor], torch.Tensor] + apply_output: Callable[[torch.Tensor], torch.Tensor] + + +def _prepare_ucpe_rope_tables( + rotary_emb_cam: torch.Tensor, + token_count: int, + half_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert complex RoPE freqs to NVlabs' interleaved cos/sin tables. + + The output convention is ``out[d] = x[d] * cos[d] + x[d ^ 1] * sin[d]``. + For every complex pair, ``sin[2i] = -imag(freq[i])`` and + ``sin[2i + 1] = imag(freq[i])``. + """ + freqs = rotary_emb_cam + while freqs.ndim > 2 and freqs.shape[0] == 1: + freqs = freqs.squeeze(0) + if freqs.ndim != 2: + raise ValueError(f"Sana-WM cam RoPE freqs must be 2D after squeeze, got {tuple(freqs.shape)}.") + if freqs.shape[0] != token_count: + raise ValueError(f"Sana-WM cam RoPE token count {freqs.shape[0]} != {token_count}.") + if freqs.shape[1] * 2 != half_dim: + raise ValueError(f"Sana-WM cam RoPE dim {freqs.shape[1] * 2} != half_dim {half_dim}.") + cos_half = freqs.real.float() + sin_half = freqs.imag.float() + rope_cos = cos_half.repeat_interleave(2, dim=-1).contiguous() + rope_sin = torch.stack((-sin_half, sin_half), dim=-1).reshape(token_count, half_dim).contiguous() + return rope_cos, rope_sin + + +def prepare_cam_geometry( + *, + camera_conditions: torch.Tensor, + spatial_shape: tuple[int, int, int], + patch_size: tuple[int, int, int], + head_dim: int, + rotary_emb: torch.Tensor | None, +) -> SanaWmCamGeometry: + """Build the per-forward UCPE geometry: ray projections, RoPE tables, apply fns. + + Call once per transformer forward and share the result across every block — + see :class:`SanaWmCamGeometry` for why. + """ + if head_dim % 4 != 0: + raise ValueError(f"UCPE head_dim must be divisible by 4, got {head_dim}.") + if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0: + raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.") + + batch = camera_conditions.shape[0] + frames, height, width = spatial_shape + token_count = frames * height * width + half_dim = head_dim // 2 + + raymats = _process_camera_conditions(camera_conditions, spatial_shape, patch_size=patch_size) + proj = raymats.reshape(batch, token_count, 4, 4).contiguous() + proj_q = proj.transpose(-1, -2).contiguous() + proj_kv = _invert_se3(proj).contiguous() + + rotary_emb_cam = _slice_rope_for_cam(rotary_emb, half_dim) + if rotary_emb_cam is None: + rope_cos = torch.ones(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32) + rope_sin = torch.zeros(token_count, half_dim, device=camera_conditions.device, dtype=torch.float32) + else: + rope_cos, rope_sin = _prepare_ucpe_rope_tables(rotary_emb_cam, token_count, half_dim) + + apply_q, apply_kv, apply_output = _prepare_ray_apply_fns(head_dim, proj, proj_q, proj_kv, rotary_emb=rotary_emb_cam) + return SanaWmCamGeometry( + proj_q=proj_q, + proj_kv=proj_kv, + rope_cos=rope_cos, + rope_sin=rope_sin, + apply_q=apply_q, + apply_kv=apply_kv, + apply_output=apply_output, + ) + + +def _validate_cam_prep_inputs( + q_raw: torch.Tensor, + k_raw: torch.Tensor, + v_raw: torch.Tensor, + proj_q: torch.Tensor, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, +) -> int: + """Validate the fused cam-prep contract and return ``head_dim // 2``.""" + if q_raw.shape != k_raw.shape or q_raw.shape != v_raw.shape: + raise ValueError(f"Sana-WM cam prep q/k/v shapes must match, got {q_raw.shape}, {k_raw.shape}, {v_raw.shape}.") + if q_raw.ndim != 4: + raise ValueError(f"Sana-WM cam prep q/k/v must be (B,N,H,D), got {tuple(q_raw.shape)}.") + batch, token_count, _num_heads, head_dim = q_raw.shape + if head_dim % 2 != 0 or (head_dim // 2) % 4 != 0: + raise ValueError(f"Sana-WM cam prep head_dim={head_dim} must satisfy D % 2 == 0 and (D/2) % 4 == 0.") + half_dim = head_dim // 2 + if proj_q.shape != (batch, token_count, 4, 4): + raise ValueError(f"Sana-WM cam prep proj_q shape {tuple(proj_q.shape)} != {(batch, token_count, 4, 4)}.") + if proj_kv.shape != (batch, token_count, 4, 4): + raise ValueError(f"Sana-WM cam prep proj_kv shape {tuple(proj_kv.shape)} != {(batch, token_count, 4, 4)}.") + if rope_cos.shape != (token_count, half_dim) or rope_sin.shape != (token_count, half_dim): + raise ValueError( + "Sana-WM cam prep rope table shapes must be " + f"{(token_count, half_dim)}, got {tuple(rope_cos.shape)} and {tuple(rope_sin.shape)}." + ) + return half_dim + + +def _apply_ray_projection(feats: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor: + batch, token_count, num_heads, half_dim = feats.shape + grouped = feats.reshape(batch, token_count, num_heads, half_dim // 4, 4) + out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), grouped) + return out.reshape(batch, token_count, num_heads, half_dim) + + +def _apply_interleaved_rope(feats: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor) -> torch.Tensor: + half_dim = feats.shape[-1] + pair_index = torch.arange(half_dim, device=feats.device) ^ 1 + paired = feats.index_select(-1, pair_index) + return feats * rope_cos.reshape(1, rope_cos.shape[0], 1, half_dim) + paired * rope_sin.reshape( + 1, rope_sin.shape[0], 1, half_dim + ) + + +def cam_prep_func( + q_normed: torch.Tensor, + k_normed: torch.Tensor, + v_raw: torch.Tensor, + *, + proj_q: torch.Tensor, + proj_kv: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + k_scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Native equivalent of NVlabs ``cam_prep_func``. + + The caller is responsible for the q/k RMSNorm, which must run through the + model's norm modules so a tensor-parallel shard still normalises over the + global channel width. This function only applies ReLU, the K scale, the + UCPE ray projection and RoPE. + + Args: + q_normed: ``(B, N, H, D)`` camera Q, already RMSNorm'd. + k_normed: ``(B, N, H, D)`` camera K, already RMSNorm'd. K must also + include the temporal short convolution. + v_raw: ``(B, N, H, D)`` raw camera V. + proj_q: ``(B, N, 4, 4)`` Q ray projection matrices. + proj_kv: ``(B, N, 4, 4)`` K/V ray projection matrices. + rope_cos: ``(N, D//2)`` interleaved-pair RoPE cosine table. + rope_sin: ``(N, D//2)`` interleaved-pair RoPE sine table. + k_scale: ``D^-0.5 * spatial_tokens^-0.5``. + + Returns: + ``q_trans``, ``k_trans``, ``v_trans`` in ``(B, H, D, N)`` layout, and + ``inflation_sq`` in ``(B, H, N)``. + """ + half_dim = _validate_cam_prep_inputs( + q_normed, + k_normed, + v_raw, + proj_q, + proj_kv, + rope_cos, + rope_sin, + ) + + q_norm = torch.relu(q_normed.float()) + k_norm = torch.relu(k_normed.float()) * float(k_scale) + value = v_raw.float() + + q_half = _apply_ray_projection(q_norm[..., :half_dim], proj_q) + k_half = _apply_ray_projection(k_norm[..., :half_dim], proj_kv) + v_half = _apply_ray_projection(value[..., :half_dim], proj_kv) + + q_rope = _apply_interleaved_rope(q_norm[..., half_dim:], rope_cos.float(), rope_sin.float()) + k_rope = _apply_interleaved_rope(k_norm[..., half_dim:], rope_cos.float(), rope_sin.float()) + v_rope = _apply_interleaved_rope(value[..., half_dim:], rope_cos.float(), rope_sin.float()) + + k_pre_sq = k_norm.square().sum(dim=-1).permute(0, 2, 1).contiguous() + k_post = torch.cat((k_half, k_rope), dim=-1) + k_post_sq = k_post.square().sum(dim=-1).permute(0, 2, 1).contiguous() + inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12) + + def to_bhdn(tensor: torch.Tensor) -> torch.Tensor: + return tensor.permute(0, 2, 3, 1).contiguous().to(v_raw.dtype) + + q_out = to_bhdn(torch.cat((q_half, q_rope), dim=-1)) + k_out = to_bhdn(k_post) + v_out = to_bhdn(torch.cat((v_half, v_rope), dim=-1)) + return q_out, k_out, v_out, inflation_sq + + +__all__ = [ + "SanaWmCamGeometry", + "cam_prep_func", + "prepare_cam_geometry", +] diff --git a/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py b/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py index 0bb8a48cd..ac7231fbc 100644 --- a/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py +++ b/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py @@ -98,7 +98,7 @@ class StableDiffusionXLPipeline( prefetch_list=sdxl_subfolders, local_files_only=local_files_only, torch_dtype=dtype, - ) + ).to(self.device) self.text_encoder_2 = from_pretrained_with_prefetch( CLIPTextModelWithProjection.from_pretrained, model, @@ -106,7 +106,7 @@ class StableDiffusionXLPipeline( prefetch_list=sdxl_subfolders, local_files_only=local_files_only, torch_dtype=dtype, - ) + ).to(self.device) self.unet = SDXLUNet2DConditionModel(od_config=od_config) self.vae = from_pretrained_with_prefetch( DistributedAutoencoderKL.from_pretrained, @@ -146,7 +146,7 @@ class StableDiffusionXLPipeline( ) text_input_ids = text_inputs.input_ids - text_encoder_device = next(text_encoder.parameters()).device + text_encoder_device = self.device outputs = text_encoder(text_input_ids.to(text_encoder_device), output_hidden_states=True) prompt_embeds = outputs.hidden_states[-2].to(dtype=self.od_config.dtype, device=self.device) diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 67e141e6a..3a828fa69 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -17,7 +17,6 @@ Key integration points: from __future__ import annotations -import json import math import os from collections.abc import Iterable @@ -28,6 +27,7 @@ import numpy as np import torch import torch.nn as nn import torchvision.transforms as T +from huggingface_hub import snapshot_download from PIL import Image from transformers import AutoTokenizer from vllm.logger import init_logger @@ -40,6 +40,9 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.transformers_utils.configs.sensenova_u1 import ( + SenseNovaU1Config, +) from .sensenova_u1_transformer import ( SenseNovaU1ForCausalLM, @@ -86,16 +89,12 @@ SYSTEM_MESSAGE_FOR_GEN = ( # --------------------------------------------------------------------------- -# Config parsed from config.json +# Config # --------------------------------------------------------------------------- - - def _resolve_model_path(model_path: str) -> str: """Resolve a HuggingFace model ID or local path to a local directory.""" if os.path.isdir(model_path): return model_path - from huggingface_hub import snapshot_download - return snapshot_download(model_path) @@ -182,73 +181,6 @@ def _load_image_native( return flatten_pv, grid_hw -def _load_sensenova_u1_config(model_path: str): - """Parse config.json and return (top_config, llm_config, vision_config) namespaces.""" - local_path = _resolve_model_path(model_path) - cfg_path = os.path.join(local_path, "config.json") - with open(cfg_path, encoding="utf-8") as f: - raw = json.load(f) - - llm_raw = raw.get("llm_config", {}) - vis_raw = raw.get("vision_config", {}) - - # LLM config needs these extra fields for 3D RoPE - llm_cfg = SimpleNamespace( - hidden_size=llm_raw.get("hidden_size", 4096), - intermediate_size=llm_raw.get("intermediate_size", 11008), - num_hidden_layers=llm_raw.get("num_hidden_layers", 32), - num_attention_heads=llm_raw.get("num_attention_heads", 32), - num_key_value_heads=llm_raw.get("num_key_value_heads", llm_raw.get("num_attention_heads", 32)), - head_dim=llm_raw.get("head_dim", llm_raw.get("hidden_size", 4096) // llm_raw.get("num_attention_heads", 32)), - hidden_act=llm_raw.get("hidden_act", "silu"), - rms_norm_eps=llm_raw.get("rms_norm_eps", 1e-6), - vocab_size=llm_raw.get("vocab_size", 152064), - rope_theta=llm_raw.get("rope_theta", 1000000.0), - rope_theta_hw=llm_raw.get("rope_theta_hw", 10000.0), - max_position_embeddings=llm_raw.get("max_position_embeddings", 40960), - max_position_embeddings_hw=llm_raw.get("max_position_embeddings_hw", 10000), - attention_bias=llm_raw.get("attention_bias", True), - layer_types=llm_raw.get("layer_types", ["full_attention"] * llm_raw.get("num_hidden_layers", 32)), - tie_word_embeddings=llm_raw.get("tie_word_embeddings", False), - ) - - vis_cfg = SimpleNamespace( - num_channels=vis_raw.get("num_channels", 3), - patch_size=vis_raw.get("patch_size", 16), - hidden_size=vis_raw.get("hidden_size", 1024), - llm_hidden_size=vis_raw.get("llm_hidden_size", [llm_cfg.hidden_size]), - downsample_ratio=vis_raw.get("downsample_ratio", [0.5]), - rope_theta_vision=vis_raw.get("rope_theta_vision", 10000.0), - max_position_embeddings_vision=vis_raw.get("max_position_embeddings_vision", 10000), - output_hidden_states=vis_raw.get("output_hidden_states", False), - use_return_dict=vis_raw.get("use_return_dict", True), - ) - - top_cfg = SimpleNamespace( - downsample_ratio=raw.get("downsample_ratio", 0.5), - template=raw.get("template", "neo1_0"), - fm_head_layers=raw.get("fm_head_layers", 2), - fm_head_dim=raw.get("fm_head_dim", 4096), - fm_head_mlp_ratio=raw.get("fm_head_mlp_ratio", 1.0), - use_pixel_head=raw.get("use_pixel_head", False), - noise_scale=raw.get("noise_scale", 1.0), - noise_scale_mode=raw.get("noise_scale_mode", "none"), - noise_scale_base_image_seq_len=raw.get("noise_scale_base_image_seq_len", 256), - noise_scale_max_value=raw.get("noise_scale_max_value", 10.0), - add_noise_scale_embedding=raw.get("add_noise_scale_embedding", False), - time_schedule=raw.get("time_schedule", "standard"), - time_shift_type=raw.get("time_shift_type", "exponential"), - base_shift=raw.get("base_shift", 0.5), - max_shift=raw.get("max_shift", 1.15), - base_image_seq_len=raw.get("base_image_seq_len", 256), - max_image_seq_len=raw.get("max_image_seq_len", 4096), - concat_time_token_num=raw.get("concat_time_token_num", 0), - t_eps=raw.get("t_eps", 0.02), - ) - - return top_cfg, llm_cfg, vis_cfg - - # --------------------------------------------------------------------------- # Vision Embedding — 2D RoPE + Conv patch embed (no transformer blocks) # --------------------------------------------------------------------------- @@ -549,7 +481,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli model_path = od_config.model self.local_model_path = _resolve_model_path(model_path) - self.top_cfg, self.llm_cfg, self.vis_cfg = _load_sensenova_u1_config(model_path) + self.model_cfg = SenseNovaU1Config.from_pretrained(self.local_model_path) + self.llm_cfg = self.model_cfg.llm_config + self.vis_cfg = self.model_cfg.vision_config self.tokenizer = AutoTokenizer.from_pretrained(self.local_model_path) self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN) @@ -573,13 +507,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli # FM modules llm_hidden_size = self.llm_cfg.hidden_size patch_size = self.vis_cfg.patch_size - merge_size = int(1 / self.top_cfg.downsample_ratio) + merge_size = int(1 / self.model_cfg.downsample_ratio) output_dim = 3 * (patch_size * merge_size) ** 2 vision_model_mot_gen = NEOVisionModel(self.vis_cfg) timestep_embedder = TimestepEmbedder(llm_hidden_size) - if self.top_cfg.use_pixel_head: + if self.model_cfg.use_pixel_head: fm_head = ConvDecoder(llm_hidden_size) else: # The reference implementation uses a fixed 2-layer GELU MLP with a 4096-d @@ -601,13 +535,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli } ) - if self.top_cfg.add_noise_scale_embedding: + if self.model_cfg.add_noise_scale_embedding: self.fm_modules["noise_scale_embedder"] = TimestepEmbedder(llm_hidden_size) # Config shortcuts self.patch_size = patch_size self.merge_size = merge_size - self.downsample_ratio = self.top_cfg.downsample_ratio + self.downsample_ratio = self.model_cfg.downsample_ratio # Weight sources for diffusers_loader self.weights_sources = [ @@ -727,6 +661,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli t, z, image_token_num, + t_eps: float, image_size=None, cache_dit_skip=False, **_kw, @@ -748,7 +683,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli ) last_hidden_state = outputs.hidden_states - if self.top_cfg.use_pixel_head: + if self.model_cfg.use_pixel_head: merge_size = self.merge_size token_h = image_size[1] // (self.patch_size * merge_size) token_w = image_size[0] // (self.patch_size * merge_size) @@ -765,7 +700,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli else: x_pred = self.fm_modules["fm_head"](last_hidden_state[:, -image_token_num:].view(B, L, -1)).view(B, L, -1) - v_pred = (x_pred - z) / (1 - t).clamp_min(self.top_cfg.t_eps) + v_pred = (x_pred - z) / (1 - t).clamp_min(t_eps) return v_pred def _apply_time_schedule(self, t, image_seq_len, timestep_shift): @@ -1055,14 +990,14 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli token_w = p.image_size[0] // (self.patch_size * merge_size) grid_hw = torch.tensor([[grid_h, grid_w]] * p.batch_size, device=self.device) - noise_scale = self.top_cfg.noise_scale - if self.top_cfg.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"): - base = float(self.top_cfg.noise_scale_base_image_seq_len) + noise_scale = self.model_cfg.noise_scale + if self.model_cfg.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"): + base = float(self.model_cfg.noise_scale_base_image_seq_len) scale = math.sqrt((grid_h * grid_w) / (merge_size**2) / base) - noise_scale = scale * float(self.top_cfg.noise_scale) - if self.top_cfg.noise_scale_mode == "dynamic_sqrt": + noise_scale = scale * float(self.model_cfg.noise_scale) + if self.model_cfg.noise_scale_mode == "dynamic_sqrt": noise_scale = math.sqrt(noise_scale) - noise_scale = min(noise_scale, self.top_cfg.noise_scale_max_value) + noise_scale = min(noise_scale, self.model_cfg.noise_scale_max_value) generator = torch.Generator(self.device).manual_seed(p.seed) image_prediction = noise_scale * torch.randn( @@ -1101,6 +1036,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli z=z, image_token_num=ns.token_h * ns.token_w, image_size=p.image_size, + t_eps=p.t_eps, ) if cache_dit_skip: kwargs["cache_dit_skip"] = True @@ -1241,7 +1177,6 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli @torch.inference_mode() def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: p = self._parse_request(req) - self.top_cfg.t_eps = p.t_eps input_images = self._extract_input_images(p.first_prompt) modalities = p.first_prompt.get("modalities", []) if isinstance(p.first_prompt, dict) else [] @@ -1438,8 +1373,8 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli ns.token_h * ns.token_w, -1, ) - if self.top_cfg.add_noise_scale_embedding: - ns_tensor = torch.full_like(t_expanded, ns.noise_scale / self.top_cfg.noise_scale_max_value) + if self.model_cfg.add_noise_scale_embedding: + ns_tensor = torch.full_like(t_expanded, ns.noise_scale / self.model_cfg.noise_scale_max_value) ns_emb = self.fm_modules["noise_scale_embedder"](ns_tensor).view( p.batch_size, ns.token_h * ns.token_w, diff --git a/vllm_omni/diffusion/models/soulx_singer/__init__.py b/vllm_omni/diffusion/models/soulx_singer/__init__.py deleted file mode 100644 index 30235d2c1..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/__init__.py +++ /dev/null @@ -1,44 +0,0 @@ -"""SoulX-Singer SVS / SVC flow-matching singing models for vLLM-Omni.""" - -__all__ = [ - "FlowMatchingAudioPipeline", - "PipelineSoulXSingerSVS", - "PipelineSoulXSingerSVC", - "get_soulxsinger_post_process_func", - "get_soulxsinger_pre_process_func", - "get_soulxsinger_svc_pre_process_func", -] - - -def __getattr__(name: str): - if name in {"FlowMatchingAudioPipeline", "get_soulxsinger_post_process_func"}: - from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( - FlowMatchingAudioPipeline, - get_soulxsinger_post_process_func, - ) - - return { - "FlowMatchingAudioPipeline": FlowMatchingAudioPipeline, - "get_soulxsinger_post_process_func": get_soulxsinger_post_process_func, - }[name] - if name in {"PipelineSoulXSingerSVC", "get_soulxsinger_svc_pre_process_func"}: - from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svc import ( - PipelineSoulXSingerSVC, - get_soulxsinger_svc_pre_process_func, - ) - - return { - "PipelineSoulXSingerSVC": PipelineSoulXSingerSVC, - "get_soulxsinger_svc_pre_process_func": get_soulxsinger_svc_pre_process_func, - }[name] - if name in {"PipelineSoulXSingerSVS", "get_soulxsinger_pre_process_func"}: - from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svs import ( - PipelineSoulXSingerSVS, - get_soulxsinger_pre_process_func, - ) - - return { - "PipelineSoulXSingerSVS": PipelineSoulXSingerSVS, - "get_soulxsinger_pre_process_func": get_soulxsinger_pre_process_func, - }[name] - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py b/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py deleted file mode 100644 index 282ab12ab..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -from .convnext import ConvNeXtV2Block -from .decoder import CFMDecoder -from .flow_matching import FlowMatchingTransformer -from .llama import DiffLlama -from .mel_transform import MelSpectrogramEncoder -from .vocoder import Vocoder -from .whisper_encoder import WhisperEncoder - -__all__ = [ - "ConvNeXtV2Block", - "FlowMatchingTransformer", - "CFMDecoder", - "DiffLlama", - "MelSpectrogramEncoder", - "Vocoder", - "WhisperEncoder", -] diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py b/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py deleted file mode 100644 index 465898e8a..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py +++ /dev/null @@ -1,46 +0,0 @@ -import torch -import torch.nn as nn - - -class GRN(nn.Module): - def __init__(self, dim): - super().__init__() - self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) - self.beta = nn.Parameter(torch.zeros(1, 1, dim)) - - def forward(self, x): - Gx = torch.norm(x, p=2, dim=1, keepdim=True) - Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) - return self.gamma * (x * Nx) + self.beta + x - - -# ref: https://github.com/SWivid/F5-TTS/blob/main/src/f5_tts/model/modules.py#L247 -class ConvNeXtV2Block(nn.Module): - def __init__( - self, - dim: int, - intermediate_dim: int, - dilation: int = 1, - ): - super().__init__() - padding = (dilation * (7 - 1)) // 2 - self.dwconv = nn.Conv1d( - dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation - ) # depthwise conv - self.norm = nn.LayerNorm(dim, eps=1e-6) - self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers - self.act = nn.GELU() - self.grn = GRN(intermediate_dim) - self.pwconv2 = nn.Linear(intermediate_dim, dim) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = x.transpose(1, 2) # b n d -> b d n - x = self.dwconv(x) - x = x.transpose(1, 2) # b d n -> b n d - x = self.norm(x) - x = self.pwconv1(x) - x = self.act(x) - x = self.grn(x) - x = self.pwconv2(x) - return residual + x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py deleted file mode 100644 index 6b1fda616..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py +++ /dev/null @@ -1,24 +0,0 @@ -import torch -import torch.nn as nn - -from vllm_omni.diffusion.models.soulx_singer.modules.flow_matching import FlowMatchingTransformer - - -class CFMDecoder(nn.Module): - def __init__(self, config): - super().__init__() - self.model = FlowMatchingTransformer(cfg=config, **config) - - def forward(self, mel, x_mask, decoder_inp, is_prompt): - outputs = self.model(mel, x_mask, decoder_inp, is_prompt) - - noise, x, flow_pred, final_mask, prompt_len = outputs["output"] - return noise, x, flow_pred, final_mask, prompt_len - - def reverse_diffusion(self, pt_mel, pt_decoder_inp, gt_decoder_inp, n_timesteps=32, cfg=1): - diffusion_cond = torch.cat([pt_decoder_inp, gt_decoder_inp], dim=1) - diffusion_cond_emb = self.model.cond_emb(diffusion_cond) - diffusion_prompt = pt_mel - - generated = self.model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps, cfg=cfg) - return generated diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py b/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py deleted file mode 100644 index 19e87b65b..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py +++ /dev/null @@ -1,423 +0,0 @@ -# https://github.com/open-mmlab/Amphion/blob/main/models/svc/flow_matching_transformer/fmt_model.py - -# Copyright (c) 2023 Amphion. -# -# This source code is licensed under the MIT license found in the -# LICENSE file in the root directory of this source tree. - -import math - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from transformers import LlamaConfig - -from vllm_omni.diffusion.models.soulx_singer.modules.llama import DiffLlama - - -class FlowMatchingTransformer(nn.Module): - def __init__( - self, - mel_dim=100, - hidden_size=1024, - num_layers=12, - num_heads=16, - cfg_drop_prob=0.2, - use_embedding=True, - cond_codebook_size=1024, - cond_scale_factor=1, - sigma=1e-5, - time_scheduler="linear", - cfg=None, - ): - super().__init__() - self.cfg = cfg - - if cfg is not None: - mel_dim = getattr(cfg, "mel_dim", mel_dim) - hidden_size = getattr(cfg, "hidden_size", hidden_size) - num_layers = getattr(cfg, "num_layers", num_layers) - num_heads = getattr(cfg, "num_heads", num_heads) - cfg_drop_prob = getattr(cfg, "cfg_drop_prob", cfg_drop_prob) - cond_codebook_size = getattr(cfg, "cond_codebook_size", cond_codebook_size) - time_scheduler = getattr(cfg, "time_scheduler", time_scheduler) - sigma = getattr(cfg, "sigma", sigma) - cond_scale_factor = getattr(cfg, "cond_scale_factor", cond_scale_factor) - - self.mel_dim = mel_dim - self.hidden_size = hidden_size - self.num_layers = num_layers - self.num_heads = num_heads - self.cfg_drop_prob = cfg_drop_prob - self.cond_codebook_size = cond_codebook_size - self.time_scheduler = time_scheduler - self.sigma = sigma - self.cond_scale_factor = cond_scale_factor - - if use_embedding: - self.cond_emb = nn.Embedding(cond_codebook_size, self.hidden_size) - else: - self.cond_emb = nn.Linear(cond_codebook_size, self.hidden_size) - - if cond_scale_factor != 1: - self.do_resampling = True - assert np.log2(cond_scale_factor).is_integer() - - up_layers = [] - for _ in range(int(np.log2(cond_scale_factor))): - up_layers.extend( - [ - nn.ConvTranspose1d(hidden_size, hidden_size, kernel_size=4, stride=2, padding=1), - nn.GELU(), - ] - ) - self.resampling_layers = nn.Sequential(*up_layers) - else: - self.do_resampling = False - - ### REPA: Use the Wav2Vec2Bert features to align. ### - self.use_repa = "repa" in cfg - self.repa_layer_index = None - if self.use_repa: - self.repa_layer_index = cfg.repa.layer_index - - self.repa_mlp_layer = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, cfg.repa.output_dim), - ) - - ### CTC: Use the ASR loss ### - self.use_ctc = "ctc" in cfg - self.ctc_layer_index = None - if self.use_ctc: - self.ctc_layer_index = cfg.ctc.layer_index - - self.ctc_mlp_layer = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, cfg.ctc.output_dim), - ) - - self.reset_parameters() - - # DiffLlama config must be overridden manually to make sure the - # internal RoPE head dimension is correct - llama_config = LlamaConfig( - vocab_size=0, - bos_token_id=None, - eos_token_id=None, - hidden_size=hidden_size, - num_attention_heads=num_heads, - num_hidden_layers=num_layers, - rms_norm_eps=1e-6, - output_hidden_states=False, - output_attentions=False, - use_cache=False, - ) - self.diff_estimator = DiffLlama( - mel_dim=mel_dim, hidden_size=hidden_size, num_heads=num_heads, num_layers=num_layers, config=llama_config - ) - - self.sigma = sigma - - @torch.no_grad() - def forward_diffusion(self, x, t, is_prompt=None): - """ - x: (B, T, mel_dim) - t: (B,) - """ - new_t = t - t = t.unsqueeze(-1).unsqueeze(-1) - z = torch.randn(x.shape, dtype=x.dtype, device=x.device, requires_grad=False) # (B, T, mel_dim) - - # get prompt len - if torch.rand(1) <= self.cfg_drop_prob: - prompt_len = torch.zeros(x.shape[0]).to(x) - is_prompt = torch.zeros_like(x[:, :, 0]) - else: - if is_prompt is None: - prompt_len = torch.randint(min(x.shape[1] // 4, 5), int(x.shape[1] * 0.4), (x.shape[0],)).to( - x.device - ) # (B,) - - # get is_prompt - is_prompt = torch.zeros_like(x[:, :, 0]) # (B, T) - col_indices = torch.arange(is_prompt.shape[1]).repeat(is_prompt.shape[0], 1).to(prompt_len) # (B, T) - is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1 # (B, T) 1 if prompt - else: - prompt_len = is_prompt.sum(dim=1) # (B,) - - mask = torch.ones_like(x[:, :, 0]) # mask if 1, not mask if 0 - mask[is_prompt.bool()] = 0 - mask = mask[:, :, None] - - # flow matching: xt = (1 - (1 - sigma) * t) * x0 + t * x; where x0 ~ N(0, 1), x is a sample - # flow gt: x - (1 - sigma) * x0 = x - (1 - sigma) * noise - xt = ((1 - (1 - self.sigma) * t) * z + t * x) * mask + x * (1 - mask) - - return xt, z, new_t, prompt_len, mask - - def loss_t(self, x, x_mask, t, cond=None, is_prompt=None): - xt, z, new_t, prompt_len, mask = self.forward_diffusion(x, t, is_prompt) - - noise = z - - # drop all condition for cfg, so if prompt_len is 0, we also drop cond - if cond is not None: - cond = cond * torch.where( - prompt_len > 0, - torch.ones_like(prompt_len), - torch.zeros_like(prompt_len), - ).to(cond.device).unsqueeze(-1).unsqueeze(-1) - - dit_output = self.diff_estimator(xt, new_t, cond, x_mask, return_dict=True) - flow_pred = dit_output["output"] # (B, T, mel_dim) - - # final mask used for loss calculation - final_mask = mask * x_mask[..., None] # (B, T, 1) - - results = {"output": (noise, x, flow_pred, final_mask, prompt_len)} - - if self.use_repa: - repa_hidden_states = dit_output["hidden_states"][self.repa_layer_index] # (B, T, hidden_size) - - repa_pred = self.repa_mlp_layer(repa_hidden_states) # (B, T, repa_dim) - results["repa"] = repa_pred - - if self.use_ctc: - ctc_hidden_states = dit_output["hidden_states"][self.ctc_layer_index] # (B, T, hidden_size) - ctc_pred = self.ctc_mlp_layer(ctc_hidden_states) # (B, T, ctc_dim) - results["ctc"] = ctc_pred - - return results - - def compute_loss(self, x, x_mask, cond=None, is_prompt=None): - # x0: (B, T, num_quantizer) - # x_mask: (B, T) mask is 0 for padding - t = torch.rand(x.shape[0], device=x.device, requires_grad=False) - t = torch.clamp(t, 1e-5, 1.0) - # generation is harder at the start than later; use a cosine - # scheduler for timestep t when time_scheduler == "cos". - if self.time_scheduler == "cos": - t = 1 - torch.cos(t * math.pi * 0.5) - else: - pass - return self.loss_t(x, x_mask, t, cond, is_prompt) - - def reset_parameters(self): - def _reset_parameters(m): - if isinstance(m, nn.MultiheadAttention): - if m._qkv_same_embed_dim: - nn.init.normal_(m.in_proj_weight, std=0.02) - else: - nn.init.normal_(m.q_proj_weight, std=0.02) - nn.init.normal_(m.k_proj_weight, std=0.02) - nn.init.normal_(m.v_proj_weight, std=0.02) - - if m.in_proj_bias is not None: - nn.init.constant_(m.in_proj_bias, 0.0) - nn.init.constant_(m.out_proj.bias, 0.0) - if m.bias_k is not None: - nn.init.xavier_normal_(m.bias_k) - if m.bias_v is not None: - nn.init.xavier_normal_(m.bias_v) - - elif ( - isinstance(m, nn.Conv1d) - or isinstance(m, nn.ConvTranspose1d) - or isinstance(m, nn.Conv2d) - or isinstance(m, nn.ConvTranspose2d) - ): - m.weight.data.normal_(0.0, 0.02) - - elif isinstance(m, nn.Linear): - m.weight.data.normal_(mean=0.0, std=0.02) - if m.bias is not None: - m.bias.data.zero_() - - elif isinstance(m, nn.Embedding): - m.weight.data.normal_(mean=0.0, std=0.02) - if m.padding_idx is not None: - m.weight.data[m.padding_idx].zero_() - - self.apply(_reset_parameters) - - @torch.no_grad() - def reverse_diffusion( - self, - cond, - prompt, - x_mask=None, - prompt_mask=None, - n_timesteps=10, - cfg=1.0, - rescale_cfg=0.75, - ): - h = 1.0 / n_timesteps - prompt_len = prompt.shape[1] - target_len = cond.shape[1] - prompt_len - trunk_dtype = cond.dtype - - if x_mask is None: - x_mask = torch.ones(cond.shape[0], target_len, dtype=trunk_dtype, device=cond.device) - if prompt_mask is None: - prompt_mask = torch.ones(cond.shape[0], prompt_len, dtype=trunk_dtype, device=cond.device) - xt_mask = torch.cat([prompt_mask, x_mask], dim=1) - z = torch.randn( - (cond.shape[0], target_len, self.mel_dim), - dtype=trunk_dtype, - device=cond.device, - requires_grad=False, - ) - xt = z - h_tensor = torch.tensor(h, dtype=trunk_dtype, device=cond.device) - - # t from 0 to 1: x0 = z ~ N(0, 1) - for i in range(n_timesteps): - xt_input = torch.cat([prompt, xt], dim=1) - t = torch.full( - (z.shape[0],), - fill_value=(i + 0.5) * h, - dtype=trunk_dtype, - device=z.device, - ) - flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) - flow_pred = flow_pred[:, prompt_len:, :] - - # cfg - if cfg > 0: - uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask) - pos_flow_pred_std = flow_pred.std() - flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) - rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() - flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg - - xt = xt + flow_pred * h_tensor - - return xt - - @torch.no_grad() - def reverse_diffusion_v2( - self, - cond, - prompt, - x_mask=None, - prompt_mask=None, - n_timesteps=10, - cfg=1.0, - rescale_cfg=0.75, - ): - h = 1.0 / n_timesteps - prompt_len = prompt.shape[1] - target_len = cond.shape[1] - prompt_len * 2 - - if x_mask is None: - x_mask = torch.ones(cond.shape[0], target_len).to(cond.device) # (B, T) - if prompt_mask is None: - prompt_mask = torch.ones(cond.shape[0], prompt_len).to(cond.device) # (B, prompt_len) - xt_mask = torch.cat([prompt_mask, x_mask, prompt_mask], dim=1) - z = torch.randn( - (cond.shape[0], target_len, self.mel_dim), - dtype=cond.dtype, - device=cond.device, - requires_grad=False, - ) - xt = z - - # t from 0 to 1: x0 = z ~ N(0, 1) - for i in range(n_timesteps): - xt_input = torch.cat([prompt, xt, prompt], dim=1) - t = (0 + (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) - flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) - flow_pred = flow_pred[:, prompt_len:-prompt_len, :] - - # cfg - if cfg > 0: - uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask) - pos_flow_pred_std = flow_pred.std() - flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) - rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() - flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg - - dxt = flow_pred * h - xt = xt + dxt - - return xt - - def forward( - self, - x: torch.Tensor, - x_mask: torch.Tensor, - cond_code: torch.Tensor, - is_prompt: torch.Tensor | None = None, - ): - """ - Args: - x: (B, T, mel_dim) - x_mask: (B, T) - cond_code: (B, T), Note that cond_code might be not at 50Hz! - """ - T = x.shape[1] - - cond = self.cond_emb(cond_code) # (B, T, hidden_size) - if self.do_resampling: - # Align to the frame rate of Mels - cond = self.resampling_layers(cond.transpose(1, 2)).transpose(1, 2) - - # print("cond_code: {}, after resampling: {}".format(cond_code.shape, cond.shape)) - - if cond.shape[1] >= T: # Check time dimension - cond = cond[:, :T, :] - else: - padding_frames = T - cond.shape[1] - last_frame = cond[:, -1:, :] - padding = last_frame.repeat(1, padding_frames, 1) - cond = torch.cat([cond, padding], dim=1) - - return self.compute_loss(x, x_mask, cond, is_prompt) - - -if __name__ == "__main__": - model_cfg = { - "mel_dim": 128, - "hidden_size": 256, - "num_layers": 8, - "num_heads": 8, - "cfg_drop_prob": 0.2, - "use_embedding": False, - "cond_codebook_size": 256, - "cond_scale_factor": 1, - "sigma": 1e-5, - "time_scheduler": "cos", - } - - device = "cuda" - x = torch.randn(2, 100, 128).to(device) - x_mask = torch.ones(2, 100).to(device) - # cond_code = torch.randint(0, 16384, (2, 25)).to(device) - cond_code = torch.randn(2, 100, 256).to(device) - - model = FlowMatchingTransformer(cfg=model_cfg, **model_cfg).to(device) - outputs = model(x, x_mask, cond_code) - print(outputs) - - noise, x, flow_pred, final_mask, prompt_len = outputs["output"] - final_mask = final_mask.squeeze(-1) - - flow_gt = x - (1 - 1e-5) * noise - - # [B, n_frames, D] - diff_loss = F.l1_loss(flow_pred, flow_gt, reduction="none").float() * final_mask.unsqueeze(-1) - diff_loss = torch.mean(diff_loss, dim=2).sum() / final_mask.sum() - - print("diff_loss:", diff_loss.item()) - - diffusion_cond = torch.randn(2, 150, 256).to(device) - diffusion_cond_emb = model.cond_emb(diffusion_cond) - diffusion_prompt = torch.randn(2, 50, 128).to(device) - n_timesteps = 32 - - generated = model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps) - print("generated:", generated.shape) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/llama.py b/vllm_omni/diffusion/models/soulx_singer/modules/llama.py deleted file mode 100644 index b0b6afce3..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/llama.py +++ /dev/null @@ -1,393 +0,0 @@ -import math - -import torch -import torch.nn as nn -from cache_dit import ForwardPattern -from transformers import LlamaConfig, LlamaModel -from transformers.modeling_outputs import BaseModelOutputWithPast -from transformers.models.llama.modeling_llama import LlamaDecoderLayer - -from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput - - -class SinusoidalPosEmb(nn.Module): - def __init__(self, dim): - super().__init__() - self.dim = dim - - def forward(self, x): - device = x.device - dtype = x.dtype - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device, dtype=dtype) * -emb) - emb = x[:, None] * emb[None, :] - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class LlamaAdaptiveRMSNorm(nn.Module): - def __init__(self, hidden_size=1024, eps=1e-6, dim_cond=1024): - super().__init__() - self.to_weight = nn.Linear(dim_cond, hidden_size) - nn.init.zeros_(self.to_weight.weight) - nn.init.ones_(self.to_weight.bias) - self.variance_epsilon = eps - self._is_hf_initialized = True - - def forward(self, hidden_states, cond_embedding): - input_dtype = hidden_states.dtype - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - - weight = self.to_weight(cond_embedding) - if len(weight.shape) == 2: - weight = weight.unsqueeze(1) - - return (weight * hidden_states).to(input_dtype) - - -class LlamaNARDecoderLayer(LlamaDecoderLayer): - def __init__(self, config: LlamaConfig, layer_idx: int): - super().__init__(config, layer_idx) - - self.input_layernorm = LlamaAdaptiveRMSNorm( - config.hidden_size, - eps=config.rms_norm_eps, - dim_cond=config.hidden_size, - ) - self.post_attention_layernorm = LlamaAdaptiveRMSNorm( - config.hidden_size, - eps=config.rms_norm_eps, - dim_cond=config.hidden_size, - ) - - def forward( - self, - hidden_states: torch.Tensor, - cond_embedding: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - output_attentions: bool | None = False, - cache_position: torch.LongTensor | None = None, - **kwargs, - ) -> tuple[torch.FloatTensor, ...]: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states, cond_embedding=cond_embedding) - - # RoPE is computed by the parent model and passed in as (cos, sin). - hidden_states, self_attn_weights = self.self_attn( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - attention_mask=attention_mask, - cache_position=cache_position, - output_attentions=output_attentions, - **kwargs, - ) - - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm( - hidden_states, - cond_embedding=cond_embedding, - ) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights,) - - return outputs - - -class DiffLlama(LlamaModel): - # Regional torch.compile in DiffusionModelRunner targets these blocks. - _repeated_blocks = ["LlamaNARDecoderLayer"] - _layerwise_offload_blocks_attrs = ["layers"] - - # cache-dit: Llama-style backbone with static condition embeddings injected - # before the decoder loop (like Stable Audio Open). Pattern_3 handles the - # cross-attention-style conditioning on the ``layers`` ModuleList. - _cache_dit_adapter_config = CacheDiTAdapterConfig( - block_forward_patterns={ - "layers": ForwardPattern.Pattern_3, - }, - ) - - # Sequence parallelism: shard mel sequence before the first decoder layer and - # gather after the output projection. Tensor parallelism still requires - # migrating attention to vLLM parallel linear layers. - _sp_plan = { - "layers.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), - }, - "mel_out_mlp": SequenceParallelOutput(gather_dim=1, expected_dims=3), - } - - @staticmethod - def _is_transformer_block(name: str, module: nn.Module) -> bool: - del module - return name.startswith("layers.") and name.split(".")[-1].isdigit() - - _hsdp_shard_conditions = [_is_transformer_block] - - def __init__( - self, - mel_dim=100, - hidden_size=1024, - num_heads=16, - num_layers=16, - dropout=0.1, - ffn_dropout=0.1, - attention_dropout=0.0, - config: LlamaConfig | None = None, - ): - if config is None: - config = LlamaConfig( - vocab_size=32000, - hidden_size=hidden_size, - intermediate_size=hidden_size * 4, - num_hidden_layers=num_layers, - num_attention_heads=num_heads, - num_key_value_heads=num_heads, - max_position_embeddings=4096, - rms_norm_eps=1e-6, - attention_dropout=attention_dropout, - use_cache=False, - attn_implementation="eager", - ) - - config.use_cache = False - - super().__init__(config) - - # Keep custom decoder layers on the same attention backend. - layer_config = LlamaConfig( - vocab_size=config.vocab_size, - hidden_size=hidden_size, - intermediate_size=hidden_size * 4, - num_hidden_layers=num_layers, - num_attention_heads=num_heads, - num_key_value_heads=num_heads, - max_position_embeddings=4096, - rms_norm_eps=config.rms_norm_eps, - attention_dropout=attention_dropout, - use_cache=False, - attn_implementation="eager", - bos_token_id=None, - eos_token_id=None, - ) - - self.layers = nn.ModuleList( - [ - LlamaNARDecoderLayer( - layer_config, - layer_idx=i, - ) - for i in range(num_layers) - ] - ) - - self.norm = LlamaAdaptiveRMSNorm(hidden_size, dim_cond=hidden_size) - - self.diff_step_embedding = SinusoidalPosEmb(hidden_size) - self.diff_step_mlp = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, hidden_size), - ) - - self.cond_mlp = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, hidden_size), - ) - - self.mel_mlp = nn.Sequential( - nn.Linear(mel_dim, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, hidden_size), - ) - - self.mel_out_mlp = nn.Sequential( - nn.Linear(hidden_size, hidden_size * 4), - nn.SiLU(), - nn.Linear(hidden_size * 4, mel_dim), - ) - - self.embed_tokens = None - - self.post_init() - - def _prepare_decoder_attention_mask( - self, - attention_mask, - input_shape, - inputs_embeds, - past_key_values_length=0, - ): - """ - Build a 4D additive padding mask. - - This model is non-autoregressive, so no causal mask is added here. - """ - - combined_attention_mask = None - - def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None): - bsz, src_len = mask.size() - tgt_len = tgt_len if tgt_len is not None else src_len - - expanded_mask = ( - mask[:, None, None, :] - .expand( - bsz, - 1, - tgt_len, - src_len, - ) - .to(dtype) - ) - - inverted_mask = 1.0 - expanded_mask - - return inverted_mask.masked_fill( - inverted_mask.to(torch.bool), - torch.finfo(dtype).min, - ) - - if attention_mask is not None: - expanded_attn_mask = _expand_mask( - attention_mask, - inputs_embeds.dtype, - tgt_len=input_shape[-1], - ).to(inputs_embeds.device) - - combined_attention_mask = ( - expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask - ) - - return combined_attention_mask - - def forward( - self, - x, - diffusion_step, - cond, - x_mask, - input_ids: torch.LongTensor = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: list[torch.FloatTensor] | None = None, - inputs_embeds: torch.FloatTensor | None = None, - use_cache: bool | None = None, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - return_dict: bool | None = False, - ) -> BaseModelOutputWithPast | torch.Tensor | dict: - batch_size, seq_length, _ = x.shape - - cond_embedding = self.cond_mlp(cond) - x = self.mel_mlp(x) - - diffusion_step = self.diff_step_embedding(diffusion_step).to(device=x.device, dtype=x.dtype) - diffusion_step = self.diff_step_mlp(diffusion_step) - - x = x + cond_embedding - - inputs_embeds = x - attention_mask = x_mask - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - - # KV cache is disabled because this diffusion decoder processes full sequences. - - if past_key_values is not None: - raise NotImplementedError("KV cache is disabled for this non-autoregressive DiffLlama implementation.") - - if position_ids is None: - device = input_ids.device if input_ids is not None else inputs_embeds.device - position_ids = torch.arange( - 0, - seq_length, - dtype=torch.long, - device=device, - ) - position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) - else: - position_ids = position_ids.view(batch_size, seq_length).long() - - if attention_mask is None: - attention_mask = torch.ones( - (batch_size, seq_length), - dtype=torch.bool, - device=inputs_embeds.device, - ) - - attention_mask = self._prepare_decoder_attention_mask( - attention_mask, - (batch_size, seq_length), - inputs_embeds, - past_key_values_length=0, - ) - - hidden_states = inputs_embeds - - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_layer_hidden_states = [] - - cache_position = torch.arange( - 0, - seq_length, - device=hidden_states.device, - dtype=torch.long, - ) - - # LlamaAttention expects precomputed rotary embeddings. - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - for idx, decoder_layer in enumerate(self.layers): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - raise NotImplementedError - - layer_outputs = decoder_layer( - hidden_states, - cond_embedding=diffusion_step, - attention_mask=attention_mask, - position_embeddings=position_embeddings, - output_attentions=output_attentions, - cache_position=cache_position, - ) - - hidden_states = layer_outputs[0] - all_layer_hidden_states.append(hidden_states.clone()) - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step) - - if output_hidden_states: - all_hidden_states += (hidden_states,) - - hidden_states = self.mel_out_mlp(hidden_states) - - if return_dict: - return { - "output": hidden_states, - "hidden_states": all_layer_hidden_states, - "attentions": all_self_attns, - } - - return hidden_states diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py b/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py deleted file mode 100644 index cc767dcf3..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py +++ /dev/null @@ -1,151 +0,0 @@ -import math -from typing import Any - -import numpy as np -import torch -import torch.nn as nn - -from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn - - -def dynamic_range_compression(x, c=1, clip_val=1e-5): - return np.log(np.clip(x, a_min=clip_val, a_max=None) * c) - - -def dynamic_range_decompression(x, c=1): - return np.exp(x) / c - - -def dynamic_range_compression_torch(x, c=1, clip_val=1e-5): - return torch.log(torch.clamp(x, min=clip_val) * c) - - -def dynamic_range_decompression_torch(x, c=1): - return torch.exp(x) / c - - -def spectral_normalize_torch(magnitudes): - output = dynamic_range_compression_torch(magnitudes) - return output - - -def spectral_de_normalize_torch(magnitudes): - output = dynamic_range_decompression_torch(magnitudes) - return output - - -class MelSpectrogram(nn.Module): - def __init__( - self, - n_fft, - num_mels, - sampling_rate, - hop_size, - win_size, - fmin, - fmax, - center=False, - ): - super().__init__() - self.n_fft = n_fft - self.hop_size = hop_size - self.win_size = win_size - self.sampling_rate = sampling_rate - self.num_mels = num_mels - self.fmin = fmin - self.fmax = fmax - self.center = center - - mel_basis = {} - hann_window = {} - - mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) - mel_basis = mel.to(torch.float32) - hann_window = torch.hann_window(win_size) - - self.register_buffer("mel_basis", mel_basis) - self.register_buffer("hann_window", hann_window) - - def forward(self, y): - y = torch.nn.functional.pad( - y.unsqueeze(1), - ( - int((self.n_fft - self.hop_size) / 2), - int((self.n_fft - self.hop_size) / 2), - ), - mode="reflect", - ) - y = y.squeeze(1) - spec = torch.stft( - y, - self.n_fft, - hop_length=self.hop_size, - win_length=self.win_size, - window=self.hann_window, - center=self.center, - pad_mode="reflect", - normalized=False, - onesided=True, - return_complex=True, - ) - spec = torch.view_as_real(spec) - - spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) - - spec = torch.matmul(self.mel_basis, spec) - spec = spectral_normalize_torch(spec) - - return spec - - -def load_mel_spectrogram(): - return load_mel_spectrogram_from_cfg(None) - - -def _get_from_mapping(cfg: Any, key: str, default: Any = None) -> Any: - """Safely read a field from a dict/OmegaConf-like object.""" - if cfg is None: - return default - if isinstance(cfg, dict): - return cfg.get(key, default) - return getattr(cfg, key, default) - - -def load_mel_spectrogram_from_cfg(audio_cfg: Any | None = None) -> MelSpectrogram: - """Build MelSpectrogram from `audio_config`-like config. - - Expected keys (either in dict or Hydra/OmegaConf object): - - hop_size, sample_rate (or sampling_rate), n_fft, num_mels, win_size, fmin, fmax - """ - # Defaults keep current behavior. - mel_cfg: dict[str, Any] = { - "hop_size": _get_from_mapping(audio_cfg, "hop_size", 480), - "sampling_rate": _get_from_mapping( - audio_cfg, - "sampling_rate", - _get_from_mapping(audio_cfg, "sample_rate", 24000), - ), - "n_fft": _get_from_mapping(audio_cfg, "n_fft", 1920), - "num_mels": _get_from_mapping(audio_cfg, "num_mels", 128), - "win_size": _get_from_mapping(audio_cfg, "win_size", 1920), - "fmin": _get_from_mapping(audio_cfg, "fmin", 0), - "fmax": _get_from_mapping(audio_cfg, "fmax", 12000), - } - - mel_model = MelSpectrogram(**mel_cfg) - mel_model.eval() - return mel_model - - -class MelSpectrogramEncoder(nn.Module): - def __init__(self, audio_config: dict | None = None): - super().__init__() - self.model = load_mel_spectrogram_from_cfg(audio_config) - audio_config = audio_config or {} - self.mel_mean = audio_config.get("mel_mean", -4.92) - self.mel_var = audio_config.get("mel_var", 8.14) - - def forward(self, x): - x = self.model(x).transpose(1, 2) - x = (x - self.mel_mean) / math.sqrt(self.mel_var) - return x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py b/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py deleted file mode 100644 index 70933dd92..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py +++ /dev/null @@ -1,958 +0,0 @@ -import math - -import numpy as np -import packaging.version as version -import torch -import torch.nn as nn - -# ---------------------- note transcription common layers ---------------------- - - -class LayerNorm(nn.LayerNorm): - def __init__(self, nout: int, dim: int = -1, eps: float = 1e-5): - """Construct an LayerNorm object.""" - super().__init__(nout, eps=eps) - self.dim = dim - - def forward(self, x): - if self.dim == -1: - return super().forward(x) - return super().forward(x.transpose(1, -1)).transpose(1, -1) - - -class Reshape(nn.Module): - def __init__(self, *args): - super().__init__() - self.shape = args - - def forward(self, x): - return x.view(self.shape) - - -class Permute(nn.Module): - def __init__(self, *args): - super().__init__() - self.args = args - - def forward(self, x): - return x.permute(self.args) - - -class LambdaLayer(nn.Module): - def __init__(self, lambd): - super().__init__() - self.lambd = lambd - - def forward(self, x): - return self.lambd(x) - - -class ResidualBlock(nn.Module): - """ - Implements a residual block pattern: - [Norm] → [Conv1d] → [Scale] → [Activation] → [Conv1d] - repeated n times, each as a residual unit. - """ - - def __init__( - self, - channels: int, - kernel_size: int, - dilation: int, - n: int = 2, - c_multiple: int = 2, - ln_eps: float = 1e-12, - ): - super().__init__() - - self.blocks = nn.ModuleList( - [ - nn.Sequential( - LayerNorm(channels, dim=1, eps=ln_eps), - nn.Conv1d( - channels, - c_multiple * channels, - kernel_size, - dilation=dilation, - padding=(dilation * (kernel_size - 1)) // 2, - ), - LambdaLayer(lambda x: x * kernel_size**-0.5), # scaling for stability - nn.LeakyReLU(), - nn.Conv1d(c_multiple * channels, channels, kernel_size=1, dilation=dilation), - ) - for _ in range(n) - ] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] - for block in self.blocks: - x_res = block(x) - x = (x + x_res) * nonpadding - return x - - -class ConvBlocks(nn.Module): - """ - Decodes the expanded phoneme encoding into spectrograms. - """ - - def __init__( - self, - hidden_size, - out_dims, - dilations, - kernel_size, - layers_in_block=2, - c_multiple=2, - ln_eps=1e-5, - is_btc=True, - post_net_kernel=3, - ): - super().__init__() - self.is_btc = is_btc - - res_blocks = [ - ResidualBlock( - hidden_size, - kernel_size=kernel_size, - dilation=1, - n=layers_in_block, - c_multiple=c_multiple, - ln_eps=ln_eps, - ) - for _ in range(1) - ] # default 1 - self.res_blocks = nn.Sequential(*res_blocks) - - # Output normalization and conv - self.last_norm = LayerNorm(hidden_size, dim=1, eps=ln_eps) - self.post_net1 = nn.Conv1d(hidden_size, out_dims, kernel_size=post_net_kernel, padding=post_net_kernel // 2) - - def forward(self, x: torch.Tensor, nonpadding: torch.Tensor | None = None) -> torch.Tensor: - if self.is_btc: - x = x.transpose(1, 2) - if nonpadding is None: - nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] - elif self.is_btc: - nonpadding = nonpadding.transpose(1, 2) - x = self.res_blocks(x) * nonpadding - x = self.last_norm(x) * nonpadding - x = self.post_net1(x) * nonpadding - if self.is_btc: - x = x.transpose(1, 2) - return x - - -def Linear(in_features: int, out_features: int, bias: bool = True, init_type: str = "xavier"): - m = nn.Linear(in_features, out_features, bias) - if init_type == "xavier": - nn.init.xavier_uniform_(m.weight) - elif init_type == "kaiming": - nn.init.kaiming_normal_(m.weight, mode="fan_in") - if bias: - nn.init.constant_(m.bias, 0.0) - return m - - -def Embedding(num_embeddings, embedding_dim, padding_idx=None, init_type="normal"): - m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx) - if init_type == "normal": - nn.init.normal_(m.weight, mean=0, std=embedding_dim**-0.5) - elif init_type == "kaiming": - nn.init.kaiming_normal_(m.weight, mode="fan_in") - if padding_idx is not None: - nn.init.constant_(m.weight[padding_idx], 0) - return m - - -# ---------------------- note transcription espnet attention ---------------------- - - -class MultiHeadedAttention(nn.Module): - """Multi-Head Attention layer. - Args: - n_head (int): The number of heads. - n_feat (int): The number of features. - dropout_rate (float): Dropout rate. - """ - - def __init__(self, n_head, n_feat, dropout_rate, flash=False): - """Construct an MultiHeadedAttention object.""" - super().__init__() - assert n_feat % n_head == 0 - # We assume d_v always equals d_k - self.d_k = n_feat // n_head - self.h = n_head - self.linear_q = nn.Linear(n_feat, n_feat) - self.linear_k = nn.Linear(n_feat, n_feat) - self.linear_v = nn.Linear(n_feat, n_feat) - self.linear_out = nn.Linear(n_feat, n_feat) - self.attn = None - self.dropout = nn.Dropout(p=dropout_rate) - self.dropout_rate = dropout_rate - self.flash = flash - - def forward_qkv(self, query, key, value): - """Transform query, key and value. - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - Returns: - torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). - torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). - torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). - """ - n_batch = query.size(0) - q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) - k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) - v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) - q = q.transpose(1, 2) # (batch, head, time1, d_k) - k = k.transpose(1, 2) # (batch, head, time2, d_k) - v = v.transpose(1, 2) # (batch, head, time2, d_k) - - return q, k, v - - def forward_attention(self, value, scores, mask): - """Compute attention context vector. - Args: - value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). - scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). - mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). - Returns: - torch.Tensor: Transformed value (#batch, time1, d_model) - weighted by the attention score (#batch, time1, time2). - """ - n_batch = value.size(0) - if mask is not None: - mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) - min_value = float(np.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) - scores = scores.masked_fill(mask, min_value) - self.attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) # (batch, head, time1, time2) - else: - self.attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) - - p_attn = self.dropout(self.attn) - x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) - x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) - - return self.linear_out(x) # (batch, time1, d_model) - - def forward(self, query, key, value, mask): - """Compute scaled dot product attention. - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - mask (torch.Tensor): Mask tensor (#batch, 1, time2) or - (#batch, time1, time2). - Returns: - torch.Tensor: Output tensor (#batch, time1, d_model). - """ - q, k, v = self.forward_qkv(query, key, value) - if version.parse(torch.__version__) >= version.parse("2.0") and self.flash: - n_batch = value.size(0) - x = torch.nn.functional.scaled_dot_product_attention( - q, k, v, attn_mask=mask.unsqueeze(1) if mask is not None else None, dropout_p=self.dropout_rate - ) - x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) - return self.linear_out(x) - else: - scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) - return self.forward_attention(v, scores, mask) - - -class RelPositionMultiHeadedAttention(MultiHeadedAttention): - """Multi-Head Attention layer with relative position encoding. - Paper: https://arxiv.org/abs/1901.02860 - Args: - n_head (int): The number of heads. - n_feat (int): The number of features. - dropout_rate (float): Dropout rate. - """ - - def __init__(self, n_head, n_feat, dropout_rate): - """Construct an RelPositionMultiHeadedAttention object.""" - super().__init__(n_head, n_feat, dropout_rate) - # linear transformation for positional encoding - self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) - # these two learnable bias are used in matrix c and matrix d - # as described in https://arxiv.org/abs/1901.02860 Section 3.3 - self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) - self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) - torch.nn.init.xavier_uniform_(self.pos_bias_u) - torch.nn.init.xavier_uniform_(self.pos_bias_v) - - def rel_shift(self, x, zero_triu=False): - """Compute relative positinal encoding. - Args: - x (torch.Tensor): Input tensor (batch, time, size). - zero_triu (bool): If true, return the lower triangular part of the matrix. - Returns: - torch.Tensor: Output tensor. - """ - zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) - x_padded = torch.cat([zero_pad, x], dim=-1) - - x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) - x = x_padded[:, :, 1:].view_as(x) - - if zero_triu: - ones = torch.ones((x.size(2), x.size(3))) - x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :] - - return x - - def forward(self, query, key, value, pos_emb, mask): - """Compute 'Scaled Dot Product Attention' with rel. positional encoding. - Args: - query (torch.Tensor): Query tensor (#batch, time1, size). - key (torch.Tensor): Key tensor (#batch, time2, size). - value (torch.Tensor): Value tensor (#batch, time2, size). - pos_emb (torch.Tensor): Positional embedding tensor (#batch, time2, size). - mask (torch.Tensor): Mask tensor (#batch, 1, time2) or - (#batch, time1, time2). - Returns: - torch.Tensor: Output tensor (#batch, time1, d_model). - """ - q, k, v = self.forward_qkv(query, key, value) - q = q.transpose(1, 2) # (batch, time1, head, d_k) - - n_batch_pos = pos_emb.size(0) - p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) - p = p.transpose(1, 2) # (batch, head, time1, d_k) - - # (batch, head, time1, d_k) - q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) - # (batch, head, time1, d_k) - q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) - - # compute attention score - # first compute matrix a and matrix c - # as described in https://arxiv.org/abs/1901.02860 Section 3.3 - # (batch, head, time1, time2) - matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) - - # compute matrix b and matrix d - # (batch, head, time1, time2) - matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) - matrix_bd = self.rel_shift(matrix_bd) - - scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) - - return self.forward_attention(v, scores, mask) - - -# ---------------------- note transcription positional encoding ---------------------- - - -class PositionalEncoding(torch.nn.Module): - """Positional encoding. - Args: - d_model (int): Embedding dimension. - max_len (int): Maximum input length. - reverse (bool): Whether to reverse the input position. - """ - - def __init__(self, d_model, max_len=5000, reverse=False): - """Construct an PositionalEncoding object.""" - super().__init__() - self.d_model = d_model - self.reverse = reverse - self.xscale = math.sqrt(self.d_model) - self.pe = None - - self.extend_pe(torch.tensor(0.0).expand(1, max_len)) - - def extend_pe(self, x): - """Reset the positional encodings.""" - if self.pe is not None: - if self.pe.size(1) >= x.size(1): - if self.pe.dtype != x.dtype or self.pe.device != x.device: - self.pe = self.pe.to(dtype=x.dtype, device=x.device) - return - pe = torch.zeros(x.size(1), self.d_model) - if self.reverse: - position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1) - else: - position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) - div_term = torch.exp( - torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model) - ) - pe[:, 0::2] = torch.sin(position * div_term) - pe[:, 1::2] = torch.cos(position * div_term) - pe = pe.unsqueeze(0) - self.pe = pe.to(device=x.device, dtype=x.dtype) - - def forward(self, x: torch.Tensor): - """Add positional encoding. - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - """ - self.extend_pe(x) - x = x * self.xscale + self.pe[:, : x.size(1)] - return x - - -class RelPositionalEncoding(PositionalEncoding): - """Relative positional encoding module. - See : Appendix B in https://arxiv.org/abs/1901.02860 - Args: - d_model (int): Embedding dimension. - max_len (int): Maximum input length. - """ - - def __init__(self, d_model, max_len=5000): - super().__init__(d_model, max_len, reverse=True) - - def forward(self, x): - """Compute positional encoding. - Args: - x (torch.Tensor): Input tensor (batch, time, `*`). - Returns: - torch.Tensor: Encoded tensor (batch, time, `*`). - torch.Tensor: Positional embedding tensor (1, time, `*`). - """ - self.extend_pe(x) - x = x * self.xscale - pos_emb = self.pe[:, : x.size(1)] - return x, pos_emb - - -# ---------------------- note transcription conformer layers ---------------------- - - -class ConvolutionModule(nn.Module): - def __init__(self, channels: int, kernel_size: int, activation: nn.Module = nn.ReLU(), bias: bool = True): - super().__init__() - if (kernel_size - 1) % 2 != 0: - raise ValueError("kernel_size should be a odd number for 'SAME' padding") - - self.pointwise_conv1 = nn.Conv1d(channels, 2 * channels, kernel_size=1, bias=bias) - self.depthwise_conv = nn.Conv1d( - channels, channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, groups=channels, bias=bias - ) - self.norm = nn.BatchNorm1d(channels) - self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1, bias=bias) - self.activation = activation - - def forward(self, x): - # (batch, time, channels) -> (batch, channels, time) - x = x.transpose(1, 2) - - x = self.pointwise_conv1(x) - x = nn.functional.glu(x, dim=1) - - x = self.depthwise_conv(x) - x = self.norm(x) - x = self.activation(x) - - x = self.pointwise_conv2(x) - x = x.transpose(1, 2) - return x - - -class MultiLayeredConv1d(nn.Module): - """ - Multi-layered Conv1d that replaces feed-forward in Transformer block. - - Reference: - FastSpeech: Fast, Robust and Controllable Text to Speech - https://arxiv.org/pdf/1905.09263.pdf - """ - - def __init__(self, in_chans, hidden_chans, kernel_size): - super().__init__() - padding = (kernel_size - 1) // 2 - self.w_1 = nn.Conv1d(in_chans, hidden_chans, kernel_size, padding=padding) - self.w_2 = nn.Conv1d(hidden_chans, in_chans, kernel_size, padding=padding) - - def forward(self, x): - """ - Args: - x (torch.Tensor): (batch, time, in_chans) - Returns: - torch.Tensor: (batch, time, in_chans) - """ - # Conv1d expects (batch, channels, time) - x = x.transpose(1, 2) - x = torch.relu(self.w_1(x)) - x = self.w_2(x) - x = x.transpose(1, 2) - return x - - -class Swish(nn.Module): - """Swish activation function.""" - - def forward(self, x): - return x * torch.sigmoid(x) - - -class EncoderLayer(nn.Module): - """ - A single encoder layer: optional macaron FFN, self-attn, conv, FFN, norms. - Args: - size (int): Input dimension. - self_attn (nn.Module): Self-attention module. - feed_forward (nn.Module): Standard feed-forward network. - feed_forward_macaron (nn.Module): Optional, second FFN in macaron structure. - conv_module (nn.Module): Optional, convolution module. - normalize_before (bool): If True, use pre-layer norm. Else, post-layer norm. - concat_after (bool): If True, concat attn in/out, else residual. - """ - - def __init__( - self, - size, - self_attn, - feed_forward, - feed_forward_macaron, - conv_module, - normalize_before=True, - concat_after=False, - ): - super().__init__() - self.size = size - self.self_attn = self_attn - self.feed_forward = feed_forward - self.feed_forward_macaron = feed_forward_macaron - self.conv_module = conv_module - - self.norm_ff = LayerNorm(size) - self.norm_mha = LayerNorm(size) - - self.ff_scale = 0.5 if feed_forward_macaron is not None else 1.0 - self.norm_ff_macaron = LayerNorm(size) if feed_forward_macaron is not None else None - - if conv_module is not None: - self.norm_conv = LayerNorm(size) - self.norm_final = LayerNorm(size) - else: - self.norm_conv = None - self.norm_final = None - - self.normalize_before = normalize_before - self.concat_after = concat_after - if self.concat_after: - self.concat_linear = nn.Linear(2 * size, size) - - def forward(self, x_input, mask, cache=None) -> tuple: - """ - Forward pass for the encoder layer. - - Args: - x_input (Tensor or Tuple): Input tensor of shape (batch, time, size) - or tuple (x, pos_emb) where pos_emb is the positional embedding. - mask (Tensor): Mask tensor of shape (batch, time) or (batch, 1, time). - cache (Tensor, optional): Input cache for streaming (can be None). - - Returns: - Tuple[Tensor, Tensor]: (outputs, mask), where outputs either includes - pos_emb or not, depending on input. - """ - # Unpack possible (x, pos_emb) tuple input - if isinstance(x_input, tuple): - x, pos_emb = x_input - else: - x, pos_emb = x_input, None - - # --- Macaron Feed-Forward (first FFN) --- - if self.feed_forward_macaron is not None: - residual = x - if self.normalize_before: - x = self.norm_ff_macaron(x) - x = residual + self.ff_scale * self.feed_forward_macaron(x) - if not self.normalize_before: - x = self.norm_ff_macaron(x) - - # --- Self-Attention --- - residual = x - if self.normalize_before: - x = self.norm_mha(x) - - # Handle cache for streaming - if cache is not None: - assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size) - x_q = x[:, -1:, :] - residual = residual[:, -1:, :] - mask = mask[:, -1:, :] if mask is not None else None - else: - x_q = x - - # Apply self-attention (with or without positional embeddings) - if pos_emb is not None: - attn_out = self.self_attn(x_q, x, x, pos_emb, mask) - else: - attn_out = self.self_attn(x_q, x, x, mask) - - if self.concat_after: - x = torch.cat((x, attn_out), dim=-1) - x = residual + self.concat_linear(x) - else: - x = residual + attn_out - if not self.normalize_before: - x = self.norm_mha(x) - - # --- Convolution Module --- - if self.conv_module is not None: - residual = x - if self.normalize_before: - x = self.norm_conv(x) - x = residual + self.conv_module(x) - if not self.normalize_before: - x = self.norm_conv(x) - - # --- Feed-Forward (second FFN) --- - residual = x - if self.normalize_before: - x = self.norm_ff(x) - x = residual + self.ff_scale * self.feed_forward(x) - if not self.normalize_before: - x = self.norm_ff(x) - - # --- Final normalization, if present --- - if self.norm_final is not None: - x = self.norm_final(x) - - # --- Append to cache if streaming --- - if cache is not None: - x = torch.cat([cache, x], dim=1) - - # --- Final return, with or without positional embedding --- - if pos_emb is not None: - return (x, pos_emb), mask - else: - return x, mask - - -class ConformerLayers(nn.Module): - def __init__(self, hidden_size, num_layers, kernel_size=9, num_heads=4, use_last_norm=True, save_hidden=False): - super().__init__() - self.use_last_norm = use_last_norm - self.layers = nn.ModuleList() - positionwise_layer = MultiLayeredConv1d - positionwise_layer_args = (hidden_size, hidden_size * 4, 1) - self.pos_embed = RelPositionalEncoding(hidden_size) - self.encoder_layers = nn.ModuleList( - [ - EncoderLayer( - hidden_size, - RelPositionMultiHeadedAttention(num_heads, hidden_size, 0.0), - positionwise_layer(*positionwise_layer_args), - positionwise_layer(*positionwise_layer_args), - ConvolutionModule(hidden_size, kernel_size, Swish()), - ) - for _ in range(num_layers) - ] - ) - if self.use_last_norm: - self.layer_norm = nn.LayerNorm(hidden_size) - else: - self.layer_norm = nn.Linear(hidden_size, hidden_size) - self.save_hidden = save_hidden - if save_hidden: - self.hiddens = [] - - def forward(self, x, padding_mask=None): - """ - - :param x: [B, T, H] - :param padding_mask: [B, T] - :return: [B, T, H] - """ - self.hiddens = [] - nonpadding_mask = x.abs().sum(-1) > 0 - x = self.pos_embed(x) - for layer in self.encoder_layers: - x, mask = layer(x, nonpadding_mask[:, None, :]) - if self.save_hidden: - self.hiddens.append(x[0]) - x = x[0] - x = self.layer_norm(x) * nonpadding_mask.float()[:, :, None] - return x - - -# ---------------------- note transcription backbone layers ---------------------- - - -class UnetDown(nn.Module): - def __init__( - self, - hidden_size, - n_layers, - kernel_size, - down_rates, - channel_multiples=None, - is_btc=True, - constant_channels=False, - ): - super().__init__() - assert n_layers == len(down_rates) # downs, down sample rate - down_rates = [int(i) for i in down_rates] - self.n_layers = n_layers - self.hidden_size = hidden_size - self.is_btc = is_btc - channel_multiples = channel_multiples if channel_multiples is not None else down_rates - self.layers = nn.ModuleList() - self.downs = nn.ModuleList() - in_channels = hidden_size - for i in range(self.n_layers): - out_channels = int(in_channels * channel_multiples[i]) if not constant_channels else in_channels - self.layers.append( - nn.Sequential( - ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), - nn.Conv1d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2), - ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), - ) - ) - self.downs.append(nn.Sequential(nn.AvgPool1d(down_rates[i]))) - in_channels = out_channels - self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6) - self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2) - - def forward(self, x, **kwargs): - # x [B, T, C] - if self.is_btc: - x = x.transpose(1, 2) # [B, C, T] - skip_xs = [] - for i in range(self.n_layers): - skip_x = self.layers[i](x) - x = self.downs[i](skip_x) - if self.is_btc: - skip_xs.append(skip_x.transpose(1, 2)) # [B, T, C] - else: - skip_xs.append(skip_x) - x = self.post_net(self.last_norm(x)) - if self.is_btc: - x = x.transpose(1, 2) - return x, skip_xs - - -class UnetMid(nn.Module): - def __init__(self, hidden_size, kernel_size, n_layers=None, in_dims=None, out_dims=None, is_btc=True, net=None): - super().__init__() - in_dims = in_dims if in_dims is not None else hidden_size - out_dims = out_dims if out_dims is not None else hidden_size - self.pre = nn.Conv1d(in_dims, hidden_size, kernel_size, padding=kernel_size // 2) - self.post = nn.Conv1d(hidden_size, out_dims, kernel_size, padding=kernel_size // 2) - self.is_btc = is_btc - if net is not None: - self.net = net - else: - self.net = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=kernel_size, - layers_in_block=2, - c_multiple=2, - post_net_kernel=3, - is_btc=is_btc, - ) - - def forward(self, x, cond=None, **kwargs): - # x [B, T, C] - if self.is_btc: - x = self.pre(x.transpose(1, 2)).transpose(1, 2) - else: - x = self.pre(x) - if cond is None: - cond = 0 - x = self.net(x + cond) - if self.is_btc: - x = self.post(x.transpose(1, 2)).transpose(1, 2) - else: - x = self.post(x) - return x - - -class UnetUp(nn.Module): - def __init__( - self, - hidden_size, - n_layers, - kernel_size, - up_rates, - channel_multiples=None, - is_btc=True, - constant_channels=False, - use_skip_layer=False, - skip_scale=1.0, - ): - super().__init__() - assert n_layers == len(up_rates) # this is reversed in up module, from the output to the interface with middle - up_rates = [int(i) for i in up_rates] - self.n_layers = n_layers - self.hidden_size = hidden_size - self.is_btc = is_btc - self.skip_scale = skip_scale - channel_multiples = channel_multiples if channel_multiples is not None else up_rates - # in_channels = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size - self.in_channels_lst = ( - (np.cumprod([1] + channel_multiples) * hidden_size).astype(int) - if not constant_channels - else [hidden_size for _ in range(self.n_layers + 1)] - ) - in_channels = self.in_channels_lst[-1] - self.ups = nn.ModuleList() - self.skip_layers = nn.ModuleList() - self.layers = nn.ModuleList() - for i in range(self.n_layers - 1, -1, -1): - out_channels = self.in_channels_lst[i] if not constant_channels else in_channels - self.ups.append( - nn.Sequential( - nn.ConvTranspose1d( - in_channels, - in_channels, - kernel_size=kernel_size, - stride=up_rates[i], - padding=kernel_size // 2, - output_padding=up_rates[i] - 1, - ), - LayerNorm(in_channels, dim=1, eps=1e-6), - nn.LeakyReLU(), - ) - ) - self.layers.append( - nn.Sequential( - nn.Conv1d(in_channels * 2, out_channels, kernel_size, padding=(kernel_size - 1) // 2), - ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), - ) - ) - if use_skip_layer: - self.skip_layers.append( - ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5) - ) - else: - self.skip_layers.append(nn.Identity()) - - in_channels = out_channels - self.out_channels = out_channels - self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6) - self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2) - - def forward(self, x, skips, **kwargs): - # x [B, T, C] - if self.is_btc: - x = x.transpose(1, 2) # [B, C, T] - for i in range(self.n_layers): - x = self.ups[i](x) - skip_x = ( - skips[self.n_layers - i - 1] if not self.is_btc else skips[self.n_layers - i - 1].transpose(1, 2) - ) # [B, T, C] -> [B, C, T] - skip_x = self.skip_layers[i](skip_x) * self.skip_scale - x = torch.cat((x, skip_x), dim=1) # [B, C, T] - x = self.layers[i](x) - x = self.post_net(self.last_norm(x)) - if self.is_btc: - x = x.transpose(1, 2) - return x - - -class Unet(nn.Module): - def __init__( - self, - hidden_size, - down_layers, - up_layers, - kernel_size, - updown_rates, - mid_layers=None, - channel_multiples=None, - is_btc=True, - constant_channels=False, - mid_net=None, - use_skip_layer=False, - skip_scale=1.0, - ): - super().__init__() - assert len(updown_rates) == down_layers == up_layers, f"{len(updown_rates)}, {down_layers}, {up_layers}" - if channel_multiples is not None: - assert len(channel_multiples) == len(updown_rates) - else: - channel_multiples = updown_rates - self.down = UnetDown( - hidden_size, down_layers, kernel_size, updown_rates, channel_multiples, is_btc, constant_channels - ) - down_out_dims = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size - self.mid = UnetMid( - hidden_size, - kernel_size, - mid_layers, - in_dims=down_out_dims, - out_dims=down_out_dims, - is_btc=is_btc, - net=mid_net, - ) - self.up = UnetUp( - hidden_size, - up_layers, - kernel_size, - updown_rates, - channel_multiples, - is_btc, - constant_channels, - use_skip_layer, - skip_scale, - ) - - def forward(self, x, mid_cond=None, **kwargs): - x, skips = self.down(x) - x = self.mid(x, mid_cond) - x = self.up(x, skips) - return x - - -class BackboneNet(nn.Module): - def __init__(self, hparams): - super().__init__() - self.hidden_size = hidden_size = hparams["hidden_size"] - updown_rates = [2, 2, 2] - channel_multiples = [1, 1, 1] - if hparams.get("updown_rates", None) is not None: - updown_rates = [int(i) for i in hparams.get("updown_rates", None).split("-")] - if hparams.get("channel_multiples", None) is not None: - channel_multiples = [float(i) for i in hparams.get("channel_multiples", None).split("-")] - assert len(updown_rates) == len(channel_multiples) - # convs - if hparams.get("bkb_net", "conv") == "conv": - self.net = Unet( - hidden_size, - down_layers=len(updown_rates), - mid_layers=hparams.get("bkb_layers", 12), - up_layers=len(updown_rates), - kernel_size=3, - updown_rates=updown_rates, - channel_multiples=channel_multiples, - is_btc=True, - constant_channels=False, - mid_net=None, - use_skip_layer=hparams.get("unet_skip_layer", False), - ) - # conformer - elif hparams.get("bkb_net", "conv") == "conformer": - mid_net = ConformerLayers( - hidden_size, - num_layers=hparams.get("bkb_layers", 12), - kernel_size=hparams.get("conformer_kernel", 9), - num_heads=4, - ) - self.net = Unet( - hidden_size, - down_layers=len(updown_rates), - up_layers=len(updown_rates), - kernel_size=3, - updown_rates=updown_rates, - channel_multiples=channel_multiples, - is_btc=True, - constant_channels=False, - mid_net=mid_net, - use_skip_layer=hparams.get("unet_skip_layer", False), - ) - - def forward(self, x): - return self.net(x) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py deleted file mode 100644 index 9e020f61b..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Preprocess neural network modules for SoulX-Singer.""" - -from .pipeline import SoulXPreprocessPipeline - -__all__ = [ - "SoulXPreprocessPipeline", -] diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py deleted file mode 100644 index 0752c73f8..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py +++ /dev/null @@ -1,326 +0,0 @@ -"""Lyric transcription (Paraformer + Parakeet) for SoulX preprocess.""" - -import os -import re -import tempfile -from pathlib import Path -from typing import Any - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -from omegaconf import open_dict -from vllm.logger import init_logger - -from .utils import get_audio_duration - -logger = init_logger(__name__) - - -def _import_funasr_automodel(): - """Import FunASR without shelling out to ffmpeg at module import time.""" - import subprocess - from functools import wraps - - real_check_output = subprocess.check_output - - @wraps(real_check_output) - def _safe_check_output(cmd, *args, **kwargs): - if cmd and Path(str(cmd[0])).name == "ffmpeg": - raise FileNotFoundError("ffmpeg unavailable for FunASR import probe") - return real_check_output(cmd, *args, **kwargs) - - subprocess.check_output = _safe_check_output - try: - from funasr import AutoModel - finally: - subprocess.check_output = real_check_output - return AutoModel - - -def build_words_with_gaps(raw_words, raw_timestamps, audio_duration: float): - words, word_durs = [], [] - prev = 0.0 - for w, t in zip(raw_words, raw_timestamps): - s, e = float(t[0]), float(t[1]) - if s > prev: - words.append("") - word_durs.append(s - prev) - words.append(w) - word_durs.append(e - s) - prev = e - - if audio_duration > prev: - if not words: - words.append("") - word_durs.append(audio_duration) - elif words[-1] != "": - words.append("") - word_durs.append(audio_duration - prev) - else: - word_durs[-1] += audio_duration - prev - return words, word_durs - - -class ParaformerASR(nn.Module): - """Mandarin/Cantonese ASR via FunASR Paraformer.""" - - def __init__(self, model_path: str, device: str): - super().__init__() - from vllm_omni.diffusion.models.soulx_singer.utils import _patch_torchaudio_load - - _patch_torchaudio_load() - AutoModel = _import_funasr_automodel() - self.asr = AutoModel(model=model_path, disable_update=True, device=device) - logger.info("Loaded Paraformer ASR from %s", model_path) - - def _generate(self, source: str) -> tuple[list[str], list[float], float]: - out = self.asr.generate(source, output_timestamp=True)[0] - raw_words = out["text"].replace("@", "").split(" ") - raw_timestamps = [[t[0] / 1000, t[1] / 1000] for t in out["timestamp"]] - duration = float(raw_timestamps[-1][1]) if raw_timestamps else 0.0 - return raw_words, raw_timestamps, duration - - def _maybe_refine_with_sidecar_f0( - self, - wav_path: str, - words: list[str], - word_durs: list[float], - *, - sample_rate: int, - hop_size: int, - ) -> tuple[list[str], list[float]]: - """Match upstream ``_ASRZhModel.process``: refine using ``{wav}_f0.npy`` if present.""" - f0_path = os.path.splitext(wav_path)[0] + "_f0.npy" - if not os.path.isfile(f0_path): - return words, word_durs - f0 = np.load(f0_path) - if f0.size == 0: - return words, word_durs - return refine_word_durs_with_f0( - words, - word_durs, - f0, - sample_rate=sample_rate, - hop_size=hop_size, - ) - - @torch.inference_mode() - def forward(self, audio: str | np.ndarray, sample_rate: int = 16000) -> tuple[list[str], list[float]]: - if isinstance(audio, str): - raw_words, raw_timestamps, _ = self._generate(audio) - audio_duration = get_audio_duration(audio) - words, word_durs = build_words_with_gaps(raw_words, raw_timestamps, audio_duration) - return self._maybe_refine_with_sidecar_f0( - audio, - words, - word_durs, - sample_rate=24000, - hop_size=480, - ) - - waveform = np.asarray(audio, dtype=np.float32).squeeze() - audio_duration = len(waveform) / float(sample_rate) - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp_path = tmp.name - try: - sf.write(tmp_path, waveform, sample_rate) - raw_words, raw_timestamps, _ = self._generate(tmp_path) - words, word_durs = build_words_with_gaps(raw_words, raw_timestamps, audio_duration) - return self._maybe_refine_with_sidecar_f0( - tmp_path, - words, - word_durs, - sample_rate=24000, - hop_size=480, - ) - finally: - Path(tmp_path).unlink(missing_ok=True) - Path(os.path.splitext(tmp_path)[0] + "_f0.npy").unlink(missing_ok=True) - - -def _disable_nemo_cuda_graphs(model) -> None: - decoding_cfg = model.cfg.decoding - with open_dict(decoding_cfg): - if hasattr(decoding_cfg, "greedy") and decoding_cfg.greedy is not None: - decoding_cfg.greedy.use_cuda_graph_decoder = False - if hasattr(decoding_cfg, "beam") and decoding_cfg.beam is not None: - decoding_cfg.beam.allow_cuda_graphs = False - model.change_decoding_strategy(decoding_cfg) - - -class ParakeetASR(nn.Module): - """English ASR via NeMo Parakeet-TDT.""" - - def __init__(self, model_path: str, device: str): - super().__init__() - try: - import nemo.collections.asr as nemo_asr # type: ignore - except Exception as e: # pragma: no cover - raise ImportError( - "NeMo (nemo_toolkit) is required for ASR English but is not available in this Python env." - ) from e - - self.asr = nemo_asr.models.ASRModel.restore_from( - restore_path=model_path, - map_location=device, - ) - self.asr.eval() - _disable_nemo_cuda_graphs(self.asr) - logger.info("Loaded Parakeet ASR from %s", model_path) - - @staticmethod - def _clean_word(word: str) -> str: - return re.sub(r"[\?\.,:]", "", word).strip() - - @staticmethod - def _extract_word_segments(output: Any) -> list[dict[str, Any]]: - ts = getattr(output, "timestamp", None) - if not ts or not isinstance(ts, dict): - return [] - word_ts = ts.get("word") - return word_ts if isinstance(word_ts, list) else [] - - @torch.inference_mode() - def forward(self, audio: str | np.ndarray, sample_rate: int = 16000) -> tuple[list[str], list[float]]: - if isinstance(audio, str): - sources = [audio] - duration = 0.0 - else: - waveform = np.asarray(audio, dtype=np.float32).squeeze() - duration = len(waveform) / float(sample_rate) - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp_path = tmp.name - sf.write(tmp_path, waveform, sample_rate) - sources = [tmp_path] - - try: - outputs = self.asr.transcribe( - sources, - timestamps=True, - batch_size=1, - num_workers=0, - ) - finally: - if not isinstance(audio, str): - Path(sources[0]).unlink(missing_ok=True) - - output = outputs[0] if outputs else None - raw_words: list[str] = [] - raw_timestamps: list[list[float]] = [] - if output is not None: - for w in self._extract_word_segments(output): - s, e = float(w.get("start", 0.0)), float(w.get("end", 0.0)) - word = self._clean_word(str(w.get("word", ""))) - if word: - raw_words.append(word) - raw_timestamps.append([s, e]) - - if raw_timestamps: - duration = max(duration, float(raw_timestamps[-1][1])) - return build_words_with_gaps(raw_words, raw_timestamps, duration) - - -def refine_word_durs_with_f0( - words: list[str], - word_durs: list[float], - f0: np.ndarray, - *, - sample_rate: int = 24000, - hop_size: int = 480, -) -> tuple[list[str], list[float]]: - boundaries = np.cumsum([0, *[int(dur * sample_rate / hop_size) for dur in word_durs]]).tolist() - sil_tolerance = 5 - ext_tolerance = 5 - new_words: list[str] = [] - new_word_durs: list[float] = [] - if words: - new_words.append(words[0]) - new_word_durs.append(word_durs[0]) - - for i in range(1, len(words)): - word = words[i] - if word == "": - start_frame = boundaries[i] - end_frame = boundaries[i + 1] - num_frames = end_frame - start_frame - frame_idx = start_frame - unvoiced_count = 0 - while frame_idx < end_frame: - if f0[frame_idx] <= 1: - unvoiced_count += 1 - if unvoiced_count >= sil_tolerance: - frame_idx -= sil_tolerance - 1 - break - else: - unvoiced_count = 0 - frame_idx += 1 - - voice_frames = frame_idx - start_frame - if voice_frames >= int(num_frames * 0.9): - new_word_durs[-1] += word_durs[i] - elif voice_frames >= ext_tolerance: - dur = voice_frames * hop_size / sample_rate - new_word_durs[-1] += dur - new_words.append("") - new_word_durs.append(word_durs[i] - dur) - else: - new_words.append(word) - new_word_durs.append(word_durs[i]) - else: - new_words.append(word) - new_word_durs.append(word_durs[i]) - return new_words, new_word_durs - - -class LyricModel(nn.Module): - def __init__( - self, - zh_model_path: str, - en_model_path: str, - device: str = "cuda", - *, - target_sr: int = 24000, - hop_size: int = 480, - ): - super().__init__() - self.device_str = device - self.en_model_path = en_model_path - self.target_sr = target_sr - self.hop_size = hop_size - self.zh_asr = ParaformerASR(zh_model_path, device=device) - self.en_asr: nn.Module | None = None - - @torch.inference_mode() - def forward( - self, - audio: str | np.ndarray, - language: str = "Mandarin", - *, - sample_rate: int = 16000, - f0: np.ndarray | None = None, - ) -> tuple[list[str], list[float]]: - if language not in {"Mandarin", "Cantonese", "English"}: - raise ValueError(f"Unsupported language: {language}") - - if language.lower() == "english": - if self.en_asr is None: - self.en_asr = ParakeetASR(self.en_model_path, device=self.device_str) - self.add_module("en_asr", self.en_asr) - words, durs = self.en_asr(audio, sample_rate=sample_rate) - else: - # Path-based ASR loads ``{wav}_f0.npy`` inside ParaformerASR (upstream parity). - if isinstance(audio, str): - words, durs = self.zh_asr(audio) - else: - words, durs = self.zh_asr(audio, sample_rate=sample_rate) - if f0 is not None and len(f0) > 0: - words, durs = refine_word_durs_with_f0( - words, - durs, - f0, - sample_rate=self.target_sr, - hop_size=self.hop_size, - ) - return words, durs diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py deleted file mode 100644 index 825c1622f..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Resample RMVPE F0 (16 kHz hop) onto the SoulX mel frame grid.""" - -from __future__ import annotations - -import math -import time -from typing import TYPE_CHECKING - -import numpy as np -from scipy.interpolate import interp1d -from vllm.logger import init_logger - -from .rmvpe import RMVPE_SAMPLE_RATE -from .utils import load_mono_audio - -if TYPE_CHECKING: - from .rmvpe import RMVPE - -logger = init_logger(__name__) - -RMVPE_HOP = 160 - - -def interpolate_f0_to_mel_grid( - f0_16k: np.ndarray, - original_length: int, - original_sr: int, - *, - target_sr: int, - hop_size: int, - max_duration: float = 300.0, -) -> np.ndarray: - batch_max_length = int(max_duration * target_sr / hop_size) - duration_in_seconds = original_length / original_sr - effective_target_length = int(duration_in_seconds * target_sr) - original_frames = math.ceil(effective_target_length / hop_size) - target_frames = min(original_frames, batch_max_length) - - t_16k = np.arange(len(f0_16k)) * (RMVPE_HOP / float(original_sr)) - t_target = np.arange(target_frames) * (hop_size / float(target_sr)) - - if len(f0_16k) > 0: - f_interp = interp1d( - t_16k, - f0_16k, - kind="linear", - bounds_error=False, - fill_value=0.0, - assume_sorted=True, - ) - f0 = f_interp(t_target) - else: - f0 = np.zeros(target_frames) - - if len(f0) != target_frames: - f0 = f0[:target_frames] if len(f0) > target_frames else np.pad(f0, (0, target_frames - len(f0))) - return f0.astype(np.float32, copy=False) - - -def extract_f0_file( - rmvpe: RMVPE, - wav_path: str, - *, - target_sr: int, - hop_size: int, - f0_path: str | None = None, - thread: float = 0.03, - max_duration: float = 300.0, - verbose: bool = False, -) -> np.ndarray: - if verbose: - logger.info("[f0] extracting from %s", wav_path) - t0 = time.perf_counter() - - audio, _ = load_mono_audio(wav_path, target_sr=RMVPE_SAMPLE_RATE) - f0_16k = rmvpe.infer_from_audio(audio, sample_rate=RMVPE_SAMPLE_RATE) - f0 = interpolate_f0_to_mel_grid( - np.asarray(f0_16k, dtype=np.float32), - original_length=audio.shape[-1], - original_sr=RMVPE_SAMPLE_RATE, - target_sr=target_sr, - hop_size=hop_size, - max_duration=max_duration, - ) - - if verbose: - voiced_ratio = float(np.mean(f0 > 0)) if len(f0) else 0.0 - logger.info( - "[f0] done: frames=%s voiced_ratio=%.3f time=%.3fs", - len(f0), - voiced_ratio, - time.perf_counter() - t0, - ) - if f0_path is not None: - np.save(f0_path, f0) - return f0 diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py deleted file mode 100644 index cf5f84e7a..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py +++ /dev/null @@ -1,431 +0,0 @@ -"""SoulX preprocess pipeline.""" - -from __future__ import annotations - -import json -import tempfile -from collections.abc import Iterable -from pathlib import Path -from typing import Any, ClassVar - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.stack import SoulXPreprocessStack -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import load_mono_audio, resample_mono -from vllm_omni.diffusion.models.soulx_singer.preprocess.metadata_utils import ( - SegmentMetadata, - _merge_group, - convert_metadata, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - load_config, - load_wav, - preprocess_weight_paths, - resolve_preprocess_weights_root, - resolve_soulx_kind, -) - -logger = init_logger(__name__) - - -def build_svs_prompt_meta( - preprocess: SoulXPreprocessPipeline, - *, - prompt_audio: str | tuple[np.ndarray, int], - prompt_meta_raw: dict[str, Any], - metadata_processor, - vocal_sep: bool | None = None, -) -> dict[str, Any]: - """Build SVS prompt_meta; in-memory prompt wav is trimmed to mel2note length.""" - if isinstance(prompt_audio, str): - return metadata_processor.process(prompt_meta_raw, prompt_audio) - prompt_meta = metadata_processor.process(prompt_meta_raw, None) - prompt_vocal, prompt_sr = preprocess._extract_vocal(prompt_audio, vocal_sep=vocal_sep) - if prompt_sr != metadata_processor.sample_rate: - prompt_vocal = resample_mono( - prompt_vocal, - orig_sr=prompt_sr, - target_sr=metadata_processor.sample_rate, - ) - max_samples = prompt_meta["mel2note"].shape[1] * metadata_processor.hop_size - segment_wav = np.asarray(prompt_vocal[:max_samples], dtype=np.float32) - prompt_meta["wav"] = torch.from_numpy(segment_wav).unsqueeze(0).float().to(metadata_processor.device) - return prompt_meta - - -class SoulXPreprocessPipeline(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - """Lazy-loaded preprocess stack integrated with vLLM-Omni diffusion lifecycle.""" - - support_audio_input: ClassVar[bool] = True - weights_sources: ClassVar[tuple] = () - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = [ - "stack.vocal_sep", - "stack.lyric", - "stack.rosvot", - ] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["stack.rmvpe"] - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - vocal_sep: bool = True, - midi_transcribe: bool = True, - max_merge_duration_ms: int = 60000, - verbose: bool = False, - extra_args: dict[str, Any] | None = None, - ) -> None: - super().__init__() - del prefix - self.od_config = od_config - extra_args = extra_args or {} - self.kind = resolve_soulx_kind(od_config.model) - hf_config = load_config(str(Path(od_config.model) / "config.yaml")) - audio_config = hf_config.audio - if self.kind == "svc" and midi_transcribe: - logger.warning("SVC model (config.json): midi_transcribe is ignored; lyric/MIDI transcription is SVS-only.") - midi_transcribe = False - - self.vocal_sep = vocal_sep - self.midi_transcribe = midi_transcribe - self.max_merge_duration_ms = max_merge_duration_ms - self.verbose = verbose - self.target_sr = int(audio_config.sample_rate) - self.hop_size = int(audio_config.hop_size) - self._metadata_processor = None - self._weights_root: Path | None = None - self.stack: SoulXPreprocessStack | None = None - self._ensure_stack(extra_args) - - def _ensure_stack(self, extra_args: dict[str, Any] | None = None) -> None: - extra_args = extra_args or {} - weights_root = resolve_preprocess_weights_root(self.od_config) - if self.stack is not None and self._weights_root == weights_root: - return - weights = preprocess_weight_paths(weights_root) - device = str(get_local_device()) - self.stack = SoulXPreprocessStack( - weights, - device, - target_sr=self.target_sr, - hop_size=self.hop_size, - verbose=self.verbose, - ) - self._weights_root = weights_root - - def forward(self, req) -> DiffusionOutput: - from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PREPROCESSED_KEY, - ) - from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( - build_metadata_processor, - build_preprocess_payload, - build_warmup_payload, - is_warmup_request, - ) - from vllm_omni.inputs.data import OmniTextPrompt - - extra_args = dict(getattr(req.sampling_params, "extra_args", None) or {}) - prompt = OmniTextPrompt(prompt=req.prompts[0]) if isinstance(req.prompts[0], str) else req.prompts[0] - info = prompt.get("additional_information") or {} - if isinstance(info, dict): - extra_args.update(dict(info.get("extra_args") or {})) - - self._ensure_stack(extra_args) - if self._metadata_processor is None: - self._metadata_processor = build_metadata_processor(self.od_config) - - device = get_local_device() - if is_warmup_request(req): - payload = build_warmup_payload( - self.kind, - metadata_processor=self._metadata_processor, - device=device, - sample_rate=self.target_sr, - ) - else: - payload = build_preprocess_payload( - self.kind, - prompt=prompt, - extra_args=extra_args, - preprocess=self, - metadata_processor=self._metadata_processor, - device=device, - sample_rate=self.target_sr, - ) - - return DiffusionOutput( - output={ - "payload": {}, - "metadata": {SOULX_PREPROCESSED_KEY: payload}, - }, - to_cpu=True, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - del weights - self._ensure_stack({}) - return set() - - def _extract_vocal( - self, - audio: str | tuple[np.ndarray, int], - *, - vocal_sep: bool | None = None, - ) -> tuple[np.ndarray, int]: - if self.stack is None: - self._ensure_stack({}) - use_sep = self.vocal_sep if vocal_sep is None else vocal_sep - if use_sep: - return self.stack.extract_vocal(audio) - return load_mono_audio(audio) - - def extract_f0(self, vocal: np.ndarray, sample_rate: int) -> np.ndarray: - if self.stack is None: - self._ensure_stack({}) - return self.stack.extract_f0(vocal, sample_rate) - - def iter_svs_metadata( - self, - audio_source: str | tuple[np.ndarray, int], - *, - language: str = "Mandarin", - vocal_sep: bool | None = None, - max_merge_duration_ms: int | None = None, - ): - """Yield merged SVS metadata dicts one segment at a time.""" - if isinstance(audio_source, tuple): - vocal, sample_rate = audio_source - origin_wav_fn = "" - else: - vocal, sample_rate = self._extract_vocal(str(audio_source), vocal_sep=vocal_sep) - origin_wav_fn = str(audio_source) - - with tempfile.TemporaryDirectory(prefix="soulx_preprocess_") as work_dir: - work = Path(work_dir) - vocal_path = work / "vocal.wav" - sf.write(vocal_path, vocal, sample_rate) - - vocal_f0_path = str(vocal_path).replace(".wav", "_f0.npy") - vocal_f0 = self.stack.extract_f0(str(vocal_path), sample_rate=sample_rate, f0_path=vocal_f0_path) - - if not self.midi_transcribe: - end_ms = int(len(vocal) / sample_rate * 1000) - item = SegmentMetadata( - item_name="full", - wav_fn=str(vocal_path), - language=language, - start_time_ms=0, - end_time_ms=end_ms, - note_text=[""], - note_dur=[end_ms / 1000.0], - note_pitch=[0], - note_type=[1], - ) - yield convert_metadata(item) - return - - segments = self.stack.ensure_segmenter().forward( - vocal, - sample_rate, - vocal_f0, - base_name="vocal", - origin_wav_fn=origin_wav_fn, - verbose=self.verbose, - ) - cut_dir = work / "cut_wavs" - cut_dir.mkdir(parents=True, exist_ok=True) - lyric = self.stack.ensure_lyric() - rosvot = self.stack.ensure_rosvot() - - long_cut_dir = work / "long_cut_wavs" - max_dur = max_merge_duration_ms or self.max_merge_duration_ms - max_gap_ms = 10000 - current_group: list[Any] = [] - current_len = 0 - prev_end = -1 - - def _flush_group() -> dict[str, Any]: - merged_item = _merge_group(vocal, sample_rate, current_group, long_cut_dir) - merged_f0_path = merged_item.wav_fn.replace(".wav", "_f0.npy") - self.stack.extract_f0(merged_item.wav_fn, sample_rate=sample_rate, f0_path=merged_f0_path) - return convert_metadata(merged_item) - - for seg in segments: - seg_key = seg["item_name"] - seg_wav_path = cut_dir / f"{seg_key}.wav" - sf.write(seg_wav_path, seg["wav"], seg["sample_rate"]) - seg_f0_path = str(seg_wav_path).replace(".wav", "_f0.npy") - self.stack.extract_f0(str(seg_wav_path), sample_rate=seg["sample_rate"], f0_path=seg_f0_path) - - words, durs = lyric.forward( - str(seg_wav_path), - language, - sample_rate=seg["sample_rate"], - ) - seg_item = { - "item_name": seg_key, - "wav_fn": str(seg_wav_path), - "start_time_ms": seg["start_time_ms"], - "end_time_ms": seg["end_time_ms"], - "origin_wav_fn": origin_wav_fn or str(vocal_path), - "words": words, - "word_durs": durs, - "language": language, - } - meta_item = rosvot.transcribe(seg_item, segment_info=seg_item, verbose=self.verbose) - - start_time = int(meta_item.get("start_time_ms", seg["start_time_ms"])) - end_time = int(meta_item.get("end_time_ms", seg["end_time_ms"])) - if current_group and ( - start_time - prev_end > max_gap_ms or current_len + end_time - start_time > max_dur - ): - yield _flush_group() - current_group = [] - current_len = 0 - - current_group.append(meta_item) - current_len += end_time - start_time - prev_end = end_time - - if current_group: - yield _flush_group() - - @staticmethod - def build_svs_payload_from_paths( - *, - prompt_metadata_path: str, - target_metadata_path: str, - audio_path: str, - metadata_processor, - ) -> dict[str, Any]: - with open(prompt_metadata_path, encoding="utf-8") as f: - prompt_metadata = json.load(f) - with open(target_metadata_path, encoding="utf-8") as f: - target_metadata = json.load(f) - if not prompt_metadata: - raise ValueError("prompt_metadata is empty") - if not target_metadata: - raise ValueError("target_metadata is empty") - - prompt_meta = metadata_processor.process(prompt_metadata[0], audio_path) - return { - "kind": "svs", - "prompt_meta": prompt_meta, - "target_meta_list": target_metadata, - } - - def build_svs_payload_from_audio( - self, - *, - prompt_audio: str | tuple[np.ndarray, int], - target_audio: str | tuple[np.ndarray, int], - metadata_processor, - language: str = "Mandarin", - vocal_sep: bool | None = None, - prompt_vocal_sep: bool | None = None, - target_vocal_sep: bool | None = None, - prompt_max_merge_duration_ms: int | None = None, - target_max_merge_duration_ms: int | None = None, - ) -> dict[str, Any]: - p_sep = prompt_vocal_sep if prompt_vocal_sep is not None else vocal_sep - t_sep = target_vocal_sep if target_vocal_sep is not None else vocal_sep - p_merge = ( - prompt_max_merge_duration_ms if prompt_max_merge_duration_ms is not None else self.max_merge_duration_ms - ) - t_merge = ( - target_max_merge_duration_ms if target_max_merge_duration_ms is not None else self.max_merge_duration_ms - ) - prompt_list = list( - self.iter_svs_metadata( - prompt_audio, - language=language, - vocal_sep=p_sep, - max_merge_duration_ms=p_merge, - ) - ) - target_list = list( - self.iter_svs_metadata( - target_audio, - language=language, - vocal_sep=t_sep, - max_merge_duration_ms=t_merge, - ) - ) - if not prompt_list or not target_list: - raise ValueError("SVS preprocess produced empty metadata") - - # Prompt wav is trimmed to mel2note length, not the metadata time window. - prompt_meta = build_svs_prompt_meta( - self, - prompt_audio=prompt_audio, - prompt_meta_raw=prompt_list[0], - metadata_processor=metadata_processor, - vocal_sep=vocal_sep, - ) - - return { - "kind": "svs", - "prompt_meta": prompt_meta, - "target_meta_list": target_list, - } - - @staticmethod - def build_svc_payload_from_paths( - *, - prompt_wav_path: str, - target_wav_path: str, - prompt_f0_path: str, - target_f0_path: str, - sample_rate: int, - device: torch.device | str, - ) -> dict[str, Any]: - prompt_wav = load_wav(prompt_wav_path, sample_rate).to(device) - target_wav = load_wav(target_wav_path, sample_rate).to(device) - prompt_f0 = torch.from_numpy(np.load(prompt_f0_path)).unsqueeze(0).to(device) - target_f0 = torch.from_numpy(np.load(target_f0_path)).unsqueeze(0).to(device) - return { - "kind": "svc", - "prompt_wav": prompt_wav, - "target_wav": target_wav, - "prompt_f0": prompt_f0, - "target_f0": target_f0, - } - - def build_svc_payload_from_audio( - self, - *, - prompt_audio: str | tuple[np.ndarray, int], - target_audio: str | tuple[np.ndarray, int], - sample_rate: int, - device: torch.device | str, - vocal_sep: bool | None = None, - ) -> dict[str, Any]: - prompt_vocal, prompt_sr = self._extract_vocal(prompt_audio, vocal_sep=vocal_sep) - target_vocal, target_sr = self._extract_vocal(target_audio, vocal_sep=vocal_sep) - - if prompt_sr != sample_rate: - prompt_vocal = resample_mono(prompt_vocal, orig_sr=prompt_sr, target_sr=sample_rate) - if target_sr != sample_rate: - target_vocal = resample_mono(target_vocal, orig_sr=target_sr, target_sr=sample_rate) - - prompt_f0 = self.extract_f0(prompt_vocal, sample_rate) - target_f0 = self.extract_f0(target_vocal, sample_rate) - - return { - "kind": "svc", - "prompt_wav": torch.from_numpy(prompt_vocal).unsqueeze(0).to(device), - "target_wav": torch.from_numpy(target_vocal).unsqueeze(0).to(device), - "prompt_f0": torch.from_numpy(prompt_f0).unsqueeze(0).to(device), - "target_f0": torch.from_numpy(target_f0).unsqueeze(0).to(device), - } diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py deleted file mode 100644 index 50ae4b2fc..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py +++ /dev/null @@ -1,593 +0,0 @@ -"""RMVPE F0 extractor. - -All model code is inlined (originally from RickyL-2000/ROSVOT -``modules/pe/rmvpe/``) with two minor substitutions: - -* ``librosa.filters.mel`` → ``torchaudio.functional.melscale_fbanks`` - (already used elsewhere in vllm-omni); -* ``librosa.sequence.viterbi`` → hand-written numpy viterbi (non-default path). -* ``pyworld`` (deprecated) → removed. -""" - -from __future__ import annotations - -import math -from typing import ClassVar - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from torchaudio.functional import melscale_fbanks -from torchaudio.transforms import Resample -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery - -logger = init_logger(__name__) - -RMVPE_SAMPLE_RATE = 16000 -_N_CLASS = 360 -_N_MELS = 128 -_MEL_FMIN = 30 -_MEL_FMAX = 8000 -_WINDOW_LENGTH = 1024 -_CONST = 1997.3794084376191 - - -class ConvBlockRes(nn.Module): - def __init__(self, in_channels: int, out_channels: int, momentum: float = 0.01): - super().__init__() - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), - nn.BatchNorm2d(out_channels, momentum=momentum), - nn.ReLU(), - nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), - nn.BatchNorm2d(out_channels, momentum=momentum), - nn.ReLU(), - ) - if in_channels != out_channels: - self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1)) - self._has_shortcut = True - else: - self._has_shortcut = False - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self._has_shortcut: - return self.conv(x) + self.shortcut(x) - return self.conv(x) + x - - -class ResEncoderBlock(nn.Module): - def __init__(self, in_channels: int, out_channels: int, kernel_size, n_blocks: int = 1, momentum: float = 0.01): - super().__init__() - self.n_blocks = n_blocks - self.conv = nn.ModuleList() - self.conv.append(ConvBlockRes(in_channels, out_channels, momentum)) - for _ in range(n_blocks - 1): - self.conv.append(ConvBlockRes(out_channels, out_channels, momentum)) - self.kernel_size = kernel_size - if self.kernel_size is not None: - self.pool = nn.AvgPool2d(kernel_size=kernel_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - for i in range(self.n_blocks): - x = self.conv[i](x) - if self.kernel_size is not None: - return x, self.pool(x) - return x - - -class ResDecoderBlock(nn.Module): - def __init__(self, in_channels: int, out_channels: int, stride, n_blocks: int = 1, momentum: float = 0.01): - super().__init__() - out_padding = (0, 1) if stride == (1, 2) else (1, 1) - self.n_blocks = n_blocks - self.conv1 = nn.Sequential( - nn.ConvTranspose2d( - in_channels, - out_channels, - kernel_size=(3, 3), - stride=stride, - padding=(1, 1), - output_padding=out_padding, - bias=False, - ), - nn.BatchNorm2d(out_channels, momentum=momentum), - nn.ReLU(), - ) - self.conv2 = nn.ModuleList() - self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum)) - for _ in range(n_blocks - 1): - self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum)) - - def forward(self, x: torch.Tensor, concat_tensor: torch.Tensor) -> torch.Tensor: - x = self.conv1(x) - x = torch.cat((x, concat_tensor), dim=1) - for i in range(self.n_blocks): - x = self.conv2[i](x) - return x - - -class Encoder(nn.Module): - def __init__( - self, - in_channels: int, - in_size: int, - n_encoders: int, - kernel_size, - n_blocks: int, - out_channels: int = 16, - momentum: float = 0.01, - ): - super().__init__() - self.n_encoders = n_encoders - self.bn = nn.BatchNorm2d(in_channels, momentum=momentum) - self.layers = nn.ModuleList() - self.latent_channels: list[list[int]] = [] - for _ in range(self.n_encoders): - self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum)) - self.latent_channels.append([out_channels, in_size]) - in_channels = out_channels - out_channels *= 2 - in_size //= 2 - self.out_size = in_size - self.out_channel = out_channels - - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: - concat_tensors: list[torch.Tensor] = [] - x = self.bn(x) - for i in range(self.n_encoders): - _c, x = self.layers[i](x) - concat_tensors.append(_c) - return x, concat_tensors - - -class Intermediate(nn.Module): - def __init__(self, in_channels: int, out_channels: int, n_inters: int, n_blocks: int, momentum: float = 0.01): - super().__init__() - self.n_inters = n_inters - self.layers = nn.ModuleList() - self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)) - for _ in range(n_inters - 1): - self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for i in range(self.n_inters): - x = self.layers[i](x) - return x - - -class Decoder(nn.Module): - def __init__(self, in_channels: int, n_decoders: int, stride, n_blocks: int, momentum: float = 0.01): - super().__init__() - self.layers = nn.ModuleList() - self.n_decoders = n_decoders - for _ in range(self.n_decoders): - out_channels = in_channels // 2 - self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)) - in_channels = out_channels - - def forward(self, x: torch.Tensor, concat_tensors: list[torch.Tensor]) -> torch.Tensor: - for i in range(self.n_decoders): - x = self.layers[i](x, concat_tensors[-1 - i]) - return x - - -class TimbreFilter(nn.Module): - def __init__(self, latent_rep_channels: list[list[int]]): - super().__init__() - self.layers = nn.ModuleList() - for ch in latent_rep_channels: - self.layers.append(ConvBlockRes(ch[0], ch[0])) - - def forward(self, x_tensors: list[torch.Tensor]) -> list[torch.Tensor]: - return [layer(t) for layer, t in zip(self.layers, x_tensors)] - - -class DeepUnet0(nn.Module): - def __init__( - self, - kernel_size, - n_blocks: int, - en_de_layers: int = 5, - inter_layers: int = 4, - in_channels: int = 1, - en_out_channels: int = 16, - ): - super().__init__() - self.encoder = Encoder(in_channels, _N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels) - self.intermediate = Intermediate( - self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks - ) - self.tf = TimbreFilter(self.encoder.latent_channels) - self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x, concat_tensors = self.encoder(x) - x = self.intermediate(x) - x = self.decoder(x, concat_tensors) - return x - - -class BiGRU(nn.Module): - def __init__(self, input_features: int, hidden_features: int, num_layers: int): - super().__init__() - self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.gru(x)[0] - - -class E2E0(nn.Module): - """E2E0 pitch estimator network (DeepUnet0 + CNN + BiGRU + classifier).""" - - def __init__( - self, - n_blocks: int, - n_gru: int, - kernel_size, - en_de_layers: int = 5, - inter_layers: int = 4, - in_channels: int = 1, - en_out_channels: int = 16, - ): - super().__init__() - self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels) - self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1)) - if n_gru: - self.fc = nn.Sequential( - BiGRU(3 * _N_MELS, 256, n_gru), - nn.Linear(512, _N_CLASS), - nn.Dropout(0.25), - nn.Sigmoid(), - ) - else: - self.fc = nn.Sequential( - nn.Linear(3 * _N_MELS, _N_CLASS), - nn.Dropout(0.25), - nn.Sigmoid(), - ) - - def forward(self, mel: torch.Tensor) -> torch.Tensor: - mel = mel.transpose(-1, -2).unsqueeze(1) - x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2) - x = self.fc(x) - return x - - -class MelSpectrogram(nn.Module): - """Mel-spectrogram extractor matching ROSVOT's MelSpectrogram. - - Log-mel spectrogram via torch.stft + adaptive filterbank. - """ - - def __init__( - self, - n_mel_channels: int = _N_MELS, - sampling_rate: int = RMVPE_SAMPLE_RATE, - win_length: int = _WINDOW_LENGTH, - hop_length: int = 160, - n_fft: int | None = None, - mel_fmin: float = _MEL_FMIN, - mel_fmax: float = _MEL_FMAX, - clamp: float = 1e-5, - ): - super().__init__() - n_fft = win_length if n_fft is None else n_fft - self.hann_window: dict[str, torch.Tensor] = {} - mel_basis = melscale_fbanks( - n_freqs=n_fft // 2 + 1, - f_min=float(mel_fmin), - f_max=float(mel_fmax) if mel_fmax is not None else 8000.0, - n_mels=n_mel_channels, - sample_rate=int(sampling_rate), - norm="slaney", - mel_scale="htk", - ).T # (n_freqs, n_mels) → (n_mels, n_freqs) to match librosa/layernorm convention - self.register_buffer("mel_basis", mel_basis) # (n_mels, n_freqs) - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - self.sampling_rate = sampling_rate - self.n_mel_channels = n_mel_channels - self.clamp = clamp - - def forward(self, audio: torch.Tensor, keyshift: float = 0, speed: float = 1, center: bool = True) -> torch.Tensor: - factor = 2 ** (keyshift / 12) - n_fft_new = int(np.round(self.n_fft * factor)) - win_length_new = int(np.round(self.win_length * factor)) - hop_length_new = int(np.round(self.hop_length * speed)) - - key = str(keyshift) + "_" + str(audio.device) - if key not in self.hann_window: - self.hann_window[key] = torch.hann_window(win_length_new).to(audio.device) - - if center: - pad_left = win_length_new // 2 - pad_right = (win_length_new + 1) // 2 - audio = F.pad(audio, (pad_left, pad_right)) - - fft = torch.stft( - audio, - n_fft=n_fft_new, - hop_length=hop_length_new, - win_length=win_length_new, - window=self.hann_window[key], - center=False, - return_complex=True, - ) - magnitude = fft.abs() - - if keyshift != 0: - size = self.n_fft // 2 + 1 - if magnitude.size(1) < size: - magnitude = F.pad(magnitude, (0, 0, 0, size - magnitude.size(1))) - magnitude = magnitude[:, :size, :] * self.win_length / win_length_new - - # register_buffer confuses type checkers; cast to Tensor explicitly - mel_basis: torch.Tensor = self.mel_basis # type: ignore[assignment] - mel_output = torch.matmul(mel_basis, magnitude) - log_mel = torch.log(torch.clamp(mel_output, min=self.clamp)) - return log_mel - - -def to_local_average_f0(hidden: torch.Tensor, center: torch.Tensor | None = None, thread: float = 0.03) -> np.ndarray: - """Decode hidden representation to F0 via local weighted average.""" - idx = torch.arange(_N_CLASS, device=hidden.device)[None, None, :] - idx_cents = idx * 20 + _CONST - if center is None: - center = torch.argmax(hidden, dim=2, keepdim=True) - start = torch.clip(center - 4, min=0) - end = torch.clip(center + 5, max=_N_CLASS) - idx_mask = (idx >= start) & (idx < end) - weights = hidden * idx_mask - product_sum = torch.sum(weights * idx_cents, dim=2) - weight_sum = torch.sum(weights, dim=2) - cents = product_sum / (weight_sum + (weight_sum == 0)) - f0 = 10 * 2 ** (cents / 1200) - uv = hidden.max(dim=2)[0] < thread - f0 = f0 * ~uv - return f0.cpu().numpy() - - -def numpy_viterbi(prob: np.ndarray, transition: np.ndarray) -> np.ndarray: - """Viterbi decoding on log-probabilities (numpy). - - Args: - prob: Emission probabilities, shape (n_states, n_obs). - transition: Transition matrix, shape (n_states, n_states), row-stochastic. - - Returns: - path: Most likely state sequence, shape (n_obs,). - """ - n_states, n_obs = prob.shape - log_emit = np.log(prob + 1e-12) - log_trans = np.log(transition + 1e-12) - - delta = log_emit[:, 0] + np.zeros(n_states) # initial: log(pi) omitted (uniform) - psi = np.zeros((n_obs, n_states), dtype=np.int64) - - for t in range(1, n_obs): - scores = delta[:, None] + log_trans # (n_states, n_states) - psi[t] = np.argmax(scores, axis=0) - delta = scores[psi[t], np.arange(n_states)] + log_emit[:, t] - - # Backtrack - path = np.zeros(n_obs, dtype=np.int64) - path[-1] = np.argmax(delta) - for t in range(n_obs - 2, -1, -1): - path[t] = psi[t + 1, path[t + 1]] - return path - - -def to_viterbi_f0(hidden: torch.Tensor, thread: float = 0.03) -> np.ndarray: - """Decode hidden representation to F0 via Viterbi. - - Uses a hand-written numpy Viterbi decoder (replacing - ``librosa.sequence.viterbi``) to avoid the librosa dependency. - """ - if not hasattr(to_viterbi_f0, "_transition"): - xx, yy = np.meshgrid(range(_N_CLASS), range(_N_CLASS)) - transition = np.maximum(30 - abs(xx - yy), 0) - transition = transition / transition.sum(axis=1, keepdims=True) - to_viterbi_f0._transition = transition # type: ignore[attr-defined] - - prob = hidden.squeeze(0).cpu().numpy() - prob = prob.T - prob = prob / prob.sum(axis=0, keepdims=True) - - path = numpy_viterbi(prob, to_viterbi_f0._transition) # type: ignore[attr-defined] - center = torch.from_numpy(path).unsqueeze(0).unsqueeze(-1).to(hidden.device) - return to_local_average_f0(hidden, center=center, thread=thread) - - -def resample_align_curve( - points: np.ndarray, original_timestep: float, target_timestep: float, align_length: int = -1 -) -> np.ndarray: - """Align a curve from one time grid to another via linear interpolation.""" - t_max = (len(points) - 1) * original_timestep - curve_interp = np.interp( - np.arange(0, t_max, target_timestep), - original_timestep * np.arange(len(points)), - points, - ).astype(points.dtype) - if align_length > 0: - if len(curve_interp) > align_length: - curve_interp = curve_interp[:align_length] - elif len(curve_interp) < align_length: - curve_interp = np.concatenate( - [curve_interp, np.full(align_length - len(curve_interp), fill_value=curve_interp[-1])], - axis=0, - ) - return curve_interp - - -class RMVPE(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - """RMVPE F0 extractor (SoulX-Singer). - - Wraps the E2E0 model with mel-spectrogram front-end and decoding - post-processing. - """ - - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = [] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["."] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = [] - - def __init__(self, model_path: str | None = None, hop_length: int = 160, device: str | None = None): - super().__init__() - self.resample_kernel: dict[str, Resample] = {} - if device is None: - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - else: - self.device = torch.device(device) - - self.model = E2E0(4, 1, (2, 2)).eval() - if model_path is not None: - ckpt = torch.load(model_path, map_location=self.device) - state_dict = ckpt if not isinstance(ckpt, dict) else ckpt.get("model", ckpt) - self.model.load_state_dict(state_dict, strict=False) - logger.info("Loaded RMVPE weights from %s", model_path) - self.model = self.model.to(self.device) - - self.mel_extractor = MelSpectrogram( - _N_MELS, - RMVPE_SAMPLE_RATE, - _WINDOW_LENGTH, - hop_length, - None, - _MEL_FMIN, - _MEL_FMAX, - ).to(self.device) - self.hop_length = hop_length - - # ---- internal helpers ---- - - @torch.no_grad() - def mel2hidden(self, mel: torch.Tensor) -> torch.Tensor: - n_frames = mel.shape[-1] - pad_len = 32 * ((n_frames - 1) // 32 + 1) - n_frames - mel = F.pad(mel, (0, pad_len), mode="constant") - hidden = self.model(mel) - return hidden[:, :n_frames] - - def decode(self, hidden: torch.Tensor, thread: float = 0.03, use_viterbi: bool = False) -> np.ndarray: - if use_viterbi: - return to_viterbi_f0(hidden, thread=thread) - return to_local_average_f0(hidden, thread=thread) - - def postprocess(self, f0: np.ndarray, fmin: float = 50, fmax: float = 1000, min_gap: int = 2) -> np.ndarray: - f0[f0 < fmin] = 0 - f0[f0 > fmax] = 0 - # Eliminate glitch: if positive f0 values span fewer than min_gap frames - # between zeros, zero them out. - for idx in range(f0.shape[0] - min_gap - 1): - if f0[idx] == 0 and f0[idx + min_gap + 1] == 0 and np.sum(f0[idx : idx + min_gap + 2]) > 0: - f0[idx : idx + min_gap + 2] = 0 - return f0 - - # ---- single audio inference ---- - - def infer_from_audio( - self, audio: np.ndarray, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False - ) -> np.ndarray: - audio_t = torch.from_numpy(audio).float().unsqueeze(0).to(self.device) - if sample_rate == RMVPE_SAMPLE_RATE: - audio_res = audio_t - else: - key_str = str(sample_rate) - if key_str not in self.resample_kernel: - self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128) - self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device) - audio_res = self.resample_kernel[key_str](audio_t) - mel = self.mel_extractor(audio_res, center=True) - hidden = self.mel2hidden(mel) - f0 = self.decode(hidden, thread=thread, use_viterbi=use_viterbi).squeeze(0) - return f0 - - def get_pitch( - self, - waveform, - sample_rate: int, - hop_size: int, - length: int, - interp_uv: bool = False, - fmin: float = 50, - fmax: float = 1000, - ) -> tuple[np.ndarray, np.ndarray]: - f0 = self.infer_from_audio(waveform, sample_rate=sample_rate) - f0 = self.postprocess(f0, fmin, fmax) - uv = f0 == 0 - time_step = hop_size / sample_rate - f0_res = resample_align_curve(f0, 0.01, time_step, length) - uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5 - if not interp_uv: - f0_res[uv_res] = 0 - return f0_res, uv_res - - # ---- batch audio inference ---- - - def infer_from_audio_batch( - self, audios, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False - ) -> list[np.ndarray]: - sizes: list[int] | None = None - if isinstance(audios, list): - audios = [torch.from_numpy(a).float() for a in audios] - sizes = [math.ceil((a.shape[0] + 1) / self.hop_length) for a in audios] - audios = nn.utils.rnn.pad_sequence(audios, batch_first=True, padding_value=0.0).to(self.device) - elif isinstance(audios, torch.Tensor): - if audios.device != self.device: - audios = audios.to(self.device) - else: - raise NotImplementedError(f"Unsupported audio type: {type(audios)}") - - if sample_rate == RMVPE_SAMPLE_RATE: - audios_res = audios - else: - key_str = str(sample_rate) - if key_str not in self.resample_kernel: - self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128) - self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device) - audios_res = self.resample_kernel[key_str](audios) - - mels = self.mel_extractor(audios_res, center=True) - hiddens = self.mel2hidden(mels) - f0 = self.decode(hiddens, thread=thread, use_viterbi=use_viterbi) - - f0s: list[np.ndarray] = [] - for i in range(f0.shape[0]): - f0s.append(f0[i, : sizes[i]] if sizes is not None else f0[i, :]) - return f0s - - def get_pitch_batch( - self, - waveforms, - sample_rate: int, - hop_size: int, - lengths: list[int], - interp_uv: bool = False, - fmin: float = 50, - fmax: float = 1000, - ) -> tuple[list[np.ndarray], list[np.ndarray]]: - f0s = self.infer_from_audio_batch(waveforms, sample_rate=sample_rate) - f0s_res: list[np.ndarray] = [] - uvs_res: list[np.ndarray] = [] - for idx, f0 in enumerate(f0s): - f0 = self.postprocess(f0, fmin, fmax, min_gap=6) - uv = f0 == 0 - length = lengths[idx] - time_step = hop_size / sample_rate - f0_res = resample_align_curve(f0, 0.01, time_step, length) - uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5 - if not interp_uv: - f0_res[uv_res] = 0 - f0s_res.append(f0_res) - uvs_res.append(uv_res) - return f0s_res, uvs_res - - def release_cuda(self): - self.model = self.model.cpu() - self.mel_extractor = self.mel_extractor.cpu() diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py deleted file mode 100644 index 0333b521e..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py +++ /dev/null @@ -1,619 +0,0 @@ -"""ROSVOT note transcription.""" - -from __future__ import annotations - -import math -from pathlib import Path -from typing import Any, ClassVar - -import numpy as np -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery -from vllm_omni.diffusion.models.soulx_singer.modules.mel_transform import spectral_normalize_torch -from vllm_omni.diffusion.models.soulx_singer.modules.note_transcription import ( - BackboneNet, - ConvBlocks, - Embedding, -) -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import ( - boundary2Interval, - denorm_f0, - f0_to_coarse, - load_model_ckpt, - load_mono_audio, - load_rosvot_config, - norm_interp_f0, - pad_or_cut_xd, - resample_mono, -) -from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn - -logger = init_logger(__name__) - - -def regulate_real_note_itv( - note_itv: np.ndarray, - note_bd: np.ndarray, - word_bd: np.ndarray, - word_durs: np.ndarray, - hop_size: int, - audio_sample_rate: int, -) -> tuple[np.ndarray, np.ndarray]: - assert note_itv.shape[0] == np.sum(note_bd) + 1 - assert np.sum(word_bd) <= np.sum(note_bd) - assert word_durs.shape[0] == np.sum(word_bd) + 1, f"{word_durs.shape[0]} {np.sum(word_bd) + 1}" - word_bd = np.cumsum(word_bd) * word_bd - word_itv = np.zeros((word_durs.shape[0], 2)) - word_offsets = np.cumsum(word_durs) - note2words = np.zeros(note_itv.shape[0], dtype=int) - for idx in range(len(word_offsets) - 1): - word_itv[idx, 1] = word_itv[idx + 1, 0] = word_offsets[idx] - word_itv[-1, 1] = word_offsets[-1] - note_itv_secs = note_itv * hop_size / audio_sample_rate - for idx, itv in enumerate(note_itv): - start_idx, end_idx = itv - if word_bd[start_idx] > 0: - word_dur_idx = word_bd[start_idx] - note_itv_secs[idx, 0] = word_itv[word_dur_idx, 0] - note2words[idx] = word_dur_idx - if word_bd[end_idx] > 0: - word_dur_idx = word_bd[end_idx] - 1 - note_itv_secs[idx, 1] = word_itv[word_dur_idx, 1] - note2words[idx] = word_dur_idx - note2words += 1 - return note_itv_secs, note2words - - -def regulate_ill_slur( - notes: np.ndarray, note_itv: np.ndarray, note2words: np.ndarray -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - res_note2words: list[int] = [] - res_note_itv: list[list[float]] = [] - res_notes: list[int] = [] - note_idx = 0 - note_idx_end = 0 - while note_idx <= len(notes) - 1: - while note_idx <= note_idx_end < len(notes) and note2words[note_idx] == note2words[note_idx_end]: - note_idx_end += 1 - res_note2words.append(note2words[note_idx]) - res_note_itv.append(note_itv[note_idx].tolist()) - res_notes.append(notes[note_idx]) - for idx in range(note_idx + 1, note_idx_end): - if notes[idx] == notes[idx - 1]: - res_note_itv[-1][1] = note_itv[idx][1] - else: - res_note_itv.append(note_itv[idx].tolist()) - res_note2words.append(note2words[idx]) - res_notes.append(notes[idx]) - note_idx = note_idx_end - return ( - np.array(res_notes, dtype=notes.dtype), - np.array(res_note_itv, dtype=note_itv.dtype), - np.array(res_note2words, dtype=note2words.dtype), - ) - - -def regulate_boundary(bd_logits, threshold, min_gap=18, ref_bd=None, ref_bd_min_gap=8, non_padding=None): - # this doesn't preserve gradient - device = bd_logits.device - bd_logits = torch.sigmoid(bd_logits).data.cpu() - # bd_logits[0] = bd_logits[-1] = 1e-5 # avoid itv invalid problem - bd = (bd_logits > threshold).long() - bd_res = torch.zeros_like(bd).long() - for i in range(bd.shape[0]): - bd_i = bd[i] - last_bd_idx = -1 - start = -1 - for j in range(bd_i.shape[0]): - if bd_i[j] == 1: - if 0 <= start < j: - continue - elif start < 0: - start = j - else: - if 0 <= start < j: - if j - 1 > start: - bd_idx = start + int(torch.argmax(bd_logits[i, start:j]).item()) - else: - bd_idx = start - if bd_idx - last_bd_idx < min_gap and last_bd_idx > 0: - bd_idx = round((bd_idx + last_bd_idx) / 2) - bd_res[i, last_bd_idx] = 0 - bd_res[i, bd_idx] = 1 - last_bd_idx = bd_idx - start = -1 - - # assert ref_bd_min_gap <= min_gap // 2 - if ref_bd is not None and ref_bd_min_gap > 0: - ref = ref_bd.data.cpu() - for i in range(bd_res.shape[0]): - ref_bd_i = ref[i] - ref_bd_i_js = [] - for j in range(ref_bd_i.shape[0]): - if ref_bd_i[j] == 1: - ref_bd_i_js.append(j) - seg_sum = torch.sum(bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap]) - if seg_sum == 0: - bd_res[i, j] = 1 - elif seg_sum == 1 and bd_res[i, j] != 1: - bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap] = ref_bd_i[ - max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap - ] - elif seg_sum > 1: - for k in range(1, ref_bd_min_gap + 1): - if bd_res[i, max(0, j - k)] == 1 and ref_bd_i[max(0, j - k)] != 1: - bd_res[i, max(0, j - k)] = 0 - break - if ( - bd_res[i, min(bd_res.shape[1] - 1, j + k)] == 1 - and ref_bd_i[min(bd_res.shape[1] - 1, j + k)] != 1 - ): - bd_res[i, min(bd_res.shape[1] - 1, j + k)] = 0 - break - bd_res[i, j] = 1 - # final check - assert torch.sum(bd_res[i, ref_bd_i_js]) == len(ref_bd_i_js), ( - f"{torch.sum(bd_res[i, ref_bd_i_js])} {len(ref_bd_i_js)}" - ) - - bd_res = bd_res.to(device) - - # force valid begin and end - bd_res[:, 0] = 0 - if non_padding is not None: - for i in range(bd_res.shape[0]): - bd_res[i, sum(non_padding[i]) - 1 :] = 0 - else: - bd_res[:, -1] = 0 - - return bd_res - - -def _align_word(word_durs: list[float], mel_len: int, hop_size: int, audio_sample_rate: int) -> np.ndarray: - mel2word = np.zeros([mel_len], int) - start_time = 0.0 - for i_word, wd in enumerate(word_durs): - start_frame = int(start_time * audio_sample_rate / hop_size + 0.5) - end_frame = int((start_time + wd) * audio_sample_rate / hop_size + 0.5) - mel2word[start_frame:end_frame] = i_word + 1 - start_time += wd - return mel2word - - -class PitchDecoder(nn.Module): - def __init__(self, hparams): - super().__init__() - self.hidden_size = hidden_size = hparams["hidden_size"] - self.dropout = hparams.get("dropout", 0.0) - self.note_bd_out = nn.Linear(hidden_size, 1) - self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0)) - - # note prediction - self.pitch_attn_num_head = hparams.get("pitch_attn_num_head", 1) - self.multihead_dot_attn = nn.Linear(hidden_size, self.pitch_attn_num_head) - self.post = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=3, - layers_in_block=1, - c_multiple=1, - post_net_kernel=3, - ) - self.pitch_out = nn.Linear(hidden_size, hparams.get("note_num", 100) + 4) - self.note_num = hparams.get("note_num", 100) - self.note_start = hparams.get("note_start", 30) - self.pitch_temperature = max(1e-7, hparams.get("note_pitch_temperature", 1.0)) - - def forward(self, feat, note_bd, train=True): - bsz, T, _ = feat.shape - - attn = torch.sigmoid(self.multihead_dot_attn(feat)) # [B, T, C] -> [B, T, num_head] - attn_feat = feat.unsqueeze(3) * attn.unsqueeze(2) # [B, T, C, 1] x [B, T, 1, num_head] -> [B, T, C, num_head] - attn_feat = torch.mean(attn_feat, dim=-1) # [B, T, C, num_head] -> [B, T, C] - mel2note = torch.cumsum(note_bd, 1) - note_length = torch.max(torch.sum(note_bd, dim=1)).item() + 1 # max length - note_lengths = torch.sum(note_bd, dim=1) + 1 # [B] - - attn = torch.mean(attn, dim=-1, keepdim=True) # [B, T, num_head] -> [B, T, 1] - denom = mel2note.new_zeros(bsz, note_length, dtype=attn.dtype).scatter_add_( - dim=1, index=mel2note, src=attn.squeeze(-1) - ) # [B, T] -> [B, note_length] count the note frames of each note (with padding excluded) - frame2note = mel2note.unsqueeze(-1).repeat(1, 1, self.hidden_size) # [B, T] -> [B, T, C], with padding included - note_aggregate = frame2note.new_zeros(bsz, note_length, self.hidden_size, dtype=attn_feat.dtype).scatter_add_( - dim=1, index=frame2note, src=attn_feat - ) # [B, T, C] -> [B, note_length, C] - note_aggregate = note_aggregate / (denom.unsqueeze(-1) + 1e-5) - note_logits = self.post(note_aggregate) - note_logits = self.pitch_out(note_logits) / self.pitch_temperature - # note_logits = torch.clamp(note_logits, min=-16., max=16.) # don't know need it or not - - note_pred = torch.softmax(note_logits, dim=-1) # [B, note_length, note_num] - note_pred = torch.argmax(note_pred, dim=-1) # [B, note_length] - # for some reason, note idx maybe 130 (why?) - note_pred[note_pred > self.note_num] = 0 - note_pred[note_pred < self.note_start] = 0 - - return note_lengths, note_logits, note_pred - - -class MidiExtractor(nn.Module): - def __init__(self, hparams): - super().__init__() - self.hparams = hparams.copy() - self.hidden_size = hidden_size = hparams["hidden_size"] - self.note_bd_threshold = hparams.get("note_bd_threshold", 0.5) - self.note_bd_min_gap = round( - hparams.get("note_bd_min_gap", 100) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"] - ) - self.note_bd_ref_min_gap = round( - hparams.get("note_bd_ref_min_gap", 50) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"] - ) - - self.mel_proj = nn.Conv1d(hparams["use_mel_bins"], hidden_size, kernel_size=3, padding=1) - self.mel_encoder = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=3, - layers_in_block=2, - c_multiple=1, - post_net_kernel=3, - ) - self.use_pitch = hparams.get("use_pitch_embed", True) - if self.use_pitch: - self.pitch_embed = Embedding(300, hidden_size, 0, "kaiming") - self.uv_embed = Embedding(3, hidden_size, 0, "kaiming") - self.use_wbd = hparams.get("use_wbd", True) - if self.use_wbd: - self.word_bd_embed = Embedding(3, hidden_size, 0, "kaiming") - self.cond_encoder = ConvBlocks( - hidden_size, - out_dims=hidden_size, - dilations=None, - kernel_size=3, - layers_in_block=1, - c_multiple=1, - post_net_kernel=3, - ) - - # backbone - self.net = BackboneNet(hparams) - - # note bd prediction - self.note_bd_out = nn.Linear(hidden_size, 1) - self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0)) - - # note prediction - self.pitch_decoder = PitchDecoder(hparams) - - def run_encoder(self, mel, word_bd=None, pitch=None, uv=None): - mel_embed = self.mel_proj(mel.transpose(1, 2)).transpose(1, 2) - mel_embed = self.mel_encoder(mel_embed) - pitch_embed = word_bd_embed = 0 - if self.use_pitch and pitch is not None and uv is not None: - pitch_embed = self.pitch_embed(pitch) + self.uv_embed(uv) # [B, T, C] - if self.use_wbd and word_bd is not None: - word_bd_embed = self.word_bd_embed(word_bd) - feat = self.cond_encoder(mel_embed + pitch_embed + word_bd_embed) - return feat - - def forward(self, mel, word_bd=None, note_bd=None, pitch=None, uv=None, non_padding=None, train=True): - ret = {} - bsz, T, _ = mel.shape - - feat = self.run_encoder(mel, word_bd, pitch, uv) - feat = self.net(feat) # [B, T, C] - - # note bd prediction - # dropout has been dropped (inference mode) - note_bd_logits = self.note_bd_out(feat).squeeze(-1) / self.note_bd_temperature - note_bd_logits = torch.clamp(note_bd_logits, min=-16.0, max=16.0) - ret["note_bd_logits"] = note_bd_logits # [B, T] - if note_bd is None or not train: - note_bd = regulate_boundary( - note_bd_logits, - self.note_bd_threshold, - self.note_bd_min_gap, - word_bd, - self.note_bd_ref_min_gap, - non_padding, - ) - ret["note_bd_pred"] = note_bd # [B, T] - - # note pitch prediction - note_lengths, note_logits, note_pred = self.pitch_decoder(feat, note_bd, train) - ret["note_lengths"], ret["note_logits"], ret["note_pred"] = note_lengths, note_logits, note_pred - - return ret - - -class MelNet(nn.Module): - def __init__(self, hparams, device="cpu") -> None: - super().__init__() - self.n_fft = hparams["fft_size"] - self.num_mels = hparams["audio_num_mel_bins"] - self.sampling_rate = hparams["audio_sample_rate"] - self.hop_size = hparams["hop_size"] - self.win_size = hparams["win_size"] - self.fmin = hparams["fmin"] - self.fmax = hparams["fmax"] - self.device = device - - mel = librosa_mel_fn( - sr=self.sampling_rate, - n_fft=self.n_fft, - n_mels=self.num_mels, - fmin=self.fmin, - fmax=self.fmax, - ) - self.mel_basis = mel.float().to(self.device) - self.hann_window = torch.hann_window(self.win_size).to(self.device) - - def to(self, device, **kwagrs): - super().to(device=device, **kwagrs) - self.mel_basis = self.mel_basis.to(device) - self.hann_window = self.hann_window.to(device) - self.device = device - - def forward(self, y, center=False, complex=False): - if isinstance(y, np.ndarray): - y = torch.FloatTensor(y) - if len(y.shape) == 1: - y = y.unsqueeze(0) - y = y.clamp(min=-1.0, max=1.0).to(self.device) - - pad_length = math.ceil(y.shape[1] / self.hop_size) * self.hop_size - y.shape[1] - y = torch.nn.functional.pad( - y.unsqueeze(1), - [int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2 + pad_length)], - mode="reflect", - ) - y = y.squeeze(1) - - spec = torch.stft( - y, - self.n_fft, - hop_length=self.hop_size, - win_length=self.win_size, - window=self.hann_window, - center=center, - pad_mode="reflect", - normalized=False, - onesided=True, - return_complex=True, - ) - if not complex: - spec = torch.view_as_real(spec) - spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) # [B, n_fft, T] - spec = torch.matmul(self.mel_basis, spec) - spec = spectral_normalize_torch(spec) - spec = spec.transpose(1, 2) # [B, T, n_fft] - else: - B, C, T, _ = spec.shape - spec = spec.transpose(1, 2) # [B, T, n_fft, 2] - return spec - - -class RosvotModel(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = ["."] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = [] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["midi.net", "midi.pitch_decoder", "midi.cond_encoder"] - - def __init__( - self, - rosvot_ckpt: str | Path, - *, - config_path: str | Path = "", - pe: nn.Module | None = None, - the: float = 0.85, - verbose: bool = False, - rosvot_source_dir: str | Path | None = None, - ): - super().__init__() - self.verbose = verbose - ckpt = Path(rosvot_ckpt) - resolved_config = Path(config_path) if config_path else ckpt.with_name("config.yaml") - self.hparams = load_rosvot_config(resolved_config, hparams_str=f"note_bd_threshold={the}") - if verbose: - logger.info("ROSVOT config: %s", resolved_config) - - self.midi = MidiExtractor(self.hparams) - self.mel_net = MelNet(self.hparams) - self.pe = pe if pe is not None and self.hparams.get("use_pitch_embed", False) else None - self._checkpoint_path = str(ckpt) - self.load_checkpoint(str(ckpt), verbose=verbose) - self.eval() - - def load_checkpoint(self, checkpoint_path: str | None = None, *, verbose: bool = False) -> None: - load_model_ckpt(self.midi, checkpoint_path or self._checkpoint_path, verbose=verbose) - - @torch.no_grad() - def forward(self, wav: torch.Tensor, word_durs: list[float]) -> dict[str, Any]: - hparams = self.hparams - if wav.ndim == 1: - wav = wav.unsqueeze(0) - - mel_len = (wav.shape[-1] + hparams["hop_size"] - 1) // hparams["hop_size"] - min_word_dur = hparams.get("min_word_dur", 20) / 1000 - wd_raw = list(word_durs) - word_durs_filtered: list[float] = [] - for i, wd in enumerate(wd_raw): - if wd < min_word_dur: - if i == 0 and len(wd_raw) > 1: - wd_raw[i + 1] += wd - elif word_durs_filtered: - word_durs_filtered[-1] += wd - else: - word_durs_filtered.append(wd) - - mel2word = _align_word(word_durs_filtered, mel_len, hparams["hop_size"], hparams["audio_sample_rate"]) - if mel2word.size > 0 and mel2word[0] == 0: - mel2word = mel2word + 1 - - real_len = min(mel_len, int(np.sum(mel2word > 0))) - T = math.ceil(min(real_len, hparams["max_frames"]) / hparams["frames_multiple"]) * hparams["frames_multiple"] - device = wav.device - self.mel_net.to(device) - wav_t = pad_or_cut_xd(wav.float(), T * hparams["hop_size"], 1) - - pitch_coarse = uv_t = None - if self.pe is not None: - f0s, uvs = self.pe.get_pitch_batch( - wav_t, - sample_rate=hparams["audio_sample_rate"], - hop_size=hparams["hop_size"], - lengths=[real_len], - fmax=hparams["f0_max"], - fmin=hparams["f0_min"], - ) - f0_1d, uv_1d = norm_interp_f0(f0s[0][:T]) - f0_t = pad_or_cut_xd(torch.as_tensor(f0_1d, device=device, dtype=torch.float32), T, 0).unsqueeze(0) - uv_t = pad_or_cut_xd(torch.as_tensor(uv_1d, device=device, dtype=torch.float32), T, 0).long().unsqueeze(0) - pitch_coarse = f0_to_coarse(denorm_f0(f0_t, uv_t)).to(device) - - mel = pad_or_cut_xd(self.mel_net(wav_t)[0], T, dim=0).unsqueeze(0) - mel_nonpadding_mask = torch.zeros(1, T, device=device) - mel_nonpadding_mask[:, :real_len] = 1.0 - mel = (mel.transpose(1, 2) * mel_nonpadding_mask.unsqueeze(1)).transpose(1, 2) - mel_nonpadding = mel.abs().sum(-1) > 0 - - mel2word_t = pad_or_cut_xd(torch.as_tensor(mel2word, device=device, dtype=torch.long), T, 0) - word_bd = torch.zeros_like(mel2word_t) - word_bd[1:] = (mel2word_t[1:] != mel2word_t[:-1]).long() - word_bd[real_len:] = 0 - word_bd = word_bd.unsqueeze(0) - - outputs = self.midi( - mel=mel[:, :, : hparams.get("use_mel_bins", 80)], - word_bd=word_bd, - pitch=pitch_coarse, - uv=uv_t, - non_padding=mel_nonpadding, - ) - outputs["word_durs_filtered"] = word_durs_filtered - outputs["real_len"] = real_len - outputs["word_bd"] = word_bd - return outputs - - @staticmethod - def _load_wav(wav_src: str | np.ndarray, sample_rate: int, *, src_sample_rate: int | None = None) -> np.ndarray: - if isinstance(wav_src, str): - wav, _ = load_mono_audio(wav_src, target_sr=sample_rate) - return wav - wav = np.asarray(wav_src, dtype=np.float32) - if src_sample_rate is not None and src_sample_rate != sample_rate: - wav = resample_mono(wav, orig_sr=src_sample_rate, target_sr=sample_rate) - return wav - - @staticmethod - def _normalize_note2words(note2words: list[int]) -> list[int]: - if not note2words: - return [] - out = [note2words[0]] - for idx in range(1, len(note2words)): - out.append(max(note2words[idx], out[-1])) - return out - - @staticmethod - def _build_ep_types(note2words: list[int], align_words: list[str]) -> list[int]: - ep_types: list[int] = [] - prev = -1 - for i, w in zip(note2words, align_words): - ep_types.append(1 if w == "" else (2 if i != prev else 3)) - prev = i - return ep_types - - @torch.no_grad() - def transcribe( - self, - item: dict[str, Any], - *, - segment_info: dict[str, Any] | None = None, - verbose: bool = False, - ) -> dict[str, Any]: - if "word_durs" not in item: - raise ValueError('item must contain "word_durs" from lyric transcription') - - if item.get("wav_fn"): - wav = self._load_wav(item["wav_fn"], self.hparams["audio_sample_rate"]) - elif item.get("wav") is not None: - wav = self._load_wav( - item["wav"], - self.hparams["audio_sample_rate"], - src_sample_rate=item.get("sample_rate"), - ) - else: - raise ValueError('item must contain "wav_fn" or "wav"') - - device = next(self.parameters()).device - outputs = self(torch.from_numpy(wav).float().to(device), list(item["word_durs"])) - real_len = int(outputs["real_len"]) - word_durs_filtered = outputs["word_durs_filtered"] - word_bd = outputs["word_bd"] - item_name = item.get("item_name", "") - - note_lengths = outputs["note_lengths"].detach().cpu().numpy() - note_bd_pred = outputs["note_bd_pred"][0].detach().cpu().numpy()[:real_len] - note_pred = outputs["note_pred"][0].detach().cpu().numpy()[: note_lengths[0]] - - if note_pred.shape == (0,): - rosvot_out = {"item_name": item_name, "pitches": [], "note_durs": [], "note2words": None} - else: - note_itv_pred = boundary2Interval(note_bd_pred) - word_bd_for_reg = word_bd[0].detach().cpu().numpy()[:real_len] - hop = self.hparams["hop_size"] - sr = self.hparams["audio_sample_rate"] - - if self.hparams.get("infer_regulate_real_note_itv", True): - try: - note_itv_pred_secs, note2words = regulate_real_note_itv( - note_itv_pred, - note_bd_pred, - word_bd_for_reg, - np.array(word_durs_filtered), - hop, - sr, - ) - note_pred, note_itv_pred_secs, note2words = regulate_ill_slur( - note_pred, note_itv_pred_secs, note2words - ) - except Exception: - if verbose: - logger.exception("ROSVOT postprocess failed") - note_itv_pred_secs = note_itv_pred * hop / sr - note2words = None - else: - note_itv_pred_secs = note_itv_pred * hop / sr - note2words = None - - rosvot_out = { - "item_name": item_name, - "pitches": note_pred.tolist(), - "note_durs": [float(itv[1] - itv[0]) for itv in note_itv_pred_secs], - "note2words": note2words.tolist() if note2words is not None else None, - } - - note2words_raw = rosvot_out.get("note2words") or [] - align_words = [item["words"][idx - 1] for idx in note2words_raw if 0 < idx <= len(item["words"])] - ep_types = self._build_ep_types(self._normalize_note2words(note2words_raw), align_words) if align_words else [] - seg = segment_info or item - - return { - "item_name": seg.get("item_name", item_name), - "wav_fn": seg.get("wav_fn", item.get("wav_fn", "")), - "origin_wav_fn": seg.get("origin_wav_fn", item.get("origin_wav_fn", "")), - "start_time_ms": seg.get("start_time_ms", ""), - "end_time_ms": seg.get("end_time_ms", ""), - "language": item.get("language", ""), - "note_text": align_words, - "note_dur": rosvot_out.get("note_durs", []), - "note_type": ep_types, - "note_pitch": rosvot_out.get("pitches", []), - } diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py deleted file mode 100644 index 0140182ab..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py +++ /dev/null @@ -1,267 +0,0 @@ -"""Rule-based vocal segmentation from F0 contour (no neural net).""" - -from __future__ import annotations - -import time -from dataclasses import dataclass -from typing import ClassVar - -import numpy as np -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery - -logger = init_logger(__name__) - - -@dataclass(frozen=True) -class VocalSegmentConfig: - hop_ms: int = 20 - smooth_ms: int = 200 - start_ms: int = 120 - end_ms: int = 200 - prepad_ms: int = 80 - postpad_ms: int = 120 - min_len_ms: int = 1000 - max_len_ms: int = 20000 - short_seg_merge_gap_ms: int = 8000 - small_gap_ms: int = 500 - lookback_ms: int = 200 - lookahead_ms: int = 100 - - -def _moving_average(x: np.ndarray, win: int) -> np.ndarray: - if win <= 1: - return x - kernel = np.ones(win, dtype=np.float32) / float(win) - return np.convolve(x, kernel, mode="same") - - -def _merge_short_segments( - segments_ms: list[list[int]], - *, - min_len_ms: int, - max_len_ms: int, - short_seg_merge_gap_ms: int, - small_gap_ms: int, -) -> list[list[int]]: - if not segments_ms: - return [] - - merged: list[list[int]] = [] - cur_start, cur_end = segments_ms[0] - for next_start, next_end in segments_ms[1:]: - cur_len = cur_end - cur_start - gap_ms = next_start - cur_end - merged_len = next_end - cur_start - should_merge = (cur_len < min_len_ms and gap_ms < short_seg_merge_gap_ms) or ( - gap_ms < small_gap_ms and merged_len < max_len_ms - ) - if should_merge: - cur_end = next_end - continue - if (cur_end - cur_start) >= min_len_ms: - merged.append([cur_start, cur_end]) - cur_start, cur_end = next_start, next_end - - if (cur_end - cur_start) >= min_len_ms: - merged.append([cur_start, cur_end]) - - if not merged: - return segments_ms - - return merged - - -def _voiced_to_segments( - voiced: np.ndarray, - *, - hop_ms: int, - smooth_ms: int, - start_ms: int, - end_ms: int, - prepad_ms: int, - postpad_ms: int, - max_len_ms: int, -) -> list[list[int]]: - smooth_frames = max(1, int(round(smooth_ms / hop_ms))) - active = _moving_average(voiced.astype(np.float32), smooth_frames) >= 0.5 - - segments: list[list[int]] = [] - start_idx = None - start_frames = max(1, int(round(start_ms / hop_ms))) - end_frames = max(1, int(round(end_ms / hop_ms))) - prepad_frames = max(0, int(round(prepad_ms / hop_ms))) - postpad_frames = max(0, int(round(postpad_ms / hop_ms))) - active_count = 0 - inactive_count = 0 - - for i, flag in enumerate(active): - if flag: - active_count += 1 - inactive_count = 0 - else: - inactive_count += 1 - active_count = 0 - - if start_idx is None: - if active_count >= start_frames: - start_idx = max(0, i - start_frames + 1 - prepad_frames) - else: - if inactive_count >= end_frames: - end_idx = min(len(active) - 1, i - end_frames + 1 + postpad_frames) - start_ms_val = start_idx * hop_ms - end_ms_val = end_idx * hop_ms + hop_ms - if end_ms_val > start_ms_val: - segments.append([int(start_ms_val), int(end_ms_val)]) - start_idx = None - - if start_idx is not None: - start_ms_val = start_idx * hop_ms - end_idx = min(len(active) - 1, len(active) - 1 + postpad_frames) - end_ms_val = end_idx * hop_ms + hop_ms - if end_ms_val > start_ms_val: - segments.append([int(start_ms_val), int(end_ms_val)]) - - def _split_segment(seg: list[int]) -> list[list[int]]: - start_ms_val, end_ms_val = seg - start_frame = int(start_ms_val // hop_ms) - end_frame = int((end_ms_val - 1) // hop_ms) - end_frame = max(start_frame, min(end_frame, len(active) - 1)) - - best_start = None - best_len = 0 - cur_start = None - cur_len = 0 - for idx in range(start_frame, end_frame + 1): - if not active[idx]: - if cur_start is None: - cur_start = idx - cur_len = 1 - else: - cur_len += 1 - else: - if cur_start is not None and cur_len > best_len: - best_start, best_len = cur_start, cur_len - cur_start = None - cur_len = 0 - if cur_start is not None and cur_len > best_len: - best_start, best_len = cur_start, cur_len - - if best_start is None: - split_frame = (start_frame + end_frame) // 2 - else: - split_frame = best_start + best_len // 2 - - split_ms = split_frame * hop_ms - if split_ms <= start_ms_val: - split_ms = start_ms_val + hop_ms - if split_ms >= end_ms_val: - split_ms = end_ms_val - hop_ms - if split_ms <= start_ms_val or split_ms >= end_ms_val: - return [seg] - return [[start_ms_val, int(split_ms)], [int(split_ms), end_ms_val]] - - queue = segments[:] - segments = [] - while queue: - seg = queue.pop(0) - if (seg[1] - seg[0]) <= max_len_ms: - segments.append(seg) - continue - parts = _split_segment(seg) - if len(parts) == 1: - segments.append(seg) - else: - queue = parts + queue - return segments - - -class VocalSegmenter(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = [] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["."] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = [] - - def __init__(self, config: VocalSegmentConfig | None = None, *, verbose: bool = False): - super().__init__() - self.config = config or VocalSegmentConfig() - self.verbose = verbose - - def forward( - self, - audio: np.ndarray, - sample_rate: int, - f0: np.ndarray, - *, - base_name: str = "vocal", - origin_wav_fn: str = "", - verbose: bool | None = None, - ) -> list[dict]: - verbose = self.verbose if verbose is None else verbose - if verbose: - logger.info("[vocal detection] segmenting %s", base_name) - t0 = time.time() - - cfg = self.config - y = np.asarray(audio, dtype=np.float32).squeeze() - sr = int(sample_rate) - - segments_ms = _voiced_to_segments( - f0 > 0, - hop_ms=cfg.hop_ms, - smooth_ms=cfg.smooth_ms, - start_ms=cfg.start_ms, - end_ms=cfg.end_ms, - prepad_ms=cfg.prepad_ms, - postpad_ms=cfg.postpad_ms, - max_len_ms=cfg.max_len_ms, - ) - segments_ms = _merge_short_segments( - segments_ms, - min_len_ms=cfg.min_len_ms, - max_len_ms=cfg.max_len_ms, - short_seg_merge_gap_ms=cfg.short_seg_merge_gap_ms, - small_gap_ms=cfg.small_gap_ms, - ) - - adjusted_segments: list[list[int]] = [] - prev_end = 0 - for start_ms, end_ms in segments_ms: - start_ms = max(0, start_ms - cfg.lookback_ms) - end_ms = min(end_ms + cfg.lookahead_ms, int(y.shape[0] / sr * 1000)) - if start_ms < prev_end and adjusted_segments: - adjusted_segments[-1][1] = start_ms - adjusted_segments.append([start_ms, end_ms]) - prev_end = end_ms - - segment_infos = [] - for idx, (start_ms, end_ms) in enumerate(adjusted_segments): - if end_ms - start_ms > cfg.max_len_ms: - start_ms = end_ms - cfg.max_len_ms - key = f"{base_name}_{idx}" - start_sample = int(start_ms * sr / 1000) - end_sample = int(end_ms * sr / 1000) - segment_infos.append( - { - "item_name": key, - "wav": y[start_sample:end_sample], - "sample_rate": sr, - "wav_fn": key, - "start_time_ms": int(start_sample * 1000 / sr), - "end_time_ms": int(end_sample * 1000 / sr), - "origin_wav_fn": origin_wav_fn, - "duration": int((end_sample - start_sample) * 1000 / sr), - } - ) - - if verbose: - logger.info( - "[vocal detection] done n_segments=%s time=%.3fs", - len(segment_infos), - time.time() - t0, - ) - return segment_infos diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py deleted file mode 100644 index e29b52a29..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Lazy-loaded SoulX preprocess model tree.""" - -import os -import tempfile -from pathlib import Path -from typing import ClassVar - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery - -from .asr import LyricModel -from .mel_grid_f0 import extract_f0_file -from .rmvpe import RMVPE -from .rosvot import RosvotModel -from .segmenter import VocalSegmenter -from .utils import resample_mono -from .vocal_sep import MelBandRoformerSeparator - -logger = init_logger(__name__) - - -class SoulXPreprocessStack(nn.Module, SupportsComponentDiscovery): - """Unified nn.Module tree for SoulX-Singer audio preprocess.""" - - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = ["vocal_sep", "lyric", "rosvot"] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = ["rmvpe"] - - def __init__( - self, - weights: dict[str, str], - device: str, - *, - target_sr: int = 24000, - hop_size: int = 480, - verbose: bool = False, - ) -> None: - super().__init__() - self.weights = weights - self.device_str = device - self.target_sr = target_sr - self.hop_size = hop_size - self.verbose = verbose - self._rmvpe: nn.Module | None = None - self._vocal_sep: nn.Module | None = None - self._segmenter: nn.Module | None = None - self._lyric: nn.Module | None = None - self._rosvot: nn.Module | None = None - - @property - def device(self) -> torch.device: - return torch.device(self.device_str) - - def ensure_rmvpe(self): - if self._rmvpe is None: - if self.verbose: - logger.info("[rmvpe] loading shared pitch model") - self._rmvpe = RMVPE(self.weights["rmvpe"], device=self.device_str) - self.add_module("rmvpe", self._rmvpe) - return self._rmvpe - - def ensure_vocal_sep(self): - if self._vocal_sep is None: - if self.verbose: - logger.info("[vocal extraction] loading model") - self._vocal_sep = MelBandRoformerSeparator( - self.weights["sep_config"], - self.weights["sep_ckpt"], - use_half=False, - ).to(self.device) - self.add_module("vocal_sep", self._vocal_sep) - if self.verbose: - logger.info("[vocal extraction] ready on %s, sr=%s", self.device_str, self._vocal_sep.sample_rate) - return self._vocal_sep - - def ensure_segmenter(self): - if self._segmenter is None: - self._segmenter = VocalSegmenter(verbose=self.verbose) - self.add_module("segmenter", self._segmenter) - return self._segmenter - - def ensure_lyric(self): - if self._lyric is None: - if self.verbose: - logger.info("[lyric transcription] loading ASR") - self._lyric = LyricModel( - self.weights["asr_zh"], - self.weights["asr_en"], - device=self.device_str, - target_sr=self.target_sr, - hop_size=self.hop_size, - ) - self.add_module("lyric", self._lyric) - return self._lyric - - def ensure_rosvot(self): - if self._rosvot is None: - if self.verbose: - logger.info("[note transcription] loading ROSVOT") - rosvot_src = os.environ.get("ROSVOT_SOURCE_DIR") - self._rosvot = RosvotModel( - self.weights["rosvot"], - pe=self.ensure_rmvpe(), - verbose=self.verbose, - rosvot_source_dir=rosvot_src, - ).to(self.device) - self.add_module("rosvot", self._rosvot) - return self._rosvot - - def extract_f0(self, vocal, sample_rate: int, *, f0_path: str | None = None) -> np.ndarray: - """Extract F0 on the mel grid; ndarray inputs are spooled to a temp wav.""" - rmvpe = self.ensure_rmvpe() - kwargs = dict( - target_sr=self.target_sr, - hop_size=self.hop_size, - f0_path=f0_path, - verbose=self.verbose, - ) - if isinstance(vocal, str): - return extract_f0_file(rmvpe, vocal, **kwargs) - tmp_path: str | None = None - try: - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - tmp_path = tmp.name - sf.write(tmp_path, np.asarray(vocal, dtype=np.float32).squeeze(), int(sample_rate)) - return extract_f0_file(rmvpe, tmp_path, **kwargs) - finally: - if tmp_path is not None: - Path(tmp_path).unlink(missing_ok=True) - - def extract_vocal(self, audio: str | tuple[np.ndarray, int]) -> tuple[np.ndarray, int]: - sep = self.ensure_vocal_sep() - if isinstance(audio, tuple): - mix, sr = audio - mix = np.asarray(mix, dtype=np.float32).squeeze() - if sr != sep.sample_rate: - mix = resample_mono(mix, orig_sr=sr, target_sr=sep.sample_rate) - return sep.separate_mono(mix, self.device), sep.sample_rate - return sep.separate_path(str(audio), self.device) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py deleted file mode 100644 index f135fd553..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py +++ /dev/null @@ -1,260 +0,0 @@ -"""Shared helpers for SoulX preprocess (I/O, config, checkpoints, pitch).""" - -from pathlib import Path -from typing import Any - -import numpy as np -import soundfile as sf -import torch -import torch.nn as nn -import torch.nn.functional as F -from omegaconf import OmegaConf -from vllm.logger import init_logger -from vllm.multimodal.audio import AudioResampler -from vllm.multimodal.media.audio import load_audio - -logger = init_logger(__name__) - - -# --- audio I/O --- - - -def resample_mono( - audio: np.ndarray, - *, - orig_sr: int, - target_sr: int, -) -> np.ndarray: - """Resample mono float waveform to ``target_sr``.""" - audio = np.asarray(audio, dtype=np.float32).squeeze() - if orig_sr == target_sr: - return audio - resampler = AudioResampler(target_sr=float(target_sr)) - return np.asarray( - resampler.resample(audio, orig_sr=float(orig_sr)), - dtype=np.float32, - ) - - -def get_audio_duration(path: str) -> float: - """Return audio duration in seconds for a file on disk.""" - return float(sf.info(path).duration) - - -def load_mono_audio( - source: str | tuple[np.ndarray, int], - *, - target_sr: int | None = None, -) -> tuple[np.ndarray, int]: - if isinstance(source, tuple): - audio, sr = source - audio = np.asarray(audio, dtype=np.float32).squeeze() - if target_sr is not None and sr != target_sr: - audio = resample_mono(audio, orig_sr=sr, target_sr=target_sr) - sr = target_sr - return audio, sr - - if target_sr is None: - audio, sr = load_audio(str(source), sr=None, mono=True) - else: - audio, sr = load_audio(str(source), sr=target_sr, mono=True) - return np.asarray(audio, dtype=np.float32).squeeze(), int(sr) - - -# --- tensor padding --- - - -def pad_or_cut_xd(values: torch.Tensor, tgt_len: int, dim: int = -1, pad_value=0) -> torch.Tensor: - ndim = values.ndim - if ndim == 1: - if values.shape[0] < tgt_len: - return F.pad(values, [0, tgt_len - values.shape[0]], value=pad_value) - return values[:tgt_len] - if ndim == 2: - if dim in (0, -2): - src_len = values.shape[0] - if src_len < tgt_len: - return F.pad(values, [0, 0, 0, tgt_len - src_len], value=pad_value) - return values[:tgt_len] - src_len = values.shape[1] - if src_len < tgt_len: - return F.pad(values, [0, tgt_len - src_len], value=pad_value) - return values[:, :tgt_len] - raise NotImplementedError(f"pad_or_cut_xd supports 1D/2D tensors, got ndim={ndim}") - - -# --- YAML config (OmegaConf) --- - - -def _apply_hparams_str(config: dict[str, Any], hparams_str: str) -> None: - if not hparams_str: - return - for item in hparams_str.split(","): - key, raw_value = item.split("=", maxsplit=1) - raw_value = raw_value.strip().strip("'\"") - node: dict[str, Any] = config - parts = key.split(".") - for part in parts[:-1]: - node = node.setdefault(part, {}) - leaf = parts[-1] - current = node.get(leaf) - if raw_value in {"True", "False"}: - node[leaf] = raw_value == "True" - elif isinstance(current, bool): - node[leaf] = raw_value.lower() == "true" - elif isinstance(current, int): - node[leaf] = int(raw_value) - elif isinstance(current, float): - node[leaf] = float(raw_value) - else: - node[leaf] = type(current)(raw_value) if current is not None else raw_value - - -def load_yaml_config( - config: str | Path, - *, - hparams_str: str = "", -) -> dict[str, Any]: - """Load YAML with optional ``base_config`` chain; returns a plain dict.""" - config_path = Path(config).expanduser() - if not config_path.is_file(): - raise FileNotFoundError(f"Config not found: {config_path}") - - merged = OmegaConf.create({}) - pending = [config_path.resolve()] - loaded: set[str] = set() - - while pending: - path = pending.pop(0) - path_key = str(path) - if path_key in loaded: - continue - loaded.add(path_key) - - current = OmegaConf.load(path_key) - base_configs = current.get("base_config") - if base_configs is not None: - entries = base_configs if OmegaConf.is_list(base_configs) else [base_configs] - for entry in entries: - base_path = Path(entry) - if not base_path.is_absolute(): - base_path = (path.parent / entry).resolve() - if not base_path.is_file(): - logger.warning("Skipping missing base config: %s", base_path) - continue - pending.insert(0, base_path) - merged = OmegaConf.merge(current, merged) - - result = OmegaConf.to_container(merged, resolve=True) - if not isinstance(result, dict): - raise TypeError(f"Expected dict config, got {type(result)}") - _apply_hparams_str(result, hparams_str) - return result - - -def load_rosvot_config(config: str | Path, *, hparams_str: str = "") -> dict[str, Any]: - return load_yaml_config(config, hparams_str=hparams_str) - - -# --- checkpoint loading --- - - -def load_model_ckpt( - model: nn.Module, - checkpoint_path: str, - *, - model_name: str = "model", - verbose: bool = True, -) -> None: - checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) - if not isinstance(checkpoint, dict) or "state_dict" not in checkpoint: - raise ValueError(f"Unexpected ROSVOT checkpoint format: {checkpoint_path}") - - state_dict = checkpoint["state_dict"] - if state_dict and "." in next(iter(state_dict)): - prefix = f"{model_name}." - state_dict = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} - elif model_name in state_dict and isinstance(state_dict[model_name], dict): - state_dict = state_dict[model_name] - - model.load_state_dict(state_dict, strict=True) - if verbose: - logger.info("Loaded ROSVOT weights from %s", checkpoint_path) - - -# --- pitch helpers (ROSVOT / RMVPE) --- - - -def f0_to_coarse(f0, f0_bin=256, f0_max=900.0, f0_min=50.0): - f0_mel_min = 1127 * np.log(1 + f0_min / 700) - f0_mel_max = 1127 * np.log(1 + f0_max / 700) - is_torch = isinstance(f0, torch.Tensor) - f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700) - f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1 - f0_mel[f0_mel <= 1] = 1 - f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1 - return (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(int) - - -def norm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100): - is_torch = isinstance(f0, torch.Tensor) - if pitch_norm == "standard": - f0 = (f0 - f0_mean) / f0_std - if pitch_norm == "log": - f0 = torch.log2(f0 + 1e-8) if is_torch else np.log2(f0 + 1e-8) - if uv is not None: - f0[uv > 0] = 0 - return f0 - - -def norm_interp_f0(f0, pitch_norm="log", f0_mean=None, f0_std=None): - is_torch = isinstance(f0, torch.Tensor) - if is_torch: - device = f0.device - f0 = f0.detach().cpu().numpy() - uv = f0 == 0 - f0 = norm_f0(f0, uv, pitch_norm, f0_mean, f0_std) - if sum(uv) == len(f0): - f0[uv] = 0 - elif sum(uv) > 0: - f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv]) - if is_torch: - uv_t = torch.FloatTensor(uv) - f0_t = torch.FloatTensor(f0) - return f0_t.to(device), uv_t.to(device) - return f0, uv - - -def denorm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100, pitch_padding=None, min=50, max=900): - is_torch = isinstance(f0, torch.Tensor) - if pitch_norm == "standard": - f0 = f0 * f0_std + f0_mean - if pitch_norm == "log": - f0 = 2**f0 - f0 = f0.clamp(min=min, max=max) if is_torch else np.clip(f0, a_min=min, a_max=max) - if uv is not None: - f0[uv > 0] = 0 - if pitch_padding is not None: - f0[pitch_padding] = 0 - return f0 - - -def boundary2Interval(bd): - is_torch = isinstance(bd, torch.Tensor) - if is_torch: - device = bd.device - bd = bd.detach().cpu().numpy() - ret = np.zeros(shape=(bd.sum() + 1, 2), dtype=int) - ret_idx = 0 - ret[0, 0] = 0 - for i, u in enumerate(bd): - if i == 0: - continue - if u == 1: - ret[ret_idx, 1] = i - ret[ret_idx + 1, 0] = i - ret_idx += 1 - ret[-1, 1] = bd.shape[0] - 1 - if is_torch: - return torch.LongTensor(ret).to(device) - return ret diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py deleted file mode 100644 index b62902d38..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py +++ /dev/null @@ -1,207 +0,0 @@ -"""BS-RoFormer vocal separation adapter (external pip package + SoulX chunking).""" - -from __future__ import annotations - -from pathlib import Path -from typing import ClassVar - -import numpy as np -import torch -import torch.nn as nn -import yaml -from omegaconf import DictConfig, OmegaConf -from vllm.logger import init_logger -from vllm.multimodal.media.audio import load_audio - -from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery -from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend - -logger = init_logger(__name__) - - -class MelBandRoformerSeparator(nn.Module, SupportAudioInput, SupportsComponentDiscovery): - support_audio_input: ClassVar[bool] = True - _dit_modules: ClassVar[list[str]] = [] - _encoder_modules: ClassVar[list[str]] = ["roformer"] - _vae_modules: ClassVar[list[str]] = [] - _resident_modules: ClassVar[list[str]] = [] - _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["roformer.layers"] - - def __init__( - self, - config_path: str | Path, - checkpoint_path: str | Path | None = None, - *, - chunk_length_sec: float = 5, - use_half: bool = True, - ): - super().__init__() - try: - from bs_roformer import MelBandRoformer as ExtModel - except ImportError as e: - raise ImportError( - 'Install BS-RoFormer: pip install "vllm-omni[soulx-preprocess]" or pip install BS-RoFormer' - ) from e - - cfg = self._load_config(Path(config_path)) - model_cfg = self._normalize_cfg(OmegaConf.to_container(cfg.model, resolve=True)) - self.roformer = ExtModel(**model_cfg) - self.config = cfg - self.sample_rate = int(cfg.audio.sample_rate) - - LayerWiseOffloadBackend.set_blocks_attr_names(self.roformer, ["layers"]) - dtype = torch.float16 if use_half else torch.float32 - self.roformer = self.roformer.to(dtype=dtype) - - chunk = int(chunk_length_sec * self.sample_rate) - self.config.inference = getattr(self.config, "inference", {}) - inf = self.config.inference - inf.chunk_size = chunk - inf.num_overlap = inf.get("num_overlap", 4) - inf.use_amp = False - inf.normalize = inf.get("normalize", True) - - fade = chunk // 10 - w = torch.ones(chunk) - w[:fade] = torch.linspace(0, 1, fade) - w[-fade:] = torch.linspace(1, 0, fade) - self.register_buffer("_window", w.to(dtype=dtype), persistent=False) - - if checkpoint_path: - self.load_checkpoint(str(checkpoint_path), use_half=use_half) - else: - self.roformer.eval() - - @staticmethod - def _normalize_cfg(cfg: dict) -> dict: - tuple_keys = {"multi_stft_resolutions_window_sizes", "freqs_per_bands", "multi_stft_window_sizes"} - - def _fix(v): - if isinstance(v, list): - return tuple(_fix(x) for x in v) - if isinstance(v, dict): - return {k: _fix(val) for k, val in v.items()} - return v - - for k in tuple_keys: - if k in cfg: - cfg[k] = _fix(cfg[k]) - return cfg - - @staticmethod - def _load_config(path: Path) -> DictConfig: - class Loader(yaml.SafeLoader): - pass - - Loader.add_constructor( - "tag:yaml.org,2002:python/tuple", - lambda loader, node: tuple(loader.construct_sequence(node)), - ) - with open(path) as f: - return OmegaConf.create(yaml.load(f, Loader=Loader)) - - def load_checkpoint(self, path: str, *, use_half: bool = True) -> None: - self.roformer.load_state_dict(torch.load(path, map_location="cpu"), strict=False) - if use_half: - self.roformer.half() - self.roformer.eval() - - @property - def chunk_size(self) -> int: - return int(self.config.inference.chunk_size) - - @property - def num_overlap(self) -> int: - return int(self.config.inference.num_overlap) - - @property - def use_amp(self) -> bool: - return bool(self.config.inference.use_amp) - - @property - def normalize_input(self) -> bool: - return bool(self.config.inference.normalize) - - @property - def num_channels(self) -> int: - return int(self.config.audio.get("num_channels", 1)) - - def _prepare_mix(self, mix: np.ndarray) -> tuple[np.ndarray, tuple[float, float] | None]: - if mix.ndim == 1: - mix = mix[np.newaxis, :] - if self.num_channels == 2 and mix.shape[0] == 1: - mix = np.repeat(mix, 2, axis=0) - if not self.normalize_input: - return mix, None - mono = mix.mean(0) - return (mix - mono.mean()) / mono.std(), (float(mono.mean()), float(mono.std())) - - @staticmethod - def _select_vocal(recon: torch.Tensor) -> torch.Tensor: - if recon.ndim == 4: - return recon[0, 0] - if recon.ndim == 3: - return recon[0] - if recon.ndim == 2: - return recon - raise ValueError(recon.shape) - - @torch.inference_mode() - def forward(self, mix: torch.Tensor) -> torch.Tensor: - if mix.ndim == 1: - mix = mix.unsqueeze(0) - if self.num_channels == 2 and mix.shape[0] == 1: - mix = mix.repeat(2, 1) - - model_dtype = next(self.roformer.parameters()).dtype - compute_dtype = torch.float32 if model_dtype == torch.float16 else model_dtype - if mix.dtype != compute_dtype: - mix = mix.to(dtype=compute_dtype) - - chunk_size = self.chunk_size - step = chunk_size // self.num_overlap - fade = chunk_size // 10 - border = chunk_size - step - n_samples = mix.shape[-1] - window = self._window.to(mix.device) - - if n_samples > 2 * border and border > 0: - mix = nn.functional.pad(mix, (border, border), mode="reflect") - - out = torch.zeros_like(mix) - weight = torch.zeros(mix.shape[-1], device=mix.device, dtype=mix.dtype) - pos = 0 - with torch.cuda.amp.autocast(enabled=self.use_amp): - while pos < mix.shape[-1]: - part = mix[:, pos : pos + chunk_size] - seg_len = part.shape[-1] - pad_mode = "reflect" if seg_len > chunk_size // 2 else "constant" - part = nn.functional.pad(part, (0, chunk_size - seg_len), mode=pad_mode, value=0) - stem = self._select_vocal(self.roformer(part.unsqueeze(0))) - win = window.clone() - if pos == 0: - win[:fade] = 1 - elif pos + step >= n_samples: - win[-fade:] = 1 - out[:, pos : pos + seg_len] += stem[:, :seg_len] * win[:seg_len] - weight[pos : pos + seg_len] += win[:seg_len] - pos += step - - out = out / weight - if n_samples > 2 * border and border > 0: - out = out[..., border:-border] - return out - - @torch.inference_mode() - def separate_mono(self, mix: np.ndarray, device: torch.device) -> np.ndarray: - m, norm = self._prepare_mix(mix) - target_dtype = next(self.roformer.parameters()).dtype - v = self.forward(torch.as_tensor(m, dtype=target_dtype, device=device)) - if norm: - v = v * norm[1] + norm[0] - v = v.mean(0) if v.ndim == 2 else v - return v.float().cpu().numpy().astype(np.float32) - - def separate_path(self, path: str, device: torch.device) -> tuple[np.ndarray, int]: - mix, _ = load_audio(path, sr=self.sample_rate, mono=False) - return self.separate_mono(np.asarray(mix, dtype=np.float32), device), self.sample_rate diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py deleted file mode 100644 index ca3ef852e..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py +++ /dev/null @@ -1,983 +0,0 @@ -import accelerate -import numpy as np -import scipy -import torch -from omegaconf import DictConfig -from torch import nn, view_as_complex, view_as_real -from torch.nn.utils import remove_weight_norm, weight_norm -from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz - - -def _deep_update_dict(base: dict, override: dict | None) -> dict: - """Recursively merge `override` into `base` (dict only).""" - if not override: - return base - for k, v in override.items(): - if isinstance(v, dict) and isinstance(base.get(k), dict): - base[k] = _deep_update_dict(base[k], v) - else: - base[k] = v - return base - - -def _get_vocoder_default_cfg() -> dict: - # Defaults keep current behavior. - return { - "preprocess": { - "hop_size": 480, - "sample_rate": 24000, - "max_length": 36000, - "n_fft": 1920, - "num_mels": 128, - "win_size": 1920, - "fmin": 0, - "fmax": 12000, - "mel_var": 8.14, - "mel_mean": -4.92, - "load_phone": False, - "load_chromagram": False, - }, - "model": { - "vocos": { - "input_channels": 128, - "dim": 1024, - "intermediate_dim": 4096, - "num_layers": 30, - "n_fft": 1920, - "hop_size": 480, - "padding": "same", - }, - "period_gan": { - "max_downsample_channels": 1024, - "channels": 64, - "channel_increasing_factor": 2, - }, - "spec_gan": { - "stft_params": { - "fft_sizes": [128, 256, 512, 1024, 2048], - "hop_sizes": [32, 64, 128, 256, 512], - "win_lengths": [128, 256, 512, 1024, 2048], - "window": "hann_window", - }, - "in_channels": 1, - "out_channels": 1, - "channels": 64, - "kernel_sizes": [5, 3], - "max_downsample_channels": 1024, - "down_scales": [2, 2, 2], - "use_weight_norm": True, - "use_complex": False, - }, - }, - "loss": { - "mel_loss": {"sample_rate": 24000}, - "disc_loss_weight": 1.0, - "mel_loss_weight": 10.0, - "adv_loss_weight": 2.0, - "fm_loss_weight": 2.0, - "spec_fm_loss_weight": 1.0, - }, - } - - -def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor: - """ - Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values. - - Args: - x (Tensor): Input tensor. - clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7. - - Returns: - Tensor: Element-wise logarithm of the input tensor with clipping applied. - """ - return torch.log(torch.clip(x, min=clip_val)) - - -def symlog(x: torch.Tensor) -> torch.Tensor: - return torch.sign(x) * torch.log1p(x.abs()) - - -def symexp(x: torch.Tensor) -> torch.Tensor: - return torch.sign(x) * (torch.exp(x.abs()) - 1) - - -class STFT(nn.Module): - def __init__( - self, - n_fft: int, - hop_length: int, - win_length: int, - center=True, - ): - super().__init__() - self.center = center - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - window = torch.hann_window(win_length) - self.register_buffer("window", window) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # x: (B, T * hop_length) - - if not self.center: - pad = self.win_length - self.hop_length - x = torch.nn.functional.pad(x, (pad // 2, pad // 2), mode="reflect") - - stft_spec = torch.stft( - x, - self.n_fft, - hop_length=self.hop_length, - win_length=self.win_length, - window=self.window, - center=self.center, - return_complex=False, - ) # (B, n_fft // 2 + 1, T, 2) - - rea = stft_spec[:, :, :, 0] # (B, n_fft // 2 + 1, T, 2) - imag = stft_spec[:, :, :, 1] # (B, n_fft // 2 + 1, T, 2) - - log_mag = torch.log( - torch.abs(torch.sqrt(torch.pow(rea, 2) + torch.pow(imag, 2))) + 1e-5 - ) # (B, n_fft // 2 + 1, T) - phase = torch.atan2(imag, rea) # (B, n_fft // 2 + 1, T) - - return log_mag, phase - - -class ISTFT(nn.Module): - """ - Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with - windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges. - See issue: https://github.com/pytorch/pytorch/issues/62323 - Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs. - The NOLA constraint is met as we trim padded samples anyway. - - Args: - n_fft (int): Size of Fourier transform. - hop_length (int): The distance between neighboring sliding window frames. - win_length (int): The size of window frame and STFT filter. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.n_fft = n_fft - self.hop_length = hop_length - self.win_length = win_length - window = torch.hann_window(win_length) - self.register_buffer("window", window) - - def forward(self, spec: torch.Tensor) -> torch.Tensor: - """ - Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram. - - Args: - spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size, - N is the number of frequency bins, and T is the number of time frames. - - Returns: - Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal. - """ - if self.padding == "center": - # Fallback to pytorch native implementation - return torch.istft( - spec, - self.n_fft, - self.hop_length, - self.win_length, - self.window, - center=True, - ) - elif self.padding == "same": - pad = (self.win_length - self.hop_length) // 2 - else: - raise ValueError("Padding must be 'center' or 'same'.") - - assert spec.dim() == 3, "Expected a 3D tensor as input" - B, N, T = spec.shape - - # Inverse FFT - ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") - ifft = ifft * self.window[None, :, None] - - # Overlap and Add - output_size = (T - 1) * self.hop_length + self.win_length - y = torch.nn.functional.fold( - ifft, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - )[:, 0, 0, pad:-pad] - - # Window envelope - window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) - window_envelope = torch.nn.functional.fold( - window_sq, - output_size=(1, output_size), - kernel_size=(1, self.win_length), - stride=(1, self.hop_length), - ).squeeze()[pad:-pad] - - # Normalize - assert (window_envelope > 1e-11).all() - y = y / window_envelope - - return y - - -class MDCT(nn.Module): - """ - Modified Discrete Cosine Transform (MDCT) module. - - Args: - frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, frame_len: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.frame_len = frame_len - N = frame_len // 2 - n0 = (N + 1) / 2 - window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() - self.register_buffer("window", window) - - pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len) - post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N) - # view_as_real: NCCL Backend does not support ComplexFloat data type - # https://github.com/pytorch/pytorch/issues/71613 - self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) - self.register_buffer("post_twiddle", view_as_real(post_twiddle)) - - def forward(self, audio: torch.Tensor) -> torch.Tensor: - """ - Apply the Modified Discrete Cosine Transform (MDCT) to the input audio. - - Args: - audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size - and T is the length of the audio. - - Returns: - Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames - and N is the number of frequency bins. - """ - if self.padding == "center": - audio = torch.nn.functional.pad(audio, (self.frame_len // 2, self.frame_len // 2)) - elif self.padding == "same": - # hop_length is 1/2 frame_len - audio = torch.nn.functional.pad(audio, (self.frame_len // 4, self.frame_len // 4)) - else: - raise ValueError("Padding must be 'center' or 'same'.") - - x = audio.unfold(-1, self.frame_len, self.frame_len // 2) - N = self.frame_len // 2 - x = x * self.window.expand(x.shape) - x = torch.fft.fft(x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1)[..., :N] - res = x * view_as_complex(self.post_twiddle).expand(x.shape) * np.sqrt(1 / N) - return torch.real(res) * np.sqrt(2) - - -class IMDCT(nn.Module): - """ - Inverse Modified Discrete Cosine Transform (IMDCT) module. - - Args: - frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, frame_len: int, padding: str = "same"): - super().__init__() - if padding not in ["center", "same"]: - raise ValueError("Padding must be 'center' or 'same'.") - self.padding = padding - self.frame_len = frame_len - N = frame_len // 2 - n0 = (N + 1) / 2 - window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() - self.register_buffer("window", window) - - pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N) - post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2)) - self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) - self.register_buffer("post_twiddle", view_as_real(post_twiddle)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients. - - Args: - x (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size, - L is the number of frames, and N is the number of frequency bins. - - Returns: - Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio. - """ - B, L, N = x.shape - Y = torch.zeros((B, L, N * 2), dtype=x.dtype, device=x.device) - Y[..., :N] = x - Y[..., N:] = -1 * torch.conj(torch.flip(x, dims=(-1,))) - y = torch.fft.ifft(Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1) - y = torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape)) * np.sqrt(N) * np.sqrt(2) - result = y * self.window.expand(y.shape) - output_size = (1, (L + 1) * N) - audio = torch.nn.functional.fold( - result.transpose(1, 2), - output_size=output_size, - kernel_size=(1, self.frame_len), - stride=(1, self.frame_len // 2), - )[:, 0, 0, :] - - if self.padding == "center": - pad = self.frame_len // 2 - elif self.padding == "same": - pad = self.frame_len // 4 - else: - raise ValueError("Padding must be 'center' or 'same'.") - - audio = audio[:, pad:-pad] - return audio - - -class FourierHead(nn.Module): - """Base class for inverse fourier modules.""" - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - raise NotImplementedError("Subclasses must implement the forward method.") - - -class ISTFTHead(FourierHead): - """ - ISTFT Head module for predicting STFT complex coefficients. - - Args: - dim (int): Hidden dimension of the model. - n_fft (int): Size of Fourier transform. - hop_length (int): The distance between neighboring sliding window frames, which should align with - the resolution of the input features. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - """ - - def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): - super().__init__() - out_dim = n_fft + 2 - self.out = torch.nn.Linear(dim, out_dim) - self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the ISTFTHead module. - - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - x = self.out(x).transpose(1, 2) - mag, p = x.chunk(2, dim=1) - mag = torch.exp(mag) - mag = torch.clip(mag, max=1e2) # safeguard to prevent excessively large magnitudes - # wrapping happens here. These two lines produce real and imaginary value - x = torch.cos(p) - y = torch.sin(p) - # recalculating phase here does not produce anything new - # only costs time - # phase = torch.atan2(y, x) - # S = mag * torch.exp(phase * 1j) - # better directly produce the complex value - - # Always compute complex values in float32 to avoid ComplexHalf warning (then cast audio back if needed) - orig_dtype = mag.dtype - mag, x, y = mag.float(), x.float(), y.float() - S = mag * (x + 1j * y) - audio = self.istft(S) - if orig_dtype != torch.float32: - audio = audio.to(orig_dtype) - return audio - - -class IMDCTSymExpHead(FourierHead): - """ - IMDCT Head module for predicting MDCT coefficients with symmetric exponential function - - Args: - dim (int): Hidden dimension of the model. - mdct_frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized - based on perceptual scaling. Defaults to None. - clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False. - """ - - def __init__( - self, - dim: int, - mdct_frame_len: int, - padding: str = "same", - sample_rate: int | None = None, - clip_audio: bool = False, - ): - super().__init__() - out_dim = mdct_frame_len // 2 - self.out = nn.Linear(dim, out_dim) - self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) - self.clip_audio = clip_audio - - if sample_rate is not None: - # optionally init the last layer following mel-scale - m_max = _hz_to_mel(sample_rate // 2) - m_pts = torch.linspace(0, m_max, out_dim) - f_pts = _mel_to_hz(m_pts) - scale = 1 - (f_pts / f_pts.max()) - - with torch.no_grad(): - self.out.weight.mul_(scale.view(-1, 1)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the IMDCTSymExpHead module. - - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - x = self.out(x) - x = symexp(x) - x = torch.clip(x, min=-1e2, max=1e2) # safeguard to prevent excessively large magnitudes - audio = self.imdct(x) - if self.clip_audio: - audio = torch.clip(x, min=-1.0, max=1.0) - - return audio - - -class IMDCTCosHead(FourierHead): - """ - IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p) - - Args: - dim (int): Hidden dimension of the model. - mdct_frame_len (int): Length of the MDCT frame. - padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". - clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False. - """ - - def __init__( - self, - dim: int, - mdct_frame_len: int, - padding: str = "same", - clip_audio: bool = False, - ): - super().__init__() - self.clip_audio = clip_audio - self.out = nn.Linear(dim, mdct_frame_len) - self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass of the IMDCTCosHead module. - - Args: - x (Tensor): Input tensor of shape (B, L, H), where B is the batch size, - L is the sequence length, and H denotes the model dimension. - - Returns: - Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. - """ - x = self.out(x) - m, p = x.chunk(2, dim=2) - m = torch.exp(m).clip(max=1e2) # safeguard to prevent excessively large magnitudes - audio = self.imdct(m * torch.cos(p)) - if self.clip_audio: - audio = torch.clip(x, min=-1.0, max=1.0) - return audio - - -class ConvNeXtBlock(nn.Module): - """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal. - - Args: - dim (int): Number of input channels. - intermediate_dim (int): Dimensionality of the intermediate layer. - layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. - Defaults to None. - adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. - None means non-conditional LayerNorm. Defaults to None. - """ - - def __init__( - self, - dim: int, - intermediate_dim: int, - layer_scale_init_value: float, - adanorm_num_embeddings: int | None = None, - ): - super().__init__() - self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv - self.adanorm = adanorm_num_embeddings is not None - if adanorm_num_embeddings: - self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) - else: - self.norm = nn.LayerNorm(dim, eps=1e-6) - self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers - self.act = nn.GELU() - self.pwconv2 = nn.Linear(intermediate_dim, dim) - self.gamma = ( - nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) - if layer_scale_init_value > 0 - else None - ) - - def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor | None = None) -> torch.Tensor: - residual = x - x = self.dwconv(x) - x = x.transpose(1, 2) # (B, C, T) -> (B, T, C) - if self.adanorm: - assert cond_embedding_id is not None - x = self.norm(x, cond_embedding_id) - else: - x = self.norm(x) - x = self.pwconv1(x) - x = self.act(x) - x = self.pwconv2(x) - if self.gamma is not None: - x = self.gamma * x - x = x.transpose(1, 2) # (B, T, C) -> (B, C, T) - - x = residual + x - return x - - -class AdaLayerNorm(nn.Module): - """ - Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes - - Args: - num_embeddings (int): Number of embeddings. - embedding_dim (int): Dimension of the embeddings. - """ - - def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.dim = embedding_dim - self.scale = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim) - self.shift = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim) - torch.nn.init.ones_(self.scale.weight) - torch.nn.init.zeros_(self.shift.weight) - - def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor: - scale = self.scale(cond_embedding_id) - shift = self.shift(cond_embedding_id) - x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps) - x = x * scale + shift - return x - - -class ResBlock1(nn.Module): - """ - ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions, - but without upsampling layers. - - Args: - dim (int): Number of input channels. - kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3. - dilation (tuple[int], optional): Dilation factors for the dilated convolutions. - Defaults to (1, 3, 5). - lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function. - Defaults to 0.1. - layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. - Defaults to None. - """ - - def __init__( - self, - dim: int, - kernel_size: int = 3, - dilation: tuple[int, int, int] = (1, 3, 5), - lrelu_slope: float = 0.1, - layer_scale_init_value: float | None = None, - ): - super().__init__() - self.lrelu_slope = lrelu_slope - self.convs1 = nn.ModuleList( - [ - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=dilation[0], - padding=self.get_padding(kernel_size, dilation[0]), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=dilation[1], - padding=self.get_padding(kernel_size, dilation[1]), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=dilation[2], - padding=self.get_padding(kernel_size, dilation[2]), - ) - ), - ] - ) - - self.convs2 = nn.ModuleList( - [ - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=1, - padding=self.get_padding(kernel_size, 1), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=1, - padding=self.get_padding(kernel_size, 1), - ) - ), - weight_norm( - nn.Conv1d( - dim, - dim, - kernel_size, - 1, - dilation=1, - padding=self.get_padding(kernel_size, 1), - ) - ), - ] - ) - - self.gamma = nn.ParameterList( - [ - ( - nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) - if layer_scale_init_value is not None - else None - ), - ( - nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) - if layer_scale_init_value is not None - else None - ), - ( - nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) - if layer_scale_init_value is not None - else None - ), - ] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma): - xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope) - xt = c1(xt) - xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope) - xt = c2(xt) - if gamma is not None: - xt = gamma * xt - x = xt + x - return x - - def remove_weight_norm(self): - for conv in self.convs1: - remove_weight_norm(conv) - for conv in self.convs2: - remove_weight_norm(conv) - - @staticmethod - def get_padding(kernel_size: int, dilation: int = 1) -> int: - return int((kernel_size * dilation - dilation) / 2) - - -class Backbone(nn.Module): - """Base class for the generator's backbone. It preserves the same temporal resolution across all layers.""" - - def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - x (Tensor): Input tensor of shape (B, C, L), where B is the batch size, - C denotes output features, and L is the sequence length. - - Returns: - Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length, - and H denotes the model dimension. - """ - raise NotImplementedError("Subclasses must implement the forward method.") - - -class VocosBackbone(Backbone): - """ - Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization - - Args: - input_channels (int): Number of input features channels. - dim (int): Hidden dimension of the model. - intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock. - num_layers (int): Number of ConvNeXtBlock layers. - layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`. - adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. - None means non-conditional model. Defaults to None. - """ - - def __init__( - self, - input_channels: int, - dim: int, - intermediate_dim: int, - num_layers: int, - layer_scale_init_value: float | None = None, - adanorm_num_embeddings: int | None = None, - ): - super().__init__() - self.input_channels = input_channels - self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3) - self.adanorm = adanorm_num_embeddings is not None - if adanorm_num_embeddings: - self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) - else: - self.norm = nn.LayerNorm(dim, eps=1e-6) - layer_scale_init_value = layer_scale_init_value or 1 / num_layers - self.convnext = nn.ModuleList( - [ - ConvNeXtBlock( - dim=dim, - intermediate_dim=intermediate_dim, - layer_scale_init_value=layer_scale_init_value, - adanorm_num_embeddings=adanorm_num_embeddings, - ) - for _ in range(num_layers) - ] - ) - self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) - self.apply(self._init_weights) - - def _init_weights(self, m): - if isinstance(m, nn.Conv1d | nn.Linear): - nn.init.trunc_normal_(m.weight, std=0.02) - nn.init.constant_(m.bias, 0) - - def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: - bandwidth_id = kwargs.get("bandwidth_id", None) - x = self.embed(x) - if self.adanorm: - assert bandwidth_id is not None - x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id) - else: - x = self.norm(x.transpose(1, 2)) - x = x.transpose(1, 2) - for conv_block in self.convnext: - x = conv_block(x, cond_embedding_id=bandwidth_id) - x = self.final_layer_norm(x.transpose(1, 2)) - return x - - -class VocosResNetBackbone(Backbone): - """ - Vocos backbone module built with ResBlocks. - - Args: - input_channels (int): Number of input features channels. - dim (int): Hidden dimension of the model. - num_blocks (int): Number of ResBlock1 blocks. - layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None. - """ - - def __init__( - self, - input_channels, - dim, - num_blocks, - layer_scale_init_value=None, - ): - super().__init__() - self.input_channels = input_channels - self.embed = weight_norm(nn.Conv1d(input_channels, dim, kernel_size=3, padding=1)) - layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3 - self.resnet = nn.Sequential( - *[ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) for _ in range(num_blocks)] - ) - - def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: - x = self.embed(x) - x = self.resnet(x) - x = x.transpose(1, 2) - return x - - -class Vocos(nn.Module): - def __init__( - self, - input_channels: int = 256, - dim: int = 384, - intermediate_dim: int = 1152, - num_layers: int = 8, - n_fft: int = 800, - hop_size: int = 200, - padding: str = "same", - adanorm_num_embeddings=None, - cfg=None, - ): - super().__init__() - - input_channels = cfg.input_channels if cfg is not None and hasattr(cfg, "input_channels") else input_channels - dim = cfg.dim if cfg is not None and hasattr(cfg, "dim") else dim - intermediate_dim = ( - cfg.intermediate_dim if cfg is not None and hasattr(cfg, "intermediate_dim") else intermediate_dim - ) - num_layers = cfg.num_layers if cfg is not None and hasattr(cfg, "num_layers") else num_layers - adanorm_num_embeddings = ( - cfg.adanorm_num_embeddings - if cfg is not None and hasattr(cfg, "adanorm_num_embeddings") - else adanorm_num_embeddings - ) - n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft - hop_size = cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size - padding = cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding - - self.backbone = VocosBackbone( - input_channels=input_channels, - dim=dim, - intermediate_dim=intermediate_dim, - num_layers=num_layers, - adanorm_num_embeddings=adanorm_num_embeddings, - ) - self.head = ISTFTHead(dim, n_fft, hop_size, padding) - - def forward(self, x): - x = self.backbone(x) - x = self.head(x) - - return x[:, None, :] - - -class JsonHParams: - def __init__(self, **kwargs): - for k, v in kwargs.items(): - if isinstance(v, dict): - v = JsonHParams(**v) - self[k] = v - - def keys(self): - return self.__dict__.keys() - - def items(self): - return self.__dict__.items() - - def values(self): - return self.__dict__.values() - - def __len__(self): - return len(self.__dict__) - - def __getitem__(self, key): - return getattr(self, key) - - def __setitem__(self, key, value): - return setattr(self, key, value) - - def __contains__(self, key): - return key in self.__dict__ - - def __repr__(self): - return self.__dict__.__repr__() - - -def build_vocoder_model(cfg): - vocoder_model = Vocos(cfg=cfg.model.vocos) - vocoder_model.eval() - return vocoder_model - - -def load_checkpoint(build_model_func, cfg, ckpt_path): - model = build_model_func(cfg) - - if ckpt_path is not None: - accelerate.load_checkpoint_and_dispatch(model, ckpt_path) - return model - - -def load_vocos_model( - ckpt_path: str | None = None, - config: DictConfig = None, - device: str = "cuda", -): - """Load Vocos vocoder. - - Args: - config: DictConfig, config for vocoder. - ckpt_path: str | None = None, path to checkpoint. - device: str, device to load model. Note: accelerate dispatch handles device placement. - """ - - merged_cfg = _deep_update_dict(_get_vocoder_default_cfg(), config) - vocoder_cfg_obj = JsonHParams(**merged_cfg) - - vocoder_model = load_checkpoint(build_vocoder_model, vocoder_cfg_obj, ckpt_path) - vocoder_model.eval() - - for param in vocoder_model.parameters(): - param.requires_grad = False - - return vocoder_model - - -class Vocoder(nn.Module): - def __init__(self, vocoder: dict | None = None, ckpt_path: str | None = None): - super().__init__() - vocoder = vocoder or {} - - model_cfg = vocoder.get("model_cfg") or vocoder.get("cfg") - - self.model = load_vocos_model( - ckpt_path=ckpt_path, - config=model_cfg, - ) - - def forward(self, x): - x = self.model(x) - return x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py deleted file mode 100644 index 379dda3a5..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Frozen Whisper encoder wrapper (wav -> encoder embeddings).""" - -import torch -import torchaudio -from transformers import WhisperFeatureExtractor, WhisperModel - -WHISPER_MEL_FRAMES = 3000 # 3000 frames at 16000 Hz - - -class WhisperEncoder: - """Auxiliary content encoder for SVC. - - Kept in FP32 for stable HF Whisper inference under vLLM-Omni's default - BF16 diffusion trunk (same idea as Stable Audio's FP32 VAE). Outputs are - cast to the trunk dtype before fusion with the CFM conditioner. - """ - - def __init__(self, device: str | torch.device | None = None) -> None: - self.fe = WhisperFeatureExtractor.from_pretrained("openai/whisper-base") - self.model = WhisperModel.from_pretrained("openai/whisper-base") - self.device = torch.device( - device or ("cuda" if torch.cuda.is_available() else "cpu"), - ) - self.model = self.model.to(device=self.device, dtype=torch.float32) - self.model.eval() - for param in self.model.parameters(): - param.requires_grad_(False) - - @property - def dtype(self) -> torch.dtype: - return torch.float32 - - def float(self) -> "WhisperEncoder": - self.model = self.model.float() - return self - - def encode( - self, - wav: torch.Tensor, - sr: int, - *, - output_dtype: torch.dtype | None = None, - ) -> torch.Tensor: - wav = ( - torchaudio.functional.resample(wav, orig_freq=sr, new_freq=self.fe.sampling_rate) - if sr != self.fe.sampling_rate - else wav - ) - wav_np = wav.detach().float().cpu().numpy().astype("float32", copy=False) - - inputs = self.fe( - wav_np, - sampling_rate=self.fe.sampling_rate, - return_tensors="pt", - padding=False, - truncation=False, - return_attention_mask=True, - ) - - input_features = inputs.input_features - num_frames = input_features.shape[-1] - if num_frames < WHISPER_MEL_FRAMES: - pad = WHISPER_MEL_FRAMES - num_frames - input_features = torch.nn.functional.pad(input_features, (0, pad)) - else: - input_features = input_features[..., :WHISPER_MEL_FRAMES] - - input_features = input_features.to(device=wav.device, dtype=self.dtype) - if self.model.device != wav.device: - self.model = self.model.to(wav.device) - - attention_mask = inputs.attention_mask.to(wav.device) if inputs.attention_mask is not None else None - - with torch.inference_mode(): - encoder_out = self.model.encoder(input_features).last_hidden_state - - if attention_mask is not None: - valid_mel_frames = attention_mask.sum(dim=1) - valid_enc_frames = (valid_mel_frames + 1) // 2 - max_valid_enc_frames = min(int(valid_enc_frames.max().item()), encoder_out.shape[1]) - encoder_out = encoder_out[:, :max_valid_enc_frames, :] - valid_len = min(int(valid_enc_frames[0].item()), max_valid_enc_frames) - if valid_len < max_valid_enc_frames: - encoder_out[0, valid_len:, :] = 0 - - if output_dtype is not None and encoder_out.dtype != output_dtype: - encoder_out = encoder_out.to(dtype=output_dtype) - return encoder_out - - -if __name__ == "__main__": - torch.manual_seed(0) - audio = torch.randn(1, 24000 * 25).float().to("cuda") - encoder = WhisperEncoder() - whisper_encoder_out = encoder.encode(audio, sr=24000) - print(whisper_encoder_out.shape) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py deleted file mode 100644 index 0961e6916..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py +++ /dev/null @@ -1,377 +0,0 @@ -import os -import time -from collections.abc import Callable -from contextlib import contextmanager -from typing import Any, ClassVar - -import numpy as np -import torch -import torch.nn as nn -from vllm.logger import init_logger -from vllm.utils.torch_utils import set_default_torch_dtype - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _unwrap, _wrap -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.interface import ( - SupportAudioInput, - SupportAudioOutput, - SupportsComponentDiscovery, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import consume_payload -from vllm_omni.diffusion.models.soulx_singer.utils import _patch_torchaudio_load, load_config -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.platforms import current_omni_platform - -logger = init_logger(__name__) - -_SOULX_PROFILER_TARGETS = [ - "forward", -] - - -def convert_soulx_audio_output_to_numpy(audio: torch.Tensor | dict[str, Any]) -> np.ndarray | dict[str, Any]: - if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): - payload = dict(audio["payload"]) - audio_payload = payload.get("audio") - if isinstance(audio_payload, torch.Tensor): - payload["audio"] = audio_payload.detach().cpu().float().numpy() - return { - "payload": payload, - "metadata": audio.get("metadata") or {}, - } - if isinstance(audio, torch.Tensor): - return audio.detach().cpu().float().numpy() - return audio - - -class FlowMatchingAudioPipeline( - nn.Module, - CFGParallelMixin, - SupportAudioInput, - SupportAudioOutput, - SupportsComponentDiscovery, - DiffusionPipelineProfilerMixin, -): - """SoulX flow-matching pipeline with CFG / CFG-parallel support.""" - - support_audio_input: ClassVar[bool] = True - support_audio_output: ClassVar[bool] = True - audio_sample_rate: ClassVar[int] = 24000 - _DEFAULT_RESCALE_CFG: ClassVar[float] = 0.75 - _dit_modules: ClassVar[list[str]] = ["cfm_decoder.model.diff_estimator"] - _encoder_modules: ClassVar[list[str]] = ["mel", "f0_encoder"] - _vae_modules: ClassVar[list[str]] = ["vocoder"] - _resident_modules: ClassVar[list[str]] = ["cfm_decoder.model.cond_emb"] - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__() - self.od_config = od_config - self.device = get_local_device() - - # 1. parse model path and yaml config - model_dir = od_config.model - assert model_dir is not None, "model_dir is not provided" - - if not os.path.exists(model_dir): - raise FileNotFoundError( - f"""Model directory {model_dir} does not exist. User should refer to - `examples/offline_inference/text_to_speech/README.md` for usage details.""" - ) - - # Empty: skip DiffusersPipelineLoader.get_all_weights() → no _prepare_weights / file glob / - # safetensors or pt iterator from the framework. Load checkpoints in __init__ instead. - self.weights_sources = () - - self.model_path = model_dir - - config_path = os.path.join(model_dir, "config.yaml") - hf_config = load_config(config_path) - - self.audio_config = hf_config.audio - self.encoder_config = hf_config.model.encoder - self.flow_matching_config = hf_config.model.flow_matching - - self.mel_dim = self.audio_config.num_mels - - self.f0_bin = self.encoder_config.f0_bin - self.f0_dim = self.encoder_config.f0_dim - self.text_dim = self.encoder_config.text_dim - self.vocab_size = self.encoder_config.vocab_size - self.pitch_dim = self.encoder_config.pitch_dim - self.type_dim = self.encoder_config.type_dim - - _patch_torchaudio_load() - - def _mel_to_audio(self, generated_mel: torch.Tensor, *, squeeze: bool = False) -> torch.Tensor: - audio = self.vocoder(generated_mel.transpose(1, 2)[0:1, ...].float()).float() - return audio.squeeze() if squeeze else audio - - def _mel_from_wav(self, wav: torch.Tensor) -> torch.Tensor: - if wav.dtype != torch.float32: - wav = wav.float() - return self.mel(wav) - - def _load_soulx_checkpoint(self, checkpoint_name: str) -> None: - weight_path = os.path.join(self.model_path, checkpoint_name) - if not os.path.isfile(weight_path): - raise FileNotFoundError( - f"Model weights not found at {weight_path}. Please check the pretrained model path." - ) - state = torch.load(weight_path, map_location=self.device) - self.mel.float() - self.load_state_dict(state["state_dict"], strict=True) - self._finalize_loaded_dtypes() - logger.info("Loaded model weights from %s", weight_path) - - def _setup_soulx_profiler(self, *, extra_targets: list[str] | None = None) -> None: - """Enable stage timing when ``od_config.enable_diffusion_pipeline_profiler`` is set.""" - targets = list(_SOULX_PROFILER_TARGETS) - if extra_targets: - targets.extend(extra_targets) - self.setup_diffusion_pipeline_profiler( - profiler_targets=targets, - enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler, - ) - - @contextmanager - def _stage_timer(self, name: str): - if not getattr(self, "enable_diffusion_pipeline_profiler", False): - yield - return - if current_omni_platform.is_available(): - current_omni_platform.synchronize() - start = time.perf_counter() - try: - yield - finally: - if current_omni_platform.is_available(): - current_omni_platform.synchronize() - duration = time.perf_counter() - start - with self._profiler_lock: - self._stage_durations[name] = self._stage_durations.get(name, 0.0) + duration - - def _profiler_stage_durations(self) -> dict[str, float] | None: - if getattr(self, "enable_diffusion_pipeline_profiler", False): - return self.stage_durations - return None - - @staticmethod - def _build_fp32_audio_modules(audio_config) -> tuple[nn.Module, nn.Module]: - """Mel/vocoder must stay FP32 even when ``od_config.dtype`` is FP16/BF16.""" - from vllm_omni.diffusion.models.soulx_singer.modules import MelSpectrogramEncoder, Vocoder - - with set_default_torch_dtype(torch.float32): - mel = MelSpectrogramEncoder(audio_config) - vocoder = Vocoder() - return mel, vocoder - - def _finalize_loaded_dtypes(self) -> None: - """Keep mel/vocoder in FP32; cast DiT trunk to ``od_config.dtype`` after load.""" - self.mel.float() - self.vocoder.float() - trunk_dtype = self.od_config.dtype - if trunk_dtype not in (torch.float16, torch.bfloat16): - return - trunk_modules = ( - "cfm_decoder", - "f0_encoder", - "preflow", - "note_text_encoder", - "note_pitch_encoder", - "note_type_encoder", - ) - for name in trunk_modules: - module = getattr(self, name, None) - if module is not None: - module.to(dtype=trunk_dtype) - whisper = getattr(self, "whisper_encoder", None) - if whisper is not None: - whisper.float() - - @property - def transformer(self) -> nn.Module: - """Alias expected by ``CFGParallelMixin`` default hooks.""" - return self.cfm_decoder.model.diff_estimator - - @property - def diffusion_trunk_dtype(self) -> torch.dtype: - """DiT trunk dtype from ``diff_estimator`` (not ``cond_emb``, which is registered first).""" - return next(self.cfm_decoder.model.diff_estimator.parameters()).dtype - - def _to_trunk_dtype(self, *tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: - """Cast tensors to ``diffusion_trunk_dtype`` at the CFM/DiT boundary.""" - trunk_dtype = self.diffusion_trunk_dtype - return tuple(t if t.dtype == trunk_dtype else t.to(dtype=trunk_dtype) for t in tensors) - - def _resolve_diffusion_generator(self, sampling_params: Any) -> torch.Generator | None: - """Build a device-local RNG for CFM noise from ``generator`` or ``seed``.""" - generator = getattr(sampling_params, "generator", None) - if generator is not None: - return generator - seed = getattr(sampling_params, "seed", None) - if seed is None: - return None - return torch.Generator(device=self.device).manual_seed(int(seed)) - - def _prepare_cfm_loop_state( - self, - prompt: torch.Tensor, - cond: torch.Tensor, - n_timesteps: int, - *, - generator: torch.Generator | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: - """Cast FP32 islands to trunk dtype and build CFM loop tensors/masks.""" - trunk_dtype = self.diffusion_trunk_dtype - prompt, cond = self._to_trunk_dtype(prompt, cond) - cond_emb = self.cfm_decoder.model.cond_emb(cond) - if cond_emb.dtype != trunk_dtype: - cond_emb = self._to_trunk_dtype(cond_emb)[0] - device = cond_emb.device - - h = 1.0 / n_timesteps - h_tensor = torch.tensor(h, device=device, dtype=trunk_dtype) - prompt_len = prompt.shape[1] - target_len = cond_emb.shape[1] - prompt_len - - x_mask = torch.ones(cond_emb.shape[0], target_len, device=device, dtype=trunk_dtype) - prompt_mask = torch.ones(cond_emb.shape[0], prompt_len, device=device, dtype=trunk_dtype) - xt_mask = torch.cat([prompt_mask, x_mask], dim=1) - xt = torch.randn( - cond_emb.shape[0], - target_len, - self.mel_dim, - dtype=trunk_dtype, - device=device, - generator=generator, - ) - return cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len - - def predict_noise( - self, - *, - x: torch.Tensor, - diffusion_step: torch.Tensor, - cond: torch.Tensor, - x_mask: torch.Tensor, - prompt_len: int | None = None, - **_: Any, - ) -> torch.Tensor: - """Forward through ``diff_estimator``, optionally slice off prompt frames.""" - x, diffusion_step, cond, x_mask = self._to_trunk_dtype(x, diffusion_step, cond, x_mask) - flow_pred = self.cfm_decoder.model.diff_estimator(x, diffusion_step, cond, x_mask) - if prompt_len is not None: - flow_pred = flow_pred[:, prompt_len:, :] - return flow_pred - - def combine_cfg_noise( - self, - positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], - negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], - true_cfg_scale: float, - cfg_normalize: bool = False, - rescale_cfg: float | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, ...]: - """SoulX CFG: pos + scale * (pos - neg) with std rescaling.""" - del cfg_normalize # SoulX uses rescale_cfg instead of cfg_normalize_function. - if rescale_cfg is None: - rescale_cfg = getattr(self, "rescale_cfg", self._DEFAULT_RESCALE_CFG) - - pos = _unwrap(_wrap(positive_noise_pred)) - neg = _unwrap(_wrap(negative_noise_pred)) - - pos_flow_pred_std = pos.std() - flow_pred_cfg = pos + true_cfg_scale * (pos - neg) - rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() - combined = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg - return combined.to(dtype=pos.dtype) - - def _run_flow_matching_loop( - self, - prompt: torch.Tensor, - cond: torch.Tensor, - n_timesteps: int, - cfg: float, - *, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - """Euler flow-matching loop aligned with ``FlowMatchingTransformer.reverse_diffusion``.""" - cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len = self._prepare_cfm_loop_state( - prompt, cond, n_timesteps, generator=generator - ) - trunk_dtype = self.diffusion_trunk_dtype - step_scale = 1.0 / n_timesteps - - for i in range(n_timesteps): - xt_input = torch.cat([prompt, xt], dim=1) - diffusion_step = torch.full( - (xt.shape[0],), - fill_value=(i + 0.5) * step_scale, - device=xt.device, - dtype=trunk_dtype, - ) - - flow_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=(cfg > 0), - true_cfg_scale=cfg, - positive_kwargs={ - "x": xt_input, - "diffusion_step": diffusion_step, - "cond": cond_emb, - "x_mask": xt_mask, - "prompt_len": prompt_len, - }, - negative_kwargs={ - "x": xt, - "diffusion_step": diffusion_step, - # CFG null cond: zero out target frames only (prompt cond stays in pos branch). - "cond": torch.zeros_like(cond_emb)[:, : xt.shape[1], :], - "x_mask": x_mask, - }, - cfg_normalize=False, - ) - - xt = xt + flow_pred * h_tensor - - return xt - - def _forward_batch_from_request( - self, - req: OmniDiffusionRequest, - *, - kind: str, - metadata_key: str, - infer_batch_fn: Callable[..., tuple[torch.Tensor, int]], - prepare_extra_args: Callable[[dict[str, Any], Any], dict[str, Any]] | None = None, - ) -> DiffusionOutput: - sampling_params = req.sampling_params - extra_args = dict(getattr(sampling_params, "extra_args", None) or {}) - if prepare_extra_args is not None: - extra_args = prepare_extra_args(extra_args, sampling_params) - sampling_params.extra_args = extra_args - - num_inference_steps = sampling_params.num_inference_steps or 32 - guidance_scale = sampling_params.guidance_scale if sampling_params.guidance_scale_provided else 3.0 - generator = self._resolve_diffusion_generator(sampling_params) - - with self._stage_timer("consume_payload"): - payload = consume_payload(req, kind, self.device) - - with self._stage_timer("infer_segments"): - audio, pitch_shift = infer_batch_fn( - payload, - extra_args=extra_args, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - ) - - return DiffusionOutput( - output={ - "payload": {"audio": audio}, - "metadata": {"audio": {metadata_key: pitch_shift}}, - }, - stage_durations=self._profiler_stage_durations() or {}, - ) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py deleted file mode 100644 index 85fed9ee7..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py +++ /dev/null @@ -1,411 +0,0 @@ -"""SoulX-Singer SVC (voice conversion) pipeline implementation.""" - -import os -from collections.abc import Iterable -from typing import Any, ClassVar - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -from vllm.logger import init_logger -from vllm.utils.torch_utils import set_default_torch_dtype - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.soulx_singer.modules import ( - CFMDecoder, - WhisperEncoder, -) -from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( - FlowMatchingAudioPipeline, - convert_soulx_audio_output_to_numpy, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PREPROCESSED_KEY, - get_soulx_preprocessed_payload, - has_precomputed, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( - attach_preprocess_for_diffusion_request, - is_warmup_request, - resolve_preprocess_audio, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - f0_to_coarse, - load_config, - resolve_pitch_shift, - validate_soulx_extra_args, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -logger = init_logger(__name__) - - -def get_soulxsinger_svc_pre_process_func(od_config: OmniDiffusionConfig): - hf_config = load_config(os.path.join(od_config.model, "config.yaml")) - sample_rate = hf_config.audio.sample_rate - device = get_local_device() - _pipeline = None - - def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - nonlocal _pipeline - - extra_args = validate_soulx_extra_args( - "svc", - dict(getattr(request.sampling_params, "extra_args", None) or {}), - ) - - # Inline build: when no warmup/no precomputed paths/no IPC payload, - # build the preprocess payload directly from audio file paths. - if not (is_warmup_request(request) or has_precomputed(extra_args, "svc")): - prompt = request.prompt - if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type] - prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type] - if prompt_audio is not None and target_audio is not None: - if _pipeline is None: - from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import ( - SoulXPreprocessPipeline, - ) - - _pipeline = SoulXPreprocessPipeline( - od_config=od_config, - vocal_sep=bool(extra_args.get("vocal_sep", False)), - verbose=bool(extra_args.get("preprocess_verbose", False)), - extra_args=extra_args, - ) - payload = _pipeline.build_svc_payload_from_audio( - prompt_audio=prompt_audio, - target_audio=target_audio, - sample_rate=sample_rate, - device=device, - vocal_sep=extra_args.get("vocal_sep"), - ) - payload.setdefault("kind", "svc") - prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment] - - return attach_preprocess_for_diffusion_request( - request, - kind="svc", - sample_rate=sample_rate, - device=device, - ) - - return pre_process_func - - -def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): - def post_process_func(audio: torch.Tensor): - return convert_soulx_audio_output_to_numpy(audio) - - return post_process_func - - -class PipelineSoulXSingerSVC(FlowMatchingAudioPipeline): - """SVC pipeline for the SoulX-Singer model.""" - - _encoder_modules: ClassVar[list[str]] = [ - "whisper_encoder", - "mel", - "f0_encoder", - ] - - EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "prompt_wav_path", - "target_wav_path", - "prompt_f0_path", - "target_f0_path", - "prompt_audio", - "target_audio", - "vocal_sep", - "preprocess_weights_dir", - "preprocess_verbose", - "auto_shift", - "pitch_shift", - } - ) - - EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"pitch_shift"}) - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__(od_config=od_config, prefix=prefix) - - self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim) - - self.cfm_decoder = CFMDecoder(self.flow_matching_config) - - self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config) - with set_default_torch_dtype(torch.float32): - self.whisper_encoder = WhisperEncoder(device=self.device) - - self._setup_soulx_profiler() - - def _encode_condition(self, *, whisper_features: torch.Tensor, f0_coarse: torch.Tensor) -> torch.Tensor: - cond = whisper_features + self.f0_encoder(f0_coarse) - return self._to_trunk_dtype(cond)[0] - - def _encode_prompt_whisper_feature(self, prompt_wav: torch.Tensor) -> torch.Tensor: - trunk_dtype = self.f0_encoder.weight.dtype - with self._stage_timer("whisper"): - return self.whisper_encoder.encode( - prompt_wav, - sr=self.audio_config.sample_rate, - output_dtype=trunk_dtype, - ) - - _encode_target_whisper_feature = _encode_prompt_whisper_feature - - def _infer_segment( - self, - *, - prompt_mel: torch.Tensor, - prompt_wav: torch.Tensor, - target_wav: torch.Tensor, - prompt_f0: torch.Tensor, - target_f0: torch.Tensor, - pitch_shift: int, - num_inference_steps: int, - guidance_scale: float, - prompt_feature: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - """Single-segment SVC inference aligned with ``SoulXSingerSVC.infer_segment``.""" - len_prompt_mel = prompt_mel.shape[1] - prompt_f0 = F.pad(prompt_f0, (0, 0, 0, max(0, len_prompt_mel - prompt_f0.shape[1])))[:, :len_prompt_mel] - - f0_coarse_prompt = f0_to_coarse(prompt_f0) - f0_coarse_target = f0_to_coarse(target_f0, f0_shift=int(pitch_shift * 5)) - f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1) - - trunk_dtype = self.f0_encoder.weight.dtype - if prompt_feature is None: - prompt_feature = self._encode_prompt_whisper_feature(prompt_wav) - with self._stage_timer("whisper"): - target_feature = self.whisper_encoder.encode( - target_wav, - sr=self.audio_config.sample_rate, - output_dtype=trunk_dtype, - ) - prompt_feature = F.pad( - prompt_feature, - (0, 0, 0, max(0, f0_coarse_prompt.shape[1] - prompt_feature.shape[1])), - )[:, : f0_coarse_prompt.shape[1], :] - target_feature = F.pad( - target_feature, - (0, 0, 0, max(0, f0_coarse_target.shape[1] - target_feature.shape[1])), - )[:, : f0_coarse_target.shape[1], :] - - whisper_features = torch.cat([prompt_feature, target_feature], dim=1) - with self._stage_timer("cond_encode"): - cond = self._encode_condition(whisper_features=whisper_features, f0_coarse=f0_coarse) - - with self._stage_timer("cfm"): - generated_mel = self._run_flow_matching_loop( - prompt=prompt_mel, - cond=cond, - n_timesteps=num_inference_steps, - cfg=guidance_scale, - generator=generator, - ) - - with self._stage_timer("vocoder"): - generated_audio = self._mel_to_audio(generated_mel, squeeze=True) - target_len = target_wav.shape[-1] - if generated_audio.shape[-1] > target_len: - generated_audio = generated_audio[:target_len] - elif generated_audio.shape[-1] < target_len: - generated_audio = F.pad(generated_audio, (0, target_len - generated_audio.shape[-1])) - return generated_audio - - def _build_vocal_segments( - self, - f0: torch.Tensor, - *, - hop_size: int, - sample_rate: int, - uv_frames_th: int = 10, - min_duration_sec: float = 15.0, - max_duration_sec: float = 30.0, - ignore_silent: bool = True, - ) -> tuple[list[tuple[float, float]], list[tuple[float, float]]]: - """Build vocal segments based on F0 contour. - - Mirrors upstream ``SoulXSingerSVC.build_vocal_segments``. - Returns (overlap_segments, segments) tuples of (start_sec, end_sec). - """ - f0_np = ( - f0.detach().float().cpu().numpy().squeeze() if isinstance(f0, torch.Tensor) else np.asarray(f0).squeeze() - ) - total_frames = int(f0_np.shape[0]) - if total_frames == 0: - return [], [] - - f0_rate = sample_rate // hop_size - min_frames = max(1, int(round(min_duration_sec * f0_rate))) - max_frames = max(1, int(round(max_duration_sec * f0_rate))) - - split_points = [0] - - def append_split_point(point: int) -> None: - point = int(max(0, min(point, total_frames))) - while point - split_points[-1] > max_frames: - split_points.append(split_points[-1] + max_frames) - if point > split_points[-1]: - split_points.append(point) - - idx = 0 - while idx < total_frames: - if f0_np[idx] == 0: - run_start = idx - while idx < total_frames and f0_np[idx] == 0: - idx += 1 - run_end = idx - if (run_end - run_start) >= uv_frames_th: - split_point = max(run_end - 5, (run_start + run_end) // 2) - append_split_point(split_point) - else: - idx += 1 - append_split_point(total_frames) - - segments: list[tuple[int, int]] = [] - overlap_segments: list[tuple[int, int]] = [] - num_overlaps = 1 - - def append_segment(start_idx: int, end_idx: int, overlaps: int = num_overlaps) -> None: - segments.append((split_points[start_idx], split_points[end_idx])) - overlap_start_idx = start_idx - if start_idx > 0 and (split_points[end_idx] - split_points[start_idx - overlaps]) <= max_frames: - overlap_start_idx = start_idx - overlaps - overlap_segments.append((split_points[overlap_start_idx], split_points[end_idx])) - - seg_start, seg_end = 0, 1 - while seg_start < len(split_points) - 1: - while seg_end < len(split_points) and (split_points[seg_end] - split_points[seg_start]) < min_frames: - seg_end += 1 - if seg_end >= len(split_points): - append_segment(seg_start, len(split_points) - 1) - break - append_segment(seg_start, seg_end) - seg_start = seg_end - seg_end = seg_start + 1 - - if ignore_silent: - filtered_idx = [] - for i, (ov_start, ov_end) in enumerate(overlap_segments): - voice_ratio = np.sum(f0_np[ov_start:ov_end] > 0) / max(1, ov_end - ov_start) - voiced_frames = np.sum(f0_np[ov_start:ov_end] > 0) - if voice_ratio > 0.05 and voiced_frames >= 10: - filtered_idx.append(i) - overlap_segments = [overlap_segments[i] for i in filtered_idx] - segments = [segments[i] for i in filtered_idx] - - # convert indices to seconds - overlap_sec = [(s / f0_rate, e / f0_rate) for s, e in overlap_segments] - seg_sec = [(s / f0_rate, e / f0_rate) for s, e in segments] - return overlap_sec, seg_sec - - def infer_svc_batch( - self, - payload: dict[str, Any], - *, - extra_args: dict[str, Any], - num_inference_steps: int, - guidance_scale: float, - generator: torch.Generator | None = None, - ) -> tuple[torch.Tensor, int]: - """Batch SVC with chunked inference for long audio (matches upstream).""" - prompt_wav = payload["prompt_wav"] - target_wav = payload["target_wav"] - prompt_f0 = payload["prompt_f0"] - target_f0 = payload["target_f0"] - - pitch_shift = resolve_pitch_shift( - auto_shift=bool(extra_args.get("auto_shift", True)), - manual_shift=int(extra_args.get("pitch_shift", 0)), - prompt_f0=prompt_f0, - target_f0=target_f0, - ) - - prompt_mel = self._mel_from_wav(prompt_wav) - - # Long audio: chunk by vocal segments (same as upstream) - max_sec = 30.0 * self.audio_config.sample_rate - if target_wav.shape[-1] < max_sec: - # Short audio: one-shot - generated_audio = self._infer_segment( - prompt_mel=prompt_mel, - prompt_wav=prompt_wav, - target_wav=target_wav, - prompt_f0=prompt_f0, - target_f0=target_f0, - pitch_shift=pitch_shift, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - ) - return generated_audio.unsqueeze(0), pitch_shift - - # Build vocal segments - overlap_segments, segments = self._build_vocal_segments( - target_f0, - hop_size=self.audio_config.hop_size, - sample_rate=self.audio_config.sample_rate, - ) - if len(segments) == 0: - segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)] - overlap_segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)] - - f0_rate = self.audio_config.sample_rate // self.audio_config.hop_size - generated_audio = torch.zeros_like(target_wav) - - for idx in range(len(segments)): - overlap_start_sec, overlap_end_sec = overlap_segments[idx] - seg_start_sec, seg_end_sec = segments[idx] - - wav_start = int(round(overlap_start_sec * self.audio_config.sample_rate)) - wav_end = int(round(overlap_end_sec * self.audio_config.sample_rate)) - f0_start = int(round(overlap_start_sec * f0_rate)) - f0_end = int(round(overlap_end_sec * f0_rate)) - - wav_start = max(0, min(wav_start, target_wav.shape[-1])) - wav_end = max(wav_start, min(wav_end, target_wav.shape[-1])) - f0_start = max(0, min(f0_start, target_f0.shape[-1])) - f0_end = max(f0_start, min(f0_end, target_f0.shape[-1])) - - segment_gt_wav = target_wav[:, wav_start:wav_end] - segment_gt_f0 = target_f0[:, f0_start:f0_end] - - segment_generated = self._infer_segment( - prompt_mel=prompt_mel, - prompt_wav=prompt_wav, - target_wav=segment_gt_wav, - prompt_f0=prompt_f0, - target_f0=segment_gt_f0, - pitch_shift=pitch_shift, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - generator=generator, - ) - - seg_start = int(round(seg_start_sec * self.audio_config.sample_rate)) - seg_end = int(round(seg_end_sec * self.audio_config.sample_rate)) - segment_generated = segment_generated[seg_start - wav_start : seg_end - wav_start] - generated_audio[:, seg_start:seg_end] = segment_generated - - return generated_audio.unsqueeze(0), pitch_shift - - @torch.inference_mode() - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - return self._forward_batch_from_request( - req, - kind="svc", - metadata_key="pitch_shift", - infer_batch_fn=self.infer_svc_batch, - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - del weights - self._load_soulx_checkpoint("model-svc.pt") - return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py deleted file mode 100644 index fb2c086f2..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py +++ /dev/null @@ -1,355 +0,0 @@ -"""SoulX-Singer SVS (score-driven) pipeline implementation.""" - -from collections.abc import Iterable -from typing import Any, ClassVar - -import torch -import torch.nn as nn -import torch.nn.functional as F -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.models.soulx_singer.modules import ( - CFMDecoder, - ConvNeXtV2Block, -) -from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( - FlowMatchingAudioPipeline, - convert_soulx_audio_output_to_numpy, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PREPROCESSED_KEY, - get_soulx_preprocessed_payload, - has_precomputed, -) -from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( - attach_preprocess_for_diffusion_request, - build_metadata_processor, - is_warmup_request, - normalize_svs_control_extra_args, - resolve_preprocess_audio, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - f0_to_coarse, - resolve_pitch_shift, - validate_soulx_extra_args, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - -logger = init_logger(__name__) - - -def get_soulxsinger_pre_process_func(od_config: OmniDiffusionConfig): - """Validate/load SVS preprocess payload for single-stage or stage-1 DiT.""" - metadata_processor = build_metadata_processor(od_config) - _pipeline = None - - def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - nonlocal _pipeline - - extra_args = validate_soulx_extra_args( - "svs", - dict(getattr(request.sampling_params, "extra_args", None) or {}), - ) - - # Inline build: when no warmup/no precomputed paths/no IPC payload, - # build the preprocess payload directly from audio file paths. - if not (is_warmup_request(request) or has_precomputed(extra_args, "svs")): - prompt = request.prompt - if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type] - prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type] - if prompt_audio is not None and target_audio is not None: - if _pipeline is None: - from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import ( - SoulXPreprocessPipeline, - ) - - _pipeline = SoulXPreprocessPipeline( - od_config=od_config, - vocal_sep=bool(extra_args.get("vocal_sep", False)), - verbose=bool(extra_args.get("preprocess_verbose", False)), - extra_args=extra_args, - ) - payload = _pipeline.build_svs_payload_from_audio( - prompt_audio=prompt_audio, - target_audio=target_audio, - metadata_processor=metadata_processor, - language=str(extra_args.get("language", "Mandarin")), - vocal_sep=extra_args.get("vocal_sep"), - prompt_vocal_sep=extra_args.get("prompt_vocal_sep"), - target_vocal_sep=extra_args.get("target_vocal_sep"), - prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"), - target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"), - ) - payload.setdefault("kind", "svs") - prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment] - - return attach_preprocess_for_diffusion_request( - request, - kind="svs", - metadata_processor=metadata_processor, - ) - - return pre_process_func - - -def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): - """Convert pipeline audio tensor output for offline consumers.""" - - def post_process_func(audio: torch.Tensor): - return convert_soulx_audio_output_to_numpy(audio) - - return post_process_func - - -def _expand_states(h: torch.Tensor, mel2token: torch.Tensor) -> torch.Tensor: - if mel2token.max() > h.size(1) - 1: - logger.warning( - "mel2token.max() (%s) is greater than h.size(1) - 1 (%s); clamping.", - mel2token.max(), - h.size(1) - 1, - ) - mel2token = torch.clamp(mel2token, 0, h.size(1) - 1) - mel2token_ = mel2token[..., None].repeat([1, 1, h.shape[-1]]) - return torch.gather(h, 1, mel2token_) - - -class PipelineSoulXSingerSVS(FlowMatchingAudioPipeline): - """Pipeline for the SoulX-Singer model (SoulX-Singer).""" - - _encoder_modules: ClassVar[list[str]] = [ - "mel", - "f0_encoder", - "preflow", - "note_text_encoder", - "note_pitch_encoder", - "note_type_encoder", - ] - - EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "prompt_metadata_path", - "target_metadata_path", - "audio_path", - "prompt_audio", - "target_audio", - "language", - "vocal_sep", - "max_merge_duration", - "preprocess_weights_dir", - "preprocess_verbose", - "control", - "auto_shift", - "pitch_shift", - } - ) - - EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"f0_shift"}) - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__(od_config=od_config, prefix=prefix) - self.metadata_processor = build_metadata_processor(od_config) - - self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim) - self.preflow = nn.Sequential( - *[ - ConvNeXtV2Block( - self.text_dim, - self.text_dim * 2, - ) - for _ in range(self.encoder_config.num_layers) - ] - ) - - self.cfm_decoder = CFMDecoder(self.flow_matching_config) - - self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config) - - self.note_text_encoder = nn.Embedding(self.vocab_size, self.text_dim) - self.note_pitch_encoder = nn.Embedding(256, self.pitch_dim) - self.note_type_encoder = nn.Embedding(256, self.type_dim) - - self._setup_soulx_profiler() - - def _encode_condition( - self, - *, - note_pitch: torch.Tensor, - note_type: torch.Tensor, - note_text: torch.Tensor, - mel2note: torch.Tensor, - f0_coarse: torch.Tensor, - ) -> torch.Tensor: - features = ( - self.note_pitch_encoder(note_pitch) + self.note_type_encoder(note_type) + self.note_text_encoder(note_text) - ) - features = self.preflow(features) - features = _expand_states(features, mel2note) - features = features + self.f0_encoder(f0_coarse) - return self._to_trunk_dtype(features)[0] - - def _infer_svs_segment( - self, - prompt_meta: dict, - target_meta: dict, - *, - pitch_shift: int, - num_inference_steps: int, - guidance_scale: float, - prompt_mel: torch.Tensor | None = None, - generator: torch.Generator | None = None, - ) -> torch.Tensor: - prompt_note_text = prompt_meta["phoneme"] - prompt_mel2note = prompt_meta["mel2note"] - prompt_note_type = prompt_meta["note_type"] - prompt_note_pitch = prompt_meta["note_pitch"] - prompt_f0 = prompt_meta["f0"] - - target_note_text = target_meta["phoneme"] - target_mel2note = target_meta["mel2note"] - target_note_type = target_meta["note_type"] - target_note_pitch = target_meta["note_pitch"] - target_f0 = target_meta["f0"] - - if target_f0 is None or prompt_f0 is None: - target_f0 = torch.zeros_like(target_mel2note).float() - prompt_f0 = torch.zeros_like(prompt_mel2note).float() - if target_note_pitch is None or prompt_note_pitch is None: - target_note_pitch = torch.zeros_like(target_note_type).int() - prompt_note_pitch = torch.zeros_like(prompt_note_type).int() - - len_prompt_note = prompt_note_pitch.shape[1] - len_prompt_mel = prompt_f0.shape[1] if prompt_f0 is not None else prompt_mel2note.shape[1] - - note_pitch = torch.cat([prompt_note_pitch, target_note_pitch], dim=1) - note_text = torch.cat([prompt_note_text, target_note_text], dim=1) - note_type = torch.cat([prompt_note_type, target_note_type], dim=1) - # Target note indices follow prompt notes in the concatenated score. - mel2note = torch.cat([prompt_mel2note, target_mel2note + len_prompt_note], dim=1) - - # pitch_shift is semitones; each coarse F0 bin is 20 cents (×5 bins per semitone). - f0_coarse_prompt = f0_to_coarse(prompt_f0) - f0_coarse_target = f0_to_coarse(target_f0, f0_shift=pitch_shift * 5) - f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1) - - note_pitch = note_pitch.clone() - note_pitch[note_pitch > 0] = note_pitch[note_pitch > 0] + pitch_shift - note_pitch = torch.clamp(note_pitch, 0, 255) - - if prompt_mel is None: - prompt_wav = prompt_meta["wav"] - with self._stage_timer("mel"): - prompt_mel = self._mel_from_wav(prompt_wav) - - if prompt_mel.shape[1] > len_prompt_mel: - prompt_mel = prompt_mel[:, :len_prompt_mel, :] - elif prompt_mel.shape[1] < len_prompt_mel: - logger.warning( - "prompt_mel length %s is shorter than metadata frames %s; padding mel.", - prompt_mel.shape[1], - len_prompt_mel, - ) - prompt_mel = F.pad(prompt_mel, (0, 0, 0, len_prompt_mel - prompt_mel.shape[1])) - - with self._stage_timer("cond_encode"): - cond = self._encode_condition( - note_pitch=note_pitch, - note_type=note_type, - note_text=note_text, - mel2note=mel2note, - f0_coarse=f0_coarse, - ) - - with self._stage_timer("cfm"): - generated_mel = self._run_flow_matching_loop( - prompt_mel, - cond, - num_inference_steps, - guidance_scale, - generator=generator, - ) - with self._stage_timer("vocoder"): - return self._mel_to_audio(generated_mel) - - def infer_svs_batch( - self, - payload: dict[str, Any], - *, - extra_args: dict[str, Any], - num_inference_steps: int, - guidance_scale: float, - generator: torch.Generator | None = None, - ) -> tuple[torch.Tensor, int]: - """Pure batch (non-streaming): run diffusion for each target segment in order and concat.""" - prompt_meta_raw = payload.get("prompt_meta") - target_meta_list = payload["target_meta_list"] - control = str(extra_args.get("control", "score")) - - prompt_meta = None - prompt_mel = None - if prompt_meta_raw: - prompt_meta = self._ensure_processed_meta(prompt_meta_raw) - prompt_meta = self._apply_control_mode(prompt_meta, control) - if prompt_meta and prompt_meta.get("wav") is not None: - prompt_mel = self._mel_from_wav(prompt_meta["wav"]) - - pieces: list[torch.Tensor] = [] - pitch_shift = 0 - - for idx, target_raw in enumerate(target_meta_list): - target_meta = self._ensure_processed_meta(target_raw) - target_meta = self._apply_control_mode(target_meta, control) - - if idx == 0: - pitch_shift = resolve_pitch_shift( - auto_shift=bool(extra_args.get("auto_shift", True)), - manual_shift=int(extra_args.get("pitch_shift", 0)), - prompt_f0=prompt_meta["f0"] if prompt_meta else None, - target_f0=target_meta["f0"], - prompt_note_pitch=prompt_meta["note_pitch"] if prompt_meta else None, - target_note_pitch=target_meta["note_pitch"], - ) - - seg = self._infer_svs_segment( - prompt_meta or {}, - target_meta, - pitch_shift=pitch_shift, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - prompt_mel=prompt_mel, - generator=generator, - ) - pieces.append(seg.squeeze()) - - if not pieces: - return torch.zeros(1, 0, device=self.device, dtype=torch.float32), pitch_shift - - full = torch.cat(pieces, dim=-1) - return full.unsqueeze(0), pitch_shift - - @staticmethod - def _apply_control_mode(meta: dict, control: str) -> dict: - meta["note_pitch"] = meta["note_pitch"] if control == "score" else None - meta["f0"] = meta["f0"] if control == "melody" else None - return meta - - def _ensure_processed_meta(self, meta: dict) -> dict: - if isinstance(meta.get("phoneme"), torch.Tensor): - return meta - return self.metadata_processor.process(meta, None) - - @torch.inference_mode() - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - return self._forward_batch_from_request( - req, - kind="svs", - metadata_key="f0_shift", - infer_batch_fn=self.infer_svs_batch, - prepare_extra_args=lambda extra_args, _sampling: normalize_svs_control_extra_args(extra_args), - ) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - del weights - self._load_soulx_checkpoint("model.pt") - return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py deleted file mode 100644 index aed87f757..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Native SoulX-Singer preprocess helpers.""" diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py deleted file mode 100644 index 5f54d47c9..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py +++ /dev/null @@ -1,86 +0,0 @@ -import re - -_EN_WORD_RE = re.compile(r"^[A-Za-z]+(?:'[A-Za-z]+)*$") -_ZH_WORD_RE = re.compile(r"[\u4e00-\u9fff]") - -EN_FLAG = "en_" -YUE_FLAG = "yue_" -ZH_FLAG = "zh_" - -_g2p_zh = None -_g2p_en = None - - -def _ensure_g2p_backends() -> None: - global _g2p_zh, _g2p_en - if _g2p_zh is None: - from g2pM import G2pM - - _g2p_zh = G2pM() - if _g2p_en is None: - from g2p_en import G2p as G2pE - - _g2p_en = G2pE() - - -def is_chinese_char(word: str) -> bool: - if len(word) != 1: - return False - return bool(_ZH_WORD_RE.fullmatch(word)) - - -def is_english_word(word: str) -> bool: - if not word: - return False - return bool(_EN_WORD_RE.fullmatch(word)) - - -def g2p_cantonese(sent): - import ToJyutping - - return ToJyutping.get_jyutping_list(sent) - - -def g2p_mandarin(sent): - _ensure_g2p_backends() - return _g2p_zh(sent, tone=True, char_split=False) - - -def g2p_english(word): - _ensure_g2p_backends() - return _g2p_en(word) - - -def g2p_transform(words, lang): - zh_words = [] - transformed_words = [0] * len(words) - - for idx, w in enumerate(words): - if w == "": - transformed_words[idx] = w - continue - - w = w.replace("?", "").replace(".", "").replace("!", "").replace(",", "") - - if is_chinese_char(w): - zh_words.append([idx, w]) - else: - if is_english_word(w): - w = EN_FLAG + "-".join(g2p_english(w.lower())) - else: - w = "" - transformed_words[idx] = w - - sent = "".join([k[1] for k in zh_words]) - - if len(sent) > 0: - if lang == "Cantonese": - g2pm_rst = g2p_cantonese(sent) - g2pm_rst = [YUE_FLAG + k[1] for k in g2pm_rst] - else: - g2pm_rst = g2p_mandarin(sent) - g2pm_rst = [ZH_FLAG + k for k in g2pm_rst] - for p, w in zip([k[0] for k in zh_words], g2pm_rst): - transformed_words[p] = w - - return transformed_words diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py deleted file mode 100644 index 8a8611121..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py +++ /dev/null @@ -1,137 +0,0 @@ -"""Metadata merge helpers for SoulX-Singer preprocess.""" - -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np -import soundfile as sf - - -@dataclass -class SegmentMetadata: - item_name: str - wav_fn: str - language: str - start_time_ms: int - end_time_ms: int - note_text: list[str] - note_dur: list[float] - note_pitch: list[int] - note_type: list[int] - origin_wav_fn: str | None = None - wav: np.ndarray | None = None - - def get(self, name: str, default: Any = None) -> Any: - return getattr(self, name, default) - - -def _merge_group( - audio: np.ndarray, - sample_rate: int, - segments: list[SegmentMetadata | dict[str, Any]], - output_dir: Path | None, - *, - end_extension_ms: int = 0, -) -> SegmentMetadata: - if not segments: - raise ValueError("segments must not be empty") - - words: list[str] = [] - durs: list[float] = [] - pitches: list[int] = [] - types: list[int] = [] - - for i, seg in enumerate(segments): - if i > 0: - prev_seg = segments[i - 1] - gap_ms = seg.get("start_time_ms", 0) - prev_seg.get("end_time_ms", 0) - if gap_ms > 0: - words.append("") - durs.append(gap_ms / 1000.0) - pitches.append(0) - types.append(1) - - words.extend(seg.get("note_text", [])) - durs.extend(seg.get("note_dur", [])) - pitches.extend(seg.get("note_pitch", [])) - types.extend(seg.get("note_type", [])) - - if end_extension_ms > 0: - words.append("") - durs.append(end_extension_ms / 1000.0) - pitches.append(0) - types.append(1) - - merged_words, merged_durs, merged_pitches, merged_types = [], [], [], [] - for word, dur, pitch, note_type in zip(words, durs, pitches, types): - if merged_words and word == "" and merged_words[-1] == "": - merged_durs[-1] += dur - else: - merged_words.append(word) - merged_durs.append(dur) - merged_pitches.append(pitch) - merged_types.append(note_type) - - languages = [s.get("language", "Mandarin") for s in segments if s.get("language")] - language = max(languages, key=languages.count) if languages else "Mandarin" - - start_ms = int(segments[0].get("start_time_ms", 0)) - end_ms = int(segments[-1].get("end_time_ms", 0)) + end_extension_ms - start_sample = int(start_ms * sample_rate // 1000) - end_sample = int(end_ms * sample_rate // 1000) - - first_item_name = segments[0].get("item_name", "segment") - song_prefix = "_".join(str(first_item_name).split("_")[:-1]) - item_name = f"{song_prefix}_{start_ms}_{end_ms}" - - wav_fn = "" - segment_audio = audio[start_sample:end_sample] - if output_dir is not None: - output_dir.mkdir(parents=True, exist_ok=True) - wav_path = output_dir / f"{item_name}.wav" - sf.write(wav_path, segment_audio, sample_rate) - wav_fn = str(wav_path) - - return SegmentMetadata( - item_name=item_name, - wav_fn=wav_fn, - language=language, - start_time_ms=start_ms, - end_time_ms=end_ms, - note_text=merged_words, - note_dur=merged_durs, - note_pitch=merged_pitches, - note_type=merged_types, - origin_wav_fn=segments[0].get("origin_wav_fn", ""), - wav=segment_audio if output_dir is None else None, - ) - - -def convert_metadata(item: SegmentMetadata, f0: np.ndarray | None = None) -> dict[str, Any]: - """Convert merged segment metadata into SoulX JSON metadata format. - - Matches upstream ``convert_metadata``: F0 always comes from ``*_f0.npy`` written - by ``extract_f0_file`` on the merged segment wav, not from an in-memory array. - """ - from vllm_omni.diffusion.models.soulx_singer.preprocess.g2p import g2p_transform - - if item.wav_fn: - f0_path = str(Path(item.wav_fn).with_suffix("")) + "_f0.npy" - f0_arr = np.load(f0_path) - elif f0 is not None: - f0_arr = f0 - else: - raise ValueError("convert_metadata requires item.wav_fn with sidecar f0 or an explicit f0 array") - - return { - "index": item.item_name, - "language": item.language, - "time": [item.start_time_ms, item.end_time_ms], - "duration": " ".join(f"{d:.2f}" for d in item.note_dur), - "text": " ".join(item.note_text), - "phoneme": " ".join(g2p_transform(item.note_text, item.language)), - "note_pitch": " ".join(map(str, item.note_pitch)), - "note_type": " ".join(map(str, item.note_type)), - "f0": " ".join(f"{x:.1f}" for x in f0_arr), - } diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py deleted file mode 100644 index ff4e4b7ca..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py +++ /dev/null @@ -1,106 +0,0 @@ -"""SoulX-Singer payload schema (guidance/precomputed tables, dummy builders).""" - -from __future__ import annotations - -from typing import Any - -import torch - -SOULX_PRECOMPUTED_KEYS_BY_KIND: dict[str, tuple[str, ...]] = { - "svs": ("prompt_metadata_path", "target_metadata_path", "audio_path"), - "svc": ("prompt_wav_path", "target_wav_path", "prompt_f0_path", "target_f0_path"), -} - - -# Key under which the SoulX preprocess payload is stored in -# prompt["additional_information"] or diffusion output metadata. -SOULX_PREPROCESSED_KEY = "soulx_preprocessed" - - -def get_soulx_preprocessed_payload(prompt: dict[str, Any]) -> dict[str, Any] | None: - """Extract attached payload from a diffusion request prompt dict.""" - additional = prompt.get("additional_information") or {} - payload = additional.get(SOULX_PREPROCESSED_KEY) - return payload if isinstance(payload, dict) else None - - -def relocate_tensors(payload: dict[str, Any], device: torch.device | str) -> dict[str, Any]: - """Recursively move all tensors inside the payload to the target device. - - Called after decode when the payload is about to be consumed by the - diffusion model (which may live on a different device / rank). - """ - if isinstance(payload, torch.Tensor): - return payload.to(device, non_blocking=True) - if isinstance(payload, dict): - return {key: relocate_tensors(value, device) for key, value in payload.items()} - if isinstance(payload, list): - return [relocate_tensors(item, device) for item in payload] - return payload - - -def consume_payload( - req: Any, - expected_kind: str, - device: torch.device | str, -) -> dict[str, Any]: - """Read the payload attached by pre_process_func, validate the kind, - and relocate tensors to the compute device. - """ - prompts = getattr(req, "prompts", None) or [] - if not prompts: - raise ValueError(f"SoulX-Singer {expected_kind} request has no prompts.") - prompt = prompts[0] - if isinstance(prompt, str): - raise ValueError(f"SoulX-Singer {expected_kind} forward requires pre_process_func output on the prompt.") - payload = get_soulx_preprocessed_payload(prompt) - if payload is None or payload.get("kind") != expected_kind: - raise ValueError( - f"SoulX-Singer {expected_kind} forward requires additional_information['soulx_preprocessed'] " - "produced by pre_process_func." - ) - return relocate_tensors(payload, device) - - -def has_precomputed(extra_args: dict[str, Any], kind: str) -> bool: - keys = SOULX_PRECOMPUTED_KEYS_BY_KIND[kind] - return all(extra_args.get(key) for key in keys) - - -def build_dummy_payload(kind: str, device: torch.device) -> dict[str, Any]: - mel_frames = 4 - hop_size = 480 - n_fft = 1920 - wav_samples = max(mel_frames * hop_size, n_fft) - if kind == "svc": - voiced_f0 = torch.full((1, mel_frames), 100.0, device=device, dtype=torch.float64) - return { - "kind": "svc", - "prompt_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), - "target_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), - "prompt_f0": voiced_f0.clone(), - "target_f0": voiced_f0.clone(), - } - return { - "kind": "svs", - "prompt_meta": { - "phoneme": torch.zeros(1, 2, device=device, dtype=torch.long), - "note_pitch": torch.zeros(1, mel_frames, device=device, dtype=torch.long), - "note_type": torch.zeros(1, mel_frames, device=device, dtype=torch.long), - "mel2note": torch.zeros(1, mel_frames, device=device, dtype=torch.long), - "f0": torch.zeros(1, mel_frames, device=device, dtype=torch.float32), - "wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), - }, - "target_meta_list": [ - { - "text": "dummy", - "phoneme": " ".join([""] * mel_frames), - "f0": " ".join(["100"] * mel_frames), - "language": "Mandarin", - "note_pitch": " ".join(["0"] * mel_frames), - "note_type": " ".join(["1"] * mel_frames), - "duration": " ".join(["1"] * mel_frames), - "time": [0, max(int(wav_samples / 24000 * 1000), 100)], - } - ], - } diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py deleted file mode 100644 index a8784fcba..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py +++ /dev/null @@ -1,223 +0,0 @@ -"""Shared ``pre_process_func`` helpers for SoulX-Singer pipelines (batch / full-payload only).""" - -from pathlib import Path -from typing import Any - -import torch -import torch.nn as nn -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import SoulXPreprocessPipeline -from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PRECOMPUTED_KEYS_BY_KIND, - SOULX_PREPROCESSED_KEY, - build_dummy_payload, - get_soulx_preprocessed_payload, - has_precomputed, -) -from vllm_omni.diffusion.models.soulx_singer.utils import ( - MetadataProcessor, - load_config, - resolve_phoneset_path, - validate_soulx_extra_args, -) -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.inputs.data import OmniTextPrompt - -logger = init_logger(__name__) - - -def is_warmup_request(request: OmniDiffusionRequest) -> bool: - return request.is_dummy_run() - - -def resolve_preprocess_audio(prompt: dict[str, Any], extra_args: dict[str, Any]) -> tuple[Any, Any]: - mm = prompt.get("multi_modal_data") or {} - prompt_audio = mm.get("audio") if isinstance(mm, dict) else None - if prompt_audio is None: - prompt_audio = extra_args.get("prompt_audio") - return prompt_audio, extra_args.get("target_audio") - - -def _preprocess_inputs_missing_error(kind: str) -> ValueError: - return ValueError( - f"SoulX-Singer {kind} requires precomputed paths " - f"{list(SOULX_PRECOMPUTED_KEYS_BY_KIND[kind])}, an attached soulx_preprocessed payload " - f"(multi-stage preprocess stage), or run with pipeline soulxsinger_{kind}." - ) - - -def build_precomputed_payload( - kind: str, - extra_args: dict[str, Any], - *, - metadata_processor, - sample_rate: int | None, - device, -) -> dict[str, Any]: - if kind == "svs": - return SoulXPreprocessPipeline.build_svs_payload_from_paths( - prompt_metadata_path=str(extra_args["prompt_metadata_path"]), - target_metadata_path=str(extra_args["target_metadata_path"]), - audio_path=str(extra_args["audio_path"]), - metadata_processor=metadata_processor, - ) - if sample_rate is None or device is None: - raise ValueError("svc precomputed payload requires sample_rate and device") - return SoulXPreprocessPipeline.build_svc_payload_from_paths( # type: ignore - prompt_wav_path=str(extra_args["prompt_wav_path"]), - target_wav_path=str(extra_args["target_wav_path"]), - prompt_f0_path=str(extra_args["prompt_f0_path"]), - target_f0_path=str(extra_args["target_f0_path"]), - sample_rate=int(sample_rate), - device=device, - ) - - -def build_warmup_payload( - kind: str, - *, - metadata_processor, - device, - sample_rate: int | None, -) -> dict[str, Any]: - if kind == "svs": - if metadata_processor is None: - raise ValueError("SVS warmup requires metadata_processor") - payload = build_dummy_payload("svs", torch.device("cpu")) - dummy_prompt = payload["prompt_meta"] - processed = metadata_processor.process(dict(payload["target_meta_list"][0])) - payload["prompt_meta"] = { - key: value.clone() if isinstance(value, torch.Tensor) else value for key, value in processed.items() - } - if isinstance(dummy_prompt.get("wav"), torch.Tensor): - payload["prompt_meta"]["wav"] = dummy_prompt["wav"].clone() - return payload - if device is None or sample_rate is None: - raise ValueError("svc warmup requires device and sample_rate") - return build_dummy_payload("svc", device) - - -def build_preprocess_payload( - kind: str, - *, - prompt: dict[str, Any], - extra_args: dict[str, Any], - preprocess: nn.Module, - metadata_processor, - device, - sample_rate: int, -) -> dict[str, Any]: - if has_precomputed(extra_args, kind): - return build_precomputed_payload( - kind, - extra_args, - metadata_processor=metadata_processor, - sample_rate=sample_rate, - device=device, - ) - - prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) - if prompt_audio is None or target_audio is None: - raise _preprocess_inputs_missing_error(kind) - - if kind == "svc": - return preprocess.build_svc_payload_from_audio( # type: ignore - prompt_audio=prompt_audio, - target_audio=target_audio, - sample_rate=sample_rate, - device=device, - vocal_sep=extra_args.get("vocal_sep"), - ) - - return preprocess.build_svs_payload_from_audio( # type: ignore - prompt_audio=prompt_audio, - target_audio=target_audio, - metadata_processor=metadata_processor, - language=str(extra_args.get("language", "Mandarin")), - vocal_sep=extra_args.get("vocal_sep"), - prompt_vocal_sep=extra_args.get("prompt_vocal_sep"), - target_vocal_sep=extra_args.get("target_vocal_sep"), - prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"), - target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"), - ) - - -def attach_preprocess_for_diffusion_request( - request: OmniDiffusionRequest, - *, - kind: str, - metadata_processor=None, - sample_rate: int | None = None, - device=None, -) -> OmniDiffusionRequest: - """Resolve preprocess payload for stage-1 diffusion (warmup / IPC / precomputed paths).""" - - extra_args = validate_soulx_extra_args( - kind, - dict(getattr(request.sampling_params, "extra_args", None) or {}), - ) - prompt = request.prompt - prompt = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt - - if is_warmup_request(request): - request.sampling_params.num_inference_steps = 1 - payload = build_warmup_payload( - kind, - metadata_processor=metadata_processor, - device=device, - sample_rate=sample_rate, - ) - elif get_soulx_preprocessed_payload(prompt): - request.prompt = prompt - if kind == "svs": - extra_args = normalize_svs_control_extra_args(extra_args) - request.sampling_params.extra_args = extra_args - return request - elif has_precomputed(extra_args, kind): - payload = build_precomputed_payload( - kind, - extra_args, - metadata_processor=metadata_processor, - sample_rate=sample_rate, - device=device, - ) - else: - raise _preprocess_inputs_missing_error(kind) - - if payload.get("kind") != kind: - raise ValueError(f"Invalid {kind} preprocess payload kind: {payload.get('kind')}") - prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload - request.prompt = prompt - - if kind == "svs": - extra_args = normalize_svs_control_extra_args(extra_args) - request.sampling_params.extra_args = extra_args - return request - - -def normalize_svs_control_extra_args(extra_args: dict[str, Any]) -> dict[str, Any]: - control = extra_args.get("control") - if control is None: - control = "score" - logger.warning("control is not provided, using 'score' as default") - elif control not in ("score", "melody"): - raise ValueError(f"Invalid control: {control}. Must be one of: ['score', 'melody']") - extra_args["control"] = control - return extra_args - - -def build_metadata_processor(od_config: OmniDiffusionConfig): - model_dir = od_config.model - config_path = Path(model_dir) / "config.yaml" - hf_config = load_config(str(config_path)) - audio_config = hf_config.audio - - return MetadataProcessor( - hop_size=audio_config.hop_size, - sample_rate=audio_config.sample_rate, - phoneset_path=resolve_phoneset_path(model_dir), - device=str(get_local_device()), - ) diff --git a/vllm_omni/diffusion/models/soulx_singer/utils.py b/vllm_omni/diffusion/models/soulx_singer/utils.py deleted file mode 100644 index 4f201c407..000000000 --- a/vllm_omni/diffusion/models/soulx_singer/utils.py +++ /dev/null @@ -1,565 +0,0 @@ -import json -from pathlib import Path -from typing import Any, Literal - -import numpy as np -import torch -import torchaudio -from huggingface_hub import snapshot_download -from omegaconf import DictConfig, OmegaConf -from vllm.logger import init_logger - -from vllm_omni.diffusion.data import OmniDiffusionConfig - -logger = init_logger(__name__) - -_MODEL_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer" -_PREPROCESS_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer-Preprocess" - - -def _patch_torchaudio_load() -> None: - """Borrowed from MOSS-TTS-Nano. Patch torchaudio.load to use soundfile - if torchcodec is unavailable. - """ - try: - import torchaudio - - torchaudio # noqa - import torchcodec # noqa: F401 - - return - except Exception: - pass - - import soundfile as sf - - def _soundfile_load(path, frame_offset=0, num_frames=-1, normalize=True, channels_first=True, format=None): - data, sr = sf.read(str(path), dtype="float32", always_2d=True) - if frame_offset > 0: - data = data[frame_offset:] - if num_frames > 0: - data = data[:num_frames] - waveform = torch.from_numpy(data) - if channels_first: - waveform = waveform.T - return waveform, sr - - def _soundfile_save(path, src, sample_rate, channels_first=True, **kwargs): - wav = src.detach().cpu().float().numpy() - if channels_first and wav.ndim == 2: - wav = wav.T - sf.write(str(path), wav, sample_rate) - - try: - import torchaudio - - torchaudio.load = _soundfile_load - torchaudio.save = _soundfile_save - logger.info("Patched torchaudio.load/save to use soundfile (torchcodec unavailable)") - except Exception as e: - logger.warning("Could not patch torchaudio: %s", e) - - -def load_config(config_path: str | Path) -> DictConfig: - """ - Load a configuration file and optionally merges it with a base configuration. - - Args: - config_path (Path): Path to the configuration file. - """ - # Load the initial configuration from the given path - config = OmegaConf.load(str(config_path)) - - # Check if there is a base configuration specified and merge if necessary - if config.get("base_config", None) is not None: - base_config = OmegaConf.load(str(config.get("base_config"))) - config = OmegaConf.merge(base_config, config) - - return config - - -SoulXKind = Literal["svs", "svc"] - -_SVS_ARCHITECTURE = "SoulXSingerPipeline" -_SVC_ARCHITECTURE = "SoulXSingerSVCPipeline" -_ARCHITECTURE_TO_KIND: dict[str, SoulXKind] = { - _SVS_ARCHITECTURE: "svs", - _SVC_ARCHITECTURE: "svc", -} - - -def load_model_config_json(model_dir: str | Path) -> dict[str, Any]: - """Load ``config.json`` from a SoulX-Singer model view directory.""" - cfg_path = Path(model_dir) / "config.json" - if not cfg_path.is_file(): - raise FileNotFoundError( - f"SoulX-Singer config.json not found at {cfg_path}. " - "Create one with architectures ['SoulXSingerPipeline'] (SVS) or " - "['SoulXSingerSVCPipeline'] (SVC); see examples/offline_inference/text_to_speech/README.md." - ) - with cfg_path.open(encoding="utf-8") as f: - return json.load(f) - - -def resolve_soulx_kind(model_dir: str | Path) -> SoulXKind: - """Resolve SVS vs SVC from ``config.json`` ``architectures[0]``.""" - cfg = load_model_config_json(model_dir) - archs = cfg.get("architectures") or [] - if not archs: - raise ValueError( - f"SoulX-Singer config.json at {Path(model_dir) / 'config.json'} is missing 'architectures'. " - "Set architectures to ['SoulXSingerPipeline'] for SVS or ['SoulXSingerSVCPipeline'] for SVC." - ) - arch = archs[0] - kind = _ARCHITECTURE_TO_KIND.get(arch) - if kind is None: - raise ValueError( - f"Unknown SoulX-Singer architecture {arch!r} in config.json. Expected one of {list(_ARCHITECTURE_TO_KIND)}." - ) - return kind - - -def validate_soulx_extra_args(kind: SoulXKind, extra_args: dict[str, Any] | None) -> dict[str, Any]: - """Validate request ``extra_args`` against the model mode from config.json.""" - from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( - SOULX_PRECOMPUTED_KEYS_BY_KIND, - has_precomputed, - ) - - extra_args = dict(extra_args or {}) - declared = extra_args.pop("kind", None) - if declared is not None and declared != kind: - logger.warning( - "Ignoring extra_args kind=%r; model config.json declares %r.", - declared, - kind, - ) - - other_kind: SoulXKind = "svc" if kind == "svs" else "svs" - if has_precomputed(extra_args, other_kind): - raise ValueError( - f"SoulX-Singer config.json declares {kind!r}, but request supplies " - f"{other_kind} precomputed paths {list(SOULX_PRECOMPUTED_KEYS_BY_KIND[other_kind])}." - ) - - if kind == "svc": - for key in ("language", "control"): - if extra_args.get(key) is not None: - logger.warning( - "Ignoring SVS-only extra_arg %r=%r on an SVC model (config.json).", - key, - extra_args.pop(key), - ) - for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svs"]: - if extra_args.get(key): - raise ValueError(f"SoulX-Singer SVC model (config.json) does not accept SVS precomputed key {key!r}.") - else: - for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svc"]: - if extra_args.get(key): - raise ValueError(f"SoulX-Singer SVS model (config.json) does not accept SVC precomputed key {key!r}.") - - return extra_args - - -# ---------------- utils for model weights ---------------- - - -def resolve_preprocess_weights_root(od_config: OmniDiffusionConfig) -> Path: - """Locate preprocess weights on disk or download from Hugging Face.""" - - # default local dir - if Path(od_config.model).parent.is_dir(): - local_dir = Path(od_config.model).parent / "SoulX-Singer-Preprocess" - else: - local_dir = Path(__file__).parents[5] / "pretrained" / "SoulX-Singer-Preprocess" - - local_dir = local_dir.expanduser() - if local_dir.is_dir(): - logger.info("Using SoulX preprocess weights from %s", local_dir) - return local_dir - - logger.info( - "SoulX preprocess weights not found locally; downloading %s", - _PREPROCESS_WEIGHTS_REPO, - ) - downloaded = snapshot_download( - _PREPROCESS_WEIGHTS_REPO, - allow_patterns=["*"], - local_dir=local_dir, - ) - return Path(downloaded) - - -def preprocess_weight_paths(weights_root: Path) -> dict[str, str]: - """Map upstream relative paths to absolute paths under ``weights_root``.""" - root = weights_root.resolve() - return { - "rmvpe": str(root / "rmvpe/rmvpe.pt"), - "sep_ckpt": str(root / "mel-band-roformer-karaoke/mel_band_roformer_karaoke_becruily.ckpt"), - "sep_config": str(root / "mel-band-roformer-karaoke/config_karaoke_becruily.yaml"), - "asr_zh": str(root / "speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"), - "asr_en": str(root / "parakeet-tdt-0.6b-v2/parakeet-tdt-0.6b-v2.nemo"), - "rosvot": str(root / "rosvot/rosvot/model.pt"), - } - - -def resolve_phoneset_path(model_dir: str) -> str: - """Resolve phoneme vocabulary; must be present in model directory.""" - candidates = ( - Path(model_dir) / "phoneme" / "phone_set.json", - Path(model_dir) / "phone_set.json", - ) - for path in candidates: - if path.is_file(): - return str(path) - raise FileNotFoundError( - "SoulX-Singer phoneset not found. Expected one of: " - f"{[str(p) for p in candidates]}. " - "Copy phone_set.json into the model directory. " - "See `examples/offline_inference/text_to_speech/README.md` for instructions." - ) - - -# ---------------- utils for data processing ---------------- - - -def load_wav(wav_path: str, sample_rate: int) -> torch.Tensor: - """Load wav file and resample to target sample rate. - - Args: - wav_path (str): Path to wav file. - sample_rate (int): Target sample rate. - - Returns: - torch.Tensor: Waveform tensor with shape (1, T). - """ - _patch_torchaudio_load() - waveform, sr = torchaudio.load(wav_path) - - if sr != sample_rate: - waveform = torchaudio.functional.resample(waveform, sr, sample_rate) - - if len(waveform.shape) > 1 and waveform.shape[0] > 1: - waveform = torch.mean(waveform, dim=0, keepdim=True) - - return waveform - - -def f0_to_coarse(f0, f0_bin=361, f0_min=32.7031956625, f0_shift=0): - """ - Convert continuous F0 values to discrete F0 bins (SIL and C1 - B6, 361 bins). - args: - f0: continuous F0 values - f0_bin: number of F0 bins - f0_min: minimum F0 value - f0_shift: shift value for F0 bins - returns: - f0_coarse: discrete F0 bins - """ - is_torch = isinstance(f0, torch.Tensor) - uv_mask = f0 <= 0 - - if is_torch: - f0_min_tensor = f0.new_tensor(f0_min) - f0_safe = torch.maximum(f0, f0_min_tensor) - f0_cents = 1200 * torch.log2(f0_safe / f0_min_tensor) - else: - f0_safe = np.maximum(f0, f0_min) - f0_cents = 1200 * np.log2(f0_safe / f0_min) - - f0_coarse = (f0_cents / 20) + 1 - - if is_torch: - f0_coarse = torch.round(f0_coarse).long() - f0_coarse = torch.clamp(f0_coarse, min=1, max=f0_bin - 1) - else: - f0_coarse = np.rint(f0_coarse).astype(int) - f0_coarse = np.clip(f0_coarse, 1, f0_bin - 1) - - f0_coarse[uv_mask] = 0 - - if f0_shift != 0: - if is_torch: - voiced = f0_coarse > 0 - if voiced.any(): - shifted = f0_coarse[voiced] + f0_shift - f0_coarse[voiced] = torch.clamp(shifted, 1, f0_bin - 1) - else: - voiced = f0_coarse > 0 - if np.any(voiced): - shifted = f0_coarse[voiced] + f0_shift - f0_coarse[voiced] = np.clip(shifted, 1, f0_bin - 1) - - return f0_coarse - - -def resolve_pitch_shift( - *, - auto_shift: bool, - manual_shift: int, - prompt_f0: torch.Tensor | None = None, - target_f0: torch.Tensor | None = None, - prompt_note_pitch: torch.Tensor | None = None, - target_note_pitch: torch.Tensor | None = None, -) -> int: - """Resolve semitone shift for auto_shift; ``f0_to_coarse(..., f0_shift=shift * 5)`` consumes it.""" - if not auto_shift or manual_shift != 0: - return int(manual_shift) - - if prompt_note_pitch is not None and target_note_pitch is not None: - target_pitched = target_note_pitch[target_note_pitch >= 1] - prompt_pitched = prompt_note_pitch[prompt_note_pitch >= 1] - if target_pitched.numel() > 0 and prompt_pitched.numel() > 0: - shift = torch.round(prompt_pitched.median() - target_pitched.median()) - if torch.isfinite(shift): - return int(shift.item()) - - if prompt_f0 is not None and target_f0 is not None: - target_voiced = target_f0[target_f0 > 0] - prompt_voiced = prompt_f0[prompt_f0 > 0] - if target_voiced.numel() > 0 and prompt_voiced.numel() > 0: - shift = torch.round(torch.log2(prompt_voiced.median() / target_voiced.median()) * 1200 / 100) - if torch.isfinite(shift): - return int(shift.item()) - - return 0 - - -class MetadataProcessor: - """Data processor for SoulX-Singer""" - - def __init__( - self, - hop_size: int, - sample_rate: int, - phoneset_path: str, - device: str = "cuda", - ): - """Initialize data processor. - - Args: - hop_size (int): Hop size in samples. - sample_rate (int): Sample rate in Hz. - phoneset_path (str): Path to phoneme set JSON file. - device (str): Device to use for tensor operations. - """ - self.hop_size = hop_size - self.sample_rate = sample_rate - self.device = device - self.load_phoneme_id_map(phoneset_path) - - def load_phoneme_id_map(self, phoneset_path: str): - with open(phoneset_path, encoding="utf-8") as f: - phoneset = json.load(f) - self.phone2idx = {ph: idx for idx, ph in enumerate(phoneset)} - - def merge_phoneme(self, meta): - merged_items = [] - - duration = [float(x) for x in meta["duration"].split()] - phoneme = [str(x).replace("", "") for x in meta["phoneme"].split()] - note_pitch = [int(x) for x in meta["note_pitch"].split()] - note_type = [int(x) if phoneme[i] != "" else 1 for i, x in enumerate(meta["note_type"].split())] - - for i in range(len(phoneme)): - if ( - i > 0 - and phoneme[i] == phoneme[i - 1] == "" - and note_type[i] == note_type[i - 1] - and note_pitch[i] == note_pitch[i - 1] - ): - merged_items[-1][1] += duration[i] - else: - merged_items.append([phoneme[i], duration[i], note_pitch[i], note_type[i]]) - - meta["phoneme"] = [x[0] for x in merged_items] - meta["duration"] = [x[1] for x in merged_items] - meta["note_pitch"] = [x[2] for x in merged_items] - meta["note_type"] = [x[3] for x in merged_items] - - return meta - - def preprocess( - self, - note_duration: list[float], - phonemes: list[str], - note_pitch: list[int], - note_type: list[int], - ): - """ - Insert and for each note. - Get aligned indices for each frame. - - Args: - note_duration: Duration of each note in seconds - phonemes: Phoneme sequence for each note - note_pitch: Pitch value for each note - note_type: Type value for each note - - """ - sample_rate = self.sample_rate - hop_size = self.hop_size - duration = sum(note_duration) * sample_rate / hop_size - mel2note = torch.zeros(int(duration), dtype=torch.long, device=self.device) - - ph_locations = [] # idx at mel scale and length - new_phonemes = [] - dur_sum = 0 - - note2origin = [] - - for ph_idx in range(len(phonemes)): - dur = int(np.round(dur_sum * sample_rate / hop_size)) - dur = min(dur, len(mel2note) - 1) - new_phonemes.append("") - note2origin.append(ph_idx) - if phonemes[ph_idx][:3] == "en_": - en_phs = ["en_" + x for x in phonemes[ph_idx][3:].split("-")] + [ - "" - ] # between en words in one note - ph_locations.append([dur, max(1, len(en_phs))]) - new_phonemes.extend(en_phs) - note2origin.extend([ph_idx] * len(en_phs)) - else: - ph_locations.append([dur, 1]) - new_phonemes.append(phonemes[ph_idx]) - note2origin.append(ph_idx) - new_phonemes.append("") - note2origin.append(ph_idx) - dur_sum += note_duration[ph_idx] - - ph_idx = 1 - for idx, (i, j) in enumerate(ph_locations): - next_phoneme_start = ph_locations[idx + 1][0] if idx < len(ph_locations) - 1 else len(mel2note) - if i >= len(mel2note) or i + j > len(mel2note): - break - if i < len(mel2note) and mel2note[i] > 0: - logger.warning(f"warning: overlap of {idx}: {mel2note[i]}") - while i < len(mel2note) and mel2note[i] > 0: - i += 1 - mel2note[i] = ph_idx - k = i + 1 - while k + j < next_phoneme_start: - mel2note[k : k + j] = torch.arange(ph_idx, ph_idx + j, device=self.device) + 1 - k += j - mel2note[next_phoneme_start - 1] = ph_idx + j + 1 - ph_idx += j + 2 # + ph repeats + - - new_phonemes = [""] + new_phonemes - new_note_pitch = [0] + [note_pitch[k] for k in note2origin] - new_note_type = [1] + [note_type[k] for k in note2origin] - - return { - "phoneme": torch.tensor([self.phone2idx[x] for x in new_phonemes], device=self.device).unsqueeze(0), - "note_pitch": torch.tensor(new_note_pitch, device=self.device).unsqueeze(0), - "note_type": torch.tensor(new_note_type, device=self.device).unsqueeze(0), - "mel2note": mel2note.unsqueeze(0), - } - - def process(self, meta: dict, wav_path: str | None = None) -> dict[str, torch.Tensor | None]: - meta = self.merge_phoneme(meta) - - item = self.preprocess( - meta["duration"], - meta["phoneme"], - meta["note_pitch"], - meta["note_type"], - ) - - f0 = [float(x) for x in meta.get("f0", "").split()] - min_frame = min(item["mel2note"].shape[1], len(f0)) if len(f0) > 0 else item["mel2note"].shape[1] - item["f0"] = torch.tensor(f0, device=self.device)[:min_frame].unsqueeze(0).float() if len(f0) > 0 else None - item["mel2note"] = item["mel2note"][:, :min_frame] - - if wav_path is not None: - waveform = load_wav(wav_path, self.sample_rate) - item["wav"] = waveform.to(self.device)[:, : min_frame * self.hop_size] - - return item - - -def build_vocal_segments( - f0, - *, - f0_rate: int = 50, - ignore_silent_frames_thresh: int = 5, - min_duration_sec_per_segment: float = 5.0, - max_duration_sec_per_segment: float = 30.0, - num_overlaps: int = 1, - ignore_silent_frames: bool = True, -) -> tuple[list[tuple[float, float]], list[tuple[float, float]]]: - """Build vocal segments from an F0 contour for chunked SVC inference.""" - if isinstance(f0, torch.Tensor): - f0_np = f0.detach().float().cpu().numpy() - else: - f0_np = np.asarray(f0, dtype=np.float32) - f0_np = np.squeeze(f0_np) - - total_frames = int(f0_np.shape[0]) - if total_frames == 0: - return [], [] - - min_frames = max(1, int(round(min_duration_sec_per_segment * f0_rate))) - max_frames = max(1, int(round(max_duration_sec_per_segment * f0_rate))) - - split_points = [0] - - def append_split_point(point: int): - point = int(max(0, min(point, total_frames))) - while point - split_points[-1] > max_frames: - split_points.append(split_points[-1] + max_frames) - if point > split_points[-1]: - split_points.append(point) - - idx = 0 - while idx < total_frames: - if f0_np[idx] == 0: - run_start = idx - while idx < total_frames and f0_np[idx] == 0: - idx += 1 - run_end = idx - if (run_end - run_start) >= ignore_silent_frames_thresh: - split_point = max(run_end - 5, (run_start + run_end) // 2) - append_split_point(split_point) - else: - idx += 1 - append_split_point(total_frames) - - segments: list[tuple[float, float]] = [] - overlap_segments: list[tuple[float, float]] = [] - - def append_segment(start_idx: int, end_idx: int, overlaps: int = num_overlaps): - segments.append((split_points[start_idx] / f0_rate, split_points[end_idx] / f0_rate)) - overlap_start_idx = start_idx - if start_idx > 0 and (split_points[end_idx] - split_points[start_idx - overlaps]) <= max_frames: - overlap_start_idx = start_idx - overlaps - overlap_segments.append((split_points[overlap_start_idx] / f0_rate, split_points[end_idx] / f0_rate)) - - segment_start, segment_end = 0, 1 - while segment_start < len(split_points) - 1: - while ( - segment_end < len(split_points) and (split_points[segment_end] - split_points[segment_start]) < min_frames - ): - segment_end += 1 - - if segment_end >= len(split_points): - append_segment(segment_start, len(split_points) - 1, overlaps=num_overlaps) - break - append_segment(segment_start, segment_end, overlaps=num_overlaps) - segment_start = segment_end - segment_end = segment_start + 1 - - if ignore_silent_frames: - filtered_idx = [] - for i, seg in enumerate(overlap_segments): - start_frame = int(seg[0] * f0_rate) - end_frame = int(seg[1] * f0_rate) - seg_frames = end_frame - start_frame - voice_frames = np.sum(f0_np[start_frame:end_frame] > 0) - if voice_frames / seg_frames > 0.05 and voice_frames >= 10: - filtered_idx.append(i) - - overlap_segments = [overlap_segments[i] for i in filtered_idx] - segments = [segments[i] for i in filtered_idx] - - return overlap_segments, segments diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py index 41f12f654..f8681408a 100644 --- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py +++ b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py @@ -81,8 +81,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov # Picked up by ``supports_audio_output`` in the diffusion engine so the # default stage metadata reports ``final_output_type="audio"`` and the - # ``multimodal_output`` payload includes the sample rate (mirrors the - # contract introduced for AudioX in #2077). + # ``multimodal_output`` payload includes the sample rate. support_audio_output: ClassVar[bool] = True audio_sample_rate: ClassVar[int] = 44100 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index f9bc9f04f..f92053e52 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -33,17 +33,19 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler -from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel +from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention, WanTransformer3DModel from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) DEBUG_PERF = False WAN_SAMPLE_SOLVER_CHOICES = {"unipc", "euler"} +FASTWAN_DMD_TIMESTEPS = (1000.0, 757.0, 522.0) +FASTWAN_DMD_SCHEDULER_SHIFT = 8.0 def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any: @@ -64,6 +66,30 @@ def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any: ) +def load_wan_weights_with_optional_gate(model: nn.Module, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load Wan weights and discard optional VSA gates absent from the checkpoint.""" + gate_param_names = {name for name, _ in model.named_parameters() if ".to_gate_compress." in name} + has_gate_compress_weights = False + + def tracked_weights(): + nonlocal has_gate_compress_weights + for name, weight in weights: + if ".to_gate_compress." in name: + has_gate_compress_weights = True + yield name, weight + + loaded_weights = AutoWeightsLoader(model).load_weights(tracked_weights()) + setattr(model, "has_gate_compress_weights", has_gate_compress_weights) + if not has_gate_compress_weights: + for module in model.modules(): + if isinstance(module, WanSelfAttention): + module.to_gate_compress = None + + # Optional gate parameters are absent from public FullAttn/DMD checkpoints. + loaded_weights.update(gate_param_names) + return loaded_weights + + def resolve_wan_sample_solver(req: OmniDiffusionRequest, default: str = "unipc") -> str: extra_args = getattr(req.sampling_params, "extra_args", {}) or {} raw = extra_args.get("sample_solver", default) @@ -204,6 +230,8 @@ def get_wan22_post_process_func( output_type: str = "np", sampling_params=None, ): + if sampling_params is not None and sampling_params.output_type is not None: + output_type = sampling_params.output_type if output_type == "latent": return video video_metadata = {} @@ -233,6 +261,8 @@ def get_wan22_pre_process_func( prompt = request.prompt multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + has_image = raw_image is not None and (not isinstance(raw_image, list) or bool(raw_image)) + request.batch_compatibility_key = ("wan22_image_condition", has_image) if isinstance(prompt, str): prompt = OmniTextPrompt(prompt=prompt) if "additional_information" not in prompt: @@ -288,6 +318,7 @@ class Wan22Pipeline( SupportsComponentDiscovery, WanLoraLoaderMixin, ): + supports_request_batch = True _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -309,6 +340,7 @@ class Wan22Pipeline( # Read model_index.json to detect expand_timesteps mode (for TI2V-5B) self.expand_timesteps = False + self.is_dmd = False self.has_transformer_2 = False if local_files_only: model_index_path = os.path.join(model, "model_index.json") @@ -316,6 +348,7 @@ class Wan22Pipeline( with open(model_index_path) as f: model_index = json.load(f) self.expand_timesteps = model_index.get("expand_timesteps", False) + self.is_dmd = model_index.get("_class_name") == "WanDMDPipeline" # Check if this is a two-stage model (MoE with transformer_2) transformer_2_path = os.path.join(model, "transformer_2") self.has_transformer_2 = os.path.exists(transformer_2_path) @@ -328,6 +361,7 @@ class Wan22Pipeline( with open(model_index_path) as f: model_index = json.load(f) self.expand_timesteps = model_index.get("expand_timesteps", False) + self.is_dmd = model_index.get("_class_name") == "WanDMDPipeline" # Check transformer_2 from model_index transformer_2_info = model_index.get("transformer_2", [None, None]) self.has_transformer_2 = transformer_2_info[0] is not None @@ -417,6 +451,10 @@ class Wan22Pipeline( else: self.transformer_2 = None + for transformer in (self.transformer, self.transformer_2): + if transformer is not None: + transformer.preserve_vsa_all_blocks = self.is_dmd + # Store the active transformer config if load_transformer: self.transformer_config = self.transformer.config @@ -425,8 +463,14 @@ class Wan22Pipeline( else: raise RuntimeError("No transformer loaded") - self._sample_solver = "unipc" - self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0 + self._sample_solver = "euler" if self.is_dmd else "unipc" + self._flow_shift = ( + FASTWAN_DMD_SCHEDULER_SHIFT + if self.is_dmd + else od_config.flow_shift + if od_config.flow_shift is not None + else 5.0 + ) self.scheduler = build_wan_scheduler(self._sample_solver, self._flow_shift) self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 @@ -475,6 +519,7 @@ class Wan22Pipeline( attention_kwargs: dict[str, Any], latent_condition: torch.Tensor | None = None, first_frame_mask: torch.Tensor | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor | AsyncLatents: if attention_kwargs is None: attention_kwargs = {} @@ -556,32 +601,51 @@ class Wan22Pipeline( cfg_normalize=False, ) - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + if self.is_dmd: + pred_clean = self.scheduler.predict_clean(noise_pred, latents, t).to(noise_pred.dtype) + if step_idx + 1 < len(timesteps): + noise = randn_tensor( + latents.shape, + generator=generator, + device=latents.device, + dtype=pred_clean.dtype, + ) + latents = self.scheduler.add_noise(pred_clean, noise, timesteps[step_idx + 1]) + else: + latents = pred_clean + else: + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) pbar.update() return latents - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - prompt: str | None = None - negative_prompt: str | None = None - prompt_embeds: torch.Tensor | None = None - negative_prompt_embeds: torch.Tensor | None = None - if len(req.prompts) > 1: - raise ValueError( - """This model only supports a single prompt, not a batched request.""", - """Please pass in a single prompt object or string, or a single-item list.""", - ) - if len(req.prompts) == 1: - first_prompt = req.prompts[0] - prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") - negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts] + negative_prompts = [ + None if isinstance(prompt, str) else prompt.get("negative_prompt") for prompt in req.prompts + ] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": None, + "negative_prompt_embeds": None, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt: list[str] | None = prompt_texts if prompt_embeds is None else None + negative_prompt: list[str] | None = None + if negative_prompt_embeds is None and any(value is not None for value in negative_prompts): + negative_prompt = [value or "" for value in negative_prompts] - if not prompt: - raise ValueError("Prompt is required for Wan2.2 generation.") + if prompt is not None and not all(prompt): + raise ValueError("Prompt is required for Wan2.2 generation when prompt_embeds are not provided.") - height = req.sampling_params.height or 480 - width = req.sampling_params.width or 832 - num_frames = req.sampling_params.num_frames or 81 + height = common.height or 480 + width = common.width or 832 + num_frames = common.num_frames or 81 # Ensure dimensions are compatible with VAE and patch size # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size) @@ -589,19 +653,28 @@ class Wan22Pipeline( mod_value = self.vae_scale_factor_spatial * patch_size[1] # 16*2=32 for TI2V, 8*2=16 for I2V height = (height // mod_value) * mod_value width = (width // mod_value) * mod_value - num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps + if self.is_dmd: + # The checkpoint was distilled for these three transitions. Ignore + # request-level step counts, including the engine's 1-step warmup. + num_steps = len(FASTWAN_DMD_TIMESTEPS) + else: + num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps - output_type = req.sampling_params.output_type or "np" + output_type = common.output_type or "np" + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 attention_kwargs: dict | None = None - guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=4.0) + guidance_low, guidance_high = resolve_wan_guidance_scales( + common, + default_guidance_scale=1.0 if self.is_dmd else 4.0, + ) # record guidance for properties self._guidance_scale = guidance_low self._guidance_scale_2 = guidance_high # Prefer engine-configured boundary_ratio, but allow per-request fallback. - boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio + boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio if boundary_ratio is None: boundary_ratio = 0.875 @@ -633,40 +706,59 @@ class Wan22Pipeline( # Fallback to text_encoder dtype if no transformer loaded dtype = self.text_encoder.dtype - # Seed / generator - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) + generator = req.collate_request_generators(num_outputs_per_prompt, None) + request_latents = req.collate_request_tensors("latents", None) if DEBUG_PERF: # Sync GPU before timing to ensure accurate measurements current_omni_platform.synchronize() _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) + do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0 + if prompt_embeds is None: + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) + else: + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + elif do_classifier_free_guidance: + _, negative_prompt_embeds = self.encode_prompt( + prompt=[""] * req.num_reqs, + negative_prompt=negative_prompt, + do_classifier_free_guidance=True, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) if DEBUG_PERF: current_omni_platform.synchronize() _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000 - sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver) - flow_shift = resolve_wan_flow_shift(req, self.od_config) - if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6: - self.scheduler = build_wan_scheduler(sample_solver, flow_shift) - self._sample_solver = sample_solver - self._flow_shift = flow_shift + if self.is_dmd: + timesteps = torch.tensor(FASTWAN_DMD_TIMESTEPS, device=device, dtype=torch.float32) + else: + first_request = req.requests[0] + sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver) + flow_shift = resolve_wan_flow_shift(first_request, self.od_config) + if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6: + self.scheduler = build_wan_scheduler(sample_solver, flow_shift) + self._sample_solver = sample_solver + self._flow_shift = flow_shift - # Timesteps - self.scheduler.set_timesteps(num_steps, device=device) - timesteps = self.scheduler.timesteps + self.scheduler.set_timesteps(num_steps, device=device) + timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) boundary_timestep = None if boundary_ratio is not None: @@ -674,37 +766,41 @@ class Wan22Pipeline( if DEBUG_PERF: _t_latent_prep_start = time.perf_counter() - multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(raw_image, list): - if len(raw_image) > 1: - logger.warning( - """Received a list of image. Only a single image is supported by this model.""" - """Taking only the first image for now.""" - ) - raw_image = raw_image[0] - if raw_image is None: - image = None - elif isinstance(raw_image, str): - image = PIL.Image.open(raw_image) - else: - image = cast(PIL.Image.Image | torch.Tensor, raw_image) + images: list[PIL.Image.Image | torch.Tensor | None] = [] + for request_prompt in req.prompts: + multi_modal_data = request_prompt.get("multi_modal_data", {}) if not isinstance(request_prompt, str) else {} + raw_image = multi_modal_data.get("image") + if isinstance(raw_image, list): + if len(raw_image) > 1: + logger.warning("Received multiple images for one Wan request; using only the first image.") + raw_image = raw_image[0] if raw_image else None + if isinstance(raw_image, str): + raw_image = PIL.Image.open(raw_image) + images.append(cast(PIL.Image.Image | torch.Tensor | None, raw_image)) latent_condition = None first_frame_mask = None - if self.expand_timesteps and image is not None: + if self.expand_timesteps and any(image is not None for image in images): + if not all(image is not None for image in images): + raise ValueError("Cannot batch Wan requests with a mix of provided and missing image conditions.") # I2V mode: encode image and prepare condition from diffusers.video_processor import VideoProcessor video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - # Preprocess image - if isinstance(image, PIL.Image.Image): - image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) - image_tensor = video_processor.preprocess(image, height=height, width=width) - else: - image_tensor = image + image_tensors = [] + for image in images: + assert image is not None + if isinstance(image, PIL.Image.Image): + image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + image_tensor = video_processor.preprocess(image, height=height, width=width) + else: + image_tensor = image.unsqueeze(0) if image.ndim == 3 else image + image_tensors.append(image_tensor) + image_tensor = DiffusionRequestBatch.collate_tensors(image_tensors, "image condition", None) + assert image_tensor is not None + image_tensor = image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0) # Use out_channels for noise latents (not in_channels which includes condition) num_channels_latents = self.transformer_config.out_channels @@ -720,7 +816,7 @@ class Wan22Pipeline( dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, ) # Encode image condition @@ -731,7 +827,6 @@ class Wan22Pipeline( image_tensor = image_tensor.unsqueeze(2) # [B, C, 1, H, W] image_tensor = image_tensor.to(device=device, dtype=self.vae.dtype) latent_condition = retrieve_latents(self.vae.encode(image_tensor), sample_mode="argmax") - latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1) # Normalize condition latents latents_mean = ( @@ -747,7 +842,7 @@ class Wan22Pipeline( # Create mask: 0 for first frame (condition), 1 for rest (to denoise) first_frame_mask = torch.ones( - 1, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device + batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device ) first_frame_mask[:, :, 0] = 0 else: @@ -762,7 +857,7 @@ class Wan22Pipeline( dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, ) if DEBUG_PERF: current_omni_platform.synchronize() @@ -785,6 +880,7 @@ class Wan22Pipeline( attention_kwargs=attention_kwargs, latent_condition=latent_condition, first_frame_mask=first_frame_mask, + generator=generator, ) # Wan2.2 is prone to out of memory errors when predicting large videos @@ -839,8 +935,13 @@ class Wan22Pipeline( _t_pipeline_wall_ms - _t_stages_sum, ) - return DiffusionOutput( - output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + return split_diffusion_output_by_request( + DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_outputs_per_prompt, ) def predict_noise( @@ -966,9 +1067,7 @@ class Wan22Pipeline( return latents def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) def check_inputs( self, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index a2e444bb0..b0cd35458 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -16,7 +16,6 @@ import torchvision.transforms.functional as TF from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel -from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.sequence import IntermediateTensors from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -36,6 +35,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( build_wan_scheduler, create_transformer_from_config, load_transformer_config, + load_wan_weights_with_optional_gate, resolve_wan_flow_shift, resolve_wan_guidance_scales, resolve_wan_sample_solver, @@ -45,7 +45,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3 from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -65,6 +65,8 @@ def get_wan22_i2v_post_process_func( output_type: str = "np", sampling_params=None, ): + if sampling_params is not None and sampling_params.output_type is not None: + output_type = sampling_params.output_type if output_type == "latent": return video video_metadata = {} @@ -84,6 +86,32 @@ def get_wan22_i2v_post_process_func( return post_process_func +def _normalize_i2v_last_image( + value: str | PIL.Image.Image | torch.Tensor | list[object] | None, +) -> PIL.Image.Image | torch.Tensor | None: + """Normalize the optional I2V last-frame condition for batching.""" + if value is None: + return None + + if isinstance(value, list): + if not value: + return None + if len(value) != 1: + raise ValueError("I2V accepts at most one last_image.") + value = value[0] + + if isinstance(value, str): + value = PIL.Image.open(value).convert("RGB") + + if not isinstance(value, (PIL.Image.Image, torch.Tensor)): + raise TypeError( + f"Unsupported last_image format {value.__class__}. " + "Expected a file path, PIL.Image.Image, torch.Tensor, or None." + ) + + return value + + def get_wan22_i2v_pre_process_func( od_config: OmniDiffusionConfig, ): @@ -109,6 +137,9 @@ def get_wan22_i2v_pre_process_func( """Please correctly set `"multi_modal_data": {"image": , …}`""", ) image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image + last_image = _normalize_i2v_last_image(multi_modal_data.get("last_image")) # type: ignore[union-attr] + prompt["multi_modal_data"]["last_image"] = last_image # type: ignore[index] + request.batch_compatibility_key = ("wan22_i2v_last_image", last_image is not None) # Calculate dimensions based on aspect ratio if not provided if request.sampling_params.height is None or request.sampling_params.width is None: @@ -157,6 +188,7 @@ class Wan22I2VPipeline( Wan2.2-style I2V (with expand_timesteps for TI2V-5B). """ + supports_request_batch = True _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -433,60 +465,67 @@ class Wan22I2VPipeline( quant_config = getattr(self.od_config, "quantization_config", None) return create_transformer_from_config(config, quant_config=quant_config) - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - prompt: str | None = None - negative_prompt: str | None = None - prompt_embeds: torch.Tensor | None = None - negative_prompt_embeds: torch.Tensor | None = None + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts] + negative_prompts = [ + None if isinstance(prompt, str) else prompt.get("negative_prompt") for prompt in req.prompts + ] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": None, + "negative_prompt_embeds": None, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] image_embeds: torch.Tensor | None = None - if len(req.prompts) > 1: - raise ValueError( - """This model only supports a single prompt, not a batched request.""", - """Please pass in a single prompt object or string, or a single-item list.""", - ) - if len(req.prompts) == 1: - first_prompt = req.prompts[0] - prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") - negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") + prompt: list[str] | None = prompt_texts if prompt_embeds is None else None + negative_prompt: list[str] | None = None + if negative_prompt_embeds is None and any(value is not None for value in negative_prompts): + negative_prompt = [value or "" for value in negative_prompts] + if prompt is not None and not all(prompt): + raise ValueError("Prompt is required for Wan2.2 I2V generation when prompt_embeds are not provided.") - if not prompt: - raise ValueError("Prompt is required for Wan2.2 generation.") + images: list[PIL.Image.Image | torch.Tensor] = [] + last_images: list[PIL.Image.Image | torch.Tensor | None] = [] + for request_prompt in req.prompts: + multi_modal_data = request_prompt.get("multi_modal_data", {}) if not isinstance(request_prompt, str) else {} + raw_image = multi_modal_data.get("image") + if raw_image is None: + raise ValueError("Image is required for I2V generation.") + if isinstance(raw_image, list): + if len(raw_image) > 1: + logger.warning("Received multiple images for one I2V request; using only the first image.") + raw_image = raw_image[0] + if isinstance(raw_image, str): + raw_image = PIL.Image.open(raw_image).convert("RGB") + images.append(cast(PIL.Image.Image | torch.Tensor, raw_image)) - # Get image from request - multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if raw_image is None: - raise ValueError("Image is required for I2V generation.") - if isinstance(raw_image, list): - if len(raw_image) > 1: - logger.warning( - """Received a list of image. Only a single image is supported by this model.""" - """Taking only the first image for now.""" - ) - raw_image = raw_image[0] - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image) - else: - image = cast(PIL.Image.Image | torch.Tensor, raw_image) + last_image = multi_modal_data.get("last_image") + if isinstance(last_image, str): + last_image = PIL.Image.open(last_image).convert("RGB") + last_images.append(cast(PIL.Image.Image | torch.Tensor | None, last_image)) + if any(image is not None for image in last_images) and not all(image is not None for image in last_images): + raise ValueError("Cannot batch I2V requests with a mix of provided and missing last_image conditions.") - last_image: PIL.Image.Image | torch.Tensor | None = None - if multi_modal_data is not None: - last_image = multi_modal_data.get("last_image", None) + height = common.height or 480 + width = common.width or 832 + num_frames = common.num_frames or 81 + num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps - height = req.sampling_params.height or 480 - width = req.sampling_params.width or 832 - num_frames = req.sampling_params.num_frames or 81 - num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps - - output_type = req.sampling_params.output_type or "np" + output_type = common.output_type or "np" + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 attention_kwargs: dict | None = None - guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=5.0) + guidance_low, guidance_high = resolve_wan_guidance_scales(common, default_guidance_scale=5.0) self._guidance_scale = guidance_low self._guidance_scale_2 = guidance_high - boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio + boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio if boundary_ratio is None: boundary_ratio = 0.875 logger.warning("boundary_ratio is required for I2V generation. using default value 0.875") @@ -495,7 +534,7 @@ class Wan22I2VPipeline( self.check_inputs( prompt=prompt, negative_prompt=negative_prompt, - image=image, + image=images, height=height, width=width, prompt_embeds=prompt_embeds, @@ -513,10 +552,8 @@ class Wan22I2VPipeline( device = self.device dtype = self.transformer.dtype - # Generator setup - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) + generator = req.collate_request_generators(num_outputs_per_prompt, None) + request_latents = req.collate_request_tensors("latents", None) if DEBUG_PERF: # Sync GPU before timing to ensure accurate measurements @@ -524,15 +561,33 @@ class Wan22I2VPipeline( _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) + do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0 + if prompt_embeds is None: + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) + else: + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + elif do_classifier_free_guidance: + _, negative_prompt_embeds = self.encode_prompt( + prompt=[""] * req.num_reqs, + negative_prompt=negative_prompt, + do_classifier_free_guidance=True, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) if DEBUG_PERF: current_omni_platform.synchronize() @@ -543,12 +598,18 @@ class Wan22I2VPipeline( if DEBUG_PERF: _t_img_enc_start = time.perf_counter() if self.has_image_encoder and self.transformer.config.image_dim is not None: - if image_embeds is None: - if last_image is None: - image_embeds = self.encode_image(image, device) - else: - image_embeds = self.encode_image([image, last_image], device) - image_embeds = image_embeds.repeat(batch_size, 1, 1) + if all(last_image is None for last_image in last_images): + image_embeds = self.encode_image(images, device) + image_embeds = image_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + else: + image_pairs = [] + for first_image, last_image in zip(images, last_images): + assert last_image is not None + image_pairs.extend([first_image, last_image]) + image_embeds = self.encode_image(image_pairs, device) + image_embeds = image_embeds.view(req.num_reqs, 2, *image_embeds.shape[1:]) + image_embeds = image_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + image_embeds = image_embeds.flatten(0, 1) image_embeds = image_embeds.to(dtype) else: image_embeds = None @@ -557,8 +618,9 @@ class Wan22I2VPipeline( current_omni_platform.synchronize() _t_img_enc_ms = (time.perf_counter() - _t_img_enc_start) * 1000 - sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver) - flow_shift = resolve_wan_flow_shift(req, self.od_config) + first_request = req.requests[0] + sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver) + flow_shift = resolve_wan_flow_shift(first_request, self.od_config) if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6: self.scheduler = build_wan_scheduler(sample_solver, flow_shift) self._sample_solver = sample_solver @@ -582,23 +644,36 @@ class Wan22I2VPipeline( video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - if isinstance(image, PIL.Image.Image): - image = TF.to_tensor(image).to(device) - image_tensor = video_processor.preprocess(image, height=height, width=width) - else: - image_tensor = image - image_tensor = image_tensor.to(device=device, dtype=torch.float32) - - # Handle last_image if provided - if last_image is not None: - if isinstance(last_image, PIL.Image.Image): - image = TF.to_tensor(last_image).to(device) - last_image_tensor = video_processor.preprocess(last_image, height=height, width=width) + image_tensors = [] + for image in images: + if isinstance(image, PIL.Image.Image): + image = TF.to_tensor(image).to(device) + image_tensor = video_processor.preprocess(image, height=height, width=width) else: - last_image_tensor = last_image - last_image_tensor = last_image_tensor.to(device=device, dtype=torch.float32) - else: - last_image_tensor = None + image_tensor = image.unsqueeze(0) if image.ndim == 3 else image + image_tensors.append(image_tensor.to(device=device, dtype=torch.float32)) + image_tensor = DiffusionRequestBatch.collate_tensors(image_tensors, "image condition", None) + assert image_tensor is not None + image_tensor = image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0) + + last_image_tensor = None + if all(last_image is not None for last_image in last_images): + last_image_tensors = [] + for last_image in last_images: + assert last_image is not None + if isinstance(last_image, PIL.Image.Image): + last_image = TF.to_tensor(last_image).to(device) + current_last_image = video_processor.preprocess(last_image, height=height, width=width) + else: + current_last_image = last_image.unsqueeze(0) if last_image.ndim == 3 else last_image + last_image_tensors.append(current_last_image.to(device=device, dtype=torch.float32)) + last_image_tensor = DiffusionRequestBatch.collate_tensors( + last_image_tensors, + "last image condition", + None, + ) + assert last_image_tensor is not None + last_image_tensor = last_image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0) latents, condition, first_frame_mask = self.prepare_latents( image=image_tensor, @@ -610,7 +685,7 @@ class Wan22I2VPipeline( dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, last_image=last_image_tensor, ) @@ -694,8 +769,13 @@ class Wan22I2VPipeline( _t_pipeline_wall_ms - _t_stages_sum, ) - return DiffusionOutput( - output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + return split_diffusion_output_by_request( + DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_outputs_per_prompt, ) def predict_noise( @@ -849,7 +929,6 @@ class Wan22I2VPipeline( # Encode through VAE latent_condition = retrieve_latents(self.vae.encode(video_condition), sample_mode="argmax") - latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1) # Normalize latents latents_mean = ( @@ -866,7 +945,7 @@ class Wan22I2VPipeline( if self.expand_timesteps: # TI2V-5B style: create mask where first frame is 0 (condition), rest is 1 (to denoise) first_frame_mask = torch.ones( - 1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device + batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device ) first_frame_mask[:, :, 0] = 0 return latents, latent_condition, first_frame_mask @@ -892,7 +971,9 @@ class Wan22I2VPipeline( condition = torch.concat([mask_lat_size, latent_condition], dim=1) # For non-expand mode, first_frame_mask is not used in the same way - first_frame_mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device) + first_frame_mask = torch.ones( + batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device + ) return latents, condition, first_frame_mask @@ -933,9 +1014,7 @@ class Wan22I2VPipeline( raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.") def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) # --------------------------------------------------------------------------- diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index 52a42cb79..bc688fbc4 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Wan2.2 Speech-to-Video (S2V) Pipeline for vLLM-Omni. @@ -22,7 +22,6 @@ from diffusers.utils.torch_utils import randn_tensor from torch import nn from torchvision import transforms from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel, Wav2Vec2ForCTC, Wav2Vec2Processor -from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.multimodal.media.audio import load_audio from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -36,6 +35,7 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( load_transformer_config, + load_wan_weights_with_optional_gate, retrieve_latents, ) from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import ( @@ -43,7 +43,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import ( ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -63,6 +63,30 @@ _S2V_DEFAULT_NEG_PROMPT = ( _AUDIO_VIDEO_RATE = 30 +def _s2v_audio_condition_key(audio: str | np.ndarray | None) -> tuple[Any, ...]: + if isinstance(audio, str): + # File paths are deliberately conservative: two different files may + # have different durations even when their request metadata matches. + return ("path", audio) + if isinstance(audio, np.ndarray): + return ("array", tuple(audio.shape), str(audio.dtype)) + return (type(audio).__name__,) + + +def _make_clip_generators( + seeds: list[int | None], + request_generators: list[torch.Generator | None], + clip_index: int, + device: torch.device, +) -> list[torch.Generator | None]: + return [ + torch.Generator(device=device).manual_seed((seed if seed is not None else 0) + clip_index) + if seed is not None or request_generator is None + else request_generator + for seed, request_generator in zip(seeds, request_generators) + ] + + def _linear_interpolation( features: torch.Tensor, input_fps: float, @@ -337,6 +361,12 @@ def get_wan22_s2v_pre_process_func( prompt["additional_information"]["init_first_frame"] = ( multi_modal_data.get("init_first_frame", False) if multi_modal_data is not None else False ) + request.batch_compatibility_key = ( + "wan22_s2v_condition", + bool(prompt["additional_information"]["init_first_frame"]), + prompt["additional_information"].get("num_repeat"), + _s2v_audio_condition_key(raw_audio), + ) request.prompt = prompt return request @@ -486,6 +516,7 @@ class Wan22S2VPipeline( """ # Default config values from Wan2.2/wan/configs/wan_s2v_14B.py + supports_request_batch = True _DEFAULT_MOTION_FRAMES = 73 _DEFAULT_INFER_FRAMES = 80 _DEFAULT_FPS = 16 @@ -942,7 +973,6 @@ class Wan22S2VPipeline( """Forward pass through the S2V transformer to predict noise. With return_dict=False, the model returns (output,) where output is [B, C, T, H, W]. - We squeeze batch dim since S2V processes one sample at a time. """ if current_model is None: current_model = self.transformer @@ -952,8 +982,8 @@ class Wan22S2VPipeline( if isinstance(result, list): return result[0] if isinstance(result, tuple): - return result[0].squeeze(0) - return result.sample.squeeze(0) + return result[0] + return result.sample def diffuse( self, @@ -962,7 +992,7 @@ class Wan22S2VPipeline( prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None, guidance_scale: float, - clip_generator: torch.Generator, + clip_generator: torch.Generator | list[torch.Generator] | None, dtype: torch.dtype, device: torch.device, max_seq_len: int, @@ -971,13 +1001,13 @@ class Wan22S2VPipeline( ref_latents: torch.Tensor, motion_frames: list[int], drop_first_motion: bool, - positive_audio_emb: torch.Tensor, - negative_audio_emb: torch.Tensor | None, + positive_audio_emb: dict[str, torch.Tensor], + negative_audio_emb: dict[str, torch.Tensor] | None, ) -> torch.Tensor: - """Denoising diffusion loop for one S2V clip. + """Denoising diffusion loop for a batch of compatible S2V clips. Args: - latents: Initial noise tensor [C, T, H, W] + latents: Initial noise tensor [B, C, T, H, W] timesteps: Denoising timesteps prompt_embeds: Text embeddings [1, seq_len, dim] negative_prompt_embeds: Negative text embeddings (optional) @@ -995,7 +1025,7 @@ class Wan22S2VPipeline( negative_audio_emb: Precomputed audio embeddings for negative prompt (optional) Returns: - Denoised latents [C, T, H, W] + Denoised latents [B, C, T, H, W] """ do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None @@ -1004,13 +1034,13 @@ class Wan22S2VPipeline( self._current_timestep = t self.record_denoise_step(step_idx, t) - latent_model_input = [latents.to(device)] - timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device) + latent_model_input = latents.to(device) + timestep = t.expand(latents.shape[0]).to(device) positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep, - "encoder_hidden_states": prompt_embeds[0:1], + "encoder_hidden_states": prompt_embeds, "cond_states": cond_latents, "motion_latents": input_motion_latents, "ref_latents": ref_latents, @@ -1023,7 +1053,7 @@ class Wan22S2VPipeline( negative_kwargs = { "hidden_states": latent_model_input, "timestep": timestep, - "encoder_hidden_states": negative_prompt_embeds[0:1], + "encoder_hidden_states": negative_prompt_embeds, "cond_states": cond_latents, "motion_latents": input_motion_latents, "ref_latents": ref_latents, @@ -1044,12 +1074,12 @@ class Wan22S2VPipeline( # Scheduler step latents = self.scheduler.step( - noise_pred.unsqueeze(0), + noise_pred, t, - latents.unsqueeze(0), + latents, return_dict=False, - generator=clip_generator, - )[0].squeeze(0) + generator=clip_generator if isinstance(clip_generator, torch.Generator) else None, + )[0] pbar.update() @@ -1080,111 +1110,178 @@ class Wan22S2VPipeline( prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, **kwargs, - ) -> DiffusionOutput: + ) -> list[DiffusionOutput]: """Run S2V generation — may produce multiple autoregressive clips. This method mirrors ``WanS2V.generate()``, reorganized into the vLLM-Omni pipeline structure. """ - # ---- Extract params from request ---- - if len(req.prompts) > 1: - raise ValueError("S2V only supports a single prompt per request.") - - if len(req.prompts) == 1: - prompt_data = req.prompts[0] - prompt = prompt_data if isinstance(prompt_data, str) else prompt_data.get("prompt") - negative_prompt = None if isinstance(prompt_data, str) else prompt_data.get("negative_prompt") + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt_texts: list[str] = [] + negative_prompts: list[str] = [] + images: list[PIL.Image.Image | None] = [] + audio_paths: list[str | np.ndarray | None] = [] + requested_repeats: list[int | None] = [] + init_first_frames: list[bool] = [] + for prompt_data in req.prompts: multi_modal_data = prompt_data.get("multi_modal_data", {}) if not isinstance(prompt_data, str) else {} additional_info = prompt_data.get("additional_information", {}) if not isinstance(prompt_data, str) else {} - if image is None: - raw_image = multi_modal_data.get("image", None) - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image).convert("RGB") - elif isinstance(raw_image, PIL.Image.Image): - image = raw_image + prompt_texts.append( + prompt_data if isinstance(prompt_data, str) else (prompt_data.get("prompt") or prompt or "") + ) + request_negative_prompt = ( + negative_prompt if isinstance(prompt_data, str) else prompt_data.get("negative_prompt", negative_prompt) + ) + negative_prompts.append(request_negative_prompt or _S2V_DEFAULT_NEG_PROMPT) - if audio_path is None: - audio_path = additional_info.get("audio_path") - if audio_path is None: - audio_path = multi_modal_data.get("audio") + raw_image = multi_modal_data.get("image", image) + if isinstance(raw_image, list): + raw_image = raw_image[0] if raw_image else None + if isinstance(raw_image, str): + raw_image = PIL.Image.open(raw_image).convert("RGB") + images.append(raw_image if isinstance(raw_image, PIL.Image.Image) else None) - if pose_video is None: - pose_video = additional_info.get("pose_video") + request_audio_path = additional_info.get("audio_path", audio_path) + if request_audio_path is None: + request_audio_path = multi_modal_data.get("audio") + audio_paths.append(request_audio_path) + requested_repeats.append(additional_info.get("num_repeat", num_repeat)) + init_first_frames.append(additional_info.get("init_first_frame", init_first_frame)) - init_first_frame = additional_info.get("init_first_frame", init_first_frame) - - if negative_prompt is None: - negative_prompt = _S2V_DEFAULT_NEG_PROMPT - - height = req.sampling_params.height or height - width = req.sampling_params.width or width + height = common.height or height + width = common.width or width # num_frames defaults to 1 (image-model sentinel); S2V needs many frames. - req_num_frames = req.sampling_params.num_frames + req_num_frames = common.num_frames infer_frames = ( infer_frames or (req_num_frames if req_num_frames and req_num_frames > 1 else None) or self._DEFAULT_INFER_FRAMES ) - num_steps = req.sampling_params.num_inference_steps or num_inference_steps + num_steps = common.num_inference_steps or num_inference_steps + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale + if common.guidance_scale_provided: + guidance_scale = common.guidance_scale self._guidance_scale = guidance_scale # ---- Validate ---- - self.check_inputs(prompt, image, audio_path, height, width, prompt_embeds) + if len(set(init_first_frames)) != 1: + raise ValueError("Batched S2V requests must have matching init_first_frame values.") + init_first_frame = init_first_frames[0] + for request_prompt, request_image, request_audio_path in zip(prompt_texts, images, audio_paths): + self.check_inputs( + request_prompt if prompt_embeds is None else None, + request_image, + request_audio_path, + height, + width, + prompt_embeds, + ) device = self.device dtype = self.transformer.dtype if hasattr(self.transformer, "dtype") else torch.bfloat16 - - if generator is None: - generator = req.sampling_params.generator - seed = req.sampling_params.seed - if generator is None and seed is not None: - generator = torch.Generator(device=device).manual_seed(seed) - if seed is None: - seed = 0 + request_generators = req.collate_request_generators(num_outputs_per_prompt, None) + batch_size = req.num_reqs * num_outputs_per_prompt + if isinstance(request_generators, list): + generators = request_generators + else: + generators = [request_generators] * batch_size + seeds = [sampling.seed for sampling in sampling_params_list for _ in range(num_outputs_per_prompt)] + request_latents = req.collate_request_tensors("latents", None) # ---- 1. Text encoding ---- if prompt_embeds is None: prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, + prompt=prompt_texts, + negative_prompt=negative_prompts, do_classifier_free_guidance=self.do_classifier_free_guidance, - num_videos_per_prompt=1, - max_sequence_length=512, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, device=device, dtype=dtype, ) else: prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) if negative_prompt_embeds is not None: negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) # ---- 2. Audio encoding ---- - audio_emb, audio_num_repeat, target_video_frames = self.encode_audio( - audio_path, infer_frames=infer_frames, device=device, dtype=dtype - ) - if num_repeat is None or num_repeat > audio_num_repeat: - num_repeat = audio_num_repeat + audio_embeddings: list[torch.Tensor] = [] + effective_repeats: list[int] = [] + target_video_frames: list[int] = [] + for request_audio_path, requested_repeat in zip(audio_paths, requested_repeats): + assert request_audio_path is not None + audio_emb, audio_num_repeat, request_target_video_frames = self.encode_audio( + request_audio_path, infer_frames=infer_frames, device=device, dtype=dtype + ) + audio_embeddings.append(audio_emb.repeat_interleave(num_outputs_per_prompt, dim=0)) + effective_repeats.append( + audio_num_repeat + if requested_repeat is None or requested_repeat > audio_num_repeat + else requested_repeat + ) + target_video_frames.append(request_target_video_frames) + + if len(set(effective_repeats)) != 1: + raise ValueError(f"Batched S2V requests must have matching num_repeat values, got {effective_repeats}.") + if len(set(target_video_frames)) != 1: + raise ValueError( + "Batched S2V requests must have matching audio embedding time lengths; " + f"got target video frames {target_video_frames}." + ) + num_repeat = effective_repeats[0] + audio_emb = DiffusionRequestBatch.collate_tensors(audio_embeddings, "S2V audio embeddings", None) + assert audio_emb is not None # Load raw audio waveform for muxing into output video - if isinstance(audio_path, np.ndarray): - raw_audio_waveform = audio_path.astype(np.float32) - raw_audio_sr = 16000 - else: - raw_audio_waveform, raw_audio_sr = load_audio(audio_path, sr=None, mono=True) + raw_audio_waveforms: list[np.ndarray] = [] + raw_audio_sample_rates: list[int] = [] + for request_audio_path in audio_paths: + assert request_audio_path is not None + if isinstance(request_audio_path, np.ndarray): + raw_audio_waveforms.append(request_audio_path.astype(np.float32, copy=False)) + raw_audio_sample_rates.append(16000) + else: + raw_audio_waveform, raw_audio_sr = load_audio(request_audio_path, sr=None, mono=True) + raw_audio_waveforms.append(raw_audio_waveform) + raw_audio_sample_rates.append(raw_audio_sr) + if len(set(raw_audio_sample_rates)) != 1 or len({waveform.shape for waveform in raw_audio_waveforms}) != 1: + raise ValueError("Batched S2V requests must have matching raw audio shapes and sample rates.") + raw_audio_sr = raw_audio_sample_rates[0] # ---- 3. Reference image encoding ---- - ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype) + ref_latents = DiffusionRequestBatch.collate_tensors( + [ + self.encode_ref_image(request_image, height, width, device=device) + .to(dtype=dtype) + .repeat_interleave(num_outputs_per_prompt, dim=0) + for request_image in images + if request_image is not None + ], + "S2V reference image latents", + None, + ) + assert ref_latents is not None # ---- 4. Initial motion latents (zeros) ---- motion_frames = min(self.motion_frames, infer_frames - 1) motion_pixels = torch.zeros( - [1, 3, motion_frames, height, width], + [batch_size, 3, motion_frames, height, width], dtype=dtype, device=device, ) @@ -1193,8 +1290,11 @@ class Wan22S2VPipeline( drop_first_motion = False # Place reference image in the last 6 frames of motion tensor_trans = transforms.ToTensor() - ref_tensor = tensor_trans(image).to(device=device, dtype=dtype) * 2 - 1.0 # [C, H, W] - motion_pixels[:, :, -6:] = ref_tensor.unsqueeze(0).unsqueeze(2).expand(-1, -1, 6, -1, -1) + ref_tensors = torch.stack( + [tensor_trans(request_image) for request_image in images if request_image is not None] + ).to(device=device, dtype=dtype) + ref_tensors = ref_tensors.repeat_interleave(num_outputs_per_prompt, dim=0) * 2 - 1.0 + motion_pixels[:, :, -6:] = ref_tensors.unsqueeze(2).expand(-1, -1, 6, -1, -1) motion_latents = self.prepare_motion_latents(motion_pixels, device=device).to(dtype=dtype) @@ -1211,24 +1311,29 @@ class Wan22S2VPipeline( clips: list[torch.Tensor] = [] # Keep a pixel-space buffer of the trailing motion_frames for the # autoregressive connection between clips. - videos_last_frames = torch.zeros([1, 3, motion_frames, height, width], dtype=dtype, device=device) + videos_last_frames = torch.zeros([batch_size, 3, motion_frames, height, width], dtype=dtype, device=device) for r in range(num_repeat): - # Per-clip seed - clip_seed = seed + r - clip_generator = torch.Generator(device=device).manual_seed(clip_seed) - + clip_generators = _make_clip_generators(seeds, generators, r, device) # -- Noise -- - latents = self.prepare_latents( - infer_frames=infer_frames, - height=height, - width=width, - motion_frames=motion_frames, - lat_motion_frames=lat_motion_frames, - dtype=dtype, - device=device, - generator=clip_generator, - ) + if r == 0 and request_latents is not None: + latents = request_latents.to(device=device, dtype=dtype) + else: + latents = torch.stack( + [ + self.prepare_latents( + infer_frames=infer_frames, + height=height, + width=width, + motion_frames=motion_frames, + lat_motion_frames=lat_motion_frames, + dtype=dtype, + device=device, + generator=clip_generator, + ) + for clip_generator in clip_generators + ] + ) # -- Scheduler -- self.scheduler.set_timesteps(num_steps, device=device) @@ -1242,11 +1347,11 @@ class Wan22S2VPipeline( # -- Pose condition for this clip -- # Default: zero condition (no pose driving) - # Shape: [B=1, C=16, T, H, W] — passed to transformer which iterates over batch dim + # Shape: [B, C=16, T, H, W] lat_h = height // self.vae_scale_factor_spatial lat_w = width // self.vae_scale_factor_spatial lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames - cond_latents = torch.zeros([1, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device) + cond_latents = torch.zeros([batch_size, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device) # -- Clone motion latents for this clip -- input_motion_latents = motion_latents.clone() @@ -1282,7 +1387,7 @@ class Wan22S2VPipeline( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, guidance_scale=guidance_scale, - clip_generator=clip_generator, + clip_generator=clip_generators, dtype=dtype, device=device, max_seq_len=max_seq_len, @@ -1300,11 +1405,10 @@ class Wan22S2VPipeline( self.transformer.to("cpu") current_omni_platform.empty_cache() - latents_for_decode = latents.unsqueeze(0) # [1, C, T, H, W] if not (drop_first_motion and r == 0): - decode_latents = torch.cat([motion_latents, latents_for_decode], dim=2) + decode_latents = torch.cat([motion_latents, latents], dim=2) else: - decode_latents = torch.cat([ref_latents, latents_for_decode], dim=2) + decode_latents = torch.cat([ref_latents, latents], dim=2) decode_latents = self._denormalize_latents(decode_latents) decode_latents = decode_latents.to(self.vae.dtype) @@ -1321,7 +1425,7 @@ class Wan22S2VPipeline( # Create buffer for broadcast clip_video = torch.empty( - (1, 3, total_frames, height, width), + (batch_size, 3, total_frames, height, width), device=decode_latents.device, dtype=decode_latents.dtype, ) @@ -1356,14 +1460,19 @@ class Wan22S2VPipeline( current_omni_platform.empty_cache() # ---- Concatenate all clips ---- - output = torch.cat(clips, dim=2) # [1, C, T_total, H, W] + output = torch.cat(clips, dim=2) # [B, C, T_total, H, W] - return DiffusionOutput( - output=(output, raw_audio_waveform, raw_audio_sr), - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + outputs = split_diffusion_output_by_request( + DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_outputs_per_prompt, ) + for request_output, raw_audio_waveform in zip(outputs, raw_audio_waveforms): + request_output.output = (request_output.output, raw_audio_waveform, raw_audio_sr) + return outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index d0cdb1250..58cd98c3e 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ VACE (Video Creation and Editing) Pipeline for WAN models. @@ -22,12 +22,12 @@ import PIL.Image import torch from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( Wan22Pipeline, + load_wan_weights_with_optional_gate, retrieve_latents, ) from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( @@ -35,13 +35,27 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( ) from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +def _vace_condition_structure(value: object) -> tuple: + if value is None: + return ("none",) + if isinstance(value, str): + return ("path", value) + if isinstance(value, PIL.Image.Image): + return ("image", value.size) + if isinstance(value, torch.Tensor): + return ("tensor", tuple(value.shape), str(value.dtype)) + if isinstance(value, list): + return ("list", tuple(_vace_condition_structure(item) for item in value)) + return (type(value).__name__,) + + def create_vace_transformer_from_config( config: dict, quant_config: QuantizationConfig | None = None, @@ -104,6 +118,7 @@ def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig): prompt["additional_information"] = {} if not multi_modal_data: + request.batch_compatibility_key = ("wan22_vace_condition", ("none",), ("none",), ("none",)) request.prompt = prompt return request @@ -155,6 +170,13 @@ def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig): mask = [mask] prompt["additional_information"]["mask"] = mask + request.batch_compatibility_key = ( + "wan22_vace_condition", + _vace_condition_structure(ref_images), + _vace_condition_structure(source_video), + _vace_condition_structure(mask), + ) + request.prompt = prompt return request @@ -482,7 +504,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): attention_kwargs: dict | None = None, vace_context_scale: float | list[float] = 1.0, **kwargs, - ) -> DiffusionOutput: + ) -> list[DiffusionOutput]: """Generate or edit video using VACE. The mode is determined by which inputs are provided in the request: @@ -507,47 +529,48 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): attention_kwargs: Additional kwargs for attention layers. vace_context_scale: VACE conditioning strength. """ - # Get parameters from request or arguments - if len(req.prompts) > 1: - raise ValueError( - "This model only supports a single prompt, not a batched request. " - "Please pass in a single prompt object or string, or a single-item list." - ) - - reference_images = None - source_video = None - source_mask = None - - if len(req.prompts) == 1: - first_prompt = req.prompts[0] - if isinstance(first_prompt, str): - prompt = first_prompt - else: - prompt = first_prompt.get("prompt") - negative_prompt = negative_prompt or first_prompt.get("negative_prompt") - prompt_embeds = prompt_embeds if prompt_embeds is not None else first_prompt.get("prompt_embeds") - negative_prompt_embeds = ( - negative_prompt_embeds - if negative_prompt_embeds is not None - else first_prompt.get("negative_prompt_embeds") - ) - - additional_info = first_prompt.get("additional_information", {}) - reference_images = additional_info.get("reference_images") - source_video = additional_info.get("source_video") - source_mask = additional_info.get("mask") - - if prompt is None and prompt_embeds is None: + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_texts = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] + negative_prompts = [None if isinstance(item, str) else item.get("negative_prompt") for item in req.prompts] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt = prompt_texts if prompt_embeds is None else None + negative_prompt = None + if negative_prompt_embeds is None and any(value is not None for value in negative_prompts): + negative_prompt = [value or "" for value in negative_prompts] + if prompt is not None and not all(prompt): raise ValueError("Prompt or prompt_embeds is required for VACE generation.") - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_frames = req.sampling_params.num_frames or frame_num - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - generator = req.sampling_params.generator or generator + reference_images_list = [] + source_videos = [] + source_masks = [] + for request_prompt in req.prompts: + additional_info = ( + request_prompt.get("additional_information", {}) if not isinstance(request_prompt, str) else {} + ) + reference_images_list.append(additional_info.get("reference_images")) + source_videos.append(additional_info.get("source_video")) + source_masks.append(additional_info.get("mask")) - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale + height = common.height or height + width = common.width or width + num_frames = common.num_frames or frame_num + num_inference_steps = common.num_inference_steps or num_inference_steps + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 + generator = req.collate_request_generators(num_outputs_per_prompt, None) + request_latents = req.collate_request_tensors("latents", None) + output_type = common.output_type or output_type + + if common.guidance_scale_provided: + guidance_scale = common.guidance_scale # Ensure dimensions are compatible with VAE and patch size mod_value = self.vae_scale_factor_spatial * 2 # 8 * 2 = 16 @@ -558,47 +581,47 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 num_frames = max(num_frames, 1) - self.check_inputs( - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - video=source_video, - mask=source_mask, - reference_images=reference_images, - ) + for index in range(req.num_reqs): + self.check_inputs( + prompt=None if prompt is None else prompt[index], + negative_prompt=None if negative_prompt is None else negative_prompt[index], + height=height, + width=width, + prompt_embeds=None if prompt_embeds is None else prompt_embeds[index], + negative_prompt_embeds=None if negative_prompt_embeds is None else negative_prompt_embeds[index], + video=source_videos[index], + mask=source_masks[index], + reference_images=reference_images_list[index], + ) device = self.device self._guidance_scale = guidance_scale dtype = self.transformer.dtype if self.transformer is not None else torch.bfloat16 - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) - # Encode prompts if prompt_embeds is None: - if prompt is None: - raise ValueError("Either prompt or prompt_embeds must be provided.") prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt=prompt, negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_scale > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, device=device, dtype=dtype, ) else: prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) if negative_prompt_embeds is not None: negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) elif guidance_scale > 1.0: _, negative_prompt_embeds = self.encode_prompt( - prompt="", - negative_prompt=None, + prompt=[""] * req.num_reqs, + negative_prompt=negative_prompt, do_classifier_free_guidance=True, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, device=device, dtype=dtype, ) @@ -608,24 +631,43 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): # so read it from whichever transformer exists. active_transformer = self.transformer if self.transformer is not None else self.transformer_2 if active_transformer.vace_patch_embedding is not None: - video, mask, ref_images_processed = self.preprocess_conditions( - video=source_video, - mask=source_mask, - reference_images=reference_images, - height=height, - width=width, - num_frames=num_frames, - dtype=dtype, - device=device, - ) + contexts = [] + reference_counts = [] + for index, (source_video, source_mask, reference_images) in enumerate( + zip(source_videos, source_masks, reference_images_list) + ): + video, mask, ref_images_processed = self.preprocess_conditions( + video=source_video, + mask=source_mask, + reference_images=reference_images, + height=height, + width=width, + num_frames=num_frames, + dtype=dtype, + device=device, + ) + request_generator = sampling_params_list[index].generator + if isinstance(request_generator, list): + request_generator = request_generator[0] + conditioning_latents = self.prepare_video_latents( + video, + mask, + ref_images_processed, + request_generator, + device, + ) + mask_encoded = self.prepare_masks(mask, ref_images_processed) + context = torch.cat([conditioning_latents, mask_encoded], dim=1) + contexts.append(context.repeat_interleave(num_outputs_per_prompt, dim=0)) + reference_counts.append(len(ref_images_processed[0]) if ref_images_processed else 0) - conditioning_latents = self.prepare_video_latents(video, mask, ref_images_processed, generator, device) - mask_encoded = self.prepare_masks(mask, ref_images_processed) - - # Unified VACE context: [video_latents, mask] along channels -> [B, C, T, H, W] - vace_context = torch.cat([conditioning_latents, mask_encoded], dim=1) - - num_reference_images = len(ref_images_processed[0]) if ref_images_processed else 0 + if len(set(reference_counts)) != 1: + raise ValueError( + f"Batched VACE conditions must have the same number of reference images; got {reference_counts}." + ) + vace_context = DiffusionRequestBatch.collate_tensors(contexts, "VACE condition", None) + assert vace_context is not None + num_reference_images = reference_counts[0] else: vace_context = None @@ -641,7 +683,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, ) # Set up scheduler @@ -654,9 +696,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): # stays None and diffuse() uses the same transformer for every step. boundary_timestep = None if self.transformer_2 is not None: - boundary_ratio = ( - self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio - ) + boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio if boundary_ratio is None: boundary_ratio = 0.875 logger.warning_once("boundary_ratio is required for Wan2.2 VACE generation. using default value 0.875") @@ -700,9 +740,11 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): latents = latents / latents_std + latents_mean output = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=output) + return split_diffusion_output_by_request( + DiffusionOutput(output=output), + req, + num_outputs_per_prompt=num_outputs_per_prompt, + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights using AutoWeightsLoader for vLLM integration.""" - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + return load_wan_weights_with_optional_gate(self, weights) diff --git a/vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py b/vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py index 25444044c..4a365c2be 100644 --- a/vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py +++ b/vllm_omni/diffusion/models/wan2_2/scheduling_wan_euler.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -110,6 +110,42 @@ class WanEulerScheduler: self._step_index = None self._begin_index = None + def sigma_for_timestep( + self, + timestep: float | torch.Tensor, + *, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + """Return the flow sigma nearest to a (possibly distilled) timestep.""" + timestep_t = torch.as_tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype) + if timestep_t.numel() != 1: + raise ValueError(f"Expected a scalar timestep, got shape {tuple(timestep_t.shape)}") + index = torch.argmin(torch.abs(self.timesteps - timestep_t)) + return self.sigmas[index].to(device=device or self.timesteps.device, dtype=dtype) + + def predict_clean( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + timestep: float | torch.Tensor, + ) -> torch.Tensor: + """Convert a flow/noise prediction into the corresponding clean latent.""" + sample_fp32 = sample.to(torch.float32) + sigma = self.sigma_for_timestep(timestep, device=sample.device, dtype=torch.float32) + return sample_fp32 - sigma * model_output.to(torch.float32) + + def add_noise( + self, + clean_sample: torch.Tensor, + noise: torch.Tensor, + timestep: float | torch.Tensor, + ) -> torch.Tensor: + """Re-noise a clean latent at a distilled flow timestep.""" + sigma = self.sigma_for_timestep(timestep, device=noise.device, dtype=torch.float32) + noisy = (1.0 - sigma) * clean_sample.to(torch.float32) + sigma * noise.to(torch.float32) + return noisy.to(noise.dtype) + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: # noqa: ARG002 return sample diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index d61d1c82b..dbc555604 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math from collections.abc import Iterable @@ -419,6 +419,24 @@ class WanSelfAttention(nn.Module): prefix=prefix, ) + # FastVideo VSA checkpoints may add a learned projection that gates + # the compressed global branch. Zero initialization makes checkpoints + # without these weights sparse-only, with no user-facing mode switch. + self.to_gate_compress: ColumnParallelLinear | None = None + if self.attn.attn_backend.get_name() == "FASTVIDEO_VSA": + self.to_gate_compress = ColumnParallelLinear( + dim, + self.inner_dim, + bias=True, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_gate_compress" if prefix else "to_gate_compress", + ) + nn.init.zeros_(self.to_gate_compress.weight) + if self.to_gate_compress.bias is not None: + nn.init.zeros_(self.to_gate_compress.bias) + def forward( self, hidden_states: torch.Tensor, @@ -441,6 +459,15 @@ class WanSelfAttention(nn.Module): key = key.unflatten(2, (self.num_kv_heads, self.head_dim)) value = value.unflatten(2, (self.num_kv_heads, self.head_dim)) + if self.to_gate_compress is not None: + gate_result = self.to_gate_compress(hidden_states) + gate_compress = gate_result[0] if isinstance(gate_result, tuple) else gate_result + gate_compress = gate_compress.unflatten(2, (self.num_heads, self.head_dim)) + if attn_metadata is None: + attn_metadata = AttentionMetadata(extra={"gate_compress": gate_compress}) + else: + attn_metadata.extra["gate_compress"] = gate_compress + # Apply rotary embeddings if rotary_emb is not None: freqs_cos, freqs_sin = rotary_emb @@ -709,6 +736,8 @@ class WanTransformerBlock(nn.Module): temb: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], hidden_states_mask: torch.Tensor | None = None, + vsa_dit_seq_shape: tuple[int, int, int] | None = None, + preserve_vsa_all_blocks: bool = False, ) -> torch.Tensor: if temb.ndim == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) @@ -729,7 +758,12 @@ class WanTransformerBlock(nn.Module): # 1. Self-attention norm_hidden_states = self.norm1(hidden_states, scale_msa, shift_msa).type_as(hidden_states) - self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask) + self_attn_extra = {} + if vsa_dit_seq_shape is not None: + self_attn_extra["vsa_dit_seq_shape"] = vsa_dit_seq_shape + if preserve_vsa_all_blocks: + self_attn_extra["preserve_vsa_all_blocks"] = True + self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask, extra=self_attn_extra) attn_output = self.attn1(norm_hidden_states, rotary_emb, self_attn_metadata) hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states) @@ -910,6 +944,9 @@ class WanTransformer3DModel(nn.Module): pos_embed_seq_len=pos_embed_seq_len, ) + # DMD/FastVideo checkpoints retain VSA semantics when top-k selects every block. + self.preserve_vsa_all_blocks = False + # 3. Transformer blocks — partitioned across PP stages via vLLM's `make_layers`. # It computes the [start_layer, end_layer) slice for this rank and fills the remaining slots # with PPMissingLayer so that weight names stay globally consistent. @@ -1044,8 +1081,19 @@ class WanTransformer3DModel(nn.Module): ) # Transformer blocks + # Preserve the post-patch (T, H, W) grid so VSA can partition + # the flattened DiT sequence into spatiotemporal blocks. + vsa_dit_seq_shape = (post_patch_num_frames, post_patch_height, post_patch_width) for block in self.blocks[self.start_layer : self.end_layer]: - hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask) + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + hidden_states_mask, + vsa_dit_seq_shape, + self.preserve_vsa_all_blocks, + ) if not is_pipeline_last_stage(): # Non-last PP stage: hand the token sequence to the caller via IntermediateTensors. diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index a13c04c99..7d04580e4 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """VACE variant of WanTransformer3DModel for conditional video generation.""" @@ -61,6 +61,7 @@ class VaceWanTransformerBlock(WanTransformerBlock): temb: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], hidden_states_mask: torch.Tensor | None = None, + vsa_dit_seq_shape: tuple[int, int, int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if self.proj_in is not None: control_hidden_states = self.proj_in(control_hidden_states) @@ -72,6 +73,7 @@ class VaceWanTransformerBlock(WanTransformerBlock): temb, rotary_emb, hidden_states_mask, + vsa_dit_seq_shape, ) conditioning_states = self.proj_out(control_hidden_states) @@ -250,6 +252,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel): ) # VACE: embed context and run conditioning blocks + vsa_dit_seq_shape = (post_patch_num_frames, post_patch_height, post_patch_width) vace_hints = None if vace_context is not None and self.vace_blocks is not None: full_seq_len = hidden_states.shape[1] * sp_size @@ -263,6 +266,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel): timestep_proj, rotary_emb, hidden_states_mask, + vsa_dit_seq_shape, ) vace_hints.append(conditioning_states) @@ -278,6 +282,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel): timestep_proj, rotary_emb, hidden_states_mask, + vsa_dit_seq_shape, ) if vace_hints is not None and self.vace_layers_mapping is not None and i in self.vace_layers_mapping: vace_idx = self.vace_layers_mapping[i] diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 9ff464494..8145ac624 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import torch from vllm.logger import init_logger from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan from vllm_omni.platforms import current_omni_platform -from .base import OffloadBackend, OffloadConfig, OffloadStrategy +from .base import OffloadBackend, OffloadConfig, OffloadStrategy, SupportsModelCpuOffload from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names from .distributed_layerwise_backend import ( DistributedLayerwiseOffloadBackend, @@ -16,13 +17,15 @@ from .distributed_layerwise_backend import ( remove_distributed_block_hook, ) from .layerwise_backend import LayerWiseOffloadBackend -from .module_residency import PinnedModuleStager -from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading +from .module_residency import BoundedAllocatorCache, PinnedModuleStager +from .offload_plan import OffloadPlan, get_offload_plan from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, remove_sequential_offload, + sequential_offload_component, ) +from .startup import OffloadStartupState, take_offload_startup_state from .tensor_utils import ( dtype_size, is_dtensor, @@ -38,18 +41,23 @@ __all__ = [ "OffloadConfig", "OffloadPlan", "OffloadStrategy", + "SupportsModelCpuOffload", "LayerWiseOffloadBackend", "DistributedLayerwiseOffloadBackend", "DistributedLayerwiseOffloadHook", "ModelLevelOffloadBackend", + "BoundedAllocatorCache", "PinnedModuleStager", "apply_sequential_offload", "remove_sequential_offload", + "sequential_offload_component", "apply_distributed_block_hook", "remove_distributed_block_hook", "get_offload_backend", + "enable_offload_backend", + "OffloadStartupState", + "take_offload_startup_state", "get_offload_plan", - "supports_mmap_loading", "get_blocks_attr_names", "get_blocks_from_dit", "set_blocks_attr_names", @@ -64,12 +72,15 @@ __all__ = [ def get_offload_backend( od_config: OmniDiffusionConfig, device: torch.device | None = None, + host_weight_plan: HostWeightPlan | None = None, ) -> OffloadBackend | None: """Create appropriate offload backend based on configuration. Args: od_config: OmniDiffusionConfig with offload settings device: Target device (auto-detected if None) + host_weight_plan: Exact loader-produced backing plan, if ordinary + weight materialization was skipped. Returns: OffloadBackend instance or None if offloading disabled @@ -82,12 +93,22 @@ def get_offload_backend( # Extract and validate configuration config = OffloadConfig.from_od_config(od_config) + if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE: + raise RuntimeError( + "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected" + ) + # Return None if no offloading requested if config.strategy == OffloadStrategy.NONE: return None # Validate platform (CUDA required for now) if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1: + if host_weight_plan is not None: + raise RuntimeError( + "The loader skipped ordinary weight materialization for DLO, " + "but this platform cannot create the required offload backend" + ) logger.warning( "Current device: %s does not support CPU offloading. Skipping offloading.", current_omni_platform.get_device_name(), @@ -99,6 +120,11 @@ def get_offload_backend( try: device = current_omni_platform.get_torch_device() except (NotImplementedError, AttributeError) as exc: + if host_weight_plan is not None: + raise RuntimeError( + "The loader skipped ordinary weight materialization for DLO, " + "but the target device could not be resolved" + ) from exc logger.error("Failed to detect device: %s. Skipping offloading.", exc) return None @@ -108,7 +134,63 @@ def get_offload_backend( elif config.strategy == OffloadStrategy.LAYER_WISE: return LayerWiseOffloadBackend(config, device) elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE: - return DistributedLayerwiseOffloadBackend(config, device) + return DistributedLayerwiseOffloadBackend( + config, + device, + host_weight_plan=host_weight_plan, + ) else: logger.error("Unknown offload strategy: %s", config.strategy) return None + + +def enable_offload_backend( + od_config: OmniDiffusionConfig, + pipeline: torch.nn.Module, + device: torch.device | None = None, +) -> tuple[torch.nn.Module, OffloadBackend | None]: + """Create and enable the loader-selected backend transactionally. + + The model runner only calls this generic offloader boundary. Loader-owned + host plans and optional fresh-model recovery callbacks stay inside the + startup state consumed here. + """ + + def enable_once(model: torch.nn.Module, state: OffloadStartupState | None) -> OffloadBackend | None: + backend: OffloadBackend | None = None + try: + backend = get_offload_backend( + od_config, + device=device, + host_weight_plan=state.host_weight_plan if state is not None else None, + ) + if backend is not None: + logger.info("Enabling offloader backend: %s", backend.__class__.__name__) + backend.enable(model) + elif state is not None: + state.close_loader_ownership() + return backend + except Exception: + if backend is not None: + try: + backend.disable() + except Exception: + logger.exception("Failed to clean up the offload backend after startup failure") + if state is not None: + state.close_loader_ownership() + raise + + startup_state = take_offload_startup_state(pipeline) + try: + return pipeline, enable_once(pipeline, startup_state) + except Exception: + if startup_state is None or not startup_state.allow_fresh_retry: + raise + assert startup_state.fresh_model_loader is not None + logger.warning( + "Loader-backed offload startup failed; retrying with a fresh canonical model", + exc_info=True, + ) + del pipeline + pipeline = startup_state.fresh_model_loader() + return pipeline, enable_once(pipeline, take_offload_startup_state(pipeline)) diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 2db80923d..c4ad93482 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -1,19 +1,40 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum +from typing import Protocol, runtime_checkable import torch from torch import nn from vllm.logger import init_logger -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig, validate_dlo_host_registration_options logger = init_logger(__name__) +@runtime_checkable +class SupportsModelCpuOffload(Protocol): + """Pipeline-owned lifecycle for model-level CPU offload. + + Pipelines with non-forward component entry points (for example VAE + ``decode_latent`` methods) need to activate those stages explicitly, so + generic forward-hook discovery cannot manage their full lifecycle. + """ + + def enable_omni_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool, + use_hsdp: bool, + ) -> None: ... + + def disable_omni_model_cpu_offload(self) -> None: ... + + class OffloadStrategy(Enum): NONE = "none" MODEL_LEVEL = "model_level" # Sequential offloading between DiT and encoders @@ -27,11 +48,13 @@ class OffloadConfig: pin_cpu_memory: bool = True use_hsdp: bool = False dp_size: int = 1 # derived from parallel_config, not user-configurable - # True: add DP sharding + AllGather. False: stream the standard loader's - # rank-local tensors (including TP-local shards) with H2D only. + # True: add DP sharding + AllGather. False: stream complete rank-local + # blocks from the loader-selected host backing with H2D only. dlo_use_allgather: bool = True dlo_resident_layers: int = 0 # leading DiT layers kept on device - model_path: str | None = None # checkpoint path for mmap weight loading + # Optional per-worker ceiling for registering an HWR mmap. Zero means no + # additional ceiling; pin_cpu_memory controls whether registration is tried. + dlo_host_registration_limit_gib: float = 0.0 @classmethod def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": @@ -58,13 +81,12 @@ class OffloadConfig: parallel_config = getattr(od_config, "parallel_config", None) use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False - # Derive dp_size from parallel_config — not user-configurable. # The offload adapts to whatever DP/SP is already configured. dp_size = 1 if parallel_config is not None: dp_size = getattr(parallel_config, "data_parallel_size", 1) - # HSDP's fully_shard_degree also contributes to effective DP + # HSDP shard and replica sizes determine the effective group size. hsdp_shard_size = getattr(parallel_config, "hsdp_shard_size", -1) if use_hsdp else -1 hsdp_replicate_size = getattr(parallel_config, "hsdp_replicate_size", 1) if use_hsdp else 1 if use_hsdp and hsdp_shard_size > 0: @@ -103,6 +125,12 @@ class OffloadConfig: # requirements (concurrent requests, dummy run skip). dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0)) + dlo_host_registration_limit_gib = validate_dlo_host_registration_options( + limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0), + enable_dlo=enable_distributed_layerwise_offload, + use_allgather=dlo_use_allgather, + hwr_mode=getattr(od_config, "host_weight_runtime_mode", "disabled"), + ) if dlo_resident_layers < 0: raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}") if dlo_resident_layers and dlo_use_allgather: @@ -116,7 +144,8 @@ class OffloadConfig: dp_size = 1 logger.info( "Distributed layerwise offload: dlo_use_allgather=False, " - "streaming standard-loader rank-local weights (no DP shard or AllGather)" + "streaming complete rank-local blocks (no DLO shard or AllGather); " + "the backend will select mmap or standard-loader host storage" ) # HSDP already shards parameters into DTensors. Running distributed @@ -137,7 +166,7 @@ class OffloadConfig: dp_size=dp_size, dlo_use_allgather=dlo_use_allgather, dlo_resident_layers=dlo_resident_layers, - model_path=getattr(od_config, "model", None), + dlo_host_registration_limit_gib=dlo_host_registration_limit_gib, ) diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py new file mode 100644 index 000000000..418b5d9fb --- /dev/null +++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA implementation of read-only host-mapping registration.""" + +from __future__ import annotations + +import ctypes +import mmap +from collections import defaultdict +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol + +import torch + +from vllm_omni.host_weight_runtime import MappedHostRegion + +from .host_registration import ( + HostRegistrationBudgetError, + HostRegistrationCleanupError, + HostRegistrationError, +) + +_CUDA_HOST_REGISTER_READ_ONLY = 0x08 +_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113 + + +class _CudaRuntime(Protocol): + def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ... + + def cudaHostUnregister(self, address: int) -> int: ... + + def cudaGetErrorString(self, error: int) -> str | bytes: ... + + def cudaGetLastError(self) -> int: ... + + +@dataclass(frozen=True) +class _AddressRange: + start: int + end: int + + @property + def size(self) -> int: + return self.end - self.start + + +def _coalesce_ranges( + ranges: Sequence[tuple[int, int]], + page_size: int = mmap.PAGESIZE, +) -> tuple[_AddressRange, ...]: + """Page-align and merge overlapping ranges from one backing mapping.""" + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + + aligned: list[_AddressRange] = [] + for start, size in ranges: + if start <= 0 or size < 0: + raise ValueError(f"invalid host range start={start}, size={size}") + if size == 0: + continue + aligned_start = start - start % page_size + end = start + size + aligned_end = ((end + page_size - 1) // page_size) * page_size + aligned.append(_AddressRange(aligned_start, aligned_end)) + + merged: list[_AddressRange] = [] + for region in sorted(aligned, key=lambda item: (item.start, item.end)): + if merged and region.start <= merged[-1].end: + previous = merged[-1] + merged[-1] = _AddressRange(previous.start, max(previous.end, region.end)) + else: + merged.append(region) + return tuple(merged) + + +def _mapped_regions(regions: Sequence[MappedHostRegion]) -> tuple[_AddressRange, ...]: + """Resolve page spans without merging unrelated artifact files.""" + by_file: dict[str, list[tuple[int, int]]] = defaultdict(list) + for region in regions: + if not region.read_only: + raise HostRegistrationError( + f"CUDA direct H2D requires an immutable mapping, but {region.file_name!r} is writable" + ) + by_file[region.file_name].append((region.address, region.length)) + return tuple(span for ranges in by_file.values() for span in _coalesce_ranges(ranges)) + + +def _error_message(runtime: _CudaRuntime, error: int) -> str: + try: + message = runtime.cudaGetErrorString(error) + except Exception: + return str(error) + if isinstance(message, bytes): + return message.decode(errors="replace") + return str(message) + + +def _consume_last_cuda_error(runtime: _CudaRuntime, expected_error: int) -> None: + """Clear one handled CUDA Runtime error before returning to PyTorch.""" + try: + get_last_error = getattr(runtime, "cudaGetLastError", None) + if get_last_error is None: + get_last_error = ctypes.CDLL(None).cudaGetLastError + get_last_error.argtypes = [] + get_last_error.restype = ctypes.c_int + pending_error = int(get_last_error()) + except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as exc: + raise HostRegistrationError("cannot clear CUDA's pending error after host-registration failure") from exc + + if pending_error not in (0, expected_error): + raise HostRegistrationError( + f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}" + ) + + +def _handled_error_message(runtime: _CudaRuntime, error: int) -> str: + error_code = int(error) + message = _error_message(runtime, error_code) + try: + _consume_last_cuda_error(runtime, error_code) + except HostRegistrationError as exc: + raise HostRegistrationError(f"{message}; {exc}") from exc + return message + + +def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool: + """Query support required by immutable artifact mappings.""" + try: + get_attribute = getattr(runtime, "cudaDeviceGetAttribute", None) + if get_attribute is None: + get_attribute = ctypes.CDLL(None).cudaDeviceGetAttribute + get_attribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int] + get_attribute.restype = ctypes.c_int + except (AttributeError, OSError) as exc: + raise HostRegistrationError("cannot query CUDA read-only host-registration support") from exc + + supported = ctypes.c_int() + try: + error = get_attribute( + ctypes.byref(supported), + _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, + torch.accelerator.current_device_index(), + ) + except Exception as exc: + raise HostRegistrationError(f"cannot query CUDA read-only host-registration support: {exc}") from exc + if int(error) != 0: + raise HostRegistrationError( + f"cudaDeviceGetAttribute(read-only host registration) failed: {_handled_error_message(runtime, error)}" + ) + return bool(supported.value) + + +class CudaHostRegistration: + """Own CUDA registrations for immutable file-backed CPU regions.""" + + def __init__(self, runtime: _CudaRuntime, regions: tuple[_AddressRange, ...]) -> None: + self._runtime = runtime + self._regions = regions + + @classmethod + def create( + cls, + regions: Sequence[MappedHostRegion], + *, + max_bytes: int | None, + ) -> CudaHostRegistration: + mapped = _mapped_regions(regions) + total_bytes = sum(region.size for region in mapped) + if max_bytes is not None and total_bytes > max_bytes: + raise HostRegistrationBudgetError( + f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget" + ) + if not mapped: + raise HostRegistrationError("no non-empty host ranges were available for registration") + if not torch.cuda.is_available(): + raise HostRegistrationError("CUDA is not available") + + try: + runtime = torch.cuda.cudart() + except Exception as exc: + raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc + if not _supports_read_only_host_registration(runtime): + raise HostRegistrationError( + "CUDA device does not support read-only host registration required by immutable HWR mappings" + ) + + registered: list[_AddressRange] = [] + try: + for region in mapped: + error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY) + if int(error) != 0: + raise HostRegistrationError( + "cudaHostRegister(read-only) failed for " + f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}" + ) + registered.append(region) + except Exception as exc: + rollback_errors: list[str] = [] + rollback_failed: list[_AddressRange] = [] + for region in reversed(registered): + try: + error = runtime.cudaHostUnregister(region.start) + if int(error) != 0: + rollback_errors.append( + f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}" + ) + rollback_failed.append(region) + except Exception as rollback_exc: + rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}") + rollback_failed.append(region) + if rollback_errors: + active_registration = cls(runtime, tuple(reversed(rollback_failed))) + raise HostRegistrationCleanupError( + f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}", + active_registration=active_registration, + ) from exc + if isinstance(exc, HostRegistrationError): + raise + raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc + + return cls(runtime, mapped) + + @property + def total_bytes(self) -> int: + return sum(region.size for region in self._regions) + + @property + def region_count(self) -> int: + return len(self._regions) + + def close(self) -> tuple[str, ...]: + """Unregister every range and retain failures for a later retry.""" + errors: list[str] = [] + failed: list[_AddressRange] = [] + for region in reversed(self._regions): + try: + error = self._runtime.cudaHostUnregister(region.start) + if int(error) != 0: + errors.append( + f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}" + ) + failed.append(region) + except Exception as exc: + errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}") + failed.append(region) + self._regions = tuple(reversed(failed)) + return tuple(errors) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 7da7c8d37..7063c2759 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1,22 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Distributed Layerwise Offload backend with H2D + AllGather overlap. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Distributed Layerwise Offload backend with double-buffered H2D. This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that: -* Shards model weights across DP ranks and stores only the local shard - (1/DP_size of full model) in host pinned memory per rank. +* Optionally shards model weights across DP ranks and reconstructs each block + with AllGather, or streams a complete rank-local block without a collective. +* Can retain compatible checkpoint tensors as node-shared mmap sources instead + of creating a persistent private host copy in every rank. * Uses a fixed double-buffer scheme that keeps only two layers' worth of weights on each device at any time. -* Asynchronously pipelines both H2D transfers and AllGather communications - on dedicated streams, fully overlapping them with computation. +* Pipelines H2D transfers and, when enabled, AllGather communications on + dedicated streams, overlapping them with computation. * Is hardware-agnostic, supporting both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms via vLLM-Omni's platform abstraction layer. """ from __future__ import annotations -import os +import threading +import time from itertools import chain from typing import Any @@ -26,12 +29,25 @@ from torch import nn from vllm.logger import init_logger from vllm_omni.diffusion.hooks import HookRegistry, ModelHook +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, +) +from vllm_omni.host_weight_runtime import HostWeightLease from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig from .block_discovery import get_blocks_from_dit +from .host_registration import ( + HostRegistration, + HostRegistrationCleanupError, + HostRegistrationError, + register_host_mappings, +) from .module_collector import ModuleDiscovery -from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading +from .offload_plan import ( + OffloadPlan, + get_offload_plan, +) from .tensor_utils import ( dtype_size as _dtype_size, ) @@ -43,6 +59,38 @@ from .tensor_utils import ( logger = init_logger(__name__) +# A backend normally owns both objects below. This process-lifetime safety +# owner prevents HostWeightLease.__del__ from unmapping storage when cleanup +# failure unwinds startup and the backend itself becomes unreachable. A clean +# retry removes the pair before closing the lease. +_ACTIVE_HWR_REGISTRATIONS: list[tuple[HostRegistration, HostWeightLease]] = [] +_ACTIVE_HWR_REGISTRATIONS_LOCK = threading.Lock() + + +def _retain_active_hwr_registration( + registration: HostRegistration, + lease: HostWeightLease, +) -> None: + with _ACTIVE_HWR_REGISTRATIONS_LOCK: + if not any( + candidate is registration and candidate_lease is lease + for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS + ): + _ACTIVE_HWR_REGISTRATIONS.append((registration, lease)) + + +def _forget_active_hwr_registration( + registration: HostRegistration, + lease: HostWeightLease, +) -> None: + with _ACTIVE_HWR_REGISTRATIONS_LOCK: + _ACTIVE_HWR_REGISTRATIONS[:] = [ + (candidate, candidate_lease) + for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS + if candidate is not registration or candidate_lease is not lease + ] + + # Threshold (in MB) for deciding whether a non-block DiT submodule should # use layerwise offload (streaming hooks) or be moved to GPU as a resident # module. Submodules larger than this are offloaded to save HBM; smaller @@ -74,6 +122,8 @@ class DistributedLayerwiseOffloadHook(ModelHook): comm_stream: Any | None = None, pin_memory: bool = True, shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + rank_local_mmap: bool = False, + tensor_transforms: dict[int, Any] | None = None, ): assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" @@ -83,6 +133,9 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.dp_size = dp_size self.rank = rank self.pin_memory = pin_memory + self.rank_local_mmap = rank_local_mmap + self.registered_mmap = False + self.tensor_transforms = tensor_transforms or {} self.copy_stream = copy_stream or current_omni_platform.Stream() self.comm_stream = comm_stream or current_omni_platform.Stream() @@ -98,8 +151,18 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Sharded host weights for the next block, keyed by dtype self.cpu_shards: dict[torch.dtype, torch.Tensor] = {} + # File-backed source tensors for rank-local mmap. Unlike cpu_shards, + # these remain immutable views of the checkpoint and are never pinned + # or flattened into a model-sized private allocation. + self.cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {} self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + # Rank-local mmap uses two host staging slots shared by every hook in + # this worker. They are assigned by the backend after all block sizes + # are known, mirroring the shared device-buffer allocation. + self.cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + self.cpu_staging_events: list[Any | None] = [None, None] + # Current slot index (0 or 1). Updated dynamically by the previous # hook's prefetch_layer call via _prefetched_slot. This ensures # correct slot tracking for ALL block counts (including odd N). @@ -154,14 +217,22 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.next_block_parameters = dict(self.next_block.named_parameters()) self.next_block_buffers = dict(self.next_block.named_buffers()) - # Shard next block's weights and store local shard in pinned CPU memory - self.cpu_shards, self.metadata = self._shard_and_pin( - self.next_block_parameters, - self.next_block_buffers, - self.dp_size, - self.rank, - self.pin_memory, - ) + if self.rank_local_mmap: + self.cpu_sources, self.metadata = self._collect_mmap_sources( + self.next_block_parameters, + self.next_block_buffers, + self.tensor_transforms, + ) + else: + # Shard next block's weights and store local shard in pinned CPU memory. + self.cpu_shards, self.metadata = self._shard_and_pin( + self.next_block_parameters, + self.next_block_buffers, + self.dp_size, + self.rank, + self.pin_memory, + self.tensor_transforms, + ) # Allocate device buffers only if not using shared buffers from backend if self._owns_buffers: @@ -178,18 +249,94 @@ class DistributedLayerwiseOffloadHook(ModelHook): if m["name"] in self.next_block_parameters else self.next_block_buffers[m["name"]] ) - repoint.append((target, dtype, m["offset"], m["numel"], m["shape"])) + repoint.append( + ( + target, + dtype, + m["offset"], + m["numel"], + m["shape"], + m["stride"], + ) + ) self._cached_repoint.append(repoint) # Pre-compute AG output sizes (avoid sum() per layer). self._ag_output_sizes: dict[torch.dtype, int] = {} for dtype, metas in self.metadata.items(): total_numel = sum(m["numel"] for m in metas) - shard_numel = self.cpu_shards[dtype].numel() - self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel + if self.rank_local_mmap: + self._ag_output_sizes[dtype] = total_numel + else: + shard_numel = self.cpu_shards[dtype].numel() + self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel return module + @staticmethod + def _collect_mmap_sources( + params: dict[str, nn.Parameter], + bufs: dict[str, torch.Tensor], + tensor_transforms: dict[int, Any] | None = None, + ) -> tuple[dict[torch.dtype, list[dict[str, Any]]], dict[torch.dtype, list[dict[str, Any]]]]: + """Retain file-backed tensors and replace module storage with placeholders. + + The returned sources preserve safetensors mmap storage. Runtime-layout + adapters are applied only while packing a bounded staging slot, so no + full-model anonymous CPU copy is retained by a worker. + """ + cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {} + metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + offsets: dict[torch.dtype, int] = {} + + for name, target in chain(params.items(), bufs.items()): + source = target.to_local() if hasattr(target, "to_local") else target + if source.device.type != "cpu": + raise ValueError( + f"Rank-local mmap storage requires CPU checkpoint views, but {name!r} is on {source.device}." + ) + + dtype = source.dtype + offset = offsets.get(dtype, 0) + transform = (tensor_transforms or {}).get(id(target)) + runtime_source = transform(source) if callable(transform) else source + if runtime_source.dtype != dtype or runtime_source.shape != source.shape: + raise ValueError( + "mmap weight transform changed tensor metadata for " + f"{name!r}: expected dtype={dtype}, shape={tuple(source.shape)}, " + f"got dtype={runtime_source.dtype}, shape={tuple(runtime_source.shape)}" + ) + stride = runtime_source.stride() + storage_numel = ( + 0 + if runtime_source.numel() == 0 + else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(runtime_source.shape, stride)) + ) + + cpu_sources.setdefault(dtype, []).append( + { + "name": name, + "tensor": source.detach(), + "transform": transform, + } + ) + metadata.setdefault(dtype, []).append( + { + "name": name, + "offset": offset, + "numel": storage_numel, + "shape": runtime_source.shape, + "stride": stride, + } + ) + offsets[dtype] = offset + storage_numel + + # The detached source above keeps the mmap storage alive while the + # module parameter/buffer is rebound to the rotating device slot. + set_tensor_storage(target, make_offload_placeholder(target)) + + return cpu_sources, metadata + @staticmethod def _shard_and_pin( params: dict[str, nn.Parameter], @@ -197,6 +344,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): dp_size: int, rank: int, pin_memory: bool, + tensor_transforms: dict[int, Any] | None = None, ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]: """Flatten params+buffers by dtype, split into DP shards, store local shard. @@ -219,8 +367,8 @@ class DistributedLayerwiseOffloadHook(ModelHook): weights_with_local = [] for name, t in name2weights.items(): local_t = t.to_local() if hasattr(t, "to_local") else t - mmap_transform = getattr(t, "mmap_weight_transform", None) - if callable(mmap_transform) and getattr(t, "mmap_weight_transform_pending", False): + mmap_transform = (tensor_transforms or {}).get(id(t)) + if callable(mmap_transform): # Some checkpoints use a layout that is converted by the # regular weight loader (for example MiniMax-H3 grouped # QKV). Apply that conversion one block at a time while @@ -228,9 +376,15 @@ class DistributedLayerwiseOffloadHook(ModelHook): # parameter as an mmap view avoids a private full-model # copy in every worker. local_t = mmap_transform(local_t) - weights_with_local.append((name, t, local_t)) + stride = local_t.stride() + storage_numel = ( + 0 + if local_t.numel() == 0 + else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride)) + ) + weights_with_local.append((name, t, local_t, storage_numel, stride)) - total_numel = sum(local.numel() for _, _, local in weights_with_local) + total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local) # Equal-sized shards (ceil division) for all_gather_into_tensor shard_size = (total_numel + dp_size - 1) // dp_size # ceil @@ -242,8 +396,13 @@ class DistributedLayerwiseOffloadHook(ModelHook): shard = torch.zeros(shard_size, dtype=dtype, device="cpu") current_offset = 0 - for name, original_tensor, local_tensor in weights_with_local: - numel = local_tensor.numel() + for ( + name, + original_tensor, + local_tensor, + storage_numel, + stride, + ) in weights_with_local: if dtype not in dtype_metadata: dtype_metadata[dtype] = [] # Offsets remain relative to the FULL flattened buffer @@ -252,27 +411,48 @@ class DistributedLayerwiseOffloadHook(ModelHook): { "name": name, "offset": current_offset, - "numel": numel, + "numel": storage_numel, "shape": local_tensor.shape, + "stride": stride, } ) # Copy ONLY the portion within [shard_start, shard_end) overlap_start = max(current_offset, shard_start) - overlap_end = min(current_offset + numel, shard_end) + overlap_end = min(current_offset + storage_numel, shard_end) if overlap_start < overlap_end: + if local_tensor.is_contiguous(): + flat_storage = local_tensor.flatten() + else: + # Online FP8 stores Cutlass weights as transposed views + # (e.g. stride=(1, K)). Flattening such a tensor in + # logical order and later rebuilding it with .view() + # changes its layout and makes scaled_mm reject it. + # Pack the physical storage order and preserve the + # original stride for zero-copy reconstruction. + flat_storage = torch.zeros( + storage_numel, + dtype=dtype, + device=local_tensor.device, + ) + physical_view = torch.as_strided( + flat_storage, + size=local_tensor.shape, + stride=stride, + ) + physical_view.copy_(local_tensor) src_start = overlap_start - current_offset src_end = overlap_end - current_offset dst_start = overlap_start - shard_start dst_end = overlap_end - shard_start - shard[dst_start:dst_end].copy_(local_tensor.flatten()[src_start:src_end]) + shard[dst_start:dst_end].copy_(flat_storage[src_start:src_end]) # Replace original tensor with placeholder (frees CPU storage) set_tensor_storage( original_tensor, make_offload_placeholder(original_tensor), ) - current_offset += numel + current_offset += storage_numel if pin_memory: shard = shard.pin_memory() @@ -306,6 +486,96 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Prefetch: H2D + AllGather (overlapped on dedicated streams) # # ------------------------------------------------------------------ # + @staticmethod + def _resolve_mmap_source( + source_info: dict[str, Any], + meta: dict[str, Any], + dtype: torch.dtype, + ) -> torch.Tensor: + source = source_info["tensor"] + transform = source_info["transform"] + if callable(transform): + source = transform(source) + if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]: + raise ValueError( + "mmap weight transform changed tensor layout for " + f"{source_info['name']!r}: expected dtype={dtype}, " + f"shape={tuple(meta['shape'])}, stride={meta['stride']}; " + f"got dtype={source.dtype}, shape={tuple(source.shape)}, " + f"stride={source.stride()}" + ) + return source + + @staticmethod + def _pack_mmap_sources( + cpu_sources: dict[torch.dtype, list[dict[str, Any]]], + metadata: dict[torch.dtype, list[dict[str, Any]]], + slot_buffers: dict[torch.dtype, torch.Tensor], + ) -> dict[torch.dtype, torch.Tensor]: + staged: dict[torch.dtype, torch.Tensor] = {} + for dtype, metas in metadata.items(): + total_numel = sum(meta["numel"] for meta in metas) + destination = slot_buffers[dtype][:total_numel] + sources = cpu_sources[dtype] + for source_info, meta in zip(sources, metas, strict=True): + source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype) + start = meta["offset"] + physical_storage = destination[start : start + meta["numel"]] + if source.is_contiguous(): + physical_storage.copy_(source.flatten()) + else: + torch.as_strided( + physical_storage, + size=source.shape, + stride=source.stride(), + ).copy_(source) + staged[dtype] = destination + return staged + + def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]: + """Pack this block's mmap views into one bounded host staging slot.""" + previous_copy = self.cpu_staging_events[slot] + if previous_copy is not None: + synchronize = getattr(previous_copy, "synchronize", None) + if callable(synchronize): + synchronize() + else: + # Platform events normally expose synchronize(). Retain a + # correctness fallback for test and non-CUDA platform shims. + current_omni_platform.synchronize() + self.cpu_staging_events[slot] = None + + slot_buffers = self.cpu_staging_buffers[slot] + if slot_buffers is None: + raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated") + return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers) + + @staticmethod + def _copy_mmap_sources_to_device( + cpu_sources: dict[torch.dtype, list[dict[str, Any]]], + metadata: dict[torch.dtype, list[dict[str, Any]]], + gpu_buffers: dict[torch.dtype, torch.Tensor], + *, + non_blocking: bool, + ) -> None: + """Copy registered source views directly into flattened device buffers.""" + for dtype, metas in metadata.items(): + destination = gpu_buffers[dtype] + sources = cpu_sources[dtype] + for source_info, meta in zip(sources, metas, strict=True): + source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype) + start = meta["offset"] + physical_storage = destination[start : start + meta["numel"]] + async_copy = non_blocking and source.is_pinned() + if source.is_contiguous(): + physical_storage.copy_(source.flatten(), non_blocking=async_copy) + else: + torch.as_strided( + physical_storage, + size=source.shape, + stride=source.stride(), + ).copy_(source, non_blocking=async_copy) + @torch.compiler.disable def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None: """Prepare next block's weights into the shared device buffer for *slot*. @@ -321,11 +591,27 @@ class DistributedLayerwiseOffloadHook(ModelHook): assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated" if self.dp_size <= 1 or self.dp_group is None: - with current_omni_platform.stream(self.copy_stream): - for dtype, cpu_shard in self.cpu_shards.items(): - gw = gpu_weights[dtype] - gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=non_blocking) - evt.record(self.copy_stream) + if self.rank_local_mmap and self.registered_mmap: + with current_omni_platform.stream(self.copy_stream): + self._copy_mmap_sources_to_device( + self.cpu_sources, + self.metadata, + gpu_weights, + non_blocking=non_blocking, + ) + evt.record(self.copy_stream) + else: + cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards + with current_omni_platform.stream(self.copy_stream): + for dtype, cpu_shard in cpu_weights.items(): + gw = gpu_weights[dtype] + async_copy = non_blocking and cpu_shard.is_pinned() + gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy) + evt.record(self.copy_stream) + if self.rank_local_mmap: + # The CPU slot may be overwritten only after this H2D copy has + # finished. The shared event protects reuse by another hook. + self.cpu_staging_events[slot] = evt else: gpu_shards: dict[torch.dtype, torch.Tensor] = {} shard_bufs = self.gpu_shard_buffers[slot] @@ -365,10 +651,14 @@ class DistributedLayerwiseOffloadHook(ModelHook): self._shared_slot_group[slot] = self._group_id # Re-point using cached metadata (avoids per-layer dict lookups). - for target, dtype, offset, numel, shape in self._cached_repoint[slot]: + for target, dtype, offset, numel, shape, stride in self._cached_repoint[slot]: set_tensor_storage( target, - gpu_weights[dtype][offset : offset + numel].view(shape), + torch.as_strided( + gpu_weights[dtype][offset : offset + numel], + size=shape, + stride=stride, + ), ) def get_weights(self, slot: int) -> dict[torch.dtype, torch.Tensor] | None: @@ -468,6 +758,8 @@ def apply_distributed_block_hook( comm_stream: Any | None = None, pin_memory: bool = True, shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + rank_local_mmap: bool = False, + tensor_transforms: dict[int, Any] | None = None, ) -> DistributedLayerwiseOffloadHook: """Register a DistributedLayerwiseOffloadHook on *module*.""" registry = HookRegistry.get_or_create(module) @@ -481,6 +773,8 @@ def apply_distributed_block_hook( comm_stream=comm_stream, pin_memory=pin_memory, shared_buffers=shared_buffers, + rank_local_mmap=rank_local_mmap, + tensor_transforms=tensor_transforms, ) registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook) return hook @@ -495,15 +789,16 @@ def remove_distributed_block_hook(module: nn.Module) -> None: class PinnedResidentLayerGroup: - """Keep selected layers in pinned host memory between requests. + """Keep selected layers available for stage-scoped device residency. TODO(offload): Extract this alongside PinnedModuleStager after the distributed shard-and-pin operation becomes a shared storage primitive. It currently remains here because it depends on DLO's local-shard layout. Unlike ``module.to(device)``/``module.to("cpu")``, this group retains a - pinned CPU master copy and never copies generated device weights back to - host. Entering the denoise stage performs one asynchronous H2D pass; + pinned CPU master copy (or mmap source plus bounded staging) and never + copies generated device weights back to host. Entering the denoise stage + performs one asynchronous H2D pass; leaving it only restores zero-sized placeholders and releases the device buffers. This lets the following VAE stage reuse the same HBM. @@ -518,10 +813,16 @@ class PinnedResidentLayerGroup: device: torch.device, copy_stream: Any, pin_memory: bool, + rank_local_mmap: bool = False, + defer_staging: bool = False, + tensor_transforms: dict[int, Any] | None = None, ) -> None: self.device = device self.copy_stream = copy_stream self.loaded = False + self.rank_local_mmap = rank_local_mmap + self.registered_mmap = False + self.pin_memory = pin_memory self._states: list[dict[str, Any]] = [] self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] @@ -529,21 +830,49 @@ class PinnedResidentLayerGroup: params = dict(block.named_parameters()) bufs = dict(block.named_buffers()) targets: dict[str, torch.Tensor] = {**params, **bufs} - cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin( - params, - bufs, - dp_size=1, - rank=0, - pin_memory=pin_memory, - ) + if rank_local_mmap: + cpu_sources, metadata = DistributedLayerwiseOffloadHook._collect_mmap_sources( + params, + bufs, + tensor_transforms, + ) + cpu_shards = {} + else: + cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin( + params, + bufs, + dp_size=1, + rank=0, + pin_memory=pin_memory, + tensor_transforms=tensor_transforms, + ) + cpu_sources = {} self._states.append( { "targets": targets, "cpu_shards": cpu_shards, + "cpu_sources": cpu_sources, "metadata": metadata, } ) + self._cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor]] = [] + self._cpu_staging_events: list[Any | None] = [None, None] + if rank_local_mmap and not defer_staging: + max_sizes: dict[torch.dtype, int] = {} + for state in self._states: + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + for _ in range(2): + buffers = {} + for dtype, total in max_sizes.items(): + buffer = torch.empty(total, dtype=dtype, device="cpu") + if pin_memory: + buffer = buffer.pin_memory() + buffers[dtype] = buffer + self._cpu_staging_buffers.append(buffers) + def load(self) -> None: if self.loaded: return @@ -554,20 +883,49 @@ class PinnedResidentLayerGroup: gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] for state in self._states: block_buffers: dict[torch.dtype, torch.Tensor] = {} - for dtype, cpu_shard in state["cpu_shards"].items(): - block_buffers[dtype] = torch.empty( - cpu_shard.shape, - dtype=dtype, - device=self.device, - ) + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + block_buffers[dtype] = torch.empty(total, dtype=dtype, device=self.device) gpu_buffers.append(block_buffers) self.copy_stream.wait_stream(current_omni_platform.current_stream()) ready = current_omni_platform.Event() with current_omni_platform.stream(self.copy_stream): - for state, block_buffers in zip(self._states, gpu_buffers): - for dtype, cpu_shard in state["cpu_shards"].items(): - block_buffers[dtype].copy_(cpu_shard, non_blocking=True) + for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)): + if self.rank_local_mmap and self.registered_mmap: + DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device( + state["cpu_sources"], + state["metadata"], + block_buffers, + non_blocking=True, + ) + continue + if self.rank_local_mmap: + slot = index % 2 + previous_copy = self._cpu_staging_events[slot] + if previous_copy is not None: + synchronize = getattr(previous_copy, "synchronize", None) + if callable(synchronize): + synchronize() + else: + current_omni_platform.synchronize() + cpu_weights = DistributedLayerwiseOffloadHook._pack_mmap_sources( + state["cpu_sources"], + state["metadata"], + self._cpu_staging_buffers[slot], + ) + else: + cpu_weights = state["cpu_shards"] + + for dtype, cpu_weight in cpu_weights.items(): + block_buffers[dtype].copy_( + cpu_weight, + non_blocking=cpu_weight.is_pinned(), + ) + if self.rank_local_mmap: + slot_ready = current_omni_platform.Event() + slot_ready.record(self.copy_stream) + self._cpu_staging_events[slot] = slot_ready ready.record(self.copy_stream) for state, block_buffers in zip(self._states, gpu_buffers): @@ -577,7 +935,11 @@ class PinnedResidentLayerGroup: for meta in metas: set_tensor_storage( targets[meta["name"]], - gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]].view(meta["shape"]), + torch.as_strided( + gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]], + size=meta["shape"], + stride=meta["stride"], + ), ) current_omni_platform.current_stream().wait_event(ready) @@ -616,9 +978,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): overlapped with computation. """ - _MMAP_PARAM_ATTRS = ("weight_loader", "mmap_weight_transform") - - def __init__(self, config: OffloadConfig, device: torch.device): + def __init__( + self, + config: OffloadConfig, + device: torch.device, + host_weight_plan: HostWeightPlan | None = None, + ): super().__init__(config, device) self.copy_stream = current_omni_platform.Stream() @@ -630,6 +995,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] self._resident_blocks: list[nn.Module] = [] self._resident_layer_group: PinnedResidentLayerGroup | None = None + self._using_mmap = False + self._using_rank_local_mmap = False + self._using_registered_mmap = False + self.host_weight_plan = host_weight_plan + self._host_weight_lease: HostWeightLease | None = None + self._host_registration: HostRegistration | None = None + self._mmap_transforms_by_tensor_id: dict[int, Any] = {} def load_resident_layers(self) -> None: """Load the model-declared leading blocks for the denoise stage.""" @@ -641,93 +1013,160 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self._resident_layer_group is not None: self._resident_layer_group.offload() - def _remember_mmap_param_attrs(self, pipeline: nn.Module) -> None: - """Save loader metadata before ``to_empty`` replaces Parameters.""" - self._mmap_param_attrs = { - name: {attr: value for attr in self._MMAP_PARAM_ATTRS if (value := getattr(param, attr, None)) is not None} - for name, param in pipeline.named_parameters() - } - - def _attach_mmap_param_attrs( + def _rank_local_source_tensors( self, - name: str, - replacement: nn.Parameter, - source: nn.Parameter | None = None, - ) -> None: - attrs = getattr(self, "_mmap_param_attrs", {}).get(name, {}) - for attr in self._MMAP_PARAM_ATTRS: - value = attrs.get(attr) - if value is None and source is not None: - value = getattr(source, attr, None) - if value is not None: - setattr(replacement, attr, value) - if attr == "mmap_weight_transform": - replacement.mmap_weight_transform_pending = True + hooks: list[DistributedLayerwiseOffloadHook], + ) -> tuple[torch.Tensor, ...]: + """Return the exact CPU sources that direct H2D would consume.""" + tensors: list[torch.Tensor] = [] + seen: set[int] = set() - def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None: - """Load DiT weights from safetensors via mmap views (no RSS). + def collect(cpu_sources: dict[torch.dtype, list[dict[str, Any]]]) -> None: + for sources in cpu_sources.values(): + for source in sources: + tensor = source["tensor"] + if id(tensor) not in seen: + seen.add(id(tensor)) + tensors.append(tensor) + + for hook in hooks: + collect(hook.cpu_sources) + if self._resident_layer_group is not None: + for state in self._resident_layer_group._states: + collect(state["cpu_sources"]) + return tuple(tensors) + + def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bool: + """Register the complete final-layout lease under pinned-memory policy.""" + lease = self._host_weight_lease + if lease is None: + return False + if not self.config.pin_cpu_memory: + logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging") + return False + if not lease.mapped_regions or not source_tensors: + logger.warning("HWR mmap registration found no mapped sources; using bounded host staging") + return False + + limit_gib = self.config.dlo_host_registration_limit_gib + max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None + started = time.perf_counter() + try: + registration = register_host_mappings( + lease.mapped_regions, + device=self.device, + max_bytes=max_bytes, + ) + try: + unpinned = [tensor for tensor in source_tensors if tensor.numel() and not tensor.is_pinned()] + except Exception as exc: + errors = registration.close() + if errors: + self._host_registration = registration + _retain_active_hwr_registration(registration, lease) + raise HostRegistrationCleanupError( + "CUDA registration succeeded but pinned-source verification failed, " + f"and rollback failed: {errors[:3]}" + ) from exc + raise HostRegistrationError(f"cannot verify registered HWR sources: {exc}") from exc + if unpinned: + errors = registration.close() + if errors: + self._host_registration = registration + _retain_active_hwr_registration(registration, lease) + raise HostRegistrationCleanupError( + "CUDA registration succeeded but PyTorch rejected mapped sources, " + f"and rollback failed: {errors[:3]}" + ) + raise HostRegistrationError( + "CUDA registration succeeded but PyTorch did not recognize " + f"{len(unpinned)} mapped source(s) as pinned" + ) + except HostRegistrationCleanupError as exc: + # Falling back could close a lease while the platform still owns + # one of its mappings. Fail startup and retain ownership for retry. + active_registration = exc.active_registration + if active_registration is not None: + self._host_registration = active_registration + _retain_active_hwr_registration(active_registration, lease) + logger.exception("HWR mmap registration rollback failed") + raise + except HostRegistrationError as exc: + logger.warning("HWR registered direct H2D unavailable (%s); using bounded host staging", exc) + return False + + self._host_registration = registration + _retain_active_hwr_registration(registration, lease) + logger.info( + "Registered %.2f GiB of HWR mmap in %d range(s) for direct H2D in %.3f s", + registration.total_bytes / 1024**3, + registration.region_count, + time.perf_counter() - started, + ) + return True + + def _configure_hwr_transfer(self, hooks: list[DistributedLayerwiseOffloadHook]) -> None: + """Select registered direct H2D or bounded staging once per backend.""" + plan = self.host_weight_plan + if ( + (not hooks and self._resident_layer_group is None) + or not self._using_rank_local_mmap + or plan is None + or plan.backing_kind != "host_weight_runtime" + ): + return + + source_tensors = self._rank_local_source_tensors(hooks) + self._using_registered_mmap = self._try_register_hwr_mmap(source_tensors) + for hook in hooks: + hook.registered_mmap = self._using_registered_mmap + if self._resident_layer_group is not None: + self._resident_layer_group.registered_mmap = self._using_registered_mmap + if self._using_registered_mmap: + self._resident_layer_group._cpu_staging_buffers.clear() + + def _release_registered_mmap(self) -> None: + """Release every platform registration before closing the HWR lease.""" + registration = self._host_registration + if registration is None: + return + errors = registration.close() + if errors: + lease = self._host_weight_lease + if lease is not None and not lease.closed: + _retain_active_hwr_registration(registration, lease) + logger.error("HWR mmap unregistration failed; retaining lease mappings for retry: %s", errors[:3]) + raise HostRegistrationCleanupError(f"failed to unregister {len(errors)} HWR mmap range(s)") + lease = self._host_weight_lease + if lease is not None: + _forget_active_hwr_registration(registration, lease) + self._host_registration = None + logger.info("Unregistered HWR mmap ranges") + + def _load_weights_via_mmap( + self, + pipeline: nn.Module, + modules, + plan: HostWeightPlan, + ) -> None: + """Load DiT checkpoint tensors as file-backed safetensors views. When the transformer is created on meta device, this method replaces meta params with mmap views of the checkpoint files. - The views point to OS page cache (shared across ranks), so - no private copy is created. _shard_and_pin then copies only - the 1/dp_size shard from the mmap view to a private buffer. + The views point to OS page cache shared across ranks. AllGather mode + copies only the rank's 1/dp_size shard to a persistent private buffer; + rank-local mode retains the views and packs one block at a time into + bounded staging storage. Non-DiT modules (VAE, encoders) are NOT affected — they were created on CPU with real weights via from_pretrained. """ - import glob - import json - from safetensors import safe_open - model_path = self.config.model_path - if not model_path: - logger.warning("No model_path for mmap weight loading, skipping") - return - - # Resolve HF repo ID to local snapshot path. - # If model_path is a local directory, use it directly; otherwise - # download the safetensors index + weight files via HuggingFace Hub. - if not os.path.isdir(model_path): - from vllm.model_executor.model_loader.weight_utils import ( - download_weights_from_hf, - ) - - logger.info("model_path %s is not local, downloading from HF", model_path) - model_path = download_weights_from_hf( - model_name_or_path=model_path, - cache_dir=None, - allow_patterns=["*.safetensors", "*.safetensors.index.json"], - ) - - # Build {checkpoint_key: file_path} from safetensors index - weight_map: dict[str, str] = {} - for idx_file in glob.glob(os.path.join(model_path, "**", "*.safetensors.index.json"), recursive=True): - idx_dir = os.path.dirname(idx_file) - with open(idx_file) as f: - idx = json.load(f) - for key, filename in idx.get("weight_map", {}).items(): - weight_map[key] = os.path.join(idx_dir, filename) - - if not weight_map: - single_files = glob.glob(os.path.join(model_path, "**", "model.safetensors"), recursive=True) - for sf in single_files: - from safetensors import safe_open as _safe_open - - f = _safe_open(sf, framework="pt", device="cpu") - for key in f.keys(): - weight_map[key] = sf - logger.info("Found single-file safetensors: %s (%d keys)", sf, len(f.keys())) - - if not weight_map: - raise RuntimeError( - "Distributed layerwise offload could not find safetensors " - f"weights at {model_path}. Expected *.safetensors.index.json " - "or model.safetensors. Ensure the checkpoint path is correct." - ) - - logger.info("Loading DiT weights via mmap (meta → page cache, no RSS): %d keys", len(weight_map)) + logger.info( + "Loading DiT weights via mmap (meta -> shared page cache): %d tensors", + len(plan.bindings), + ) # --- Convert DiT modules to meta device --- # The transformer was created normally (with random weights and @@ -740,36 +1179,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # __init__. Save them before meta conversion and restore after # mmap loading, so we don't need model-specific buffer rebuild code. # - # Fail closed for unsupported configurations BEFORE to_empty. - # The mmap path bypasses AutoWeightsLoader, which means TP-aware - # weight_loader callbacks (fused QKV, row-parallel, etc.) are not - # invoked. When TP > 1, these callbacks shard weights across TP - # ranks — skipping them produces incorrect weights. - # Check the actual TP world size from parallel_state (initialized - # before enable()). Note: params may have custom weight_loader - # attributes even at TP=1 (e.g. QKVParallelLinear), so we cannot - # reject based on attribute presence alone. - try: - from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_world_size, - ) - - tp_world = get_tensor_model_parallel_world_size() - except Exception: - tp_world = 1 - if tp_world > 1: - raise ValueError( - "Distributed layerwise offload with mmap loading does not " - "support Tensor Parallel (TP > 1). TP-aware weight_loader " - "callbacks are bypassed by the mmap path. Use DP or SP " - "instead of TP, or disable distributed layerwise offload." - ) - - # ``Module.to_empty`` creates new Parameter objects and drops custom - # attributes installed by vLLM weight loaders. Preserve them by full - # pipeline name so mmap replacements can reproduce checkpoint layout - # transforms such as MiniMax-H3 grouped QKV. - self._remember_mmap_param_attrs(pipeline) + # The loader proved topology, coverage, source metadata, and adapter + # compatibility before it skipped ordinary weight materialization. + # This method only realizes that exact plan; it does not select or + # rediscover a checkpoint layout. # # Important: when DiT modules are nested (e.g. transformer contains @@ -812,189 +1225,63 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): len(saved_buffers.get(id(dit_module), {})), ) - # Build reverse mapping {model_param_name: (ckpt_key, file_path)} using - # the pipeline's _remap_ckpt_key (if available). This handles the - # model-specific checkpoint key remapping (e.g., Cosmos3's - # layers.0.mlp_moe_gen.gate_proj → transformer.gen_layers.0.mlp.gate_proj). - remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) - model_to_ckpt: dict[str, tuple[str, str]] = {} - if remap_fn is not None: - for ckpt_key, file_path in weight_map.items(): - model_name = remap_fn(ckpt_key) - if model_name is not None: - model_to_ckpt[model_name] = (ckpt_key, file_path) - logger.info( - "Built reverse remap: %d model params from %d checkpoint keys", len(model_to_ckpt), len(weight_map) - ) - else: - # No remap function — use checkpoint keys directly - for ckpt_key, file_path in weight_map.items(): - model_to_ckpt[ckpt_key] = (ckpt_key, file_path) - - # Collect all meta params in DiT modules - dit_ids = set() - for dit_module in modules.dits: - dit_ids.update(id(m) for m in dit_module.modules()) - # Cache open file handles file_cache: dict[str, Any] = {} - loaded = 0 - skipped = 0 loaded_names: set[str] = set() - # Discover blocks first so we can distinguish block params (deferred - # to _shard_and_pin) from non-block params (loaded here as mmap views). - block_module_ids: set[int] = set() - for dit_module in modules.dits: - blocks_attr_names, blocks = get_blocks_from_dit(dit_module) - for block in blocks: - block_module_ids.update(id(m) for m in block.modules()) + # Realize the loader's exact runtime-name bindings. In particular, + # do not reconstruct names from DLO block discovery: storage planning + # and transfer topology are separate contracts. + for runtime_name, binding in plan.bindings.items(): + parent_path, _, leaf_name = runtime_name.rpartition(".") + try: + parent = pipeline.get_submodule(parent_path) + except AttributeError as exc: + raise RuntimeError(f"Host-weight plan target module {parent_path!r} no longer exists") from exc - for name, param in pipeline.named_parameters(): - if not (hasattr(param, "is_meta") and param.is_meta): - continue + target = parent._parameters.get(leaf_name) + is_parameter = target is not None + if target is None: + target = parent._buffers.get(leaf_name) + if target is None: + raise RuntimeError(f"Host-weight plan target tensor {runtime_name!r} no longer exists") - # Check if this param is inside a block (will be handled by - # _shard_and_pin via mmap views in section 2 below). - parent = pipeline - parts = name.split(".") - is_block_param = False - for part in parts[:-1]: - try: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - except (AttributeError, IndexError, TypeError): - break - if id(parent) in block_module_ids: - is_block_param = True - break - - if is_block_param: - # Block param — will be loaded as mmap view in section 2 - continue - - # Non-block param (merger, patch_embed, norm, proj_out, etc.) - # Look up in reverse remap mapping (handles Cosmos3 key remapping). - entry = model_to_ckpt.get(name) - if entry is None: - skipped += 1 - continue - - ckpt_key, file_path = entry - if file_path not in file_cache: - file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") - f = file_cache[file_path] - tensor = f.get_tensor(ckpt_key) - - # Replace meta param with mmap view (no copy!) - parent = pipeline - for part in parts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) - self._attach_mmap_param_attrs(name, replacement, param) - parent._parameters[parts[-1]] = replacement - loaded += 1 - loaded_names.add(name) - - logger.info("Mmap weight loading: %d non-block params loaded, %d skipped", loaded, skipped) - - # Now load DiT block params from safetensors using mmap views. - # Each block param is assigned an mmap view (no RSS). - # _shard_and_pin will later copy the shard portion to a private buffer. - block_loaded = 0 - for dit_idx, dit_module in enumerate(modules.dits): - dit_name = modules.dit_names[dit_idx] - blocks_attr_names, blocks = get_blocks_from_dit(dit_module) - if not blocks: - continue - - # Determine the blocks_attr for this DiT module (e.g. "layers" or "gen_layers") - blocks_attr = None - for attr_name in blocks_attr_names: - if hasattr(dit_module, attr_name): - candidate = getattr(dit_module, attr_name) - if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: - blocks_attr = attr_name - break - - for block_idx, block in enumerate(blocks): - # Build the full model param prefix for this block - # e.g. "transformer.language_model.layers.0" or "transformer.gen_layers.0" - block_full_prefix = f"{dit_name}.{blocks_attr}.{block_idx}" - - for bname, bparam in block.named_parameters(): - if not (hasattr(bparam, "is_meta") and bparam.is_meta): - continue - - # Use reverse remap to find checkpoint key - full_param_name = f"{block_full_prefix}.{bname}" - entry = model_to_ckpt.get(full_param_name) - if entry is None: - skipped += 1 - continue - - ckpt_key, file_path = entry - if file_path not in file_cache: - file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") - f = file_cache[file_path] - tensor = f.get_tensor(ckpt_key) - - # Replace meta param with mmap view - parent = block - bparts = bname.split(".") - for part in bparts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - replacement = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) - self._attach_mmap_param_attrs(full_param_name, replacement, bparam) - parent._parameters[bparts[-1]] = replacement - block_loaded += 1 - loaded_names.add(full_param_name) - - logger.info("Mmap block loading: %d params loaded, %d skipped", block_loaded, skipped) - - # Strict validation: check that all meta params were loaded. - # This replaces the validation that AutoWeightsLoader.load_weights() - # performs in the regular load path. - remaining_meta = [ - name for name, param in pipeline.named_parameters() if hasattr(param, "is_meta") and param.is_meta - ] - if remaining_meta: - # Filter out params inside submodules that will be loaded later - # by _load_module_weights_from_mmap (non-block submodules) - dit_param_names = set() - for dit_module in modules.dits: - for pname, _ in dit_module.named_parameters(): - dit_param_names.add(pname) - truly_missing = [n for n in remaining_meta if n in dit_param_names] - if truly_missing: - logger.warning( - "Mmap loading: %d params still on meta device after " - "loading (first 5: %s). These may be loaded later by " - "_load_module_weights_from_mmap or are expected to be " - "meta (e.g. DTensor placeholders).", - len(truly_missing), - truly_missing[:5], + if binding.file_path not in file_cache: + file_cache[binding.file_path] = safe_open( + binding.file_path, + framework="pt", + device="cpu", ) + tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key) + if is_parameter: + replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad) + parent._parameters[leaf_name] = replacement + else: + replacement = tensor + parent._buffers[leaf_name] = replacement + loaded_names.add(runtime_name) + + logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names)) # Keep file handles open — _shard_and_pin will read from the mmap views. # They will be released after _shard_and_pin completes (when params are # replaced with offload placeholders). self._mmap_file_cache = file_cache - self._mmap_model_to_ckpt = model_to_ckpt - # The regular loader runs model-specific post-load transforms after # assigning checkpoint tensors. The mmap path bypasses that loader, so # preserve the same lifecycle for transforms such as Cosmos3's fp32 # timestep embedder. - for dit_module in modules.dits: + for dit_name, dit_module in zip(modules.dit_names, modules.dits): + # Restore non-persistent buffers before post-load hooks and strict + # validation. They are constructor-derived and intentionally have + # no checkpoint binding. + bufs = saved_buffers.get(id(dit_module), {}) + for name, buf in bufs.items(): + parent_path, _, leaf_name = name.rpartition(".") + parent = dit_module.get_submodule(parent_path) + restore_device = torch.device("cpu") if self._using_rank_local_mmap else self.device + parent._buffers[leaf_name] = buf.to(restore_device) + post_load_weights = getattr(dit_module, "post_load_weights", None) if callable(post_load_weights): post_load_weights() @@ -1005,73 +1292,45 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # load path (e.g. Cosmos3 checks for missing audio/action weights). validate = getattr(dit_module, "validate_loaded_weights", None) if callable(validate): - validate(loaded_names) + local_loaded_names = { + name.removeprefix(f"{dit_name}.") for name in loaded_names if name.startswith(f"{dit_name}.") + } + validate(local_loaded_names) - # Restore non-persistent buffers that were saved before meta - # conversion. These buffers (e.g. RoPE inv_freq, timestep - # freqs) are computed from formulas in __init__ and are not - # present in the checkpoint. By saving and restoring them - # generically, we avoid model-specific buffer rebuild code. - bufs = saved_buffers.get(id(dit_module), {}) - for name, buf in bufs.items(): - parent = dit_module - parts = name.split(".") - for part in parts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - # Place on the same device as the module's parameters - # Place on the offloader's target device (not the module's - # current device, which may still be meta at this point). - parent._buffers[parts[-1]] = buf.to(self.device) - - def _load_module_weights_from_mmap(self, module: nn.Module, dit_name: str, child_name: str) -> None: - """Load a non-block DiT submodule's weights from safetensors via mmap.""" - from safetensors import safe_open - - if not hasattr(self, "_mmap_file_cache") or not self._mmap_file_cache: - return - - file_cache = self._mmap_file_cache - model_to_ckpt = getattr(self, "_mmap_model_to_ckpt", {}) - loaded = 0 - - for pname, param in module.named_parameters(): - if not (hasattr(param, "is_meta") and param.is_meta): + # Post-load hooks may rebind parameters (for example while casting a + # submodule), so associate deferred bounded transforms with the final + # runtime tensor objects rather than the initial mmap replacements. + self._mmap_transforms_by_tensor_id.clear() + for runtime_name, binding in plan.bindings.items(): + if binding.transform is None: continue + parent_path, _, leaf_name = runtime_name.rpartition(".") + parent = pipeline.get_submodule(parent_path) + target = parent._parameters.get(leaf_name) + if target is None: + target = parent._buffers.get(leaf_name) + if target is None: + raise RuntimeError( + f"Host-weight transform target {runtime_name!r} no longer exists after post-load processing" + ) + self._mmap_transforms_by_tensor_id[id(target)] = binding.transform - # Build full model param name and look up in reverse remap - full_name = f"{dit_name}.{child_name}.{pname}" - entry = model_to_ckpt.get(full_name) - if entry is None: - # Try without dit_name prefix - full_name = f"{child_name}.{pname}" - entry = model_to_ckpt.get(full_name) - if entry is None: - continue - - ckpt_key, file_path = entry - if file_path not in file_cache: - file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") - f = file_cache[file_path] - tensor = f.get_tensor(ckpt_key) - - # Replace meta param with mmap view - parent = module - parts = pname.split(".") - for part in parts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) - self._attach_mmap_param_attrs(full_name, replacement, param) - parent._parameters[parts[-1]] = replacement - loaded += 1 - - if loaded > 0: - logger.info("Loaded %d params for submodule %s.%s via mmap", loaded, dit_name, child_name) + remaining_meta: list[str] = [] + for dit_name, dit_module in zip(modules.dit_names, modules.dits): + remaining_meta.extend( + f"{dit_name}.{name}" for name, tensor in dit_module.named_parameters() if tensor.is_meta + ) + for name, tensor in dit_module.named_buffers(): + parent_path, _, leaf_name = name.rpartition(".") + owner = dit_module.get_submodule(parent_path) + if leaf_name not in owner._non_persistent_buffers_set and tensor.is_meta: + remaining_meta.append(f"{dit_name}.{name}") + if remaining_meta: + raise RuntimeError( + "The prevalidated host-weight plan left " + f"{len(remaining_meta)} DiT tensors on the meta device " + f"(first 5: {remaining_meta[:5]})." + ) def _init_dp_group(self) -> None: """Reuse the process group initialized by parallel_state. @@ -1265,6 +1524,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) sub_hooks = [last_hook] for i, block in enumerate(blocks[:-1]): @@ -1280,6 +1541,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) sub_hooks.append(hook) @@ -1322,9 +1585,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): logger.debug("Skipped blocks module %s", name) continue - has_meta = any(getattr(param, "is_meta", False) for param in module.parameters()) - if has_meta: - self._load_module_weights_from_mmap(module, dit_name, name) module_mb = ( sum( param.nelement() * param.element_size() if not getattr(param, "is_meta", False) else 0 @@ -1345,7 +1605,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): try: module.to(self.device) except (NotImplementedError, RuntimeError): - self._load_module_weights_from_mmap(module, dit_name, name) # Non-persistent buffers such as RoPE frequencies do not # exist in the checkpoint and must be reconstructed. has_meta_buffer = any(getattr(buffer, "is_meta", False) for buffer in module.buffers(recurse=True)) @@ -1375,6 +1634,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): buffer.data = buffer.data.to(self.device, non_blocking=True) def enable(self, pipeline: nn.Module) -> None: + """Enable DLO and make partial startup failures transactional.""" + try: + self._enable(pipeline) + except BaseException: + try: + self.disable() + except BaseException: + logger.exception("DLO cleanup failed while handling an enable failure") + raise + + def _enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("DistributedLayerwiseOffloadBackend already enabled") return @@ -1388,6 +1658,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): modules = ModuleDiscovery.discover(pipeline) if not modules.dits: + if self.host_weight_plan is not None: + raise RuntimeError( + "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it" + ) logger.warning("No DiT/transformer modules found, skipping distributed layer-wise offloading") return @@ -1402,33 +1676,57 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.config.dlo_resident_layers, ) - # Load weights via mmap for DLO+AllGather. - # TODO(offload): Add a rank-local mmap path for dlo_no_use_allgather. - # It must apply declarative checkpoint-to-runtime adapters before the - # standard TP loader shards each block, and must cover staged encoders - # and VAEs as well as DiT blocks. The current AllGather mmap path is - # deliberately not reused because it changes the weight layout. - # Gate condition MUST match diffusers_loader.py: - # supports_mmap_loading(pipeline) and not _has_online_quant - # When the gate is False, the loader has already loaded weights via - # regular load_weights() + _process_weights_after_loading(). - if self.config.dlo_use_allgather and self.dp_size > 1: - _has_online_quant = any( - getattr(getattr(module, "quant_method", None), "uses_meta_device", False) - for module in pipeline.modules() - ) - if _has_online_quant: - raise ValueError( - "Online quantization is incompatible with DLO+AllGather: " - "the sharding + AllGather mechanism flattens weights by " - "dtype, which breaks quantized weight/scale layouts. " - "Please use --dlo-no-use-allgather or disable online " - "quantization." + # Storage selection belongs to the loader. DLO consumes the exact + # prevalidated plan that caused the loader to skip materialization; + # without a plan, all weights must already come from the ordinary + # loader. The transfer protocol is selected independently below. + self._using_mmap = self.host_weight_plan is not None + self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1 + if self._using_mmap: + if self.host_weight_plan.backing_kind == "host_weight_runtime": + carrier = self.host_weight_plan.lease_carrier + if carrier is None: + raise RuntimeError("DLO received a Host Weight Runtime plan without a lease carrier") + self._host_weight_lease = carrier.take() + if self._host_weight_lease.closed: + raise RuntimeError("DLO received a closed Host Weight Runtime lease") + # The final-layout restorer has already rebound the model to + # immutable host tensors. Treat those tensors as mmap-like + # sources; transport setup below selects registered direct + # H2D or the bounded two-slot staging fallback. + self._using_rank_local_mmap = True + logger.info( + "DLO consuming final-layout Host Weight Runtime lease %s", + self._host_weight_lease.provenance.resolution_id, + ) + elif self.host_weight_plan.backing_kind == "checkpoint_mmap": + self._load_weights_via_mmap( + pipeline, + modules, + self.host_weight_plan, ) - if supports_mmap_loading(pipeline): - self._load_weights_via_mmap(pipeline, modules) else: - logger.info("Weights loaded via regular loader — skipping mmap (model does not support mmap)") + raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}") + if self._using_rank_local_mmap: + logger.info( + "DLO rank-local host storage enabled: source pages are " + "node-shared; transfer setup will select registered direct H2D or bounded host staging" + ) + else: + remaining_meta = [ + name + for dit_name, dit_module in zip(modules.dit_names, modules.dits) + for name, tensor in chain( + dit_module.named_parameters(), + dit_module.named_buffers(), + ) + if getattr(tensor, "is_meta", False) + ] + if remaining_meta: + raise RuntimeError( + f"DLO received meta tensors without a loader-owned host-weight plan (first 5: {remaining_meta[:5]})" + ) + logger.info("DLO is using host tensors materialized by the ordinary loader") # Keep VAE/encoders on CPU; move to GPU on-demand via hooks. # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9) @@ -1528,6 +1826,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook] @@ -1544,6 +1844,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) block_hooks.append(hook) @@ -1572,24 +1874,43 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.device, self.copy_stream, self.config.pin_cpu_memory, + rank_local_mmap=self._using_rank_local_mmap, + defer_staging=bool(self._all_hook_groups), + tensor_transforms=self._mmap_transforms_by_tensor_id, ) pipeline._dlo_residency_controller = self + all_hooks = [hook for group in self._all_hook_groups for hook in group] + self._configure_hwr_transfer(all_hooks) + if not self._all_hook_groups: self.enabled = bool(self._resident_blocks) + if self._using_mmap and not self.enabled: + self._release_mmap_handles() return # Unified allocation: 2 shared output buffers + 2 shared shard buffers # sized to the max block across ALL module groups (gen_layers + # language_model). Groups execute sequentially, so 2 buffers suffice. - all_hooks: list[DistributedLayerwiseOffloadHook] = [] - for group in self._all_hook_groups: - all_hooks.extend(group) - unified_buffers = self._allocate_shared_buffers(all_hooks) unified_shard_buffers = None if self.dp_size > 1: unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) + unified_cpu_staging = None + cpu_staging_events = None + if self._using_rank_local_mmap and not self._using_registered_mmap: + unified_cpu_staging = self._allocate_shared_cpu_staging_buffers( + all_hooks, + self._resident_layer_group, + ) + cpu_staging_events = [None, None] + if self._resident_layer_group is not None: + # Resident and streamed layers execute in the same stage and + # reuse the same host slots. Events serialize slot reuse. + self._resident_layer_group._cpu_staging_buffers = [ + buffers for buffers in unified_cpu_staging if buffers is not None + ] + self._resident_layer_group._cpu_staging_events = cpu_staging_events # Shared slot-group tracker: _shared_slot_group[slot] = group_id # that last wrote to that slot. Group-first hooks use this to @@ -1602,6 +1923,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): hook._owns_buffers = False if unified_shard_buffers is not None: hook.gpu_shard_buffers = unified_shard_buffers + if unified_cpu_staging is not None and cpu_staging_events is not None: + hook.cpu_staging_buffers = unified_cpu_staging + hook.cpu_staging_events = cpu_staging_events hook._group_id = group_idx hook._shared_slot_group = shared_slot_group @@ -1628,9 +1952,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.enabled = True - # Release mmap file handles — _shard_and_pin has copied all shards, - # params now point to offload placeholders (not mmap views). - self._release_mmap_handles() + if self._using_mmap and not self._using_rank_local_mmap: + # AllGather mode copied each rank's persistent shard, so the source + # mappings are no longer needed. Rank-local mode retains them as + # the node-shared host master until disable(). + self._release_mmap_handles() # Assign GPU buffers to sharded on-demand modules (VAE/encoders). # Each module gets a dedicated input (shard-sized) and output @@ -1652,13 +1978,21 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._cleanup_after_loading() def _release_mmap_handles(self) -> None: - """Release safetensors mmap file handles.""" + """Release source handles and the transport-owned HWR lease.""" + if self._host_registration is not None: + raise HostRegistrationCleanupError("cannot close HWR mappings while host registration is still active") + self._mmap_transforms_by_tensor_id.clear() if hasattr(self, "_mmap_file_cache"): self._mmap_file_cache.clear() del self._mmap_file_cache - if hasattr(self, "_mmap_model_to_ckpt"): - del self._mmap_model_to_ckpt logger.info("Released safetensors mmap file handles") + lease = self._host_weight_lease + self._host_weight_lease = None + if lease is not None and not lease.closed: + lease.close() + logger.info("Released Host Weight Runtime lease %s", lease.provenance.resolution_id) + if self.host_weight_plan is not None and self.host_weight_plan.lease_carrier is not None: + self.host_weight_plan.lease_carrier.close() def _cleanup_after_loading(self) -> None: """Synchronize and release freed device/CPU memory after sharding.""" @@ -1674,9 +2008,32 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): pass def disable(self) -> None: - if not self.enabled: + has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed + has_registration = self._host_registration is not None + has_carrier = ( + self.host_weight_plan is not None + and self.host_weight_plan.lease_carrier is not None + and not self.host_weight_plan.lease_carrier.closed + ) + has_partial_hooks = bool( + self._blocks + or self._all_hook_groups + or getattr(self, "_on_demand_handles", ()) + or self._resident_layer_group is not None + ) + if ( + not self.enabled + and not hasattr(self, "_mmap_file_cache") + and not has_open_lease + and not has_registration + and not has_carrier + and not has_partial_hooks + ): return + if self._using_rank_local_mmap or has_registration: + current_omni_platform.synchronize() + for blocks in self._blocks: for block in blocks: remove_distributed_block_hook(block) @@ -1691,6 +2048,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups.clear() self._resident_blocks.clear() self._resident_layer_group = None + self._release_registered_mmap() + self._release_mmap_handles() + self._using_mmap = False + self._using_rank_local_mmap = False + self._using_registered_mmap = False self.enabled = False logger.info("Distributed layer-wise offloading disabled") @@ -1730,6 +2092,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) return shared_buffers + @staticmethod + def _allocate_shared_cpu_staging_buffers( + hooks: list[DistributedLayerwiseOffloadHook], + resident_group: PinnedResidentLayerGroup | None = None, + ) -> list[dict[torch.dtype, torch.Tensor] | None]: + """Allocate two bounded host slots for rank-local mmap -> device copies.""" + max_sizes: dict[torch.dtype, int] = {} + for hook in hooks: + for dtype, metas in hook.metadata.items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + if resident_group is not None: + for state in resident_group._states: + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + + pin_memory = hooks[0].pin_memory if hooks else bool(resident_group and resident_group.pin_memory) + shared_staging: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + for slot in range(2): + buffers: dict[torch.dtype, torch.Tensor] = {} + for dtype, total_numel in max_sizes.items(): + buffer = torch.empty(total_numel, dtype=dtype, device="cpu") + if pin_memory: + buffer = buffer.pin_memory() + buffers[dtype] = buffer + shared_staging[slot] = buffers + + logger.info( + "Allocated 2 shared host staging buffers for rank-local mmap (max block size: %s, pinned=%s)", + {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()}, + pin_memory, + ) + return shared_staging + @staticmethod def _allocate_shared_shard_buffers( hooks: list[DistributedLayerwiseOffloadHook], diff --git a/vllm_omni/diffusion/offloader/host_registration.py b/vllm_omni/diffusion/offloader/host_registration.py new file mode 100644 index 000000000..c9f225122 --- /dev/null +++ b/vllm_omni/diffusion/offloader/host_registration.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Platform-neutral lifecycle for registering existing host mappings.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +import torch + +from vllm_omni.host_weight_runtime import MappedHostRegion + + +class HostRegistrationError(RuntimeError): + """Raised when existing host ranges cannot be registered safely.""" + + +class HostRegistrationBudgetError(HostRegistrationError): + """Raised before mutation when the complete mapping exceeds its budget.""" + + +class HostRegistrationCleanupError(HostRegistrationError): + """Raised when registered ranges cannot be released safely.""" + + def __init__(self, message: str, *, active_registration: HostRegistration | None = None) -> None: + super().__init__(message) + self.active_registration = active_registration + + +class HostRegistration(Protocol): + """Transport-owned registration for one process-local host mapping.""" + + @property + def total_bytes(self) -> int: ... + + @property + def region_count(self) -> int: ... + + def close(self) -> tuple[str, ...]: ... + + +def register_host_mappings( + regions: Sequence[MappedHostRegion], + *, + device: torch.device, + max_bytes: int | None, +) -> HostRegistration: + """Register immutable ranges with the active platform implementation.""" + if device.type == "cuda": + from .cuda_host_registration import CudaHostRegistration + + return CudaHostRegistration.create(regions, max_bytes=max_bytes) + raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}") + + +__all__ = [ + "HostRegistration", + "HostRegistrationBudgetError", + "HostRegistrationCleanupError", + "HostRegistrationError", + "register_host_mappings", +] diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 9979d01b1..662e1304f 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -139,29 +139,57 @@ class LayerwiseOffloadHook(ModelHook): weights_with_local = [] for name, t in name2weights.items(): local_t = t.to_local() if hasattr(t, "to_local") else t - weights_with_local.append((name, t, local_t)) - total_numel = sum(local.numel() for _, _, local in weights_with_local) + stride = local_t.stride() + storage_numel = ( + 0 + if local_t.numel() == 0 + else 1 + sum((size - 1) * axis_stride for size, axis_stride in zip(local_t.shape, stride)) + ) + weights_with_local.append((name, t, local_t, storage_numel, stride)) + total_numel = sum(storage_numel for _, _, _, storage_numel, _ in weights_with_local) cpu_tensor = torch.empty(total_numel, dtype=dtype, device="cpu", pin_memory=pin_memory) current_offset = 0 - for name, original_tensor, local_tensor in weights_with_local: - numel = local_tensor.numel() - cpu_tensor[current_offset : current_offset + numel].copy_(local_tensor.flatten()) + for ( + name, + original_tensor, + local_tensor, + storage_numel, + stride, + ) in weights_with_local: + if local_tensor.is_contiguous(): + flat_storage = local_tensor.flatten() + else: + # Cutlass FP8 weights use a transposed physical layout. + # Preserve it across the flattened CPU staging buffer. + flat_storage = torch.zeros( + storage_numel, + dtype=dtype, + device=local_tensor.device, + ) + physical_view = torch.as_strided( + flat_storage, + size=local_tensor.shape, + stride=stride, + ) + physical_view.copy_(local_tensor) + cpu_tensor[current_offset : current_offset + storage_numel].copy_(flat_storage) if dtype not in dtype_metadata: dtype_metadata[dtype] = [] dtype_metadata[dtype].append( { "name": name, "offset": current_offset, - "numel": numel, + "numel": storage_numel, "shape": local_tensor.shape, + "stride": stride, } ) LayerwiseOffloadHook._set_tensor_storage( original_tensor, LayerwiseOffloadHook._make_offload_placeholder(original_tensor) ) - current_offset += numel + current_offset += storage_numel dtype_cpu_flattened_weights[dtype] = cpu_tensor @@ -210,7 +238,11 @@ class LayerwiseOffloadHook(ModelHook): LayerwiseOffloadHook._set_tensor_storage( target_param_or_buf, - gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]].view(metadata["shape"]), + torch.as_strided( + gpu_weight[metadata["offset"] : metadata["offset"] + metadata["numel"]], + size=metadata["shape"], + stride=metadata["stride"], + ), ) self._prefetch_done = evt diff --git a/vllm_omni/diffusion/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py index da32ab520..ab584c790 100644 --- a/vllm_omni/diffusion/offloader/module_residency.py +++ b/vllm_omni/diffusion/offloader/module_residency.py @@ -1,89 +1,271 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""On-demand module staging backed by an immutable pinned CPU snapshot.""" +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""On-demand module staging backed by immutable pinned CPU storage.""" from __future__ import annotations +from collections.abc import Iterable +from dataclasses import dataclass from itertools import chain from typing import Any import torch from torch import nn +from vllm.logger import init_logger from vllm_omni.platforms import current_omni_platform -from .tensor_utils import set_tensor_storage +from .tensor_utils import is_dtensor, set_tensor_storage + +logger = init_logger(__name__) -class PinnedModuleStager: - """Stage an immutable module without copying device weights back to CPU. +class BoundedAllocatorCache: + """Retain reusable allocator blocks without monopolizing device memory. - ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter - after each forward. In inference the weights are immutable, so retain one - pinned CPU master instead. ``load`` materializes fresh device tensors from - that master; ``offload`` only rebinds the Parameters/Buffers to it and - releases the device allocations. + Component offload normally calls ``empty_cache`` after every stage. That + makes the next stage return to the device allocator even though PyTorch's + cached blocks are immediately reusable. This policy keeps the cache while + both of these bounds hold: + + * cached-but-unallocated memory is at most 25% of device capacity; and + * at least 5% of device capacity is physically free. + + Missing memory telemetry is handled conservatively by releasing the cache. + Failure paths can force release; normal executor shutdown keeps its own + unconditional device-cache cleanup because this policy is not global. """ def __init__( self, - module: nn.Module, + device: torch.device, + *, + max_cached_fraction: float = 0.25, + min_free_fraction: float = 0.05, + ) -> None: + if not 0.0 <= max_cached_fraction <= 1.0: + raise ValueError(f"max_cached_fraction must be in [0, 1], got {max_cached_fraction}") + if not 0.0 <= min_free_fraction <= 1.0: + raise ValueError(f"min_free_fraction must be in [0, 1], got {min_free_fraction}") + self.device = device + self.max_cached_fraction = max_cached_fraction + self.min_free_fraction = min_free_fraction + + def _should_release(self) -> bool: + reserved = int(torch.accelerator.memory_reserved(self.device)) + allocated = int(torch.accelerator.memory_allocated(self.device)) + free, total = current_omni_platform.get_device_memory(self.device) + cached = max(0, reserved - allocated) + return cached > int(total * self.max_cached_fraction) or free < int(total * self.min_free_fraction) + + def release_if_needed(self, *, force: bool = False) -> bool: + """Release cached blocks when a bound is crossed or release is forced.""" + if not force: + try: + if not self._should_release(): + return False + except Exception as exc: + # Preserve the pre-retention behavior on platforms that do not + # expose allocator telemetry through torch.accelerator. + logger.debug("Allocator cache telemetry unavailable; releasing cache: %s", exc) + current_omni_platform.empty_cache() + return True + + +@dataclass(frozen=True) +class _TensorBinding: + target: torch.Tensor + dtype: torch.dtype + shape: tuple[int, ...] + stride: tuple[int, ...] + storage_offset: int + + +@dataclass +class _StorageGroup: + master: torch.Tensor + bindings: list[_TensorBinding] + + +class PinnedModuleStager: + """Stage immutable module groups without copying device weights to CPU. + + ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter + after each forward. In inference the weights are immutable, so retain one + pinned CPU master instead. ``load`` materializes device storage from that + master; ``offload`` only rebinds Parameters and buffers to the master. + + A module iterable is treated as one staging group. It uses one copy stream + and one reusable completion event. Tensors sharing storage keep their + shapes, strides, offsets, dtypes, and aliases across every transition. + """ + + def __init__( + self, + module: nn.Module | Iterable[nn.Module], device: torch.device, *, pin_memory: bool = True, copy_stream: Any | None = None, + cache_retention: BoundedAllocatorCache | None = None, ) -> None: - self.device = device - self.copy_stream = copy_stream or current_omni_platform.Stream() - self.loaded = False - self._entries: list[tuple[torch.Tensor, torch.Tensor]] = [] - self._device_tensors: list[torch.Tensor] = [] + modules = (module,) if isinstance(module, nn.Module) else tuple(module) + if not modules or not all(isinstance(item, nn.Module) for item in modules): + raise ValueError("PinnedModuleStager requires at least one nn.Module") - # parameters()/buffers() remove duplicate objects by default, which - # preserves tied weights and buffer aliases when their storage moves. - for target in chain(module.parameters(), module.buffers()): - master = target.detach() - if master.device.type != "cpu": - master = master.to("cpu") + self.device = device + self.copy_stream = copy_stream if copy_stream is not None else current_omni_platform.Stream() + self._ready_event = current_omni_platform.Event() + self.cache_retention = cache_retention + self.loaded = False + self._groups = self._snapshot_groups(modules, pin_memory=pin_memory) + self._device_storages: list[torch.Tensor] = [] + self._restore_masters() + + @staticmethod + def _local_tensor(target: torch.Tensor) -> torch.Tensor: + return target.to_local() if is_dtensor(target) else target + + @classmethod + def _snapshot_groups( + cls, + modules: tuple[nn.Module, ...], + *, + pin_memory: bool, + ) -> list[_StorageGroup]: + targets: list[torch.Tensor] = [] + seen_targets: set[int] = set() + for target in chain.from_iterable(chain(item.parameters(), item.buffers()) for item in modules): + if id(target) not in seen_targets: + seen_targets.add(id(target)) + targets.append(target) + + grouped: dict[tuple[Any, ...], tuple[torch.Tensor, list[_TensorBinding]]] = {} + for target in targets: + local = cls._local_tensor(target) + if local.is_meta: + raise ValueError("PinnedModuleStager cannot snapshot a meta tensor") + storage = local.untyped_storage() + if storage.nbytes() == 0: + storage_key: tuple[Any, ...] = ("empty", id(target)) + else: + storage_key = ( + local.device.type, + local.device.index, + storage.data_ptr(), + storage.nbytes(), + ) + binding = _TensorBinding( + target=target, + dtype=local.dtype, + shape=tuple(local.shape), + stride=tuple(local.stride()), + storage_offset=local.storage_offset(), + ) + if storage_key not in grouped: + grouped[storage_key] = (local, []) + grouped[storage_key][1].append(binding) + + groups: list[_StorageGroup] = [] + for source, bindings in grouped.values(): + storage = source.untyped_storage() + storage_view = torch.empty(0, dtype=torch.uint8, device=source.device).set_( + storage, + 0, + (storage.nbytes(),), + (1,), + ) + master = storage_view.detach() if storage_view.device.type == "cpu" else storage_view.to("cpu") if pin_memory and not master.is_pinned(): master = master.pin_memory() - set_tensor_storage(target, master) - self._entries.append((target, master)) + groups.append(_StorageGroup(master=master, bindings=bindings)) + return groups + + @staticmethod + def _view(backing: torch.Tensor, binding: _TensorBinding) -> torch.Tensor: + return torch.empty(0, dtype=binding.dtype, device=backing.device).set_( + backing.untyped_storage(), + binding.storage_offset, + binding.shape, + binding.stride, + ) + + def _bind(self, storages: list[torch.Tensor]) -> None: + for storage, group in zip(storages, self._groups): + for binding in group.bindings: + set_tensor_storage(binding.target, self._view(storage, binding)) + + def _restore_masters(self) -> None: + self._bind([group.master for group in self._groups]) + + def set_cache_retention(self, cache_retention: BoundedAllocatorCache | None) -> None: + self.cache_retention = cache_retention + + def _release_cache(self, *, force: bool = False) -> None: + if self.cache_retention is None: + current_omni_platform.empty_cache() + else: + self.cache_retention.release_if_needed(force=force) + + def _load_once(self) -> None: + device_storages = [torch.empty_like(group.master, device=self.device) for group in self._groups] + compute_stream = current_omni_platform.current_stream() + self.copy_stream.wait_stream(compute_stream) + with current_omni_platform.stream(self.copy_stream): + for device_storage, group in zip(device_storages, self._groups): + device_storage.copy_(group.master, non_blocking=group.master.is_pinned()) + self._ready_event.record(self.copy_stream) + + self._bind(device_storages) + compute_stream.wait_event(self._ready_event) + self._device_storages = device_storages + self.loaded = True + + def _cleanup_failed_load(self) -> None: + try: + current_omni_platform.synchronize() + except Exception: + logger.debug("Device synchronization failed while cleaning up module staging", exc_info=True) + self._restore_masters() + self._device_storages.clear() + self.loaded = False + self._release_cache(force=True) def load(self) -> None: if self.loaded: return - - # Allocate on the compute stream so the caching allocator can reuse - # memory released by the preceding encoder/DiT stage. - device_tensors = [torch.empty_like(master, device=self.device) for _, master in self._entries] - compute_stream = current_omni_platform.current_stream() - self.copy_stream.wait_stream(compute_stream) - ready = current_omni_platform.Event() - with current_omni_platform.stream(self.copy_stream): - for device_tensor, (_, master) in zip(device_tensors, self._entries): - device_tensor.copy_(master, non_blocking=master.is_pinned()) - ready.record(self.copy_stream) - - for device_tensor, (target, _) in zip(device_tensors, self._entries): - set_tensor_storage(target, device_tensor) - compute_stream.wait_event(ready) - self._device_tensors = device_tensors - self.loaded = True + try: + self._load_once() + except torch.OutOfMemoryError: + # A retained cache is normally reusable by this process, but an + # explicit flush gives external memory pressure one bounded retry. + self._cleanup_failed_load() + try: + self._load_once() + except BaseException: + self._cleanup_failed_load() + raise + except BaseException: + self._cleanup_failed_load() + raise def offload(self) -> None: if not self.loaded: return # The module has completed on the compute stream. Synchronize once at - # the stage boundary, then release HBM by rebinding to the unchanged - # host masters. No device-to-host transfer is necessary. - current_omni_platform.synchronize() - for target, master in self._entries: - set_tensor_storage(target, master) - self._device_tensors.clear() + # the stage boundary, then discard device storage without any D2H copy. + try: + current_omni_platform.synchronize() + self._restore_masters() + except BaseException: + self._device_storages.clear() + self.loaded = False + self._release_cache(force=True) + raise + self._device_storages.clear() self.loaded = False - current_omni_platform.empty_cache() + self._release_cache() -__all__ = ["PinnedModuleStager"] +__all__ = ["BoundedAllocatorCache", "PinnedModuleStager"] diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index f8943757e..2e1443df7 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -55,20 +55,3 @@ class OffloadPlan: def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None: """Retrieve the OffloadPlan declared by the pipeline, if any.""" return getattr(pipeline, "_offload_plan", None) - - -def supports_mmap_loading(pipeline: nn.Module) -> bool: - """Whether the pipeline supports mmap weight loading. - - A pipeline supports mmap if any module defines ``_remap_ckpt_key``, - which maps checkpoint keys to model parameter names. Models without - this method must use the regular ``load_weights()`` path. - - This check is shared between ``diffusers_loader.py`` (to decide whether - to skip ``load_weights``) and ``DistributedLayerwiseOffloadBackend.enable()`` - (to decide whether to call ``_load_weights_via_mmap``), ensuring both - paths use the same gate condition. - """ - return bool(getattr(pipeline, "_supports_mmap_loading", True)) and any( - callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules() - ) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 0b031be3a..36e0f7b22 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator +from contextlib import contextmanager + import torch from torch import nn from torch.distributed._tensor import DTensor # type: ignore[attr-defined] @@ -9,7 +12,7 @@ from vllm.logger import init_logger from vllm_omni.diffusion.hooks import HookRegistry, ModelHook from vllm_omni.platforms import current_omni_platform -from .base import OffloadBackend, OffloadConfig +from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload from .module_collector import ModuleDiscovery logger = init_logger(__name__) @@ -123,6 +126,7 @@ def apply_sequential_offload( device: torch.device, pin_memory: bool = True, use_hsdp: bool = False, + offload_initial_dits: bool = False, ) -> None: """Apply sequential offloading hooks to DiT and encoder modules. @@ -136,6 +140,7 @@ def apply_sequential_offload( device: Target GPU device for loading pin_memory: Whether to pin CPU memory for faster transfers use_hsdp: Whether HSDP is enabled (affects non_blocking behavior) + offload_initial_dits: Whether to begin with all DiT modules on CPU. Example: >>> apply_sequential_offload( @@ -170,6 +175,14 @@ def apply_sequential_offload( registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) logger.debug("Registered offload hook for %s", enc.__class__.__name__) + if offload_initial_dits: + try: + for dit_mod in dit_modules: + _get_sequential_offload_hook(dit_mod)._to_cpu(dit_mod) + except Exception: + remove_sequential_offload([*dit_modules, *encoder_modules]) + raise + def remove_sequential_offload(modules: list[nn.Module]) -> None: """Remove sequential offloading hooks from modules. @@ -188,6 +201,31 @@ def remove_sequential_offload(modules: list[nn.Module]) -> None: logger.debug("Removed offload hook from %s", module.__class__.__name__) +def _get_sequential_offload_hook(module: nn.Module) -> SequentialOffloadHook: + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + hook = registry.get_hook(SequentialOffloadHook._HOOK_NAME) if registry is not None else None + if not isinstance(hook, SequentialOffloadHook): + raise RuntimeError(f"{module.__class__.__name__} has no sequential offload hook") + return hook + + +@contextmanager +def sequential_offload_component(module: nn.Module) -> Iterator[None]: + """Activate and release a hooked component called outside ``forward``.""" + hook = _get_sequential_offload_hook(module) + try: + hook.pre_forward(module) + yield + except BaseException: + try: + hook._to_cpu(module) + except Exception: + logger.exception("Failed to release %s after component failure", module.__class__.__name__) + raise + else: + hook._to_cpu(module) + + class ModelLevelOffloadBackend(OffloadBackend): """Model-level (sequential) offloading backend. @@ -197,23 +235,17 @@ class ModelLevelOffloadBackend(OffloadBackend): def __init__(self, config: OffloadConfig, device: torch.device): super().__init__(config, device) self._offload_modules: list[nn.Module] = [] # Track modules with hooks - self._custom_pipeline: nn.Module | None = None + self._custom_pipeline: SupportsModelCpuOffload | None = None def enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("ModelLevelOffloadBackend already enabled") return - # Pipelines with a nested transformer (e.g. Cosmos3's reasoner/generator - # pathways) own their own mutual-exclusion swaps and cannot be offloaded - # with the generic encoder<->DiT hooks. Delegate to the custom hook. - # TODO: Cosmos3 is the only implementer today, so we duck-type the - # enable/disable_omni_model_cpu_offload pair via getattr. Once a second - # pipeline needs it, formalize this as a Protocol (like - # SupportsComponentDiscovery) instead of getattr detection. - custom_enable = getattr(pipeline, "enable_omni_model_cpu_offload", None) - if callable(custom_enable): - custom_enable( + # Pipelines with non-forward component entry points own their complete + # mutual-exclusion lifecycle. Delegate through the explicit protocol. + if isinstance(pipeline, SupportsModelCpuOffload): + pipeline.enable_omni_model_cpu_offload( device=self.device, pin_memory=self.config.pin_cpu_memory, use_hsdp=self.config.use_hsdp, @@ -283,9 +315,7 @@ class ModelLevelOffloadBackend(OffloadBackend): return if self._custom_pipeline is not None: - custom_disable = getattr(self._custom_pipeline, "disable_omni_model_cpu_offload", None) - if callable(custom_disable): - custom_disable() + self._custom_pipeline.disable_omni_model_cpu_offload() self._custom_pipeline = None self.enabled = False logger.info("Model-level offloading disabled") diff --git a/vllm_omni/diffusion/offloader/startup.py b/vllm_omni/diffusion/offloader/startup.py new file mode 100644 index 000000000..7861ddfd2 --- /dev/null +++ b/vllm_omni/diffusion/offloader/startup.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generic startup handoff between model loaders and offload backends.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +from torch import nn + +from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan + +_STARTUP_STATE_ATTR = "_omni_offload_startup_state" + + +@dataclass +class OffloadStartupState: + """Loader-owned state consumed by the offloader startup boundary.""" + + host_weight_plan: HostWeightPlan | None = None + fresh_model_loader: Callable[[], nn.Module] | None = None + allow_fresh_retry: bool = False + + def __post_init__(self) -> None: + if self.allow_fresh_retry != (self.fresh_model_loader is not None): + raise ValueError("fresh-model retry policy requires exactly one retry callback") + + def close_loader_ownership(self) -> None: + """Release a plan that never reached a backend.""" + if self.host_weight_plan is None: + return + carrier = self.host_weight_plan.lease_carrier + if carrier is not None: + carrier.close() + + +def attach_offload_startup_state(model: nn.Module, state: OffloadStartupState) -> None: + """Attach one process-local startup handoff to a loaded pipeline.""" + previous = getattr(model, _STARTUP_STATE_ATTR, None) + if previous is not None: + previous.close_loader_ownership() + setattr(model, _STARTUP_STATE_ATTR, state) + + +def take_offload_startup_state(model: nn.Module) -> OffloadStartupState | None: + """Take and remove the loader handoff from a pipeline exactly once.""" + state = getattr(model, _STARTUP_STATE_ATTR, None) + if state is None: + return None + delattr(model, _STARTUP_STATE_ATTR) + return state + + +__all__ = [ + "OffloadStartupState", + "attach_offload_startup_state", + "take_offload_startup_state", +] diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 4e68bfd25..62c2330d5 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -8,7 +8,7 @@ from dataclasses import dataclass, field from typing import TypeGuard, cast from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.io_support import supports_audio_output +from vllm_omni.diffusion.io_support import get_diffusion_output_type, supports_audio_output from vllm_omni.diffusion.registry import DiffusionModelRegistry from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniPromptType @@ -146,6 +146,11 @@ def format_diffusion_outputs( is_text_output = postprocess_output.primary_key == "text" is_audio_output = supports_audio_output(od_config.model_class_name) + final_output_type = ( + postprocess_output.primary_key + if postprocess_output.primary_key in {"image", "video"} + else get_diffusion_output_type(od_config.model_class_name) + ) audio_sample_rate = _metadata_audio_sample_rate(postprocess_output.metadata) if is_audio_output and audio_sample_rate is None: model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name) @@ -167,6 +172,7 @@ def format_diffusion_outputs( postprocess_output=postprocess_output, is_text_output=is_text_output, is_audio_output=is_audio_output, + final_output_type=final_output_type, audio_sample_rate=audio_sample_rate, finished=diffusion_output.finished, ) @@ -272,6 +278,7 @@ def _format_single_prompt_output( postprocess_output: DiffusionPostprocessOutput, is_text_output: bool, is_audio_output: bool, + final_output_type: str, audio_sample_rate: int | None, finished: bool = True, ) -> list[OmniRequestOutput]: @@ -339,6 +346,7 @@ def _format_single_prompt_output( trajectory_log_probs=trajectory_log_probs, trajectory_decoded=trajectory_decoded, multimodal_output=mm_output, + final_output_type=final_output_type, stage_durations=diffusion_output.stage_durations, peak_memory_mb=diffusion_output.peak_memory_mb, finished=finished, diff --git a/vllm_omni/diffusion/quantization/hsdp_fp8.py b/vllm_omni/diffusion/quantization/hsdp_fp8.py index 7c4666e29..6ffa8f2b4 100644 --- a/vllm_omni/diffusion/quantization/hsdp_fp8.py +++ b/vllm_omni/diffusion/quantization/hsdp_fp8.py @@ -33,9 +33,17 @@ import types from torch import nn from vllm.logger import init_logger from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod +from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, +) logger = init_logger(__name__) +_FP8_TRANSPOSED_WEIGHT_METHODS = ( + Fp8LinearMethod, + Fp8PerTensorOnlineLinearMethod, +) + def _build_transposed_get_layer_params(original_bound_method): """Return a ``_get_layer_params`` replacement that transposes ``w``. @@ -55,9 +63,8 @@ def _build_transposed_get_layer_params(original_bound_method): def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int: """Make online-FP8 linear layers in ``model`` FSDP2-compatible. - For every layer whose ``quant_method`` is an :class:`Fp8LinearMethod` - and whose weight is currently a non-contiguous transpose view, this - function: + For every layer whose quantization method stores FP8 weights as a + non-contiguous transpose view, this function: 1. Replaces ``layer.weight`` with the underlying ``(out, in)`` row-major contiguous storage so FSDP2 ``fully_shard`` accepts it. @@ -76,7 +83,7 @@ def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int: patched_kernel_ids: set[int] = set() for module in model.modules(): qm = getattr(module, "quant_method", None) - if not isinstance(qm, Fp8LinearMethod): + if not isinstance(qm, _FP8_TRANSPOSED_WEIGHT_METHODS): continue weight = getattr(module, "weight", None) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 75ef7391a..ed6a816a9 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import importlib @@ -9,7 +9,7 @@ from vllm.model_executor.model_loader.utils import configure_quant_config from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry from vllm_omni.diffusion.config import set_current_diffusion_config -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig, uses_diffusers_adapter from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model from vllm_omni.diffusion.forward_context import get_forward_context @@ -61,6 +61,11 @@ _DIFFUSION_MODELS = { "pipeline_wan2_2", "Wan22Pipeline", ), + "WanDMDPipeline": ( + "wan2_2", + "pipeline_wan2_2", + "Wan22Pipeline", + ), "WanVACEPipeline": ( "wan2_2", "pipeline_wan2_2_vace", @@ -71,6 +76,21 @@ _DIFFUSION_MODELS = { "pipeline_ltx2", "LTX2Pipeline", ), + "LTX2TwoStagePipeline": ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2TwoStagePipeline", + ), + "LTX2DistilledOneStagePipeline": ( + "ltx2", + "pipeline_ltx2", + "LTX2DistilledOneStagePipeline", + ), + "LTX2DistilledTwoStagePipeline": ( + "ltx2", + "pipeline_ltx2_two_stage", + "LTX2DistilledTwoStagePipeline", + ), "LTX2DistilledPipeline": ( "ltx2", "pipeline_ltx2_two_stage", @@ -121,6 +141,11 @@ _DIFFUSION_MODELS = { "pipeline_wan2_2_i2v", "WanI2VDMD2Pipeline", ), + "LingBotWorldCausalDMDPipeline": ( + "lingbot_world", + "pipeline", + "LingBotWorldCausalDMDPipeline", + ), "LongCatImagePipeline": ( "longcat_image", "pipeline_longcat_image", @@ -151,6 +176,11 @@ _DIFFUSION_MODELS = { "pipeline_ming_imagegen", "MingImagePipeline", ), + "SanaWmPipeline": ( + "sana_wm", + "pipeline_sana_wm", + "SanaWmPipeline", + ), "InternVLAA1Pipeline": ( "internvla_a1", "pipeline_internvla_a1", @@ -161,6 +191,11 @@ _DIFFUSION_MODELS = { "pipeline_gr00t", "Gr00tN1d7Pipeline", ), + "Pi0Pipeline": ( + "pi0", + "pipeline_pi0", + "Pi0Pipeline", + ), "LongCatImageEditPipeline": ( "longcat_image", "pipeline_longcat_image_edit", @@ -231,21 +266,11 @@ _DIFFUSION_MODELS = { "pipeline_flux2", "Flux2Pipeline", ), - "DreamIDOmniPipeline": ( - "dreamid_omni", - "pipeline_dreamid_omni", - "DreamIDOmniPipeline", - ), "SenseNovaU1Pipeline": ( "sensenova_u1", "pipeline_sensenova_u1", "SenseNovaU1Pipeline", ), - "AudioXPipeline": ( - "audiox", - "pipeline_audiox", - "AudioXPipeline", - ), "HunyuanVideo15Pipeline": ( "hunyuan_video", "pipeline_hunyuan_video_1_5", @@ -261,10 +286,15 @@ _DIFFUSION_MODELS = { "pipeline_lingbot_video", "LingBotVideoPipeline", ), - "MagiHumanPipeline": ( - "magi_human", - "pipeline_magi_human", - "MagiHumanPipeline", + "SanaVideoPipeline": ( + "sana_video", + "pipeline_sana_video", + "SanaVideoPipeline", + ), + "SanaImageToVideoPipeline": ( + "sana_video", + "pipeline_sana_video_i2v", + "SanaImageToVideoPipeline", ), "OmniVoicePipeline": ( "omnivoice", @@ -286,16 +316,6 @@ _DIFFUSION_MODELS = { "pipeline_cosmos3", "Cosmos3OmniDiffusersPipeline", ), - "SoulXSingerPipeline": ( - "soulx_singer", - "pipeline_soulx_singer_svs", - "PipelineSoulXSingerSVS", - ), - "SoulXSingerSVCPipeline": ( - "soulx_singer", - "pipeline_soulx_singer_svc", - "PipelineSoulXSingerSVC", - ), "DiffusersAdapterPipeline": ( "diffusers_adapter", "pipeline_diffusers_adapter", @@ -306,6 +326,11 @@ _DIFFUSION_MODELS = { "pipeline_hidream_image", "HiDreamImagePipeline", ), + "HiDreamO1ImagePipeline": ( + "hidream_o1_image", + "pipeline_hidream_o1_image", + "HiDreamO1ImagePipeline", + ), "DreamZeroPipeline": ( "dreamzero", "pipeline_dreamzero", @@ -337,7 +362,12 @@ DiffusionModelRegistry = _ModelRegistry( _NO_CACHE_ACCELERATION = { # Pipelines that do not support cache acceleration (cache_dit / tea_cache). "NextStep11Pipeline", - "AudioXPipeline", + # π0 is a flow-matching VLA with a self-contained sample_actions loop and no + # DiT-style ``.transformer`` block list, so cache_dit / tea_cache cannot apply + # to it; list it here so a stray cache_backend override disables gracefully + # instead of erroring. + "Pi0Pipeline", + "LingBotWorldCausalDMDPipeline", } @@ -434,6 +464,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - try: sp_size = od_config.parallel_config.sequence_parallel_size + assert sp_size is not None if sp_size <= 1: return @@ -514,27 +545,30 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "OvisImagePipeline": "get_ovis_image_post_process_func", "BooguImagePipeline": "get_boogu_image_post_process_func", "WanPipeline": "get_wan22_post_process_func", + "WanDMDPipeline": "get_wan22_post_process_func", "WanVACEPipeline": "get_wan22_vace_post_process_func", "LTX2Pipeline": "get_ltx2_post_process_func", + "LTX2TwoStagePipeline": "get_ltx2_post_process_func", + "LTX2DistilledOneStagePipeline": "get_ltx2_post_process_func", + "LTX2DistilledTwoStagePipeline": "get_ltx2_post_process_func", "LTX2DistilledPipeline": "get_ltx2_post_process_func", "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func", "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func", "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", - "SoulXSingerPipeline": "get_soulxsinger_post_process_func", - "SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func", - "AudioXPipeline": "get_audiox_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanS2VPipeline": "get_wan22_s2v_post_process_func", "WanT2VDMD2Pipeline": "get_wan22_post_process_func", "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func", + "LingBotWorldCausalDMDPipeline": "get_lingbot_world_post_process_func", "LongCatImagePipeline": "get_longcat_image_post_process_func", "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_post_process_func", "BagelPipeline": "get_bagel_post_process_func", "LancePipeline": "get_lance_post_process_func", "MingImagePipeline": "get_ming_image_post_process_func", "InternVLAA1Pipeline": "get_internvla_a1_post_process_func", + "Pi0Pipeline": "get_pi0_post_process_func", "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", "FluxKontextPipeline": "get_flux_kontext_post_process_func", @@ -552,13 +586,14 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func", "HunyuanImage3Pipeline": "get_hunyuan_image3_post_process_func", "LingBotVideoPipeline": "get_lingbot_video_post_process_func", - "MagiHumanPipeline": "get_magi_human_post_process_func", + "SanaVideoPipeline": "get_sana_video_post_process_func", + "SanaImageToVideoPipeline": "get_sana_video_i2v_post_process_func", "OmniVoicePipeline": "get_omnivoice_post_process_func", - "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func", "Cosmos3OmniPipeline": "get_cosmos3_post_process_func", "HiDreamImagePipeline": "get_hidream_image_post_process_func", + "HiDreamO1ImagePipeline": "get_hidream_o1_image_post_process_func", "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func", "Krea2Pipeline": "get_krea2_post_process_func", "HunyuanImage3ForCausalMM": "get_hunyuan_image3_post_process_func", @@ -584,22 +619,23 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func", "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func", "WanPipeline": "get_wan22_pre_process_func", + "WanDMDPipeline": "get_wan22_pre_process_func", "WanVACEPipeline": "get_wan22_vace_pre_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func", "WanS2VPipeline": "get_wan22_s2v_pre_process_func", "WanT2VDMD2Pipeline": "get_wan22_pre_process_func", "WanI2VDMD2Pipeline": "get_wan22_i2v_pre_process_func", + "LingBotWorldCausalDMDPipeline": "get_lingbot_world_pre_process_func", "OmniGen2Pipeline": "get_omnigen2_pre_process_func", "HeliosPipeline": "get_helios_pre_process_func", "HeliosPyramidPipeline": "get_helios_pre_process_func", "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_pre_process_func", "LingBotVideoPipeline": "get_lingbot_video_pre_process_func", + "SanaImageToVideoPipeline": "get_sana_video_i2v_pre_process_func", "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func", - "MagiHumanPipeline": "get_magi_human_pre_process_func", + "SanaWmPipeline": "get_sana_wm_pre_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func", - "SoulXSingerPipeline": "get_soulxsinger_pre_process_func", - "SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func", } @@ -675,6 +711,7 @@ def register_diffusion_model( def _load_process_func(od_config: OmniDiffusionConfig, func_name: str): """Load and return a process function from the appropriate module.""" + assert od_config.model_class_name is not None mod_folder, mod_relname, _ = _DIFFUSION_MODELS[od_config.model_class_name] if mod_relname == "": # Full module path (registered via register_diffusion_model) @@ -688,6 +725,10 @@ def _load_process_func(od_config: OmniDiffusionConfig, func_name: str): def get_diffusion_post_process_func(od_config: OmniDiffusionConfig): + # Keep the checkpoint's native class name for modality/capability metadata, + # but do not run its tensor postprocessor on Diffusers' decoded outputs. + if uses_diffusers_adapter(od_config): + return None if od_config.model_class_name not in _DIFFUSION_POST_PROCESS_FUNCS: return None func_name = _DIFFUSION_POST_PROCESS_FUNCS[od_config.model_class_name] @@ -702,6 +743,9 @@ def get_diffusion_ir_op_priority_func(od_config: OmniDiffusionConfig): def get_diffusion_pre_process_func(od_config: OmniDiffusionConfig): + # The adapter translates requests to Diffusers call arguments itself. + if uses_diffusers_adapter(od_config): + return None if od_config.model_class_name not in _DIFFUSION_PRE_PROCESS_FUNCS: return None # Return None if no pre-processing function is registered (for backward compatibility) func_name = _DIFFUSION_PRE_PROCESS_FUNCS[od_config.model_class_name] diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 8834e408e..b2eae3fca 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -1,6 +1,6 @@ # adapted from sglang and fastvideo # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -14,6 +14,26 @@ if TYPE_CHECKING: from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id" +_DUMMY_DIFFUSION_REQUEST_ID_PREFIX = f"{DUMMY_DIFFUSION_REQUEST_ID}/" + + +def resolve_video_num_frames( + num_frames: int | None, + *, + default_num_frames: int, + is_dummy_run: bool, +) -> int: + """Resolve the shared image-model frame sentinel for a video pipeline. + + ``OmniDiffusionSamplingParams`` defaults ``num_frames`` to one for image + models, so an omitted video API field reaches model code as ``1``. Video + pipelines with a different default must materialize it at their boundary. + Startup profiling intentionally requests one frame, however, and must stay + lightweight. + """ + if num_frames is None or (num_frames == 1 and not is_dummy_run): + return default_num_frames + return num_frames @dataclass @@ -43,6 +63,15 @@ class OmniDiffusionRequest: # BaseScheduler takes ownership and clears this field before the request is # sent to a Worker. diffusion_kv_requests: tuple[DiffusionKVRequest, ...] | None = None + # Optional model-specific condition structure used by request-batch admission. + # This is populated by a pipeline preprocessor before the request reaches + # the scheduler; ``None`` keeps the default behavior for other pipelines. + batch_compatibility_key: tuple[Any, ...] | None = None + # KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward + # and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric. + kv_recv_ms: float = 0.0 + # Time spent waiting for initial admission by the diffusion scheduler. + scheduler_queue_wait_ms: float | None = None def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" @@ -86,4 +115,6 @@ class OmniDiffusionRequest: @classmethod def is_dummy_run_request_id(cls, request_id: str | None) -> bool: - return request_id == DUMMY_DIFFUSION_REQUEST_ID + return request_id == DUMMY_DIFFUSION_REQUEST_ID or ( + isinstance(request_id, str) and request_id.startswith(_DUMMY_DIFFUSION_REQUEST_ID_PREFIX) + ) diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index b558c1ee6..8fe79a11c 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -3,14 +3,19 @@ from __future__ import annotations +import time from abc import ABC, abstractmethod from collections import deque from dataclasses import fields +from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.v1.kv_cache_interface import KVCacheConfig from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager +from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( CachedRequestData, @@ -49,9 +54,17 @@ class BaseScheduler(ABC): self._finished_req_ids: set[str] = set() self.max_num_running_reqs: int = 1 self._prefetch_enabled: bool = False - self._diffusion_kv_enabled = False + self._diffusion_kv_manager: DiffusionKVCacheManager | None = None - def initialize(self, od_config: OmniDiffusionConfig) -> None: + def initialize( + self, + od_config: OmniDiffusionConfig, + *, + kv_cache_config: KVCacheConfig | None = None, + scheduler_block_size: int | None = None, + hash_block_size: int | None = None, + kv_vllm_config: VllmConfig | None = None, + ) -> None: self.od_config = od_config self._request_states.clear() self._step_id = 0 @@ -66,10 +79,36 @@ class BaseScheduler(ABC): self.max_num_running_reqs = 1 omni_kv = getattr(od_config, "omni_kv_config", None) or {} self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False)) - self._diffusion_kv_enabled = ( + diffusion_kv_enabled = ( getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) is DiffusionKVCacheMode.PAGED_SCHEDULER ) + if diffusion_kv_enabled: + if kv_cache_config is None: + raise ValueError("paged_scheduler Diffusion KV requires a native Scheduler KVCacheConfig") + if scheduler_block_size is None or hash_block_size is None: + raise ValueError("paged_scheduler Diffusion KV requires native scheduler/hash block sizes") + if kv_vllm_config is None: + raise ValueError("paged_scheduler Diffusion KV requires the native VllmConfig used for cache sizing") + self._diffusion_kv_manager = DiffusionKVCacheManager( + kv_cache_config, + max_model_len=kv_vllm_config.model_config.max_model_len, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + max_in_flight_tokens=kv_vllm_config.max_in_flight_tokens, + ) + else: + if any( + value is not None + for value in ( + kv_cache_config, + scheduler_block_size, + hash_block_size, + kv_vllm_config, + ) + ): + raise ValueError("dense_legacy Scheduler received unexpected Diffusion KV cache initialization state") + self._diffusion_kv_manager = None self._reset_scheduler_state() def add_request(self, request: OmniDiffusionRequest) -> str: @@ -79,6 +118,7 @@ class BaseScheduler(ABC): if request_id in self._request_states: raise ValueError(f"request_id {request_id!r} is already active.") state = self._make_request_state(request_id, request) + state.queued_at = time.perf_counter() self._request_states[request_id] = state self._waiting.append(request_id) logger.debug("%s add_request: %s (waiting=%d)", self.__class__.__name__, request_id, len(self._waiting)) @@ -104,6 +144,35 @@ class BaseScheduler(ABC): if not self._can_schedule_waiting(state): break + diffusion_kv_metadata: DiffusionKVMetadata | None = None + if self._diffusion_kv_manager is not None: + if self._diffusion_kv_manager.has_request(request_id): + diffusion_kv_metadata = self._diffusion_kv_manager.get_metadata(request_id) + else: + try: + allocation = self._diffusion_kv_manager.reserve_request( + request_id, + state.diffusion_kv_requests, + ) + except Exception as exc: + # schedule() runs outside the Engine execution-error + # wrapper. Convert both expected admission failures and + # unexpected native allocator failures into a terminal + # request error so the busy loop can wake its stream. + if not isinstance(exc, DiffusionKVAdmissionError): + logger.exception( + "Unexpected Diffusion KV allocation failure for request %s", + request_id, + ) + self._finish_requests( + {request_id: DiffusionRequestStatus.FINISHED_ERROR}, + {request_id: str(exc)}, + ) + continue + if allocation is None: + break + diffusion_kv_metadata = allocation + self._waiting.popleft() was_new_request = state.status == DiffusionRequestStatus.WAITING if not self._running: @@ -111,7 +180,13 @@ class BaseScheduler(ABC): state.status = DiffusionRequestStatus.RUNNING self._running.append(request_id) if was_new_request: - scheduled_new_reqs.append(NewRequestData.from_state(state)) + state.req.scheduler_queue_wait_ms = max((time.perf_counter() - state.queued_at) * 1000.0, 0.0) + scheduled_new_reqs.append( + NewRequestData.from_state( + state, + diffusion_kv_metadata=diffusion_kv_metadata, + ) + ) else: scheduled_cached_request_ids.append(request_id) @@ -184,6 +259,8 @@ class BaseScheduler(ABC): return self._request_states.get(request_id) def pop_request_state(self, request_id: str) -> SchedulerRequestState | None: + if self._diffusion_kv_manager is not None: + self._diffusion_kv_manager.free_request(request_id) self._pop_extra_request_state(request_id) return self._request_states.pop(request_id, None) @@ -206,6 +283,9 @@ class BaseScheduler(ABC): self._finish_requests({request_id: status for request_id in request_ids}) def close(self) -> None: + if self._diffusion_kv_manager is not None: + self._diffusion_kv_manager.close() + self._diffusion_kv_manager = None self._request_states.clear() self._waiting.clear() self._running.clear() @@ -245,6 +325,8 @@ class BaseScheduler(ABC): self._waiting = deque(request_id for request_id in self._waiting if request_id not in waiting_to_remove) for request_id in finished_req_ids: + if self._diffusion_kv_manager is not None: + self._diffusion_kv_manager.free_request(request_id) state = self._request_states[request_id] status = statuses[request_id] state.status = status @@ -266,7 +348,13 @@ class BaseScheduler(ABC): # update_from_output() processes the runner output. It is already # marked finished at that point, but we still need to surface its id # in this update so the engine can observe the terminal state. + # Also surface admission failures recorded while schedule() built this + # output. Older finished ids retained only for Worker cleanup have + # already been popped by the Engine and are deliberately ignored. finished_req_ids = { + request_id for request_id in sched_output.finished_req_ids if request_id in self._request_states + } + finished_req_ids |= { request_id for request_id in sched_output.scheduled_request_ids if request_id in self._finished_req_ids } finished_req_ids |= self._finish_requests(statuses, errors) @@ -280,7 +368,7 @@ class BaseScheduler(ABC): def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> SchedulerRequestState: kv_requests = request.diffusion_kv_requests or () - if self._diffusion_kv_enabled: + if self._diffusion_kv_manager is not None: self._reject_legacy_dense_kv(request) if not kv_requests: raise ValueError("paged_scheduler request preprocessing did not produce DiffusionKVRequest state") diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index 8fb36f405..62e5eddef 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -88,9 +88,10 @@ class RequestBatchSamplingParamsKey: """Request level Batch-compatibility key derived from ``OmniDiffusionSamplingParams``. Only request-batch-wide fields belong here. Request-local values such as - seeds, generators, latent tensors, timesteps, and pipeline-specific + seeds, generators, latent tensors, timesteps, and request-local ``extra_args`` are read per request from - ``DiffusionRequestBatch.sampling_params_list``. + ``DiffusionRequestBatch.sampling_params_list``. Structural ``extra_args`` + that affect batching are listed explicitly below. """ # Spatial / temporal shape. @@ -131,6 +132,16 @@ class RequestBatchSamplingParamsKey: # the model and must remain distinct from explicit ``lossless``. quality: str | None = None + # Wan scheduler structure is carried through extra_args. Requests using + # different solvers or flow shifts must not share a request batch. + sample_solver: str | None = None + flow_shift: float | None = None + + # Pipeline-specific condition structure populated during preprocessing. + # It prevents independently valid requests with incompatible conditions + # from being admitted to the same request batch. + condition_key: tuple[Any, ...] | None = None + # LoRA identity. lora_int_id: int | None = None lora_scale: float = 1.0 @@ -146,6 +157,7 @@ class SchedulerRequestState: diffusion_kv_requests: tuple[DiffusionKVRequest, ...] = () status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING error: str | None = None + queued_at: float = 0.0 def is_finished(self) -> bool: return DiffusionRequestStatus.is_finished(self.status) diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py index 968360183..4168ae974 100644 --- a/vllm_omni/diffusion/sched/request_scheduler.py +++ b/vllm_omni/diffusion/sched/request_scheduler.py @@ -22,7 +22,21 @@ if TYPE_CHECKING: # on sampling params, so it must be resolved separately from the bulk lookup. _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset( field.name for field in fields(RequestBatchSamplingParamsKey) -) - {"lora_int_id"} +) - {"condition_key", "flow_shift", "lora_int_id", "sample_solver"} + + +def _normalize_explicit_sample_solver(value: object | None) -> str | None: + """Normalize an explicitly provided solver without selecting a default.""" + if value is None: + return None + return str(value).strip().lower() + + +def _normalize_explicit_flow_shift(value: object | None) -> float | None: + """Normalize an explicitly provided flow shift without selecting a default.""" + if value is None: + return None + return float(value) def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey: @@ -31,6 +45,13 @@ def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> Re # LoRA identity is optional on sampling params (and on test stubs). lora_request = getattr(sampling, "lora_request", None) key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES} + extra_args = sampling.extra_args or {} + # Match pipeline resolution for explicit overrides, but preserve None: + # pipeline/engine defaults are configuration-dependent and must not be + # inferred while building the request-batch key. + key_kwargs["sample_solver"] = _normalize_explicit_sample_solver(extra_args.get("sample_solver")) + key_kwargs["flow_shift"] = _normalize_explicit_flow_shift(extra_args.get("flow_shift")) + key_kwargs["condition_key"] = getattr(request, "batch_compatibility_key", None) key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None return RequestBatchSamplingParamsKey(**key_kwargs) diff --git a/vllm_omni/diffusion/sched/sigma_schedule.py b/vllm_omni/diffusion/sched/sigma_schedule.py index bea6164d3..7eb3ed38c 100644 --- a/vllm_omni/diffusion/sched/sigma_schedule.py +++ b/vllm_omni/diffusion/sched/sigma_schedule.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -37,8 +37,11 @@ class DMD2SigmaSchedule: # have to be rejected before the monotonicity check rather than after it. if any(not math.isfinite(value) for value in values): raise ValueError("DMD2 base_schedule entries must be finite") - if values[0] != 1.0 or values[-1] != 0.0: - raise ValueError("DMD2 base_schedule must start at 1.0 and end at 0.0") + # A student whose training capped the noise level (``max_timestep_ratio``) + # starts just below 1.0 rather than at it, so the opening position is + # bounded rather than pinned; the terminal one is always clean latents. + if not 0.0 < values[0] <= 1.0 or values[-1] != 0.0: + raise ValueError("DMD2 base_schedule must start in (0.0, 1.0] and end at 0.0") if any(curr <= nxt for curr, nxt in zip(values, values[1:], strict=False)): raise ValueError("DMD2 base_schedule must be strictly decreasing") diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index a8733b8d0..fe7b6c408 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Stage Diffusion Client for vLLM-Omni multi-stage runtime. Owns the frontend-side ZMQ sockets for StageDiffusionProc and exposes the @@ -30,6 +33,7 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import ( ) from vllm_omni.engine.stage_client import StageClientBase from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -45,14 +49,13 @@ def create_diffusion_client( od_config: OmniDiffusionConfig, metadata: StageMetadata, stage_init_timeout: int, - batch_size: int = 1, use_inline: bool = False, ) -> Any: """Factory to create either an inline or out-of-process diffusion client.""" if use_inline: from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient - return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size) + return InlineStageDiffusionClient(model, od_config, metadata) proc_manager = StageDiffusionProcManager( model=model, od_config=od_config, @@ -63,7 +66,6 @@ def create_diffusion_client( request_address=proc_manager.addresses.inputs[0], response_address=proc_manager.addresses.outputs[0], proc_manager=proc_manager, - batch_size=batch_size, ) @@ -87,14 +89,12 @@ class StageDiffusionClient(StageClientBase): response_address: str, *, proc_manager: StageDiffusionProcManager | None = None, - batch_size: int = 1, ) -> None: self._initialize_client( metadata, request_address, response_address, proc_manager=proc_manager, - batch_size=batch_size, ) @classmethod @@ -105,7 +105,6 @@ class StageDiffusionClient(StageClientBase): response_address: str, *, proc_manager: StageDiffusionProcManager | None = None, - batch_size: int = 1, ) -> StageDiffusionClient: """Create a client for an already-running diffusion subprocess.""" return cls( @@ -113,7 +112,6 @@ class StageDiffusionClient(StageClientBase): request_address, response_address, proc_manager=proc_manager, - batch_size=batch_size, ) def _initialize_client( @@ -123,7 +121,6 @@ class StageDiffusionClient(StageClientBase): response_address: str, *, proc_manager: StageDiffusionProcManager | None = None, - batch_size: int, ) -> None: self._set_stage_metadata(metadata) self._proc_manager = proc_manager @@ -140,11 +137,10 @@ class StageDiffusionClient(StageClientBase): self._start_proc_monitor() logger.info( - "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)", + "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s)", self.stage_id, self.replica_id, self._proc_manager is not None, - batch_size, ) def _set_stage_metadata(self, metadata: StageMetadata) -> None: @@ -154,6 +150,7 @@ class StageDiffusionClient(StageClientBase): self.final_output_type = metadata.final_output_type self.model_stage = metadata.model_stage self.default_sampling_params = metadata.default_sampling_params + self.prompt_transform_func = metadata.prompt_transform_func self.prompt_expand_func = metadata.prompt_expand_func self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False) self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None) @@ -240,6 +237,14 @@ class StageDiffusionClient(StageClientBase): if msg_type == "result": self._output_queue.put_nowait(msg["output"]) + elif msg_type == "metrics": + self._output_queue.put_nowait( + OmniRequestOutput( + request_id=DIFFUSION_METRICS_ONLY_REQUEST_ID, + finished=True, + metrics=msg.get("metrics") or {}, + ) + ) elif msg_type == "rpc_result": self._rpc_results[msg["rpc_id"]] = msg["result"] elif msg_type == "error": @@ -264,14 +269,14 @@ class StageDiffusionClient(StageClientBase): # Route request errors as error outputs so the Orchestrator # sees the request complete (instead of hanging forever). if req_id is not None: - self._output_queue.put_nowait( - OmniRequestOutput.from_error( - req_id, - error_msg, - status_code=status_code, - error_type=error_type, - ) + error_output = OmniRequestOutput.from_error( + req_id, + error_msg, + status_code=status_code, + error_type=error_type, ) + error_output.metrics.update(msg.get("metrics") or {}) + self._output_queue.put_nowait(error_output) # Fields that are subprocess-local and cannot be serialized across # process boundaries. They are recreated in the subprocess with diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index f01348c74..dad453faa 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Subprocess entry point for the diffusion engine. StageDiffusionProc runs DiffusionEngine in a child process, @@ -22,7 +25,12 @@ from vllm.logger import init_logger from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx from vllm.utils.system_utils import get_mp_context from vllm.v1.engine.core import EngineCoreProc -from vllm.v1.engine.utils import CoreEngine, EngineZmqAddresses, wait_for_engine_startup +from vllm.v1.engine.utils import ( + CoreEngine, + CoreEngineLaunch, + EngineZmqAddresses, + wait_for_engine_startup, +) from vllm.v1.utils import shutdown from vllm_omni.diffusion.data import DiffusionRequestAbortedError @@ -36,6 +44,7 @@ from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage from vllm_omni.engine.stage_init_utils import set_death_signal from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.metrics.utils import diffusion_exception_metrics from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -361,6 +370,16 @@ class StageDiffusionProc: request_id, str(e), ) + metrics = diffusion_exception_metrics(e) + if metrics: + await response_socket.send( + encoder.encode( + { + "type": "metrics", + "metrics": metrics, + } + ) + ) except Exception as e: logger.exception("Diffusion request %s failed: %s", request_id, e) status_code, error_type = client_error_metadata(e) @@ -372,6 +391,7 @@ class StageDiffusionProc: "error": str(e), "status_code": status_code, "error_type": error_type, + "metrics": diffusion_exception_metrics(e), } ) ) @@ -426,9 +446,8 @@ class StageDiffusionProc: elif msg_type == "abort": for rid in msg.get("request_ids", []): - task = tasks.pop(rid, None) - if task: - task.cancel() + # Let the request task consume the terminal abort + # output so it can publish the scheduler snapshot. self._engine.abort(rid) elif msg_type == "collective_rpc": @@ -579,6 +598,8 @@ class StageDiffusionProc: - ``omni_replica_id``: cluster-unique replica id within the stage (logging / metrics only). """ + from vllm_omni.plugins import load_omni_general_plugins + shutdown_requested = False set_death_signal(signal.SIGTERM) @@ -592,6 +613,11 @@ class StageDiffusionProc: signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) + # ``spawn`` starts this process with a fresh interpreter, so plugin + # side effects (for example, custom diffusion loader hooks) must be + # applied again before the engine is constructed. + load_omni_general_plugins() + proc = cls(model, od_config) coord_client: OmniCoordClientForStage | None = None handshake_ctx: zmq.Context | None = None @@ -750,7 +776,6 @@ class StageDiffusionProcManager: with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: wait_for_engine_startup( handshake_socket, - self.addresses, [CoreEngine(index=0, local=True)], SimpleNamespace( data_parallel_size_local=1, @@ -759,8 +784,17 @@ class StageDiffusionProcManager: ), False, None, - self, - None, + CoreEngineLaunch( + engine_manager=self, + coordinator=None, + addresses=self.addresses, + tensor_queue=None, + # StageDiffusionProcManager is not a CoreEngineProcManager, + # so upstream's isinstance-gated sentinel registration would + # be skipped; watch the subprocess directly so a proc death + # during handshake is still detected. + watched_frontend_processes=[self.proc], + ), ) except Exception: shutdown([self.proc]) diff --git a/vllm_omni/diffusion/utils/future_utils.py b/vllm_omni/diffusion/utils/future_utils.py new file mode 100644 index 000000000..5492c6a21 --- /dev/null +++ b/vllm_omni/diffusion/utils/future_utils.py @@ -0,0 +1,27 @@ +"""Helpers for resolving ``concurrent.futures.Future`` objects shared across threads.""" + +import concurrent.futures +from typing import Any, TypeVar + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_ResultT = TypeVar("_ResultT") + + +def try_set_result(fut: concurrent.futures.Future[_ResultT], result: _ResultT) -> None: + # fut may be cancelled concurrently (e.g. asyncio.wait_for timeout) between + # the caller's fut.done() check and this call; drop the late result instead + # of crashing the pump thread. + try: + fut.set_result(result) + except concurrent.futures.InvalidStateError: + logger.debug("Dropping late result for already-resolved/cancelled future") + + +def try_set_exception(fut: concurrent.futures.Future[Any], exc: BaseException) -> None: + try: + fut.set_exception(exc) + except concurrent.futures.InvalidStateError: + logger.debug("Dropping late exception for already-resolved/cancelled future") diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py index 78d8c6536..5c41685f0 100644 --- a/vllm_omni/diffusion/utils/hf_utils.py +++ b/vllm_omni/diffusion/utils/hf_utils.py @@ -71,6 +71,30 @@ def _looks_like_dreamzero(model_name: str) -> bool: return False +HIDREAM_O1_SIGNATURE_WEIGHTS = ( + "model.final_layer2.linear.weight", + "final_layer2.linear.weight", +) + + +def _looks_like_hidream_o1(model_name: str, config: Mapping | None = None) -> bool: + """Detect HiDream-O1 without matching regular Qwen3-VL checkpoints.""" + try: + cfg = config if config is not None else get_hf_file_to_dict("config.json", model_name) + if not isinstance(cfg, Mapping) or cfg.get("model_type") != "qwen3_vl": + return False + + index = get_hf_file_to_dict("model.safetensors.index.json", model_name) + if not isinstance(index, Mapping): + return False + weight_map = index.get("weight_map") + if not isinstance(weight_map, Mapping): + return False + return any(key in weight_map for key in HIDREAM_O1_SIGNATURE_WEIGHTS) + except (AttributeError, OSError, TypeError, ValueError): + return False + + @lru_cache def is_diffusion_model(model_name: str) -> bool: """Check if a model is a diffusion model. @@ -115,7 +139,4 @@ def is_diffusion_model(model_name: str) -> bool: except Exception as e: logger.debug("Failed to load diffusers config via DiffusionPipeline: %s", e) - # Bagel and DreamZero are not diffusers pipelines (no model_index.json), - # but are still diffusion-style models in vllm-omni. Detect them via - # config.json. - return _looks_like_bagel(model_name) or _looks_like_dreamzero(model_name) + return _looks_like_bagel(model_name) or _looks_like_dreamzero(model_name) or _looks_like_hidream_o1(model_name) diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index 1ace493db..df67f9d99 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -5,6 +5,7 @@ from __future__ import annotations import io +from collections.abc import Iterable from fractions import Fraction from typing import Any, cast @@ -200,3 +201,71 @@ def mux_video_audio_bytes( container.close() return buf.getvalue() + + +def mux_av_video_audio_bytes( + video_frames: Iterable[av.VideoFrame], + width: int, + height: int, + audio_waveform: np.ndarray | None = None, + *, + fps: float = 25.0, + audio_sample_rate: int | None = None, + video_codec: str = "h264", + audio_codec: str = "aac", + crf: str = "18", + video_codec_options: dict[str, str] | None = None, +) -> bytes: + """Mux preconstructed video frames and optional audio into MP4 bytes.""" + buf = io.BytesIO() + with cast(Any, av.open(buf, mode="w", format="mp4")) as container: + v_stream = cast( + av.VideoStream, + container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)), + ) + v_stream.width = width + v_stream.height = height + v_stream.pix_fmt = "yuv420p" + + options: dict[str, object] = {"crf": str(crf)} + if video_codec_options: + options.update(video_codec_options) + v_stream.options = options + + a_stream: av.AudioStream | None = None + samples: np.ndarray | None = None + layout: str | None = None + if audio_waveform is not None: + effective_audio_sample_rate = 44100 if audio_sample_rate is None else audio_sample_rate + samples = audio_waveform.astype(np.float32) + if samples.ndim == 1: + samples = samples.reshape(1, -1) + elif samples.ndim == 2 and samples.shape[0] > samples.shape[1]: + samples = np.ascontiguousarray(samples.T) + num_channels = samples.shape[0] + layout = "stereo" if num_channels >= 2 else "mono" + a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=effective_audio_sample_rate)) + a_stream.layout = layout + + for frame in video_frames: + for packet in v_stream.encode(frame): + container.mux(packet) + for packet in v_stream.encode(): + container.mux(packet) + + if a_stream is not None and audio_waveform is not None: + if samples is None or layout is None: + raise ValueError("Audio samples were not prepared for muxing.") + audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout) + audio_frame.sample_rate = effective_audio_sample_rate + # AAC has a one-frame encoder delay. Mark the input waveform as + # starting at t=0 so the MP4 muxer writes the corresponding negative + # priming timestamp instead of exposing the delay as leading silence. + audio_frame.pts = 0 + audio_frame.time_base = Fraction(1, effective_audio_sample_rate) + for packet in a_stream.encode(audio_frame): + container.mux(packet) + for packet in a_stream.encode(): + container.mux(packet) + + return buf.getvalue() diff --git a/vllm_omni/diffusion/vllm_config.py b/vllm_omni/diffusion/vllm_config.py new file mode 100644 index 000000000..ac591d687 --- /dev/null +++ b/vllm_omni/diffusion/vllm_config.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Build native vLLM configuration views for the Omni diffusion runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +from vllm.config import CompilationConfig, DeviceConfig, VllmConfig +from vllm.transformers_utils.config import get_hf_text_config + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.platforms import current_omni_platform + + +def resolve_diffusion_max_model_len(od_config: OmniDiffusionConfig) -> int: + """Resolve the native KV admission bound for a diffusion stage.""" + + configured = getattr(od_config, "max_model_len", None) + if configured is not None: + configured = int(configured) + if configured == -1: + configured = None + elif configured <= 0: + raise ValueError(f"max_model_len must be positive for Scheduler-managed Diffusion KV, got {configured}") + else: + return configured + + # ``max_model_len`` keeps the native vLLM meaning: the per-sequence KV + # admission ceiling consumed by ``KVCacheManager``. HunyuanImage3 exposes + # that ceiling directly as ``max_position_embeddings`` because its text, + # reference-image, and generated-image tokens share one Transformer + # sequence. Other diffusion models commonly express their serving limits + # through resolution, duration, or frame-count constraints instead. When + # paged KV support expands to those models, the model integration must + # derive ``max_model_len`` from its serving limits, while model-owned + # preprocessing validates each request and derives its matching + # ``seq_len``. Text-encoder length fields must not be mistaken for the DiT + # KV limit. + hf_config = getattr(od_config, "tf_model_config", None) + for field_name in ("max_position_embeddings", "max_sequence_length", "seq_length"): + value = getattr(hf_config, field_name, None) + if value is None and isinstance(hf_config, dict): + value = hf_config.get(field_name) + if value is not None and int(value) > 0: + return int(value) + + mode = getattr(getattr(od_config, "diffusion_kv_mode", None), "value", None) + if mode == "paged_scheduler": + raise ValueError( + "paged_scheduler Diffusion KV requires max_model_len or a model config " + "with max_position_embeddings/max_sequence_length/seq_length" + ) + # Dense execution does not consume the native KV admission bound, but the + # shared Worker VllmConfig still requires a positive placeholder. + return 1 + + +@dataclass +class _DiffusionVllmModelConfig: + """Private adapter for the native ``VllmConfig.model_config`` surface.""" + + model: str | None + dtype: torch.dtype + max_model_len: int + original_max_model_len: int | None + runner_type: str = "generate" + quantization: str | None = None + quantization_config: Any | None = None + hf_config: Any | None = None + hf_text_config: Any | None = None + multimodal_config: Any | None = None + enforce_eager: bool = False + disable_cascade_attn: bool = False + enable_return_routed_experts: bool = False + enable_prompt_embeds: bool = False + logits_processors: tuple[str, ...] = () + use_mla: bool = False + is_moe: bool = False + # Native vLLM attention layers/builders read these fields directly. The + # diffusion model does not have a full vLLM ModelConfig, so keep the + # narrow compatibility surface here instead of manufacturing a second + # model configuration object in the Worker backend. + is_mm_prefix_lm: bool = False + rswa_window: int | None = None + _attention_num_heads: int | None = field(default=None, init=False, repr=False) + _attention_num_kv_heads: int | None = field(default=None, init=False, repr=False) + _attention_head_size: int | None = field(default=None, init=False, repr=False) + + # Needed for models that bundle things like LogitsProcessors, e.g., SenseNova + head_dtype: torch.dtype | None = None + + @property + def is_quantized(self) -> bool: + return self.quantization is not None + + def is_nvfp4_quantized(self) -> bool: + return self.quantization == "modelopt_fp4" + + def set_attention_geometry( + self, + *, + num_heads: int, + num_kv_heads: int, + head_size: int, + ) -> None: + """Record rank-local geometry needed by native metadata builders. + + Diffusion attention modules are loaded outside vLLM's model loader, so + their dimensions are not available when this lightweight config is + created. The paged Worker backend fills them in after model loading + and before calling ``init_attn_backend``. + """ + + values = { + "num_heads": num_heads, + "num_kv_heads": num_kv_heads, + "head_size": head_size, + } + invalid = {name: value for name, value in values.items() if type(value) is not int or value <= 0} + if invalid: + raise ValueError(f"Diffusion attention geometry must contain positive integers: {invalid!r}") + self._attention_num_heads = num_heads + self._attention_num_kv_heads = num_kv_heads + self._attention_head_size = head_size + + def _require_attention_geometry(self) -> tuple[int, int, int]: + geometry = ( + self._attention_num_heads, + self._attention_num_kv_heads, + self._attention_head_size, + ) + if any(value is None for value in geometry): + raise RuntimeError("Diffusion attention geometry is unavailable before paged cache layer registration") + num_heads, num_kv_heads, head_size = geometry + assert num_heads is not None and num_kv_heads is not None and head_size is not None + return num_heads, num_kv_heads, head_size + + def get_num_attention_heads(self, parallel_config: Any) -> int: + del parallel_config + return self._require_attention_geometry()[0] + + def get_num_kv_heads(self, parallel_config: Any) -> int: + del parallel_config + return self._require_attention_geometry()[1] + + def get_head_size(self) -> int: + return self._require_attention_geometry()[2] + + @property + def is_diffusion(self) -> bool: + # This config is also installed as the current VllmConfig for the + # existing Omni diffusion runtime, which does not use vLLM's native + # diffusion runner. Preserve that established backend-selection flag. + return False + + +def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig: + quant_config = getattr(od_config, "quantization_config", None) + quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None + hf_config = getattr(od_config, "tf_model_config", None) + hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config + configured_max_model_len = getattr(od_config, "max_model_len", None) + max_model_len = resolve_diffusion_max_model_len(od_config) + paged_scheduler = ( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ) + return _DiffusionVllmModelConfig( + model=od_config.model, + dtype=od_config.dtype, + max_model_len=max_model_len, + # Preserve the user's -1 sentinel only for the native paged-cache + # sizing flow. Dense keeps its established resolved model-config view. + original_max_model_len=configured_max_model_len if paged_scheduler else max_model_len, + quantization=quantization, + quantization_config=quant_config, + hf_config=hf_config, + hf_text_config=hf_text_config, + enforce_eager=getattr(od_config, "enforce_eager", False), + is_moe=bool(getattr(od_config, "is_moe", False)), + ) + + +def create_base_diffusion_vllm_config( + device: torch.device, + od_config: OmniDiffusionConfig, +) -> VllmConfig: + """Create the native vLLM 0.27 config used by diffusion.""" + + return VllmConfig( + compilation_config=CompilationConfig(), + device_config=DeviceConfig(device=device), + additional_config=od_config.additional_config, + ) + + +def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiffusionConfig) -> VllmConfig: + """Populate native model, parallel, cache, and Scheduler configuration.""" + + parallel_config = od_config.parallel_config + vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size + vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size + if parallel_config.enable_expert_parallel and od_config.is_moe: + vllm_config.parallel_config.data_parallel_size = ( + parallel_config.data_parallel_size * parallel_config.cfg_parallel_size + ) + vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size + vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel + + vllm_config.model_config = _make_diffusion_vllm_model_config(od_config) # type: ignore[assignment] + vllm_config.quant_config = od_config.quantization_config + vllm_config.profiler_config = od_config.profiler_config + if ( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ): + vllm_config.parallel_config.pipeline_parallel_size = parallel_config.pipeline_parallel_size + vllm_config.cache_config.gpu_memory_utilization = float(getattr(od_config, "gpu_memory_utilization", 0.9)) + vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None) + # Prefix hashes/publication are intentionally deferred to the next PR. + vllm_config.cache_config.enable_prefix_caching = False + vllm_config.scheduler_config.max_num_seqs = int(od_config.max_num_seqs) + max_num_batched_tokens = getattr(od_config, "max_num_batched_tokens", None) + if max_num_batched_tokens is not None: + vllm_config.scheduler_config.max_num_batched_tokens = int(max_num_batched_tokens) + current_omni_platform.configure_diffusion_vllm_config(vllm_config, od_config) + return vllm_config + + +def create_diffusion_vllm_config( + device: torch.device, + od_config: OmniDiffusionConfig, +) -> VllmConfig: + """Build the same fully populated native config in Engine and Workers.""" + + vllm_config = create_base_diffusion_vllm_config( + device, + od_config, + ) + return configure_diffusion_vllm_config(vllm_config, od_config) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index d118a9833..ea0dff90a 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Diffusion Model Runner for vLLM-Omni. @@ -11,6 +11,7 @@ model-related operations. from __future__ import annotations import copy +import gc import time from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext @@ -21,7 +22,9 @@ from torch.profiler import record_function from vllm.config import LoadConfig, VllmConfig from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec +from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import cache_summary from vllm_omni.diffusion.cache.prompt_embed_cache import ( install_prompt_embed_cache, @@ -32,6 +35,12 @@ from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata +from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend +from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( + DiffusionPagedAttentionMetadata, + DiffusionPagedAttentionRow, +) +from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_rank from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import ( @@ -41,12 +50,14 @@ from vllm_omni.diffusion.models.interface import ( supports_prompt_update, supports_step_execution, ) -from vllm_omni.diffusion.offloader import get_offload_backend +from vllm_omni.diffusion.offloader import enable_offload_backend from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( + CachedRequestData, DiffusionSchedulerOutput, KVPrefetchJob, + NewRequestData, validate_new_request_data_identity, ) from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents @@ -71,6 +82,31 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _dit_any_rank_failed(local_failed: bool) -> bool: + """All-reduce a per-request failure flag across the DiT process group. + + Every DiT rank must reach this point in lockstep; the caller wraps + ``pipeline.prepare_encode`` so both the success and the exception paths + call the helper. In single-rank execution this collapses to the local + flag with no collectives issued. + """ + if not torch.distributed.is_initialized(): + return local_failed + try: + from vllm_omni.diffusion.distributed.parallel_state import get_dit_group + + group = get_dit_group() + except (AssertionError, ImportError): + group = None + if group is None: + return local_failed + signal = torch.tensor(1 if local_failed else 0, dtype=torch.int32) + if current_omni_platform.is_available(): + signal = signal.to(device=current_omni_platform.device_type) + torch.distributed.all_reduce(signal, op=torch.distributed.ReduceOp.MAX, group=group) + return bool(signal.item()) + + def _normalize_pipeline_outputs( outputs: object, *, @@ -139,6 +175,15 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): self.offload_backend: Any | None = None self.prompt_embed_cache: Any | None = None self.input_batch: InputBatch | None = None + self.model_memory_usage = 0 + self.diffusion_kv_backend = DiffusionKVModelRunnerBackend( + vllm_config=vllm_config, + od_config=od_config, + device=device, + ) + # Compatibility view for callers that inspect the installed config; + # physical ownership remains in ``diffusion_kv_backend``. + self.kv_cache_config: KVCacheConfig | None = None # Cache for per-request stepwise state. self.state_cache: dict[str, StepRequestState] = {} @@ -167,8 +212,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): metadata: DiffusionKVMetadata | None, ) -> None: cache_mode = getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) - # W0 does not wire the Scheduler allocator or Worker installer yet, so - # validate an allocation snapshot when present without requiring one. + if cache_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is None: + raise ValueError(f"paged_scheduler request {request_id!r} requires Diffusion KV metadata") if cache_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER and metadata is not None: raise ValueError(f"{cache_mode.value} request {request_id!r} must not carry Diffusion KV metadata") @@ -273,6 +318,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): m.consumed_memory / GiB_bytes, time_after_load - time_before_load, ) + self.model_memory_usage = int(m.consumed_memory) logger.info("Model runner: Model loaded successfully.") if self.od_config.streaming_output and not getattr(self.od_config, "step_execution", False): @@ -291,11 +337,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): f"{self.od_config.model_class_name} does not support that contract." ) - # Apply CPU offloading - self.offload_backend = get_offload_backend(self.od_config, device=self.device) - if self.offload_backend is not None: - logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}") - self.offload_backend.enable(self.pipeline) + # The offloader owns loader-plan handoff and startup recovery. The + # runner only receives the pipeline/backend pair that is ready to use. + self.pipeline, self.offload_backend = enable_offload_backend( + self.od_config, + self.pipeline, + device=self.device, + ) # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: @@ -364,6 +412,105 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): logger.info("Model runner: Initialization complete.") + def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: + """Collect native specs from cache-enabled loaded attention modules.""" + + if self.od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER: + return {} + if self.pipeline is None: + raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs") + + cache_layers: dict[str, tuple[Attention, KVCacheSpec]] = {} + for module_path, module in self.pipeline.named_modules(): + if not isinstance(module, Attention): + continue + spec = module.get_kv_cache_spec(self.vllm_config) + if spec is not None: + layer_name = module.prefix + if not isinstance(layer_name, str) or not layer_name: + raise RuntimeError( + "Paged Diffusion Attention must expose a non-empty canonical prefix; " + f"module_path={module_path!r}" + ) + if layer_name in cache_layers: + raise RuntimeError(f"Duplicate canonical paged Diffusion Attention prefix {layer_name!r}") + cache_layers[layer_name] = (module, spec) + if not cache_layers: + raise RuntimeError( + "paged_scheduler Diffusion KV found no cache-enabled Attention modules " + f"in {type(self.pipeline).__name__}" + ) + return self.diffusion_kv_backend.register_kv_cache_layers(cache_layers) + + def set_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None: + """Physically initialize the Engine-generated rank-local config.""" + + self.diffusion_kv_backend.initialize_kv_cache(kv_cache_config) + self.kv_cache_config = self.diffusion_kv_backend.kv_cache_config + + def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool: + return self.diffusion_kv_backend.install_diffusion_kv_metadata(metadata) + + def get_diffusion_kv_row( + self, + request_id: str, + sequence_id: int | None, + context_id: str | None = None, + ) -> int: + return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id) + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int: + return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids) + + def refresh_diffusion_kv_block_table_layout(self) -> None: + self.diffusion_kv_backend.refresh_block_table_layout() + + def _build_paged_attention_metadata( + self, + metadata: list[DiffusionKVMetadata], + ) -> DiffusionPagedAttentionMetadata: + """Translate Scheduler sequence state into ordered request-level rows.""" + + cfg_size = int(getattr(self.od_config.parallel_config, "cfg_parallel_size", 1) or 1) + cfg_rank = get_classifier_free_guidance_rank() if cfg_size > 1 else None + prefill_rows: list[DiffusionPagedAttentionRow] = [] + denoise_rows: list[DiffusionPagedAttentionRow] = [] + for request_metadata in metadata: + sequences = request_metadata.sequences + if cfg_rank is not None and len(sequences) > 1: + if len(sequences) != cfg_size: + raise ValueError( + "Paged CFG parallel execution requires one Scheduler sequence per CFG rank: " + f"request={request_metadata.request_id!r}, cfg_size={cfg_size}, rows={len(sequences)}" + ) + sequences = (sequences[cfg_rank],) + for sequence in sequences: + active_seq_len = sequence.prefix_len + sequence.target_len + if active_seq_len > sequence.seq_len: + raise ValueError( + "Paged denoise span exceeds its Scheduler allocation: " + f"request={request_metadata.request_id!r}, sequence={sequence.sequence_id}, " + f"active={active_seq_len}, allocated={sequence.seq_len}" + ) + prefill_rows.append( + DiffusionPagedAttentionRow( + request_id=request_metadata.request_id, + sequence_id=sequence.sequence_id, + query_len=sequence.seq_len, + seq_len=sequence.seq_len, + ) + ) + denoise_rows.append( + DiffusionPagedAttentionRow( + request_id=request_metadata.request_id, + sequence_id=sequence.sequence_id, + query_len=sequence.target_len, + seq_len=active_seq_len, + kv_start_pos=sequence.prefix_len, + ) + ) + return DiffusionPagedAttentionMetadata(tuple(prefill_rows), tuple(denoise_rows)) + def clear_prompt_embed_cache(self) -> None: """Evict all cached text-encoder outputs (e.g. between training epochs). @@ -434,6 +581,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): target_device=self._target_device if use_prefetch else getattr(self.pipeline, "device", None), ) kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000 + req.kv_recv_ms = kv_recv_ms logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms) # Kick off the next request's prefetch (+ H2D) to overlap this forward. @@ -491,6 +639,10 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): reqs: list[OmniDiffusionRequest], outputs: list[DiffusionOutput], ) -> BatchRunnerOutput: + for i in range(len(reqs)): + # Carry the runner-measured KV-recv timing onto the output so the + # engine's step_streaming can surface it as diffusion_kv_load_s. + outputs[i].kv_recv_ms = reqs[i].kv_recv_ms return BatchRunnerOutput.from_list( [ RunnerOutput( @@ -512,6 +664,9 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): require_request_batch_support: bool, kv_prefetch_job: KVPrefetchJob | None = None, record_name: str, + record_output_peak_memory: bool = True, + in_diffusion_kv_memory_profile: bool = False, + diffusion_kv_metadata: list[DiffusionKVMetadata] | None = None, ) -> BatchRunnerOutput: assert self.pipeline is not None, "Model not loaded. Call load_model() first." if not reqs: @@ -541,10 +696,30 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): batch = DiffusionRequestBatch(requests=reqs) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 - if is_primary: + if is_primary and record_output_peak_memory: current_omni_platform.reset_peak_memory_stats() - with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config): + paged_kv_runtime = None + paged_kv_context: AbstractContextManager[Any] = nullcontext() + if diffusion_kv_metadata is not None: + if len(diffusion_kv_metadata) != len(reqs): + raise ValueError( + "Diffusion KV metadata count must match the request batch: " + f"metadata={len(diffusion_kv_metadata)}, requests={len(reqs)}" + ) + paged_metadata = self._build_paged_attention_metadata(diffusion_kv_metadata) + paged_kv_runtime, paged_kv_context = self.diffusion_kv_backend.activate_paged_attention_metadata( + paged_metadata + ) + with ( + set_forward_context( + vllm_config=self.vllm_config, + omni_diffusion_config=od_config, + paged_kv_runtime=paged_kv_runtime, + in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, + ), + paged_kv_context, + ): with record_function(record_name): raw_outputs = self.pipeline.forward(batch) outputs = _normalize_pipeline_outputs( @@ -554,7 +729,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): pipeline_name=type(self.pipeline).__name__, ) - if is_primary and outputs: + if is_primary and outputs and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() for output in outputs: output.peak_memory_mb = max(output.peak_memory_mb, batch_peak_memory_mb) @@ -610,17 +785,81 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): request_id=req.request_id, metadata=diffusion_kv_metadata, ) - runner_output = self._execute_request_list( - [req], - od_config=self.od_config, - allow_single_output=True, - require_request_batch_support=False, - kv_prefetch_job=kv_prefetch_job, - record_name="pipeline_forward", - ) - output = runner_output.runner_outputs[0].result - assert output is not None - return output + installed_request = False + if diffusion_kv_metadata is not None: + self.install_diffusion_kv_metadata(diffusion_kv_metadata) + installed_request = True + try: + runner_output = self._execute_request_list( + [req], + od_config=self.od_config, + allow_single_output=True, + require_request_batch_support=False, + kv_prefetch_job=kv_prefetch_job, + record_name="pipeline_forward", + diffusion_kv_metadata=[diffusion_kv_metadata] if diffusion_kv_metadata is not None else None, + ) + output = runner_output.runner_outputs[0].result + assert output is not None + return output + finally: + if installed_request: + self.remove_diffusion_kv_requests([req.request_id]) + + def profile_run(self, requests: list[OmniDiffusionRequest]) -> None: + """Run the maximum per-rank request batch for memory profiling. + + This deliberately bypasses Scheduler admission and Diffusion KV + metadata validation because cache capacity has not been sized yet. + It otherwise uses the normal execution-mode path so model inputs, + collective communication, backend workspaces, denoising, and decode + allocations contribute to the observed peak. Step execution profiles + one fused ``InputBatch`` instead of sequential single-request forwards. + """ + + if not requests: + raise ValueError("Diffusion memory profiling requires at least one request.") + + runner_output: BatchRunnerOutput | None = None + request_ids = [request.request_id for request in requests] + try: + if getattr(self.od_config, "step_execution", False): + scheduler_output = DiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=[ + NewRequestData(request_id=request.request_id, req=request) for request in requests + ], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=len(requests), + num_waiting_reqs=0, + ) + runner_output = self._execute_stepwise( + scheduler_output, + validate_kv_metadata=False, + record_output_peak_memory=False, + in_diffusion_kv_memory_profile=True, + ) + else: + runner_output = self._execute_request_list( + requests, + od_config=self.od_config, + allow_single_output=len(requests) == 1, + require_request_batch_support=len(requests) > 1, + record_name="pipeline_memory_profile", + # The enclosing native memory_profiling context owns the + # peak counters. Resetting them here would discard request + # preparation allocations and understate the budget. + record_output_peak_memory=False, + in_diffusion_kv_memory_profile=True, + ) + current_omni_platform.synchronize() + finally: + for request_id in request_ids: + self.state_cache.pop(request_id, None) + self.input_batch = None + del runner_output + gc.collect() def execute_model_batch( self, @@ -639,14 +878,29 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): request_id=new_req.req.request_id, metadata=new_req.diffusion_kv_metadata, ) - reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs] - return self._execute_request_list( - reqs, - od_config=od_config, - allow_single_output=False, - require_request_batch_support=True, - record_name="pipeline_forward_batch", - ) + installed_request_ids: list[str] = [] + try: + for new_req in scheduler_output.scheduled_new_reqs: + if new_req.diffusion_kv_metadata is not None: + self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata) + installed_request_ids.append(new_req.request_id) + reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs] + return self._execute_request_list( + reqs, + od_config=od_config, + allow_single_output=False, + require_request_batch_support=True, + record_name="pipeline_forward_batch", + diffusion_kv_metadata=[ + new_req.diffusion_kv_metadata + for new_req in scheduler_output.scheduled_new_reqs + if new_req.diffusion_kv_metadata is not None + ] + or None, + ) + finally: + if installed_request_ids: + self.remove_diffusion_kv_requests(installed_request_ids) # ------------------------------------------------------------------ # Step-wise execution @@ -700,25 +954,70 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): return resolved, new_request_ids - def _prepare_batch_inputs(self, states: list[StepRequestState], new_request_ids: list[str]) -> InputBatch: + def _prepare_batch_inputs( + self, + states: list[StepRequestState], + new_request_ids: list[str], + ) -> tuple[list[StepRequestState], InputBatch | None, list[RunnerOutput]]: # process new reqs + prepared_states: list[StepRequestState] = [] + error_outputs: list[RunnerOutput] = [] for state in states: if state.request_id in new_request_ids: - self._initialize_generator(state.sampling) - clear_pipeline_stage_durations(self.pipeline) - # encode - self.pipeline.prepare_encode(state) - merge_stage_durations( - state, - consume_pipeline_stage_durations(self.pipeline), - ) + # Everything that runs before ``_dit_any_rank_failed`` must be + # inside the try: an exception in ``_initialize_generator`` or + # ``clear_pipeline_stage_durations`` on one rank would skip the + # all-reduce here while every peer proceeds into it, and the + # peers then hang on the NCCL collective until timeout. + per_req_exc: BaseException | None = None + try: + self._initialize_generator(state.sampling) + clear_pipeline_stage_durations(self.pipeline) + # encode + self.pipeline.prepare_encode(state) + merge_stage_durations( + state, + consume_pipeline_stage_durations(self.pipeline), + ) + except Exception as exc: + per_req_exc = exc + # Pipelines that do rank-0-only work (e.g. MiniMax H3 + # reference-video prep) must broadcast per-request failures + # internally so downstream collectives stay in step; even so, + # cross-check that every DiT rank agrees so a rank-local error + # (or a future pipeline that omits the guard) does not leave + # the process group half-way through a new request. + if _dit_any_rank_failed(per_req_exc is not None): + self.state_cache.pop(state.request_id, None) + if per_req_exc is None: + per_req_exc = RuntimeError( + f"Stepwise preparation failed on another DiT rank for {state.request_id}" + ) + logger.error( + "Stepwise request preparation failed for %s: %s", + state.request_id, + per_req_exc, + exc_info=isinstance(per_req_exc, Exception), + ) + error_outputs.append( + RunnerOutput( + request_id=state.request_id, + step_index=state.step_index, + finished=True, + result=DiffusionOutput.from_exception(per_req_exc), + ) + ) + continue + prepared_states.append(state) + if not prepared_states: + return prepared_states, None, error_outputs input_batch = InputBatch.make_batch( - states, + prepared_states, cached_batch=getattr(self, "input_batch", None), ) self.input_batch = input_batch - return input_batch + return prepared_states, input_batch, error_outputs def _update_states_after( self, @@ -746,13 +1045,30 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput: """Execute one step for one scheduled request and return runner output.""" + return self._execute_stepwise( + scheduler_output, + validate_kv_metadata=True, + record_output_peak_memory=True, + ) + + def _execute_stepwise( + self, + scheduler_output: DiffusionSchedulerOutput, + *, + validate_kv_metadata: bool, + record_output_peak_memory: bool, + in_diffusion_kv_memory_profile: bool = False, + ) -> BatchRunnerOutput: + """Execute one step with explicit validation and profiling policy.""" + assert self.pipeline is not None, "Model not loaded. Call load_model() first." for new_req in scheduler_output.scheduled_new_reqs: validate_new_request_data_identity(new_req) - self._validate_diffusion_kv_metadata( - request_id=new_req.req.request_id, - metadata=new_req.diffusion_kv_metadata, - ) + if validate_kv_metadata: + self._validate_diffusion_kv_metadata( + request_id=new_req.req.request_id, + metadata=new_req.diffusion_kv_metadata, + ) if not self._supports_step_mode(): raise ValueError("Current pipeline does not support step execution.") # Stepwise mode only supports the basic state-driven denoise path for now. @@ -761,21 +1077,69 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if self.od_config.cache_backend not in (None, "none"): raise ValueError("Step mode does not support cache_backend yet.") + # Scheduler metadata is installed only for newly admitted requests; + # cached requests continue to use the row installed on their first + # step. Paged rows are retired here for scheduler-finished requests so + # a subsequent wave can reuse the native table slots. Dense execution + # must not issue Worker KV cleanup side effects. + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + and scheduler_output.finished_req_ids + ): + self.remove_diffusion_kv_requests(list(scheduler_output.finished_req_ids)) + installed_request_ids: list[str] = [] + try: + for new_req in scheduler_output.scheduled_new_reqs: + if new_req.diffusion_kv_metadata is not None: + self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata) + installed_request_ids.append(new_req.request_id) + return self._execute_stepwise_core( + scheduler_output, + record_output_peak_memory=record_output_peak_memory, + in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, + ) + except Exception: + if installed_request_ids: + self.remove_diffusion_kv_requests(installed_request_ids) + raise + + def _execute_stepwise_core( + self, + scheduler_output: DiffusionSchedulerOutput, + *, + record_output_peak_memory: bool, + in_diffusion_kv_memory_profile: bool = False, + ) -> BatchRunnerOutput: + """Run the denoise step after metadata admission and row installation.""" + use_hsdp = self.od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: had_active_states = bool(self.state_cache) states, new_request_ids = self._update_states(scheduler_output) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 - if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available(): + if ( + record_output_peak_memory + and new_request_ids + and not had_active_states + and is_primary + and current_omni_platform.is_available() + ): current_omni_platform.reset_peak_memory_stats() - input_batch = self._prepare_batch_inputs(states, new_request_ids) + states, input_batch, runner_output_list = self._prepare_batch_inputs(states, new_request_ids) + if input_batch is None: + return BatchRunnerOutput.from_list(runner_output_list) attn_metadata = {} + kv_backend = getattr(self, "diffusion_kv_backend", None) + paged_kv_runtime = kv_backend if getattr(kv_backend, "paged_attention_adapter", None) is not None else None with set_forward_context( vllm_config=self.vllm_config, omni_diffusion_config=self.od_config, attn_metadata=attn_metadata, + paged_kv_runtime=paged_kv_runtime, + in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, ): clear_pipeline_stage_durations(self.pipeline) noise_pred = self.pipeline.denoise_step(input_batch, states=states) @@ -786,7 +1150,6 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): denoise_stage_durations, ) - runner_output_list = [] pipeline_interrupted = getattr(self.pipeline, "interrupt", False) if noise_pred is None and pipeline_interrupted: for state in states: @@ -839,6 +1202,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): offset = offset + row_num except Exception as per_req_exc: offset = offset + row_num + self.state_cache.pop(req.request_id, None) logger.error( "Stepwise per-request error for %s: %s", req.request_id, @@ -850,7 +1214,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): request_id=req.request_id, step_index=req.step_index, finished=True, - result=DiffusionOutput(error=str(per_req_exc)), + result=DiffusionOutput.from_exception(per_req_exc), ) ) @@ -860,7 +1224,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): f"but batched noise_pred has {noise_pred.shape[0]} rows." ) - if is_primary: + if is_primary and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() states_by_id = {state.request_id: state for state in states} for state in states: @@ -876,7 +1240,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): state.peak_memory_mb, ) + terminal_request_ids = [ + runner_output.request_id for runner_output in runner_output_list if runner_output.finished + ] self._update_states_after(states, input_batch, pipeline_interrupted) + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + and terminal_request_ids + ): + self.remove_diffusion_kv_requests(terminal_request_ids) return BatchRunnerOutput.from_list(runner_output_list) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 0820d82b4..5f8fa3e82 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Diffusion Worker for vLLM-Omni. @@ -13,23 +13,27 @@ import multiprocessing as mp import os import queue import signal +import sys import threading import traceback import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext -from dataclasses import dataclass from typing import Any import torch +import torch.distributed as dist import zmq -from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config +from vllm.config import VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue +from vllm.distributed.parallel_state import get_ep_group, get_tp_group from vllm.logger import init_logger from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler -from vllm.transformers_utils.config import get_hf_text_config from vllm.utils.import_utils import resolve_obj_by_qualname -from vllm.utils.mem_utils import GiB_bytes +from vllm.utils.mem_utils import GiB_bytes, MemorySnapshot, format_gib, memory_profiling +from vllm.utils.system_utils import decorate_logs, set_process_title +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec +from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import ( @@ -41,18 +45,32 @@ from vllm_omni.diffusion.data import ( OmniSleepTask, OmniWakeTask, ) +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, + get_cfg_group, + get_dp_group, + get_fs_group, + get_hsdp_replicate_group, + get_pp_group, + get_sp_group, init_distributed_environment, initialize_model_parallel, + model_parallel_is_initialized, ) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob +from vllm_omni.diffusion.sched.interface import ( + DiffusionSchedulerOutput, + KVPrefetchJob, + NewRequestData, + validate_new_request_data_identity, +) +from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput from vllm_omni.engine.stage_init_utils import set_death_signal @@ -65,51 +83,90 @@ from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory logger = init_logger(__name__) _ASYNC_OUTPUT_THREAD_JOIN_TIMEOUT_S = 10.0 +# Maximum time (in seconds) to wait for pending background D2H / SHM packing +# to drain before the worker executes memory-releasing lifecycle tasks +# (e.g. during sleep/wake transitions). This barrier prevents device tensors +# from being freed while the side CUDA stream is still actively reading them. +_ASYNC_OUTPUT_DRAIN_TIMEOUT_S = 10.0 + +# Worker entry points that release device memory. Background D2H/SHM packing +# still reads model output tensors, so it must finish before these run. +_MEMORY_RELEASING_METHODS = frozenset({"sleep", "handle_sleep_task"}) -@dataclass -class _DiffusionVllmModelConfig: - model: str - dtype: torch.dtype - quantization: str | None = None - quantization_config: Any | None = None - hf_config: Any | None = None - hf_text_config: Any | None = None - multimodal_config: Any | None = None - enforce_eager: bool = False - disable_cascade_attn: bool = False - enable_return_routed_experts: bool = False - is_moe: bool = False - - def is_quantized(self) -> bool: - return self.quantization is not None - - def is_model_moe(self) -> bool: - return self.is_moe - - def is_nvfp4_quantized(self) -> bool: - return self.quantization == "modelopt_fp4" - - @property - def is_diffusion(self) -> bool: - return False +def _cleanup_after_execution_error(exc: Exception) -> None: + """Release device tensors retained by a failed execution traceback.""" + exc.__traceback__ = None + try: + gc.collect() + current_omni_platform.empty_cache() + except Exception: + logger.warning("Failed to release device memory after an execution error", exc_info=True) -def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig: - quant_config = getattr(od_config, "quantization_config", None) - quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None - hf_config = getattr(od_config, "tf_model_config", None) - hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config - return _DiffusionVllmModelConfig( - model=od_config.model, - dtype=od_config.dtype, - quantization=quantization, - quantization_config=quant_config, - hf_config=hf_config, - hf_text_config=hf_text_config, - enforce_eager=getattr(od_config, "enforce_eager", False), - is_moe=bool(getattr(od_config, "is_moe", False)), - ) +def _all_gather_rank_values(value: Any) -> list[Any]: + if not dist.is_available() or not dist.is_initialized(): + return [value] + values: list[Any] = [None] * dist.get_world_size() + dist.all_gather_object(values, value) + return values + + +def _run_and_gather_rank_values(operation: str, func: Callable[[], Any]) -> list[Any]: + """Run one rank-local probe without stranding peers on local failure.""" + + try: + local_result = (True, func()) + except Exception as exc: + logger.exception("%s failed on this Worker rank", operation) + local_result = (False, f"{type(exc).__name__}: {exc}") + + rank_results = _all_gather_rank_values(local_result) + failures = [f"rank {rank}: {result}" for rank, (ok, result) in enumerate(rank_results) if not ok] + if failures: + raise RuntimeError(f"{operation} failed on " + "; ".join(failures)) + return [result for _, result in rank_results] + + +def _setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep: bool, + use_hsdp: bool, + hsdp_replicate_size: int = 1, +) -> None: + """Set the worker process title and log prefix from initialized groups.""" + process_name = "DiffusionWorker" + if model_parallel_is_initialized(): + dp_group = get_dp_group() + pp_group = get_pp_group() + sp_group = get_sp_group() + cfg_group = get_cfg_group() + tp_group = get_tp_group() + + if dp_group.world_size > 1: + process_name += f"_DP{dp_group.rank_in_group}" + if pp_group.world_size > 1: + process_name += f"_PP{pp_group.rank_in_group}" + if sp_group.world_size > 1: + process_name += f"_SP{sp_group.rank_in_group}" + if cfg_group.world_size > 1: + process_name += f"_CFG{cfg_group.rank_in_group}" + if tp_group.world_size > 1: + process_name += f"_TP{tp_group.rank_in_group}" + if use_hsdp: + fs_group = get_fs_group() + if fs_group.world_size > 1: + process_name += f"_FS{fs_group.rank_in_group}" + if hsdp_replicate_size > 1: + replicate_group = get_hsdp_replicate_group() + if replicate_group.world_size > 1: + process_name += f"_RP{replicate_group.rank_in_group}" + if enable_ep: + ep_group = get_ep_group() + if ep_group.world_size > 1: + process_name += f"_EP{ep_group.rank_in_group}" + + set_process_title(name=process_name, prefix="vLLM-Omni") + decorate_logs(process_name) @contextmanager @@ -143,37 +200,6 @@ def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[No yield -def _is_unexpected_additional_config_type_error(exc: TypeError) -> bool: - """Return True only for constructor rejections of the additional_config kwarg.""" - message = str(exc) - return "unexpected keyword argument" in message and "additional_config" in message - - -def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDiffusionConfig) -> VllmConfig: - """Create a worker-local VllmConfig while preserving additional_config when supported.""" - config_kwargs: dict[str, Any] = { - "compilation_config": CompilationConfig(), - "device_config": DeviceConfig(device=device), - } - if od_config.additional_config: - config_kwargs["additional_config"] = od_config.additional_config - - try: - return VllmConfig(**config_kwargs) - except TypeError as exc: - if not _is_unexpected_additional_config_type_error(exc): - raise - - logger.debug("Worker-local VllmConfig does not accept additional_config in constructor: %s", exc) - config_kwargs.pop("additional_config", None) - vllm_config = VllmConfig(**config_kwargs) - try: - setattr(vllm_config, "additional_config", dict(od_config.additional_config)) - except Exception as set_exc: # pragma: no cover - defensive for older vLLM builds - logger.warning("Failed to attach additional_config to worker VllmConfig: %s", set_exc) - return vllm_config - - def _get_cumem_allocator_class() -> type: from vllm.device_allocator.cumem import CuMemAllocator @@ -214,6 +240,8 @@ class DiffusionWorker: self.device: torch.device | None = None self.vllm_config: VllmConfig | None = None self.model_runner: DiffusionModelRunner | None = None + self.init_snapshot: MemorySnapshot | None = None + self.requested_memory: int | None = None self._sleep_saved_buffers: dict[str, torch.Tensor] = {} self.lora_manager: DiffusionLoRAManager | None = None # Worker-side cache of (lora_request, lora_scale) per scheduled @@ -278,22 +306,7 @@ class DiffusionWorker: # Create vllm_config for parallel configuration. Pass explicit device_config # so DeviceConfig does not rely on current_platform in worker subprocesses. - vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config) - parallel_config = self.od_config.parallel_config - vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size - vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size - if parallel_config.enable_expert_parallel and self.od_config.is_moe: - # Diffusion uses its own DP/CFG/SP groups normally. vLLM groups are - # only remapped for expert-parallel runtimes that consume vLLM's - # FusedMoE/EP semantics. - vllm_config.parallel_config.data_parallel_size = ( - parallel_config.data_parallel_size * parallel_config.cfg_parallel_size - ) - vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size - vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel - vllm_config.profiler_config = self.od_config.profiler_config - vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config) # type: ignore[assignment] - vllm_config.quant_config = self.od_config.quantization_config + vllm_config = create_diffusion_vllm_config(self.device, self.od_config) # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping. # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion. vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config) @@ -328,9 +341,31 @@ class DiffusionWorker: tensor_parallel_size=parallel_config.tensor_parallel_size, pipeline_parallel_size=parallel_config.pipeline_parallel_size, fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1, - hsdp_replicate_size=parallel_config.hsdp_replicate_size if parallel_config.use_hsdp else 1, enable_expert_parallel=parallel_config.enable_expert_parallel, + use_hsdp=parallel_config.use_hsdp, ) + _setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep=parallel_config.enable_expert_parallel, + use_hsdp=parallel_config.use_hsdp, + hsdp_replicate_size=parallel_config.hsdp_replicate_size, + ) + if ( + getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ): + gc.collect() + current_omni_platform.empty_cache() + self.init_snapshot = MemorySnapshot(device=self.device) + self.requested_memory = request_memory( + self.init_snapshot, + vllm_config.cache_config, + ) + logger.debug( + "Worker %d: Diffusion KV initial memory snapshot: %r; requested=%s GiB", + self.rank, + self.init_snapshot, + format_gib(self.requested_memory), + ) init_workspace_manager(self.device) def _create_profiler(self) -> WorkerProfiler | None: @@ -384,9 +419,130 @@ class DiffusionWorker: if load_format != "dummy": assert self.model_runner.pipeline is not None + def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]: + """Return native rank-local specs for every diffusion Worker.""" + + assert self.model_runner is not None + return _run_and_gather_rank_values( + "Diffusion KV cache spec discovery", + self.model_runner.get_kv_cache_spec, + ) + + def determine_available_kv_memory(self, profile_requests: list[OmniDiffusionRequest]) -> list[int]: + """Profile and return each rank's safe Diffusion KV memory budget.""" + + def determine_local_memory() -> int: + assert self.vllm_config is not None + assert self.model_runner is not None + if self.init_snapshot is None or self.requested_memory is None: + raise RuntimeError("Diffusion KV memory snapshot was not captured before model loading") + override = self.vllm_config.cache_config.kv_cache_memory_bytes + if override: + # Match native vLLM: an explicit cache budget skips automatic + # capacity derivation, but still runs the maximum-shape model + # request so lazy kernels and communication buffers initialize. + self.model_runner.profile_run(profile_requests) + logger.info( + "Worker %d: Initial free memory %s GiB, reserved %s GiB memory for " + "Diffusion KV Cache as specified by kv_cache_memory_bytes config and " + "skipped automatic memory profiling. This does not respect the " + "gpu_memory_utilization config. A profile warmup was still executed.", + self.rank, + format_gib(self.init_snapshot.free_memory), + format_gib(int(override)), + ) + return int(override) + + with memory_profiling( + self.init_snapshot, + weights_memory=self.model_runner.model_memory_usage, + ) as profile_result: + self.model_runner.profile_run(profile_requests) + + available_memory = self.requested_memory - profile_result.non_kv_cache_memory + if available_memory <= 0: + raise RuntimeError( + "No memory remains for Diffusion KV cache after profiling: " + f"requested_memory={self.requested_memory} bytes, " + f"non_kv_cache_memory={profile_result.non_kv_cache_memory} bytes. " + "Increase gpu_memory_utilization or reduce the maximum profile request shape." + ) + free_gpu_memory = profile_result.after_profile.free_memory + unrequested_memory = self.init_snapshot.free_memory - self.requested_memory + logger.debug( + "Worker %d: Initial free memory: %s GiB; Requested memory: %f (util), %s GiB", + self.rank, + format_gib(self.init_snapshot.free_memory), + self.vllm_config.cache_config.gpu_memory_utilization, + format_gib(self.requested_memory), + ) + logger.debug( + "Worker %d: Free memory after profiling: %s GiB (total), %s GiB (within requested)", + self.rank, + format_gib(free_gpu_memory), + format_gib(free_gpu_memory - unrequested_memory), + ) + logger.debug("Worker %d: %r", self.rank, profile_result) + logger.info( + "Worker %d: Available Diffusion KV cache memory: %s GiB", + self.rank, + format_gib(available_memory), + ) + return int(available_memory) + + return [ + int(value) + for value in _run_and_gather_rank_values( + "Diffusion KV memory discovery", + determine_local_memory, + ) + ] + + def set_kv_cache_configs( + self, + kv_cache_configs: list[KVCacheConfig], + resolved_max_model_len: int, + ) -> None: + """Select this rank's config and initialize its physical KV pages.""" + + assert self.model_runner is not None + assert self.vllm_config is not None + if len(kv_cache_configs) != self.od_config.num_gpus: + raise ValueError( + "Diffusion KVCacheConfig rank count mismatch: " + f"expected={self.od_config.num_gpus}, got={len(kv_cache_configs)}" + ) + if not 0 <= self.rank < len(kv_cache_configs): + raise ValueError(f"Diffusion Worker rank {self.rank} has no rank-local KVCacheConfig") + if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0: + raise ValueError("resolved Diffusion KV max_model_len must be a positive integer") + + # Native cache sizing may resolve an explicit ``-1`` model-length + # sentinel to the capacity that actually fits the profiled pool. + self.vllm_config.model_config.max_model_len = resolved_max_model_len + kv_cache_config = kv_cache_configs[self.rank] + self.vllm_config.cache_config.num_gpu_blocks = kv_cache_config.num_blocks + with self._maybe_get_memory_pool_context("kv_cache"): + self.model_runner.set_kv_cache_config(kv_cache_config) + + def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int: + """Clear Worker-local rows without freeing Scheduler-owned blocks.""" + + assert self.model_runner is not None, "Model runner not initialized" + return self.model_runner.remove_diffusion_kv_requests(request_ids) + def init_lora_manager(self) -> None: """Initialize the LoRA manager for this worker.""" - if self.model_runner.pipeline is None: + pipeline = self.model_runner.pipeline + if pipeline is None: + return + + # A release whose weights cannot be expressed as switchable LoRA layers + # is fused into the checkpoint while the pipeline loads. There is then + # no adapter left to register, and handing the same path to the manager + # would only fail on a format it does not accept. + if getattr(pipeline, "lora_is_fused", False): + logger.info("LoRA was fused into the checkpoint at load time; skipping the dynamic LoRA manager.") return lora_path = self.od_config.lora_path @@ -435,20 +591,40 @@ class DiffusionWorker: else: profiler.stop() + def _run_ar_diffusion_session_lifecycle(self, method: str, session_id: str) -> bool: + """Delegate an optional AR session lifecycle call to the model runner.""" + if not isinstance(session_id, str) or not session_id.strip(): + raise ValueError("session_id must be a non-empty string.") + assert self.model_runner is not None, "Model runner not initialized" + lifecycle_method = getattr(self.model_runner, method, None) + if not callable(lifecycle_method): + return False + lifecycle_method(session_id) + return True + + def reset_ar_diffusion_session(self, session_id: str) -> bool: + """Reset runner-owned AR state through the collective RPC boundary.""" + return self._run_ar_diffusion_session_lifecycle("reset_session", session_id) + + def close_ar_diffusion_session(self, session_id: str) -> bool: + """Close runner-owned AR state through the collective RPC boundary.""" + return self._run_ar_diffusion_session_lifecycle("close_session", session_id) + def execute_model( self, - req: OmniDiffusionRequest | list[OmniDiffusionRequest], + req: OmniDiffusionRequest | list[NewRequestData], od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, diffusion_kv_metadata: DiffusionKVMetadata | None = None, ) -> DiffusionOutput: """Execute a forward pass by delegating to the model runner. - If *req* is a list (DP multi-concurrency), each rank picks one - request based on its distributed rank. AllGather in the layerwise - offload only gathers weight shards (request-independent), so all - ranks stay synchronised at each AllGather call while computing - different activations. + If *req* is a list (DP multi-concurrency), each rank picks one complete + NewRequestData envelope based on its distributed rank. AllGather in + the layerwise offload only gathers weight shards (request-independent), + so all ranks stay synchronised at each AllGather call while computing + different activations. Selecting the envelope keeps Scheduler-issued + KV metadata bound to the request that owns its block tables. Each rank returns its OWN DiffusionOutput (no gather). The executor collects N responses via the per-worker result queues. @@ -464,7 +640,10 @@ class DiffusionWorker: dp_rank = get_data_parallel_rank() idx = dp_rank % len(req) - req = req[idx] + new_req = req[idx] + validate_new_request_data_identity(new_req) + req = new_req.req + diffusion_kv_metadata = new_req.diffusion_kv_metadata if self.lora_manager is not None: try: @@ -564,11 +743,15 @@ class DiffusionWorker: logger.warning("LoRA activation skipped: %s", exc) def remove_lora(self, adapter_id: int) -> bool: + if self.lora_manager is None: + return False return self.lora_manager.remove_adapter(adapter_id) def add_lora(self, lora_request: LoRARequest) -> bool: # NOTE (Alex): We have not implemented the API routing # for the frontend server yet. + if self.lora_manager is None: + return False return self.lora_manager.add_adapter(lora_request) def submit_interaction( @@ -581,9 +764,13 @@ class DiffusionWorker: self.model_runner.submit_interaction(request_id, interaction) def list_loras(self) -> list[int]: + if self.lora_manager is None: + return [] return self.lora_manager.list_adapters() def pin_lora(self, adapter_id: int) -> bool: + if self.lora_manager is None: + return False return self.lora_manager.pin_adapter(adapter_id) def sleep(self, level: int = 1) -> int: @@ -649,6 +836,8 @@ class DiffusionWorker: allocator = CuMemAllocator.get_instance() allocator.wake_up(tags) current_omni_platform.synchronize() + if self.model_runner is not None and (tags is None or "kv_cache" in tags): + self.model_runner.refresh_diffusion_kv_block_table_layout() if len(self._sleep_saved_buffers) and self.model_runner is not None: model = self.model_runner.pipeline for name, buffer in model.named_buffers(): @@ -777,11 +966,25 @@ class DiffusionWorker: def shutdown(self) -> None: """Shutdown the worker and cleanup distributed environment.""" - if self.model_runner is not None: - mgr = getattr(self.model_runner, "kv_transfer_manager", None) - if mgr is not None: - mgr.shutdown_prefetch() - destroy_distributed_env() + try: + if self.model_runner is not None: + mgr = getattr(self.model_runner, "kv_transfer_manager", None) + try: + offload_backend = getattr(self.model_runner, "offload_backend", None) + if offload_backend is not None: + offload_backend.disable() + finally: + if mgr is not None: + mgr.shutdown_prefetch() + finally: + try: + a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute") + if a2a_permute is not None: + a2a_permute.clear_a2a_permute_workspaces() + except Exception: + logger.exception("Failed to release fused Ulysses symmetric-memory workspaces") + finally: + destroy_distributed_env() class CustomPipelineWorkerExtension: @@ -850,6 +1053,13 @@ class WorkerProc: self._async_output_queue: queue.Queue | None = None self._async_output_thread: threading.Thread | None = None + self._async_output_done = threading.Condition() + self._async_output_pending = 0 + # MessageQueue.acquire_write() is single-writer: it reads current_idx, + # writes that block, then advances the index. The async output thread + # enqueues OUTPUT_READY while the main loop enqueues COMPUTE_DONE, so + # unsynchronized writers can target the same block and drop a message. + self._result_mq_lock = threading.Lock() if not self.od_config.step_execution: self._async_output_queue = queue.Queue() self._async_output_thread = threading.Thread( @@ -882,25 +1092,32 @@ class WorkerProc: ) return wrapper + def _enqueue_result(self, msg: Any) -> None: + """Serialize writes to the single-writer result queue.""" + with self._result_mq_lock: + self.result_mq.enqueue(msg) + def _return_result(self, output: Any, rpc_id: str | None = None) -> None: """Reply to client, only on rank 0.""" if self.result_mq is None: return if isinstance(output, OmniACK): - self.result_mq.enqueue(output) + self._enqueue_result(output) return # Async path: enqueue compute_done immediately, bg thread does D2H+SHM. if not self.od_config.step_execution and isinstance(output, (DiffusionOutput, BatchRunnerOutput)): async_output_id = WorkerProc._generate_async_output_id() gpu_event = current_omni_platform.record_device_event() + with self._async_output_done: + self._async_output_pending += 1 self._async_output_queue.put((output, async_output_id, gpu_event)) msg = AsyncDiffusionOutput( kind=AsyncOutputKind.COMPUTE_DONE, rpc_id=rpc_id, async_output_id=async_output_id, ) - self.result_mq.enqueue(msg) + self._enqueue_result(msg) return # Sync path (original, or async fallback). @@ -909,7 +1126,7 @@ class WorkerProc: except Exception as e: if hasattr(output, "output"): logger.warning("SHM pack failed for model output: %s", e) - self.result_mq.enqueue(output) + self._enqueue_result(output) def _async_output_loop(self): """Background thread: D2H + SHM packing for async diffusion output. @@ -932,7 +1149,7 @@ class WorkerProc: pack_diffusion_output_shm(output, d2h_stream=d2h_stream) d2h_stream.synchronize() - self.result_mq.enqueue( + self._enqueue_result( AsyncDiffusionOutput( kind=AsyncOutputKind.OUTPUT_READY, async_output_id=async_output_id, @@ -944,13 +1161,38 @@ class WorkerProc: "Async output packing failed for id '%s'; sending error", async_output_id, ) - self.result_mq.enqueue( + self._enqueue_result( AsyncDiffusionOutput( kind=AsyncOutputKind.OUTPUT_READY, async_output_id=async_output_id, error="Background D2H/SHM packing failed", ) ) + finally: + with self._async_output_done: + # Clamped: only items enqueued by _return_result are counted. + self._async_output_pending = max(0, self._async_output_pending - 1) + self._async_output_done.notify_all() + + def drain_async_outputs(self, timeout: float = _ASYNC_OUTPUT_DRAIN_TIMEOUT_S) -> bool: + """Block until background D2H/SHM packing has no work left. + + Returns False if outputs are still in flight when ``timeout`` expires. + """ + with self._async_output_done: + if self._async_output_pending == 0: + return True + drained = self._async_output_done.wait_for(lambda: self._async_output_pending == 0, timeout=timeout) + pending = self._async_output_pending + if not drained: + logger.warning( + "Worker %d: %d async output(s) still in flight after %.1fs; " + "releasing device memory now may drop OUTPUT_READY messages", + self.gpu_id, + pending, + timeout, + ) + return drained def shutdown(self) -> None: """Stop background work and release worker-owned IPC resources.""" @@ -1038,7 +1280,6 @@ class WorkerProc: return None, False result = None - rpc_exception: Exception | None = None status: dict[str, Any] = { "rank": self.gpu_id, "ok": True, @@ -1049,11 +1290,12 @@ class WorkerProc: } try: + if method in _MEMORY_RELEASING_METHODS: + self.drain_async_outputs() # Use execute_method from WorkerWrapperBase for consistent method resolution result = self.worker.execute_method(method, *args, **kwargs) except Exception as e: logger.error(f"Error executing RPC: {e}", exc_info=True) - rpc_exception = e status.update( { "ok": False, @@ -1062,6 +1304,9 @@ class WorkerProc: "traceback": traceback.format_exc(), } ) + if not collect_rank_status: + raise + _cleanup_after_execution_error(e) if isinstance(result, bool): status["bool_result"] = result @@ -1081,12 +1326,14 @@ class WorkerProc: ) return None, False - if rpc_exception is not None: - raise rpc_exception if isinstance(result, dict) and wave_id is not None: result["wave_id"] = wave_id return result, should_reply + def recv_message(self) -> Any: + """Receive one complete broadcast message without dropping overflow data.""" + return self.mq.dequeue(indefinite=True) + def _worker_busy_loop(self) -> None: """Main busy loop for Multiprocessing Workers.""" logger.info(f"Worker {self.gpu_id} ready to receive requests via shared memory") @@ -1094,7 +1341,7 @@ class WorkerProc: while self._running: msg = None try: - msg = self.mq.dequeue(timeout=1.0) + msg = self.recv_message() except Exception: if self.wake_event and self.wake_event.is_set(): self.wake_event.clear() @@ -1110,6 +1357,7 @@ class WorkerProc: continue if isinstance(msg, dict) and msg.get("type") == "sleep": + self.drain_async_outputs() task = OmniSleepTask(level=msg.get("level", 2), task_id=msg.get("task_id", "local")) ack = self.worker.handle_sleep_task(task) self._return_result(ack) @@ -1126,6 +1374,8 @@ class WorkerProc: self._return_result(result, rpc_id=rpc_id) except Exception as e: logger.error(f"Error processing RPC: {e}", exc_info=True) + error = str(e) + _cleanup_after_execution_error(e) # Apply the same reply gate as the success path so # non-output ranks don't enqueue stale error replies # that compete with the expected responder's message. @@ -1136,11 +1386,11 @@ class WorkerProc: if rpc_id is not None: # Async RPC: must complete the executor's pending # future so collective_rpc() doesn't hang. - self.result_mq.enqueue( + self._enqueue_result( AsyncDiffusionOutput( kind=AsyncOutputKind.RPC_RESULT, rpc_id=rpc_id, - error=str(e), + error=error, ) ) elif output_rank is None and exec_all_ranks: @@ -1168,11 +1418,11 @@ class WorkerProc: except Exception: dp_rank = self.gpu_id self._return_result( - {"status": "error", "error": str(e), "dp_rank": dp_rank, "wave_id": wave_id} + {"status": "error", "error": error, "dp_rank": dp_rank, "wave_id": wave_id} ) elif output_rank is None or output_rank == self.gpu_id: # Normal RPC: only the expected rank replies - self._return_result({"status": "error", "error": str(e), "wave_id": wave_id}) + self._return_result({"status": "error", "error": error, "wave_id": wave_id}) elif isinstance(msg, dict) and msg.get("type") == "shutdown": logger.info("Worker %s: Received shutdown message", self.gpu_id) @@ -1189,6 +1439,7 @@ class WorkerProc: exc_info=True, ) output = DiffusionOutput.from_exception(e) + _cleanup_after_execution_error(e) try: self._return_result(output) @@ -1224,13 +1475,11 @@ class WorkerProc: set_death_signal(signal.SIGTERM) - # Set process title for visibility in nvidia-smi / htop (optional, non-fatal) - try: - import setproctitle - - setproctitle.setproctitle(f"vLLM-Omni::DiffusionWorker-{rank}") - except ImportError: - pass # setproctitle not installed, skip process title setting + _setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep=od_config.parallel_config.enable_expert_parallel, + use_hsdp=od_config.parallel_config.use_hsdp, + hsdp_replicate_size=od_config.parallel_config.hsdp_replicate_size, + ) load_omni_general_plugins() worker_proc = None @@ -1369,7 +1618,7 @@ class WorkerWrapperBase: def execute_model( self, - req: OmniDiffusionRequest, + req: OmniDiffusionRequest | list[NewRequestData], od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, diffusion_kv_metadata: DiffusionKVMetadata | None = None, diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index f1e224b15..d585b3ba1 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import time from collections.abc import Callable from typing import Any +from vllm.v1.request import Request + from vllm_omni.metrics import OrchestratorAggregator from .utils.logging import get_connector_logger @@ -216,38 +218,159 @@ def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: return sum_user_len + assistant_len -def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], request: Any) -> None: +def construct_next_stage_streaming_input_prompt( + payload_data: dict[str, Any], + request: Request, + *, + max_model_len: int | None = None, + previous_condition_len: int | None = None, + previous_condition_seq: int | None = None, + condition_seq: int | None = None, + recompute_previous_chunks: int = 0, +) -> bool: """Update a downstream streaming request prompt from connector payload ids. Async-chunk downstream stages are prewarmed before the real Talker prompt is known. When a Thinker payload carries ``ids.prompt``, this helper: - * Preserves ``num_computed_tokens`` (the scheduler token watermark). + * Preserves ``num_computed_tokens`` while extending a non-window prompt. + Explicit replacements and one-condition window recomputes reset it. * Moves already-computed output tokens into ``prompt_token_ids``. * Appends a new placeholder prompt slice sized from the upstream ids. * Refreshes block hashes so the scheduler allocates KV slots for the extended prompt without discarding prior computed state. """ - ids = payload_data.get("ids", {}) - meta = payload_data.get("meta", {}) - replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True - next_stage_prompt_len = meta.get("next_stage_prompt_len") if isinstance(meta, dict) else None - if replace_prompt and isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0: + ids = payload_data.get("ids") + if not isinstance(ids, dict): + ids = {} + payload_data["ids"] = ids + meta = payload_data.get("meta") + if not isinstance(meta, dict): + meta = {} + payload_data["meta"] = meta + # This flag is transported through a merge-based runtime buffer. Set it on + # every condition so a prior rollover cannot be replayed by a later append. + meta["streaming_prompt_recompute"] = False + next_stage_prompt_len = meta.get("next_stage_prompt_len") + if "next_stage_generation_tokens" in meta: + next_generation_tokens = meta["next_stage_generation_tokens"] + if ( + not isinstance(next_generation_tokens, int) + or isinstance(next_generation_tokens, bool) + or next_generation_tokens <= 0 + ): + raise ValueError("next_stage_generation_tokens must be a positive integer") + generation_reserve = next_generation_tokens + else: + generation_reserve = 0 + capacity_limit = ( + max_model_len + if isinstance(max_model_len, int) and not isinstance(max_model_len, bool) and max_model_len > 0 + else None + ) + managed_prompt_len: int | None = None + if generation_reserve > 0 and capacity_limit is not None: + if ( + not isinstance(next_stage_prompt_len, int) + or isinstance(next_stage_prompt_len, bool) + or next_stage_prompt_len <= 0 + ): + raise ValueError("capacity-managed streaming prompt requires a positive next_stage_prompt_len") + managed_prompt_len = next_stage_prompt_len + if managed_prompt_len + generation_reserve > capacity_limit: + raise ValueError( + "fresh streaming prompt plus generation reserve exceeds max_model_len: " + f"prompt={managed_prompt_len}, reserve={generation_reserve}, limit={capacity_limit}" + ) + explicit_replacement = meta.get("replace_streaming_prompt") is True + if isinstance(recompute_previous_chunks, bool) or recompute_previous_chunks not in (0, 1): + raise ValueError("streaming prompt recompute supports exactly one previous chunk") + has_window_contract = recompute_previous_chunks == 1 + window_recompute = not explicit_replacement and has_window_contract and previous_condition_seq is not None + exceeds_accumulated_capacity = ( + managed_prompt_len is not None + and capacity_limit is not None + and request.num_computed_tokens + managed_prompt_len + generation_reserve > capacity_limit + ) + if not explicit_replacement and exceeds_accumulated_capacity and not window_recompute: + raise ValueError("capacity-managed streaming prompt rollover requires window_size=1") + replacement_prompt_len = next_stage_prompt_len + if window_recompute: + if managed_prompt_len is None or capacity_limit is None: + raise ValueError("streaming prompt window requires a positive generation reserve and max_model_len") + if ( + not isinstance(previous_condition_len, int) + or isinstance(previous_condition_len, bool) + or previous_condition_len <= 0 + or not isinstance(previous_condition_seq, int) + or isinstance(previous_condition_seq, bool) + or previous_condition_seq < 0 + or not isinstance(condition_seq, int) + or isinstance(condition_seq, bool) + or condition_seq < 0 + ): + raise ValueError("streaming prompt recompute requires condition lengths and monotonic sequence metadata") + if condition_seq != previous_condition_seq + 1: + raise ValueError( + "streaming prompt recompute skipped a Talker condition: " + f"previous={previous_condition_seq}, current={condition_seq}" + ) + + num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0) + confirmed_num_computed_tokens = max(0, int(request.num_computed_tokens) - num_placeholders) + if confirmed_num_computed_tokens < request.num_prompt_tokens: + raise ValueError( + "confirmed streaming output ends before the current condition: " + f"confirmed={confirmed_num_computed_tokens}, prompt={request.num_prompt_tokens}" + ) + previous_codec_ids = list(request._all_token_ids[request.num_prompt_tokens : confirmed_num_computed_tokens]) + max_confirmed_codec_tokens = generation_reserve - 1 + if len(previous_codec_ids) > max_confirmed_codec_tokens: + raise ValueError( + "streaming prompt recompute retained too many codec tokens: " + f"retained={len(previous_codec_ids)}, limit={max_confirmed_codec_tokens}" + ) + + replacement_prompt_len = previous_condition_len + len(previous_codec_ids) + managed_prompt_len + if replacement_prompt_len + generation_reserve > capacity_limit: + raise ValueError( + "sliding streaming prompt plus generation reserve exceeds max_model_len: " + f"previous={previous_condition_len}, codec={len(previous_codec_ids)}, " + f"current={managed_prompt_len}, reserve={generation_reserve}, limit={capacity_limit}" + ) + ids["streaming_prompt_previous_codes"] = previous_codec_ids + meta.update( + { + "streaming_prompt_recompute": True, + "streaming_condition_seq": condition_seq, + } + ) + logger.debug( + "Recomputing a one-chunk streaming prompt window: previous=%d, codec=%d, current=%d, reserve=%d, limit=%d", + previous_condition_len, + len(previous_codec_ids), + managed_prompt_len, + generation_reserve, + capacity_limit, + ) + + replace_prompt = explicit_replacement or window_recompute + if replace_prompt and isinstance(replacement_prompt_len, int) and replacement_prompt_len > 0: # Some downstream stages consume complete, independently conditioned # segments instead of extending an existing KV prefix. The producer # declares that transport behavior explicitly in payload metadata. - new_prompt = [0] * next_stage_prompt_len + new_prompt = [0] * replacement_prompt_len request._output_token_ids.clear() request._all_token_ids.clear() request._all_token_ids.extend(new_prompt) request.prompt_token_ids = new_prompt request.num_computed_tokens = 0 - request.num_prompt_tokens = next_stage_prompt_len + request.num_prompt_tokens = replacement_prompt_len request.update_block_hashes() - return + return True prompt_token_ids = ids.get("prompt", None) if not prompt_token_ids: - return + return False num_computed_tokens = request.num_computed_tokens kept_output_tokens = request._all_token_ids[request.num_prompt_tokens : num_computed_tokens] del request._all_token_ids[num_computed_tokens:] @@ -264,3 +387,4 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re request.prompt_token_ids.extend(new_prompt or ()) request.update_block_hashes() request.num_prompt_tokens = len(request.prompt_token_ids) + return False diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py index bbaca48e0..3420dcb94 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import threading from collections import deque @@ -35,6 +35,10 @@ class OmniTransferAdapterBase: self.stop_event = threading.Event() self._recv_cond = threading.Condition() self._save_cond = threading.Condition() + self._send_failures: dict[str, str] = {} + self._send_failure_lock = threading.Lock() + self._receive_failures: dict[str, str] = {} + self._receive_failure_lock = threading.Lock() self.recv_thread = threading.Thread(target=self.recv_loop, daemon=True) self.recv_thread.start() @@ -82,6 +86,35 @@ class OmniTransferAdapterBase: elif not any_success and not self.stop_event.is_set(): self._recv_cond.wait(timeout=0.001) + def record_send_failure(self, request_id: str | None, reason: str) -> None: + """Note that a chunk for *request_id* will never be delivered. + + Called from the sender thread; drained by the scheduler thread, hence + the lock. ``_waiting_since`` needs none because only the scheduler + touches it. + """ + if not request_id: + return + with self._send_failure_lock: + self._send_failures.setdefault(request_id, reason) + + def collect_failed_send_request_ids(self) -> dict[str, str]: + """Drain and return the requests whose send gave up.""" + with self._send_failure_lock: + failures, self._send_failures = self._send_failures, {} + return failures + + def record_receive_failure(self, request_id: str, reason: str) -> None: + """Record an invalid received chunk that cannot be retried.""" + with self._receive_failure_lock: + self._receive_failures.setdefault(request_id, reason) + + def collect_failed_receive_request_ids(self) -> dict[str, str]: + """Drain and return requests whose received chunk was invalid.""" + with self._receive_failure_lock: + failures, self._receive_failures = self._receive_failures, {} + return failures + def save_loop(self): """Loop to send outgoing data.""" while not self.stop_event.is_set(): @@ -90,7 +123,18 @@ class OmniTransferAdapterBase: try: self._send_single_request(task) except Exception as e: - logger.warning(f"Error saving data for {task.get('request_id')}: {e}") + # R1.2 of #4855. The task was popleft()-ed and is not re-queued, + # so this is a give-up, not a retry: the consumer will wait for a + # chunk that is never coming. Record it so the scheduler can fail + # the request now instead of leaving it to the wait deadline. + # The task dict carries the Request object, not a bare id -- + # `task.get("request_id")` is always None (it was that way in the + # original warning too, which logged "None" for every failure). + # Use the scheduler-side id: that is what `self.requests` is keyed + # by, and what `finish_requests` expects. + failed = getattr(task.get("request"), "request_id", None) + logger.error("Send gave up for %s: %s", failed, e) + self.record_send_failure(failed, f"{type(e).__name__}: {e}") with self._save_cond: if not self._pending_save_reqs and not self.stop_event.is_set(): diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index d722adad2..b9b03c9c5 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -1,14 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import copy import importlib +import inspect +import threading +import time from collections import defaultdict, deque -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping from typing import Any import torch from vllm.v1.metrics.stats import PrefillStats from vllm.v1.request import Request, RequestStatus +from vllm.v1.utils import ConstantList from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload @@ -21,6 +26,58 @@ from .base import OmniTransferAdapterBase logger = get_connector_logger(__name__) +class _SenderGeneration: + """Fence one external request generation without blocking cleanup.""" + + def __init__(self) -> None: + self.in_flight = False + self.cancelled = False + + +class _LoadEntry: + """Identify one receiver registration across queue and I/O boundaries.""" + + __slots__ = ("request",) + + def __init__(self, request: Request) -> None: + self.request = request + + @property + def request_id(self) -> str: + return self.request.request_id + + @property + def external_req_id(self) -> str: + return self.request.external_req_id + + +def _resolve_talker_streaming_prompt_config(model_config: Any) -> tuple[int, int]: + """Return the effective Talker context limit and recompute window.""" + max_model_len = int(getattr(model_config, "max_model_len", 0) or 0) + hf_config = getattr(model_config, "hf_config", None) + tts_config = getattr(hf_config, "tts_config", None) + if tts_config is None and getattr(hf_config, "model_type", None) == "minicpmtts": + tts_config = hf_config + tts_max_model_len = ( + tts_config.get("max_position_embeddings", 0) + if isinstance(tts_config, Mapping) + else getattr(tts_config, "max_position_embeddings", 0) + ) + tts_max_model_len = int(tts_max_model_len or 0) + if tts_max_model_len > 0: + max_model_len = min(max_model_len, tts_max_model_len) if max_model_len else tts_max_model_len + attention_type = ( + tts_config.get("attention_type") + if isinstance(tts_config, Mapping) + else getattr(tts_config, "attention_type", None) + ) + # Official MiniCPMTTS supports both full_attention and + # sliding_recompute. Only an explicit Talker-stage policy enables the + # latter; native duplex alone must not override checkpoint semantics. + previous_chunks = 1 if attention_type == "sliding_recompute" else 0 + return max_model_len, previous_chunks + + class OmniChunkTransferAdapter(OmniTransferAdapterBase): """Chunk-level transfer adapter for Omni connector pipelines. @@ -41,7 +98,24 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): def __init__(self, vllm_config: Any): model_config = vllm_config.model_config + # The base constructor starts the save thread, so sender-generation + # state must exist before it can observe an enqueued task. The global + # lock only protects this map and short state transitions. Slow payload + # construction and connector writes never hold it. + self._sender_state_lock = threading.Lock() + # Serialize receiver registration/cleanup with the commit of a chunk + # returned by connector.get(). The get itself stays outside the lock so + # scheduler cleanup never waits for connector I/O. + self._receiver_state_lock = threading.Lock() + # Includes requests queued for polling and the request currently being + # polled. Per-registration identity lets the receiver reject an old + # segment's queued or in-flight work after the same request id resumes. + self._registered_load_entries: dict[str, _LoadEntry] = {} + self._sender_tokens: dict[str, _SenderGeneration] = {} self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs + self._max_model_len, self._streaming_prompt_previous_chunks = _resolve_talker_streaming_prompt_config( + model_config + ) active_stream_window = int(getattr(model_config, "active_stream_window", 0) or 0) model_max_num_seqs = int(getattr(model_config, "max_num_seqs", self.scheduler_max_num_seqs) or 0) if model_max_num_seqs <= 0: @@ -67,6 +141,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1) module = importlib.import_module(module_path) self.custom_process_next_stage_input_func = getattr(module, func_name) + self._processor_accepts_step_tokens: dict[Callable[..., Any], bool] = {} # mapping for request id and chunk id self.put_req_chunk: dict[str, int] = defaultdict(int) self.get_req_chunk: dict[str, int] = defaultdict(int) @@ -74,6 +149,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # but popped at segment boundaries (unlike put_req_chunk which is # request-global for connector key continuity). self.ramp_chunk_count: dict[str, int] = defaultdict(int) + self._adaptive_states: dict[str, Any] = {} self.upstream_exhausted_requests: set[str] = set() self.segment_finished_requests: set[str] = set() self.request_payload = {} @@ -83,6 +159,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.waiting_for_chunk_waiting_requests: deque[Any] = deque() self.waiting_for_chunk_running_requests: deque[Any] = deque() self.requests_with_ready_chunks = set() + self.replaced_streaming_prompt_ids: set[str] = set() self.requests_origin_status = {} self._active_streams: dict[str, Any] = {} # Private hold-queue for non-active running requests. Restored to @@ -93,7 +170,25 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # been read"). self._held_non_active: deque[Any] = deque() self.requests_num_chunks_sent: dict[str, int] = defaultdict(int) + # Boundary tasks advance this before the background sender runs, so a + # late old-segment frame cannot restore the previous watermark. + self._segment_generation: dict[str, int] = defaultdict(int) self._pending_streaming_prefills: dict[str, dict] = {} + # The recv thread only records AR payloads. Prompt mutation is finalized + # by the scheduler after it observes the ready marker, where the request + # token counters form a consistent snapshot. + self._pending_ar_prompt_updates: dict[ + str, + tuple[Request, dict[str, Any], bool, int | None, bool], + ] = {} + self._streaming_condition_lengths: dict[str, int] = {} + self._streaming_condition_seqs: dict[str, int] = {} + # Monotonic timestamp of when each request last began waiting for a + # chunk, refreshed every time one arrives. Read by + # collect_timed_out_request_ids() so a stream that stops advancing + # becomes a client-visible error instead of parking forever. Mirrors + # OmniSchedulingCoordinator._waiting_since on the full-payload path. + self._waiting_since: dict[str, float] = {} @staticmethod def _is_truthy_scalar(value: Any) -> bool: @@ -110,6 +205,38 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0) return max(0, num_computed - num_placeholders) + @staticmethod + def _snapshot_processor_request(request: Request) -> Request: + """Snapshot mutable processor inputs at save-queue admission time.""" + + def snapshot_container(value: Any) -> Any: + if isinstance(value, dict): + return {key: snapshot_container(item) for key, item in value.items()} + if isinstance(value, list): + return [snapshot_container(item) for item in value] + if isinstance(value, tuple): + return tuple(snapshot_container(item) for item in value) + return value + + snapshot = copy.copy(request) + for name in ( + "additional_information", + "prompt_token_ids", + ): + if hasattr(request, name): + setattr(snapshot, name, snapshot_container(getattr(request, name))) + + for private_name, public_name in ( + ("_all_token_ids", "all_token_ids"), + ("_output_token_ids", "output_token_ids"), + ): + token_ids = getattr(request, private_name, None) + if isinstance(token_ids, list): + frozen_ids = token_ids.copy() + setattr(snapshot, private_name, frozen_ids) + setattr(snapshot, public_name, ConstantList(frozen_ids)) + return snapshot + @staticmethod def _refresh_generation_chunk_prefill_state(request: Request) -> None: request.num_prompt_tokens = len(request.prompt_token_ids) @@ -151,16 +278,41 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): return if not hasattr(request, "additional_information"): request.additional_information = None - self._cancelled_load_reqs.discard(request.request_id) - self._pending_load_reqs.append(request) + with self._receiver_state_lock: + request_id = request.request_id + if request_id in self._registered_load_entries: + return + entry = _LoadEntry(request) + self._cancelled_load_reqs.discard(request_id) + self._registered_load_entries[request_id] = entry + self._pending_load_reqs.append(entry) with self._recv_cond: self._recv_cond.notify() + def _accepts_new_token_ids(self, processor: Callable[..., Any]) -> bool: + cached = self._processor_accepts_step_tokens.get(processor) + if cached is not None: + return cached + try: + parameters = inspect.signature(processor).parameters.values() + except (TypeError, ValueError): + accepts = False + else: + accepts = any( + parameter.name == "new_token_ids" or parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in parameters + ) + self._processor_accepts_step_tokens[processor] = accepts + return accepts + def save_async( self, multimodal_output: dict[str, Any] | None = None, request: Request | None = None, is_segment_finished: bool = False, + new_token_ids: Iterable[int] | None = None, + confirmed_num_computed_tokens: int | None = None, + segment_generation: int | None = None, ): """Build and enqueue one chunk for asynchronous sending. @@ -178,38 +330,104 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): multimodal_output: Per-request multimodal output dictionary request: Request object is_segment_finished: whether the segment of request is finished + new_token_ids: immutable snapshot of tokens generated by this + scheduler step, before resumable cleanup mutates the request + confirmed_num_computed_tokens: committed token count captured + before a streaming transition can mutate ``request`` + segment_generation: generation captured before a resumable stop + can apply a queued update to the mutable request """ is_finished = request.is_finished() and not request.resumable - - confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request) - - # If the request is preempted, skip the already saved chunks. - if confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0): + if not hasattr(self, "_segment_generation"): + self._segment_generation = defaultdict(int) + external_req_id = request.external_req_id + if segment_generation is None: + raw_generation = getattr(request, "_omni_segment_generation", 0) + try: + generation = int(raw_generation) + except (TypeError, ValueError): + generation = 0 + else: + generation = int(segment_generation) + expected_generation = self._segment_generation.get(external_req_id, generation) + if generation < expected_generation: logger.warning( - f"Enqueue save_async for request {request.external_req_id}, " - f"request.num_computed_tokens={request.num_computed_tokens}, " - f"request.num_output_placeholders={getattr(request, 'num_output_placeholders', 0)}, " - f"previous_chunks_sent={self.requests_num_chunks_sent.get(request.external_req_id, 0)}" + "Skip late save_async for request %s, segment_generation=%s, expected=%s", + external_req_id, + generation, + expected_generation, ) return + self._segment_generation[external_req_id] = generation - self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens + if confirmed_num_computed_tokens is None: + confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request) + processor_request = ( + self._snapshot_processor_request(request) if self.custom_process_next_stage_input_func else request + ) task = { "multimodal_output": multimodal_output, - "request": request, + "request": processor_request, "is_finished": is_finished, "is_segment_finished": is_segment_finished, + "new_token_ids": tuple(int(token_id) for token_id in (new_token_ids or ())), + "segment_generation": generation, } - self._pending_save_reqs.append(task) + + reject_reason = None + with self._sender_state_lock: + sender_token = self._sender_tokens.get(external_req_id) + if sender_token is not None and sender_token.cancelled: + reject_reason = "previous sender generation is still draining" + # If the request is preempted, skip the already saved chunks. + elif not is_segment_finished and confirmed_num_computed_tokens < self.requests_num_chunks_sent.get( + external_req_id, 0 + ): + logger.warning( + f"Enqueue save_async for request {external_req_id}, " + f"request.num_computed_tokens={request.num_computed_tokens}, " + f"request.num_output_placeholders={getattr(request, 'num_output_placeholders', 0)}, " + f"previous_chunks_sent={self.requests_num_chunks_sent.get(external_req_id, 0)}" + ) + return + + else: + self.requests_num_chunks_sent[external_req_id] = confirmed_num_computed_tokens + if sender_token is None: + sender_token = _SenderGeneration() + self._sender_tokens[external_req_id] = sender_token + task["sender_token"] = sender_token + self._pending_save_reqs.append(task) + if is_segment_finished: + # The queued FIFO item now owns the old segment. Start the next + # segment's deduplication watermark before the worker sends it. + # Only advance past the current generation when the producer + # maintains a request-side counter that can catch up. On paths + # where the request carries no `_omni_segment_generation`, a + # speculative +1 makes every subsequent frame look stale and + # starves the stream permanently (see #6816). + if getattr(request, "_omni_segment_generation", None) is not None: + self._segment_generation[external_req_id] = generation + 1 + else: + self._segment_generation[external_req_id] = generation + self.requests_num_chunks_sent.pop(external_req_id, None) + if reject_reason is not None: + logger.error("Cannot enqueue %s: %s", external_req_id, reject_reason) + self.record_send_failure(request.request_id, reject_reason) + return with self._save_cond: self._save_cond.notify() - def _poll_single_request(self, request: Request): + def _poll_single_request(self, entry: _LoadEntry): + request = entry.request stage_id = self.connector.stage_id target_stage_id = stage_id - 1 req_id = request.request_id - chunk_id = self.get_req_chunk[req_id] - external_req_id = self.request_ids_mapping.get(req_id, req_id) + with self._receiver_state_lock: + if self._registered_load_entries.get(req_id) is not entry: + return True + chunk_id = self.get_req_chunk[req_id] + external_req_id = self.request_ids_mapping.get(req_id, req_id) connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" # Use timeout=0 for non-blocking poll @@ -221,10 +439,46 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): ) except Exception as e: logger.error(f"SharedMemoryConnector get failed for req {connector_get_key}: {e}") + with self._receiver_state_lock: + if self._registered_load_entries.get(req_id) is not entry: + return True return False if result is None: + with self._receiver_state_lock: + if self._registered_load_entries.get(req_id) is not entry: + return True return False + + with self._receiver_state_lock: + # cleanup_receiver() can run while connector.get() is in flight. + # Treat cleanup as a commit barrier: a late chunk must not recreate + # prompt/window state for the removed request. + if self._registered_load_entries.get(req_id) is not entry: + return True + is_success = self._commit_received_chunk( + request, + result, + stage_id=stage_id, + req_id=req_id, + chunk_id=chunk_id, + connector_get_key=connector_get_key, + ) + if is_success: + self._registered_load_entries.pop(req_id, None) + return is_success + + def _commit_received_chunk( + self, + request: Request, + result: tuple[dict[str, Any], int], + *, + stage_id: int, + req_id: str, + chunk_id: int, + connector_get_key: str, + ) -> bool: + """Commit a received connector chunk while receiver state is locked.""" payload_data, size = result if payload_data: @@ -232,14 +486,38 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.get_req_chunk[req_id] += 1 meta = payload_data.get("meta", {}) + if not isinstance(meta, dict): + meta = {} + payload_data["meta"] = meta payload_finished = self._is_truthy_scalar(meta.get("finished")) payload_segment_finished = self._is_truthy_scalar(meta.get("is_segment_finished")) if self.model_mode == "ar": - request.additional_information = payload_data + was_resumable = bool(getattr(request, "resumable", False)) + meta["streaming_prompt_recompute"] = False + prompt_len = meta.get("next_stage_prompt_len") replace_prompt = meta.get("replace_streaming_prompt") is True - if getattr(request, "resumable", False) and (chunk_id > 0 or replace_prompt): - # For new streaming input segment, we should update prompt from payload - construct_next_stage_streaming_input_prompt(payload_data, request) + payload_ids = payload_data.get("ids") + has_prompt_payload = bool( + replace_prompt + or (isinstance(prompt_len, int) and not isinstance(prompt_len, bool) and prompt_len > 0) + or (isinstance(payload_ids, dict) and payload_ids.get("prompt")) + ) + window_condition_len = ( + prompt_len + if self._streaming_prompt_previous_chunks == 1 + and isinstance(prompt_len, int) + and not isinstance(prompt_len, bool) + and prompt_len > 0 + else None + ) + update_prompt = bool(was_resumable and has_prompt_payload and (chunk_id > 0 or replace_prompt)) + self._pending_ar_prompt_updates[req_id] = ( + request, + payload_data, + update_prompt, + window_condition_len, + was_resumable, + ) if payload_finished: self.upstream_exhausted_requests.add(req_id) @@ -268,15 +546,26 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): else: prompt_token_ids = new_ids request.prompt_token_ids = prompt_token_ids + # Full-snapshot producers opt in explicitly; generation and + # diffusion models keep their existing incremental merge. prev_info = getattr(request, "additional_information", None) - info = dict(prev_info) if isinstance(prev_info, dict) else {} + replace_snapshot = meta.get("replace_runtime_additional_information") is True + info = {} if replace_snapshot else (dict(prev_info) if isinstance(prev_info, dict) else {}) for key, value in payload_data.items(): if key == "codes": - if use_tensor_codes and isinstance(value, dict): + if isinstance(value, dict): existing_sub = info.get(key) merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {} - merged_sub.update(value) - info[key] = merged_sub + for subkey, subvalue in value.items(): + # A 1-D audio tensor is represented by the + # placeholder prompt above, but sibling fields + # such as the reference voice still belong in + # the current runtime snapshot. + if subkey == "audio" and not use_tensor_codes: + continue + merged_sub[subkey] = subvalue + if merged_sub: + info[key] = merged_sub continue if isinstance(value, dict): existing_sub = info.get(key) @@ -313,6 +602,44 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): return False def _send_single_request(self, task: dict): + request = task["request"] + external_req_id = request.external_req_id + sender_token = task.get("sender_token") + if sender_token is None: + self._send_single_request_for_generation(task) + return + with self._sender_state_lock: + is_current = self._sender_tokens.get(external_req_id) is sender_token + if is_current and not sender_token.cancelled: + sender_token.in_flight = True + else: + is_current = False + if not is_current: + logger.debug("Discarding stale queued chunk for aborted request %s", external_req_id) + return + try: + self._send_single_request_for_generation(task, sender_token) + finally: + with self._sender_state_lock: + if self._sender_tokens.get(external_req_id) is sender_token: + sender_token.in_flight = False + if sender_token.cancelled: + self._sender_tokens.pop(external_req_id, None) + self._clear_sender_state_locked(external_req_id) + + def _sender_generation_is_active( + self, + external_req_id: str, + sender_token: _SenderGeneration, + ) -> bool: + with self._sender_state_lock: + return self._sender_tokens.get(external_req_id) is sender_token and not sender_token.cancelled + + def _send_single_request_for_generation( + self, + task: dict, + sender_token: _SenderGeneration | None = None, + ): raw_mm = task["multimodal_output"] multimodal_output = unflatten_payload(raw_mm) if isinstance(raw_mm, Mapping) else raw_mm request = task["request"] @@ -327,10 +654,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): payload_data: OmniPayloadStruct | None = None if self.custom_process_next_stage_input_func: try: - payload_data = self.custom_process_next_stage_input_func( - transfer_manager=self, - multimodal_output=multimodal_output, - request=request, + processor = self.custom_process_next_stage_input_func + processor_kwargs = { + "transfer_manager": self, + "multimodal_output": multimodal_output, + "request": request, # Existing processors use is_finished as a flush signal. # Terminal stops no longer count as segment boundaries # (is_segment_finished is False when the request finishes, @@ -338,8 +666,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # accumulated tail on the terminal chunk — otherwise the # downstream stage receives the finished marker without # the final payload (#5413). - is_finished=is_segment_finished or is_finished, - ) + "is_finished": is_segment_finished or is_finished, + } + if self._accepts_new_token_ids(processor): + processor_kwargs["new_token_ids"] = task.get("new_token_ids", ()) + payload_data = processor(**processor_kwargs) except Exception as e: logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}") @@ -352,10 +683,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): payload_data = OmniPayloadStruct() if payload_data.meta is None: payload_data.meta = MetaStruct() - payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool) + # A model-native codec stream may span many resumable scheduler + # requests. Its processor is the only layer that knows whether a + # request finish is a transport wake or the lifetime of the codec + # state itself. Preserve an explicit streaming lifetime decision; + # all ordinary producers keep the scheduler-derived default. + processor_controls_codec_lifetime = payload_data.meta.finished is not None and self._is_truthy_scalar( + payload_data.meta.codec_streaming + ) + if not processor_controls_codec_lifetime: + payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool) if payload_data.meta.is_segment_finished is None: payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool) + if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): + logger.debug("Skipping cancelled chunk for request %s before connector put", external_req_id) + return + success, size, metadata = self.connector.put( from_stage=str(stage_id), to_stage=str(next_stage_id), @@ -363,6 +707,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): data=payload_data, ) + if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): + # cleanup_sender() may cancel this generation while put() is + # blocked. Do not let that stale completion recreate state for a + # request whose cleanup is already in progress. + logger.debug("Ignoring completed put for cancelled request %s", external_req_id) + return + if success: self.put_req_chunk[external_req_id] += 1 self.ramp_chunk_count[external_req_id] += 1 @@ -379,15 +730,36 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): elif finished_flag is not None: is_payload_finished = bool(finished_flag) - # Reclaim per-request async state only after the terminal payload - # has been sent successfully. This avoids cleanup->save races. - if is_payload_finished: + # Processor-controlled codec ``finished`` can mark the end of one + # native-duplex turn while the persistent request remains live. + # Only the scheduler's whole-request terminal may reclaim shared + # sender/receiver state; otherwise a late old-turn put can erase a + # newly admitted turn for the same request. + if is_payload_finished and is_finished: self.cleanup(request.request_id, external_req_id) + else: + # R1.2 of #4855. connector.put returning False is a silent drop: no + # exception, no retry, and the caller's `if success:` block simply + # does not run. /dev/shm exhaustion in SharedMemoryConnector.put is + # one way to get here. Record it so the request fails now rather + # than parking in WAITING_FOR_CHUNK until the deadline. + logger.error( + "Chunk send failed for %s (stage %s -> %s); giving up on this chunk", + external_req_id, + stage_id, + next_stage_id, + ) + # Key on the scheduler-side id. `external_req_id` is the user-facing id + # (InputProcessor renames request_id to an internal UUID and keeps the + # original in external_req_id, see async_omni_engine.py), while + # `self.requests` -- and therefore `finish_requests` -- is keyed by the + # internal one. + self.record_send_failure(request.request_id, "connector.put reported failure") if is_segment_finished: self.code_prompt_token_ids.pop(external_req_id, None) - self.requests_num_chunks_sent.pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) + self._adaptive_states.pop(external_req_id, None) cached_ic = getattr(self, "_cached_ic", None) if cached_ic is not None: cached_ic.pop(external_req_id, None) @@ -409,9 +781,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): def cleanup_receiver(self, request_id: str) -> None: """Reclaim receiver-side per-request state (keyed by internal id). - Safe to call from the scheduler even when ``save_async()`` has - enqueued work that the background thread has not yet processed, - because it only touches receiver-side dictionaries. + Safe to call from the scheduler while ``connector.get()`` is in + flight: cleanup and the post-get state commit are serialized, and a + cancelled request cannot publish the late chunk. Must also purge the request from the chunk-parking deques (``waiting_for_chunk_waiting_requests`` / ``_running_requests`` / @@ -425,19 +797,30 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): Idempotent: calling with an already-cleaned or unknown id is safe. """ + with self._receiver_state_lock: + self._clear_receiver_state_locked(request_id) + + def _clear_receiver_state_locked(self, request_id: str) -> None: + """Clear receiver state while ``_receiver_state_lock`` is held.""" self._active_streams.pop(request_id, None) self.upstream_exhausted_requests.discard(request_id) self.segment_finished_requests.discard(request_id) self.get_req_chunk.pop(request_id, None) self.requests_with_ready_chunks.discard(request_id) + self.replaced_streaming_prompt_ids.discard(request_id) self.request_ids_mapping.pop(request_id, None) self.requests_origin_status.pop(request_id, None) + self._pending_ar_prompt_updates.pop(request_id, None) + self._streaming_condition_lengths.pop(request_id, None) + self._streaming_condition_seqs.pop(request_id, None) + self._registered_load_entries.pop(request_id, None) self._discard_from_chunk_deque(self.waiting_for_chunk_waiting_requests, request_id) self._discard_from_chunk_deque(self.waiting_for_chunk_running_requests, request_id) self._discard_from_chunk_deque(self._held_non_active, request_id) self._cancelled_load_reqs.add(request_id) self._finished_load_reqs.discard(request_id) + self._waiting_since.pop(request_id, None) @staticmethod def _discard_from_chunk_deque(deque_list: deque[Any], request_id: str) -> None: @@ -451,16 +834,36 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): def cleanup_sender(self, external_req_id: str) -> None: """Reclaim sender-side per-request state (keyed by external id). - Must only be called after the terminal chunk has actually been - sent (i.e. from ``_send_single_request``), not before. + Called after a terminal chunk is sent or when the scheduler aborts the + request before a terminal chunk can be produced. In-flight sends are + cancelled here and reclaim their own state in ``finally``; cleanup + never waits for connector I/O on the scheduler thread. Idempotent: calling with an already-cleaned or unknown id is safe. """ + with self._sender_state_lock: + sender_token = self._sender_tokens.get(external_req_id) + if sender_token is None: + self._clear_sender_state_locked(external_req_id) + return + sender_token.cancelled = True + if sender_token.in_flight: + # The sender's finally block is the first point where no + # in-flight code can mutate this request's state. Let it own + # cleanup without making the scheduler wait for connector I/O. + return + self._sender_tokens.pop(external_req_id, None) + self._clear_sender_state_locked(external_req_id) + + def _clear_sender_state_locked(self, external_req_id: str) -> None: + """Clear sender state while ``_sender_state_lock`` is held.""" self.put_req_chunk.pop(external_req_id, None) self.request_payload.pop(external_req_id, None) self.code_prompt_token_ids.pop(external_req_id, None) self.requests_num_chunks_sent.pop(external_req_id, None) + self._segment_generation.pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) + self._adaptive_states.pop(external_req_id, None) self._pending_streaming_prefills.pop(external_req_id, None) cached_ic = getattr(self, "_cached_ic", None) @@ -538,6 +941,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): RequestStatus.RUNNING, self._finished_load_reqs, ) + self._apply_pending_ar_prompt_updates(scheduler_requests) + self._requeue_replaced_prompts(waiting_queue, running_queue) while len(running_queue) > self.scheduler_max_num_seqs: request = running_queue.pop() request.status = RequestStatus.PREEMPTED @@ -552,9 +957,80 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._process_chunk_queue( running_queue, self.waiting_for_chunk_running_requests, RequestStatus.RUNNING, self._finished_load_reqs ) + self._apply_pending_ar_prompt_updates(scheduler_requests) + self._requeue_replaced_prompts(waiting_queue, running_queue) self._promote_active_streams(waiting_queue) self._preempt_non_active_running(waiting_queue, running_queue) + def _apply_pending_ar_prompt_updates(self, scheduler_requests: dict[str, Request] | None) -> None: + """Finalize ready AR prompt updates on the scheduler thread.""" + for request_id in tuple(self.requests_with_ready_chunks): + with self._receiver_state_lock: + pending = self._pending_ar_prompt_updates.pop(request_id, None) + if pending is None: + continue + ( + fallback_request, + payload_data, + update_prompt, + window_condition_len, + was_resumable, + ) = pending + request = scheduler_requests.get(request_id) if scheduler_requests is not None else fallback_request + if request is None: + continue + + request.additional_information = payload_data + meta = payload_data.get("meta") + meta = meta if isinstance(meta, dict) else {} + previous_condition_seq = self._streaming_condition_seqs.get(request_id) + condition_seq = previous_condition_seq + 1 if previous_condition_seq is not None else 0 + is_window_condition = window_condition_len is not None + if is_window_condition: + meta["streaming_condition_seq"] = condition_seq + if is_window_condition: + update_prompt = bool( + was_resumable + and (previous_condition_seq is not None or meta.get("replace_streaming_prompt") is True) + ) + try: + if update_prompt: + replaced = construct_next_stage_streaming_input_prompt( + payload_data, + request, + max_model_len=self._max_model_len, + previous_condition_len=self._streaming_condition_lengths.get(request_id), + previous_condition_seq=previous_condition_seq, + condition_seq=condition_seq, + recompute_previous_chunks=self._streaming_prompt_previous_chunks, + ) + if replaced: + self.replaced_streaming_prompt_ids.add(request_id) + except ValueError as exc: + # The connector chunk was consumed, so retrying would skip + # this transition and desynchronize Talker. Let the + # scheduler fail it. + self.record_receive_failure(request_id, str(exc)) + continue + + if window_condition_len is not None: + self._streaming_condition_lengths[request_id] = window_condition_len + self._streaming_condition_seqs[request_id] = condition_seq + + def _requeue_replaced_prompts(self, waiting_queue: Any, running_queue: list[Request]) -> None: + """Move a replaced running prompt back through scheduler admission.""" + for request in list(running_queue): + request_id = request.request_id + if ( + request_id not in self.replaced_streaming_prompt_ids + or request_id not in self.requests_with_ready_chunks + ): + continue + running_queue.remove(request) + request.status = RequestStatus.WAITING + self.requests_origin_status[request_id] = RequestStatus.WAITING + waiting_queue.add_request(request) + def _promote_active_streams(self, queue: Any) -> None: if len(self._active_streams) >= self._active_window: return @@ -579,6 +1055,42 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._active_streams[request_id] = request return True + def collect_timed_out_request_ids(self, timeout_s: float) -> set[str]: + """Return IDs whose chunk wait has exceeded *timeout_s*. + + The async-chunk path had no deadline of any kind: a request parks in + ``WAITING_FOR_CHUNK`` and the receiver re-queues it on every failed + poll (``transfer_adapter/base.py``) with no attempt counter, so a + producer that crashed, gave up after its send retries, or simply never + emitted a terminal chunk left the request waiting forever + (vllm-project/vllm-omni#3833). The full-payload path has had a net + since ``OmniSchedulingCoordinator.collect_timed_out_request_ids``; this + is its async-chunk counterpart, and both are driven by the same + ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``. + + The clock measures *stall* time, not stream lifetime: it starts when a + request begins waiting for a chunk and resets each time one arrives, so + a long but healthy stream is never failed. + + Clears ``_waiting_since`` for the expired IDs. The caller marks them + ``FINISHED_ERROR`` via the scheduler's ``finish_requests``, which routes + back through this adapter's own ``finish_requests`` and releases the + rest of the per-request state. + """ + if timeout_s <= 0 or not self._waiting_since: + return set() + now = time.monotonic() + timed_out_ids = {req_id for req_id, started in self._waiting_since.items() if now - started > timeout_s} + for req_id in timed_out_ids: + self._waiting_since.pop(req_id, None) + logger.warning( + "[Stage-%s] Request %s timed out waiting for a chunk (stalled > %.0fs)", + self.connector.stage_id, + req_id, + timeout_s, + ) + return timed_out_ids + @property def num_running_waiting_for_chunk(self) -> int: """Count running requests temporarily removed while awaiting a chunk.""" @@ -624,11 +1136,15 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # Requests that waiting for chunk self.load_async(request) request.status = RequestStatus.WAITING_FOR_CHUNK + self._waiting_since.setdefault(request.request_id, time.monotonic()) else: if request.request_id in finished_load_reqs: request.status = target_status finished_load_reqs.remove(request.request_id) self.requests_with_ready_chunks.add(request.request_id) + # A chunk landed: restart the clock for the next one, so the + # deadline measures stall time rather than total stream time. + self._waiting_since.pop(request.request_id, None) continue queue.remove(request) self.requests_origin_status[request.request_id] = target_status @@ -767,11 +1283,15 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # Requests that waiting for chunk self.load_async(request) request.status = RequestStatus.WAITING_FOR_CHUNK + self._waiting_since.setdefault(request.request_id, time.monotonic()) else: if request.request_id in finished_load_reqs: request.status = target_status finished_load_reqs.remove(request.request_id) self.requests_with_ready_chunks.add(request.request_id) + # A chunk landed: restart the clock for the next one, so the + # deadline measures stall time rather than total stream time. + self._waiting_since.pop(request.request_id, None) continue queue.remove(request) self.requests_origin_status[request.request_id] = target_status @@ -782,6 +1302,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): for req_data in scheduler_output.scheduled_new_reqs: if req_data.req_id in self.requests_with_ready_chunks: self.requests_with_ready_chunks.remove(req_data.req_id) + if req_data.req_id in self.replaced_streaming_prompt_ids: + external_req_id = self.request_ids_mapping.get(req_data.req_id) + if external_req_id is not None: + self.requests_num_chunks_sent.pop(external_req_id, None) + self.replaced_streaming_prompt_ids.remove(req_data.req_id) if scheduler_output.scheduled_cached_reqs: for req_id in scheduler_output.scheduled_cached_reqs.req_ids: @@ -799,16 +1324,45 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): else: request_ids = requests.keys() + connector_owned_ids = { + request.request_id + for queue in ( + self.waiting_for_chunk_waiting_requests, + self.waiting_for_chunk_running_requests, + self._held_non_active, + ) + for request in queue + } + # First pass: collect requests to remove from queues + request_ids = set(request_ids) for req_id in request_ids: request = requests.get(req_id) if requests else None - if request is None or request.is_finished(): + if request is None: # Invalid request ID. continue - if req_id in self.requests_origin_status: + resumable_segment_stop = bool( + getattr(request, "resumable", False) and request.status == RequestStatus.FINISHED_STOPPED + ) + if request.is_finished() and not resumable_segment_stop: + continue + # Once restored to a scheduler queue, the saved origin is stale and + # must not overwrite statuses such as WAITING_FOR_STREAMING_REQ. + if req_id in self.requests_origin_status and req_id in connector_owned_ids: request.status = self.requests_origin_status.pop(req_id) - request_ids = set(request_ids) + # An abort can terminate a long-lived native codec stream before it + # emits a terminal payload. Reclaim both sides of the adapter so a + # later request reusing the same external id starts with empty state. + for req_id in request_ids: + request = requests.get(req_id) if requests else None + external_req_id = ( + getattr(request, "external_req_id", None) + if request is not None + else self.request_ids_mapping.get(req_id) + ) or req_id + self.cleanup_receiver(req_id) + self.cleanup_sender(external_req_id) self.waiting_for_chunk_waiting_requests = deque( request for request in self.waiting_for_chunk_waiting_requests if request.request_id not in request_ids @@ -821,10 +1375,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): ) for req_id in request_ids: - self._active_streams.pop(req_id, None) - self.requests_with_ready_chunks.discard(req_id) - self.upstream_exhausted_requests.discard(req_id) - self._finished_load_reqs.discard(req_id) - self._cancelled_load_reqs.add(req_id) + self.cleanup_receiver(req_id) return [] diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a3aa96f60..0eb18c686 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,9 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import argparse import json import os import tempfile +from collections.abc import Callable from dataclasses import dataclass, field, fields -from typing import Any +from typing import Any, cast from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.logger import init_logger @@ -12,6 +16,7 @@ from vllm_omni.config import OmniModelConfig from vllm_omni.outputs.output_modality import OutputModality from vllm_omni.platforms import current_omni_platform from vllm_omni.plugins import load_omni_general_plugins +from vllm_omni.worker.omni_connector_validation import validate_worker_omni_connector logger = init_logger(__name__) @@ -158,6 +163,7 @@ class OmniEngineArgs(EngineArgs): If None, default processing is used. stage_connector_spec: Extra configuration for stage connector async_chunk: If set to True, perform async chunk + session_mode: Request lifecycle mode, either turn-based or duplex worker_type: Model Type, e.g., "ar" or "generation" task_type: Model-defined startup task type. Consumers validate the supported values and decide whether it selects request behavior, @@ -167,9 +173,6 @@ class OmniEngineArgs(EngineArgs): Required when single-stage mode is active. omni_master_port: TCP port for the OmniMasterServer registration socket. Required when single-stage mode is active. - stage_configs_path: Optional path to a JSON/YAML file containing - stage configurations for the multi-stage pipeline. If None, - stage configs are resolved from the model's default configuration. output_modalities: Optional list of output modality names to enable (e.g. ["text", "audio"]). If None, all modalities supported by the model are used. @@ -185,9 +188,12 @@ class OmniEngineArgs(EngineArgs): engine_output_type: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None + requires_full_payload_input: bool = False stage_connector_spec: dict[str, Any] = field(default_factory=dict) subtalker_sampling_params: dict[str, Any] | None = None + silence_ban_frames: int = 0 async_chunk: bool = False + session_mode: str = "turn" retains_state_across_chunks: bool = False # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. # Must be declared here so engine_args dict propagation does not silently @@ -231,19 +237,26 @@ class OmniEngineArgs(EngineArgs): omni_dp_size_local: int = 1 omni_lb_policy: str = "random" omni_heartbeat_timeout: float = 30.0 - stage_configs_path: str | None = None output_modalities: list[str] | None = None log_stats: bool = False custom_pipeline_args: dict[str, Any] | None = None has_sampling_extra_args: bool = False + sampling_extra_args_keys: tuple[str, ...] = () def __post_init__(self) -> None: - if self.worker_cls is None: + if self.worker_cls in (None, "auto"): if self.worker_type == "ar": self.worker_cls = current_omni_platform.get_omni_ar_worker_cls() elif self.worker_type == "generation": self.worker_cls = current_omni_platform.get_omni_generation_worker_cls() load_omni_general_plugins() + + connector_extra = self.stage_connector_spec.get("extra") + connector_role = connector_extra.get("role") if isinstance(connector_extra, dict) else None + needs_connector = bool( + self.requires_full_payload_input or self.custom_process_next_stage_input_func or connector_role is not None + ) + validate_worker_omni_connector(self.worker_cls, needs_connector) super().__post_init__() def _ensure_omni_models_registered(self): @@ -300,24 +313,26 @@ class OmniEngineArgs(EngineArgs): } stage_connector_config["extra"]["stage_id"] = self.stage_id + hf_overrides = cast(dict[str, Any] | Callable[[Any], Any] | None, getattr(self, "hf_overrides", None)) # If model_arch is specified, inject it into hf_overrides so vLLM can # resolve the architecture even when config.json lacks 'architectures'. # Also inject model_type so AutoConfig can resolve the correct config # class for models with empty or missing config.json (e.g. CosyVoice3). if self.model_arch: - if self.hf_overrides is None: - self.hf_overrides = {} - if isinstance(self.hf_overrides, dict): - self.hf_overrides.setdefault("architectures", [self.model_arch]) - if "model_type" not in self.hf_overrides: + if hf_overrides is None: + hf_overrides = {} + self.hf_overrides = hf_overrides + if isinstance(hf_overrides, dict): + hf_overrides.setdefault("architectures", [self.model_arch]) + if "model_type" not in hf_overrides: model_type = _ARCH_TO_MODEL_TYPE.get(self.model_arch) if model_type is not None: - self.hf_overrides.setdefault("model_type", model_type) + hf_overrides.setdefault("model_type", model_type) # Stage wrappers (e.g. Code2Wav) may need max_model_len larger # than the base checkpoint's text max_position_embeddings. if self.model_arch == "Qwen3TTSCode2Wav" and self.max_model_len is not None: - self.hf_overrides.setdefault("talker_config", {}).setdefault( + hf_overrides.setdefault("talker_config", {}).setdefault( "max_position_embeddings", int(self.max_model_len) ) @@ -333,33 +348,37 @@ class OmniEngineArgs(EngineArgs): if model_type is not None: self._patch_empty_hf_config(model_type) + tokenizer = cast(str | None, getattr(self, "tokenizer", None)) + # NemotronVoiceChat: the checkpoint's only tokenizer subfolder is the # auxiliary rnnt_tokenizer/ (ASR loss vocab), which the generic # subfolder auto-detect below would wrongly pick up. The text tokenizer # is the Nemotron backbone's, resolved from a local override or its # HF id (cache-friendly). - if not self.tokenizer and self.model_arch in ( + if not tokenizer and self.model_arch in ( "NemotronVoiceChatThinkerForConditionalGeneration", "NemotronVoiceChatTalkerForConditionalGeneration", "NemotronVoiceChatCode2Wav", ): self.tokenizer = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or "nvidia/NVIDIA-Nemotron-Nano-9B-v2" - logger.info("NemotronVoiceChat: using text tokenizer from %s", self.tokenizer) + tokenizer = self.tokenizer + logger.info("NemotronVoiceChat: using text tokenizer from %s", tokenizer) # Auto-detect tokenizer for models that store it in a subdirectory # rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/). - if not self.tokenizer and self.model: + if not tokenizer and self.model: model_path = self.model if os.path.isdir(model_path) and not os.path.isfile(os.path.join(model_path, "tokenizer_config.json")): for subfolder in sorted(os.listdir(model_path)): candidate = os.path.join(model_path, subfolder) if os.path.isdir(candidate) and os.path.isfile(os.path.join(candidate, "tokenizer_config.json")): self.tokenizer = candidate + tokenizer = candidate logger.info("Auto-detected tokenizer at %s", candidate) break elif not os.path.isdir(model_path): - subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch) - if subfolder: + tokenizer_subfolder = _TOKENIZER_SUBFOLDER_MAP.get(self.model_arch or "") + if tokenizer_subfolder: # Download just the tokenizer files from the subfolder try: from huggingface_hub import snapshot_download @@ -367,18 +386,18 @@ class OmniEngineArgs(EngineArgs): local_dir = snapshot_download( model_path, allow_patterns=[ - f"{subfolder}/tokenizer*", - f"{subfolder}/tokenization_*", - f"{subfolder}/special_tokens*", - f"{subfolder}/vocab*", - f"{subfolder}/merges*", - f"{subfolder}/added_tokens*", + f"{tokenizer_subfolder}/tokenizer*", + f"{tokenizer_subfolder}/tokenization_*", + f"{tokenizer_subfolder}/special_tokens*", + f"{tokenizer_subfolder}/vocab*", + f"{tokenizer_subfolder}/merges*", + f"{tokenizer_subfolder}/added_tokens*", ], ) - candidate = os.path.join(local_dir, subfolder) + candidate = os.path.join(local_dir, tokenizer_subfolder) if os.path.isdir(candidate): self.tokenizer = candidate - logger.info("Downloaded tokenizer from %s/%s", model_path, subfolder) + logger.info("Downloaded tokenizer from %s/%s", model_path, tokenizer_subfolder) except Exception as e: logger.warning("Failed to download tokenizer subfolder: %s", e) @@ -398,6 +417,7 @@ class OmniEngineArgs(EngineArgs): # All kwargs below are Omni specific stage_id=self.stage_id, async_chunk=self.async_chunk, + session_mode=self.session_mode, retains_state_across_chunks=self.retains_state_across_chunks, active_stream_window=self.active_stream_window, duplex_max_sessions=self.duplex_max_sessions, @@ -408,11 +428,14 @@ class OmniEngineArgs(EngineArgs): engine_output_type=self.engine_output_type, hf_config_name=self.hf_config_name, custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, + requires_full_payload_input=self.requires_full_payload_input, stage_connector_config=stage_connector_config, subtalker_sampling_params=self.subtalker_sampling_params, + silence_ban_frames=self.silence_ban_frames, omni_kv_config=self.omni_kv_config, task_type=self.task_type, has_sampling_extra_args=self.has_sampling_extra_args, + sampling_extra_args_keys=tuple(self.sampling_extra_args_keys or ()), ) return omni_config @@ -490,7 +513,6 @@ class OrchestratorArgs: ray_address: str | None = None # === Config Files === - stage_configs_path: str | None = None deploy_config: str | None = None stage_overrides: str | None = None # raw JSON string; parsed downstream # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays @@ -526,6 +548,7 @@ class OrchestratorArgs: diffusers_call_kwargs: str = "{}" ulysses_degree: int | None = None ulysses_mode: str = "strict" + ulysses_a2a_permute: bool | None = None ring_degree: int | None = None allgather_degree: int | None = None diffusion_quantization_config: str | None = None @@ -533,6 +556,7 @@ class OrchestratorArgs: hsdp_shard_size: int = -1 hsdp_replicate_size: int = 1 diffusion_attention_backend: str | None = None + fastvideo_vsa_topk: int | None = None diffusion_attention_config: str | None = None diffusion_compile_granularity: str | None = None diffusion_compile_dynamic: bool | None = None @@ -549,6 +573,9 @@ class OrchestratorArgs: enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True dlo_resident_layers: int = 0 + host_weight_runtime_mode: str = "disabled" + host_weight_runtime_root: str | None = None + dlo_host_registration_limit_gib: float = 0.0 boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None @@ -557,7 +584,6 @@ class OrchestratorArgs: cfg_parallel_size: int = 1 vae_patch_parallel_size: int = 1 vae_parallel_mode: str = "tile" - text_encoder_tp_size: int = 1 default_sampling_params: str | None = None max_generated_image_size: int | None = None tts_max_instructions_length: int | None = None @@ -582,7 +608,6 @@ SHARED_FIELDS: frozenset[str] = frozenset( "model", # orch: detect model_type; engine: load weights "stage_id", # orch: route (headless); engine: identity "log_stats", # both want the flag - "stage_configs_path", # orch: load legacy YAML; engine: may reference for validation "async_chunk", # orch: read from CLI, redistribute; engine: per-stage flag "tokenizer", # orch: detect model type; engine: tokenization } diff --git a/vllm_omni/engine/async_engine_utils.py b/vllm_omni/engine/async_engine_utils.py index 6f544239b..fbe4632ab 100644 --- a/vllm_omni/engine/async_engine_utils.py +++ b/vllm_omni/engine/async_engine_utils.py @@ -24,6 +24,24 @@ logger = init_logger(__name__) SHUTDOWN_ENQUEUE_TIMEOUT_S = 1.0 SHUTDOWN_JOIN_TIMEOUT_S = 30.0 _WEAK_SHUTDOWN_JOIN_TIMEOUT_S = 1.0 +_JANUS_SYNC_QUEUE_SHUTDOWN = getattr(janus, "SyncQueueShutDown", None) +_LEGACY_JANUS_QUEUE_CLOSED_MESSAGE = "Operation on the closed queue is forbidden" +_RPC_RESULT_ROUTER_CLOSED_MESSAGES = { + "RPC result router closed", + "RPC result router is closed", +} + + +def is_janus_sync_queue_shutdown(exc: Exception) -> bool: + if _JANUS_SYNC_QUEUE_SHUTDOWN is not None: + return isinstance(exc, _JANUS_SYNC_QUEUE_SHUTDOWN) + return isinstance(exc, RuntimeError) and str(exc) == _LEGACY_JANUS_QUEUE_CLOSED_MESSAGE + + +def is_abort_transport_shutdown(exc: Exception) -> bool: + return is_janus_sync_queue_shutdown(exc) or ( + isinstance(exc, RuntimeError) and str(exc) in _RPC_RESULT_ROUTER_CLOSED_MESSAGES + ) def inject_global_id(target: Any, request_id: str) -> None: diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index aa284067b..1efafbbab 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Async Omni Engine for vLLM-Omni multi-stage runtime. @@ -9,9 +12,10 @@ from __future__ import annotations import asyncio import concurrent.futures -import dataclasses +import copy import json import queue +import shutil import threading import time import uuid @@ -24,7 +28,6 @@ import janus import torch from omegaconf import OmegaConf from vllm import envs as vllm_envs -from vllm.engine.arg_utils import EngineArgs from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest @@ -34,32 +37,42 @@ from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remot from vllm_omni.config.stage_config import ( DuplexSessionRuntimeConfig, load_deploy_config, - strip_parent_engine_args, ) -from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config -from vllm_omni.diffusion.diffusion_engine import supports_audio_output +from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY, TRANSFORM_OWNED_META_KEYS +from vllm_omni.diffusion.data import ( + DiffusionParallelConfig, + parse_attention_config, + resolve_model_class_name, +) +from vllm_omni.diffusion.io_support import get_diffusion_output_type +from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_engine_utils import ( SHUTDOWN_ENQUEUE_TIMEOUT_S, SHUTDOWN_JOIN_TIMEOUT_S, apply_omni_final_stage_metadata, enqueue_orchestrator_shutdown, inject_global_id, + is_abort_transport_shutdown, + is_janus_sync_queue_shutdown, shutdown_runtime_after_orchestrator, upgrade_to_omni_request, weak_shutdown_async_omni_engine, ) from vllm_omni.engine.messages import ( AbortRequestMessage, + AbortResultMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, InteractionMessage, + OutputMessage, StageSubmissionMessage, ) from vllm_omni.engine.orchestrator import Orchestrator from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.engine.stage_client import StageClient from vllm_omni.engine.stage_init_utils import build_stage0_input_processor from vllm_omni.engine.stage_pool import StagePool @@ -69,7 +82,6 @@ from vllm_omni.engine.stage_runtime import ( ) from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import ( - is_new_format_deploy_config, load_and_resolve_stage_configs, parse_stage_overrides, ) @@ -85,39 +97,6 @@ if TYPE_CHECKING: _STARTUP_POLL_INTERVAL_S = 1.0 _REQUEST_QUEUE_MAXSIZE = 256 -# ============================================================================ -# Parent-EngineArgs field-routing contracts (consumed by -# AsyncOmniEngine._strip_parent_engine_args when ``stage_configs_path`` is set). -# ============================================================================ - -# Fields that must survive the "equal to default → strip" filter because -# diffusion stages need them even when equal to vllm's default value -# (e.g. colocate worker setup relies on worker_extension_cls being forwarded). -_PARENT_ARGS_KEEP: frozenset[str] = frozenset( - { - "worker_extension_cls", - "allowed_local_media_path", - "allowed_media_domains", - # Legacy stage-config YAMLs may intentionally leave parallel or - # distributed knobs unspecified at the stage level and rely on - # top-level CLI values to fill them in during the per-stage merge. - # Keep these fields so stages that omit them can inherit CLI values, - # while stages with explicit YAML values still win because the legacy - # stage-config loader prefers stage-local engine args. - "tensor_parallel_size", - } -) - -# Omni orchestrator-level fields consumed by ``_resolve_stage_configs`` that -# must never leak into per-stage EngineArgs (``stage_configs_path`` would -# trigger the ``create_model_config`` guard). -_PARENT_ARGS_STRIP: frozenset[str] = frozenset({"stage_configs_path"}) - - -# Fields always populated by callers (via ``from_cli_args`` / ``asdict``) so -# their presence as an override is never a surprise — suppress the -# "override ignored" warning for these. -_PARENT_ARGS_NO_WARN: frozenset[str] = frozenset({"model"}) class AsyncOmniEngine: @@ -140,16 +119,19 @@ class AsyncOmniEngine: _log_stats: bool = False _coordinator_runtime: Any = None _transfer_emitter: Any = None + _prom_metrics: Any = None _enable_orch_monitor: bool = False + # Lazily created by get_output_blocking_async(). + _output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None def __init__( self, model: str, stage_init_timeout: int = 300, init_timeout: int = 600, - diffusion_batch_size: int = 1, single_stage_mode: bool = False, transfer_emitter: Any = None, + prom_metrics: Any = None, log_stats: bool = False, tokenizer: str | None = None, trust_remote_code: bool | None = None, @@ -157,7 +139,6 @@ class AsyncOmniEngine: ) -> None: self.model = model self.tokenizer = tokenizer - self.diffusion_batch_size = diffusion_batch_size # Cached by get_diffusion_od_config(). self._diffusion_od_config_view: Any = None startup_timeout = int(init_timeout) @@ -166,6 +147,7 @@ class AsyncOmniEngine: # Optional: when None, Orchestrator silently skips TX emit (existing # RX path still works via OrchestratorAggregator). self._transfer_emitter = transfer_emitter + self._prom_metrics = prom_metrics # Drives upstream EngineCore + scheduler stats production. When False # the engine skips SchedulerStats / IterationStats; the per-(stage, # replica) vllm:* wrap stays registered but reads zero. Respects the @@ -210,10 +192,8 @@ class AsyncOmniEngine: self._omni_master_port, ) - # Stage resolution pops deploy_config, so get pipeline-wide settings - # beforehand. The stage CLI exposes the same deploy YAML through - # stage_configs_path. - deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path") + # Stage resolution pops deploy_config, so get pipeline-wide settings beforehand. + deploy_config_path = kwargs.get("deploy_config") # ``trust_remote_code`` is tri-state (bool | None): ``None`` means "not # specified" so stage-config resolution can defer to the deploy yaml's # per-stage value (see ``with_trust_remote_code_override``). The @@ -254,6 +234,7 @@ class AsyncOmniEngine: self.stage_pools: list[StagePool] = [] self.stage_clients: list[StageClient] = [] # logical-stage view for external readers self.input_processor: InputProcessor | None = None + self.prompt_transform_func: Any | None = None self.prompt_expand_func: Any | None = None self.supported_tasks: tuple[str, ...] = ("generate",) self.default_sampling_params_list: list[OmniSamplingParams] = [] @@ -271,6 +252,7 @@ class AsyncOmniEngine: self._correlated_rpc_client: CorrelatedRpcClient | None = None self._duplex_control_client: DuplexControlClient | None = None self._running_counter = OmniRequestCounter() + self._engines_waiting_counter = OmniRequestCounter() logger.info(f"[AsyncOmniEngine] Launching Orchestrator thread with {self.num_stages} stages") @@ -336,7 +318,6 @@ class AsyncOmniEngine: config_path=self.config_path, single_stage_mode=self.single_stage_mode, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=self.diffusion_batch_size, async_chunk=self.async_chunk, tokenizer=self.tokenizer, single_stage_id_filter=self._single_stage_id_filter, @@ -362,6 +343,9 @@ class AsyncOmniEngine: if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None else None ) + self.prompt_transform_func = ( + getattr(self.stage_clients[0], "prompt_transform_func", None) if self.stage_clients else None + ) self.prompt_expand_func = next( ( getattr(client, "prompt_expand_func", None) @@ -414,9 +398,13 @@ class AsyncOmniEngine: getattr(self, "_duplex_runtime_extension_path", None) ) if duplex_runtime_extension is not None: + stage_clients = [] + for pool in self.stage_pools: + assert pool.stage_client is not None + stage_clients.append(pool.stage_client) validate_duplex_runtime_extension( duplex_runtime_extension, - sampling_defaults=tuple(pool.stage_client.default_sampling_params for pool in self.stage_pools), + sampling_defaults=tuple(client.default_sampling_params for client in stage_clients), ) orchestrator = Orchestrator( @@ -428,7 +416,9 @@ class AsyncOmniEngine: pd_config=pd_config, membership_controller=membership_controller, running_counter=self._running_counter, + engines_waiting_counter=self._engines_waiting_counter, transfer_emitter=self._transfer_emitter, + prom_metrics=self._prom_metrics, log_stats=self._log_stats, enable_orch_monitor=self._enable_orch_monitor, duplex_runtime_extension=duplex_runtime_extension, @@ -704,14 +694,45 @@ class AsyncOmniEngine: # Keep the original prompt for downstream stages (they need the raw # dict, e.g. for multi_modal_data). + if isinstance(prompt, dict): + raw_info = prompt.get("additional_information") + if isinstance(raw_info, dict): + raw_meta = raw_info.get("meta") + if isinstance(raw_meta, dict): + for key in TRANSFORM_OWNED_META_KEYS: + raw_meta.pop(key, None) original_prompt = prompt preselected_stage0_replica: int | None = None + request_artifact_dirs: list[str] = [] stage_type = self.stage_metadata[0].stage_type output_prompt_text: Any = None _preprocess_ms = 0.0 if stage_type != "diffusion" and not isinstance(prompt, EngineCoreRequest): - # Inject global_request_id into the raw prompt. + # Stage transforms and downstream stages must share the same + # request identity, including when the transform replaces the + # prompt object. + if isinstance(prompt, dict): + inject_global_id(prompt, request_id) + elif isinstance(prompt, list): + for item in prompt: + inject_global_id(item, request_id) + + prompt_transform_func = getattr(self, "prompt_transform_func", None) + if prompt_transform_func is not None: + if isinstance(prompt, dict): + prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) + prompt = prompt_transform_func( + copy.copy(prompt), + effective_sampling_params_list, + ) + if isinstance(prompt, dict): + raw_dirs = prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) + if isinstance(raw_dirs, list) and all(isinstance(path, str) for path in raw_dirs): + request_artifact_dirs = list(raw_dirs) + if isinstance(original_prompt, dict): + original_prompt[REQUEST_ARTIFACT_DIRS_KEY] = request_artifact_dirs + if isinstance(prompt, dict): inject_global_id(prompt, request_id) elif isinstance(prompt, list): @@ -724,6 +745,7 @@ class AsyncOmniEngine: ) # Full input processing (tokenization, multimodal, etc.) + assert self.input_processor is not None _t_preprocess = time.perf_counter() try: request = self.input_processor.process_inputs( @@ -742,12 +764,25 @@ class AsyncOmniEngine: except Exception: if preselected_stage0_replica is not None and self.stage_pools: self.stage_pools[0].release_binding(request_id) + for artifact_dir in request_artifact_dirs: + shutil.rmtree(artifact_dir, ignore_errors=True) raise _preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0 # TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have # additional_information field in the prompt. request = upgrade_to_omni_request(request, prompt) + if isinstance(request, OmniEngineCoreRequest) and request.additional_information is not None: + processed_info = deserialize_additional_information(request.additional_information) + processed_meta = processed_info.get("meta") + if isinstance(processed_meta, dict): + if isinstance(original_prompt, dict): + original_info = dict(original_prompt.get("additional_information") or {}) + original_meta = dict(original_info.get("meta") or {}) + original_meta.update(processed_meta) + original_info["meta"] = original_meta + original_prompt["additional_information"] = original_info + if reasoning_ended is not None: request.reasoning_ended = reasoning_ended @@ -767,6 +802,8 @@ class AsyncOmniEngine: if output_prompt_text is None and isinstance(original_prompt, dict): output_prompt_text = original_prompt.get("prompt") prompt = request + else: + request_artifact_dirs = [] return StageSubmissionMessage( type=message_type, @@ -780,25 +817,36 @@ class AsyncOmniEngine: preprocess_ms=_preprocess_ms, request_timestamp=request_timestamp, enqueue_ts=time.perf_counter(), + request_artifact_dirs=request_artifact_dirs or None, ) - def _enqueue_cfg_companions( + def _build_cfg_companions( self, parent_id: str, original_prompt: Any, stage0_params: Any, sampling_params_list: list[Any], - ) -> None: - """Expand prompt into CFG companions, process through InputProcessor, and enqueue.""" - try: - expanded = self.prompt_expand_func(original_prompt, stage0_params) - except Exception: - logger.exception("[AsyncOmniEngine] prompt_expand_func failed for req %s", parent_id) - return + ) -> list[AddCompanionRequestMessage]: + """Expand a prompt into its CFG companions, without enqueueing any. + Construction is separated from admission so a guided request is + all-or-nothing. A model whose guidance is mandatory cannot decode a + request whose companion never arrived: the pair never completes, the + request occupies scheduler and KV capacity for the scheduler's whole + hold budget, and then produces no audio. Raising here instead means the + caller learns immediately and nothing was admitted. + + Raises: + Exception: Whatever prompt expansion or input processing raised. + The caller is expected to let it reach the client. + """ + assert self.prompt_expand_func is not None + expanded = self.prompt_expand_func(original_prompt, stage0_params) if not expanded: - return + return [] + companions: list[AddCompanionRequestMessage] = [] + assert self.input_processor is not None for ep in expanded: cid = f"{parent_id}{ep.request_id_suffix}" companion_prompt = ep.prompt @@ -814,6 +862,14 @@ class AsyncOmniEngine: params=companion_params, supported_tasks=self.supported_tasks, ) + # Same restore the parent request gets: the upstream input + # processor drops omni-only prompt fields, so without this the + # companion reaches the worker with no additional_information at + # all. That is where ``global_request_id`` lives, and it is what + # was just injected above, so skipping it silently undoes the + # injection: the model sees the companion row with no id and + # cannot match it to its conditioned partner. + request = upgrade_to_omni_request(request, companion_prompt) request.external_req_id = cid # Companions are stage-0-final for ordinary downstream payloads, # but diffusion still needs their CFG KV caches. @@ -822,7 +878,7 @@ class AsyncOmniEngine: # Registration of this companion on stage-0's output processor is # deferred to Orchestrator._handle_add_companion, which routes # admission through StagePool.submit_initial(..., affinity_request_id=...). - self.request_queue.sync_q.put( + companions.append( AddCompanionRequestMessage( companion_id=cid, parent_id=parent_id, @@ -832,11 +888,35 @@ class AsyncOmniEngine: sampling_params_list=companion_spl, ) ) + return companions + + def _enqueue_cfg_companions( + self, + parent_id: str, + original_prompt: Any, + stage0_params: Any, + sampling_params_list: list[Any], + ) -> None: + """Build and enqueue CFG companions, tolerating a build failure. + + Kept for callers that admit the parent first and cannot roll it back. + Prefer building with :meth:`_build_cfg_companions` before the parent is + admitted, so the pair is atomic. + """ + try: + companions = self._build_cfg_companions(parent_id, original_prompt, stage0_params, sampling_params_list) + except Exception: + logger.exception("[AsyncOmniEngine] CFG companion build failed for req %s", parent_id) + return + for companion in companions: + self.request_queue.sync_q.put(companion) + if not companions: + return logger.info( "[AsyncOmniEngine] CFG expansion for req %s: %d companions", parent_id, - len(expanded), + len(companions), ) @staticmethod @@ -863,7 +943,7 @@ class AsyncOmniEngine: "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1, } - if cache_backend in ("step_cache"): + if cache_backend == "step_cache": return { "step_cache_dit_enabled": True, "velocity_sim_thresholds": [0.95, 0.93], @@ -968,9 +1048,10 @@ class AsyncOmniEngine: ring_degree = normalized_kwargs.get("ring_degree") or 1 allgather_degree = normalized_kwargs.get("allgather_degree") or 1 ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" + ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False)) sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 - data_parallel_size = normalized_kwargs.get("data_parallel_size") or 1 + data_parallel_size = normalized_kwargs.get("data_parallel_size") tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 @@ -994,6 +1075,7 @@ class AsyncOmniEngine: ring_degree=ring_degree, allgather_degree=allgather_degree, ulysses_mode=ulysses_mode, + ulysses_a2a_permute=ulysses_a2a_permute, cfg_parallel_size=cfg_parallel_size, vae_patch_parallel_size=vae_patch_parallel_size, vae_parallel_mode=vae_parallel_mode, @@ -1003,24 +1085,45 @@ class AsyncOmniEngine: hsdp_replicate_size=hsdp_replicate_size, ) + num_gpus = normalized_kwargs.get("num_gpus") + if num_gpus is not None: + num_gpus = int(num_gpus) + parallel_config.resolve_data_parallel_size(num_gpus) + num_devices = max(1, int(parallel_config.world_size)) devices = ",".join(str(i) for i in range(num_devices)) model_class_name = kwargs.get("model_class_name", None) - final_output_type = "audio" if model_class_name and supports_audio_output(model_class_name) else "image" + final_output_type = get_diffusion_output_type(model_class_name) attention_config = None if ( kwargs.get("diffusion_attention_config") is not None or kwargs.get("diffusion_attention_backend") is not None + or kwargs.get("fastvideo_vsa_topk") is not None ): attention_config = parse_attention_config( kwargs.get("diffusion_attention_config"), attention_backend=kwargs.get("diffusion_attention_backend"), + fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), ) + extras = dict(kwargs.get("extras") or {}) + for key, default in ( + ("auxiliary_text_encoder", None), + ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), + ): + top_level_value = kwargs.get(key) + if top_level_value is not None: + extras[key] = top_level_value + else: + extras.setdefault(key, default) + stage_engine_args = { "max_num_seqs": kwargs.get("max_num_seqs") or 1, "parallel_config": parallel_config, + # Default-stage construction bypasses the structured projection. + # Runner selection remains owned by the selected engine/platform. + "engine_backend": kwargs.get("engine_backend", "default"), "model_class_name": kwargs.get("model_class_name", None), "task_type": kwargs.get("task_type", None), "model_config": kwargs.get("model_config", None), @@ -1037,6 +1140,9 @@ class AsyncOmniEngine: "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), + "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), + "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), + "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( "regional" @@ -1056,9 +1162,7 @@ class AsyncOmniEngine: "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), "worker_extension_cls": kwargs.get("worker_extension_cls", None), "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), - "distributed_executor_backend": ( - "mp" if kwargs.get("distributed_executor_backend") is None else kwargs["distributed_executor_backend"] - ), + "distributed_executor_backend": kwargs.get("distributed_executor_backend"), "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32), @@ -1074,10 +1178,7 @@ class AsyncOmniEngine: "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), "streaming_output": kwargs.get("diffusion_streaming_output", False), "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), - "extras": { - "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None), - "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), - }, + "extras": extras, **( { "profiler_config": asdict(kwargs["profiler_config"]) @@ -1088,6 +1189,8 @@ class AsyncOmniEngine: else {} ), } + if num_gpus is not None: + stage_engine_args["num_gpus"] = num_gpus # Only set dtype if it was already explicitly passed and normalized if "dtype" in normalized_kwargs: stage_engine_args["dtype"] = normalized_kwargs["dtype"] @@ -1116,40 +1219,6 @@ class AsyncOmniEngine: default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" return default_stage_cfg - @staticmethod - def _strip_single_engine_args(kwargs: dict[str, Any]) -> dict[str, Any]: - """Remove parent ``EngineArgs`` fields from *kwargs*. - - When ``stage_configs_path`` is set, per-stage engine args are defined - in the YAML. Top-level single-engine fields (``compilation_config``, - ``tensor_parallel_size``, …) must not leak into per-stage configs via - the ``base_engine_args`` merge in ``load_stage_configs_from_yaml`` — - they can cause type errors (e.g. ``compilation_config`` as a JSON - string rejected by ``VllmConfig``) or silently override YAML values. - - Logs a warning for any parent field whose value differs from the - dataclass default, so users know their explicit overrides are ignored. - See the module-level ``_PARENT_ARGS_*`` constants for the routing - contracts this method enforces. - """ - parent_fields: dict[str, dataclasses.Field] = {f.name: f for f in dataclasses.fields(EngineArgs)} - result, overridden = strip_parent_engine_args( - kwargs, - parent_fields=parent_fields, - keep_keys=_PARENT_ARGS_KEEP, - strip_keys=_PARENT_ARGS_STRIP, - no_warn_keys=_PARENT_ARGS_NO_WARN, - ) - - if overridden: - logger.warning( - "stage_configs_path is set — the following top-level engine " - "args are ignored (per-stage YAML takes precedence): %s", - ", ".join(sorted(overridden)), - ) - - return result - def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: """Apply a strategy-derived ``omni_lb_policy`` to the engine. @@ -1188,36 +1257,51 @@ class AsyncOmniEngine: ) -> tuple[str, list[Any]]: """Resolve stage configs and inject defaults shared by orchestrator/headless.""" - stage_configs_path = kwargs.get("stage_configs_path", None) + for legacy_arg in ("stage_configs_path", "stage_configs"): + if legacy_arg in kwargs: + raise ValueError(f"`{legacy_arg}` is no longer supported; use `deploy_config` instead.") + deploy_config_path = kwargs.pop("deploy_config", None) strategy_config_path = kwargs.pop("strategy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) - explicit_stage_configs = kwargs.pop("stage_configs", None) - if explicit_stage_configs is not None: - logger.warning( - "`stage_configs` is not part of the public API. " - "Ignoring it and resolving stages from stage_configs_path/model factory." - ) - - # Legacy ``stage_args`` YAMLs own per-stage EngineArgs completely, so - # strip parent CLI fields. New-format deploy YAMLs (also accepted via - # ``--stage-configs-path``) merge CLI overrides the same way as - # ``--deploy-config`` — do not strip, or flags like - # ``interleave_mm_strings`` / ``media_io_kwargs`` are silently dropped. - if stage_configs_path is not None and not is_new_format_deploy_config(stage_configs_path): - base_kwargs = self._strip_single_engine_args(kwargs) - else: - base_kwargs = kwargs # Parse --stage-overrides JSON string if provided stage_overrides = parse_stage_overrides(stage_overrides_json) + # Unregistered diffusion checkpoints use the single-stage fallback + # below instead of StageConfigFactory, so fold stage-0 model extras + # into the fallback's input as well. Registered pipelines still get + # the complete override mapping through load_and_resolve_stage_configs. + default_stage_kwargs = kwargs + stage_zero_overrides = (stage_overrides or {}).get("0", {}) + if "extras" in stage_zero_overrides: + override_extras = stage_zero_overrides["extras"] + if not isinstance(override_extras, Mapping): + raise TypeError("stage 0 extras must be a mapping") + default_stage_kwargs = { + **kwargs, + "extras": { + **(kwargs.get("extras") or {}), + **override_extras, + }, + } + + def create_default_stage_config() -> list: + fallback_kwargs = dict(default_stage_kwargs) + if not fallback_kwargs.get("model_class_name"): + model_class_name = resolve_model_class_name( + model, + fallback_kwargs.get("diffusion_load_format", "default"), + ) + if model_class_name is not None: + fallback_kwargs["model_class_name"] = model_class_name + return self._create_default_diffusion_stage_cfg(fallback_kwargs) + config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs( model, - stage_configs_path, - base_kwargs, + kwargs, trust_remote_code=trust_remote_code, - default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs), + default_stage_cfg_factory=create_default_stage_config, deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, strategy_config_path=strategy_config_path, @@ -1262,6 +1346,7 @@ class AsyncOmniEngine: if ( kwargs.get("diffusion_attention_config") is not None or kwargs.get("diffusion_attention_backend") is not None + or kwargs.get("fastvideo_vsa_topk") is not None ): has_stage_attention = ( hasattr(cfg.engine_args, "diffusion_attention_config") @@ -1271,6 +1356,7 @@ class AsyncOmniEngine: cfg.engine_args.diffusion_attention_config = parse_attention_config( kwargs.get("diffusion_attention_config"), attention_backend=kwargs.get("diffusion_attention_backend"), + fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), ) quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config") if quantization_config is not None: @@ -1347,32 +1433,61 @@ class AsyncOmniEngine: a queue + coroutine-switch round-trip. The Orchestrator receives a ready-to-submit OmniEngineCoreRequest. """ - msg = self._build_add_request_message( - request_id=request_id, - prompt=prompt, - prompt_text=prompt_text, - sampling_params_list=sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - tokenization_kwargs=tokenization_kwargs, - trace_headers=trace_headers, - priority=priority, - data_parallel_rank=data_parallel_rank, - reasoning_ended=reasoning_ended, - resumable=resumable, - ) - self.request_queue.sync_q.put(msg) + try: + msg = self._build_add_request_message( + request_id=request_id, + prompt=prompt, + prompt_text=prompt_text, + sampling_params_list=sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + tokenization_kwargs=tokenization_kwargs, + trace_headers=trace_headers, + priority=priority, + data_parallel_rank=data_parallel_rank, + reasoning_ended=reasoning_ended, + resumable=resumable, + ) + except BaseException: + if isinstance(prompt, dict): + for artifact_dir in prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) or (): + if isinstance(artifact_dir, str): + shutil.rmtree(artifact_dir, ignore_errors=True) + raise + # CFG companions are built before the parent is admitted, so the group + # is all-or-nothing: a build failure raises here, nothing is enqueued, + # and the caller sees the error. Admitting the parent first would leave + # an orphan holding scheduler and KV capacity that can never complete, + # because a model whose guidance is mandatory cannot decode a request + # whose companion never arrived. + companions: list[AddCompanionRequestMessage] = [] + try: + if self.prompt_expand_func is not None and final_stage_id > 0: + effective_spl = msg.sampling_params_list + stage0_params = effective_spl[0] if effective_spl else None + if stage0_params is not None: + companions = self._build_cfg_companions( + request_id, msg.original_prompt, stage0_params, effective_spl + ) - # CFG companion expansion: create and enqueue companion requests - # so the AR stage also generates their KV caches. - if self.prompt_expand_func is not None and final_stage_id > 0: - original_prompt = msg.original_prompt - effective_spl = msg.sampling_params_list - stage0_params = effective_spl[0] if effective_spl else None - if stage0_params is not None: - self._enqueue_cfg_companions(request_id, original_prompt, stage0_params, effective_spl) + self.request_queue.sync_q.put(msg) + except BaseException: + for artifact_dir in msg.request_artifact_dirs or (): + shutil.rmtree(artifact_dir, ignore_errors=True) + raise + finally: + if isinstance(msg.original_prompt, dict): + msg.original_prompt.pop(REQUEST_ARTIFACT_DIRS_KEY, None) + for companion in companions: + self.request_queue.sync_q.put(companion) + if companions: + logger.info( + "[AsyncOmniEngine] CFG expansion for req %s: %d companions", + request_id, + len(companions), + ) async def add_request_async( self, @@ -1744,19 +1859,111 @@ class AsyncOmniEngine: raise RuntimeError("Orchestrator died unexpectedly. See logs above.") return None + async def get_output_blocking_async(self, timeout: float = 1.0) -> EngineQueueMessage | None: + """Blocking-wait read from the Orchestrator output queue. + + Waits up to ``timeout`` seconds in a dedicated drain thread for the + next message (condition-variable wakeup instead of a poll cadence); + returns ``None`` on timeout so the caller keeps its liveness check, + mirroring ``try_get_output_async``'s contract. Used by the serving + final-output drain when ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` is on. + """ + executor = self._output_drain_executor + if executor is None: + executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="omni-output-drain", + ) + self._output_drain_executor = executor + + sync_q = self.output_queue.sync_q + + def _drain_get() -> EngineQueueMessage | None: + # Exceptions are swallowed to a None sentinel: the queue may be + # closed mid-shutdown, and an exception left on an executor future + # after task cancellation would warn as never-retrieved. + try: + return sync_q.get(timeout=timeout) + except queue.Empty: + return None + except Exception: + return None + + loop = asyncio.get_running_loop() + msg = await loop.run_in_executor(executor, _drain_get) + if msg is None and not self.is_alive(): + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") + return msg + def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: """Get cached metadata for a stage.""" return self.stage_metadata[stage_id] def abort(self, request_ids: list[str]) -> None: - """Send abort message to the Orchestrator.""" + """Fire-and-forget abort: enqueue and return without waiting. + + Prefer :meth:`abort_async` when the caller needs acknowledgment that + stage aborts, binding release, and orchestrator request cleanup finished. + """ + if not request_ids or getattr(self, "_shutdown_called", False): + return if self.request_queue is None: raise RuntimeError("request_queue is not initialized") - self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) + try: + self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_janus_sync_queue_shutdown(exc): + return + raise - async def abort_async(self, request_ids: list[str]) -> None: - """Async abort API.""" - self.abort(request_ids) + async def abort_async( + self, + request_ids: list[str], + timeout: float | None = None, + ) -> list[OutputMessage]: + """Abort requests and wait for orchestrator acknowledgment. + + Unlike :meth:`abort`, this generates an ``rpc_id``, correlates the + :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and + raises if the orchestrator reports failure or times out. + + Returns: + Final-stage AR abort ``OutputMessage`` list carrying partial + tokens generated before abort (empty for diffusion / no OP state). + """ + if not request_ids or getattr(self, "_shutdown_called", False): + return + if self.request_queue is None: + raise RuntimeError("request_queue is not initialized") + transport = self._correlated_rpc_client + if transport is None: + raise RuntimeError("correlated RPC client is not initialized") + + rpc_id = uuid.uuid4().hex + msg = AbortRequestMessage(request_ids=request_ids, rpc_id=rpc_id) + + def _wait() -> AbortResultMessage: + result_msg = transport.execute( + ("abort", rpc_id), + msg, + timeout=timeout, + timeout_message=f"abort timed out after {timeout} seconds", + block_on_submit=True, + ) + if not isinstance(result_msg, AbortResultMessage): + raise RuntimeError(f"unexpected abort result type: {type(result_msg).__name__}") + return result_msg + + loop = asyncio.get_running_loop() + try: + result_msg = await loop.run_in_executor(None, _wait) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc): + return + raise + if not result_msg.success: + raise RuntimeError(result_msg.error or "abort failed") + return list(result_msg.abort_outputs or []) def submit_interaction( self, @@ -1896,6 +2103,12 @@ class AsyncOmniEngine: except Exception: pass + if self._output_drain_executor is not None: + # Any in-flight blocking get bails out within its ≤1 s timeout + # (or immediately via the queue close above), so don't wait. + self._output_drain_executor.shutdown(wait=False) + self._output_drain_executor = None + if hasattr(self, "_runtime") and self._runtime is not None and orchestrator_stopped: try: self._runtime.shutdown() diff --git a/vllm_omni/engine/membership_controller.py b/vllm_omni/engine/membership_controller.py index a7676de54..8f008e286 100644 --- a/vllm_omni/engine/membership_controller.py +++ b/vllm_omni/engine/membership_controller.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MembershipController: distributed replica lifecycle management. @@ -49,10 +49,18 @@ class MembershipController: self._stage_pools = stage_pools self._remote_replica_factory = remote_replica_factory self._membership_tasks: set[asyncio.Task[None]] = set() + self._attaching_remote_replicas: set[tuple[int, int]] = set() + self._attached_remote_replicas: set[tuple[int, int]] = set() + self._remote_replica_keys_by_input_addr: dict[str, tuple[int, int]] = {} + self._watch_generation = 0 + self._attach_start_generation: dict[tuple[int, int], int] = {} + self._up_generation_by_input_addr: dict[tuple[int, str], int] = {} + self._observed_up_attached_addrs: set[tuple[int, str]] = set() self._shutdown_event = asyncio.Event() self._watcher_task: asyncio.Task[None] | None = None self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None + self._replica_removed_callback: Callable[[int, int], None] | None = None self._hub = OmniCoordClientForHub(coordinator_pub_address) factory = load_balancer_factory @@ -67,8 +75,34 @@ class MembershipController: async def handle_register(self, stage_id: int, replica_id: int) -> None: """Handle a register_remote_replica message (fire-and-forget).""" + if self._shutdown_event.is_set(): + logger.info( + "[MembershipController] ignoring register during shutdown stage=%d replica=%d", + stage_id, + replica_id, + ) + return + + key = (stage_id, replica_id) + if key in self._attached_remote_replicas: + logger.info( + "[MembershipController] ignoring register for attached replica stage=%d replica=%d", + stage_id, + replica_id, + ) + return + if key in self._attaching_remote_replicas: + logger.info( + "[MembershipController] ignoring register while attach is in flight stage=%d replica=%d", + stage_id, + replica_id, + ) + return + + self._attaching_remote_replicas.add(key) + self._attach_start_generation[key] = self._watch_generation self._spawn_task( - self._do_register(stage_id, replica_id), + self._register_once(stage_id, replica_id), label=f"register-s{stage_id}-r{replica_id}", ) @@ -85,8 +119,12 @@ class MembershipController: return effective_output_queue = output_queue if output_queue is not None else self._output_queue effective_cleanup_callback = cleanup_callback if cleanup_callback is not None else self._cleanup_callback + replica_id = pool.get_replica_id_by_addr(input_addr) + client = pool.get_client_by_addr(input_addr) affected = pool.invalidate_addr(input_addr) self._detach_replica(stage_id, input_addr) + if client is not None and replica_id is not None and self._replica_removed_callback is not None: + self._replica_removed_callback(stage_id, replica_id) if affected and effective_cleanup_callback is not None: await effective_cleanup_callback(affected) if affected and effective_output_queue is not None: @@ -131,8 +169,24 @@ class MembershipController: while not self._shutdown_event.is_set(): try: snap = self._hub.get_replica_list() + self._watch_generation += 1 current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP} - for stage_id, addr in last_up - current: + attached = { + (stage_id, input_addr) + for input_addr, (stage_id, _replica_id) in self._remote_replica_keys_by_input_addr.items() + } + self._observed_up_attached_addrs.update(attached & current) + attaching_stages = {stage_id for stage_id, _replica_id in self._attaching_remote_replicas} + for addr_key in current: + if addr_key[0] in attaching_stages: + self._up_generation_by_input_addr[addr_key] = self._watch_generation + + # Preserve watcher cleanup for untracked/static clients. A + # controller-attached client is eligible only after this + # attachment has been observed UP at least once. + untracked_disappeared = (last_up - current) - attached + confirmed_attached_disappeared = self._observed_up_attached_addrs - current + for stage_id, addr in untracked_disappeared | confirmed_attached_disappeared: self._spawn_task( self.handle_unregister(stage_id, addr), label=f"unregister-s{stage_id}", @@ -148,24 +202,75 @@ class MembershipController: except asyncio.CancelledError: raise - async def _do_register(self, stage_id: int, replica_id: int) -> None: + async def _register_once(self, stage_id: int, replica_id: int) -> None: + key = (stage_id, replica_id) + attach_start_generation = self._attach_start_generation[key] + try: + input_addr = await self._do_register(stage_id, replica_id) + if input_addr is None: + return + self._attached_remote_replicas.add(key) + self._remote_replica_keys_by_input_addr[input_addr] = key + addr_key = (stage_id, input_addr) + if self._up_generation_by_input_addr.get(addr_key, -1) > attach_start_generation: + self._observed_up_attached_addrs.add(addr_key) + finally: + self._attaching_remote_replicas.discard(key) + self._attach_start_generation.pop(key, None) + if not any(attaching_stage_id == stage_id for attaching_stage_id, _ in self._attaching_remote_replicas): + self._up_generation_by_input_addr = { + addr_key: generation + for addr_key, generation in self._up_generation_by_input_addr.items() + if addr_key[0] != stage_id + } + + async def _do_register(self, stage_id: int, replica_id: int) -> str | None: pool = self._pool_for_stage_id(stage_id) if pool is None: logger.warning("[MembershipController] register: stage_id %d out of range", stage_id) - return + return None client = await asyncio.to_thread(self._remote_replica_factory, stage_id, replica_id) - input_addr = StagePool._client_input_addr(client) - if input_addr is None: - raise RuntimeError(f"remote replica factory for stage {stage_id} produced a client without input address") - pool.add_client(input_addr, client, replica_id=replica_id) + if self._shutdown_event.is_set(): + try: + client.shutdown() + except Exception: + logger.exception( + "[MembershipController] failed to shutdown late replica client stage=%d replica=%d", + stage_id, + replica_id, + ) + return None + try: + input_addr = StagePool._client_input_addr(client) + if input_addr is None: + raise RuntimeError( + f"remote replica factory for stage {stage_id} produced a client without input address" + ) + pool.add_client(input_addr, client, replica_id=replica_id) + except Exception: + try: + client.shutdown() + except Exception: + logger.exception( + "[MembershipController] failed to shutdown rejected replica client stage=%d replica=%d", + stage_id, + replica_id, + ) + raise logger.info( "[MembershipController] attached remote replica stage=%d replica=%d addr=%s", stage_id, replica_id, input_addr, ) + return input_addr def _detach_replica(self, stage_id: int, input_addr: str) -> None: + key = self._remote_replica_keys_by_input_addr.pop(input_addr, None) + if key is not None: + self._attached_remote_replicas.discard(key) + self._attaching_remote_replicas.discard(key) + self._observed_up_attached_addrs.discard((stage_id, input_addr)) pool = self._pool_for_stage_id(stage_id) if pool is None: return @@ -183,10 +288,12 @@ class MembershipController: *, output_queue: asyncio.Queue[EngineQueueMessage], cleanup_callback: Callable[[list[str]], Awaitable[None]], + replica_removed_callback: Callable[[int, int], None] | None = None, ) -> None: """Install shared cleanup sinks for watcher-driven unregister events.""" self._output_queue = output_queue self._cleanup_callback = cleanup_callback + self._replica_removed_callback = replica_removed_callback def _pool_for_stage_id(self, stage_id: int) -> StagePool | None: if not (0 <= stage_id < len(self._stage_pools)): diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 80bf42b0e..f6a91dbcc 100644 --- a/vllm_omni/engine/messages.py +++ b/vllm_omni/engine/messages.py @@ -27,6 +27,7 @@ class StageSubmissionMessage(EngineQueueMessage, kw_only=True): request_timestamp: float enqueue_ts: float final_output_stage_ids: list[int] | None = None + request_artifact_dirs: list[str] | None = None class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True): @@ -42,6 +43,9 @@ class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True): class AbortRequestMessage(EngineQueueMessage, kw_only=True): type: Literal["abort"] = "abort" request_ids: list[str] + # When set, Orchestrator emits AbortResultMessage on rpc_async_queue so + # callers can await acknowledgment. Omit for fire-and-forget abort(). + rpc_id: str | None = None class InteractionMessage(EngineQueueMessage, kw_only=True): @@ -98,6 +102,20 @@ class OutputMessage(EngineQueueMessage, kw_only=True): stage_submit_ts: float | None = None +class AbortResultMessage(EngineQueueMessage, kw_only=True): + type: Literal["abort_result"] = "abort_result" + rpc_id: str + success: bool + error: str | None = None + # Final-stage AR abort outputs (partial tokens) for frontend delivery. + # Empty for diffusion / requests with no output-processor state. + abort_outputs: list[OutputMessage] | None = None + + @property + def rpc_correlation_key(self) -> tuple[str, str]: + return ("abort", self.rpc_id) + + class StageMetricsMessage(EngineQueueMessage, kw_only=True): type: Literal["stage_metrics"] = "stage_metrics" request_id: str diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 901bdb980..c0240c4bc 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Orchestrator for vLLM-Omni multi-stage runtime. @@ -12,9 +15,12 @@ handled by :class:`MembershipController`, which is injected optionally. from __future__ import annotations import asyncio +import os +import shutil import time as _time from collections.abc import Awaitable, Callable from dataclasses import dataclass, field +from http import HTTPStatus from typing import TYPE_CHECKING, Any import janus @@ -35,6 +41,7 @@ from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, + AbortResultMessage, AddCompanionRequestMessage, CollectiveRPCRequestMessage, CollectiveRPCResultMessage, @@ -51,12 +58,38 @@ from vllm_omni.engine.messages import ( from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError +from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE +from vllm_omni.metrics import definitions as metric_defs from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger +from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) + +def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None: + for artifact_dir in artifact_dirs: + shutil.rmtree(artifact_dir, ignore_errors=True) + + +# VLLM_OMNI_EVENT_DRIVEN_ORCH=1 switches the orchestration loop (and the +# serving-side final-output drain in entrypoints/async_omni.py) from the legacy +# 1 ms poll cadence to event-driven wakeups: one reader task per live LLM stage +# replica awaits `client.get_output_async()` directly — the same pattern vLLM's +# own AsyncLLM output handler uses — and feeds a single serial dispatch queue. +# Default off; the legacy poll loop remains the fallback. +_EVENT_DRIVEN_ORCH_ENV = "VLLM_OMNI_EVENT_DRIVEN_ORCH" + +# How often the event-driven loop reconciles its reader-task set against +# `available_replica_ids()` (elastic membership, replica eviction) while idle. +_ORCH_READER_RECONCILE_INTERVAL_S = 0.5 + + +def _event_driven_orch_enabled() -> bool: + return os.environ.get(_EVENT_DRIVEN_ORCH_ENV, "0").strip().lower() in ("1", "true", "yes", "on") + + if TYPE_CHECKING: from vllm_omni.experimental.fullduplex.engine.contracts import ( DuplexControlPlanePort, @@ -125,8 +158,16 @@ def build_engine_core_request_from_tokens( pooling_params = None if isinstance(params, SamplingParams): sampling_params = params.clone() - if sampling_params.max_tokens is None and model_config is not None: - sampling_params.max_tokens = model_config.max_model_len - len(prompt_token_ids) + if model_config is not None: + remaining = model_config.max_model_len - len(prompt_token_ids) + if sampling_params.max_tokens is None: + sampling_params.max_tokens = remaining + # ``check_stop`` returns early while ``min_tokens`` is unmet, so a + # min_tokens wider than the leftover context window never reaches + # the length cap: the request stalls once the scheduler can grant + # no further tokens. Stage prompts are built here, so clamp now. + if sampling_params.min_tokens > remaining: + sampling_params.min_tokens = max(0, remaining) else: pooling_params = params.clone() @@ -186,20 +227,41 @@ class OrchestratorRequestState: duplex_stage_fences: dict[int, DuplexFence] = field(default_factory=dict) duplex_config_generation: int = -1 running_counter_registered: bool = False + request_artifact_dirs: set[str] = field(default_factory=set) + + +@dataclass +class StreamingSegmentState: + """Streaming segment-boundary state for one stage of a request. + + Every stage of a multi-stage pipeline reaches its own segment boundaries + independently, so this is tracked per stage rather than per request. + """ + + # Flag of segment of streaming input finished + finished: bool = False + # Tokens from the current raw segment boundary. The vLLM output processor + # does not guarantee that EngineCoreOutput.new_token_ids survives on the + # processed RequestOutput used by the routing layer. + token_ids: list[int] = field(default_factory=list) + output_metadata: dict[str, Any] = field(default_factory=dict) @dataclass class StreamingInputState: # Flag of streaming input request enabled: bool = False - # Flag of segment of streaming input finished - segment_finished: bool = False - # Tokens from the current raw segment boundary. The vLLM output processor - # does not guarantee that EngineCoreOutput.new_token_ids survives on the - # processed RequestOutput used by the routing layer. - segment_token_ids: list[int] = field(default_factory=list) - segment_output_metadata: dict[str, Any] = field(default_factory=dict) - # Streaming update prompt length + # Keyed by stage. ``_orchestration_loop`` records a boundary per stage poll + # into this one ``req_state``, while every consumer reads it for a specific + # ``stage_id``, so a flat slot is only correct as long as each read happens in + # the same iteration as its own stage's write. It is not for a stage that + # never records one (the ``poll_diffusion_output`` branch reaches + # ``_handle_processed_outputs`` without the raw-output loop) or for any + # consumer reading outside the writing iteration. + segments: dict[int, StreamingSegmentState] = field(default_factory=dict) + # Streaming update prompt length. NOT stage-keyed: its only consumer reads it + # through the streaming context on the stage-input-bridge path, not by stage + # id. Do not assume the rest of this struct is stage-safe. new_prompt_len_snapshot: int | None = None # Model/bridge-specific runtime states (e.g., thinker->talker) bridge_states: dict[str, Any] = field(default_factory=dict) @@ -207,6 +269,14 @@ class StreamingInputState: # while the downstream input processor consumes upstream token output. source_token_decoder: Callable[..., str] | None = None + def segment(self, stage_id: int | None) -> StreamingSegmentState: + """Return ``stage_id``'s segment state, or a fresh empty one if unreported. + + Fresh rather than shared: ``output_metadata`` is handed out by reference. + """ + segment = self.segments.get(stage_id) + return segment if segment is not None else StreamingSegmentState() + class _OrchestratorDuplexStagePort: """Adapts generic stage pools to the model-neutral duplex control plane.""" @@ -221,7 +291,7 @@ class _OrchestratorDuplexStagePort: async_chunk: bool, prewarm_async_chunk_stages: Callable[ [str, Any, OrchestratorRequestState], - Awaitable[None], + Awaitable[bool], ], ) -> None: self._stage_pools = stage_pools @@ -248,7 +318,15 @@ class _OrchestratorDuplexStagePort: duplex_state = {} request_state.streaming.bridge_states["duplex"] = duplex_state previous_epoch = duplex_state.get("epoch") - if not isinstance(duplex_state.get("model_turn_id"), int) or previous_epoch != context.fence.epoch: + current_model_turn_id = duplex_state.get("model_turn_id") + if ( + not isinstance(current_model_turn_id, int) + or previous_epoch != context.fence.epoch + or current_model_turn_id < context.fence.turn_id + ): + # A serving-side safety boundary can close a Realtime response + # before the model emits its normal turn_eos. Catch up the + # engine-owned identity when the next fenced append starts. duplex_state["model_turn_id"] = context.fence.turn_id duplex_state.update( { @@ -307,11 +385,23 @@ class _OrchestratorDuplexStagePort: else: replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None) if self._async_chunk and context.stage_id == 0: - await self._prewarm_async_chunk_stages( + prewarmed = await self._prewarm_async_chunk_stages( context.request_id, request, request_state, ) + if not prewarmed: + # The prewarm already failed the request, aborted it, and + # popped its state. Falling through would write fences and + # a submit timestamp onto that orphaned object and + # re-increment the running counter the cleanup just + # released, then hand the control plane a success result + # for a request that no longer exists. Raise instead: + # ``DuplexControlPlane.handle_append`` turns this into an + # error result for the append. + raise RuntimeError( + f"async-chunk prewarm failed for duplex request {context.request_id}; the request was aborted" + ) request_state.duplex_stage_fences[context.stage_id] = context.fence request_state.stage_submit_ts[context.stage_id] = _time.time() if not request_state.running_counter_registered and self._running_counter is not None: @@ -333,22 +423,26 @@ class Orchestrator: # Class-level defaults so tests that bypass __init__ via object.__new__ # don't AttributeError when transfer / counter emit paths access them. _running_counter: OmniRequestCounter | None = None + _engines_waiting_counter: OmniRequestCounter | None = None _transfer_emitter: Any = None + _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None duplex_control_plane: DuplexControlPlanePort | None = None def __init__( self, request_async_queue: janus.AsyncQueue[EngineQueueMessage], - output_async_queue: janus.AsyncQueue[dict[str, Any]], - rpc_async_queue: janus.AsyncQueue[dict[str, Any]], + output_async_queue: janus.AsyncQueue[EngineQueueMessage], + rpc_async_queue: janus.AsyncQueue[EngineQueueMessage], stage_pools: list[StagePool], *, async_chunk: bool = False, pd_config: dict[str, Any] | None = None, membership_controller: MembershipController | None = None, running_counter: OmniRequestCounter | None = None, + engines_waiting_counter: OmniRequestCounter | None = None, transfer_emitter: Any = None, + prom_metrics: Any = None, log_stats: bool = False, enable_orch_monitor: bool = False, duplex_runtime_extension: DuplexRuntimeExtension | None = None, @@ -363,6 +457,8 @@ class Orchestrator: self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools self.log_stats = log_stats + self._prom_metrics = prom_metrics + self._stage_replica_waiting: dict[tuple[int, int], int] = {} self._orch_monitor = create_orch_monitor( enabled=enable_orch_monitor, replica_sampler=self._sample_replica_metrics, @@ -381,7 +477,13 @@ class Orchestrator: self._pd_bootstrap_addr = pd_config.get("bootstrap_addr") self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") self.request_states: dict[str, OrchestratorRequestState] = {} - self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats) + self._init_metrics_state( + stage_pools, + running_counter, + transfer_emitter, + engines_waiting_counter=engines_waiting_counter, + log_stats=log_stats, + ) self._cfg_tracker = CfgCompanionTracker() self._stage_input_processors: dict[int, Any] = {} @@ -417,6 +519,7 @@ class Orchestrator: self._shutdown_event = asyncio.Event() self._stages_shutdown = False + self._event_driven_orch = _event_driven_orch_enabled() # Distributed membership (optional, injected by DistStageRuntime) self._membership = membership_controller @@ -426,6 +529,7 @@ class Orchestrator: stage_pools: list[StagePool], running_counter: OmniRequestCounter | None, transfer_emitter: Any, + engines_waiting_counter: OmniRequestCounter | None = None, log_stats: bool = False, ) -> None: """Wire up all metric-related orchestrator state. @@ -451,6 +555,7 @@ class Orchestrator: from iteration_stats independently of scheduler gauges. """ self._running_counter = running_counter + self._engines_waiting_counter = engines_waiting_counter self._transfer_emitter = transfer_emitter # Flat engine_idx ↔ (stage, replica) maps. The reverse map is @@ -516,6 +621,7 @@ class Orchestrator: self._membership.install_unregister_handlers( output_queue=self.output_async_queue, cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True), + replica_removed_callback=self._remove_stage_replica_waiting, ) membership_watcher = self._membership.start() @@ -671,6 +777,7 @@ class Orchestrator: final_output_stage_ids=final_output_stage_ids, request_timestamp=float(msg.request_timestamp or _time.time()), mm_features=getattr(prompt, "mm_features", None), + request_artifact_dirs=set(msg.request_artifact_dirs or ()), ) self.request_states[request_id] = req_state self._register_running_request(req_state) @@ -795,13 +902,44 @@ class Orchestrator: ) async def _handle_abort(self, msg: AbortRequestMessage) -> None: - """Handle an abort message from the main thread.""" + """Handle an abort message from the main thread. + + When ``msg.rpc_id`` is set, emit :class:`AbortResultMessage` after + stage aborts, binding release, and request cleanup complete so the + caller can await acknowledgment. Fire-and-forget aborts omit rpc_id. + + AR final-stage abort outputs (partial tokens) are attached to the + result message when ``rpc_id`` is set, or enqueued on + ``output_async_queue`` for fire-and-forget aborts. + """ request_ids = msg.request_ids - # _cleanup_request_ids is CFG-aware: it expands aborted parents to - # their companions and fails a deferred parent whose companion is - # aborted before its output arrived. - await self._cleanup_request_ids(list(request_ids), abort=True) - logger.info("[Orchestrator] Aborted request(s) %s", request_ids) + error: str | None = None + abort_outputs: list[OutputMessage] = [] + try: + # _cleanup_request_ids is CFG-aware: it expands aborted parents to + # their companions and fails a deferred parent whose companion is + # aborted before its output arrived. + abort_outputs = await self._cleanup_request_ids(list(request_ids), abort=True) + logger.info("[Orchestrator] Aborted request(s) %s", request_ids) + except Exception as exc: + error = str(exc) + logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids) + if msg.rpc_id is None: + raise + if msg.rpc_id is not None: + # Always emit a result when rpc_id is set so CorrelatedRpcClient + # waiters are unblocked even on failure. + await self.rpc_async_queue.put( + AbortResultMessage( + rpc_id=msg.rpc_id, + success=error is None, + error=error, + abort_outputs=abort_outputs or None, + ) + ) + elif abort_outputs: + for output_msg in abort_outputs: + await self.output_async_queue.put(output_msg) async def _handle_interaction(self, msg: InteractionMessage) -> None: """Handle a midway interaction for an active streaming diffusion request.""" @@ -841,13 +979,53 @@ class Orchestrator: ) ) - async def _abort_request_ids(self, request_ids: list[str]) -> None: - """Forward abort requests to all stage pools.""" + async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage]: + """Forward abort requests to all stage pools. + + Collects final-stage AR abort outputs (partial tokens) while request + state is still available for stage-id filtering. Diffusion pools do + not contribute abort outputs. + """ if not request_ids: - return + return [] + abort_outputs: list[OutputMessage] = [] for pool in self.stage_pools: - await pool.abort_requests(request_ids) + stage_outputs = await pool.abort_requests(request_ids) or [] + if bool(getattr(pool, "final_output", False)) and getattr(pool, "stage_type", None) != "diffusion": + final_output_type = getattr(pool.stage_client, "final_output_type", None) or "text" + for orch_req_id, request_output in stage_outputs: + req_state = self.request_states.get(orch_req_id) + if req_state is None: + continue + final_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} + if pool.stage_id not in final_stage_ids: + continue + engine_output = OmniRequestOutput.from_stage_output( + request_output, + request_id=orch_req_id, + finished=True, + stage_id=pool.stage_id, + replica_id=pool.get_bound_replica_id(orch_req_id), + final_output_type=final_output_type, + ) + abort_outputs.append( + OutputMessage( + request_id=orch_req_id, + stage_id=pool.stage_id, + replica_id=pool.get_bound_replica_id(orch_req_id), + engine_outputs=engine_output, + metrics=None, + finished=False, + stage_submit_ts=req_state.stage_submit_ts.get(pool.stage_id), + ) + ) pool.release_bindings(request_ids) + last_index_by_req: dict[str, int] = {} + for index, output_msg in enumerate(abort_outputs): + last_index_by_req[output_msg.request_id] = index + for index, output_msg in enumerate(abort_outputs): + output_msg.finished = index == last_index_by_req[output_msg.request_id] + return abort_outputs def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" @@ -909,11 +1087,102 @@ class Orchestrator: async def _orchestration_output_handler(self) -> None: """Poll all stages, handle transfers, send final outputs to main.""" try: - await self._orchestration_loop() + if self._event_driven_orch: + await self._orchestration_loop_event_driven() + else: + await self._orchestration_loop() except asyncio.CancelledError: logger.debug("[Orchestrator] _orchestration_output_handler cancelled") return + async def _process_llm_stage_outputs( + self, + stage_id: int, + replica_id: int, + raw_outputs: Any, + raw_terminal_request_ids: set[str], + ) -> list[Any]: + """Process one raw LLM poll result; returns processed request outputs. + + Shared by the legacy poll loop and the event-driven dispatcher so both + run the exact same per-output handling. Callers own the + ``EngineDeadError`` catch, since eviction needs the replica id, and + supply ``raw_terminal_request_ids`` — the per-poll accumulator this + fills for the caller to drain through + ``_finish_raw_terminal_requests`` once routing is done. + """ + pool = self.stage_pools[stage_id] + await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) + for eco in raw_outputs.outputs: + # Emit kv_wait_s before _handle_kv_ready_raw_outputs' + # async_chunk early-return so it lands in all modes. + kv_params = getattr(eco, "kv_transfer_params", None) + if ( + self._prom_metrics is not None + and isinstance(kv_params, dict) + and (kv_wait_s := kv_params.get("kv_wait_s")) is not None + ): + self._prom_metrics.observe_kv_wait( + kv_params.get("connector_type") or "unknown", + float(kv_wait_s), + ) + req_state = self.request_states.get(getattr(eco, "request_id", None)) + if req_state is None or not req_state.streaming.enabled: + continue + segment_finished = bool(getattr(eco, "is_segment_finished", False)) + raw_mm = self._completion_multimodal_output(eco, None) + req_state.streaming.segments[stage_id] = StreamingSegmentState( + finished=segment_finished, + token_ids=(self._coerce_int_list(getattr(eco, "new_token_ids", None)) if segment_finished else []), + output_metadata=(dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}), + ) + req_state.streaming.new_prompt_len_snapshot = getattr( + eco, + "new_prompt_len_snapshot", + None, + ) + if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state): + raw_terminal_request_ids.add(req_state.request_id) + iteration_stats = IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None + processed = await pool.process_llm_raw_outputs( + replica_id, + raw_outputs, + iteration_stats=iteration_stats, + ) + if self._stat_logger is not None and (raw_outputs.scheduler_stats is not None or iteration_stats is not None): + self._stat_logger.record( + raw_outputs.scheduler_stats, + iteration_stats, + engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)], + ) + _sched_stats = raw_outputs.scheduler_stats + if ( + self._prom_metrics is not None + and _sched_stats is not None + and getattr(_sched_stats, "num_waiting_reqs", None) is not None + ): + self._update_stage_replica_waiting( + stage_id, + replica_id, + int(_sched_stats.num_waiting_reqs), + ) + return processed + + def _absorb_diffusion_metrics(self, stage_id: int, replica_id: int, diffusion_output: Any) -> bool: + """Drain a diffusion output's piggybacked metrics. + + Returns ``True`` when the output carried nothing but metrics and must + not be routed. Shared by both orchestration loops: the event-driven + poller would otherwise hand a metrics-only sentinel to + ``_handle_processed_outputs`` as if it were a real request output. + """ + output_metrics = getattr(diffusion_output, "metrics", None) + if isinstance(output_metrics, dict): + n_waiting = output_metrics.pop(metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY, None) + if n_waiting is not None: + self._update_stage_replica_waiting(stage_id, replica_id, int(n_waiting)) + return diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID + async def _orchestration_loop(self) -> None: """Poll stage pools and route logical outputs.""" while not self._shutdown_event.is_set(): @@ -923,6 +1192,7 @@ class Orchestrator: for replica_id in pool.available_replica_ids(): if self._shutdown_event.is_set(): return + raw_terminal_request_ids: set[str] = set() # Shared catch so a dead replica on either poll path is # evicted rather than tearing down every stage (#4285). @@ -932,6 +1202,10 @@ class Orchestrator: if diffusion_output is None: continue + if self._absorb_diffusion_metrics(stage_id, replica_id, diffusion_output): + idle = False + continue + pool.record_output_timestamps([diffusion_output]) processed = [diffusion_output] else: @@ -939,46 +1213,9 @@ class Orchestrator: if raw_outputs is None: continue - await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) - for eco in raw_outputs.outputs: - req_state = self.request_states.get(getattr(eco, "request_id", None)) - if req_state is None or not req_state.streaming.enabled: - continue - req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False)) - req_state.streaming.segment_token_ids = ( - self._coerce_int_list(getattr(eco, "new_token_ids", None)) - if req_state.streaming.segment_finished - else [] - ) - raw_mm = self._completion_multimodal_output(eco, None) - req_state.streaming.segment_output_metadata = ( - dict(raw_mm) - if req_state.streaming.segment_finished and isinstance(raw_mm, dict) - else {} - ) - req_state.streaming.new_prompt_len_snapshot = getattr( - eco, - "new_prompt_len_snapshot", - None, - ) - if req_state.streaming.enabled: - await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state) - iteration_stats = ( - IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None + processed = await self._process_llm_stage_outputs( + stage_id, replica_id, raw_outputs, raw_terminal_request_ids ) - processed = await pool.process_llm_raw_outputs( - replica_id, - raw_outputs, - iteration_stats=iteration_stats, - ) - if self._stat_logger is not None and ( - raw_outputs.scheduler_stats is not None or iteration_stats is not None - ): - self._stat_logger.record( - raw_outputs.scheduler_stats, - iteration_stats, - engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)], - ) except asyncio.CancelledError: raise except EngineDeadError as e: @@ -995,6 +1232,7 @@ class Orchestrator: raise await self._handle_processed_outputs(stage_id, replica_id, processed) + await self._finish_raw_terminal_requests(stage_id, replica_id, raw_terminal_request_ids) idle = False self._orch_monitor.note_loop(idle=idle) @@ -1003,6 +1241,258 @@ class Orchestrator: else: await asyncio.sleep(0) + async def _orchestration_loop_event_driven(self) -> None: + """Event-driven variant of ``_orchestration_loop``. + + Selected by ``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``. One reader task per + available LLM stage replica awaits ``client.get_output_async()`` + directly — the same pattern vLLM's own ``AsyncLLM`` output handler uses + — and feeds a single dispatch queue. This coroutine consumes that queue + serially, so routing/handling semantics are identical to the legacy + loop; only the 1 ms poll cadence (and its per-tick ``asyncio.wait_for`` + task churn) is removed. Diffusion stages keep their nowait-poll contract + via a per-pool poller task feeding the same queue. + """ + ready_q: asyncio.Queue[tuple[str, int, int, Any]] = asyncio.Queue() + readers: dict[tuple[int, int], tuple[asyncio.Task, Any]] = {} + pollers: dict[int, asyncio.Task] = {} + reaped: list[asyncio.Task] = [] + + async def _llm_replica_reader(stage_id: int, replica_id: int, client: Any) -> None: + try: + while not self._shutdown_event.is_set(): + raw_outputs = await client.get_output_async() + # Same keep/drop rule as StagePool._poll_stage_raw: a batch + # with no request outputs still carries SchedulerStats on + # throttled ticks, and dropping it loses the KV/queue gauges + # for that interval. Only a fully empty batch is dropped. A + # poll-style client returns one instead of blocking, so keep + # the legacy 1 ms cadence for those; a blocking client only + # lands here rarely, where 1 ms is irrelevant. + if ( + not raw_outputs.outputs + and raw_outputs.scheduler_stats is None + and not raw_outputs.finished_requests + ): + await asyncio.sleep(0.001) + continue + await ready_q.put(("llm", stage_id, replica_id, raw_outputs)) + except asyncio.CancelledError: + raise + except BaseException as e: # noqa: BLE001 - routed to the dispatcher + await ready_q.put(("error", stage_id, replica_id, e)) + + async def _diffusion_poller(stage_id: int, pool: StagePool) -> None: + try: + while not self._shutdown_event.is_set(): + got = False + for replica_id in pool.available_replica_ids(): + try: + output = pool.poll_diffusion_output(replica_id) + except EngineDeadError as e: + await ready_q.put(("error", stage_id, replica_id, e)) + continue + if output is None: + continue + await ready_q.put(("diffusion", stage_id, replica_id, output)) + got = True + await asyncio.sleep(0 if got else 0.001) + except asyncio.CancelledError: + raise + except BaseException as e: # noqa: BLE001 - routed to the dispatcher + await ready_q.put(("error", stage_id, -1, e)) + + def _reconcile_readers() -> None: + """Track available replicas: spawn/reap LLM readers and diffusion pollers. + + Walks ``available_replica_ids()`` for parity with the legacy loop, + so a replica evicted by ``_handle_dead_replica`` (or marked + unavailable by membership) stops being read here too. + + Diffusion pollers are reconciled here rather than spawned once at + startup: a pool with no live replica yet reports ``stage_type is + None``, so a stage whose first replica registers at runtime would + otherwise never get a poller and its outputs would never drain. + """ + live: set[tuple[int, int]] = set() + live_pools: set[int] = set() + for stage_id in range(self.num_stages): + pool = self.stage_pools[stage_id] + stage_type = pool.stage_type + if stage_type is None: + # No live replica yet; nothing to attach to. A later + # reconcile picks the stage up once one registers. + continue + if stage_type == "diffusion": + live_pools.add(stage_id) + existing_poller = pollers.get(stage_id) + if existing_poller is None or existing_poller.done(): + pollers[stage_id] = asyncio.create_task( + _diffusion_poller(stage_id, pool), + name=f"orch-diffusion-poller-s{stage_id}", + ) + continue + for replica_id in pool.available_replica_ids(): + client = pool.clients[replica_id] + if client is None: + continue + key = (stage_id, replica_id) + live.add(key) + existing = readers.get(key) + if existing is None: + readers[key] = ( + asyncio.create_task( + _llm_replica_reader(stage_id, replica_id, client), + name=f"orch-reader-s{stage_id}r{replica_id}", + ), + client, + ) + continue + if existing[0].done(): + # The reader exited, which means it already queued its + # failure. Leave the slot alone until the dispatcher + # handles that error and evicts the replica; respawning + # now just re-raises the same error against the same + # dead client and queues a duplicate eviction. + continue + if existing[1] is client: + continue + # Client swapped underneath us: retire the old reader. + existing[0].cancel() + reaped.append(existing[0]) + readers[key] = ( + asyncio.create_task( + _llm_replica_reader(stage_id, replica_id, client), + name=f"orch-reader-s{stage_id}r{replica_id}", + ), + client, + ) + for key in list(readers): + if key not in live: + task, _ = readers.pop(key) + task.cancel() + reaped.append(task) + for stage_id in list(pollers): + if stage_id not in live_pools: + task = pollers.pop(stage_id) + task.cancel() + reaped.append(task) + + _reconcile_readers() + logger.info( + "[Orchestrator] Event-driven orchestration loop enabled (%s): %d LLM replica readers, %d diffusion pollers", + _EVENT_DRIVEN_ORCH_ENV, + len(readers), + len(pollers), + ) + + shutdown_task = asyncio.create_task(self._shutdown_event.wait(), name="orch-shutdown-wait") + pending_get: asyncio.Task | None = None + next_reconcile = _time.monotonic() + _ORCH_READER_RECONCILE_INTERVAL_S + try: + while not self._shutdown_event.is_set(): + if pending_get is None: + pending_get = asyncio.create_task(ready_q.get(), name="orch-dispatch-get") + done, _ = await asyncio.wait( + {pending_get, shutdown_task}, + timeout=max(0.0, next_reconcile - _time.monotonic()), + return_when=asyncio.FIRST_COMPLETED, + ) + if shutdown_task in done: + return + # Reconcile on a wall-clock schedule, not only when the queue + # goes idle: under continuous traffic `asyncio.wait` returns on + # every output and a timeout-only reconcile would never fire, + # so a replica that registers at runtime would never get a + # reader and its outputs would never drain. + if _time.monotonic() >= next_reconcile: + next_reconcile = _time.monotonic() + _ORCH_READER_RECONCILE_INTERVAL_S + _reconcile_readers() + if not done: + self._orch_monitor.note_loop(idle=True) + continue + kind, stage_id, replica_id, payload = pending_get.result() + pending_get = None + + if kind == "error": + # replica_id < 0 means the failure was raised by a poller + # task itself rather than attributed to one replica, so + # there is nothing to evict -- and evict_replica() rejects + # an out-of-range id, which would turn a survivable death + # into a ValueError out of the dispatcher. + if isinstance(payload, EngineDeadError) and replica_id >= 0: + # Same per-replica isolation as the legacy loop (#4285): + # evict and keep serving. Reconcile immediately so the + # dead replica's reader is reaped now rather than at the + # next scheduled tick. Skip a replica already evicted by + # an earlier error from the same slot -- evict_replica() + # would shut a None client down a second time. + if replica_id in self.stage_pools[stage_id].live_replica_ids(): + await self._handle_dead_replica(stage_id, replica_id, payload) + _reconcile_readers() + continue + if self._shutdown_event.is_set(): + return + logger.error( + "[Orchestrator] Stage-%s replica-%s reader failed: %r", + stage_id, + replica_id, + payload, + ) + raise payload + + raw_terminal_request_ids: set[str] = set() + try: + if kind == "diffusion": + pool = self.stage_pools[stage_id] + if self._absorb_diffusion_metrics(stage_id, replica_id, payload): + self._orch_monitor.note_loop(idle=False) + continue + pool.record_output_timestamps([payload]) + processed = [payload] + else: + processed = await self._process_llm_stage_outputs( + stage_id, replica_id, payload, raw_terminal_request_ids + ) + except asyncio.CancelledError: + raise + except EngineDeadError as e: + if replica_id in self.stage_pools[stage_id].live_replica_ids(): + await self._handle_dead_replica(stage_id, replica_id, e) + _reconcile_readers() + continue + except Exception: + if self._shutdown_event.is_set(): + return + logger.exception( + "[Orchestrator] Stage-%s replica-%s processing failed", + stage_id, + replica_id, + ) + raise + + await self._handle_processed_outputs(stage_id, replica_id, processed) + await self._finish_raw_terminal_requests(stage_id, replica_id, raw_terminal_request_ids) + # Mirrors the legacy loop, which sets idle=False only after a + # poll produced outputs that were routed -- an evicted replica + # `continue`s above without marking the tick active. + self._orch_monitor.note_loop(idle=False) + finally: + shutdown_task.cancel() + if pending_get is not None: + pending_get.cancel() + reader_tasks = [task for task, _ in readers.values()] + poller_tasks = list(pollers.values()) + for task in (*reader_tasks, *poller_tasks): + task.cancel() + # `reaped` holds readers/pollers retired mid-run (client swap, + # eviction, stage teardown). They were cancelled but never awaited, + # so gather them here too rather than leaving dangling tasks. + cleanup = [shutdown_task, *reader_tasks, *poller_tasks, *reaped] + if pending_get is not None: + cleanup.append(pending_get) + await asyncio.gather(*cleanup, return_exceptions=True) + async def _handle_processed_outputs(self, stage_id: int, replica_id: int, outputs: list[Any]) -> None: """Route processed stage outputs produced by one stage poll.""" pool = self.stage_pools[stage_id] @@ -1022,7 +1512,7 @@ class Orchestrator: continue stage_metrics = None - segment_finished = req_state.streaming.enabled and req_state.streaming.segment_finished + segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished if output.finished or segment_finished: stage_metrics = pool.build_stage_metrics( [output], @@ -1056,6 +1546,35 @@ class Orchestrator: close_duplex_sessions=True, ) + def _update_stage_replica_waiting(self, stage_id: int, replica_id: int, n_waiting: int) -> None: + """Update one replica snapshot and expose the stage-wide total.""" + self._stage_replica_waiting[(stage_id, replica_id)] = max(int(n_waiting), 0) + self._set_stage_waiting_total(stage_id) + self._sync_engines_waiting_counter() + + def _remove_stage_replica_waiting(self, stage_id: int, replica_id: int) -> None: + """Drop a dead replica's snapshot and refresh the stage total.""" + self._stage_replica_waiting.pop((stage_id, replica_id), None) + self._set_stage_waiting_total(stage_id) + self._sync_engines_waiting_counter() + + def _sync_engines_waiting_counter(self) -> None: + """Aggregate per-replica waiting into the counter shared with the + frontend so engine-queued requests show as waiting, not running.""" + counter = self._engines_waiting_counter + if counter is not None: + counter.value = sum(self._stage_replica_waiting.values()) + + def _set_stage_waiting_total(self, stage_id: int) -> None: + if self._prom_metrics is None: + return + total = sum( + n_waiting + for (snapshot_stage_id, _), n_waiting in self._stage_replica_waiting.items() + if snapshot_stage_id == stage_id + ) + self._prom_metrics.set_stage_waiting_requests(stage_id, total) + async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: EngineDeadError) -> None: """Evict a dead stage replica and fail the requests stranded on it (#4285). @@ -1074,6 +1593,7 @@ class Orchestrator: error, ) pool.evict_replica(replica_id) + self._remove_stage_replica_waiting(stage_id, replica_id) stage_has_live = bool(pool.live_replica_ids()) failed_ids: list[str] = [] for req_id, req_state in list(self.request_states.items()): @@ -1128,6 +1648,35 @@ class Orchestrator: abort=True, ) + async def _fail_request_client_error( + self, + req_id: str, + stage_id: int, + error: str, + *, + close_duplex_sessions: bool = False, + ) -> None: + """Fail one request with a non-fatal 400 (bad input, engine survives). + + The non-fatal counterpart of `_fail_request_dead_stage`: emits a + client-error ErrorMessage (default `fatal=False`) so the engine keeps + serving, then releases the request's state across every stage pool. + """ + await self.output_async_queue.put( + ErrorMessage( + error=error, + status_code=HTTPStatus.BAD_REQUEST.value, + error_type=DEFAULT_CLIENT_ERROR_TYPE, + request_id=req_id, + stage_id=stage_id, + ) + ) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + abort=True, + close_duplex_sessions=close_duplex_sessions, + ) + async def _dispatch_or_fail_request( self, dispatch: Callable[[], Awaitable[Any]], @@ -1149,7 +1698,9 @@ class Orchestrator: StageUnavailableError (no live replica / evicted slot) or EngineDeadError (replica died mid-submit before the poll loop evicted it). Both mean "this request cannot be placed", not "the server is - broken": fail the request and keep serving (#4285). Unrelated errors + broken": fail the request and keep serving (#4285). An OverflowError + from encoding an out-of-range request value is likewise failed as a + client error (400) rather than killing the loop. Other unrelated errors propagate. Returns False when the request was failed. ``req_id`` is the request the failure is attributed to; ``dispatch_req_id`` @@ -1192,6 +1743,24 @@ class Orchestrator: if dead_replica is not None and dead_replica in pool.live_replica_ids(): await self._handle_dead_replica(stage_id, dead_replica, e) return False + except OverflowError as e: + # Catch overflow errors in the request payload to ensure that msgpack + # encoding can't kill the engine; instead, we just fail this request. + # The frontend should generally validate this as well, but this is needed + # in case we get values like seed > 2**64-1 from sampling params. + logger.warning( + "[Orchestrator] %s dispatch for req=%s hit an unserializable value on stage-%s: %s", + operation, + req_id, + stage_id, + e, + ) + await self._fail_request_client_error( + req_id, + stage_id, + f"Request contains a value that cannot be serialized: {e}", + ) + return False # ---- Shared helpers ---- @@ -1201,7 +1770,7 @@ class Orchestrator: *, abort: bool = False, close_duplex_sessions: bool = False, - ) -> None: + ) -> list[OutputMessage]: """Release pool bindings and logical request state for the given ids. CFG-aware: cleaning a parent releases its tracker state and pulls its @@ -1210,9 +1779,13 @@ class Orchestrator: can never complete). Every teardown path (stage error, abort, replica loss, membership unregister) funnels through here, so tracker state cannot outlive its requests. + + Returns: + Final-stage AR abort ``OutputMessage`` list when ``abort=True``; + otherwise an empty list. """ if not request_ids: - return + return [] cleanup_ids = list(dict.fromkeys(request_ids)) batch = set(cleanup_ids) @@ -1254,13 +1827,16 @@ class Orchestrator: cleanup_ids.extend(stale_request_ids) cleanup_ids = list(dict.fromkeys(cleanup_ids)) + abort_outputs: list[OutputMessage] = [] try: if abort: - await self._abort_request_ids(cleanup_ids) + abort_outputs = await self._abort_request_ids(cleanup_ids) self._release_request_bindings(cleanup_ids) for request_id in cleanup_ids: self._pd_kv_params.pop(request_id, None) req_state = self.request_states.pop(request_id, None) + if req_state is not None: + cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ())) if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: self._running_counter.decrement() req_state.running_counter_registered = False @@ -1270,14 +1846,15 @@ class Orchestrator: raise if closing_session_ids and self.duplex_control_plane is not None: self.duplex_control_plane.finalize_closed_sessions(closing_session_ids) + return abort_outputs async def _apply_raw_terminal_stage_finish( self, stage_id: int, eco: Any, req_state: OrchestratorRequestState, - ) -> None: - """Record session-level finish markers dropped by the streaming output processor. + ) -> bool: + """Record a session-level finish marker and report whether it was terminal. Streaming segment stops set ``is_segment_finished=True`` and are handled via processed outputs. Session termination (e.g. ``finish_requests`` after @@ -1290,14 +1867,60 @@ class Orchestrator: outputs after stage-0 session end. """ if getattr(eco, "finish_reason", None) is None: - return + return False if getattr(eco, "is_segment_finished", False): - return + return False final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} if stage_id not in final_output_stage_ids: - return + return False req_state.finished_final_output_stage_ids.add(stage_id) + return True + + async def _finish_raw_terminal_requests( + self, + stage_id: int, + replica_id: int, + request_ids: set[str], + ) -> None: + """Finish streaming requests whose raw terminal had no processed output.""" + pool = self.stage_pools[stage_id] + if not pool.final_output: + return + + for request_id in request_ids: + req_state = self.request_states.get(request_id) + if req_state is None or self._is_duplex_session_request(req_state): + continue + + final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} + if not final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids): + continue + + final_output_type = getattr(pool.stage_client, "final_output_type", None) + logger.info( + "[Orchestrator] req=%s stage-%s raw terminal produced no processed terminal; returning empty %s output", + request_id, + stage_id, + final_output_type or "text", + ) + terminal_output = _build_terminal_empty_output( + request_id, + final_output_type=final_output_type, + audio_sample_rate=pool._infer_audio_sample_rate(), + ) + await self.output_async_queue.put( + OutputMessage( + request_id=request_id, + stage_id=stage_id, + replica_id=replica_id, + engine_outputs=terminal_output, + metrics=None, + finished=True, + stage_submit_ts=req_state.stage_submit_ts.get(stage_id), + ) + ) + await self._cleanup_request_ids([request_id, *self._cfg_tracker.cleanup_parent(request_id)]) def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" @@ -1356,6 +1979,7 @@ class Orchestrator: req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) + segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished # CFG companion: stash output so parent can bundle [parent, *companions] # into source_outputs for the bridge (e.g. thinker2imagegen). if finished and self._cfg_tracker.is_companion(req_id): @@ -1365,11 +1989,7 @@ class Orchestrator: return request_finished = False - if ( - finished - and self.stage_pools[stage_id].final_output - and not (req_state.streaming.enabled and req_state.streaming.segment_finished) - ): + if finished and self.stage_pools[stage_id].final_output and not segment_finished: req_state.finished_final_output_stage_ids.add(stage_id) final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) @@ -1381,7 +2001,7 @@ class Orchestrator: # filter out again (the official implementation returns exactly one # result per audio chunk). is_duplex_stage0_segment = ( - stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished + stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment(0).finished ) if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment: await self.output_async_queue.put( @@ -1393,7 +2013,10 @@ class Orchestrator: metrics=stage_metrics, finished=( request_finished - or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished) + or ( + self._is_duplex_session_request(req_state) + and req_state.streaming.segment(stage_id).finished + ) ), stage_submit_ts=submit_ts, ) @@ -1428,7 +2051,7 @@ class Orchestrator: return if ( - (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished)) + (finished or segment_finished) and stage_id < req_state.final_stage_id and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1)) and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled) @@ -1567,7 +2190,9 @@ class Orchestrator: def _duplex_output_context( req_state: OrchestratorRequestState, *, - stage_id: int | None = None, + # Required: segment state is stage-keyed, so a caller that omitted this + # would silently read "no boundary" rather than the intended stage's. + stage_id: int | None, ) -> DuplexOutputContext | None: identity = req_state.duplex_identity if identity is None: @@ -1578,15 +2203,16 @@ class Orchestrator: ) fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence + segment = req_state.streaming.segment(stage_id) return DuplexOutputContext( identity=DuplexRequestIdentity( session_id=identity.session_id, fence=fence, ), final_stage_id=req_state.final_stage_id, - segment_finished=req_state.streaming.enabled and req_state.streaming.segment_finished, - segment_token_ids=tuple(req_state.streaming.segment_token_ids), - segment_output_metadata=req_state.streaming.segment_output_metadata, + segment_finished=req_state.streaming.enabled and segment.finished, + segment_token_ids=tuple(segment.token_ids), + segment_output_metadata=segment.output_metadata, ) @staticmethod @@ -2122,13 +2748,28 @@ class Orchestrator: req_state.prompt, streaming_context=req_state.streaming, ) - except Exception: + except Exception as exc: logger.exception( "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s", req_id, next_logical, ) - raise + if not self._is_duplex_session_request(req_state): + raise + await self.output_async_queue.put( + ErrorMessage( + request_id=req_id, + stage_id=next_logical, + error=f"Stage-{next_logical} input processor failed: {type(exc).__name__}: {exc}", + error_type=type(exc).__name__, + ) + ) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + abort=True, + close_duplex_sessions=True, + ) + return finally: req_state.streaming.source_token_decoder = previous_decoder @@ -2217,18 +2858,44 @@ class Orchestrator: request_id: str, stage0_request: Any, req_state: OrchestratorRequestState, - ) -> None: - """Pre-submit downstream stages for async-chunk mode.""" + ) -> bool: + """Pre-submit downstream stages for async-chunk mode. + + Returns False when the request was failed and cleaned up in here, so a + caller still holding ``req_state`` stops instead of recording state on + an object the cleanup already popped from ``request_states``. + """ if req_state.final_stage_id <= 0: - return + return True prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None) if prompt_token_ids is None: - logger.warning( - "[Orchestrator] async_chunk prewarm skipped for req=%s: stage0 prompt_token_ids missing", + # R1.3 of #4855. Skipping the prewarm leaves every downstream stage + # unsubmitted, so the request produces no output and no error -- it + # just stops. An embeds-only prompt in async-chunk mode is a config + # mistake, and the client should hear about it now rather than wait + # for the stage-input deadline to notice nothing ever arrived. + # + # Non-fatal, with a 4xx: `fatal=True` means "the engine is dead", and + # the entrypoints act on it -- `OmniBase._handle_output_message` + # raises `OmniEngineDeadError` out of `OmniLLM.generate`'s batch loop, + # failing every other request in that batch, and the serving layer + # calls `terminate_if_errored`. One client's bad prompt must not do + # either. This is the same shape `_handle_stage_error` uses for a + # request-scoped client error. + logger.error( + "[Orchestrator] req=%s: async_chunk prewarm needs stage0 prompt_token_ids " + "and none were provided; failing the request", request_id, ) - return + await self._fail_request_client_error( + request_id, + 0, + "async_chunk requires prompt_token_ids on the stage-0 request; " + "an embeds-only prompt cannot prewarm downstream stages", + close_duplex_sessions=True, + ) + return False for next_stage_id in range(1, req_state.final_stage_id + 1): next_pool = self.stage_pools[next_stage_id] @@ -2303,7 +2970,7 @@ class Orchestrator: # Attribute the failure to the stage that failed, not stage 0; the # request is already cleaned up, so stop prewarming. if not submitted: - return + return False # The guard swallows submit_initial's return value; the pool binding # it recorded carries the same replica. @@ -2329,6 +2996,8 @@ class Orchestrator: tx_ms=_tx_ms, ) + return True + def _build_kv_sender_info( self, sender_stage_ids: list[int], diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 74f65d8c2..06110e005 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -35,6 +35,7 @@ class StageClient(Protocol): final_output: bool final_output_type: FinalOutputModalityType | None default_sampling_params: OmniSamplingParams + prompt_transform_func: Any | None prompt_expand_func: Any | None requires_multimodal_data: bool custom_process_input_func: Any | None diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 5fbbbc46d..cef4dbf62 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -143,6 +143,7 @@ class StageEngineCoreClientBase(StageClientBase): self.final_output = metadata.final_output self.final_output_type = metadata.final_output_type self.default_sampling_params = metadata.default_sampling_params + self.prompt_transform_func = metadata.prompt_transform_func self.prompt_expand_func = metadata.prompt_expand_func self.custom_process_input_func = metadata.custom_process_input_func diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 3ac719a6b..4bcf47d3b 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -31,7 +31,7 @@ from vllm.v1.engine.utils import ( from vllm_omni.distributed.omni_coordinator import create_stage_coord_client from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.stage_init_utils import ( - maybe_apply_audex_cfg_patches, + maybe_apply_cfg_scheduler_patches, set_death_signal, ) @@ -134,9 +134,10 @@ class StageEngineCoreProc(EngineCoreProc): _vllm_engine_core_module.EngineCoreRequest, ) - # Audex CFG scheduler patches must land before EngineCore builds - # its Scheduler; gated on the stage's logits_processors config. - maybe_apply_audex_cfg_patches(kwargs.get("vllm_config")) + # CFG pairing scheduler patches must land before EngineCore builds + # its Scheduler; gated on the stage's logits_processors and its + # default sampling extra_args. + maybe_apply_cfg_scheduler_patches(kwargs.get("vllm_config")) engine_core = StageEngineCoreProc( *args, diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 5dacc31df..5abcbd249 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Helpers for launching and handshaking omni engine cores.""" from __future__ import annotations @@ -22,6 +25,7 @@ from vllm.utils.network_utils import get_open_ports_list, zmq_socket_ctx from vllm.v1.engine.coordinator import DPCoordinator from vllm.v1.engine.utils import ( CoreEngine, + CoreEngineLaunch, CoreEngineProcManager, EngineZmqAddresses, get_engine_zmq_addresses, @@ -833,13 +837,16 @@ def connect_remote_engine_cores( ) wait_for_engine_startup( handshake_socket, - addresses, engines_to_handshake, vllm_config.parallel_config, False, # coordinated_dp vllm_config.cache_config, - None, # proc_manager (remote — no local procs) - None, # coord_process + CoreEngineLaunch( + engine_manager=None, # remote — no local procs + coordinator=None, + addresses=addresses, + tensor_queue=None, + ), ) finally: omni_master_server.release_route_port_reservations( @@ -875,13 +882,16 @@ def connect_remote_diffusion_proc( yield StageReplicaResources(addresses=addresses) wait_for_engine_startup( handshake_socket, - addresses, [CoreEngine(index=0, local=False)], _single_diffusion_parallel_config(local_client=False), False, None, - None, - None, + CoreEngineLaunch( + engine_manager=None, # remote — no local procs + coordinator=None, + addresses=addresses, + tensor_queue=None, + ), ) finally: omni_master_server.release_route_port_reservations( @@ -1077,13 +1087,16 @@ def _launch_omni_core_engines( yield local_engine_manager, coordinator, addresses wait_for_engine_startup( handshake_socket, - addresses, engines_to_handshake, parallel_config, parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe, vllm_config.cache_config, - local_engine_manager, - coordinator.proc if coordinator else None, + CoreEngineLaunch( + engine_manager=local_engine_manager, + coordinator=coordinator, + addresses=addresses, + tensor_queue=None, + ), ) finally: omni_master_server.release_route_port_reservations( @@ -1166,13 +1179,16 @@ def launch_stage_replica( ) wait_for_engine_startup( handshake_socket, - addresses, engines_to_handshake, vllm_config.parallel_config, False, # coordinated_dp vllm_config.cache_config, - engine_manager, - None, # coordinator_proc + CoreEngineLaunch( + engine_manager=engine_manager, + coordinator=None, + addresses=addresses, + tensor_queue=None, + ), ) @@ -1530,7 +1546,6 @@ def launch_diffusion_stage_replica( stage_config: Any, metadata: Any, stage_init_timeout: int, - batch_size: int, use_inline: bool, replica_id: int = 0, omni_master_server: OmniMasterServer | None = None, @@ -1549,7 +1564,6 @@ def launch_diffusion_stage_replica( stage_config, metadata, stage_init_timeout=stage_init_timeout, - batch_size=batch_size, use_inline=use_inline, ) return client, StageReplicaResources() @@ -1608,7 +1622,6 @@ def launch_diffusion_stage_replica( request_address=proc_manager.addresses.inputs[0], response_address=proc_manager.addresses.outputs[0], proc_manager=proc_manager, - batch_size=batch_size, ) return client, StageReplicaResources( manager=proc_manager, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index c2e1ab5bc..6fdcac81b 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Stage initialization helpers for vLLM-Omni multi-stage runtime. @@ -11,24 +14,37 @@ from __future__ import annotations import copy import fcntl import importlib +import json import multiprocessing as mp import os import time -from collections.abc import Callable, Generator, Sequence +from collections.abc import Callable, Collection, Generator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass, fields, replace +from pathlib import Path from typing import Any, Literal, cast +import regex as re from vllm.logger import init_logger from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams from vllm.tokenizers import cached_tokenizer_from_config +from vllm.transformers_utils.repo_utils import hf_api +from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.input_processor import InputProcessor from vllm.v1.executor import Executor from vllm_omni.config.omni_config import ( + _CACHE_STAGE_ENGINE_FIELD_MAP, + _DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP, + _DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP, + _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS, + _DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP, + _LOAD_STAGE_ENGINE_FIELD_MAP, + _PARALLEL_CONFIG_ENGINE_FIELD_MAP, + _SCHEDULER_STAGE_ENGINE_FIELD_MAP, BaseVllmOmniStageConfig, VllmOmniDiffusionStageConfig, ) @@ -71,22 +87,144 @@ class LogicalStageInitPlan: replicas: list[ReplicaInitPlan] -def _resolve_model_to_local_path(model: str) -> str: - """Resolve an HF Hub model ID to a local cache path.""" +def _missing_stage_subdirs(base: str, subdirs: Sequence[str]) -> list[str]: + """Return the entries of ``subdirs`` that are not directories under ``base``.""" + return [subdir for subdir in subdirs if not os.path.isdir(os.path.join(base, subdir))] + + +# Artifacts that make a snapshot subfolder trustworthy. A directory that +# exists but holds none of these is an interrupted download, not a snapshot; +# treating it as complete strips the Hub fallback vLLM would need later. +_WEIGHT_ARTIFACT_PATTERNS = ("*.safetensors", "*.bin", "*.pt", "*.gguf") +# Vocabulary-bearing files. Configs and chat templates are small and download +# first, so their presence alone cannot distinguish a tokenizer folder from an +# interrupted download. +_TOKENIZER_ARTIFACT_NAMES = ( + "tokenizer.json", + "tokenizer.model", + "spiece.model", + "sentencepiece.bpe.model", + "vocab.json", + "vocab.txt", +) +# HF sharded checkpoints name their pieces `-NNNNN-of-NNNNN.` and +# always ship an index; a shard-named file without one is a partial download. +_SHARD_NAME_RE = re.compile(r"-\d+-of-\d+\.(safetensors|bin)$") + + +def _indexed_shards_complete(folder: Path) -> bool | None: + """Check sharded weights against their index; ``None`` when no index exists.""" + indexes = list(folder.rglob("*.index.json")) + if not indexes: + return None + for index in indexes: + try: + weight_map = json.loads(index.read_text()).get("weight_map") or {} + except (OSError, ValueError): + return False + shards = set(weight_map.values()) + if not shards or any(not (index.parent / shard).is_file() for shard in shards): + return False + return True + + +def _subdir_is_populated(base: str, subdir: str, needs_weights: bool) -> bool: + folder = Path(base) / subdir + if not folder.is_dir(): + return False + if needs_weights: + indexed = _indexed_shards_complete(folder) + if indexed is not None: + return indexed + weights = [path for pattern in _WEIGHT_ARTIFACT_PATTERNS for path in folder.rglob(pattern)] + if not weights: + return False + return not any(_SHARD_NAME_RE.search(path.name) for path in weights) + return any((folder / name).is_file() for name in _TOKENIZER_ARTIFACT_NAMES) + + +def _incomplete_stage_subdirs( + base: str, + subdirs: Sequence[str], + weight_subdirs: Collection[str] = (), +) -> list[str]: + """Return the entries of ``subdirs`` without a POPULATED directory under ``base``. + + Hub-snapshot paths use this stricter check: ``os.path.isdir`` alone accepts + a subfolder holding only config.json from an interrupted download, and the + warm-cache early return would then convert the Hub ID into a local path + vLLM cannot fetch missing weights for. ``weight_subdirs`` names the entries + that must contain a weight artifact, not merely any file. + """ + return [ + subdir for subdir in subdirs if not _subdir_is_populated(base, subdir, needs_weights=subdir in weight_subdirs) + ] + + +def _resolve_model_to_local_path( + model: str, + required_subdirs: Sequence[str] = (), + weight_subdirs: Collection[str] = (), + *, + revision: str | None = None, + download_dir: str | None = None, +) -> str: + """Resolve an HF Hub model ID to a local path that holds ``required_subdirs``. + + ``snapshot_download(local_files_only=True)`` returns the snapshot root as + soon as *any* file of the repo is cached, even when the subfolders this + stage needs were never materialized. Joining a stage subdir onto such a + root produces a path that exists nowhere, and upstream ``EngineArgs`` + forwards a non-directory ``model`` to HuggingFace as a repo id, which fails + with an ``HFValidationError`` about the cache path. Verify the subfolders + here and pull just the missing ones, so the join always lands on a real + directory or raises an error that names what is missing. + + ``revision`` and ``download_dir`` mirror the engine args of the same name: + once the repo ID is replaced by a local path, downstream ModelConfig can no + longer correct either, so they must shape the snapshot selection here. + """ if os.path.isdir(model): return model + # Keep the warm-cache path offline-friendly: no Hub round trip when the + # stage's subfolders are already there. try: - from huggingface_hub import snapshot_download - - # Keep init path resolution offline-friendly. - return snapshot_download(model, local_files_only=True) - except Exception: - logger.warning( - "[stage_init] Could not resolve %s to local snapshot; using as-is", - model, + cached_root: str | None = hf_api().snapshot_download( + model, local_files_only=True, revision=revision, cache_dir=download_dir ) - return model + except Exception: + cached_root = None + if cached_root is not None and not _incomplete_stage_subdirs(cached_root, required_subdirs, weight_subdirs): + return cached_root + + # Cold cache, or a snapshot root whose stage subfolders were never (or + # only partially) downloaded: pull exactly the subfolders this stage asked + # for. snapshot_download resumes a partial subfolder for free. + allow_patterns = [f"{subdir.strip('/')}/*" for subdir in required_subdirs] or None + try: + resolved = hf_api().snapshot_download( + model, allow_patterns=allow_patterns, revision=revision, cache_dir=download_dir + ) + except Exception as exc: + raise RuntimeError( + f"[stage_init] Could not resolve {model!r} to a local snapshot containing " + f"{sorted(required_subdirs)}: the download failed and " + + ( + f"the cached snapshot {cached_root!r} is missing or incomplete for " + f"{sorted(_incomplete_stage_subdirs(cached_root, required_subdirs, weight_subdirs))}." + if cached_root is not None + else "nothing is cached locally." + ) + ) from exc + + missing = _incomplete_stage_subdirs(resolved, required_subdirs, weight_subdirs) + if missing: + raise RuntimeError( + f"[stage_init] Snapshot {resolved!r} for {model!r} has no populated {sorted(missing)} " + "subfolder; the stage cannot be initialized from it." + ) + return resolved def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> str: @@ -96,14 +234,66 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s if model_subdir is None and tokenizer_subdir is None: return model - resolved_base = _resolve_model_to_local_path(model) + revision = engine_args.get("revision") + tokenizer_revision = engine_args.get("tokenizer_revision") + download_dir = engine_args.get("download_dir") + # A tokenizer pinned to a different revision cannot come from the model's + # snapshot; resolve it against its own. An empty subdir means the snapshot + # root and still needs its own revision. + split_tokenizer = tokenizer_subdir is not None and tokenizer_revision is not None and tokenizer_revision != revision + + required_subdirs = [subdir for subdir in (model_subdir, tokenizer_subdir) if subdir] + model_required = [subdir for subdir in (model_subdir,) if subdir] if split_tokenizer else required_subdirs + weight_subdirs = frozenset(subdir for subdir in (model_subdir,) if subdir) + if is_runai_obj_uri(model): + # Object-storage URIs stay opaque until each stage builds its own + # ModelConfig, so the joins below are resolved by vLLM's streamer + # rather than by the local filesystem. + resolved_base = model + tokenizer_base = model + else: + resolved_base = _resolve_model_to_local_path( + model, model_required, weight_subdirs, revision=revision, download_dir=download_dir + ) + # Reachable for a local model directory; the Hub branch above has + # already failed closed on a missing subfolder. + missing = _missing_stage_subdirs(resolved_base, model_required) + if missing: + raise RuntimeError( + f"[stage_init] Model directory {resolved_base!r} has no {sorted(missing)} " + "subfolder; the stage cannot be initialized from it." + ) + if split_tokenizer: + # An empty subdir targets the snapshot root, which cannot be + # subset by allow_patterns; resolve the whole revision. + tokenizer_required = [tokenizer_subdir] if tokenizer_subdir else [] + tokenizer_base = _resolve_model_to_local_path( + model, tokenizer_required, revision=tokenizer_revision, download_dir=download_dir + ) + missing = _missing_stage_subdirs(tokenizer_base, tokenizer_required) + if missing: + raise RuntimeError( + f"[stage_init] Tokenizer directory {tokenizer_base!r} has no {sorted(missing)} " + "subfolder; the stage cannot be initialized from it." + ) + if not tokenizer_required and not _subdir_is_populated(tokenizer_base, "", False): + # An empty subdir means the tokenizer lives at the snapshot + # root, so there is no subfolder for the check above to look + # at and any resolved root would otherwise pass. Require the + # vocabulary artifacts themselves. + raise RuntimeError( + f"[stage_init] Tokenizer directory {tokenizer_base!r} holds none of " + f"{list(_TOKENIZER_ARTIFACT_NAMES)}; the stage cannot be initialized from it." + ) + else: + tokenizer_base = resolved_base if model_subdir: model = os.path.join(resolved_base, model_subdir) logger.info("[stage_init] Using model subdirectory: %s", model) if tokenizer_subdir is not None: - tokenizer_path = os.path.join(resolved_base, tokenizer_subdir) if tokenizer_subdir else resolved_base + tokenizer_path = os.path.join(tokenizer_base, tokenizer_subdir) if tokenizer_subdir else tokenizer_base engine_args["tokenizer"] = tokenizer_path logger.info("[stage_init] Using tokenizer from: %s", tokenizer_path) elif model_subdir and "tokenizer" not in engine_args: @@ -114,6 +304,23 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s return model +def _resolve_model_path(model: str, engine_args: dict[str, Any]) -> str: + """Apply a model-owned path resolver from stage engine args.""" + resolver_path = engine_args.pop("model_path_resolver", None) + if resolver_path is None: + return model + resolver = _resolve_omni_metadata_hook(str(resolver_path)) + if resolver is None: + return model + return str( + resolver( + model, + engine_args.get("revision"), + engine_args.get("task_type"), + ) + ) + + def apply_cli_tokenizer( engine_args: dict[str, Any], *, @@ -347,6 +554,7 @@ class StageMetadata: custom_process_input_func: Callable | None model_stage: str | None runtime_cfg: Any + prompt_transform_func: Callable | None = None prompt_expand_func: Callable | None = None cfg_kv_collect_func: Callable | None = None # Multi-replica: replica_id distinguishes replicas of the same stage. @@ -416,6 +624,12 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: mod_path, fn_name = _cpif_path.rsplit(".", 1) custom_process_input_func = getattr(importlib.import_module(mod_path), fn_name) + prompt_transform_func: Callable | None = None + _ptf_path = _get_attr_or_item(stage_config, "prompt_transform_func") + if _ptf_path: + _mod, _fn = _ptf_path.rsplit(".", 1) + prompt_transform_func = getattr(importlib.import_module(_mod), _fn) + prompt_expand_func: Callable | None = None _pef_path = _get_attr_or_item(stage_config, "prompt_expand_func") if _pef_path: @@ -444,6 +658,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: custom_process_input_func=custom_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, + prompt_transform_func=prompt_transform_func, cfg_kv_collect_func=cfg_kv_collect_func, ) @@ -464,6 +679,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: custom_process_input_func=custom_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, + prompt_transform_func=prompt_transform_func, prompt_expand_func=prompt_expand_func, ) @@ -510,6 +726,7 @@ def extract_stage_metadata_from_omni_stage_config( custom_process_input_func=custom_process_input_func, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, + prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), cfg_kv_collect_func=_resolve_omni_metadata_hook(stage_config.cfg_kv_collect_func), ) @@ -526,6 +743,7 @@ def extract_stage_metadata_from_omni_stage_config( custom_process_input_func=custom_process_input_func, model_stage=stage_config.model_stage, runtime_cfg=stage_config.runtime_config, + prompt_transform_func=_resolve_omni_metadata_hook(stage_config.prompt_transform_func), prompt_expand_func=_resolve_omni_metadata_hook(stage_config.prompt_expand_func), ) @@ -557,27 +775,32 @@ def _maybe_set_qwen3_omni_moe_env(engine_args_dict: dict[str, Any]) -> None: def split_devices_for_replicas( devices_str: str | None, num_replicas: int, - tp_size: int, + devices_per_replica: int, stage_id: int, -) -> list[str]: +) -> list[str | None]: """Split a devices string into per-replica subsets. + The result always has one entry per replica, because callers index it by + replica id. When ``devices_str`` is ``None`` the stage declares no explicit + placement, so every replica gets ``None`` and inherits the launcher's + ``CUDA_VISIBLE_DEVICES``. + When ``num_replicas`` is 1, returns ``[devices_str]`` unchanged. Otherwise, two YAML shapes are accepted: - 1. **Legacy / pool mode** — ``len(devices) == num_replicas * tp_size``: + 1. **Legacy / pool mode** — ``len(devices) == num_replicas * devices_per_replica``: the string enumerates the full per-stage pool. Each replica gets - ``tp_size`` consecutive entries. The values are logical indices + ``devices_per_replica`` consecutive entries. The values are logical indices into the launcher's ``CUDA_VISIBLE_DEVICES``. ``split_devices_for_replicas("1,2,3,4", 2, 2, 1) → ["1,2", "3,4"]`` - 2. **Template mode** — ``len(devices) == tp_size``: the YAML declares + 2. **Template mode** — ``len(devices) == devices_per_replica``: the YAML declares a single per-replica template (the same shape one replica would - use), and is **dp-independent**. Each replica r gets the offsets - ``[r*tp_size + a for a in template]`` of the launcher's + use), and is **replica-count-independent**. Each replica r gets the offsets + ``[r*devices_per_replica + a for a in template]`` of the launcher's ``CUDA_VISIBLE_DEVICES``. The template's entries must lie in - ``[0, tp_size)``. + ``[0, devices_per_replica)``. ``split_devices_for_replicas("0,1", 2, 2, 1) → ["0,1", "2,3"]`` ``split_devices_for_replicas("0,1", 4, 2, 1) → ["0,1", "2,3", "4,5", "6,7"]`` @@ -589,32 +812,40 @@ def split_devices_for_replicas( Any other length raises ``ValueError`` (the two modes are length-disjoint for ``num_replicas > 1``). """ - if num_replicas <= 1 or devices_str is None: - return [devices_str] if devices_str is not None else [devices_str] + if devices_str is None: + # No explicit placement: hand back one empty slot per replica, matching + # ``get_headless_replica_devices``. Returning a single-element list here + # made callers index past the end for every replica after the first. + return [None] * max(1, num_replicas) + + if num_replicas <= 1: + return [devices_str] device_list = [d.strip() for d in devices_str.split(",") if d.strip()] - if len(device_list) == num_replicas * tp_size: - return [",".join(device_list[r * tp_size : (r + 1) * tp_size]) for r in range(num_replicas)] + if len(device_list) == num_replicas * devices_per_replica: + return [ + ",".join(device_list[r * devices_per_replica : (r + 1) * devices_per_replica]) for r in range(num_replicas) + ] - if len(device_list) == tp_size: + if len(device_list) == devices_per_replica: try: offsets = [int(a) for a in device_list] except ValueError as e: raise ValueError(f"Stage {stage_id}: template-mode devices must be ints, got {devices_str!r}") from e - bad = [a for a in offsets if not (0 <= a < tp_size)] + bad = [a for a in offsets if not (0 <= a < devices_per_replica)] if bad: raise ValueError( f"Stage {stage_id}: template-mode device offset(s) {bad} " - f"out of range [0, {tp_size}); devices={devices_str!r}" + f"out of range [0, {devices_per_replica}); devices={devices_str!r}" ) - return [",".join(str(r * tp_size + a) for a in offsets) for r in range(num_replicas)] + return [",".join(str(r * devices_per_replica + a) for a in offsets) for r in range(num_replicas)] raise ValueError( f"Stage {stage_id}: devices={devices_str!r} has {len(device_list)} id(s); " - f"need either {tp_size} (template, dp-independent) or " - f"{num_replicas * tp_size} (pool / legacy). " - f"num_replicas={num_replicas}, tensor_parallel_size={tp_size}." + f"need either {devices_per_replica} (per-replica template) or " + f"{num_replicas * devices_per_replica} (pool / legacy). " + f"num_replicas={num_replicas}, devices_per_replica={devices_per_replica}." ) @@ -626,9 +857,32 @@ def get_stage_tp_size(stage_cfg: Any) -> int: return int(getattr(engine_args, "tensor_parallel_size", 1) or 1) -def get_stage_devices_per_replica(stage_cfg: Any) -> int: +def _get_local_llm_parallel_sizes( + stage_cfg: Any, + engine_args: Any | None = None, +) -> tuple[int, int, int]: + """Return ``(tp, local_dp, pp)`` for one local LLM replica. + + ``data_parallel_size`` is cluster-wide, whereas ``runtime.devices`` is + local to this process. Prefer an explicitly resolved + ``data_parallel_size_local`` (including zero for a head process that owns + no local engines), and only fall back to the global DP width when it is + unset. + """ + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) + tp_size = int(_get_attr_or_item(engine_args, "tensor_parallel_size", 1) or 1) + pp_size = int(_get_attr_or_item(engine_args, "pipeline_parallel_size", 1) or 1) + local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size_local", None) + if local_dp_size is None: + local_dp_size = _get_attr_or_item(engine_args, "data_parallel_size", 1) + return tp_size, int(local_dp_size if local_dp_size is not None else 1), pp_size + + +def get_stage_devices_per_replica(stage_cfg: Any, engine_args: Any | None = None) -> int: """Return the number of devices consumed by one replica of *stage_cfg*.""" - engine_args = getattr(stage_cfg, "engine_args", {}) + if engine_args is None: + engine_args = getattr(stage_cfg, "engine_args", {}) if getattr(stage_cfg, "stage_type", "llm") == "diffusion": parallel_config = _get_attr_or_item(engine_args, "parallel_config") if parallel_config is None: @@ -645,14 +899,15 @@ def get_stage_devices_per_replica(stage_cfg: Any) -> int: except Exception: return 1 - return get_stage_tp_size(stage_cfg) + tp_size, local_dp_size, pp_size = _get_local_llm_parallel_sizes(stage_cfg, engine_args) + return tp_size * max(1, local_dp_size) * pp_size def compute_replica_layout( stage_configs: Sequence[Any], *, allow_zero: bool = False, -) -> tuple[list[int], dict[int, list[str]]]: +) -> tuple[list[int], dict[int, list[str | None]]]: """Compute per-stage replica counts and device assignments. Args: @@ -680,7 +935,7 @@ def compute_replica_layout( raise ValueError(f"num_replicas must be >= 0, got {num_replicas}") replicas_per_stage.append(num_replicas if allow_zero else max(1, num_replicas)) - replica_devices_map: dict[int, list[str]] = {} + replica_devices_map: dict[int, list[str | None]] = {} for stage_id, stage_cfg in enumerate(stage_configs): num_replicas = replicas_per_stage[stage_id] if num_replicas <= 1: @@ -781,38 +1036,53 @@ def stage_runtime_env(stage_id: int, runtime_cfg: Any) -> Generator[None, None, def _project_omni_config_fields( config: Any, *, + field_map: Mapping[str, str] | None = None, exclude: frozenset[str] = frozenset(), ) -> dict[str, Any]: """Copy defined typed config fields into backend adapter kwargs.""" projected: dict[str, Any] = {} for config_field in fields(config): name = config_field.name - if name in exclude: + if name in exclude or (field_map is not None and name not in field_map): continue value = getattr(config, name) if value is not None: - projected[name] = copy.deepcopy(value) + projected[field_map.get(name, name) if field_map is not None else name] = copy.deepcopy(value) return projected +def _project_upstream_config_fields( + config: Any, + field_map: Mapping[str, str], +) -> dict[str, Any]: + """Project every explicit upstream input, including newly added fields.""" + explicit_fields: frozenset[str] = getattr(config, "_omni_explicit_fields", frozenset()) + unprojected_fields = explicit_fields - frozenset(field_map) + if unprojected_fields: + names = ", ".join(sorted(unprojected_fields)) + raise ValueError(f"{type(config).__name__} has explicit field(s) with no EngineArgs projection: {names}") + return _project_omni_config_fields( + config, + field_map=field_map, + exclude=frozenset(field_map) - explicit_fields, + ) + + def _project_omni_stage_engine_args( stage_config: BaseVllmOmniStageConfig, ) -> dict[str, Any]: """Read backend inputs from one structured stage config.""" engine_args: dict[str, Any] = {} is_diffusion = isinstance(stage_config, VllmOmniDiffusionStageConfig) - if is_diffusion: - engine_args.update(_project_omni_config_fields(stage_config.diffusion_config)) + diffusion_stage = cast(VllmOmniDiffusionStageConfig, stage_config) + engine_args.update(_project_omni_config_fields(diffusion_stage.diffusion_config)) for config, excluded_fields in ( ( stage_config.model_config, frozenset({"default_sampling_params", "has_sampling_extra_args"}), ), - (stage_config.load_config, frozenset()), - (stage_config.cache_config, frozenset()), - (stage_config.scheduler_config, frozenset()), ( stage_config.runtime_config, frozenset({"devices", "num_replicas", "env", "num_gpus"}), @@ -825,6 +1095,26 @@ def _project_omni_stage_engine_args( ) ) + if is_diffusion: + for config, field_map in ( + (stage_config.load_config, _DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP), + (stage_config.cache_config, _DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP), + (stage_config.scheduler_config, _DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP), + ): + engine_args.update(_project_omni_config_fields(config, field_map=field_map)) + else: + for config, field_map in ( + (stage_config.load_config, _LOAD_STAGE_ENGINE_FIELD_MAP), + (stage_config.cache_config, _CACHE_STAGE_ENGINE_FIELD_MAP), + (stage_config.scheduler_config, _SCHEDULER_STAGE_ENGINE_FIELD_MAP), + ): + engine_args.update(_project_upstream_config_fields(config, field_map)) + + for name in ("compilation_config", "profiler_config"): + value = getattr(stage_config, name) + if value is not None: + engine_args[name] = copy.deepcopy(value) + # The legacy builder always emits this key, including for pipelines such # as Audex that intentionally defer architecture discovery to HF config. engine_args["model_arch"] = copy.deepcopy(stage_config.model_config.model_arch) @@ -837,6 +1127,7 @@ def _project_omni_stage_engine_args( "hf_config_name": stage_config.hf_config_name, "engine_output_type": stage_config.engine_output_type, "custom_process_next_stage_input_func": stage_config.custom_process_next_stage_input_func, + "model_path_resolver": topology.model_path_resolver, "retains_state_across_chunks": topology.retains_state_across_chunks, } engine_args.update( @@ -851,13 +1142,14 @@ def _project_omni_stage_engine_args( if is_diffusion: engine_args["parallel_config"] = _project_omni_config_fields( stage_config.parallel_config, + field_map={name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS}, exclude=frozenset({"world_size"}), ) else: engine_args.update( - _project_omni_config_fields( + _project_upstream_config_fields( stage_config.parallel_config, - exclude=frozenset({"world_size"}), + _PARALLEL_CONFIG_ENGINE_FIELD_MAP, ) ) @@ -871,6 +1163,20 @@ def _project_omni_stage_engine_args( return engine_args +def _sampling_extra_args_keys(default_sampling_params: Any) -> tuple[str, ...]: + """Key names of a stage's default sampling ``extra_args``, sorted. + + Only the keys travel into the engine config: engine-core code needs to know + which request-shaping conventions a stage uses (e.g. CFG request pairing) + before any request exists, while the values stay a serving-layer concern. + """ + extra_args = _to_dict(default_sampling_params or {}).get("extra_args") or {} + try: + return tuple(sorted(str(key) for key in extra_args)) + except TypeError: + return () + + def _finalize_engine_args_dict( engine_args_dict: dict[str, Any], *, @@ -880,6 +1186,7 @@ def _finalize_engine_args_dict( stage_connector_spec: dict[str, Any] | None, cli_tokenizer: str | None, has_sampling_extra_args: bool, + sampling_extra_args_keys: tuple[str, ...] = (), ) -> dict[str, Any]: """Apply representation-independent engine adapter behavior.""" pipeline_model_root = model @@ -887,6 +1194,7 @@ def _finalize_engine_args_dict( stage_defines_tokenizer = ( engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None ) + model = _resolve_model_path(model, engine_args_dict) audex_stage = str(engine_args_dict.get("model_stage") or "") if audex_stage == "audex_xcodec": # TTA stage 1 decodes with the external XCodec1 checkpoint, not a @@ -960,6 +1268,7 @@ def _finalize_engine_args_dict( engine_args_dict.setdefault("enable_prefix_caching", False) engine_args_dict["has_sampling_extra_args"] = has_sampling_extra_args + engine_args_dict["sampling_extra_args_keys"] = sampling_extra_args_keys # TODO: Remove this after the performance regression is fixed # Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni to avoid performance regression @@ -990,6 +1299,7 @@ def build_legacy_engine_args_dict( stage_connector_spec=stage_connector_spec, cli_tokenizer=cli_tokenizer, has_sampling_extra_args=bool(default_sp.get("extra_args")), + sampling_extra_args_keys=_sampling_extra_args_keys(default_sp), ) @@ -1036,9 +1346,78 @@ def build_engine_args_dict_from_omni_stage_config( stage_connector_spec=stage_connector_spec, cli_tokenizer=cli_tokenizer, has_sampling_extra_args=stage_config.model_config.has_sampling_extra_args, + sampling_extra_args_keys=_sampling_extra_args_keys(stage_config.model_config.default_sampling_params), ) +def _count_stage_devices(devices: Any) -> int | None: + if devices is None: + return None + if isinstance(devices, (list, tuple)): + return len(devices) + values = [device for device in str(devices).split(",") if device.strip()] + return len(values) or None + + +def _check_stage_device_layout(stage_config: Any, engine_args_dict: dict[str, Any]) -> None: + """Fail early when a stage's world size cannot fit its assigned ``devices``. + + Re-runs :func:`check_device_layout` (normally only reached on the + ``--strategy-config`` path) against the fully resolved per-stage layout, so + an inconsistent ``tensor_parallel_size`` vs ``devices`` (issue #5003) is + reported here with a clear message instead of surfacing later as an opaque + worker-side ``local rank ... out of bounds`` assertion. + """ + from vllm_omni.config.composable_parallel import StrategyApplyError, check_device_layout + + runtime = getattr(stage_config, "runtime", None) + devices = _get_attr_or_item(runtime, "devices", None) if runtime is not None else None + if devices is None: + # No explicit placement -> vLLM assigns devices itself; nothing to check. + return + + num_replicas = _get_attr_or_item(runtime, "num_replicas", 1) if runtime is not None else 1 + stage_id = getattr(stage_config, "stage_id", "?") + tp_size, local_dp_size, pp_size = _get_local_llm_parallel_sizes(stage_config, engine_args_dict) + if local_dp_size == 0: + # This process hosts no local DP engines, so its local device list does + # not describe the cluster-wide DP layout and must not be validated. + return + + try: + check_device_layout( + devices, + tensor_parallel_size=tp_size, + data_parallel_size=local_dp_size, + pipeline_parallel_size=pp_size, + num_replicas=int(num_replicas or 1), + role=f"stage-{stage_id}", + ) + except StrategyApplyError as e: + message = ( + f"Stage {stage_id}: device layout is inconsistent — {e} " + "Set devices and the per-stage TP, local DP, PP, and replica counts " + "so the declared device count matches the local world size." + ) + device_count = _count_stage_devices(devices) + world_without_tp = local_dp_size * pp_size + valid_without_tp = {world_without_tp, int(num_replicas or 1) * world_without_tp} + if tp_size > 1 and device_count in valid_without_tp: + message += ( + " This layout is consistent with issue #5003: a top-level " + "--tensor-parallel-size is applied to every stage, but each stage's " + "`devices` is not adjusted automatically. Pass --stage-overrides " + "to set tensor_parallel_size and devices together on every stage, " + "so single-GPU stages get tensor_parallel_size=1, e.g. " + '\'{"0": {"tensor_parallel_size": 4, "devices": "0,1,2,3"}, ' + '"1": {"tensor_parallel_size": 1, "devices": "0"}, ' + '"2": {"tensor_parallel_size": 1, "devices": "1"}}\'. ' + "Or omit the top-level --tensor-parallel-size and set it only in " + "stage-0's override." + ) + raise ValueError(message) from e + + def build_vllm_config( stage_config: Any, model: str, @@ -1074,6 +1453,16 @@ def build_vllm_config( omni_engine_args = OmniEngineArgs(**filtered_engine_args_dict) + # Guard against a per-stage world size that its assigned ``devices`` cannot + # satisfy (issue #5003). A top-level ``--tensor-parallel-size`` is broadcast + # to every stage, but ``devices`` is not, so a stage can end up with e.g. + # tensor_parallel_size=4 while still holding a single-GPU deploy default. + # Without --strategy-config the strategy-path device check never runs, so + # the mismatch used to surface only as an opaque worker-side assertion + # ("DP adjusted local rank N is out of bounds for M devices."). Re-run the + # same check here, before workers spawn, to fail early with a clear message. + _check_stage_device_layout(stage_config, filtered_engine_args_dict) + # Multi-stage pipelines (qwen3_tts code2wav, etc.) set max_model_len # larger than HF max_position_embeddings by design. vLLM's validator # rejects that without the env flag. @@ -1400,6 +1789,7 @@ def build_diffusion_config( num_devices_per_stage = od_config.parallel_config.world_size device_control_env = current_omni_platform.device_control_env_var visible_devices_str = os.environ.get(device_control_env) if device_control_env else None + physical_devices: list[str | int] if visible_devices_str: physical_devices = [device.strip() for device in visible_devices_str.split(",") if device.strip()] else: @@ -1423,7 +1813,6 @@ def initialize_diffusion_stage( stage_cfg: Any, metadata: StageMetadata, stage_init_timeout: int, - batch_size: int = 1, use_inline: bool = False, ) -> Any: """Build a diffusion stage client. @@ -1433,35 +1822,53 @@ def initialize_diffusion_stage( stage_cfg: Stage configuration. metadata: Extracted stage metadata. stage_init_timeout: Timeout in seconds for stage initialization handshake - batch_size: Maximum number of requests to batch together in the - diffusion engine. Passed through to ``StageDiffusionClient`` - and ultimately to ``AsyncOmni``. use_inline: If True, uses the inline diffusion client instead of subprocess. """ from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client od_config = build_diffusion_config(model, stage_cfg, metadata) - return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline) + return create_diffusion_client(model, od_config, metadata, stage_init_timeout, use_inline) -def maybe_apply_audex_cfg_patches(vllm_config: Any) -> None: - """Install the Audex CFG scheduler patches for CFG-configured engines. +def _stage_declares_cfg_pairs(model_config: Any) -> bool: + """Whether this stage submits classifier-free-guidance request pairs. + + Two independent declarations, because a model may own either side of the + mechanism without the other: + + * a CFG logits processor is configured (blending implies pairing), or + * the stage's default sampling ``extra_args`` carry ``cfg_role``, which is + how a model declares paired requests without owning a logits processor. + Only the *key set* survives into the engine config (see + ``sampling_extra_args_keys``); the values stay in the serving layer. + """ + processors = getattr(model_config, "logits_processors", None) or [] + if any("CFGLogitsProcessor" in getattr(proc, "__name__", str(proc)) for proc in processors): + return True + return "cfg_role" in (getattr(model_config, "sampling_extra_args_keys", None) or ()) + + +def maybe_apply_cfg_scheduler_patches(vllm_config: Any) -> None: + """Install the CFG pairing scheduler patches for CFG-configured engines. Must run in the engine-core process BEFORE ``Scheduler`` is constructed: the patch wraps ``Scheduler.__init__`` to add the pair registry, so a - scheduler built earlier would never become pair-aware. Gated on the - engine's ``logits_processors`` so non-CFG stages stay untouched. + scheduler built earlier would never become pair-aware. Gated so non-CFG + stages stay untouched. """ model_config = getattr(vllm_config, "model_config", None) - processors = getattr(model_config, "logits_processors", None) or [] - if not any("AudexCFGLogitsProcessor" in getattr(proc, "__name__", str(proc)) for proc in processors): + if model_config is None or not _stage_declares_cfg_pairs(model_config): return - from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches + from vllm_omni.model_executor.models.common.cfg_pairing import apply_cfg_patches apply_cfg_patches() from vllm.v1.core.sched.scheduler import Scheduler - if not getattr(Scheduler.schedule, "_audex_cfg_patched", False): - raise RuntimeError("Audex CFG scheduler patches failed to install before Scheduler construction") + if not getattr(Scheduler.schedule, "_cfg_pairing_patched", False): + raise RuntimeError("CFG pairing scheduler patches failed to install before Scheduler construction") + + +# Name kept for callers written against the Audex-only gate. +maybe_apply_audex_cfg_patches = maybe_apply_cfg_scheduler_patches diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 7d1bfc60a..e7de5f51c 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unified stage-local runtime abstraction for vLLM-Omni.""" from __future__ import annotations @@ -94,6 +97,16 @@ class StagePool: DISPATCH_WAIT_TIMEOUT_S: float = 10.0 DISPATCH_RETRY_INTERVAL_S: float = 0.1 + # Only these EngineCore helpers may skip collective_rpc_async. A generic + # ``{method}_async`` on AsyncMPClient must not silently drop timeout. + _ENGINE_CORE_CONTROL_ASYNC_METHODS = frozenset( + { + "pause_scheduler", + "resume_scheduler", + "sleep", + "wake_up", + } + ) def __init__( self, @@ -620,6 +633,15 @@ class StagePool: if callable(pop_native_text_metrics): native_text_metrics = pop_native_text_metrics(request_id) native_generation_tokens = native_text_metrics.get("num_generation_tokens") + finish_reason = next( + ( + str(reason) + for request_output in reversed(request_outputs) + for completion in reversed(getattr(request_output, "outputs", []) or []) + if (reason := getattr(completion, "finish_reason", None)) is not None + ), + None, + ) num_tokens_out = ( max(int(native_generation_tokens), 0) if isinstance(native_generation_tokens, int) and not isinstance(native_generation_tokens, bool) @@ -643,11 +665,6 @@ class StagePool: ) image_pixels = self._count_image_pixels(request_outputs) if output_unit_type == "image" else 0 num_inference_steps = coerce_positive_int_scalar(getattr(sampling_params, "num_inference_steps", None)) or 0 - denoise_step_latency_ms = ( - defs.compute_denoise_step_latency(stage_gen_time_ms, num_inference_steps) - if output_unit_type == "image" - else 0.0 - ) has_output_timestamps = bool(output_timestamps) first_ts = output_timestamps[0] if has_output_timestamps else now serving_time_to_first_output_ms = ( @@ -655,6 +672,11 @@ class StagePool: if non_empty_first_output_ts is not None else 0.0 ) + image_time_to_first_output_ms = ( + max((non_empty_first_output_ts - submit_ts) * 1000.0, 0.0) + if non_empty_first_output_ts is not None and output_unit_type == "image" + else 0.0 + ) remaining_ms = max((now - first_ts) * 1000.0, 0.0) if output_unit_count > 1 and self._is_streaming_output_unit_type(output_unit_type): time_per_output_unit_ms = remaining_ms / float(output_unit_count - 1) @@ -688,6 +710,7 @@ class StagePool: # happens inside the model runner and is not observable here. batch_size=1, replica_id=replica_id, + finish_reason=finish_reason, rx_decode_time_ms=0.0, rx_transfer_bytes=0, rx_in_flight_time_ms=0.0, @@ -699,10 +722,11 @@ class StagePool: audio_sample_rate=audio_sample_rate, audio_duration_s=audio_duration_s, image_pixels=image_pixels, - denoise_step_latency_ms=denoise_step_latency_ms, + num_inference_steps=num_inference_steps, output_unit_type=output_unit_type, output_unit_count=output_unit_count, serving_time_to_first_output_ms=serving_time_to_first_output_ms, + image_time_to_first_output_ms=image_time_to_first_output_ms, time_per_output_unit_ms=time_per_output_unit_ms, inter_output_latency_ms=inter_output_latency_ms, inter_output_latencies_ms=inter_output_latencies_ms, @@ -1175,14 +1199,19 @@ class StagePool: # ---- Stage-local control plane ---- - async def abort_requests(self, request_ids: list[str]) -> None: + async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]: """Abort the given requests in this stage pool. Request-bound abort routing stays inside the pool because route affinity (``request_id -> replica_id``) is pool-owned. + + For AR stages, collect terminal abort ``RequestOutput`` objects from + output-processor state *before* EngineCore abort so partial tokens + generated so far are preserved. Diffusion stages return no abort + outputs (whole-sample retry remains the recovery path). """ if not request_ids: - return + return [] request_ids_by_replica: dict[int, list[str]] = {} for request_id in request_ids: @@ -1192,18 +1221,45 @@ class StagePool: continue request_ids_by_replica.setdefault(replica_id, []).append(request_id) + abort_outputs: list[tuple[str, Any]] = [] + is_diffusion = self.stage_type == "diffusion" for replica_id, replica_request_ids in request_ids_by_replica.items(): + # Orchestrator ids are OP external ids; EngineCore may use a + # different internal request_id assigned by InputProcessor. + # Always key collected outputs by the orchestrator id we passed + # (internal=False), never by RequestOutput.request_id — some stages + # still surface EngineCore/internal ids and would drop the prefix. + engine_abort_ids = list(replica_request_ids) + if not is_diffusion and self._output_processor is not None: + collect = getattr(self._output_processor, "abort_requests_collecting_outputs", None) + if collect is not None: + engine_abort_ids = [] + for orch_req_id in replica_request_ids: + collected_ids, stage_outputs = collect( + [orch_req_id], + internal=False, + commit_state=False, + ) + for req_out in stage_outputs: + abort_outputs.append((orch_req_id, req_out)) + if collected_ids: + engine_abort_ids.extend(collected_ids) + else: + engine_abort_ids.append(orch_req_id) + else: + aborted = self._output_processor.abort_requests(replica_request_ids, internal=False) + if aborted: + engine_abort_ids = list(aborted) client = self.clients[replica_id] if client is None: continue - await client.abort_requests_async(replica_request_ids) + await client.abort_requests_async(engine_abort_ids) + if not is_diffusion and self._output_processor is not None: + commit = getattr(self._output_processor, "commit_aborted_request_state", None) + if callable(commit): + commit(replica_request_ids, internal=False) - # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that - # would otherwise leak — aborted requests never produce a final - # EngineCoreOutput, so process_outputs() never fires its cleanup path. - all_aborted = [rid for ids in request_ids_by_replica.values() for rid in ids] - if all_aborted and self._output_processor is not None: - self._output_processor.abort_requests(all_aborted, internal=True) + return abort_outputs async def collective_rpc( self, @@ -1213,7 +1269,13 @@ class StagePool: args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, ) -> dict[str, Any] | Any: - """Dispatch a stage-scoped control-plane RPC to one physical route.""" + """Dispatch a stage-scoped control-plane RPC to one physical route. + + EngineCore control methods (``pause_scheduler``, ``resume_scheduler``, + ``sleep``, ``wake_up``) propagate worker exceptions. Capability probes + and other RPCs still return ``{"supported": False, "error": ...}``. + """ + args = tuple(args or ()) kwargs = dict(kwargs or {}) client = self.clients[replica_id] if client is None: @@ -1222,6 +1284,14 @@ class StagePool: "error": f"stage {self.stage_id} replica {replica_id} is not attached", } try: + if self.stage_type != "diffusion" and method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS: + client_method = getattr(client, f"{method}_async", None) + if callable(client_method): + result = client_method(*args, **kwargs) + if timeout is not None: + return await asyncio.wait_for(result, timeout=timeout) + return await result + return await client.collective_rpc_async( method=method, timeout=timeout, @@ -1235,9 +1305,15 @@ class StagePool: replica_id, method, ) + if method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS: + raise + if isinstance(exc, TimeoutError): + error = f"{type(exc).__name__}: {method} timed out after {timeout}s" + else: + error = str(exc) or repr(exc) return { "supported": False, - "error": str(exc), + "error": error, } def shutdown_replica(self, replica_id: int) -> None: diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index bf84f7793..049d09688 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage Runtime implementations for single-node and distributed omni stages.""" @@ -12,7 +12,7 @@ import threading from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass -from typing import Any +from typing import Any, cast import janus from omegaconf import OmegaConf @@ -58,6 +58,7 @@ from vllm_omni.engine.stage_init_utils import ( load_omni_transfer_config_for_model, prepare_engine_environment, release_device_locks, + stage_runtime_env, ) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.entrypoints.stage_utils import resolve_stage_physical_devices @@ -86,7 +87,6 @@ class StageRemoteFactoryContext: base_metadata: Any vllm_config: Any | None = None executor_class: type | None = None - diffusion_batch_size: int = 1 def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: @@ -123,7 +123,6 @@ class StageRuntime: config_path: str, *, stage_init_timeout: int, - diffusion_batch_size: int, async_chunk: bool, tokenizer: str | None = None, log_stats: bool = False, @@ -132,7 +131,6 @@ class StageRuntime: self._model = model self._config_path = config_path self._stage_init_timeout = stage_init_timeout - self._diffusion_batch_size = diffusion_batch_size self._async_chunk = async_chunk self._tokenizer = tokenizer self._log_stats = log_stats @@ -336,7 +334,7 @@ class StageRuntime: self, omni_transfer_config: Any, replicas_per_stage: Sequence[int], - replica_devices_map: Mapping[int, Sequence[str]], + replica_devices_map: Mapping[int, Sequence[str | None]], ) -> list[LogicalStageInitPlan]: """Build startup plans for every logical stage and replica.""" stage_plans: list[LogicalStageInitPlan] = [] @@ -544,7 +542,7 @@ class StageRuntime: self, plan: ReplicaInitPlan, stage_init_timeout: int, - ) -> StageEngineCoreClientBase: + ) -> StagePoolClient: """Initialize one local LLM replica using vLLM's launch/attach pattern.""" resources: StageReplicaResources | None = None stage_client = None @@ -576,7 +574,7 @@ class StageRuntime: ) # Serialize engine-core spawning across all LLM replicas to avoid # ZMQ port-allocation races and simultaneous CUDA context init. - with self._replica_launch_lock: + with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): with launch_stage_replica( vllm_config=vllm_config, executor_class=executor_class, @@ -607,7 +605,7 @@ class StageRuntime: ) logger.info("[StageRuntime] Stage %s initialized", plan.metadata.stage_id) - return stage_client + return cast(StagePoolClient, stage_client) except Exception: if stage_client is not None: try: @@ -640,35 +638,50 @@ class StageRuntime: self, plan: ReplicaInitPlan, stage_init_timeout: int, - ) -> Any: + ) -> StagePoolClient: """Initialize one local diffusion replica end-to-end.""" client = None resources = None try: - with self._stage_device_scope(plan.metadata.stage_id, plan.metadata.runtime_cfg): + with ( + stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg), + self._stage_device_scope(plan.metadata.stage_id, plan.metadata.runtime_cfg), + ): omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector if omni_conn_cfg: + if omni_from is None or omni_to is None: + raise RuntimeError("Omni KV connector requires source and destination stages") inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs) + engine_args = getattr(plan.stage_cfg, "engine_args", {}) + inline_diffusion = ( + engine_args.get("inline_diffusion", False) + if hasattr(engine_args, "get") + else getattr(engine_args, "inline_diffusion", False) + ) + custom_pipeline_args = ( + engine_args.get("custom_pipeline_args") + if hasattr(engine_args, "get") + else getattr(engine_args, "custom_pipeline_args", None) + ) client, resources = launch_diffusion_stage_replica( model=self._model, stage_config=plan.stage_cfg, metadata=plan.metadata, stage_init_timeout=stage_init_timeout, - batch_size=self._diffusion_batch_size, - use_inline=self._num_stages == 1 and plan.num_replicas == 1, + use_inline=plan.num_replicas == 1 + and bool(self._num_stages == 1 or inline_diffusion or custom_pipeline_args), replica_id=plan.replica_id, omni_master_server=self._get_omni_master_server(), omni_coordinator_address=self._get_coordinator_address(), ) logger.info( - "[StageRuntime] Stage %s replica %s initialized (diffusion, batch_size=%d)", + "[StageRuntime] Stage %s replica %s initialized (diffusion)", plan.metadata.stage_id, plan.replica_id, - self._diffusion_batch_size, ) - return client + return cast(StagePoolClient, client) except Exception: if client is not None: try: @@ -750,7 +763,6 @@ class DistStageRuntime(StageRuntime): config_path: str, *, stage_init_timeout: int, - diffusion_batch_size: int, async_chunk: bool, single_stage_id_filter: int | None, omni_master_address: str, @@ -767,7 +779,6 @@ class DistStageRuntime(StageRuntime): model=model, config_path=config_path, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, log_stats=log_stats, @@ -896,7 +907,6 @@ class DistStageRuntime(StageRuntime): base_metadata=metadata, vllm_config=plan.stage_vllm_config, executor_class=plan.executor_class, - diffusion_batch_size=self._diffusion_batch_size, ) return self._create_remote_replica_client(ctx, plan.replica_id) @@ -959,7 +969,6 @@ class DistStageRuntime(StageRuntime): base_metadata=template.metadata, vllm_config=template.stage_vllm_config, executor_class=template.executor_class, - diffusion_batch_size=self._diffusion_batch_size, ) return contexts @@ -1027,14 +1036,13 @@ class DistStageRuntime(StageRuntime): metadata, request_address=resources.addresses.inputs[0], response_address=resources.addresses.outputs[0], - batch_size=ctx.diffusion_batch_size, ) logger.info( "[DistStageRuntime] Remote diffusion replica attached stage=%d replica=%d", stage_id, replica_id, ) - return client + return cast(StagePoolClient, client) if ctx.vllm_config is None: raise RuntimeError(f"Remote LLM stage {stage_id} is missing vllm_config") @@ -1091,7 +1099,6 @@ def create_stage_runtime( *, single_stage_mode: bool, stage_init_timeout: int, - diffusion_batch_size: int, async_chunk: bool, tokenizer: str | None = None, # Distributed-only params: @@ -1113,7 +1120,6 @@ def create_stage_runtime( model=model, config_path=config_path, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, log_stats=log_stats, @@ -1130,7 +1136,6 @@ def create_stage_runtime( model=model, config_path=config_path, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, log_stats=log_stats, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index f6899deaa..8ac0fde56 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -17,7 +17,7 @@ from vllm import TokensPrompt from vllm.engine.protocol import EngineClient, StreamingInput from vllm.logger import init_logger from vllm.lora.request import LoRARequest -from vllm.outputs import PoolingRequestOutput +from vllm.outputs import CompletionOutput, PoolingRequestOutput from vllm.plugins.io_processors import get_io_processor from vllm.pooling_params import PoolingParams from vllm.renderers.inputs.preprocess import extract_prompt_components @@ -37,7 +37,6 @@ from vllm_omni.errors import client_error_metadata from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.inputs.preprocess import InputPreprocessor @@ -54,6 +53,11 @@ if TYPE_CHECKING: logger = init_logger(__name__) _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 +# Blocking-wait interval for the event-driven final-output drain +# (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): a message wakes the drain immediately via +# the janus queue's condition variable; this timeout only bounds how often the +# orchestrator liveness check runs while the pipeline is idle. +_FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 class AsyncEventResolver: @@ -123,9 +127,8 @@ class AsyncOmni(EngineClient, OmniBase): Args: model: Model name or path to load. **kwargs: Additional keyword arguments. - - stage_configs_path: Optional path to YAML file containing stage - configurations. If None, configurations are resolved from model - pipeline factory. + - deploy_config: Optional path to a deploy YAML. If None, + configurations are resolved from the model pipeline factory. - log_stats: Whether to enable statistics logging. - stage_init_timeout: Timeout for per-stage initialization. - init_timeout: Total timeout for orchestrator startup. @@ -147,7 +150,19 @@ class AsyncOmni(EngineClient, OmniBase): OmniBase.__init__(self, model=model, **kwargs) self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False + # In-flight EngineCore submits (non-streaming add_request, or each + # streaming ADD/update/final marker). sleep() waits for this to hit + # zero so a pipelined request cannot race into EngineCore during + # drain/offload. Streaming generate does not hold a slot while + # waiting for the next client chunk. + self._admitting: int = 0 + # True after pause_generation() or AR EngineCore sleep; wake_up must + # not reopen generate() until resume_generation(). Diffusion-only + # sleep uses _paused as a temporary admission gate and clears it + # on wake so sleep → wake → generate keeps working. + self._hold_admission_until_resume: bool = False self._sleeping_tags: set[str] = set() + self._stage_sleeping_tags: dict[int, set[str]] = {} self._level2_sleeping: bool = False self._duplex_request_client: DuplexRequestClient | None = None self.duplex_lifecycle_events: asyncio.Queue[DuplexSessionLifecycleMessage] = asyncio.Queue() @@ -505,31 +520,40 @@ class AsyncOmni(EngineClient, OmniBase): external_request_id = request_id request_id = self._get_unique_request_id(external_request_id) - # Wait until generation is resumed if the engine is paused + # Wait until generation is resumed if the engine is paused. Non-streaming + # generate holds an admission slot until add_request completes so sleep() + # cannot race a just-unblocked generate into EngineCore during offload. + # Streaming generate does **not** hold the slot while waiting for the + # first client chunk; the input pump acquires it immediately before each + # EngineCore ADD/update. + streaming_input = isinstance(prompt, AsyncGenerator) async with self._pause_cond: await self._pause_cond.wait_for(lambda: not self._paused) + if not streaming_input: + self._admitting = getattr(self, "_admitting", 0) + 1 + admitting = not streaming_input logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}") - _sleeping_tags = getattr(self, "_sleeping_tags", None) - if _sleeping_tags: - raise RuntimeError( - f"Generation rejected: Engine is partially or fully asleep. " - f"Currently sleeping tags: {list(_sleeping_tags)}. " - f"Please perform a full wake_up before generating." - ) - - # Reject diffusion list-prompt early with a clear API error. - if isinstance(prompt, list) and any( - getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", []) - ): - raise ValueError( - "Diffusion stages accept only a single prompt per request. " - "Submit multiple independent requests to use scheduler batching." - ) - input_stream_task: asyncio.Task | None = None try: + _sleeping_tags = getattr(self, "_sleeping_tags", None) + if _sleeping_tags: + raise RuntimeError( + f"Generation rejected: Engine is partially or fully asleep. " + f"Currently sleeping tags: {list(_sleeping_tags)}. " + f"Please perform a full wake_up before generating." + ) + + # Reject diffusion list-prompt early with a clear API error. + if isinstance(prompt, list) and any( + getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", []) + ): + raise ValueError( + "Diffusion stages accept only a single prompt per request. " + "Submit multiple independent requests to use scheduler batching." + ) + # Start final output dispatcher on the first call to generate() self._final_output_handler() @@ -590,8 +614,11 @@ class AsyncOmni(EngineClient, OmniBase): req_sp_list[p_id] = self._prepare_prefill_sampling_params(request_id, req_sp_list[p_id]) # Add request(s) to stage 0. For streaming inputs, submit - # chunks incrementally through streaming_update. - if isinstance(prompt, AsyncGenerator): + # chunks incrementally through streaming_update. The helper + # returns as soon as the pump task is created; each ADD/update + # takes its own admission slot inside the pump. + if streaming_input: + first_chunk_submitted = asyncio.get_running_loop().create_future() input_stream_task = await self._add_streaming_input_request( request_id=request_id, input_stream=prompt, @@ -600,7 +627,9 @@ class AsyncOmni(EngineClient, OmniBase): final_output_stage_ids=final_output_stage_ids, arrival_time=wall_start_ts, lora_request=lora_request, + first_chunk_submitted=first_chunk_submitted, ) + await first_chunk_submitted else: await self.engine.add_request_async( request_id=request_id, @@ -614,6 +643,11 @@ class AsyncOmni(EngineClient, OmniBase): submit_ts = time.time() req_state.metrics.stage_first_ts[0] = submit_ts req_start_ts[request_id] = submit_ts + if admitting: + await self._release_generate_admission() + admitting = False + # Refresh gauges on arrival. + self._publish_request_gauges(len(self.request_states)) # Process results based on mode # Both sequential and async_chunk modes read the same message stream @@ -630,27 +664,38 @@ class AsyncOmni(EngineClient, OmniBase): logger.debug(f"[AsyncOmni] Request {request_id} completed") - # The input pump can outlive a normally-completed generation - # (e.g. a realtime client that keeps streaming after the final - # output); a live pump keeps issuing streaming updates for a - # request the orchestrator no longer tracks. - if input_stream_task is not None and not input_stream_task.done(): - input_stream_task.cancel() - - self._log_summary_and_cleanup(request_id) - except (asyncio.CancelledError, GeneratorExit): - if input_stream_task is not None and not input_stream_task.done(): - input_stream_task.cancel() - self._fire_failure_counter_if_alive(request_id) + self._record_request_failure_once(request_id, reason="client_disconnect") await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: - self._fire_failure_counter_if_alive(request_id) + self._record_request_failure_once(request_id, reason="stage_error") await self._abort_internal_requests(request_id) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise + finally: + if input_stream_task is not None and not input_stream_task.done(): + input_stream_task.cancel() + if admitting: + await self._release_generate_admission() + self._log_summary_and_cleanup(request_id) + + async def _release_generate_admission(self) -> None: + """Drop one in-flight generate admission slot held across add_request.""" + async with self._pause_cond: + self._admitting = max(getattr(self, "_admitting", 1) - 1, 0) + self._pause_cond.notify_all() + + async def _submit_with_admission(self, awaitable): + """Wait for resume, hold one admission slot for a single EngineCore submit.""" + async with self._pause_cond: + await self._pause_cond.wait_for(lambda: not self._paused) + self._admitting = getattr(self, "_admitting", 0) + 1 + try: + return await awaitable + finally: + await self._release_generate_admission() async def _add_streaming_input_request( self, @@ -662,6 +707,7 @@ class AsyncOmni(EngineClient, OmniBase): final_output_stage_ids: Sequence[int], arrival_time: float, lora_request: Any = None, + first_chunk_submitted: asyncio.Future[None] | None = None, ) -> asyncio.Task: """Submit a streaming input generator as incremental stage-0 updates.""" if not sampling_params_list: @@ -678,6 +724,10 @@ class AsyncOmni(EngineClient, OmniBase): stage0_params = stage0_params.clone() stage0_params.skip_clone = True + def _mark_first_chunk_submitted() -> None: + if first_chunk_submitted is not None and not first_chunk_submitted.done(): + first_chunk_submitted.set_result(None) + async def handle_inputs() -> None: nonlocal has_submitted_first_chunk cancelled = False @@ -691,29 +741,34 @@ class AsyncOmni(EngineClient, OmniBase): prompt_text, _, _ = extract_prompt_components(self.model_config, chunk_prompt) if not has_submitted_first_chunk: - await self.engine.add_request_async( - request_id=request_id, - prompt=chunk_prompt, - prompt_text=prompt_text, - sampling_params_list=chunk_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=True, + await self._submit_with_admission( + self.engine.add_request_async( + request_id=request_id, + prompt=chunk_prompt, + prompt_text=prompt_text, + sampling_params_list=chunk_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=True, + ) ) has_submitted_first_chunk = True + _mark_first_chunk_submitted() else: - await self.engine.add_streaming_update_async( - request_id=request_id, - prompt=chunk_prompt, - prompt_text=prompt_text, - sampling_params_list=chunk_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=True, + await self._submit_with_admission( + self.engine.add_streaming_update_async( + request_id=request_id, + prompt=chunk_prompt, + prompt_text=prompt_text, + sampling_params_list=chunk_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=True, + ) ) except (asyncio.CancelledError, GeneratorExit): cancelled = True @@ -728,37 +783,47 @@ class AsyncOmni(EngineClient, OmniBase): ) ) finally: - if not cancelled: - # Send empty final request to indicate that inputs have - # finished. Don't send if canceled (session was aborted). - final_sampling_params_list = list(sampling_params_list) - final_sampling_params_list[0] = stage0_params - final_prompt = TokensPrompt(prompt_token_ids=[0]) + try: + if not cancelled: + # Send empty final request to indicate that inputs have + # finished. Don't send if canceled (session was aborted). + final_sampling_params_list = list(sampling_params_list) + final_sampling_params_list[0] = stage0_params + final_prompt = TokensPrompt(prompt_token_ids=[0]) - if has_submitted_first_chunk: - await self.engine.add_streaming_update_async( - request_id=request_id, - prompt=final_prompt, - prompt_text=None, - sampling_params_list=final_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=False, - ) - else: - await self.engine.add_request_async( - request_id=request_id, - prompt=final_prompt, - prompt_text=None, - sampling_params_list=final_sampling_params_list, - final_stage_id=final_stage_id, - final_output_stage_ids=final_output_stage_ids, - arrival_time=arrival_time, - lora_request=lora_request, - resumable=False, - ) + if has_submitted_first_chunk: + await self._submit_with_admission( + self.engine.add_streaming_update_async( + request_id=request_id, + prompt=final_prompt, + prompt_text=None, + sampling_params_list=final_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=False, + ) + ) + else: + await self._submit_with_admission( + self.engine.add_request_async( + request_id=request_id, + prompt=final_prompt, + prompt_text=None, + sampling_params_list=final_sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=final_output_stage_ids, + arrival_time=arrival_time, + lora_request=lora_request, + resumable=False, + ) + ) + has_submitted_first_chunk = True + finally: + # Unblock generate() even on cancel / empty stream / submit + # failure so it can observe a terminal abort or empty result. + _mark_first_chunk_submitted() input_stream_task = asyncio.create_task(handle_inputs()) req_state.input_stream_task = input_stream_task @@ -877,14 +942,29 @@ class AsyncOmni(EngineClient, OmniBase): engine = self.engine + # Event-driven drain (VLLM_OMNI_EVENT_DRIVEN_ORCH=1): block on the + # queue's condition variable in a dedicated thread instead of the + # get_nowait + 1 ms sleep cadence. Same flag as the orchestrator-side + # event-driven loop (vllm_omni/engine/orchestrator.py). + from vllm_omni.engine.orchestrator import _event_driven_orch_enabled + + event_driven_drain = _event_driven_orch_enabled() and hasattr(engine, "get_output_blocking_async") + async def _final_output_loop(): """Background coroutine that dispatches final outputs to request queues.""" try: while True: - msg = await engine.try_get_output_async() - if msg is None: - await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) - continue + if event_driven_drain: + msg = await engine.get_output_blocking_async(timeout=_FINAL_OUTPUT_BLOCKING_WAIT_S) + if msg is None: + # Timed out with the orchestrator alive; loop for + # the periodic liveness check. + continue + else: + msg = await engine.try_get_output_async() + if msg is None: + await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) + continue if isinstance(msg, dict) and msg.get("type") == "ack": ack_data = msg.get("ack") @@ -1010,6 +1090,30 @@ class AsyncOmni(EngineClient, OmniBase): return results + async def _engine_core_rpc( + self, + method: str, + *, + stage_ids: list[int], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> list[Any]: + """Call an AR EngineCore helper via collective_rpc (orchestrator loop). + + StagePool resolves ``{method}_async`` on the AR client when present + (vLLM AsyncMPClient convention). Raises if any replica reports failure. + """ + results = await self.collective_rpc( + method=method, + args=args, + kwargs=kwargs, + stage_ids=stage_ids, + ) + for result in results: + if isinstance(result, dict) and result.get("error"): + raise RuntimeError(f"{method} failed: {result['error']}") + return results + @staticmethod def _coerce_stage_bool(result: Any) -> bool: """Reduce a stage RPC result to a boolean. @@ -1080,33 +1184,188 @@ class AsyncOmni(EngineClient, OmniBase): await self._abort(internal_req_ids) async def _abort(self, request_ids: list[str]) -> None: - """Submit request IDs to be aborted to the engine.""" - await self.engine.abort_async(request_ids) + """Abort request IDs via the engine and enqueue terminal abort outputs. + + Waits for orchestrator abort acknowledgment, enqueues any AR terminal + abort outputs (partial tokens) into each request's asyncio queue, then + cancels the input pump. Frontend ``request_states`` stay registered so + ``generate()`` can consume the terminal message in + ``_process_orchestrator_results`` and run normal cleanup. + + When ``abort_async`` returns no output for an active request (OP not + registered yet, unbound replica, or orchestrator id drop), enqueue a + synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. + """ + abort_outputs = await self.engine.abort_async(request_ids) or [] + delivered: set[str] = set() + for output_msg in abort_outputs: + req_id = getattr(output_msg, "request_id", None) + if req_id is None: + continue + state = self.request_states.get(req_id) + if state is None: + logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) + continue + await state.queue.put(output_msg) + delivered.add(req_id) for rid in request_ids: - state = self.request_states.pop(rid, None) + state = self.request_states.get(rid) + if state is not None and rid not in delivered: + queue = getattr(state, "queue", None) + if queue is not None: + await state.queue.put(self._synthetic_abort_output_message(rid)) + delivered.add(rid) + for rid in request_ids: + self._record_request_failure_once(rid, reason="client_abort") + state = self.request_states.get(rid) input_stream_task = getattr(state, "input_stream_task", None) if input_stream_task is not None and not input_stream_task.done(): input_stream_task.cancel() if self.log_stats: logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) + @staticmethod + def _synthetic_abort_output_message(request_id: str) -> OutputMessage: + """Terminal abort OutputMessage used when the engine returned none.""" + engine_output = OmniRequestOutput( + request_id=request_id, + finished=True, + stage_id=0, + final_output_type="text", + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="abort", + stop_reason=None, + ) + ], + ) + return OutputMessage( + request_id=request_id, + stage_id=0, + replica_id=None, + engine_outputs=engine_output, + metrics=None, + finished=True, + stage_submit_ts=None, + ) + + def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]: + """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC).""" + n_stages = len(self.engine.stage_clients) + if stage_ids is None: + stage_ids = list(range(n_stages)) + else: + invalid = [sid for sid in stage_ids if not isinstance(sid, int) or sid < 0 or sid >= n_stages] + if invalid: + raise ValueError( + f"Invalid stage_ids {invalid}; valid range is 0..{n_stages - 1}" + if n_stages + else f"Invalid stage_ids {invalid}; this engine has no stages" + ) + ar_stage_ids: list[int] = [] + diffusion_stage_ids: list[int] = [] + for sid in stage_ids: + client = self.engine.stage_clients[sid] + if getattr(client, "stage_type", "llm") == "diffusion": + diffusion_stage_ids.append(sid) + else: + ar_stage_ids.append(sid) + return ar_stage_ids, diffusion_stage_ids + + def _sleeping_tags_for_stages(self, stage_ids: list[int]) -> set[str]: + """Union of sleeping tags recorded for the given stages.""" + per_stage = getattr(self, "_stage_sleeping_tags", None) or {} + tags: set[str] = set() + for sid in stage_ids: + tags.update(per_stage.get(sid, set())) + return tags + + def _refresh_union_sleeping_tags(self) -> None: + """Keep ``_sleeping_tags`` as the engine-wide union of per-stage tags.""" + per_stage = getattr(self, "_stage_sleeping_tags", None) or {} + union: set[str] = set() + for stage_tags in per_stage.values(): + union.update(stage_tags) + self._sleeping_tags = union + + def _record_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None: + per_stage = getattr(self, "_stage_sleeping_tags", None) + if per_stage is None: + per_stage = {} + self._stage_sleeping_tags = per_stage + tag_set = set(tags) + for sid in stage_ids: + per_stage.setdefault(sid, set()).update(tag_set) + self._refresh_union_sleeping_tags() + + def _clear_stage_sleep(self, stage_ids: list[int], tags: Iterable[str]) -> None: + per_stage = getattr(self, "_stage_sleeping_tags", None) + if per_stage is None: + self._sleeping_tags = set() + return + tag_set = set(tags) + for sid in stage_ids: + remaining = per_stage.get(sid) + if remaining is None: + continue + remaining.difference_update(tag_set) + if not remaining: + per_stage.pop(sid, None) + self._refresh_union_sleeping_tags() + async def pause_generation( self, *, mode: PauseMode = "abort", wait_for_inflight_requests: bool = False, clear_cache: bool = True, + stage_ids: list[int] | None = None, ) -> None: - """Pause generation.""" + """Pause generation, mirroring vLLM AsyncLLM.pause_generation. + + 1. Stop frontend admission (``_paused``). + 2. For AR/LLM stages, call EngineCore.pause_scheduler via the + Orchestrator loop (abort/wait/keep + optional cache clear). + 3. Diffusion stages have no EngineCore scheduler — only frontend + admission is paused for them. + + Note: ``sleep()`` already pauses the AR scheduler internally (same as + vLLM EngineCore.sleep). Call this API when you need pause *without* + freeing GPU memory (e.g. weight sync). + """ + if wait_for_inflight_requests: + mode = "wait" + async with self._pause_cond: - if self._paused: - return + # Keep running EngineCore pause + cache clear even when frontend + # admission is already paused (sleep or a prior pause_generation). self._paused = True + self._hold_admission_until_resume = True - # TODO: Implement request draining if wait_for_inflight_requests + ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + if ar_stage_ids: + logger.info( + "[%s] Pausing AR stage(s) %s via EngineCore.pause_scheduler(mode=%s)", + self._name, + ar_stage_ids, + mode, + ) + # Same API name as vLLM AsyncMPClient.pause_scheduler_async; routed + # through collective_rpc so it runs on the orchestrator event loop. + await self._engine_core_rpc( + "pause_scheduler", + stage_ids=ar_stage_ids, + kwargs={"mode": mode, "clear_cache": clear_cache}, + ) + # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler + # already clears AR-side caches when clear_cache=True. if clear_cache: - # Clear caches for all stages. await self.reset_prefix_cache( reset_running_requests=not wait_for_inflight_requests, reset_connector=True, @@ -1114,14 +1373,20 @@ class AsyncOmni(EngineClient, OmniBase): await self.reset_mm_cache() await self.reset_encoder_cache() - async def resume_generation(self) -> None: - """Resume generation.""" + async def resume_generation(self, stage_ids: list[int] | None = None) -> None: + """Resume generation after :meth:`pause_generation`.""" + ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + if ar_stage_ids: + logger.info("[%s] Resuming AR stage(s) %s via EngineCore", self._name, ar_stage_ids) + await self._engine_core_rpc("resume_scheduler", stage_ids=ar_stage_ids) + async with self._pause_cond: self._paused = False + self._hold_admission_until_resume = False self._pause_cond.notify_all() async def is_paused(self) -> bool: - """Check if paused.""" + """Check if frontend admission is paused.""" async with self._pause_cond: return self._paused @@ -1151,11 +1416,24 @@ class AsyncOmni(EngineClient, OmniBase): return await self.collective_rpc(method="profile", args=(False, None), stage_ids=stages) async def reset_mm_cache(self) -> None: - """Reset the multi-modal cache for all stages. + """Reset the frontend (P0) multimodal processor cache. - TODO: Forward to Orchestrator process via message. + ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache. + Clearing P0 avoids hash-only follow-up requests after that reset. """ - logger.warning("[AsyncOmni] reset_mm_cache not yet supported with Orchestrator process") + processor = getattr(self, "input_processor", None) + if processor is None: + processor = getattr(self.engine, "input_processor", None) + cache = getattr(processor, "mm_processor_cache", None) + if cache is None: + logger.debug("[AsyncOmni] reset_mm_cache: no frontend mm_processor_cache") + return + for name in ("clear", "reset", "clear_cache"): + fn = getattr(cache, name, None) + if callable(fn): + fn() + return + logger.debug("[AsyncOmni] reset_mm_cache: cache has no clear/reset method") async def reset_encoder_cache(self) -> None: """Reset the encoder cache for all stages. @@ -1179,41 +1457,108 @@ class AsyncOmni(EngineClient, OmniBase): async def sleep( self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort" ) -> list[OmniACK]: - self._final_output_handler() - if stage_ids is None: - stage_ids = list(range(len(self.engine.stage_clients))) - total_workers = 0 - for sid in stage_ids: - client = self.engine.stage_clients[sid] - # During the Diffusion phase, regardless of the TP amount, - # currently only a summary ACK is reported at Rank 0. - if getattr(client, "stage_type", "") == "diffusion": - total_workers += 1 - else: - config = self.engine.stage_vllm_configs[sid] - actual_tp = config.parallel_config.tensor_parallel_size if config else 1 - total_workers += actual_tp + """Put stages to sleep. + AR/LLM stages use EngineCore.sleep (pause scheduler, wait idle, then + offload/discard memory) — matching vLLM AsyncLLM.sleep. + + Diffusion stages keep the existing worker-level handle_sleep_task RPC + because StageDiffusionProc does not expose EngineCore.pause_scheduler. + + Frontend admission is blocked at the start of this call (``_paused``) + so pipelined :meth:`generate` cannot race into stages while sleep is + in flight. This does **not** invoke EngineCore.pause_scheduler again + (sleep already pauses the AR scheduler). + + For AR / mixed engines, ``wake_up`` does **not** clear ``_paused``; + callers must :meth:`resume_generation` when ready (typical trainer + order: pause → abort → sleep → train → wake → resume). Diffusion-only + engines have no EngineCore pause to hold, so ``wake_up`` restores + admission and ``sleep → wake → generate`` keeps working. + """ + # Block admission before any sleep RPC so generate() waits on + # _pause_cond during the drain/offload window. Wait until generate() + # coroutines that already passed the pause check have submitted (or + # failed) so EngineCore does not see ADD frames while sleeping. + async with self._pause_cond: + self._paused = True + await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0) + + # P0 sender cache must drop hashes before EngineCore.sleep clears P1. + await self.reset_mm_cache() + + self._final_output_handler() + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + final_acks: list[OmniACK] = [] + if ar_stage_ids: + self._hold_admission_until_resume = True + logger.info( + "[%s] Sleeping AR stage(s) %s via EngineCore.sleep(level=%s, mode=%s)", + self._name, + ar_stage_ids, + level, + mode, + ) + await self._engine_core_rpc( + "sleep", + stage_ids=ar_stage_ids, + args=(level, mode), + ) + # EngineCore.sleep has no OmniACK handshake; emit stage-level SUCCESS + # markers so callers/tests that count ACKs keep a stable API. + task_id = f"engine_core-sleep-{uuid.uuid4().hex[:8]}" + final_acks.extend( + OmniACK( + task_id=task_id, + status="SUCCESS", + stage_id=sid, + rank=0, + metadata={"path": "engine_core", "level": level, "mode": mode}, + ) + for sid in ar_stage_ids + ) + + if diffusion_stage_ids: + final_acks.extend(await self._sleep_diffusion(diffusion_stage_ids, level)) + + self._record_stage_sleep( + ar_stage_ids + diffusion_stage_ids, + [CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value], + ) + if level == 2: + self._level2_sleeping = True + return final_acks + + async def _sleep_diffusion(self, stage_ids: list[int], level: int) -> list[OmniACK]: + """Worker-level sleep RPC for diffusion stages only.""" + # Diffusion reports one summary ACK at rank 0 regardless of TP. + total_workers = len(stage_ids) task_id = str(uuid.uuid4()) self.event_resolver.watch_task(task_id, expected_count=total_workers) - logger.info(f"[{self._name}] Sleep initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") + logger.info("[%s] Sleep (diffusion) initiated (Task: %s).", self._name, task_id) task = OmniSleepTask(level=level, task_id=task_id) rpc_results = await self.collective_rpc(method="handle_sleep_task", args=(task,), stage_ids=stage_ids) - final_acks = [] + final_acks: list[OmniACK] = [] for stage_res in rpc_results: worker_acks = stage_res if isinstance(stage_res, list) else [stage_res] for ack in worker_acks: if ack is not None: await self.event_resolver.resolve(ack) final_acks.append(ack) - if not hasattr(self, "_sleeping_tags"): - self._sleeping_tags = set() - self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value]) - if level == 2: - self._level2_sleeping = True return final_acks async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]: + """Wake stages after sleep. + + AR/LLM stages use EngineCore.wake_up (restore memory, auto-resume + scheduler). Diffusion stages keep the worker-level wake RPC. + + Does **not** clear the frontend ``_paused`` admission gate when + :meth:`pause_generation` ran or AR stages were slept — call + :meth:`resume_generation` when the trainer is ready to admit new + requests. Diffusion-only ``sleep`` uses ``_paused`` only as a race + guard; this method restores admission after a successful wake. + """ self._final_output_handler() if getattr(self, "_level2_sleeping", False): @@ -1223,7 +1568,12 @@ class AsyncOmni(EngineClient, OmniBase): "Use sleep(level=1) instead, which offloads weights to CPU RAM " "and supports fast DMA restore." ) - _current_tags = getattr(self, "_sleeping_tags", set()) + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + target_stage_ids = ar_stage_ids + diffusion_stage_ids + _current_tags = self._sleeping_tags_for_stages(target_stage_ids) + per_stage = getattr(self, "_stage_sleeping_tags", None) or {} + if not _current_tags and not per_stage: + _current_tags = set(getattr(self, "_sleeping_tags", set())) if tags is None: requested_tags = list(_current_tags) else: @@ -1232,39 +1582,64 @@ class AsyncOmni(EngineClient, OmniBase): logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.") return [] - if stage_ids is None: - stage_ids = list(range(len(self.engine.stage_clients))) - total_workers = 0 - for sid in stage_ids: - client = self.engine.stage_clients[sid] - if getattr(client, "stage_type", "") == "diffusion": - total_workers += 1 - else: - config = self.engine.stage_vllm_configs[sid] - total_workers += config.parallel_config.tensor_parallel_size if config else 1 + final_acks: list[OmniACK] = [] + if ar_stage_ids: + logger.info("[%s] Waking AR stage(s) %s via EngineCore", self._name, ar_stage_ids) + await self._engine_core_rpc( + "wake_up", + stage_ids=ar_stage_ids, + kwargs={"tags": requested_tags}, + ) + task_id = f"engine_core-wake-{uuid.uuid4().hex[:8]}" + final_acks.extend( + OmniACK( + task_id=task_id, + status="SUCCESS", + stage_id=sid, + rank=0, + metadata={"path": "engine_core", "tags": list(requested_tags)}, + ) + for sid in ar_stage_ids + ) + + if diffusion_stage_ids: + final_acks.extend(await self._wake_diffusion(diffusion_stage_ids, requested_tags)) + + self._clear_stage_sleep(target_stage_ids, requested_tags) + # Only clear the level-2 flag once all tags are warm, in case partial + # wake support (e.g. tags=["kv_cache"] only) is added in the future. + if not getattr(self, "_sleeping_tags", None): + self._level2_sleeping = False + logger.info( + "[%s] Wake-up complete for stage(s) %s.", + self._name, + ar_stage_ids + diffusion_stage_ids, + ) + # Diffusion-only sleep uses `_paused` as a race guard. Restore + # generate() admission after memory is back. AR/mixed sleep and + # pause_generation keep the trainer hold until resume_generation. + if not getattr(self, "_hold_admission_until_resume", False): + async with self._pause_cond: + self._paused = False + self._pause_cond.notify_all() + return final_acks + + async def _wake_diffusion(self, stage_ids: list[int], requested_tags: list[str]) -> list[OmniACK]: + """Worker-level wake RPC for diffusion stages only.""" + total_workers = len(stage_ids) task_id = str(uuid.uuid4()) self.event_resolver.watch_task(task_id, expected_count=total_workers) - logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") + logger.info("[%s] Wake-up (diffusion) initiated (Task: %s).", self._name, task_id) task = OmniWakeTask(tags=requested_tags, task_id=task_id) rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids) - final_acks = [] + final_acks: list[OmniACK] = [] for stage_res in rpc_results: worker_acks = stage_res if isinstance(stage_res, list) else [stage_res] for ack in worker_acks: if ack is not None: await self.event_resolver.resolve(ack) final_acks.append(ack) - current_omni_platform.synchronize() await asyncio.sleep(0.1) - - for t in requested_tags: - if hasattr(self, "_sleeping_tags"): - self._sleeping_tags.discard(t) - # Only clear the level-2 flag once all tags are warm, in case partial - # wake support (e.g. tags=["kv_cache"] only) is added in the future. - if not getattr(self, "_sleeping_tags", None): - self._level2_sleeping = False - logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.") return final_acks async def is_sleeping(self) -> bool: diff --git a/vllm_omni/entrypoints/cli/__init__.py b/vllm_omni/entrypoints/cli/__init__.py index 2ffba6130..c13ac94be 100644 --- a/vllm_omni/entrypoints/cli/__init__.py +++ b/vllm_omni/entrypoints/cli/__init__.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """CLI helpers for vLLM-Omni entrypoints.""" # To ensure patch imports work properly, disable unused import checks @@ -6,8 +9,9 @@ from vllm_omni.benchmarks.patch import patch # isort: on +from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand from .serve import OmniServeCommand -__all__ = ["OmniServeCommand", "OmniBenchmarkServingSubcommand"] +__all__ = ["OmniServeCommand", "OmniBenchmarkServingSubcommand", "OmniDuplexEvalSubcommand"] diff --git a/vllm_omni/entrypoints/cli/benchmark/cli_args.py b/vllm_omni/entrypoints/cli/benchmark/cli_args.py index c7f220870..2cf668769 100644 --- a/vllm_omni/entrypoints/cli/benchmark/cli_args.py +++ b/vllm_omni/entrypoints/cli/benchmark/cli_args.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """vLLM-Omni extensions for the ``vllm bench serve`` CLI. Core functions: @@ -16,6 +19,63 @@ by ``add_omni_args``. """ import argparse +import math +from pathlib import Path + +_DEFAULT_OMNIINTERACT_NUM_PROMPTS = 3 + + +def _positive_finite_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError(f"must be a finite positive number, got {value!r}") + return parsed + + +def _existing_file(value: str) -> str: + path = Path(value).expanduser() + if not path.is_file(): + raise argparse.ArgumentTypeError(f"file does not exist: {value!r}") + return str(path) + + +def add_omniinteract_cli_args(parser: argparse.ArgumentParser) -> None: + from vllm_omni.benchmarks.data_modules.omniinteract_dataset import ( + DEFAULT_MAX_VIDEO_DURATION_S, + OMNIINTERACT_SUBSETS, + ) + + group = parser.add_argument_group("OmniInteract Benchmark Options") + group.add_argument( + "--omniinteract-subsets", nargs="+", choices=OMNIINTERACT_SUBSETS, default=list(OMNIINTERACT_SUBSETS) + ) + group.add_argument( + "--omniinteract-timeout-s", type=_positive_finite_float, default=900.0, help="Complete session timeout." + ) + group.add_argument( + "--omniinteract-media-timeout-s", + type=_positive_finite_float, + default=600.0, + help="Per-command media timeout.", + ) + group.add_argument( + "--omniinteract-max-video-duration-s", + type=_positive_finite_float, + default=DEFAULT_MAX_VIDEO_DURATION_S, + help="Reject media longer than this safety limit before decoding.", + ) + group.add_argument( + "--omniinteract-ref-audio", type=_existing_file, help="Reference WAV for native-duplex audio output." + ) + group.add_argument( + "--omniinteract-require-response", action="store_true", help="Fail LISTEN-only functional E2E cases." + ) + group.add_argument( + "--omniinteract-output-dir", + type=Path, + default=Path("omniinteract-output"), + help="Directory for case and evaluator artifacts.", + ) def add_multi_stage_cli_args(parser: argparse.ArgumentParser) -> None: @@ -188,6 +248,7 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None: _OMNI_BENCH_DATASET_CHOICES = ( "daily-omni", + "omniinteract", "seed-tts", "seed-tts-text", "seed-tts-design", @@ -218,6 +279,11 @@ def extend_omni_choices(parser: argparse.ArgumentParser) -> None: def update_omni_help(parser: argparse.ArgumentParser) -> None: """Update upstream argument help text to describe Omni-specific behavior.""" for action in parser._actions: + if action.dest == "num_prompts": + action.help = ( + f"{action.help} OmniInteract uses {_DEFAULT_OMNIINTERACT_NUM_PROMPTS} when this option is omitted; " + "0 selects all available cases." + ) if action.dest == "percentile_metrics": action.help = ( "Comma-separated list of selected metrics to report percentiles. " @@ -254,6 +320,7 @@ def update_omni_help(parser: argparse.ArgumentParser) -> None: def add_omni_args(parser: argparse.ArgumentParser) -> None: """Register all vLLM-Omni serving benchmark arguments.""" add_daily_omni_cli_args(parser) + add_omniinteract_cli_args(parser) add_seed_tts_cli_args(parser) add_multi_stage_cli_args(parser) add_diffusion_cli_args(parser) @@ -261,6 +328,28 @@ def add_omni_args(parser: argparse.ArgumentParser) -> None: def preprocess_serve_args(args: argparse.Namespace) -> None: """Apply serving benchmark CLI transformations after parsing.""" + if getattr(args, "dataset_name", None) == "omniinteract": + if getattr(args, "backend", None) != "openai-realtime-duplex": + raise ValueError("OmniInteract requires --backend openai-realtime-duplex") + if getattr(args, "endpoint", None) != "/v1/realtime": + raise ValueError("OmniInteract requires --endpoint /v1/realtime") + if not getattr(args, "omniinteract_ref_audio", None): + raise ValueError("OmniInteract requires --omniinteract-ref-audio") + if getattr(args, "ignore_eos", False): + raise ValueError("OmniInteract does not support --ignore-eos") + if getattr(args, "profile", False): + raise ValueError("OmniInteract does not support --profile") + if getattr(args, "skip_tokenizer_init", False): + raise ValueError("OmniInteract does not support --skip-tokenizer-init") + if float(getattr(args, "probe_request_rate", 0.0) or 0.0) > 0: + raise ValueError("OmniInteract does not support --probe-request-rate") + if "num_prompts" not in getattr(args, "explicit_keys", ()): + args.num_prompts = _DEFAULT_OMNIINTERACT_NUM_PROMPTS + max_concurrency = getattr(args, "max_concurrency", None) + if max_concurrency is None: + args.max_concurrency = 1 + elif max_concurrency <= 0: + raise ValueError("OmniInteract requires --max-concurrency to be positive") extra_body = dict(getattr(args, "extra_body", None) or {}) bot_task = getattr(args, "bot_task", None) if getattr(args, "backend", None) == "openai-image-edits-omni" and bot_task is not None: diff --git a/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py b/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py new file mode 100644 index 000000000..a43a5ec71 --- /dev/null +++ b/vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import argparse +import asyncio +import concurrent.futures +import json +from pathlib import Path + +from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DEFAULT_DATASET, load_samples +from vllm_omni.benchmarks.duplex.omni_duplex_eval_eval import evaluate_sample, summarize_scores +from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge +from vllm_omni.benchmarks.duplex.omni_duplex_eval_runner import generate_sample +from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase + + +def _common(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--dataset", default=DEFAULT_DATASET) + parser.add_argument("--split", default="all") + parser.add_argument("--family", choices=("all", "rtd", "pr"), default="all") + parser.add_argument("--media-root") + parser.add_argument("--limit", type=int) + parser.add_argument("--ids", nargs="*") + + +def add_cli_args(parser: argparse.ArgumentParser) -> None: + actions = parser.add_subparsers(dest="action", required=True) + generate = actions.add_parser("generate") + _common(generate) + generate.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1") + generate.add_argument("--model", required=True) + generate.add_argument("--ref-audio", required=True) + generate.add_argument("--response-root", required=True) + generate.add_argument("--fps", type=float, default=1.0) + generate.add_argument("--mix", choices=("question",), default="question") + generate.add_argument("--pace", choices=("realtime", "as-fast-as-possible"), default="realtime") + generate.add_argument("--clock", choices=("media",), default="media") + generate.add_argument("--concurrency", type=int, default=1) + generate.add_argument("--overwrite", action="store_true") + evaluate = actions.add_parser("evaluate") + _common(evaluate) + evaluate.add_argument("--response-root", required=True) + evaluate.add_argument("--score-root", required=True) + evaluate.add_argument("--judge-base-url", default="http://127.0.0.1:8000") + evaluate.add_argument("--judge-model", required=True) + evaluate.add_argument("--judge-api-key", default="EMPTY") + evaluate.add_argument("--judge-video-mode", choices=("video_url", "frame-sample"), default="video_url") + evaluate.add_argument("--judge-fps", type=int, default=2) + evaluate.add_argument("--window-size", type=float, default=10.0) + evaluate.add_argument("--allow-invalid-clock", action="store_true") + evaluate.add_argument("--eval-workers", type=int, default=1) + evaluate.add_argument("--overwrite", action="store_true") + summarize = actions.add_parser("summarize") + summarize.add_argument("--score-root", required=True) + + +def run(args: argparse.Namespace) -> int: + if args.action == "summarize": + print(json.dumps(summarize_scores(args.score_root), ensure_ascii=False, indent=2)) + return 0 + + samples = load_samples( + args.dataset, + split=args.split, + family=args.family, + media_root=args.media_root, + limit=args.limit, + ids=args.ids, + ) + if args.action == "generate": + if args.concurrency < 1: + raise ValueError("--concurrency must be at least 1") + + async def generate() -> None: + semaphore = asyncio.Semaphore(args.concurrency) + + async def generate_one(sample) -> None: + async with semaphore: + await generate_sample( + sample, + url=args.url, + model=args.model, + ref_audio=args.ref_audio, + output_root=args.response_root, + fps=args.fps, + mix=args.mix, + pace=args.pace, + clock=args.clock, + overwrite=args.overwrite, + ) + + await asyncio.gather(*(generate_one(sample) for sample in samples)) + + asyncio.run(generate()) + return 0 + + if args.eval_workers < 1: + raise ValueError("--eval-workers must be at least 1") + judge = DuplexJudge(args.judge_base_url, args.judge_model, api_key=args.judge_api_key) + + def evaluate_one(sample) -> None: + response_path = Path(args.response_root) / sample.split / f"{sample.id}.json" + score_path = Path(args.score_root) / sample.split / f"{sample.id}.json" + if score_path.exists() and not args.overwrite: + return + evaluate_sample( + sample, + response_path, + score_path, + judge, + judge_fps=args.judge_fps, + judge_video_mode=args.judge_video_mode, + window_size=args.window_size, + allow_invalid_clock=args.allow_invalid_clock, + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=args.eval_workers) as executor: + list(executor.map(evaluate_one, samples)) + return 0 + + +class OmniDuplexEvalSubcommand(OmniBenchmarkSubcommandBase): + """Run the Omni-DuplexEval generation or scoring workflow.""" + + name = "omni-duplex-eval" + help = "Generate, evaluate, or summarize Omni-DuplexEval artifacts." + + @classmethod + def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: + add_cli_args(parser) + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + run(args) diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py index 5e46aecd2..7c7b05221 100644 --- a/vllm_omni/entrypoints/cli/benchmark/serve.py +++ b/vllm_omni/entrypoints/cli/benchmark/serve.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Online serving benchmark subcommand for vLLM-Omni. ``OmniBenchmarkServingSubcommand`` starts from vLLM's serving benchmark @@ -24,7 +27,7 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase): """The `serve` subcommand for vllm bench.""" name = "serve" - help = "Benchmark the online serving throughput. Supports Daily-Omni and Seed-TTS datasets." + help = "Benchmark online serving. Supports Daily-Omni, OmniInteract, and Seed-TTS datasets." @classmethod def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index dfa222cd1..f4c890816 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Omni serve command for vLLM-Omni. @@ -537,6 +540,15 @@ class OmniServeCommand(CLISubcommand): "'strict' keeps the original divisibility requirements; " "'advanced_uaa' enables the experimental UAA path for uneven sequence/head shapes.", ) + omni_config_group.add_argument( + "--ulysses-a2a-permute", + action=argparse.BooleanOptionalAction, + default=None, + help=( + "Enable fused permute-free Ulysses all-to-all over NCCL symmetric memory. " + "Only strict Ulysses layouts are eligible. Defaults to disabled." + ), + ) omni_config_group.add_argument( "--ring", "--ring-degree", @@ -606,6 +618,12 @@ class OmniServeCommand(CLISubcommand): "May be combined with --diffusion-attention-config.per_role.* overrides, " "but mutually exclusive with --diffusion-attention-config.default.backend.", ) + omni_config_group.add_argument( + "--fastvideo-vsa-topk", + type=int, + default=None, + help="Number of key/value blocks selected per query block by FASTVIDEO_VSA.", + ) omni_config_group.add_argument( "--diffusion-attention-config", "-dac", @@ -729,6 +747,38 @@ class OmniServeCommand(CLISubcommand): help="Keep this many leading main-DiT blocks resident on the device " "while distributed layerwise offload streams the remaining blocks.", ) + omni_config_group.add_argument( + "--host-weight-runtime-mode", + choices=("disabled", "preferred", "required"), + default="disabled", + help=( + "Host Weight Runtime policy for eligible no-AllGather DLO: " + "disabled does not consult HWR; preferred restores an exact hit " + "or canonically loads and publishes on a miss; required restores " + "an exact hit or fails startup. Populate a required store with " + "preferred first." + ), + ) + omni_config_group.add_argument( + "--host-weight-runtime-root", + type=str, + default=None, + help=( + "Writable node-local Host Weight Runtime store shared by workers " + "in one storage domain. Required for preferred and required; use " + "the same persistent path for population and serving." + ), + ) + omni_config_group.add_argument( + "--dlo-host-registration-limit-gib", + type=float, + default=0.0, + help=( + "Optional per-worker GiB ceiling for registering an HWR mmap for direct H2D. " + "Zero applies no additional ceiling. Eligible no-AllGather HWR hits attempt registration " + "under the existing pinned-memory policy and fall back to bounded staging when unavailable." + ), + ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio", @@ -780,10 +830,9 @@ class OmniServeCommand(CLISubcommand): omni_config_group.add_argument( "--text-encoder-tp-size", type=int, - default=1, + default=None, help="Tensor-parallel degree for the diffusion text encoder. " - "Shards the Qwen3-VL encoder across the first N DiT ranks, " - "removing the rank-0 encoder memory hotspot in no-offload runs. " + "Shards the encoder across the first N DiT ranks. " "Equivalent to setting DiffusionParallelConfig.text_encoder_tp_size.", ) omni_config_group.add_argument( @@ -823,13 +872,12 @@ class OmniServeCommand(CLISubcommand): default=False, help="Enable chunked streaming output for diffusion (mainly video generation) models that support it.", ) - # TTS-specific parameters omni_config_group.add_argument( "--tts-max-instructions-length", type=int, default=None, - help="Maximum length for TTS voice style instructions (overrides stage config, default: 500).", + help="Maximum length for TTS voice style instructions (overrides the pipeline default, default: 500).", ) # Disable safety guardrails for this server (currently only applicable for Cosmos3) @@ -945,8 +993,6 @@ def run_headless(args: TrackingNamespace) -> None: config_path, stage_configs, _ = load_and_resolve_stage_configs( model, - # The serve CLI no longer accepts legacy stage_args YAMLs. - None, args_dict, # store_true cannot express an explicit False: absent maps to None # ("not specified") so the deploy yaml's per-stage value applies. diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 9526df36a..0aa0c7c5a 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -17,6 +17,8 @@ class ClientRequestState: self.stage_id: int | None = None self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None + # Request-scoped idempotency guard for Prometheus failure counters. + self.failure_recorded = False # Wall-clock time at which the user's request arrived in the engine # entrypoint. Set in async_omni.generate() before the orchestrator # accepts the request. Used as the t0 anchor for audio_ttfp. diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 14bdc10be..28bbe9219 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -200,6 +200,8 @@ class Omni(OmniBase): self._log_summary_and_cleanup(req_id) except Exception: if "active_reqs" in locals() and active_reqs: + for req_id in active_reqs: + self._record_request_failure_once(req_id, reason="stage_error") self.abort(list(active_reqs)) raise finally: @@ -210,6 +212,7 @@ class Omni(OmniBase): request_ids = [request_id] if isinstance(request_id, str) else list(request_id) self.engine.abort(request_ids) for req_id in request_ids: + self._record_request_failure_once(req_id, reason="client_abort") self.request_states.pop(req_id, None) if self.log_stats: logger.info("[Omni] Aborted request(s) %s", ",".join(request_ids)) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index b787c44a4..0aa6832d0 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -1,15 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os import time import weakref -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from dataclasses import fields, is_dataclass from typing import TYPE_CHECKING, Any, Literal import huggingface_hub import vllm.envs as envs +from omegaconf import OmegaConf from vllm.logger import init_logger from vllm.transformers_utils.repo_utils import file_or_path_exists +from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -27,6 +33,11 @@ from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_ from vllm_omni.metrics.prometheus import OmniPrometheusMetrics from vllm_omni.metrics.stats import OrchestratorAggregator from vllm_omni.metrics.transfer import OmniTransferMetrics +from vllm_omni.metrics.utils import ( + extract_queue_wait_s, + normalize_failure_reason, + observe_stage_workload_metrics, +) from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.tracking_parser import TrackingNamespace @@ -66,6 +77,14 @@ def omni_snapshot_download(model_id: str) -> str: if os.path.exists(model_id): return model_id + # Object-storage models must remain URIs until each stage constructs its + # ModelConfig. vLLM then materializes only config/tokenizer files locally + # and keeps the URI in model_weights for Run:AI streaming. Treating the URI + # as a Hugging Face repo here either fails validation or downloads through + # the wrong backend before the stage processes are created. + if is_runai_obj_uri(model_id): + return model_id + # TODO: this is just a workaround for quickly use modelscope, we should support # modelscope in weight loading feature instead of using `snapshot_download` # Read through ``vllm.envs`` so the flag keeps vLLM's semantics (only @@ -168,7 +187,9 @@ class OmniBase(PDDisaggregationMixin): # override the deploy YAML's ``async_chunk: true`` default. async_chunk = kwargs.get("async_chunk") output_modalities = kwargs.pop("output_modalities", None) - diffusion_batch_size: int = kwargs.pop("diffusion_batch_size", 1) + + if "diffusion_batch_size" in kwargs: + raise TypeError("`diffusion_batch_size` has been removed in Omni/AsyncOmni. Use `max_num_seqs`.") if "log_requests" in kwargs: raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.") @@ -188,13 +209,14 @@ class OmniBase(PDDisaggregationMixin): # (which forwards it to the Orchestrator background thread for # TX-side emit; see Orchestrator._forward_to_next_stage). self.transfer_metrics = OmniTransferMetrics(model_name=model, log_stats=log_stats) + self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats) st = time.time() self.engine = AsyncOmniEngine( model=model, init_timeout=init_timeout, stage_init_timeout=stage_init_timeout, - diffusion_batch_size=diffusion_batch_size, transfer_emitter=self.transfer_metrics, + prom_metrics=self.prom_metrics, log_stats=log_stats, **kwargs, ) @@ -210,10 +232,10 @@ class OmniBase(PDDisaggregationMixin): self.request_states: dict[str, ClientRequestState] = {} self._consumed_metric_messages: dict[str, set[int]] = {} - self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats) self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats) self.default_sampling_params_list = self.engine.default_sampling_params_list + self.sampling_constraints_list = self._get_sampling_constraints_list(self.engine.stage_configs) if not self.output_modalities: self.output_modalities = [ self.engine.get_stage_metadata(i).final_output_type for i in range(self.engine.num_stages) @@ -316,6 +338,7 @@ class OmniBase(PDDisaggregationMixin): sampling_params_list: Sequence[Any] | Any | None, allow_delta_coercion: bool = False, ) -> Sequence[Any]: + """Resolve request parameters; pipeline sampling constraints override caller values.""" if sampling_params_list is None: normalized = self.default_sampling_params_list # Set the output kind to delta since no params were specified @@ -330,31 +353,87 @@ class OmniBase(PDDisaggregationMixin): raise ValueError(f"Expected {self.num_stages} sampling params, got a single sampling params object") if len(normalized) != self.num_stages: raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}") + + if sampling_params_list is not None: + normalized = [ + self._apply_sampling_constraints(params, constraints) + for params, constraints in zip(normalized, self.sampling_constraints_list, strict=True) + ] return normalized - def _fire_failure_counter_if_alive(self, request_id: str) -> None: - """Fire the abort/exception bucket of requests_success_total. + @staticmethod + def _get_sampling_constraints_list(stage_configs: Sequence[Any]) -> list[dict[str, Any]]: + """Extract each stage's required sampling settings from runtime configs.""" + constraints_list = [] + for stage_config in stage_configs: + constraints = getattr(stage_config, "sampling_constraints", {}) + if not isinstance(constraints, Mapping): + constraints_list.append({}) + continue + if OmegaConf.is_config(constraints): + constraints = OmegaConf.to_container(constraints, resolve=True) + constraints_list.append(dict(constraints)) + return constraints_list - Called from cancel / exception paths in async_omni.generate() BEFORE - _abort_internal_requests pops request_states — that method resolves - the internal id by dict lookup, so popping first would no-op it. We - keep this counter fire separate from _log_summary_and_cleanup (which - pops) so the abort path can still find the state to clean up. - """ + @staticmethod + def _apply_sampling_constraints(params: Any, constraints: Mapping[str, Any]) -> Any: + """Rebuild params with pipeline-required settings without mutating caller input.""" + if not constraints: + return params + if isinstance(params, Mapping): + return {**params, **constraints} + if is_dataclass(params): + values = {field.name: getattr(params, field.name) for field in fields(params) if field.init} + elif struct_fields := getattr(params, "__struct_fields__", None): + values = { + name: getattr(params, name) + for name in struct_fields + if not name.startswith("_") and name != "output_text_buffer_length" + } + else: + raise TypeError(f"Expected a mapping, dataclass, or msgspec struct, got {type(params).__name__}") + return type(params)(**{**values, **constraints}) + + def _record_request_failure_once(self, request_id: str, reason: str) -> None: req_state = self.request_states.get(request_id) prom = getattr(self, "prom_metrics", None) - if req_state is None or req_state.metrics is None or prom is None: + metrics = getattr(req_state, "metrics", None) + if metrics is None or prom is None: + return + if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False): return - if str(request_id) not in req_state.metrics.e2e_done: - prom.request_failed() - def _log_summary_and_cleanup(self, request_id: str) -> None: + req_state.failure_recorded = True + prom.request_failed() + prom.inc_requests_failed(normalize_failure_reason(reason)) + + def _publish_request_gauges(self, total: int) -> None: + """Publish num_requests_running / num_requests_waiting for ``total`` in-flight requests. + + ``_running_counter`` counts requests from the moment the orchestrator + dispatches them to a stage engine, which includes requests still queued + inside a stage scheduler. The orchestrator reports those via + ``_engines_waiting_counter``. + """ + prom = getattr(self, "prom_metrics", None) + if prom is None: + return + engine = getattr(self, "engine", None) + counter = getattr(engine, "_running_counter", None) + waiting_counter = getattr(engine, "_engines_waiting_counter", None) + dispatched = counter.value if counter is not None else total + in_engine_waiting = waiting_counter.value if waiting_counter is not None else 0 + in_engine_waiting = min(max(0, in_engine_waiting), dispatched) + prom.set_running(max(0, dispatched - in_engine_waiting)) + prom.set_waiting(max(0, total - dispatched) + in_engine_waiting) + + def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") -> None: req_state = self.request_states.get(request_id) try: if req_state is None or req_state.metrics is None: return if str(request_id) not in req_state.metrics.e2e_done: - self.prom_metrics.request_failed() + self._record_request_failure_once(request_id, reason) if self.log_stats: # Emit per-request orchestrator timing (including e2e_total_ms) # before dropping request state. @@ -373,13 +452,7 @@ class OmniBase(PDDisaggregationMixin): # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. - prom = getattr(self, "prom_metrics", None) - counter = getattr(getattr(self, "engine", None), "_running_counter", None) - if prom is not None: - total = len(self.request_states) - running = counter.value if counter is not None else total - prom.set_running(running) - prom.set_waiting(max(0, total - running)) + self._publish_request_gauges(len(self.request_states)) def _compute_final_stage_id(self, output_modalities: list[str] | None) -> int: return get_final_stage_id_for_e2e( @@ -571,6 +644,30 @@ class OmniBase(PDDisaggregationMixin): metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs) metrics.on_stage_metrics(stage_id, req_id, _m, output_type) consumed.add(msg_id) + if peak_memory_mb > 0: + self.prom_metrics.set_peak_memory(stage_id, peak_memory_mb) + self.prom_metrics.observe_stage_gen_time( + stage_id, + stage_meta.stage_type, + _m.stage_gen_time_ms / 1000.0, + ) + observe_stage_workload_metrics( + self.prom_metrics, + stage_type=stage_meta.stage_type, + stage_metrics=_m, + ) + # Only meaningful for diffusion stages (text stages have no exec key). + _diff_m = _m.diffusion_metrics or {} + stage_in_queue_s = _diff_m.get("scheduler_queue_wait_s") + if stage_in_queue_s is not None: + self.prom_metrics.observe_stage_in_queue(stage_id, float(stage_in_queue_s)) + if _m.output_unit_type == "image": + if result.replica_id is not None: + self.mod_metrics.observe_image_ttfp( + str(stage_id), + str(result.replica_id), + _m.image_time_to_first_output_ms / 1000.0, + ) if not stage_meta.final_output: return None @@ -604,6 +701,10 @@ class OmniBase(PDDisaggregationMixin): _gen_tok += int(evt.num_tokens_out) self.prom_metrics.observe_tokens(_prompt_tok, _gen_tok) + queue_wait_s = extract_queue_wait_s(_m.pipeline_timings if _m is not None else None) + if queue_wait_s is not None: + self.prom_metrics.observe_queue_wait(queue_wait_s) + # Modality observe inside the same finalize guard so it fires # once per request and inherits the try/except isolation. observe_modality_at_finalize( @@ -622,12 +723,8 @@ class OmniBase(PDDisaggregationMixin): # hasn't popped self.request_states yet — exclude the finalizing # request from `total` so waiting doesn't read 1 and stay stuck # there until the next request arrives. - counter = getattr(self.engine, "_running_counter", None) is_finalizing = finished and stage_id == final_stage_id_for_e2e - total = max(0, len(self.request_states) - (1 if is_finalizing else 0)) - running = counter.value if counter is not None else total - self.prom_metrics.set_running(running) - self.prom_metrics.set_waiting(max(0, total - running)) + self._publish_request_gauges(max(0, len(self.request_states) - (1 if is_finalizing else 0))) response_metrics: dict[str, Any] = {} stage_metrics: dict[str, dict[str, Any]] = {} diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 0a8daa055..1e55d4476 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import asyncio import base64 import dataclasses @@ -43,6 +43,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, ) +from vllm.entrypoints.openai.cli_args import make_arg_parser # yapf conflicts with isort for this block # yapf: disable @@ -64,6 +65,14 @@ from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.pooling.scoring.serving import ServingScores from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens +# vLLM < 0.28 keeps create_error_response under serve.utils (it is not +# re-exported from vllm.entrypoints.serve); 0.28+ moved it under +# serve.exception_handling and re-exports it from the package root. +try: + from vllm.entrypoints.serve import create_error_response +except ImportError: + from vllm.entrypoints.serve.utils.error_response import create_error_response + # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. # Keep a fallback for older/newer upstream layouts during rebase windows. from vllm.entrypoints.serve.instrumentator.basic import base @@ -74,7 +83,6 @@ from vllm.entrypoints.serve.utils.api_utils import ( validate_json_request, with_cancellation, ) -from vllm.entrypoints.serve.utils.error_response import create_error_response from vllm.entrypoints.serve.utils.orca_metrics import metrics_header from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config @@ -539,6 +547,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None if _should_enable_profiler_endpoints(stage_configs): logger.warning("Profiler endpoints are enabled. This should ONLY be used for local development!") + _remove_route_from_app(app, "/start_profile", frozenset({"POST"})) + _remove_route_from_app(app, "/stop_profile", frozenset({"POST"})) app.include_router(profiler_router) vllm_config = await _get_vllm_config(engine_client) @@ -612,6 +622,9 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, await shutdown_task finally: state = getattr(app, "state", None) + serving_video = getattr(state, "openai_serving_video", None) if state is not None else None + if serving_video is not None: + serving_video.shutdown() serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None if serving_speech is not None: serving_speech.shutdown() @@ -679,7 +692,7 @@ async def build_async_omni_from_stage_config( EngineClient instance (AsyncOmni) ready for use Note: - Stage configurations are loaded from args.stage_configs_path if provided, + Stage configurations are loaded from ``args.deploy_config`` when provided, otherwise from the model's default configuration. """ @@ -822,6 +835,8 @@ async def omni_init_app_state( diffusion_engine=engine_client, model_name=model_name, stage_configs=diffusion_stage_configs, + allowed_local_media_path=getattr(args, "allowed_local_media_path", ""), + allowed_media_domains=getattr(args, "allowed_media_domains", None), ) state.openai_serving_duplex = None state.openai_streaming_speech = None @@ -996,11 +1011,9 @@ async def omni_init_app_state( ) # Warm up chat template processing to avoid first-request latency - # Upstream f5ffc59b6a moved warmup onto OnlineRenderer. Accelerator - # images can temporarily lag that renderer API, where warmup is optional. - renderer_warmup = getattr(state.online_renderer, "warmup", None) - if renderer_warmup is not None: - renderer_warmup() + # Upstream f5ffc59b6a removed OpenAIServingChat.warmup() and moved the + # warmup onto the renderer (OnlineRenderer.warmup()); mirror upstream. + state.online_renderer.warmup() state.openai_serving_completion = ( OpenAIServingCompletion( @@ -1156,7 +1169,7 @@ async def omni_init_app_state( state.openai_serving_duplex = None if state.openai_serving_chat is not None and should_enable_duplex_endpoint( state.stage_configs, - config_path=getattr(args, "stage_configs_path", None) or getattr(args, "deploy_config", None), + config_path=getattr(args, "deploy_config", None), ): from vllm_omni.experimental.fullduplex.openai.serving import OmniDuplexSessionHandler @@ -1481,8 +1494,7 @@ async def list_voices(raw_request: Request): status_code=HTTPStatus.NOT_FOUND, ) - # Get all speakers (both model built-in and uploaded) - speakers = sorted(handler.supported_speakers) if handler.supported_speakers else [] + speakers = sorted(handler._get_available_voices()) # Get uploaded speakers details uploaded_speakers = [] @@ -2377,7 +2389,7 @@ async def edit_images( status_code=generation_result.error.code if generation_result.error else 400, detail=generation_result.message, ) - images, _, _, cot_output = generation_result + images, stage_durations, peak_memory_mb, cot_output = generation_result else: # Single-stage diffusion: use the direct path. result = await _generate_with_async_omni( @@ -2388,6 +2400,8 @@ async def edit_images( request_id=request_id, ) images = _extract_images_from_result(result) + stage_durations = getattr(result, "stage_durations", None) + peak_memory_mb = getattr(result, "peak_memory_mb", None) logger.debug(f"Successfully generated {len(images)} image(s)") @@ -2408,6 +2422,10 @@ async def edit_images( output_format=output_format, size=size_str, cot_output=cot_output, + metrics={ + "stage_durations": stage_durations or None, + "peak_memory_mb": float(peak_memory_mb) if peak_memory_mb else None, + }, ) except (EngineGenerateError, EngineDeadError) as exc: @@ -3826,26 +3844,16 @@ async def omni_wakeup(request: OmniWakeupRequest, raw_request: Request): if __name__ == "__main__": - import asyncio - - from vllm.entrypoints.openai.cli_args import make_arg_parser - parser = TrackingArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server") parser = make_arg_parser(parser) - registered_flags = set() - for action in parser._actions: - registered_flags.update(action.option_strings) - - # FIXME - this is broken on `main`. `make_arg_parser` does not handle omni engine args properly. - if "--omni" not in registered_flags: - parser.add_argument("--omni", action="store_true", default=False, help="Enable vLLM-Omni mode.") - if "--enable-sleep-mode" not in registered_flags: - parser.add_argument( - "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode." - ) + # Ensure that passing --omni won't crash the server. + # NOTE: the value here does not matter since we are always running the Omni server + # when __main__ is called, i.e., --omni is only used when called through the entrypoints. + parser.add_argument("--omni", action="store_true", default=False) args = parser.parse_args() - if not hasattr(args, "model_tag"): - setattr(args, "model_tag", args.model) - if hasattr(args, "model_tag") and args.model_tag is None: - args.model_tag = args.model + # sync args.model to model_tag, because if we pass the model positionally, + # args.model will be the default from vLLM's ModelConfig (currently + # Qwen/Qwen3-0.6B) and crash cryptically. + if args.model_tag is not None: + args.model = args.model_tag asyncio.run(omni_run_server(args)) diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index b2b55bc52..1c556fc79 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -1,9 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import math from typing import Any, Literal import numpy as np from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator +# Bound int request fields to avoid overflow issues. +_INT64_MIN = -(2**63) +_INT64_MAX = 2**63 - 1 + _MAX_EMBEDDING_DIM = 8192 SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"}) @@ -59,6 +65,16 @@ class OpenAICreateSpeechRequest(BaseModel): validation_alias=AliasChoices("voice", "speaker"), description="Speaker/voice to use. For Qwen3-TTS: vivian, ryan, aiden, etc.", ) + + @field_validator("voice", mode="before") + @classmethod + def _normalize_voice(cls, voice: str | dict[str, str]) -> str: + if isinstance(voice, dict): + if "id" not in voice: + raise ValueError("voice dict must contain 'id'") + return voice["id"] + return voice + instructions: str | None = Field( default=None, description="Instructions for voice style/emotion (maps to 'instruct' for Qwen3-TTS)", @@ -132,18 +148,21 @@ class OpenAICreateSpeechRequest(BaseModel): ) max_new_tokens: int | None = Field( default=None, + ge=1, + le=_INT64_MAX, description="Maximum tokens to generate", ) seed: int | None = Field( default=None, - ge=0, - le=2**63 - 1, + ge=_INT64_MIN, + le=_INT64_MAX, description="Random seed for reproducible generation. When set, ensures " "deterministic output for the same input text and seed value.", ) initial_codec_chunk_frames: int | None = Field( default=None, ge=0, + le=_INT64_MAX, description="Per-request initial chunk size override. If null, computed dynamically based on server load.", ) non_streaming_mode: bool | None = Field( @@ -355,10 +374,14 @@ class OpenAICreateAudioGenerateRequest(BaseModel): ) num_inference_steps: int | None = Field( default=None, + ge=1, + le=_INT64_MAX, description="Number of inference steps", ) seed: int | None = Field( default=None, + ge=_INT64_MIN, + le=_INT64_MAX, description="Random seed for reproducibility", ) @@ -403,8 +426,8 @@ class SpeechBatchItem(BaseModel): ref_audio: str | None = None ref_text: str | None = None x_vector_only_mode: bool | None = None - max_new_tokens: int | None = None - initial_codec_chunk_frames: int | None = Field(default=None, ge=0) + max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX) + initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX) non_streaming_mode: bool | None = None @@ -423,8 +446,8 @@ class BatchSpeechRequest(BaseModel): ref_audio: str | None = None ref_text: str | None = None x_vector_only_mode: bool | None = None - max_new_tokens: int | None = None - initial_codec_chunk_frames: int | None = Field(default=None, ge=0) + max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX) + initial_codec_chunk_frames: int | None = Field(default=None, ge=0, le=_INT64_MAX) non_streaming_mode: bool | None = None @@ -507,11 +530,12 @@ class StreamingSpeechSessionConfig(BaseModel): instructions: str | None = None response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT speed: float | None = Field(default=1.0, ge=0.25, le=4.0) - max_new_tokens: int | None = Field(default=None, ge=1) + max_new_tokens: int | None = Field(default=None, ge=1, le=_INT64_MAX) initial_codec_chunk_frames: int | None = Field( default=None, ge=0, - description="Initial chunk size for reduced TTFA. Overrides stage config for this session.", + le=_INT64_MAX, + description="Initial chunk size for reduced TTFA. Overrides the deploy config for this session.", ) non_streaming_mode: bool | None = Field( default=None, diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 0a20ba1c8..2f28b6199 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ OpenAI-compatible protocol definitions for image generation. @@ -21,6 +21,10 @@ from pydantic import BaseModel, Field, field_validator from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers +# Bound int request fields to avoid overflow issues. +_INT64_MIN = -(2**63) +_INT64_MAX = 2**63 - 1 + _IMAGE_FILE_METADATA = { "jpg": ("jpg", "image/jpeg"), "jpeg": ("jpeg", "image/jpeg"), @@ -144,7 +148,7 @@ class ImageGenerationRequest(BaseModel): default=None, description="Optional model-specific parameters passed directly to the model's extra_args.", ) - seed: int | None = Field(default=None, description="Random seed for reproducibility") + seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility") generator_device: str | None = Field( default=None, description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.", diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index 87959b123..f71e9feed 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/vllm_omni/entrypoints/openai/protocol/videos.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ OpenAI-compatible protocol definitions for video generation. @@ -20,6 +20,10 @@ from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_vali from vllm_omni.entrypoints.openai.image_api_utils import parse_size from vllm_omni.inputs.data import DIFFUSION_QUALITY_LEVELS +# Bound int request fields to avoid overflow issues. +_INT64_MIN = -(2**63) +_INT64_MAX = 2**63 - 1 + class VideoGenerationStatus(str, Enum): QUEUED = "queued" @@ -48,10 +52,10 @@ def file_extension(media_type: str): class VideoParams(BaseModel): """Optional block for video-specific parameters.""" - width: int | None = Field(default=None, ge=1, description="Video width in pixels") - height: int | None = Field(default=None, ge=1, description="Video height in pixels") - num_frames: int | None = Field(default=None, ge=1, description="Number of frames") - fps: int | None = Field(default=None, ge=1, description="Frames per second for output video") + width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels") + height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels") + num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames") + fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video") @property def size(self) -> str | None: @@ -140,10 +144,10 @@ class VideoGenerationRequest(BaseModel): user: str | None = Field(default=None, description="User identifier for tracking") # Video-specific fields (top-level for OpenAI-style compatibility) - width: int | None = Field(default=None, ge=1, description="Video width in pixels") - height: int | None = Field(default=None, ge=1, description="Video height in pixels") - fps: int | None = Field(default=None, ge=1, description="Frames per second for output video") - num_frames: int | None = Field(default=None, ge=1, description="Number of frames to generate") + width: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video width in pixels") + height: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Video height in pixels") + fps: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Frames per second for output video") + num_frames: int | None = Field(default=None, ge=1, le=_INT64_MAX, description="Number of frames to generate") aspect_ratio: str | None = Field( default=None, description=( @@ -151,7 +155,12 @@ class VideoGenerationRequest(BaseModel): "FL2VA follows the input image; Ref2VA defaults to 16:9." ), ) - short_edge: int | None = Field(default=None, ge=1, description="MiniMax H3 output short edge in pixels") + short_edge: int | None = Field( + default=None, + ge=1, + le=_INT64_MAX, + description="MiniMax H3 output short edge in pixels", + ) num_outputs_per_prompt: int = Field( default=1, ge=1, @@ -218,7 +227,7 @@ class VideoGenerationRequest(BaseModel): le=20.0, description="True CFG scale (model-specific parameter, may be ignored if not supported)", ) - seed: int | None = Field(default=None, description="Random seed for reproducibility") + seed: int | None = Field(default=None, ge=_INT64_MIN, le=_INT64_MAX, description="Random seed for reproducibility") generate_sound: bool = Field( default=False, description="Request model-generated audio for video models that support sound generation.", @@ -237,6 +246,7 @@ class VideoGenerationRequest(BaseModel): frame_interpolation_exp: int = Field( default=1, ge=1, + le=_INT64_MAX, description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.", ) frame_interpolation_scale: float = Field( diff --git a/vllm_omni/entrypoints/openai/serving_audio_generate.py b/vllm_omni/entrypoints/openai/serving_audio_generate.py index 594b8c137..1975f1f16 100644 --- a/vllm_omni/entrypoints/openai/serving_audio_generate.py +++ b/vllm_omni/entrypoints/openai/serving_audio_generate.py @@ -95,11 +95,14 @@ class OmniOpenAIServingAudioGenerate(OpenAIServing, AudioMixin): "audio_end_in_s": audio_end_in_s, } - logger.info( - "Audio generation request %s: prompt=%r", - request_id, - request.input[:50] + "..." if len(request.input) > 50 else request.input, - ) + logger.info("Audio generation request %s", request_id) + _rl = getattr(self, "request_logger", None) + if _rl: + base_len = len(f"Audio generation request {request_id}: prompt=") + raw_max = getattr(_rl, "max_log_len", None) + cap = raw_max if isinstance(raw_max, int) else 200 + text = request.input[: max(cap - base_len, 0)] + logger.debug("Audio generation request %s: prompt=%r", request_id, text) generator = self.engine_client.generate( prompt=prompt, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 423ca6d82..e799fce81 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 import json @@ -19,9 +22,11 @@ from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateContentFormatOption, ConversationMessage, - get_history_tool_calls_cnt, make_tool_call_id, ) +from vllm.parser.utils import ( + count_chat_history_tool_calls as get_history_tool_calls_cnt, +) from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -37,6 +42,7 @@ from vllm_omni.metrics.modality import ( observe_audio_first_packet, observe_audio_streaming_finalize, ) +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY from vllm_omni.model_extras import ( get_extra_body_params, get_extra_output_params, @@ -58,7 +64,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionResponseStreamChoice, ChatMessage, ) -from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.openai.chat_completion.serving import ( + OpenAIServingChat, + _get_mm_token_counts, + _make_prompt_tokens_details, +) from vllm.entrypoints.openai.engine.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -67,7 +77,6 @@ from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, FunctionCall, FunctionDefinition, - PromptTokenUsageInfo, RequestResponseMetadata, ToolCall, UsageInfo, @@ -83,14 +92,7 @@ from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.multimodal.media.connector import MediaConnector from vllm.outputs import RequestOutput - -try: - from vllm.parser.mistral import MistralToolCall -except ImportError: - # Accelerator images can lag the core vLLM release by one version. - from vllm.tool_parsers.mistral_tool_parser import MistralToolCall - - +from vllm.parser.mistral import MistralToolCall from vllm.reasoning import ReasoningParser from vllm.renderers import BaseRenderer, merge_kwargs from vllm.renderers.inputs import TokPrompt @@ -248,6 +250,44 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return True return False + async def _attach_minicpmo45_reference_audio( + self, + engine_prompt: TokPrompt, + request: ChatLikeRequest | ResponsesRequest, + ) -> None: + if not self._has_minicpmo45_stage(): + return + + reference_audio_source = getattr(request, "ref_audio", None) + extra_body = getattr(request, "extra_body", None) + model_extra = getattr(request, "model_extra", None) + if reference_audio_source is None and isinstance(extra_body, dict): + reference_audio_source = extra_body.get("ref_audio") + if reference_audio_source is None and isinstance(model_extra, dict): + nested_extra_body = model_extra.get("extra_body") + if isinstance(nested_extra_body, dict): + reference_audio_source = nested_extra_body.get("ref_audio") + if reference_audio_source is None: + reference_audio_source = model_extra.get("ref_audio") + if reference_audio_source is None: + return + if not isinstance(reference_audio_source, str): + raise ValueError("MiniCPM-o 4.5 chat ref_audio must be a URI string") + + model_config = self.model_config + media_connector = MediaConnector( + media_io_kwargs=getattr(request, "media_io_kwargs", None), + allowed_local_media_path=getattr(model_config, "allowed_local_media_path", "") or "", + allowed_media_domains=getattr(model_config, "allowed_media_domains", None), + ) + reference_waveform, reference_sample_rate = await media_connector.fetch_audio_async(reference_audio_source) + # The renderer keeps processed Stage-0 features, not the source waveform. + # Keep the decoded audio only on the original prompt for the downstream bridge. + cast(dict[str, object], engine_prompt)[MINICPMO45_REFERENCE_AUDIO_KEY] = ( + reference_waveform, + reference_sample_rate, + ) + def _fix_minicpmo45_audio_stream_output_kinds( self, sampling_params_list: list[Any], @@ -730,8 +770,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Schedule the request and get the result generator. generators: list[AsyncGenerator[RequestOutput, None]] = [] + mm_token_counts: dict[str, int] | None = None try: for i, engine_prompt in enumerate(engine_prompts): + if self.enable_prompt_tokens_details: + mm_token_counts = _get_mm_token_counts(engine_prompt) if hasattr(request, "sampling_params_list"): sampling_params_list = self._to_sampling_params_list(request.sampling_params_list) else: @@ -814,6 +857,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer, request_metadata, reasoning_parser, + mm_token_counts=mm_token_counts, raw_request=raw_request, ) @@ -827,6 +871,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer, request_metadata, reasoning_parser, + mm_token_counts=mm_token_counts, ) except ValueError as e: return self.create_error_response(e) @@ -935,6 +980,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) prompt_additional_information["deferred_multi_modal_data"] = deferred_multi_modal_data + await self._attach_minicpmo45_reference_audio(engine_prompt, request) + speaker = getattr(request, "voice", None) or getattr(request, "speaker", None) normalized = validate_requested_speaker(speaker, self._get_supported_speakers()) if normalized is not None: @@ -1272,7 +1319,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): max_tokens is NOT computed dynamically — it uses the deploy YAML default. Args: - default_params: Default SamplingParams from stage config YAML. + default_params: Default SamplingParams from the resolved pipeline and deploy configuration. request: The chat completion request containing user-provided values. Returns: @@ -1387,6 +1434,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer: AnyTokenizer, request_metadata: RequestResponseMetadata, reasoning_parser: ReasoningParser | None = None, + mm_token_counts: dict[str, int] | None = None, raw_request: Request | None = None, ): created_time = int(time.time()) @@ -1407,6 +1455,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): stop_reason_emitted: list[bool] = [False] * num_choices num_prompt_tokens = 0 num_cached_tokens = None + num_cache_creation_tokens = None if self.use_harmony: harmony_parsers = [get_streamable_parser_for_assistant() for _ in range(num_choices)] harmony_tools_streamed = [False] * num_choices @@ -1503,6 +1552,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): num_prompt_tokens += len(res.encoder_prompt_token_ids) num_cached_tokens = res.num_cached_tokens + num_cache_creation_tokens = res.num_cache_creation_tokens # Send first response for each choice with role # NOTE: num_choices defaults to 1 so this usually executes once per request for i in range(num_choices): @@ -2234,8 +2284,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): completion_tokens=completion_tokens, total_tokens=num_prompt_tokens + completion_tokens, ) - if self.enable_prompt_tokens_details and num_cached_tokens: - final_usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=num_cached_tokens) + final_usage.prompt_tokens_details = _make_prompt_tokens_details( + self.enable_prompt_tokens_details, + num_cached_tokens, + num_cache_creation_tokens, + mm_token_counts, + ) final_usage_chunk = OmniChatCompletionStreamResponse( id=request_id, @@ -2307,6 +2361,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, reasoning_parser: ReasoningParser | None = None, + mm_token_counts: dict[str, int] | None = None, ) -> ErrorResponse | OmniChatCompletionResponse: created_time = int(time.time()) final_res: RequestOutput | None = None @@ -2361,6 +2416,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation, role, reasoning_parser, + mm_token_counts=mm_token_counts, ) final_res = omni_outputs else: @@ -2462,6 +2518,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation: list[ConversationMessage], role: str, reasoning_parser: ReasoningParser | None = None, + mm_token_counts: dict[str, int] | None = None, ): final_res = omni_outputs if self.tool_call_id_type == "kimi_k2": @@ -2698,8 +2755,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, ) - if self.enable_prompt_tokens_details and final_res.num_cached_tokens: - usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens) + usage.prompt_tokens_details = _make_prompt_tokens_details( + self.enable_prompt_tokens_details, + final_res.num_cached_tokens, + final_res.num_cache_creation_tokens, + mm_token_counts, + ) prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs) prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None @@ -3528,12 +3589,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return self._create_error_response(str(exc), status_code=400) logger.info( - "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s", + "Diffusion chat request %s: ref_images=%d, params=%s", request_id, - prompt[:50] + "..." if len(prompt) > 50 else prompt, len(reference_images), {key: value for key, value in extra_body.items() if value is not None}, ) + _rl = getattr(self, "request_logger", None) + if _rl: + base_len = len(f"Diffusion chat request {request_id}: prompt=") + raw_max = getattr(_rl, "max_log_len", None) + cap = raw_max if isinstance(raw_max, int) else 200 + text = prompt[: max(cap - base_len, 0)] + logger.debug("Diffusion chat request %s: prompt=%r", request_id, text) # Decode reference images if provided pil_images: list[Image.Image] = [] @@ -3621,7 +3688,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if reference_audios: gen_params.extra_args["audio_path"] = reference_audios[0] - # Generate image or audio (e.g. AudioX) via AsyncOmni + # Generate image or audio via AsyncOmni diffusion_engine = cast(AsyncOmni, self._diffusion_engine) stage_configs = list(getattr(diffusion_engine, "stage_configs", []) or []) default_sampling_params_list = get_default_sampling_params_list(diffusion_engine) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 109777282..9ed81a862 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import base64 import hashlib @@ -9,18 +12,18 @@ import re import struct import time from collections import OrderedDict -from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor from http import HTTPStatus from pathlib import Path from typing import Any +from urllib.parse import urlparse +from urllib.request import url2pathname import numpy as np import soundfile as sf import torch from fastapi import HTTPException, Request, UploadFile from fastapi.responses import Response, StreamingResponse -from transformers.utils.hub import cached_file from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing from vllm.entrypoints.launcher import terminate_if_errored from vllm.entrypoints.openai.engine.protocol import ( @@ -53,6 +56,7 @@ from vllm_omni.entrypoints.openai.speech_usage import ( ) from vllm_omni.entrypoints.openai.tts_adapters import ( SpeechServingContext, + TTSGenerationError, all_tts_stage_keys, detect_tts_model_type, resolve_adapter, @@ -70,16 +74,18 @@ from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( create_instruction as ming_create_instruction, ) -from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.utils.speaker_cache import ( - get_speaker_cache, - iter_custom_voice_profiles, - load_validated_profile_tensors, -) +from vllm_omni.utils.speaker_cache import get_speaker_cache logger = init_logger(__name__) + +_SPEECH_USAGE_INPUT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TOKENS" +_SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS" +_SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS" +_SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS" +_SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS" + # TTS Configuration # # The stage-key -> model-type mapping is NOT declared here: it is derived from @@ -92,40 +98,15 @@ logger = init_logger(__name__) # reference (issue #4644). Matches the offline example/test and upstream demo. _COSYVOICE3_PROMPT_DELIMITER = "<|endofprompt|>" _COSYVOICE3_PROMPT_PREFIX = f"You are a helpful assistant.{_COSYVOICE3_PROMPT_DELIMITER}" -_SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = { - "fish_tts", - "qwen3_tts", - "voxtral_tts", - "cosyvoice3", - "voxcpm2", - "higgs_audio_v2", - "higgs_audio_v3", - "indextts2", - "indextts2_5", - "audex", - "audex_tta", -} + # Audex contract: zero-codec / invalid generations arrive as empty terminal # payloads and must fail the request, never serialize as a successful empty # WAV. Covers both the TTS ("audex") and TTA ("audex_tta") pipelines. _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES = frozenset({"audex", "audex_tta"}) -_TTS_LANGUAGES = frozenset( - { - "Auto", - "Chinese", - "English", - "Japanese", - "Korean", - "German", - "French", - "Russian", - "Portuguese", - "Spanish", - "Italian", - } -) _REF_AUDIO_MIN_DURATION = 1.0 # seconds _REF_AUDIO_MAX_DURATION = 30.0 # seconds +_REF_AUDIO_METADATA_FETCH_ATTEMPTS = 3 +_REMOTE_REF_AUDIO_SCHEMES = frozenset({"http", "https", "data"}) _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256 _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024 _HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES = 256 @@ -134,6 +115,13 @@ _QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" _TTS_MAX_INSTRUCTIONS_LENGTH = 500 _TTS_MAX_NEW_TOKENS_MAX = 4096 _MING_DEFAULT_PROMPT = MING_DEFAULT_PROMPT +_DEFAULT_VOICE_NAME = "default" + + +def _is_default_voice(voice, supported_speakers): + """Check if a lowercased voice name is the placeholder default and not + an actual registered/built-in speaker.""" + return voice == _DEFAULT_VOICE_NAME and voice not in supported_speakers def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16) -> bytes: @@ -249,6 +237,13 @@ def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str: must also be folded in. ``voice_created_at`` bumps on every (re-)upload, which uniquely identifies the resolved reference artifact together with the voice name; the decoded ref_audio array itself need not be hashed. + + Local ``ref_audio`` locators are the same class of gap: the raw request + hashes the path/URI string, so an on-disk rewrite is invisible to the + salt. Callers stash the content-aware key from ``_resolve_ref_audio`` as + ``ref_audio_cache_key`` (and ``ref_audio_2_cache_key`` / emotion keys + when those clips exist) so a file edit cannot reuse KV from the previous + reference. """ h = hashlib.sha256() for part in ( @@ -266,13 +261,15 @@ def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str: if part is not None: h.update(repr(part).encode("utf-8")) # Fold resolved conditioning that is auto-derived for uploaded voices and - # absent from the raw request. + # absent from the raw request (or that a locator string cannot see). for key in ( "voice_created_at", "task_type", "speaker", "ref_text", "x_vector_only_mode", + "ref_audio_cache_key", + "ref_audio_2_cache_key", ): h.update(b"\x00") value = tts_params.get(key) if tts_params is not None else None @@ -283,6 +280,8 @@ def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str: class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _diffusion_mode: bool = False + _media_connector: MediaConnector | None = None + _allowed_local_media_path: str = "" _tts_executor: ThreadPoolExecutor | None = None def _init_speaker_storage(self) -> None: @@ -296,9 +295,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): except ValueError: logger.warning("Invalid SPEAKER_MAX_UPLOADED=%r; using default 1000", _raw_cap) self._max_uploaded_speakers = 1000 - self.uploaded_speakers: dict[str, dict] = {} - self.precomputed_speakers: dict[str, dict[str, Any]] = {} - self.supported_speakers: set[str] = set() + self.uploaded_speakers: dict[str, dict[str, Any]] = {} self._ref_audio_data_url_cache: dict[str, str] = {} self._ref_audio_resolve_cache: OrderedDict[str, tuple[list[float], int, int, str]] = OrderedDict() self._ref_audio_resolve_cache_bytes = 0 @@ -390,7 +387,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): speaker_data.setdefault("name", voice_name_lower) speaker_data.setdefault("file_size", int(path.stat().st_size)) self.uploaded_speakers[voice_name_lower] = speaker_data - self.supported_speakers.add(voice_name_lower) self._last_upload_ts = max(self._last_upload_ts, int(speaker_data.get("created_at", 0))) restored += 1 if restored: @@ -402,6 +398,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): diffusion_engine: "Any", model_name: str, stage_configs: "list[Any] | None" = None, + allowed_local_media_path: str = "", + allowed_media_domains: list[str] | None = None, ) -> "OmniOpenAIServingSpeech": """Create a speech serving instance for pure diffusion TTS models. @@ -413,15 +411,23 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): instance._diffusion_engine = diffusion_engine instance._diffusion_model_name = model_name instance._diffusion_stage_configs = stage_configs + instance._allowed_local_media_path = allowed_local_media_path + instance._media_connector = MediaConnector( + allowed_local_media_path=allowed_local_media_path, + allowed_media_domains=allowed_media_domains, + ) instance._tts_model_type = "omnivoice" instance._is_tts = False # Diffusion-only instances don't have a TTS stage; set None so any # ``_is_tts_model()`` / ``_tts_stage`` access doesn't raise AttributeError. instance._tts_stage = None + instance._adapter = None instance._init_speaker_storage() return instance def __init__(self, *args, **kwargs): + self._media_connector = None + self._allowed_local_media_path = "" self.model_name = kwargs.pop("model_name", None) # True when the server was launched with --forced-aligner (a pooling # aligner stage is appended to the pipeline). Gates word_timestamps. @@ -437,12 +443,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Cached per process: the CosyVoice3 Qwen tokenizer + resolved model # path used for dynamic-token sizing. Without this, every request # re-ran snapshot_download + reloaded the tokenizer (~100 ms on the - # TTFP critical path) in _apply_cosyvoice3_dynamic_tokens. + # TTFP critical path) in the CosyVoice3 adapter's sampling override. self._cosyvoice3_tokenizer = None # Determine TTS model type or None self._tts_model_type = self._detect_tts_model_type() - self.precomputed_speakers = self._load_precomputed_speakers() + # Resolve the per-model serving adapter (RFC #4327), keyed on the + # detected model-type. Every dedicated TTS model has an adapter; the + # adapter owns request validation, prompt/param building, capability + # metadata, and sampling overrides. The model-type label remains in the + # orchestrator for compatibility during this incremental migration. + self._adapter = None + if self._tts_stage is not None: + adapter_cls = resolve_adapter(self._tts_model_type) + if adapter_cls is not None: + ctx = SpeechServingContext(server=self, engine_client=self.engine_client) + self._adapter = adapter_cls(ctx) + logger.info("Resolved TTS serving adapter: %s", adapter_cls.__name__) + + adapter = self._adapter + if adapter is not None: + adapter.load_capabilities() + available_speakers = self._get_available_speakers() + logger.info("Loaded %d supported speakers: %s", len(available_speakers), sorted(available_speakers)) # Sub-variant inside the full MOSS-TTS family. We collapse all five # variants onto the same _tts_model_type="moss_tts" because they share @@ -458,26 +481,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Cache TTS configuration values (computed once, reused per request) self._max_instructions_length = self._compute_max_instructions_length() - # Merge built-in speakers into the set initialized by _init_speaker_storage. - self.supported_speakers |= self._load_supported_speakers() - self.supported_speakers |= set(self.precomputed_speakers) - - self.supported_languages = self._load_supported_languages() - self._tts_tokenizer = None self._voxcpm2_tokenizer = None self._audex_tokenizer = None self._audex_tta_rvq = None self._voxcpm2_split_map: dict[int, list[int]] = {} - logger.info("Loaded %d supported speakers: %s", len(self.supported_speakers), sorted(self.supported_speakers)) - # Batch configuration self._batch_max_items: int = getattr(self.engine_client, "tts_batch_max_items", 32) - # Load speech tokenizer codec parameters for prompt length estimation - self._codec_frame_rate: float | None = self._load_codec_frame_rate() - # Shared thread pool executor for blocking TTS preprocessing # operations. max_workers=1 serializes tokenizer access to avoid # Rust RefCell "Already borrowed" errors from concurrent use. @@ -486,19 +498,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._build_fish_speech_prompt_async = make_async(self._build_fish_speech_prompt, executor=self._tts_executor) self._estimate_prompt_len_async = make_async(self._estimate_prompt_len, executor=self._tts_executor) - # Resolve the per-model serving adapter (RFC #4327), keyed on the - # detected model-type. Every dedicated TTS model has an adapter; the - # adapter owns request validation + prompt/param building. Sampling - # overrides and the model-type label remain in the orchestrator tail - # (keyed on ``_tts_model_type``) during this incremental migration. - self._adapter = None - if self._tts_stage is not None: - adapter_cls = resolve_adapter(self._tts_model_type) - if adapter_cls is not None: - ctx = SpeechServingContext(server=self, engine_client=self.engine_client) - self._adapter = adapter_cls(ctx) - logger.info("Resolved TTS serving adapter: %s", adapter_cls.__name__) - def _get_tts_adapter(self): """Return the per-model serving adapter for the current ``_tts_model_type``. @@ -520,145 +519,44 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self._adapter is None or type(self._adapter) is not adapter_cls: ctx = SpeechServingContext(server=self, engine_client=self.engine_client) self._adapter = adapter_cls(ctx) + self._adapter.load_capabilities() return self._adapter def _uses_native_speed_control(self) -> bool: adapter = self._get_tts_adapter() return bool(adapter is not None and adapter.native_speed_control) + def _get_available_speakers(self) -> set[str]: + """Return all built-in, precomputed, and runtime-uploaded speakers.""" + available_speakers = set(self.uploaded_speakers) + if self._adapter is not None: + available_speakers.update(self._adapter.capabilities.supported_speakers) + available_speakers.update(self._adapter.capabilities.precomputed_speakers) + return available_speakers + def _audio_encode_speed(self, request: OpenAICreateSpeechRequest) -> float: if self._uses_native_speed_control(): return 1.0 return float(request.speed or 1.0) async def warmup(self) -> None: - """Run a synthetic speech request to trigger all first-request warmup. + """Run model-specific startup warmup through the resolved adapter. - Unlike qwen3-tts, whose CUDA Graph warmup targets a standalone tokenizer - decoder (no vLLM dependencies) and can complete entirely at model-init - time, VoxCPM2 needs to warm up PagedAttention scaffold/residual LLMs. - Their CUDA Graph capture requires a vLLM ``ForwardContext`` - (attn_metadata, slot_mapping, etc.) that only exists during real - inference steps. The same request also pays the one-time torch.compile - JIT tax for the LocDiT estimator, feat_encoder, AudioVAE decoder, and - projection helpers. - - For VoxCPM2 this shifts ~15s of torch.compile + CUDA Graph capture from - the first user request to server startup. + Warmup requirements differ by model. For example, Qwen3-TTS can warm up + its standalone tokenizer decoder during model initialization, while + VoxCPM2 requires a real inference request with an active vLLM + ``ForwardContext``. The adapter owns that model-specific lifecycle logic. """ - if self._tts_model_type != "voxcpm2": + if self._adapter is None: return - - t0 = time.time() - logger.info("Running warmup speech request for model_type=%s", self._tts_model_type) - # VoxCPM2 has no predefined speaker presets — "default" means zero-shot - # mode (no voice cloning). The voice field is required by the OpenAI - # API schema but semantically ignored by the model. - warmup_req = OpenAICreateSpeechRequest( - input="Warmup.", - voice="default", - response_format="wav", - speed=1.0, - stream=False, - model=self.model_name, - ) - try: - _audio_bytes, _media_type = await self._generate_audio_bytes(warmup_req, request_id="speech-warmup") - except Exception as exc: - logger.warning("Speech warmup failed (non-fatal): %s", exc) - return - - elapsed = time.time() - t0 - logger.info("Speech warmup complete in %.1fs", elapsed) - - def _get_qwen_tts_expected_speaker_embedding_dim(self) -> int | None: - """Return the loaded Qwen3-TTS speaker embedding dim, if known. - - The user-provided speaker embedding is concatenated directly with - talker codec embeddings, so the real compatibility requirement is the - talker hidden size. - """ - if self._tts_model_type != "qwen3_tts": - return None - hf_config = self.engine_client.model_config.hf_config - talker_config = hf_config.talker_config - return int(talker_config.hidden_size) - - def _validate_qwen_tts_speaker_embedding_dim(self, emb_dim: int) -> str | None: - expected_dim = self._get_qwen_tts_expected_speaker_embedding_dim() - if expected_dim is None: - return None - if emb_dim != expected_dim: - return f"speaker_embedding has {emb_dim} dimensions; expected {expected_dim} for the loaded Qwen3-TTS model" - return None - - def _load_codec_frame_rate(self) -> float | None: - """Load codec frame rate from speech tokenizer config for prompt length estimation.""" - if self._tts_model_type == "ming_tts": - try: - from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingTTSConfig - - hf_config = self.engine_client.model_config.hf_config - ming_cfg = MingTTSConfig.from_hf_config(hf_config) - patch_size = int(ming_cfg.patch_size) - audio_frame_hop = int(ming_cfg.audio_frame_hop) - sample_rate = int(ming_cfg.sample_rate) - if patch_size <= 0 or audio_frame_hop <= 0 or sample_rate <= 0: - raise ValueError( - "Ming config has invalid tokenizer timing values: " - f"patch_size={patch_size}, audio_frame_hop={audio_frame_hop}, sample_rate={sample_rate}" - ) - rate = float(sample_rate) / float(audio_frame_hop * patch_size) - logger.info( - "Derived Ming codec frame rate: %.1f Hz (sample_rate=%s, audio_frame_hop=%s, patch_size=%s)", - rate, - sample_rate, - audio_frame_hop, - patch_size, - ) - return rate - except Exception as e: - logger.warning(f"Failed to derive Ming codec frame rate from hf_config: {e}") - - try: - model_path = self.engine_client.model_config.model - st_config_path = os.path.join(model_path, "speech_tokenizer", "config.json") - if not os.path.exists(st_config_path): - st_config_path = cached_file(model_path, "speech_tokenizer/config.json") - if st_config_path is not None and os.path.exists(st_config_path): - with open(st_config_path) as f: - st_config = json.load(f) - output_sr = st_config.get("output_sample_rate") - downsample = st_config.get("encode_downsample_rate") - if output_sr and downsample and downsample > 0: - rate = float(output_sr) / float(downsample) - logger.info( - "Loaded codec frame rate: %.1f Hz (output_sample_rate=%s, encode_downsample_rate=%s)", - rate, - output_sr, - downsample, - ) - return rate - except Exception as e: - logger.warning("Failed to load codec frame rate from speech tokenizer config: %s", e) - - # Fallback: try codec_frame_rate_hz from hf_config - try: - hf_config = self.engine_client.model_config.hf_config - rate = getattr(hf_config, "codec_frame_rate_hz", None) - if rate is not None: - logger.info("Using codec frame rate from hf_config: %s Hz", rate) - return float(rate) - except Exception: - pass - return None + await self._adapter.warmup() def shutdown(self) -> None: """Shut down the TTS thread pool executor.""" if self._tts_executor is not None: self._tts_executor.shutdown(wait=False, cancel_futures=True) self._tts_executor = None - for name in list(self.uploaded_speakers.keys()): + for name in list(self.uploaded_speakers): self._speaker_cache.clear(name) def _find_tts_stage(self): @@ -701,45 +599,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): getattr(self._tts_stage.engine_args, "model_arch", None), ) - def _get_custom_voice_dir(self) -> str | None: - try: - value = getattr(self.engine_client.model_config.hf_config, "custom_voice_dir", None) - except AttributeError: - return None - if isinstance(value, os.PathLike): - return os.fspath(value) - if isinstance(value, str) and value: - return value - return None - - def _load_precomputed_speakers(self) -> dict[str, dict[str, Any]]: - """Load precomputed voice names from ``custom_voice_dir`` for API validation.""" - if self._tts_model_type not in ("qwen3_tts", "voxcpm2"): - return {} - custom_voice_dir = self._get_custom_voice_dir() - if not custom_voice_dir: - return {} - - profiles: dict[str, dict[str, Any]] = {} - qwen3_embedding_dim = self._get_qwen_tts_expected_speaker_embedding_dim() - for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type=self._tts_model_type): - tensors = load_validated_profile_tensors( - profile, - expected_model_type=self._tts_model_type, - qwen3_embedding_dim=qwen3_embedding_dim, - ) - if tensors is None: - continue - profiles[profile["voice_name_lower"]] = profile - if profiles: - logger.info( - "Loaded %d precomputed %s voice profile(s) from %s", - len(profiles), - self._tts_model_type, - custom_voice_dir, - ) - return profiles - def _compute_max_instructions_length(self) -> int: """Compute max instructions length with precedence: CLI > stage config > default. @@ -759,62 +618,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # 3. Default fallback return _TTS_MAX_INSTRUCTIONS_LENGTH - def _load_supported_speakers(self) -> set[str]: - """Load supported speakers (case-insensitive) from the model configuration.""" - if self._tts_model_type == "ming_flash_omni_tts": - # Ming-flash-omni drives speaker selection via the caption JSON - # (audio_sequence[0]["说话人"]) rather than a spk_id table, so there - # is no static speaker list to surface here. - return set() - try: - if self._tts_model_type == "glm_tts": - return set() - if self._tts_model_type == "ming_tts": - return set() - if self._tts_model_type == "voxcpm2": - return {"default"} - if self._tts_model_type == "voxtral_tts": - config = self.engine_client.model_config.hf_config.audio_config - else: - # Default is qwen3_tts path - config = getattr(self.engine_client.model_config.hf_config, "talker_config", None) - if config is None: - return set() - - # Check for speakers in either spk_id or speaker_id - for attr_name in ["spk_id", "speaker_id"]: - if isinstance(config, dict): - speakers_dict = config.get(attr_name) - else: - speakers_dict = getattr(config, attr_name, None) - if speakers_dict and isinstance(speakers_dict, dict): - return {speaker.lower() for speaker in speakers_dict.keys()} - - logger.warning("No speakers found in config (checked spk_id and speaker_id)") - except Exception as e: - logger.warning("Could not load speakers from model config: %s", e) - - return set() - - def _load_supported_languages(self) -> frozenset[str]: - """Load supported languages (title-cased) from the model configuration""" - if self._tts_model_type != "qwen3_tts": - return _TTS_LANGUAGES - try: - config = self.engine_client.model_config.hf_config.talker_config - - if isinstance(config, dict): - codec_language_id = config.get("codec_language_id") - else: - codec_language_id = getattr(config, "codec_language_id", None) - - if codec_language_id and isinstance(codec_language_id, Mapping): - return frozenset(str(language).title() for language in codec_language_id) | {"Auto"} - - logger.warning("No codec_language_id found in talker_config; falling back to default languages") - except Exception as e: - logger.warning("Could not load languages from model config: %s", e) - return _TTS_LANGUAGES + def _get_available_voices(self) -> set[str]: + """Get all voice names accepted by the API, including the placeholder default.""" + return self._get_available_speakers() | {_DEFAULT_VOICE_NAME} def _estimate_ref_code_len(self, ref_audio: object) -> int | None: """Estimate ref_code length from ref_audio waveform without running the codec. @@ -822,7 +628,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): The codec produces one frame per (output_sample_rate / encode_downsample_rate) audio samples, so ref_code_len = ceil(duration_seconds * codec_frame_rate). """ - if self._codec_frame_rate is None: + codec_frame_rate = self._adapter.capabilities.codec_frame_rate + if codec_frame_rate is None: return None try: # ref_audio comes from tts_params as [[wav_array, sr]] or similar nested structure @@ -847,7 +654,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if sr <= 0 or n_samples <= 0: return None duration = n_samples / sr - return math.ceil(duration * self._codec_frame_rate) + return math.ceil(duration * codec_frame_rate) except Exception: return None @@ -945,6 +752,38 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): details = self._compute_speech_input_details(request, tts_params) return build_speech_usage(details, output_tokens) + def _validate_tts_generation( + self, + tts_params: dict[str, Any], + usage_acc: SpeechOutputTokenCounter, + ) -> None: + adapter = self._get_tts_adapter() + if adapter is None: + return + adapter.validate_generation( + tts_params, + stage0_finish_reason=usage_acc.stage0_finish_reason, + output_tokens=usage_acc.total(), + ) + + @staticmethod + def _build_speech_usage_headers(usage: SpeechTokenUsage | None) -> dict[str, str]: + """Map speech usage into non-streaming response headers. + + Returns an empty dict when usage is unavailable, allowing callers to + merge these headers with other optional response headers. + """ + if usage is None: + return {} + + return { + _SPEECH_USAGE_INPUT_TOKENS_HEADER: str(usage.input_tokens), + _SPEECH_USAGE_OUTPUT_TOKENS_HEADER: str(usage.output_tokens), + _SPEECH_USAGE_TOTAL_TOKENS_HEADER: str(usage.total_tokens), + _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER: str(usage.input_token_details.text_tokens), + _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER: str(usage.input_token_details.audio_tokens), + } + def _estimate_fish_ref_code_len(self, ref_audio: object) -> int | None: """Estimate Fish Speech semantic token length from raw reference audio.""" from vllm_omni.model_executor.models.fish_speech.dac_utils import ( @@ -1007,12 +846,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ref_sr = None voice_profile = None if request.ref_audio is not None: - ref_audio, ref_sr = await self._resolve_ref_audio(request.ref_audio) + ref_audio, ref_sr, _ = await self._resolve_ref_audio(request.ref_audio) elif uploaded_ref is not None: wav_np, ref_sr = uploaded_ref ref_audio = wav_np.tolist() elif request.voice is not None: - voice_profile = self.precomputed_speakers.get(request.voice.lower()) + voice_profile = self._adapter.capabilities.precomputed_speakers.get(request.voice.lower()) return build_voxcpm2_prompt( hf_config=self.engine_client.model_config.hf_config, tokenizer=self._voxcpm2_tokenizer, @@ -1175,7 +1014,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if voice_name_lower not in self.uploaded_speakers: return old = self.uploaded_speakers.pop(voice_name_lower) - self.supported_speakers.discard(voice_name_lower) self._ref_audio_data_url_cache.pop(voice_name_lower, None) old_path = old.get("file_path") if old_path: @@ -1317,7 +1155,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raise ValueError(f"Failed to save voice file: {e}") self.uploaded_speakers[voice_name_lower] = speaker_data - self.supported_speakers.add(voice_name_lower) logger.info("Uploaded new voice '%s' with consent ID '%s'", name, consent) @@ -1368,13 +1205,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)") emb_dim = len(embedding) - if self._tts_model_type == "ming_tts": - if emb_dim != SPEAKER_EMBEDDING_DIM: - raise ValueError(f"Ming speaker embedding must have {SPEAKER_EMBEDDING_DIM} dims, got {emb_dim}") - else: - dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim) - if dim_err is not None: - raise ValueError(dim_err) + dim_err = self._adapter.validate_tts_embedding_dim(emb_dim) if self._adapter is not None else None + if dim_err is not None: + raise ValueError(dim_err) async with self._upload_lock: voice_name_lower = name.lower() @@ -1414,7 +1247,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): speaker_data["file_size"] = file_path.stat().st_size self.uploaded_speakers[voice_name_lower] = speaker_data - self.supported_speakers.add(voice_name_lower) logger.info("Uploaded voice '%s' from speaker embedding (%d-dim)", name, emb_dim) @@ -1444,7 +1276,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return False speaker_info = self.uploaded_speakers.pop(voice_name_lower) - self.supported_speakers.discard(voice_name_lower) self._ref_audio_data_url_cache.pop(voice_name_lower, None) file_path = speaker_info.get("file_path") @@ -1502,11 +1333,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """Validate ref_audio is a supported URI format. Returns error or None.""" if not isinstance(ref_audio, str): return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" - if not ( - ref_audio.startswith(("http://", "https://")) - or ref_audio.startswith("data:") - or ref_audio.startswith("file://") - ): + scheme = (urlparse(ref_audio).scheme or "").lower() + if scheme not in {"http", "https", "data", "file"}: return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" return None @@ -1555,7 +1383,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._moss_processor_cache = proc return proc - async def _build_moss_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + async def _build_moss_tts_params( + self, + request: OpenAICreateSpeechRequest, + *, + has_inline_ref_audio: bool = False, + ) -> dict[str, Any]: """Build the talker prompt + ``additional_information`` payload for any MOSS-TTS-family request (nano + 5 full variants). @@ -1583,8 +1416,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): } if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) params["prompt_audio_array"] = [[wav_list, sr]] + params["ref_audio_cache_key"] = cache_key return params # ---- MOSS-TTS-Realtime: keep the old prompt_audio_array path ---- @@ -1602,8 +1436,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): } if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) params["prompt_audio_array"] = [[wav_list, sr]] + params["ref_audio_cache_key"] = cache_key return params # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator) @@ -1634,30 +1469,47 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # only pulls it on the delay-family path (alongside the upstream proc). from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes - _voice = getattr(request, "voice", None) - _voice = _voice.strip() if isinstance(_voice, str) else "" - _voice_created = self._voice_created_at(_voice.lower()) if _voice else 0 + # Named-voice speaker cache is only valid for uploaded speakers + # without an inline ref_audio. ``request.voice`` plus a file/URL + # otherwise keys on (name, created_at=0) and skips the content-aware + # resolve — the adapter already applies this guard when tagging + # ``voice_name`` on tts_params. + _raw_voice = getattr(request, "voice", None) + _raw_voice = _raw_voice.strip() if isinstance(_raw_voice, str) else "" + _voice_lower = _raw_voice.lower() if _raw_voice else "" + _use_named_voice = bool(_voice_lower) and _voice_lower in self.uploaded_speakers and not has_inline_ref_audio + _voice = _voice_lower if _use_named_voice else "" + _voice_created = self._voice_created_at(_voice) if _voice else 0 + _resolve_keys: list[str] = [] - async def _encode_ref(ref_str: str) -> torch.Tensor: + async def _tracked_resolve(ref_str: str) -> tuple[list, int, str]: + wav_list, sr, cache_key = await self._resolve_ref_audio(ref_str) + _resolve_keys.append(cache_key) + return wav_list, sr, cache_key + + async def _encode_ref(ref_str: str, *, named_voice: bool) -> torch.Tensor: + # Named-voice cache is (voice_name, created_at). TTSD speaker 2 is a + # different clip, so it must stay anonymous or it reuses speaker 1. + use_named = named_voice and bool(_voice) return await encode_reference_codes( ref_str, processor=proc, - resolve_ref_audio=self._resolve_ref_audio, + resolve_ref_audio=_tracked_resolve, speaker_cache=self._speaker_cache, variant=v, n_vq=n_vq, sr_target=sr_target, - voice_name=_voice or None, - voice_created_at=_voice_created, + voice_name=_voice if use_named else None, + voice_created_at=_voice_created if use_named else 0, ) user_kwargs: dict[str, Any] = {"text": request.input or ""} if v in ("tts", "local"): - user_kwargs["reference"] = [await _encode_ref(request.ref_audio)] + user_kwargs["reference"] = [await _encode_ref(request.ref_audio, named_voice=True)] elif v == "ttsd": - refs = [await _encode_ref(request.ref_audio)] + refs = [await _encode_ref(request.ref_audio, named_voice=True)] if request.ref_audio_2: - refs.append(await _encode_ref(request.ref_audio_2)) + refs.append(await _encode_ref(request.ref_audio_2, named_voice=False)) user_kwargs["reference"] = refs elif v == "sound_effect": user_kwargs["text"] = request.input or "" # may be empty @@ -1691,6 +1543,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): } if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] + if _resolve_keys: + params["ref_audio_cache_key"] = _resolve_keys[0] + if len(_resolve_keys) > 1: + params["ref_audio_2_cache_key"] = _resolve_keys[1] return params async def _build_higgs_audio_v2_params(self, request: OpenAICreateSpeechRequest): @@ -1719,7 +1575,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt_token_ids = input_ids_to_python_list(inputs) return tokens_input(prompt_token_ids=prompt_token_ids) - wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + wav_list, sr, _ = await self._resolve_ref_audio(request.ref_audio) wav = np.asarray(wav_list, dtype=np.float32) out = await asyncio.to_thread( build_voice_clone_prompt, @@ -1788,8 +1644,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): encode_reference_audio, ) - wav_list, sr = await self._resolve_ref_audio(request.ref_audio) - artifact_key = self._get_resolved_ref_audio_artifact_key(request.ref_audio) + wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) + artifact_key = self._get_resolved_ref_audio_artifact_key(cache_key) wav = np.asarray(wav_list, dtype=np.float32) ref_codes_delayed, cache_hit, inflight_wait = await self._resolve_higgs_audio_v3_ref_codes( artifact_key, @@ -1811,7 +1667,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prompt["additional_information"] = { "audio_input_ids": ref_codes_delayed.to(torch.long), "audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool), + "ref_audio_cache_key": cache_key, } + # Placeholder prompt ids do not change when a same-path file is + # rewritten at the same duration; fold the content-aware resolve key + # into cache_salt so prefix caching cannot reuse the previous KV. + prompt["cache_salt"] = _conditioning_cache_salt(request, prompt["additional_information"]) return prompt async def _resolve_higgs_audio_v3_ref_codes( @@ -1902,43 +1763,100 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._higgs_audio_v3_adapter = adapter return adapter - async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int]: - """Resolve ref_audio to (wav_samples, sample_rate). + @staticmethod + def _local_ref_audio_stat_path(ref_audio_str: str) -> str | None: + """Filesystem path to stat for a local locator, else ``None``. - Delegates to upstream vLLM's MediaConnector which handles http(s) - URLs, ``data:`` base64 URIs, and ``file:`` local paths (the latter - gated by ``--allowed-local-media-path``). + Scheme comparison is case-insensitive (``urlparse`` lowercases it), so + ``FILE:///...`` is treated as a local file the same way ``file://`` is. + Remote ``http`` / ``https`` / ``data`` locators, including ``HTTP://``, + return ``None`` and stay string-keyed. """ - cache_key = hashlib.sha1(ref_audio_str.encode("utf-8")).hexdigest() - cached = self._ref_audio_resolve_cache.get(cache_key) - if cached is not None: - self._ref_audio_resolve_cache.move_to_end(cache_key) - wav_list, sr, _, _ = cached - logger.debug( - "Resolved ref_audio from cache: samples=%d sr=%d duration_s=%.3f", - len(wav_list), - sr, - len(wav_list) / sr if sr > 0 else 0.0, - ) - return wav_list, sr + parsed = urlparse(ref_audio_str) + scheme = (parsed.scheme or "").lower() + if scheme in _REMOTE_REF_AUDIO_SCHEMES: + return None + if scheme == "file": + netloc = parsed.netloc or "" + if netloc.lower() not in ("", "localhost"): + raise OSError(f"file:// URI with non-local authority {netloc!r} cannot be stat'd") + return url2pathname(parsed.path or "") + if scheme: + return None + return ref_audio_str - # In diffusion mode, model_config may not be available - if self._diffusion_mode: - connector = MediaConnector() - else: - model_config = self.model_config - connector = MediaConnector( - allowed_local_media_path=model_config.allowed_local_media_path, - allowed_media_domains=model_config.allowed_media_domains, - ) - fetch_start_s = time.perf_counter() - wav_np, sr = await connector.fetch_audio_async(ref_audio_str) - fetch_decode_ms = (time.perf_counter() - fetch_start_s) * 1000.0 + @staticmethod + def _get_ref_audio_cache_key( + ref_audio_str: str, + allowed_local_media_path: str | None = None, + ) -> str: + """Compute a cache key hash for *ref_audio_str*. + + For local files (bare paths and ``file://`` URIs) the key folds in + ``st_mtime_ns`` and ``st_size`` so that an on-disk edit automatically + invalidates the cached waveform without a server restart. Remote URLs + and ``data:`` URIs are keyed on the raw string alone — the server does + not re-fetch them to check for changes. + + ``os.stat`` runs only when *allowed_local_media_path* is a non-empty + path. vLLM's default is ``""`` (not ``None``); both mean local media + is refused, so a stat would be exposure for a request that cannot + load the file. + + Note: when the key changes the *previous* entry stays in + ``_ref_audio_resolve_cache`` until LRU eviction, so + ``_discard_ref_audio_artifact_ready_if_unreferenced`` will not fire for + the replaced file and its stale artifact key remains in + ``_ref_audio_model_artifact_ready``. This is functionally correct + (the stale entry is never *used*) but doubles memory until eviction. + """ + cache_key_source = ref_audio_str + try: + path = OmniOpenAIServingSpeech._local_ref_audio_stat_path(ref_audio_str) + except OSError as exc: + path = None + logger.debug("Skipping local-file cache metadata for %s: %s", ref_audio_str[:80], exc) + if path is not None: + try: + # Only stat paths the server operator has explicitly permitted + # via --allowed-local-media-path. The default is "" (and + # diffusion mode passes None); MediaConnector refuses local + # file loads in both cases, so a stat is exposure for zero + # benefit. ``if not`` covers "" and None; ``is None`` would + # treat "" as an allowlist of the process cwd. + if not allowed_local_media_path: + raise OSError("no allowed_local_media_path configured; skipping stat") + resolved = os.path.realpath(path) + allowed_base = os.path.realpath(allowed_local_media_path) + if os.path.commonpath([resolved, allowed_base]) != allowed_base: + raise OSError("path outside allowed_local_media_path; skipping stat") + st = os.stat(path) + cache_key_source = f"{ref_audio_str}:{st.st_mtime_ns}:{st.st_size}" + except (OSError, ValueError): + # Truncate the value to avoid dumping huge base64 blobs into + # the server log when a client omits the ``data:`` prefix. + display = ref_audio_str[:80] + if len(ref_audio_str) > 80: + display += f"... ({len(ref_audio_str)} chars)" + logger.debug( + "Failed to stat ref_audio path %s; falling back to string-only cache key (stale cache possible)", + display, + ) + return hashlib.sha1(cache_key_source.encode("utf-8")).hexdigest() + + async def _ref_audio_cache_key(self, ref_audio_str: str, allowed_local_media_path: str | None) -> str: + """Stat local files off the event loop; remote locators stay a cheap hash.""" + return await asyncio.to_thread( + self._get_ref_audio_cache_key, + ref_audio_str, + allowed_local_media_path, + ) + + def _finalize_fetched_ref_audio(self, wav_np: np.ndarray, sr: int) -> tuple[list[float], int, str, float]: wav_np = np.asarray(wav_np, dtype=np.float32) if wav_np.ndim > 1: wav_np = np.mean(wav_np, axis=-1) sr = int(sr) - artifact_key = self._make_ref_audio_artifact_cache_key(wav_np, sr) duration = len(wav_np) / sr if sr > 0 else 0.0 if duration < _REF_AUDIO_MIN_DURATION: raise ValueError( @@ -1950,19 +1868,89 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): f"Reference audio too long ({duration:.1f}s). " f"Maximum {_REF_AUDIO_MAX_DURATION:.0f}s supported — use a shorter clip." ) - tolist_start_s = time.perf_counter() - wav_list = wav_np.tolist() - tolist_ms = (time.perf_counter() - tolist_start_s) * 1000.0 - logger.debug( - "Resolved ref_audio: fetch_decode_ms=%.3f tolist_ms=%.3f samples=%d sr=%d duration_s=%.3f", - fetch_decode_ms, - tolist_ms, - len(wav_np), - sr, - duration, + artifact_key = self._make_ref_audio_artifact_cache_key(wav_np, sr) + return wav_np.tolist(), sr, artifact_key, duration + + async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int, str]: + """Resolve ref_audio to (wav_samples, sample_rate, cache_key). + + Delegates to upstream vLLM's MediaConnector which handles http(s) + URLs, ``data:`` base64 URIs, and ``file:`` local paths (the latter + gated by ``--allowed-local-media-path``). + + Local file references incorporate mtime and size into the cache key + so that modified files are automatically reloaded without a server + restart. Remote URLs remain cached by their original string locator. + + The returned *cache_key* should be passed to + ``_get_resolved_ref_audio_artifact_key`` when the caller needs the + artifact key, avoiding a redundant ``os.stat`` and the TOCTOU window. + After a cache miss the file is re-stat'd; a metadata change retries + the fetch so the stored waveform matches the key. + """ + # Pass the allowed-local-media-path so the stat is restricted to + # paths the server operator has explicitly permitted. + allowed_path: str | None + if self._diffusion_mode: + allowed_path = self._allowed_local_media_path + else: + allowed_path = getattr(self.model_config, "allowed_local_media_path", None) + + wav_list: list[float] | None = None + sr = 0 + artifact_key = "" + post_key = "" + for attempt in range(_REF_AUDIO_METADATA_FETCH_ATTEMPTS): + cache_key = await self._ref_audio_cache_key(ref_audio_str, allowed_path) + cached = self._ref_audio_resolve_cache.get(cache_key) + if cached is not None: + self._ref_audio_resolve_cache.move_to_end(cache_key) + wav_list, sr, _, _ = cached + logger.debug( + "Resolved ref_audio from cache: samples=%d sr=%d duration_s=%.3f", + len(wav_list), + sr, + len(wav_list) / sr if sr > 0 else 0.0, + ) + return wav_list, sr, cache_key + + if self._media_connector is None: + model_config = self.model_config + self._media_connector = MediaConnector( + allowed_local_media_path=model_config.allowed_local_media_path, + allowed_media_domains=model_config.allowed_media_domains, + ) + + fetch_start_s = time.perf_counter() + wav_np, fetched_sr = await self._media_connector.fetch_audio_async(ref_audio_str) + fetch_decode_ms = (time.perf_counter() - fetch_start_s) * 1000.0 + tolist_start_s = time.perf_counter() + wav_list, sr, artifact_key, duration = self._finalize_fetched_ref_audio(wav_np, fetched_sr) + tolist_ms = (time.perf_counter() - tolist_start_s) * 1000.0 + logger.debug( + "Resolved ref_audio: fetch_decode_ms=%.3f tolist_ms=%.3f samples=%d sr=%d duration_s=%.3f", + fetch_decode_ms, + tolist_ms, + len(wav_list), + sr, + duration, + ) + post_key = await self._ref_audio_cache_key(ref_audio_str, allowed_path) + if post_key == cache_key: + self._put_resolved_ref_audio(cache_key, wav_list, sr, artifact_key) + return wav_list, sr, cache_key + logger.debug( + "ref_audio metadata changed during fetch (attempt %d/%d); retrying", + attempt + 1, + _REF_AUDIO_METADATA_FETCH_ATTEMPTS, + ) + + logger.warning( + "ref_audio file changed during fetch after %d attempts; skipping resolve cache", + _REF_AUDIO_METADATA_FETCH_ATTEMPTS, ) - self._put_resolved_ref_audio(cache_key, wav_list, sr, artifact_key) - return wav_list, sr + assert wav_list is not None and post_key + return wav_list, sr, post_key @staticmethod def _make_ref_audio_artifact_cache_key(wav: np.ndarray, sr: int) -> str: @@ -1973,12 +1961,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): h.update(wav_f32.tobytes(order="C")) return h.hexdigest() - def _get_resolved_ref_audio_artifact_key(self, ref_audio_str: str) -> str | None: - source_key = hashlib.sha1(ref_audio_str.encode("utf-8")).hexdigest() - cached = self._ref_audio_resolve_cache.get(source_key) + def _get_resolved_ref_audio_artifact_key(self, cache_key: str) -> str | None: + """Look up the artifact key for a previously resolved ref_audio. + + *cache_key* must be the exact key returned by the preceding + ``_resolve_ref_audio`` call. Passing the key explicitly avoids a + second ``os.stat`` syscall and eliminates the TOCTOU window that + would arise from recomputing the key independently. + """ + cached = self._ref_audio_resolve_cache.get(cache_key) if cached is None: return None - self._ref_audio_resolve_cache.move_to_end(source_key) + self._ref_audio_resolve_cache.move_to_end(cache_key) return cached[3] def _put_resolved_ref_audio(self, cache_key: str, wav_list: list[float], sr: int, artifact_key: str) -> None: @@ -2043,7 +2037,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): async def _resolve_ref_audio_many(self, ref_audio_list: list[str]) -> list[tuple[list[float], int]]: resolved = [] for ref_audio in ref_audio_list: - resolved.append(await self._resolve_ref_audio(ref_audio)) + wav_list, sr, _ = await self._resolve_ref_audio(ref_audio) + resolved.append((wav_list, sr)) return resolved # ---- Ming TTS helpers ---- @@ -2194,6 +2189,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_start_s: float | None = None, include_sample_rate: bool = False, usage_acc: SpeechOutputTokenCounter | None = None, + tts_params: dict[str, Any] | None = None, collect: dict | None = None, ): """Generate audio chunks for streaming response. @@ -2219,6 +2215,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): stream_start_s = request_start_s if request_start_s is not None else time.perf_counter() artifact_ready = False + # SSE supplies an accumulator for usage output. Raw-audio and WebSocket + # streams retain terminal metrics only when their model adapter needs + # generation validation. + adapter = self._get_tts_adapter() + if tts_params is not None and usage_acc is None and adapter is not None and adapter.validates_generation: + usage_acc = SpeechOutputTokenCounter() + try: async for res in generator: # Tally generated codec tokens for usage (reads per-stage metrics @@ -2296,6 +2299,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Audex contract: zero codec tokens must abort the stream, not # complete it cleanly with zero audio bytes. raise ValueError("Audex produced no audio (the thinker emitted zero or invalid codec tokens)") + # Check before committing the reference-audio artifact or logging + # success. Streaming protocols may already have emitted partial + # bytes, but they must terminate as an error rather than cleanly. + if tts_params is not None and usage_acc is not None: + self._validate_tts_generation(tts_params, usage_acc) self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True total_ms = (time.perf_counter() - stream_start_s) * 1000.0 @@ -2380,6 +2388,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ``input_tokens`` is computed from the request text + reference audio. """ usage_acc = SpeechOutputTokenCounter() + emitted_audio = False try: async for chunk in self._generate_audio_chunks( generator, @@ -2388,6 +2397,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raw_request=raw_request, request_start_s=request_start_s, usage_acc=usage_acc, + tts_params=tts_params, ): payload = { "type": "speech.audio.delta", @@ -2395,6 +2405,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "response_format": response_format, } data = json.dumps(payload, separators=(",", ":")) + emitted_audio = True yield f"event: speech.audio.delta\ndata: {data}\n\n" done_payload: dict[str, Any] = {"type": "speech.audio.done"} if request is not None: @@ -2406,16 +2417,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): except asyncio.CancelledError: raise except Exception as e: - payload = { - "type": "speech.audio.error", - "error": { - "message": str(e), - "type": "server_error", - "param": None, - "code": HTTPStatus.INTERNAL_SERVER_ERROR.value, - }, + error: dict[str, Any] = { + "message": str(e), + "type": "server_error", + "param": None, + "code": HTTPStatus.INTERNAL_SERVER_ERROR.value, } - data = json.dumps(payload, separators=(",", ":")) + if emitted_audio: + error.update(partial_audio=True, action="discard") + error_payload: dict[str, Any] = { + "type": "speech.audio.error", + "error": error, + } + data = json.dumps(error_payload, separators=(",", ":")) yield f"event: speech.audio.error\ndata: {data}\n\n" @staticmethod @@ -2478,6 +2492,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Speaker (voice) if request.voice is not None: voice_lower = request.voice.lower() + precomputed_speakers = self._adapter.capabilities.precomputed_speakers if self._adapter is not None else {} params["speaker"] = [request.voice] params["voice_created_at"] = [self._voice_created_at(voice_lower)] @@ -2511,8 +2526,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): logger.info( "Auto-set ref_audio for uploaded voice: %s (icl=%s)", request.voice, bool(stored_ref_text) ) - elif voice_lower in self.precomputed_speakers and request.ref_audio is None: - profile = self.precomputed_speakers[voice_lower] + elif voice_lower in precomputed_speakers and request.ref_audio is None: + profile = precomputed_speakers[voice_lower] mode = str(profile.get("mode") or "xvec").lower() params["speaker"] = [voice_lower] params["task_type"] = ["Base"] @@ -2713,7 +2728,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): + mm_processor_kwargs with prompt_text. """ # Resolve reference audio - wav_samples, sr = await self._resolve_ref_audio(request.ref_audio) + wav_samples, sr, _ = await self._resolve_ref_audio(request.ref_audio) audio_data = (np.asarray(wav_samples, dtype=np.float32), sr) # Wrap the reference transcript in the CosyVoice3 instruction template @@ -2807,154 +2822,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._audex_tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, folder)) return self._audex_tokenizer - def _inject_audex_tta_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any: - """Complete the Audex TTA contract on a CLONE of the stage-0 params. - - Always attaches the RVQ phase-mask contract (required for token - validity) and the CFG pair contract at the official default scale - 3.0 unless the caller passed ``cfg_scale`` (1.0 disables guidance). - - Returns the cloned params; the input is never mutated. On the - no-extra_params online path the input is an element of the engine's - SHARED ``default_sampling_params_list`` — writing per-request pair - state (``cfg_pair_id``/``cfg_null_prompt``) into it would race - concurrent requests and leak stale CFG metadata into later ones. - """ - import copy - - from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt - from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids - - cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt - if not isinstance(cond_prompt, str) or not cond_prompt: - raise ValueError("Audex TTA requires the adapter-built caption prompt") - - stage0_params = copy.deepcopy(stage0_params) - tokenizer = self._get_audex_tokenizer() - if self._audex_tta_rvq is None: - phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer) - self._audex_tta_rvq = { - "phase_token_ids": phase_token_ids, - "start_tid": start_tid, - "end_tid": end_tid, - # Official generation cap (decode truncates at 500 frames). - "codec_cap": 4000, - # The TTA prompt ends with . - "start_in_prompt": True, - } - - extra_args = getattr(stage0_params, "extra_args", None) - if extra_args is None: - extra_args = {} - stage0_params.extra_args = extra_args - extra_args["tta_rvq"] = self._audex_tta_rvq - - cfg_scale = extra_args.get("cfg_scale") - cfg_scale = 3.0 if cfg_scale is None else float(cfg_scale) - if cfg_scale <= 1.0: - extra_args.pop("cfg_scale", None) - return stage0_params - extra_args.update( - { - "cfg_scale": cfg_scale, - "cfg_role": "cond", - "cfg_pair_id": request_id, - "cfg_null_prompt": build_tta_null_prompt(cond_prompt, tokenizer), - } - ) - return stage0_params - - def _inject_audex_cfg_pair_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any: - """Complete the Audex CFG contract on a CLONE of the stage-0 params. - - The adapter validated ``cfg_scale``; here (where the final request id - exists) the cond role, pair id, and the length-matched null prompt are - attached. ``cfg_scale`` absent/1.0 returns a clone value-identical to - the non-CFG flow (the key is dropped rather than forwarded). - - Returns the cloned params; the input is never mutated (it may be an - element of the engine's shared ``default_sampling_params_list`` — - see ``_inject_audex_tta_args``). - """ - import copy - - stage0_params = copy.deepcopy(stage0_params) - extra_args = getattr(stage0_params, "extra_args", None) - if not extra_args or extra_args.get("cfg_scale") is None: - if extra_args: - extra_args.pop("cfg_scale", None) - return stage0_params - cfg_scale = float(extra_args["cfg_scale"]) - if cfg_scale <= 1.0: - extra_args.pop("cfg_scale", None) - return stage0_params - - from vllm_omni.model_executor.models.audex.prompt import build_null_prompt - - cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt - if not isinstance(cond_prompt, str) or not cond_prompt: - raise ValueError("Audex CFG requires the adapter-built text prompt") - null_prompt = build_null_prompt(cond_prompt, self._get_audex_tokenizer()) - extra_args.update( - { - "cfg_scale": cfg_scale, - "cfg_role": "cond", - "cfg_pair_id": request_id, - "cfg_null_prompt": null_prompt, - } - ) - # Guided decoding sharpens the distribution; the unguided default - # temperature (0.1) adds excess sampling noise under CFG. Measured - # on the en-24 gate corpus at cfg 1.5: temp 0.1 -> CER 7.31%, - # temp 0.05 -> CER 6.87% (unguided baseline 7.24%). - stage0_params.temperature = 0.05 - return stage0_params - - def _apply_cosyvoice3_dynamic_tokens( - self, - sampling_params_list: list, - request: OpenAICreateSpeechRequest, - ) -> list: - """Set min/max tokens from tokenized text length (ratios target tokens, not chars).""" - import copy - - from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer - from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token - - sampling_params_list = copy.deepcopy(sampling_params_list) - hf_cfg = self.model_config.hf_config - # Build the Qwen tokenizer once per process (resolving the model dir via - # snapshot_download at most once) and reuse it across requests. - tokenizer = self._cosyvoice3_tokenizer - if tokenizer is None: - model_path = self.engine_client.model_config.model - if not os.path.isdir(model_path): - from huggingface_hub import snapshot_download - - model_path = snapshot_download(model_path) - tokenizer = get_qwen_tokenizer( - token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path), - skip_special_tokens=hf_cfg.skip_special_tokens, - version=hf_cfg.version, - ) - self._cosyvoice3_tokenizer = tokenizer - _, text_token_len = extract_text_token( - request.input, - tokenizer, - hf_cfg.allowed_special, - ) - min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2) - max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20) - sampling_params_list[0].min_tokens = max(1, int(text_token_len * min_ratio)) - sampling_params_list[0].max_tokens = min(2048, int(text_token_len * max_ratio)) - logger.info( - "CosyVoice3 dynamic tokens: text_tokens=%d, min_tokens=%d, max_tokens=%d", - text_token_len, - sampling_params_list[0].min_tokens, - sampling_params_list[0].max_tokens, - ) - return sampling_params_list - # ---- GLM-TTS helpers ---- async def _build_glm_tts_prompt( @@ -2975,7 +2842,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): """ # Voice cloning requires ref_audio + ref_text if request.ref_audio is not None and request.ref_text: - wav_samples, sr = await self._resolve_ref_audio(request.ref_audio) + wav_samples, sr, _ = await self._resolve_ref_audio(request.ref_audio) audio_data = (np.asarray(wav_samples, dtype=np.float32), int(sr)) mm_kwargs: dict[str, Any] = { @@ -3107,8 +2974,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Uploaded voice: ref_audio was auto-set as [base64_data_url] ref_audio_source = tts_params["ref_audio"][0] if ref_audio_source is not None and isinstance(ref_audio_source, str): - wav_list, sr = await self._resolve_ref_audio(ref_audio_source) - artifact_key = self._get_resolved_ref_audio_artifact_key(ref_audio_source) + wav_list, sr, cache_key = await self._resolve_ref_audio(ref_audio_source) + tts_params["ref_audio_cache_key"] = cache_key + artifact_key = self._get_resolved_ref_audio_artifact_key(cache_key) if artifact_key: tts_params[_QWEN3_TTS_REF_AUDIO_CACHE_KEY] = [artifact_key] ref_code_length = self._estimate_ref_code_len([wav_list, sr]) @@ -3154,9 +3022,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Build prompt + tts_params via the per-model adapter (RFC #4327). Every # dedicated TTS model resolves to an adapter that owns its validation, - # uploaded-speaker handling, and prompt/param building. Sampling - # overrides and the model-type label remain in the orchestrator tail - # below (keyed on ``_tts_model_type``) during this incremental migration. + # uploaded-speaker handling, prompt/param building, and sampling + # overrides. The model-type label remains available for compatibility. # Non-TTS deployments (no adapter) fall through to the rejection below. # Capture inline-ref-audio status BEFORE validate(): several adapters # apply uploaded speakers inside validate(), which sets request.ref_audio @@ -3200,69 +3067,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = tts_params.get("task_type", ["unknown"])[0] else: model_type = "generic" - logger.info( - "TTS speech request %s: text=%r, model=%s", - request_id, - request.input[:50] + "..." if len(request.input) > 50 else request.input, - model_type, - ) - - # CosyVoice3: set dynamic min/max tokens based on text length. - # The official model requires min_token_text_ratio to prevent early - # EOS and max_token_text_ratio to cap generation length. - if self._tts_model_type == "cosyvoice3" and sampling_params_list: - sampling_params_list = self._apply_cosyvoice3_dynamic_tokens(sampling_params_list, request) - - # GLM-TTS: set dynamic min/max tokens based on text length. - if self._tts_model_type == "glm_tts" and sampling_params_list: - import copy - - sampling_params_list = copy.deepcopy(sampling_params_list) - glm_metadata = prompt.get("additional_information") if isinstance(prompt, dict) else None - text_len_value = None - if isinstance(glm_metadata, dict): - text_len_value = glm_metadata.get("glm_tts_text_token_len") - if isinstance(text_len_value, list) and text_len_value: - text_len_value = text_len_value[0] - text_token_len = ( - int(text_len_value) - if text_len_value is not None - else self._estimate_glm_tts_text_token_len(request.input) - ) - hf_cfg = self.model_config.hf_config - min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2) - max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20) - stage_min_tokens = getattr(sampling_params_list[0], "min_tokens", None) - stage_max_tokens = getattr(sampling_params_list[0], "max_tokens", None) - cap_candidates = [int(cap) for cap in (stage_max_tokens, request.max_new_tokens) if cap is not None] - hard_cap = min(cap_candidates) if cap_candidates else None - - min_tokens = max(1, int(text_token_len * min_ratio)) - if stage_min_tokens is not None: - min_tokens = max(min_tokens, int(stage_min_tokens)) - if hard_cap is not None: - min_tokens = min(min_tokens, hard_cap) - - max_tokens = max(min_tokens, int(text_token_len * max_ratio)) - if hard_cap is not None: - max_tokens = min(max_tokens, hard_cap) - sampling_params_list[0].min_tokens = min_tokens - sampling_params_list[0].max_tokens = max_tokens - seed = getattr(request, "seed", None) - if seed is not None: - sampling_params_list[0].seed = seed - logger.info( - "GLM-TTS dynamic tokens: text_tokens=%d, min_ratio=%s, max_ratio=%s, " - "stage_min=%s, stage_max=%s, request_max=%s, min_tokens=%d, max_tokens=%d", - text_token_len, - min_ratio, - max_ratio, - stage_min_tokens, - stage_max_tokens, - request.max_new_tokens, - min_tokens, - max_tokens, - ) + logger.info("TTS speech request %s: model=%s", request_id, model_type) + _rl = getattr(self, "request_logger", None) + if _rl: + base_len = len(f"TTS speech request {request_id}: text=") + raw_max = getattr(_rl, "max_log_len", None) + cap = raw_max if isinstance(raw_max, int) else 200 + text = request.input[: max(cap - base_len, 0)] + logger.debug("TTS speech request %s: text=%r", request_id, text) # Apply model-specific extra parameters if request.extra_params is not None and sampling_params_list: @@ -3282,52 +3094,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): sampling_params_list[0].extra_args.update(request.extra_params) logger.info("Applied extra_params: %s", request.extra_params) - # Audex CFG: turn an adapter-validated cfg_scale into the engine-side - # pair contract. This must run here (not in the adapter) because the - # pair id is the final request_id, which the adapter never sees. - # The injectors return a CLONE: without extra_params above, - # sampling_params_list still holds the engine's shared default params, - # which must never carry per-request CFG state. - if self._tts_model_type == "audex" and sampling_params_list: - sampling_params_list[0] = self._inject_audex_cfg_pair_args(request_id, prompt, sampling_params_list[0]) - elif self._tts_model_type == "audex_tta" and sampling_params_list: - sampling_params_list[0] = self._inject_audex_tta_args(request_id, prompt, sampling_params_list[0]) - - # Some TTS model defaults come from deploy YAML. Their AR - # generation length is controlled by SamplingParams.max_tokens, so only - # override it when the caller explicitly requests max_new_tokens. - if ( - self._tts_model_type in _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES - and request.max_new_tokens is not None - and sampling_params_list - ): - import copy - - sampling_params_list = copy.deepcopy(sampling_params_list) - sampling_params_list[0].max_tokens = request.max_new_tokens - if self._tts_model_type == "cosyvoice3": - sampling_params_list[0].min_tokens = min( - getattr(sampling_params_list[0], "min_tokens", 0), - request.max_new_tokens, - ) - elif self._tts_model_type == "ming_tts" and sampling_params_list: - import copy - - from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( - MOE_TEXT_EOS_TOKEN_ID, - TEXT_EOS_TOKEN_ID, - ) - - hf_config = self.engine_client.model_config.hf_config - is_moe = getattr(hf_config, "model_type", "") == "bailingmm" - stop_token_id = MOE_TEXT_EOS_TOKEN_ID if is_moe else TEXT_EOS_TOKEN_ID - - sampling_params_list = copy.deepcopy(sampling_params_list) - sampling_params_list[0].stop_token_ids = [int(stop_token_id)] - if request.max_new_tokens is not None: - # Ming emits TEXT_EOS after the latent decode budget is exhausted, so - # Stage-0 needs one extra token beyond ming_max_decode_steps. - sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1 + # Apply adapter-owned sampling overrides, including request-level token + # limits and model-specific dynamic token or stop-token configuration. + if sampling_params_list and (adapter := self._get_tts_adapter()) is not None: + sampling_params_list = adapter.apply_sampling_overrides(sampling_params_list, request, prompt, request_id) if request.seed is not None and sampling_params_list: import copy @@ -3374,7 +3144,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return request_id, generator, tts_params async def _generate_pcm_chunks( - self, generator, request_id: str, *, include_sample_rate: bool = False, collect: dict | None = None + self, + generator, + request_id: str, + *, + include_sample_rate: bool = False, + tts_params: dict[str, Any] | None = None, + collect: dict | None = None, ): """Yield raw PCM byte chunks from the engine generator. @@ -3388,15 +3164,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id, response_format="pcm", include_sample_rate=include_sample_rate, + tts_params=tts_params, collect=collect, ): yield chunk async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest): """Yield raw PCM bytes for a speech request as soon as chunks are decoded.""" - request_id, generator, _ = await self._prepare_speech_generation(request) + request_id, generator, tts_params = await self._prepare_speech_generation(request) try: - async for chunk in self._generate_pcm_chunks(generator, request_id): + async for chunk in self._generate_pcm_chunks(generator, request_id, tts_params=tts_params): yield chunk finally: self._discard_ref_audio_artifact_warmup(request_id) @@ -3413,7 +3190,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # ``usage_out`` is an opt-in output channel: when a list is passed, the # computed SpeechTokenUsage is appended to it. The return stays a # 2-tuple so existing callers (and their test mocks) are unaffected; - # only the batch path, which surfaces per-item usage, opts in. + # batch and non-streaming response-header paths opt in when surfacing + # usage outside the raw audio body. request_id, generator, bytes_tts_params = await self._prepare_speech_generation( request, request_id=request_id, has_inline_ref_audio=has_inline_ref_audio ) @@ -3467,6 +3245,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if final_output is None: raise ValueError("No output generated from the model.") + self._validate_tts_generation(bytes_tts_params or {}, usage_acc) + # Extract audio from the audio-bearing res (not necessarily the last # yielded one, which may be the aligner's timestamps output). audio_source = audio_res if audio_res is not None else final_output @@ -3557,6 +3337,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if not artifact_ready: self._discard_ref_audio_artifact_warmup(request_id) + def _get_normalized_voice(self, voice: str | None) -> str | None: + """Get the normalized voice to be used; currently this means that + the voice is a: + - lowercase str if it's a valid supported/uploaded speaker + - None if the voice is the placeholder default or not provided + """ + if voice is not None: + voice = voice.lower() + available_speakers = self._get_available_speakers() + if voice not in available_speakers: + raise ValueError( + f"Invalid voice '{voice}'. Supported: {', '.join(sorted(self._get_available_voices()))}" + ) + return voice + async def _create_diffusion_speech( self, request: OpenAICreateSpeechRequest, @@ -3573,11 +3368,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if fmt_err: return self._diffusion_error_response(fmt_err, status_code=400) - if request.voice: - voice_lower = request.voice.lower() - if voice_lower not in self.uploaded_speakers and voice_lower not in self.supported_speakers: - all_voices = sorted(self.uploaded_speakers.keys() | self.supported_speakers) - raise ValueError(f"Invalid voice '{request.voice}'. Supported: {', '.join(all_voices) or 'none'}") + request.voice = self._get_normalized_voice(request.voice) has_inline_ref_audio = request.ref_audio is not None err = self._apply_uploaded_speaker(request) @@ -3587,26 +3378,31 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id = f"speech-{random_uuid()}" prompt: dict[str, Any] = {"input": request.input} if request.ref_audio: - wav, sr = await self._resolve_ref_audio(request.ref_audio) + wav, sr, _ = await self._resolve_ref_audio(request.ref_audio) prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr) if request.ref_text: prompt["ref_text"] = request.ref_text if request.voice: - voice_lower = request.voice.lower() - if voice_lower in self.uploaded_speakers and not has_inline_ref_audio: - prompt["voice_name"] = voice_lower - prompt["voice_created_at"] = self._voice_created_at(voice_lower) + if request.voice in self.uploaded_speakers and not has_inline_ref_audio: + prompt["voice_name"] = request.voice + prompt["voice_created_at"] = self._voice_created_at(request.voice) if request.language: prompt["lang"] = request.language if request.instructions: prompt["instruct"] = request.instructions logger.info( - "Diffusion TTS speech request %s: text=%r, voice_clone=%s", + "Diffusion TTS speech request %s: voice_clone=%s", request_id, - request.input[:50] + "..." if len(request.input) > 50 else request.input, "ref_audio" in prompt, ) + _rl = getattr(self, "request_logger", None) + if _rl: + base_len = len(f"Diffusion TTS speech request {request_id}: text=") + raw_max = getattr(_rl, "max_log_len", None) + cap = raw_max if isinstance(raw_max, int) else 200 + text = request.input[: max(cap - base_len, 0)] + logger.debug("Diffusion TTS speech request %s: text=%r", request_id, text) if request.extra_params is not None and not isinstance(request.extra_params, dict): raise ValueError("extra_params must be a JSON object/dict.") extra = dict(request.extra_params or {}) @@ -3734,6 +3530,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Raw audio streaming yields each Code2Wav chunk as raw bytes as soon as it is decoded. Raw WAV streaming emits a header with placeholder size values first. """ + if request.voice is not None: + if _is_default_voice(request.voice.lower(), self._get_available_speakers()): + request.voice = None + if self._diffusion_mode: return await self._create_diffusion_speech(request) @@ -3771,7 +3571,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return error media_type = "audio/wav" if response_format == "wav" else "audio/pcm" - _, generator, _ = await self._prepare_speech_generation(request, request_id=request_id) + _, generator, raw_tts_params = await self._prepare_speech_generation(request, request_id=request_id) return StreamingResponse( self._generate_audio_chunks( generator, @@ -3779,6 +3579,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): response_format, raw_request=raw_request, request_start_s=request_start_s, + tts_params=raw_tts_params, ), media_type=media_type, ) @@ -3806,7 +3607,36 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) collect: dict = {} - audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id, collect=collect) + usage_box: list[SpeechTokenUsage] = [] + try: + audio_bytes, media_type = await self._generate_audio_bytes( + request, request_id=request_id, usage_out=usage_box, collect=collect + ) + except TTSGenerationError as error: + # An adapter can reject otherwise completed audio. Retry only + # retryable, stochastic, server-budgeted non-streaming + # requests: changing an explicit seed would violate + # reproducibility, and changing a caller-provided token limit + # would violate the requested budget. + if not error.retryable or request.seed is not None or request.max_new_tokens is not None: + raise + retry_seed = int(random_uuid()[:8], 16) + retry_request = request.model_copy(update={"seed": retry_seed}) + retry_request_id = f"speech-{random_uuid()}" + collect.clear() + usage_box.clear() + logger.warning( + "TTS request %s failed generation validation; retrying once as %s with seed=%d", + request_id, + retry_request_id, + retry_seed, + ) + audio_bytes, media_type = await self._generate_audio_bytes( + retry_request, + request_id=retry_request_id, + usage_out=usage_box, + collect=collect, + ) total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.info( "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d", @@ -3814,7 +3644,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): total_ms, len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)), ) - headers = {} + headers = self._build_speech_usage_headers(usage_box[0] if usage_box else None) if collect.get("word_timestamps") is not None: # Default ensure_ascii keeps the header latin-1 encodable (non-ASCII words \uXXXX-escaped). ts_json = json.dumps(collect["word_timestamps"]) @@ -3859,6 +3689,16 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): e, ) return self.create_error_response(e) + except TTSGenerationError as e: + total_ms = (time.perf_counter() - request_start_s) * 1000.0 + logger.error( + "[SpeechE2E] request_id=%s stream=%s status=invalid_generation total_ms=%.2f error=%s", + request_id, + bool(request.stream), + total_ms, + e, + ) + return self._diffusion_error_response(str(e), status_code=500) except Exception as e: total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.exception( diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 2b71ad6c7..0f883f179 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """WebSocket handler for streaming text input TTS. Accepts text incrementally via WebSocket, buffers it until input.done, and @@ -329,7 +332,7 @@ class OmniStreamingSpeechHandler: request_id = None try: if config.stream_audio: - request_id, generator, _ = await self._speech_service._prepare_speech_generation(request) + request_id, generator, tts_params = await self._speech_service._prepare_speech_generation(request) if config.word_timestamps: total_bytes = await self._stream_audio_with_alignments( websocket=websocket, @@ -339,9 +342,16 @@ class OmniStreamingSpeechHandler: utterance_index=utterance_index, sentence_index=sentence_index, language=config.language, + tts_params=tts_params, ) else: - async with aclosing(self._speech_service._generate_pcm_chunks(generator, request_id)) as stream: + async with aclosing( + self._speech_service._generate_pcm_chunks( + generator, + request_id, + tts_params=tts_params, + ) + ) as stream: async for chunk in stream: total_bytes += len(chunk) await websocket.send_bytes(chunk) @@ -367,6 +377,7 @@ class OmniStreamingSpeechHandler: await self._send_error( websocket, f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}", + partial_audio=total_bytes > 0, ) finally: try: @@ -392,6 +403,7 @@ class OmniStreamingSpeechHandler: utterance_index: int, sentence_index: int, language: str | None = None, + tts_params: dict | None = None, ) -> int: """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence. @@ -438,6 +450,7 @@ class OmniStreamingSpeechHandler: generator, request_id, include_sample_rate=True, + tts_params=tts_params, collect=collect, ) ) as stream: @@ -463,14 +476,15 @@ class OmniStreamingSpeechHandler: return total_bytes @staticmethod - async def _send_error(websocket: WebSocket, message: str) -> None: + async def _send_error(websocket: WebSocket, message: str, *, partial_audio: bool = False) -> None: """Send an error message to the client.""" try: - await websocket.send_json( - { - "type": "error", - "message": message, - } - ) + payload: dict[str, object] = { + "type": "error", + "message": message, + } + if partial_audio: + payload.update(partial_audio=True, action="discard") + await websocket.send_json(payload) except Exception: pass # Connection may already be closed; safe to ignore diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 9d511ba6d..398c7b81f 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -7,8 +7,10 @@ import copy import time from collections.abc import Mapping from dataclasses import dataclass +from functools import cached_property from http import HTTPStatus -from typing import Any, cast +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, cast from fastapi import HTTPException from PIL import Image @@ -27,9 +29,14 @@ from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, get_default_sampling_params_list, ) -from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request -from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64 +from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request +from vllm_omni.entrypoints.openai.video_api_utils import ( + _encode_video_bytes, + _PlanarFrameConverter, + encode_video_base64, +) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.model_extras import should_preserve_reference_image_size from vllm_omni.outputs.output_metadata import ( DiffusionMetadataMapping, DiffusionMultimodalOutput, @@ -38,6 +45,11 @@ from vllm_omni.outputs.output_metadata import ( logger = init_logger(__name__) +_VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS = 8 + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig + @dataclass class ReferenceImage: @@ -87,6 +99,17 @@ class OmniOpenAIServingVideo: self._engine_client = engine_client self._model_name = model_name self._stage_configs = stage_configs + self._video_frame_converter = _PlanarFrameConverter(max_workers=_VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS) + logger.info( + "Video response frame conversion pool configured: workers=%d", + self._video_frame_converter.max_workers, + ) + + def _resolve_diffusion_od_config(self) -> OmniDiffusionConfig | SimpleNamespace | None: + get_od_config = getattr(self._engine_client, "get_diffusion_od_config", None) + if callable(get_od_config): + return get_od_config() + return getattr(self._engine_client, "od_config", None) @property def model_name(self) -> str | None: @@ -100,20 +123,47 @@ class OmniOpenAIServingVideo: if self._stage_configs is None and stage_configs is not None: self._stage_configs = stage_configs + @cached_property + def preserves_reference_image_size(self) -> bool: + """Return whether the active pipeline owns reference-image resizing.""" + od_config = self._resolve_diffusion_od_config() + model_class_name = None if od_config is None else getattr(od_config, "model_class_name", None) + model = getattr(od_config, "model", None) if od_config is not None else None + model = model or self.model_name + revision = getattr(od_config, "revision", None) if od_config is not None else None + return should_preserve_reference_image_size( + model_class_name, + model=None if model is None else str(model), + revision=revision, + ) + @property def supports_mixed_reference_inputs(self) -> bool: """Return whether the configured diffusion model accepts mixed refs.""" - get_od_config = getattr(self._engine_client, "get_diffusion_od_config", None) - od_config = get_od_config() if callable(get_od_config) else getattr(self._engine_client, "od_config", None) + od_config = self._resolve_diffusion_od_config() if od_config is None: return False capability = getattr(od_config, "supports_mixed_reference_inputs", None) - if isinstance(capability, bool): - return capability - model_class_name = getattr(od_config, "model_class_name", None) - return get_diffusion_model_metadata(model_class_name).supports_mixed_reference_inputs + model_archs = [model_class_name] + for stage_config in self.stage_configs or (): + stage_get = ( + stage_config.get if isinstance(stage_config, Mapping) else lambda key: getattr(stage_config, key, None) + ) + engine_args = stage_get("engine_args") or {} + model_archs.extend( + ( + stage_get("model_arch"), + engine_args.get("model_class_name") + if isinstance(engine_args, Mapping) + else getattr(engine_args, "model_class_name", None), + ) + ) + metadata_capability = any( + get_diffusion_model_metadata(model_arch).supports_mixed_reference_inputs for model_arch in model_archs + ) + return capability is True or metadata_capability @classmethod def for_diffusion( @@ -128,6 +178,9 @@ class OmniOpenAIServingVideo: stage_configs=stage_configs, ) + def shutdown(self) -> None: + self._video_frame_converter.shutdown() + async def _run_and_extract( self, request: VideoGenerationRequest, @@ -154,7 +207,12 @@ class OmniOpenAIServingVideo: provided_fields = request.model_fields_set fps_provided = self._request_fps_provided(request) vp = request.resolve_video_params() - if input_image is not None and vp.width is not None and vp.height is not None: + if ( + input_image is not None + and vp.width is not None + and vp.height is not None + and not self.preserves_reference_image_size + ): target_size = (vp.width, vp.height) image_items = input_image if isinstance(input_image, list) else [input_image] resized_images = [ @@ -310,7 +368,12 @@ class OmniOpenAIServingVideo: video_data = [ VideoData( b64_json=( - encode_video_base64(video, fps=artifacts.output_fps, video_codec_options=video_codec_options) + encode_video_base64( + video, + fps=artifacts.output_fps, + video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, + ) if artifacts.audios[idx] is None else encode_video_base64( video, @@ -318,6 +381,7 @@ class OmniOpenAIServingVideo: audio=artifacts.audios[idx], audio_sample_rate=artifacts.audio_sample_rate, video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, ) ), action=artifacts.actions[idx], @@ -374,6 +438,7 @@ class OmniOpenAIServingVideo: fps=artifacts.output_fps, **({"audio": audio, "audio_sample_rate": artifacts.audio_sample_rate} if audio is not None else {}), video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, ) _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms) @@ -441,14 +506,13 @@ class OmniOpenAIServingVideo: detail="Stage configs not found. Start server with an omni diffusion model.", ) - # Video generation endpoint only supports diffusion stages. - for stage in stage_configs: - stage_type = get_stage_type(stage) - if stage_type != "diffusion": - raise HTTPException( - status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, - detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.", - ) + # Video pipelines may use preparatory AR stages (for example, a + # vLLM-hosted text encoder) before the diffusion stage. + if not is_video_generation_pipeline(stage_configs): + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="No final video output stage found in video generation pipeline.", + ) # Common generation logic for both paths engine_client = cast(AsyncOmni, self._engine_client) diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index 770890faf..a9abad146 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -51,7 +51,7 @@ from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, get_default_sampling_params_list, ) -from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request +from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request from vllm_omni.entrypoints.openai.video_api_utils import ( StreamingVideoFormat, create_streaming_video_encoder, @@ -616,13 +616,11 @@ class OmniStreamingVideoOutputHandler: detail="Stage configs not found. Start server with an omni diffusion model.", ) - for stage in stage_configs: - stage_type = get_stage_type(stage) - if stage_type != "diffusion": - raise HTTPException( - status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, - detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.", - ) + if not is_video_generation_pipeline(stage_configs): + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="No final video output stage found in video generation pipeline.", + ) sampling_params_list = build_stage_sampling_params_list( list(stage_configs), diff --git a/vllm_omni/entrypoints/openai/speech_usage.py b/vllm_omni/entrypoints/openai/speech_usage.py index 5e38405fd..cb7667b48 100644 --- a/vllm_omni/entrypoints/openai/speech_usage.py +++ b/vllm_omni/entrypoints/openai/speech_usage.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Token-usage accounting for the Speech (``/v1/audio/speech``) API. Why this module exists (issue #4646) @@ -160,6 +161,7 @@ class SpeechOutputTokenCounter: """ output_tokens: int = 0 + stage0_finish_reason: str | None = None def observe(self, res: Any) -> None: metrics = getattr(res, "metrics", None) @@ -168,6 +170,11 @@ class SpeechOutputTokenCounter: stage_metrics = metrics.get("stage_metrics") if not isinstance(stage_metrics, dict): return + stage0 = stage_metrics.get("0", stage_metrics.get(0)) + if isinstance(stage0, dict): + reason = stage0.get("finish_reason") + if reason is not None: + self.stage0_finish_reason = str(reason) for stage in stage_metrics.values(): if not isinstance(stage, dict): continue diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index cd44f0ff5..dcf8c4eaa 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Registry of TTS serving adapters, and the model detection built on it. Adapters register themselves by their ``name`` (the model-type discriminator) @@ -9,8 +10,6 @@ key alone cannot distinguish. That metadata is the single source of truth. :func:`detect_tts_model_type` and :func:`all_tts_stage_keys` derive the whole stage -> model-type mapping from it, so a new TTS model is one adapter file and needs no edit to ``serving_speech.py``. -Models still awaiting an adapter are listed explicitly in -:data:`LEGACY_TTS_DETECTORS`. """ from vllm.logger import init_logger @@ -18,10 +17,10 @@ from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters.base import ( ARTTSAdapter, DiffusionTTSAdapter, - LegacyDetector, OutputPolicy, PreparedRequest, SpeechServingContext, + TTSGenerationError, TTSModelAdapter, ) @@ -29,10 +28,6 @@ logger = init_logger(__name__) TTS_ADAPTER_REGISTRY: dict[str, type[TTSModelAdapter]] = {} -#: Model types the serving layer detects that have no adapter yet. See -#: :class:`LegacyDetector`. This list must only ever shrink. -LEGACY_TTS_DETECTORS: tuple[LegacyDetector, ...] = () - def register_tts_adapter(cls: type[TTSModelAdapter]) -> type[TTSModelAdapter]: """Class decorator: index ``cls`` under its ``name`` (model-type).""" @@ -57,16 +52,13 @@ def resolve_adapter(model_type: str | None) -> type[TTSModelAdapter] | None: return TTS_ADAPTER_REGISTRY.get(model_type) -def iter_tts_detectors() -> list[type[TTSModelAdapter] | LegacyDetector]: +def iter_tts_detectors() -> list[type[TTSModelAdapter]]: """Every detector, in resolution order. Sorted by ``detect_priority`` then ``name``, so the order is total and does not depend on adapter import order. """ - detectors: list[type[TTSModelAdapter] | LegacyDetector] = [ - *TTS_ADAPTER_REGISTRY.values(), - *LEGACY_TTS_DETECTORS, - ] + detectors = list(TTS_ADAPTER_REGISTRY.values()) detectors.sort(key=lambda d: (d.detect_priority, d.name)) return detectors @@ -121,6 +113,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 indextts2, ming_flash_omni_tts, ming_tts, + minimax_music3, moss_tts, omnivoice, qwen3_tts, @@ -132,12 +125,11 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 __all__ = [ "ARTTSAdapter", "DiffusionTTSAdapter", - "LegacyDetector", "OutputPolicy", "PreparedRequest", "SpeechServingContext", + "TTSGenerationError", "TTSModelAdapter", - "LEGACY_TTS_DETECTORS", "TTS_ADAPTER_REGISTRY", "all_tts_model_types", "all_tts_stage_keys", diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex.py b/vllm_omni/entrypoints/openai/tts_adapters/audex.py index 4533a420b..ab7c7c25e 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/audex.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/audex.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 """Audex (Nemotron-Labs-Audex-2B) TTS serving adapter.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt if TYPE_CHECKING: @@ -17,8 +17,8 @@ if TYPE_CHECKING: AUDEX_CFG_SCALE_MIN = 1.0 AUDEX_CFG_SCALE_MAX = 10.0 -# Internal CFG pair-plumbing keys; injected by the serving layer once the -# request id exists, never accepted from callers. +# Internal CFG pair-plumbing keys; injected by this adapter once it receives +# the final request id, never accepted from callers. _AUDEX_INTERNAL_CFG_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt") @@ -74,3 +74,65 @@ class AudexAdapter(ARTTSAdapter): ) -> PreparedRequest: prompt = {"prompt": build_cond_prompt(request.input)} return PreparedRequest(prompt=prompt, tts_params={}, model_type="audex") + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + if request_id is None: + raise ValueError("Audex CFG requires a request id") + stage0_params = sampling_params_list[0] + sampling_params_list[0] = self._inject_audex_cfg_pair_args( + request_id=request_id, prompt=prompt, stage0_params=stage0_params + ) + return apply_max_new_tokens(sampling_params_list, request) + + def _inject_audex_cfg_pair_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any: + """Complete the Audex CFG contract on a CLONE of the stage-0 params. + + The adapter validated ``cfg_scale``; here (where the final request id + exists) the cond role, pair id, and the length-matched null prompt are + attached. ``cfg_scale`` absent/1.0 returns a clone value-identical to + the non-CFG flow (the key is dropped rather than forwarded). + + Returns the cloned params; the input is never mutated (it may be an + element of the engine's shared ``default_sampling_params_list`` — + see ``_inject_audex_tta_args``). + """ + server = self.ctx.server + import copy + + stage0_params = copy.deepcopy(stage0_params) + extra_args = getattr(stage0_params, "extra_args", None) + if not extra_args or extra_args.get("cfg_scale") is None: + if extra_args: + extra_args.pop("cfg_scale", None) + return stage0_params + cfg_scale = float(extra_args["cfg_scale"]) + if cfg_scale <= 1.0: + extra_args.pop("cfg_scale", None) + return stage0_params + + from vllm_omni.model_executor.models.audex.prompt import build_null_prompt + + cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt + if not isinstance(cond_prompt, str) or not cond_prompt: + raise ValueError("Audex CFG requires the adapter-built text prompt") + null_prompt = build_null_prompt(cond_prompt, server._get_audex_tokenizer()) + extra_args.update( + { + "cfg_scale": cfg_scale, + "cfg_role": "cond", + "cfg_pair_id": request_id, + "cfg_null_prompt": null_prompt, + } + ) + # Guided decoding sharpens the distribution; the unguided default + # temperature (0.1) adds excess sampling noise under CFG. Measured + # on the en-24 gate corpus at cfg 1.5: temp 0.1 -> CER 7.31%, + # temp 0.05 -> CER 6.87% (unguided baseline 7.24%). + stage0_params.temperature = 0.05 + return stage0_params diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py b/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py index f754f0158..33dd2294a 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/audex_tta.py @@ -2,16 +2,16 @@ """Audex (Nemotron-Labs-Audex-2B) text-to-audio serving adapter. Caption → general audio through the ``audex_tta`` pipeline. -The adapter builds the TTA prompt; the serving layer completes the request -after the final request id exists: the RVQ phase-mask contract +The adapter builds the TTA prompt and, after receiving the final request id, +completes the sampling parameters with the RVQ phase-mask contract (``extra_args["tta_rvq"]``) and the CFG pair contract (official default scale 3.0 — guidance is effectively mandatory for TTA quality). """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt if TYPE_CHECKING: @@ -21,8 +21,7 @@ AUDEX_TTA_DEFAULT_CFG_SCALE = 3.0 AUDEX_TTA_CFG_SCALE_MIN = 1.0 AUDEX_TTA_CFG_SCALE_MAX = 10.0 -# Internal plumbing keys; injected by the serving layer, never accepted -# from callers. +# Internal plumbing keys; injected by this adapter, never accepted from callers. _INTERNAL_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt", "tta_rvq") @@ -67,3 +66,76 @@ class AudexTTAAdapter(ARTTSAdapter): ) -> PreparedRequest: prompt = {"prompt": build_tta_cond_prompt(request.input)} return PreparedRequest(prompt=prompt, tts_params={}, model_type="audex_tta") + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + if request_id is None: + raise ValueError("Audex TTA requires a request id") + stage0_params = sampling_params_list[0] + sampling_params_list[0] = self._inject_audex_tta_args( + request_id=request_id, prompt=prompt, stage0_params=stage0_params + ) + return apply_max_new_tokens(sampling_params_list, request) + + def _inject_audex_tta_args(self, request_id: str, prompt: Any, stage0_params: Any) -> Any: + """Complete the Audex TTA contract on a CLONE of the stage-0 params. + + Always attaches the RVQ phase-mask contract (required for token + validity) and the CFG pair contract at the official default scale + 3.0 unless the caller passed ``cfg_scale`` (1.0 disables guidance). + + Returns the cloned params; the input is never mutated. On the + no-extra_params online path the input is an element of the engine's + SHARED ``default_sampling_params_list`` — writing per-request pair + state (``cfg_pair_id``/``cfg_null_prompt``) into it would race + concurrent requests and leak stale CFG metadata into later ones. + """ + server = self.ctx.server + import copy + + from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt + from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids + + cond_prompt = prompt.get("prompt") if isinstance(prompt, dict) else prompt + if not isinstance(cond_prompt, str) or not cond_prompt: + raise ValueError("Audex TTA requires the adapter-built caption prompt") + + stage0_params = copy.deepcopy(stage0_params) + tokenizer = server._get_audex_tokenizer() + if server._audex_tta_rvq is None: + phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer) + server._audex_tta_rvq = { + "phase_token_ids": phase_token_ids, + "start_tid": start_tid, + "end_tid": end_tid, + # Official generation cap (decode truncates at 500 frames). + "codec_cap": 4000, + # The TTA prompt ends with . + "start_in_prompt": True, + } + + extra_args = getattr(stage0_params, "extra_args", None) + if extra_args is None: + extra_args = {} + stage0_params.extra_args = extra_args + extra_args["tta_rvq"] = server._audex_tta_rvq + + cfg_scale = extra_args.get("cfg_scale") + cfg_scale = 3.0 if cfg_scale is None else float(cfg_scale) + if cfg_scale <= 1.0: + extra_args.pop("cfg_scale", None) + return stage0_params + extra_args.update( + { + "cfg_scale": cfg_scale, + "cfg_role": "cond", + "cfg_pair_id": request_id, + "cfg_null_prompt": build_tta_null_prompt(cond_prompt, tokenizer), + } + ) + return stage0_params diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index b8679041c..4842acfd7 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Base contract for per-model TTS serving adapters. This package factors the per-model ``if self._tts_model_type == ...`` dispatch @@ -7,19 +8,34 @@ model's request normalization, validation, prompt/param building, sampling overrides, and output policy, so adding a model means writing one adapter file instead of editing the shared serving module in ~10 scattered places. -See the RFC for the full design (issue #4327). This is the foundation landed in -the first migration PR; Qwen3-TTS is the first model routed through it while the -remaining models stay on the legacy path until individually migrated. +See the RFC for the full design (issue #4327). """ from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar +from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_codec_frame_rate, load_supported_speakers + if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +DEFAULT_TTS_LANGUAGES = frozenset( + { + "Auto", + "Chinese", + "English", + "Japanese", + "Korean", + "German", + "French", + "Russian", + "Portuguese", + "Spanish", + "Italian", + } +) _conditioning_cache_salt_fn: "Callable[..., str] | None" = None @@ -39,6 +55,29 @@ def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: di return _conditioning_cache_salt_fn(request, tts_params) +def apply_max_new_tokens( + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", +) -> list: + """Apply a request-level ``max_new_tokens`` limit.""" + if request.max_new_tokens is None: + return sampling_params_list + + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + sampling_params_list[0].max_tokens = request.max_new_tokens + return sampling_params_list + + +class TTSGenerationError(RuntimeError): + """A completed TTS generation that must not be returned as valid audio.""" + + def __init__(self, message: str, *, retryable: bool = False) -> None: + super().__init__(message) + self.retryable = retryable + + @dataclass class OutputPolicy: """How the orchestrator aggregates engine output for a model. @@ -88,6 +127,14 @@ class SpeechServingContext: diffusion_engine: Any | None = None +@dataclass(frozen=True) +class TTSCapabilities: + precomputed_speakers: dict[str, dict[str, Any]] = field(default_factory=dict) + supported_speakers: frozenset[str] = frozenset() + supported_languages: frozenset[str] = DEFAULT_TTS_LANGUAGES + codec_frame_rate: float | None = None + + class TTSModelAdapter(ABC): """Mandatory base class for a TTS model served via ``/v1/audio/speech``. @@ -115,6 +162,8 @@ class TTSModelAdapter(ABC): detect_priority: ClassVar[int] = 100 #: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``. backend: ClassVar[str] = "ar" + #: Whether streaming entrypoints must retain terminal metrics for validation. + validates_generation: ClassVar[bool] = False #: Whether the model consumes ``request.speed`` in its native parameters. native_speed_control: ClassVar[bool] = False @@ -124,6 +173,7 @@ class TTSModelAdapter(ABC): def __init__(self, ctx: SpeechServingContext) -> None: self.ctx = ctx + self.capabilities = TTSCapabilities() @classmethod def matches(cls, model_stage: str | None, model_arch: str | None) -> bool: @@ -184,6 +234,8 @@ class TTSModelAdapter(ABC): self, sampling_params_list: list, request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, ) -> list: """Apply model-specific sampling mutations. @@ -192,6 +244,49 @@ class TTSModelAdapter(ABC): """ return sampling_params_list + def validate_generation( + self, + tts_params: Mapping[str, object], + *, + stage0_finish_reason: str | None, + output_tokens: int, + ) -> None: + """Reject a completed generation that is invalid for this model. + + Adapters that override this hook must also set + :attr:`validates_generation` so streaming entrypoints retain the + terminal metrics needed by the validation without charging that cost + to unrelated TTS models. + """ + + async def warmup(self) -> None: + return + + def validate_tts_embedding_dim(self, emb_dim: int) -> str | None: + return None + + def load_capabilities(self) -> TTSCapabilities: + self.capabilities = TTSCapabilities( + precomputed_speakers=self._load_precomputed_speakers(), + supported_speakers=frozenset(self._load_supported_speakers()), + supported_languages=self._load_supported_languages(), + codec_frame_rate=self._load_codec_frame_rate(), + ) + return self.capabilities + + def _load_precomputed_speakers(self) -> dict[str, dict[str, Any]]: + return {} + + def _load_supported_speakers(self) -> set[str]: + # Preserve the legacy default path, which reads talker_config. + return load_supported_speakers(self.ctx.engine_client) + + def _load_supported_languages(self) -> frozenset[str]: + return DEFAULT_TTS_LANGUAGES + + def _load_codec_frame_rate(self) -> float | None: + return load_codec_frame_rate(self.ctx.engine_client) + class ARTTSAdapter(TTSModelAdapter): """Adapter for models served through the AR ``engine_client`` path. @@ -228,39 +323,13 @@ class DiffusionTTSAdapter(TTSModelAdapter): return frozenset(params) if params is not None else frozenset() -@dataclass(frozen=True) -class LegacyDetector: - """A model type the serving layer detects but has no adapter for yet. - - Detection has to keep naming these models so ``serving_speech.py`` can route - them down its legacy path, but they own none of the adapter contract. Each - entry is a migration debt: delete it when the model gets an adapter. The - pre-commit gate (``tools/pre_commit/check_tts_adapter.py``) fails if the list - grows. - - Exposes the subset of the :class:`TTSModelAdapter` class surface that - detection reads, so both kinds live in one sorted sequence. - """ - - name: str - stage_keys: frozenset[str] = frozenset() - model_archs: frozenset[str] = frozenset() - arch_identifies_entry_stage: bool = False - detect_priority: int = 100 - - def matches(self, model_stage: str | None, model_arch: str | None) -> bool: - if model_arch is not None and model_arch in self.model_archs: - return True - return model_stage is not None and model_stage in self.stage_keys - - # Re-exported here to avoid import cycles at call sites. __all__ = [ "ARTTSAdapter", "DiffusionTTSAdapter", - "LegacyDetector", "OutputPolicy", "PreparedRequest", "SpeechServingContext", + "TTSGenerationError", "TTSModelAdapter", ] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/capabilities.py b/vllm_omni/entrypoints/openai/tts_adapters/capabilities.py new file mode 100644 index 000000000..310161d82 --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/capabilities.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared loading mechanisms for TTS model capability metadata.""" + +import json +import os +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch +from transformers.utils.hub import cached_file +from vllm.logger import init_logger + +from vllm_omni.utils.speaker_cache import iter_custom_voice_profiles, load_validated_profile_tensors + +logger = init_logger(__name__) + + +def load_supported_speakers(engine_client: Any, config: Any = None) -> set[str]: + """Extract supported speaker names from a model configuration.""" + try: + if config is None: + config = getattr(engine_client.model_config.hf_config, "talker_config", None) + if config is None: + return set() + + for attr_name in ("spk_id", "speaker_id"): + speakers = config.get(attr_name) if isinstance(config, dict) else getattr(config, attr_name, None) + if speakers and isinstance(speakers, dict): + return {speaker.lower() for speaker in speakers} + + logger.warning("No speakers found in config (checked spk_id and speaker_id)") + except Exception as exc: + logger.warning("Could not load speakers from model config: %s", exc) + return set() + + +def load_codec_frame_rate(engine_client: Any) -> float | None: + """Load codec frame rate from speech-tokenizer or HF configuration.""" + try: + model_path = engine_client.model_config.model + config_path = os.path.join(model_path, "speech_tokenizer", "config.json") + if not os.path.exists(config_path): + config_path = cached_file(model_path, "speech_tokenizer/config.json") + if config_path is not None and os.path.exists(config_path): + with open(config_path) as file: + config = json.load(file) + output_sr = config.get("output_sample_rate") + downsample = config.get("encode_downsample_rate") + if output_sr and downsample and downsample > 0: + rate = float(output_sr) / float(downsample) + logger.info( + "Loaded codec frame rate: %.1f Hz (output_sample_rate=%s, encode_downsample_rate=%s)", + rate, + output_sr, + downsample, + ) + return rate + except Exception as exc: + logger.warning("Failed to load codec frame rate from speech tokenizer config: %s", exc) + + try: + rate = getattr(engine_client.model_config.hf_config, "codec_frame_rate_hz", None) + if rate is not None: + logger.info("Using codec frame rate from hf_config: %s Hz", rate) + return float(rate) + except Exception: + pass + return None + + +def load_precomputed_speakers( + engine_client: Any, + *, + expected_model_type: str, + validate_profile: Callable[[dict[str, Any], dict[str, torch.Tensor]], str | None], +) -> dict[str, dict[str, Any]]: + """Load and validate precomputed profiles from ``custom_voice_dir``.""" + try: + custom_voice_dir = getattr(engine_client.model_config.hf_config, "custom_voice_dir", None) + except AttributeError: + return {} + if isinstance(custom_voice_dir, os.PathLike): + custom_voice_dir = os.fspath(custom_voice_dir) + if not isinstance(custom_voice_dir, str) or not custom_voice_dir: + return {} + + profiles: dict[str, dict[str, Any]] = {} + for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type=expected_model_type): + tensors = load_validated_profile_tensors( + profile, + expected_model_type=expected_model_type, + validate_profile=validate_profile, + ) + if tensors is not None: + profiles[profile["voice_name_lower"]] = profile + if profiles: + logger.info( + "Loaded %d precomputed %s voice profile(s) from %s", + len(profiles), + expected_model_type, + Path(custom_voice_dir).expanduser(), + ) + return profiles diff --git a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py index 85648442f..ec1c145c7 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 """CosyVoice3 TTS serving adapter.""" -from typing import TYPE_CHECKING +import os +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest @@ -9,6 +12,8 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, Prepare if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +logger = init_logger(__name__) + @register_tts_adapter class CosyVoice3Adapter(ARTTSAdapter): @@ -47,6 +52,63 @@ class CosyVoice3Adapter(ARTTSAdapter): self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: prompt = await self.ctx.server._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio) - # NOTE: CosyVoice3 dynamic-token sampling stays in the orchestrator tail - # (keyed on _tts_model_type) during this incremental migration. return PreparedRequest(prompt=prompt, tts_params={}, model_type="cosyvoice3") + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + """Set dynamic min/max tokens based on tokenized text length. + + The official model uses ``min_token_text_ratio`` to prevent early EOS + and ``max_token_text_ratio`` to cap the generated token count. + """ + import copy + + from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer + from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token + + server = self.ctx.server + sampling_params_list = copy.deepcopy(sampling_params_list) + hf_cfg = server.model_config.hf_config + # Build the Qwen tokenizer once per process (resolving the model dir via + # snapshot_download at most once) and reuse it across requests. + tokenizer = server._cosyvoice3_tokenizer + if tokenizer is None: + model_path = server.engine_client.model_config.model + if not os.path.isdir(model_path): + from huggingface_hub import snapshot_download + + model_path = snapshot_download(model_path) + tokenizer = get_qwen_tokenizer( + token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path), + skip_special_tokens=hf_cfg.skip_special_tokens, + version=hf_cfg.version, + ) + server._cosyvoice3_tokenizer = tokenizer + _, text_token_len = extract_text_token( + request.input, + tokenizer, + hf_cfg.allowed_special, + ) + min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2) + max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20) + sampling_params_list[0].min_tokens = max(1, int(text_token_len * min_ratio)) + sampling_params_list[0].max_tokens = min(2048, int(text_token_len * max_ratio)) + + if request.max_new_tokens is not None: + sampling_params_list[0].max_tokens = request.max_new_tokens + sampling_params_list[0].min_tokens = min( + getattr(sampling_params_list[0], "min_tokens", 0), + request.max_new_tokens, + ) + logger.info( + "CosyVoice3 dynamic tokens: text_tokens=%d, min_tokens=%d, max_tokens=%d", + text_token_len, + sampling_params_list[0].min_tokens, + sampling_params_list[0].max_tokens, + ) + return sampling_params_list diff --git a/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py b/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py index e8080a427..a6426a13d 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py @@ -1,10 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 """Fish Speech serving adapter (retires the legacy ``_is_fish_speech`` flag).""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + conditioning_cache_salt, +) if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -44,12 +49,22 @@ class FishSpeechAdapter(ARTTSAdapter): ) -> PreparedRequest: server = self.ctx.server ref_audio_data = None + tts_params: dict = {} if request.ref_audio is not None: - wav_list, sr = await server._resolve_ref_audio(request.ref_audio) + wav_list, sr, cache_key = await server._resolve_ref_audio(request.ref_audio) ref_audio_data = (wav_list, sr) + tts_params["ref_audio_cache_key"] = cache_key prompt = await server._build_fish_speech_prompt_async( request, ref_audio_data=ref_audio_data, has_inline_ref_audio=has_inline_ref_audio ) - tts_params = {} prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="fish_speech") + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py index ddfd194e2..e65edb01b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """GLM-TTS serving adapter.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any + +from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest @@ -9,6 +11,8 @@ from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, Prepare if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +logger = init_logger(__name__) + @register_tts_adapter class GlmTTSAdapter(ARTTSAdapter): @@ -43,6 +47,66 @@ class GlmTTSAdapter(ARTTSAdapter): self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: prompt = await self.ctx.server._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio) - # GLM-TTS dynamic-token sampling stays in the orchestrator tail - # (keyed on _tts_model_type) during this incremental migration. return PreparedRequest(prompt=prompt, tts_params={}, model_type="glm_tts") + + def _load_supported_speakers(self) -> set[str]: + return set() + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + # GLM-TTS: set dynamic min/max tokens based on text length. + import copy + + server = self.ctx.server + sampling_params_list = copy.deepcopy(sampling_params_list) + glm_metadata = prompt.get("additional_information") if isinstance(prompt, dict) else None + text_len_value = None + if isinstance(glm_metadata, dict): + text_len_value = glm_metadata.get("glm_tts_text_token_len") + if isinstance(text_len_value, list) and text_len_value: + text_len_value = text_len_value[0] + text_token_len = ( + int(text_len_value) + if text_len_value is not None + else server._estimate_glm_tts_text_token_len(request.input) + ) + hf_cfg = server.model_config.hf_config + min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2) + max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20) + stage_min_tokens = getattr(sampling_params_list[0], "min_tokens", None) + stage_max_tokens = getattr(sampling_params_list[0], "max_tokens", None) + cap_candidates = [int(cap) for cap in (stage_max_tokens, request.max_new_tokens) if cap is not None] + hard_cap = min(cap_candidates) if cap_candidates else None + + min_tokens = max(1, int(text_token_len * min_ratio)) + if stage_min_tokens is not None: + min_tokens = max(min_tokens, int(stage_min_tokens)) + if hard_cap is not None: + min_tokens = min(min_tokens, hard_cap) + + max_tokens = max(min_tokens, int(text_token_len * max_ratio)) + if hard_cap is not None: + max_tokens = min(max_tokens, hard_cap) + sampling_params_list[0].min_tokens = min_tokens + sampling_params_list[0].max_tokens = max_tokens + seed = getattr(request, "seed", None) + if seed is not None: + sampling_params_list[0].seed = seed + logger.info( + "GLM-TTS dynamic tokens: text_tokens=%d, min_ratio=%s, max_ratio=%s, " + "stage_min=%s, stage_max=%s, request_max=%s, min_tokens=%d, max_tokens=%d", + text_token_len, + min_ratio, + max_ratio, + stage_min_tokens, + stage_max_tokens, + request.max_new_tokens, + min_tokens, + max_tokens, + ) + return sampling_params_list diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py index 99949d8b3..74ec591cc 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 """Higgs-Audio v2 serving adapter.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -99,3 +99,12 @@ class HiggsAudioV2Adapter(ARTTSAdapter): additional["voice_name"] = voice_lower additional["voice_created_at"] = server._voice_created_at(voice_lower) return PreparedRequest(prompt=prompt, tts_params={}, model_type="higgs_audio_v2") + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py index 9cae9aa2d..411d97f31 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 """Higgs-Audio v3 serving adapter.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -36,3 +36,12 @@ class HiggsAudioV3Adapter(ARTTSAdapter): ) -> PreparedRequest: prompt = await self.ctx.server._build_higgs_audio_v3_params(request) return PreparedRequest(prompt=prompt, tts_params={}, model_type="higgs_audio_v3") + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py b/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py index d1cb2a469..500287e1c 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py @@ -14,7 +14,7 @@ from vllm.inputs import tokens_input from vllm.utils import random_uuid from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens from vllm_omni.model_executor.models.indextts2.configuration_indextts2 import ( INDEXTTS25_MAX_DURATION_FACTOR, INDEXTTS25_MIN_DURATION_FACTOR, @@ -108,6 +108,8 @@ def indextts2_conditioning_cache_salt( "use_random", "lang", "text_normalization", + "ref_audio_cache_key", + "emo_audio_cache_key", ): h.update(b"\x00") h.update(key.encode("utf-8")) @@ -241,6 +243,15 @@ class IndexTTS2Adapter(ARTTSAdapter): prompt["cache_salt"] = indextts2_conditioning_cache_salt(request, tts_params) return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: OpenAICreateSpeechRequest, + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) + async def _build_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: server = self.ctx.server params: dict[str, Any] = {"text": [request.input]} @@ -251,8 +262,9 @@ class IndexTTS2Adapter(ARTTSAdapter): ref_audio_source = server._get_uploaded_audio_data(voice_lower) using_uploaded_voice = ref_audio_source is not None if ref_audio_source is not None and isinstance(ref_audio_source, str): - wav_list, sr = await server._resolve_ref_audio(ref_audio_source) + wav_list, sr, cache_key = await server._resolve_ref_audio(ref_audio_source) params["voice"] = [[wav_list, sr]] + params["ref_audio_cache_key"] = [cache_key] if using_uploaded_voice and voice_lower: params["voice_name"] = [voice_lower] params["voice_created_at"] = [server._voice_created_at(voice_lower)] @@ -262,8 +274,9 @@ class IndexTTS2Adapter(ARTTSAdapter): if key not in extras: continue if key == "emo_audio": - wav_list, sr = await server._resolve_ref_audio(extras[key]) + wav_list, sr, emo_key = await server._resolve_ref_audio(extras[key]) params[key] = [[wav_list, sr]] + params["emo_audio_cache_key"] = [emo_key] else: params[key] = [extras[key]] return params diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py index f351e7531..bc084d6b4 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_flash_omni_tts.py @@ -16,6 +16,9 @@ logger = init_logger(__name__) @register_tts_adapter class MingFlashOmniTTSAdapter(ARTTSAdapter): + # Ming-flash-omni drives speaker selection via the caption JSON + # (audio_sequence[0]["说话人"]) rather than a spk_id table, so there + # is no static speaker list to surface here. stage_keys = frozenset({"ming_tts"}) name = "ming_flash_omni_tts" @@ -61,3 +64,7 @@ class MingFlashOmniTTSAdapter(ARTTSAdapter): server = self.ctx.server prompt = server._build_ming_flash_omni_prompt(request) return PreparedRequest(prompt=prompt, tts_params={}, model_type="ming_flash_omni_tts") + + def _load_supported_speakers(self) -> set[str]: + # Speaker selection is driven by caption JSON rather than a static table. + return set() diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py index dd8b8b621..613968b8c 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 """Ming-TTS (dense) serving adapter.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm.logger import init_logger @@ -64,7 +64,7 @@ class MingTTSAdapter(ARTTSAdapter): if isinstance(ref_audio_source, list): ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source) elif ref_audio_source is not None and isinstance(ref_audio_source, str): - wav_list, sr = await server._resolve_ref_audio(ref_audio_source) + wav_list, sr, _ = await server._resolve_ref_audio(ref_audio_source) ref_audio_data = (wav_list, sr) prompt = server._build_ming_dense_prompt( request, @@ -73,9 +73,36 @@ class MingTTSAdapter(ARTTSAdapter): voice_created_at=uploaded_audio_created_at, ) tts_params = prompt.get("additional_information", {}) - # Ming stop-token / max_tokens sampling stays in the orchestrator tail. return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="ming_tts") + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + import copy + + server = self.ctx.server + + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + MOE_TEXT_EOS_TOKEN_ID, + TEXT_EOS_TOKEN_ID, + ) + + hf_config = server.engine_client.model_config.hf_config + is_moe = getattr(hf_config, "model_type", "") == "bailingmm" + stop_token_id = MOE_TEXT_EOS_TOKEN_ID if is_moe else TEXT_EOS_TOKEN_ID + + sampling_params_list = copy.deepcopy(sampling_params_list) + sampling_params_list[0].stop_token_ids = [int(stop_token_id)] + if request.max_new_tokens is not None: + # Ming emits TEXT_EOS after the latent decode budget is exhausted, so + # Stage-0 needs one extra token beyond ming_max_decode_steps. + sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1 + return sampling_params_list + def _validate_ming_tts_single_speaker_request(self, request: "OpenAICreateSpeechRequest") -> str | None: server = self.ctx.server if request.ref_audio is not None: @@ -157,3 +184,44 @@ class MingTTSAdapter(ARTTSAdapter): return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" return None + + def validate_tts_embedding_dim(self, emb_dim: int) -> str | None: + if emb_dim != SPEAKER_EMBEDDING_DIM: + return f"Ming speaker embedding must have {SPEAKER_EMBEDDING_DIM} dims, got {emb_dim}" + return None + + def _load_ming_tts_codec_frame_rate(self) -> float | None: + server = self.ctx.server + try: + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingTTSConfig + + hf_config = server.engine_client.model_config.hf_config + ming_cfg = MingTTSConfig.from_hf_config(hf_config) + patch_size = int(ming_cfg.patch_size) + audio_frame_hop = int(ming_cfg.audio_frame_hop) + sample_rate = int(ming_cfg.sample_rate) + if patch_size <= 0 or audio_frame_hop <= 0 or sample_rate <= 0: + raise ValueError( + "Ming config has invalid tokenizer timing values: " + f"patch_size={patch_size}, audio_frame_hop={audio_frame_hop}, sample_rate={sample_rate}" + ) + rate = float(sample_rate) / float(audio_frame_hop * patch_size) + logger.info( + "Derived Ming codec frame rate: %.1f Hz (sample_rate=%s, audio_frame_hop=%s, patch_size=%s)", + rate, + sample_rate, + audio_frame_hop, + patch_size, + ) + return rate + except Exception as e: + logger.warning(f"Failed to derive Ming codec frame rate from hf_config: {e}") + + def _load_supported_speakers(self) -> set[str]: + return set() + + def _load_codec_frame_rate(self) -> float | None: + codec_frame_rate = self._load_ming_tts_codec_frame_rate() + if codec_frame_rate is None: + codec_frame_rate = super()._load_codec_frame_rate() + return codec_frame_rate diff --git a/vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py b/vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py new file mode 100644 index 000000000..50b9c2d71 --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax Music 3 serving adapter for ``/v1/audio/speech``. + +This is a text-to-music model on a speech endpoint, so most of the speech +contract does not apply. ``input`` carries the lyrics and ``instructions`` the +caption that describes genre, instrumentation, tempo and production. There is +no speaker to select, no reference audio, and no temperature: sampling is +fixed by the checkpoint. Unsupported parameters are rejected rather than +silently ignored, so a mistaken request fails immediately instead of quietly +returning something the caller did not ask for. +""" + +from typing import TYPE_CHECKING, Any + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AUDIO_FRAME_RATE, + MAX_AUDIO_FRAMES, + MAX_PROMPT_TOKENS, +) +from vllm_omni.model_executor.models.minimax_music3.prompt import ( + build_prompt, + validate_tokenizer_ids, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + +# Captions are long by design: the checkpoint was trained on structured +# descriptions covering global attributes, emotional progression, vocal detail +# and arrangement. The reference captions run to roughly 5,000 characters, so +# the generic 500-character instructions cap does not apply here. The real +# limit is the tokenized prompt cap, enforced downstream. +_MAX_CAPTION_CHARS = 20_000 + +# Sampling is fixed by the checkpoint: guidance at 1.5, then a seeded top-k 50 +# draw. Accepting these would imply a control the model does not have. +_UNSUPPORTED_SAMPLING_PARAMS = ("temperature", "top_p", "top_k", "repetition_penalty") + +# CFG pair plumbing is derived from the request id by the scheduler. A caller +# supplying these would desync the guided pair. +_INTERNAL_CFG_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt") + + +@register_tts_adapter +class MiniMaxMusic3Adapter(ARTTSAdapter): + """Lyrics plus caption in, one 32 kHz stereo song out.""" + + name = "minimax_music3" + stage_keys = frozenset({"minimax_music3_ar"}) + # 9,000 frames at 25 fps is six minutes, the checkpoint's own ceiling. + max_new_tokens_min = 1 + max_new_tokens_max = MAX_AUDIO_FRAMES + + def __init__(self, ctx: Any) -> None: + super().__init__(ctx) + self._cached_tokenizer: Any = None + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + if not request.input or not request.input.strip(): + return "MiniMax Music 3 requires non-empty 'input' carrying the lyrics" + + instructions = request.instructions + if not isinstance(instructions, str) or not instructions.strip(): + return ( + "MiniMax Music 3 requires 'instructions' describing the music " + "(genre, instrumentation, tempo, mood); it is what decides the arrangement" + ) + if len(instructions) > _MAX_CAPTION_CHARS: + return f"'instructions' exceeds the maximum length of {_MAX_CAPTION_CHARS} characters" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"'max_new_tokens' must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return ( + f"'max_new_tokens' counts audio frames at {AUDIO_FRAME_RATE} per second " + f"and cannot exceed {self.max_new_tokens_max}" + ) + + voice = (request.voice or "").strip().lower() + if voice not in ("", "default"): + return ( + "MiniMax Music 3 has no speaker to select; the vocal comes from " + f"'instructions'. Omit 'voice' or pass 'default'; got {request.voice!r}" + ) + if request.ref_audio is not None or request.ref_text is not None: + return "MiniMax Music 3 does not support reference-audio conditioning" + if request.language is not None: + return "MiniMax Music 3 has no language tag; the language follows the lyrics" + if request.task_type is not None: + return "MiniMax Music 3 does not support 'task_type'" + if request.speed is not None and float(request.speed) != 1.0: + return "MiniMax Music 3 only supports speed=1.0; put the tempo in 'instructions', e.g. 'at 92 BPM'" + + extra: dict[str, Any] = request.extra_params or {} + for key in _INTERNAL_CFG_KEYS: + if key in extra: + return f"extra_params.{key} is managed internally by the server and cannot be set" + rejected = sorted(key for key in _UNSUPPORTED_SAMPLING_PARAMS if extra.get(key) is not None) + if rejected: + return "MiniMax Music 3 has fixed sampling and does not accept: " + ", ".join(rejected) + return None + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del sampling_params_list, has_inline_ref_audio + prompt = build_prompt(request.instructions, request.input) + # Submit token ids, not text. Guidance needs a companion whose prompt is + # the same length with the caption-and-lyrics span replaced, and that + # length match can only be guaranteed on ids. Handing over text would + # leave the expansion with nothing to build the twin from, and a request + # that decodes unguided produces a different song. + token_ids = self._tokenizer().encode(prompt) + if len(token_ids) > MAX_PROMPT_TOKENS: + raise ValueError( + f"MiniMax Music 3 prompt is {len(token_ids)} tokens; " + f"the maximum is {MAX_PROMPT_TOKENS}. Shorten the caption or the lyrics." + ) + # The frame budget is enforced by the model, not by the engine's + # max_tokens. The model needs one decode step beyond the last frame to + # deliver the song's final window, so capping the engine at exactly + # this many tokens would cut the tail off. + max_frames = int(request.max_new_tokens or MAX_AUDIO_FRAMES) + # Carried on the prompt rather than through tts_params: the adapter + # path hands PreparedRequest.prompt to the engine untouched, so only + # what is attached here reaches the model's per-request info. + return PreparedRequest( + prompt={ + "prompt_token_ids": token_ids, + "additional_information": {"max_audio_frames": [max_frames]}, + }, + tts_params={"max_audio_frames": [max_frames]}, + model_type=self.name, + ) + + def _tokenizer(self) -> Any: + """The music tokenizer, loaded once per process. + + Its special-token ids are pinned: a checkpoint whose tokenizer + disagrees is not this model, and quietly producing noise is worse than + failing at the first request. + """ + if self._cached_tokenizer is None: + from transformers import AutoTokenizer + + stage = getattr(self.ctx.server, "_tts_stage", None) + model_path = getattr(getattr(stage, "engine_args", None), "tokenizer", None) + if not model_path: + model_path = self.ctx.server.engine_client.model_config.tokenizer + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + validate_tokenizer_ids(tokenizer) + self._cached_tokenizer = tokenizer + return self._cached_tokenizer + + +__all__ = ["MiniMaxMusic3Adapter"] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index ad43a1c01..3d6e99552 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MOSS-TTS serving adapters (Nano + full family). Both variants share the same build/validate flow (``_build_moss_tts_params`` @@ -10,7 +11,12 @@ from typing import TYPE_CHECKING from vllm.inputs import tokens_input from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + conditioning_cache_salt, +) if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -86,15 +92,19 @@ class _MossTTSAdapterBase(ARTTSAdapter): self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: server = self.ctx.server - tts_params = await server._build_moss_tts_params(request) + tts_params = await server._build_moss_tts_params(request, has_inline_ref_audio=has_inline_ref_audio) if request.voice: voice_lower = request.voice.lower() if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: tts_params["voice_name"] = [voice_lower] tts_params["voice_created_at"] = [server._voice_created_at(voice_lower)] - # MOSS reads the resolved seed at build time (it samples internally). - if sampling_params_list and getattr(sampling_params_list[0], "seed", None) is not None: - tts_params["seed"] = [sampling_params_list[0].seed] + # MOSS samples internally from additional_information. build() runs + # before the shared path applies request.seed to SamplingParams. + seed = request.seed + if seed is None and sampling_params_list: + seed = getattr(sampling_params_list[0], "seed", None) + if seed is not None: + tts_params["seed"] = [int(seed)] if isinstance(tts_params.get("prompt_token_ids"), list): prompt_token_ids = tts_params.pop("prompt_token_ids") prompt = tokens_input(prompt_token_ids=prompt_token_ids) @@ -104,6 +114,15 @@ class _MossTTSAdapterBase(ARTTSAdapter): prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) + @register_tts_adapter class MossTTSNanoAdapter(_MossTTSAdapterBase): diff --git a/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py b/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py index aa3ddd3b4..deafab685 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py @@ -33,7 +33,7 @@ class OmniVoiceAdapter(ARTTSAdapter): server = self.ctx.server prompt: dict = {"input": request.input} if request.ref_audio: - wav, sr = await server._resolve_ref_audio(request.ref_audio) + wav, sr, _ = await server._resolve_ref_audio(request.ref_audio) prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr) if request.ref_text: prompt["ref_text"] = request.ref_text diff --git a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py index c0be8c685..c2fd4faf8 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py @@ -1,24 +1,45 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Qwen3-TTS serving adapter.""" +from collections.abc import Mapping from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + DEFAULT_TTS_LANGUAGES, + ARTTSAdapter, + PreparedRequest, + TTSGenerationError, +) +from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_precomputed_speakers +from vllm_omni.utils.speaker_cache import validate_qwen3_tts_profile if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest logger = init_logger(__name__) +QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY = "_qwen3_tts_effective_max_tokens" +_MIN_CODEC_FRAMES = 192 +_MAX_CODEC_FRAMES_PER_TEXT_TOKEN = 12 + + +class Qwen3TTSCodecLimitError(TTSGenerationError): + """Qwen3-TTS Base exhausted its codec budget without emitting EOS.""" + + def __init__(self, message: str) -> None: + super().__init__(message, retryable=True) + @register_tts_adapter class Qwen3TTSAdapter(ARTTSAdapter): """Adapter for Qwen3-TTS (AR ``engine_client`` backend).""" + validates_generation = True stage_keys = frozenset({"qwen3_tts"}) name = "qwen3_tts" @@ -37,7 +58,7 @@ class Qwen3TTSAdapter(ARTTSAdapter): # Normalize voice to lowercase for case-insensitive matching if request.voice is not None: request.voice = request.voice.lower() - if request.task_type is None and request.voice in server.precomputed_speakers: + if request.task_type is None and request.voice in self.capabilities.precomputed_speakers: request.task_type = "Base" task_type = request.task_type or "CustomVoice" @@ -48,21 +69,23 @@ class Qwen3TTSAdapter(ARTTSAdapter): # Validate language (case-insensitive; normalized to the title-cased config form) if request.language is not None: request.language = request.language.title() - if request.language not in server.supported_languages: + if request.language not in self.capabilities.supported_languages: return ( - f"Invalid language '{request.language}'. Supported: {', '.join(sorted(server.supported_languages))}" + f"Invalid language '{request.language}'. Supported: " + f"{', '.join(sorted(self.capabilities.supported_languages))}" ) # Validate speaker for CustomVoice task if task_type == "CustomVoice": - if not server.supported_speakers: + available_speakers = server._get_available_speakers() + if not available_speakers: return ( "This model does not support CustomVoice task (no speakers configured). " "Use task_type='Base' with ref_audio/ref_text for voice cloning, " "or use a CustomVoice model." ) - if request.voice is not None and request.voice not in server.supported_speakers: - return f"Invalid voice '{request.voice}'. Supported: {', '.join(sorted(server.supported_speakers))}" + if request.voice is not None and request.voice not in available_speakers: + return f"Invalid voice '{request.voice}'. Supported: {', '.join(sorted(available_speakers))}" # Validate speaker_embedding constraints if request.speaker_embedding is not None: @@ -74,7 +97,7 @@ class Qwen3TTSAdapter(ARTTSAdapter): # Base task validation so callers don't need to pass it explicitly. request.x_vector_only_mode = True emb_len = len(request.speaker_embedding) - dim_err = server._validate_qwen_tts_speaker_embedding_dim(emb_len) + dim_err = self.validate_tts_embedding_dim(emb_len) if dim_err is not None: return dim_err # Validate Base task requirements @@ -103,8 +126,8 @@ class Qwen3TTSAdapter(ARTTSAdapter): file_path = Path(speaker_info["file_path"]) if not file_path.exists(): return f"Data file for uploaded speaker '{request.voice}' not found on disk" - elif voice_lower in server.precomputed_speakers: - profile = server.precomputed_speakers[voice_lower] + elif voice_lower in self.capabilities.precomputed_speakers: + profile = self.capabilities.precomputed_speakers[voice_lower] mode = str(profile.get("mode") or "xvec").lower() ref_text = request.ref_text or profile.get("ref_text") if mode == "icl" and (not isinstance(ref_text, str) or not ref_text.strip()): @@ -163,3 +186,148 @@ class Qwen3TTSAdapter(ARTTSAdapter): model_type=tts_params.get("task_type", ["unknown"])[0], warmup_artifact_key=warmup_key, ) + + def _get_expected_speaker_embedding_dim(self) -> int: + """Return the loaded Qwen3-TTS speaker embedding dim, if known. + + The user-provided speaker embedding is concatenated directly with + talker codec embeddings, so the real compatibility requirement is the + talker hidden size. + """ + hf_config = self.ctx.engine_client.model_config.hf_config + talker_config = hf_config.talker_config + return int(talker_config.hidden_size) + + def _load_precomputed_speakers(self) -> dict[str, dict]: + return load_precomputed_speakers( + self.ctx.engine_client, + expected_model_type=self.name, + validate_profile=lambda profile, tensors: validate_qwen3_tts_profile( + profile, + tensors, + expected_embedding_dim=self._get_expected_speaker_embedding_dim(), + ), + ) + + def _load_supported_languages(self) -> frozenset[str]: + try: + config = self.ctx.engine_client.model_config.hf_config.talker_config + + if isinstance(config, dict): + codec_language_id = config.get("codec_language_id") + else: + codec_language_id = getattr(config, "codec_language_id", None) + + if codec_language_id and isinstance(codec_language_id, Mapping): + return frozenset(str(language).title() for language in codec_language_id) | {"Auto"} + + logger.warning("No codec_language_id found in talker_config; falling back to default languages") + except Exception as exc: + logger.warning("Could not load languages from model config: %s", exc) + return DEFAULT_TTS_LANGUAGES + + def validate_tts_embedding_dim(self, emb_dim: int) -> str | None: + expected_dim = self._get_expected_speaker_embedding_dim() + if emb_dim != expected_dim: + return f"speaker_embedding has {emb_dim} dimensions; expected {expected_dim} for the loaded Qwen3-TTS model" + return None + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + """Apply a text-scaled safety ceiling to Base codec generation. + + Qwen3-TTS can rarely enter a repetitive state in which codec EOS is no + longer reachable through top-k sampling. A fixed 4096-frame ceiling + turns that into several minutes of unusable audio. Bound the default + Base-task budget by text length, while preserving an explicit caller + ``max_new_tokens`` override and the configured budget for other tasks. + """ + import copy + + del request_id + server = self.ctx.server + # Only scalar fields on stage 0 are changed below. Shallow-copy each + # stage so the shared defaults stay immutable without deep-copying the + # complete sampling configuration on every request. + sampling_params_list = [copy.copy(params) for params in sampling_params_list] + configured_cap = getattr(sampling_params_list[0], "max_tokens", None) + task_type = request.task_type or "CustomVoice" + text_tokens = None + dynamic_cap = None + effective_cap: int | None + + if request.max_new_tokens is not None: + # An explicit request budget is an opt-out from the automatic + # ceiling. It remains an upper bound, and a length finish is still + # surfaced as an incomplete generation rather than valid audio. + effective_cap = int(request.max_new_tokens) + elif task_type == "Base": + counted_text_tokens = server._count_usage_text_tokens(request.input) + if counted_text_tokens > 0: + text_tokens = counted_text_tokens + dynamic_cap = max(_MIN_CODEC_FRAMES, text_tokens * _MAX_CODEC_FRAMES_PER_TEXT_TOKEN) + effective_cap = min(dynamic_cap, int(configured_cap)) if configured_cap is not None else dynamic_cap + else: + # Token counting is best-effort. If the tokenizer is missing or + # rejects the input, preserve the configured budget instead of + # truncating an otherwise valid request at the minimum ceiling. + effective_cap = int(configured_cap) if configured_cap is not None else None + else: + effective_cap = int(configured_cap) if configured_cap is not None else None + + if effective_cap is not None: + effective_cap = max(1, effective_cap) + sampling_params_list[0].max_tokens = effective_cap + sampling_params_list[0].min_tokens = min( + int(getattr(sampling_params_list[0], "min_tokens", 0) or 0), + effective_cap, + ) + + if isinstance(prompt, dict): + additional_information = prompt.get("additional_information") + if isinstance(additional_information, dict) and effective_cap is not None: + additional_information[QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] = [effective_cap] + + logger.debug( + "Qwen3-TTS codec budget: task_type=%s text_tokens=%s dynamic_cap=%s " + "configured_cap=%s request_cap=%s effective_cap=%s", + task_type, + text_tokens, + dynamic_cap, + configured_cap, + request.max_new_tokens, + effective_cap, + ) + return sampling_params_list + + def validate_generation( + self, + tts_params: Mapping[str, object], + *, + stage0_finish_reason: str | None, + output_tokens: int, + ) -> None: + if QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY not in tts_params: + return + task_type = tts_params.get("task_type") + if isinstance(task_type, (list, tuple)): + task_type = task_type[0] if task_type else None + if task_type != "Base" or stage0_finish_reason != "length": + return + + raw_limit = tts_params.get(QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY) + if isinstance(raw_limit, (list, tuple)): + raw_limit = raw_limit[0] if raw_limit else None + try: + limit = int(raw_limit) if isinstance(raw_limit, (str, bytes, bytearray, int, float)) else 0 + except (TypeError, ValueError): + limit = 0 + raise Qwen3TTSCodecLimitError( + "Qwen3-TTS Base did not emit codec EOS before its token budget " + f"({output_tokens}/{limit} codec tokens); the generated audio is incomplete." + ) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py index b0f188fbb..760bb3339 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py @@ -1,13 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 """VoxCPM2 serving adapter (AR base-LM + diffusion side-computation).""" -from typing import TYPE_CHECKING +import time +from typing import Any +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens +from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_precomputed_speakers +from vllm_omni.utils.speaker_cache import validate_voxcpm2_profile -if TYPE_CHECKING: - from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +logger = init_logger(__name__) @register_tts_adapter @@ -31,7 +36,7 @@ class VoxCPM2Adapter(ARTTSAdapter): if request.voice is not None: request.voice = request.voice.lower() - available_voices = set(server.uploaded_speakers) | set(server.precomputed_speakers) | {"default"} + available_voices = server._get_available_speakers() if request.voice not in available_voices: supported = ", ".join(sorted(available_voices)) or "none" return f"Invalid voice '{request.voice}'. Supported: {supported}" @@ -65,8 +70,59 @@ class VoxCPM2Adapter(ARTTSAdapter): tts_params = {} if request.voice: voice_lower = request.voice.lower() - if voice_lower in server.uploaded_speakers or voice_lower in server.precomputed_speakers: + if voice_lower in server.uploaded_speakers or voice_lower in self.capabilities.precomputed_speakers: additional = prompt.setdefault("additional_information", {}) additional["voice_name"] = voice_lower additional["voice_created_at"] = server._voice_created_at(voice_lower) return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="voxcpm2") + + async def warmup(self) -> None: + """Warm up VoxCPM2 through a synthetic serving request. + + VoxCPM2 needs to warm up its PagedAttention scaffold/residual LLMs. CUDA + Graph capture requires a vLLM ``ForwardContext`` containing attention + metadata and slot mappings, which is only available during inference. The + request also pays the one-time torch.compile cost for LocDiT, feat_encoder, + AudioVAE, and projection helpers. + """ + server = self.ctx.server + t0 = time.time() + logger.info("Running warmup speech request for model_type=%s", self.name) + # VoxCPM2 has no predefined speaker presets — "default" means zero-shot + # mode (no voice cloning). The voice field is required by the OpenAI + # API schema but semantically ignored by the model. + warmup_req = OpenAICreateSpeechRequest( + input="Warmup.", + voice="default", + response_format="wav", + speed=1.0, + stream=False, + model=server.model_name, + ) + try: + _audio_bytes, _media_type = await server._generate_audio_bytes(warmup_req, request_id="speech-warmup") + except Exception as exc: + logger.warning("Speech warmup failed (non-fatal): %s", exc) + return + + elapsed = time.time() - t0 + logger.info("Speech warmup complete in %.1fs", elapsed) + + def _load_precomputed_speakers(self) -> dict[str, dict]: + return load_precomputed_speakers( + self.ctx.engine_client, + expected_model_type=self.name, + validate_profile=validate_voxcpm2_profile, + ) + + def _load_supported_speakers(self) -> set[str]: + return {"default"} + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py index 0b2d400b1..c9bea7bb0 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 """Voxtral TTS serving adapter.""" -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter -from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, apply_max_new_tokens +from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -32,8 +33,9 @@ class VoxtralTTSAdapter(ARTTSAdapter): if request.voice is not None: request.voice = request.voice.lower() - if server.supported_speakers and request.voice not in server.supported_speakers: - return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(server.supported_speakers))}" + available_speakers = server._get_available_speakers() + if available_speakers and request.voice not in available_speakers: + return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(available_speakers))}" if request.max_new_tokens is not None: if request.max_new_tokens < self.max_new_tokens_min: @@ -48,3 +50,16 @@ class VoxtralTTSAdapter(ARTTSAdapter): ) -> PreparedRequest: prompt = await self.ctx.server._build_voxtral_prompt_async(request) return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts") + + def _load_supported_speakers(self) -> set[str]: + config = self.ctx.engine_client.model_config.hf_config.audio_config + return load_supported_speakers(self.ctx.engine_client, config) + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + return apply_max_new_tokens(sampling_params_list, request) diff --git a/vllm_omni/entrypoints/openai/utils.py b/vllm_omni/entrypoints/openai/utils.py index 5cca49129..4aac073d4 100644 --- a/vllm_omni/entrypoints/openai/utils.py +++ b/vllm_omni/entrypoints/openai/utils.py @@ -20,6 +20,27 @@ def get_stage_type(stage_cfg: Any) -> str: return getattr(stage_cfg, "stage_type", "llm") +def is_video_generation_pipeline(stage_configs: list[Any] | None) -> bool: + """Return whether a pipeline declares a final video output stage.""" + for stage in stage_configs or (): + if isinstance(stage, dict): + final_output = stage.get("final_output", False) + final_output_type = stage.get("final_output_type") + elif hasattr(stage, "get"): + try: + final_output = stage.get("final_output", False) + final_output_type = stage.get("final_output_type") + except Exception: + final_output = getattr(stage, "final_output", False) + final_output_type = getattr(stage, "final_output_type", None) + else: + final_output = getattr(stage, "final_output", False) + final_output_type = getattr(stage, "final_output_type", None) + if final_output and final_output_type in {"video", "videos"}: + return True + return False + + def parse_lora_request(lora_body: Any) -> tuple[LoRARequest | None, float | None]: """Parse a request-level LoRA object into a LoRARequest and optional scale. diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 5a4f9f826..34bc23384 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Shared helper utilities for OpenAI-compatible video generation API. """ @@ -10,13 +10,19 @@ import base64 import binascii import os import tempfile +import threading +from collections import deque +from collections.abc import Generator +from concurrent.futures import Future, ThreadPoolExecutor from io import BytesIO -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal, TypeAlias import httpx import numpy as np import torch from PIL import Image, UnidentifiedImageError +from vllm import envs +from vllm.logger import init_logger from vllm.multimodal.video import ( VIDEO_LOADER_REGISTRY, VideoBackend, @@ -34,6 +40,20 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoReference, ) +if TYPE_CHECKING: + import av + + +logger = init_logger(__name__) + + +DEFAULT_AUDIO_SAMPLE_RATE = 24_000 + + +VideoInput: TypeAlias = torch.Tensor | np.ndarray | list[torch.Tensor | np.ndarray | Image.Image] +AudioSample = int | float +AudioInput: TypeAlias = torch.Tensor | np.ndarray | list[AudioSample] | list[list[AudioSample]] + class VideoFrames(list[Image.Image]): """Decoded video frames plus source metadata.""" @@ -172,11 +192,21 @@ async def decode_image_url(image_url: str) -> Image.Image: return _decode_base64_image(image_url, source="image_reference.image_url") if image_url.startswith(("http://", "https://")): - async with httpx.AsyncClient(timeout=60) as client: + allow_redirects = envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS + async with httpx.AsyncClient(timeout=60, follow_redirects=allow_redirects) as client: try: response = await client.get(image_url) response.raise_for_status() - except httpx.HTTPError as exc: + except httpx.HTTPStatusError as exc: + if exc.response.has_redirect_location and not allow_redirects: + raise InvalidInputReferenceError( + "Invalid image_reference.image_url: redirect response was rejected because " + "VLLM_MEDIA_URL_ALLOW_REDIRECTS is disabled." + ) from exc + raise InvalidInputReferenceError( + f"Invalid image_reference.image_url: server returned HTTP {exc.response.status_code}." + ) from exc + except httpx.RequestError as exc: raise InvalidInputReferenceError( "Invalid image_reference.image_url: failed to download image." ) from exc @@ -366,7 +396,7 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray: # Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no # numpy dtype, so ``.numpy()`` below raises on it. video_tensor = video_tensor.float().clamp(-1, 1) * 0.5 + 0.5 - else: + elif video_tensor.dtype != torch.uint8: video_tensor = video_tensor.to(torch.float32) / 255.0 video_array = video_tensor.numpy() return _normalize_single_video_array(video_array) @@ -387,9 +417,9 @@ def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray: video_array = np.transpose(video_array, (0, 2, 3, 1)) if np.issubdtype(video_array.dtype, np.floating): - if video_array.min() < 0.0 or video_array.max() > 1.0: + if video_array.size and (video_array.min() < 0.0 or video_array.max() > 1.0): video_array = np.clip(video_array, -1.0, 1.0) * 0.5 + 0.5 - elif np.issubdtype(video_array.dtype, np.integer): + elif video_array.dtype != np.uint8 and np.issubdtype(video_array.dtype, np.integer): video_array = video_array.astype(np.float32) / 255.0 return video_array @@ -405,7 +435,7 @@ def _normalize_video_array(video_array: np.ndarray) -> list[np.ndarray] | np.nda def _normalize_frames(frames: list[Any]) -> list[np.ndarray]: - """Normalize a list of frames into numpy arrays with values in [0,1].""" + """Normalize a list of frames into numpy arrays, uint8 ones unchanged.""" normalized: list[np.ndarray] = [] for frame in frames: if isinstance(frame, torch.Tensor): @@ -421,9 +451,9 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]: frame_array = np.transpose(frame_array, (1, 2, 0)) if np.issubdtype(frame_array.dtype, np.floating): - if frame_array.min() < 0.0 or frame_array.max() > 1.0: + if frame_array.size and (frame_array.min() < 0.0 or frame_array.max() > 1.0): frame_array = np.clip(frame_array, -1.0, 1.0) * 0.5 + 0.5 - elif np.issubdtype(frame_array.dtype, np.integer): + elif frame_array.dtype != np.uint8 and np.issubdtype(frame_array.dtype, np.integer): frame_array = frame_array.astype(np.float32) / 255.0 normalized.append(frame_array) @@ -431,7 +461,14 @@ def _normalize_frames(frames: list[Any]) -> list[np.ndarray]: def _coerce_video_to_frames(video: Any) -> list[np.ndarray]: - """Convert a video payload into a list of normalized float32 frames.""" + """Convert a video payload into a list of normalized frames. + + Frames are float32 in [0, 1], except uint8 ones, which pass through: the + muxer's own dtype is uint8, so normalising here would only pay for a + full-size conversion each way. The direct planar path additionally needs + contiguous channel planes and falls back for interleaved RGB whatever the + dtype; the standard muxer takes the uint8 frames as they are. + """ if isinstance(video, torch.Tensor): video_array = _normalize_video_tensor(video) return list(video_array) @@ -475,17 +512,29 @@ def _coerce_audio_to_numpy(audio: Any) -> np.ndarray: return arr.astype(np.float32) -def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray: - """Convert a video payload into contiguous uint8 frames shaped (F, H, W, 3).""" +def _prepare_video_frames(video: Any) -> tuple[list[np.ndarray], tuple[int, ...], np.dtype]: + """Normalize and validate frames for the common video encoding dispatcher.""" frames = _coerce_video_to_frames(video) if not frames: raise ValueError("No frames found to encode.") frame_shape = frames[0].shape + if any(frame.shape != frame_shape for frame in frames[1:]): + raise ValueError("All video frames must have the same shape.") + + common_dtype = np.result_type(*(frame.dtype for frame in frames)) + return frames, frame_shape, common_dtype + + +def _coerce_prepared_video_to_uint8_frames( + frames: list[np.ndarray], + frame_shape: tuple[int, ...], + common_dtype: np.dtype, +) -> np.ndarray: + """Convert prepared frames into contiguous uint8 frames for the legacy muxer.""" has_alpha = len(frame_shape) == 3 and frame_shape[-1] == 4 output_shape = (*frame_shape[:-1], 3) if has_alpha else frame_shape frames_u8 = np.empty((len(frames), *output_shape), dtype=np.uint8) - common_dtype = np.result_type(*(frame.dtype for frame in frames)) # Convert one frame at a time instead of stacking the normalized float # payload first. Long videos can otherwise require another full-size @@ -509,27 +558,305 @@ def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray: return frames_u8 -def _encode_video_bytes( +def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray: + """Convert a video payload into contiguous uint8 frames shaped (F, H, W, 3).""" + if ( + isinstance(video, np.ndarray) + and video.dtype == np.uint8 + and video.ndim == 4 + and video.shape[-1] == 3 + and video.flags.c_contiguous + ): + return video + frames, frame_shape, common_dtype = _prepare_video_frames(video) + return _coerce_prepared_video_to_uint8_frames(frames, frame_shape, common_dtype) + + +def _direct_planar_fallback_reason( + frames: list[np.ndarray], + frame_shape: tuple[int, ...], + common_dtype: np.dtype, + *, + allow_strided_rgb_planes: bool = False, +) -> str | None: + """Return a stable reason when direct planar muxing cannot consume frames.""" + if len(frame_shape) != 3 or frame_shape[0] <= 0 or frame_shape[1] <= 0 or frame_shape[2] not in (3, 4): + return "unsupported_shape" + + if not ( + common_dtype == np.dtype(np.uint8) + or np.issubdtype(common_dtype, np.bool_) + or np.issubdtype(common_dtype, np.floating) + ): + return "unsupported_dtype" + + if not allow_strided_rgb_planes and not all( + frame[..., channel].flags.c_contiguous for frame in frames for channel in range(3) + ): + return "non_contiguous_rgb_planes" + + return None + + +class _PlanarFrameConverter: + """Convert response frames with one reusable, bounded executor.""" + + def __init__(self, max_workers: int) -> None: + self._max_workers = max_workers + if self._max_workers < 1: + raise ValueError("max_workers must be positive") + self._executor = ( + ThreadPoolExecutor(max_workers=self._max_workers, thread_name_prefix="video-planar") + if self._max_workers > 1 + else None + ) + self._scratch_local = threading.local() + self._shutdown = False + + @property + def max_workers(self) -> int: + return self._max_workers + + def _channel_scratch(self, height: int, width: int, common_dtype: np.dtype) -> np.ndarray: + scratch_dtype = np.dtype(common_dtype) + if not np.issubdtype(scratch_dtype, np.floating): + scratch_dtype = np.dtype(np.float64) + scratch = getattr(self._scratch_local, "channel", None) + if scratch is None or scratch.shape != (height, width) or scratch.dtype != scratch_dtype: + scratch = np.empty((height, width), dtype=scratch_dtype) + self._scratch_local.channel = scratch + return scratch + + def _build_frame(self, frame: np.ndarray, common_dtype: np.dtype) -> av.VideoFrame: + """Build one quantized GBR PyAV frame.""" + import av + + height, width = frame.shape[:2] + scratch = None if frame.dtype == np.uint8 else self._channel_scratch(height, width, common_dtype) + av_frame = av.VideoFrame(width, height, format="gbrp") + for plane, channel in zip(av_frame.planes, (1, 2, 0)): + if plane.height < height or plane.line_size < width: + raise ValueError("PyAV video plane is smaller than the requested frame dimensions.") + plane_view = np.frombuffer( + memoryview(plane), + dtype=np.uint8, + count=plane.height * plane.line_size, + ).reshape(plane.height, plane.line_size) + plane_view.fill(0) + if frame.dtype == np.uint8: + plane_view[:height, :width] = frame[..., channel] + else: + assert scratch is not None + np.copyto(scratch, frame[..., channel], casting="unsafe") + np.clip(scratch, 0.0, 1.0, out=scratch) + scratch *= 255.0 + np.rint(scratch, out=scratch) + plane_view[:height, :width] = scratch + return av_frame + + def iter_frames( + self, + frames: list[np.ndarray], + common_dtype: np.dtype, + ) -> Generator[av.VideoFrame, None, None]: + """Yield converted frames in order with at most two batches in flight.""" + effective_workers = min(len(frames), self._max_workers) + if effective_workers <= 1: + for frame in frames: + yield self._build_frame(frame, common_dtype) + return + + executor = self._executor + assert executor is not None + pending: deque[Future[av.VideoFrame]] = deque() + next_index = 0 + max_pending = 2 * effective_workers + try: + while next_index < len(frames) and len(pending) < max_pending: + pending.append(executor.submit(self._build_frame, frames[next_index], common_dtype)) + next_index += 1 + + while pending: + converted_frame = pending.popleft().result() + if next_index < len(frames): + pending.append(executor.submit(self._build_frame, frames[next_index], common_dtype)) + next_index += 1 + yield converted_frame + finally: + for future in pending: + future.cancel() + + def shutdown(self) -> None: + if not self._shutdown: + self._shutdown = True + if self._executor is not None: + self._executor.shutdown(wait=True, cancel_futures=True) + + +def _iter_planar_video_frames( + frames: list[np.ndarray], + common_dtype: np.dtype, +) -> Generator[av.VideoFrame, None, None]: + """Yield serial planar frames for direct utility callers.""" + converter = _PlanarFrameConverter(max_workers=1) + try: + yield from converter.iter_frames(frames, common_dtype) + finally: + converter.shutdown() + + +def _log_video_encoding_path( + *, + selected_path: str, + frames: list[np.ndarray], + frame_shape: tuple[int, ...], + common_dtype: np.dtype, + fps: int, + audio: AudioInput | None, + audio_sample_rate: int | None, + effective_frame_conversion_workers: int, + reason: str | None = None, +) -> None: + reason_field = "" if reason is None else f" reason={reason}" + logger.info( + "Video response encoding route selected: selected_path=%s%s frames=%d frame_shape=%s dtype=%s fps=%s " + "audio_present=%s effective_audio_sample_rate=%s effective_frame_conversion_workers=%s", + selected_path, + reason_field, + len(frames), + frame_shape, + np.dtype(common_dtype).name, + fps, + audio is not None, + audio_sample_rate, + effective_frame_conversion_workers, + ) + + +def _resolve_audio_sample_rate(audio: AudioInput | None, audio_sample_rate: int | None) -> int: + if audio is not None and audio_sample_rate is None: + logger.info_once( + "Audio sample rate was not provided; using default sample rate of %s Hz.", + DEFAULT_AUDIO_SAMPLE_RATE, + ) + return audio_sample_rate or DEFAULT_AUDIO_SAMPLE_RATE + + +def _encode_prepared_video_bytes_legacy( + frames: list[np.ndarray], + frame_shape: tuple[int, ...], + common_dtype: np.dtype, + fps: int, + audio: Any | None = None, + audio_sample_rate: int | None = None, + video_codec_options: dict[str, str] | None = None, +) -> bytes: + """Encode validated frames through the compatibility path used before planar encoding.""" + from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes + + audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None + return mux_video_audio_bytes( + _coerce_prepared_video_to_uint8_frames(frames, frame_shape, common_dtype), + audio_np, + fps=float(fps), + audio_sample_rate=audio_sample_rate or DEFAULT_AUDIO_SAMPLE_RATE, + video_codec_options=video_codec_options, + ) + + +def _encode_video_bytes_legacy( video: Any, fps: int, audio: Any | None = None, audio_sample_rate: int | None = None, video_codec_options: dict[str, str] | None = None, ) -> bytes: - """Encode a video payload into MP4 bytes, optionally muxing audio.""" - from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes - - audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None - - return mux_video_audio_bytes( - _coerce_video_to_uint8_frames(video), - audio_np, - fps=float(fps), - audio_sample_rate=audio_sample_rate or 24000, + """Encode through the compatibility path used before planar encoding.""" + frames, frame_shape, common_dtype = _prepare_video_frames(video) + return _encode_prepared_video_bytes_legacy( + frames, + frame_shape, + common_dtype, + fps, + audio=audio, + audio_sample_rate=_resolve_audio_sample_rate(audio, audio_sample_rate), video_codec_options=video_codec_options, ) +def _encode_video_bytes( + video: Any, + fps: int, + audio: Any | None = None, + audio_sample_rate: int | None = None, + video_codec_options: dict[str, str] | None = None, + frame_converter: _PlanarFrameConverter | None = None, +) -> bytes: + """Encode a video payload through the direct planar or legacy path.""" + from vllm_omni.diffusion.utils.media_utils import mux_av_video_audio_bytes + + # Prepare once so validation is shared by both paths and malformed common + # input is reported before any muxer is opened. + frames, frame_shape, common_dtype = _prepare_video_frames(video) + effective_audio_sample_rate = _resolve_audio_sample_rate(audio, audio_sample_rate) if audio is not None else None + fallback_reason = _direct_planar_fallback_reason( + frames, + frame_shape, + common_dtype, + allow_strided_rgb_planes=frame_converter is not None and frame_converter.max_workers > 1, + ) + if fallback_reason is not None: + _log_video_encoding_path( + selected_path="legacy_fallback", + reason=fallback_reason, + frames=frames, + frame_shape=frame_shape, + common_dtype=common_dtype, + fps=fps, + audio=audio, + audio_sample_rate=effective_audio_sample_rate, + effective_frame_conversion_workers=0, + ) + return _encode_prepared_video_bytes_legacy( + frames, + frame_shape, + common_dtype, + fps, + audio=audio, + audio_sample_rate=effective_audio_sample_rate, + video_codec_options=video_codec_options, + ) + + owns_frame_converter = frame_converter is None + converter = frame_converter or _PlanarFrameConverter(max_workers=1) + _log_video_encoding_path( + selected_path="direct_planar", + frames=frames, + frame_shape=frame_shape, + common_dtype=common_dtype, + fps=fps, + audio=audio, + audio_sample_rate=effective_audio_sample_rate, + effective_frame_conversion_workers=min(len(frames), converter.max_workers), + ) + audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None + video_frames = converter.iter_frames(frames, common_dtype) + try: + return mux_av_video_audio_bytes( + video_frames, + width=frame_shape[1], + height=frame_shape[0], + audio_waveform=audio_np, + fps=float(fps), + audio_sample_rate=effective_audio_sample_rate, + video_codec_options=video_codec_options, + ) + finally: + video_frames.close() + if owns_frame_converter: + converter.shutdown() + + class FragmentedMP4VideoEncoder: """Normalize video chunks and append them to one fragmented MP4 stream.""" @@ -585,9 +912,15 @@ def encode_video_base64( audio: Any | None = None, audio_sample_rate: int | None = None, video_codec_options: dict[str, str] | None = None, + frame_converter: _PlanarFrameConverter | None = None, ) -> str: """Encode a video (frames/array/tensor) to base64 MP4.""" video_bytes = _encode_video_bytes( - video, fps=fps, audio=audio, audio_sample_rate=audio_sample_rate, video_codec_options=video_codec_options + video, + fps=fps, + audio=audio, + audio_sample_rate=audio_sample_rate, + video_codec_options=video_codec_options, + frame_converter=frame_converter, ) return base64.b64encode(video_bytes).decode("utf-8") diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index c899b2697..281cbc55c 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -10,19 +10,26 @@ from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +# Default for ``visible_baseline``: read the live device-control env var. +# Explicit ``None`` means the caller captured an unset env and must not fall +# back to a later thread's temporary CUDA_VISIBLE_DEVICES. +_LIVE_DEVICE_ENV = object() + def resolve_stage_physical_devices( stage_id: int, devices: str | int | None, *, - visible_baseline: str | None = None, + visible_baseline: str | None | object = _LIVE_DEVICE_ENV, ) -> str | None: """Map logical stage devices to physical IDs without mutating process env. - When ``visible_baseline`` is provided it is used for logical-to-physical - mapping instead of the current device-control env var. This keeps parallel - stage initialization correct even if another thread has temporarily - narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage. + ``visible_baseline`` is the device-control env captured at the start of + replica init. Pass the captured string, or ``None`` when that env was + unset. Omitting the argument reads the live env (single-threaded callers). + + Parallel stage init must pass the captured value: another thread may have + already narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage. """ if devices in (None, "cpu"): return None @@ -32,7 +39,7 @@ def resolve_stage_physical_devices( device_list = _parse_device_list(devices) env_var = current_omni_platform.device_control_env_var - visible_devices = visible_baseline if visible_baseline is not None else os.environ.get(env_var) + visible_devices = os.environ.get(env_var) if visible_baseline is _LIVE_DEVICE_ENV else visible_baseline if visible_devices is not None: visible_device_list = _parse_device_list(visible_devices) device_list = _map_device_list(stage_id, device_list, visible_device_list) @@ -43,7 +50,7 @@ def set_stage_devices( stage_id: int, devices: str | int | None, *, - visible_baseline: str | None = None, + visible_baseline: str | None | object = _LIVE_DEVICE_ENV, ) -> str | None: """Configure per-stage device visibility and current device (CUDA or NPU). diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index ff0cd4f9e..d1c4c8aba 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import json import os import types @@ -6,16 +9,20 @@ from dataclasses import fields, is_dataclass from pathlib import Path from typing import Any, get_args, get_origin -import yaml from vllm.logger import init_logger from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.transformers_utils.config import get_config, get_hf_file_to_dict from vllm.transformers_utils.repo_utils import file_or_path_exists -from vllm_omni.config.config_factory import StageConfigFactory, with_trust_remote_code_override +from vllm_omni.config.config_factory import ( + StageConfigFactory, + _materialize_object_storage_configs, + _name_match_candidate, + with_trust_remote_code_override, +) from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import _DEPLOY_DIR -from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs +from vllm_omni.config.yaml_util import create_config, load_yaml_config from vllm_omni.diffusion.utils.hf_utils import ( _looks_like_dreamzero, get_diffusion_model_index, @@ -30,15 +37,6 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent logger = init_logger(__name__) -def is_new_format_deploy_config(path: str | None) -> bool: - """Return True when *path* is a deploy YAML (``stages:`` without legacy ``stage_args:``).""" - if not path or not os.path.exists(path): - return False - with open(path, encoding="utf-8") as f: - peek = yaml.safe_load(f) or {} - return isinstance(peek, dict) and "stages" in peek and "stage_args" not in peek - - _DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = { "GlmImagePipeline": "glm_image", } @@ -220,12 +218,13 @@ def _try_resolve_omni_model_type(model: str) -> str | None: Searches both the legacy ``stage_configs/*.yaml`` directory and the migrated ``deploy/*.yaml`` directory for a stem that substring-matches - the model path (e.g. ``cosyvoice3`` in - ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). The longest match wins so - ``cosyvoice3`` beats ``cosyvoice`` and ``bagel_single_stage`` beats - ``bagel``. + the model path basename (e.g. ``cosyvoice3`` in + ``FunAudioLLM/Fun-CosyVoice3-0.5B-2512``). Only the basename is scanned + so URI segments such as the bucket name cannot select an unrelated + pipeline. The longest match wins so ``cosyvoice3`` beats ``cosyvoice`` + and ``bagel_single_stage`` beats ``bagel``. """ - model_lower = model.lower().replace("-", "").replace("_", "") + model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "") best_match: str | None = None best_len = 0 for subdir in ("model_executor/stage_configs", "deploy"): @@ -284,24 +283,29 @@ def resolve_model_config_path(model: str) -> str | None: Raises: ValueError: If model_type cannot be determined """ + # Object-storage URIs are streamed by vLLM's Run:AI streamer only after + # each stage builds its ModelConfig, so config resolution here reads the + # local copy that vLLM-Omni materializes for such URIs. Name-based + # fallbacks keep the original string (the materialized path is a hash). + config_source = _materialize_object_storage_configs(model) # Try to get config from standard transformers format first try: - hf_config = get_config(model, trust_remote_code=True) + hf_config = get_config(config_source, trust_remote_code=True) model_type = hf_config.model_type except (ValueError, Exception): # If standard transformers format fails, try diffusers format - if get_diffusion_model_index(model) is not None: - model_type = _try_get_class_name_from_diffusers_config(model) + if get_diffusion_model_index(config_source) is not None: + model_type = _try_get_class_name_from_diffusers_config(config_source) if model_type is None: raise ValueError( f"Could not determine model_type for diffusers model: {model}. " "Please ensure its Diffusers pipeline index contains '_class_name'" ) - elif file_or_path_exists(model, "config.json", revision=None): + elif file_or_path_exists(config_source, "config.json", revision=None): # Try to read config.json manually for custom models like Bagel that fail get_config # but have a valid config.json with model_type try: - config_dict = get_hf_file_to_dict("config.json", model, revision=None) + config_dict = get_hf_file_to_dict("config.json", config_source, revision=None) if config_dict and "model_type" in config_dict: model_type = config_dict["model_type"] else: @@ -325,7 +329,7 @@ def resolve_model_config_path(model: str) -> str | None: ) default_config_path = current_omni_platform.get_default_stage_config_path() - if model_type == "vla" and _looks_like_dreamzero(model): + if model_type == "vla" and _looks_like_dreamzero(config_source): model_type = "dreamzero" if model_type in _DIFFUSERS_CLASS_TO_CONFIG: @@ -368,11 +372,12 @@ def load_stage_configs_from_model( ) -> tuple[list, str | None]: """Load stage configurations from model's default config file. - For models registered in the pipeline registry (new path), uses + For models registered in the pipeline registry, uses ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges PipelineConfig + DeployConfig + CLI overrides. - For other models (legacy path), loads stage configs from YAML. + Models that cannot be resolved through the registry return no stages so the + caller can apply its default stage configuration, when available. Args: model: Model name or path (used to determine model_type) @@ -382,7 +387,7 @@ def load_stage_configs_from_model( stage_overrides: Per-stage overrides from --stage-overrides. strategy_config_path: Optional path to a composable-parallel ``strategy.yaml`` whose derived sizing is overlaid onto the - registry-merged stages (opt-in; ignored on the legacy YAML path). + registry-merged stages. Returns: ``(stage_configs, omni_lb_policy)``: the list of stage configuration @@ -426,69 +431,17 @@ def load_stage_configs_from_model( # Convert StageConfig objects to OmegaConf for backward compat return [stage.to_omegaconf() for stage in stages], omni_lb_policy - # Legacy fallback: load from YAML. A composable-parallel strategy cannot be - # applied here (it overlays onto registry-merged stages), so warn rather than - # silently dropping the operator's --strategy-config. + strategy_note = "" if strategy_config_path is not None: - logger.warning( - "--strategy-config (%s) was provided but model %r resolves via the " - "legacy stage_configs YAML path, which does not support " - "composable-parallel strategies; the strategy is ignored. Use a " - "registry-based model to apply it.", - strategy_config_path, - model, - ) - stage_config_path = resolve_model_config_path(model) - if stage_config_path is None: - return [], None - stage_configs = load_stage_configs_from_yaml( - config_path=stage_config_path, - base_engine_args=base_engine_args, - prefer_stage_engine_args=True, + strategy_note = f" Strategy config {strategy_config_path!r} was not applied." + logger.warning( + "No registered PipelineConfig resolved for model %r. Legacy `stage_args` " + "YAMLs are no longer supported; register the pipeline and provide deployment " + "overrides through `deploy_config`.%s", + model, + strategy_note, ) - return stage_configs, None - - -def load_stage_configs_from_yaml( - config_path: str, - base_engine_args: dict | None = None, - prefer_stage_engine_args: bool = True, -) -> list: - """Load stage configurations from a YAML file (legacy OmegaConf path). - - TODO(@lishunyang12): remove once all models use PipelineConfig + DeployConfig. - - Args: - config_path: Path to the YAML configuration file - base_engine_args: Engine args supplied by the caller. - prefer_stage_engine_args: When True, YAML stage args override caller - engine args. When False, caller engine args override YAML defaults. - - Returns: - List of stage configuration dictionaries from the file's stage_args - """ - if base_engine_args is None: - base_engine_args = {} - config_data = load_yaml_config(config_path) - stage_args = config_data.stage_args - global_async_chunk = config_data.get("async_chunk", False) - # Convert any nested dataclass objects to dicts before creating DictConfig - base_engine_args = _convert_dataclasses_to_dict(base_engine_args) - base_engine_args = create_config(base_engine_args) - for stage_arg in stage_args: - base_engine_args_tmp = base_engine_args.copy() - # Update base_engine_args with stage-specific engine_args if they exist - if hasattr(stage_arg, "engine_args") and stage_arg.engine_args is not None: - if prefer_stage_engine_args: - merged_engine_args = merge_configs(base_engine_args_tmp, stage_arg.engine_args) - else: - merged_engine_args = merge_configs(stage_arg.engine_args, base_engine_args_tmp) - base_engine_args_tmp = create_config(merged_engine_args) - stage_type = getattr(stage_arg, "stage_type", "llm") - if hasattr(stage_arg, "runtime") and stage_arg.runtime is not None and stage_type != "diffusion": - base_engine_args_tmp.async_chunk = global_async_chunk - stage_arg.engine_args = base_engine_args_tmp - return stage_args + return [], None def filter_stages( @@ -606,7 +559,6 @@ def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None: def load_and_resolve_stage_configs( model: str, - stage_configs_path: str | None, kwargs: dict | None, *, trust_remote_code: bool | None, @@ -615,17 +567,15 @@ def load_and_resolve_stage_configs( stage_overrides: dict[str, dict[str, Any]] | None = None, strategy_config_path: str | None = None, ) -> tuple[str, list, str | None]: - """Load stage configurations from model or YAML file with fallback to defaults. + """Load stage configurations from a deploy YAML or model defaults. Args: model: Model name or path - stage_configs_path: Optional path to legacy YAML (stage_args format) kwargs: Engine arguments to merge with stage configs trust_remote_code: Whether to trust remote code while resolving the model config. default_stage_cfg_factory: Optional callable that takes no args and returns default stage config list when no configs are found - deploy_config_path: Optional path to deploy YAML (new format). - Mutually exclusive with ``stage_configs_path``. + deploy_config_path: Optional path to a deploy YAML. stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON. Keys are stage_id strings, values are dicts of overrides. strategy_config_path: Optional path to a composable-parallel @@ -636,67 +586,35 @@ def load_and_resolve_stage_configs( the strategy-derived pipeline-wide load-balancer policy (``None`` when no strategy set one), returned for the engine to apply. """ - if stage_configs_path is not None and deploy_config_path is not None: - raise ValueError( - "--stage-configs-path and --deploy-config are mutually exclusive: " - "they use different path resolution rules and loading paths. " - "Use --deploy-config for new-format YAMLs (preferred); " - "--stage-configs-path is kept only for the legacy `stage_args` format " - "and will be removed in a future release." - ) - if stage_configs_path is not None and deploy_config_path is None: - if not os.path.exists(stage_configs_path): - raise FileNotFoundError( - f"--stage-configs-path {stage_configs_path!r} does not exist. " - "Legacy `stage_configs/` yamls were replaced by `vllm_omni/deploy/.yaml`; " - "use --deploy-config. See docs/configuration/stage_configs.md." - ) - if is_new_format_deploy_config(stage_configs_path): - deploy_config_path = stage_configs_path - stage_configs_path = None + config_path = deploy_config_path if deploy_config_path is not None else resolve_model_config_path(model) + stage_configs, omni_lb_policy = load_stage_configs_from_model( + model, + trust_remote_code=trust_remote_code, + base_engine_args=kwargs, + deploy_config_path=deploy_config_path, + stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, + ) + if not stage_configs: + if default_stage_cfg_factory is not None: + default_stage_cfg = default_stage_cfg_factory() + stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg)) else: - logger.warning( - "--stage-configs-path is deprecated; migrate %r and use --deploy-config.", - stage_configs_path, - ) - - omni_lb_policy: str | None = None - if deploy_config_path is not None: - config_path = deploy_config_path - stage_configs, omni_lb_policy = load_stage_configs_from_model( - model, - trust_remote_code=trust_remote_code, - base_engine_args=kwargs, - deploy_config_path=deploy_config_path, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ) - if not stage_configs: - if default_stage_cfg_factory is not None: - default_stage_cfg = default_stage_cfg_factory() - stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg)) - else: - stage_configs = [] - elif stage_configs_path is None: - config_path = resolve_model_config_path(model) - stage_configs, omni_lb_policy = load_stage_configs_from_model( - model, - trust_remote_code=trust_remote_code, - base_engine_args=kwargs, - stage_overrides=stage_overrides, - strategy_config_path=strategy_config_path, - ) - if not stage_configs: - if default_stage_cfg_factory is not None: - default_stage_cfg = default_stage_cfg_factory() - stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg)) - else: - stage_configs = [] - else: - config_path = stage_configs_path - stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=kwargs) + stage_configs = [] stage_configs = filter_stages(config_path, stage_configs, kwargs) + + # ``revision`` selects the model source for every stage. Keep the global + # CLI pin on each stage's engine args so diffusion config construction can + # pass it through to all component loaders. The registry-to-legacy-stage + # conversion currently drops this inherited vLLM EngineArgs field. + revision = (kwargs or {}).get("revision") + if revision is not None: + for stage_config in stage_configs: + engine_args = stage_config.get("engine_args") if hasattr(stage_config, "get") else None + if engine_args is not None and engine_args.get("revision") is None: + engine_args["revision"] = revision + logger.debug(f"stage_configs: {stage_configs}") return config_path, stage_configs, omni_lb_policy diff --git a/vllm_omni/experimental/ar_diffusion/__init__.py b/vllm_omni/experimental/ar_diffusion/__init__.py index b87ac1ccc..bd0f2eac2 100644 --- a/vllm_omni/experimental/ar_diffusion/__init__.py +++ b/vllm_omni/experimental/ar_diffusion/__init__.py @@ -13,12 +13,32 @@ from vllm_omni.experimental.ar_diffusion.capability import ( SupportsARDiffusionPipeline, SupportsARDiffusionWarmup, ) +from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine +from vllm_omni.experimental.ar_diffusion.session import ( + ARDiffusionSession, + ARDiffusionSessionManager, + ARDiffusionWorkerLifecycle, +) +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + AR_DIFFUSION_TICK_KEY, + ARDiffusionChunkMetadata, + ARDiffusionControlInput, + ARDiffusionTickRequest, +) __all__ = [ "ARDiffusionKVBranchSpec", "ARDiffusionCrossAttentionKVSpec", "ARDiffusionEngine", + "ARDiffusionOmniTickConsumer", + "ARDiffusionSession", + "ARDiffusionSessionManager", + "ARDiffusionWorkerLifecycle", + "AR_DIFFUSION_TICK_KEY", + "ARDiffusionChunkMetadata", + "ARDiffusionControlInput", + "ARDiffusionTickRequest", "ARDiffusionKVCacheSpec", "SupportsARDiffusionPipeline", "SupportsARDiffusionWarmup", diff --git a/vllm_omni/experimental/ar_diffusion/capability.py b/vllm_omni/experimental/ar_diffusion/capability.py index b30d10c1e..64eb62cf0 100644 --- a/vllm_omni/experimental/ar_diffusion/capability.py +++ b/vllm_omni/experimental/ar_diffusion/capability.py @@ -58,6 +58,9 @@ class ARDiffusionKVCacheSpec: ``max_scratch_tokens_per_branch`` is the maximum non-video KV (for example, action/state registers) that must coexist with an uncommitted video block. + ``model_owned_state_bytes_per_session`` is a conservative upper bound for + persistent per-session CUDA state outside runner-owned self/cross-attention + KV and scratch pools. ``session_capacity`` is the pipeline's upper bound; a runner may retain fewer sessions when its current memory and execution policy is stricter. @@ -79,6 +82,7 @@ class ARDiffusionKVCacheSpec: cross_attention: tuple[ARDiffusionCrossAttentionKVSpec, ...] = () max_model_len: int = 1 << 20 max_scratch_tokens_per_branch: int = 0 + model_owned_state_bytes_per_session: int = 0 def __post_init__(self) -> None: positive_fields = { @@ -101,6 +105,11 @@ class ARDiffusionKVCacheSpec: "AR-Diffusion max_scratch_tokens_per_branch must be non-negative, " f"got {self.max_scratch_tokens_per_branch}" ) + if self.model_owned_state_bytes_per_session < 0: + raise ValueError( + "AR-Diffusion model_owned_state_bytes_per_session must be non-negative, " + f"got {self.model_owned_state_bytes_per_session}" + ) if not self.kv_branches: raise ValueError("AR-Diffusion requires at least one KV branch") kv_branch_names = [kv_branch.name for kv_branch in self.kv_branches] diff --git a/vllm_omni/experimental/ar_diffusion/consumer.py b/vllm_omni/experimental/ar_diffusion/consumer.py new file mode 100644 index 000000000..a819e0b31 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/consumer.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +"""AsyncOmni consumer for typed realtime AR-Diffusion ticks.""" + +from __future__ import annotations + +import copy +from collections.abc import AsyncIterator, Callable, Mapping, Sequence +from typing import Any, Protocol + +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionChunkMetadata, + ARDiffusionTickRequest, +) +from vllm_omni.inputs.data import ( + OmniDiffusionSamplingParams, + OmniPromptType, + OmniSamplingParams, +) +from vllm_omni.outputs import OmniRequestOutput + +AR_DIFFUSION_OUTPUT_METADATA_KEY = "ar_diffusion" + + +class ARDiffusionGenerateClient(Protocol): + """Internal subset of AsyncOmni used to execute one tick.""" + + def generate( + self, + prompt: OmniPromptType, + *, + request_id: str, + sampling_params_list: Sequence[OmniSamplingParams], + output_modalities: list[str] | None = None, + ) -> AsyncIterator[OmniRequestOutput]: ... + + +def _metadata_mapping(output: OmniRequestOutput) -> Mapping[str, Any]: + multimodal_output = output.multimodal_output + if not isinstance(multimodal_output, Mapping): + raise ValueError("AR-Diffusion output must contain a multimodal output mapping.") + metadata = multimodal_output.get("metadata") + if not isinstance(metadata, Mapping): + raise ValueError("AR-Diffusion output must contain multimodal_output.metadata.") + value = metadata.get(AR_DIFFUSION_OUTPUT_METADATA_KEY) + if not isinstance(value, Mapping): + raise ValueError(f"AR-Diffusion output metadata must contain {AR_DIFFUSION_OUTPUT_METADATA_KEY!r}.") + return value + + +def _parse_chunk_metadata(output: OmniRequestOutput) -> ARDiffusionChunkMetadata: + value = _metadata_mapping(output) + session_id = value.get("session_id") + request_id = value.get("request_id") + chunk_index = value.get("chunk_index") + applied_event_ids = value.get("applied_event_ids", ()) + if not isinstance(session_id, str) or not session_id.strip(): + raise ValueError("AR-Diffusion output metadata session_id must be a non-empty string.") + if not isinstance(request_id, str) or not request_id.strip(): + raise ValueError("AR-Diffusion output metadata request_id must be a non-empty string.") + if isinstance(chunk_index, bool) or not isinstance(chunk_index, int) or chunk_index < 0: + raise ValueError("AR-Diffusion output metadata chunk_index must be a non-negative integer.") + if not isinstance(applied_event_ids, Sequence) or isinstance(applied_event_ids, (str, bytes, bytearray)): + raise ValueError("AR-Diffusion output metadata applied_event_ids must be a sequence.") + event_ids = tuple(applied_event_ids) + if any(isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0 for event_id in event_ids): + raise ValueError("AR-Diffusion output metadata applied_event_ids must contain non-negative integers.") + if event_ids != tuple(sorted(set(event_ids))): + raise ValueError("AR-Diffusion output metadata applied_event_ids must be unique and strictly increasing.") + return ARDiffusionChunkMetadata( + session_id=session_id, + request_id=request_id, + chunk_index=chunk_index, + applied_event_ids=event_ids, + ) + + +class ARDiffusionOmniTickConsumer: + """Converts a typed tick into one AsyncOmni diffusion request. + + ``prompt_provider`` packages session initialization data such as an initial + image into the normal Omni prompt. Control ``track/schema/data`` remains in + the typed tick and is not interpreted by this class. + """ + + def __init__( + self, + engine_client: ARDiffusionGenerateClient, + *, + prompt_provider: Callable[[ARDiffusionTickRequest], OmniPromptType], + sampling_params_list: Sequence[OmniSamplingParams], + diffusion_stage_id: int, + output_modalities: Sequence[str] | None = None, + ) -> None: + templates = tuple(sampling_params_list) + if not templates: + raise ValueError("sampling_params_list must not be empty.") + if ( + isinstance(diffusion_stage_id, bool) + or not isinstance(diffusion_stage_id, int) + or not 0 <= diffusion_stage_id < len(templates) + ): + raise ValueError("diffusion_stage_id must index sampling_params_list.") + if not isinstance(templates[diffusion_stage_id], OmniDiffusionSamplingParams): + raise TypeError("The selected diffusion stage must use OmniDiffusionSamplingParams.") + stage_pools = getattr(getattr(engine_client, "engine", None), "stage_pools", None) + if stage_pools is None or not 0 <= diffusion_stage_id < len(stage_pools): + raise ValueError( + "AR-Diffusion realtime execution requires an initialized AsyncOmni " + "engine with a matching diffusion stage pool." + ) + replica_count = getattr(stage_pools[diffusion_stage_id], "num_replicas", None) + if isinstance(replica_count, bool) or not isinstance(replica_count, int) or replica_count < 1: + raise ValueError("The selected AR-Diffusion stage must expose a positive num_replicas.") + if replica_count != 1: + raise ValueError( + "AR-Diffusion realtime sessions currently require exactly one " + "replica for the selected diffusion stage; session-affine " + f"routing is not implemented (got num_replicas={replica_count})." + ) + self._engine_client = engine_client + self._prompt_provider = prompt_provider + self._sampling_params_templates = templates + self._diffusion_stage_id = diffusion_stage_id + self._output_modalities = list(output_modalities) if output_modalities is not None else None + + def _sampling_params_for_tick( + self, + tick: ARDiffusionTickRequest, + ) -> list[OmniSamplingParams]: + params = [copy.copy(template) for template in self._sampling_params_templates] + diffusion_params = params[self._diffusion_stage_id] + assert isinstance(diffusion_params, OmniDiffusionSamplingParams) + diffusion_params.extra_args = { + **(diffusion_params.extra_args or {}), + **tick.to_extra_args(), + } + return params + + async def execute_tick(self, tick: ARDiffusionTickRequest) -> OmniRequestOutput: + prompt = self._prompt_provider(tick) + final_output: OmniRequestOutput | None = None + async for output in self._engine_client.generate( + prompt, + request_id=tick.request_id, + sampling_params_list=self._sampling_params_for_tick(tick), + output_modalities=self._output_modalities, + ): + if output.stage_id != self._diffusion_stage_id: + continue + if output.finished: + final_output = output + + if final_output is None: + raise RuntimeError( + f"AR-Diffusion tick {tick.request_id!r} completed without a final " + f"output from stage {self._diffusion_stage_id}." + ) + if final_output.error: + raise RuntimeError(f"AR-Diffusion tick {tick.request_id!r} failed: {final_output.error}") + if final_output.request_id != tick.request_id: + raise ValueError("AR-Diffusion output request_id does not match the submitted tick request_id.") + # Parse here so malformed/missing standard metadata fails at the engine + # boundary, before the session attempts to commit accepted events. + self.chunk_metadata(final_output) + return final_output + + def chunk_metadata(self, output: OmniRequestOutput) -> ARDiffusionChunkMetadata: + return _parse_chunk_metadata(output) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py index b407d24a3..94e6edf0d 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py @@ -32,7 +32,6 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( ARDiffusionPagedLayerInputs, ar_diffusion_paged_attention, paged_write_attn, - set_current_paged_kv_cache, ) __all__ = [ @@ -53,5 +52,4 @@ __all__ = [ "paged_write_attn", "pool_write_chunk", "resident_block_ids", - "set_current_paged_kv_cache", ] diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py index 07be2546e..ca91a8bb6 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py @@ -23,7 +23,8 @@ class ARDiffusionKVConfig: sink_chunks: int = 0 # Boundary reset vs. sliding replacement. reset_at_boundary: bool = False - # Fraction of free device memory budgeted for the AR-Diffusion KV pool. + # Fraction of free device memory used to admit additional resident + # sessions. One session is admitted whenever it fits actual free memory. gpu_memory_fraction: float = 0.1 # When CUDA graph / torch.compile is on (not enforce_eager), pre-capture the # DiT graphs for every window-fill shape at load time via a synthetic rollout, diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index 9ef3cca0c..da61afe4d 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -174,6 +174,7 @@ class ARDiffusionKVCache: device: torch.device | None = None, frames_per_block: int = 1, max_scratch_tokens_per_branch: int = 0, + model_owned_state_bytes_per_session: int = 0, ) -> None: if not config.enable: raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig") @@ -197,13 +198,17 @@ class ARDiffusionKVCache: self.config = config self.kv_branches = kv_branches - self.session_capacity = session_capacity + self.requested_session_capacity = session_capacity self._kv_branch_local_indices = {kv_branch.name: kv_branch.local_index for kv_branch in kv_branches} self.num_local_kv_branches = max(local_indices) + 1 if frames_per_block <= 0: raise ValueError(f"frames_per_block must be positive, got {frames_per_block}") if max_scratch_tokens_per_branch < 0: raise ValueError(f"max_scratch_tokens_per_branch must be non-negative, got {max_scratch_tokens_per_branch}") + if model_owned_state_bytes_per_session < 0: + raise ValueError( + f"model_owned_state_bytes_per_session must be non-negative, got {model_owned_state_bytes_per_session}" + ) self.frames_per_block = int(frames_per_block) self.block_size = block_size self.num_layers = num_layers @@ -217,28 +222,11 @@ class ARDiffusionKVCache: self.device = device or torch.device("cpu") self._allocate_tensors = device is not None self._adapters: dict[str, ARDiffusionRequestAdapter] = {} - - # Named cross-attention K/V is allocated lazily per session and released - # with that session. Reserve its worst-case capacity when sizing the - # paged self-attention pool so the two stores share one memory budget. self._cross_sessions: dict[ str, dict[str, dict[str, tuple[list[torch.Tensor], list[torch.Tensor]]]], ] = {} - def _cross_pool_bytes(length: int) -> int: - return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers) - - cross_bytes_per_session = sum(_cross_pool_bytes(length) for length in self.cross_attention_lengths.values()) - cross_total_bytes = cross_bytes_per_session * session_capacity - if cross_total_bytes: - _log.info( - "AR-Diffusion cross-attn reservation: %.1f MiB/session × %d sessions = %.1f MiB", - cross_bytes_per_session / (1024 * 1024), - session_capacity, - cross_total_bytes / (1024 * 1024), - ) - self.spec = ChunkWindowSpec( block_size=block_size, num_kv_heads=num_kv_heads, @@ -250,41 +238,7 @@ class ARDiffusionKVCache: sink_chunks=config.sink_chunks, reset_at_boundary=config.reset_at_boundary, ) - # Each pool block spans all layers' K/V, so size against the per-layer - # page size times the layer count. - # Size the self-attn pool against memory reserved for the maximum number - # of lazily allocated cross-attention sessions. - num_blocks = compute_num_blocks( - max(0, available_bytes - cross_total_bytes), - config.gpu_memory_fraction, - self.spec.page_size_bytes * num_layers, - ) - # Floor: one forward needs the resident window plus the in-flight chunk - # (frames_per_block frame-blocks) for every KV branch THIS rank runs, - # with a little eviction-transient headroom. The memory-fraction - # heuristic can under-size this once block_size grows — e.g. frame-granular - # paging at the true frame_seqlen makes each block larger and the pool - # fewer-blocks — so guarantee the minimum the rollout cannot run without, - # otherwise allocate_chunk hits an exhausted pool mid-forward. - resident_blocks = config.sink_chunks + config.window_chunks - min_blocks = self.num_local_kv_branches * (resident_blocks + self.frames_per_block) + 2 - if num_blocks < min_blocks: - _log.warning( - "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block " - "floor (%d local KV branch(es) x (sink_chunks=%d + window_chunks=%d " - "+ frames_per_block=%d) + 2 headroom)", - num_blocks, - min_blocks, - self.num_local_kv_branches, - config.sink_chunks, - config.window_chunks, - self.frames_per_block, - ) - num_blocks = min_blocks - layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)] - self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len) - self.managed_num_blocks = num_blocks - self.num_blocks = num_blocks + # Scratch blocks are outside KVCacheManager ownership. A non-committing # forward needs one block per current frame plus space for any # model-declared action/state tokens that coexist with video KV. @@ -297,6 +251,101 @@ class ARDiffusionKVCache: scratch_per_kv_branch = max(minimum_scratch_blocks, override_blocks) self.scratch_blocks_per_kv_branch = scratch_per_kv_branch self.scratch_num_blocks = self.num_local_kv_branches * scratch_per_kv_branch + + # The self-attention pool, scratch pool, and lazily materialized + # cross-attention caches share one hard memory budget. Select the largest + # feasible resident-session count rather than raising a block-count floor + # past that budget. All resident sessions need a complete sink + window; + # only one request can be in flight, so frames_per_block is counted once. + page_size_bytes = self.spec.page_size_bytes * num_layers + self.available_memory_bytes = available_bytes + self.configured_memory_budget_bytes = int(available_bytes * config.gpu_memory_fraction) + # Reuse the public helper's validation for the memory fraction/page size. + compute_num_blocks(available_bytes, config.gpu_memory_fraction, page_size_bytes) + self.scratch_reserved_bytes = self.scratch_num_blocks * page_size_bytes + self.model_owned_state_bytes_per_session = model_owned_state_bytes_per_session + + def _cross_pool_bytes(length: int) -> int: + return int(2 * len(self.kv_branches) * length * num_kv_heads * head_size * dtype.itemsize * num_layers) + + self.cross_attention_bytes_per_session = sum( + _cross_pool_bytes(length) for length in self.cross_attention_lengths.values() + ) + + def _required_managed_blocks(capacity: int) -> int: + resident_per_session = config.sink_chunks + config.window_chunks + return self.num_local_kv_branches * (capacity * resident_per_session + self.frames_per_block) + 2 + + def _required_bytes(capacity: int) -> int: + return ( + self.scratch_reserved_bytes + + capacity * self.cross_attention_bytes_per_session + + capacity * self.model_owned_state_bytes_per_session + + _required_managed_blocks(capacity) * page_size_bytes + ) + + one_session_bytes = _required_bytes(1) + if one_session_bytes > available_bytes: + raise ValueError( + "AR-Diffusion available device memory cannot fit one session: " + f"available={available_bytes} bytes, required={one_session_bytes} bytes " + "(managed self-attention + cross-attention + scratch + model-owned state)." + ) + self.memory_budget_bytes = max(self.configured_memory_budget_bytes, one_session_bytes) + if self.memory_budget_bytes > self.configured_memory_budget_bytes: + _log.warning( + "AR-Diffusion raised the configured memory budget from %d to %d bytes " + "to admit one session that fits actual free device memory", + self.configured_memory_budget_bytes, + self.memory_budget_bytes, + ) + + effective_capacity = 0 + required_managed_blocks = 0 + for candidate in range(session_capacity, 0, -1): + candidate_managed_blocks = _required_managed_blocks(candidate) + if _required_bytes(candidate) <= self.memory_budget_bytes: + effective_capacity = candidate + required_managed_blocks = candidate_managed_blocks + break + assert effective_capacity > 0 + + self.session_capacity = effective_capacity + self.cross_attention_reserved_bytes = self.cross_attention_bytes_per_session * effective_capacity + self.model_owned_state_reserved_bytes = self.model_owned_state_bytes_per_session * effective_capacity + self_attn_budget_bytes = ( + self.memory_budget_bytes + - self.scratch_reserved_bytes + - self.cross_attention_reserved_bytes + - self.model_owned_state_reserved_bytes + ) + num_blocks = self_attn_budget_bytes // page_size_bytes + assert num_blocks >= required_managed_blocks + if effective_capacity < session_capacity: + _log.warning( + "AR-Diffusion resident session capacity reduced from %d to %d by the KV memory budget", + session_capacity, + effective_capacity, + ) + if self.cross_attention_reserved_bytes: + _log.info( + "AR-Diffusion cross-attn reservation: %.1f MiB/session × %d sessions = %.1f MiB", + self.cross_attention_bytes_per_session / (1024 * 1024), + effective_capacity, + self.cross_attention_reserved_bytes / (1024 * 1024), + ) + if self.model_owned_state_reserved_bytes: + _log.info( + "AR-Diffusion model-owned state reservation: %.1f MiB/session × %d sessions = %.1f MiB", + self.model_owned_state_bytes_per_session / (1024 * 1024), + effective_capacity, + self.model_owned_state_reserved_bytes / (1024 * 1024), + ) + + layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)] + self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len) + self.managed_num_blocks = num_blocks + self.num_blocks = num_blocks self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks self.null_block_id = self.manager.block_pool.null_block.block_id @@ -561,6 +610,11 @@ class ARDiffusionKVCache: """ _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks) adapter.on_chunk_committed() + self.manager.remove_skipped_blocks( + adapter.request_id, + adapter.num_computed_tokens, + num_prompt_tokens=adapter.num_prompt_tokens, + ) _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks) # -- pool-backed K/V access -------------------------------------------- diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py index bcd06b63a..b2e19a95f 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py @@ -10,19 +10,6 @@ import torch from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping -# Set by ARDiffusionPagedForwardContext.prepare() before each KV branch forward and -# read by the fused write+attend custom op below. The pools are process-lifetime -# allocations owned by one ARDiffusionKVCache per worker process (CFG-parallel / -# TP ranks are separate processes), mirroring vLLM's forward-context pattern for -# keeping multi-GiB cache tensors out of the compiled graph's inputs. -_CURRENT_PAGED_KV_CACHE: Any = None - - -def set_current_paged_kv_cache(kv_cache: Any) -> None: - global _CURRENT_PAGED_KV_CACHE - _CURRENT_PAGED_KV_CACHE = kv_cache - - _LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {} @@ -50,6 +37,9 @@ class ARDiffusionPagedLayerInputs(NamedTuple): """ layer_idx: torch.Tensor + key_pool: torch.Tensor + value_pool: torch.Tensor + block_size: int seq_len: int video_slots: torch.Tensor action_slots: torch.Tensor @@ -228,7 +218,6 @@ class ARDiffusionPagedForwardContext: ) = self.build_block_table(action_len=action_len, query_len=query_len, device=device) if self.action_slot_mapping is None: self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device) - set_current_paged_kv_cache(self.kv_cache) self._prepared = True def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs: @@ -236,6 +225,9 @@ class ARDiffusionPagedForwardContext: raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()") return ARDiffusionPagedLayerInputs( layer_idx=_layer_idx_tensor(layer_idx), + key_pool=self.kv_cache._k_pools[layer_idx], + value_pool=self.kv_cache._v_pools[layer_idx], + block_size=int(self.kv_cache.block_size), seq_len=int(self.seq_len), video_slots=self.current_video_slot_mapping, action_slots=self.action_slot_mapping, @@ -365,6 +357,22 @@ def _resolve_fa_version(head_size: int) -> int: return version +def _rocm_flash_attn_varlen_func(): + """Resolve a ROCm varlen kernel, preferring AITER when available. + + The caller gathers paged KV into packed tensors before invoking this + function, avoiding AITER releases whose ``block_table`` kernel is broken. + """ + try: + from aiter import flash_attn_varlen_func + + return flash_attn_varlen_func + except ImportError: + from flash_attn import flash_attn_varlen_func + + return flash_attn_varlen_func + + def ar_diffusion_paged_attention( query: torch.Tensor, key_cache: torch.Tensor, @@ -401,6 +409,34 @@ def ar_diffusion_paged_attention( softmax_scale, causal=causal, ) + elif torch.version.hip is not None: + # vllm.vllm_flash_attn contains CUDA-only extensions. ROCm's AITER and + # upstream flash-attn expose the standard cu_seqlens_k API instead of + # vLLM's seqused_k/fa_version API. The ROCm flash-attn paged kernel also + # requires 128-token blocks, while AR-Diffusion uses frame-aligned + # 16-token blocks, so gather the visible blocks on-device first. + flash_attn_varlen_func = _rocm_flash_attn_varlen_func() + cu_seqlens_k = torch.cat([seq_lens.new_zeros(1), torch.cumsum(seq_lens, dim=0, dtype=torch.int32)]) + positions = torch.arange(int(max_seq_len), device=query_flat.device) + logical_blocks = torch.div(positions, key_cache.shape[1], rounding_mode="floor") + offsets = positions % key_cache.shape[1] + physical_blocks = block_table[:, logical_blocks].long() + gathered_k = key_cache[physical_blocks, offsets] + gathered_v = value_cache[physical_blocks, offsets] + valid = positions.unsqueeze(0) < seq_lens.unsqueeze(1) + packed_k = gathered_k[valid] + packed_v = gathered_v[valid] + out = flash_attn_varlen_func( + q=query_flat, + k=packed_k, + v=packed_v, + cu_seqlens_q=query_start_loc, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=int(max_query_len), + max_seqlen_k=int(max_seq_len), + softmax_scale=float(softmax_scale), + causal=causal, + ) else: from vllm.vllm_flash_attn import flash_attn_varlen_func @@ -432,17 +468,18 @@ def ar_diffusion_paged_attention( # One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats # it as a single graph node (no eager island, no graph breaks), and the K/V slot # writes happen inside the op so write→read ordering with the block-table kernel -# is internal — no hidden-mutation ordering hazards against a separate reader op -# and no multi-GiB pool tensors as graph inputs (the flat and paged pool views -# alias the same storage, which AOTAutograd handles poorly as inputs). Pattern -# follows sage_attn3.py in-repo and vLLM's own unified attention ops. +# is internal. The flat pools are explicit mutable inputs: Inductor/CUDA Graph +# must track their storage lifetime instead of observing an undeclared mutation +# through a process-global registry. def _paged_write_attn_impl( query: torch.Tensor, k_curr: torch.Tensor, v_curr: torch.Tensor, k_act: torch.Tensor | None, v_act: torch.Tensor | None, - layer_idx: torch.Tensor, + key_pool: torch.Tensor, + value_pool: torch.Tensor, + block_size: int, video_slots: torch.Tensor, action_slots: torch.Tensor, block_table: torch.Tensor, @@ -452,21 +489,17 @@ def _paged_write_attn_impl( max_seq_len: int, softmax_scale: float, ) -> torch.Tensor: - kv = _CURRENT_PAGED_KV_CACHE - if kv is None: - raise RuntimeError("ar_diffusion_paged_write_attn called before prepare() set the KV pool registry") - layer_idx = int(layer_idx) - k_pool = kv._k_pools[layer_idx] - v_pool = kv._v_pools[layer_idx] - k_pool[video_slots] = k_curr.to(k_pool.dtype) - v_pool[video_slots] = v_curr.to(v_pool.dtype) + key_pool[video_slots] = k_curr.to(key_pool.dtype) + value_pool[video_slots] = v_curr.to(value_pool.dtype) if k_act is not None and v_act is not None and k_act.shape[0] > 0: - k_pool[action_slots] = k_act.to(k_pool.dtype) - v_pool[action_slots] = v_act.to(v_pool.dtype) + key_pool[action_slots] = k_act.to(key_pool.dtype) + value_pool[action_slots] = v_act.to(value_pool.dtype) + key_cache = key_pool.unflatten(0, (-1, block_size)) + value_cache = value_pool.unflatten(0, (-1, block_size)) return ar_diffusion_paged_attention( query, - kv.key_cache(layer_idx), - kv.value_cache(layer_idx), + key_cache, + value_cache, block_table=block_table, query_start_loc=query_start_loc, seq_lens=seq_lens, @@ -481,14 +514,19 @@ def _paged_write_attn_impl( # the module from sys.modules (same as sage_attn3.py). if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"): - @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=()) + @torch.library.custom_op( + "vllm_omni::ar_diffusion_paged_write_attn", + mutates_args=("key_pool", "value_pool"), + ) def _paged_write_attn_op( query: torch.Tensor, k_curr: torch.Tensor, v_curr: torch.Tensor, k_act: torch.Tensor | None, v_act: torch.Tensor | None, - layer_idx: torch.Tensor, + key_pool: torch.Tensor, + value_pool: torch.Tensor, + block_size: int, video_slots: torch.Tensor, action_slots: torch.Tensor, block_table: torch.Tensor, @@ -504,7 +542,9 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"): v_curr, k_act, v_act, - layer_idx, + key_pool, + value_pool, + block_size, video_slots, action_slots, block_table, @@ -522,7 +562,9 @@ if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"): v_curr, k_act, v_act, - layer_idx, + key_pool, + value_pool, + block_size, video_slots, action_slots, block_table, @@ -545,7 +587,9 @@ def paged_write_attn( v_curr, k_act, v_act, - inputs.layer_idx, + inputs.key_pool, + inputs.value_pool, + inputs.block_size, inputs.video_slots, inputs.action_slots, inputs.block_table, diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py index d114212bb..ec14a9177 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py @@ -130,7 +130,7 @@ class ARDiffusionKVState: if ctx.commit_current and ctx._allocated_video: n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size for _ in range(n_chunks): - ctx.adapter.on_chunk_committed() + self.kv_cache.commit_chunk(ctx.adapter) self._committed[kv_branch] += ctx.seq_len _log.debug( "AR-Diffusion COMMIT [%s] new_tokens=%d chunks=%d resident=%d/%d", @@ -173,6 +173,17 @@ class ARDiffusionKVState: for i in range(self.num_layers) ] + def clear_cross_attention(self) -> None: + """Invalidate all named cross-attention caches without dropping self-KV. + + Prompt changes may preserve the autoregressive world history while + requiring fresh text projections. This operation deliberately leaves + branch adapters and their resident self-attention blocks untouched. + """ + if self._closed: + raise RuntimeError(f"AR-Diffusion session {self.session_id!r} is closed") + self.kv_cache.release_cross_attention(self.session_id) + def close(self) -> None: """Release all self- and cross-attention storage owned by this session.""" if self._closed: diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py index 31a23eaf7..410c45fb1 100644 --- a/vllm_omni/experimental/ar_diffusion/runner.py +++ b/vllm_omni/experimental/ar_diffusion/runner.py @@ -21,6 +21,7 @@ from vllm_omni.experimental.ar_diffusion.capability import ( from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState +from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest logger = init_logger(__name__) @@ -51,7 +52,6 @@ class ARDiffusionModelRunner(DiffusionModelRunner): """ _WARMUP_SID = "__ardiffusion_warmup__" - _MAX_RESIDENT_SESSIONS = 1 def __init__(self, vllm_config: object, od_config: OmniDiffusionConfig, device: torch.device) -> None: super().__init__(vllm_config, od_config, device) @@ -123,7 +123,6 @@ class ARDiffusionModelRunner(DiffusionModelRunner): self.ar_diffusion_kv_config = config self._ar_diffusion_capability = capability self._ar_diffusion_kv_cache_spec = spec - self._session_capacity = min(spec.session_capacity, self._MAX_RESIDENT_SESSIONS) self.kv_cache = ARDiffusionKVCache( config, num_layers=spec.num_layers, @@ -134,12 +133,14 @@ class ARDiffusionModelRunner(DiffusionModelRunner): max_model_len=spec.max_model_len, available_bytes=self._available_memory_bytes() if available_bytes is None else available_bytes, kv_branches=spec.kv_branches, - session_capacity=self._session_capacity, + session_capacity=spec.session_capacity, cross_attention_lengths=spec.cross_attention_lengths, frames_per_block=spec.frames_per_block, max_scratch_tokens_per_branch=spec.max_scratch_tokens_per_branch, + model_owned_state_bytes_per_session=spec.model_owned_state_bytes_per_session, device=self.device, ) + self._session_capacity = self.kv_cache.session_capacity logger.info( "AR-Diffusion KV cache: blocks=%d layers=%d local_kv_heads=%d head_size=%d " "tokens/frame=%d frames/block=%d window=%d sink=%d kv_branches=%s cross=%s " @@ -230,9 +231,14 @@ class ARDiffusionModelRunner(DiffusionModelRunner): return state @staticmethod - def _request_session(req: OmniDiffusionRequest) -> tuple[str, dict]: + def _request_session( + req: OmniDiffusionRequest, + ) -> tuple[str, dict, ARDiffusionTickRequest | None]: extra_args = req.sampling_params.extra_args or {} - return str(extra_args.get("session_id") or "default"), extra_args + tick = ARDiffusionTickRequest.from_extra_args(extra_args) + if tick is not None: + return tick.session_id, extra_args, tick + return str(extra_args.get("session_id") or "default"), extra_args, None def execute_model( self, @@ -245,8 +251,10 @@ class ARDiffusionModelRunner(DiffusionModelRunner): if capability is None: raise RuntimeError("AR-Diffusion capability missing after KV cache initialization") - session_id, extra_args = self._request_session(req) - if extra_args.get("reset", False): + session_id, extra_args, tick = self._request_session(req) + reset = tick.reset if tick is not None else bool(extra_args.get("reset", False)) + close_session = tick.close_session if tick is not None else bool(extra_args.get("close_session", False)) + if reset: self.reset_session(session_id) state = self._get_or_create_session(session_id) started = time.perf_counter() @@ -268,7 +276,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner): ) raise self._perf_e2e_times.append(time.perf_counter() - started) - if extra_args.get("close_session", False): + if close_session: self.close_session(session_id) return output diff --git a/vllm_omni/experimental/ar_diffusion/session.py b/vllm_omni/experimental/ar_diffusion/session.py new file mode 100644 index 000000000..6bb4b27d9 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/session.py @@ -0,0 +1,578 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Internal realtime session control plane for AR-Diffusion. + +This module owns transport-level identity, event ordering, backpressure, and +chunk-boundary snapshots. Model-defined control payloads remain opaque here. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections import OrderedDict +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Any, Generic, Protocol, TypeVar + +from vllm_omni.experimental.ar_diffusion.tick_protocol import ( + ARDiffusionChunkMetadata, + ARDiffusionControlInput, + ARDiffusionTickRequest, +) + + +class ARDiffusionSessionError(RuntimeError): + """Base error for realtime AR-Diffusion session operations.""" + + +class ARDiffusionSessionStateError(ARDiffusionSessionError): + """Raised when an operation is invalid for the current session state.""" + + +class ARDiffusionEventQueueFullError(ARDiffusionSessionError): + """Raised when accepting an event would exceed the bounded queue.""" + + +class ARDiffusionStaleEventError(ARDiffusionSessionError): + """Raised when an event is older than the current monotonic boundary.""" + + +class ARDiffusionChunkMetadataError(ARDiffusionSessionError): + """Raised when a tick consumer returns metadata for a different chunk.""" + + +class ARDiffusionEventAcceptanceStatus(str, Enum): + ACCEPTED = "accepted" + DUPLICATE = "duplicate" + + +class ARDiffusionSessionStatus(str, Enum): + ACTIVE = "active" + RESETTING = "resetting" + FAILED = "failed" + CLOSING = "closing" + CLEANUP_FAILED = "cleanup_failed" + CLOSED = "closed" + + +def _copy_control(control: ARDiffusionControlInput) -> ARDiffusionControlInput: + return ARDiffusionControlInput( + track=control.track, + schema=control.schema, + data=control.data, + ) + + +@dataclass(frozen=True) +class ARDiffusionSessionEvent: + """One atomically accepted prompt/control update. + + ``prompt=None`` means that this event does not update the prompt. Controls + are keyed by ``track`` when the chunk snapshot is assembled; their schema + and data are otherwise opaque to this layer. + """ + + event_id: int + prompt: str | None = None + controls: tuple[ARDiffusionControlInput, ...] = () + + def __post_init__(self) -> None: + if isinstance(self.event_id, bool) or not isinstance(self.event_id, int) or self.event_id < 0: + raise ValueError("event_id must be a non-negative integer.") + if self.prompt is not None and not isinstance(self.prompt, str): + raise ValueError("prompt must be a string or None.") + controls = tuple(self.controls) + if any(not isinstance(control, ARDiffusionControlInput) for control in controls): + raise ValueError("controls must contain ARDiffusionControlInput values.") + tracks = [control.track for control in controls] + if len(tracks) != len(set(tracks)): + raise ValueError("A composite event may update each control track at most once.") + if self.prompt is None and not controls: + raise ValueError("An event must update the prompt, one or more controls, or both.") + object.__setattr__(self, "controls", controls) + + +@dataclass(frozen=True) +class ARDiffusionEventAcceptance: + """Immediate queue acceptance result; this does not mean applied.""" + + session_id: str + event_id: int + status: ARDiffusionEventAcceptanceStatus + pending_event_count: int + + +TickOutputT = TypeVar("TickOutputT") + + +class ARDiffusionTickConsumer(Protocol[TickOutputT]): + """Consumes one immutable chunk-boundary tick.""" + + async def execute_tick(self, tick: ARDiffusionTickRequest) -> TickOutputT: ... + + def chunk_metadata(self, output: TickOutputT) -> ARDiffusionChunkMetadata: ... + + +@dataclass(frozen=True) +class ARDiffusionPreparedControls: + """One reducer-owned, transactional control snapshot. + + ``private_state`` is opaque to the session layer and becomes durable only + after the corresponding model tick succeeds. + """ + + controls: tuple[ARDiffusionControlInput, ...] + private_state: Any = None + + +class ARDiffusionControlReducer(Protocol): + """Optional model adapter for stateful controls sampled per AR block.""" + + def prepare( + self, + *, + current_controls: Mapping[str, ARDiffusionControlInput], + events: Sequence[ARDiffusionSessionEvent], + chunk_index: int, + ) -> ARDiffusionPreparedControls: ... + + def commit(self, prepared: ARDiffusionPreparedControls) -> None: ... + + def reset(self) -> None: ... + + +class ARDiffusionSessionLifecycle(Protocol): + """Worker-facing session lifecycle operations.""" + + async def reset_session(self, session_id: str) -> None: ... + + async def close_session(self, session_id: str) -> None: ... + + +class ARDiffusionCollectiveRPCClient(Protocol): + """Subset of the engine control RPC surface used by session lifecycle.""" + + async def collective_rpc( + self, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + stage_ids: list[int] | None = None, + ) -> list[Any]: ... + + +def _collect_rpc_support_results(value: Any, *, errors: list[str], supported: list[bool]) -> None: + if isinstance(value, bool): + supported.append(value) + return + if isinstance(value, (list, tuple)): + for item in value: + _collect_rpc_support_results(item, errors=errors, supported=supported) + return + if isinstance(value, Mapping): + if value.get("supported") is False: + errors.append(str(value.get("error") or "target stage does not support this lifecycle RPC")) + return + if "result" in value: + _collect_rpc_support_results(value["result"], errors=errors, supported=supported) + + +class ARDiffusionWorkerLifecycle: + """Routes reset/close through the existing orchestrator collective RPC.""" + + def __init__( + self, + rpc_client: ARDiffusionCollectiveRPCClient, + *, + stage_ids: Sequence[int], + timeout: float | None = None, + ) -> None: + normalized_stage_ids = tuple(stage_ids) + if not normalized_stage_ids: + raise ValueError("stage_ids must select at least one AR-Diffusion stage.") + if any(isinstance(stage_id, bool) or not isinstance(stage_id, int) or stage_id < 0 for stage_id in stage_ids): + raise ValueError("stage_ids must contain non-negative integers.") + self._rpc_client = rpc_client + self._stage_ids = normalized_stage_ids + self._timeout = timeout + + async def _invoke(self, method: str, session_id: str) -> None: + results = await self._rpc_client.collective_rpc( + method=method, + timeout=self._timeout, + args=(session_id,), + stage_ids=list(self._stage_ids), + ) + errors: list[str] = [] + supported: list[bool] = [] + _collect_rpc_support_results(results, errors=errors, supported=supported) + if errors: + raise ARDiffusionSessionError(f"{method} failed: {'; '.join(errors)}") + if not supported or not all(supported): + raise ARDiffusionSessionError(f"{method} is not supported by every selected worker.") + + async def reset_session(self, session_id: str) -> None: + await self._invoke("reset_ar_diffusion_session", session_id) + + async def close_session(self, session_id: str) -> None: + await self._invoke("close_ar_diffusion_session", session_id) + + +def _default_request_id(session_id: str, chunk_index: int) -> str: + return f"ar-{session_id}-{chunk_index}-{uuid.uuid4().hex}" + + +class ARDiffusionSession(Generic[TickOutputT]): + """Transactional event queue and chunk snapshot for one realtime world.""" + + def __init__( + self, + session_id: str, + *, + tick_consumer: ARDiffusionTickConsumer[TickOutputT], + lifecycle: ARDiffusionSessionLifecycle, + max_pending_events: int, + request_id_factory: Callable[[str, int], str] | None = None, + control_reducer: ARDiffusionControlReducer | None = None, + ) -> None: + if not isinstance(session_id, str) or not session_id.strip(): + raise ValueError("session_id must be a non-empty string.") + if isinstance(max_pending_events, bool) or not isinstance(max_pending_events, int) or max_pending_events <= 0: + raise ValueError("max_pending_events must be a positive integer.") + self.session_id = session_id + self._tick_consumer = tick_consumer + self._lifecycle = lifecycle + self._max_pending_events = max_pending_events + self._request_id_factory = request_id_factory or _default_request_id + self._control_reducer = control_reducer + + self._status = ARDiffusionSessionStatus.ACTIVE + self._pending_events: dict[int, ARDiffusionSessionEvent] = {} + self._recent_event_ids: OrderedDict[int, None] = OrderedDict() + self._event_floor = -1 + self._inflight_event_floor: int | None = None + self._chunk_index = 0 + self._prompt: str | None = None + self._controls: dict[str, ARDiffusionControlInput] = {} + + self._state_lock = asyncio.Lock() + self._tick_lock = asyncio.Lock() + + @property + def status(self) -> ARDiffusionSessionStatus: + return self._status + + @property + def chunk_index(self) -> int: + return self._chunk_index + + @property + def pending_event_count(self) -> int: + return len(self._pending_events) + + def _require_active(self) -> None: + if self._status is not ARDiffusionSessionStatus.ACTIVE: + raise ARDiffusionSessionStateError(f"AR-Diffusion session {self.session_id!r} is {self._status.value}.") + + def _require_resettable(self) -> None: + if self._status not in { + ARDiffusionSessionStatus.ACTIVE, + ARDiffusionSessionStatus.FAILED, + }: + raise ARDiffusionSessionStateError(f"AR-Diffusion session {self.session_id!r} is {self._status.value}.") + + async def accept_event(self, event: ARDiffusionSessionEvent) -> ARDiffusionEventAcceptance: + """Atomically enqueue an event or return an idempotent duplicate result.""" + if not isinstance(event, ARDiffusionSessionEvent): + raise TypeError("event must be an ARDiffusionSessionEvent.") + async with self._state_lock: + self._require_active() + if event.event_id in self._pending_events: + return ARDiffusionEventAcceptance( + session_id=self.session_id, + event_id=event.event_id, + status=ARDiffusionEventAcceptanceStatus.DUPLICATE, + pending_event_count=len(self._pending_events), + ) + if event.event_id in self._recent_event_ids: + return ARDiffusionEventAcceptance( + session_id=self.session_id, + event_id=event.event_id, + status=ARDiffusionEventAcceptanceStatus.DUPLICATE, + pending_event_count=len(self._pending_events), + ) + + event_floor = self._event_floor + if self._inflight_event_floor is not None: + event_floor = max(event_floor, self._inflight_event_floor) + if event.event_id <= event_floor: + raise ARDiffusionStaleEventError( + f"event_id {event.event_id} is stale; the session event floor is {event_floor}." + ) + if len(self._pending_events) >= self._max_pending_events: + raise ARDiffusionEventQueueFullError( + f"AR-Diffusion session {self.session_id!r} event queue is full " + f"({self._max_pending_events} pending)." + ) + + stored_event = ARDiffusionSessionEvent( + event_id=event.event_id, + prompt=event.prompt, + controls=tuple(_copy_control(control) for control in event.controls), + ) + self._pending_events[event.event_id] = stored_event + self._recent_event_ids[event.event_id] = None + while len(self._recent_event_ids) > self._max_pending_events: + self._recent_event_ids.popitem(last=False) + return ARDiffusionEventAcceptance( + session_id=self.session_id, + event_id=event.event_id, + status=ARDiffusionEventAcceptanceStatus.ACCEPTED, + pending_event_count=len(self._pending_events), + ) + + def _build_tick_locked( + self, + ) -> tuple[ + ARDiffusionTickRequest, + str | None, + dict[str, ARDiffusionControlInput], + ARDiffusionPreparedControls | None, + ]: + event_ids = tuple(sorted(self._pending_events)) + events = tuple(self._pending_events[event_id] for event_id in event_ids) + prompt = self._prompt + for event in events: + if event.prompt is not None: + prompt = event.prompt + + prepared_controls: ARDiffusionPreparedControls | None = None + if self._control_reducer is None: + controls = {track: _copy_control(control) for track, control in self._controls.items()} + for event in events: + for control in event.controls: + controls[control.track] = _copy_control(control) + else: + prepared_controls = self._control_reducer.prepare( + current_controls={track: _copy_control(control) for track, control in self._controls.items()}, + events=events, + chunk_index=self._chunk_index, + ) + if not isinstance(prepared_controls, ARDiffusionPreparedControls): + raise TypeError("control reducer prepare() must return ARDiffusionPreparedControls.") + controls = {} + for control in prepared_controls.controls: + if not isinstance(control, ARDiffusionControlInput): + raise TypeError("prepared controls must contain ARDiffusionControlInput values.") + if control.track in controls: + raise ValueError("prepared controls must contain at most one value per track.") + controls[control.track] = _copy_control(control) + + chunk_index = self._chunk_index + request_id = self._request_id_factory(self.session_id, chunk_index) + tick = ARDiffusionTickRequest( + session_id=self.session_id, + request_id=request_id, + chunk_index=chunk_index, + applied_event_ids=event_ids, + prompt=prompt, + controls=tuple(_copy_control(controls[track]) for track in sorted(controls)), + ) + self._inflight_event_floor = event_ids[-1] if event_ids else None + return tick, prompt, controls, prepared_controls + + async def _commit_tick( + self, + tick: ARDiffusionTickRequest, + prompt: str | None, + controls: dict[str, ARDiffusionControlInput], + prepared_controls: ARDiffusionPreparedControls | None, + ) -> None: + async with self._state_lock: + # Reducer state and the session snapshot form one logical commit. + # Run the only fallible operation before mutating session fields so + # a reducer failure cannot acknowledge events or advance the chunk. + if prepared_controls is not None: + assert self._control_reducer is not None + self._control_reducer.commit(prepared_controls) + for event_id in tick.applied_event_ids: + self._pending_events.pop(event_id, None) + if tick.applied_event_ids: + self._event_floor = tick.applied_event_ids[-1] + self._prompt = prompt + self._controls = controls + self._chunk_index += 1 + self._inflight_event_floor = None + + async def _fail_tick(self, tick_error: BaseException) -> None: + async with self._state_lock: + self._inflight_event_floor = None + self._status = ARDiffusionSessionStatus.FAILED + try: + # A model forward may already have released its state. This + # idempotent close also covers successful forwards whose metadata + # or local reducer commit could not be accepted. + await self._lifecycle.close_session(self.session_id) + except Exception as cleanup_error: + raise ARDiffusionSessionError( + "AR-Diffusion tick failed " + f"({type(tick_error).__name__}: {tick_error}) and worker " + f"session cleanup also failed: {cleanup_error}" + ) from tick_error + + async def next_chunk(self) -> TickOutputT: + """Execute and atomically commit one chunk-boundary snapshot.""" + async with self._tick_lock: + async with self._state_lock: + self._require_active() + tick, prompt, controls, prepared_controls = self._build_tick_locked() + + expected_metadata = ARDiffusionChunkMetadata.from_tick(tick) + try: + output = await self._tick_consumer.execute_tick(tick) + metadata = self._tick_consumer.chunk_metadata(output) + if not isinstance(metadata, ARDiffusionChunkMetadata) or metadata != expected_metadata: + raise ARDiffusionChunkMetadataError( + "Tick consumer metadata does not match the submitted AR-Diffusion tick." + ) + except BaseException as tick_error: + await self._fail_tick(tick_error) + raise + + commit_task = asyncio.create_task( + self._commit_tick( + tick, + prompt, + controls, + prepared_controls, + ) + ) + try: + await asyncio.shield(commit_task) + except asyncio.CancelledError: + # Once a consumer reports success, the local chunk/event state + # must commit even if the transport task is being cancelled. + try: + await commit_task + except BaseException as commit_error: + await self._fail_tick(commit_error) + raise + raise + except BaseException as commit_error: + await self._fail_tick(commit_error) + raise + return output + + async def reset(self) -> None: + """Reset worker state and discard queued/current transport state.""" + async with self._tick_lock: + async with self._state_lock: + self._require_resettable() + self._status = ARDiffusionSessionStatus.RESETTING + try: + await self._lifecycle.reset_session(self.session_id) + except BaseException: + async with self._state_lock: + self._status = ARDiffusionSessionStatus.FAILED + raise + async with self._state_lock: + if self._pending_events: + self._event_floor = max(self._event_floor, max(self._pending_events)) + self._pending_events.clear() + self._prompt = None + self._controls.clear() + if self._control_reducer is not None: + self._control_reducer.reset() + self._inflight_event_floor = None + self._chunk_index = 0 + self._status = ARDiffusionSessionStatus.ACTIVE + + async def close(self) -> None: + """Close worker state, retaining a tombstone until cleanup succeeds.""" + async with self._tick_lock: + async with self._state_lock: + if self._status is ARDiffusionSessionStatus.CLOSED: + return + self._status = ARDiffusionSessionStatus.CLOSING + try: + await self._lifecycle.close_session(self.session_id) + except BaseException: + async with self._state_lock: + self._status = ARDiffusionSessionStatus.CLEANUP_FAILED + raise + else: + async with self._state_lock: + self._pending_events.clear() + self._prompt = None + self._controls.clear() + self._inflight_event_floor = None + self._status = ARDiffusionSessionStatus.CLOSED + + async def disconnect(self) -> None: + """Disconnect uses the same owning-worker cleanup as explicit close.""" + await self.close() + + +class ARDiffusionSessionManager(Generic[TickOutputT]): + """Creates and owns transport sessions sharing one tick/lifecycle backend.""" + + def __init__( + self, + *, + tick_consumer: ARDiffusionTickConsumer[TickOutputT], + lifecycle: ARDiffusionSessionLifecycle, + max_pending_events: int, + request_id_factory: Callable[[str, int], str] | None = None, + session_id_factory: Callable[[], str] | None = None, + control_reducer_factory: Callable[[], ARDiffusionControlReducer] | None = None, + ) -> None: + self._tick_consumer = tick_consumer + self._lifecycle = lifecycle + self._max_pending_events = max_pending_events + self._request_id_factory = request_id_factory + self._session_id_factory = session_id_factory or (lambda: uuid.uuid4().hex) + self._control_reducer_factory = control_reducer_factory + self._sessions: dict[str, ARDiffusionSession[TickOutputT]] = {} + self._lock = asyncio.Lock() + + async def create_session(self, session_id: str | None = None) -> ARDiffusionSession[TickOutputT]: + resolved_session_id = session_id or self._session_id_factory() + async with self._lock: + if resolved_session_id in self._sessions: + raise ARDiffusionSessionStateError(f"AR-Diffusion session {resolved_session_id!r} already exists.") + session = ARDiffusionSession( + resolved_session_id, + tick_consumer=self._tick_consumer, + lifecycle=self._lifecycle, + max_pending_events=self._max_pending_events, + request_id_factory=self._request_id_factory, + control_reducer=( + self._control_reducer_factory() if self._control_reducer_factory is not None else None + ), + ) + self._sessions[resolved_session_id] = session + return session + + async def get_session(self, session_id: str) -> ARDiffusionSession[TickOutputT]: + async with self._lock: + try: + return self._sessions[session_id] + except KeyError as exc: + raise ARDiffusionSessionStateError(f"AR-Diffusion session {session_id!r} does not exist.") from exc + + async def reset_session(self, session_id: str) -> None: + session = await self.get_session(session_id) + await session.reset() + + async def close_session(self, session_id: str) -> None: + session = await self.get_session(session_id) + await session.close() + async with self._lock: + if self._sessions.get(session_id) is session: + self._sessions.pop(session_id) + + async def disconnect(self, session_id: str) -> None: + """Release the owning worker route when its transport disconnects.""" + await self.close_session(session_id) diff --git a/vllm_omni/experimental/ar_diffusion/tick_protocol.py b/vllm_omni/experimental/ar_diffusion/tick_protocol.py new file mode 100644 index 000000000..aa8315118 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/tick_protocol.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Typed request boundary shared by AR-Diffusion models and serving APIs. + +The transport layer owns event ordering and snapshots. Model adapters consume +the resulting immutable tick without learning about WebSocket messages. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +AR_DIFFUSION_TICK_KEY = "ar_diffusion_tick" + + +def _non_empty_string(value: object, *, field: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field} must be a non-empty string.") + return value + + +def _non_negative_int(value: object, *, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{field} must be a non-negative integer.") + return value + + +def _freeze_transport_value(value: Any, *, path: str) -> Any: + """Copy JSON-like transport data into recursively immutable containers.""" + if value is None or isinstance(value, (str, bool, int, float)): + return value + if isinstance(value, Mapping): + frozen: dict[str, Any] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise ValueError(f"{path} mapping keys must be strings.") + frozen[key] = _freeze_transport_value(item, path=f"{path}.{key}") + return MappingProxyType(frozen) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return tuple(_freeze_transport_value(item, path=f"{path}[]") for item in value) + raise ValueError( + f"{path} must contain only transport-safe mappings, sequences, and scalar values; got {type(value).__name__}." + ) + + +def _thaw_transport_value(value: Any) -> Any: + """Return a detached JSON-serializable copy of an immutable snapshot.""" + if isinstance(value, Mapping): + return {key: _thaw_transport_value(item) for key, item in value.items()} + if isinstance(value, tuple): + return [_thaw_transport_value(item) for item in value] + return value + + +@dataclass(frozen=True) +class ARDiffusionControlInput: + """One model-defined control track in a chunk-boundary snapshot.""" + + track: str + schema: str + data: Mapping[str, Any] + + def __post_init__(self) -> None: + _non_empty_string(self.track, field="control.track") + _non_empty_string(self.schema, field="control.schema") + if not isinstance(self.data, Mapping): + raise ValueError("control.data must be a mapping.") + object.__setattr__( + self, + "data", + _freeze_transport_value(self.data, path=f"control[{self.track!r}].data"), + ) + + def to_dict(self) -> dict[str, Any]: + return { + "track": self.track, + "schema": self.schema, + "data": _thaw_transport_value(self.data), + } + + @classmethod + def from_dict(cls, value: object) -> ARDiffusionControlInput: + if not isinstance(value, Mapping): + raise ValueError("Each control must be a mapping.") + return cls( + track=value.get("track"), # type: ignore[arg-type] + schema=value.get("schema"), # type: ignore[arg-type] + data=value.get("data"), # type: ignore[arg-type] + ) + + +@dataclass(frozen=True) +class ARDiffusionTickRequest: + """A single atomic AR-block computation. + + ``request_id`` identifies this computation, ``session_id`` identifies the + persistent rollout, and ``applied_event_ids`` records which accepted + interaction events are represented by this immutable snapshot. + """ + + session_id: str + request_id: str + chunk_index: int + applied_event_ids: tuple[int, ...] = () + prompt: str | None = None + controls: tuple[ARDiffusionControlInput, ...] = () + reset: bool = False + close_session: bool = False + + def __post_init__(self) -> None: + _non_empty_string(self.session_id, field="session_id") + _non_empty_string(self.request_id, field="request_id") + _non_negative_int(self.chunk_index, field="chunk_index") + if self.prompt is not None and not isinstance(self.prompt, str): + raise ValueError("prompt must be a string or None.") + if not isinstance(self.reset, bool): + raise ValueError("reset must be a boolean.") + if not isinstance(self.close_session, bool): + raise ValueError("close_session must be a boolean.") + + event_ids = tuple( + _non_negative_int(event_id, field="applied_event_ids item") for event_id in self.applied_event_ids + ) + if event_ids != tuple(sorted(set(event_ids))): + raise ValueError("applied_event_ids must be unique and strictly increasing.") + controls = tuple(self.controls) + if any(not isinstance(control, ARDiffusionControlInput) for control in controls): + raise ValueError("controls must contain ARDiffusionControlInput values.") + object.__setattr__(self, "applied_event_ids", event_ids) + object.__setattr__(self, "controls", controls) + + def to_extra_args(self) -> dict[str, Any]: + """Serialize the tick under one namespaced ``extra_args`` key.""" + return { + AR_DIFFUSION_TICK_KEY: { + "session_id": self.session_id, + "request_id": self.request_id, + "chunk_index": self.chunk_index, + "applied_event_ids": list(self.applied_event_ids), + "prompt": self.prompt, + "controls": [control.to_dict() for control in self.controls], + "reset": self.reset, + "close_session": self.close_session, + } + } + + @classmethod + def from_extra_args( + cls, + extra_args: Mapping[str, Any] | None, + ) -> ARDiffusionTickRequest | None: + """Parse a typed tick, returning ``None`` for legacy AR requests.""" + if extra_args is None or AR_DIFFUSION_TICK_KEY not in extra_args: + return None + value = extra_args[AR_DIFFUSION_TICK_KEY] + if not isinstance(value, Mapping): + raise ValueError(f"{AR_DIFFUSION_TICK_KEY} must be a mapping.") + raw_event_ids = value.get("applied_event_ids", ()) + raw_controls = value.get("controls", ()) + if not isinstance(raw_event_ids, Sequence) or isinstance(raw_event_ids, (str, bytes)): + raise ValueError("applied_event_ids must be a sequence.") + if not isinstance(raw_controls, Sequence) or isinstance(raw_controls, (str, bytes)): + raise ValueError("controls must be a sequence.") + return cls( + session_id=value.get("session_id"), # type: ignore[arg-type] + request_id=value.get("request_id"), # type: ignore[arg-type] + chunk_index=value.get("chunk_index"), # type: ignore[arg-type] + applied_event_ids=tuple(raw_event_ids), + prompt=value.get("prompt"), # type: ignore[arg-type] + controls=tuple(ARDiffusionControlInput.from_dict(control) for control in raw_controls), + reset=value.get("reset", False), # type: ignore[arg-type] + close_session=value.get("close_session", False), # type: ignore[arg-type] + ) + + +@dataclass(frozen=True) +class ARDiffusionChunkMetadata: + """Identity metadata returned with one generated AR block.""" + + session_id: str + request_id: str + chunk_index: int + applied_event_ids: tuple[int, ...] = () + + @classmethod + def from_tick(cls, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata: + return cls( + session_id=tick.session_id, + request_id=tick.request_id, + chunk_index=tick.chunk_index, + applied_event_ids=tick.applied_event_ids, + ) + + def to_dict(self) -> dict[str, Any]: + return { + "session_id": self.session_id, + "request_id": self.request_id, + "chunk_index": self.chunk_index, + "applied_event_ids": list(self.applied_event_ids), + } diff --git a/vllm_omni/experimental/fullduplex/DESIGN.md b/vllm_omni/experimental/fullduplex/DESIGN.md index 256214299..916a2a096 100644 --- a/vllm_omni/experimental/fullduplex/DESIGN.md +++ b/vllm_omni/experimental/fullduplex/DESIGN.md @@ -58,13 +58,13 @@ The checkpoint keeps these verified contracts: - scheduler data-plane append over a resumable request; - stale epoch/turn/response fencing; - `/v1/duplex` and OpenAI Realtime projection; +- optional OpenAI `server_vad` turn detection backed by per-session Silero VAD; - existing JoyVL behavior. The checkpoint does not claim: - scheduler-native KV append; -- deterministic VAD-triggered interruption (the browser intentionally does not - run VAD; MiniCPM owns listen/speak decisions at model-unit boundaries); +- an acoustic-onset-to-cancel latency target across arbitrary client chunk sizes; - production multi-session admission, fairness, capacity, or failure recovery; - bounded long-session KV; - video input or audio/video synchronization. @@ -640,9 +640,8 @@ multi-turn runtime evidence. During auto-response overlap, `preserve_realtime_input` distinguishes "do not append this silent chunk to the native buffer" from "clear the open Realtime -item". Silent overlap no longer discards earlier user PCM. This is an input -ownership correction. It does not add a VAD policy: overlap is admitted to -Stage0 and the model decides whether to listen or speak. +item". Silent overlap no longer discards earlier user PCM. A separate explicit +`server_vad` policy is documented in the Realtime contract above. The first chunk of one overlapping input item also reserves its target model turn. A later Realtime commit uses that reserved identity even if response EOS @@ -1024,6 +1023,6 @@ Passing this checkpoint supports the statement: > Single-session, model-owned MiniCPM-o 4.5 native duplex is reviewable on the > validated H20 configuration. -It does not support claims for deterministic VAD-triggered interruption, -multi-session production concurrency, bounded long-session KV, -scheduler-native append, or video input. +It does not support claims for an acoustic-onset barge-in latency target across +arbitrary client chunk sizes, multi-session production concurrency, bounded +long-session KV, scheduler-native append, or video input. diff --git a/vllm_omni/experimental/fullduplex/client.py b/vllm_omni/experimental/fullduplex/client.py index 97391a5c9..7c114f6c8 100644 --- a/vllm_omni/experimental/fullduplex/client.py +++ b/vllm_omni/experimental/fullduplex/client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Reusable Realtime WebSocket, PCM, and event helpers for MiniCPM-o demos.""" from __future__ import annotations @@ -8,7 +11,7 @@ import json import math import time import wave -from collections.abc import Callable +from collections.abc import Callable, Iterator, Sequence from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -25,11 +28,32 @@ except ImportError as exc: # pragma: no cover - example dependency PCM16_SAMPLE_RATE = 16_000 PCM16_BYTES_PER_SAMPLE = 2 +# Server-side model unit boundaries, in cumulative appended audio. +# Stage0 configures the streaming mel processor with first_chunk_ms=1035 and +# chunk_ms=1000; the processor aligns the first chunk down to a hop_length (160 +# samples) multiple, so unit 0 closes at 16480 samples and every later unit +# closes 16000 samples after it. Camera frames must ride the append that closes +# a unit, otherwise Stage0 cannot bind them to that unit's audio. +DUPLEX_FIRST_UNIT_MS = 1030 +DUPLEX_UNIT_MS = 1000 + + +def duplex_unit_boundary_ms(unit_index: int) -> int: + """Cumulative appended audio, in ms, that closes model unit ``unit_index``.""" + return DUPLEX_FIRST_UNIT_MS + max(0, int(unit_index)) * DUPLEX_UNIT_MS + def _rounded_ms(value: float) -> float: return round(float(value), 3) +def _finite_number(value: object, *, nonnegative: bool = False) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + parsed = float(value) + return parsed if math.isfinite(parsed) and (not nonnegative or parsed >= 0) else None + + def _interval_summary(values: list[float]) -> dict[str, float | int]: clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0) if not clean: @@ -56,11 +80,7 @@ def summarize_session_request_metrics( """Average client-observed metrics across turns that emitted audio.""" def mean(metric: str, *, digits: int = 3) -> float | None: - values = [ - float(request[metric]) - for request in request_metrics - if isinstance(request.get(metric), int | float) and math.isfinite(float(request[metric])) - ] + values = [value for request in request_metrics if (value := _finite_number(request.get(metric))) is not None] return round(sum(values) / len(values), digits) if values else None return { @@ -93,24 +113,57 @@ def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None: def build_realtime_url( url: str, - model: str, + model: str | None, *, autostart: bool | None = None, + native_duplex: bool | None = True, session_id: str | None = None, ) -> str: """Add the explicit native-duplex query parameters to a Realtime URL.""" parts = urlsplit(url) + if parts.scheme in {"http", "https"}: + parts = parts._replace(scheme="ws" if parts.scheme == "http" else "wss") + if parts.scheme not in {"ws", "wss"} or not parts.netloc: + raise ValueError(f"Unsupported Realtime URL: {url!r}") query = dict(parse_qsl(parts.query, keep_blank_values=True)) - query.setdefault("duplex", "1") - query.setdefault("model", model) - query.setdefault("minicpmo45_native_duplex", "1") + query["duplex"] = "1" + if model: + query["model"] = model + if native_duplex is not None: + query["minicpmo45_native_duplex"] = "1" if native_duplex else "0" if autostart is not None: - query.setdefault("autostart", "1" if autostart else "0") + query["autostart"] = "1" if autostart else "0" if session_id: - query.setdefault("session_id", session_id) + query["session_id"] = session_id return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)) +def reference_audio_data_url(path: str | None) -> str | None: + """Encode a local reference WAV for a Realtime session update.""" + if path is None: + return None + audio = Path(path).expanduser().resolve() + if not audio.is_file(): + raise FileNotFoundError(f"Reference audio does not exist: {audio}") + return "data:audio/wav;base64," + base64.b64encode(audio.read_bytes()).decode("ascii") + + +def chunk_period_ms(events: list[dict[str, object]], *, default: int = 1000) -> int: + """Read the negotiated native-duplex model-unit duration.""" + for event in reversed(events): + session = event.get("session") + capabilities = session.get("capabilities") if isinstance(session, dict) else None + value = capabilities.get("chunk_period_ms") if isinstance(capabilities, dict) else None + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value + return default + + +def has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool: + unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000 + return bool(unit_bytes and len(pcm16) % unit_bytes) + + def read_pcm16_wav(path: Path) -> bytes: """Read a mono, uncompressed, 16 kHz PCM16 WAV file.""" with wave.open(str(path), "rb") as wav_file: @@ -202,6 +255,25 @@ class RealtimeEventCollector: chunk for response_id in self.response_ids for chunk in self.response_audio.get(response_id, ()) ) + def response_text(self, response_id: str) -> str: + """Join all text/transcript deltas for one response identity.""" + return "".join( + str(event.get("delta") or "") + for event in self.events + if self.response_id(event) == response_id + and event.get("type") + in { + "response.audio_transcript.delta", + "response.output_text.delta", + "response.text.delta", + } + ) + + def response_is_done(self, response_id: str) -> bool: + return any( + event.get("type") == "response.done" and self.response_id(event) == response_id for event in self.events + ) + def errors(self) -> list[dict[str, object]]: return [event for event in self.events if event.get("type") == "error"] @@ -282,15 +354,16 @@ class RealtimeEventCollector: if stage0_metrics is not None: raw_itls = stage0_metrics.get("vllm_itls_ms") itls = ( - [float(value) for value in raw_itls if isinstance(value, int | float)] + [parsed for value in raw_itls if (parsed := _finite_number(value, nonnegative=True)) is not None] if isinstance(raw_itls, list) else [] ) result["stage0_tokens"] = { "source": "engine_stage_metrics", - "output_token_count": int(stage0_metrics.get("num_tokens_out") or 0), - "ttft_ms": float(stage0_metrics.get("vllm_ttft_ms") or 0.0), - "tpot_ms": float(stage0_metrics.get("vllm_tpot_ms") or 0.0), + "output_token_count": int(_finite_number(stage0_metrics.get("num_tokens_out"), nonnegative=True) or 0), + "ttft_ms": _finite_number(stage0_metrics.get("vllm_ttft_ms"), nonnegative=True) or 0.0, + "tpot_ms": _finite_number(stage0_metrics.get("vllm_tpot_ms"), nonnegative=True), + "itls_ms": itls, "inter_token_interval_ms": _interval_summary(itls), } @@ -368,15 +441,27 @@ class RealtimeEventCollector: class RealtimeDuplexClient: """Small async client used by the user demo and reusable smoke probes.""" - def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None: + def __init__( + self, + url: str, + *, + max_size: int = 64 * 1024 * 1024, + additional_headers: dict[str, str] | None = None, + ) -> None: self.url = url self.max_size = max_size + self.additional_headers = additional_headers self.events = RealtimeEventCollector() self._ws: Any = None self._reader_task: asyncio.Task[None] | None = None + self._media_clock_ms = 0 async def __aenter__(self) -> RealtimeDuplexClient: - self._ws = await websockets.connect(self.url, max_size=self.max_size) + self._ws = await websockets.connect( + self.url, + max_size=self.max_size, + additional_headers=self.additional_headers, + ) self._reader_task = asyncio.create_task(self._read_events()) return self @@ -398,6 +483,7 @@ class RealtimeDuplexClient: continue event = json.loads(raw) if isinstance(event, dict): + event.setdefault("_media_clock_ms", self._media_clock_ms) self.events.add(event) except ConnectionClosed: return @@ -405,6 +491,19 @@ class RealtimeDuplexClient: async def send(self, event: dict[str, object]) -> None: await self._ws.send(json.dumps(event)) + def raise_if_reader_stopped(self) -> None: + """Fail a caller waiting on events after the WebSocket reader exits.""" + + task = self._reader_task + if task is None or not task.done(): + return + if task.cancelled(): + raise ConnectionError("Realtime WebSocket reader was cancelled") + error = task.exception() + if error is not None: + raise ConnectionError("Realtime WebSocket reader failed") from error + raise ConnectionError("Realtime WebSocket closed before the requested event arrived") + async def configure( self, model: str, @@ -417,7 +516,9 @@ class RealtimeDuplexClient: auto_response: bool = True, temperature: float | None = None, extra_body: dict[str, object] | None = None, + turn_detection: dict[str, object] | None = None, session_id: str | None = None, + idle_timeout_s: float | None = None, timeout_s: float = 20.0, ) -> None: session_extra_body = dict(extra_body or {}) @@ -436,8 +537,12 @@ class RealtimeDuplexClient: "modalities": ["audio", "text"], "input_audio_format": "pcm16", "output_audio_format": output_audio_format, - "turn_detection": None, - "overlap_policy": "listen_only", + "turn_detection": dict(turn_detection) if turn_detection is not None else None, + "overlap_policy": ( + "barge_in_on_speech" + if turn_detection is not None and turn_detection.get("interrupt_response", True) is True + else "listen_only" + ), "playback_commit_policy": "ack_only", "extra_body": session_extra_body, } @@ -448,14 +553,23 @@ class RealtimeDuplexClient: if instructions is not None: session["instructions"] = instructions if initial_user_text is not None: - extra_body = session["extra_body"] - assert isinstance(extra_body, dict) - extra_body["duplex_initial_user_text"] = initial_user_text + session_extra = session["extra_body"] + assert isinstance(session_extra, dict) + session_extra["duplex_initial_user_text"] = initial_user_text if session_id: session["session_id"] = session_id + if idle_timeout_s is not None: + session["idle_timeout_s"] = idle_timeout_s await self.send({"type": "session.update", "session": session}) + + def session_created() -> bool: + if self.events.count("session.created") > 0: + return True + self.raise_if_reader_stopped() + return False + await wait_for( - lambda: self.events.count("session.created") > 0, + session_created, timeout_s=timeout_s, label="session.created", ) @@ -466,28 +580,97 @@ class RealtimeDuplexClient: *, chunk_ms: int = 200, realtime: bool = True, - ) -> None: + video_frames: Sequence[str] | None = None, + stacked_video_frames: Sequence[str | None] | None = None, + ) -> int: + """Append PCM16 audio, optionally interleaving omni camera frames. + + ``video_frames`` holds base64 JPEG/PNG frames in capture order, one per + second of the clip. Frame ``k`` rides the append that closes model unit + ``k`` (see ``duplex_unit_boundary_ms``), which reproduces the official + ``streaming_prefill(audio_waveform=<1 s>, frame_list=[frame])`` pairing: + a second of audio and the picture captured during it enter the same + unit. Sending on whole-second boundaries instead would strand frame 0 on + an append that cannot close a unit yet, and shift every later frame one + unit ahead of its audio. + + ``stacked_video_frames`` is the optional parallel track of composites + built by ``video_stacking.concat_frames``: entry ``k`` tiles the + sub-frames captured *inside* unit ``k``, and rides the same append right + after the base frame, giving the official ``frame_list=[base, + composite]``. The audio is untouched — a unit stays one second however + many sub-frames the composite carries. ``None`` entries send the base + frame alone. + + Returns the number of base frames actually sent (a clip shorter than the + frame list leaves the tail unsent). + """ chunk_bytes = max( PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_ms // 1000, PCM16_BYTES_PER_SAMPLE, ) + frames = list(video_frames or []) + stacked = list(stacked_video_frames or []) + + def units() -> Iterator[tuple[bytes, list[str] | None]]: + audio_end_ms = 0 + frames_sent = 0 + for offset in range(0, len(pcm16), chunk_bytes): + chunk = pcm16[offset : offset + chunk_bytes] + audio_end_ms += len(chunk) * 1000 // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE) + if not frames or audio_end_ms < duplex_unit_boundary_ms(frames_sent): + yield chunk, None + continue + # A video advances one frame per unit and holds its last frame + # when the audio outlives the clip; a still image is a + # one-element list and therefore repeats. + index = min(frames_sent, len(frames) - 1) + composite = stacked[index] if index < len(stacked) else None + frames_sent += 1 + yield chunk, [frames[index]] if composite is None else [frames[index], composite] + + return await self.stream_av_units(units(), realtime=realtime) + + async def stream_av_units( + self, + units: Any, + *, + realtime: bool = True, + ) -> int: + """Stream PCM16 units, optionally attaching camera frames to each unit. + + A unit's frame slot takes a single JPEG/PNG (raw bytes or base64) or a + sequence of them, which the Realtime wire caps at two per append. + Returns the number of appends that carried at least one frame. + """ audio_end_ms = 0 - for offset in range(0, len(pcm16), chunk_bytes): - chunk = pcm16[offset : offset + chunk_bytes] + frames_sent = 0 + for chunk, frame in units: + if not chunk: + continue duration_ms = len(chunk) * 1000 // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE) audio_end_ms += duration_ms - await self.send( - { - "type": "input_audio_buffer.append", - "audio": base64.b64encode(chunk).decode("ascii"), - "input_audio_format": "pcm16", - "sample_rate_hz": PCM16_SAMPLE_RATE, - "duration_ms": duration_ms, - "audio_end_ms": audio_end_ms, - } - ) + self._media_clock_ms = audio_end_ms + event: dict[str, object] = { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(chunk).decode("ascii"), + "input_audio_format": "pcm16", + "sample_rate_hz": PCM16_SAMPLE_RATE, + "duration_ms": duration_ms, + "audio_end_ms": audio_end_ms, + } + encoded_frames = [ + base64.b64encode(item).decode("ascii") if isinstance(item, bytes | bytearray) else item + for item in (frame if isinstance(frame, list | tuple) else [frame]) + if item is not None + ] + if encoded_frames: + event["video_frames"] = encoded_frames + frames_sent += 1 + await self.send(event) if realtime: await asyncio.sleep(duration_ms / 1000) + return frames_sent async def commit(self) -> None: await self.send({"type": "input_audio_buffer.commit", "final": True}) @@ -495,23 +678,34 @@ class RealtimeDuplexClient: async def acknowledge_playback(self) -> None: for response_id in self.events.response_ids: pcm16 = self.events.audio_bytes(response_id) - if not pcm16: + if not pcm16 and self.events.response_is_done(response_id): continue played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE) - await self.send( - { - "type": "playback.ack", - "response_id": response_id, - "item_id": f"item_{response_id}", - "played_ms": played_ms, - "committed_ms": played_ms, - } - ) + await self.send_playback_ack(response_id, played_ms) + + async def send_playback_ack(self, response_id: str, played_ms: int) -> None: + await self.send( + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": played_ms, + "committed_ms": played_ms, + } + ) async def close_session(self, *, timeout_s: float = 20.0) -> None: + from_index = len(self.events.events) await self.send({"type": "session.close"}) + + def session_closed() -> bool: + if any(event.get("type") == "session.closed" for event in self.events.events[from_index:]): + return True + self.raise_if_reader_stopped() + return False + await wait_for( - lambda: self.events.count("session.closed") > 0, + session_closed, timeout_s=timeout_s, label="session.closed", ) diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py b/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py index 83ef48b1e..0afffa03c 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/adapter.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import base64 @@ -11,6 +14,7 @@ from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45Duplex from vllm_omni.experimental.fullduplex.openai.protocol import DuplexSessionConfig from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ( ServingRuntimeConfigError, + reject_changed_runtime_value, ) @@ -66,7 +70,13 @@ class MiniCPMO45NativeDuplexServingAdapter: current: object, ) -> dict[str, object]: runtime_config = deepcopy(dict(current)) if isinstance(current, dict) else {} - runtime_config["instructions"] = config.instructions + reject_changed_runtime_value( + config.instructions, + runtime_config.get("instructions"), + message="instructions cannot be changed after the session is created", + code="instructions_update_unsupported", + error_cls=MiniCPMO45ClientRuntimeConfigError, + ) stage_max_tokens = runtime_config.get("duplex_stage_max_tokens") stage_max_tokens = deepcopy(stage_max_tokens) if isinstance(stage_max_tokens, dict) else {} stage_max_tokens["0"] = ( @@ -97,10 +107,14 @@ class MiniCPMO45NativeDuplexServingAdapter: cls._apply_default_scheduler_policy(runtime_config, config=config, model_config=model_config) ref_audio = config.ref_audio - if ref_audio is None and isinstance(extra_body.get("ref_audio"), str): - ref_audio = extra_body.pop("ref_audio") - if ref_audio is None and isinstance(extra_body.get("tts_ref_audio"), str): - ref_audio = extra_body.pop("tts_ref_audio") + extra_ref_audio = extra_body.get("ref_audio") + if ref_audio is None and isinstance(extra_ref_audio, str): + extra_body.pop("ref_audio") + ref_audio = extra_ref_audio + extra_tts_ref_audio = extra_body.get("tts_ref_audio") + if ref_audio is None and isinstance(extra_tts_ref_audio, str): + extra_body.pop("tts_ref_audio") + ref_audio = extra_tts_ref_audio if ref_audio is None: if any(str(modality).lower() == "audio" for modality in config.modalities): @@ -161,7 +175,11 @@ class MiniCPMO45NativeDuplexServingAdapter: stop_token_ids = cls._native_stage0_stop_token_ids(model_config) if stop_token_ids: stage0_params["stop_token_ids"] = stop_token_ids - runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params} + # Stage 1 keeps the deploy YAML's codec knobs, minus upstream's + # min_new_token=50: that floor is for whole-utterance chat TTS, while a + # duplex chunk is 26 codec samples, so check_stop would never release + # the Talker and Thinker would stall on the next model turn. + runtime_config["duplex_stage_sampling_params"] = {"0": stage0_params, "1": {"min_tokens": 0}} scheduler_token_id = cls._native_scheduler_token_id(model_config) if scheduler_token_id is not None: runtime_config["duplex_scheduler_token_id"] = scheduler_token_id @@ -236,6 +254,8 @@ class MiniCPMO45NativeDuplexServingAdapter: if token_id is not None: return token_id eos_id = getattr(tokenizer, "eos_token_id", None) + if eos_id is None: + return None try: return int(eos_id) except (TypeError, ValueError): @@ -265,10 +285,13 @@ class MiniCPMO45NativeDuplexServingAdapter: value = convert(token) if isinstance(value, list): value = value[0] if len(value) == 1 else None - try: - token_id = int(value) - except (TypeError, ValueError): + if value is None: token_id = -1 + else: + try: + token_id = int(value) + except (TypeError, ValueError): + token_id = -1 unk_token_id = getattr(tokenizer, "unk_token_id", None) if token_id >= 0 and token_id != unk_token_id: return token_id diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/policy.py b/vllm_omni/experimental/fullduplex/minicpmo45/policy.py index d869e54c9..7a1779617 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/policy.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/policy.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from typing import Any @@ -24,7 +27,9 @@ class MiniCPMO45DuplexPolicy: # Vision framing contract (omni duplex). Official streaming_prefill feeds # each frame as + 64 resampler embeddings + inside the # unit, ahead of the unit's audio embeddings (max_slice_nums=1 in - # streaming, so exactly one 64-token block per frame). + # streaming, so exactly one 64-token block per frame). A unit may carry + # its base frame plus an optional stacked composite of the sub-frames + # captured inside that unit (2 blocks). VISION_EMBEDS_PER_FRAME = 64 VISION_TOKENS_PER_FRAME = VISION_EMBEDS_PER_FRAME + 2 # + embeds + DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK = 20 @@ -77,6 +82,8 @@ class MiniCPMO45DuplexPolicy: "audio_placeholder_token_id": "<|audio|>", "image_start_token_id": "", "image_end_token_id": "", + "slice_start_token_id": "", + "slice_end_token_id": "", } @classmethod @@ -92,7 +99,7 @@ class MiniCPMO45DuplexPolicy: value = value[0] if len(value) == 1 else None unk_token_id = getattr(tokenizer, "unk_token_id", None) try: - candidate = int(value) + candidate = int(value) if value is not None else -1 except (TypeError, ValueError): candidate = -1 if candidate >= 0 and candidate != unk_token_id: @@ -118,7 +125,8 @@ class MiniCPMO45DuplexPolicy: if resolved.get("listen_token_id", -1) < 0: eos_id = getattr(tokenizer, "eos_token_id", None) try: - resolved["listen_token_id"] = int(eos_id) + if eos_id is not None: + resolved["listen_token_id"] = int(eos_id) except (TypeError, ValueError): pass return resolved diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/runtime.py b/vllm_omni/experimental/fullduplex/minicpmo45/runtime.py index b5fbd37a5..ecd634f73 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/runtime.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/runtime.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations from base64 import b64decode from binascii import Error as BinasciiError -from typing import Any +from typing import Any, cast from vllm.sampling_params import SamplingParams @@ -22,6 +22,9 @@ _DUPLEX_SAMPLES_PER_AUDIO_TOKEN = 1600 # + 64 resampler embeddings + per frame (max_slice_nums=1), # matching MiniCPMO45DuplexPolicy.VISION_TOKENS_PER_FRAME. _DUPLEX_VISION_TOKENS_PER_FRAME = 66 +# Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD +# sliced (1 source + 2 patches on 960x540) and the composite is not. +_DUPLEX_HD_SLICES_PER_BASE_FRAME = 3 def _duplex_frame_count(payload: object) -> int: @@ -33,6 +36,22 @@ def _duplex_frame_count(payload: object) -> int: return sum(1 for frame in frames if isinstance(frame, str) and frame) +def _duplex_vision_tokens(payload: object) -> int: + """Scheduler slots for this append's camera track. + + Audio is never stacked: a unit still carries one second of soundtrack. + ``stack_frames`` only adds a second *image*. Official HD on that pair is + ``[2, 1]``, so the base frame reserves three 66-token blocks and every + extra frame reserves one. + """ + count = _duplex_frame_count(payload) + if count <= 0: + return 0 + if count >= 2: + return (_DUPLEX_HD_SLICES_PER_BASE_FRAME + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME + return count * _DUPLEX_VISION_TOKENS_PER_FRAME + + def _duplex_pcm_sample_count(payload: object) -> int | None: if not isinstance(payload, dict): return None @@ -48,7 +67,7 @@ def _duplex_pcm_sample_count(payload: object) -> int | None: def duplex_payload_is_exact_chunks(payload: object) -> bool: sample_count = _duplex_pcm_sample_count(payload) - return bool(sample_count) and sample_count % _DUPLEX_CHUNK_SAMPLES == 0 + return sample_count is not None and sample_count != 0 and sample_count % _DUPLEX_CHUNK_SAMPLES == 0 def duplex_first_append_unit_count(payload: object) -> int | None: @@ -59,7 +78,7 @@ def duplex_first_append_unit_count(payload: object) -> int | None: def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int: - vision_tokens = _duplex_frame_count(payload) * _DUPLEX_VISION_TOKENS_PER_FRAME + vision_tokens = _duplex_vision_tokens(payload) sample_count = _duplex_pcm_sample_count(payload) if sample_count is None: return max(1, int(default)) + vision_tokens @@ -114,9 +133,7 @@ def build_duplex_data_plane_prompt( token_budget += context_reserve first_units = duplex_first_append_unit_count(payload) if first_units is not None: - token_budget = ( - context_reserve + first_units * 12 - 1 + _duplex_frame_count(payload) * _DUPLEX_VISION_TOKENS_PER_FRAME - ) + token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload) if seq > 1 and duplex_payload_is_exact_chunks(payload): token_budget += 1 if final and duplex_payload_is_exact_chunks(payload): @@ -124,7 +141,7 @@ def build_duplex_data_plane_prompt( extra_body = session_config.get("extra_body") raw_token_id = runtime_config.get("duplex_scheduler_token_id") try: - token_id = max(0, int(raw_token_id)) + token_id = 0 if raw_token_id is None else max(0, int(raw_token_id)) except (TypeError, ValueError): token_id = 0 force_listen_count = _duplex_force_listen_count(extra_body) @@ -163,16 +180,17 @@ def build_duplex_data_plane_prompt( def _coerce_int(value: object) -> int | None: - if hasattr(value, "detach"): + detach = getattr(value, "detach", None) + if callable(detach): try: - value = value.detach().cpu().reshape(-1) - if value.numel() == 0: + flat: Any = detach().cpu().reshape(-1) + if flat.numel() == 0: return None - value = value[0].item() + value = flat[0].item() except Exception: return None try: - return int(value) + return int(cast(Any, value)) except (TypeError, ValueError): return None diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/session.py b/vllm_omni/experimental/fullduplex/minicpmo45/session.py index 763bcf8b3..1848ee81b 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/session.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/session.py @@ -16,6 +16,7 @@ class MiniCPMO45ServingSessionState: audio_buffer: MiniCPMO45PcmAppendBuffer = field(default_factory=MiniCPMO45PcmAppendBuffer) input_since_commit: bool = False speech_since_commit: bool = False + native_context_locked: bool = False committed_audio_payload: dict[str, object] | None = None committed_audio_operation_id: str | None = None committed_audio_reserved_bytes: int = 0 diff --git a/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py b/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py index 392484c98..fbea82188 100644 --- a/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py +++ b/vllm_omni/experimental/fullduplex/minicpmo45/stage0.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import base64 @@ -42,6 +45,20 @@ class _MiniCPMO45Stage0SessionState: class MiniCPMO45Stage0DuplexRuntime: """Build scheduler-owned MiniCPM-o 4.5 Stage0 duplex inputs.""" + unit_token_id: int = -1 + unit_end_token_id: int = -1 + listen_token_id: int = -1 + speak_token_id: int = -1 + tts_bos_token_id: int = -1 + tts_eos_token_id: int = -1 + tts_pad_token_id: int = -1 + chunk_eos_token_id: int = -1 + chunk_tts_eos_token_id: int = -1 + turn_eos_token_id: int = -1 + audio_placeholder_token_id: int = -1 + image_start_token_id: int = -1 + image_end_token_id: int = -1 + def __init__(self, stage_model: Any, *, model_path: str | None = None, device: str = "cuda") -> None: self.stage_model = stage_model self.model_path = model_path @@ -81,7 +98,7 @@ class MiniCPMO45Stage0DuplexRuntime: processor._streaming_mel_processor = copy.deepcopy(shared_mel) state.streaming_processor = processor if processor is None: - return + return None set_streaming_mode = getattr(processor, "set_streaming_mode", None) if callable(set_streaming_mode): set_streaming_mode( @@ -175,24 +192,40 @@ class MiniCPMO45Stage0DuplexRuntime: self._require_special_token_ids() if audio_waveform is None or len(audio_waveform) == 0: return self._stage_prefill_result(False, start_time, "empty audio") - # Omni duplex: encode this append's camera frames up front so the unit - # loop can interleave one block per unit, mirroring official - # streaming_prefill (feed , then image embeds, then audio). - frame_blocks: list[Any] = [] - if video_frames: - self._require_vision_token_ids() - frame_blocks = self._stage_vision_embeddings(video_frames) - if frame_blocks is None or len(frame_blocks) != len(video_frames): - return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") state.audio_buffer = np.concatenate([state.audio_buffer, np.asarray(audio_waveform, dtype=np.float32)]) chunk_size = self._streaming_chunk_size(processor) self._pad_first_audio_chunk_if_needed(state, processor) if len(state.audio_buffer) < chunk_size: + if video_frames: + # Frames belong to the unit this append closes, and the + # scheduler already reserved their vision tokens against this + # append. A frame arriving before its unit can close has + # nowhere to go, so report it instead of dropping it under a + # generic buffering result. + return self._stage_prefill_result( + False, + start_time, + f"{len(video_frames)} video frame(s) on an append that closes no model unit " + f"(need {chunk_size} samples, only {len(state.audio_buffer)}); " + "send frames on unit boundaries", + ) return self._stage_prefill_result( False, start_time, f"audio not enough: need {chunk_size} samples, only {len(state.audio_buffer)}", ) + # Omni duplex: encode this append's camera frames so the first unit can + # carry them, mirroring official streaming_prefill (feed , then + # every frame_list entry as its own block, then the audio). + frame_blocks: list[Any] = [] + if video_frames: + self._require_vision_token_ids() + encoded_frames = self._stage_vision_embeddings(video_frames) + if encoded_frames is None or len(encoded_frames) != len(video_frames): + return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") + if any(not slices for slices in encoded_frames): + return self._stage_prefill_result(False, start_time, "streaming vision embedding failed") + frame_blocks = encoded_frames embed_parts: list[Any] = [] token_ids: list[int] = [] @@ -246,16 +279,22 @@ class MiniCPMO45Stage0DuplexRuntime: embed_parts.append(self._embed_token(self.unit_token_id)) token_ids.append(self.unit_token_id) if frame_blocks: - # Official order inside a unit: + 64 resampler - # embeddings + ahead of the unit's audio embeddings - # (max_slice_nums=1 in streaming, one frame per unit). - vision_block = self._as_2d_tensor(frame_blocks.pop(0)) - embed_parts.append(self._embed_token(self.image_start_token_id)) - token_ids.append(int(self.image_start_token_id)) - embed_parts.append(vision_block) - token_ids.extend([self._vision_embedding_placeholder_token_id()] * int(vision_block.shape[0])) - embed_parts.append(self._embed_token(self.image_end_token_id)) - token_ids.append(int(self.image_end_token_id)) + # Official order inside a unit: every frame_list entry as + # + 64 embeds + , then any HD patches as + # + 64 embeds + , all ahead of this unit's + # *one* second of audio. stack_frames never stacks audio. + vision_placeholder = self._vision_embedding_placeholder_token_id() + for frame_slices in frame_blocks: + for slice_index, frame_block in enumerate(frame_slices): + vision_block = self._as_2d_tensor(frame_block) + start_id, end_id = self._vision_slice_token_ids(is_source=slice_index == 0) + embed_parts.append(self._embed_token(start_id)) + token_ids.append(int(start_id)) + embed_parts.append(vision_block) + token_ids.extend([vision_placeholder] * int(vision_block.shape[0])) + embed_parts.append(self._embed_token(end_id)) + token_ids.append(int(end_id)) + frame_blocks = [] embed_parts.append(audio_embeds) token_ids.extend( [self._audio_embedding_placeholder_token_id()] * int(self._as_2d_tensor(audio_embeds).shape[0]) @@ -264,14 +303,6 @@ class MiniCPMO45Stage0DuplexRuntime: state.audio_chunk_idx += 1 units_built += 1 chunk_size = self._streaming_chunk_size(processor) - if frame_blocks: - # Each frame reserves one image block in the append plan. A frame - # without a matching audio unit would desynchronize that plan. - return self._stage_prefill_result( - False, - start_time, - f"{len(frame_blocks)} video frame(s) left without a matching audio unit", - ) # Match official streaming_prefill: per chunk feed ONLY +audio. The assistant # turn is opened once at session init; re-emitting the turn-open prefix per chunk # re-opened the turn each chunk -> degenerate repetition. tts_bos/listen/turn_eos are @@ -331,6 +362,20 @@ class MiniCPMO45Stage0DuplexRuntime: embeds = embedder(token) return self._as_2d_tensor(embeds) + def _token_embedding_dtype(self) -> Any: + """dtype of the decoder token embeddings, the unit's reference dtype. + + Official ``get_vllm_embedding`` casts vision hidden states to the token + embedding dtype before they enter the sequence, so the vision tower's + own dtype must not leak into the concatenated unit. + """ + embedder = self._token_embedder() + weight = getattr(embedder, "weight", None) + dtype = getattr(weight, "dtype", None) + if dtype is not None: + return dtype + return getattr(next(self.thinker.parameters()), "dtype", None) + def _token_embedder(self) -> Any: nested_embed = getattr(getattr(getattr(self.thinker, "llm", None), "model", None), "embed_tokens", None) if callable(nested_embed): @@ -392,7 +437,7 @@ class MiniCPMO45Stage0DuplexRuntime: ) if len(state.audio_buffer) >= first_chunk_samples: return - padding = np.zeros(first_chunk_samples - len(state.audio_buffer), dtype=np.float32) + padding: np.ndarray = np.zeros(first_chunk_samples - len(state.audio_buffer), dtype=np.float32) state.audio_buffer = np.concatenate([padding, state.audio_buffer]) def _stage_param(self, name: str, default: Any) -> Any: @@ -454,7 +499,7 @@ class MiniCPMO45Stage0DuplexRuntime: previous_audio_past_key_values = ( getattr(self.thinker, "audio_past_key_values", None) if has_audio_cache else None ) - if has_audio_cache: + if has_audio_cache and state is not None: self.thinker.audio_past_key_values = state.audio_past_key_values try: for target in (self.stage_model, self.thinker): @@ -469,18 +514,18 @@ class MiniCPMO45Stage0DuplexRuntime: suffix_extra_frames=2, ) ) - if has_audio_cache: + if has_audio_cache and state is not None: state.audio_past_key_values = getattr(self.thinker, "audio_past_key_values", None) return result except TypeError: result = self._cat_nested_tensors(get_streaming(batch_feature)) - if has_audio_cache: + if has_audio_cache and state is not None: state.audio_past_key_values = getattr(self.thinker, "audio_past_key_values", None) return result get_hidden = getattr(target, "get_audio_hidden_states", None) if callable(get_hidden): result = self._cat_nested_tensors(get_hidden(batch_feature)) - if has_audio_cache: + if has_audio_cache and state is not None: state.audio_past_key_values = getattr(self.thinker, "audio_past_key_values", None) return result return None @@ -627,14 +672,14 @@ class MiniCPMO45Stage0DuplexRuntime: return original_register(config_class, *args, **kwargs) with _MINICPMO45_PROCESSOR_LOAD_LOCK: - AutoImageProcessor.register = staticmethod(register_image_processor) + setattr(AutoImageProcessor, "register", staticmethod(register_image_processor)) try: processor = AutoProcessor.from_pretrained( model_path, trust_remote_code=True, ) finally: - AutoImageProcessor.register = staticmethod(original_register) + setattr(AutoImageProcessor, "register", staticmethod(original_register)) except Exception as exc: raise RuntimeError(f"Failed to load MiniCPM-o duplex processor from {model_path!r}") from exc if getattr(processor, "tokenizer", None) is None: @@ -764,21 +809,24 @@ class MiniCPMO45Stage0DuplexRuntime: f"missing or unknown: {', '.join(missing)}" ) - def _stage_vision_embeddings(self, frames: list[Any]) -> list[Any] | None: - """Encode camera frames for omni duplex via the loaded vision tower. + def _vision_slice_token_ids(self, *, is_source: bool) -> tuple[int, int]: + if is_source: + return int(self.image_start_token_id), int(self.image_end_token_id) + start = getattr(self, "slice_start_token_id", None) + end = getattr(self, "slice_end_token_id", None) + if isinstance(start, int) and start >= 0 and isinstance(end, int) and end >= 0: + return start, end + return int(self.image_start_token_id), int(self.image_end_token_id) - Semantics mirror ``MiniCPMODuplex.streaming_prefill`` (one 64-embedding - block per frame at ``max_slice_nums=1``), executed through the vLLM - wrapper's ``get_vision_hidden_states`` (vpm + resampler) since the - stage model does not expose the remote-code ``get_vision_embedding``. - """ - process_image = getattr(self.processor, "process_image", None) - if not callable(process_image): - return None - try: - processed = process_image(frames, max_slice_nums=1) - except Exception: # noqa: BLE001 - prefill fails with a reason - return None + def _official_max_slice_nums(self, frame_count: int) -> list[int]: + """Official stacked pair uses HD on the current frame only: ``[2, 1]``.""" + if frame_count <= 0: + return [] + if frame_count == 1: + return [1] + return [2] + [1] * (frame_count - 1) + + def _encode_processed_vision(self, processed: Any) -> list[Any] | None: targets = (self.stage_model, self.thinker, getattr(self.stage_model, "model", None)) for target in targets: if target is None: @@ -801,17 +849,48 @@ class MiniCPMO45Stage0DuplexRuntime: flat_pixels.append(slice_pixels.to(device=device, dtype=dtype)) tgt_tensor = image_tgt if hasattr(image_tgt, "reshape") else torch.tensor(image_tgt) flat_tgt.append(tgt_tensor.reshape(-1, 2)) + if not flat_pixels: + return None tgt_sizes = torch.cat(flat_tgt, dim=0).to(device=device, dtype=torch.int64) with torch.no_grad(): hidden = get_hidden({"pixel_values": flat_pixels, "tgt_sizes": tgt_sizes}) except Exception: # noqa: BLE001 - prefill fails with a reason return None expected = MiniCPMO45DuplexPolicy.VISION_EMBEDS_PER_FRAME + embed_dtype = self._token_embedding_dtype() out: list[Any] = [] for block in hidden: block_2d = self._as_2d_tensor(block) if int(block_2d.shape[0]) != expected: return None + if embed_dtype is not None and block_2d.dtype != embed_dtype: + block_2d = block_2d.to(dtype=embed_dtype) out.append(block_2d) return out return None + + def _stage_vision_embeddings(self, frames: list[Any]) -> list[list[Any]] | None: + """Encode camera frames for omni duplex via the loaded vision tower. + + Official ``streaming_prefill`` encodes each ``frame_list`` entry as its + own image (source + optional HD slices) and never stacks audio: the + unit still carries one second of soundtrack. A stacked pair uses + ``max_slice_nums=[2, 1]`` so the current frame keeps the elevator + display readable and the composite stays a single 448-normalized tile. + Frames are processed one at a time because ``process_image([a, b])`` + packs both PILs into one batch item. + """ + process_image = getattr(self.processor, "process_image", None) + if not callable(process_image): + return None + out: list[list[Any]] = [] + for frame, max_slices in zip(frames, self._official_max_slice_nums(len(frames))): + try: + processed = process_image([frame], max_slice_nums=max_slices) + except Exception: # noqa: BLE001 - prefill fails with a reason + return None + encoded = self._encode_processed_vision(processed) + if encoded is None: + return None + out.append(encoded) + return out diff --git a/vllm_omni/experimental/fullduplex/model_executor.py b/vllm_omni/experimental/fullduplex/model_executor.py index eebe01876..58ba597df 100644 --- a/vllm_omni/experimental/fullduplex/model_executor.py +++ b/vllm_omni/experimental/fullduplex/model_executor.py @@ -4,7 +4,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import torch + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.sample.metadata import SamplingMetadata @dataclass(frozen=True, slots=True) @@ -103,4 +108,55 @@ class DuplexSamplingHelper: self.hook_active = False -__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow"] +class DuplexSamplingRunnerMixin: + """Runner-side wiring for the experimental duplex sampling hook. + + GPU and NPU AR runners are siblings, not parent and child, so neither + inherits the other's ``_sample``. Keep the hook here and have each runner + opt in, instead of copying the GPU ``_sample`` path and dropping the + ``prepare_duplex_sampling`` call. + """ + + def _init_duplex_sampling_state(self) -> None: + """Reset hook resolution. Call from ``__init__``, before ``load_model``.""" + self._duplex_sampling_hook = None + self._duplex_sampling_hook_resolved = False + + def _resolve_duplex_sampling_hook(self, *, force: bool = False): + """Bind ``model.prepare_duplex_sampling``. Call with ``force`` after load.""" + if not force and getattr(self, "_duplex_sampling_hook_resolved", False): + return self._duplex_sampling_hook + candidate = getattr(getattr(self, "model", None), "prepare_duplex_sampling", None) + self._duplex_sampling_hook = candidate if callable(candidate) else None + self._duplex_sampling_hook_resolved = True + if self._duplex_sampling_hook is not None and not hasattr(self, "_duplex_sampling_helper"): + self._duplex_sampling_helper = DuplexSamplingHelper() + return self._duplex_sampling_hook + + def _update_duplex_sampling_states(self, scheduler_output: SchedulerOutput) -> None: + """Refresh which batch rows are duplex rows. Call from ``_update_states``.""" + if self._resolve_duplex_sampling_hook() is None: + return + helper = getattr(self, "_duplex_sampling_helper", None) + if helper is not None: + helper.update_states(self, scheduler_output) + + def _apply_duplex_sampling(self, logits: torch.Tensor, prepared_sampling_metadata: SamplingMetadata) -> None: + """Hand the model its duplex rows. Call from ``_sample``, before the sampler.""" + prepare_duplex_sampling = self._resolve_duplex_sampling_hook() + if prepare_duplex_sampling is None: + return + helper = getattr(self, "_duplex_sampling_helper", None) + rows = helper.rows(self) if helper is not None and helper.active_request_ids else () + if rows or (helper is not None and helper.hook_active): + prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) + if helper is not None: + helper.hook_active = bool(rows) + + def _clear_duplex_sampling(self) -> None: + helper = getattr(self, "_duplex_sampling_helper", None) + if helper is not None: + helper.clear() + + +__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow", "DuplexSamplingRunnerMixin"] diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/__init__.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/__init__.py new file mode 100644 index 000000000..ce738cce6 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Native Full-Duplex integration for NVIDIA NemotronLabs VoiceChat.""" diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/data_plane.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/data_plane.py new file mode 100644 index 000000000..7cc2a43ef --- /dev/null +++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/data_plane.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Output projection for Nemotron VoiceChat's independent channels.""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field +from typing import Any +from uuid import uuid4 + +import numpy as np +import torch +from vllm.logger import init_logger + +from vllm_omni.experimental.fullduplex.engine.contracts import ( + duplex_resource_request_belongs_to_session, +) +from vllm_omni.experimental.fullduplex.output import ( + get_duplex_output_decision, +) + +logger = init_logger(__name__) + +EncodeAudio = Callable[[object, int, str, float | None], str | None] + + +@dataclass(frozen=True, slots=True) +class NemotronVoiceChatDataPlaneContext: + epoch: int = 0 + turn_id: int = 0 + auto_responds: bool = True + response_format: str = "wav" + speed: float | None = None + modalities: tuple[str, ...] = ("text", "audio") + + +@dataclass(slots=True) +class _RequestState: + text_frames: int = 0 + audio_frames: int = 0 + pending_speech_end_frames: list[int] = field(default_factory=list) + function_active: bool = False + function_tokens: list[int] = field(default_factory=list) + function_call_id: str | None = None + terminal: bool = False + + +def _coerce_ints(value: object) -> list[int]: + if value is None: + return [] + if isinstance(value, torch.Tensor): + value = value.detach().cpu().reshape(-1).tolist() + elif isinstance(value, np.ndarray): + value = value.reshape(-1).tolist() + if isinstance(value, int): + return [value] + if not isinstance(value, list | tuple): + return [] + values: list[int] = [] + for item in value: + if isinstance(item, torch.Tensor): + if item.numel() != 1: + continue + item = item.item() + elif isinstance(item, np.generic): + item = item.item() + try: + values.append(int(item)) + except (TypeError, ValueError): + continue + return values + + +def _unwrap(output: object) -> tuple[object, object | None, int | None]: + stage_id = getattr(output, "stage_id", None) + inner = getattr(output, "request_output", None) + if inner is not None and inner is not output: + output = inner + outputs = getattr(output, "outputs", None) + completion = outputs[0] if isinstance(outputs, list) and outputs else None + if stage_id is None: + stage_id = getattr(output, "stage_id", None) + return output, completion, int(stage_id) if isinstance(stage_id, int) else None + + +def _multimodal(output: object, completion: object | None) -> dict[str, object]: + decision = get_duplex_output_decision(output) + metadata = getattr(decision, "metadata", None) + if isinstance(metadata, Mapping): + return dict(metadata) + for candidate in ( + getattr(output, "multimodal_output", None), + getattr(completion, "multimodal_output", None) if completion is not None else None, + ): + if isinstance(candidate, Mapping): + return dict(candidate) + return {} + + +def _audio_value(metadata: Mapping[str, object]) -> object | None: + value = next((metadata[key] for key in ("audio", "model_outputs", "latent") if key in metadata), None) + if isinstance(value, list) and len(value) == 1: + return value[0] + return value + + +def _sample_rate(metadata: Mapping[str, object]) -> int: + value = metadata.get("sr", metadata.get("sample_rate_hz", 22050)) + if isinstance(value, list) and value: + value = value[0] + if hasattr(value, "item"): + value = value.item() + return int(value) if isinstance(value, int | float) else 22050 + + +def _audio_samples(audio: object | None) -> int: + if audio is None: + return 0 + if isinstance(audio, torch.Tensor): + return int(audio.numel()) + try: + return int(np.asarray(audio, dtype=np.float32).size) + except (TypeError, ValueError): + return 0 + + +def _speech_end_event(request_id: str) -> dict[str, object]: + return { + "stage_role": "tts", + "is_listen": False, + "data_plane_request_id": request_id, + "text": "", + "end_of_turn": True, + } + + +class NemotronVoiceChatDataPlaneSession: + """Join frame-locked text/function outputs with Stage-2 audio.""" + + def __init__(self, encode_audio: EncodeAudio) -> None: + self._encode_audio = encode_audio + self._requests: dict[str, _RequestState] = {} + self._tokenizer = None + self._tokenizer_ref: str | None = None + self._special_ids: dict[str, int] | None = None + + def configure_runtime(self, runtime_config: Mapping[str, object], *, tokenizer: Any | None = None) -> None: + """Install the serving-resolved tokenizer contract as the sole truth.""" + try: + special_ids = { + "bos": int(runtime_config["nvc_text_bos_id"]), + "eos": int(runtime_config["nvc_text_eos_id"]), + "pad": int(runtime_config["nvc_text_pad_id"]), + "sotc": int(runtime_config["nvc_function_sotc_id"]), + "eotc": int(runtime_config["nvc_function_eotc_id"]), + } + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Nemotron VoiceChat data plane requires resolved special-token ids") from exc + tokenizer_ref = runtime_config.get("nvc_tokenizer_ref") + if not isinstance(tokenizer_ref, str) or not tokenizer_ref: + raise ValueError("Nemotron VoiceChat data plane requires nvc_tokenizer_ref") + self._special_ids = special_ids + self._tokenizer_ref = tokenizer_ref + self._tokenizer = tokenizer + + def begin_request(self, request_id: str) -> None: + self._requests.setdefault(request_id, _RequestState()).terminal = False + + def is_terminal(self, request_id: str | None) -> bool: + return bool(request_id and self._requests.get(request_id) and self._requests[request_id].terminal) + + def mark_terminal(self, request_id: str) -> None: + self._requests.setdefault(request_id, _RequestState()).terminal = True + + def close_stream(self, request_id: str) -> None: + self._requests.pop(request_id, None) + + def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None: + if active_request_id: + self._requests.pop(active_request_id, None) + for request_id in tuple(self._requests): + if duplex_resource_request_belongs_to_session(request_id, session_id): + self._requests.pop(request_id, None) + + def _load_tokenizer(self): + if self._tokenizer is None: + from transformers import AutoTokenizer + + if self._tokenizer_ref is None: + raise RuntimeError("Nemotron VoiceChat data plane was used before runtime configuration") + self._tokenizer = AutoTokenizer.from_pretrained(self._tokenizer_ref, trust_remote_code=False) + return self._tokenizer + + def _ids(self) -> dict[str, int]: + if self._special_ids is None: + raise RuntimeError("Nemotron VoiceChat data plane was used before runtime configuration") + return self._special_ids + + def _decode(self, token_ids: list[int]) -> str: + if not token_ids: + return "" + return str(self._load_tokenizer().decode(token_ids, skip_special_tokens=True)) + + def project( + self, + result: object, + *, + context: NemotronVoiceChatDataPlaneContext | None = None, + ) -> Iterator[dict[str, object]]: + if not isinstance(result, dict): + return + outputs = result.get("data_plane_outputs") + if not isinstance(outputs, list): + return + for output in outputs: + yield from self.project_output(output, context=context) + + def project_output( + self, + output: object, + *, + context: NemotronVoiceChatDataPlaneContext | None = None, + ) -> Iterator[dict[str, object]]: + context = context or NemotronVoiceChatDataPlaneContext() + outer = output + output, completion, stage_id = _unwrap(output) + request_id = getattr(output, "request_id", None) or getattr(outer, "request_id", None) + request_id = str(request_id) if request_id is not None else "" + state = self._requests.setdefault(request_id, _RequestState()) + metadata = _multimodal(outer, completion) + modalities = {modality.lower() for modality in context.modalities} + wants_text = "text" in modalities + wants_audio = "audio" in modalities + if stage_id == 0 or "nvc_text_token_ids" in metadata: + ids = self._ids() + text_ids = _coerce_ints(metadata.get("nvc_text_token_ids")) + if not text_ids and completion is not None: + text_ids = _coerce_ints( + getattr(completion, "token_ids", None) or getattr(completion, "cumulative_token_ids", None) + ) + for token_id in text_ids[-1:]: + state.text_frames += 1 + if token_id == ids["eos"]: + if state.audio_frames >= state.text_frames: + yield _speech_end_event(request_id) + else: + state.pending_speech_end_frames.append(state.text_frames) + elif token_id == ids["pad"]: + yield { + "stage_role": "llm", + "is_listen": True, + "model_listen": True, + "listen_source": "model_listen", + "data_plane_request_id": request_id, + "end_of_turn": False, + } + elif token_id != ids["bos"] and wants_text: + text = self._decode([token_id]) + if text: + yield { + "stage_role": "llm", + "is_listen": False, + "data_plane_request_id": request_id, + "text": text, + "end_of_turn": False, + } + + for function_id in _coerce_ints(metadata.get("nvc_function_token"))[-1:]: + yield from self._project_function_token( + function_id, + state=state, + request_id=request_id, + ids=ids, + ) + return + + audio = _audio_value(metadata) + sample_rate = _sample_rate(metadata) + sample_count = _audio_samples(audio) + encoded = ( + self._encode_audio(audio, sample_rate, context.response_format, context.speed) + if wants_audio and audio is not None + else None + ) + audio_chunk_complete = sample_count > 0 + if audio_chunk_complete: + # A native-duplex Stage-2 segment can contain one or more 80 ms + # codec frames. Count them against the frame-locked Stage-0 text + # channel. The long-lived codec request deliberately keeps + # ``finished=False`` across scheduler wakes; a non-empty decoded + # chunk, not request lifetime, is the audio-coverage signal. + duration_ms = sample_count * 1000 / max(1, sample_rate) + state.audio_frames += max(1, round(duration_ms / 80)) + end_of_turn = bool(state.pending_speech_end_frames) and state.audio_frames >= state.pending_speech_end_frames[0] + if encoded: + yield { + "stage_role": "tts", + "is_listen": False, + "data_plane_request_id": request_id, + "text": "", + "audio_data": encoded, + "audio_format": context.response_format, + "sample_rate_hz": sample_rate, + "audio_duration_ms": round(sample_count * 1000 / max(1, sample_rate)), + "audio_text_mark": True, + "end_of_turn": end_of_turn, + } + elif end_of_turn: + yield _speech_end_event(request_id) + if end_of_turn: + state.pending_speech_end_frames.pop(0) + + def _project_function_token( + self, + token_id: int, + *, + state: _RequestState, + request_id: str, + ids: Mapping[str, int], + ) -> Iterator[dict[str, object]]: + if token_id == ids["sotc"]: + state.function_active = True + state.function_tokens.clear() + state.function_call_id = f"call_{uuid4().hex}" + return + if token_id == ids["eotc"] and state.function_active: + state.function_active = False + raw = self._decode(state.function_tokens) + state.function_tokens.clear() + try: + calls = self._parse_calls(raw) + except ValueError as exc: + state.function_call_id = None + yield { + "stage_role": "function", + "data_plane_request_id": request_id, + "error_code": "nemotron_function_call_parse_error", + "error": str(exc), + } + return + for index, call in enumerate(calls): + yield { + "stage_role": "function", + "data_plane_request_id": request_id, + "function_call": True, + "call_id": ( + state.function_call_id if index == 0 and state.function_call_id else f"call_{uuid4().hex}" + ), + "name": str(call["name"]), + "arguments": ( + call.get("arguments") + if isinstance(call.get("arguments"), str) + else json.dumps(call.get("arguments", {}), separators=(",", ":")) + ), + } + state.function_call_id = None + return + if state.function_active and token_id != ids["pad"]: + state.function_tokens.append(token_id) + + @staticmethod + def _parse_calls(value: str) -> list[dict[str, object]]: + text = value.strip() + if "" in text: + text = text.split("", 1)[1] + if "" in text: + text = text.split("", 1)[0] + try: + parsed = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError("Nemotron VoiceChat produced malformed function-call JSON") from exc + if isinstance(parsed, dict): + parsed = [parsed] + if not isinstance(parsed, list) or not parsed: + raise ValueError("Nemotron VoiceChat function-call payload must be a non-empty object or list") + calls: list[dict[str, object]] = [] + for item in parsed: + if not isinstance(item, dict): + raise ValueError("Nemotron VoiceChat function-call entries must be objects") + name = item.get("name") + if not isinstance(name, str) or not name.strip(): + raise ValueError("Nemotron VoiceChat function-call entries require a non-empty name") + arguments = item.get("arguments", {}) + if isinstance(arguments, str): + try: + decoded_arguments = json.loads(arguments) + except json.JSONDecodeError as exc: + raise ValueError("Nemotron VoiceChat function-call arguments contain malformed JSON") from exc + if not isinstance(decoded_arguments, dict): + raise ValueError("Nemotron VoiceChat function-call arguments must encode a JSON object") + elif not isinstance(arguments, dict): + raise ValueError("Nemotron VoiceChat function-call arguments must be a JSON object") + calls.append({"name": name.strip(), "arguments": arguments}) + return calls + + +__all__ = [ + "NemotronVoiceChatDataPlaneContext", + "NemotronVoiceChatDataPlaneSession", +] diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/input.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/input.py new file mode 100644 index 000000000..e69ebf4f4 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/input.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""80 ms PCM packetization for Nemotron VoiceChat native duplex.""" + +from __future__ import annotations + +import binascii + +import numpy as np +import pybase64 as base64 + +NEMOTRON_VOICECHAT_SAMPLE_RATE = 16000 +NEMOTRON_VOICECHAT_FRAME_SAMPLES = 1280 +_SAMPLE_BYTES = 4 +_FRAME_BYTES = NEMOTRON_VOICECHAT_FRAME_SAMPLES * _SAMPLE_BYTES + + +def _require_terminal_tail(buffered_bytes: int) -> None: + if buffered_bytes > _FRAME_BYTES: + raise ValueError( + "Nemotron VoiceChat manual/deferred commit cannot contain more than one 80 ms frame; " + "enable extra_body.auto_response for native full-duplex streaming" + ) + + +def decode_pcm_f32le(payload: object, *, exact_frame: bool = False) -> bytes: + if not isinstance(payload, dict): + raise ValueError("Nemotron VoiceChat duplex audio payload must be a mapping") + if payload.get("format") != "pcm_f32le": + raise ValueError("Nemotron VoiceChat duplex audio format must be pcm_f32le") + if payload.get("sample_rate_hz") != NEMOTRON_VOICECHAT_SAMPLE_RATE: + raise ValueError("Nemotron VoiceChat duplex audio sample_rate_hz must be 16000") + encoded = payload.get("audio") + if not isinstance(encoded, str): + raise ValueError("Nemotron VoiceChat duplex audio must be base64 pcm_f32le") + try: + raw = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError("Nemotron VoiceChat duplex audio is not valid base64") from exc + if len(raw) % _SAMPLE_BYTES: + raise ValueError("Nemotron VoiceChat pcm_f32le payload has a partial sample") + values = np.frombuffer(raw, dtype=" dict[str, object]: + return { + "type": "audio", + "audio": base64.b64encode(raw).decode("ascii"), + "format": "pcm_f32le", + "sample_rate_hz": NEMOTRON_VOICECHAT_SAMPLE_RATE, + "final": final, + } + + +class NemotronVoiceChatPcmAppendReservation: + def __init__( + self, + owner: NemotronVoiceChatPcmAppendBuffer, + *, + operation_id: str, + payload: dict[str, object] | None, + raw: bytes, + ) -> None: + self._owner = owner + self.operation_id = operation_id + self.payload = payload + self.raw = raw + self._active = True + + @property + def active(self) -> bool: + return self._active + + @property + def byte_count(self) -> int: + return len(self.raw) + + def commit(self) -> None: + if not self._active: + return + self._active = False + if self in self._owner._reservations: + self._owner._reservations.remove(self) + + def rollback(self) -> None: + if not self._active: + return + try: + index = self._owner._reservations.index(self) + except ValueError: + self._active = False + return + restore = bytearray() + for reservation in self._owner._reservations[index:]: + if reservation._active: + restore.extend(reservation.raw) + reservation._active = False + del self._owner._reservations[index:] + self._owner._buffer[:0] = restore + + +class NemotronVoiceChatPcmAppendBuffer: + """Adapt Realtime packets to the current full-duplex PCM buffer contract. + + The model consumes exactly one 1280-sample frame per scheduler append. + Realtime packets may split a frame arbitrarily, but a single packet is + capped at one frame so the serving contract never silently drops a second + ready frame behind a single-reservation API. + """ + + def __init__(self) -> None: + self._buffer = bytearray() + self._reservations: list[NemotronVoiceChatPcmAppendReservation] = [] + + @property + def pending_byte_count(self) -> int: + return len(self._buffer) + + def clear(self) -> None: + for reservation in self._reservations: + reservation._active = False + self._reservations.clear() + self._buffer.clear() + + def clear_force_listen(self) -> None: + # Nemotron VoiceChat has no serving-side force-listen packet bit. + return + + def has_pending(self) -> bool: + return bool(self._buffer) + + def has_reserved(self) -> bool: + return any(reservation.active for reservation in self._reservations) + + def _reserve_frame( + self, + *, + operation_id: str, + final: bool, + ) -> NemotronVoiceChatPcmAppendReservation | None: + if len(self._buffer) < _FRAME_BYTES: + return None + raw = bytes(self._buffer[:_FRAME_BYTES]) + del self._buffer[:_FRAME_BYTES] + reservation = NemotronVoiceChatPcmAppendReservation( + self, + operation_id=operation_id, + payload=_frame_payload(raw, final=final), + raw=raw, + ) + self._reservations.append(reservation) + return reservation + + def prepare_append( + self, + payload: dict[str, object], + *, + operation_id: str, + chunk_period_ms: int, + allow_emit: bool, + ) -> NemotronVoiceChatPcmAppendReservation | None: + if chunk_period_ms != 80: + raise ValueError("Nemotron VoiceChat native duplex chunk_period_ms must be 80") + raw = decode_pcm_f32le(payload) + if len(raw) > _FRAME_BYTES: + raise ValueError("Nemotron VoiceChat Realtime append must contain at most 1280 samples") + self._buffer.extend(raw) + if not allow_emit: + return None + return self._reserve_frame(operation_id=operation_id, final=False) + + def prepare_commit( + self, + *, + operation_id: str, + chunk_period_ms: int, + ) -> NemotronVoiceChatPcmAppendReservation: + if chunk_period_ms != 80: + raise ValueError("Nemotron VoiceChat native duplex chunk_period_ms must be 80") + _require_terminal_tail(len(self._buffer)) + if not self._buffer: + reservation = NemotronVoiceChatPcmAppendReservation( + self, + operation_id=operation_id, + payload=None, + raw=b"", + ) + else: + raw = bytes(self._buffer) + self._buffer.clear() + reservation = NemotronVoiceChatPcmAppendReservation( + self, + operation_id=operation_id, + payload=_frame_payload(raw + bytes(_FRAME_BYTES - len(raw)), final=True), + raw=raw, + ) + self._reservations.append(reservation) + return reservation + + def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None: + if chunk_period_ms != 80: + raise ValueError("Nemotron VoiceChat native duplex chunk_period_ms must be 80") + _require_terminal_tail(len(self._buffer)) + if not self._buffer: + return None + raw = bytes(self._buffer) + self._buffer.clear() + return _frame_payload(raw + bytes(_FRAME_BYTES - len(raw)), final=True) + + +__all__ = [ + "NEMOTRON_VOICECHAT_FRAME_SAMPLES", + "NEMOTRON_VOICECHAT_SAMPLE_RATE", + "NemotronVoiceChatPcmAppendBuffer", + "NemotronVoiceChatPcmAppendReservation", + "decode_pcm_f32le", +] diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/runtime.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/runtime.py new file mode 100644 index 000000000..99de30aa7 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/runtime.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Engine policy for Nemotron VoiceChat's frame-locked duplex timeline.""" + +from __future__ import annotations + +from typing import Any + +from vllm.sampling_params import RequestOutputKind, SamplingParams + +from vllm_omni.experimental.fullduplex.engine.contracts import ( + DuplexAppendPlan, + DuplexInputMode, + DuplexOutputAction, + DuplexOutputDecision, +) +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence +from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import ( + decode_pcm_f32le, +) + + +def _plain_token_ids(value: object, *, name: str) -> list[int]: + if not isinstance(value, list | tuple) or not value: + raise ValueError(f"Nemotron VoiceChat runtime requires non-empty {name}") + try: + return [int(token_id) for token_id in value] + except (TypeError, ValueError) as exc: + raise ValueError(f"Nemotron VoiceChat {name} must contain token ids") from exc + + +def _positive_int(value: object, *, name: str) -> int: + try: + parsed = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Nemotron VoiceChat requires a positive {name}") from exc + if parsed <= 0: + raise ValueError(f"Nemotron VoiceChat requires a positive {name}") + return parsed + + +class NemotronVoiceChatDuplexRuntimeExtension: + """Append one 80 ms acoustic frame to one resumable thinker request.""" + + def configure_sampling_params( + self, + *, + runtime_config: dict[str, Any], + defaults: tuple[object, ...], + ) -> tuple[object, ...]: + del runtime_config + if not defaults: + return defaults + # Every scheduler segment is one transport wake of a long-lived + # request. CUMULATIVE output makes the multimodal output processor + # concatenate all earlier Stage-2 waveforms and resend them on every + # wake (80, 160, 240, ... ms), causing O(n^2) audio replay. Duplex + # consumers require per-wake deltas at every stage. + configured = [params.clone() if isinstance(params, SamplingParams) else params for params in defaults] + for params in configured: + if isinstance(params, SamplingParams): + params.output_kind = RequestOutputKind.DELTA + stage0 = defaults[0] + if isinstance(stage0, SamplingParams): + stage0 = stage0.clone() + stage0.temperature = 0.0 + stage0.top_p = 1.0 + stage0.top_k = 0 + stage0.max_tokens = 1 + stage0.ignore_eos = True + stage0.output_kind = RequestOutputKind.DELTA + configured[0] = stage0 + return tuple(configured) + + def plan_append( + self, + *, + request_id: str, + fence: DuplexFence, + session_config: dict[str, Any], + runtime_config: dict[str, Any], + seq: int, + turn_seq: int, + mode: DuplexInputMode, + payload: object, + final: bool, + sampling_params: object, + ) -> DuplexAppendPlan: + del sampling_params, turn_seq + if mode is not DuplexInputMode.APPEND_AUDIO_CHUNK: + raise ValueError(f"Nemotron VoiceChat does not support duplex input mode {mode.value!r}") + decode_pcm_f32le(payload, exact_frame=True) + normalized_payload = dict(payload) + prompt_ids = _plain_token_ids( + runtime_config.get("nvc_prompt_token_ids"), + name="nvc_prompt_token_ids", + ) + max_model_len = _positive_int( + runtime_config.get("nvc_max_model_len"), + name="nvc_max_model_len", + ) + required_model_len = len(prompt_ids) + seq + 1 + if required_model_len > max_model_len: + raise ValueError( + "Nemotron VoiceChat native duplex session exceeds the Stage-0 " + f"max_model_len: prompt_tokens={len(prompt_ids)} + input_frames={seq} + " + f"sampled_token=1 gives {required_model_len} > {max_model_len}; " + "start a new session or raise Stage 0 max_model_len" + ) + try: + pad_id = int(runtime_config["nvc_text_pad_id"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Nemotron VoiceChat runtime requires nvc_text_pad_id") from exc + scheduler_prompt = prompt_ids + [pad_id] if seq <= 1 else [pad_id] + return DuplexAppendPlan( + prompt={ + "prompt_token_ids": scheduler_prompt, + "model_intermediate_buffer": { + "request_id": request_id, + "global_request_id": [fence.session_id], + "duplex": { + "data_plane": True, + "fence": fence, + "session_id": fence.session_id, + "incarnation": fence.incarnation, + "epoch": fence.epoch, + "source_input_seq": seq, + "seq": seq, + "mode": mode.value, + "payload": normalized_payload, + "final": final, + "session_config": dict(session_config), + "runtime_config": dict(runtime_config), + "scheduler_token_budget": len(scheduler_prompt), + }, + }, + } + ) + + def decide_output( + self, + *, + stage_id: int, + final_stage_id: int, + segment_finished: bool, + segment_token_ids: tuple[int, ...], + segment_output_metadata: dict[str, Any], + output: object, + ) -> DuplexOutputDecision | None: + del final_stage_id, segment_finished, output + if stage_id != 0: + return None + metadata = dict(segment_output_metadata) + metadata["nvc_text_token_ids"] = list(segment_token_ids) + # Stage 0 is a client-visible side channel even though Stage 2 is the + # configured audio response stage. Mark it explicitly so the shared + # collector does not discard it while waiting for the final stage. + metadata["duplex_direct_response"] = True + return DuplexOutputDecision( + action=DuplexOutputAction.DIRECT_RESPONSE, + metadata=metadata, + final_output_type="text", + ) + + +__all__ = ["NemotronVoiceChatDuplexRuntimeExtension"] diff --git a/vllm_omni/experimental/fullduplex/nemotron_voicechat/serving_adapter.py b/vllm_omni/experimental/fullduplex/nemotron_voicechat/serving_adapter.py new file mode 100644 index 000000000..2505b3766 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/nemotron_voicechat/serving_adapter.py @@ -0,0 +1,363 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Realtime serving adapter for Nemotron VoiceChat native duplex.""" + +from __future__ import annotations + +import asyncio +import json +import os +from collections.abc import Callable, Mapping +from copy import deepcopy +from typing import Any + +from vllm_omni.experimental.fullduplex.minicpmo45.session import ( + MiniCPMO45ServingSessionState, +) +from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import ( + NemotronVoiceChatDataPlaneContext, + NemotronVoiceChatDataPlaneSession, +) +from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import ( + NemotronVoiceChatPcmAppendBuffer, +) +from vllm_omni.experimental.fullduplex.openai.protocol import DuplexCapabilities +from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ServingRuntimeConfigError + +EncodeAudio = Callable[[object, int, str, float | None], str | None] + +_DEFAULT_SYSTEM_PROMPT = ( + "You are an AI voice assistant developed by NVIDIA. " + "Your name is NVIDIA Voice Chat. " + "Answer in a spoken, conversational style rather than a written one. " + "Do not repeat the same sentence over and over again. " + "Start the conversation by greeting the user." +) +_PRIVATE_KEYS = frozenset( + { + "nvc_prompt_token_ids", + "nvc_text_bos_id", + "nvc_text_eos_id", + "nvc_text_pad_id", + "nvc_function_sotc_id", + "nvc_function_eotc_id", + "nvc_function_eotr_id", + "nvc_max_model_len", + "nvc_tokenizer_ref", + "nvc_tools_signature", + "nvc_function_response_generation", + "nvc_function_response_token_ids", + "nvc_function_response_call_id", + "nvc_function_response_batches", + } +) + + +def _stt_config(model_config: Any) -> dict[str, Any]: + hf_config = getattr(model_config, "hf_config", None) + stt_cfg = getattr(hf_config, "stt_cfg", None) + if not isinstance(stt_cfg, dict): + raise ServingRuntimeConfigError("Nemotron VoiceChat checkpoint STT configuration is unavailable") + return stt_cfg + + +def _normalized_tools(config: object) -> tuple[list[dict[str, object]], str]: + extra_body = getattr(config, "extra_body", None) + raw_tools = extra_body.get("realtime_tools") if isinstance(extra_body, dict) else None + if raw_tools is None: + return [], "[]" + if not isinstance(raw_tools, list): + raise ServingRuntimeConfigError("Nemotron VoiceChat tools must be a list") + if len(raw_tools) > 5: + raise ServingRuntimeConfigError("Nemotron VoiceChat supports at most 5 tools per session") + + normalized: list[dict[str, object]] = [] + for index, tool in enumerate(raw_tools): + if not isinstance(tool, dict): + raise ServingRuntimeConfigError(f"Nemotron VoiceChat tool {index} must be an object") + function = tool.get("function", tool) + if not isinstance(function, dict): + raise ServingRuntimeConfigError(f"Nemotron VoiceChat tool {index} has no function definition") + definition = {key: value for key, value in function.items() if key != "type"} + if not isinstance(definition.get("name"), str) or not str(definition["name"]).strip(): + raise ServingRuntimeConfigError(f"Nemotron VoiceChat tool {index} requires a name") + normalized.append(definition) + signature = json.dumps(normalized, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + return normalized, signature + + +def _render_tool_prompt(instructions: str, tools: list[dict[str, object]]) -> str: + if not tools: + return instructions + available = json.dumps(tools, ensure_ascii=False, sort_keys=True) + return ( + f"{instructions}\n\n" + "You can use the following tools to assist the user if required:" + f"\n{available}\n\n" + "If you decide to call any tool(s), use the following format:\n" + '[{"name": "tool_name1", "arguments": "tool_args1"}, ' + '{"name": "tool_name2", "arguments": "tool_args2"}]\n\n' + "The user will execute tool-calls and return responses from tool(s) in this format:\n" + '[{"tool_response1"}, {"tool_response2"}]\n\n' + "Based on the tool responses, you can call additional tools if needed, correct tool calls if any " + "errors are found, or just respond to the user." + ) + + +def _render_tool_response(output: str) -> str: + """Serialize one Realtime function result using NVIDIA's function channel.""" + try: + value = json.loads(output) + except json.JSONDecodeError: + value = output + payload = json.dumps([value], ensure_ascii=True, separators=(",", ":")) + return f"{payload}" + + +def _require_native_full_duplex(config: object) -> None: + extra_body = getattr(config, "extra_body", None) + enabled = isinstance(extra_body, dict) and ( + extra_body.get("auto_response") is True or extra_body.get("full_duplex") is True + ) + if not enabled: + raise ServingRuntimeConfigError( + "Nemotron VoiceChat currently supports model-native full-duplex streaming only; " + "set extra_body.auto_response=true", + code="unsupported_nemotron_duplex_mode", + ) + + +def _tokenize_runtime(model_config: Any, instructions: str) -> tuple[dict[str, object], Any]: + from transformers import AutoTokenizer + + stt_cfg = _stt_config(model_config) + tokenizer_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get( + "pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + ) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_ref, trust_remote_code=False) + + def token(name: str, default: str) -> int: + value = tokenizer.convert_tokens_to_ids(stt_cfg.get(name, default)) + if value is None: + raise ServingRuntimeConfigError(f"Nemotron VoiceChat tokenizer does not define {name}") + return int(value) + + bos_id = token("bos_token", "") + eos_id = token("eos_token", "") + pad_id = token("pad_token", "") + prompt_ids = [bos_id] + list(tokenizer.encode(instructions, add_special_tokens=False)) + [eos_id] + runtime = { + "instructions": instructions, + "nvc_prompt_token_ids": prompt_ids, + "nvc_text_bos_id": bos_id, + "nvc_text_eos_id": eos_id, + "nvc_text_pad_id": pad_id, + "nvc_function_sotc_id": int(tokenizer.convert_tokens_to_ids("")), + "nvc_function_eotc_id": int(tokenizer.convert_tokens_to_ids("")), + "nvc_function_eotr_id": int(tokenizer.convert_tokens_to_ids("")), + "nvc_tokenizer_ref": str(tokenizer_ref), + } + return runtime, tokenizer + + +class NemotronVoiceChatServingRuntimeAdapter: + adapter_id = "nemotron_voicechat" + silence_continuation_samples = 1280 + # The session-owned persistent drain is the sole data-plane output + # consumer. Collecting on append would race that drain and can consume a + # Stage2 audio output before the realtime projector observes it. + collect_outputs_on_append = False + clean_response_done_prefix = "" + interrupted_tts_prefix = "" + private_runtime_config_keys = _PRIVATE_KEYS + + def __init__(self, encode_audio: EncodeAudio) -> None: + self.session_states: dict[str, MiniCPMO45ServingSessionState] = {} + self.data_plane = NemotronVoiceChatDataPlaneSession(encode_audio) + self._tokenizer: Any | None = None + + def create_session_state(self) -> MiniCPMO45ServingSessionState: + return MiniCPMO45ServingSessionState( + audio_buffer=NemotronVoiceChatPcmAppendBuffer(), + ) + + def session_state(self, session_id: str) -> MiniCPMO45ServingSessionState: + return self.session_states.setdefault(session_id, self.create_session_state()) + + def remove_session_state(self, session_id: str) -> None: + self.session_states.pop(session_id, None) + + @staticmethod + def is_enabled(config: object) -> bool: + del config + return True + + @staticmethod + def capabilities(*, max_sessions: int) -> DuplexCapabilities: + supports_multi_session = max_sessions > 1 + return DuplexCapabilities( + supports_model_native_turn_policy=True, + supports_external_turn_signal=False, + supports_client_commit=True, + supports_barge_in=False, + supports_playback_ack=True, + supports_input_append=True, + supports_replace_latest_chunk=False, + supports_reencode_context=False, + supports_rollback_to_checkpoint=False, + supports_turn_commit_only=False, + supports_kv_lease=False, + supports_core_kv_lease=False, + supports_model_internal_state=True, + supports_stage_resumption=True, + supports_scheduler_native_append=False, + supports_core_resumable_request=True, + supports_stage_connector_handoff=True, + supports_independent_io_streams=True, + supports_realtime_endpoint=True, + supports_multi_session=supports_multi_session, + supports_multi_session_same_replica=False, + supports_session_lease=True, + supports_session_resume=True, + session_admission_mode="engine_managed", + supports_audio_truncate=False, + requires_model_runner_kv=True, + requires_native_stage_role=True, + implementation_level="model_native_duplex", + adapter_patterns=["scheduler_data_plane"], + input_modes=["append_audio_chunk"], + signal_sources=["model_native", "client_event"], + stage_handoff_transport="scheduler_data_plane", + chunk_period_ms=80, + target_barge_in_latency_ms=None, + ) + + @staticmethod + def validate_client_extra_body(extra_body: object) -> None: + if not isinstance(extra_body, dict): + return + private = sorted(_PRIVATE_KEYS.intersection(extra_body)) + if private: + raise ServingRuntimeConfigError( + "Nemotron VoiceChat runtime configuration is server-owned: " + ", ".join(private) + ) + + async def prepare_runtime_config( + self, + config: object, + *, + model_config: Any, + ) -> dict[str, object]: + self.validate_client_extra_body(getattr(config, "extra_body", None)) + _require_native_full_duplex(config) + instructions = str(getattr(config, "instructions", None) or _DEFAULT_SYSTEM_PROMPT) + tools, tools_signature = _normalized_tools(config) + rendered_prompt = _render_tool_prompt(instructions, tools) + runtime, tokenizer = await asyncio.to_thread(_tokenize_runtime, model_config, rendered_prompt) + max_model_len = getattr(model_config, "max_model_len", None) + if not isinstance(max_model_len, int) or max_model_len <= 0: + raise ServingRuntimeConfigError("Nemotron VoiceChat requires a positive Stage-0 max_model_len") + runtime["nvc_max_model_len"] = max_model_len + # Keep raw session values for immutable-in-incarnation update checks; + # only the rendered prompt token ids enter Stage-0 KV. + runtime["instructions"] = instructions + runtime["nvc_tools_signature"] = tools_signature + self.data_plane.configure_runtime(runtime, tokenizer=tokenizer) + self._tokenizer = tokenizer + return runtime + + def runtime_config_for_function_output( + self, + item: Mapping[str, object], + current: Mapping[str, object], + ) -> dict[str, object]: + """Queue a client tool result for frame-locked function-channel injection.""" + call_id = item.get("call_id") + output = item.get("output") + if not isinstance(call_id, str) or not call_id: + raise ServingRuntimeConfigError( + "function_call_output requires call_id", + code="invalid_function_call_output", + ) + if not isinstance(output, str): + raise ServingRuntimeConfigError( + "function_call_output requires a string output", + code="invalid_function_call_output", + ) + if self._tokenizer is None: + raise ServingRuntimeConfigError( + "Nemotron VoiceChat tokenizer is unavailable for function output", + code="invalid_function_call_output", + ) + response_token_ids = list( + self._tokenizer.encode( + _render_tool_response(output), + add_special_tokens=False, + ) + ) + if not response_token_ids: + raise ServingRuntimeConfigError( + "function_call_output tokenized to an empty response", + code="invalid_function_call_output", + ) + runtime = deepcopy(dict(current)) + generation = int(runtime.get("nvc_function_response_generation", 0)) + 1 + token_ids = [int(token_id) for token_id in response_token_ids] + batches = runtime.get("nvc_function_response_batches") + if not isinstance(batches, list): + batches = [] + batches = [*batches, {"generation": generation, "call_id": call_id, "token_ids": token_ids}] + runtime["nvc_function_response_generation"] = generation + runtime["nvc_function_response_token_ids"] = token_ids + runtime["nvc_function_response_call_id"] = call_id + runtime["nvc_function_response_batches"] = batches + return runtime + + @staticmethod + def runtime_config_for_update( + config: object, + current: Mapping[str, object], + ) -> dict[str, object]: + NemotronVoiceChatServingRuntimeAdapter.validate_client_extra_body(getattr(config, "extra_body", None)) + _require_native_full_duplex(config) + runtime = deepcopy(dict(current)) + instructions = str(getattr(config, "instructions", None) or _DEFAULT_SYSTEM_PROMPT) + _, tools_signature = _normalized_tools(config) + # The system/tool prompt is already resident in Stage-0 KV. Updating it + # without a new incarnation would make the advertised config disagree + # with model state, so reject such updates explicitly. + if runtime and instructions != runtime.get("instructions"): + raise ServingRuntimeConfigError( + "Nemotron VoiceChat instructions cannot change inside an active duplex incarnation" + ) + if runtime and tools_signature != runtime.get("nvc_tools_signature", "[]"): + raise ServingRuntimeConfigError( + "Nemotron VoiceChat tools cannot change inside an active duplex incarnation" + ) + return runtime + + @staticmethod + def data_plane_context( + *, + epoch: int, + turn_id: int, + active_response_turn_id: int | None, + active_response_id: str | None, + auto_responds: bool, + response_format: str, + speed: float | None, + modalities: tuple[str, ...], + ) -> NemotronVoiceChatDataPlaneContext: + del active_response_turn_id, active_response_id + return NemotronVoiceChatDataPlaneContext( + epoch=epoch, + turn_id=turn_id, + auto_responds=auto_responds, + response_format=response_format, + speed=speed, + modalities=modalities, + ) + + +__all__ = ["NemotronVoiceChatServingRuntimeAdapter"] diff --git a/vllm_omni/experimental/fullduplex/openai/chat_fallback.py b/vllm_omni/experimental/fullduplex/openai/chat_fallback.py index 76c45fc73..4edba1646 100644 --- a/vllm_omni/experimental/fullduplex/openai/chat_fallback.py +++ b/vllm_omni/experimental/fullduplex/openai/chat_fallback.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -43,20 +46,22 @@ class ChatFallbackProjectorMixin: else: await self._emit_full_response(session, result, epoch, response_id, send_json) if session.epoch == epoch: + final_stage_metrics = session.accumulate_response_stage_metrics(None) should_commit = self._should_commit_response_to_history(session, response_id) committed_message = session.end_response(commit_text=should_commit) if should_commit: session.register_history_item(f"item_{response_id}", committed_message) - await send_json( - { - "type": "response.done", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "committed": committed_message is not None, - "playback": session.playback.as_dict(), - } - ) + done_payload: dict[str, object] = { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": epoch, + "committed": committed_message is not None, + "playback": session.playback.as_dict(), + } + if final_stage_metrics: + done_payload["vllm_omni"] = {"stage_metrics": final_stage_metrics} + await send_json(done_payload) except asyncio.CancelledError: raise except Exception as exc: @@ -177,6 +182,11 @@ class ChatFallbackProjectorMixin: response_id: str, send_json, ) -> None: + metrics = payload.get("metrics") + stage_metrics = metrics.get("stage_metrics") if isinstance(metrics, dict) else None + if isinstance(stage_metrics, dict): + session.replace_response_stage_metric_snapshots(stage_metrics) + modality = payload.get("modality") if modality not in {None, "text", "audio"}: await send_json( diff --git a/vllm_omni/experimental/fullduplex/openai/protocol.py b/vllm_omni/experimental/fullduplex/openai/protocol.py index 0a87d55b4..8b3812bc3 100644 --- a/vllm_omni/experimental/fullduplex/openai/protocol.py +++ b/vllm_omni/experimental/fullduplex/openai/protocol.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import copy @@ -9,7 +12,6 @@ from uuid import uuid4 class DuplexOverlapPolicy(str, Enum): - AUTO = "auto" LISTEN_ONLY = "listen_only" BARGE_IN_ON_SPEECH = "barge_in_on_speech" @@ -101,7 +103,7 @@ class DuplexCapabilities: supports_multi_session = max_sessions > 1 return cls( supports_model_native_turn_policy=True, - supports_barge_in=False, + supports_barge_in=True, supports_input_append=True, supports_replace_latest_chunk=False, supports_reencode_context=False, @@ -129,6 +131,8 @@ class DuplexCapabilities: signal_sources=["model_native", "client_event", "server_policy"], stage_handoff_transport="scheduler_data_plane", chunk_period_ms=1000, + # Barge-in latency depends on client chunking and lacks hardware E2E + # measurement, so do not advertise an invented target. target_barge_in_latency_ms=None, ) @@ -385,6 +389,8 @@ class ResponseState: active_request_id: str | None = None active_response_id: str | None = None active_response_turn_id: int | None = None + active_response_input_commit_seq: int | None = None + active_response_awaits_input_commit: bool = False last_response_id: str | None = None assistant_text_buffer: list[str] = field(default_factory=list) assistant_audio_text_marks: list[DuplexAssistantAudioTextMark] = field(default_factory=list) @@ -406,12 +412,18 @@ class PlaybackLedger: class ConversationHistory: messages: list[dict[str, object]] = field(default_factory=list) item_ids: dict[str, dict[str, object]] = field(default_factory=dict) + history_item_placeholders: dict[str, dict[str, object]] = field(default_factory=dict) item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) pending_item_ids: dict[str, dict[str, object]] = field(default_factory=dict) pending_item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) + pending_item_input_commit_seqs: dict[str, int] = field(default_factory=dict) pending_truncations_ms: dict[str, int] = field(default_factory=dict) + hard_truncations_ms: dict[str, int] = field(default_factory=dict) last_assistant_full_message: dict[str, object] | None = None last_assistant_audio_text_marks: list[DuplexAssistantAudioTextMark] = field(default_factory=list) + assistant_response_snapshots: dict[str, tuple[dict[str, object], tuple[DuplexAssistantAudioTextMark, ...], int]] = ( + field(default_factory=dict) + ) @dataclass @@ -461,7 +473,8 @@ class DuplexSession: @property def history(self) -> tuple[dict[str, object], ...]: - return tuple(dict(message) for message in self._conversation.messages) + placeholders = {id(message) for message in self._conversation.history_item_placeholders.values()} + return tuple(dict(message) for message in self._conversation.messages if id(message) not in placeholders) @property def active_request_id(self) -> str | None: @@ -500,6 +513,81 @@ class DuplexSession: def last_assistant_audio_text_marks(self) -> tuple[DuplexAssistantAudioTextMark, ...]: return tuple(self._conversation.last_assistant_audio_text_marks) + def has_assistant_response_item(self, response_id: str, item_id: str) -> bool: + if item_id != f"item_{response_id}": + return False + if self.active_response_id == response_id: + return True + message = self._conversation.item_ids.get(item_id) or self._conversation.pending_item_ids.get(item_id) + if isinstance(message, dict) and message.get("role") == "assistant": + return True + return response_id in self._conversation.assistant_response_snapshots + + def playback_ack_is_too_late(self, response_id: str, item_id: str) -> bool: + # An earlier ACK reserved this item's exact history position. Later + # ACKs update that slot in place, so they cannot append an old + # assistant turn after a newer user input. + if item_id in self._conversation.item_ids or item_id in self._conversation.history_item_placeholders: + return False + input_commit_seq = self._conversation.pending_item_input_commit_seqs.get(item_id) + if input_commit_seq is None: + snapshot = self._conversation.assistant_response_snapshots.get(response_id) + if snapshot is not None: + input_commit_seq = snapshot[2] + if input_commit_seq is None and self.active_response_id == response_id: + input_commit_seq = self._response.active_response_input_commit_seq + return input_commit_seq is not None and self.input_commit_seq > input_commit_seq + + def reserve_history_item(self, item_id: str) -> None: + """Reserve the current history position for a response-owned item.""" + if item_id in self._conversation.item_ids or item_id in self._conversation.history_item_placeholders: + return + placeholder: dict[str, object] = {} + self._conversation.history_item_placeholders[item_id] = placeholder + self._conversation.messages.append(placeholder) + + def _store_history_item_message( + self, + item_id: str, + message: dict[str, object], + ) -> dict[str, object]: + """Materialize or update a response item without changing its order.""" + existing = self._conversation.item_ids.get(item_id) + if existing is not None: + if existing is not message: + existing.clear() + existing.update(message) + return existing + + placeholder = self._conversation.history_item_placeholders.pop(item_id, None) + if placeholder is not None: + for index, candidate in enumerate(self._conversation.messages): + if candidate is placeholder: + self._conversation.messages[index] = message + break + else: + self._conversation.messages.append(message) + elif not any(candidate is message for candidate in self._conversation.messages): + self._conversation.messages.append(message) + self._conversation.item_ids[item_id] = message + return message + + def _discard_history_item_placeholder(self, item_id: str) -> bool: + placeholder = self._conversation.history_item_placeholders.pop(item_id, None) + if placeholder is None: + return False + self._conversation.messages = [ + candidate for candidate in self._conversation.messages if candidate is not placeholder + ] + return True + + def release_response_history_snapshot(self, response_id: str | None) -> None: + """Release a final response snapshot after its playback fully commits.""" + if response_id is None or response_id == self.active_response_id: + return + self._conversation.assistant_response_snapshots.pop(response_id, None) + self._conversation.hard_truncations_ms.pop(f"item_{response_id}", None) + @property def playback(self) -> DuplexPlaybackView: return self._playback.current.snapshot() @@ -645,6 +733,7 @@ class DuplexSession: else: content = text + self._bind_active_response_to_input_commit(self._input.commit_seq + 1) self._input.commit_seq += 1 message = {"role": "user", "content": content} self._conversation.messages.append(message) @@ -671,6 +760,7 @@ class DuplexSession: } if transcript: input_audio_part["transcript"] = transcript + self._bind_active_response_to_input_commit(self._input.commit_seq + 1) self._input.commit_seq += 1 message = {"role": "user", "content": [input_audio_part]} if transcript: @@ -692,6 +782,11 @@ class DuplexSession: if completed_turn_id >= self.turn_id: self.turn_id = completed_turn_id + 1 + def _bind_active_response_to_input_commit(self, input_commit_seq: int) -> None: + if self.active_response_id is not None and self._response.active_response_awaits_input_commit: + self._response.active_response_input_commit_seq = int(input_commit_seq) + self._response.active_response_awaits_input_commit = False + def reserve_response_options(self, options: ResponseCreateOptions) -> None: if self._response.active_response_id is not None: raise RuntimeError("response options cannot be reserved while a response is active") @@ -720,6 +815,8 @@ class DuplexSession: response_id = f"resp-{self.session_id}-{self.epoch}-{uuid4().hex[:8]}" self._response.active_response_id = response_id self._response.active_response_turn_id = self.turn_id if turn_id is None else int(turn_id) + self._response.active_response_input_commit_seq = self.input_commit_seq + self._response.active_response_awaits_input_commit = self.turn_state == DuplexTurnState.USER_SPEAKING self._response.last_response_id = response_id self._response.assistant_text_buffer.clear() self._response.assistant_audio_text_marks.clear() @@ -817,6 +914,23 @@ class DuplexSession: return copy.deepcopy(self._response.stage_metrics) + def replace_response_stage_metric_snapshots( + self, + stage_metrics: Mapping[object, object] | None, + ) -> dict[str, dict[str, object]]: + """Merge cumulative chat snapshots by replacing each stage's latest value.""" + if self.active_response_id is None or not isinstance(stage_metrics, Mapping): + return copy.deepcopy(self._response.stage_metrics) + + for raw_stage_id, raw_values in stage_metrics.items(): + if not isinstance(raw_values, Mapping): + continue + stage_id = str(raw_stage_id) + self._response.stage_metrics[stage_id] = copy.deepcopy(dict(raw_values)) + self._response.stage_metric_tpot_weighted_ms.pop(stage_id, None) + self._response.stage_metric_tpot_weight.pop(stage_id, None) + return copy.deepcopy(self._response.stage_metrics) + def accumulate_overlap_speech(self, duration_ms: int) -> int: self._input.overlap_speech_ms += max(0, int(duration_ms)) return self._input.overlap_speech_ms @@ -916,12 +1030,23 @@ class DuplexSession: playback_commit_policy: str | None = None, preserve_request: bool = False, ) -> dict[str, object] | None: + response_id = self._response.active_response_id + response_input_commit_seq = self._response.active_response_input_commit_seq + if response_input_commit_seq is None: + response_input_commit_seq = self.input_commit_seq + response_history_is_late = self.input_commit_seq > response_input_commit_seq assistant_text = "".join(self._response.assistant_text_buffer).strip() message = None if assistant_text: self._conversation.last_assistant_full_message = {"role": "assistant", "content": assistant_text} self._conversation.last_assistant_audio_text_marks = list(self._response.assistant_audio_text_marks) - if commit_text and assistant_text: + if response_id is not None: + self._conversation.assistant_response_snapshots[response_id] = ( + copy.deepcopy(self._conversation.last_assistant_full_message), + tuple(copy.deepcopy(self._response.assistant_audio_text_marks)), + response_input_commit_seq, + ) + if commit_text and assistant_text and not response_history_is_late: committed_text = self._playback_committed_text( assistant_text, playback_commit_policy=playback_commit_policy, @@ -930,12 +1055,28 @@ class DuplexSession: committed_text = "" if commit_text and committed_text: message = {"role": "assistant", "content": committed_text} - self._conversation.messages.append(message) + item_id = f"item_{response_id}" if response_id is not None else None + if item_id is not None and item_id in self._conversation.history_item_placeholders: + self._store_history_item_message(item_id, message) + else: + self._conversation.messages.append(message) + effective_playback_policy = playback_commit_policy or self.config.playback_commit_policy + if ( + response_id is not None + and assistant_text + and message is None + and effective_playback_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value + ): + self.register_history_item(f"item_{response_id}", None) + elif response_id is not None and not assistant_text: + self._discard_history_item_placeholder(f"item_{response_id}") self._response.assistant_text_buffer.clear() if not preserve_request: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None + self._response.active_response_input_commit_seq = None + self._response.active_response_awaits_input_commit = False self._clear_response_metrics() self.turn_state = DuplexTurnState.IDLE self._restore_response_config() @@ -944,16 +1085,20 @@ class DuplexSession: def register_history_item(self, item_id: str | None, message: dict[str, object] | None) -> None: if not item_id: return + response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None + response_snapshot = ( + self._conversation.assistant_response_snapshots.get(response_id) if response_id is not None else None + ) + response_audio_text_marks = list(response_snapshot[1]) if response_snapshot is not None else None if message is None: - last_message = self._conversation.last_assistant_full_message - if last_message is None: + if response_snapshot is None: return - self._conversation.pending_item_ids[item_id] = dict(last_message) - if self._conversation.last_assistant_audio_text_marks: - self._conversation.pending_item_audio_text_marks[item_id] = list( - self._conversation.last_assistant_audio_text_marks - ) - pending_audio_ms = self._conversation.pending_truncations_ms.pop(item_id, None) + last_message, audio_text_marks, _ = response_snapshot + self._conversation.pending_item_ids[item_id] = copy.deepcopy(last_message) + self._conversation.pending_item_input_commit_seqs[item_id] = response_snapshot[2] + if audio_text_marks: + self._conversation.pending_item_audio_text_marks[item_id] = list(copy.deepcopy(audio_text_marks)) + pending_audio_ms = self._conversation.pending_truncations_ms.get(item_id) if pending_audio_ms is not None: self.truncate_history_item( item_id, @@ -961,36 +1106,40 @@ class DuplexSession: playback=self._playback_cursor_for_item_id(item_id), ) return - pending_audio_ms = self._conversation.pending_truncations_ms.pop(item_id, None) - if pending_audio_ms is not None: - self._truncate_message_to_audio_ms( - message, - audio_end_ms=pending_audio_ms, - marks=self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks, - playback=self._playback_cursor_for_item_id(item_id), - ) - if self._message_text_len(message) <= 0: - try: - self._conversation.messages.remove(message) - except ValueError: - pass - return - self._conversation.item_ids[item_id] = message + message = self._store_history_item_message(item_id, message) + pending_audio_ms = self._conversation.pending_truncations_ms.get(item_id) self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) if message.get("role") == "assistant": - marks = self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks + marks = ( + response_audio_text_marks + if response_audio_text_marks is not None + else self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks + ) if marks: self._conversation.item_audio_text_marks[item_id] = list(marks) + if pending_audio_ms is not None: + self.truncate_history_item( + item_id, + audio_end_ms=pending_audio_ms, + playback=self._playback_cursor_for_item_id(item_id), + ) def delete_history_item(self, item_id: str) -> bool: + response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None message = self._conversation.item_ids.pop(item_id, None) self._conversation.item_audio_text_marks.pop(item_id, None) pending = self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) self._conversation.pending_truncations_ms.pop(item_id, None) + self._conversation.hard_truncations_ms.pop(item_id, None) + removed_placeholder = self._discard_history_item_placeholder(item_id) + if response_id is not None: + self._conversation.assistant_response_snapshots.pop(response_id, None) if message is None: - return pending is not None + return pending is not None or removed_placeholder try: self._conversation.messages.remove(message) except ValueError: @@ -1003,8 +1152,43 @@ class DuplexSession: *, audio_end_ms: int, playback: DuplexPlaybackCursor | DuplexPlaybackView | None = None, + hard: bool = False, ) -> bool: playback = playback or self._playback_cursor_for_item_id(item_id) + audio_end_ms = max(0, int(audio_end_ms)) + if hard: + previous_cap_ms = self._conversation.hard_truncations_ms.get(item_id) + self._conversation.hard_truncations_ms[item_id] = ( + audio_end_ms if previous_cap_ms is None else min(previous_cap_ms, audio_end_ms) + ) + hard_cap_ms = self._conversation.hard_truncations_ms.get(item_id) + if hard_cap_ms is not None: + audio_end_ms = min(audio_end_ms, hard_cap_ms) + response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None + response_snapshot = ( + self._conversation.assistant_response_snapshots.get(response_id) if response_id is not None else None + ) + if response_snapshot is not None: + full_message, full_marks, _ = response_snapshot + message = copy.deepcopy(full_message) + changed = self._truncate_message_to_audio_ms( + message, + audio_end_ms=audio_end_ms, + marks=list(full_marks), + playback=playback, + ) + if not changed or self._message_text_len(message) <= 0: + self._conversation.pending_truncations_ms[item_id] = max(0, int(audio_end_ms)) + return False + self._store_history_item_message(item_id, message) + if full_marks: + self._conversation.item_audio_text_marks[item_id] = list(copy.deepcopy(full_marks)) + self._conversation.pending_item_ids.pop(item_id, None) + self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) + self._conversation.pending_truncations_ms.pop(item_id, None) + return True + message = self._conversation.item_ids.get(item_id) if message is None: pending = self._conversation.pending_item_ids.get(item_id) @@ -1022,7 +1206,10 @@ class DuplexSession: if changed: self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) self._conversation.pending_truncations_ms.pop(item_id, None) + if item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) return changed self._conversation.messages.append(message) self._conversation.item_ids[item_id] = message @@ -1032,7 +1219,10 @@ class DuplexSession: ) self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) self._conversation.pending_truncations_ms.pop(item_id, None) + if item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) return True changed = self._truncate_message_to_audio_ms( message, @@ -1047,6 +1237,10 @@ class DuplexSession: self._conversation.messages.remove(message) except ValueError: pass + if item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) + elif changed and item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) return changed def _truncate_message_to_audio_ms( @@ -1193,6 +1387,8 @@ class DuplexSession: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None + self._response.active_response_input_commit_seq = None + self._response.active_response_awaits_input_commit = False self._clear_response_metrics() self._restore_response_config() self.turn_state = DuplexTurnState.BARGE_IN @@ -1206,6 +1402,8 @@ class DuplexSession: self.state = DuplexSessionState.CLOSED self.turn_state = DuplexTurnState.IDLE self._response.active_response_turn_id = None + self._response.active_response_input_commit_seq = None + self._response.active_response_awaits_input_commit = False self._clear_response_metrics() self._restore_response_config() diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_input.py b/vllm_omni/experimental/fullduplex/openai/realtime_input.py index e3fe7f3ec..3d77a4172 100644 --- a/vllm_omni/experimental/fullduplex/openai/realtime_input.py +++ b/vllm_omni/experimental/fullduplex/openai/realtime_input.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio import base64 import binascii +from collections.abc import Callable from typing import Any from uuid import uuid4 @@ -14,12 +18,53 @@ from vllm_omni.experimental.fullduplex.openai.audio import ( from vllm_omni.experimental.fullduplex.openai.realtime_state import ( REALTIME_INPUT_AUDIO_FORMATS, REALTIME_OUTPUT_AUDIO_FORMATS, + RealtimeStateOwner, +) +from vllm_omni.experimental.fullduplex.openai.vad import ( + SILERO_VAD_MIN_THRESHOLD, + ServerVADUnavailableError, + SileroStreamingVAD, + SileroVADConfig, + StreamingVADResult, ) -class RealtimeInputTranslator: +class RealtimeInputTranslator(RealtimeStateOwner): """Translate and validate client Realtime events for the duplex core.""" + _opened: bool + _initial_session_update: bool + _input_speech_started: bool + _input_audio_buffer_has_audio: bool + _input_audio_buffer_had_non_speech: bool + _input_audio_buffer_transcript_parts: list[str] + _active_input_item_id: str | None + _last_conversation_item_id: str | None + _active_response_id: str | None + _last_response_id: str | None + _input_audio_format: str + _output_audio_format: str + _input_sample_rate_hz: int + _conversation_items: dict[str, dict[str, object]] + _item_truncation_cursors: dict[str, tuple[int, int]] + _pending_outbound: asyncio.Queue[dict[str, object]] + _pending_commit_item_ids: asyncio.Queue[str] + + async def _send_realtime_payload(self, payload: dict[str, object]) -> None: + raise NotImplementedError + + def _realtime_error_payload( + self, + code: str, + message: str, + *, + event_id: object | None = None, + param: object | None = None, + ) -> dict[str, object]: + raise NotImplementedError + + _response_is_done: Callable[[object], bool] + async def discard_pending_input_audio( self, *, @@ -67,7 +112,8 @@ class RealtimeInputTranslator: async def _to_duplex_event(self, event: dict[str, object]) -> dict[str, object] | None: event_type = event.get("type") if event_type == "session.update": - session_payload = event.get("session") if isinstance(event.get("session"), dict) else event + session = event.get("session") + session_payload: dict[str, object] = session if isinstance(session, dict) else event format_error = self._validate_realtime_session_audio_formats(session_payload) if format_error is not None: await self._send_realtime_payload( @@ -89,6 +135,7 @@ class RealtimeInputTranslator: ) ) return None + self._prepare_realtime_turn_detection_update(session_payload) self._apply_realtime_session_defaults(session_payload) session_payload.update(self._realtime_overlap_fields(session_payload)) if not self._opened: @@ -99,6 +146,7 @@ class RealtimeInputTranslator: "type": "turn.signal", "event": "session.update", "payload": session_payload, + "realtime_event_id": event.get("event_id"), } if event_type == "conversation.item.create": item = event.get("item") @@ -189,7 +237,7 @@ class RealtimeInputTranslator: ) return None item = self._conversation_items.get(item_id) - if item is None: + if not isinstance(item, dict): await self._send_realtime_payload( self._realtime_error_payload( "item_not_found", @@ -275,15 +323,56 @@ class RealtimeInputTranslator: ) ) return None - looks_like_speech = self._input_looks_like_speech(event, audio=audio, fmt=fmt) + detector = self._server_vad + vad_result: StreamingVADResult | None = None + if isinstance(detector, SileroStreamingVAD): + try: + vad_result = detector.process_base64(audio, fmt=fmt, sample_rate_hz=sample_rate_hz) + except (ServerVADUnavailableError, ValueError) as exc: + unavailable = isinstance(exc, ServerVADUnavailableError) + await self._send_realtime_payload( + self._realtime_error_payload( + "server_vad_unavailable" if unavailable else "bad_audio", + str(exc), + event_id=event.get("event_id"), + param="turn_detection" if unavailable else "audio", + ) + ) + return None + looks_like_speech = ( + vad_result.is_speech + if vad_result is not None + else self._input_looks_like_speech(event, audio=audio, fmt=fmt) + ) self._input_audio_buffer_has_audio = self._input_audio_buffer_has_audio or ( looks_like_speech and isinstance(audio, str) and bool(audio) ) self._input_audio_buffer_had_non_speech = self._input_audio_buffer_had_non_speech or ( not looks_like_speech and isinstance(audio, str) and bool(audio) ) - if looks_like_speech: + stopped_event: dict[str, object] | None = None + if vad_result is not None and vad_result.speech_stopped: + stopped_event = dict(event) + if vad_result.speech_end_ms is not None: + stopped_event["audio_end_ms"] = vad_result.speech_end_ms + if stopped_event is not None and vad_result is not None and vad_result.speech_active: + await self._emit_input_speech_stopped( + stopped_event, + item_id=self._active_input_item_id or f"item_{uuid4().hex}", + ) + if vad_result is not None and vad_result.speech_started: + started_event = dict(event) + if vad_result.speech_start_ms is not None: + started_event["audio_start_ms"] = vad_result.speech_start_ms + await self._emit_input_speech_started(started_event) + elif vad_result is None and looks_like_speech: await self._emit_input_speech_started(event) + if stopped_event is not None and vad_result is not None and not vad_result.speech_active: + await self._emit_input_speech_stopped( + stopped_event, + item_id=self._active_input_item_id or f"item_{uuid4().hex}", + ) + if looks_like_speech: self._remember_input_transcript_hint(event) payload = { "type": "input_audio_buffer.append", @@ -304,16 +393,28 @@ class RealtimeInputTranslator: ) ) return None - payload["video_frames"] = [frame for frame in video_frames if isinstance(frame, str) and frame] + if isinstance(video_frames, list): + payload["video_frames"] = [frame for frame in video_frames if isinstance(frame, str) and frame] self._copy_realtime_input_hints(event, payload) - if not looks_like_speech: - payload["is_speech"] = False + payload["is_speech"] = looks_like_speech + if vad_result is not None: + payload["vad"] = { + "backend": "silero", + "is_speech": vad_result.is_speech, + "speech_active": vad_result.speech_active, + "speech_started": vad_result.speech_started, + "speech_stopped": vad_result.speech_stopped, + "speech_probability": vad_result.speech_probability, + } + if vad_result.speech_active: + payload["force_listen"] = True return payload if event_type == "input_audio_buffer.commit": if not self._input_audio_buffer_has_audio: if self._input_audio_buffer_had_non_speech: self._input_audio_buffer_had_non_speech = False self._active_input_item_id = None + self._reset_server_vad() return { "type": "input_audio_buffer.commit", "final": event.get("final", True), @@ -335,6 +436,7 @@ class RealtimeInputTranslator: self._active_input_item_id = None self._input_audio_buffer_has_audio = False self._input_audio_buffer_had_non_speech = False + self._reset_server_vad() payload = { "type": "input_audio_buffer.commit", "final": event.get("final", True), @@ -350,6 +452,7 @@ class RealtimeInputTranslator: self._input_audio_buffer_has_audio = False self._input_audio_buffer_had_non_speech = False self._input_audio_buffer_transcript_parts.clear() + self._reset_server_vad() return {"type": "input_audio_buffer.clear", "reason": event_type} if event_type == "output_audio_buffer.clear": payload = {"type": "output_audio_buffer.clear", "reason": event_type} @@ -561,7 +664,8 @@ class RealtimeInputTranslator: def _remove_conversation_item(self, item_id: str) -> bool: removed = self._conversation_items.pop(item_id, None) is not None if self._last_conversation_item_id == item_id: - self._last_conversation_item_id = next(reversed(self._conversation_items), None) + remaining_ids = list(self._conversation_items) + self._last_conversation_item_id = remaining_ids[-1] if remaining_ids else None self._item_truncation_cursors.pop(item_id, None) return removed @@ -647,29 +751,34 @@ class RealtimeInputTranslator: audio_config = session_payload.get("audio") audio_input = audio_config.get("input") if isinstance(audio_config, dict) else None audio_output = audio_config.get("output") if isinstance(audio_config, dict) else None - extra_body = ( - dict(session_payload.get("extra_body")) if isinstance(session_payload.get("extra_body"), dict) else {} - ) + extra_body_payload = session_payload.get("extra_body") + extra_body = dict(extra_body_payload) if isinstance(extra_body_payload, dict) else {} extra_body["realtime_session_payload"] = self._json_safe_realtime_payload(session_payload) if isinstance(session_payload.get("tools"), list): extra_body["realtime_tools"] = session_payload["tools"] if isinstance(session_payload.get("tool_choice"), str | dict): extra_body["realtime_tool_choice"] = session_payload["tool_choice"] - if isinstance(session_payload.get("metadata"), dict): - extra_body["realtime_metadata"] = dict(session_payload["metadata"]) - if isinstance(session_payload.get("include"), list): - extra_body["realtime_include"] = list(session_payload["include"]) - if isinstance(session_payload.get("prompt"), dict): - extra_body["realtime_prompt"] = dict(session_payload["prompt"]) + metadata = session_payload.get("metadata") + if isinstance(metadata, dict): + extra_body["realtime_metadata"] = dict(metadata) + include = session_payload.get("include") + if isinstance(include, list): + extra_body["realtime_include"] = list(include) + prompt = session_payload.get("prompt") + if isinstance(prompt, dict): + extra_body["realtime_prompt"] = dict(prompt) input_audio_transcription = self._input_audio_transcription_config(session_payload) if isinstance(input_audio_transcription, dict): extra_body["realtime_input_audio_transcription"] = dict(input_audio_transcription) - if isinstance(session_payload.get("input_audio_noise_reduction"), dict): - extra_body["realtime_input_audio_noise_reduction"] = dict(session_payload["input_audio_noise_reduction"]) - if isinstance(audio_input, dict) and isinstance(audio_input.get("noise_reduction"), dict): - extra_body["realtime_input_audio_noise_reduction"] = dict(audio_input["noise_reduction"]) - if isinstance(session_payload.get("audio"), dict): - extra_body["realtime_audio"] = dict(session_payload["audio"]) + noise_reduction = session_payload.get("input_audio_noise_reduction") + if isinstance(noise_reduction, dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(noise_reduction) + audio_noise_reduction = audio_input.get("noise_reduction") if isinstance(audio_input, dict) else None + if isinstance(audio_noise_reduction, dict): + extra_body["realtime_input_audio_noise_reduction"] = dict(audio_noise_reduction) + audio = session_payload.get("audio") + if isinstance(audio, dict): + extra_body["realtime_audio"] = dict(audio) if isinstance(session_payload.get("tracing"), str | dict): extra_body["realtime_tracing"] = session_payload["tracing"] response_format = self._duplex_response_format(self._output_audio_format) @@ -966,10 +1075,14 @@ class RealtimeInputTranslator: def _validate_realtime_video_frames(video_frames: object, max_slice_nums: object) -> str | None: """Validate omni-duplex camera frames on input_audio_buffer.append. - Wire contract matches the official MiniCPM-o-Demo omni client: one - base64 JPEG per ~1 s audio chunk. HD slicing (max_slice_nums > 1) is - not implemented by the duplex adapter yet and is rejected explicitly - rather than silently ignored. + Wire contract matches the official MiniCPM-o duplex loop: one base + base64 JPEG per ~1 s audio chunk, optionally followed by that unit's + stacked composite tiling the sub-frames captured inside it (at most 2 + images either way). HD slicing (max_slice_nums > 1) is not implemented + by the duplex adapter yet and is rejected explicitly rather than + silently ignored; note official suggests it for composites + (``max_slice_nums=[2, 1]``), so stacked detail is capped at + scale_resolution until that lands. """ if max_slice_nums not in (None, 1): return "max_slice_nums > 1 (HD slicing) is not implemented by the duplex Realtime adapter" @@ -992,25 +1105,124 @@ class RealtimeInputTranslator: return None @staticmethod - def _validate_realtime_turn_detection(session_payload: dict[str, object]) -> str | None: - configured_values: list[tuple[str, object]] = [] - if "turn_detection" in session_payload: - configured_values.append(("turn_detection", session_payload["turn_detection"])) + def _configured_realtime_turn_detection(session_payload: dict[str, object]) -> tuple[str | None, object]: + field = "turn_detection" if "turn_detection" in session_payload else None + selected = session_payload.get("turn_detection") audio_config = session_payload.get("audio") if isinstance(audio_config, dict): audio_input = audio_config.get("input") if isinstance(audio_input, dict) and "turn_detection" in audio_input: - configured_values.append(("audio.input.turn_detection", audio_input["turn_detection"])) - for field_path, turn_detection in configured_values: - if turn_detection is None: - continue - turn_detection_type = turn_detection.get("type") if isinstance(turn_detection, dict) else turn_detection - return ( - f"{field_path}={turn_detection_type!r} is not implemented by the duplex Realtime adapter; " - "set turn_detection to null and commit input explicitly, or use the model-owned duplex policy" - ) + field, selected = "audio.input.turn_detection", audio_input["turn_detection"] + return field, selected + + @classmethod + def _validate_realtime_turn_detection(cls, session_payload: dict[str, object]) -> str | None: + field, turn_detection = cls._configured_realtime_turn_detection(session_payload) + if field is None: + if session_payload.get("overlap_policy") == "barge_in_on_speech": + return ( + "overlap_policy='barge_in_on_speech' requires " + "turn_detection.type='server_vad' on the Realtime endpoint" + ) + return None + if turn_detection is not None and not isinstance(turn_detection, dict): + return f"{field} must be null or an object with type='server_vad'" + if isinstance(turn_detection, dict): + if turn_detection.get("type") != "server_vad": + return f"{field}.type must be 'server_vad'" + threshold = turn_detection.get("threshold", 0.5) + if not cls._valid_vad_number(threshold, minimum=SILERO_VAD_MIN_THRESHOLD, strict=True, maximum=1): + return f"{field}.threshold must be greater than {SILERO_VAD_MIN_THRESHOLD} and at most 1" + for name in ("prefix_padding_ms", "silence_duration_ms", "min_speech_duration_ms"): + value = turn_detection.get(name) + if value is not None and not cls._valid_vad_number(value, minimum=0): + return f"{field}.{name} must be a non-negative number" + if turn_detection.get("interrupt_response", True) is not True: + return ( + f"{field}.interrupt_response=false is unsupported; " + "use turn_detection=null for model-owned listen/speak" + ) + desired_policy = "barge_in_on_speech" if turn_detection is not None else "listen_only" + overlap_policy = session_payload.get("overlap_policy") + if isinstance(overlap_policy, str) and overlap_policy != desired_policy: + return f"overlap_policy={overlap_policy!r} conflicts with turn_detection; expected {desired_policy!r}" return None + @staticmethod + def _valid_vad_number(value: object, *, minimum: float, maximum: float | None = None, strict: bool = False) -> bool: + return ( + isinstance(value, int | float) + and not isinstance(value, bool) + and np.isfinite(float(value)) + and (float(value) > minimum if strict else float(value) >= minimum) + and (maximum is None or float(value) <= maximum) + ) + + def _prepare_realtime_turn_detection_update(self, session_payload: dict[str, object]) -> None: + field, configured = self._configured_realtime_turn_detection(session_payload) + if field is None: + return + if self._pending_turn_detection_update is not None: + raise RuntimeError("A Realtime turn-detection update is already pending") + + turn_detection: dict[str, object] | None = None + detector: SileroStreamingVAD | None = None + if configured is None: + session_payload["overlap_policy"] = "listen_only" + else: + assert isinstance(configured, dict) + turn_detection = { + "type": "server_vad", + "interrupt_response": True, + "threshold": 0.5, + "prefix_padding_ms": 300, + "silence_duration_ms": 500, + "min_speech_duration_ms": 96, + **configured, + } + detector = SileroStreamingVAD( + SileroVADConfig( + threshold=float(turn_detection["threshold"]), + prefix_padding_ms=int(turn_detection["prefix_padding_ms"]), + silence_duration_ms=int(turn_detection["silence_duration_ms"]), + min_speech_duration_ms=max(32, int(turn_detection["min_speech_duration_ms"])), + ) + ) + session_payload["turn_detection"] = dict(turn_detection) + session_payload["overlap_policy"] = "barge_in_on_speech" + + self._pending_turn_detection_update = (turn_detection, detector, asyncio.Event()) + + async def wait_for_realtime_turn_detection_update(self) -> None: + candidate = self._pending_turn_detection_update + if candidate is not None: + await candidate[2].wait() + + def commit_realtime_turn_detection_update(self) -> None: + self._resolve_realtime_turn_detection_update(commit=True) + + def reject_realtime_turn_detection_update(self) -> None: + self._resolve_realtime_turn_detection_update(commit=False) + + def _resolve_realtime_turn_detection_update(self, *, commit: bool) -> None: + candidate = self._pending_turn_detection_update + if candidate is None: + return + turn_detection, detector, resolved = candidate + discarded = self._server_vad if commit else detector + if commit: + self._turn_detection, self._server_vad = turn_detection, detector + self._pending_turn_detection_update = None + resolved.set() + reset = getattr(discarded, "reset", None) + if callable(reset): + reset() + + def _reset_server_vad(self) -> None: + reset = getattr(self._server_vad, "reset", None) + if callable(reset): + reset() + @staticmethod def _input_audio_transcription_config(session_payload: dict[str, object]) -> dict[str, object] | None: transcription = session_payload.get("input_audio_transcription") @@ -1052,6 +1264,46 @@ class RealtimeInputTranslator: item_id = str(item["id"]) item_type = item.get("type") role = item.get("role") + if item_type == "function_call_output": + call_id = item.get("call_id") + output = item.get("output") + matching_call = any( + known.get("type") == "function_call" and known.get("call_id") == call_id + for known in self._conversation_items.values() + ) + duplicate_output = any( + known.get("type") == "function_call_output" and known.get("call_id") == call_id + for known in self._conversation_items.values() + ) + if not isinstance(call_id, str) or not call_id or not matching_call: + await self._send_realtime_payload( + self._realtime_error_payload( + "invalid_function_call_output", + "function_call_output requires the call_id of a completed function call", + event_id=event.get("event_id"), + param="item.call_id", + ) + ) + return None + if not isinstance(output, str) or duplicate_output: + await self._send_realtime_payload( + self._realtime_error_payload( + "invalid_function_call_output", + ( + "function_call_output requires a string output" + if not isinstance(output, str) + else f"function_call_output already exists for call_id {call_id}" + ), + event_id=event.get("event_id"), + param="item.output", + ) + ) + return None + # Do not retain the result until SessionRunner has delivered its + # runtime update. On success, the resulting + # ``conversation.item.created`` event records it in the shared + # input/output protocol state; on failure the client may retry the + # same call_id instead of being rejected by provisional history. if item_type != "message" or role in {"assistant", "system"}: return { "type": "turn.signal", @@ -1211,10 +1463,11 @@ class RealtimeInputTranslator: self._input_speech_started = True if self._active_input_item_id is None: self._active_input_item_id = f"item_{uuid4().hex}" + audio_start_ms = event.get("audio_start_ms", 0) await self._send_realtime_payload( { "type": "input_audio_buffer.speech_started", - "audio_start_ms": int(event.get("audio_start_ms", 0) or 0), + "audio_start_ms": int(audio_start_ms) if isinstance(audio_start_ms, int | float) else 0, "item_id": self._active_input_item_id, } ) diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_output.py b/vllm_omni/experimental/fullduplex/openai/realtime_output.py index 57c07c852..2a68f79ad 100644 --- a/vllm_omni/experimental/fullduplex/openai/realtime_output.py +++ b/vllm_omni/experimental/fullduplex/openai/realtime_output.py @@ -262,15 +262,8 @@ class RealtimeOutputProjector: content_index=0, audio_end_ms=committed_audio_ms, ) - payloads.append( - { - "type": "conversation.item.truncated", - "item_id": item_id, - "content_index": 0, - "audio_end_ms": committed_audio_ms, - "event": event, - } - ) + # Align server history with acknowledged playback; only explicit + # client truncation emits ``conversation.item.truncated``. payloads.extend(self._realtime_audio_done_events(event, response_id)) payloads.extend( self._realtime_response_terminal_events( @@ -346,13 +339,74 @@ class RealtimeOutputProjector: "reason": event.get("finish_reason") or "stop", }, ) + if event_type == "function_call.done": + call_id = event.get("call_id") + name = event.get("name") + arguments = event.get("arguments", "") + if not isinstance(call_id, str) or not call_id or not isinstance(name, str) or not name: + return [{"type": "duplex.function_call.done", "event": event}] + if not isinstance(arguments, str): + arguments = str(arguments) + response_id = f"resp_{uuid4().hex}" + item = { + "id": f"item_{uuid4().hex}", + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "name": name, + "call_id": call_id, + "arguments": arguments, + } + self._conversation_items[str(item["id"])] = item + response = { + "id": response_id, + "object": "realtime.response", + "status": "completed", + "status_details": {"type": "completed", "reason": "completed"}, + "output": [dict(item)], + "metadata": {}, + } + return [ + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + *self._conversation_item_added_events(item), + { + "type": "response.output_item.added", + "response_id": response_id, + "output_index": 0, + "item": dict(item), + }, + { + "type": "response.function_call_arguments.delta", + "response_id": response_id, + "item_id": item["id"], + "output_index": 0, + "call_id": call_id, + "delta": arguments, + }, + { + "type": "response.function_call_arguments.done", + "response_id": response_id, + "item_id": item["id"], + "output_index": 0, + "call_id": call_id, + "arguments": arguments, + }, + { + "type": "response.output_item.done", + "response_id": response_id, + "output_index": 0, + "item": dict(item), + }, + self._conversation_item_done_event(item), + {"type": "response.done", "response": response}, + ] if event_type == "error": raw_error = event.get("error") if isinstance(raw_error, dict): return [event] message = str(raw_error or event.get("message") or "Duplex runtime error") code = str(event.get("code") or "duplex_error") - return [self._realtime_error_payload(code, message)] + return [self._realtime_error_payload(code, message, event_id=event.get("realtime_event_id"))] if event_type == "session.closed": return [{"type": "session.closed", "event": event}] return [{"type": f"duplex.{event_type}", "event": event}] diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_session.py b/vllm_omni/experimental/fullduplex/openai/realtime_session.py index 519c071cf..a85ab1354 100644 --- a/vllm_omni/experimental/fullduplex/openai/realtime_session.py +++ b/vllm_omni/experimental/fullduplex/openai/realtime_session.py @@ -58,6 +58,7 @@ class NativeRealtimeSessionProtocol( while True: if not self._pending_outbound.empty(): return json.dumps(await self._pending_outbound.get()) + await self.wait_for_realtime_turn_detection_update() raw = await websocket.receive_text() try: event = json.loads(raw) diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_state.py b/vllm_omni/experimental/fullduplex/openai/realtime_state.py index ac3b2a8de..6a2c8f347 100644 --- a/vllm_omni/experimental/fullduplex/openai/realtime_state.py +++ b/vllm_omni/experimental/fullduplex/openai/realtime_state.py @@ -48,6 +48,9 @@ REALTIME_ERROR_TYPES_BY_CODE = { "input_audio_buffer_empty": "invalid_request_error", "missing_item_id": "invalid_request_error", "item_not_found": "invalid_request_error", + "playback_item_mismatch": "invalid_request_error", + "playback_item_not_found": "invalid_request_error", + "playback_ack_too_late": "invalid_request_error", "unsupported_audio_format": "invalid_request_error", "unsupported_ref_audio_path": "invalid_request_error", "ref_audio_required": "invalid_request_error", @@ -57,6 +60,7 @@ REALTIME_ERROR_TYPES_BY_CODE = { "native_text_append_unsupported": "invalid_request_error", "runtime_native_stage_role_required": "server_error", "runtime_native_runner_kv_required": "server_error", + "server_vad_unavailable": "server_error", } @@ -105,6 +109,9 @@ class RealtimeSessionState: _input_sample_rate_hz: int = 16000 _output_audio_format: str = "pcm16" _overlap_silence_rms: float = 0.003 + _turn_detection: dict[str, object] | None = None + _server_vad: object | None = None + _pending_turn_detection_update: tuple[dict[str, object] | None, object | None, asyncio.Event] | None = None _send_realtime_json: Any = None _initial_session_update: bool = False _input_speech_started: bool = False @@ -180,6 +187,11 @@ class RealtimeStateOwner: _input_sample_rate_hz: int = _RealtimeStateField() _output_audio_format: str = _RealtimeStateField() _overlap_silence_rms: float = _RealtimeStateField() + _turn_detection: dict[str, object] | None = _RealtimeStateField() + _server_vad: object | None = _RealtimeStateField() + _pending_turn_detection_update: tuple[dict[str, object] | None, object | None, asyncio.Event] | None = ( + _RealtimeStateField() + ) _send_realtime_json: Any = _RealtimeStateField() _initial_session_update: bool = _RealtimeStateField() _input_speech_started: bool = _RealtimeStateField() diff --git a/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py b/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py index f88cffd3b..d05b709ff 100644 --- a/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py +++ b/vllm_omni/experimental/fullduplex/openai/runtime_adapter.py @@ -14,6 +14,18 @@ class ServingRuntimeConfigError(ValueError): self.code = code +def reject_changed_runtime_value( + new_value: object, + current_value: object, + *, + message: str, + code: str, + error_cls: type[ServingRuntimeConfigError] = ServingRuntimeConfigError, +) -> None: + if new_value != current_value: + raise error_cls(message, code=code) + + class PcmAppendReservation(Protocol): operation_id: str payload: dict[str, object] | None @@ -64,6 +76,7 @@ class ServingRuntimeSessionState(Protocol): audio_buffer: PcmAppendBuffer input_since_commit: bool speech_since_commit: bool + native_context_locked: bool committed_audio_payload: dict[str, object] | None committed_audio_operation_id: str | None committed_audio_reserved_bytes: int @@ -111,6 +124,7 @@ class ServingRuntimeAdapter(Protocol): clean_response_done_prefix: str interrupted_tts_prefix: str private_runtime_config_keys: frozenset[str] + collect_outputs_on_append: bool def create_session_state(self) -> ServingRuntimeSessionState: ... diff --git a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py index 258f512b6..98e0aaaf5 100644 --- a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py +++ b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -135,7 +138,9 @@ class NativeRuntimeBridgeMixin: incarnation=session.incarnation, ) if self._callable_accepts_keyword(append_input, "collect_outputs"): - append_kwargs["collect_outputs"] = False + append_kwargs["collect_outputs"] = bool( + getattr(self._serving_runtime_adapter, "collect_outputs_on_append", False) + ) result = await append_input(session.session_id, **append_kwargs) except Exception as exc: logger.exception("Failed to append duplex runtime input: %s", exc) @@ -315,27 +320,86 @@ class NativeRuntimeBridgeMixin: native.data_plane_task = task return True - # One model unit (1 s at 16 kHz) of pcm_f32le silence, matching the - # official full-duplex behavior where the microphone keeps streaming - # silence while the assistant speaks; replies span multiple units. + # One MiniCPM model unit (1 s at 16 kHz) is the compatibility default. + # Other native-duplex adapters may use a different frame-locked unit. _NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO = base64.b64encode(bytes(16000 * 4)).decode("ascii") _NATIVE_RESPONSE_MAX_CONTINUATION_UNITS = 8 + # Each unit is one second of silence. Auto-response needs substantially + # more room than the eight decision beats above, but must still terminate + # well before the 300-second benchmark playback horizon. The final unit + # asks the model to LISTEN; a protocol fallback closes the response if that + # logit hint is ignored. + _NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS = 64 def _native_silence_unit_payload(self) -> dict[str, object]: + samples = int(getattr(self._serving_runtime_adapter, "silence_continuation_samples", 16000)) + audio = ( + self._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO + if samples == 16000 + else base64.b64encode(bytes(samples * 4)).decode("ascii") + ) return { "type": "audio", - "audio": self._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO, + "audio": audio, "format": "pcm_f32le", "sample_rate_hz": 16000, } def _native_response_continuations_remaining(self, session: DuplexSession, response_id: str) -> bool: - if self._session_auto_responds(session): - return True native = self._runtime_session_state(session) owner_id = f"response:{response_id}" count = native.continuation_units if native.continuation_owner_id == owner_id else 0 - return count < self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS + limit = ( + self._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS + if self._session_auto_responds(session) + else self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS + ) + return count < limit + + async def _finish_bounded_native_auto_response( + self, + send_json, + *, + session: DuplexSession, + expected_epoch: int | None, + model_turn_id: int | None = None, + ) -> None: + """Close an auto-response whose final silence unit did not end it.""" + native = self._runtime_session_state(session) + response_id = session.active_response_id + response_epoch = session.epoch + response_turn_id = model_turn_id if model_turn_id is not None else session.active_response_turn_id + if expected_epoch is not None and response_epoch != expected_epoch: + return + native.clear_continuation() + await send_json( + { + "type": "response.listen", + "session_id": session.session_id, + "epoch": response_epoch, + "reason": "continuation_limit", + "model_listen": False, + } + ) + if response_id is None: + if session.epoch == response_epoch and response_turn_id is not None: + session.complete_model_turn(response_turn_id) + return + if session.epoch != response_epoch or session.active_response_id != response_id: + return + if response_turn_id is not None: + session.complete_model_turn(response_turn_id) + session.end_response(commit_text=False, preserve_request=True) + await send_json( + { + "type": "response.done", + "session_id": session.session_id, + "response_id": response_id, + "epoch": response_epoch, + "committed": False, + "playback": session.playback.as_dict(), + } + ) def _native_silence_continuation_is_stale( self, @@ -400,9 +464,23 @@ class NativeRuntimeBridgeMixin: owner_id = f"model-turn:{expected_model_turn_id}" payload_turn_id = expected_model_turn_id count = native.continuation_units if native.continuation_owner_id == owner_id else 0 - if not auto_response and count >= self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS: + continuation_limit = ( + self._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS + if auto_response + else self._NATIVE_RESPONSE_MAX_CONTINUATION_UNITS + ) + if count >= continuation_limit: + if auto_response: + await self._finish_bounded_native_auto_response( + send_json, + session=session, + expected_epoch=expected_epoch, + model_turn_id=payload_turn_id, + ) return payload = self._native_silence_unit_payload() + if auto_response and count + 1 == continuation_limit: + payload["force_listen"] = True payload["duplex_turn_id"] = payload_turn_id scheduler = native.silence_continuation_scheduler @@ -826,6 +904,18 @@ class NativeRuntimeBridgeMixin: } ) return close_reason, True + if native_result.get("function_call") is True: + await send_json( + { + "type": "function_call.done", + "session_id": session.session_id, + "epoch": session.epoch, + "call_id": native_result.get("call_id"), + "name": native_result.get("name"), + "arguments": native_result.get("arguments", ""), + } + ) + return close_reason, True if native_result.get("requires_stage_handoff") is True or native_result.get("requires_tts_stage") is True: # Stage0 announces a talker handoff before Stage1 has produced # client-visible audio. A single client input stream may contain @@ -861,11 +951,17 @@ class NativeRuntimeBridgeMixin: and not session.active_response_accepts_model_turn(model_turn_id) ): return close_reason, emitted_response + auto_response = self._session_auto_responds(session) + active_response_id = session.active_response_id + auto_continuations_remaining = active_response_id is None or self._native_response_continuations_remaining( + session, active_response_id + ) non_terminal_auto_listen = ( - self._session_auto_responds(session) - and session.active_response_id is not None + auto_response + and active_response_id is not None and session.active_request_id is not None and native_result.get("end_of_turn") is not True + and auto_continuations_remaining ) if non_terminal_auto_listen: await self._maybe_continue_native_response( @@ -874,7 +970,6 @@ class NativeRuntimeBridgeMixin: expected_epoch=expected_epoch, ) return close_reason, emitted_response - auto_response = self._session_auto_responds(session) if not auto_response and data_plane_request_id == session.active_request_id: session.clear_request() model_listen = native_result.get("model_listen") @@ -913,6 +1008,14 @@ class NativeRuntimeBridgeMixin: expected_epoch=expected_epoch, ) return close_reason, emitted_response + if auto_response: + self._runtime_session_state(session).clear_continuation() + if not auto_continuations_remaining: + completed_turn_id = model_turn_id + if completed_turn_id is None: + completed_turn_id = session.active_response_turn_id + if completed_turn_id is not None: + session.complete_model_turn(completed_turn_id) session.end_response(commit_text=False, preserve_request=auto_response) await send_json( { diff --git a/vllm_omni/experimental/fullduplex/openai/serving.py b/vllm_omni/experimental/fullduplex/openai/serving.py index 472df0256..0b3bd246c 100644 --- a/vllm_omni/experimental/fullduplex/openai/serving.py +++ b/vllm_omni/experimental/fullduplex/openai/serving.py @@ -494,6 +494,22 @@ class OmniDuplexSessionHandler( # playback only controls history ACKs, not model admission. if is_speech: session.accumulate_overlap_speech(duration_ms) + vad_speech_started = self._vad_speech_started(event, payload) + if ( + self._uses_native_input_append(session) + and session.capabilities.supports_barge_in + and is_speech + and session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + and vad_speech_started is not False + ): + return { + "action": "barge_in", + "reason": ("server_vad_speech_started" if vad_speech_started is True else "barge_in_on_speech"), + "cancel_reason": "turn_detected" if vad_speech_started is True else "barge_in", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + } return { "action": "listen", "reason": "auto_response_continuous", @@ -501,7 +517,7 @@ class OmniDuplexSessionHandler( "overlap_speech_ms": session.overlap_speech_ms, "buffer_audio": True, "defer_runtime_append": False, - "force_listen": False, + "force_listen": event.get("force_listen") is True or payload.get("force_listen") is True, "preserve_realtime_input": True, } if bool(event.get("force_listen", False)): @@ -553,9 +569,22 @@ class OmniDuplexSessionHandler( session.accumulate_overlap_speech(duration_ms) if policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value: + vad_speech_started = self._vad_speech_started(event, payload) + if vad_speech_started is False: + return { + "action": "listen", + "reason": "server_vad_utterance_active", + "duration_ms": duration_ms, + "overlap_speech_ms": session.overlap_speech_ms, + "buffer_audio": True, + "defer_runtime_append": False, + "force_listen": True, + "preserve_realtime_input": True, + } return { "action": "barge_in", - "reason": "policy_barge_in_on_speech", + "reason": ("server_vad_speech_started" if vad_speech_started is True else "policy_barge_in_on_speech"), + "cancel_reason": "turn_detected" if vad_speech_started is True else "barge_in", "duration_ms": duration_ms, "overlap_speech_ms": session.overlap_speech_ms, "buffer_audio": True, @@ -599,6 +628,17 @@ class OmniDuplexSessionHandler( "defer_runtime_append": True, } + @staticmethod + def _vad_speech_started( + event: Mapping[str, object], + payload: Mapping[str, object], + ) -> bool | None: + for source in (event, payload): + vad = source.get("vad") + if isinstance(vad, Mapping) and isinstance(vad.get("speech_started"), bool): + return bool(vad["speech_started"]) + return None + @staticmethod def _event_requests_barge_in(event: Mapping[str, object]) -> bool: if event.get("force_barge_in") is True: @@ -1155,6 +1195,24 @@ class OmniDuplexSessionHandler( candidate_config: DuplexSessionConfig, ) -> dict[str, object] | None: if not self._uses_native_input_append(session): + if candidate_config.overlap_policy != DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value: + return None + return { + "type": "error", + "code": "server_vad_requires_native_duplex", + "error": "Realtime server_vad requires a model-native duplex runtime", + } + if ( + candidate_config.instructions != session.config.instructions + and self._runtime_session_state(session).native_context_locked + ): + return { + "type": "error", + "session_id": session.session_id, + "code": "instructions_update_unsupported", + "error": "session.update cannot change instructions after the native duplex context is initialized", + } + if session.config.extra_body.get("minicpmo45_native_duplex") is not True: return None if not self._config_requests_audio_output(candidate_config): return None @@ -1326,6 +1384,19 @@ class OmniDuplexSessionHandler( "code": "ref_audio_update_unsupported", "error": "session.update cannot change ref_audio after the native duplex runtime is open", } + extra_body_payload = payload.get("extra_body") + if isinstance(extra_body_payload, dict) and "minicpmo45_native_duplex" in extra_body_payload: + current_native_duplex = session.config.extra_body.get("minicpmo45_native_duplex") + if ( + current_native_duplex is not None + and extra_body_payload["minicpmo45_native_duplex"] != current_native_duplex + ): + return { + "type": "error", + "session_id": session.session_id, + "code": "native_duplex_mode_update_unsupported", + "error": "session.update cannot change minicpmo45_native_duplex after the session is created", + } if isinstance(payload.get("instructions"), str): session.config.instructions = str(payload["instructions"]) elif "instructions" in payload and payload.get("instructions") is None: @@ -1601,7 +1672,60 @@ class OmniDuplexSessionHandler( if response_id is None and item_id is None and session.active_response_id is not None: response_id = session.active_response_id item_id = f"item_{response_id}" - if event.get("truncate") is True: + if response_id is not None: + expected_item_id = f"item_{response_id}" + if item_id is None: + item_id = expected_item_id + elif item_id != expected_item_id: + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_item_mismatch", + "error": "playback.ack item_id must match item_.", + } + ) + return + if not session.has_assistant_response_item(response_id, item_id): + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_item_not_found", + "error": f"No assistant response item is registered for {response_id}.", + } + ) + return + if session.playback_ack_is_too_late(response_id, item_id): + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_ack_too_late", + "error": "playback.ack arrived after a later user input was committed.", + } + ) + return + # Reserve the response's current history position before any later + # input commit can append a user item. A 0 ms ACK is intentional: + # the response may be active but have no audio delta yet. + session.reserve_history_item(item_id) + elif item_id is not None: + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_item_not_found", + "error": "playback.ack requires a response-owned assistant item.", + } + ) + return + hard_truncate = event.get("truncate") is True + if hard_truncate: playback = session.acknowledge_playback( int(played_ms), committed_cursor, @@ -1619,10 +1743,21 @@ class OmniDuplexSessionHandler( ) committed_history = False if isinstance(item_id, str) and item_id: + expected_item_id = f"item_{response_id}" if response_id is not None else None + if ( + expected_item_id == item_id + and item_id not in session.history_item_ids + and item_id not in session.pending_history_item_ids + ): + # ACK_ONLY responses are normally registered as pending when + # the response ends. Recover the response-local item if that + # registration was lost across the response/ack boundary. + session.register_history_item(item_id, None) committed_history = session.truncate_history_item( item_id, audio_end_ms=committed_cursor, playback=playback, + hard=hard_truncate, ) elif session.pending_history_item_ids: # A plain playback ack has no OpenAI item id. Commit the only @@ -1635,6 +1770,7 @@ class OmniDuplexSessionHandler( item_id, audio_end_ms=committed_cursor, playback=playback, + hard=hard_truncate, ) elif session.active_response_id is not None: item_id = f"item_{session.active_response_id}" @@ -1642,6 +1778,7 @@ class OmniDuplexSessionHandler( item_id, audio_end_ms=committed_cursor, playback=playback, + hard=hard_truncate, ) elif session.last_assistant_full_message is not None: if item_id is None and session.history_item_ids: @@ -1657,6 +1794,7 @@ class OmniDuplexSessionHandler( item_id, audio_end_ms=committed_cursor, playback=playback, + hard=hard_truncate, ) await send_json( { @@ -1673,6 +1811,7 @@ class OmniDuplexSessionHandler( ) if committed_history and committed_cursor >= max(playback.sent_ms, playback.generated_ms): session.release_response_playback(response_id) + session.release_response_history_snapshot(response_id) async def _cancel_active_response( self, @@ -1699,10 +1838,14 @@ class OmniDuplexSessionHandler( item_id = f"item_{old_response_id}" if committed_message is not None: session.register_history_item(item_id, committed_message) - elif committed_ms > 0: + elif committed_ms > 0 and not session.playback_ack_is_too_late(old_response_id, item_id): session.truncate_history_item(item_id, audio_end_ms=committed_ms) new_epoch, old_playback = self._advance_barge_in_epoch(session) if old_request_id is not None: + # Epoch-scoped request ids prevent state reuse, but explicitly + # release projector/parser cursors so cancelled epochs do not + # accumulate until the whole session closes. + self._serving_runtime_adapter.data_plane.close_stream(old_request_id) await self._abort_request_background( session, old_request_id, diff --git a/vllm_omni/experimental/fullduplex/openai/session_runner.py b/vllm_omni/experimental/fullduplex/openai/session_runner.py index 4b5e4f2e6..e3a8a8cdf 100644 --- a/vllm_omni/experimental/fullduplex/openai/session_runner.py +++ b/vllm_omni/experimental/fullduplex/openai/session_runner.py @@ -18,6 +18,7 @@ from vllm_omni.experimental.fullduplex.openai.commit_policy import ( decide_commit_action, ) from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexOverlapPolicy, DuplexPlaybackCommitPolicy, DuplexSession, DuplexSessionState, @@ -28,6 +29,7 @@ from vllm_omni.experimental.fullduplex.openai.realtime_session import ( ) from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ( PcmAppendReservation, + ServingRuntimeConfigError, ServingRuntimeSessionState, payload_turn_id, ) @@ -247,6 +249,9 @@ class DuplexSessionRunnerMixin: await actor.enqueue_event(event) except WebSocketDisconnect: await actor.enqueue_event({"type": "__disconnect__"}) + except Exception as exc: + await emit_event({"type": "error", "error": str(exc), "code": "realtime_input_failed"}) + await actor.enqueue_event({"type": "__disconnect__"}) async def next_actor_event() -> dict[str, object]: nonlocal pending_turn_reservations, transport_detached @@ -344,6 +349,13 @@ class DuplexSessionRunnerMixin: ) precreated_response_id = session.active_response_id if precreate_response else None + def _discard_retained_committed_audio() -> None: + if ( + retained_committed_payload is not None + and native.committed_audio_payload is retained_committed_payload + ): + session.release_input_bytes(native.clear_committed_audio()) + async def _run() -> bool: nonlocal runtime_closed try: @@ -357,6 +369,7 @@ class DuplexSessionRunnerMixin: expected_epoch=append_epoch, ) if append_ok: + native.native_context_locked = True if pcm_reservation is not None: pcm_reservation.commit() session.release_input_bytes(pcm_reservation.byte_count) @@ -365,8 +378,10 @@ class DuplexSessionRunnerMixin: and native.committed_audio_payload is retained_committed_payload ): session.release_input_bytes(native.clear_committed_audio()) - elif pcm_reservation is not None: - pcm_reservation.rollback() + else: + if pcm_reservation is not None: + pcm_reservation.rollback() + _discard_retained_committed_audio() if ( not append_ok and precreated_response_id is not None @@ -400,10 +415,13 @@ class DuplexSessionRunnerMixin: ) return append_ok except asyncio.CancelledError: + if pcm_reservation is not None: + pcm_reservation.rollback() raise except Exception as exc: if pcm_reservation is not None: pcm_reservation.rollback() + _discard_retained_committed_audio() logger.exception("Native duplex append task failed: %s", exc) await self._send_runtime_error(emit_event, "runtime_append_task_failed", exc, session=session) if session.state != DuplexSessionState.CLOSED: @@ -438,19 +456,34 @@ class DuplexSessionRunnerMixin: if not predecessor_ok: if pcm_reservation is not None: pcm_reservation.rollback() + _discard_retained_committed_audio() return False if actor.closing or runtime_closed or session.state != DuplexSessionState.OPEN: if pcm_reservation is not None: pcm_reservation.rollback() + _discard_retained_committed_audio() return False if before_append is not None and not before_append(): if pcm_reservation is not None: pcm_reservation.rollback() + _discard_retained_committed_audio() return True if pcm_reservation is not None and not pcm_reservation.active: + _discard_retained_committed_audio() return False return await _run() + def _release_cancelled_retained_audio(done: asyncio.Task[bool]) -> None: + if done.cancelled(): + _discard_retained_committed_audio() + return + try: + append_ok = done.result() + except Exception: + append_ok = False + if not append_ok: + _discard_retained_committed_audio() + predecessor = actor.native_append_tail if predecessor is not None and predecessor.done(): try: @@ -463,6 +496,7 @@ class DuplexSessionRunnerMixin: # is an explicit retry and starts a new chain. predecessor = None task = asyncio.create_task(_run_in_wire_order(predecessor)) + task.add_done_callback(_release_cancelled_retained_audio) actor.native_append_tail = task actor.track_append_task( task, @@ -642,6 +676,9 @@ class DuplexSessionRunnerMixin: if handshake is None: return session = handshake.session + if initial_update_error := self._runtime_session_candidate_update_error(session, session.config): + await emit_event(initial_update_error) + return if handshake.resumed: native = self._serving_runtime_adapter.session_states[session.session_id] actor.tasks = self._session_tasks[session.session_id] @@ -693,6 +730,8 @@ class DuplexSessionRunnerMixin: ) if isinstance(open_result, dict): created_payload["runtime_control"] = self._redact_runtime_control_result(open_result) + if realtime_protocol is not None: + realtime_protocol.commit_realtime_turn_detection_update() await emit_event(created_payload) resume_credential_delivered = session.capabilities.supports_session_resume reader_task = asyncio.create_task(read_event_loop(), name="duplex-session-reader") @@ -875,7 +914,11 @@ class DuplexSessionRunnerMixin: if event_type in {"input.cancel", "response.cancel", "barge_in", "output_audio_buffer.clear"}: cancel_reason = ( - "output_audio_buffer_clear" if event_type == "output_audio_buffer.clear" else "barge_in" + "output_audio_buffer_clear" + if event_type == "output_audio_buffer.clear" + else "client_cancelled" + if event_type == "response.cancel" + else "barge_in" ) cancelled_fence = DuplexFence( session.session_id, @@ -927,6 +970,8 @@ class DuplexSessionRunnerMixin: had_native_append = await actor.cancel_append_tasks( response_bound_only=event_type in {"response.cancel", "output_audio_buffer.clear"}, ) + if event_type == "response.cancel": + session.release_input_bytes(native.clear_committed_audio()) had_native_stream = native.data_plane_task is not None cancelled = await self._cancel_active_response( session, @@ -1038,9 +1083,20 @@ class DuplexSessionRunnerMixin: await emit_event(self._barge_in_unsupported_error(session)) continue if turn_event == "session.update": + realtime_event_id = event.get("realtime_event_id") + + async def emit_update_event(update_event: dict[str, object]) -> None: + if isinstance(realtime_event_id, str) and realtime_event_id: + update_event = {**update_event, "realtime_event_id": realtime_event_id} + await emit_event(update_event) + + def reject_update() -> None: + if realtime_protocol is not None: + realtime_protocol.reject_realtime_turn_detection_update() + payload = event.get("payload") if not isinstance(payload, dict): - await emit_event( + await emit_update_event( { "type": "error", "session_id": session.session_id, @@ -1048,12 +1104,22 @@ class DuplexSessionRunnerMixin: "error": "session.update requires a session payload", } ) + reject_update() continue if not await wait_for_native_append_tail(): + await emit_update_event( + { + "type": "error", + "code": "session_update_aborted", + "error": "session.update was not applied because the preceding append failed", + } + ) + reject_update() continue runtime_update_error = self._runtime_session_update_error(session, payload) if runtime_update_error is not None: - await emit_event(runtime_update_error) + await emit_update_event(runtime_update_error) + reject_update() continue previous_config = session.config candidate_config = deepcopy(previous_config) @@ -1063,30 +1129,47 @@ class DuplexSessionRunnerMixin: finally: session.replace_config(previous_config) if update_error is not None: - await emit_event(update_error) + await emit_update_event(update_error) + reject_update() continue runtime_update_error = self._runtime_session_candidate_update_error( session, candidate_config, ) if runtime_update_error is not None: - await emit_event(runtime_update_error) + await emit_update_event(runtime_update_error) + reject_update() + continue + try: + candidate_runtime_config = self._runtime_config_for_session_update( + session, + candidate_config, + ) + except ServingRuntimeConfigError as exc: + await emit_update_event( + { + "type": "error", + "session_id": session.session_id, + "code": exc.code, + "error": str(exc), + } + ) + reject_update() continue - candidate_runtime_config = self._runtime_config_for_session_update( - session, - candidate_config, - ) if not await self._signal_runtime_session( session, turn_event, - emit_event, + emit_update_event, session_config=candidate_config.as_dict(), runtime_config=candidate_runtime_config, ): + reject_update() continue session.replace_config(candidate_config) session.replace_runtime_config(candidate_runtime_config) - await emit_event( + if realtime_protocol is not None: + realtime_protocol.commit_realtime_turn_detection_update() + await emit_update_event( { "type": "session.updated", "session": session.as_public_dict(), @@ -1096,6 +1179,53 @@ class DuplexSessionRunnerMixin: if turn_event == "conversation.item.create": payload = event.get("payload") item = payload.get("item") if isinstance(payload, dict) else None + item_type = item.get("type") if isinstance(item, dict) else None + prepare_function_output = getattr( + self._serving_runtime_adapter, + "runtime_config_for_function_output", + None, + ) + if item_type == "function_call_output" and callable(prepare_function_output): + if not await wait_for_native_append_tail(): + continue + try: + candidate_runtime_config = prepare_function_output( + item, + dict(session.runtime_config), + ) + except ServingRuntimeConfigError as exc: + await emit_event( + { + "type": "error", + "session_id": session.session_id, + "code": exc.code, + "error": str(exc), + } + ) + continue + if not await self._signal_runtime_session( + session, + "session.update", + emit_event, + runtime_config=candidate_runtime_config, + ): + continue + session.replace_runtime_config(candidate_runtime_config) + await emit_event( + { + "type": "conversation.item.created", + "session_id": session.session_id, + "item": item, + "created": True, + } + ) + await self._maybe_continue_native_response( + emit_event, + session=session, + expected_epoch=session.epoch, + expected_model_turn_id=session.turn_id, + ) + continue message = self._realtime_item_to_history_message(item) item_id = item.get("id") if isinstance(item, dict) else None if message is not None: @@ -1131,6 +1261,7 @@ class DuplexSessionRunnerMixin: session.truncate_history_item( item_id, audio_end_ms=int(audio_end_ms) if isinstance(audio_end_ms, int | float) else 0, + hard=True, ) if isinstance(item_id, str) else False @@ -1249,7 +1380,16 @@ class DuplexSessionRunnerMixin: buffer_overlap_audio = True if self._uses_native_input_append(session): mark_pending_silence_superseded() - overlap_active = not self._session_auto_responds(session) and native_response_in_progress() + overlap_active = native_response_in_progress() and ( + not self._session_auto_responds(session) + or ( + session.capabilities.supports_barge_in + and ( + session.config.overlap_policy == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value + or self._event_requests_barge_in(event) + ) + ) + ) if overlap_active: decision = self._overlap_decision(session, event, payload) await self._emit_overlap_decision(emit_event, session, decision) @@ -1290,11 +1430,12 @@ class DuplexSessionRunnerMixin: native.speech_since_commit = False await actor.cancel_append_tasks() had_native_stream = native.data_plane_task is not None + cancel_reason = str(decision.get("cancel_reason") or "barge_in") cancelled = await self._cancel_active_response( session, actor.active_response_task, emit_event, - reason="barge_in", + reason=cancel_reason, ) had_native_stream = ( await self._cancel_native_data_plane_stream(session) or had_native_stream @@ -1310,7 +1451,7 @@ class DuplexSessionRunnerMixin: "type": "audio.cancelled", "session_id": session.session_id, "response_id": old_response_id, - "reason": "barge_in", + "reason": cancel_reason, "cancelled_epoch": old_epoch, "epoch": new_epoch, "committed_ms": committed_ms, @@ -1330,7 +1471,7 @@ class DuplexSessionRunnerMixin: "type": "audio.cancelled", "session_id": session.session_id, "response_id": session.last_response_id, - "reason": "barge_in", + "reason": cancel_reason, "cancelled_epoch": old_epoch, "epoch": new_epoch, "committed_ms": committed_ms, @@ -1862,6 +2003,8 @@ class DuplexSessionRunnerMixin: with suppress(Exception): await emit_event({"type": "error", "error": str(exc), "code": "internal_error"}) finally: + if realtime_protocol is not None: + realtime_protocol.reject_realtime_turn_detection_update() if session is not None: if reader_task is not None and not reader_task.done(): reader_task.cancel() diff --git a/vllm_omni/experimental/fullduplex/openai/vad.py b/vllm_omni/experimental/fullduplex/openai/vad.py new file mode 100644 index 000000000..f81559613 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/openai/vad.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import binascii +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np +import pybase64 as base64 + +SILERO_VAD_MIN_THRESHOLD = 0.15 + + +class ServerVADUnavailableError(RuntimeError): + pass + + +@dataclass(frozen=True, slots=True) +class SileroVADConfig: + threshold: float = 0.5 + prefix_padding_ms: int = 300 + silence_duration_ms: int = 500 + min_speech_duration_ms: int = 96 + + +@dataclass(frozen=True, slots=True) +class StreamingVADResult: + is_speech: bool + speech_active: bool + speech_started: bool = False + speech_stopped: bool = False + speech_probability: float = 0.0 + speech_start_ms: int | None = None + speech_end_ms: int | None = None + + +class SileroStreamingVAD: + _SAMPLE_RATE_HZ = 16_000 + _WINDOW_SAMPLES = 512 + + def __init__( + self, + config: SileroVADConfig, + *, + frame_scorer: Callable[[np.ndarray], float] | None = None, + ) -> None: + self.config = config + self._frame_scorer = frame_scorer + self._model: object | None = None + self._clear_stream_state() + + def _clear_stream_state(self) -> None: + self._pending = np.empty(0, dtype=np.float32) + self._speech_active = False + self._candidate_samples = 0 + self._candidate_start_sample: int | None = None + self._silence_samples = 0 + self._processed_samples = 0 + self._resample_rate_hz: int | None = None + self._resample_remainder = 0 + + def reset(self) -> None: + self._clear_stream_state() + reset_states = getattr(self._model, "reset_states", None) + if callable(reset_states): + reset_states() + + def process_base64( + self, + audio: object, + *, + fmt: object, + sample_rate_hz: object, + ) -> StreamingVADResult: + if fmt != "pcm_f32le" or not isinstance(audio, str): + raise ValueError("Silero server VAD requires decoded pcm_f32le audio") + try: + raw = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError("Silero server VAD received invalid base64 audio") from exc + if len(raw) % 4: + raise ValueError("Silero server VAD received an incomplete pcm_f32le frame") + samples = np.frombuffer(raw, dtype=" StreamingVADResult: + samples = np.asarray(samples, dtype=np.float32).reshape(-1) + if samples.size: + self._pending = np.concatenate((self._pending, samples)) + + started = False + stopped = False + start_ms: int | None = None + end_ms: int | None = None + max_probability = 0.0 + contained_speech = self._speech_active + min_speech_samples = self.config.min_speech_duration_ms * 16 + min_silence_samples = self.config.silence_duration_ms * 16 + negative_threshold = max(0.0, self.config.threshold - SILERO_VAD_MIN_THRESHOLD) + + while self._pending.size >= self._WINDOW_SAMPLES: + frame = self._pending[: self._WINDOW_SAMPLES] + self._pending = self._pending[self._WINDOW_SAMPLES :] + frame_start = self._processed_samples + self._processed_samples += self._WINDOW_SAMPLES + probability = min(1.0, max(0.0, float(self._score_frame(frame)))) + max_probability = max(max_probability, probability) + + if self._speech_active: + contained_speech = True + if probability < negative_threshold: + self._silence_samples += self._WINDOW_SAMPLES + if self._silence_samples >= min_silence_samples: + speech_end_sample = self._processed_samples - self._silence_samples + self._speech_active = False + self._silence_samples = 0 + stopped = True + end_ms = max(0, int(round(speech_end_sample / 16))) + else: + self._silence_samples = 0 + continue + + if probability >= self.config.threshold: + if self._candidate_samples == 0: + self._candidate_start_sample = frame_start + self._candidate_samples += self._WINDOW_SAMPLES + if self._candidate_samples >= min_speech_samples: + candidate_start = self._candidate_start_sample or 0 + prefix_samples = self.config.prefix_padding_ms * 16 + self._speech_active = True + self._candidate_samples = self._silence_samples = 0 + self._candidate_start_sample = None + contained_speech = True + started = True + start_ms = max(0, int(round((candidate_start - prefix_samples) / 16))) + else: + self._candidate_samples = 0 + self._candidate_start_sample = None + + return StreamingVADResult( + contained_speech, self._speech_active, started, stopped, max_probability, start_ms, end_ms + ) + + def _score_frame(self, frame: np.ndarray) -> float: + if self._frame_scorer is not None: + return float(self._frame_scorer(frame)) + if self._model is None: + try: + import torch + from silero_vad import load_silero_vad + except ImportError as exc: + raise ServerVADUnavailableError( + "server_vad requires the optional 'silero-vad' package; install vllm-omni[server-vad]" + ) from exc + self._model = load_silero_vad(onnx=False) + self._frame_scorer = lambda value: float( + self._model(torch.from_numpy(np.ascontiguousarray(value)), self._SAMPLE_RATE_HZ).item() + ) + return float(self._frame_scorer(frame)) diff --git a/vllm_omni/experimental/fullduplex/openai/websocket.py b/vllm_omni/experimental/fullduplex/openai/websocket.py index 18cd11c8d..3b3dc4d5d 100644 --- a/vllm_omni/experimental/fullduplex/openai/websocket.py +++ b/vllm_omni/experimental/fullduplex/openai/websocket.py @@ -121,12 +121,15 @@ class DuplexSessionTasks: ] if not tasks: return False + cancelled_tail = self.native_append_tail if self.native_append_tail in tasks else None for task in tasks: task.cancel() try: await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=timeout_s) except TimeoutError: pass + if cancelled_tail is not None and self.native_append_tail is cancelled_tail: + self.native_append_tail = None return True diff --git a/vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py b/vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py index 66244ca68..d4f920d44 100644 --- a/vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py +++ b/vllm_omni/experimental/fullduplex/personaplex/serving_adapter.py @@ -15,6 +15,7 @@ from vllm_omni.experimental.fullduplex.openai.protocol import ( ) from vllm_omni.experimental.fullduplex.openai.runtime_adapter import ( ServingRuntimeConfigError, + reject_changed_runtime_value, ) from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA from vllm_omni.experimental.fullduplex.personaplex.data_plane import ( @@ -31,6 +32,7 @@ _PRIVATE_RUNTIME_CONFIG_KEYS = frozenset( { "personaplex_prefill_slots", "personaplex_model_path", + "minicpmo45_native_duplex", } ) @@ -40,6 +42,7 @@ class PersonaPlexServingSessionState: audio_buffer: PersonaPlexPcmAppendBuffer = field(default_factory=PersonaPlexPcmAppendBuffer) input_since_commit: bool = False speech_since_commit: bool = False + native_context_locked: bool = False committed_audio_payload: dict[str, object] | None = None committed_audio_operation_id: str | None = None committed_audio_reserved_bytes: int = 0 @@ -190,9 +193,23 @@ class PersonaPlexServingRuntimeAdapter: current: Mapping[str, object], ) -> dict[str, object]: cls.validate_client_extra_body(getattr(config, "extra_body", None)) + new_persona = str(getattr(config, "instructions", None) or DEFAULT_PERSONA) + reject_changed_runtime_value( + new_persona, + current.get("personaplex_persona"), + message="PersonaPlex persona (instructions) cannot be changed after the session is created", + code="persona_update_unsupported", + ) + new_voice = cls._voice_name(getattr(config, "voice", None)) + reject_changed_runtime_value( + new_voice, + current.get("personaplex_voice_prompt"), + message="PersonaPlex voice cannot be changed after the session is created", + code="voice_update_unsupported", + ) runtime_config = deepcopy(dict(current)) - runtime_config["personaplex_voice_prompt"] = cls._voice_name(getattr(config, "voice", None)) - runtime_config["personaplex_persona"] = str(getattr(config, "instructions", None) or DEFAULT_PERSONA) + runtime_config["personaplex_voice_prompt"] = new_voice + runtime_config["personaplex_persona"] = new_persona return runtime_config @staticmethod diff --git a/vllm_omni/experimental/fullduplex/video_stacking.py b/vllm_omni/experimental/fullduplex/video_stacking.py new file mode 100644 index 000000000..2a7706759 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/video_stacking.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Intra-unit frame stacking for the omni duplex video track. + +Official MiniCPM-o duplex raises the visual refresh rate without touching the +audio cadence: ``get_video_frame_audio_segments(stack_frames=N)`` samples the +``N-1`` sub-frames captured inside each second (skipping the one that would +duplicate the second's base frame) and tiles them into a *single* composite +image. The duplex loop then feeds ``frame_list=[base, composite]`` alongside +that second's audio, so a unit still carries one second of audio and at most +two images however large ``N`` is. + +Tiling therefore belongs to media preprocessing, not to the wire: a caller +turns its sub-frames into one composite here and sends it as the second frame +of the append that closes the unit. +""" + +from __future__ import annotations + +import io +from collections.abc import Sequence +from typing import Any + +import pybase64 as base64 + +_BG_COLOR = (255, 255, 255) +_LINE_COLOR = (0, 0, 0) +_LINE_WIDTH = 6 + + +def _as_image(frame: Any) -> Any: + """Coerce a PIL image, encoded bytes, or base64/data-URL string to PIL RGB.""" + from PIL import Image + + if isinstance(frame, Image.Image): + return frame + if isinstance(frame, (bytes, bytearray)): + return Image.open(io.BytesIO(bytes(frame))).convert("RGB") + if isinstance(frame, str): + payload = frame.split(",")[-1] if ";base64," in frame else frame + return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB") + raise TypeError(f"unsupported frame type for stacking: {type(frame)}") + + +def _grid_shape(images: Sequence[Any], *, cell_w: int, cell_h: int) -> tuple[int, int]: + """Rows and columns for ``images``, matching official ``concat_images``. + + 4 frames tile 2x2; 2 and 3 frames pick the row/column arrangement whose + canvas is closest to square, so a landscape clip stacks vertically and a + portrait one horizontally; anything else falls back to a single row. + """ + count = len(images) + if count == 4: + return 2, 2 + if count == 1: + return 1, 1 + if count not in (2, 3): + return 1, count + + candidates = [(1, count), (count, 1)] + + def squareness(rows: int, cols: int) -> float: + width = cols * cell_w + (cols - 1) * _LINE_WIDTH + height = rows * cell_h + (rows - 1) * _LINE_WIDTH + return abs(width / max(1, height) - 1.0) + + ratios = [squareness(rows, cols) for rows, cols in candidates] + if count == 2 and ratios[0] == ratios[1]: + average_aspect = sum(im.width / max(1, im.height) for im in images) / count + return (1, 2) if average_aspect >= 1.0 else (2, 1) + return candidates[min(range(len(candidates)), key=lambda index: ratios[index])] + + +def _letterbox(image: Any, target_w: int, target_h: int) -> Any: + from PIL import Image + + image = image.convert("RGB") + width, height = image.size + scale = min(target_w / width, target_h / height) + new_size = (max(1, round(width * scale)), max(1, round(height * scale))) + resized = image.resize(new_size, Image.Resampling.BICUBIC) + canvas = Image.new("RGB", (target_w, target_h), _BG_COLOR) + canvas.paste(resized, ((target_w - new_size[0]) // 2, (target_h - new_size[1]) // 2)) + return canvas + + +def concat_frames(frames: Sequence[Any]) -> Any: + """Tile a unit's sub-frames into one PIL image, as official stacking does. + + Cells are the largest input width and height, each frame is letterboxed to + keep its aspect ratio, and a separator band is drawn only in the interior + seams so the model can tell the sub-frames apart. + """ + from PIL import Image + + images = [_as_image(frame) for frame in frames] + if not images: + raise ValueError("cannot stack an empty frame list") + + cell_w = max(image.width for image in images) + cell_h = max(image.height for image in images) + rows, cols = _grid_shape(images, cell_w=cell_w, cell_h=cell_h) + + canvas = Image.new( + "RGB", + (cols * cell_w + (cols - 1) * _LINE_WIDTH, rows * cell_h + (rows - 1) * _LINE_WIDTH), + _LINE_COLOR, + ) + for index, image in enumerate(images[: rows * cols]): + row, col = divmod(index, cols) + canvas.paste( + _letterbox(image, cell_w, cell_h), + (col * (cell_w + _LINE_WIDTH), row * (cell_h + _LINE_WIDTH)), + ) + return canvas + + +def concat_frames_b64(frames: Sequence[Any], *, quality: int = 95) -> str: + """``concat_frames`` re-encoded as a base64 JPEG, ready for the wire.""" + buf = io.BytesIO() + concat_frames(frames).save(buf, format="JPEG", quality=quality) + return base64.b64encode(buf.getvalue()).decode("ascii") + + +def unit_subframe_offsets(stack_frames: int, *, unit_s: float = 1.0) -> list[float]: + """Sub-frame capture offsets inside a unit, in seconds from its start. + + ``stack_frames=5`` yields 0.2/0.4/0.6/0.8 s for a 1 s unit. Offset 0 is + skipped because official skips it: it is the unit's base frame, already + sent as ``frame_list[0]``. + """ + count = max(1, int(stack_frames)) + return [index / count * unit_s for index in range(1, count)] diff --git a/vllm_omni/host_weight_runtime/__init__.py b/vllm_omni/host_weight_runtime/__init__.py new file mode 100644 index 000000000..5a504adfc --- /dev/null +++ b/vllm_omni/host_weight_runtime/__init__.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Neutral contracts for immutable runtime-ready host weight artifacts.""" + +from .config import ( + CapacityPolicy, + HostWeightRuntimeConfig, + IntegrityPolicy, + ProductionPolicy, + RemoteImportPolicy, + RemoteOnMiss, + RuntimeMode, + StorageClass, + StorageDomainPolicy, + StorageScope, + ValidationLevel, + WaitPolicy, +) +from .errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionAction, ResolutionStage +from .identity import ( + AdaptationIdentity, + CanonicalJson, + ComponentIdentity, + IdentityValidationError, + ProducerIdentity, + RuntimeWeightLayout, + WeightArtifactIdentity, + WeightRepresentation, + WeightSourceIdentity, +) +from .inspection import ( + ArtifactInspection, + ArtifactInventoryEntry, + ArtifactInventoryState, + DomainInspection, + HostWeightInspector, +) +from .lease import HostWeightLease, HostWeightLeaseCarrier, LeaseProvenance, MappedHostRegion +from .manifest import ( + FileManifestEntry, + ManifestValidationError, + PublicationMarker, + TensorKind, + TensorManifestEntry, + WeightManifest, +) +from .outcomes import ( + AttemptResult, + HostWeightResolution, + PostLoadPublicationOutcome, + PostLoadPublicationReport, + ResolutionAttempt, + ResolutionOutcome, + ResolutionReport, + StoreResult, + StoreStatus, +) +from .protocols import ( + ArtifactWriter, + BuildRequest, + CoordinationScope, + HostWeightStore, + LookupPhase, + ProductionMetadata, + ProductionSourceMode, + RemoteArtifactDescriptor, + RemoteArtifactProvider, + TensorFileWriter, + TensorWriteSpec, + WeightProducer, + WeightProductionSpec, + WeightRestorePlan, + WeightRestorer, +) +from .runtime import HostWeightRuntime + +__all__ = [ + "AdaptationIdentity", + "ArtifactInspection", + "ArtifactInventoryEntry", + "ArtifactInventoryState", + "ArtifactWriter", + "AttemptResult", + "BuildRequest", + "CanonicalJson", + "CapacityPolicy", + "ComponentIdentity", + "CoordinationScope", + "DomainInspection", + "FailureCode", + "FileManifestEntry", + "HostWeightError", + "HostWeightFailure", + "HostWeightLease", + "HostWeightLeaseCarrier", + "HostWeightResolution", + "HostWeightInspector", + "HostWeightRuntime", + "HostWeightRuntimeConfig", + "HostWeightStore", + "IdentityValidationError", + "IntegrityPolicy", + "LeaseProvenance", + "LookupPhase", + "ManifestValidationError", + "MappedHostRegion", + "ProducerIdentity", + "PostLoadPublicationOutcome", + "PostLoadPublicationReport", + "ProductionMetadata", + "ProductionPolicy", + "ProductionSourceMode", + "PublicationMarker", + "RemoteArtifactDescriptor", + "RemoteArtifactProvider", + "RemoteImportPolicy", + "RemoteOnMiss", + "ResolutionAction", + "ResolutionAttempt", + "ResolutionOutcome", + "ResolutionReport", + "ResolutionStage", + "RuntimeMode", + "RuntimeWeightLayout", + "StorageClass", + "StorageDomainPolicy", + "StorageScope", + "StoreResult", + "StoreStatus", + "TensorFileWriter", + "TensorKind", + "TensorManifestEntry", + "TensorWriteSpec", + "ValidationLevel", + "WaitPolicy", + "WeightArtifactIdentity", + "WeightManifest", + "WeightProducer", + "WeightProductionSpec", + "WeightRepresentation", + "WeightRestorePlan", + "WeightRestorer", + "WeightSourceIdentity", +] diff --git a/vllm_omni/host_weight_runtime/config.py b/vllm_omni/host_weight_runtime/config.py new file mode 100644 index 000000000..0bacd69d2 --- /dev/null +++ b/vllm_omni/host_weight_runtime/config.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed policy for the loader-adjacent Host Weight Runtime.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path + + +class RuntimeMode(str, Enum): + DISABLED = "disabled" + PREFERRED = "preferred" + REQUIRED = "required" + + +class StorageScope(str, Enum): + NODE = "node" + NUMA = "numa" + + +class StorageClass(str, Enum): + DISK = "disk" + TMPFS = "tmpfs" + + +class ValidationLevel(str, Enum): + MANIFEST_AND_METADATA = "manifest_and_metadata" + FULL_CHECKSUM = "full_checksum" + FS_VERITY = "fs_verity" + + +class RemoteOnMiss(str, Enum): + DISABLED = "disabled" + TRY = "try" + REQUIRE = "require" + + +@dataclass(frozen=True) +class StorageDomainPolicy: + root: Path + scope: StorageScope = StorageScope.NODE + domain_id: str = "node" + storage_class: StorageClass = StorageClass.DISK + numa_prefault: bool = False + verify_numa_residency: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.root, Path) or not str(self.root): + raise ValueError("host weight store root must not be empty") + if not isinstance(self.scope, StorageScope) or not isinstance(self.storage_class, StorageClass): + raise ValueError("storage scope and class must use their typed enum values") + if not isinstance(self.domain_id, str) or not self.domain_id: + raise ValueError("host weight store domain_id must not be empty") + if not isinstance(self.numa_prefault, bool) or not isinstance(self.verify_numa_residency, bool): + raise ValueError("NUMA placement options must be booleans") + if self.scope is StorageScope.NODE and (self.numa_prefault or self.verify_numa_residency): + raise ValueError("NUMA placement options require scope='numa'") + + +@dataclass(frozen=True) +class CapacityPolicy: + max_artifact_bytes: int | None = None + max_store_bytes: int | None = None + min_free_bytes: int = 0 + eviction: str = "none" + + def __post_init__(self) -> None: + for name, value in ( + ("max_artifact_bytes", self.max_artifact_bytes), + ("max_store_bytes", self.max_store_bytes), + ): + if value is not None and (isinstance(value, bool) or not isinstance(value, int) or value <= 0): + raise ValueError(f"{name} must be positive when configured") + if isinstance(self.min_free_bytes, bool) or not isinstance(self.min_free_bytes, int) or self.min_free_bytes < 0: + raise ValueError("min_free_bytes must not be negative") + if not isinstance(self.eviction, str) or self.eviction != "none": + raise ValueError("automatic host weight eviction is not implemented") + + +@dataclass(frozen=True) +class IntegrityPolicy: + local_lookup: ValidationLevel = ValidationLevel.MANIFEST_AND_METADATA + require_trusted_permissions: bool = True + background_scrub: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.local_lookup, ValidationLevel): + raise ValueError("local lookup validation must use ValidationLevel") + if not isinstance(self.require_trusted_permissions, bool) or not isinstance(self.background_scrub, bool): + raise ValueError("integrity switches must be booleans") + + +@dataclass(frozen=True) +class RemoteImportPolicy: + on_local_miss: RemoteOnMiss = RemoteOnMiss.DISABLED + providers: tuple[str, ...] = () + + def __post_init__(self) -> None: + if not isinstance(self.on_local_miss, RemoteOnMiss): + raise ValueError("remote on-miss policy must use RemoteOnMiss") + if not isinstance(self.providers, tuple): + raise ValueError("remote providers must be an immutable tuple") + if self.on_local_miss is not RemoteOnMiss.DISABLED and not self.providers: + raise ValueError("remote providers are required when remote import is enabled") + if any(not isinstance(provider, str) or not provider for provider in self.providers): + raise ValueError("remote provider identifiers must not be empty") + + +@dataclass(frozen=True) +class ProductionPolicy: + allow_local_build: bool = True + allow_post_load_publish: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.allow_local_build, bool) or not isinstance(self.allow_post_load_publish, bool): + raise ValueError("production switches must be booleans") + + +@dataclass(frozen=True) +class WaitPolicy: + coordination_timeout_seconds: float = 600.0 + + def __post_init__(self) -> None: + timeout = self.coordination_timeout_seconds + if ( + isinstance(timeout, bool) + or not isinstance(timeout, int | float) + or not math.isfinite(timeout) + or timeout <= 0 + ): + raise ValueError("coordination_timeout_seconds must be positive") + + +@dataclass(frozen=True) +class HostWeightRuntimeConfig: + mode: RuntimeMode = RuntimeMode.DISABLED + domain: StorageDomainPolicy | None = None + remote: RemoteImportPolicy = field(default_factory=RemoteImportPolicy) + production: ProductionPolicy = field(default_factory=ProductionPolicy) + capacity: CapacityPolicy = field(default_factory=CapacityPolicy) + wait: WaitPolicy = field(default_factory=WaitPolicy) + integrity: IntegrityPolicy = field(default_factory=IntegrityPolicy) + schema_version: int = 1 + + def __post_init__(self) -> None: + if isinstance(self.schema_version, bool) or self.schema_version != 1: + raise ValueError("only Host Weight Runtime config schema_version=1 is supported") + if not isinstance(self.mode, RuntimeMode): + raise ValueError("runtime mode must use RuntimeMode") + if self.domain is not None and not isinstance(self.domain, StorageDomainPolicy): + raise ValueError("runtime domain must use StorageDomainPolicy") + for name, value, expected in ( + ("remote", self.remote, RemoteImportPolicy), + ("production", self.production, ProductionPolicy), + ("capacity", self.capacity, CapacityPolicy), + ("wait", self.wait, WaitPolicy), + ("integrity", self.integrity, IntegrityPolicy), + ): + if not isinstance(value, expected): + raise ValueError(f"runtime {name} policy has the wrong type") + if self.mode is not RuntimeMode.DISABLED and self.domain is None: + raise ValueError("enabled Host Weight Runtime requires a storage domain") + + +__all__ = [ + "CapacityPolicy", + "HostWeightRuntimeConfig", + "IntegrityPolicy", + "ProductionPolicy", + "RemoteImportPolicy", + "RemoteOnMiss", + "RuntimeMode", + "StorageClass", + "StorageDomainPolicy", + "StorageScope", + "ValidationLevel", + "WaitPolicy", +] diff --git a/vllm_omni/host_weight_runtime/errors.py b/vllm_omni/host_weight_runtime/errors.py new file mode 100644 index 000000000..d3ab4f3c2 --- /dev/null +++ b/vllm_omni/host_weight_runtime/errors.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stable Host Weight Runtime failures and actions.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum + +from .identity import CanonicalJson + + +class ResolutionStage(str, Enum): + CONFIGURATION = "configuration" + DOMAIN = "domain" + LOOKUP = "lookup" + VALIDATION = "validation" + REMOTE = "remote" + PRODUCTION = "production" + CAPACITY = "capacity" + RESTORATION = "restoration" + LIFECYCLE = "lifecycle" + CANONICAL_LOADING = "canonical_loading" + + +class FailureCode(str, Enum): + INVALID_CONFIG = "invalid_config" + COHORT_POLICY_MISMATCH = "cohort_policy_mismatch" + DOMAIN_POLICY_MISMATCH = "domain_policy_mismatch" + DOMAIN_UNAVAILABLE = "domain_unavailable" + STORAGE_CLASS_MISMATCH = "storage_class_mismatch" + NUMA_PLACEMENT_INVALID = "numa_placement_invalid" + LOCAL_MISS = "local_miss" + ARTIFACT_DENIED = "artifact_denied" + ACTIVE_BUILD_TIMEOUT = "active_build_timeout" + MANIFEST_INCOMPATIBLE = "manifest_incompatible" + MANIFEST_CORRUPT = "manifest_corrupt" + CHECKSUM_MISMATCH = "checksum_mismatch" + IDENTITY_COLLISION = "identity_collision" + REMOTE_MISS = "remote_miss" + REMOTE_UNAVAILABLE = "remote_unavailable" + REMOTE_AUTH_FAILED = "remote_auth_failed" + REMOTE_TIMEOUT = "remote_timeout" + REMOTE_ARTIFACT_INVALID = "remote_artifact_invalid" + PRODUCER_UNAVAILABLE = "producer_unavailable" + PRODUCER_UNSUPPORTED = "producer_unsupported" + PRODUCER_FAILED = "producer_failed" + PRODUCTION_TIMEOUT = "production_timeout" + ARTIFACT_TOO_LARGE = "artifact_too_large" + STORE_LIMIT_EXCEEDED = "store_limit_exceeded" + INSUFFICIENT_CAPACITY = "insufficient_capacity" + ENOSPC = "enospc" + RESTORER_UNAVAILABLE = "restorer_unavailable" + RESTORE_PREFLIGHT_FAILED = "restore_preflight_failed" + RESTORE_COMMIT_FAILED = "restore_commit_failed" + ACTIVE_LEASE = "active_lease" + PUBLICATION_FAILED = "publication_failed" + QUARANTINE_FAILED = "quarantine_failed" + CANONICAL_SOURCE_FAILED = "canonical_source_failed" + UNSUPPORTED = "unsupported" + + +class ResolutionAction(str, Enum): + RETURN_LEASE = "return_lease" + TRY_REMOTE = "try_remote" + TRY_PRODUCER = "try_producer" + QUARANTINE_THEN_PRODUCE = "quarantine_then_produce" + CANONICAL_FALLBACK = "canonical_fallback" + FAIL_STARTUP = "fail_startup" + SKIP = "skip" + + +@dataclass(frozen=True) +class HostWeightFailure: + stage: ResolutionStage + code: FailureCode + retryable: bool + message: str + details: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + if not isinstance(self.stage, ResolutionStage) or not isinstance(self.code, FailureCode): + raise ValueError("host weight failure stage and code must use their typed enums") + if not isinstance(self.retryable, bool): + raise ValueError("host weight failure retryable must be a boolean") + if not isinstance(self.message, str) or not self.message: + raise ValueError("host weight failure requires a message") + if not isinstance(self.details, CanonicalJson): + raise ValueError("host weight failure details must use CanonicalJson") + + +class HostWeightError(RuntimeError): + """Internal exception carrying a stable, policy-neutral failure.""" + + def __init__(self, failure: HostWeightFailure): + super().__init__(failure.message) + self.failure = failure + + +__all__ = [ + "FailureCode", + "HostWeightError", + "HostWeightFailure", + "ResolutionAction", + "ResolutionStage", +] diff --git a/vllm_omni/host_weight_runtime/filesystem/__init__.py b/vllm_omni/host_weight_runtime/filesystem/__init__.py new file mode 100644 index 000000000..2083c1e85 --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/__init__.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Local filesystem Host Weight Store implementation.""" + +from .store import FilesystemHostWeightStore, detect_filesystem_type, detect_storage_class + +__all__ = [ + "FilesystemHostWeightStore", + "detect_filesystem_type", + "detect_storage_class", +] diff --git a/vllm_omni/host_weight_runtime/filesystem/locks.py b/vllm_omni/host_weight_runtime/filesystem/locks.py new file mode 100644 index 000000000..70a400dac --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/locks.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stable POSIX advisory locks used by the filesystem store.""" + +from __future__ import annotations + +import errno +import fcntl +import os +import time +from pathlib import Path + + +class FileLockTimeoutError(TimeoutError): + pass + + +class FileLock: + """One held flock whose inode is stable outside artifact directories.""" + + def __init__(self, path: Path, *, exclusive: bool, deadline: float | None, nonblocking: bool = False) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + self.path = path + self._fd = os.open(path, os.O_CREAT | os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW, 0o600) + self._creator_pid = os.getpid() + self._closed = False + operation = fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH + if nonblocking: + operation |= fcntl.LOCK_NB + try: + if nonblocking: + self._acquire_once(operation) + elif deadline is None: + self._acquire_blocking(operation) + else: + self._acquire_until(operation | fcntl.LOCK_NB, deadline) + except BaseException: + os.close(self._fd) + self._closed = True + raise + + def _acquire_once(self, operation: int) -> None: + try: + fcntl.flock(self._fd, operation) + except OSError as exc: + if exc.errno in (errno.EACCES, errno.EAGAIN): + raise FileLockTimeoutError(f"lock is active: {self.path}") from exc + raise + + def _acquire_blocking(self, operation: int) -> None: + while True: + try: + fcntl.flock(self._fd, operation) + return + except OSError as exc: + if exc.errno != errno.EINTR: + raise + + def _acquire_until(self, operation: int, deadline: float) -> None: + while True: + try: + fcntl.flock(self._fd, operation) + return + except OSError as exc: + if exc.errno == errno.EINTR: + continue + if exc.errno not in (errno.EACCES, errno.EAGAIN): + raise + remaining = deadline - time.monotonic() + if remaining <= 0: + raise FileLockTimeoutError(f"timed out waiting for {self.path}") from exc + time.sleep(min(0.05, remaining)) + + def close(self) -> None: + if self._closed: + return + self._closed = True + if os.getpid() != self._creator_pid: + # flock state belongs to the inherited open-file description. A + # child must only drop its descriptor reference; LOCK_UN would + # release the creator process's active lock as well. + os.close(self._fd) + return + try: + fcntl.flock(self._fd, fcntl.LOCK_UN) + finally: + os.close(self._fd) + + def __enter__(self) -> FileLock: + return self + + def __exit__(self, exc_type: object, exc: object, _traceback: object) -> None: + self.close() + + +def lock_is_active(path: Path) -> bool: + """Return a point-in-time contention snapshot without creating a lock file.""" + try: + fd = os.open(path, os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW) + except FileNotFoundError: + return False + try: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as exc: + if exc.errno in (errno.EACCES, errno.EAGAIN): + return True + raise + fcntl.flock(fd, fcntl.LOCK_UN) + return False + finally: + os.close(fd) + + +__all__ = ["FileLock", "FileLockTimeoutError", "lock_is_active"] diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py new file mode 100644 index 000000000..0f01e36da --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -0,0 +1,1593 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Crash-safe node-local filesystem implementation of HostWeightStore.""" + +from __future__ import annotations + +import contextlib +import errno +import hashlib +import json +import logging +import os +import shutil +import stat +import time +import uuid +from collections import defaultdict +from collections.abc import Callable, Iterator +from pathlib import Path +from typing import Protocol + +import regex as re +import torch +from safetensors import SafetensorError, safe_open + +from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel +from ..errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionStage +from ..identity import CanonicalJson, WeightArtifactIdentity, canonical_json +from ..inspection import ( + ArtifactInspection, + ArtifactInventoryEntry, + ArtifactInventoryState, + DomainInspection, +) +from ..lease import HostWeightLease, LeaseProvenance, MappedHostRegion +from ..manifest import ManifestValidationError, PublicationMarker, TensorManifestEntry, WeightManifest +from ..outcomes import StoreResult, StoreStatus +from ..protocols import BuildRequest, CoordinationScope, WeightProducer +from .locks import FileLock, FileLockTimeoutError, lock_is_active +from .writer import FilesystemArtifactWriter + +logger = logging.getLogger(__name__) + +_DOMAIN_FILE = "domain.json" +_DOMAIN_POLICY_FILE = "domain-policy.json" +_MANIFEST_FILE = "manifest.json" +_READY_FILE = "READY.json" +_MAX_METADATA_BYTES = 64 * 1024**2 +_FILE_HASH_CHUNK_BYTES = 8 * 1024**2 +_ARTIFACT_KEY_RE = re.compile(r"[0-9a-f]{64}") +_ARTIFACT_INVALIDATION_MODE = stat.S_ISVTX +_ARTIFACT_STRUCTURE_ERRNOS = frozenset( + { + errno.EACCES, + errno.EISDIR, + errno.ELOOP, + errno.ENOENT, + errno.ENOTDIR, + errno.EPERM, + } +) +_LOCAL_DISK_FILESYSTEM_TYPES = frozenset( + { + "aufs", + "bcachefs", + "btrfs", + "erofs", + "ext2", + "ext3", + "ext4", + "f2fs", + "fuseblk", + "jfs", + "ntfs3", + "overlay", + "reiserfs", + "rootfs", + "xfs", + "zfs", + } +) +_REMOTE_FILESYSTEM_TYPES = frozenset( + { + "9p", + "afs", + "beegfs", + "ceph", + "cifs", + "fuse.gcsfuse", + "fuse.s3fs", + "fuse.sshfs", + "gcsfuse", + "glusterfs", + "gpfs", + "lustre", + "nfs", + "nfs4", + "orangefs", + "panfs", + "s3fs", + "smb3", + "sshfs", + "virtiofs", + "wekafs", + } +) + + +class _ExitContext(Protocol): + def __exit__(self, exc_type: None, exc: None, _traceback: None) -> object: ... + + +def _make_fd_closer(fd: int) -> Callable[[], None]: + def close() -> None: + os.close(fd) + + return close + + +def _make_context_closer(context: _ExitContext) -> Callable[[], None]: + def close() -> None: + context.__exit__(None, None, None) + + return close + + +def _failure( + stage: ResolutionStage, + code: FailureCode, + message: str, + *, + retryable: bool = False, + details: object | None = None, +) -> HostWeightFailure: + return HostWeightFailure( + stage=stage, + code=code, + retryable=retryable, + message=message, + details=CanonicalJson.from_value({} if details is None else details), + ) + + +def _fsync_directory(path: Path) -> None: + fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _storage_io_failure( + exc: OSError, + *, + operation: str, + details: object | None = None, +) -> HostWeightFailure: + if exc.errno in (errno.ENOSPC, errno.EDQUOT): + return _failure( + ResolutionStage.CAPACITY, + FailureCode.ENOSPC, + f"{operation}: {exc}", + retryable=True, + details=details, + ) + return _failure( + ResolutionStage.DOMAIN, + FailureCode.DOMAIN_UNAVAILABLE, + f"{operation}: {exc}", + retryable=True, + details=details, + ) + + +def _path_exists_without_following(path: Path) -> bool: + try: + path.lstat() + except FileNotFoundError: + return False + return True + + +def _os_error_details(exc: OSError) -> dict[str, object]: + return { + "exception_type": type(exc).__name__, + "errno": exc.errno, + "message": str(exc), + } + + +@contextlib.contextmanager +def _release_lock_after_authoritative_commit( + lock: FileLock, + *, + key: str, + committed: Callable[[], bool], +) -> Iterator[None]: + try: + yield + finally: + try: + lock.close() + except OSError as exc: + if not committed(): + raise + logger.warning( + "artifact %s was authoritative before its build-lock release error: %s", + key, + exc, + ) + + +def _harden_open_directory(fd: int, path: Path, mode: int) -> None: + """Remove write bits from an open directory and make that change durable.""" + try: + os.fchmod(fd, mode) + except OSError as fd_error: + try: + os.chmod(path, mode, follow_symlinks=False) + except OSError as path_error: + raise path_error from fd_error + actual_mode = stat.S_IMODE(os.fstat(fd).st_mode) + if actual_mode != mode: + raise OSError( + errno.EIO, + f"artifact directory {path} has mode {actual_mode:#o}, expected {mode:#o}", + ) + os.fsync(fd) + + +def _artifact_has_invalidation_marker(path: Path) -> bool: + return bool(path.stat(follow_symlinks=False).st_mode & _ARTIFACT_INVALIDATION_MODE) + + +def _mark_open_artifact_invalidated(fd: int, path: Path) -> tuple[int, int]: + """Persist a fail-closed marker and return the open inode identity.""" + marked_mode = stat.S_IMODE(os.fstat(fd).st_mode) | _ARTIFACT_INVALIDATION_MODE + os.fchmod(fd, marked_mode) + marked_stat = os.fstat(fd) + if not marked_stat.st_mode & _ARTIFACT_INVALIDATION_MODE: + raise OSError(errno.EIO, f"artifact directory {path} was not marked invalid") + os.fsync(fd) + return marked_stat.st_dev, marked_stat.st_ino + + +def _mark_artifact_invalidated(path: Path) -> tuple[int, int]: + """Persist a fail-closed marker without allocating a new directory entry.""" + fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW) + try: + return _mark_open_artifact_invalidated(fd, path) + finally: + os.close(fd) + + +def _move_artifact_locked(source: Path, destination: Path) -> None: + """Move an artifact directory while its exclusive artifact lock is held. + + Supported local filesystems may reject renaming the mode-0555 directory + used for a READY artifact. Widen the inode only while the lock excludes + readers, then restore an immutable mode through the same open descriptor. + """ + source_stat = source.lstat() + if stat.S_ISDIR(source_stat.st_mode): + fd = os.open(source, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW) + try: + original_mode = stat.S_IMODE(os.fstat(fd).st_mode) + source_immutable_mode = original_mode & ~0o222 + transition_immutable_mode = source_immutable_mode | _ARTIFACT_INVALIDATION_MODE + destination_immutable_mode = source_immutable_mode & ~_ARTIFACT_INVALIDATION_MODE + writable_mode = transition_immutable_mode | stat.S_IWUSR + current_path = source + try: + _harden_open_directory(fd, source, transition_immutable_mode) + os.fchmod(fd, writable_mode) + os.replace(source, destination) + current_path = destination + except BaseException as move_error: + try: + _harden_open_directory(fd, current_path, source_immutable_mode) + except BaseException as hardening_error: + raise hardening_error from move_error + raise + # Keep the invalidation marker durable until both directory + # entries are durable. If either parent fsync fails (or the + # process crashes first), a rename rollback must not expose an + # authoritative artifact with the marker already cleared. + _harden_open_directory(fd, destination, transition_immutable_mode) + _fsync_directory(source.parent) + _fsync_directory(destination.parent) + _harden_open_directory(fd, destination, destination_immutable_mode) + finally: + os.close(fd) + else: + os.replace(source, destination) + _fsync_directory(source.parent) + _fsync_directory(destination.parent) + + +def _write_atomic_json(path: Path, value: object, *, mode: int = 0o444) -> None: + data = canonical_json(value) + temporary = path.parent / f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp" + fd = os.open(temporary, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o600) + try: + offset = 0 + while offset < len(data): + count = os.write(fd, data[offset:]) + if count <= 0: + raise OSError(errno.EIO, f"short write for {temporary}") + offset += count + os.fsync(fd) + os.fchmod(fd, mode) + os.fsync(fd) + except BaseException: + os.close(fd) + temporary.unlink(missing_ok=True) + raise + else: + os.close(fd) + os.replace(temporary, path) + _fsync_directory(path.parent) + + +def _read_json_file(path: Path) -> dict[str, object]: + fd = os.open(path, os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW) + try: + return _read_json_fd(fd, path.name) + finally: + os.close(fd) + + +def _read_json_at(directory_fd: int, name: str, *, require_read_only: bool = False) -> dict[str, object]: + fd = os.open(name, os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, dir_fd=directory_fd) + try: + return _read_json_fd(fd, name, require_read_only=require_read_only) + finally: + os.close(fd) + + +def _read_json_fd(fd: int, name: str, *, require_read_only: bool = False) -> dict[str, object]: + file_stat = os.fstat(fd) + if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_nlink != 1: + raise ManifestValidationError(f"metadata file {name!r} is not one regular file") + if file_stat.st_size > _MAX_METADATA_BYTES: + raise ManifestValidationError(f"metadata file {name!r} exceeds {_MAX_METADATA_BYTES} bytes") + if require_read_only and (file_stat.st_uid != os.geteuid() or file_stat.st_mode & 0o222): + raise ManifestValidationError(f"published metadata file {name!r} is not immutable") + chunks: list[bytes] = [] + offset = 0 + while offset < file_stat.st_size: + chunk = os.pread(fd, min(1024**2, file_stat.st_size - offset), offset) + if not chunk: + raise ManifestValidationError(f"metadata file {name!r} ended unexpectedly") + chunks.append(chunk) + offset += len(chunk) + try: + value = json.loads(b"".join(chunks)) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ManifestValidationError(f"metadata file {name!r} is not valid JSON") from exc + if not isinstance(value, dict): + raise ManifestValidationError(f"metadata file {name!r} must contain an object") + return value + + +def _sha256_fd(fd: int, size: int) -> str: + digest = hashlib.sha256() + offset = 0 + while offset < size: + chunk = os.pread(fd, min(_FILE_HASH_CHUNK_BYTES, size - offset), offset) + if not chunk: + raise ManifestValidationError("artifact file ended while hashing") + digest.update(chunk) + offset += len(chunk) + return digest.hexdigest() + + +def _mount_unescape(value: str) -> str: + return value.replace("\\040", " ").replace("\\011", "\t").replace("\\012", "\n").replace("\\134", "\\") + + +def detect_filesystem_type(path: Path) -> str: + """Return the kernel filesystem type for the most specific mount.""" + resolved = path.resolve() + mountinfo = Path("/proc/self/mountinfo") + if not mountinfo.is_file(): # pragma: no cover - vLLM-Omni production is Linux + raise ValueError("cannot verify filesystem locality without /proc/self/mountinfo") + best_length = -1 + best_type = "" + for line in mountinfo.read_text(encoding="utf-8").splitlines(): + left, separator, right = line.partition(" - ") + if not separator: + continue + left_fields = left.split() + right_fields = right.split() + if len(left_fields) < 5 or not right_fields: + continue + mount_point = Path(_mount_unescape(left_fields[4])) + try: + resolved.relative_to(mount_point) + except ValueError: + continue + if len(str(mount_point)) > best_length: + best_length = len(str(mount_point)) + best_type = right_fields[0] + if not best_type: + raise ValueError(f"cannot identify the filesystem containing {resolved}") + return best_type + + +def _storage_class_for_filesystem_type(filesystem_type: str) -> StorageClass: + if filesystem_type in {"tmpfs", "ramfs"}: + return StorageClass.TMPFS + if filesystem_type in _LOCAL_DISK_FILESYSTEM_TYPES: + return StorageClass.DISK + if filesystem_type in _REMOTE_FILESYSTEM_TYPES: + raise ValueError(f"filesystem type {filesystem_type!r} is not node-local") + raise ValueError(f"filesystem type {filesystem_type!r} is not recognized as a local filesystem") + + +def detect_storage_class(path: Path) -> StorageClass: + """Classify a verified node-local mount without spawning platform tools.""" + return _storage_class_for_filesystem_type(detect_filesystem_type(path)) + + +class FilesystemHostWeightStore: + """A cooperative local store with stable locks and immutable publication.""" + + def __init__( + self, + domain: StorageDomainPolicy, + capacity: CapacityPolicy, + integrity: IntegrityPolicy, + ) -> None: + self.domain_policy = domain + self.capacity_policy = capacity + self.integrity_policy = integrity + self.root = domain.root.expanduser().resolve() + self.locks_dir = self.root / "locks" + self.artifacts_dir = self.root / "artifacts" + self.tmp_dir = self.root / "tmp" + self.quarantine_dir = self.root / "quarantine" + self.deny_dir = self.root / "deny" + if integrity.local_lookup is ValidationLevel.FS_VERITY: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.UNSUPPORTED, + "filesystem-verity validation is not implemented", + ) + ) + if domain.numa_prefault or domain.verify_numa_residency: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.UNSUPPORTED, + "enforced NUMA placement and residency verification are not implemented", + ) + ) + if integrity.background_scrub: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.UNSUPPORTED, + "background integrity scrubbing is not implemented", + ) + ) + try: + self._initialize_domain() + except HostWeightError: + raise + except OSError as exc: + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.DOMAIN_UNAVAILABLE, + f"cannot initialize storage domain {self.root}: {exc}", + retryable=True, + ) + ) from exc + except (ValueError, ManifestValidationError) as exc: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"storage domain metadata is invalid under {self.root}: {exc}", + ) + ) from exc + + def _initialize_domain(self) -> None: + try: + self.filesystem_type = detect_filesystem_type(self.root) + detected_storage_class = _storage_class_for_filesystem_type(self.filesystem_type) + except ValueError as exc: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"local filesystem storage is unavailable for {self.root}: {exc}", + ) + ) from exc + if detected_storage_class is not self.domain_policy.storage_class: + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.STORAGE_CLASS_MISMATCH, + "configured storage class " + f"{self.domain_policy.storage_class.value} differs from detected {detected_storage_class.value}", + ) + ) + + try: + self.root.mkdir(parents=True, exist_ok=True, mode=0o700) + self.locks_dir.mkdir(exist_ok=True, mode=0o700) + for directory in (self.artifacts_dir, self.tmp_dir, self.quarantine_dir, self.deny_dir): + directory.mkdir(exist_ok=True, mode=0o700) + except OSError as exc: + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.DOMAIN_UNAVAILABLE, + f"cannot initialize {self.root}: {exc}", + retryable=True, + ) + ) from exc + + if self.integrity_policy.require_trusted_permissions: + for directory in ( + self.root, + self.locks_dir, + self.artifacts_dir, + self.tmp_dir, + self.quarantine_dir, + self.deny_dir, + ): + directory_stat = directory.lstat() + if ( + not stat.S_ISDIR(directory_stat.st_mode) + or directory_stat.st_uid != os.geteuid() + or directory_stat.st_mode & 0o022 + ): + raise HostWeightError( + _failure( + ResolutionStage.DOMAIN, + FailureCode.DOMAIN_UNAVAILABLE, + f"host weight store directory is not exclusively controlled by this user: {directory}", + ) + ) + + with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=None): + domain_path = self.root / _DOMAIN_FILE + if domain_path.exists(): + domain_metadata = _read_json_file(domain_path) + expected = { + "schema_version": 1, + "scope": self.domain_policy.scope.value, + "domain_id": self.domain_policy.domain_id, + "backend": "local_filesystem", + "storage_class": self.domain_policy.storage_class.value, + "filesystem_type": self.filesystem_type, + } + if set(domain_metadata) != {*expected, "domain_uuid"}: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"storage domain metadata has an incompatible schema under {self.root}", + ) + ) + actual = {key: domain_metadata.get(key) for key in expected} + domain_uuid = domain_metadata.get("domain_uuid") + if actual != expected or not isinstance(domain_uuid, str) or not domain_uuid: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"storage domain metadata does not match configured policy under {self.root}", + ) + ) + else: + domain_metadata = { + "schema_version": 1, + "domain_uuid": uuid.uuid4().hex, + "scope": self.domain_policy.scope.value, + "domain_id": self.domain_policy.domain_id, + "backend": "local_filesystem", + "storage_class": self.domain_policy.storage_class.value, + "filesystem_type": self.filesystem_type, + } + _write_atomic_json(domain_path, domain_metadata) + self.domain_uuid = str(domain_metadata["domain_uuid"]) + + policy_path = self.root / _DOMAIN_POLICY_FILE + requested_policy = self._capacity_document(policy_version=1) + if policy_path.exists(): + actual_policy = _read_json_file(policy_path) + if set(actual_policy) != set(requested_policy): + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"capacity policy has an incompatible schema under {self.root}", + ) + ) + comparable_actual = {key: actual_policy.get(key) for key in requested_policy if key != "policy_version"} + comparable_requested = { + key: value for key, value in requested_policy.items() if key != "policy_version" + } + if comparable_actual != comparable_requested: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"capacity policy differs from the authoritative policy under {self.root}", + ) + ) + policy_version = actual_policy.get("policy_version") + if isinstance(policy_version, bool) or not isinstance(policy_version, int) or policy_version < 1: + raise HostWeightError( + _failure( + ResolutionStage.CONFIGURATION, + FailureCode.DOMAIN_POLICY_MISMATCH, + f"capacity policy has an invalid version under {self.root}", + ) + ) + self.policy_version = policy_version + else: + _write_atomic_json(policy_path, requested_policy) + self.policy_version = 1 + + def _capacity_document(self, *, policy_version: int) -> dict[str, object]: + return { + "schema_version": 1, + "policy_version": policy_version, + "max_artifact_bytes": self.capacity_policy.max_artifact_bytes, + "max_store_bytes": self.capacity_policy.max_store_bytes, + "min_free_bytes": self.capacity_policy.min_free_bytes, + "eviction": self.capacity_policy.eviction, + } + + def _build_lock_path(self, key: str) -> Path: + return self.locks_dir / f"{key}.build.lock" + + def _artifact_lock_path(self, key: str) -> Path: + return self.locks_dir / f"{key}.artifact.lock" + + def _entry_path(self, key: str) -> Path: + return self.artifacts_dir / key + + def _deny_path(self, key: str) -> Path: + return self.deny_dir / f"{key}.json" + + def lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None = None, + ) -> StoreResult: + return self._lookup(identity, validation=validation, deadline=deadline, record_invalid=True) + + def _lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None, + record_invalid: bool, + ) -> StoreResult: + if validation is ValidationLevel.FS_VERITY: + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.VALIDATION, + FailureCode.UNSUPPORTED, + "filesystem-verity validation is modeled but not implemented", + ), + ) + try: + artifact_lock = FileLock(self._artifact_lock_path(identity.key), exclusive=False, deadline=deadline) + except FileLockTimeoutError as exc: + return StoreResult( + StoreStatus.TIMEOUT, + failure=_failure( + ResolutionStage.LOOKUP, + FailureCode.ACTIVE_BUILD_TIMEOUT, + str(exc), + retryable=True, + ), + ) + except OSError as exc: + return StoreResult( + StoreStatus.FAILED, + failure=_failure(ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE, str(exc), retryable=True), + ) + + entry_path = self._entry_path(identity.key) + try: + entry_exists = _path_exists_without_following(entry_path) + deny_exists = entry_exists and _path_exists_without_following(self._deny_path(identity.key)) + except OSError as exc: + artifact_lock.close() + return StoreResult( + StoreStatus.FAILED, + failure=_storage_io_failure( + exc, + operation=f"failed to inspect artifact state for {identity.key}", + ), + ) + if not entry_exists: + artifact_lock.close() + return StoreResult(StoreStatus.MISS) + if deny_exists: + artifact_lock.close() + return StoreResult( + StoreStatus.INVALID, + failure=_failure( + ResolutionStage.LOOKUP, + FailureCode.ARTIFACT_DENIED, + f"artifact {identity.key} is denied pending quarantine", + ), + ) + try: + lease = self._open_lease(identity, entry_path, artifact_lock, validation) + try: + denied_during_lookup = _path_exists_without_following( + self._deny_path(identity.key) + ) or _artifact_has_invalidation_marker(entry_path) + except OSError: + lease.close() + raise + if denied_during_lookup: + lease.close() + return StoreResult( + StoreStatus.INVALID, + failure=_failure( + ResolutionStage.LOOKUP, + FailureCode.ARTIFACT_DENIED, + f"artifact {identity.key} was denied during lookup", + ), + ) + return StoreResult(StoreStatus.HIT, lease=lease) + except HostWeightError as exc: + failure = exc.failure + except OSError as exc: + if exc.errno not in _ARTIFACT_STRUCTURE_ERRNOS: + artifact_lock.close() + return StoreResult( + StoreStatus.FAILED, + failure=_storage_io_failure( + exc, + operation=f"failed to read or validate artifact {identity.key}", + ), + ) + failure = _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_CORRUPT, + f"artifact {identity.key} failed validation: {exc}", + ) + except (SafetensorError, ValueError, ManifestValidationError) as exc: + failure = _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_CORRUPT, + f"artifact {identity.key} failed validation: {exc}", + ) + + try: + if record_invalid: + self._write_deny(identity.key, failure) + except OSError as exc: + containment = self._contain_invalid_after_deny_failure( + identity.key, + artifact_lock, + deadline=deadline, + ) + return StoreResult( + StoreStatus.FAILED, + failure=_storage_io_failure( + exc, + operation=f"failed to deny invalid artifact {identity.key}", + details={"containment": containment}, + ), + ) + finally: + artifact_lock.close() + return StoreResult(StoreStatus.INVALID, failure=failure) + + def _open_lease( + self, + identity: WeightArtifactIdentity, + entry_path: Path, + artifact_lock: FileLock, + validation: ValidationLevel, + ) -> HostWeightLease: + resources: list[Callable[[], None]] = [] + try: + entry_fd = os.open(entry_path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW) + resources.append(_make_fd_closer(entry_fd)) + entry_stat = os.fstat(entry_fd) + if entry_stat.st_mode & _ARTIFACT_INVALIDATION_MODE: + raise HostWeightError( + _failure( + ResolutionStage.LOOKUP, + FailureCode.ARTIFACT_DENIED, + f"artifact {identity.key} is marked invalid pending quarantine", + ) + ) + if self.integrity_policy.require_trusted_permissions and ( + entry_stat.st_uid != os.geteuid() or entry_stat.st_mode & 0o222 + ): + raise ManifestValidationError("published artifact directory is not immutable") + names = set(os.listdir(entry_fd)) + marker_document = _read_json_at( + entry_fd, + _READY_FILE, + require_read_only=self.integrity_policy.require_trusted_permissions, + ) + manifest_document = _read_json_at( + entry_fd, + _MANIFEST_FILE, + require_read_only=self.integrity_policy.require_trusted_permissions, + ) + marker = PublicationMarker.from_dict(marker_document) + manifest = WeightManifest.from_dict(manifest_document) + if marker.schema_version != 1 or manifest.schema_version != 1 or marker.integrity_mode != "full_checksum": + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_INCOMPATIBLE, + "artifact publication schema or integrity mode is unsupported", + ) + ) + expected_names = {_READY_FILE, _MANIFEST_FILE, *(entry.file_name for entry in manifest.files)} + if names != expected_names: + raise ManifestValidationError( + f"artifact directory entries differ: missing={sorted(expected_names - names)}, " + f"unexpected={sorted(names - expected_names)}" + ) + manifest_digest = hashlib.sha256(manifest.canonical_bytes).hexdigest() + if marker.manifest_sha256 != manifest_digest: + raise HostWeightError( + _failure(ResolutionStage.VALIDATION, FailureCode.CHECKSUM_MISMATCH, "manifest digest mismatch") + ) + if marker.artifact_content_sha256 != manifest.artifact_content_sha256: + raise HostWeightError( + _failure(ResolutionStage.VALIDATION, FailureCode.CHECKSUM_MISMATCH, "content digest mismatch") + ) + if manifest.artifact_key != identity.key or manifest.identity.canonical_bytes != identity.canonical_bytes: + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_INCOMPATIBLE, + "artifact semantic identity differs from the exact request", + ) + ) + if ( + manifest.producer_schema != identity.producer.manifest_schema + or manifest.restorer_schema != identity.producer.restorer_schema + ): + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.MANIFEST_INCOMPATIBLE, + "artifact producer or restorer schema differs from the exact request", + ) + ) + + tensors_by_file: dict[str, list[TensorManifestEntry]] = defaultdict(list) + for tensor_entry in manifest.tensors: + tensors_by_file[tensor_entry.file_name].append(tensor_entry) + tensors: dict[str, torch.Tensor] = {} + regions: list[MappedHostRegion] = [] + backing_digest = hashlib.sha256() + for file_entry in manifest.files: + fd = os.open( + file_entry.file_name, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=entry_fd, + ) + resources.append(_make_fd_closer(fd)) + file_stat = os.fstat(fd) + if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_nlink != 1: + raise ManifestValidationError(f"payload {file_entry.file_name!r} is not one regular file") + if file_stat.st_size != file_entry.size: + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.CHECKSUM_MISMATCH, + f"payload size differs for {file_entry.file_name!r}", + ) + ) + if self.integrity_policy.require_trusted_permissions and ( + file_stat.st_uid != os.geteuid() or file_stat.st_mode & 0o222 + ): + raise ManifestValidationError(f"published payload {file_entry.file_name!r} is not immutable") + if ( + validation is ValidationLevel.FULL_CHECKSUM + and _sha256_fd(fd, file_stat.st_size) != file_entry.sha256 + ): + raise HostWeightError( + _failure( + ResolutionStage.VALIDATION, + FailureCode.CHECKSUM_MISMATCH, + f"payload digest differs for {file_entry.file_name!r}", + ) + ) + backing_digest.update( + canonical_json( + { + "file": file_entry.file_name, + "device": file_stat.st_dev, + "inode": file_stat.st_ino, + "size": file_stat.st_size, + } + ) + ) + handle_context = safe_open(f"/proc/self/fd/{fd}", framework="pt", device="cpu") + handle = handle_context.__enter__() + resources.append(_make_context_closer(handle_context)) + expected_entries = tensors_by_file[file_entry.file_name] + if set(handle.keys()) != {entry.storage_key for entry in expected_entries}: + raise ManifestValidationError(f"payload tensor keys differ for {file_entry.file_name!r}") + addresses: list[tuple[int, int]] = [] + for tensor_entry in expected_entries: + tensor = handle.get_tensor(tensor_entry.storage_key) + actual_nbytes = tensor.numel() * tensor.element_size() + if ( + tuple(tensor.shape) != tensor_entry.shape + or str(tensor.dtype) != tensor_entry.dtype + or tuple(tensor.stride()) != tensor_entry.stride + or actual_nbytes != tensor_entry.nbytes + or tensor.device.type != "cpu" + ): + raise ManifestValidationError(f"mapped tensor metadata differs for {tensor_entry.name!r}") + tensors[tensor_entry.name] = tensor + if actual_nbytes: + addresses.append((tensor.data_ptr(), tensor.data_ptr() + actual_nbytes)) + if addresses: + merged: list[list[int]] = [] + for start, end in sorted(addresses): + if merged and start <= merged[-1][1]: + merged[-1][1] = max(merged[-1][1], end) + else: + merged.append([start, end]) + regions.extend( + MappedHostRegion( + file_name=file_entry.file_name, + address=start, + length=end - start, + ) + for start, end in merged + ) + + backing_digest.update(self.domain_uuid.encode("ascii")) + backing_digest.update(marker.publication_id.encode("ascii")) + provenance = LeaseProvenance( + resolution_id=uuid.uuid4().hex, + identity_digest=identity.key, + artifact_content_sha256=manifest.artifact_content_sha256, + domain_uuid=self.domain_uuid, + publication_id=marker.publication_id, + resolution_source="local", + artifact_origin="producer", + producer_id=identity.producer.producer_id, + provider_id=None, + validation_level=validation.value, + backing_fingerprint=backing_digest.hexdigest(), + ) + return HostWeightLease( + identity=identity, + manifest=manifest, + tensors=tensors, + mapped_regions=tuple(regions), + provenance=provenance, + resource_closers=tuple(resources), + release_artifact_lock=artifact_lock.close, + ) + except BaseException: + for closer in reversed(resources): + with contextlib.suppress(Exception): + closer() + raise + + def get_or_build( + self, + request: BuildRequest, + producer: WeightProducer, + *, + validation: ValidationLevel, + deadline: float, + ) -> StoreResult: + identity = request.identity + if producer.spec.coordination_scope is not CoordinationScope.SINGLE_PROCESS: + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.PRODUCER_UNSUPPORTED, + f"coordination scope {producer.spec.coordination_scope.value} is not implemented", + ), + ) + if tuple(output.key for output in producer.spec.outputs) != (identity.key,): + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.PRODUCER_UNAVAILABLE, + "producer does not declare exactly the requested artifact identity", + ), + ) + + initial = self.lookup(identity, validation=validation, deadline=deadline) + if initial.status is StoreStatus.HIT: + return initial + if initial.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return initial + + try: + build_lock = FileLock(self._build_lock_path(identity.key), exclusive=True, deadline=deadline) + except FileLockTimeoutError as exc: + return StoreResult( + StoreStatus.TIMEOUT, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.ACTIVE_BUILD_TIMEOUT, + str(exc), + retryable=True, + ), + ) + except OSError as exc: + return StoreResult( + StoreStatus.FAILED, + failure=_failure(ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE, str(exc), retryable=True), + ) + + ready_outcome_authoritative = False + with _release_lock_after_authoritative_commit( + build_lock, + key=identity.key, + committed=lambda: ready_outcome_authoritative, + ): + recheck = self.lookup(identity, validation=validation, deadline=deadline) + if recheck.status is StoreStatus.HIT: + ready_outcome_authoritative = True + return StoreResult(StoreStatus.JOINED, lease=recheck.lease) + if recheck.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return recheck + + try: + with FileLock(self._artifact_lock_path(identity.key), exclusive=True, deadline=deadline): + self._repair_quarantine_transitions_locked(identity.key) + if _path_exists_without_following(self._entry_path(identity.key)): + self._quarantine_locked(identity.key, reason="validation_failure") + self._remove_cleanup_tombstones_locked(identity.key) + self._remove_deny_locked(identity.key) + self._remove_stale_temps_locked(identity.key) + except FileLockTimeoutError as exc: + return StoreResult( + StoreStatus.TIMEOUT, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.ACTIVE_BUILD_TIMEOUT, + str(exc), + retryable=True, + ), + ) + except OSError as exc: + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.LIFECYCLE, + FailureCode.QUARANTINE_FAILED, + f"failed to prepare artifact rebuild: {exc}", + retryable=True, + ), + ) + + staging = self.tmp_dir / f"{identity.key}.{os.getpid()}.{uuid.uuid4().hex}" + writer: FilesystemArtifactWriter | None = None + manifest: WeightManifest | None = None + producer_completed = False + try: + staging.mkdir(mode=0o700) + writer = FilesystemArtifactWriter(staging, capacity_check=self._check_capacity) + metadata = producer.produce(writer) + producer_completed = True + manifest, _ = writer.finalize(identity, metadata) + artifact_bytes = self._tree_bytes(staging) + self._check_capacity(artifact_bytes, 0) + # Keep staging writable for the atomic rename. The sticky bit + # is an internal invalidation marker until publication hardens + # the destination. + os.chmod(staging, 0o755 | _ARTIFACT_INVALIDATION_MODE) + _fsync_directory(staging) + except HostWeightError as exc: + if writer is not None: + writer.abort() + self._discard_temp(identity.key, staging) + return StoreResult(StoreStatus.FAILED, failure=exc.failure) + except OSError as exc: + if writer is not None: + writer.abort() + self._discard_temp(identity.key, staging) + if exc.errno in (errno.ENOSPC, errno.EDQUOT): + code = FailureCode.ENOSPC + stage = ResolutionStage.CAPACITY + elif writer is None or producer_completed: + code = FailureCode.DOMAIN_UNAVAILABLE + stage = ResolutionStage.DOMAIN + else: + code = FailureCode.PRODUCER_FAILED + stage = ResolutionStage.PRODUCTION + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + stage, + code, + f"artifact production failed: {exc}", + retryable=stage in {ResolutionStage.CAPACITY, ResolutionStage.DOMAIN}, + ), + ) + except Exception as exc: + if writer is not None: + writer.abort() + self._discard_temp(identity.key, staging) + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.PRODUCTION, + FailureCode.PRODUCER_FAILED, + f"artifact producer failed: {exc}", + ), + ) + + assert manifest is not None + publication_error: OSError | None = None + publication_containment: dict[str, object] | None = None + outcome_uncertain_error: OSError | None = None + competing_entry_error: OSError | None = None + try: + with FileLock(self._artifact_lock_path(identity.key), exclusive=True, deadline=None): + entry_path = self._entry_path(identity.key) + if _path_exists_without_following(entry_path): + try: + existing_document = _read_json_file(entry_path / _MANIFEST_FILE) + except OSError as exc: + competing_entry_error = exc + raise + existing = WeightManifest.from_dict(existing_document) + same_semantic_content = ( + existing.identity.canonical_bytes == manifest.identity.canonical_bytes + and existing.artifact_content_sha256 == manifest.artifact_content_sha256 + and existing.producer_schema == manifest.producer_schema + and existing.restorer_schema == manifest.restorer_schema + and existing.format_metadata == manifest.format_metadata + ) + if not same_semantic_content: + collision = self.quarantine_dir / f"{identity.key}.identity-collision.{uuid.uuid4().hex}" + _move_artifact_locked(staging, collision) + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.VALIDATION, + FailureCode.IDENTITY_COLLISION, + "the same semantic identity produced different artifact content", + ), + ) + ready_outcome_authoritative = True + self._remove_tree(staging) + self._remove_deny_locked(identity.key) + else: + published = False + try: + os.replace(staging, entry_path) + published = True + os.chmod(entry_path, 0o555) + _fsync_directory(entry_path) + _fsync_directory(self.tmp_dir) + self._remove_deny_locked(identity.key) + # The destination-parent sync is the filesystem + # publication commit. + _fsync_directory(self.artifacts_dir) + ready_outcome_authoritative = True + except OSError as exc: + publication_error = exc + if published: + try: + publication_containment = self._contain_failed_publication_locked(identity.key) + except OSError as containment_error: + outcome_uncertain_error = containment_error + raise + raise + except (OSError, ValueError, ManifestValidationError) as exc: + if not ready_outcome_authoritative: + self._discard_temp(identity.key, staging) + if competing_entry_error is not None: + return StoreResult( + StoreStatus.FAILED, + failure=_storage_io_failure( + competing_entry_error, + operation=f"failed to inspect competing publication {identity.key}", + details={ + "outcome_uncertain": True, + "competing_entry": "preserved", + }, + ), + ) + if outcome_uncertain_error is not None: + assert publication_error is not None + return StoreResult( + StoreStatus.FAILED, + failure=_storage_io_failure( + outcome_uncertain_error, + operation=f"failed to contain outcome-uncertain publication {identity.key}", + details={ + "outcome_uncertain": True, + "publication_error": _os_error_details(publication_error), + "containment_error": _os_error_details(outcome_uncertain_error), + }, + ), + ) + return StoreResult( + StoreStatus.FAILED, + failure=_failure( + ResolutionStage.LIFECYCLE, + FailureCode.PUBLICATION_FAILED, + f"atomic artifact publication failed: {exc}", + details={"containment": publication_containment} + if publication_containment is not None + else None, + ), + ) + logger.warning("artifact %s was authoritative before a post-authority error: %s", identity.key, exc) + self._discard_temp(identity.key, staging) + + opened = self.lookup(identity, validation=validation, deadline=deadline) + if opened.status is not StoreStatus.HIT: + return opened + return StoreResult(StoreStatus.BUILT, lease=opened.lease) + + def _check_capacity(self, projected_artifact_bytes: int, additional_bytes: int) -> None: + policy = self.capacity_policy + if policy.max_artifact_bytes is not None and projected_artifact_bytes > policy.max_artifact_bytes: + raise HostWeightError( + _failure( + ResolutionStage.CAPACITY, + FailureCode.ARTIFACT_TOO_LARGE, + f"projected artifact size {projected_artifact_bytes} exceeds {policy.max_artifact_bytes}", + retryable=True, + ) + ) + store_bytes = self._tree_bytes(self.root) + if policy.max_store_bytes is not None and store_bytes + additional_bytes > policy.max_store_bytes: + raise HostWeightError( + _failure( + ResolutionStage.CAPACITY, + FailureCode.STORE_LIMIT_EXCEEDED, + f"store size plus allocation would exceed {policy.max_store_bytes}", + retryable=True, + ) + ) + free_bytes = shutil.disk_usage(self.root).free + if free_bytes - additional_bytes < policy.min_free_bytes: + raise HostWeightError( + _failure( + ResolutionStage.CAPACITY, + FailureCode.INSUFFICIENT_CAPACITY, + f"allocation would leave less than {policy.min_free_bytes} free bytes", + retryable=True, + ) + ) + + def _write_deny(self, key: str, failure: HostWeightFailure) -> None: + path = self._deny_path(key) + document = { + "schema_version": 1, + "stage": failure.stage.value, + "reason": failure.code.value, + "message": failure.message, + } + try: + fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o444) + except FileExistsError: + return + try: + data = canonical_json(document) + offset = 0 + while offset < len(data): + written = os.write(fd, data[offset:]) + if written == 0: + raise OSError("short write while publishing deny marker") + offset += written + os.fsync(fd) + finally: + os.close(fd) + _fsync_directory(self.deny_dir) + + def _remove_deny_locked(self, key: str) -> None: + path = self._deny_path(key) + if _path_exists_without_following(path): + path.unlink() + _fsync_directory(self.deny_dir) + + def _quarantine_locked(self, key: str, *, reason: str) -> Path | None: + entry = self._entry_path(key) + if not _path_exists_without_following(entry): + return None + destination = self.quarantine_dir / f"{key}.{reason}.{uuid.uuid4().hex}" + _move_artifact_locked(entry, destination) + return destination + + def _contain_failed_publication_locked(self, key: str) -> dict[str, object]: + """Keep an outcome-uncertain publication from becoming a later HIT.""" + containment: dict[str, object] = {} + marker_error: OSError | None = None + try: + _mark_artifact_invalidated(self._entry_path(key)) + except OSError as exc: + marker_error = exc + containment["invalidation_marker_error"] = _os_error_details(exc) + else: + containment["invalidation_marker"] = "durable" + try: + destination = self._quarantine_locked(key, reason="publication_failure") + except OSError as quarantine_error: + containment["quarantine"] = "failed" + containment["quarantine_error"] = _os_error_details(quarantine_error) + if marker_error is not None: + raise quarantine_error from marker_error + else: + containment["quarantine"] = "already_absent" if destination is None else "completed" + return containment + + def _contain_invalid_after_deny_failure( + self, + key: str, + artifact_lock: FileLock, + *, + deadline: float | None, + ) -> dict[str, object]: + """Attempt persistent exclusion when deny-marker publication fails.""" + containment: dict[str, object] = {} + + entry = self._entry_path(key) + entry_fd: int | None = None + invalidated_inode: tuple[int, int] | None = None + try: + entry_fd = os.open(entry, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW) + entry_stat = os.fstat(entry_fd) + invalidated_inode = (entry_stat.st_dev, entry_stat.st_ino) + _mark_open_artifact_invalidated(entry_fd, entry) + except OSError as exc: + containment["invalidation_marker_error"] = _os_error_details(exc) + else: + containment["invalidation_marker"] = "durable" + + try: + try: + artifact_lock.close() + except OSError as exc: + containment["shared_lock_release_error"] = _os_error_details(exc) + + try: + exclusive_lock = FileLock( + self._artifact_lock_path(key), + exclusive=True, + deadline=deadline, + nonblocking=deadline is None, + ) + except FileLockTimeoutError as exc: + containment["quarantine"] = "deferred_for_active_lease" + containment["exclusive_lock_error"] = _os_error_details(exc) + return containment + except OSError as exc: + containment["quarantine"] = "deferred_for_lock_error" + containment["exclusive_lock_error"] = _os_error_details(exc) + return containment + + try: + with exclusive_lock: + if not _path_exists_without_following(entry): + containment["quarantine"] = "already_absent" + elif invalidated_inode is None: + containment["quarantine"] = "deferred_for_unverified_artifact" + else: + current_stat = entry.lstat() + current_inode = (current_stat.st_dev, current_stat.st_ino) + if current_inode != invalidated_inode: + containment["quarantine"] = "already_replaced" + else: + destination = self._quarantine_locked(key, reason="validation_failure") + containment["quarantine"] = "already_absent" if destination is None else "completed" + except OSError as exc: + containment["quarantine"] = "failed" + containment["quarantine_error"] = _os_error_details(exc) + return containment + finally: + if entry_fd is not None: + try: + os.close(entry_fd) + except OSError as exc: + containment["artifact_fd_close_error"] = _os_error_details(exc) + + def _remove_stale_temps_locked(self, key: str) -> None: + for path in self.tmp_dir.glob(f"{key}.*"): + self._remove_tree(path) + _fsync_directory(self.tmp_dir) + + def _remove_cleanup_tombstones_locked(self, key: str) -> None: + tombstones = tuple(self.quarantine_dir.glob(f"{key}.cleanup.*")) + for tombstone in tombstones: + self._remove_tree(tombstone) + if tombstones: + _fsync_directory(self.quarantine_dir) + + def _repair_quarantine_transitions_locked(self, key: str) -> None: + repaired = False + for path in self.quarantine_dir.glob(f"{key}.*"): + path_stat = path.lstat() + if not stat.S_ISDIR(path_stat.st_mode) or not path_stat.st_mode & _ARTIFACT_INVALIDATION_MODE: + continue + fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW) + try: + mode = stat.S_IMODE(os.fstat(fd).st_mode) + _harden_open_directory( + fd, + path, + mode & ~0o222 & ~_ARTIFACT_INVALIDATION_MODE, + ) + finally: + os.close(fd) + repaired = True + if repaired: + _fsync_directory(self.quarantine_dir) + + def _discard_temp(self, key: str, path: Path) -> None: + with contextlib.suppress(OSError, FileLockTimeoutError): + with FileLock(self._artifact_lock_path(key), exclusive=True, deadline=None): + self._remove_tree(path) + _fsync_directory(self.tmp_dir) + + @staticmethod + def _remove_tree(path: Path) -> None: + if path.is_symlink() or path.is_file(): + path.unlink(missing_ok=True) + elif _path_exists_without_following(path): + for directory, subdirectories, _ in os.walk(path, topdown=False, followlinks=False): + for name in subdirectories: + child = Path(directory) / name + if not child.is_symlink(): + os.chmod(child, 0o700) + os.chmod(directory, 0o700) + shutil.rmtree(path) + + @staticmethod + def _tree_bytes(root: Path) -> int: + with contextlib.suppress(OSError): + root_stat = root.lstat() + if stat.S_ISREG(root_stat.st_mode): + return root_stat.st_size + total = 0 + for directory, _, files in os.walk(root, followlinks=False): + for name in files: + path = Path(directory) / name + with contextlib.suppress(OSError): + file_stat = path.lstat() + if stat.S_ISREG(file_stat.st_mode): + total += file_stat.st_size + return total + + def cleanup(self, identity: WeightArtifactIdentity) -> HostWeightFailure | None: + """Explicitly remove an inactive artifact without waiting for leases.""" + try: + try: + build_lock = FileLock( + self._build_lock_path(identity.key), + exclusive=True, + deadline=None, + nonblocking=True, + ) + except FileLockTimeoutError: + return _failure( + ResolutionStage.LIFECYCLE, + FailureCode.ACTIVE_BUILD_TIMEOUT, + f"artifact {identity.key} has an active builder", + retryable=True, + ) + with build_lock: + try: + artifact_lock = FileLock( + self._artifact_lock_path(identity.key), + exclusive=True, + deadline=None, + nonblocking=True, + ) + except FileLockTimeoutError: + return _failure( + ResolutionStage.LIFECYCLE, + FailureCode.ACTIVE_LEASE, + f"artifact {identity.key} has an active lease", + retryable=True, + ) + with artifact_lock: + self._repair_quarantine_transitions_locked(identity.key) + self._quarantine_locked(identity.key, reason="cleanup") + self._remove_cleanup_tombstones_locked(identity.key) + self._remove_deny_locked(identity.key) + except OSError as exc: + return _failure( + ResolutionStage.LIFECYCLE, + FailureCode.QUARANTINE_FAILED, + f"failed to clean up artifact {identity.key}: {exc}", + retryable=True, + ) + return None + + def inspect_domain(self) -> DomainInspection: + usage = shutil.disk_usage(self.root) + inventory: list[ArtifactInventoryEntry] = [] + for state, directory in ( + (ArtifactInventoryState.READY, self.artifacts_dir), + (ArtifactInventoryState.DENIED, self.deny_dir), + (ArtifactInventoryState.QUARANTINED, self.quarantine_dir), + (ArtifactInventoryState.TEMPORARY, self.tmp_dir), + ): + for path in directory.iterdir(): + artifact_key = path.name[:64] + if _ARTIFACT_KEY_RE.fullmatch(artifact_key) is None: + continue + reason = None + if state is ArtifactInventoryState.DENIED: + with contextlib.suppress(OSError, ValueError, ManifestValidationError): + reason_value = _read_json_file(path).get("reason") + if isinstance(reason_value, str): + reason = reason_value + inventory.append( + ArtifactInventoryEntry( + artifact_key=artifact_key, + state=state, + storage_name=path.name, + size_bytes=self._tree_bytes(path), + reason=reason, + ) + ) + inventory.sort(key=lambda item: (item.artifact_key, item.state.value, item.storage_name)) + return DomainInspection( + domain_uuid=self.domain_uuid, + scope=self.domain_policy.scope.value, + domain_id=self.domain_policy.domain_id, + storage_class=self.domain_policy.storage_class.value, + filesystem_type=self.filesystem_type, + policy_version=self.policy_version, + store_bytes=self._tree_bytes(self.root), + filesystem_free_bytes=usage.free, + max_artifact_bytes=self.capacity_policy.max_artifact_bytes, + max_store_bytes=self.capacity_policy.max_store_bytes, + min_free_bytes=self.capacity_policy.min_free_bytes, + inventory=tuple(inventory), + ) + + def inspect_artifact( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + ) -> ArtifactInspection: + result = self._lookup( + identity, + validation=validation, + deadline=time.monotonic(), + record_invalid=False, + ) + manifest = None + publication_id = None + backing_fingerprint = None + if result.lease is not None: + manifest = result.lease.manifest + publication_id = result.lease.provenance.publication_id + backing_fingerprint = result.lease.provenance.backing_fingerprint + result.lease.close() + entry = self._entry_path(identity.key) + return ArtifactInspection( + artifact_key=identity.key, + status=result.status, + size_bytes=self._tree_bytes(entry), + manifest=manifest, + publication_id=publication_id, + backing_fingerprint=backing_fingerprint, + failure=result.failure, + build_lock_active=lock_is_active(self._build_lock_path(identity.key)), + artifact_lock_active=lock_is_active(self._artifact_lock_path(identity.key)), + ) + + +__all__ = ["FilesystemHostWeightStore", "detect_storage_class"] diff --git a/vllm_omni/host_weight_runtime/filesystem/writer.py b/vllm_omni/host_weight_runtime/filesystem/writer.py new file mode 100644 index 000000000..283ee2a9e --- /dev/null +++ b/vllm_omni/host_weight_runtime/filesystem/writer.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Store-scoped streaming safetensors publication.""" + +from __future__ import annotations + +import contextlib +import errno +import hashlib +import json +import os +import struct +import uuid +from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from types import TracebackType + +import torch + +from ..identity import WeightArtifactIdentity, canonical_json +from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest +from ..protocols import ProductionMetadata, TensorWriteSpec + +_HASH_CHUNK_BYTES = 8 * 1024**2 +_MAX_HASH_WORKERS = 8 +_MANIFEST_FILE = "manifest.json" +_READY_FILE = "READY.json" + +_SAFETENSORS_DTYPES: dict[torch.dtype, str] = { + torch.bool: "BOOL", + torch.uint8: "U8", + torch.int8: "I8", + torch.int16: "I16", + torch.int32: "I32", + torch.int64: "I64", + torch.float16: "F16", + torch.bfloat16: "BF16", + torch.float32: "F32", + torch.float64: "F64", +} +for _torch_name, _safe_name in ( + ("float8_e4m3fn", "F8_E4M3"), + ("float8_e5m2", "F8_E5M2"), + ("float8_e4m3fnuz", "F8_E4M3FNUZ"), + ("float8_e5m2fnuz", "F8_E5M2FNUZ"), +): + if (_dtype := getattr(torch, _torch_name, None)) is not None: + _SAFETENSORS_DTYPES[_dtype] = _safe_name + + +class ArtifactWriterError(RuntimeError): + pass + + +def _safe_relative_name(value: str) -> str: + path = PurePosixPath(value) + if not value or path.is_absolute() or len(path.parts) != 1 or path.name != value or value in {".", ".."}: + raise ArtifactWriterError(f"artifact payload name must be one safe relative basename: {value!r}") + if not value.endswith(".safetensors"): + raise ArtifactWriterError(f"artifact payload must use the .safetensors format: {value!r}") + return value + + +def _pwrite_all(fd: int, data: bytes | bytearray | memoryview, offset: int) -> None: + view = memoryview(data).cast("B") + written = 0 + while written < len(view): + count = os.pwrite(fd, view[written:], offset + written) + if count <= 0: + raise OSError(errno.EIO, "short pwrite while publishing host weights") + written += count + + +def _fsync_directory(path: Path) -> None: + fd = os.open(path, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _fsync_payload(fd: int) -> None: + try: + os.fsync(fd) + finally: + os.close(fd) + + +@dataclass(frozen=True) +class _PayloadDigests: + file: str + tensors: dict[str, str] + + +def _scan_payload(path: Path, specs: tuple[TensorWriteSpec, ...]) -> _PayloadDigests: + """Calculate file and tensor digests in one ordered read.""" + file_digest = hashlib.sha256() + tensor_digests: dict[str, str] = {} + with path.open("rb") as handle: + prefix = handle.read(8) + if len(prefix) != 8: + raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}") + header_size = struct.unpack(" None: + self._size = size + self._ranges: list[tuple[int, int]] = [] + + def add(self, start: int, end: int) -> None: + if start < 0 or end < start or end > self._size: + raise ArtifactWriterError(f"write range [{start}, {end}) exceeds tensor storage {self._size}") + if start == end: + return + for existing_start, existing_end in self._ranges: + if start < existing_end and existing_start < end: + raise ArtifactWriterError(f"overlapping tensor write ranges [{start}, {end}) and {self._ranges}") + self._ranges.append((start, end)) + + def complete(self) -> bool: + if self._size == 0: + return True + cursor = 0 + for start, end in sorted(self._ranges): + if start != cursor: + return False + cursor = end + return cursor == self._size + + +class StreamingSafeTensorFile: + """One preallocated safetensors file filled through bounded CPU writes.""" + + def __init__( + self, + owner: FilesystemArtifactWriter, + relative_name: str, + specs: tuple[TensorWriteSpec, ...], + *, + ordered: bool, + ) -> None: + self._owner = owner + self.relative_name = _safe_relative_name(relative_name) + if not specs: + raise ArtifactWriterError("a tensor payload file must contain at least one tensor") + by_name = {spec.name: spec for spec in specs} + if len(by_name) != len(specs): + raise ArtifactWriterError("tensor payload contains duplicate storage keys") + unsupported = [spec.dtype for spec in specs if spec.dtype not in _SAFETENSORS_DTYPES] + if unsupported: + raise ArtifactWriterError(f"unsupported safetensors dtypes: {unsupported}") + self.specs = tuple(sorted(specs, key=lambda spec: spec.name)) + self._specs = {spec.name: spec for spec in self.specs} + self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs} + self._ordered = ordered + self._closed = False + + offsets: dict[str, int] = {} + cursor = 0 + for spec in self.specs: + offsets[spec.name] = cursor + cursor += spec.nbytes + header: dict[str, object] = {"__metadata__": {"format": "pt", "host_weight_runtime": "1"}} + for spec in self.specs: + start = offsets[spec.name] + header[spec.name] = { + "dtype": _SAFETENSORS_DTYPES[spec.dtype], + "shape": list(spec.shape), + "data_offsets": [start, start + spec.nbytes], + } + encoded_header = json.dumps(header, sort_keys=True, separators=(",", ":")).encode("utf-8") + encoded_header += b" " * (-len(encoded_header) % 8) + self._payload_offset = 8 + len(encoded_header) + self._offsets = offsets + self._path = owner.staging_dir / self.relative_name + total_size = self._payload_offset + cursor + owner._reserve_file(total_size) + self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600) + prefix = struct.pack(" memoryview: + if tensor.device.type != "cpu" or tensor.is_meta: + raise ArtifactWriterError("artifact producers must provide materialized CPU tensor chunks") + contiguous = tensor.detach().contiguous() + return memoryview(contiguous.view(torch.uint8).numpy()).cast("B") + + def write_tensor(self, name: str, tensor: torch.Tensor) -> None: + spec = self._get_spec(name) + if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype: + raise ArtifactWriterError( + f"tensor {name!r} differs from its declared shape/dtype: " + f"got {tuple(tensor.shape)}/{tensor.dtype}, expected {spec.shape}/{spec.dtype}" + ) + self._write(name, 0, self._bytes(tensor)) + + def write_rows(self, name: str, start_row: int, tensor: torch.Tensor) -> None: + spec = self._get_spec(name) + if not spec.shape: + if start_row != 0 or tuple(tensor.shape) != (): + raise ArtifactWriterError(f"scalar tensor {name!r} only accepts one complete scalar write") + self.write_tensor(name, tensor) + return + if tensor.dtype != spec.dtype or tuple(tensor.shape[1:]) != spec.shape[1:]: + raise ArtifactWriterError(f"row chunk for {name!r} differs from its declared dtype or trailing shape") + row_count = tensor.shape[0] + if start_row < 0 or start_row + row_count > spec.shape[0]: + raise ArtifactWriterError(f"row chunk for {name!r} exceeds its declared first dimension") + row_bytes = spec.nbytes // spec.shape[0] if spec.shape[0] else 0 + self._write(name, start_row * row_bytes, self._bytes(tensor)) + + def _get_spec(self, name: str) -> TensorWriteSpec: + if self._closed: + raise ArtifactWriterError("cannot write through a closed tensor file") + try: + return self._specs[name] + except KeyError as exc: + raise ArtifactWriterError(f"tensor {name!r} was not declared in this payload file") from exc + + def _write(self, name: str, byte_offset: int, data: memoryview) -> None: + file_offset = self._payload_offset + self._offsets[name] + byte_offset + if self._ordered and file_offset != self._write_cursor: + raise ArtifactWriterError( + f"ordered tensor payload expected offset {self._write_cursor}, got {file_offset} for {name!r}" + ) + coverage = self._coverage[name] + coverage.add(byte_offset, byte_offset + len(data)) + _pwrite_all(self._fd, data, file_offset) + if self._ordered: + assert self._file_digest is not None + self._file_digest.update(data) + self._tensor_digests[name].update(data) + self._write_cursor += len(data) + + def close(self) -> None: + if self._closed: + return + missing = [name for name, coverage in self._coverage.items() if not coverage.complete()] + if missing: + self.abort() + raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}") + self._closed = True + digests = None + if self._ordered: + assert self._file_digest is not None + digests = _PayloadDigests( + self._file_digest.hexdigest(), + {name: digest.hexdigest() for name, digest in self._tensor_digests.items()}, + ) + try: + os.fchmod(self._fd, 0o444) + self._owner._closed(self, digests, self._fd) + except BaseException: + os.close(self._fd) + raise + + def abort(self) -> None: + if self._closed: + return + self._closed = True + os.close(self._fd) + size = self._path.stat().st_size if self._path.exists() else 0 + self._path.unlink(missing_ok=True) + self._owner._aborted(self, size) + + def __enter__(self) -> StreamingSafeTensorFile: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + _traceback: TracebackType | None, + ) -> None: + if exc_type is None: + self.close() + else: + self.abort() + + +class FilesystemArtifactWriter: + """Scoped writer that cannot publish outside one invisible staging dir.""" + + def __init__( + self, + staging_dir: Path, + *, + capacity_check: Callable[[int, int], None], + ) -> None: + self.staging_dir = staging_dir + self._capacity_check = capacity_check + self._active: set[StreamingSafeTensorFile] = set() + self._files: dict[str, tuple[TensorWriteSpec, ...]] = {} + self._digests: dict[str, _PayloadDigests] = {} + self._reserved_bytes = 0 + self._finalized = False + self._fsync_executor: ThreadPoolExecutor | None = None + self._fsync_futures: list[Future[None]] = [] + + def open_tensor_file( + self, + relative_name: str, + specs: tuple[TensorWriteSpec, ...], + *, + ordered: bool = False, + ) -> StreamingSafeTensorFile: + if self._finalized: + raise ArtifactWriterError("cannot add payloads after artifact finalization") + relative_name = _safe_relative_name(relative_name) + if relative_name in self._files or any(item.relative_name == relative_name for item in self._active): + raise ArtifactWriterError(f"artifact payload {relative_name!r} was declared more than once") + return StreamingSafeTensorFile(self, relative_name, specs, ordered=ordered) + + def _reserve_file(self, size: int) -> None: + self._capacity_check(self._reserved_bytes + size, size) + self._reserved_bytes += size + + def _forget_reserved_file(self, size: int) -> None: + self._reserved_bytes = max(0, self._reserved_bytes - size) + + def _opened(self, writer: StreamingSafeTensorFile) -> None: + self._active.add(writer) + + def _closed(self, writer: StreamingSafeTensorFile, digests: _PayloadDigests | None, fd: int) -> None: + self._active.remove(writer) + self._files[writer.relative_name] = writer.specs + if digests is not None: + self._digests[writer.relative_name] = digests + if self._fsync_executor is None: + self._fsync_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="hwr-fsync") + self._fsync_futures.append(self._fsync_executor.submit(_fsync_payload, fd)) + + def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None: + self._active.discard(writer) + self._forget_reserved_file(size) + + def abort(self) -> None: + for writer in tuple(self._active): + with contextlib.suppress(Exception): + writer.abort() + with contextlib.suppress(Exception): + self._wait_for_payload_fsync() + + def _wait_for_payload_fsync(self) -> None: + if self._fsync_executor is None: + return + executor = self._fsync_executor + futures = tuple(self._fsync_futures) + self._fsync_executor = None + self._fsync_futures.clear() + executor.shutdown(wait=True) + for future in futures: + future.result() + + def finalize( + self, + identity: WeightArtifactIdentity, + metadata: ProductionMetadata, + ) -> tuple[WeightManifest, PublicationMarker]: + if self._finalized: + raise ArtifactWriterError("artifact writer was finalized more than once") + if self._active: + raise ArtifactWriterError("artifact producer left tensor payload files open") + if not self._files: + raise ArtifactWriterError("artifact producer emitted no tensor payload files") + if metadata.producer_schema != identity.producer.manifest_schema: + raise ArtifactWriterError("production metadata does not match the identity's manifest schema") + if metadata.restorer_schema != identity.producer.restorer_schema: + raise ArtifactWriterError("production metadata does not match the identity's restorer schema") + + self._wait_for_payload_fsync() + + tensor_entries: list[TensorManifestEntry] = [] + file_entries: list[FileManifestEntry] = [] + seen_tensors: set[str] = set() + file_names = sorted(self._files) + unhashed = [file_name for file_name in file_names if file_name not in self._digests] + if unhashed: + with ThreadPoolExecutor(max_workers=min(_MAX_HASH_WORKERS, len(unhashed))) as executor: + scans = executor.map( + lambda file_name: _scan_payload(self.staging_dir / file_name, self._files[file_name]), + unhashed, + ) + self._digests.update(zip(unhashed, scans, strict=True)) + + for file_name in file_names: + path = self.staging_dir / file_name + specs = self._files[file_name] + digests = self._digests[file_name] + file_entries.append( + FileManifestEntry( + file_name=file_name, + size=path.stat().st_size, + sha256=digests.file, + ) + ) + for spec in specs: + if spec.name in seen_tensors: + raise ArtifactWriterError(f"artifact tensor {spec.name!r} occurs in multiple payload files") + seen_tensors.add(spec.name) + tensor_entries.append( + TensorManifestEntry( + name=spec.name, + kind=spec.kind, + role=spec.role, + file_name=file_name, + storage_key=spec.name, + shape=spec.shape, + dtype=str(spec.dtype), + stride=tuple(torch.empty(spec.shape, dtype=spec.dtype, device="meta").stride()), + nbytes=spec.nbytes, + sha256=digests.tensors[spec.name], + ) + ) + + ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name)) + content_digest = hashlib.sha256() + for entry in ordered_tensor_entries: + content_digest.update(canonical_json(entry.semantic_dict())) + + manifest = WeightManifest( + schema_version=1, + artifact_key=identity.key, + identity=identity, + producer_schema=metadata.producer_schema, + restorer_schema=metadata.restorer_schema, + tensors=ordered_tensor_entries, + files=tuple(file_entries), + artifact_content_sha256=content_digest.hexdigest(), + format_metadata=metadata.format_metadata, + ) + manifest_bytes = manifest.canonical_bytes + self._write_metadata_file(_MANIFEST_FILE, manifest_bytes) + marker = PublicationMarker( + schema_version=1, + publication_id=uuid.uuid4().hex, + manifest_sha256=hashlib.sha256(manifest_bytes).hexdigest(), + artifact_content_sha256=manifest.artifact_content_sha256, + integrity_mode="full_checksum", + ) + self._write_metadata_file(_READY_FILE, marker.canonical_bytes) + _fsync_directory(self.staging_dir) + self._finalized = True + return manifest, marker + + def _write_metadata_file(self, name: str, data: bytes) -> None: + path = self.staging_dir / name + fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o600) + try: + _pwrite_all(fd, data, 0) + os.fsync(fd) + os.fchmod(fd, 0o444) + os.fsync(fd) + finally: + os.close(fd) + + +__all__ = [ + "ArtifactWriterError", + "FilesystemArtifactWriter", + "StreamingSafeTensorFile", +] diff --git a/vllm_omni/host_weight_runtime/identity.py b/vllm_omni/host_weight_runtime/identity.py new file mode 100644 index 000000000..77b769728 --- /dev/null +++ b/vllm_omni/host_weight_runtime/identity.py @@ -0,0 +1,407 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Deterministic semantic identity for runtime-ready host weights.""" + +from __future__ import annotations + +import hashlib +import json +import math +from dataclasses import dataclass, field +from typing import TypeAlias, cast + +JsonScalar: TypeAlias = None | bool | int | float | str +JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"] + + +class IdentityValidationError(ValueError): + """Raised when an artifact identity is not deterministic JSON.""" + + +def _validate_json(value: object, path: str = "$") -> JsonValue: + if value is None or isinstance(value, (bool, int, str)): + return value + if isinstance(value, float): + if not math.isfinite(value): + raise IdentityValidationError(f"{path} contains a non-finite float") + return value + if isinstance(value, list | tuple): + return [_validate_json(item, f"{path}[{index}]") for index, item in enumerate(value)] + if isinstance(value, dict): + normalized: dict[str, JsonValue] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise IdentityValidationError(f"{path} contains a non-string mapping key") + normalized[key] = _validate_json(item, f"{path}.{key}") + return {key: normalized[key] for key in sorted(normalized)} + raise IdentityValidationError(f"{path} contains unsupported value type {type(value).__qualname__}") + + +def canonical_json(value: object) -> bytes: + """Return one stable JSON encoding, rejecting process-specific objects.""" + return json.dumps( + _validate_json(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + + +@dataclass(frozen=True) +class CanonicalJson: + """Immutable snapshot of deterministic JSON metadata.""" + + encoded: bytes + + def __post_init__(self) -> None: + if not isinstance(self.encoded, bytes): + raise IdentityValidationError("canonical metadata must be immutable bytes") + try: + value = json.loads(self.encoded) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise IdentityValidationError("canonical metadata is not valid JSON") from exc + if canonical_json(value) != self.encoded: + raise IdentityValidationError("canonical metadata does not use the required stable encoding") + + @classmethod + def from_value(cls, value: object) -> CanonicalJson: + return cls(canonical_json(value)) + + @classmethod + def empty(cls) -> CanonicalJson: + return cls(b"{}") + + def to_value(self) -> JsonValue: + return json.loads(self.encoded) + + +def _require_text(name: str, value: object) -> str: + if not isinstance(value, str) or not value: + raise IdentityValidationError(f"{name} must not be empty") + return value + + +def _require_int(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise IdentityValidationError(f"{name} must be an integer") + return value + + +def _require_metadata(name: str, value: object) -> None: + if not isinstance(value, CanonicalJson): + raise IdentityValidationError(f"{name} must use CanonicalJson") + + +def _require_exact_keys(value: dict[str, object], expected: set[str], kind: str) -> None: + keys = set(value) + if keys != expected: + raise IdentityValidationError( + f"{kind} fields differ: missing={sorted(expected - keys)}, unexpected={sorted(keys - expected)}" + ) + + +@dataclass(frozen=True) +class WeightSourceIdentity: + model_id: str + revision: str + fingerprint: str + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("source.model_id", self.model_id) + _require_text("source.revision", self.revision) + _require_text("source.fingerprint", self.fingerprint) + _require_metadata("source.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "model_id": self.model_id, + "revision": self.revision, + "fingerprint": self.fingerprint, + "metadata": self.metadata.to_value(), + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightSourceIdentity: + _require_exact_keys(value, {"model_id", "revision", "fingerprint", "metadata"}, "source identity") + return cls( + model_id=_require_text("source.model_id", value["model_id"]), + revision=_require_text("source.revision", value["revision"]), + fingerprint=_require_text("source.fingerprint", value["fingerprint"]), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class ComponentIdentity: + name: str + ownership: str + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("component.name", self.name) + _require_text("component.ownership", self.ownership) + _require_metadata("component.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return {"name": self.name, "ownership": self.ownership, "metadata": self.metadata.to_value()} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> ComponentIdentity: + _require_exact_keys(value, {"name", "ownership", "metadata"}, "component identity") + return cls( + name=_require_text("component.name", value["name"]), + ownership=_require_text("component.ownership", value["ownership"]), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class WeightRepresentation: + name: str + dtype: str + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("representation.name", self.name) + _require_text("representation.dtype", self.dtype) + _require_metadata("representation.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return {"name": self.name, "dtype": self.dtype, "metadata": self.metadata.to_value()} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightRepresentation: + _require_exact_keys(value, {"name", "dtype", "metadata"}, "weight representation") + return cls( + name=_require_text("representation.name", value["name"]), + dtype=_require_text("representation.dtype", value["dtype"]), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class RuntimeWeightLayout: + name: str + tensor_parallel_size: int = 1 + tensor_parallel_rank: int = 0 + sequence_parallel_size: int = 1 + sequence_parallel_backend: str = "none" + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("layout.name", self.name) + tensor_parallel_size = _require_int("layout.tensor_parallel_size", self.tensor_parallel_size) + tensor_parallel_rank = _require_int("layout.tensor_parallel_rank", self.tensor_parallel_rank) + sequence_parallel_size = _require_int("layout.sequence_parallel_size", self.sequence_parallel_size) + if tensor_parallel_size < 1: + raise IdentityValidationError("tensor_parallel_size must be positive") + if not 0 <= tensor_parallel_rank < tensor_parallel_size: + raise IdentityValidationError("tensor_parallel_rank must be within tensor_parallel_size") + if sequence_parallel_size < 1: + raise IdentityValidationError("sequence_parallel_size must be positive") + _require_text("layout.sequence_parallel_backend", self.sequence_parallel_backend) + _require_metadata("layout.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "name": self.name, + "tensor_parallel_size": self.tensor_parallel_size, + "tensor_parallel_rank": self.tensor_parallel_rank, + "sequence_parallel_size": self.sequence_parallel_size, + "sequence_parallel_backend": self.sequence_parallel_backend, + "metadata": self.metadata.to_value(), + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> RuntimeWeightLayout: + expected = { + "name", + "tensor_parallel_size", + "tensor_parallel_rank", + "sequence_parallel_size", + "sequence_parallel_backend", + "metadata", + } + _require_exact_keys(value, expected, "runtime weight layout") + return cls( + name=_require_text("layout.name", value["name"]), + tensor_parallel_size=_require_int("layout.tensor_parallel_size", value["tensor_parallel_size"]), + tensor_parallel_rank=_require_int("layout.tensor_parallel_rank", value["tensor_parallel_rank"]), + sequence_parallel_size=_require_int("layout.sequence_parallel_size", value["sequence_parallel_size"]), + sequence_parallel_backend=_require_text( + "layout.sequence_parallel_backend", value["sequence_parallel_backend"] + ), + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class AdaptationIdentity: + kind: str = "base" + fingerprint: str | None = None + metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + _require_text("adaptation.kind", self.kind) + if self.kind != "base" and not self.fingerprint: + raise IdentityValidationError("non-base adaptation requires a fingerprint") + if self.fingerprint is not None and (not isinstance(self.fingerprint, str) or not self.fingerprint): + raise IdentityValidationError("adaptation fingerprint must be a non-empty string or null") + _require_metadata("adaptation.metadata", self.metadata) + + def to_dict(self) -> dict[str, JsonValue]: + return {"kind": self.kind, "fingerprint": self.fingerprint, "metadata": self.metadata.to_value()} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> AdaptationIdentity: + _require_exact_keys(value, {"kind", "fingerprint", "metadata"}, "adaptation identity") + fingerprint = value["fingerprint"] + if fingerprint is not None and not isinstance(fingerprint, str): + raise IdentityValidationError("adaptation fingerprint must be a string or null") + return cls( + kind=_require_text("adaptation.kind", value["kind"]), + fingerprint=fingerprint, + metadata=CanonicalJson.from_value(value["metadata"]), + ) + + +@dataclass(frozen=True) +class ProducerIdentity: + producer_id: str + version: str + implementation_fingerprint: str + manifest_schema: str + restorer_schema: str + + def __post_init__(self) -> None: + for name, value in ( + ("producer.producer_id", self.producer_id), + ("producer.version", self.version), + ("producer.implementation_fingerprint", self.implementation_fingerprint), + ("producer.manifest_schema", self.manifest_schema), + ("producer.restorer_schema", self.restorer_schema), + ): + _require_text(name, value) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "producer_id": self.producer_id, + "version": self.version, + "implementation_fingerprint": self.implementation_fingerprint, + "manifest_schema": self.manifest_schema, + "restorer_schema": self.restorer_schema, + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> ProducerIdentity: + expected = { + "producer_id", + "version", + "implementation_fingerprint", + "manifest_schema", + "restorer_schema", + } + _require_exact_keys(value, expected, "producer identity") + return cls( + producer_id=_require_text("producer.producer_id", value["producer_id"]), + version=_require_text("producer.version", value["version"]), + implementation_fingerprint=_require_text( + "producer.implementation_fingerprint", value["implementation_fingerprint"] + ), + manifest_schema=_require_text("producer.manifest_schema", value["manifest_schema"]), + restorer_schema=_require_text("producer.restorer_schema", value["restorer_schema"]), + ) + + +@dataclass(frozen=True) +class WeightArtifactIdentity: + schema_version: int + source: WeightSourceIdentity + component: ComponentIdentity + representation: WeightRepresentation + layout: RuntimeWeightLayout + adaptation: AdaptationIdentity + producer: ProducerIdentity + + def __post_init__(self) -> None: + if _require_int("identity.schema_version", self.schema_version) < 1: + raise IdentityValidationError("identity schema_version must be positive") + for name, value, expected in ( + ("source", self.source, WeightSourceIdentity), + ("component", self.component, ComponentIdentity), + ("representation", self.representation, WeightRepresentation), + ("layout", self.layout, RuntimeWeightLayout), + ("adaptation", self.adaptation, AdaptationIdentity), + ("producer", self.producer, ProducerIdentity), + ): + if not isinstance(value, expected): + raise IdentityValidationError(f"identity {name} has the wrong contract type") + + def to_dict(self) -> dict[str, JsonValue]: + return { + "schema_version": self.schema_version, + "source": self.source.to_dict(), + "component": self.component.to_dict(), + "representation": self.representation.to_dict(), + "layout": self.layout.to_dict(), + "adaptation": self.adaptation.to_dict(), + "producer": self.producer.to_dict(), + } + + @property + def canonical_bytes(self) -> bytes: + return canonical_json(self.to_dict()) + + @property + def key(self) -> str: + return hashlib.sha256(self.canonical_bytes).hexdigest() + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightArtifactIdentity: + expected = { + "schema_version", + "source", + "component", + "representation", + "layout", + "adaptation", + "producer", + } + _require_exact_keys(value, expected, "weight artifact identity") + source = value["source"] + component = value["component"] + representation = value["representation"] + layout = value["layout"] + adaptation = value["adaptation"] + producer = value["producer"] + if any( + not isinstance(section, dict) + for section in (source, component, representation, layout, adaptation, producer) + ): + raise IdentityValidationError("identity sections must be JSON objects") + return cls( + schema_version=_require_int("identity.schema_version", value["schema_version"]), + source=WeightSourceIdentity.from_dict(cast(dict[str, object], source)), + component=ComponentIdentity.from_dict(cast(dict[str, object], component)), + representation=WeightRepresentation.from_dict(cast(dict[str, object], representation)), + layout=RuntimeWeightLayout.from_dict(cast(dict[str, object], layout)), + adaptation=AdaptationIdentity.from_dict(cast(dict[str, object], adaptation)), + producer=ProducerIdentity.from_dict(cast(dict[str, object], producer)), + ) + + +__all__ = [ + "AdaptationIdentity", + "CanonicalJson", + "ComponentIdentity", + "IdentityValidationError", + "JsonValue", + "ProducerIdentity", + "RuntimeWeightLayout", + "WeightArtifactIdentity", + "WeightRepresentation", + "WeightSourceIdentity", + "canonical_json", +] diff --git a/vllm_omni/host_weight_runtime/inspection.py b/vllm_omni/host_weight_runtime/inspection.py new file mode 100644 index 000000000..c61760502 --- /dev/null +++ b/vllm_omni/host_weight_runtime/inspection.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed, read-only snapshots for local Host Weight Store inspection.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Protocol + +from .config import ValidationLevel +from .errors import HostWeightFailure +from .identity import WeightArtifactIdentity +from .manifest import WeightManifest +from .outcomes import StoreStatus + + +class ArtifactInventoryState(str, Enum): + READY = "ready" + DENIED = "denied" + QUARANTINED = "quarantined" + TEMPORARY = "temporary" + + +@dataclass(frozen=True) +class ArtifactInventoryEntry: + artifact_key: str + state: ArtifactInventoryState + storage_name: str + size_bytes: int + reason: str | None = None + + +@dataclass(frozen=True) +class DomainInspection: + domain_uuid: str + scope: str + domain_id: str + storage_class: str + filesystem_type: str + policy_version: int + store_bytes: int + filesystem_free_bytes: int + max_artifact_bytes: int | None + max_store_bytes: int | None + min_free_bytes: int + inventory: tuple[ArtifactInventoryEntry, ...] + + +@dataclass(frozen=True) +class ArtifactInspection: + artifact_key: str + status: StoreStatus + size_bytes: int + manifest: WeightManifest | None + publication_id: str | None + backing_fingerprint: str | None + failure: HostWeightFailure | None + build_lock_active: bool + artifact_lock_active: bool + + +class HostWeightInspector(Protocol): + def inspect_domain(self) -> DomainInspection: ... + + def inspect_artifact( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + ) -> ArtifactInspection: ... + + +__all__ = [ + "ArtifactInspection", + "ArtifactInventoryEntry", + "ArtifactInventoryState", + "DomainInspection", + "HostWeightInspector", +] diff --git a/vllm_omni/host_weight_runtime/lease.py b/vllm_omni/host_weight_runtime/lease.py new file mode 100644 index 000000000..f3a62e067 --- /dev/null +++ b/vllm_omni/host_weight_runtime/lease.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Process-local ownership for immutable mmap-backed host tensors.""" + +from __future__ import annotations + +import contextlib +import threading +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import NoReturn, SupportsIndex + +import torch + +from .identity import WeightArtifactIdentity +from .manifest import WeightManifest + + +@dataclass(frozen=True) +class MappedHostRegion: + """One process-local virtual range belonging to an artifact payload file.""" + + file_name: str + address: int + length: int + read_only: bool = True + + def __post_init__(self) -> None: + if not self.file_name or self.address < 0 or self.length < 0: + raise ValueError("mapped host region contains invalid fields") + + +@dataclass(frozen=True) +class LeaseProvenance: + resolution_id: str + identity_digest: str + artifact_content_sha256: str + domain_uuid: str + publication_id: str + resolution_source: str + artifact_origin: str + producer_id: str | None + provider_id: str | None + validation_level: str + backing_fingerprint: str + + +class HostWeightLease: + """Stable tensor views and resources retained until transport teardown.""" + + def __init__( + self, + *, + identity: WeightArtifactIdentity, + manifest: WeightManifest, + tensors: Mapping[str, torch.Tensor], + mapped_regions: tuple[MappedHostRegion, ...], + provenance: LeaseProvenance, + resource_closers: tuple[Callable[[], None], ...], + release_artifact_lock: Callable[[], None], + ) -> None: + if manifest.identity.canonical_bytes != identity.canonical_bytes: + raise ValueError("lease manifest identity differs from the requested identity") + if set(tensors) != {entry.name for entry in manifest.tensors}: + raise ValueError("lease tensor names differ from the validated manifest") + entries = {entry.name: entry for entry in manifest.tensors} + for name, tensor in tensors.items(): + entry = entries[name] + if ( + not isinstance(tensor, torch.Tensor) + or tensor.device.type != "cpu" + or tuple(tensor.shape) != entry.shape + or str(tensor.dtype) != entry.dtype + or tuple(tensor.stride()) != entry.stride + or tensor.numel() * tensor.element_size() != entry.nbytes + ): + raise ValueError(f"lease tensor {name!r} differs from the validated host metadata") + if provenance.identity_digest != identity.key: + raise ValueError("lease provenance identity digest differs from the requested identity") + if provenance.artifact_content_sha256 != manifest.artifact_content_sha256: + raise ValueError("lease provenance content digest differs from the validated manifest") + if not isinstance(mapped_regions, tuple) or any( + not isinstance(region, MappedHostRegion) for region in mapped_regions + ): + raise ValueError("lease mapped regions must be an immutable tuple") + if not isinstance(resource_closers, tuple) or any(not callable(closer) for closer in resource_closers): + raise ValueError("lease resource closers must be an immutable tuple of callables") + if not callable(release_artifact_lock): + raise ValueError("lease artifact lock release must be callable") + self.identity = identity + self.manifest = manifest + self._tensor_storage = dict(tensors) + self.tensors = MappingProxyType(self._tensor_storage) + self.mapped_regions = mapped_regions + self.provenance = provenance + self._resource_closers = resource_closers + self._release_artifact_lock = release_artifact_lock + self._state_lock = threading.RLock() + self._closed = False + + @property + def closed(self) -> bool: + with self._state_lock: + return self._closed + + def close(self) -> None: + with self._state_lock: + if self._closed: + return + self._closed = True + + errors: list[BaseException] = [] + self._tensor_storage.clear() + for closer in reversed(self._resource_closers): + try: + closer() + except BaseException as exc: # pragma: no cover - best-effort teardown aggregation + errors.append(exc) + try: + self._release_artifact_lock() + except BaseException as exc: # pragma: no cover - best-effort teardown aggregation + errors.append(exc) + if errors: + raise RuntimeError(f"failed to close {len(errors)} host weight lease resources") from errors[0] + + def __enter__(self) -> HostWeightLease: + if self.closed: + raise RuntimeError("cannot re-enter a closed HostWeightLease") + return self + + def __exit__(self, exc_type: object, exc: object, _traceback: object) -> None: + self.close() + + def __reduce_ex__(self, _protocol: SupportsIndex) -> NoReturn: + raise TypeError("HostWeightLease is process-local and cannot be serialized") + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.close() + + +class HostWeightLeaseCarrier: + """Process-local, single-take ownership carrier for a host-weight lease. + + The loader creates the carrier after a restore transaction has committed. + The runner/backend takes the lease exactly once before starting any + asynchronous transport work. If setup fails before that handoff, closing + the carrier releases the loader-owned lease. + """ + + def __init__(self, lease: HostWeightLease) -> None: + if not isinstance(lease, HostWeightLease): + raise TypeError("HostWeightLeaseCarrier requires one HostWeightLease") + if lease.closed: + raise ValueError("cannot carry a closed HostWeightLease") + self._lease: HostWeightLease | None = lease + self._taken = False + self._lock = threading.RLock() + + @property + def taken(self) -> bool: + with self._lock: + return self._taken + + @property + def closed(self) -> bool: + with self._lock: + return self._lease is None or self._lease.closed + + def take(self) -> HostWeightLease: + """Transfer the lease to the transport owner exactly once.""" + with self._lock: + if self._taken: + raise RuntimeError("HostWeightLeaseCarrier was already taken") + lease = self._lease + if lease is None or lease.closed: + raise RuntimeError("HostWeightLeaseCarrier does not own an open lease") + self._lease = None + self._taken = True + return lease + + def close(self) -> None: + """Close the lease if ownership has not been transferred.""" + with self._lock: + lease = self._lease + self._lease = None + if lease is not None: + lease.close() + + def __reduce_ex__(self, _protocol: SupportsIndex) -> NoReturn: + raise TypeError("HostWeightLeaseCarrier is process-local and cannot be serialized") + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.close() + + +__all__ = ["HostWeightLease", "HostWeightLeaseCarrier", "LeaseProvenance", "MappedHostRegion"] diff --git a/vllm_omni/host_weight_runtime/manifest.py b/vllm_omni/host_weight_runtime/manifest.py new file mode 100644 index 000000000..dfeaf2b95 --- /dev/null +++ b/vllm_omni/host_weight_runtime/manifest.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Versioned immutable manifests for runtime-ready host weights.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from pathlib import PurePosixPath +from typing import cast + +import regex as re + +from .identity import CanonicalJson, JsonValue, WeightArtifactIdentity, canonical_json + +_SHA256_RE = re.compile(r"[0-9a-f]{64}") + + +class ManifestValidationError(ValueError): + """Raised when a published manifest violates its schema.""" + + +class TensorKind(str, Enum): + PARAMETER = "parameter" + BUFFER = "buffer" + + +def _require_text(name: str, value: object) -> str: + if not isinstance(value, str) or not value: + raise ManifestValidationError(f"{name} must be a non-empty string") + return value + + +def _require_int(name: str, value: object) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise ManifestValidationError(f"{name} must be an integer") + return value + + +def _require_sha256(name: str, value: object) -> str: + if not isinstance(value, str) or _SHA256_RE.fullmatch(value) is None: + raise ManifestValidationError(f"{name} must be a lowercase SHA-256 digest") + return value + + +def _require_safe_file_name(value: object) -> str: + if not isinstance(value, str): + raise ManifestValidationError("artifact file name must be a string") + path = PurePosixPath(value) + if not value or path.is_absolute() or len(path.parts) != 1 or path.name != value or value in {".", ".."}: + raise ManifestValidationError(f"artifact file name must be one safe relative basename: {value!r}") + return value + + +def _require_int_tuple(name: str, value: object) -> tuple[int, ...]: + if not isinstance(value, list): + raise ManifestValidationError(f"{name} must be an array") + return tuple(_require_int(f"{name}[{index}]", item) for index, item in enumerate(value)) + + +def _exact_keys(value: dict[str, object], expected: set[str], kind: str) -> None: + keys = set(value) + if keys != expected: + raise ManifestValidationError( + f"{kind} fields differ: missing={sorted(expected - keys)}, unexpected={sorted(keys - expected)}" + ) + + +@dataclass(frozen=True) +class TensorManifestEntry: + name: str + kind: TensorKind + role: str + file_name: str + storage_key: str + shape: tuple[int, ...] + dtype: str + stride: tuple[int, ...] + nbytes: int + sha256: str + + def __post_init__(self) -> None: + _require_text("tensor.name", self.name) + _require_text("tensor.storage_key", self.storage_key) + _require_text("tensor.role", self.role) + _require_text("tensor.dtype", self.dtype) + if not isinstance(self.kind, TensorKind): + raise ManifestValidationError("tensor kind must use TensorKind") + _require_safe_file_name(self.file_name) + if not isinstance(self.shape, tuple) or any( + isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 0 for dimension in self.shape + ): + raise ManifestValidationError(f"tensor {self.name!r} has a negative shape dimension") + if ( + not isinstance(self.stride, tuple) + or len(self.stride) != len(self.shape) + or any(isinstance(value, bool) or not isinstance(value, int) or value < 0 for value in self.stride) + ): + raise ManifestValidationError(f"tensor {self.name!r} has an invalid stride") + if _require_int(f"tensors[{self.name!r}].nbytes", self.nbytes) < 0: + raise ManifestValidationError(f"tensor {self.name!r} has negative nbytes") + _require_sha256(f"tensors[{self.name!r}].sha256", self.sha256) + + def semantic_dict(self) -> dict[str, JsonValue]: + """Return content identity independent of physical file sharding.""" + return { + "name": self.name, + "kind": self.kind.value, + "role": self.role, + "shape": list(self.shape), + "dtype": self.dtype, + "stride": list(self.stride), + "nbytes": self.nbytes, + "sha256": self.sha256, + } + + def to_dict(self) -> dict[str, JsonValue]: + return { + "name": self.name, + "kind": self.kind.value, + "role": self.role, + "file_name": self.file_name, + "storage_key": self.storage_key, + "shape": list(self.shape), + "dtype": self.dtype, + "stride": list(self.stride), + "nbytes": self.nbytes, + "sha256": self.sha256, + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> TensorManifestEntry: + expected = { + "name", + "kind", + "role", + "file_name", + "storage_key", + "shape", + "dtype", + "stride", + "nbytes", + "sha256", + } + _exact_keys(value, expected, "tensor manifest entry") + kind_value = _require_text("tensor.kind", value["kind"]) + try: + kind = TensorKind(kind_value) + except ValueError as exc: + raise ManifestValidationError(f"unknown tensor kind {value['kind']!r}") from exc + return cls( + name=_require_text("tensor.name", value["name"]), + kind=kind, + role=_require_text("tensor.role", value["role"]), + file_name=_require_safe_file_name(value["file_name"]), + storage_key=_require_text("tensor.storage_key", value["storage_key"]), + shape=_require_int_tuple("tensor.shape", value["shape"]), + dtype=_require_text("tensor.dtype", value["dtype"]), + stride=_require_int_tuple("tensor.stride", value["stride"]), + nbytes=_require_int("tensor.nbytes", value["nbytes"]), + sha256=_require_sha256("tensor.sha256", value["sha256"]), + ) + + +@dataclass(frozen=True) +class FileManifestEntry: + file_name: str + size: int + sha256: str + + def __post_init__(self) -> None: + _require_safe_file_name(self.file_name) + if _require_int(f"files[{self.file_name!r}].size", self.size) < 0: + raise ManifestValidationError(f"artifact file {self.file_name!r} has negative size") + _require_sha256(f"files[{self.file_name!r}].sha256", self.sha256) + + def to_dict(self) -> dict[str, JsonValue]: + return {"file_name": self.file_name, "size": self.size, "sha256": self.sha256} + + @classmethod + def from_dict(cls, value: dict[str, object]) -> FileManifestEntry: + _exact_keys(value, {"file_name", "size", "sha256"}, "file manifest entry") + return cls( + file_name=_require_safe_file_name(value["file_name"]), + size=_require_int("file.size", value["size"]), + sha256=_require_sha256("file.sha256", value["sha256"]), + ) + + +@dataclass(frozen=True) +class WeightManifest: + schema_version: int + artifact_key: str + identity: WeightArtifactIdentity + producer_schema: str + restorer_schema: str + tensors: tuple[TensorManifestEntry, ...] + files: tuple[FileManifestEntry, ...] + artifact_content_sha256: str + format_metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + if _require_int("manifest.schema_version", self.schema_version) < 1: + raise ManifestValidationError("manifest schema_version must be positive") + _require_sha256("artifact_key", self.artifact_key) + _require_sha256("artifact_content_sha256", self.artifact_content_sha256) + if not isinstance(self.identity, WeightArtifactIdentity): + raise ManifestValidationError("manifest identity must use WeightArtifactIdentity") + if self.artifact_key != self.identity.key: + raise ManifestValidationError("manifest artifact_key does not match its semantic identity") + _require_text("manifest.producer_schema", self.producer_schema) + _require_text("manifest.restorer_schema", self.restorer_schema) + if ( + not isinstance(self.tensors, tuple) + or not isinstance(self.files, tuple) + or not self.tensors + or not self.files + or any(not isinstance(entry, TensorManifestEntry) for entry in self.tensors) + or any(not isinstance(entry, FileManifestEntry) for entry in self.files) + ): + raise ManifestValidationError("an artifact must contain at least one tensor and one file") + if not isinstance(self.format_metadata, CanonicalJson): + raise ManifestValidationError("manifest format_metadata must use CanonicalJson") + tensor_names = [entry.name for entry in self.tensors] + if tensor_names != sorted(tensor_names) or len(set(tensor_names)) != len(tensor_names): + raise ManifestValidationError("manifest tensors must have unique names in canonical order") + file_names = [entry.file_name for entry in self.files] + if file_names != sorted(file_names) or len(set(file_names)) != len(file_names): + raise ManifestValidationError("manifest files must have unique names in canonical order") + available_files = set(file_names) + if any(entry.file_name not in available_files for entry in self.tensors): + raise ManifestValidationError("tensor manifest references an undeclared artifact file") + storage_bindings = [(entry.file_name, entry.storage_key) for entry in self.tensors] + if len(set(storage_bindings)) != len(storage_bindings): + raise ManifestValidationError("manifest tensor storage bindings must be unique") + + @property + def canonical_bytes(self) -> bytes: + return canonical_json(self.to_dict()) + + @property + def total_tensor_bytes(self) -> int: + return sum(entry.nbytes for entry in self.tensors) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "schema_version": self.schema_version, + "artifact_key": self.artifact_key, + "identity": self.identity.to_dict(), + "producer_schema": self.producer_schema, + "restorer_schema": self.restorer_schema, + "tensors": [entry.to_dict() for entry in self.tensors], + "files": [entry.to_dict() for entry in self.files], + "artifact_content_sha256": self.artifact_content_sha256, + "format_metadata": self.format_metadata.to_value(), + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> WeightManifest: + expected = { + "schema_version", + "artifact_key", + "identity", + "producer_schema", + "restorer_schema", + "tensors", + "files", + "artifact_content_sha256", + "format_metadata", + } + _exact_keys(value, expected, "weight manifest") + identity = value["identity"] + if not isinstance(identity, dict): + raise ManifestValidationError("manifest identity must be an object") + tensors = value["tensors"] + files = value["files"] + if not isinstance(tensors, list) or not all(isinstance(item, dict) for item in tensors): + raise ManifestValidationError("manifest tensors must be an array of objects") + if not isinstance(files, list) or not all(isinstance(item, dict) for item in files): + raise ManifestValidationError("manifest files must be an array of objects") + return cls( + schema_version=_require_int("manifest.schema_version", value["schema_version"]), + artifact_key=_require_sha256("manifest.artifact_key", value["artifact_key"]), + identity=WeightArtifactIdentity.from_dict(cast(dict[str, object], identity)), + producer_schema=_require_text("manifest.producer_schema", value["producer_schema"]), + restorer_schema=_require_text("manifest.restorer_schema", value["restorer_schema"]), + tensors=tuple(TensorManifestEntry.from_dict(cast(dict[str, object], item)) for item in tensors), + files=tuple(FileManifestEntry.from_dict(cast(dict[str, object], item)) for item in files), + artifact_content_sha256=_require_sha256( + "manifest.artifact_content_sha256", value["artifact_content_sha256"] + ), + format_metadata=CanonicalJson.from_value(value["format_metadata"]), + ) + + +@dataclass(frozen=True) +class PublicationMarker: + schema_version: int + publication_id: str + manifest_sha256: str + artifact_content_sha256: str + integrity_mode: str + + def __post_init__(self) -> None: + if _require_int("publication.schema_version", self.schema_version) < 1: + raise ManifestValidationError("publication schema_version must be positive") + _require_text("publication.publication_id", self.publication_id) + _require_text("publication.integrity_mode", self.integrity_mode) + _require_sha256("publication.manifest_sha256", self.manifest_sha256) + _require_sha256("publication.artifact_content_sha256", self.artifact_content_sha256) + + @property + def canonical_bytes(self) -> bytes: + return canonical_json(self.to_dict()) + + def to_dict(self) -> dict[str, JsonValue]: + return { + "schema_version": self.schema_version, + "publication_id": self.publication_id, + "manifest_sha256": self.manifest_sha256, + "artifact_content_sha256": self.artifact_content_sha256, + "integrity_mode": self.integrity_mode, + } + + @classmethod + def from_dict(cls, value: dict[str, object]) -> PublicationMarker: + expected = { + "schema_version", + "publication_id", + "manifest_sha256", + "artifact_content_sha256", + "integrity_mode", + } + _exact_keys(value, expected, "publication marker") + return cls( + schema_version=_require_int("publication.schema_version", value["schema_version"]), + publication_id=_require_text("publication.publication_id", value["publication_id"]), + manifest_sha256=_require_sha256("publication.manifest_sha256", value["manifest_sha256"]), + artifact_content_sha256=_require_sha256( + "publication.artifact_content_sha256", value["artifact_content_sha256"] + ), + integrity_mode=_require_text("publication.integrity_mode", value["integrity_mode"]), + ) + + +__all__ = [ + "FileManifestEntry", + "ManifestValidationError", + "PublicationMarker", + "TensorKind", + "TensorManifestEntry", + "WeightManifest", +] diff --git a/vllm_omni/host_weight_runtime/outcomes.py b/vllm_omni/host_weight_runtime/outcomes.py new file mode 100644 index 000000000..28711f636 --- /dev/null +++ b/vllm_omni/host_weight_runtime/outcomes.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Typed store, resolution, and post-load publication outcomes.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING + +from .errors import HostWeightFailure, ResolutionAction, ResolutionStage + +if TYPE_CHECKING: + from .lease import HostWeightLease + + +class StoreStatus(str, Enum): + HIT = "hit" + MISS = "miss" + BUILT = "built" + JOINED = "joined" + INVALID = "invalid" + TIMEOUT = "timeout" + FAILED = "failed" + + +@dataclass(frozen=True) +class StoreResult: + status: StoreStatus + lease: HostWeightLease | None = None + failure: HostWeightFailure | None = None + + def __post_init__(self) -> None: + has_lease = self.status in {StoreStatus.HIT, StoreStatus.BUILT, StoreStatus.JOINED} + if has_lease != (self.lease is not None): + raise ValueError(f"store result {self.status.value} has inconsistent lease ownership") + if self.status in {StoreStatus.HIT, StoreStatus.MISS, StoreStatus.BUILT, StoreStatus.JOINED}: + if self.failure is not None: + raise ValueError(f"successful store result {self.status.value} must not carry a failure") + if self.status in {StoreStatus.INVALID, StoreStatus.TIMEOUT, StoreStatus.FAILED} and self.failure is None: + raise ValueError(f"store result {self.status.value} requires a typed failure") + + +class AttemptResult(str, Enum): + HIT = "hit" + MISS = "miss" + SUCCESS = "success" + TIMEOUT = "timeout" + ERROR = "error" + SKIPPED = "skipped" + + +class ResolutionOutcome(str, Enum): + LOCAL_HIT = "local_hit" + REMOTE_IMPORT = "remote_import" + LOCAL_PRODUCTION = "local_production" + CANONICAL_DIRECT = "canonical_direct" + CANONICAL_FALLBACK = "canonical_fallback" + FAILED = "failed" + + +class PostLoadPublicationOutcome(str, Enum): + """Terminal outcome for one explicit post-load publication attempt.""" + + RUNTIME_DISABLED = "runtime_disabled" + POLICY_DISABLED = "policy_disabled" + ALREADY_PRESENT = "already_present" + PUBLISHED = "published" + JOINED = "joined" + FAILED = "failed" + + +@dataclass(frozen=True) +class ResolutionAttempt: + stage: ResolutionStage + result: AttemptResult + action: ResolutionAction + elapsed_seconds: float + failure: HostWeightFailure | None = None + + def __post_init__(self) -> None: + if not math.isfinite(self.elapsed_seconds) or self.elapsed_seconds < 0: + raise ValueError("resolution attempt elapsed_seconds must be finite and non-negative") + if self.result in {AttemptResult.HIT, AttemptResult.MISS, AttemptResult.SUCCESS} and self.failure is not None: + raise ValueError(f"resolution attempt {self.result.value} must not carry a failure") + if self.result in {AttemptResult.TIMEOUT, AttemptResult.ERROR} and self.failure is None: + raise ValueError(f"resolution attempt {self.result.value} requires a typed failure") + + +@dataclass(frozen=True) +class ResolutionReport: + resolution_id: str + outcome: ResolutionOutcome + attempts: tuple[ResolutionAttempt, ...] + elapsed_seconds: float + + def __post_init__(self) -> None: + if not self.resolution_id: + raise ValueError("resolution report requires a resolution_id") + if not isinstance(self.attempts, tuple): + raise ValueError("resolution report attempts must be an immutable tuple") + if not math.isfinite(self.elapsed_seconds) or self.elapsed_seconds < 0: + raise ValueError("resolution report elapsed_seconds must be finite and non-negative") + + +@dataclass(frozen=True) +class HostWeightResolution: + report: ResolutionReport + lease: HostWeightLease | None = None + + def __post_init__(self) -> None: + lease_outcomes = { + ResolutionOutcome.LOCAL_HIT, + ResolutionOutcome.REMOTE_IMPORT, + ResolutionOutcome.LOCAL_PRODUCTION, + } + if (self.report.outcome in lease_outcomes) != (self.lease is not None): + raise ValueError("resolution outcome has inconsistent lease ownership") + + +@dataclass(frozen=True) +class PostLoadPublicationReport: + operation_id: str + outcome: PostLoadPublicationOutcome + identity_digest: str | None + elapsed_seconds: float + failure: HostWeightFailure | None = None + + def __post_init__(self) -> None: + if not self.operation_id: + raise ValueError("post-load publication report requires an operation_id") + if not isinstance(self.outcome, PostLoadPublicationOutcome): + raise ValueError("post-load publication outcome must use PostLoadPublicationOutcome") + if self.identity_digest is not None and not self.identity_digest: + raise ValueError("post-load publication identity digest must not be empty") + requires_identity = self.outcome not in { + PostLoadPublicationOutcome.RUNTIME_DISABLED, + PostLoadPublicationOutcome.POLICY_DISABLED, + } + if requires_identity != (self.identity_digest is not None): + raise ValueError("post-load publication identity ownership is inconsistent") + if not math.isfinite(self.elapsed_seconds) or self.elapsed_seconds < 0: + raise ValueError("post-load publication elapsed_seconds must be finite and non-negative") + if (self.outcome is PostLoadPublicationOutcome.FAILED) != (self.failure is not None): + raise ValueError("post-load publication failure ownership is inconsistent") + + +__all__ = [ + "AttemptResult", + "HostWeightResolution", + "PostLoadPublicationOutcome", + "PostLoadPublicationReport", + "ResolutionAttempt", + "ResolutionOutcome", + "ResolutionReport", + "StoreResult", + "StoreStatus", +] diff --git a/vllm_omni/host_weight_runtime/protocols.py b/vllm_omni/host_weight_runtime/protocols.py new file mode 100644 index 000000000..f3738c1f3 --- /dev/null +++ b/vllm_omni/host_weight_runtime/protocols.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Provider and consumer protocols for Host Weight Runtime extensions.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field +from enum import Enum +from types import TracebackType +from typing import TYPE_CHECKING, Protocol + +import torch + +from .identity import CanonicalJson, WeightArtifactIdentity +from .manifest import TensorKind + +if TYPE_CHECKING: + from .config import ValidationLevel + from .lease import HostWeightLease + from .outcomes import StoreResult + + +class ProductionSourceMode(str, Enum): + FINALIZED_MODEL = "finalized_model" + CHECKPOINT_STREAM = "checkpoint_stream" + + +class CoordinationScope(str, Enum): + SINGLE_PROCESS = "single_process" + GROUP_COHORT = "group_cohort" + ARTIFACT_BUNDLE = "artifact_bundle" + + +class LookupPhase(str, Enum): + PRE_LOAD_SAFE = "pre_load_safe" + POST_LOAD_ONLY = "post_load_only" + + +@dataclass(frozen=True) +class TensorWriteSpec: + name: str + shape: tuple[int, ...] + dtype: torch.dtype + kind: TensorKind = TensorKind.PARAMETER + role: str = "weight" + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or not self.name or not isinstance(self.role, str) or not self.role: + raise ValueError("tensor write name and role must not be empty") + if not isinstance(self.shape, tuple) or any( + isinstance(dimension, bool) or not isinstance(dimension, int) or dimension < 0 for dimension in self.shape + ): + raise ValueError(f"tensor {self.name!r} has a negative shape dimension") + if not isinstance(self.dtype, torch.dtype): + raise ValueError(f"tensor {self.name!r} must declare one torch dtype") + if not isinstance(self.kind, TensorKind): + raise ValueError(f"tensor {self.name!r} must use TensorKind") + + @property + def numel(self) -> int: + return math.prod(self.shape) if self.shape else 1 + + @property + def nbytes(self) -> int: + return self.numel * torch.empty((), dtype=self.dtype).element_size() + + +class TensorFileWriter(Protocol): + def write_tensor(self, name: str, tensor: torch.Tensor) -> None: ... + + def write_rows(self, name: str, start_row: int, tensor: torch.Tensor) -> None: ... + + def close(self) -> None: ... + + def __enter__(self) -> TensorFileWriter: ... + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + _traceback: TracebackType | None, + ) -> bool | None: ... + + +class ArtifactWriter(Protocol): + def open_tensor_file( + self, + relative_name: str, + specs: tuple[TensorWriteSpec, ...], + *, + ordered: bool = False, + ) -> TensorFileWriter: ... + + +@dataclass(frozen=True) +class WeightProductionSpec: + producer_id: str + outputs: tuple[WeightArtifactIdentity, ...] + source_mode: ProductionSourceMode + coordination_scope: CoordinationScope = CoordinationScope.SINGLE_PROCESS + lookup_phase: LookupPhase = LookupPhase.PRE_LOAD_SAFE + + def __post_init__(self) -> None: + if not isinstance(self.producer_id, str) or not self.producer_id or not self.outputs: + raise ValueError("producer_id and outputs must not be empty") + if not isinstance(self.outputs, tuple) or any( + not isinstance(output, WeightArtifactIdentity) for output in self.outputs + ): + raise ValueError("producer outputs must be an immutable tuple of artifact identities") + if not isinstance(self.source_mode, ProductionSourceMode): + raise ValueError("producer source mode must use ProductionSourceMode") + if not isinstance(self.coordination_scope, CoordinationScope): + raise ValueError("producer coordination scope must use CoordinationScope") + if not isinstance(self.lookup_phase, LookupPhase): + raise ValueError("producer lookup phase must use LookupPhase") + if len({output.key for output in self.outputs}) != len(self.outputs): + raise ValueError("producer outputs must have unique semantic identities") + if any(output.producer.producer_id != self.producer_id for output in self.outputs): + raise ValueError("producer outputs must name the declaring producer_id") + + +@dataclass(frozen=True) +class ProductionMetadata: + producer_schema: str + restorer_schema: str + format_metadata: CanonicalJson = field(default_factory=CanonicalJson.empty) + + def __post_init__(self) -> None: + if ( + not isinstance(self.producer_schema, str) + or not self.producer_schema + or not isinstance(self.restorer_schema, str) + or not self.restorer_schema + ): + raise ValueError("producer and restorer schema identifiers must not be empty") + if not isinstance(self.format_metadata, CanonicalJson): + raise ValueError("production format metadata must use CanonicalJson") + + +class WeightProducer(Protocol): + """Synchronous V1 producer; callers do not preempt an active build.""" + + @property + def spec(self) -> WeightProductionSpec: ... + + def produce(self, writer: ArtifactWriter) -> ProductionMetadata: + """Produce one exact artifact synchronously or raise.""" + ... + + +@dataclass(frozen=True) +class BuildRequest: + identity: WeightArtifactIdentity + + +class WeightRestorePlan(Protocol): + """A validated, one-shot model mutation plan.""" + + def commit(self) -> None: + """Mutate the target model exactly once or raise.""" + ... + + +class WeightRestorer(Protocol): + """Validate a lease-to-model restore without mutating either input.""" + + @property + def schema(self) -> str: ... + + def plan_restore(self, model: object, lease: HostWeightLease) -> WeightRestorePlan: + """Return a validated plan; all model mutation is deferred to commit().""" + ... + + +class HostWeightStore(Protocol): + """Exact artifact access whose deadlines bound lock acquisition only.""" + + def lookup( + self, + identity: WeightArtifactIdentity, + *, + validation: ValidationLevel, + deadline: float | None = None, + ) -> StoreResult: ... + + def get_or_build( + self, + request: BuildRequest, + producer: WeightProducer, + *, + validation: ValidationLevel, + deadline: float, + ) -> StoreResult: ... + + +@dataclass(frozen=True) +class RemoteArtifactDescriptor: + identity: WeightArtifactIdentity + provider_id: str + manifest_sha256: str + content_sha256: str + + +class RemoteArtifactProvider(Protocol): + @property + def provider_id(self) -> str: ... + + def probe(self, identity: WeightArtifactIdentity) -> RemoteArtifactDescriptor | None: ... + + def fetch(self, descriptor: RemoteArtifactDescriptor, writer: ArtifactWriter) -> ProductionMetadata: ... + + +__all__ = [ + "ArtifactWriter", + "BuildRequest", + "CoordinationScope", + "HostWeightStore", + "LookupPhase", + "ProductionMetadata", + "ProductionSourceMode", + "RemoteArtifactDescriptor", + "RemoteArtifactProvider", + "TensorFileWriter", + "TensorWriteSpec", + "WeightProducer", + "WeightProductionSpec", + "WeightRestorePlan", + "WeightRestorer", +] diff --git a/vllm_omni/host_weight_runtime/runtime.py b/vllm_omni/host_weight_runtime/runtime.py new file mode 100644 index 000000000..fbc726669 --- /dev/null +++ b/vllm_omni/host_weight_runtime/runtime.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Policy-driven exact resolution over HostWeightStore contracts.""" + +from __future__ import annotations + +import logging +import time +import uuid +from collections.abc import Callable + +from .config import HostWeightRuntimeConfig, RemoteOnMiss, RuntimeMode +from .errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionAction, ResolutionStage +from .identity import CanonicalJson, WeightArtifactIdentity +from .lease import HostWeightLease +from .outcomes import ( + AttemptResult, + HostWeightResolution, + PostLoadPublicationOutcome, + PostLoadPublicationReport, + ResolutionAttempt, + ResolutionOutcome, + ResolutionReport, + StoreResult, + StoreStatus, +) +from .protocols import BuildRequest, HostWeightStore, LookupPhase, WeightProducer + +logger = logging.getLogger(__name__) + + +class HostWeightRuntime: + """Resolve or publish one exact representation without owning restoration.""" + + def __init__( + self, + config: HostWeightRuntimeConfig, + store: HostWeightStore | None = None, + *, + observer: Callable[[ResolutionReport], None] | None = None, + publication_observer: Callable[[PostLoadPublicationReport], None] | None = None, + _initialization_failure: HostWeightFailure | None = None, + ) -> None: + if observer is not None and not callable(observer): + raise ValueError("Host Weight Runtime observer must be callable") + if publication_observer is not None and not callable(publication_observer): + raise ValueError("Host Weight Runtime publication observer must be callable") + if config.mode is RuntimeMode.DISABLED and (store is not None or _initialization_failure is not None): + raise ValueError("disabled Host Weight Runtime must not initialize a store") + if config.mode is not RuntimeMode.DISABLED and (store is None) == (_initialization_failure is None): + raise ValueError("enabled Host Weight Runtime requires exactly one store or initialization failure") + self.config = config + self.store = store + self._initialization_failure = _initialization_failure + self._observer = observer + self._publication_observer = publication_observer + + @classmethod + def from_config( + cls, + config: HostWeightRuntimeConfig, + *, + observer: Callable[[ResolutionReport], None] | None = None, + publication_observer: Callable[[PostLoadPublicationReport], None] | None = None, + ) -> HostWeightRuntime: + """Construct lazily so disabled mode performs no filesystem probe.""" + if config.mode is RuntimeMode.DISABLED: + return cls(config, observer=observer, publication_observer=publication_observer) + assert config.domain is not None + from .filesystem import FilesystemHostWeightStore # noqa: PLC0415 + + try: + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity) + except HostWeightError as exc: + return cls( + config, + observer=observer, + publication_observer=publication_observer, + _initialization_failure=exc.failure, + ) + return cls(config, store, observer=observer, publication_observer=publication_observer) + + @property + def enabled(self) -> bool: + return self.config.mode is not RuntimeMode.DISABLED + + def resolve( + self, + identity: WeightArtifactIdentity | None = None, + *, + producer: WeightProducer | None = None, + ) -> HostWeightResolution: + started = time.monotonic() + if not self.enabled: + return self._finish( + ResolutionReport( + resolution_id=uuid.uuid4().hex, + outcome=ResolutionOutcome.CANONICAL_DIRECT, + attempts=(), + elapsed_seconds=time.monotonic() - started, + ) + ) + if identity is None: + raise ValueError("enabled Host Weight Runtime requires an exact artifact identity") + if self._initialization_failure is not None: + failure = self._initialization_failure + action = self._action_for_failure(failure) + return self._finish_terminal( + ( + ResolutionAttempt( + stage=failure.stage, + result=AttemptResult.ERROR, + action=action, + elapsed_seconds=time.monotonic() - started, + failure=failure, + ), + ), + started, + ) + assert self.store is not None + deadline = started + self.config.wait.coordination_timeout_seconds + attempts: list[ResolutionAttempt] = [] + + lookup_started = time.monotonic() + local = self.store.lookup( + identity, + validation=self.config.integrity.local_lookup, + deadline=deadline, + ) + if local.status is StoreStatus.HIT: + assert local.lease is not None + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.LOOKUP, + result=AttemptResult.HIT, + action=ResolutionAction.RETURN_LEASE, + elapsed_seconds=time.monotonic() - lookup_started, + ) + ) + return self._finish( + ResolutionReport( + resolution_id=local.lease.provenance.resolution_id, + outcome=ResolutionOutcome.LOCAL_HIT, + attempts=tuple(attempts), + elapsed_seconds=time.monotonic() - started, + ), + lease=local.lease, + ) + + local_action = self._next_after_local(local, producer) + attempts.append( + self._store_attempt( + ResolutionStage.LOOKUP, + local, + time.monotonic() - lookup_started, + action=local_action, + ) + ) + if local.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return self._finish_terminal(tuple(attempts), started) + + producer_allowed = self._producer_allowed(producer) + if self.config.remote.on_local_miss is not RemoteOnMiss.DISABLED: + remote_failure = HostWeightFailure( + stage=ResolutionStage.REMOTE, + code=FailureCode.UNSUPPORTED, + retryable=False, + message="remote artifact providers are modeled but not implemented", + details=CanonicalJson.empty(), + ) + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.REMOTE, + result=AttemptResult.SKIPPED, + action=self._action_for_failure(remote_failure), + elapsed_seconds=0.0, + failure=remote_failure, + ) + ) + return self._finish_terminal(tuple(attempts), started) + + if producer_allowed: + assert producer is not None + production_started = time.monotonic() + produced = self.store.get_or_build( + BuildRequest(identity), + producer, + validation=self.config.integrity.local_lookup, + deadline=deadline, + ) + if produced.status in {StoreStatus.HIT, StoreStatus.BUILT, StoreStatus.JOINED}: + assert produced.lease is not None + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.PRODUCTION, + result=AttemptResult.SUCCESS, + action=ResolutionAction.RETURN_LEASE, + elapsed_seconds=time.monotonic() - production_started, + ) + ) + return self._finish( + ResolutionReport( + resolution_id=produced.lease.provenance.resolution_id, + outcome=ResolutionOutcome.LOCAL_PRODUCTION, + attempts=tuple(attempts), + elapsed_seconds=time.monotonic() - started, + ), + lease=produced.lease, + ) + attempts.append( + self._store_attempt( + ResolutionStage.PRODUCTION, + produced, + time.monotonic() - production_started, + action=self._terminal_action_for_store_result(produced), + ) + ) + else: + unavailable = HostWeightFailure( + stage=ResolutionStage.PRODUCTION, + code=FailureCode.PRODUCER_UNAVAILABLE, + retryable=True, + message="no allowed producer can build the exact requested representation", + details=CanonicalJson.empty(), + ) + attempts.append( + ResolutionAttempt( + stage=ResolutionStage.PRODUCTION, + result=AttemptResult.SKIPPED, + action=self._action_for_failure(unavailable), + elapsed_seconds=0.0, + failure=unavailable, + ) + ) + + return self._finish_terminal(tuple(attempts), started) + + def publish_after_load( + self, + identity: WeightArtifactIdentity | None = None, + *, + producer: WeightProducer | None = None, + ) -> PostLoadPublicationReport: + """Synchronously warm a future startup without revising resolution. + + A successful store operation returns a validated lease. The runtime + closes that lease before returning because POST_LOAD_ONLY production + must not mutate or rebind the model serving the current startup. + """ + started = time.monotonic() + operation_id = uuid.uuid4().hex + if not self.enabled: + return self._finish_publication( + PostLoadPublicationReport( + operation_id=operation_id, + outcome=PostLoadPublicationOutcome.RUNTIME_DISABLED, + identity_digest=None, + elapsed_seconds=time.monotonic() - started, + ) + ) + if not self.config.production.allow_post_load_publish: + return self._finish_publication( + PostLoadPublicationReport( + operation_id=operation_id, + outcome=PostLoadPublicationOutcome.POLICY_DISABLED, + identity_digest=None, + elapsed_seconds=time.monotonic() - started, + ) + ) + if identity is None: + raise ValueError("enabled post-load publication requires an exact artifact identity") + if producer is None: + raise ValueError("enabled post-load publication requires one producer") + if producer.spec.lookup_phase is not LookupPhase.POST_LOAD_ONLY: + raise ValueError("publish_after_load requires a POST_LOAD_ONLY producer") + + identity_digest = identity.key + if self._initialization_failure is not None: + return self._finish_publication( + PostLoadPublicationReport( + operation_id=operation_id, + outcome=PostLoadPublicationOutcome.FAILED, + identity_digest=identity_digest, + elapsed_seconds=time.monotonic() - started, + failure=self._initialization_failure, + ) + ) + + assert self.store is not None + result = self.store.get_or_build( + BuildRequest(identity), + producer, + validation=self.config.integrity.local_lookup, + deadline=started + self.config.wait.coordination_timeout_seconds, + ) + successful_outcomes = { + StoreStatus.HIT: PostLoadPublicationOutcome.ALREADY_PRESENT, + StoreStatus.BUILT: PostLoadPublicationOutcome.PUBLISHED, + StoreStatus.JOINED: PostLoadPublicationOutcome.JOINED, + } + outcome = successful_outcomes.get(result.status) + if outcome is not None: + assert result.lease is not None + try: + result.lease.close() + except Exception as exc: + return self._finish_publication( + PostLoadPublicationReport( + operation_id=operation_id, + outcome=PostLoadPublicationOutcome.FAILED, + identity_digest=identity_digest, + elapsed_seconds=time.monotonic() - started, + failure=HostWeightFailure( + stage=ResolutionStage.LIFECYCLE, + code=FailureCode.PUBLICATION_FAILED, + retryable=False, + message="post-load publication lease cleanup failed", + details=CanonicalJson.from_value( + { + "exception_type": type(exc).__name__, + "store_outcome": result.status.value, + } + ), + ), + ) + ) + return self._finish_publication( + PostLoadPublicationReport( + operation_id=operation_id, + outcome=outcome, + identity_digest=identity_digest, + elapsed_seconds=time.monotonic() - started, + ) + ) + + failure = result.failure + if failure is None: + failure = HostWeightFailure( + stage=ResolutionStage.PRODUCTION, + code=FailureCode.PRODUCER_FAILED, + retryable=False, + message=f"post-load publication returned unexpected store status {result.status.value}", + details=CanonicalJson.empty(), + ) + return self._finish_publication( + PostLoadPublicationReport( + operation_id=operation_id, + outcome=PostLoadPublicationOutcome.FAILED, + identity_digest=identity_digest, + elapsed_seconds=time.monotonic() - started, + failure=failure, + ) + ) + + def _mode_terminal_action(self) -> ResolutionAction: + if self.config.mode is RuntimeMode.PREFERRED: + return ResolutionAction.CANONICAL_FALLBACK + return ResolutionAction.FAIL_STARTUP + + def _action_for_failure(self, failure: HostWeightFailure | None) -> ResolutionAction: + if failure is None: + raise ValueError("failed store results require a typed failure") + if self.config.mode is RuntimeMode.PREFERRED and failure.retryable: + return ResolutionAction.CANONICAL_FALLBACK + return ResolutionAction.FAIL_STARTUP + + def _terminal_action_for_store_result(self, result: StoreResult) -> ResolutionAction: + if result.status in {StoreStatus.MISS, StoreStatus.INVALID}: + return self._mode_terminal_action() + return self._action_for_failure(result.failure) + + def _next_after_local( + self, + result: StoreResult, + producer: WeightProducer | None, + ) -> ResolutionAction: + if result.status in {StoreStatus.TIMEOUT, StoreStatus.FAILED}: + return self._terminal_action_for_store_result(result) + if self.config.remote.on_local_miss is not RemoteOnMiss.DISABLED: + return ResolutionAction.TRY_REMOTE + if self._producer_allowed(producer): + if result.status is StoreStatus.INVALID: + return ResolutionAction.QUARANTINE_THEN_PRODUCE + return ResolutionAction.TRY_PRODUCER + return self._mode_terminal_action() + + def _producer_allowed(self, producer: WeightProducer | None) -> bool: + return ( + producer is not None + and self.config.production.allow_local_build + and producer.spec.lookup_phase is LookupPhase.PRE_LOAD_SAFE + ) + + @staticmethod + def _store_attempt( + stage: ResolutionStage, + result: StoreResult, + elapsed: float, + *, + action: ResolutionAction, + ) -> ResolutionAttempt: + if result.status is StoreStatus.MISS: + return ResolutionAttempt( + stage=stage, + result=AttemptResult.MISS, + action=action, + elapsed_seconds=elapsed, + ) + if result.status is StoreStatus.TIMEOUT: + attempt_result = AttemptResult.TIMEOUT + elif result.status is StoreStatus.INVALID: + attempt_result = AttemptResult.ERROR + else: + attempt_result = AttemptResult.ERROR + return ResolutionAttempt( + stage=stage, + result=attempt_result, + action=action, + elapsed_seconds=elapsed, + failure=result.failure, + ) + + def _finish_terminal( + self, + attempts: tuple[ResolutionAttempt, ...], + started: float, + ) -> HostWeightResolution: + action = attempts[-1].action + if action is ResolutionAction.CANONICAL_FALLBACK: + outcome = ResolutionOutcome.CANONICAL_FALLBACK + elif action is ResolutionAction.FAIL_STARTUP: + outcome = ResolutionOutcome.FAILED + else: + raise ValueError(f"resolution cannot terminate with action {action.value}") + return self._finish( + ResolutionReport( + resolution_id=uuid.uuid4().hex, + outcome=outcome, + attempts=attempts, + elapsed_seconds=time.monotonic() - started, + ) + ) + + def _finish( + self, + report: ResolutionReport, + *, + lease: HostWeightLease | None = None, + ) -> HostWeightResolution: + if self._observer is not None: + try: + self._observer(report) + except Exception: + logger.exception( + "Host Weight Runtime observer failed", + extra={"resolution_id": report.resolution_id}, + ) + return HostWeightResolution(report=report, lease=lease) + + def _finish_publication( + self, + report: PostLoadPublicationReport, + ) -> PostLoadPublicationReport: + if self._publication_observer is not None: + try: + self._publication_observer(report) + except Exception: + logger.exception( + "Host Weight Runtime publication observer failed", + extra={"operation_id": report.operation_id}, + ) + return report + + +__all__ = ["HostWeightRuntime"] diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index b5d5217f5..68adb44ca 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Single source of truth for vLLM-Omni Prometheus + bench CLI metric naming. Consumed by: @@ -61,6 +64,7 @@ IMAGE_COUNT = "image_count" IMAGE_GENERATION = "image_generation" IMAGE_GENERATION_TIME_MS = f"{IMAGE_GENERATION}_time_ms" IMAGE_PIXELS = "image_pixels" +DIFFUSION_SCHEDULER_WAITING_KEY = "scheduler_num_waiting_reqs" TOTAL_IMAGES = "total_images" IMAGE_THROUGHPUT = "image_throughput" AVERAGE_PIXELS_PER_IMAGE = "average_pixels_per_image" @@ -91,6 +95,7 @@ POSTPROCESS_TIME_MS = f"{POSTPROCESS_TIME}_ms" POSTPROCESS_TIMES_MS = f"{POSTPROCESS_TIME}s_ms" OUTPUT_UNIT_COUNT = "output_unit_count" SERVING_TIME_TO_FIRST_OUTPUT_MS = "serving_time_to_first_output_ms" +IMAGE_TIME_TO_FIRST_OUTPUT_MS = "image_time_to_first_output_ms" SERVING_TIME_TO_FIRST_OUTPUTS_MS = "serving_time_to_first_outputs_ms" TIME_PER_OUTPUT_UNIT_MS = "time_per_output_unit_ms" TIME_PER_OUTPUT_UNITS_MS = "time_per_output_units_ms" @@ -152,6 +157,30 @@ TRANSFER_RX_S = METRIC_PREFIX + "transfer_rx_s" TRANSFER_IN_FLIGHT_S = METRIC_PREFIX + "transfer_in_flight_s" +# ============================================================================ +# Image / diffusion service-level families +# ============================================================================ +STAGE_GEN_TIME_S = METRIC_PREFIX + "stage_gen_time_s" +REQUEST_QUEUE_WAIT_S = METRIC_PREFIX + "request_queue_wait_s" +STAGE_WAITING_REQUESTS = METRIC_PREFIX + "stage_waiting_requests" +NUM_INFERENCE_STEPS = METRIC_PREFIX + "num_inference_steps" +IMAGE_COUNT_METRIC = METRIC_PREFIX + IMAGE_COUNT +IMAGE_PIXELS_METRIC = METRIC_PREFIX + IMAGE_PIXELS +PEAK_MEMORY_MB = METRIC_PREFIX + "peak_memory_mb" +REQUESTS_FAILED = METRIC_PREFIX + "requests_failed" +KV_WAIT_S = METRIC_PREFIX + "kv_wait_s" +DIFFUSION_FORWARD_S = METRIC_PREFIX + "diffusion_forward_s" +DIFFUSION_EXEC_S = METRIC_PREFIX + "diffusion_exec_s" +DIFFUSION_EXEC_PER_STEP_S = METRIC_PREFIX + "diffusion_exec_per_step_s" +DIFFUSION_PREPROCESS_S = METRIC_PREFIX + "diffusion_preprocess_s" +DIFFUSION_POSTPROCESS_S = METRIC_PREFIX + "diffusion_postprocess_s" +VAE_DECODE_S = METRIC_PREFIX + "vae_decode_s" +DENOISE_STEP_LATENCY_S = METRIC_PREFIX + DENOISE_STEP_LATENCY + "_s" +DIFFUSION_KV_LOAD_S = METRIC_PREFIX + "diffusion_kv_load_s" +IMAGE_TTFP_S = METRIC_PREFIX + "image_ttfp_s" +STAGE_IN_QUEUE_S = METRIC_PREFIX + "stage_in_queue_s" + + # ============================================================================ # Label sets # ============================================================================ @@ -163,6 +192,11 @@ SUCCESS_LABELS = ("model_name", "finished_reason") # ``stage`` + ``replica``. STAGE_LABELS = ("model_name", "stage", "replica") +STAGE_GEN_TIME_LABELS = ("model_name", "stage", "stage_type") +DIFFUSION_LABELS = ("model_name", "stage") +FAILED_LABELS = ("model_name", "reason") +KV_WAIT_LABELS = ("model_name", "connector_type") + # Audio continuity Counter carries an extra ``threshold_ms`` label so multiple # threshold buckets can be tracked simultaneously. The ``_ms`` suffix names a # numeric threshold *value* in ms, not a time-bearing metric. @@ -311,16 +345,6 @@ def compute_audio_frames( return pcm_nbytes // frame_width -def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int) -> float: - """Mean denoise step latency = image stage generation time / step count. - - The returned value uses the same time unit as ``stage_gen_time``. - """ - if num_inference_steps <= 0: - return 0.0 - return stage_gen_time / float(num_inference_steps) - - # ============================================================================ # Audio sample-rate resolution # ============================================================================ diff --git a/vllm_omni/metrics/modality.py b/vllm_omni/metrics/modality.py index b83912ed3..e0d51e2b0 100644 --- a/vllm_omni/metrics/modality.py +++ b/vllm_omni/metrics/modality.py @@ -27,7 +27,7 @@ from typing import Any from prometheus_client import Counter, Histogram from vllm_omni.metrics import definitions as defs -from vllm_omni.metrics.utils import count_audio_frames, extract_mm_output +from vllm_omni.metrics.utils import observe_audio_finalize, observe_diffusion_finalize _stage_labels = list(defs.STAGE_LABELS) @@ -105,6 +105,40 @@ _diffusion_postprocess_family = Histogram( labelnames=_stage_labels, buckets=defs.SECONDS_FAST_BUCKETS, ) +_vae_decode_family = Histogram( + defs.VAE_DECODE_S, + "VAE decode latency in seconds (latents -> pixels/audio/video).", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_diffusion_forward_family = Histogram( + defs.DIFFUSION_FORWARD_S, + "Diffusion forward-only latency in seconds (denoise loop; excludes " + "preprocess / postprocess / VAE decode / KV load). Absent when the " + "pipeline profiler is off.", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_diffusion_kv_load_family = Histogram( + defs.DIFFUSION_KV_LOAD_S, + "Diffusion KV-recv latency in seconds (AR→diffusion KV fetch; absent when " + "the stage has no upstream KV to receive).", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) +_image_ttfp_family = Histogram( + defs.IMAGE_TTFP_S, + "Image time-to-first-output in seconds (stage submit → first image materialized; non-streaming single-image). " + "Stage-level, not e2e — excludes API queue and inter-stage transfer before the image stage.", + labelnames=_stage_labels, + buckets=defs.SECONDS_BUCKETS, +) +_denoise_step_latency_family = Histogram( + defs.DENOISE_STEP_LATENCY_S, + "Mean per-step denoise forward latency in seconds (forward_time / num_inference_steps).", + labelnames=_stage_labels, + buckets=defs.SECONDS_FAST_BUCKETS, +) class OmniModalityMetrics: @@ -206,6 +240,31 @@ class OmniModalityMetrics: replica=replica, ).observe(seconds) + def observe_vae_decode(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _vae_decode_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_diffusion_forward(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _diffusion_forward_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_diffusion_kv_load(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _diffusion_kv_load_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_image_ttfp(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds < 0: + return + _image_ttfp_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + + def observe_denoise_step_latency(self, stage: str, replica: str, seconds: float) -> None: + if not self._log_stats or seconds <= 0: + return + _denoise_step_latency_family.labels(model_name=self._model_name, stage=stage, replica=replica).observe(seconds) + def observe_modality_at_finalize( mod_metrics: OmniModalityMetrics, @@ -216,61 +275,34 @@ def observe_modality_at_finalize( stage_metrics: Any, engine_outputs: Any, ) -> None: - """Route audio-path observations for a finalized request. + """Route per-modality observations for a finalized request. Used by ``omni_base._process_single_result`` inside the e2e_done finalize - guard so it fires once per request. Skips text path (covered by upstream - ``vllm:*{stage="thinker", ...}``) and any case where required inputs are - missing — caller should not need to pre-validate. + guard so it fires once per request. Text path falls through — covered by + upstream ``vllm:*{stage="thinker", ...}``. Caller should not need to + pre-validate; missing inputs are silently skipped. audio_ttfp is intentionally NOT observed here; it's emitted by the streaming hook at first-packet time, not at finalize. """ - if replica_id is None or stage_metrics is None or output_type is None: + if replica_id is None or stage_metrics is None: return - stage_label = str(stage_id) - replica_label = str(replica_id) - gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0 - mm_out = extract_mm_output(engine_outputs) + observe_diffusion_finalize( + mod_metrics, + stage_id=stage_id, + replica_id=replica_id, + stage_metrics=stage_metrics, + ) if output_type == "audio": - sample_rate = defs.resolve_audio_sample_rate(mm_out) - n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0) - if n_frames == 0: - n_frames = count_audio_frames(mm_out) - mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames) - duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0 - if duration_s > 0: - mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s) - mod_metrics.observe_audio_rtf( - stage_label, - replica_label, - defs.compute_audio_rtf(gen_time_s, duration_s), - ) - else: - mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") - - dm = getattr(stage_metrics, "diffusion_metrics", None) - if dm: - _key_map = { - "diffusion_engine_exec_time_s": mod_metrics.observe_diffusion_exec, - "preprocess_time_s": mod_metrics.observe_diffusion_preprocess, - "postprocess_time_s": mod_metrics.observe_diffusion_postprocess, - } - for key, observe_fn in _key_map.items(): - val = dm.get(key) - if val is not None: - observe_fn(stage_label, replica_label, float(val)) - - exec_time = dm.get("diffusion_engine_exec_time_s") - num_steps = dm.get("num_inference_steps") - if exec_time is not None and num_steps and num_steps > 0: - mod_metrics.observe_diffusion_exec_per_step( - stage_label, - replica_label, - float(exec_time) / int(num_steps), - ) + observe_audio_finalize( + mod_metrics, + stage_id=stage_id, + replica_id=replica_id, + stage_metrics=stage_metrics, + engine_outputs=engine_outputs, + ) def observe_audio_first_packet( diff --git a/vllm_omni/metrics/prometheus.py b/vllm_omni/metrics/prometheus.py index d372c3d6b..911d21e64 100644 --- a/vllm_omni/metrics/prometheus.py +++ b/vllm_omni/metrics/prometheus.py @@ -38,6 +38,66 @@ _generation_tokens_family = Counter( labelnames=_labelnames, ) +_stage_gen_time_family = Histogram( + defs.STAGE_GEN_TIME_S, + "Per-stage generation time in seconds for diffusion / image stages.", + labelnames=list(defs.STAGE_GEN_TIME_LABELS), + buckets=defs.SECONDS_BUCKETS, +) +_request_queue_wait_family = Histogram( + defs.REQUEST_QUEUE_WAIT_S, + "Per-request queue wait between submit arrival and stage submit, in seconds.", + labelnames=_labelnames, + buckets=defs.SECONDS_BUCKETS, +) +_stage_waiting_requests_family = Gauge( + defs.STAGE_WAITING_REQUESTS, + "Sum of the latest scheduler waiting snapshots across a stage's replicas; " + "diffusion snapshots update when outputs arrive.", + labelnames=list(defs.DIFFUSION_LABELS), +) +_num_inference_steps_family = Histogram( + defs.NUM_INFERENCE_STEPS, + "Diffusion step count distribution per request.", + labelnames=_labelnames, +) +_image_count_family = Counter( + defs.IMAGE_COUNT_METRIC, + "Cumulative image count; per-model rate. Throughput recovered via rate().", + labelnames=_labelnames, +) +_image_pixels_family = Histogram( + defs.IMAGE_PIXELS_METRIC, + "Image pixel count distribution per request.", + labelnames=_labelnames, +) +_peak_memory_family = Gauge( + defs.PEAK_MEMORY_MB, + "Peak GPU memory in MB observed during a stage's generation.", + labelnames=list(defs.DIFFUSION_LABELS), +) +_stage_in_queue_family = Histogram( + defs.STAGE_IN_QUEUE_S, + "Time in seconds from diffusion scheduler enqueue to initial execution admission. " + "Zero is the healthy low-load baseline and is counted.", + labelnames=list(defs.DIFFUSION_LABELS), + buckets=defs.SECONDS_FAST_BUCKETS, +) + +_requests_failed_family = Counter( + defs.REQUESTS_FAILED, + "Total requests by failure reason. Pairs with requests_success_total " + "(abort bucket) to give the full failure surface. Reasons are bounded to " + "client_abort, client_disconnect, stage_error, and unknown.", + labelnames=list(defs.FAILED_LABELS), +) +_kv_wait_s_family = Histogram( + defs.KV_WAIT_S, + "KV cache transfer wait time per request, sliced by connector backend.", + labelnames=list(defs.KV_WAIT_LABELS), + buckets=defs.SECONDS_BUCKETS, +) + class OmniPrometheusMetrics: """Label-bound wrapper around the raw Prometheus metrics. @@ -54,6 +114,10 @@ class OmniPrometheusMetrics: self._e2e_latency = _e2e_latency_family.labels(model_name=model_name) self._prompt_tokens = _prompt_tokens_family.labels(model_name=model_name) self._generation_tokens = _generation_tokens_family.labels(model_name=model_name) + self._request_queue_wait = _request_queue_wait_family.labels(model_name=model_name) + self._num_inference_steps = _num_inference_steps_family.labels(model_name=model_name) + self._image_count = _image_count_family.labels(model_name=model_name) + self._image_pixels = _image_pixels_family.labels(model_name=model_name) def set_running(self, n: int) -> None: if not self._log_stats: @@ -89,13 +153,80 @@ class OmniPrometheusMetrics: def request_failed(self) -> None: if not self._log_stats: return - # Pipeline-level "fail" maps to the upstream FinishReason.ABORT bucket; - # a single counter family now covers both normal stops and aborts. _completion_family.labels( model_name=self._model_name, finished_reason="abort", ).inc() + def observe_stage_gen_time(self, stage: int, stage_type: str, gen_time_s: float) -> None: + if not self._log_stats: + return + _stage_gen_time_family.labels( + model_name=self._model_name, + stage=str(stage), + stage_type=stage_type, + ).observe(max(gen_time_s, 0.0)) + + def observe_stage_in_queue(self, stage: int, in_queue_s: float) -> None: + if not self._log_stats: + return + _stage_in_queue_family.labels( + model_name=self._model_name, + stage=str(stage), + ).observe(max(in_queue_s, 0.0)) + + def observe_queue_wait(self, queue_wait_s: float) -> None: + if not self._log_stats: + return + self._request_queue_wait.observe(max(queue_wait_s, 0.0)) + + def set_stage_waiting_requests(self, stage: int, n_waiting: int) -> None: + if not self._log_stats: + return + _stage_waiting_requests_family.labels( + model_name=self._model_name, + stage=str(stage), + ).set(max(n_waiting, 0)) + + def observe_num_inference_steps(self, n_steps: int) -> None: + if not self._log_stats or n_steps <= 0: + return + self._num_inference_steps.observe(n_steps) + + def inc_image_count(self, n_images: int = 1) -> None: + if not self._log_stats or n_images <= 0: + return + self._image_count.inc(n_images) + + def observe_image_pixels(self, n_pixels: int) -> None: + if not self._log_stats or n_pixels <= 0: + return + self._image_pixels.observe(n_pixels) + + def set_peak_memory(self, stage: int, peak_memory_mb: float) -> None: + if not self._log_stats: + return + _peak_memory_family.labels( + model_name=self._model_name, + stage=str(stage), + ).set(max(peak_memory_mb, 0.0)) + + def inc_requests_failed(self, reason: str) -> None: + if not self._log_stats: + return + _requests_failed_family.labels( + model_name=self._model_name, + reason=reason or "unknown", + ).inc() + + def observe_kv_wait(self, connector_type: str, kv_wait_s: float) -> None: + if not self._log_stats or kv_wait_s < 0: + return + _kv_wait_s_family.labels( + model_name=self._model_name, + connector_type=connector_type or "unknown", + ).observe(kv_wait_s) + class OmniRequestCounter: """Running-request counter written by the orchestrator thread, read by the client thread.""" diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 95198ef6a..9e9191eaa 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import time @@ -42,6 +45,7 @@ class StageRequestStats: stage_id: int | None = None replica_id: int | None = None final_output_type: str | None = None + finish_reason: str | None = None request_id: str | None = None postprocess_time_ms: float = 0.0 diffusion_metrics: dict[str, float] = None @@ -49,11 +53,13 @@ class StageRequestStats: audio_sample_rate: int = 0 audio_duration_s: float = 0.0 image_pixels: int = 0 + num_inference_steps: int = 0 denoise_step_latency_ms: float = 0.0 pipeline_timings: dict[str, float] | None = None output_unit_type: str | None = None output_unit_count: int = 0 serving_time_to_first_output_ms: float = 0.0 + image_time_to_first_output_ms: float = 0.0 time_per_output_unit_ms: float = 0.0 inter_output_latency_ms: float = 0.0 inter_output_latencies_ms: list[float] | None = None @@ -121,6 +127,7 @@ STAGE_EXCLUDE = { "rx_decode_time_ms", "rx_in_flight_time_ms", "final_output_type", + "finish_reason", "pipeline_timings", } TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"} @@ -468,6 +475,7 @@ class OrchestratorAggregator: current = { "stage_id": sid, "final_output_type": evt.final_output_type, + "finish_reason": evt.finish_reason, defs.NUM_TOKENS_IN: int(evt.num_tokens_in), defs.NUM_TOKENS_OUT: int(evt.num_tokens_out), defs.STAGE_GEN_TIME_MS: float(evt.stage_gen_time_ms), @@ -480,6 +488,7 @@ class OrchestratorAggregator: "output_unit_type": evt.output_unit_type, defs.OUTPUT_UNIT_COUNT: int(evt.output_unit_count), defs.SERVING_TIME_TO_FIRST_OUTPUT_MS: float(evt.serving_time_to_first_output_ms), + defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS: float(evt.image_time_to_first_output_ms), defs.TIME_PER_OUTPUT_UNIT_MS: float(evt.time_per_output_unit_ms), defs.INTER_OUTPUT_LATENCY_MS: float(evt.inter_output_latency_ms), defs.INTER_OUTPUT_LATENCIES_MS: list(evt.inter_output_latencies_ms or []), @@ -492,6 +501,8 @@ class OrchestratorAggregator: current[defs.NUM_TOKENS_IN] = int(current.get(defs.NUM_TOKENS_IN, 0)) + int(evt.num_tokens_in) current[defs.NUM_TOKENS_OUT] = int(current.get(defs.NUM_TOKENS_OUT, 0)) + int(evt.num_tokens_out) + if evt.finish_reason is not None: + current["finish_reason"] = evt.finish_reason current[defs.STAGE_GEN_TIME_MS] = float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) + float(evt.stage_gen_time_ms) current[defs.POSTPROCESS_TIME_MS] = float(current.get(defs.POSTPROCESS_TIME_MS, 0.0)) + float( evt.postprocess_time_ms @@ -513,6 +524,11 @@ class OrchestratorAggregator: if current_first_output_ms <= 0 < first_output_ms: current[defs.SERVING_TIME_TO_FIRST_OUTPUT_MS] = first_output_ms + image_first_output_ms = float(evt.image_time_to_first_output_ms) + current_image_first_output_ms = float(current.get(defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS, 0.0)) + if current_image_first_output_ms <= 0 < image_first_output_ms: + current[defs.IMAGE_TIME_TO_FIRST_OUTPUT_MS] = image_first_output_ms + if evt.output_unit_type: current["output_unit_type"] = evt.output_unit_type if evt.final_output_type: @@ -578,11 +594,20 @@ class OrchestratorAggregator: stats.request_id = req_id if final_output_type is not None: stats.final_output_type = final_output_type + diffusion_metrics_by_request = getattr(self, "diffusion_metrics", None) stats.diffusion_metrics = ( - {k: float(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} - if req_id in self.diffusion_metrics + {k: float(v) for k, v in diffusion_metrics_by_request.pop(req_id, {}).items()} + if diffusion_metrics_by_request is not None and req_id in diffusion_metrics_by_request else None ) + forward_time_s = (stats.diffusion_metrics or {}).get("forward_time_s") + num_inference_steps = int(getattr(stats, "num_inference_steps", 0) or 0) + if num_inference_steps > 0 and getattr(stats, "output_unit_type", None) == "image": + if forward_time_s is not None: + denoise_time_ms = float(forward_time_s) * 1000.0 + else: + denoise_time_ms = float(getattr(stats, "stage_gen_time_ms", 0.0) or 0.0) + stats.denoise_step_latency_ms = denoise_time_ms / num_inference_steps return stats def on_stage_metrics( @@ -628,10 +653,24 @@ class OrchestratorAggregator: _postproc_ms = (time.perf_counter() - _t0) * 1000.0 self.record_stage_postprocess_time(stage_id, req_id, _postproc_ms) + _MS_TO_S: dict[str, str] = { + "preprocess_time_ms": "preprocess_time_s", + "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s", + "postprocess_time_ms": "postprocess_time_s", + "vae_decode_time_ms": "vae_decode_time_s", + "forward_time_ms": "forward_time_s", + "scheduler_queue_wait_ms": "scheduler_queue_wait_s", + "kv_recv_time_ms": "kv_recv_time_s", + } + def accumulate_diffusion_metrics(self, stage_type: str, req_id: Any, engine_outputs: Any) -> None: """Accumulate diffusion metrics for a request. - Handles extraction and accumulation of diffusion stage metrics. + Engine emits ``*_ms`` timings; the accumulator converts them to ``_s`` + keys via ``_MS_TO_S`` so downstream observers read a uniform + seconds-bearing dict. Per-chunk timing keys are summed; non-timing + keys (e.g. ``image_num`` / ``resolution`` from ``format_diffusion_outputs``) + preserve the existing ``+=`` semantics. Args: req_id: Request ID @@ -643,20 +682,16 @@ class OrchestratorAggregator: diffusion_metrics: dict = getattr(engine_output, "metrics", {}) if isinstance(diffusion_metrics, list): diffusion_metrics = diffusion_metrics[0] - if diffusion_metrics: - _MS_TO_S = { - "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s", - "preprocess_time_ms": "preprocess_time_s", - "postprocess_time_ms": "postprocess_time_s", - "diffusion_engine_total_time_ms": "diffusion_engine_total_time_s", - } - for key, value in diffusion_metrics.items(): - if value is None: - continue - if key in _MS_TO_S: - self.diffusion_metrics[req_id][_MS_TO_S[key]] = float(value) / 1000.0 - else: - self.diffusion_metrics[req_id][key] = value + if not diffusion_metrics: + return + bucket = self.diffusion_metrics[req_id] + for key, value in diffusion_metrics.items(): + if value is None: + continue + if key in self._MS_TO_S: + bucket[self._MS_TO_S[key]] += float(value) / 1000.0 + else: + bucket[key] += float(value) def on_forward( self, diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py index ef8f69ed1..10cdab94a 100644 --- a/vllm_omni/metrics/utils.py +++ b/vllm_omni/metrics/utils.py @@ -1,9 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + from collections.abc import Callable, Mapping, Sequence from dataclasses import fields -from typing import Any +from typing import TYPE_CHECKING, Any from prettytable import PrettyTable +from vllm_omni.metrics import definitions as metric_defs + +if TYPE_CHECKING: + from vllm_omni.metrics.modality import OmniModalityMetrics + from vllm_omni.metrics.prometheus import OmniPrometheusMetrics + from vllm_omni.metrics.stats import StageRequestStats + +FAILURE_REASONS = frozenset({"client_abort", "client_disconnect", "stage_error", "unknown"}) +DIFFUSION_METRICS_ONLY_REQUEST_ID = "__vllm_omni_diffusion_metrics__" + + +def extract_queue_wait_s(pipeline_timings: Mapping[str, float] | None) -> float | None: + if pipeline_timings is None or "queue_wait_ms" not in pipeline_timings: + return None + return float(pipeline_timings["queue_wait_ms"] or 0.0) / 1000.0 + + +def observe_stage_workload_metrics( + prom_metrics: OmniPrometheusMetrics, + *, + stage_type: str, + stage_metrics: StageRequestStats, +) -> None: + if stage_type == "diffusion": + prom_metrics.observe_num_inference_steps(stage_metrics.num_inference_steps) + + if stage_metrics.output_unit_type == "image": + prom_metrics.observe_image_pixels(stage_metrics.image_pixels) + prom_metrics.inc_image_count(stage_metrics.output_unit_count) + + +def observe_audio_finalize( + mod_metrics: OmniModalityMetrics, + *, + stage_id: int, + replica_id: int, + stage_metrics: Any, + engine_outputs: Any, +) -> None: + stage_label = str(stage_id) + replica_label = str(replica_id) + gen_time_s = float(getattr(stage_metrics, "stage_gen_time_ms", 0.0)) / 1000.0 + mm_out = extract_mm_output(engine_outputs) + + sample_rate = metric_defs.resolve_audio_sample_rate(mm_out) + n_frames = int(getattr(stage_metrics, "audio_generated_frames", 0) or 0) + if n_frames == 0: + n_frames = count_audio_frames(mm_out) + mod_metrics.inc_audio_frames(stage_label, replica_label, n_frames) + duration_s = n_frames / sample_rate if sample_rate > 0 else 0.0 + if duration_s > 0: + mod_metrics.observe_audio_duration(stage_label, replica_label, duration_s) + mod_metrics.observe_audio_rtf( + stage_label, + replica_label, + metric_defs.compute_audio_rtf(gen_time_s, duration_s), + ) + else: + mod_metrics.inc_audio_skipped(stage_label, replica_label, "no_audio_data") + + +def observe_diffusion_finalize( + mod_metrics: OmniModalityMetrics, + *, + stage_id: int, + replica_id: int, + stage_metrics: Any, +) -> None: + diff_metrics = getattr(stage_metrics, "diffusion_metrics", None) + if not diff_metrics: + return + + stage_label = str(stage_id) + replica_label = str(replica_id) + + exec_s = diff_metrics.get("diffusion_engine_exec_time_s") + if exec_s is not None: + mod_metrics.observe_diffusion_exec(stage_label, replica_label, float(exec_s)) + num_steps = int( + getattr(stage_metrics, "num_inference_steps", 0) or diff_metrics.get("num_inference_steps") or 0 + ) + if num_steps > 0: + mod_metrics.observe_diffusion_exec_per_step(stage_label, replica_label, float(exec_s) / num_steps) + + pre_s = diff_metrics.get("preprocess_time_s") + if pre_s is not None: + mod_metrics.observe_diffusion_preprocess(stage_label, replica_label, float(pre_s)) + + post_s = diff_metrics.get("postprocess_time_s") + if post_s is not None: + mod_metrics.observe_diffusion_postprocess(stage_label, replica_label, float(post_s)) + + vae_s = diff_metrics.get("vae_decode_time_s") + if vae_s is not None: + mod_metrics.observe_vae_decode(stage_label, replica_label, float(vae_s)) + + forward_s = diff_metrics.get("forward_time_s") + if forward_s is not None: + mod_metrics.observe_diffusion_forward(stage_label, replica_label, float(forward_s)) + num_steps = int( + getattr(stage_metrics, "num_inference_steps", 0) or diff_metrics.get("num_inference_steps") or 0 + ) + if num_steps > 0 and getattr(stage_metrics, "output_unit_type", None) == "image": + mod_metrics.observe_denoise_step_latency(stage_label, replica_label, float(forward_s) / num_steps) + + kv_load_s = diff_metrics.get("kv_recv_time_s") + if kv_load_s is not None: + mod_metrics.observe_diffusion_kv_load(stage_label, replica_label, float(kv_load_s)) + + +def normalize_failure_reason(reason: str | None) -> str: + """Map request failure details to the bounded metrics taxonomy.""" + return reason if reason in FAILURE_REASONS else "unknown" + + +def diffusion_scheduler_waiting_metrics(n_waiting: int) -> dict[str, int]: + """Build the diffusion scheduler snapshot consumed by the orchestrator.""" + return {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: max(int(n_waiting), 0)} + + +def diffusion_exception_metrics(exc: BaseException) -> dict[str, Any]: + """Return metrics attached to a terminal diffusion error.""" + metrics = getattr(exc, "diffusion_metrics", None) + return dict(metrics) if isinstance(metrics, dict) else {} + + +def sum_diffusion_stage_durations_ms(output: Any, suffix: str) -> float | None: + """Sum matching diffusion stage durations in milliseconds.""" + durations = getattr(output, "stage_durations", None) + if not isinstance(durations, dict): + return None + + total = 0.0 + found = False + for key, value in durations.items(): + if not key.endswith(suffix): + continue + try: + total += float(value) + found = True + except (TypeError, ValueError): + continue + return total * 1000.0 if found else None + + +def extract_diffusion_vae_decode_ms(output: Any) -> float | None: + return sum_diffusion_stage_durations_ms(output, ".vae.decode") + + +def extract_diffusion_denoise_ms(output: Any) -> float | None: + """Extract denoise-loop time (transformer and per-step scheduler).""" + return sum_diffusion_stage_durations_ms(output, ".diffuse") + def coerce_positive_int_scalar(value: object) -> int | None: """Coerce a value to a positive int without importing tensor libs. diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py index c5225ec02..59c2cd860 100644 --- a/vllm_omni/model_executor/model_loader/weight_utils.py +++ b/vllm_omni/model_executor/model_loader/weight_utils.py @@ -1,3 +1,4 @@ +import os import time from pathlib import Path @@ -85,3 +86,46 @@ def download_weights_from_hf_specific( time_taken, ) return hf_folder + + +def resolve_model_to_local_path( + model: str, + *, + allow_download: bool = False, + allow_patterns: list[str] | None = None, + cache_dir: str | None = None, +) -> str: + """Resolve a model reference to a local directory. + Failure is always raised: whether a failure is fatal is the caller's decision. + + Args: + model: A local directory or an HF Hub repo ID. + allow_download: Fetch missing files instead of requiring a warm cache. + Downloads go through :func:`download_weights_from_hf_specific` + allow_patterns: Restrict the fetch to these glob patterns. + Defaults to the whole repository. Ignored when model is already a directory. + cache_dir: Download cache location; None uses the HF default. + + Returns: + Path to a local directory containing the model files. + """ + if os.path.isdir(model): + return model + + if allow_download: + return download_weights_from_hf_specific( + model_name_or_path=model, + cache_dir=cache_dir, + allow_patterns=allow_patterns or ["*"], + require_all=True, + ) + + # Cache-only lookups stay on huggingface_hub even under VLLM_USE_MODELSCOPE: + # ModelScope's snapshot_download has no local_files_only equivalent, + # so a ModelScope deployment that needs resolution here must pass allow_download. + return huggingface_hub.snapshot_download( + model, + cache_dir=cache_dir, + allow_patterns=allow_patterns, + local_files_only=True, + ) diff --git a/vllm_omni/model_executor/models/audex/cfg.py b/vllm_omni/model_executor/models/audex/cfg.py index 488097216..3c1837140 100644 --- a/vllm_omni/model_executor/models/audex/cfg.py +++ b/vllm_omni/model_executor/models/audex/cfg.py @@ -3,8 +3,7 @@ # # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0): # inference_scripts_vllm/audiogen_scripts/cfg_logits_processor.py -# inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py -# Divergences from the official files are commented at the divergence site. +# Divergences from the official file are commented at the divergence site. """Classifier-free guidance for the Audex thinker stage. CFG pairs a conditional request with an unconditional (null-prompt) @@ -20,16 +19,13 @@ Requests opt in via ``SamplingParams.extra_args``: uncond: {"cfg_scale": 1.5, "cfg_role": "uncond", "cfg_pair_id": } Blending is only correct when both pair members decode the same position in -the same engine step, so ``apply_cfg_patches()`` makes the vLLM scheduler -pair-aware: a lone member waits for its partner, partners stay adjacent, -their ``num_computed_tokens`` are equalized after every schedule, and they -finish together. The patch is a no-op for engines that never see a CFG -request (the pair registry stays empty). +the same engine step. Keeping the pair step-locked is model-neutral and lives +in :mod:`vllm_omni.model_executor.models.common.cfg_pairing`; the scheduler +patches are re-exported here so existing imports keep working. """ from __future__ import annotations -from collections import deque from typing import Any import torch @@ -42,18 +38,20 @@ from vllm.v1.sample.logits_processor import ( MoveDirectionality, ) +from vllm_omni.model_executor.models.common.cfg_pairing import ( + _COND, + _UNCOND, + apply_cfg_patches, + cfg_patches_applied, +) + logger = init_logger(__name__) -_COND = "cond" -_UNCOND = "uncond" - - -def _cfg_extra(request: Any, key: str) -> Any: - sampling_params = getattr(request, "sampling_params", None) - extra_args = getattr(sampling_params, "extra_args", None) - if extra_args: - return extra_args.get(key) - return None +__all__ = [ + "AudexCFGLogitsProcessor", + "apply_cfg_patches", + "cfg_patches_applied", +] class AudexCFGLogitsProcessor(LogitsProcessor): @@ -208,360 +206,3 @@ class AudexCFGLogitsProcessor(LogitsProcessor): logits[uncond_idx] = blended return logits - - -# --------------------------------------------------------------------------- -# Scheduler pair synchronization -# --------------------------------------------------------------------------- - - -def _wait_queues(scheduler: Any) -> list[Any]: - # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks - # structurally not-ready requests in a second ``skipped_waiting`` queue; - # pair-hold must cover both queues or a lone member parked there would - # be admitted without its partner. - queues = [scheduler.waiting] - skipped = getattr(scheduler, "skipped_waiting", None) - if skipped is not None: - queues.append(skipped) - return queues - - -def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None = None) -> bool: - pair_id = scheduler._cfg_req_to_pair.get(request_id) - if pair_id is None: - return True - roles = scheduler._cfg_pairs.get(pair_id, {}) - if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()): - return False - if blocked_ids: - # A partner parked in skipped_waiting (e.g. waiting for remote KVs) - # may fail promotion this step while this member gets scheduled from - # waiting — hold this member until the partner is schedulable. - for rid in roles.values(): - if rid != request_id and rid in blocked_ids: - return False - return True - - -# Schedule steps a lone pair member may wait for its partner. Partners arrive -# within a handful of steps in practice; exhausting this budget means the -# companion was never created (e.g. prompt expansion failed), and the request -# is released to decode unguided instead of hanging forever. -_MAX_PAIR_HOLD_STEPS = 512 - - -def _drop_stale_pair(scheduler: Any, request_id: str) -> None: - pair_id = scheduler._cfg_req_to_pair.get(request_id) - if pair_id is None: - return - for rid in scheduler._cfg_pairs.pop(pair_id, {}).values(): - scheduler._cfg_req_to_pair.pop(rid, None) - logger.warning( - "CFG pair %s: partner of %s never arrived after %d schedule steps; releasing the request unguided", - pair_id, - request_id, - _MAX_PAIR_HOLD_STEPS, - ) - - -def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]: - """Pull CFG requests whose partner is not yet admitted out of the wait queues. - - The engine schedules as soon as the first member of a pair arrives over - IPC; without the hold, the lone member prefills one step early and the - pair is permanently offset by one token. Held requests are re-prepended - after the schedule step. - """ - hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {} - scheduler._cfg_hold_counts = hold_counts - - skipped = getattr(scheduler, "skipped_waiting", None) - blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set() - - held: list[tuple[Any, Any]] = [] - for queue in _wait_queues(scheduler): - # A member is only "blocked" for its partner's sake when it sits in - # the OTHER queue; members of the queue being scanned move together. - partner_blockers = blocked_ids if queue is not skipped else set() - held_here = [] - for request in list(queue): - if _pair_complete(scheduler, request.request_id, partner_blockers): - hold_counts.pop(request.request_id, None) - continue - count = hold_counts.get(request.request_id, 0) + 1 - if count > _MAX_PAIR_HOLD_STEPS: - hold_counts.pop(request.request_id, None) - _drop_stale_pair(scheduler, request.request_id) - continue - hold_counts[request.request_id] = count - held_here.append(request) - if held_here: - queue.remove_requests(held_here) - held.extend((queue, req) for req in held_here) - return held - - -def _release_held(held: list[tuple[Any, Any]]) -> None: - for queue, request in reversed(held): - queue.prepend_request(request) - - -def _reorder_waiting_for_cfg(scheduler: Any) -> None: - """Move CFG pair partners adjacent in the FCFS waiting queue.""" - waiting = scheduler.waiting - # Only the FCFS queue (a deque) has caller-controlled ordering; priority - # queues order themselves and pair sync then relies on hold + equalize. - if not isinstance(waiting, deque) or len(waiting) < 2: - return - - requests = list(waiting) - waiting.clear() - - seen: set[str] = set() - result: list[Any] = [] - for request in requests: - rid = request.request_id - if rid in seen: - continue - seen.add(rid) - result.append(request) - - pair_id = scheduler._cfg_req_to_pair.get(rid) - if pair_id is None: - continue - for partner_id in scheduler._cfg_pairs.get(pair_id, {}).values(): - if partner_id != rid and partner_id not in seen: - for candidate in requests: - if candidate.request_id == partner_id: - seen.add(partner_id) - result.append(candidate) - break - - waiting.extend(result) - - -def _equalize_cfg_pair_progress(scheduler: Any, scheduler_output: Any) -> None: - """Bring both members of every CFG pair to the same ``num_computed_tokens``. - - Prefix-cache hits or unequal chunked-prefill budget can leave one member - ahead after a schedule step; reduce the faster member's allocation so - both land on the same position. Over-allocated KV blocks remain and are - consumed next step — nothing leaks. - """ - for roles in scheduler._cfg_pairs.values(): - cond_id = roles.get(_COND) - uncond_id = roles.get(_UNCOND) - if not cond_id or not uncond_id: - continue - - cond_sched = scheduler_output.num_scheduled_tokens.get(cond_id, 0) - uncond_sched = scheduler_output.num_scheduled_tokens.get(uncond_id, 0) - if cond_sched == 0 or uncond_sched == 0: - continue - - cond_req = scheduler.requests.get(cond_id) - uncond_req = scheduler.requests.get(uncond_id) - if not cond_req or not uncond_req: - continue - - if cond_req.num_computed_tokens == uncond_req.num_computed_tokens: - continue - - target = min(cond_req.num_computed_tokens, uncond_req.num_computed_tokens) - - feasible = all( - req.num_computed_tokens - target < sched - for req, sched in ((cond_req, cond_sched), (uncond_req, uncond_sched)) - ) - if not feasible: - continue - - for req_id, req, orig_sched in ( - (cond_id, cond_req, cond_sched), - (uncond_id, uncond_req, uncond_sched), - ): - diff = req.num_computed_tokens - target - if diff > 0: - req.num_computed_tokens = target - scheduler_output.num_scheduled_tokens[req_id] = orig_sched - diff - scheduler_output.total_num_scheduled_tokens -= diff - logger.debug( - "CFG equalize %s: %d -> %d scheduled, computed -> %d", - req_id, - orig_sched, - orig_sched - diff, - target, - ) - - -def _drop_split_pairs(scheduler: Any) -> None: - """Deregister a desynced CFG pair (one member preempted / left behind). - - Manual re-preemption surgery is unsafe against the upstream scheduler. - Instead, a pair that split anyway (KV-pressure preemption of exactly one - member, or unrecoverable progress skew) is dropped from the registry: - the cond request continues UNGUIDED (its logits are never blended again) - rather than being re-paired against a desynced partner, and the - companion's discarded output no longer matters. - """ - if not scheduler._cfg_pairs: - return - - running_ids = {req.request_id for req in scheduler.running} - to_drop: list[tuple[str, str, str, str]] = [] - for pair_id, roles in scheduler._cfg_pairs.items(): - cond_id = roles.get(_COND) - uncond_id = roles.get(_UNCOND) - if cond_id is None or uncond_id is None: - continue - if cond_id not in scheduler.requests or uncond_id not in scheduler.requests: - continue - - cond_running = cond_id in running_ids - uncond_running = uncond_id in running_ids - if cond_running != uncond_running: - # Split only counts once decoding started: a lone member that has - # computed tokens while its partner was pushed back is desynced. - runner = scheduler.requests.get(cond_id if cond_running else uncond_id) - if runner is not None and runner.num_computed_tokens > 0: - to_drop.append((pair_id, cond_id, uncond_id, "one member preempted")) - continue - - if cond_running and uncond_running: - cond_req = scheduler.requests.get(cond_id) - uncond_req = scheduler.requests.get(uncond_id) - if cond_req and uncond_req and cond_req.num_computed_tokens != uncond_req.num_computed_tokens: - logger.warning( - "CFG pair progress mismatch (equalize will retry): %s@%d vs %s@%d", - cond_id, - cond_req.num_computed_tokens, - uncond_id, - uncond_req.num_computed_tokens, - ) - - for pair_id, cond_id, uncond_id, reason in to_drop: - scheduler._cfg_pairs.pop(pair_id, None) - scheduler._cfg_req_to_pair.pop(cond_id, None) - scheduler._cfg_req_to_pair.pop(uncond_id, None) - logger.warning( - "CFG pair %s split (%s); releasing cond=%s unguided, uncond=%s output is discarded", - pair_id, - reason, - cond_id, - uncond_id, - ) - - -_patches_applied = False - - -def cfg_patches_applied() -> bool: - """True once :func:`apply_cfg_patches` has patched the scheduler in this process.""" - return _patches_applied - - -def _patch_scheduler() -> None: - from vllm.v1.core.sched.scheduler import Scheduler - - orig_init = Scheduler.__init__ - orig_add = Scheduler.add_request - orig_finish = Scheduler.finish_requests - orig_schedule = Scheduler.schedule - - def _init(self, *args: Any, **kwargs: Any) -> None: - orig_init(self, *args, **kwargs) - self._cfg_pairs: dict[str, dict[str, str]] = {} - self._cfg_req_to_pair: dict[str, str] = {} - - def _add(self, request: Any) -> None: - orig_add(self, request) - pair_id = _cfg_extra(request, "cfg_pair_id") - role = _cfg_extra(request, "cfg_role") - if pair_id and role: - self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id - self._cfg_req_to_pair[request.request_id] = pair_id - # An asymmetric prefix-cache hit (the cond prompt shares cached - # branches; the null prompt does not) can advance one - # member past what post-schedule equalization can claw back. - # Skip prefix-cache reads for CFG requests entirely. - if hasattr(request, "skip_reading_prefix_cache"): - request.skip_reading_prefix_cache = True - - def _finish(self, request_ids: Any, finished_status: Any) -> Any: - if request_ids is None: - result = orig_finish(self, request_ids, finished_status) - self._cfg_pairs.clear() - self._cfg_req_to_pair.clear() - return result - - if isinstance(request_ids, str): - request_ids = {request_ids} - else: - request_ids = set(request_ids) - - # Finish both members together: a surviving partner would decode - # unpaired garbage (its logits are never blended again). - partner_ids: set[str] = set() - for req_id in request_ids: - pair_id = self._cfg_req_to_pair.get(req_id) - if pair_id is None: - continue - for rid in self._cfg_pairs.get(pair_id, {}).values(): - if rid != req_id and rid in self.requests: - partner_ids.add(rid) - - all_ids = request_ids | partner_ids - result = orig_finish(self, all_ids, finished_status) - - for req_id in all_ids: - pair_id = self._cfg_req_to_pair.pop(req_id, None) - if pair_id: - self._cfg_pairs.pop(pair_id, None) - return result - - def _schedule(self, *args: Any, **kwargs: Any) -> Any: - if not self._cfg_pairs: - return orig_schedule(self, *args, **kwargs) - - _reorder_waiting_for_cfg(self) - held = _hold_incomplete_pairs(self) - - # Chunked prefill can split one pair member's prompt across steps - # while the other fits in a single step; halving the long-prefill - # threshold makes both chunk, and equalize aligns the remainder. - scheduler_config = getattr(self, "scheduler_config", None) - orig_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None) - if scheduler_config is not None and orig_threshold is not None: - scheduler_config.long_prefill_token_threshold = self.max_num_scheduled_tokens // 2 - try: - scheduler_output = orig_schedule(self, *args, **kwargs) - finally: - if scheduler_config is not None and orig_threshold is not None: - scheduler_config.long_prefill_token_threshold = orig_threshold - _release_held(held) - - _equalize_cfg_pair_progress(self, scheduler_output) - _drop_split_pairs(self) - return scheduler_output - - _schedule._audex_cfg_patched = True # type: ignore[attr-defined] - - Scheduler.__init__ = _init - Scheduler.add_request = _add - Scheduler.finish_requests = _finish - Scheduler.schedule = _schedule - - -def apply_cfg_patches() -> None: - """Make the vLLM v1 scheduler CFG-pair-aware. Idempotent, per process. - - Must run in the engine-core process before ``Scheduler`` is constructed - (the ``__init__`` wrapper installs the pair registry). - """ - global _patches_applied - if _patches_applied: - return - _patches_applied = True - logger.info("Applying Audex CFG scheduler patches to vLLM v1") - _patch_scheduler() diff --git a/vllm_omni/model_executor/models/audex/pipeline.py b/vllm_omni/model_executor/models/audex/pipeline.py index 8723a2138..a7a4d7960 100644 --- a/vllm_omni/model_executor/models/audex/pipeline.py +++ b/vllm_omni/model_executor/models/audex/pipeline.py @@ -66,6 +66,7 @@ AUDEX_TTS_PIPELINE = PipelineConfig( tokenizer_subdir="checkpoint_folder_audiogen", sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only", sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) @@ -127,6 +128,7 @@ AUDEX_S2S_PIPELINE = PipelineConfig( tokenizer_subdir="checkpoint_folder_audiogen", sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only", sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) @@ -190,6 +192,7 @@ AUDEX_TTA_PIPELINE = PipelineConfig( # payload ships via thinker2xcodec_full_payload. sync_process_input_func=f"{_PROC}.thinker2code2wav_token_only", sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py index 14aa5ff30..de163778d 100644 --- a/vllm_omni/model_executor/models/aura_omni/pipeline.py +++ b/vllm_omni/model_executor/models/aura_omni/pipeline.py @@ -80,6 +80,7 @@ AURA_OMNI_PIPELINE = PipelineConfig( sync_process_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_token_only", sampling_constraints={"detokenize": True}, extras={"tts_args": {"max_instructions_length": 500}}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/common/cfg_pairing.py b/vllm_omni/model_executor/models/common/cfg_pairing.py new file mode 100644 index 000000000..ae4b3ac38 --- /dev/null +++ b/vllm_omni/model_executor/models/common/cfg_pairing.py @@ -0,0 +1,460 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0): +# inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py +# Divergences from the official file are commented at the divergence site. +"""Model-neutral classifier-free-guidance request pairing for the vLLM v1 scheduler. + +A guided request decodes together with an unconditional companion so their +logits can be blended every step. Blending is only correct when both members +decode the same position in the same engine step, so :func:`apply_cfg_patches` +makes the scheduler pair-aware: a lone member waits for its partner, partners +stay adjacent in the waiting queue, their ``num_computed_tokens`` are equalized +after every schedule, and they finish together. + +Requests opt in through ``SamplingParams.extra_args``:: + + cond: {"cfg_role": "cond", "cfg_pair_id": } + uncond: {"cfg_role": "uncond", "cfg_pair_id": } + +``cfg_pair_id`` is optional. A model whose serving layer cannot mint a pair id +(the companion request id is only known inside the engine) may send ``cfg_role`` +alone; the pair id is then derived from the request id, stripping the +companion's registered suffix (see :func:`register_cfg_uncond_suffix`). + +Logit blending itself is the caller's concern: this module only keeps the two +rows step-locked. The patch is a no-op for engines that never see a CFG request +(the pair registry stays empty). +""" + +from __future__ import annotations + +from collections import deque +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +__all__ = [ + "apply_cfg_patches", + "cfg_patches_applied", + "cfg_uncond_suffixes", + "register_cfg_uncond_suffix", +] + +_COND = "cond" +_UNCOND = "uncond" + +# Request-id suffix appended to the companion request by the Audex stage input +# processor (``CFG_UNCOND_SUFFIX`` in +# vllm_omni/model_executor/stage_input_processors/audex.py). Models that mint +# the companion request id differently register their own suffix. +_UNCOND_REQUEST_ID_SUFFIXES: set[str] = {"__cfg_uncond"} + + +def register_cfg_uncond_suffix(suffix: str) -> None: + """Register a companion request-id suffix used to derive a missing pair id.""" + if not suffix or not isinstance(suffix, str): + raise ValueError(f"CFG uncond request-id suffix must be a non-empty string, got {suffix!r}") + _UNCOND_REQUEST_ID_SUFFIXES.add(suffix) + + +def cfg_uncond_suffixes() -> frozenset[str]: + """Companion request-id suffixes registered in this process.""" + return frozenset(_UNCOND_REQUEST_ID_SUFFIXES) + + +def _derive_pair_id(request_id: str, role: str) -> str: + """Pair id for a request that carries ``cfg_role`` but no ``cfg_pair_id``. + + The cond request id IS the pair id; its companion carries the same id plus + a registered suffix, so stripping the suffix recovers it. Longest suffix + first, so a registered suffix extending another still resolves. + """ + if role == _UNCOND: + for suffix in sorted(_UNCOND_REQUEST_ID_SUFFIXES, key=len, reverse=True): + if request_id.endswith(suffix): + return request_id[: -len(suffix)] + return request_id + + +def _cfg_extra(request: Any, key: str) -> Any: + sampling_params = getattr(request, "sampling_params", None) + extra_args = getattr(sampling_params, "extra_args", None) + if extra_args: + return extra_args.get(key) + return None + + +def _wait_queues(scheduler: Any) -> list[Any]: + # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks + # structurally not-ready requests in a second ``skipped_waiting`` queue; + # pair-hold must cover both queues or a lone member parked there would + # be admitted without its partner. + queues = [scheduler.waiting] + skipped = getattr(scheduler, "skipped_waiting", None) + if skipped is not None: + queues.append(skipped) + return queues + + +def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None = None) -> bool: + pair_id = scheduler._cfg_req_to_pair.get(request_id) + if pair_id is None: + return True + roles = scheduler._cfg_pairs.get(pair_id, {}) + if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()): + return False + if blocked_ids: + # A partner parked in skipped_waiting (e.g. waiting for remote KVs) + # may fail promotion this step while this member gets scheduled from + # waiting; hold this member until the partner is schedulable. + for rid in roles.values(): + if rid != request_id and rid in blocked_ids: + return False + return True + + +# Schedule steps a lone pair member may wait for its partner. Partners arrive +# within a handful of steps in practice; exhausting this budget means the +# companion was never created (e.g. prompt expansion failed), and the request +# is released to decode unguided instead of hanging forever. +_MAX_PAIR_HOLD_STEPS = 512 + + +def _drop_stale_pair(scheduler: Any, request_id: str) -> None: + pair_id = scheduler._cfg_req_to_pair.get(request_id) + if pair_id is None: + return + for rid in scheduler._cfg_pairs.pop(pair_id, {}).values(): + scheduler._cfg_req_to_pair.pop(rid, None) + logger.warning( + "CFG pair %s: partner of %s never arrived after %d schedule steps; releasing the request unguided", + pair_id, + request_id, + _MAX_PAIR_HOLD_STEPS, + ) + + +def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]: + """Pull CFG requests whose partner is not yet admitted out of the wait queues. + + The engine schedules as soon as the first member of a pair arrives over + IPC; without the hold, the lone member prefills one step early and the + pair is permanently offset by one token. Held requests are re-prepended + after the schedule step. + """ + hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {} + scheduler._cfg_hold_counts = hold_counts + + skipped = getattr(scheduler, "skipped_waiting", None) + blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set() + + held: list[tuple[Any, Any]] = [] + for queue in _wait_queues(scheduler): + # A member is only "blocked" for its partner's sake when it sits in + # the OTHER queue; members of the queue being scanned move together. + partner_blockers = blocked_ids if queue is not skipped else set() + held_here = [] + for request in list(queue): + if _pair_complete(scheduler, request.request_id, partner_blockers): + hold_counts.pop(request.request_id, None) + continue + count = hold_counts.get(request.request_id, 0) + 1 + if count > _MAX_PAIR_HOLD_STEPS: + hold_counts.pop(request.request_id, None) + _drop_stale_pair(scheduler, request.request_id) + continue + hold_counts[request.request_id] = count + held_here.append(request) + if held_here: + queue.remove_requests(held_here) + held.extend((queue, req) for req in held_here) + return held + + +def _release_held(held: list[tuple[Any, Any]]) -> None: + for queue, request in reversed(held): + queue.prepend_request(request) + + +def _reorder_waiting_for_cfg(scheduler: Any) -> None: + """Move CFG pair partners adjacent in the FCFS waiting queue.""" + waiting = scheduler.waiting + # Only the FCFS queue (a deque) has caller-controlled ordering; priority + # queues order themselves and pair sync then relies on hold + equalize. + if not isinstance(waiting, deque) or len(waiting) < 2: + return + + requests = list(waiting) + waiting.clear() + + seen: set[str] = set() + result: list[Any] = [] + for request in requests: + rid = request.request_id + if rid in seen: + continue + seen.add(rid) + result.append(request) + + pair_id = scheduler._cfg_req_to_pair.get(rid) + if pair_id is None: + continue + for partner_id in scheduler._cfg_pairs.get(pair_id, {}).values(): + if partner_id != rid and partner_id not in seen: + for candidate in requests: + if candidate.request_id == partner_id: + seen.add(partner_id) + result.append(candidate) + break + + waiting.extend(result) + + +def _equalize_cfg_pair_progress(scheduler: Any, scheduler_output: Any) -> None: + """Bring both members of every CFG pair to the same ``num_computed_tokens``. + + Prefix-cache hits or unequal chunked-prefill budget can leave one member + ahead after a schedule step; reduce the faster member's allocation so + both land on the same position. Over-allocated KV blocks remain and are + consumed next step, so nothing leaks. + """ + for roles in scheduler._cfg_pairs.values(): + cond_id = roles.get(_COND) + uncond_id = roles.get(_UNCOND) + if not cond_id or not uncond_id: + continue + + cond_sched = scheduler_output.num_scheduled_tokens.get(cond_id, 0) + uncond_sched = scheduler_output.num_scheduled_tokens.get(uncond_id, 0) + if cond_sched == 0 or uncond_sched == 0: + continue + + cond_req = scheduler.requests.get(cond_id) + uncond_req = scheduler.requests.get(uncond_id) + if not cond_req or not uncond_req: + continue + + if cond_req.num_computed_tokens == uncond_req.num_computed_tokens: + continue + + target = min(cond_req.num_computed_tokens, uncond_req.num_computed_tokens) + + feasible = all( + req.num_computed_tokens - target < sched + for req, sched in ((cond_req, cond_sched), (uncond_req, uncond_sched)) + ) + if not feasible: + continue + + for req_id, req, orig_sched in ( + (cond_id, cond_req, cond_sched), + (uncond_id, uncond_req, uncond_sched), + ): + diff = req.num_computed_tokens - target + if diff > 0: + req.num_computed_tokens = target + scheduler_output.num_scheduled_tokens[req_id] = orig_sched - diff + scheduler_output.total_num_scheduled_tokens -= diff + logger.debug( + "CFG equalize %s: %d -> %d scheduled, computed -> %d", + req_id, + orig_sched, + orig_sched - diff, + target, + ) + + +def _drop_split_pairs(scheduler: Any) -> None: + """Deregister a desynced CFG pair (one member preempted / left behind). + + Manual re-preemption surgery is unsafe against the upstream scheduler. + Instead, a pair that split anyway (KV-pressure preemption of exactly one + member, or unrecoverable progress skew) is dropped from the registry: + the cond request continues UNGUIDED (its logits are never blended again) + rather than being re-paired against a desynced partner, and the + companion's discarded output no longer matters. + """ + if not scheduler._cfg_pairs: + return + + running_ids = {req.request_id for req in scheduler.running} + to_drop: list[tuple[str, str, str, str]] = [] + for pair_id, roles in scheduler._cfg_pairs.items(): + cond_id = roles.get(_COND) + uncond_id = roles.get(_UNCOND) + if cond_id is None or uncond_id is None: + continue + if cond_id not in scheduler.requests or uncond_id not in scheduler.requests: + continue + + cond_running = cond_id in running_ids + uncond_running = uncond_id in running_ids + if cond_running != uncond_running: + # Split only counts once decoding started: a lone member that has + # computed tokens while its partner was pushed back is desynced. + runner = scheduler.requests.get(cond_id if cond_running else uncond_id) + if runner is not None and runner.num_computed_tokens > 0: + to_drop.append((pair_id, cond_id, uncond_id, "one member preempted")) + continue + + if cond_running and uncond_running: + cond_req = scheduler.requests.get(cond_id) + uncond_req = scheduler.requests.get(uncond_id) + if cond_req and uncond_req and cond_req.num_computed_tokens != uncond_req.num_computed_tokens: + logger.warning( + "CFG pair progress mismatch (equalize will retry): %s@%d vs %s@%d", + cond_id, + cond_req.num_computed_tokens, + uncond_id, + uncond_req.num_computed_tokens, + ) + + for pair_id, cond_id, uncond_id, reason in to_drop: + scheduler._cfg_pairs.pop(pair_id, None) + scheduler._cfg_req_to_pair.pop(cond_id, None) + scheduler._cfg_req_to_pair.pop(uncond_id, None) + logger.warning( + "CFG pair %s split (%s); releasing cond=%s unguided, uncond=%s output is discarded", + pair_id, + reason, + cond_id, + uncond_id, + ) + + +_patches_applied = False + + +def cfg_patches_applied() -> bool: + """True once :func:`apply_cfg_patches` has patched the scheduler in this process.""" + return _patches_applied + + +def _patch_scheduler(scheduler_cls: type | None = None) -> None: + """Wrap the scheduler methods that keep a CFG pair step-locked. + + ``scheduler_cls`` defaults to vLLM's v1 ``Scheduler``; accepting a class + keeps the wrappers exercisable without building a real engine. + """ + if scheduler_cls is None: + from vllm.v1.core.sched.scheduler import Scheduler + + scheduler_cls = Scheduler + + orig_init = scheduler_cls.__init__ + orig_add = scheduler_cls.add_request + orig_finish = scheduler_cls.finish_requests + orig_schedule = scheduler_cls.schedule + + def _init(self, *args: Any, **kwargs: Any) -> None: + orig_init(self, *args, **kwargs) + self._cfg_pairs: dict[str, dict[str, str]] = {} + self._cfg_req_to_pair: dict[str, str] = {} + + def _add(self, request: Any) -> None: + orig_add(self, request) + pair_id = _cfg_extra(request, "cfg_pair_id") + role = _cfg_extra(request, "cfg_role") + if role and not pair_id: + # Models that cannot inject a pair id from the serving layer send + # the role alone and let the request id carry the pairing. It has + # to be the user-facing id: the engine renames every request to an + # internal UUID on admission, and the companion suffix survives + # only on ``external_req_id``. Deriving from the UUID would give + # the two members unrelated pair ids, so they would never be held + # for each other and could be scheduled apart. + external_id = getattr(request, "external_req_id", None) or request.request_id + pair_id = _derive_pair_id(str(external_id), role) + if pair_id and role: + self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id + self._cfg_req_to_pair[request.request_id] = pair_id + # An asymmetric prefix-cache hit (the cond prompt shares cached + # branches; the null prompt does not) can advance one member past + # what post-schedule equalization can claw back. Skip prefix-cache + # reads for CFG requests entirely. + if hasattr(request, "skip_reading_prefix_cache"): + request.skip_reading_prefix_cache = True + + def _finish(self, request_ids: Any, finished_status: Any) -> Any: + if request_ids is None: + result = orig_finish(self, request_ids, finished_status) + self._cfg_pairs.clear() + self._cfg_req_to_pair.clear() + return result + + if isinstance(request_ids, str): + request_ids = {request_ids} + else: + request_ids = set(request_ids) + + # Finish both members together: a surviving partner would decode + # unpaired garbage (its logits are never blended again). + partner_ids: set[str] = set() + for req_id in request_ids: + pair_id = self._cfg_req_to_pair.get(req_id) + if pair_id is None: + continue + for rid in self._cfg_pairs.get(pair_id, {}).values(): + if rid != req_id and rid in self.requests: + partner_ids.add(rid) + + all_ids = request_ids | partner_ids + result = orig_finish(self, all_ids, finished_status) + + for req_id in all_ids: + pair_id = self._cfg_req_to_pair.pop(req_id, None) + if pair_id: + self._cfg_pairs.pop(pair_id, None) + return result + + def _schedule(self, *args: Any, **kwargs: Any) -> Any: + if not self._cfg_pairs: + return orig_schedule(self, *args, **kwargs) + + _reorder_waiting_for_cfg(self) + held = _hold_incomplete_pairs(self) + + # Chunked prefill can split one pair member's prompt across steps + # while the other fits in a single step; halving the long-prefill + # threshold makes both chunk, and equalize aligns the remainder. + scheduler_config = getattr(self, "scheduler_config", None) + orig_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None) + if scheduler_config is not None and orig_threshold is not None: + scheduler_config.long_prefill_token_threshold = self.max_num_scheduled_tokens // 2 + try: + scheduler_output = orig_schedule(self, *args, **kwargs) + finally: + if scheduler_config is not None and orig_threshold is not None: + scheduler_config.long_prefill_token_threshold = orig_threshold + _release_held(held) + + _equalize_cfg_pair_progress(self, scheduler_output) + _drop_split_pairs(self) + return scheduler_output + + _schedule._cfg_pairing_patched = True # type: ignore[attr-defined] + # Legacy marker name kept for callers written against the Audex-only patch. + _schedule._audex_cfg_patched = True # type: ignore[attr-defined] + + scheduler_cls.__init__ = _init + scheduler_cls.add_request = _add + scheduler_cls.finish_requests = _finish + scheduler_cls.schedule = _schedule + + +def apply_cfg_patches() -> None: + """Make the vLLM v1 scheduler CFG-pair-aware. Idempotent, per process. + + Must run in the engine-core process before ``Scheduler`` is constructed + (the ``__init__`` wrapper installs the pair registry). + """ + global _patches_applied + if _patches_applied: + return + _patches_applied = True + logger.info("Applying CFG pairing scheduler patches to vLLM v1") + _patch_scheduler() diff --git a/vllm_omni/model_executor/models/common/ming/__init__.py b/vllm_omni/model_executor/models/common/ming/__init__.py index 208f01a7c..5134c7cf9 100644 --- a/vllm_omni/model_executor/models/common/ming/__init__.py +++ b/vllm_omni/model_executor/models/common/ming/__init__.py @@ -1,2 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Components shared by the Ming-family model packages. + +Naming: The flow-matching pieces all use the cfm_ prefix, +because "flow matching", "fm" and "cfm" were three names for one concept here: + +* dit_blocks - DiT building blocks (DiTBlock, FinalLayer, CondEmbedder) +* cfm_dit - the DiT used as the CFM velocity model +* cfm_solver - timestep schedules and the ODE/SDE integrator +* cfm_head - inference-only CFM wrapper preserving cfm.model.* names +* cfm_cudagraph - batch-1 CUDA-graph executor for the sampling loop +""" diff --git a/vllm_omni/model_executor/models/ming_tts/aggregator.py b/vllm_omni/model_executor/models/common/ming/aggregator.py similarity index 96% rename from vllm_omni/model_executor/models/ming_tts/aggregator.py rename to vllm_omni/model_executor/models/common/ming/aggregator.py index 3233c9659..e32e8b47b 100644 --- a/vllm_omni/model_executor/models/ming_tts/aggregator.py +++ b/vllm_omni/model_executor/models/common/ming/aggregator.py @@ -7,7 +7,7 @@ import torch import torch.nn as nn from x_transformers.x_transformers import RotaryEmbedding -from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer +from vllm_omni.model_executor.models.common.ming.dit_blocks import DiTBlock, FinalLayer class Aggregator(nn.Module): diff --git a/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py b/vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py similarity index 98% rename from vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py rename to vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py index 7a7280034..1eb3a9db0 100644 --- a/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_cudagraph.py @@ -17,7 +17,7 @@ from __future__ import annotations import torch import torch.nn as nn -from vllm_omni.model_executor.models.common.ming.fm import build_timesteps +from vllm_omni.model_executor.models.common.ming.cfm_solver import build_timesteps class CFMSampler(nn.Module): diff --git a/vllm_omni/model_executor/models/ming_tts/fm/dit.py b/vllm_omni/model_executor/models/common/ming/cfm_dit.py similarity index 97% rename from vllm_omni/model_executor/models/ming_tts/fm/dit.py rename to vllm_omni/model_executor/models/common/ming/cfm_dit.py index 3eb13e987..c29ddf48d 100644 --- a/vllm_omni/model_executor/models/ming_tts/fm/dit.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_dit.py @@ -7,7 +7,7 @@ import torch.nn as nn from x_transformers.x_transformers import RotaryEmbedding from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding -from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer +from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer class DiT(nn.Module): diff --git a/vllm_omni/model_executor/models/ming_tts/flowloss_head.py b/vllm_omni/model_executor/models/common/ming/cfm_head.py similarity index 94% rename from vllm_omni/model_executor/models/ming_tts/flowloss_head.py rename to vllm_omni/model_executor/models/common/ming/cfm_head.py index 3f300cec0..14e63557e 100644 --- a/vllm_omni/model_executor/models/ming_tts/flowloss_head.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_head.py @@ -5,9 +5,8 @@ import torch import torch.nn as nn -from vllm_omni.model_executor.models.common.ming.fm import Solver, build_timesteps - -from .fm.dit import DiT +from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT +from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, build_timesteps class FlowLoss(nn.Module): diff --git a/vllm_omni/model_executor/models/common/ming/fm.py b/vllm_omni/model_executor/models/common/ming/cfm_solver.py similarity index 77% rename from vllm_omni/model_executor/models/common/ming/fm.py rename to vllm_omni/model_executor/models/common/ming/cfm_solver.py index d2cbc6211..5d9db4634 100644 --- a/vllm_omni/model_executor/models/common/ming/fm.py +++ b/vllm_omni/model_executor/models/common/ming/cfm_solver.py @@ -3,7 +3,21 @@ import torch -from .dit import get_epss_timesteps + +def get_epss_timesteps(n, device, dtype): + dt = 1 / 32 + predefined_timesteps = { + 5: [0, 2, 4, 8, 16, 32], + 6: [0, 2, 4, 6, 8, 16, 32], + 7: [0, 2, 4, 6, 8, 16, 24, 32], + 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], + 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], + 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], + } + t = predefined_timesteps.get(n, []) + if not t: + return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) + return dt * torch.tensor(t, device=device, dtype=dtype) def apply_sway_sampling(t, sway_sampling_coef): diff --git a/vllm_omni/model_executor/models/common/ming/dit.py b/vllm_omni/model_executor/models/common/ming/dit_blocks.py similarity index 93% rename from vllm_omni/model_executor/models/common/ming/dit.py rename to vllm_omni/model_executor/models/common/ming/dit_blocks.py index 56e65f47f..8a093811e 100644 --- a/vllm_omni/model_executor/models/common/ming/dit.py +++ b/vllm_omni/model_executor/models/common/ming/dit_blocks.py @@ -216,19 +216,3 @@ class CondEmbedder(nn.Module): def forward(self, llm_cond): return self.cond_embedder(llm_cond) - - -def get_epss_timesteps(n, device, dtype): - dt = 1 / 32 - predefined_timesteps = { - 5: [0, 2, 4, 8, 16, 32], - 6: [0, 2, 4, 6, 8, 16, 32], - 7: [0, 2, 4, 6, 8, 16, 24, 32], - 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], - 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], - 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], - } - t = predefined_timesteps.get(n, []) - if not t: - return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) - return dt * torch.tensor(t, device=device, dtype=dtype) diff --git a/vllm_omni/model_executor/models/common/ming/speaker_extractor.py b/vllm_omni/model_executor/models/common/ming/speaker_extractor.py new file mode 100644 index 000000000..bb06a3b58 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ming/speaker_extractor.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py +"""CAM++ speaker-embedding extraction, shared by the Ming model packages.""" + +from __future__ import annotations + +import os + +import torch +import torchaudio +from vllm.multimodal.media.audio import load_audio + +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path +from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor + + +class SpeakerEmbeddingExtractor: + """CAM++ ONNX speaker embeddings, resampled to the extractor's rate.""" + + def __init__(self, model, target_sr=16000, *, allow_download: bool = False): + local_model_path = resolve_model_to_local_path( + model, + allow_download=allow_download, + allow_patterns=["campplus.onnx"], + ) + campplus_path = os.path.join(local_model_path, "campplus.onnx") + if not os.path.exists(campplus_path): + raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}") + + self.target_sr = int(target_sr) + self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr) + + def extract_from_waveform(self, waveform, sample_rate): + if not isinstance(waveform, torch.Tensor): + waveform = torch.as_tensor(waveform) + + tensor = waveform.detach().to(torch.float32) + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + if int(sample_rate) != self.target_sr: + tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor) + + embedding = self._core._extract_spk_embedding(tensor) + return embedding.squeeze(0).to(dtype=torch.float32) + + def extract_from_file(self, audio_path): + audio, sample_rate = load_audio(audio_path, sr=None, mono=False) + waveform = torch.from_numpy(audio).contiguous() + return self.extract_from_waveform(waveform, sample_rate) + + def extract_many(self, audio_paths): + return [self.extract_from_file(path) for path in audio_paths] + + +__all__ = ["SpeakerEmbeddingExtractor"] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py b/vllm_omni/model_executor/models/common/ming/spk_embedding.py similarity index 100% rename from vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py rename to vllm_omni/model_executor/models/common/ming/spk_embedding.py diff --git a/vllm_omni/model_executor/models/common/ops/__init__.py b/vllm_omni/model_executor/models/common/ops/__init__.py new file mode 100644 index 000000000..d9b522bdf --- /dev/null +++ b/vllm_omni/model_executor/models/common/ops/__init__.py @@ -0,0 +1,17 @@ +"""Common operators shared across models and hardware backends. + +This module contains fused operators that can be reused across multiple models +and compiled for different hardware targets (CUDA, ROCm, etc.). +""" + +from vllm_omni.model_executor.models.common.ops.fused_adaptive_group_norm_silu import ( + fused_adaptive_group_norm_silu, +) +from vllm_omni.model_executor.models.common.ops.fused_group_norm_silu import ( + fused_group_norm_silu, +) + +__all__ = [ + "fused_group_norm_silu", + "fused_adaptive_group_norm_silu", +] diff --git a/vllm_omni/model_executor/models/common/ops/fused_adaptive_group_norm_silu.py b/vllm_omni/model_executor/models/common/ops/fused_adaptive_group_norm_silu.py new file mode 100644 index 000000000..2755c3ece --- /dev/null +++ b/vllm_omni/model_executor/models/common/ops/fused_adaptive_group_norm_silu.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# ruff: noqa: N803 + +"""Fused Adaptive Group Normalization (AdaGN) + SiLU operator. + +This module implements the fused AdaGN+SiLU pattern commonly used in Diffusion +Transformer ResBlocks: + + output = SiLU(GroupNorm(x, weight, bias) * (1 + scale) + shift) + +where ``scale`` and ``shift`` are per-(batch, channel) conditioning signals +derived from the timestep embedding. Eager PyTorch spends four kernels and three +full-size intermediates on this; the fused version does it in one pass. + +Falls back to native PyTorch ops when Triton is unavailable. +""" + +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON, tl, triton + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=1), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=2), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=4), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=6), + ], + key=["spatial_size", "C"], +) +@triton.jit +def _adaptive_group_norm_silu_kernel( + # Input/output pointers + x_ptr, + out_ptr, + # Affine parameters + weight_ptr, + bias_ptr, + # Conditioning signals, both (B, C) contiguous + scale_ptr, + shift_ptr, + # Shape info; x is contiguous (B, C, spatial_size) + C, + spatial_size, + num_groups: tl.constexpr, + eps: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + num_stages: tl.constexpr, +): + """ + Fused AdaGN + SiLU kernel: SiLU(GroupNorm(x) * (1 + scale) + shift). + Mean and variance use a parallel Welford reduction (block-level register + reduction merged with the Chan formula), which avoids the catastrophic + cancellation of E[x^2] - E[x]^2 on large-offset inputs. x is read once for + the statistics and once for normalization (2 reads total). + One program handles each (batch, group) pair. Channels are processed serially, + while the spatial axis is vectorized for diffusion workloads. + Moments are accumulated in fp32 to match PyTorch numerics. + """ + pid = tl.program_id(0) + + group_size = C // num_groups + n_idx = pid // num_groups + g_idx = pid % num_groups + + # === Pass 1: Welford reduction over the whole group (fp32) === + # Each block is reduced in registers (block mean, then centered M2), so the + # loaded values are reused instead of re-reading memory; blocks are merged + # with the Chan et al. parallel-Welford formula. This keeps both the mean and + # the variance accurate for inputs like ``10000 +- 0.1``, where the naive + # ``E[x^2] - E[x]^2`` cancels catastrophically, and reads x only once for the + # statistics (2 reads total including normalize). + n_total = tl.zeros([1], dtype=tl.float32) + mean_total = tl.zeros([1], dtype=tl.float32) + m2_total = tl.zeros([1], dtype=tl.float32) + + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # block-level reduction in registers (no extra memory traffic) + n = tl.sum(tl.where(mask, 1.0, 0.0), axis=0) + bsum = tl.sum(x_val, axis=0) + bmean = bsum / n + bm2 = tl.sum(tl.where(mask, (x_val - bmean) * (x_val - bmean), 0.0), axis=0) + + # Chan et al. merge into the running (n, mean, m2) + delta = bmean - mean_total + new_n = n_total + n + mean_total = mean_total + delta * (n / new_n) + m2_total = m2_total + bm2 + delta * delta * (n_total * n / new_n) + n_total = new_n + + mean = mean_total + var = m2_total / n_total + rstd = 1.0 / tl.sqrt(var + eps) + + # === Pass 2: normalize, affine, then adaptive modulation === + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + weight_val = tl.load(weight_ptr + c_idx).to(tl.float32) + bias_val = tl.load(bias_ptr + c_idx).to(tl.float32) + scale_val = tl.load(scale_ptr + n_idx * C + c_idx).to(tl.float32) + shift_val = tl.load(shift_ptr + n_idx * C + c_idx).to(tl.float32) + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # GroupNorm: (x - mean) * rstd * weight + bias + norm_val = (x_val - mean) * rstd * weight_val + bias_val + + # AdaGN: norm * (1 + scale) + shift + out_val = norm_val * (1.0 + scale_val) + shift_val + + # SiLU: out * sigmoid(out) + out_val = out_val * tl.sigmoid(out_val) + + # ``tl.store`` casts to the output pointer's dtype, which the + # caller picked to match eager GroupNorm's autocast behaviour. + tl.store(out_ptr + base + offsets, out_val, mask=mask) + + +def fused_adaptive_group_norm_silu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + num_groups: int, + eps: float = 1e-6, +) -> torch.Tensor: + """ + Fused Adaptive GroupNorm + SiLU. + + Computes SiLU(GroupNorm(x) * (1 + scale) + shift) in one fused operation, avoiding intermediate tensors. + + - x: (B, C, *spatial), any spatial rank ≥ 1. + - weight, bias: per-channel parameters, shape (C,). + - scale, shift: adaptive parameters with B * C elements, broadcastable per channel. + - num_groups: GroupNorm group count. + - eps: numerical stability epsilon. + + Returns a tensor with the same shape as x, matching eager F.group_norm dtype behavior. + + Spatial dimensions are flattened during computation and restored afterward. + GroupNorm statistics cover each channel group across all spatial positions, so the result is exact. + """ + assert x.ndim >= 3, f"Expected at least 3D input (B, C, *spatial), got {x.ndim}D" + B, C = x.shape[:2] + assert C % num_groups == 0, f"num_channels ({C}) must be divisible by num_groups ({num_groups})" + assert weight.ndim == 1 and weight.size(0) == C, f"Weight shape {tuple(weight.shape)} doesn't match channels {C}" + assert bias.ndim == 1 and bias.size(0) == C, f"Bias shape {tuple(bias.shape)} doesn't match channels {C}" + assert scale.numel() == B * C, f"scale has {scale.numel()} elements, expected B*C = {B * C}" + assert shift.numel() == B * C, f"shift has {shift.numel()} elements, expected B*C = {B * C}" + + # Fallback if Triton not available (NPU, CPU, ...) + if not HAS_TRITON: + broadcast = (B, C) + (1,) * (x.ndim - 2) + normed = F.group_norm(x, num_groups, weight, bias, eps) + return F.silu(normed * (1.0 + scale.reshape(broadcast)) + shift.reshape(broadcast)) + + orig_shape = x.shape + + # The kernel indexes x as a dense (B, C, spatial) block, so normalize the + # layout here. ``reshape``/``contiguous`` are no-ops for the common case of + # a contiguous activation coming out of a conv. + x_flat = x.contiguous().reshape(B, C, -1) + spatial_size = x_flat.size(2) + + # ``reshape`` rather than ``view``: scale/shift typically arrive as halves + # of a torch.chunk along dim 1, which is not contiguous when B > 1. + scale_2d = scale.reshape(B, C).contiguous() + shift_2d = shift.reshape(B, C).contiguous() + + # Allocate with the dtype eager GroupNorm would return, so the fused path + # stays a drop-in replacement inside autocast regions. + out_dtype = x_flat.dtype + if torch.is_autocast_enabled(x_flat.device.type): + out_dtype = torch.float32 + out_flat = torch.empty_like(x_flat, dtype=out_dtype) + + # Only B*num_groups programs are launched, which is well under the SM count + # for typical diffusion batches. A memory-bound kernel can still saturate + # HBM from few CTAs, but only with enough loads in flight, so widen the CTA + # for the large activations instead of leaving it at the 4-warp default. + + # One program per (batch, group) pair. + grid = (B * num_groups,) + + _adaptive_group_norm_silu_kernel[grid]( + x_flat, + out_flat, + weight, + bias, + scale_2d, + shift_2d, + C, + spatial_size, + num_groups=num_groups, + eps=eps, + ) + + return out_flat.reshape(orig_shape) + + +__all__ = ["fused_adaptive_group_norm_silu"] diff --git a/vllm_omni/model_executor/models/common/ops/fused_group_norm_silu.py b/vllm_omni/model_executor/models/common/ops/fused_group_norm_silu.py new file mode 100644 index 000000000..21fb245c6 --- /dev/null +++ b/vllm_omni/model_executor/models/common/ops/fused_group_norm_silu.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# ruff: noqa: N803 + +"""Fused GroupNorm + SiLU operator. + +This operator fuses GroupNorm followed by SiLU activation into a single kernel, +reducing memory traffic and kernel launch overhead. The implementation uses +Triton for CUDA/ROCm compatibility, and falls back to native PyTorch ops when +Triton is unavailable (NPU, CPU, ...), so callers never need a platform check. + +Measured against eager ``F.silu(F.group_norm(...))`` on one L20X, bf16, 32 +groups: 1.1-1.5x at the DiT ResBlock's activation sizes, where both paths are +dominated by launch overhead, and 2.2-2.9x at the VAE's decode-resolution +activations, where the saved memory traffic is what pays. +""" + +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON, tl, triton + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=1), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=2), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=4), + triton.Config({"BLOCK_SIZE": 4096}, num_warps=16, num_stages=6), + ], + key=["spatial_size", "C"], +) +@triton.jit +def _group_norm_silu_kernel( + # Input/Output pointers + x_ptr, + out_ptr, + # Normalization parameters + weight_ptr, + bias_ptr, + # Shape info; x is contiguous (N, C, spatial_size) + C, + spatial_size, + num_groups: tl.constexpr, + eps: tl.constexpr, + # Block sizes + BLOCK_SIZE: tl.constexpr, + num_stages: tl.constexpr, +): + """ + Fused GroupNorm + SiLU kernel. + Computes SiLU(GroupNorm(x)) in one kernel, avoiding intermediate tensors. + Mean and variance use a parallel Welford reduction (block-level register + reduction merged with the Chan formula), which avoids the catastrophic + cancellation of E[x^2] - E[x]^2 on large-offset inputs. x is read once for + the statistics and once for normalization (2 reads total). + One program handles each (batch, group) pair. Channels are processed serially, + while spatial positions are vectorized for diffusion workloads. + Moments are accumulated in fp32 to match PyTorch numerics. + """ + pid = tl.program_id(0) + + group_size = C // num_groups + n_idx = pid // num_groups + g_idx = pid % num_groups + + # === Pass 1: Welford reduction over the whole group (fp32) === + # Each block is reduced in registers (block mean, then centered M2), so the + # loaded values are reused instead of re-reading memory; blocks are merged + # with the Chan et al. parallel-Welford formula. This keeps both the mean and + # the variance accurate for inputs like ``10000 +- 0.1``, where the naive + # ``E[x^2] - E[x]^2`` cancels catastrophically, and reads x only once for the + # statistics (2 reads total including normalize). + n_total = tl.zeros([1], dtype=tl.float32) + mean_total = tl.zeros([1], dtype=tl.float32) + m2_total = tl.zeros([1], dtype=tl.float32) + + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # block-level reduction in registers (no extra memory traffic) + n = tl.sum(tl.where(mask, 1.0, 0.0), axis=0) + bsum = tl.sum(x_val, axis=0) + bmean = bsum / n + bm2 = tl.sum(tl.where(mask, (x_val - bmean) * (x_val - bmean), 0.0), axis=0) + + # Chan et al. merge into the running (n, mean, m2) + delta = bmean - mean_total + new_n = n_total + n + mean_total = mean_total + delta * (n / new_n) + m2_total = m2_total + bm2 + delta * delta * (n_total * n / new_n) + n_total = new_n + + mean = mean_total + var = m2_total / n_total + rstd = 1.0 / tl.sqrt(var + eps) + + # === Pass 2: normalize, apply affine, and SiLU === + for c_offset in range(group_size): + c_idx = g_idx * group_size + c_offset + base = n_idx * C * spatial_size + c_idx * spatial_size + + weight_val = tl.load(weight_ptr + c_idx).to(tl.float32) + bias_val = tl.load(bias_ptr + c_idx).to(tl.float32) + + for s_start in tl.range(0, spatial_size, BLOCK_SIZE, num_stages=num_stages): + offsets = s_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < spatial_size + + x_val = tl.load(x_ptr + base + offsets, mask=mask, other=0.0) + x_val = x_val.to(tl.float32) + + # Normalize and apply affine + norm_val = (x_val - mean) * rstd * weight_val + bias_val + + # Apply SiLU: x * sigmoid(x) + out_val = norm_val * tl.sigmoid(norm_val) + + # ``tl.store`` casts to the output pointer's dtype, which the + # caller picked to match eager GroupNorm's autocast behaviour. + tl.store(out_ptr + base + offsets, out_val, mask=mask) + + +def fused_group_norm_silu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int = 32, + eps: float = 1e-6, +) -> torch.Tensor: + """ + Fused GroupNorm + SiLU activation. + Computes SiLU(GroupNorm(x, num_groups, weight, bias, eps)) in a single Triton kernel, + avoiding intermediate tensors and reducing memory traffic and launch overhead. + + - x: (N, C, *spatial), any spatial rank. + - weight, bias: per-channel parameters, shape (C,). + - num_groups: GroupNorm group count, default 32. + - eps: numerical stability epsilon, default 1e-6. + Uses fp32 accumulation for numeric alignment with PyTorch. + Returns the same shape and eager F.group_norm-compatible dtype. + Spatial dimensions are flattened during computation and restored afterward; + this is exact since GroupNorm reduces across each channel group and all spatial positions. + Non-contiguous inputs are materialized before launch. + """ + # Fallback if Triton not available (NPU, CPU, ...) + if not HAS_TRITON: + return F.silu(F.group_norm(x, num_groups, weight, bias, eps)) + + # Validate inputs + assert x.ndim >= 3, f"Expected at least 3D input (N, C, *spatial), got {x.ndim}D" + assert x.size(1) % num_groups == 0, f"Channels {x.size(1)} must be divisible by num_groups {num_groups}" + assert weight.ndim == 1 and weight.size(0) == x.size(1), ( + f"Weight shape {weight.shape} doesn't match channels {x.size(1)}" + ) + assert bias.ndim == 1 and bias.size(0) == x.size(1), f"Bias shape {bias.shape} doesn't match channels {x.size(1)}" + + # Collapse arbitrary spatial ranks into a single axis so one kernel serves + # both the 2D DiT blocks and the 3D VAE blocks. + # + # ``contiguous()`` is not just defensive. HunyuanImage3's UNetUp feeds this + # op straight out of ``rearrange(x, "b (h w) c -> b c h w")``, i.e. a + # permuted view whose *channel* stride is 1. Indexing that layout directly + # from the kernel makes every warp's spatial load stride by C elements, and + # the resulting uncoalesced traffic turned the op into 0.44x of eager at + # (1, 4096, 64, 64). One coalesced pre-pass costs far less than that, so + # normalize the layout here and let the kernel assume a dense block. + orig_shape = x.shape + B, C = orig_shape[0], orig_shape[1] + x_flat = x.contiguous().reshape(B, C, -1) + spatial_size = x_flat.size(2) + + # Allocate output with the dtype eager GroupNorm would return, so that the + # fused path stays a drop-in replacement inside autocast regions. + out_dtype = x_flat.dtype + if torch.is_autocast_enabled(x_flat.device.type): + out_dtype = torch.float32 + out_flat = torch.empty_like(x_flat, dtype=out_dtype) + + # Only B*num_groups programs are launched, which is well under the SM count + # for typical diffusion batches. A memory-bound kernel can still saturate + # HBM from few CTAs, but only with enough loads in flight, so widen the CTA + # for the large activations instead of leaving it at the 4-warp default. + + # One program per (batch, group) pair. + grid = (B * num_groups,) + + _group_norm_silu_kernel[grid]( + x_flat, + out_flat, + weight, + bias, + C, + spatial_size, + num_groups=num_groups, + eps=eps, + ) + + return out_flat.reshape(orig_shape) diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 8c9c134cf..3a47e2912 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -3,7 +3,8 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models. * SDPA attention (F.scaled_dot_product_attention) with native GQA support -* HF-compatible CPU/CUDA numerics with NPU-only fused norm/RoPE fast paths +* HF-compatible CPU/CUDA numerics with packed qkv/gate_up projections + and NPU-only fused norm/RoPE fast paths * Per-call embedding buffer to avoid cross-request aliasing * Pre-allocated position_ids (read-only, safe to persist) * torch.compile (epilogue_fusion=False) on inner transformer by default @@ -45,8 +46,9 @@ if current_omni_platform.is_npu(): # HuggingFace reference on CPU/CUDA. vLLM's fused kernels (RMSNorm, # QKVParallel, get_rope) introduce small precision differences that compound # across the autoregressive steps of the code predictor, causing severe audio -# quality degradation. The Ascend fused norm/RoPE kernels below are dispatched -# by the current device platform. +# quality degradation. Packing qkv/gate_up into plain ``nn.Linear`` keeps the +# same algebra but can still change GPU bit-level accumulation. The Ascend +# fused norm/RoPE kernels below are dispatched by the current device platform. # # See: https://github.com/vllm-project/vllm-omni/issues/2274 @@ -175,11 +177,23 @@ class CodePredictorAttention(nn.Module): self.scaling = self.head_dim**-0.5 self.max_seq = int(config.num_code_groups) + 1 - # Separate q/k/v projections matching HF (no fused packing) + # Fused QKV projection. + # + # The HF reference uses three separate ``nn.Linear`` modules for + # q/k/v. Packing those weights row-wise into one ``nn.Linear`` keeps + # the same matmul at the algebra level; see the module-level + # HF-numerics note for the GPU bit-level drift caveat. The goal here + # is narrower: reduce eager host launch overhead while keeping the + # downstream q_norm / k_norm / SDPA math and HF checkpoint loading + # contract intact. bias = getattr(config, "attention_bias", False) - self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=bias) - self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias) - self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias) + self._q_size = self.num_heads * self.head_dim + self._kv_size = self.num_kv_heads * self.head_dim + self.qkv_proj = nn.Linear( + self.hidden_size, + self._q_size + 2 * self._kv_size, + bias=bias, + ) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) @@ -197,6 +211,15 @@ class CodePredictorAttention(nn.Module): ) self.register_buffer("_fusion_causal_mask", fusion_mask, persistent=False) + def _split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_raw = qkv[..., : self._q_size] + kv_raw = qkv[..., self._q_size :] + # ``split`` would leave the non-leading k/v slices strided like the + # packed qkv buffer. Materialize them once here so SDPA sees the same + # dense k/v layout as the unfused projection path. + k_raw, v_raw = (t.contiguous() for t in kv_raw.split(self._kv_size, dim=-1)) + return q_raw, k_raw, v_raw + def _forward_npu_attention( self, q: torch.Tensor, @@ -257,9 +280,14 @@ class CodePredictorAttention(nn.Module): hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) - q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2) - k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2) - v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2) + # Single fused matmul. ``q`` is immediately consumed by RMSNorm; for + # ``k/v`` we re-densify the packed tail once so SDPA does not inherit + # the packed-buffer stride pattern from the fused projection. + qkv = self.qkv_proj(hidden_states) + q_raw, k_raw, v_raw = self._split_qkv(qkv) + q = self.q_norm(q_raw.view(hidden_shape_q)).transpose(1, 2) + k = self.k_norm(k_raw.view(hidden_shape_kv)).transpose(1, 2) + v = v_raw.view(hidden_shape_kv).transpose(1, 2) cos, sin = position_embeddings # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads @@ -295,12 +323,22 @@ class CodePredictorMLP(nn.Module): def __init__(self, config, *, prefix: str = "") -> None: super().__init__() - self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) - self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + # Fused gate_up projection. Same packing idea as ``qkv_proj`` above: + # one matmul instead of two, with the same algebraic result as + # separate gate/up linears. See the module-level HF-numerics note for + # the GPU bit-level drift caveat. + self._intermediate_size = int(config.intermediate_size) + self.gate_up_proj = nn.Linear( + config.hidden_size, + 2 * config.intermediate_size, + bias=False, + ) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) + gate_up = self.gate_up_proj(hidden_states) + gate, up = gate_up.split([self._intermediate_size, self._intermediate_size], dim=-1) + return self.down_proj(F.silu(gate) * up) # =================================================================== @@ -411,18 +449,140 @@ class CodePredictorBaseModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states.to(input_dtype) + # Mapping from HF checkpoint shard name to the (fused parameter, + # slice index) pair used to assemble the fused tensor. See + # ``load_weights`` below. + _FUSED_QKV_SHARDS: tuple[tuple[str, str, int], ...] = ( + ("q_proj", "qkv_proj", 0), + ("k_proj", "qkv_proj", 1), + ("v_proj", "qkv_proj", 2), + ) + _FUSED_GATE_UP_SHARDS: tuple[tuple[str, str, int], ...] = ( + ("gate_proj", "gate_up_proj", 0), + ("up_proj", "gate_up_proj", 1), + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights, transparently re-packing HF q/k/v and gate/up shards + into the fused ``qkv_proj`` and ``gate_up_proj`` parameters. + + HF checkpoints ship q_proj / k_proj / v_proj (and gate_proj / + up_proj) as separate tensors. Our model defines a single + ``qkv_proj`` (``gate_up_proj``) whose weight is the row-wise + concatenation of those shards in q/k/v (gate/up) order. We + buffer the shards per layer prefix and write the fused parameter + in one go once all shards have arrived. + + Both ``.weight`` and ``.bias`` shards are packed this way, so a + code-predictor config with ``attention_bias=True`` is handled + correctly: the q/k/v bias vectors are concatenated into + ``qkv_proj.bias`` instead of being dropped (which would leave the + fused bias randomly initialized). ``gate_up_proj`` has no bias, + so only its ``.weight`` shards are ever routed. + """ params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() + fused_layer_prefixes = { + name.split(marker, 1)[0] + for name in params_dict + for marker in (".self_attn.qkv_proj.", ".mlp.gate_up_proj.") + if marker in name + } + touched_fused_layers: set[str] = set() + + # layer_prefix -> {fused_name: {kind: {slice_idx: shard_tensor}}} + # where ``kind`` is "weight" or "bias". + pending: dict[str, dict[str, dict[str, dict[int, torch.Tensor]]]] = {} + + def _try_finalize_fused( + layer_prefix: str, + fused_name: str, + kind: str, + num_shards: int, + ) -> None: + bucket = pending.get(layer_prefix, {}).get(fused_name, {}).get(kind) + if bucket is None or len(bucket) < num_shards: + return + target_param_name = f"{layer_prefix}.{fused_name}.{kind}" + target_param = params_dict.get(target_param_name) + if target_param is None: + return + ordered = [bucket[i] for i in range(num_shards)] + fused_weight = torch.cat(ordered, dim=0) + weight_loader = getattr(target_param, "weight_loader", default_weight_loader) + weight_loader(target_param, fused_weight) + loaded_params.add(target_param_name) + del pending[layer_prefix][fused_name][kind] + if not pending[layer_prefix][fused_name]: + del pending[layer_prefix][fused_name] + if not pending[layer_prefix]: + del pending[layer_prefix] + + def _route_shard( + name: str, + loaded_weight: torch.Tensor, + shard_table: tuple[tuple[str, str, int], ...], + parent_attr: str, + ) -> bool: + """Stash a HF ``.weight``/``.bias`` shard if its fused counterpart + exists in this module; return True iff the shard was consumed.""" + for shard_name, fused_name, slice_idx in shard_table: + for kind in ("weight", "bias"): + suffix = f".{parent_attr}.{shard_name}.{kind}" + if not name.endswith(suffix): + continue + layer_prefix = name[: -len(f".{shard_name}.{kind}")] # ...self_attn or ...mlp + fused_param_name = f"{layer_prefix}.{fused_name}.{kind}" + if fused_param_name not in params_dict: + # e.g. a gate/up ``.bias`` shard when the fused + # projection has no bias: nothing to pack into. + return False + pending.setdefault(layer_prefix, {}).setdefault(fused_name, {}).setdefault(kind, {})[slice_idx] = ( + loaded_weight + ) + _try_finalize_fused(layer_prefix, fused_name, kind, num_shards=len(shard_table)) + return True + return False + for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue + touched_fused_layers.update( + layer_prefix for layer_prefix in fused_layer_prefixes if name.startswith(f"{layer_prefix}.") + ) + if _route_shard(name, loaded_weight, self._FUSED_QKV_SHARDS, parent_attr="self_attn"): + continue + if _route_shard(name, loaded_weight, self._FUSED_GATE_UP_SHARDS, parent_attr="mlp"): + continue param = params_dict.get(name) if param is None: continue weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) + + if pending: + unfinished = { + lp: {fn: list(kinds.keys()) for fn, kinds in buckets.items()} for lp, buckets in pending.items() + } + raise RuntimeError( + f"CodePredictor load_weights: incomplete fused shards for layers {unfinished}; " + f"check that the checkpoint contains all q/k/v and gate/up tensors." + ) + + missing_fused = sorted( + name + for name in params_dict + if (".self_attn.qkv_proj." in name or ".mlp.gate_up_proj." in name) + and name not in loaded_params + and any(name.startswith(f"{layer_prefix}.") for layer_prefix in touched_fused_layers) + ) + if missing_fused: + raise RuntimeError( + f"CodePredictor load_weights: missing fused parameters {missing_fused}; " + f"check that the checkpoint contains the q/k/v and gate/up shards, " + f"or the already-fused qkv_proj/gate_up_proj tensors." + ) return loaded_params @@ -997,18 +1157,47 @@ class CodePredictorWrapper(nn.Module): logger.info("Prepared NPU code predictor weights: linear=%d", linear_count) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights directly (no fused projection remapping needed).""" + """Load wrapper weights, delegating inner-model tensors to the fused loader. + + Qwen3-TTS checkpoints use the current ``model.*`` wrapper layout, while + some Qwen3-Omni artifacts expose the shared body directly as + ``layers.*`` / ``codec_embedding.*``. Normalize both forms before + calling ``CodePredictorBaseModel.load_weights`` so its q/k/v and gate/up + repacking guards see the HF shards instead of treating them as unrelated + wrapper-level tensors. + """ loaded: set[str] = set() model_weights: list[tuple[str, torch.Tensor]] = [] other_weights: list[tuple[str, torch.Tensor]] = [] + inner_model_prefixes = ("codec_embedding.", "layers.", "norm.", "rotary_emb.") + wrapper_prefixes = tuple( + dict.fromkeys( + prefix + for prefix in ( + f"{self.prefix}.", + f"{self.prefix.rsplit('.', 1)[-1]}." if self.prefix else "", + "code_predictor.", + ) + if prefix + ) + ) for name, w in weights: if "rotary_emb.inv_freq" in name: continue - if name.startswith("model."): - model_weights.append((name[len("model.") :], w)) + + normalized_name = name + for wrapper_prefix in wrapper_prefixes: + if normalized_name.startswith(wrapper_prefix): + normalized_name = normalized_name[len(wrapper_prefix) :] + break + + if normalized_name.startswith("model."): + model_weights.append((normalized_name[len("model.") :], w)) + elif normalized_name.startswith(inner_model_prefixes): + model_weights.append((normalized_name, w)) else: - other_weights.append((name, w)) + other_weights.append((normalized_name, w)) loaded_model = self.model.load_weights(model_weights) loaded |= {f"model.{n}" for n in loaded_model} diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py index 78c1ea66f..2ec12a251 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py @@ -148,9 +148,9 @@ class ConditionalCFM(BASECFM): # ``.contiguous().data_ptr()`` could free the temp -> dangling ptr). io_dtype = getattr(self.estimator, "io_dtype", x.dtype) [estimator, stream], trt_engine = self.estimator.acquire_estimator() - # NOTE need to synchronize when switching stream - torch.cuda.current_stream().synchronize() - with stream: + caller_stream = torch.cuda.current_stream(x.device) + stream.wait_stream(caller_stream) + with torch.cuda.stream(stream): x_e = x.to(io_dtype).contiguous() mask_e = mask.to(io_dtype).contiguous() mu_e = mu.to(io_dtype).contiguous() @@ -176,8 +176,13 @@ class ConditionalCFM(BASECFM): for i, j in enumerate(data_ptrs): estimator.set_tensor_address(trt_engine.get_tensor_name(i), j) # run trt engine - assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True - torch.cuda.current_stream().synchronize() + assert estimator.execute_async_v3(stream.cuda_stream) is True + for tensor in (x_e, mask_e, mu_e, t_e, spks_e, cond_e, out_e): + if tensor.is_cuda: + tensor.record_stream(stream) + caller_stream.wait_stream(stream) + if out_e.is_cuda: + out_e.record_stream(caller_stream) self.estimator.release_estimator(estimator, stream) return out_e.to(x.dtype) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index d8d99aa21..2a5d472a0 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/hifigan # Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) @@ -494,7 +494,7 @@ class HiFTGenerator(nn.Module): self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], - window=self.stft_window, + window=self._get_stft_window(x), return_complex=True, ) spec = torch.view_as_real(spec) # [B, F, TT, 2] @@ -509,10 +509,15 @@ class HiFTGenerator(nn.Module): self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], - window=self.stft_window, + window=self._get_stft_window(magnitude), ) return inverse_transform + def _get_stft_window(self, tensor: torch.Tensor) -> torch.Tensor: + if self.stft_window.device != tensor.device: + self.stft_window = self.stft_window.to(tensor.device) + return self.stft_window + def _decode_pre_istft( self, x: torch.Tensor, @@ -696,7 +701,11 @@ class CausalHiFTGenerator(HiFTGenerator): self.ups.apply(init_weights) self.conv_post.apply(init_weights) self.reflection_pad = nn.ReflectionPad1d((1, 0)) - self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) + self.register_buffer( + "stft_window", + torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)), + persistent=False, + ) self.conv_pre_look_right = conv_pre_look_right self.f0_predictor = f0_predictor diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index f2f9a6a0f..81bae6d6c 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import os from collections.abc import Iterable, Mapping, Sequence from functools import partial @@ -35,6 +35,7 @@ from vllm.multimodal.processing import ( from vllm.sequence import IntermediateTensors from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.ops.topk_topp_sampler import random_sample from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct @@ -566,8 +567,8 @@ class CosyVoice3Model( return float(value) @staticmethod - def _multinomial_sample(probs: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: - return torch.multinomial(probs, 1, replacement=True, generator=generator).reshape(()) + def _random_sample_one(probs: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + return random_sample(probs.unsqueeze(0), {} if generator is None else {0: generator}).reshape(()) @classmethod def _nucleus_sample_one( @@ -602,7 +603,7 @@ class CosyVoice3Model( # First token always passes (cum_before[0] = 0 < top_p for any top_p > 0), # so ``weights`` is guaranteed to have at least one nonzero entry. The # final ``.item()`` is the ONLY D2H sync per call. - sample_idx = torch.multinomial(weights, 1, replacement=True, generator=generator) + sample_idx = cls._random_sample_one(weights, generator=generator) return int(sorted_idx[sample_idx].item()) @classmethod @@ -632,9 +633,14 @@ class CosyVoice3Model( rep_num = int((recent == top_id).sum().item()) if rep_num >= win_size * tau_r: weighted_scores = weighted_scores.clone() + original_score = weighted_scores[top_id].clone() weighted_scores[top_id] = float("-inf") - fallback_probs = weighted_scores.softmax(dim=0) - top_id = int(cls._multinomial_sample(fallback_probs, generator=generator).item()) + weighted_scores[top_id] = torch.where( + torch.isfinite(weighted_scores).any(), + weighted_scores[top_id], + original_score, + ) + top_id = int(cls._random_sample_one(weighted_scores.softmax(dim=0), generator=generator).item()) return top_id def _cosyvoice3_ras_enabled(self, sampling_metadata: SamplingMetadata) -> bool: @@ -677,6 +683,10 @@ class CosyVoice3Model( logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf")) for processor in sampling_metadata.logitsprocs.non_argmax_invariant: logits = processor.apply(logits) + finite_logits = torch.isfinite(logits) + if not finite_logits.any(dim=-1).all().item(): + raise ValueError("CosyVoice3 sampling received a row with no finite logits") + logits.masked_fill_(~finite_logits, float("-inf")) sampling_cfg = dict(self.config.llm.get("sampling", {})) default_top_p = float(sampling_cfg.get("top_p", 0.8)) diff --git a/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py b/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py index 230fb14fb..ea2f84e39 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py +++ b/vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py @@ -121,7 +121,7 @@ class TrtContextWrapper: for _ in range(trt_concurrent): ctx = engine.create_execution_context() assert ctx is not None, "failed to create TRT execution context (out of memory?)" - stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device))) + stream = torch.cuda.Stream(torch.device(device)) self._pool.put([ctx, stream]) def acquire_estimator(self): diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py index 413dc1dc3..d6bb0e5bb 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py +++ b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py @@ -48,6 +48,7 @@ COSYVOICE3_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="latent", sync_process_input_func=f"{_PROC}.text2flow_token_only", + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/covo_audio/pipeline.py b/vllm_omni/model_executor/models/covo_audio/pipeline.py index c2e1911ef..b3bd9c0c6 100644 --- a/vllm_omni/model_executor/models/covo_audio/pipeline.py +++ b/vllm_omni/model_executor/models/covo_audio/pipeline.py @@ -47,6 +47,7 @@ COVO_AUDIO_PIPELINE = PipelineConfig( engine_output_type="audio", sync_process_input_func=f"{_PROC}.llm2code2wav_token_only", sampling_constraints={"detokenize": False}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/dots_tts/__init__.py b/vllm_omni/model_executor/models/dots_tts/__init__.py new file mode 100644 index 000000000..fb6738a29 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py new file mode 100644 index 000000000..086ac109a --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_dit.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone +# +# Merged from two upstream files (preserved as-is, inference behaviour +# unchanged; only internal imports were folded inline): +# - modules/backbone/layers.py -> Dropout, Mlp, rotate_half, +# apply_rotary_pos_emb, RotaryEmbedding, +# MultiHeadAttention +# (Conv1d / ConvTranspose1d are NOT +# re-vendored here — they already live in +# dots_tts_vocoder.py) +# - modules/backbone/dit.py -> modulate, TimestepEmbedder, FinalLayer, +# DiTBlock, DiT +# +# Per-section upstream attributions are preserved at each section header. +# Nothing in this file diverges from upstream behaviour. + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +# ============================================================================ +# Building blocks (adapted from modules/backbone/layers.py) +# ---------------------------------------------------------------------------- +# Conv1d / ConvTranspose1d from layers.py are already vendored in +# dots_tts_vocoder.py (they are used by the BigVGAN-style AudioVAE). The +# pieces below are the ones consumed by the DiT flow-matching head. +# ============================================================================ + + +class Dropout(nn.Module): + def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **_kwargs): + super().__init__() + if p < 0.0 or p > 1.0: + raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}") + self.p = p + self.inplace = inplace + self.force_drop = force_drop + + def forward(self, x, **_kwargs): + return F.dropout( + x, + p=self.p, + training=True if self.force_drop else self.training, + inplace=self.inplace, + ) + + +class Mlp(nn.Module): + def __init__( + self, + hidden_size, + ffn_hidden_size=4096, + act_layer=nn.GELU, + dropout=0.0, + **_kwargs, + ): + super().__init__() + self.fc1 = nn.Linear(hidden_size, ffn_hidden_size) + self.act = act_layer() + self.fc2 = nn.Linear(ffn_hidden_size, hidden_size) + self.drop = Dropout(dropout) + + def forward(self, x, _mask=None): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + return self.drop(x) + + +def rotate_half(x): + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +@torch.autocast(enabled=False, device_type="cuda") +def apply_rotary_pos_emb(pos, t): + if pos.dim() == 3: + pos = pos.unsqueeze(1) + return t * pos.cos() + rotate_half(t) * pos.sin() + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, theta=50000): + super().__init__() + self.register_buffer( + "inv_freq", + 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)), + persistent=False, + ) + self._theta = float(theta) + + def _apply(self, fn): + inv_freq = self.inv_freq + super()._apply(fn) + self.inv_freq = inv_freq.to(device=self.inv_freq.device, dtype=torch.float32) + return self + + @torch.autocast(enabled=False, device_type="cuda") + def forward(self, t): + inv_freq = self.inv_freq + if inv_freq.device != t.device: + raise RuntimeError(f"RotaryEmbedding buffer device mismatch: inv_freq={inv_freq.device} input={t.device}.") + t = t.to(dtype=inv_freq.dtype) + if t.dim() == 1: + freqs = torch.einsum("i , j -> i j", t, inv_freq) + else: + freqs = torch.einsum("bi, j -> bij", t, inv_freq) + return torch.cat((freqs, freqs), dim=-1) + + +class MultiHeadAttention(nn.Module): + """Multi-head attention""" + + def __init__( + self, + hidden_size: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_norm: bool = False, + attn_drop: float = 0.0, + dropout: float = 0.0, + norm_layer: str = "LayerNorm", + rotary_bias: bool = False, + rotary_theta: float | None = 50000, + **_kwargs, + ): + super().__init__() + assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads" + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scale = self.head_dim**-0.5 + self.rotary_bias = rotary_bias + + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + + norm_layer = getattr(nn, norm_layer) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + + self.attn_drop = Dropout(attn_drop) + self.o_proj = nn.Linear(hidden_size, hidden_size) + self.o_dropout = Dropout(dropout) + + if self.rotary_bias: + self.rotary = RotaryEmbedding(self.head_dim, theta=rotary_theta) + + def forward(self, q, k=None, v=None, mask=None, pos_ids=None, **_kwargs): + k = k or q + v = v or q + B, L, _ = q.shape + _, S, _ = v.shape + if mask is not None: + if mask.ndim == 2: # [B, L] + assert L == S + mask = rearrange(mask, "b j -> b 1 1 j") + mask = mask.expand(-1, self.num_heads, L, -1) + elif mask.ndim == 3: # [B, L, S] + assert mask.size(1) == L and mask.size(2) == S + mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1) + + q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v) + q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) + k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads) + v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads) + q, k = self.q_norm(q), self.k_norm(k) + + # Apply rotary + if self.rotary_bias: + if L == S: + if pos_ids is None: + rotary_emb = self.rotary(torch.arange(L, device=q.device)) + else: + rotary_emb = self.rotary(pos_ids) + q, k = (apply_rotary_pos_emb(rotary_emb, tensor) for tensor in (q, k)) + else: + q_rotary_emb = self.rotary(torch.arange(L, device=q.device)) + k_rotary_emb = self.rotary(torch.arange(S, device=k.device)) + q = apply_rotary_pos_emb(q_rotary_emb, q) + k = apply_rotary_pos_emb(k_rotary_emb, k) + + attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device) + + if mask is not None: + attn_bias.masked_fill_(mask.logical_not(), float("-inf")) + + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_bias, + dropout_p=self.attn_drop.p if self.training else 0.0, + ) + + out = rearrange(out, "b h n d -> b n (h d)") + return self.o_dropout(self.o_proj(out)) + + def decode_step(self, x, *, cache, positions: torch.Tensor): + if x.size(1) <= 0: + raise ValueError("MultiHeadAttention.decode_step expects a non-empty input.") + if positions.ndim != 1 or positions.size(0) != x.size(1): + raise ValueError("MultiHeadAttention.decode_step positions must match the decode block length.") + + q = self.q_proj(x) + k = self.k_proj(x) + v = self.v_proj(x) + + q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads) + k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads) + v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads) + q, k = self.q_norm(q), self.k_norm(k) + block_len = q.size(2) + + if self.rotary_bias: + rotary_emb = self.rotary(positions) + q = apply_rotary_pos_emb(rotary_emb, q) + k = apply_rotary_pos_emb(rotary_emb, k) + + cached_k, cached_v = cache + cached_k.index_copy_(2, positions, k) + cached_v.index_copy_(2, positions, v) + + cache_capacity = cached_k.size(2) + key_positions = torch.arange( + cache_capacity, + device=x.device, + dtype=torch.long, + ).unsqueeze(0) + query_positions = positions.unsqueeze(1) + causal_mask = key_positions <= query_positions + valid_mask = key_positions <= positions[-1] + attn_bias = torch.zeros( + q.size(0), + self.num_heads, + block_len, + cache_capacity, + dtype=q.dtype, + device=q.device, + ) + attn_bias.masked_fill_( + (causal_mask & valid_mask).unsqueeze(0).unsqueeze(0).logical_not(), + float("-inf"), + ) + + out = F.scaled_dot_product_attention( + q, + cached_k, + cached_v, + attn_mask=attn_bias, + dropout_p=self.attn_drop.p if self.training else 0.0, + ) + out = rearrange(out, "b h n d -> b n (h d)") + return self.o_dropout(self.o_proj(out)), cache + + +# ============================================================================ +# DiT flow-matching head (adapted from modules/backbone/dit.py) +# ---------------------------------------------------------------------------- +# DiT.__init__ takes a ``transformer_config`` object exposing ``.to_dict()``, +# ``.hidden_size`` and ``.num_layers``. Upstream constructs it from a pydantic +# config class higher up the stack; vLLM-Omni callers can pass anything that +# satisfies the same minimal interface (e.g. a stdlib @dataclass with a +# ``to_dict`` method or SimpleNamespace). See dots_tts_talker.py for the +# concrete soar-checkpoint config factory. +# ============================================================================ + + +def modulate(x, shift, scale, **_kwargs): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + return self.mlp(t_freq) + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size, output_size): + super().__init__() + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 2 * hidden_size, bias=True), + ) + self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-5) + self.linear = nn.Linear(hidden_size, output_size, bias=True) + + def forward(self, x, c, **_kwargs): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm(x), shift, scale) + return self.linear(x) + + +class DiTBlock(nn.Module): + def __init__( + self, + attention: nn.Module, + ffn: nn.Module, + hidden_size: int = 1024, + modulation: bool = False, + eps: float = 1e-5, + **_kwargs, + ): + super().__init__() + self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps) + self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps) + self.attn = attention + self.ffn = ffn + self.modulation = modulation + if modulation: + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 6 * hidden_size, bias=True), + ) + + def forward(self, x, condition=None, mask=None, **kwargs): + if condition is None: + assert not self.modulation, "Without global condition, must set modulation to False" + else: + assert self.modulation, "With global condition, must set modulation to True" + shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(condition).chunk( + 6, dim=1 + ) + + if condition is not None: + pack_indices = kwargs.get("pack_indices") + if pack_indices is not None: + gate_attn = gate_attn[pack_indices] + gate_ffn = gate_ffn[pack_indices] + else: + gate_attn = gate_attn.unsqueeze(1) + gate_ffn = gate_ffn.unsqueeze(1) + + if condition is not None: + x = x + gate_attn * self.attn( + modulate(self.norm1(x), shift_attn, scale_attn, **kwargs), + mask=mask, + **kwargs, + ) + else: + x = x + self.attn(self.norm1(x), mask=mask, **kwargs) + + if condition is not None: + x = x + gate_ffn * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn, **kwargs)) + else: + x = x + self.ffn(self.norm2(x), mask=mask) + return x + + +class DiT(nn.Module): + def __init__( + self, + in_dim, + out_dim, + transformer_config, + *, + mode: str = "flow_matching", + ): + super().__init__() + if mode not in {"flow_matching", "meanflow"}: + raise ValueError(f"DiT mode must be 'flow_matching' or 'meanflow', got {mode!r}.") + + transformer_kwargs = transformer_config.to_dict() + model_dim = transformer_config.hidden_size + self.mode = mode + self.num_layers = transformer_config.num_layers + + self.input_layer = nn.Linear(in_dim, model_dim) + self.time_embedder = TimestepEmbedder(model_dim) + if mode == "meanflow": + self.duration_embedder = TimestepEmbedder(model_dim) + + self.blocks = nn.ModuleList() + for i in range(self.num_layers): + attn_block = MultiHeadAttention(**transformer_kwargs, name=f"layer_{i}") + ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **transformer_kwargs) + self.blocks.append(DiTBlock(attention=attn_block, ffn=ffn_block, **transformer_kwargs)) + + self.output_layer = FinalLayer(model_dim, out_dim) + self.initialize_weights() + + def initialize_weights(self): + def _basic_init(module): + if isinstance(module, nn.Linear): + torch.nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + self.apply(_basic_init) + + nn.init.normal_(self.time_embedder.mlp[0].weight, std=0.02) + nn.init.normal_(self.time_embedder.mlp[2].weight, std=0.02) + + for block in self.blocks: + if hasattr(block, "adaLN_modulation"): + nn.init.constant_(block.adaLN_modulation[-1].weight, 0) + nn.init.constant_(block.adaLN_modulation[-1].bias, 0) + + nn.init.constant_(self.output_layer.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.output_layer.adaLN_modulation[-1].bias, 0) + nn.init.constant_(self.output_layer.linear.weight, 0) + nn.init.constant_(self.output_layer.linear.bias, 0) + + def forward( + self, + x, + timesteps, + duration: torch.Tensor | None = None, + mask=None, + attn_mask=None, + g_cond: torch.Tensor | None = None, + **kwargs, + ): + t = self.time_embedder(timesteps) + c = t + duration_embedder = getattr(self, "duration_embedder", None) + if duration_embedder is not None and duration is not None: + c = c + duration_embedder(duration) + if g_cond is not None: + c = c + g_cond + + x = self.input_layer(x) + for block in self.blocks: + x = block(x, c, mask=attn_mask, **kwargs) + return self.output_layer(x, c, **kwargs) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py new file mode 100644 index 000000000..39d1a4973 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_patch_encoder.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/backbone +# +# Single upstream source (preserved as-is, inference behaviour unchanged; +# only the internal `from dots_tts.modules.backbone.layers import ...` line +# was folded to point at our sibling vendor files): +# - modules/backbone/semantic_encoder.py -> SemanticEncoderDecodeState, +# TransformerEncoderLayer, +# SuperviseEncoder, +# VAESemanticEncoder +# +# Building blocks (Conv1d / Mlp / MultiHeadAttention) are imported from the +# sibling vendor files where they already live: +# - Conv1d <- dots_tts_vocoder.py +# - Mlp <- dots_tts_dit.py +# - MultiHeadAttention <- dots_tts_dit.py +# +# Nothing else in this file diverges from upstream behaviour. + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import ( + Mlp, + MultiHeadAttention, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import Conv1d + + +@dataclass +class SemanticEncoderDecodeState: + conv_tail: torch.Tensor + layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...] + seq_len: int + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + hidden_size, + num_heads=16, + ffn_hidden_size=4096, + attn_dropout=0.0, + ffn_dropout=0.0, + norm_layer="LayerNorm", + **kwargs, + ): + super().__init__() + self.attn = MultiHeadAttention( + hidden_size, + num_heads, + attn_drop=attn_dropout, + norm_layer=norm_layer, + **kwargs, + ) + norm_cls = getattr(nn, norm_layer) + self.attn_norm = norm_cls(hidden_size) + self.ffn = Mlp(hidden_size, ffn_hidden_size, dropout=ffn_dropout, act_layer=nn.SiLU) + self.ffn_norm = norm_cls(hidden_size) + self.hidden_size = hidden_size + + def _build_causal_mask(self, T: int, device): # noqa: N803 + return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device)) + + def _build_padding_mask(self, x_lens, max_len: int, device): + B = x_lens.size(0) + positions = torch.arange(max_len, device=device).unsqueeze(0).expand(B, -1) + return positions < x_lens.unsqueeze(1) + + def _fuse_attn_mask(self, causal_mask, padding_mask): + if causal_mask is None and padding_mask is None: + return None + if causal_mask is None: + row = padding_mask.unsqueeze(2) + col = padding_mask.unsqueeze(1) + return row & col + if padding_mask is None: + return causal_mask.unsqueeze(0) + + _B, _T = padding_mask.shape + causal = causal_mask.unsqueeze(0) + row = padding_mask.unsqueeze(2) + col = padding_mask.unsqueeze(1) + pad_2d = row & col + return causal & pad_2d + + def forward( + self, + x, + x_lens=None, + causal=True, + ): + _B, T, C = x.shape + assert self.hidden_size == C + device = x.device + + causal_mask = self._build_causal_mask(T, device) if causal else None + if x_lens is not None: + padding_mask = self._build_padding_mask(x_lens, T, device) + else: + padding_mask = None + fused_mask = self._fuse_attn_mask(causal_mask, padding_mask) + + h = self.attn_norm(x) + h = self.attn( + q=h, + mask=fused_mask, + ) + x = x + h + + h = self.ffn_norm(x) + h = self.ffn(h) + return x + h + + def decode_step( + self, + x, + *, + cache: tuple[torch.Tensor, torch.Tensor], + positions: torch.Tensor, + ): + if x.size(1) <= 0: + raise ValueError("TransformerEncoderLayer.decode_step expects a non-empty input.") + + h = self.attn_norm(x) + h, cache = self.attn.decode_step(h, cache=cache, positions=positions) + x = x + h + + h = self.ffn_norm(x) + h = self.ffn(h) + return x + h, cache + + +class SuperviseEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.get("hidden_size", 1024) + self.layers = nn.ModuleList( + [ + TransformerEncoderLayer( + hidden_size=self.hidden_size, + num_heads=config.get("num_heads", 16), + ffn_hidden_size=config.get("ffn_hidden_size", 4096), + norm_layer=config.get("norm_layer", "LayerNorm"), + ) + for _ in range(config.get("num_layers", 6)) + ] + ) + self.causal = config.get("causal", False) + + def forward(self, x, x_lens=None): + batch_size, seq_len, _ = x.shape + if x_lens is None: + x_lens = torch.full((batch_size,), seq_len, device=x.device, dtype=torch.long) + for layer in self.layers: + x = layer(x, x_lens=x_lens, causal=self.causal) + return x + + def init_decode_state( + self, + *, + batch_size: int, + max_seq_len: int, + device: torch.device, + dtype: torch.dtype, + ): + layer_caches = [] + for layer in self.layers: + cache_shape = ( + batch_size, + layer.attn.num_heads, + max_seq_len, + layer.attn.head_dim, + ) + layer_caches.append( + ( + torch.zeros(cache_shape, dtype=dtype, device=device), + torch.zeros(cache_shape, dtype=dtype, device=device), + ) + ) + return tuple(layer_caches) + + def reset_decode_state( + self, + layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...], + ) -> None: + if len(layer_caches) != len(self.layers): + raise ValueError("Layer cache count does not match encoder depth.") + for key_cache, value_cache in layer_caches: + key_cache.zero_() + value_cache.zero_() + + def decode_step(self, x, *, layer_caches, positions: torch.Tensor): + if len(layer_caches) != len(self.layers): + raise ValueError("Layer cache count does not match encoder depth.") + + for layer, cache in zip(self.layers, layer_caches, strict=True): + x, _ = layer.decode_step(x, cache=cache, positions=positions) + return x + + +class VAESemanticEncoder(nn.Module): + def __init__(self, in_dim, out_dim, config): + super().__init__() + in_ds_rate = 2 + self.patch_size = int(config.patch_size) + self.in_ds_rate = in_ds_rate + self.ds_proj = Conv1d(in_dim, in_dim, kernel_size=in_ds_rate, stride=in_ds_rate, causal=True) + self.in_proj = nn.Linear(in_dim, config.PatchEncoder.hidden_size) + self.encoder = SuperviseEncoder(config.PatchEncoder) + self.out_ds_rate = self.patch_size // in_ds_rate + self.out_proj = nn.Linear(config.PatchEncoder.hidden_size * self.out_ds_rate, out_dim) + + def forward(self, x, x_lens=None): + x = self._downsample(x) + x = self.in_proj(x) + z = self.encoder(x, x_lens=x_lens) + return self._project_embeddings(z) + + def init_decode_state( + self, + *, + max_audio_patch_count: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> SemanticEncoderDecodeState: + return SemanticEncoderDecodeState( + conv_tail=torch.zeros( + (batch_size, self.ds_proj.in_channels, self.ds_proj.left_padding), + dtype=dtype, + device=device, + ), + layer_caches=self.encoder.init_decode_state( + batch_size=batch_size, + max_seq_len=max_audio_patch_count * self.out_ds_rate, + device=device, + dtype=dtype, + ), + seq_len=0, + ) + + def reset_decode_state(self, state: SemanticEncoderDecodeState) -> None: + state.conv_tail.zero_() + self.encoder.reset_decode_state(state.layer_caches) + state.seq_len = 0 + + def prefill( + self, + x, + state: SemanticEncoderDecodeState, + ) -> tuple[torch.Tensor, SemanticEncoderDecodeState]: + if x.ndim != 3: + raise ValueError(f"VAESemanticEncoder.prefill expects rank-3 input, got {tuple(x.shape)}.") + if x.size(1) % self.patch_size != 0: + raise ValueError(f"Prompt latent length {x.size(1)} must be divisible by patch_size={self.patch_size}.") + + if x.size(1) == 0: + return ( + x.new_zeros((x.size(0), 0, self.out_proj.out_features)), + state, + ) + if state.conv_tail.size(0) != x.size(0): + raise ValueError("VAESemanticEncoder.prefill batch size does not match decode state.") + + step_inputs = self.in_proj(self._downsample(x)) + expected_token_count = (x.size(1) // self.patch_size) * self.out_ds_rate + if step_inputs.size(1) != expected_token_count: + raise RuntimeError( + "Patch encoder prefill produced an unexpected token count: " + f"expected={expected_token_count} actual={step_inputs.size(1)}." + ) + + current_seq_len = state.seq_len + next_seq_len = current_seq_len + step_inputs.size(1) + cache_capacity = state.layer_caches[0][0].size(2) + if next_seq_len > cache_capacity: + raise ValueError( + "Patch encoder prefill exceeds decode-state capacity: " + f"required={next_seq_len} capacity={cache_capacity}." + ) + + positions = torch.arange(step_inputs.size(1), device=x.device, dtype=torch.long) + current_seq_len + encoded = self.encoder.decode_step( + step_inputs, + layer_caches=state.layer_caches, + positions=positions, + ) + embedding = self._project_embeddings(encoded) + raw = x.transpose(1, 2) + state.conv_tail.copy_(raw[..., -self.ds_proj.left_padding :]) + state.seq_len = next_seq_len + return embedding, state + + def decode_patch( + self, + latent_patch, + conv_tail: torch.Tensor, + layer_caches: tuple[tuple[torch.Tensor, torch.Tensor], ...], + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if latent_patch.ndim != 3: + raise ValueError(f"VAESemanticEncoder.decode_patch expects rank-3 input, got {tuple(latent_patch.shape)}.") + if latent_patch.size(1) != self.patch_size: + raise ValueError(f"decode_patch expects patch length {self.patch_size}, got {latent_patch.size(1)}.") + if positions.ndim != 1 or positions.size(0) != self.out_ds_rate: + raise ValueError("decode_patch positions must be a rank-1 tensor matching out_ds_rate.") + + step_inputs, conv_tail = self._downsample_step( + latent_patch, + conv_tail=conv_tail, + ) + if step_inputs.size(1) != self.out_ds_rate: + raise RuntimeError(f"Downsample step produced {step_inputs.size(1)} tokens, expected {self.out_ds_rate}.") + + encoded = self.encoder.decode_step( + step_inputs, + layer_caches=layer_caches, + positions=positions, + ) + embedding = self._project_embeddings(encoded) + return embedding, conv_tail + + def _downsample(self, x): + return self.ds_proj(x.transpose(1, 2)).transpose(1, 2) + + def _project_embeddings(self, z): + if self.out_ds_rate > 1: + z = rearrange(z, "b (s d) h -> b s (d h)", d=self.out_ds_rate) + return self.out_proj(z) + + def _downsample_step(self, latent_patch, *, conv_tail): + raw = latent_patch.transpose(1, 2) + conv_input = torch.cat([conv_tail, raw], dim=-1) + + projected = F.conv1d( + conv_input, + self.ds_proj.weight, + self.ds_proj.bias, + stride=self.ds_proj.stride[0], + padding=0, + dilation=self.ds_proj.dilation[0], + groups=self.ds_proj.groups, + ).transpose(1, 2) + new_conv_tail = raw[..., -self.ds_proj.left_padding :] + return self.in_proj(projected), new_conv_tail diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py new file mode 100644 index 000000000..a2c4ac9f9 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_prompt.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dots.tts prompt builder for vllm-omni's Omni engine + serving layer. + +Mirrors upstream ``build_generation_schedule`` (rednote-hilab/dots.tts @ +a393d2e data/pipelines/tokenizing.py:169) for the +no-prompt-audio + audio-placeholder path — text-only zero-shot TTS. + +V1 returns real Qwen2 token IDs (not placeholders ``[1] * N`` like +voxcpm2's builder). Reason: our talker's preprocess() runs +``self.model.embed_tokens(input_ids)`` directly, so real IDs flow +through naturally. voxcpm2 uses placeholders because its prefill mixes +reference-audio embeddings (no token IDs); dots.tts V1 zero-shot is +purely text + audio_gen_start, all real IDs. + +Voice clone (reference audio + speaker x-vector) is intentionally not +wired here — V1 is zero-shot only. Add reference_audio / +voice_profile / prompt_text branches in a follow-up (mirror voxcpm2 +``build_voxcpm2_prompt``). +""" + +from __future__ import annotations + +from typing import Any + +from vllm.inputs import tokens_input + +# Upstream rednote-hilab/dots.tts @ a393d2e +# data/pipelines/tts_pipeline.py:17-18 +_TTS_TEXT_PREFIX = "[文本]" +_TTS_AUDIO_PREFIX = "[文本对应语音]" +_AUDIO_GEN_START_TOKEN = "<|audio_gen_start|>" + + +def build_dots_tts_prompt( + tokenizer: Any, + text: str, +) -> dict[str, Any]: + """Build a dots.tts prefill prompt dict for ``Omni.generate()``. + + Returns ``{"prompt_token_ids": [...], "additional_information": {...}}`` + where ``prompt_token_ids`` is the real Qwen2 token sequence: + + [text_tokens of "[文本][文本对应语音]", audio_gen_start_id] + + talker's preprocess() embeds these into ``hidden_states`` for the AR + base LM. Decode then drives via ``state.curr_embed_for_next`` (DiT + → patch_encoder AR loopback). + """ + if not isinstance(text, str) or not text.strip(): + raise ValueError(f"text must be a non-empty string, got {text!r}") + + prompt_text = f"{_TTS_TEXT_PREFIX}{text}{_TTS_AUDIO_PREFIX}" + text_ids = tokenizer.encode(prompt_text, add_special_tokens=False) + + audio_gen_start_id = tokenizer.convert_tokens_to_ids(_AUDIO_GEN_START_TOKEN) + if not isinstance(audio_gen_start_id, int) or audio_gen_start_id == tokenizer.unk_token_id: + raise ValueError( + f"Tokenizer does not know {_AUDIO_GEN_START_TOKEN!r} " + f"(got id={audio_gen_start_id!r}). Did you load the " + f"dots.tts tokenizer (with added_tokens.json)?" + ) + prompt_token_ids = list(text_ids) + [audio_gen_start_id] + + prompt = tokens_input(prompt_token_ids=prompt_token_ids) + prompt["additional_information"] = { + # Reserved for future voice clone (ref_audio / speaker x-vector). + # V1 zero-shot leaves this empty; preprocess() ignores it. + } + return prompt diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py new file mode 100644 index 000000000..09961d0ab --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_speaker_encoder.py @@ -0,0 +1,732 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Portions Copyright 3D-Speaker +# (https://github.com/alibaba-damo-academy/3D-Speaker), Apache-2.0. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules/speaker +# (whose CAM++ implementation is itself adapted from 3D-Speaker). +# +# Merged from four upstream files (preserved as-is, inference behaviour +# unchanged; only internal cross-module imports were folded inline): +# - modules/speaker/encoder.py -> SpeakerXVectorFeatures +# - modules/speaker/fbank.py -> fbank constants + extract_speaker_fbank +# - modules/speaker/campplus.py -> FCM, CAMPPlus +# - modules/speaker/campplus_layers.py -> get_nonlinear, statistics_pooling, +# StatsPool, TDNNLayer, CAMLayer, +# CAMDenseTDNNLayer, CAMDenseTDNNBlock, +# TransitLayer, DenseLayer, BasicResBlock +# +# Plus two tiny helpers folded from utils/audio.py (used by +# extract_speaker_fbank): +# - high_quality_resample, extract_fbank +# +# Nothing in this file diverges from upstream behaviour. + +from __future__ import annotations + +import math +import random +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +import torchaudio +import torchaudio.compliance.kaldi as Kaldi +import torchaudio.functional as AF +from torch.nn.utils.rnn import pad_sequence + +# ============================================================================ +# Constants (from modules/speaker/fbank.py) +# ============================================================================ + +_SPEAKER_FBANK_SAMPLE_RATE = 16000 +_SPEAKER_FBANK_N_MELS = 80 +_SPEAKER_FBANK_MEAN_NORM = True +_SPEAKER_FBANK_DITHER = 0.0 + + +# ============================================================================ +# Audio helpers (folded from utils/audio.py — used by extract_speaker_fbank) +# ============================================================================ + + +def high_quality_resample(x, orig_sr, target_sr): + return AF.resample( + x, + orig_freq=orig_sr, + new_freq=target_sr, + lowpass_filter_width=64, + rolloff=0.95, + resampling_method="sinc_interp_kaiser", + ) + + +def extract_fbank( + waveform: torch.Tensor, + *, + sample_rate: int, + n_mels: int, + dither: float = 0.0, + mean_norm: bool = False, +) -> torch.Tensor: + if waveform.ndim == 1: + feature_input = waveform.unsqueeze(0) + elif waveform.ndim == 2: + feature_input = waveform if waveform.size(0) == 1 else waveform[0:1, :] + else: + raise ValueError(f"FBank expects a 1D or 2D waveform, got shape {tuple(waveform.shape)}.") + features = Kaldi.fbank( + feature_input, + num_mel_bins=n_mels, + sample_frequency=sample_rate, + dither=dither, + ) + if mean_norm: + features = features - features.mean(dim=0, keepdim=True) + return features + + +# ============================================================================ +# Speaker fbank entry (from modules/speaker/fbank.py) +# ============================================================================ + + +def extract_speaker_fbank( + waveform: torch.Tensor, + *, + sample_rate: int, +) -> torch.Tensor: + feature_input = waveform + if sample_rate != _SPEAKER_FBANK_SAMPLE_RATE: + feature_input = high_quality_resample( + waveform, + orig_sr=sample_rate, + target_sr=_SPEAKER_FBANK_SAMPLE_RATE, + ) + return extract_fbank( + feature_input, + sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, + n_mels=_SPEAKER_FBANK_N_MELS, + dither=_SPEAKER_FBANK_DITHER, + mean_norm=_SPEAKER_FBANK_MEAN_NORM, + ) + + +# ============================================================================ +# CAM++ building blocks (from modules/speaker/campplus_layers.py) +# ---------------------------------------------------------------------------- +# 3D-Speaker, Apache-2.0. +# ============================================================================ + + +def get_nonlinear(config_str, channels): + nonlinear = nn.Sequential() + for name in config_str.split("-"): + if name == "relu": + nonlinear.add_module("relu", nn.ReLU(inplace=True)) + elif name == "prelu": + nonlinear.add_module("prelu", nn.PReLU(channels)) + elif name == "batchnorm": + nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels)) + elif name == "batchnorm_": + nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False)) + else: + raise ValueError(f"Unexpected module ({name}).") + return nonlinear + + +def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, _eps=1e-2): + mean = x.mean(dim=dim) + std = x.std(dim=dim, unbiased=unbiased) + stats = torch.cat([mean, std], dim=-1) + if keepdim: + stats = stats.unsqueeze(dim=dim) + return stats + + +class StatsPool(nn.Module): + def forward(self, x): + return statistics_pooling(x) + + +class TDNNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + bias=False, + config_str="batchnorm-relu", + ): + super().__init__() + if padding < 0: + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.linear = nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + x = self.linear(x) + return self.nonlinear(x) + + +class CAMLayer(nn.Module): + def __init__( + self, + bn_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + bias, + reduction=2, + ): + super().__init__() + self.linear_local = nn.Conv1d( + bn_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1) + self.relu = nn.ReLU(inplace=True) + self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + y = self.linear_local(x) + context = x.mean(-1, keepdim=True) + self.seg_pooling(x) + context = self.relu(self.linear1(context)) + m = self.sigmoid(self.linear2(context)) + return y * m + + def seg_pooling(self, x, seg_len=100, stype="avg"): + if stype == "avg": + seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + elif stype == "max": + seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + else: + raise ValueError("Wrong segment pooling type.") + shape = seg.shape + seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1) + return seg[..., : x.shape[-1]] + + +class CAMDenseTDNNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.memory_efficient = memory_efficient + self.nonlinear1 = get_nonlinear(config_str, in_channels) + self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False) + self.nonlinear2 = get_nonlinear(config_str, bn_channels) + self.cam_layer = CAMLayer( + bn_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + + def bn_function(self, x): + return self.linear1(self.nonlinear1(x)) + + def forward(self, x): + if self.training and self.memory_efficient: + x = cp.checkpoint(self.bn_function, x) + else: + x = self.bn_function(x) + return self.cam_layer(self.nonlinear2(x)) + + +class CAMDenseTDNNBlock(nn.ModuleList): + def __init__( + self, + num_layers, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + for i in range(num_layers): + layer = CAMDenseTDNNLayer( + in_channels=in_channels + i * out_channels, + out_channels=out_channels, + bn_channels=bn_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + bias=bias, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.add_module(f"tdnnd{i + 1}", layer) + + def forward(self, x): + for layer in self: + x = torch.cat([x, layer(x)], dim=1) + return x + + +class TransitLayer(nn.Module): + def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"): + super().__init__() + self.nonlinear = get_nonlinear(config_str, in_channels) + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + + def forward(self, x): + x = self.nonlinear(x) + return self.linear(x) + + +class DenseLayer(nn.Module): + def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"): + super().__init__() + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + if len(x.shape) == 2: + x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1) + else: + x = self.linear(x) + return self.nonlinear(x) + + +class BasicResBlock(nn.Module): + expansion = 1 + + def __init__(self, in_planes, planes, stride=1): + super().__init__() + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d( + in_planes, + self.expansion * planes, + kernel_size=1, + stride=(stride, 1), + bias=False, + ), + nn.BatchNorm2d(self.expansion * planes), + ) + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out += self.shortcut(x) + return F.relu(out) + + +# ============================================================================ +# CAM++ model (from modules/speaker/campplus.py) +# ---------------------------------------------------------------------------- +# 3D-Speaker, Apache-2.0. +# ============================================================================ + + +class FCM(nn.Module): + def __init__( + self, + block=BasicResBlock, + num_blocks=(2, 2), + m_channels=32, + feat_dim=_SPEAKER_FBANK_N_MELS, + ): + super().__init__() + self.in_planes = m_channels + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(m_channels) + + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2) + self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2) + + self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(m_channels) + self.out_channels = m_channels * (feat_dim // 8) + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + x = x.unsqueeze(1) + out = F.relu(self.bn1(self.conv1(x))) + out = self.layer1(out) + out = self.layer2(out) + out = F.relu(self.bn2(self.conv2(out))) + + shape = out.shape + return out.reshape(shape[0], shape[1] * shape[2], shape[3]) + + +class CAMPPlus(nn.Module): + _TDNN_KERNEL_SIZE = 5 + _TDNN_STRIDE = 2 + _TDNN_PADDING = 2 + + def __init__( + self, + feat_dim=_SPEAKER_FBANK_N_MELS, + embedding_size=512, + growth_rate=32, + bn_size=4, + init_channels=128, + config_str="batchnorm-relu", + memory_efficient=True, + ): + super().__init__() + + self.head = FCM(feat_dim=feat_dim) + channels = self.head.out_channels + + self.xvector = nn.Sequential( + OrderedDict( + [ + ( + "tdnn", + TDNNLayer( + channels, + init_channels, + self._TDNN_KERNEL_SIZE, + stride=self._TDNN_STRIDE, + dilation=1, + padding=-1, + config_str=config_str, + ), + ), + ] + ) + ) + channels = init_channels + for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2), strict=True)): + block = CAMDenseTDNNBlock( + num_layers=num_layers, + in_channels=channels, + out_channels=growth_rate, + bn_channels=bn_size * growth_rate, + kernel_size=kernel_size, + dilation=dilation, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.xvector.add_module(f"block{i + 1}", block) + channels = channels + num_layers * growth_rate + self.xvector.add_module( + f"transit{i + 1}", + TransitLayer(channels, channels // 2, bias=False, config_str=config_str), + ) + channels //= 2 + + self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels)) + + self.xvector.add_module("stats", StatsPool()) + self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")) + + for m in self.modules(): + if isinstance(m, (nn.Conv1d, nn.Linear)): + nn.init.kaiming_normal_(m.weight.data) + if m.bias is not None: + nn.init.zeros_(m.bias) + + @staticmethod + def _conv_output_lengths(lengths, kernel_size, stride=1, padding=0, dilation=1): + return ( + torch.div( + lengths + 2 * padding - dilation * (kernel_size - 1) - 1, + stride, + rounding_mode="floor", + ) + + 1 + ) + + @staticmethod + def _make_length_mask(lengths, max_len, device): + lengths = lengths.to(device=device, dtype=torch.long).clamp(min=0, max=max_len) + return torch.arange(max_len, device=device).unsqueeze(0) < lengths.unsqueeze(1) + + def _masked_stats_pooling(self, x, lengths, unbiased=True, eps=1e-2): + lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1, max=x.size(-1)) + mask = self._make_length_mask(lengths, x.size(-1), x.device).unsqueeze(1) + mask = mask.to(dtype=x.dtype) + + denom = lengths.to(dtype=x.dtype).view(-1, 1).clamp_min(1.0) + mean = (x * mask).sum(dim=-1) / denom + + centered = (x - mean.unsqueeze(-1)) * mask + var_denom = (lengths - 1).clamp_min(1).to(dtype=x.dtype).view(-1, 1) if unbiased else denom + var = centered.pow(2).sum(dim=-1) / var_denom + std = torch.sqrt(var.clamp_min(eps)) + return torch.cat([mean, std], dim=1) + + def forward(self, x, lengths=None): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + x = self.head(x) + if lengths is not None: + lengths = lengths.to(device=x.device, dtype=torch.long).clamp(min=1) + + for name, module in self.xvector.named_children(): + if name == "stats": + x = self._masked_stats_pooling(x, lengths) if lengths is not None else module(x) + continue + + x = module(x) + if name == "tdnn" and lengths is not None: + lengths = self._conv_output_lengths( + lengths, + kernel_size=self._TDNN_KERNEL_SIZE, + stride=self._TDNN_STRIDE, + padding=self._TDNN_PADDING, + ) + + return x + + +# ============================================================================ +# SpeakerXVectorFeatures wrapper (from modules/speaker/encoder.py) +# ============================================================================ + + +class SpeakerXVectorFeatures(nn.Module): + """Speaker embedding extractor based on 3D-Speaker CAM++.""" + + def __init__( + self, + sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, + campplus_embedding_size=512, + max_audio_seconds=10.0, + ): + super().__init__() + + self.sample_rate = sample_rate + self.max_audio_seconds = float(max_audio_seconds) + self.model = CAMPPlus( + feat_dim=_SPEAKER_FBANK_N_MELS, + embedding_size=campplus_embedding_size, + ) + self.resample = None + if self.sample_rate != _SPEAKER_FBANK_SAMPLE_RATE: + self.resample = torchaudio.transforms.Resample( + orig_freq=sample_rate, + new_freq=_SPEAKER_FBANK_SAMPLE_RATE, + ) + + for param in self.model.parameters(): + param.requires_grad = False + + @staticmethod + def _normalize_lengths(lengths, batch_size, max_length, device, *, min_length): + if lengths is None: + return torch.full( + (batch_size,), + max_length, + device=device, + dtype=torch.long, + ) + return lengths.to(device=device, dtype=torch.long).clamp( + min=min_length, + max=max_length, + ) + + def _crop_audio(self, audio, audio_lengths=None): + original_lengths = self._normalize_lengths( + audio_lengths, + audio.size(0), + audio.size(-1), + audio.device, + min_length=0, + ) + if self.max_audio_seconds <= 0: + return audio, original_lengths, original_lengths, torch.zeros_like(original_lengths) + + max_input_length = round(self.sample_rate * self.max_audio_seconds) + cropped_audio = [] + cropped_lengths = [] + starts = [] + + for index, total_length_tensor in enumerate(original_lengths): + total_length = int(total_length_tensor.item()) + cropped_length = min(total_length, max_input_length) + start = random.randint(0, total_length - cropped_length) if total_length > cropped_length else 0 + cropped_audio.append(audio[index, start : start + cropped_length]) + cropped_lengths.append(cropped_length) + starts.append(start) + + return ( + pad_sequence( + cropped_audio, + batch_first=True, + padding_value=0.0, + ), + original_lengths, + torch.tensor( + cropped_lengths, + device=audio.device, + dtype=torch.long, + ), + torch.tensor(starts, device=audio.device, dtype=torch.long), + ) + + def _crop_fbank( + self, + fbank, + fbank_lengths, + original_audio_lengths, + cropped_audio_lengths, + starts, + ): + original_fbank_lengths = self._normalize_lengths( + fbank_lengths, + fbank.size(0), + fbank.size(1), + fbank.device, + min_length=1, + ) + cropped_fbank = [] + cropped_fbank_lengths = [] + + for index, total_feat_length_tensor in enumerate(original_fbank_lengths): + total_audio_length = int(original_audio_lengths[index].item()) + total_feat_length = int(total_feat_length_tensor.item()) + start_audio = int(starts[index].item()) + end_audio = start_audio + int(cropped_audio_lengths[index].item()) + + if total_audio_length > 0: + start_feat = math.floor(start_audio * total_feat_length / total_audio_length) + end_feat = math.ceil(end_audio * total_feat_length / total_audio_length) + else: + start_feat = 0 + end_feat = 1 + + start_feat = min(start_feat, total_feat_length - 1) + end_feat = min(max(end_feat, start_feat + 1), total_feat_length) + cropped_fbank.append(fbank[index, start_feat:end_feat]) + cropped_fbank_lengths.append(end_feat - start_feat) + + return pad_sequence( + cropped_fbank, + batch_first=True, + padding_value=0.0, + ), torch.tensor( + cropped_fbank_lengths, + device=fbank.device, + dtype=torch.long, + ) + + def _extract_fbank_batch(self, audio, audio_lengths): + if self.resample is not None: + audio = self.resample(audio) + audio_lengths = torch.ceil(audio_lengths.float() * (_SPEAKER_FBANK_SAMPLE_RATE / self.sample_rate)).long() + + audio_cpu = audio.detach().cpu() + features = [] + + for index, valid_length_tensor in enumerate(audio_lengths): + valid_length = int(valid_length_tensor.item()) + waveform = audio_cpu[index, :valid_length] + if waveform.numel() == 0: + waveform = audio_cpu.new_zeros(1) + features.append( + extract_speaker_fbank( + waveform, + sample_rate=_SPEAKER_FBANK_SAMPLE_RATE, + ) + ) + + fbank_lengths = torch.tensor( + [feature.size(0) for feature in features], + device=audio.device, + dtype=torch.long, + ) + fbank = pad_sequence( + features, + batch_first=True, + padding_value=0.0, + ).to(device=audio.device, dtype=audio.dtype) + return fbank, fbank_lengths + + @torch.no_grad() + @torch.autocast(enabled=False, device_type="cuda") + def forward(self, audio, audio_lengths=None, fbank=None, fbank_lengths=None, **_kwargs): + self.model.eval() + audio = audio.float() + if audio.dim() == 3: + if audio.size(1) != 1: + raise ValueError(f"Speaker encoder expects mono audio, got shape {tuple(audio.shape)}.") + audio = audio[:, 0] + elif audio.dim() != 2: + raise ValueError(f"Speaker encoder expects a 2D or 3D audio tensor, got shape {tuple(audio.shape)}.") + + audio, original_audio_lengths, cropped_audio_lengths, starts = self._crop_audio( + audio, + audio_lengths=audio_lengths, + ) + + if fbank is None: + fbank, fbank_lengths = self._extract_fbank_batch( + audio, + cropped_audio_lengths, + ) + else: + if not isinstance(fbank, torch.Tensor): + raise TypeError("Speaker encoder expects `fbank` to be a torch.Tensor.") + if fbank.dim() != 3 or fbank.size(0) != audio.size(0): + raise ValueError( + f"Speaker encoder expects `fbank` with shape (B, T, F) and matching batch size, got {tuple(fbank.shape)}." # noqa: E501 + ) + fbank, fbank_lengths = self._crop_fbank( + fbank.to(device=audio.device, dtype=torch.float32), + fbank_lengths, + original_audio_lengths, + cropped_audio_lengths, + starts, + ) + + return self.model(fbank, lengths=fbank_lengths) diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py new file mode 100644 index 000000000..368a3fd1e --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py @@ -0,0 +1,1426 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dots.tts talker — vLLM-native AR base LM + audio side path. + +Mirrors upstream rednote-hilab/dots.tts (pinned @ a393d2e): + + * Qwen2.5-1.5B base LM runs under vLLM PagedAttention (stock + ``Qwen2Model``); weights come from the checkpoint's ``llm.model.*`` + namespace. ``DotsTTSConfig`` hoists the LM fields to top-level + attributes. + * Everything else is side computation, invisible to the vLLM + scheduler and driven once per step by ``_finish_decode``: + DiT flow-matching head (N-step Euler, fp32 integration state, + deterministic per-request noise) → io_helper denormalize → + patch_encoder AR loopback (produces the next step's LLM input + embedding) + streaming AudioVAE decode (per-request + ``BigVGANStreamState`` sliding window, so every latent patch is + decoded with real left conv context) → 48 kHz wav chunks pushed to + ``_audio_queue`` → eos_proj stop signal via ``_results_queue`` + (threshold 0.8, upstream default; the stop step also drains the + vocoder's 2-frame lookahead tail via ``stream_flush``). + * Cross-step state is isolated per request in ``_RequestState`` + (keyed by request_id); eviction is deferred to the end of + ``forward()`` so a finishing request's audio still drains. + * The CAM++ speaker encoder is constructed and its weights load, but + voice cloning is not exposed yet — generation is zero-shot and the + DiT falls back to its null conditioning (``fm_null_g_cond``). + +Debugging: set ``DOTS_TTS_BETA_TRACE=1`` for an env-gated per-patch +rms/max trace across the full side path (zero overhead when unset). + +Reference implementations: ``vllm_omni/model_executor/models/voxcpm2/`` +(architectural template: vLLM-native base LM + side path + per-request +state) and ``ming_flash_omni`` (AudioVAE weight-loading pattern). +""" + +from __future__ import annotations + +import dataclasses +import hashlib +import os +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.qwen2 import Qwen2Model +from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.dots_tts.dots_tts_dit import DiT +from vllm_omni.model_executor.models.dots_tts.dots_tts_patch_encoder import ( + VAESemanticEncoder, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_speaker_encoder import ( + SpeakerXVectorFeatures, +) +from vllm_omni.model_executor.models.dots_tts.dots_tts_vocoder import ( + AudioVAE, + AudioVAEConfig, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +def _require_config_block(hf_config: Any, name: str) -> dict[str, Any]: + """Fetch a sub-config dict (``vocoder`` / ``DiT`` / ``PatchEncoder``) + from the checkpoint's config.json, failing loudly when absent — + silently falling back to another checkpoint's values would produce + garbage audio, not an error.""" + block = getattr(hf_config, name, None) + if not isinstance(block, dict) or not block: + raise ValueError( + f"dots.tts checkpoint config.json has no '{name}' block — cannot " + "build the audio side path. Expected the upstream dots.tts config " + "layout (e.g. rednote-hilab/dots.tts-soar)." + ) + return block + + +def _build_audio_vae_config(hf_config: Any) -> AudioVAEConfig: + # AudioVAEConfig declares exactly the checkpoint's 17 ``vocoder`` keys, + # so an unknown / missing key raises TypeError instead of guessing. + return AudioVAEConfig(**_require_config_block(hf_config, "vocoder")) + + +# Architecture constants used by the per-request FM workspace math. All +# released dots.tts checkpoints (soar / base / mf) share these values; +# _validate_architecture_constants refuses a checkpoint that disagrees +# instead of silently generating garbage. +_FM_HIDDEN = 1024 # DiT.hidden_size (also LLM↔DiT projection target) +_LATENT_DIM = 128 # AudioVAE.latent_dim +_LATENT_PATCH_SIZE = 4 # config.patch_size — DiT samples 4 latent frames / audio patch +_HIDDEN_PATCH_SIZE = 1 # upstream core.py hardcodes self.hidden_patch_size = 1 +_MAX_AUDIO_PATCHES = 1024 # ~164 s @ 4×1920/48k; bounds per-request FM static buffer +_DIT_NUM_STEPS = int( + os.environ.get("DOTS_TTS_DIT_NUM_STEPS", "10") +) # env-gated; upstream default 10 for fixed-step Euler +_DIT_GUIDANCE_SCALE = 1.2 # upstream default guidance_scale for soar +_DIT_NOISE_SEED = 20260601 # base seed for per-request FM noise (voxcpm2 parity) +_PATCH_ENCODER_OUT_DS_RATE = 2 # patch_size / in_ds_rate = 4 / 2 (VAESemanticEncoder hardcodes in_ds_rate=2) + + +class _IOHelper: + """Latent stats wrapper. Holds (mean, var) over the 128-dim VAE latent + space and normalizes / denormalizes between DiT-internal space (which + operates on normalized latents) and AudioVAE / patch_encoder input space + (which expects raw latents). + + Mirrors upstream IOHelper (rednote-hilab/dots.tts @ a393d2e + models/dots_tts/core.py:684). Plain class (not nn.Module) so it stays + out of state_dict() and load_weights() doesn't have to mute its buffers. + Stats are loaded from latent_stats.pt — independent of safetensors. + """ + + def __init__(self, latent_stats_path: str | None = None) -> None: + if latent_stats_path is None: + self.global_mean: torch.Tensor | None = None + self.global_var: torch.Tensor | None = None + return + stats = torch.load(latent_stats_path, weights_only=False) + self.global_mean = torch.as_tensor(stats["mean"]) + self.global_var = torch.as_tensor(stats["var"]) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + if self.global_mean is None or self.global_var is None: + return x + mean = self.global_mean.to(device=x.device, dtype=x.dtype) + var = self.global_var.to(device=x.device, dtype=x.dtype) + return (x - mean) / torch.sqrt(var) + + def denormalize(self, x: torch.Tensor) -> torch.Tensor: + if self.global_mean is None or self.global_var is None: + return x + mean = self.global_mean.to(device=x.device, dtype=x.dtype) + var = self.global_var.to(device=x.device, dtype=x.dtype) + return x * torch.sqrt(var) + mean + + +# Per-request state container for AR-loop continuity. All cross-step +# side-path state lives here, keyed by request_id, so concurrent requests +# stay fully isolated (the voxcpm2 _RequestState pattern). +@dataclass +class _RequestState: + request_id: str + # AR loopback: patch_encoder's previous decode output, fed to the LM + # as next-step inputs_embeds. Shape: [1, llm_hidden] = [1, 1536]. + # (patch_encoder collapses its internal _PATCH_ENCODER_OUT_DS_RATE + # positions into one LLM token via out_proj; see + # VAESemanticEncoder._project_embeddings.) _finish_decode writes a + # fresh tensor every call. + curr_embed_for_next: torch.Tensor | None = None + # Stop-signal cache: _finish_decode populates precomputed_stop_logits; + # compute_logits drains it. + precomputed_stop_logits: torch.Tensor | None = None + is_stopping: bool = False + prefill_completed: bool = False + # Per-request FM noise counter: draw #n of this request hashes to a + # deterministic Generator seed (see _run_dit_n_step_euler), so outputs + # are reproducible run-to-run and concurrent requests cannot perturb + # each other's noise streams (voxcpm2 _fill_deterministic_cfm_noise). + noise_step: int = 0 + # Per-request FM static workspace (lazy-allocated by + # _initialize_request_fm_state on first _finish_decode call). Sized for + # _MAX_AUDIO_PATCHES × (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) = 1024 × 5 + # = 5120 positions. + fm_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] + fm_cfg_sequence: torch.Tensor | None = None # [1, fm_capacity, _FM_HIDDEN] + fm_null_g_cond: torch.Tensor | None = None # [1, _FM_HIDDEN] + fm_seq_len: int = 0 + fm_capacity: int = 0 + # Speaker x-vector after _xvec_proj. Zero-shot generation leaves this + # None → DiT falls back to fm_null_g_cond; voice-clone wiring is a + # follow-up. + g_cond: torch.Tensor | None = None + # AR-loop bookkeeping. + # patch_encoder_state holds conv_tail + per-layer KV caches for the + # patch_encoder's streaming decode (mirrors upstream state.patch_encoder_state). + # Lazy-allocated by _run_patch_encoder_loopback on first audio step. + patch_encoder_state: Any | None = None + # Review M2: streaming vocoder state (BigVGANStreamState — LSTM hidden + # + decoder sliding latent window). Gives each patch real left conv + # context; output lags input by decoder.stream_lookahead (2 frames = + # 3840 samples). Lazy-allocated by _run_vocoder_stream_step. + vocoder_stream_state: Any | None = None + + +# DiT transformer hyperparameters from rednote-hilab/dots.tts-soar config.json +# (the ``DiT`` block). Mirrors the shape of an upstream pydantic config: DiT's +# constructor calls ``transformer_config.to_dict()`` and reads ``.hidden_size`` +# / ``.num_layers``, so this @dataclass plus the ``to_dict`` helper satisfies +# that minimal interface without dragging in upstream's config system. +# +# Note: ``attn_dropout`` in the JSON spells differently from MultiHeadAttention's +# ``attn_drop`` kwarg — the value gets absorbed by ``**_kwargs`` and the MHA +# default 0.0 applies. Same numeric outcome for the soar checkpoint (both 0.0), +# so this is benign; flagging it because it would mask non-default values if +# upstream ever raised this field. +@dataclass +class _DiTConfig: + num_layers: int = 18 + num_heads: int = 16 + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + modulation: bool = True + qkv_bias: bool = False + qk_norm: bool = True + attn_dropout: float = 0.0 + dropout: float = 0.0 + norm_layer: str = "RMSNorm" + alibi_bias: bool = False + rotary_bias: bool = True + rotary_theta: float = 10000.0 + + def to_dict(self) -> dict[str, Any]: + return dataclasses.asdict(self) + + +def _build_dit_config(hf_config: Any) -> _DiTConfig: + # _DiTConfig declares exactly the checkpoint's 13 ``DiT`` keys, so an + # unknown key raises TypeError instead of being silently dropped. + return _DiTConfig(**_require_config_block(hf_config, "DiT")) + + +# Patch encoder hyperparameters from rednote-hilab/dots.tts-soar config.json +# (the top-level ``patch_size`` and the ``PatchEncoder`` sub-block). +# +# Two layers of config: +# * Outer (``_PatchEncoderConfig``): VAESemanticEncoder reads ``.patch_size`` +# and forwards ``.PatchEncoder`` to its inner SuperviseEncoder. +# * Inner (``_PatchEncoderInner``): the constructor uses BOTH attribute +# access (``config.PatchEncoder.hidden_size``) and dict-style +# ``config.get("hidden_size", 1024)``. Upstream's pydantic ConfigBase +# supports both natively; we add a tiny ``.get()`` helper to a dataclass. +# +# Upstream's ``_EncoderConfig`` also defines qk_norm / rotary_bias / +# rotary_theta / input_dim, but ``SuperviseEncoder.__init__`` only reads six +# keys (num_layers / num_heads / hidden_size / ffn_hidden_size / norm_layer / +# causal). The rest are swallowed and MultiHeadAttention falls back to its +# own defaults — matching upstream behaviour means *not* plumbing them through. +@dataclass +class _PatchEncoderInner: + num_layers: int = 24 + num_heads: int = 16 + hidden_size: int = 1024 + ffn_hidden_size: int = 4096 + norm_layer: str = "RMSNorm" + causal: bool = True + + def get(self, key: str, default: Any = None) -> Any: + return getattr(self, key, default) + + +@dataclass +class _PatchEncoderConfig: + patch_size: int = 4 + PatchEncoder: _PatchEncoderInner = dataclasses.field(default_factory=_PatchEncoderInner) + + +def _build_patch_encoder_config(hf_config: Any) -> _PatchEncoderConfig: + block = _require_config_block(hf_config, "PatchEncoder") + # SuperviseEncoder reads only the six declared fields; upstream's + # pydantic config swallows the other checkpoint keys, so filtering + # here matches upstream behaviour exactly. + inner_fields = {f.name for f in dataclasses.fields(_PatchEncoderInner)} + inner = _PatchEncoderInner(**{k: v for k, v in block.items() if k in inner_fields}) + return _PatchEncoderConfig( + patch_size=int(hf_config.patch_size), + PatchEncoder=inner, + ) + + +def _validate_architecture_constants(hf_config: Any, dit_config: _DiTConfig) -> None: + """The FM workspace helpers size buffers off module-level constants; + a checkpoint that disagrees would corrupt the DiT conditioning, so + refuse it up front (making these fully config-driven is deferred + until a released checkpoint actually differs).""" + checks = [ + ("latent_dim", getattr(hf_config, "latent_dim", None), _LATENT_DIM), + ("patch_size", getattr(hf_config, "patch_size", None), _LATENT_PATCH_SIZE), + ("DiT.hidden_size", dit_config.hidden_size, _FM_HIDDEN), + ] + mismatches = [f"{name}={got!r} (expected {want})" for name, got, want in checks if got != want] + if mismatches: + raise ValueError( + "dots.tts checkpoint config disagrees with this integration's " + "architecture constants: " + ", ".join(mismatches) + ) + + +class DotsTTSForConditionalGeneration(nn.Module): + """dots.tts AR talker: Qwen2.5 base LM + DiT / AudioVAE / CAM++ side path.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + self.vllm_config = vllm_config + self.config = vllm_config.model_config.hf_config + + # vllm-omni capability flags — gpu_model_runner checks these to decide + # whether to call preprocess() / postprocess() / accept multimodal + # outputs. Without these, the runner bypasses our preprocess hook + # and _pending_requests stays empty (forward never dispatches the + # side-path). Voxcpm2 sets the same three. + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + + self.model = Qwen2Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Side-path module hyperparameters come off the checkpoint's + # config.json blocks (fail loudly if a block is absent); the FM + # workspace constants are cross-checked against the same config. + dit_config = _build_dit_config(self.config) + _validate_architecture_constants(self.config, dit_config) + llm_hidden = self.config.hidden_size + fm_hidden = dit_config.hidden_size + latent_dim = int(self.config.latent_dim) + xvec_dim = int(getattr(self.config, "campplus_embedding_size", 512)) + + self._audio_vae = AudioVAE(_build_audio_vae_config(self.config)) + # Upstream serializes vocoder.safetensors with weight_norm folded into + # plain `weight` on the decoder (encoder kept weight_norm — it's not in + # the synthesis hot path). Match that layout before load_weights runs, + # so checkpoint keys align with our state_dict 1:1. + self._audio_vae.remove_weight_norm() + # AudioVAE / speaker encoder stay fp32 — upstream only casts + # `model.core` to bf16 (runtime.py:81), leaving vocoder + xvector + # extractor in fp32 because their kernels (conv1d / fbank) need + # higher precision. Without this _audio_vae would inherit talker's + # bf16 dtype via vLLM's auto-cast and dtype-mismatch under + # stream_step (input .float() vs bf16 weights). + + # DiT flow-matching head. Dimensions per upstream core.py:101-104: + # in_dim = config.DiT.hidden_size (DiT internal space) + # out_dim = config.latent_dim (AudioVAE input space) + # mode is "flow_matching" for soar/base; "meanflow" only for the mf + # checkpoint (handled by a separate factory in a later step). + self._head = DiT( + in_dim=fm_hidden, + out_dim=latent_dim, + transformer_config=dit_config, + mode="flow_matching", + ) + + # Patch encoder: closes the AR loop by mapping each DiT-emitted audio + # latent patch back into the LLM hidden space, so Qwen2 "hears" what it + # has already generated. Upstream core.py:75 calls this with + # in_dim=latent_dim and out_dim=llm.config.hidden_size. + self._patch_encoder = VAESemanticEncoder( + in_dim=latent_dim, + out_dim=llm_hidden, + config=_build_patch_encoder_config(self.config), + ) + + # Five thin projectors (upstream core.py:82-113). Soar dimensions: + # llm_hidden=1536, fm_hidden=1024, latent_dim=128, xvec_dim=512. + # These are kept as direct attributes (no vendor file) — they are flat + # nn.Linear / nn.Sequential, and live at the top of the checkpoint + # namespace under their own short prefixes (see load_weights). + self._hidden_proj = nn.Linear(llm_hidden, fm_hidden) + self._latent_proj = nn.Linear(latent_dim, fm_hidden) + self._coordinate_proj = nn.Linear(latent_dim, fm_hidden) + self._xvec_proj = nn.Sequential( + nn.Linear(xvec_dim, fm_hidden), + nn.LayerNorm(fm_hidden), + ) + # eos_proj is the stop predictor — feeds LLM hidden through a 2-layer + # MLP and produces a 2-way logit (continue / stop). + self._eos_proj = nn.Sequential( + nn.Linear(llm_hidden, llm_hidden), + nn.SiLU(), + nn.Linear(llm_hidden, 2), + ) + + # CAM++ speaker encoder (3D-Speaker), produces the x-vector from + # reference audio for voice cloning. sample_rate matches AudioVAE's + # input space (48 kHz); the wrapper resamples to 16k internally for + # CAM++. Weights are frozen (upstream does `requires_grad = False`). + self._speaker_encoder = SpeakerXVectorFeatures( + sample_rate=self._audio_vae.sample_rate, + campplus_embedding_size=xvec_dim, + max_audio_seconds=float(getattr(self.config, "xvec_max_audio_seconds", 10.0)), + ) + + # Pin AudioVAE + speaker encoder to fp32 (mirror upstream + # runtime.py:81 where only `model.core` is cast to bf16). vLLM's + # auto-cast would otherwise lower these to bf16, breaking + # stream_step (input .float() vs bf16 weights) and the fbank + # pipeline. + self._audio_vae.float() + self._speaker_encoder.float() + + # Latent stats — independent of safetensors (torch.load on + # latent_stats.pt sitting next to the checkpoint files). Used by + # _finish_decode to bridge between DiT's normalized latent space and + # the raw latent space expected by patch_encoder + AudioVAE. + latent_stats_path = self._resolve_latent_stats_path(vllm_config.model_config.model) + if latent_stats_path is not None: + self._io_helper = _IOHelper(latent_stats_path) + else: + logger.warning( + "latent_stats.pt not resolvable for %s; running with " + "identity normalize/denormalize. Wav will be ~10x quiet " + "and distorted (no sqrt(var) denormalize).", + vllm_config.model_config.model, + ) + self._io_helper = _IOHelper() + + # (β) diagnostic: env-gated per-patch rms/max trace across the full + # side path (LLM hidden → DiT → denormalize → patch_encoder → wav → + # eos). Zero overhead when DOTS_TTS_BETA_TRACE is unset (single + # bool check per step); kept as the model's debugging instrument + # (cf. voxcpm2's enable_profiling runtime knob). + self._beta_trace = bool(os.environ.get("DOTS_TTS_BETA_TRACE")) + + # Per-request state plumbing. + # _active_states: per-request state keyed by request_id (created in + # preprocess prefill, evicted by _flush_deferred_cleanup). + # _pending_requests: per-step list of (req_id, is_prefill, embeds, + # span_len) populated by preprocess(), consumed by forward(). + # _deferred_cleanup_ids: requests vLLM has marked finished but whose + # side-path audio still needs draining in the current forward step; + # evicted from _active_states at the end of forward(). + self._active_states: dict[str, _RequestState] = {} + self._pending_requests: list[tuple[str, bool, torch.Tensor, int]] = [] + self._deferred_cleanup_ids: set[str] = set() + + # Audio side-channel: + # _audio_queue: _finish_decode pushes (req_id, audio_tensor) tuples + # here; make_omni_output drains them into multimodal_outputs each + # step. + # _sample_rate: cached from AudioVAE config (= 48000 for soar). + self._audio_queue: list[tuple[str, torch.Tensor]] = [] + self._sample_rate = self._audio_vae.sample_rate + + # Stop-signal side-channel: + # _results_queue: _finish_decode pushes (req_id, stop_logits) tuples + # here; compute_logits drains them into the [bsz, vocab_size] + # logits tensor each step. + self._results_queue: list[tuple[str, torch.Tensor | None]] = [] + + logger.info( + "DotsTTS talker built (base_lm=Qwen2[28L,12H,1536d], " + "audio_vae=AudioVAE, dit=DiT[18L,16H,1024d], " + "patch_encoder=VAESemanticEncoder[24L,16H,1024d], " + "5 projectors, speaker_encoder=CAM++[7.2M], " + "io_helper=%s, model=%s, dit_steps=%d); " + "_finish_decode: DiT → denormalize → patch_encoder AR + " + "streaming AudioVAE decode → wav + eos_proj → stop signal.", + "loaded" if self._io_helper.global_mean is not None else "identity", + vllm_config.model_config.model, + _DIT_NUM_STEPS, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: + """Drive the base LM, then the per-request audio side path. + + One vLLM step: base LM forward under PagedAttention, then slice + the scaffold hidden states per request (in _pending_requests + order) and run _finish_decode on each — DiT sampling, streaming + AudioVAE decode, AR loopback, and stop signal all happen there. + vLLM only ever sees the returned hidden states; audio leaves + through the _audio_queue side channel. + """ + output = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + if isinstance(output, IntermediateTensors): + return output + if isinstance(output, tuple): + output = output[0] + + # Per-request side-path: slice scaffold_hidden by token span_len + # in the order preprocess() appended to _pending_requests. + token_offset = 0 + for req_id, is_prefill, _embeds, span_len in self._pending_requests: + req_hidden = output[token_offset : token_offset + span_len] + token_offset += span_len + self._finish_decode(req_id, req_hidden, is_prefill) + + # End-of-step cleanup + self._pending_requests.clear() + self._flush_deferred_cleanup() + + return output + + # ── vllm-omni protocol methods ── + + @staticmethod + def _resolve_latent_stats_path(model_arg: str) -> str | None: + """Resolve latent_stats.pt to a local filesystem path. + + ``vllm_config.model_config.model`` is typically an HF repo ID like + ``"rednote-hilab/dots.tts-soar"`` — joining it with ``"latent_ + stats.pt"`` gives a non-existent path. Try the local-dir form + first (covers ``--model /local/path``), then fall back to HF + cache lookup with ``local_files_only=True`` so we never trigger + a download from this side (the file was fetched alongside the + safetensors at load time). + """ + local = os.path.join(model_arg, "latent_stats.pt") + if os.path.exists(local): + return local + try: + from huggingface_hub import hf_hub_download + + return hf_hub_download( + repo_id=model_arg, + filename="latent_stats.pt", + local_files_only=True, + ) + except Exception: + return None + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + """Map token ids to embeddings via the base LM's embed_tokens. + + Required by vLLM's ``is_text_generation_model`` Protocol check + (interfaces_base._check_vllm_model_embed_input_ids). Without this + the model fails ``--runner generate`` validation at engine startup. + """ + return self.model.embed_tokens(input_ids) + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None = None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Per-step embedding + per-request state registration. + + Prefill: embed the wrapped prompt + ``[文本][文本对应语音]`` via the LM's + embed_tokens. Caller (serving layer or test harness) is + responsible for the wrap — talker just embeds whatever it gets. + + Decode (span_len == 1): use ``state.curr_embed_for_next`` from + the previous step's AR loopback (DiT → patch_encoder output); + zero fallback only if the loopback hasn't produced one yet. + + Voice clone is intentionally not wired here — zero-shot only. + + Returns ``(input_ids unchanged, embeds [span_len, hidden_size], + {})``. ``forward()`` consumes ``self._pending_requests`` to + slice scaffold_hidden per request. + """ + # voxcpm2 protocol compatibility: flatten additional_information. + additional = info_dict.get("additional_information") + if isinstance(additional, dict): + merged = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional.items(): + merged.setdefault(k, v) + info_dict = merged + + span_len = int(input_ids.shape[0]) + dev = input_ids.device + req_id = info_dict.get("request_id", "default") + state = self._active_states.get(req_id) + # Prefill is detected by state lifecycle (not span_len), so decode + # can use span_len=1 with a multi-position curr_embed_for_next. + is_prefill = state is None or not state.prefill_completed + + if is_prefill: + if state is None: + state = _RequestState(request_id=req_id) + self._active_states[req_id] = state + # Reset per-step state on every prefill — previous run is stale. + state.curr_embed_for_next = None + state.precomputed_stop_logits = None + state.is_stopping = False + state.prefill_completed = False + state.noise_step = 0 + # Reset AR-loop state on every prefill. + state.patch_encoder_state = None + state.vocoder_stream_state = None + embeds = self.model.embed_tokens(input_ids) + else: + curr = state.curr_embed_for_next + if curr is not None: + embeds = curr.to(dev) # [1, llm_hidden] + else: + # First decode step before AR loop closed — zero fallback. + embed_dtype = self.model.embed_tokens.weight.dtype + embeds = torch.zeros( + 1, + self.config.hidden_size, + device=dev, + dtype=embed_dtype, + ) + + self._pending_requests.append((req_id, is_prefill, embeds, span_len)) + return input_ids, embeds, {} + + def postprocess(self, *args: Any, **kwargs: Any) -> dict: + return {} + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Mark finished requests for deferred eviction. + + vLLM scheduler calls this BEFORE forward() to notify that certain + requests have completed. We can't drop their _RequestState yet — + the current forward step still needs to drain side-path audio + chunks. Real eviction happens at the end of forward() via + _flush_deferred_cleanup. + """ + for req_id in finished_req_ids: + if req_id in self._active_states: + self._deferred_cleanup_ids.add(req_id) + + def _flush_deferred_cleanup(self) -> None: + """Evict states marked by on_requests_finished. Called at the + tail of forward().""" + for req_id in self._deferred_cleanup_ids: + self._active_states.pop(req_id, None) + self._deferred_cleanup_ids.clear() + + def _finish_decode( + self, + req_id: str, + req_hidden: torch.Tensor, + is_prefill: bool, + ) -> None: + """One request's per-step side path: DiT → AR loopback → wav → stop. + + One vLLM step = one audio patch. patch_encoder.decode_patch + collapses its internal _PATCH_ENCODER_OUT_DS_RATE positions into + a single LLM input embedding (see VAESemanticEncoder._project_ + embeddings), so vLLM's span_len=1 decode matches naturally — no + slot bookkeeping needed. + + Per-call flow (prefill or decode): + 1. _hidden_proj(last_hidden) → fm_sequence (+1 position) + 2. DiT N-step Euler → audio_patch [1, 4, 128] (normalized) + 3. _latent_proj(audio_patch) → fm_sequence (+4 positions, normalized) + 4. io_helper.denormalize(audio_patch) → audio_patch_raw + 5. patch_encoder(audio_patch_raw) → next-step embed [1, 1536] + 6. AudioVAE stream_step(audio_patch_raw) → wav chunk pushed to + _audio_queue (prefill's patch #0 included; output lags 2 + frames behind input) + 7. (decode only) eos_proj(last_hidden) → stop_logits pushed to + _results_queue; threshold > 0.8 flips state.is_stopping and + stream_flush drains the decoder's lookahead tail. + """ + state = self._active_states.get(req_id) + if state is None: + return + + if state.fm_sequence is None: + self._initialize_request_fm_state( + state, + device=req_hidden.device, + dtype=req_hidden.dtype, + ) + + # This step's append would overflow the per-request FM buffer. + # Stop gracefully (same path as a model-decided stop below) instead + # of raising past _append_hidden_chunk/_append_history_chunk — + # those raises are engine-fatal, not request-fatal, and would take + # every other in-flight and future request down with this one. + if state.fm_seq_len + _HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE > state.fm_capacity: + state.is_stopping = True + stop_logits = torch.tensor([[0.0, 1.0]], device=req_hidden.device, dtype=req_hidden.dtype) + state.precomputed_stop_logits = stop_logits + tail = self._audio_vae.stream_flush(state.vocoder_stream_state) + if tail.size(-1) > 0: + self._audio_queue.append((req_id, tail.reshape(-1))) + self._results_queue.append((req_id, stop_logits)) + return + + # 1. Append last LLM hidden to fm_sequence (+1 position). + last_hidden = req_hidden[-_HIDDEN_PATCH_SIZE:].unsqueeze(0) + + self._append_hidden_chunk(state, last_hidden) + + # 2. DiT N-step Euler → audio latent patch [1, 4, 128] (normalized). + audio_patch = self._run_dit_n_step_euler(state) + + # 3. Append latent to fm_sequence (+4 positions). Stays in + # normalized space — DiT's KV cache lives in that space. + self._append_history_chunk(state, audio_patch) + + # 4. Denormalize for downstream consumers (patch_encoder, AudioVAE). + audio_patch_raw = self._io_helper.denormalize(audio_patch) + + # 5. patch_encoder AR loopback → next-step inputs_embeds [1, 1536]. + next_embeds = self._run_patch_encoder_loopback(state, audio_patch_raw) + state.curr_embed_for_next = next_embeds.squeeze(0).detach() + + # 6. Streaming AudioVAE decode → wav chunk. Runs on prefill too: + # upstream zero-shot emits every DiT patch including the first + # (_generate_latents_stream drops one patch only under prompt + # prefill / voice clone), and our prefill patch #0 corresponds + # to upstream's first decode-loop patch. + wav = self._run_vocoder_stream_step(state, audio_patch_raw) + if wav.size(-1) > 0: + self._audio_queue.append((req_id, wav.reshape(-1))) + + if is_prefill: + if self._beta_trace: + self._beta_trace_log( + state, + is_prefill=True, + last_hidden=last_hidden, + audio_patch=audio_patch, + audio_patch_raw=audio_patch_raw, + next_embeds=next_embeds, + wav=wav if wav.size(-1) > 0 else None, + prob_stop=None, + ) + state.prefill_completed = True + # Keep len(_results_queue) == bsz: compute_logits pairs queue + # entries with batch rows positionally, so every request in the + # batch must push exactly one entry per step (voxcpm2_talker.py + # prefill placeholder pattern). None → forced continue via the + # stop_logits-is-None branch in compute_logits. + self._results_queue.append((req_id, None)) + return + + # 7. eos_proj → stop probability. Mirrors upstream + # model.py:1681 — softmax([continue, stop]) on detached LLM + # hidden, threshold the stop slot at 0.8 (upstream default + # eos_threshold). When stop fires for this step, compute_logits + # drains the high-stop logit on the next dispatch and vLLM + # terminates the request. + eos_logits = self._eos_proj(last_hidden.detach()).softmax(dim=-1) + stop_logits = eos_logits.squeeze(0) + state.precomputed_stop_logits = stop_logits + if eos_logits[0, -1, 1].item() > 0.8: + state.is_stopping = True + # Upstream flushes after its decode loop (model.py:1898); our + # request ends when vLLM sees this step's stop token, so drain + # the decoder's 2-frame lookahead tail in the same step. + # Requests cut by max_tokens/abort never reach this path and + # lose the 40 ms tail — no engine hook fires late enough to + # flush after the final step. + tail = self._audio_vae.stream_flush(state.vocoder_stream_state) + if tail.size(-1) > 0: + self._audio_queue.append((req_id, tail.reshape(-1))) + self._results_queue.append((req_id, stop_logits)) + + if self._beta_trace: + self._beta_trace_log( + state, + is_prefill=False, + last_hidden=last_hidden, + audio_patch=audio_patch, + audio_patch_raw=audio_patch_raw, + next_embeds=next_embeds, + wav=wav if wav.size(-1) > 0 else None, + prob_stop=eos_logits[0, -1, 1].item(), + ) + + def _beta_trace_log( + self, + state: _RequestState, + *, + is_prefill: bool, + last_hidden: torch.Tensor, + audio_patch: torch.Tensor, + audio_patch_raw: torch.Tensor, + next_embeds: torch.Tensor, + wav: torch.Tensor | None, + prob_stop: float | None, + ) -> None: + """(β) diagnostic — patch label matches output.wav indexing in + decode rows; prefill row labels itself ``pre``.""" + + def rmsmax(t: torch.Tensor) -> tuple[float, float]: + f = t.detach().float() + return f.pow(2).mean().sqrt().item(), f.abs().max().item() + + patch_label = "pre" if is_prefill else f"{state.fm_seq_len // 5 - 2:2d}" + req_tail = state.request_id[-8:] + hid_rms, hid_max = rmsmax(last_hidden) + dit_rms, dit_max = rmsmax(audio_patch) + raw_rms, raw_max = rmsmax(audio_patch_raw) + nxt_rms, nxt_max = rmsmax(next_embeds) + if wav is not None: + w_rms, w_max = rmsmax(wav) + wav_part = f"wav:rms={w_rms:.3f} max={w_max:.3f}" + else: + wav_part = "wav:rms= - max= - " + stop_part = f"stop={prob_stop:.3f}" if prob_stop is not None else "stop= - " + logger.info( + f"[β req={req_tail} patch={patch_label}] " + f"hid:rms={hid_rms:.3f} max={hid_max:.3f} | " + f"dit:rms={dit_rms:.3f} max={dit_max:.3f} | " + f"raw:rms={raw_rms:.3f} max={raw_max:.3f} | " + f"next:rms={nxt_rms:.3f} max={nxt_max:.3f} | " + f"{wav_part} | {stop_part}" + ) + + # ── FM helpers (per-request workspace + DiT N-step Euler) ── + + def _initialize_request_fm_state( + self, + state: _RequestState, + *, + device: torch.device, + dtype: torch.dtype, + ) -> None: + """Lazy-allocate per-request FM workspace. + + Mirrors upstream _allocate_generate_state (model.py:413) but per- + request (no workspace sharing) so concurrent requests stay + isolated. + """ + fm_capacity = _MAX_AUDIO_PATCHES * (_HIDDEN_PATCH_SIZE + _LATENT_PATCH_SIZE) + state.fm_capacity = fm_capacity + state.fm_sequence = torch.zeros( + (1, fm_capacity, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_cfg_sequence = torch.zeros( + (1, fm_capacity, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_null_g_cond = torch.zeros( + (1, _FM_HIDDEN), + device=device, + dtype=dtype, + ) + state.fm_seq_len = 0 + + def _append_hidden_chunk( + self, + state: _RequestState, + hidden_chunk: torch.Tensor, + ) -> None: + """Project last _HIDDEN_PATCH_SIZE LLM hiddens via _hidden_proj into fm_sequence. + + Mirrors upstream _append_hidden_chunk (model.py:1308): cfg buffer + gets the null-projected counterpart (``_hidden_proj`` of zero) so + CFG dropout sees no past hidden, only past audio history. + """ + last_hidden = hidden_chunk[:, -_HIDDEN_PATCH_SIZE:, :] + projected = self._hidden_proj(last_hidden) + null_projected = self._hidden_proj(torch.zeros_like(last_hidden)) + end = state.fm_seq_len + projected.size(1) + if end > state.fm_capacity: + raise RuntimeError(f"FM buffer overflow on hidden append: end={end} cap={state.fm_capacity}") + state.fm_sequence[:, state.fm_seq_len : end].copy_(projected.to(state.fm_sequence.dtype)) + state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(null_projected.to(state.fm_cfg_sequence.dtype)) + state.fm_seq_len = end + + def _append_history_chunk( + self, + state: _RequestState, + latent_chunk: torch.Tensor, + ) -> None: + """Project audio latent via _latent_proj into fm_sequence (history append). + + Mirrors upstream _append_history_chunk (model.py:1325): cfg + buffer gets the same value (latent history is shared between + conditional and unconditional CFG paths — only the hidden chunk + gets nulled). + """ + history_latent = self._latent_proj(latent_chunk) + end = state.fm_seq_len + history_latent.size(1) + if end > state.fm_capacity: + raise RuntimeError(f"FM buffer overflow on latent append: end={end} cap={state.fm_capacity}") + state.fm_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_sequence.dtype)) + state.fm_cfg_sequence[:, state.fm_seq_len : end].copy_(history_latent.to(state.fm_cfg_sequence.dtype)) + state.fm_seq_len = end + + def _build_fm_attn_mask( + self, + state: _RequestState, + total_len: int, + device: torch.device, + ) -> torch.Tensor: + """Build DiT attention mask for one decode step. + + Mirrors upstream _build_fm_attn_mask (model.py:1203). Layout: + * [0, fm_seq_len) — committed history. Causal over the past + up to a non-causal tail block of size _HIDDEN_PATCH_SIZE that + sees itself and the new noise patch (the just-appended hidden). + * [latent_start, total_len) — new noise patch, bidirectional + with the committed history and itself. + """ + attn_mask = torch.zeros((1, total_len, total_len), device=device, dtype=torch.bool) + latent_start = total_len - _LATENT_PATCH_SIZE + block_start = state.fm_seq_len - _HIDDEN_PATCH_SIZE + if block_start > 0: + causal_mask = torch.ones((block_start, block_start), device=device, dtype=torch.bool).triu(1).logical_not() + attn_mask[:, :block_start, :block_start] = causal_mask + attn_mask[:, block_start : state.fm_seq_len, : state.fm_seq_len] = True + attn_mask[:, block_start : state.fm_seq_len, latent_start:] = True + attn_mask[:, latent_start:, : state.fm_seq_len] = True + attn_mask[:, latent_start:, latent_start:] = True + if latent_start > state.fm_seq_len: + padding_indices = torch.arange(state.fm_seq_len, latent_start, device=device) + attn_mask[:, padding_indices, padding_indices] = True + return attn_mask + + def _build_fm_pos_ids( + self, + state: _RequestState, + total_len: int, + device: torch.device, + ) -> torch.Tensor: + """Build DiT position ids for one decode step. + + Mirrors upstream _build_fm_pos_ids (model.py:1236). Position ids + stay fp32 — DiT's RotaryEmbedding does its own dtype handling. + """ + pos_ids = torch.zeros((1, total_len), device=device, dtype=torch.float32) + latent_start = total_len - _LATENT_PATCH_SIZE + if state.fm_seq_len > 0: + pos_ids[:, : state.fm_seq_len] = torch.arange( + state.fm_seq_len, + device=device, + dtype=pos_ids.dtype, + ) + pos_ids[:, latent_start:] = torch.arange( + state.fm_seq_len, + state.fm_seq_len + _LATENT_PATCH_SIZE, + device=device, + dtype=pos_ids.dtype, + ) + return pos_ids + + def _run_dit_n_step_euler( + self, + state: _RequestState, + *, + num_steps: int = _DIT_NUM_STEPS, + guidance_scale: float = _DIT_GUIDANCE_SCALE, + ) -> torch.Tensor: + """DiT N-step Euler integration → audio latent patch [1, 4, 128]. + + Manual Euler loop (upstream uses torchdyn.odeint; we manualize to + avoid the dep and keep control flow explicit). Each step: + 1. z_proj = _coordinate_proj(z), placed into the last + _LATENT_PATCH_SIZE positions of input_sequence / cfg_sequence. + 2. Batch [conditional, unconditional] → DiT velocity field. + 3. CFG blend: v_c + guidance_scale * (v_c - v_u). + 4. Euler: z += v * dt. + + Mirrors upstream fm_solver_step (core.py:295) + _flow_matching_step_fm + (core.py:463), reduced to a fixed-step Euler integrator. + """ + assert state.fm_sequence is not None + assert state.fm_cfg_sequence is not None + assert state.fm_null_g_cond is not None + device = state.fm_sequence.device + dtype = state.fm_sequence.dtype + total_len = state.fm_seq_len + _LATENT_PATCH_SIZE + + # Workspace: committed history + 4 noise slots (overwritten per step). + input_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) + input_sequence[:, : state.fm_seq_len] = state.fm_sequence[:, : state.fm_seq_len] + cfg_sequence = torch.zeros((1, total_len, _FM_HIDDEN), device=device, dtype=dtype) + cfg_sequence[:, : state.fm_seq_len] = state.fm_cfg_sequence[:, : state.fm_seq_len] + + attn_mask = self._build_fm_attn_mask(state, total_len, device) + pos_ids = self._build_fm_pos_ids(state, total_len, device) + + g_cond = state.g_cond if state.g_cond is not None else state.fm_null_g_cond + g_cond = g_cond.to(device=device, dtype=dtype) + g_cond_batched = torch.cat([g_cond, torch.zeros_like(g_cond)], dim=0) + + latent_start = total_len - _LATENT_PATCH_SIZE + # fp32 ODE integration (same fix as Ming's CFM sampler, PR #4341): + # bf16's 7-bit mantissa loses the small per-step increments in the + # z += v*dt accumulation (hidden-state cos drift vs upstream past + # step ~7) and quantizes the linspace timesteps. Integration + # state stays fp32; DiT matmuls still run bf16 under the autocast + # block below, and the result is cast back to the sequence dtype + # on return. + # Deterministic per-request noise (voxcpm2 _fill_deterministic_cfm_ + # noise pattern): hash seed:request_key:draw# into a private + # Generator instead of the global CUDA RNG, so outputs reproduce + # run-to-run and concurrent requests cannot perturb each other's + # noise streams. request_key strips the engine's per-run "_" + # uuid suffix so replay across runs keys on the stable batch index. + request_key = state.request_id.split("_", 1)[0] + if not request_key.isdigit(): + request_key = state.request_id + noise_key = f"{_DIT_NOISE_SEED}:{request_key}:{state.noise_step}".encode() + digest = hashlib.blake2b(noise_key, digest_size=8).digest() + gen = torch.Generator(device=device) + gen.manual_seed(int.from_bytes(digest, "little") & 0x7FFF_FFFF_FFFF_FFFF) + z = torch.empty((1, _LATENT_PATCH_SIZE, _LATENT_DIM), device=device, dtype=torch.float32) + z.normal_(generator=gen) + state.noise_step += 1 + dt = 1.0 / num_steps + times = torch.linspace(0.0, 1.0, num_steps + 1, device=device, dtype=torch.float32) + + # Match upstream's autocast wrapper (model.py:280). DiT's + # TimestepEmbedder internally forces fp32 (`.float()` on freqs); + # without autocast the fp32 output hits a bf16 Linear → dtype + # mismatch. autocast bridges this for matmuls inside DiT. + use_amp = device.type == "cuda" and dtype in (torch.float16, torch.bfloat16) + with torch.autocast( + device_type=device.type if device.type == "cuda" else "cuda", + dtype=dtype if use_amp else torch.float32, + enabled=use_amp, + ): + for step in range(num_steps): + t = times[step].reshape(1) + z_proj = self._coordinate_proj(z) + z_c = input_sequence.clone() + z_c[:, latent_start:] = z_proj + z_u = cfg_sequence.clone() + z_u[:, latent_start:] = z_proj + z_batched = torch.cat([z_c, z_u], dim=0) + t_batched = t.repeat(2) + vt = self._head( + x=z_batched, + timesteps=t_batched, + attn_mask=attn_mask, + pos_ids=pos_ids, + g_cond=g_cond_batched, + ) + # Upcast DiT output before CFG blend + Euler update so the + # accumulation arithmetic stays fp32 (autocast only affects + # matmuls; these pointwise ops keep their input dtype). + vt = vt[:, latent_start:].float() + vt_c, vt_u = vt[0:1], vt[1:2] + velocity = vt_c + guidance_scale * (vt_c - vt_u) + z = z + velocity * dt + return z.to(dtype) + + def _run_patch_encoder_loopback( + self, + state: _RequestState, + audio_patch_raw: torch.Tensor, + ) -> torch.Tensor: + """Decode one audio latent patch → 1 LLM input embed. + + Mirrors upstream _consume_audio_patch (model.py:1557) — patch_encoder + side only. Lazy-allocates state.patch_encoder_state on first call. + Returns ``llm_embedding`` of shape ``[1, 1, llm_hidden]`` — the + _PATCH_ENCODER_OUT_DS_RATE encoder positions fold into the feature + axis before ``out_proj``, not into separate output positions. + + ``audio_patch_raw`` must already be denormalized (raw VAE latent + space). Caller (_finish_decode) runs io_helper.denormalize on + DiT's normalized output before passing in. + """ + if state.patch_encoder_state is None: + state.patch_encoder_state = self._patch_encoder.init_decode_state( + max_audio_patch_count=_MAX_AUDIO_PATCHES, + batch_size=1, + device=audio_patch_raw.device, + dtype=audio_patch_raw.dtype, + ) + pe_state = state.patch_encoder_state + positions = ( + torch.arange( + _PATCH_ENCODER_OUT_DS_RATE, + device=audio_patch_raw.device, + dtype=torch.long, + ) + + pe_state.seq_len + ) + llm_embedding, conv_tail = self._patch_encoder.decode_patch( + audio_patch_raw, + pe_state.conv_tail, + pe_state.layer_caches, + positions, + ) + pe_state.conv_tail.copy_(conv_tail) + pe_state.seq_len += _PATCH_ENCODER_OUT_DS_RATE + return llm_embedding + + def _run_vocoder_stream_step( + self, + state: _RequestState, + audio_patch_raw: torch.Tensor, + ) -> torch.Tensor: + """Streaming AudioVAE decode of one latent patch (review M2). + + Mirrors upstream generate_audio_stream (model.py:1881-1900): a + per-request BigVGANStreamState carries the decoder's LSTM hidden + and sliding latent window across steps, so each patch is decoded + with real left context instead of inference_from_latents' + zero-padded isolated boundary. Output lags input by + decoder.stream_lookahead (2 frames): the first call emits 2 frames + (3840 samples), steady state 4 (7680); the stop path in + _finish_decode drains the tail via stream_flush. + + audio_patch_raw: [1, T=4, D=128] denormalized latents. + Returns wav chunk [1, 1, n_samples] (n_samples may be 0). + """ + if state.vocoder_stream_state is None: + state.vocoder_stream_state = self._audio_vae.init_stream_state( + batch_size=1, + chunk_size=_LATENT_PATCH_SIZE, + ) + # stream_step expects [B, latent_dim, T]; the vocoder is pinned + # fp32 (__init__), so cast the (possibly bf16) latents up. + return self._audio_vae.stream_step( + audio_patch_raw.transpose(1, 2).float(), + state.vocoder_stream_state, + ) + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + """Drain self._audio_queue into the OmniOutput's multimodal_outputs. + + Output contract is **delta** (SKILL I1): each call emits only + audio produced since the previous step; the engine's + _consolidate_multimodal_tensors concatenates across steps for + offline consumers. + + Same-request multi-chunk: chunks for the same req_id within one + step are cat'd together. Multi-request: one entry per req_id + in multimodal_outputs["model_outputs"]. + + Note: ``sr`` is not a declared OmniPayload key but the engine / + serving consumer reads it (voxcpm2 convention). + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + + mm: dict[str, Any] = {} + if self._audio_queue: + audio_by_req: dict[str, torch.Tensor] = {} + for req_id, audio in self._audio_queue: + if audio is None: + continue + if req_id in audio_by_req: + audio_by_req[req_id] = torch.cat( + [audio_by_req[req_id].reshape(-1), audio.reshape(-1)], + dim=0, + ) + else: + audio_by_req[req_id] = audio + if audio_by_req: + sr = torch.tensor(self._sample_rate, dtype=torch.int32) + ready_req_ids = list(audio_by_req) + chunks = [audio_by_req[req_id].reshape(-1) for req_id in ready_req_ids] + mm["model_outputs"] = chunks + mm["sr"] = [sr for _ in ready_req_ids] + # sparse_audio: ["1"] flags GPUARModelRunner to skip the default + # payload["hidden"] = scaffold_hidden injection (gpu_ar_model_ + # runner.py:1114). Otherwise prefill scaffold (24 hidden × 1536 + # = 29184 numbers observed) leaks into mm["audio"] via the + # output_processor "hidden" → target_key="audio" rename + # (output_processor.py:84), prepending ~0.61 s of noise. + mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} + self._audio_queue.clear() + else: + # Empty-audio step (prefill) still needs the sparse marker so the + # engine doesn't bleed scaffold_hidden into mm["audio"]. + mm["model_outputs"] = [] + mm["sr"] = [] + mm["meta"] = {"req_id": [], "sparse_audio": ["1"]} + + if self._beta_trace: + payload = mm.get("model_outputs") + if payload: + for i, t in enumerate(payload): + logger.info( + f"[β-emit] mm[model_outputs][{i}]: shape={tuple(t.shape)} " + f"dtype={t.dtype} samples={t.numel()} " + f"rms={t.float().pow(2).mean().sqrt().item():.4f} " + f"max={t.abs().max().item():.4f}" + ) + else: + logger.info("[β-emit] mm empty / no model_outputs") + mo_kind = type(model_outputs).__name__ + mo_info = f"shape={tuple(model_outputs.shape)}" if hasattr(model_outputs, "shape") else "no_shape" + logger.info(f"[β-emit] text_hidden_states arg: {mo_kind} {mo_info}") + + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + """Per-step stop signal logits (vLLM sampling protocol). + + Encodes continue/stop into ``logits[i, 0]`` / ``logits[i, 1]`` + (voxcpm2 convention — rest of vocab stays at -inf so the sampler + picks between the two slots). Source of stop logits is + ``self._results_queue``, populated by forward()'s _finish_decode. + When the queue is empty (preprocess-only steps), defaults to + all-continue. + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + + bsz = hidden_states.shape[0] + logits = torch.full( + (bsz, self.config.vocab_size), + float("-inf"), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + if self._results_queue: + for i, (req_id, stop_logits) in enumerate(self._results_queue): + if i >= bsz: + break + state = self._active_states.get(req_id) + if stop_logits is not None: + if state is not None and state.is_stopping: + logits[i, 0] = 0.0 + logits[i, 1] = 1.0 + else: + # is_stopping=False means prob_stop <= 0.8 — force + # continue. Feeding the raw softmax (continue, stop) + # pair as logits would let the greedy sampler compare + # them directly, silently lowering the effective stop + # threshold to 0.5 (early stop, swallowed endings). + logits[i, 0] = 1.0 + if state is not None: + state.precomputed_stop_logits = None + else: + # No stop signal pushed this step (prefill placeholder). + # Default to continue — stop is driven by eos_proj. + logits[i, 0] = 1.0 # continue + self._results_queue.clear() + else: + logits[:, 0] = 1.0 # all continue + + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load all dots.tts checkpoint weights — step 6b. + + Routing branches in a single pass over the input iterator (it can only + be consumed once): + + * **AudioVAE**: keys from ``vocoder.safetensors`` (no parent prefix). + Matched via set membership against ``self._audio_vae.state_dict()``; + acts as the final catch-all so it must run last. + * **DiT**: keys from ``model.safetensors`` under ``velocity_field_predictor.``. + * **patch_encoder**: keys from ``model.safetensors`` under ``patch_encoder.``. + * **Five projectors** (``hidden_proj`` / ``latent_proj`` / + ``coordinate_proj`` / ``xvec_proj`` / ``eos_proj``): each lives at + its own top-level prefix; dispatched through a shared table. + * **Qwen2 base LM**: keys under ``llm.model.``, handed off to vLLM's + Qwen2Model fuser (q/k/v_proj -> qkv_proj, gate/up_proj -> + gate_up_proj). ``llm.lm_head.*`` is silently dropped (soar uses + ``tie_word_embeddings=True`` and our ``self.model`` is a Qwen2Model + with no lm_head module). + * **CAM++ speaker encoder**: keys from ``speaker_encoder.safetensors`` + (no parent prefix, same convention as AudioVAE). Routed by set + membership against ``self._speaker_encoder.state_dict()`` (keys are + ``model.*`` / ``resample.kernel``, no distinguishing top-level + prefix). **Placed before the AudioVAE fallback** so its keys don't + fall through to the catch-all. + """ + vae_state_keys = set(self._audio_vae.state_dict().keys()) + dit_state_keys = set(self._head.state_dict().keys()) + patch_state_keys = set(self._patch_encoder.state_dict().keys()) + speaker_state_keys = set(self._speaker_encoder.state_dict().keys()) + + # (top-level prefix in checkpoint, local attr name, local nn.Module) + projector_specs: list[tuple[str, str, nn.Module]] = [ + ("hidden_proj.", "_hidden_proj", self._hidden_proj), + ("latent_proj.", "_latent_proj", self._latent_proj), + ("coordinate_proj.", "_coordinate_proj", self._coordinate_proj), + ("xvec_proj.", "_xvec_proj", self._xvec_proj), + ("eos_proj.", "_eos_proj", self._eos_proj), + ] + projector_state_keys = {prefix: set(mod.state_dict().keys()) for prefix, _, mod in projector_specs} + projector_matched: dict[str, list[tuple[str, torch.Tensor]]] = {prefix: [] for prefix, _, _ in projector_specs} + + matched_vae: list[tuple[str, torch.Tensor]] = [] + matched_dit: list[tuple[str, torch.Tensor]] = [] + matched_patch: list[tuple[str, torch.Tensor]] = [] + matched_llm: list[tuple[str, torch.Tensor]] = [] + matched_speaker: list[tuple[str, torch.Tensor]] = [] + skipped_lm_head = 0 + + DIT_PREFIX = "velocity_field_predictor." + PATCH_PREFIX = "patch_encoder." + LLM_MODEL_PREFIX = "llm.model." + LLM_LM_HEAD_PREFIX = "llm.lm_head." + for name, tensor in weights: + if name.startswith(DIT_PREFIX): + candidate = name[len(DIT_PREFIX) :] + if candidate in dit_state_keys: + matched_dit.append((candidate, tensor)) + continue + if name.startswith(PATCH_PREFIX): + candidate = name[len(PATCH_PREFIX) :] + if candidate in patch_state_keys: + matched_patch.append((candidate, tensor)) + continue + if name.startswith(LLM_LM_HEAD_PREFIX): + # Qwen2.5-1.5B-Base has tie_word_embeddings=True, and our + # ``self.model`` is a Qwen2Model (no lm_head submodule). Drop. + skipped_lm_head += 1 + continue + if name.startswith(LLM_MODEL_PREFIX): + candidate = name[len(LLM_MODEL_PREFIX) :] + # Hand off untouched — Qwen2Model.load_weights does its own + # fusion (q/k/v_proj -> qkv_proj, gate/up_proj -> gate_up_proj) + # and tolerates unknown keys, so no local key set to gate. + matched_llm.append((candidate, tensor)) + continue + # Projectors: 5 candidate top-level prefixes, first hit wins. + hit_projector = False + for prefix, _attr, _mod in projector_specs: + if name.startswith(prefix): + candidate = name[len(prefix) :] + if candidate in projector_state_keys[prefix]: + projector_matched[prefix].append((candidate, tensor)) + hit_projector = True + break + if hit_projector: + continue + # Speaker encoder: set-membership match (no distinguishing prefix + # in checkpoint; placed before the AudioVAE fallback so its keys + # don't get silently dropped by the catch-all). + if name in speaker_state_keys: + matched_speaker.append((name, tensor)) + continue + candidate = name[len("vocoder.") :] if name.startswith("vocoder.") else name + if candidate in vae_state_keys: + matched_vae.append((candidate, tensor)) + + loaded: set[str] = set() + + if matched_vae: + vae_loader = AutoWeightsLoader(self._audio_vae) + loaded_vae = vae_loader.load_weights(iter(matched_vae)) + loaded.update(f"_audio_vae.{name}" for name in loaded_vae) + logger.info( + "DotsTTS load_weights: loaded %d/%d AudioVAE tensors.", + len(loaded_vae), + len(vae_state_keys), + ) + + if matched_dit: + dit_loader = AutoWeightsLoader(self._head) + loaded_dit = dit_loader.load_weights(iter(matched_dit)) + loaded.update(f"_head.{name}" for name in loaded_dit) + logger.info( + "DotsTTS load_weights: loaded %d/%d DiT tensors.", + len(loaded_dit), + len(dit_state_keys), + ) + + if matched_patch: + patch_loader = AutoWeightsLoader(self._patch_encoder) + loaded_patch = patch_loader.load_weights(iter(matched_patch)) + loaded.update(f"_patch_encoder.{name}" for name in loaded_patch) + logger.info( + "DotsTTS load_weights: loaded %d/%d patch_encoder tensors.", + len(loaded_patch), + len(patch_state_keys), + ) + + any_projector_matched = False + for prefix, attr_name, mod in projector_specs: + matched = projector_matched[prefix] + if not matched: + continue + any_projector_matched = True + proj_loader = AutoWeightsLoader(mod) + loaded_proj = proj_loader.load_weights(iter(matched)) + loaded.update(f"{attr_name}.{name}" for name in loaded_proj) + logger.info( + "DotsTTS load_weights: loaded %d/%d %s tensors.", + len(loaded_proj), + len(projector_state_keys[prefix]), + prefix[:-1], + ) + + if matched_llm: + # vLLM's Qwen2Model.load_weights does its own fusion + tolerates + # unknown keys. It returns the set of param names it actually + # accepted; re-prefix with ``model.`` for our return value (our + # local attr is ``self.model``). + loaded_llm = self.model.load_weights(iter(matched_llm)) + loaded.update(f"model.{name}" for name in loaded_llm) + logger.info( + "DotsTTS load_weights: loaded %d Qwen2 tensors (fed %d; %d llm.lm_head.* skipped — tied embeddings).", + len(loaded_llm), + len(matched_llm), + skipped_lm_head, + ) + + if matched_speaker: + speaker_loader = AutoWeightsLoader(self._speaker_encoder) + loaded_speaker = speaker_loader.load_weights(iter(matched_speaker)) + loaded.update(f"_speaker_encoder.{name}" for name in loaded_speaker) + logger.info( + "DotsTTS load_weights: loaded %d/%d CAM++ speaker_encoder tensors.", + len(loaded_speaker), + len(speaker_state_keys), + ) + + if ( + not matched_vae + and not matched_dit + and not matched_patch + and not any_projector_matched + and not matched_llm + and not matched_speaker + ): + logger.warning( + "DotsTTS load_weights: no AudioVAE / DiT / " + "patch_encoder / projector / Qwen2 / speaker_encoder keys matched." + ) + + return loaded diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py new file mode 100644 index 000000000..fe17b5c0b --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_vocoder.py @@ -0,0 +1,1434 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 The vLLM-Omni team. +# Copyright (c) rednote-hilab. All rights reserved. +# Adapted from: +# https://github.com/rednote-hilab/dots.tts/tree/a393d2e/src/dots_tts/modules +# +# Merged from six upstream files (preserved as-is, inference behaviour +# unchanged; only internal imports were folded inline): +# - modules/vocoder/config.py -> AudioVAEConfig +# - modules/backbone/layers.py -> Conv1d, ConvTranspose1d +# (Mlp/MultiHeadAttention/RotaryEmbedding +# are not vendored here; they belong +# to the DiT side-path module added later) +# - modules/vocoder/alias_free_filter.py -> sinc, kaiser_sinc_filter1d, +# LowPassFilter1d +# - modules/vocoder/alias_free_resample.py-> UpSample1d, DownSample1d +# - modules/vocoder/alias_free_act.py -> Snake, SnakeBeta, Activation1d +# - modules/vocoder/bigvgan.py -> AudioVAE and all its sub-modules +# +# Per-section upstream attributions are preserved at each section header. +# AudioVAEConfig is the only piece that diverges from upstream: pydantic +# ConfigBase replaced with a stdlib @dataclass (extra fields not exercised by +# any shipped checkpoint). + +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, field +from fractions import Fraction + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import pow, sin +from torch.nn import Parameter +from torch.nn.utils import remove_weight_norm, weight_norm + +# ============================================================================ +# Configuration (adapted from modules/vocoder/config.py) +# ---------------------------------------------------------------------------- +# Upstream subclasses a pydantic ``ConfigBase`` with ``extra="allow"``; both the +# dots.tts-soar and dots.tts-mf checkpoint ``config.json`` pass only the 17 +# declared fields below, so the escape hatch is never exercised. ``.get()`` +# helper preserves bigvgan.py's ``h.get(...)`` call sites without modification. +# ``num_decoder_lookahead`` / ``num_encoder_lookahead`` are looked up via +# ``h.get(...)`` but not present in any shipped config; they fall back to 2. +# ============================================================================ + + +@dataclass +class AudioVAEConfig: + sample_rate: int = 24000 + upsample_rates: list[int] = field(default_factory=list) + upsample_kernel_sizes: list[int] = field(default_factory=list) + upsample_initial_channel: int = 1536 + resblock: str = "1" + resblock_kernel_sizes: list[int] = field(default_factory=list) + resblock_dilation_sizes: list[list[int]] = field(default_factory=list) + downsample_rates: list[int] = field(default_factory=list) + downsample_channels: list[int] = field(default_factory=list) + activation: str = "snakebeta" + snake_logscale: bool = True + latent_dim: int = 128 + causal: bool = False + mi_num_layers: int = 4 + causal_encoder: bool = False + use_bias_at_final: bool = True + use_tanh_at_final: bool = True + + num_decoder_lookahead: int = 2 + num_encoder_lookahead: int = 2 + + def get(self, key: str, default=None): + return getattr(self, key, default) + + +# ============================================================================ +# Conv layers (vendored from modules/backbone/layers.py) +# ---------------------------------------------------------------------------- +# Only the Conv1d / ConvTranspose1d helpers are needed by the AudioVAE. +# Mlp / MultiHeadAttention / RotaryEmbedding from the same upstream file are +# reserved for the DiT flow-matching head, which gets its own vendored module +# in a later step. +# ============================================================================ + + +class Conv1d(nn.Conv1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 1, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + padding_mode: str = "zeros", + bias: bool = True, + padding=None, + causal: bool = False, + **_kwargs, + ): + self.causal = causal + if padding is None: + if causal: + padding = 0 + self.left_padding = dilation * (kernel_size - 1) + else: + padding = int((kernel_size * dilation - dilation) / 2) + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + padding_mode=padding_mode, + bias=bias, + ) + + self.in_channels = in_channels + + def forward(self, x): + if self.causal: + x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2) + return super().forward(x) + + +class ConvTranspose1d(nn.ConvTranspose1d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + output_padding: int = 0, + groups: int = 1, + bias: bool = True, + dilation: int = 1, + padding=None, + padding_mode: str = "zeros", + causal: bool = False, + **_kwargs, + ): + if padding is None: + padding = 0 if causal else (kernel_size - stride) // 2 + if causal: + assert padding == 0, "padding is not allowed in causal ConvTranspose1d." + assert kernel_size == 2 * stride, "kernel_size must be equal to 2*stride in Causal ConvTranspose1d." + + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + output_padding=output_padding, + groups=groups, + bias=bias, + dilation=dilation, + padding_mode=padding_mode, + ) + + self.causal = causal + self.stride = stride + + def forward(self, x): + x = super().forward(x) + if self.causal: + x = x[:, :, : -self.stride] + return x + + +# ============================================================================ +# Alias-free low-pass filter (vendored from modules/vocoder/alias_free_filter.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch under +# the Apache License 2.0. The sinc fallback is adopted from adefossez's +# julius.core.sinc (MIT). +# ============================================================================ + + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + # Normalize filter to have sum = 1, otherwise we will have a small leakage + # of the constant component in the input signal. + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + channels: int = 1, + causal: bool = True, + fixed_filter: bool = False, + ): + # kernel_size should be even number for stylegan3 setup, + # in this implementation, odd number is also possible. + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + if causal: + self.pad_left = kernel_size - 1 + self.pad_right = 0 + else: + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + self.fixed_filter = fixed_filter + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + if fixed_filter: + self.register_buffer("filter", filter) + else: + self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) + + # input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + if self.fixed_filter: + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + else: + out = F.conv1d(x, self.filter, stride=self.stride, groups=C) + return out + + +# ============================================================================ +# Alias-free up/down sampling (vendored from modules/vocoder/alias_free_resample.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). +# ============================================================================ + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.channels = channels + self.causal = causal + self.fixed_filter = fixed_filter + if causal: + self.pad = 0 + else: + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) + if self.fixed_filter: + self.register_buffer("filter", filter) + else: + self.filter = nn.Parameter(filter.expand(channels, -1, -1).clone()) + + # x: [B, C, T] + def forward(self, x): + _, C, _ = x.shape + x = F.pad(x, (self.pad, self.pad), mode="replicate") + if self.fixed_filter: + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + else: + x = self.ratio * F.conv_transpose1d(x, self.filter, stride=self.stride, groups=C) + if self.causal: + x = x[..., : -(self.kernel_size - self.stride)] + else: + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None, channels=None, causal=True, fixed_filter=False): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + channels=channels, + causal=causal, + fixed_filter=fixed_filter, + ) + + def forward(self, x): + return self.lowpass(x) + + +# ============================================================================ +# Snake activations + alias-free Activation1d (vendored from +# modules/vocoder/alias_free_act.py) +# ---------------------------------------------------------------------------- +# Adapted from https://github.com/junjun3518/alias-free-torch (Apache 2.0). +# Snake / SnakeBeta paper: https://arxiv.org/abs/2006.08195 +# ============================================================================ + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + causal=True, + fixed_filter=False, + ): + super().__init__() + # causal=False + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d( + up_ratio, + up_kernel_size, + activation.in_features, + causal=causal, + fixed_filter=fixed_filter, + ) + self.downsample = DownSample1d( + down_ratio, + down_kernel_size, + activation.in_features, + causal=causal, + fixed_filter=fixed_filter, + ) + + # x: [B,C,T] + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + return self.downsample(x) + + +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + # Snake := x + 1/a * sin^2 (xa) + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + return x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude + of the periodic components. + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + # SnakeBeta := x + 1/b * sin^2 (xa) + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + +# ============================================================================ +# BigVGAN AudioVAE (vendored from modules/vocoder/bigvgan.py) +# ---------------------------------------------------------------------------- +# Copyright (c) 2022 NVIDIA CORPORATION. +# Licensed under the MIT license. +# Upstream layout flattened: the only modification is that the three +# original ``from dots_tts.modules.{backbone,vocoder}.* import ...`` lines +# are dropped, since those names are already defined above in this file. +# ============================================================================ + + +@dataclass(slots=True) +class BigVGANStreamState: + lstm_hidden: tuple[torch.Tensor, torch.Tensor] + decoder: DecoderStreamState + + +@dataclass(slots=True) +class DecoderStreamState: + window: torch.Tensor + chunk_size: int + total_frames: int = 0 + emitted_frames: int = 0 + + +def _empty_chunk( + ref: torch.Tensor, + *, + channels: int | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + return ref.new_zeros( + (ref.size(0), channels or ref.size(1), 0), + dtype=dtype or ref.dtype, + ) + + +def _module_state_device_dtype(module: nn.Module) -> tuple[torch.device, torch.dtype]: + for name in ("weight", "bias", "filter"): + tensor = getattr(module, name, None) + if tensor is not None: + return tensor.device, tensor.dtype + for tensor in itertools.chain(module.parameters(), module.buffers()): + return tensor.device, tensor.dtype + raise RuntimeError(f"Unable to infer state dtype/device for {type(module).__name__}.") + + +def _stream_state_zeros( + batch_size: int, + channels: int, + length: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + return torch.zeros( + (batch_size, channels, max(0, int(length))), + device=device, + dtype=dtype, + ) + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +class Conv1d_S(nn.Module): + "Conv1d for spectral normalisation and orthogonal initialisation" + + def __init__( + self, + in_channels, + out_channels, + kernel_size=1, + stride=1, + dilation=1, + groups=1, + causal=False, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.stride = stride + self.dilation = dilation + self.groups = groups + self.causal = causal + pad = 0 if causal else dilation * (kernel_size - 1) // 2 + self.causal_pad = dilation * (kernel_size - 1) if causal else 0 + + self.layer = weight_norm( + nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=pad, + dilation=dilation, + groups=groups, + ) + ) + + def forward(self, inputs): + if self.causal and self.causal_pad > 0: + inputs = F.pad(inputs, (self.causal_pad, 0)) + return self.layer(inputs) + + +class SLSTM(nn.Module): + """ + LSTM without worrying about the hidden state, nor the layout of the data. + Expects input as convolutional layout. + """ + + def __init__( + self, + dimension: int, + num_layers: int = 2, + skip: bool = True, + bidirectional: bool = False, + ): + super().__init__() + self.skip = skip + self.bidirectional = bidirectional + self.lstm = nn.LSTM( + input_size=dimension, + hidden_size=dimension, + num_layers=num_layers, + bidirectional=bidirectional, + batch_first=True, + ) + self._stream_num_layers = num_layers + self._stream_weight_ih = tuple(getattr(self.lstm, f"weight_ih_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_weight_hh = tuple(getattr(self.lstm, f"weight_hh_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_bias_ih = tuple(getattr(self.lstm, f"bias_ih_l{layer_idx}") for layer_idx in range(num_layers)) + self._stream_bias_hh = tuple(getattr(self.lstm, f"bias_hh_l{layer_idx}") for layer_idx in range(num_layers)) + if self.bidirectional: + self.proj_out = nn.Linear(dimension * 2, dimension) + + def forward(self, x): + y, _ = self.lstm(x) + if self.bidirectional: + y = self.proj_out(y) + if self.skip: + y = y + x + return y + + def stream_step( + self, + x: torch.Tensor, + hidden: tuple[torch.Tensor, torch.Tensor] | None, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + if self.bidirectional: + raise RuntimeError("Streaming only supports unidirectional SLSTM.") + + residual = x + if hidden is None: + hidden = self.init_stream_state(x.size(0)) + + hidden_h, hidden_c = hidden + next_hidden_h = [] + next_hidden_c = [] + for layer_idx in range(self._stream_num_layers): + layer_input = x + hx = hidden_h[layer_idx] + cx = hidden_c[layer_idx] + outputs = [] + weight_ih = self._stream_weight_ih[layer_idx] + weight_hh = self._stream_weight_hh[layer_idx] + bias_ih = self._stream_bias_ih[layer_idx] + bias_hh = self._stream_bias_hh[layer_idx] + + for frame_idx in range(x.size(1)): + gates = F.linear(layer_input[:, frame_idx, :], weight_ih, bias_ih) + gates = gates + F.linear(hx, weight_hh, bias_hh) + input_gate, forget_gate, cell_gate, output_gate = gates.chunk(4, dim=-1) + input_gate = torch.sigmoid(input_gate) + forget_gate = torch.sigmoid(forget_gate) + cell_gate = torch.tanh(cell_gate) + output_gate = torch.sigmoid(output_gate) + cx = forget_gate * cx + input_gate * cell_gate + hx = output_gate * torch.tanh(cx) + outputs.append(hx) + + x = torch.stack(outputs, dim=1) + next_hidden_h.append(hx) + next_hidden_c.append(cx) + + y = x + if self.skip: + y = y + residual + return y, (torch.stack(next_hidden_h, dim=0), torch.stack(next_hidden_c, dim=0)) + + def init_stream_state( + self, + batch_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_directions = 2 if self.bidirectional else 1 + state_shape = ( + self.lstm.num_layers * num_directions, + batch_size, + self.lstm.hidden_size, + ) + weight = self.lstm.weight_hh_l0 + return ( + weight.new_zeros(state_shape), + weight.new_zeros(state_shape), + ) + + +class ResStack(nn.Module): + def __init__(self, channel, kernel_size=3, base=3, nums=4, causal=False): + super().__init__() + + self.layers = nn.ModuleList([]) + for i in range(nums): + dil = base**i + pad1 = dil * (kernel_size - 1) if causal else dil + pad2 = (kernel_size - 1) if causal else 1 + block = [ + nn.LeakyReLU(), + ] + if causal and pad1 > 0: + block.append(nn.ConstantPad1d((pad1, 0), 0.0)) + block.append( + nn.utils.weight_norm( + nn.Conv1d( + channel, + channel, + kernel_size=kernel_size, + dilation=dil, + padding=0 if causal else pad1, + ) + ) + ) + block.append(nn.LeakyReLU()) + if causal and pad2 > 0: + block.append(nn.ConstantPad1d((pad2, 0), 0.0)) + block.append( + nn.utils.weight_norm( + nn.Conv1d( + channel, + channel, + kernel_size=kernel_size, + dilation=1, + padding=0 if causal else pad2, + ) + ) + ) + self.layers.append(nn.Sequential(*block)) + + def forward(self, x): + for layer in self.layers: + x = x + layer(x) + return x + + +class Encoder(nn.Module): + def __init__( + self, + in_channels=1, + out_channels=100, + base_channels=12, + proj_kernel_size=3, + stack_kernel_size=3, + stack_dilation_base=2, + stacks=6, + channels=(12, 24, 48, 96, 192, 384, 768), + down_sample_factors=(2, 2, 2, 2, 4, 4), + causal=False, + lookahead=0, + ): + super().__init__() + + act_slope = 0.2 + layers = [] + # pre proj_layer + layers += [ + Conv1d_S( + in_channels, + base_channels, + kernel_size=proj_kernel_size, + stride=1, + causal=causal, + ), + nn.LeakyReLU(act_slope, True), + ] + + # channels: [512, 256, 128, 64], upsample_factors: [5, 2, 2] + for (in_c, out_c), down_f in zip(itertools.pairwise(channels), down_sample_factors, strict=True): + layers += [ + Conv1d_S(in_c, out_c, kernel_size=down_f * 2, stride=down_f, causal=causal), + ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks, causal=causal), + nn.LeakyReLU(act_slope, True), + ] + + # post layers + if lookahead > 0: + layers += [ + Conv1d_S( + channels[-1], + out_channels, + kernel_size=lookahead * 2 + 1, + stride=1, + causal=False, + ), + ] + else: + layers += [ + Conv1d_S( + channels[-1], + out_channels, + kernel_size=proj_kernel_size, + stride=1, + causal=causal, + ), + ] + self.generator = nn.Sequential(*layers) + + def forward(self, conditions, _z_inputs=None): + return self.generator(conditions) + + +class AMPBlock1(torch.nn.Module): + def __init__( + self, + h, + channels, + kernel_size=3, + dilation=(1, 3, 5), + activation=None, + causal=True, + ): + super().__init__() + self.h = h + + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + causal=causal, + ) + ), + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)), + ] + ) + self.convs2.apply(init_weights) + + self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers + + if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=Snake(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs1: + remove_weight_norm(layer) + for layer in self.convs2: + remove_weight_norm(layer) + + +class AMPBlock2(torch.nn.Module): + def __init__(self, h, channels, kernel_size=3, dilation=(1, 3), activation=None, causal=True): + super().__init__() + self.h = h + + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + causal=causal, + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + causal=causal, + ) + ), + ] + ) + self.convs.apply(init_weights) + + self.num_layers = len(self.convs) # total number of conv layers + + if activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=Snake(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + self.activations = nn.ModuleList( + [ + Activation1d( + activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale), + causal=causal, + fixed_filter=True, + ) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + for c, a in zip(self.convs, self.activations, strict=True): + xt = a(x) + xt = c(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs: + remove_weight_norm(layer) + + +class Decoder(nn.Module): + def __init__(self, h): + super().__init__() + self.h = h + causal = h.causal + self._stream_window_sizes: dict[int, int] = {} + + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + + num_decoder_lookahead = h.get("num_decoder_lookahead", 2) + # pre conv + self.conv_pre = weight_norm( + Conv1d( + h.latent_dim, + h.upsample_initial_channel, + kernel_size=2 * num_decoder_lookahead + 1, + stride=1, + causal=False, + ) + ) + + # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default + resblock = AMPBlock1 if h.resblock == "1" else AMPBlock2 + + # transposed conv-based upsamplers. does not apply anti-aliasing + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes, strict=True)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + causal=causal, + ) + ) + ] + ) + ) + + # residual blocks using anti-aliased multi-periodicity composition modules (AMP) + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for k, d in zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes, strict=True): + self.resblocks.append(resblock(h, ch, k, d, activation=h.activation, causal=causal)) + + # post conv + if h.activation == "snake": # periodic nonlinearity with snake function and anti-aliasing + activation_post = Snake(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) + elif h.activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing + activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale) + self.activation_post = Activation1d(activation=activation_post, causal=causal, fixed_filter=False) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal, bias=h.get("use_bias_at_final", True))) + + # weight initialization + for i in range(len(self.ups)): + self.ups[i].apply(init_weights) + self.conv_post.apply(init_weights) + + def forward(self, z): + # pre conv + x = self.conv_pre(z) + + for i in range(self.num_upsamples): + # upsampling + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + # AMP blocks + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + # post conv + x = self.activation_post(x) + x = self.conv_post(x) + if self.h.get("use_tanh_at_final", True): + x = torch.tanh(x) + else: + x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1] + + return x + + @property + def stream_lookahead(self) -> int: + return (self.conv_pre.kernel_size[0] - 1) // 2 + + @staticmethod + def _conv1d_left_context(layer) -> int: + dilation = layer.dilation[0] if isinstance(layer.dilation, tuple) else layer.dilation + kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size + if getattr(layer, "causal", False): + return dilation * (kernel_size - 1) + return layer.padding[0] if isinstance(layer.padding, tuple) else layer.padding + + @staticmethod + def _convtranspose1d_left_context(layer) -> int: + stride = layer.stride[0] if isinstance(layer.stride, tuple) else layer.stride + kernel_size = layer.kernel_size[0] if isinstance(layer.kernel_size, tuple) else layer.kernel_size + if not getattr(layer, "causal", False): + raise NotImplementedError("Streaming only supports causal ConvTranspose1d.") + if kernel_size != 2 * stride: + raise ValueError( + "Streaming ConvTranspose1d expects kernel_size == 2 * stride, got " + f"kernel_size={kernel_size} stride={stride}." + ) + return 1 + + @classmethod + def _activation_left_context(cls, activation: Activation1d) -> int: + upsample = activation.upsample + downsample = activation.downsample.lowpass + if not upsample.causal or not downsample.padding or downsample.pad_right != 0: + raise NotImplementedError("Streaming only supports causal alias-free activations.") + ratio = int(upsample.ratio) + if ratio != int(downsample.stride): + raise ValueError( + "Alias-free activation expects matched up/down ratios, got " + f"up_ratio={ratio} down_ratio={downsample.stride}." + ) + total_left = (upsample.kernel_size - 1) + (downsample.kernel_size - 1) + return (total_left + ratio - 1) // ratio + + @classmethod + def _ampblock_left_context(cls, block) -> int: + if isinstance(block, AMPBlock1): + left_context = 0 + acts1 = block.activations[::2] + acts2 = block.activations[1::2] + for conv1, conv2, act1, act2 in zip( + block.convs1, + block.convs2, + acts1, + acts2, + strict=True, + ): + left_context += ( + cls._activation_left_context(act1) + + cls._conv1d_left_context(conv1) + + cls._activation_left_context(act2) + + cls._conv1d_left_context(conv2) + ) + return left_context + if isinstance(block, AMPBlock2): + left_context = 0 + for conv, activation in zip(block.convs, block.activations, strict=True): + left_context += cls._activation_left_context(activation) + cls._conv1d_left_context(conv) + return left_context + raise TypeError(f"Unsupported resblock type: {type(block).__name__}.") + + def _stream_left_context(self) -> int: + left_context = Fraction(self._conv1d_left_context(self.conv_pre), 1) + current_scale = Fraction(1, 1) + for stage_idx, upsample_layers in enumerate(self.ups): + for upsample in upsample_layers: + left_context += current_scale * self._convtranspose1d_left_context(upsample) + stride = upsample.stride[0] if isinstance(upsample.stride, tuple) else upsample.stride + current_scale /= int(stride) + + stage_start = stage_idx * self.num_kernels + stage_end = stage_start + self.num_kernels + stage_context = max(self._ampblock_left_context(block) for block in self.resblocks[stage_start:stage_end]) + left_context += current_scale * stage_context + + left_context += current_scale * self._activation_left_context(self.activation_post) + left_context += current_scale * self._conv1d_left_context(self.conv_post) + return int(left_context.__ceil__()) + + def stream_window_size(self, chunk_size: int) -> int: + if chunk_size < 1: + raise ValueError(f"chunk_size must be >= 1, got {chunk_size}.") + cached = self._stream_window_sizes.get(chunk_size) + if cached is not None: + return cached + + window_size = chunk_size + self.stream_lookahead + self._stream_left_context() + self._stream_window_sizes[chunk_size] = window_size + return window_size + + def remove_weight_norm(self): + for upsample_layers in self.ups: + for upsample_layer in upsample_layers: + remove_weight_norm(upsample_layer) + for resblock in self.resblocks: + resblock.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + + +class AudioVAE(nn.Module): + def __init__(self, h: AudioVAEConfig): + super().__init__() + self.config = h + + self.h = h + self.hop_size = int(np.prod(h.downsample_rates)) + self.sample_rate = h.sample_rate + self.decoder_lookahead = int(h.get("num_decoder_lookahead", 2)) + + self.audio_encoder = Encoder( + out_channels=h.latent_dim, + down_sample_factors=h.downsample_rates, + channels=h.downsample_channels, + causal=h.causal_encoder, + lookahead=h.get("num_encoder_lookahead", 2), + ) + + intermediate_size = h.latent_dim * 4 + self.enc_mi_layer = nn.Sequential( + nn.Linear(h.latent_dim, intermediate_size), + SLSTM(intermediate_size, num_layers=h.mi_num_layers), + nn.Linear(intermediate_size, h.latent_dim), + ) + self.dec_mi_layer = nn.Sequential( + nn.Linear(h.latent_dim, intermediate_size), + SLSTM(intermediate_size, num_layers=h.mi_num_layers), + nn.Linear(intermediate_size, h.latent_dim), + ) + self.pre_proj = Conv1d( + in_channels=h.latent_dim, + out_channels=h.latent_dim * 2, + kernel_size=1, + stride=1, + ) + self.post_proj = Conv1d(in_channels=h.latent_dim, out_channels=h.latent_dim, kernel_size=1, stride=1) + + self.decoder = Decoder(h) + + def inference(self, data): + latents = self.extract_latents(data["sample"]) + return {"sample": self.inference_from_latents(latents)} + + @torch.autocast(enabled=False, device_type="cuda") + def extract_latents(self, x, do_sample=False): + x = x.float() + x = self.audio_encoder(x) + x = x.permute(0, 2, 1) + x = self.enc_mi_layer(x) + x = x.permute(0, 2, 1) + x = self.pre_proj(x) + if do_sample: + m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) + x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) + return x + + def inference_from_latents(self, x, do_sample=True, noise_scale=1.0): + if do_sample: + assert x.size(1) == self.h.latent_dim * 2, f"Input must be like [B, D, H], got {x.shape}" + m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1) + x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) * noise_scale + else: + assert x.size(1) == self.h.latent_dim, f"Input must be like [B, D, H], got {x.shape}" + x = self.post_proj(x) + x = x.permute(0, 2, 1) + x = self.dec_mi_layer(x) + x = x.permute(0, 2, 1) + return self.decoder(x) + + def _validate_stream_latents(self, latents: torch.Tensor) -> None: + if latents.ndim != 3: + raise ValueError( + f"Streaming latents must have shape [batch, latent_dim, frames], got {tuple(latents.shape)}." + ) + if latents.size(1) != self.h.latent_dim: + raise ValueError(f"Streaming latent_dim must be {self.h.latent_dim}, got {latents.size(1)}.") + + def init_stream_state( + self, + batch_size: int = 1, + chunk_size: int = 8, + ) -> BigVGANStreamState: + if not self.h.causal: + raise RuntimeError("Strict streaming requires a causal vocoder.") + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}.") + window_size = self.decoder.stream_window_size(chunk_size) + device, dtype = _module_state_device_dtype(self.decoder.conv_pre) + return BigVGANStreamState( + lstm_hidden=self.dec_mi_layer[1].init_stream_state(batch_size), + decoder=DecoderStreamState( + window=_stream_state_zeros( + batch_size, + self.h.latent_dim, + window_size, + device=device, + dtype=dtype, + ), + chunk_size=int(chunk_size), + ), + ) + + def _decode_stream_latents( + self, + latents: torch.Tensor, + lstm_hidden: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + self._validate_stream_latents(latents) + latents = latents.float() + x = self.post_proj(latents) + x = x.permute(0, 2, 1) + x = self.dec_mi_layer[0](x) + x, lstm_hidden = self.dec_mi_layer[1].stream_step(x, lstm_hidden) + x = self.dec_mi_layer[2](x) + decoder_dtype = self.decoder.conv_pre.weight.dtype + return x.permute(0, 2, 1).to(dtype=decoder_dtype), lstm_hidden + + def _prepare_stream_decoder_input( + self, + latents: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_input, state.lstm_hidden = self._decode_stream_latents( + latents, + state.lstm_hidden, + ) + return decoder_input + + def _append_stream_decoder_input_tensor( + self, + decoder_input: torch.Tensor, + window: torch.Tensor, + valid_frames: torch.Tensor, + ) -> torch.Tensor: + if window.dtype != decoder_input.dtype: + window = window.to(dtype=decoder_input.dtype) + chunk_size = int(decoder_input.size(-1)) + if chunk_size >= window.size(-1): + raise ValueError(f"decoder window size {window.size(-1)} must be larger than chunk_size {chunk_size}.") + positions = torch.arange( + window.size(-1), + device=window.device, + dtype=valid_frames.dtype, + ) + clipped_valid = valid_frames.clamp(min=0, max=window.size(-1)) + combined = torch.cat( + [window, decoder_input.new_zeros(window.size(0), window.size(1), chunk_size)], + dim=-1, + ) + insert_index = clipped_valid + torch.arange( + chunk_size, + device=window.device, + dtype=valid_frames.dtype, + ) + combined.scatter_( + -1, + insert_index.view(1, 1, -1).expand_as(decoder_input), + decoder_input, + ) + new_valid = (clipped_valid + chunk_size).clamp(max=window.size(-1)) + start = (clipped_valid + chunk_size - window.size(-1)).clamp(min=0) + gather_index = (start + positions).clamp(max=combined.size(-1) - 1) + gathered = combined.gather( + -1, + gather_index.view(1, 1, -1).expand_as(window), + ) + mask = (positions < new_valid).to(dtype=window.dtype).view(1, 1, -1) + return gathered * mask + + def _append_stream_decoder_input( + self, + decoder_input: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_state = state.decoder + chunk_size = int(decoder_input.size(-1)) + if chunk_size != decoder_state.chunk_size: + raise ValueError(f"Streaming chunk_size must stay fixed at {decoder_state.chunk_size}, got {chunk_size}.") + window = decoder_state.window + valid_frames = min(decoder_state.total_frames, window.size(-1)) + valid_frames_tensor = window.new_tensor(valid_frames, dtype=torch.int64) + new_window = self._append_stream_decoder_input_tensor( + decoder_input, + window, + valid_frames_tensor, + ) + decoder_state.window = new_window + decoder_state.total_frames += chunk_size + return new_window + + def compiled_stream_step( + self, + latents: torch.Tensor, + hidden_h: torch.Tensor, + hidden_c: torch.Tensor, + window: torch.Tensor, + valid_frames: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + decoder_input, (hidden_h, hidden_c) = self._decode_stream_latents( + latents, + (hidden_h, hidden_c), + ) + new_window = self._append_stream_decoder_input_tensor( + decoder_input, + window, + valid_frames, + ) + audio_window = self.decode_stream_window(new_window) + return audio_window, hidden_h, hidden_c, new_window + + def decode_stream_window(self, window: torch.Tensor) -> torch.Tensor: + decoder_dtype = self.decoder.conv_pre.weight.dtype + if window.dtype != decoder_dtype: + window = window.to(dtype=decoder_dtype) + return self.decoder(window) + + def _slice_stream_audio_window( + self, + audio_window: torch.Tensor, + state: BigVGANStreamState, + *, + final: bool, + ) -> torch.Tensor: + decoder_state = state.decoder + stable_end = ( + decoder_state.total_frames if final else max(0, decoder_state.total_frames - self.decoder.stream_lookahead) + ) + if stable_end <= decoder_state.emitted_frames: + return _empty_chunk(audio_window, channels=1) + + window_size = decoder_state.window.size(-1) + valid_frames = min(decoder_state.total_frames, window_size) + window_start = decoder_state.total_frames - valid_frames + if decoder_state.emitted_frames < window_start: + raise RuntimeError("Decoder stream window is too short for fixed-graph decoding.") + + local_start = decoder_state.emitted_frames - window_start + local_end = stable_end - window_start + sample_start = local_start * self.hop_size + sample_end = local_end * self.hop_size + decoder_state.emitted_frames = stable_end + return audio_window[..., sample_start:sample_end] + + def stream_step( + self, + latents: torch.Tensor, + state: BigVGANStreamState, + ) -> torch.Tensor: + decoder_input = self._prepare_stream_decoder_input(latents, state) + window = self._append_stream_decoder_input(decoder_input, state) + audio_window = self.decode_stream_window(window) + return self._slice_stream_audio_window(audio_window, state, final=False) + + def stream_flush(self, state: BigVGANStreamState) -> torch.Tensor: + audio_window = self.decode_stream_window(state.decoder.window) + return self._slice_stream_audio_window(audio_window, state, final=True) + + def remove_weight_norm(self): + self.decoder.remove_weight_norm() + + +__all__ = ["AudioVAEConfig", "AudioVAE"] diff --git a/vllm_omni/model_executor/models/dots_tts/pipeline.py b/vllm_omni/model_executor/models/dots_tts/pipeline.py new file mode 100644 index 000000000..22b50c629 --- /dev/null +++ b/vllm_omni/model_executor/models/dots_tts/pipeline.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dots.tts pipeline topology (frozen). + +Single-stage AR TTS: text → 48 kHz waveform in one pass. +Architecture: Qwen2.5-1.5B base LM (vLLM-native) → DiTAR / MeanFlow head +→ AudioVAE.decode → BigVGAN → 48 kHz mono. + +Speaker conditioning via CAM++ x-vector injected in preprocess(). + +Modeled after voxcpm2's single-stage pipeline (same "vLLM-native base LM ++ side-path computation" pattern), but simplified: no residual LM, no FSQ. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +DOTS_TTS_PIPELINE = PipelineConfig( + model_type="dots_tts", + model_arch="DotsTTSForConditionalGeneration", + default_deploy_config_name="dots_tts.yaml", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="latent_generator", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="audio", + owns_tokenizer=True, + engine_output_type="audio", + sampling_constraints={ + "detokenize": False, + # Stop signal slot in compute_logits (voxcpm2 convention — + # talker emits logits[i, 1] = 1.0 when prob_stop > 0.8, vLLM + # matches against this list to terminate the request). Not + # Qwen2.5 EOS (151643) — that ID would never come out of the + # 2-slot continue/stop softmax our compute_logits builds. + "stop_token_ids": [1], + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/dynin_omni/pipeline.py b/vllm_omni/model_executor/models/dynin_omni/pipeline.py index 584b6ab5c..6f59ebb36 100644 --- a/vllm_omni/model_executor/models/dynin_omni/pipeline.py +++ b/vllm_omni/model_executor/models/dynin_omni/pipeline.py @@ -50,6 +50,7 @@ DYNIN_OMNI_PIPELINE = PipelineConfig( engine_output_type="latent", custom_process_input_func=f"{_PROC}.token2text_to_token2image", custom_process_next_stage_input_func=f"{_PROC}.token2image_to_token2audio_full_payload", + requires_full_payload_input=True, ), StagePipelineConfig( stage_id=2, @@ -60,6 +61,7 @@ DYNIN_OMNI_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="latent", custom_process_input_func=f"{_PROC}.token2image_to_token2audio", + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/gepard/__init__.py b/vllm_omni/model_executor/models/gepard/__init__.py new file mode 100644 index 000000000..fb6738a29 --- /dev/null +++ b/vllm_omni/model_executor/models/gepard/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/gepard/configuration_gepard.py b/vllm_omni/model_executor/models/gepard/configuration_gepard.py new file mode 100644 index 000000000..d6696f52b --- /dev/null +++ b/vllm_omni/model_executor/models/gepard/configuration_gepard.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Config for Gepard-1.0, a single-stage autoregressive TTS. + +Text tokens -> one 32-code FSQ audio frame per step -> NeMo NanoCodec -> +waveform. The backbone is a vLLM-native ``Qwen3_5ForCausalLM``; the codebook +heads, binary stop head and voice-clone ref_compressor are Gepard additions. + +Parses the model's ``gepard_config.json`` sidecar, which nests the LM +parameters under ``backbone_config`` and carries audio-head cardinalities, +special tokens, codec settings and the short-text repetition layout. +""" + +from __future__ import annotations + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +_BACKBONE_MODEL_TYPE = "qwen3_5_text" + +# Stripped from the backbone config: the model uses standard 1D RoPE, and these +# would flip vLLM's uses_mrope() to True. +_MROPE_KEYS = ("mrope_section", "mrope_interleaved") + + +class GepardConfig(PretrainedConfig): + """Configuration for the Gepard-1.0 native-AR TTS model. + + Args mirror ``gepard_config.json``. Defaults match the trained + ``nineninesix/gepard-1.0`` checkpoint so an instance built with no + arguments (e.g. dummy/profiling loads) is still self-consistent. + """ + + model_type = "gepard" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + backbone_config: dict | None = None, + audio_heads: dict | None = None, + audio_embed_dim: int = 32, + special_tokens: dict | None = None, + text_repetition: dict | None = None, + codec: dict | None = None, + voice_cloning: dict | None = None, + stop_threshold: float = 0.5, + stop_loss_weight: float = 2.0, + stop_pos_weight: float = 25.0, + temperature: float = 0.3, + **kwargs, + ): + # Everything is assigned before super().__init__() at the bottom: + # transformers 5.x runs validators there that call get_text_config(). + self.backbone_config = self._normalize_backbone(backbone_config) + + self.audio_head_levels = self._parse_audio_heads(audio_heads) + self.num_audio_heads = len(self.audio_head_levels) + self.audio_embed_dim = audio_embed_dim + # head0 is the token vLLM samples; STOP is a synthetic sentinel one past + # its valid range. + self.head0_vocab_size = self.audio_head_levels[0] if self.audio_head_levels else 8 + self.stop_token = self.head0_vocab_size + + st = special_tokens or {} + self.start_of_text = st.get("start_of_text", 248073) + self.end_of_text = st.get("end_of_text", 248074) + self.start_of_speech = st.get("start_of_speech", 248070) + self.end_of_speech = st.get("end_of_speech", 248071) + self.tts_pad = st.get("tts_pad", 248076) + # The speaker placeholder slots begin at tokeniser_length. + self.tokeniser_length = st.get("tokeniser_length", 248077) + self.speaker_token_base = self.tokeniser_length + + # Prompt layout, read by prompt.py. These must match the training + # layout, so they come from the checkpoint rather than a literal. + tr = text_repetition or {} + self.text_repetition_enabled = tr.get("enabled", True) + self.text_repetition_target_tokens = tr.get("target_text_tokens", 16) + self.text_repetition_apply_below = tr.get("apply_below", 13) + self.text_repetition_max_repeats = tr.get("max_repeats", 8) + + # NeMo NanoCodec, which runs outside vLLM. + cc = codec or {} + self.codec_id = cc.get("codec_id", "nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps") + self.codec_sample_rate = cc.get("sample_rate", 22050) + self.codec_frame_rate_hz = cc.get("frame_rate_hz", 21.5) + # Per-group FSQ levels; num_codec_groups * len(fsq_levels) == 32. + self.fsq_levels = cc.get("fsq_levels", [8, 7, 6, 6]) + self.num_codec_groups = cc.get("num_layers", 8) + self.codec_do_unfold = cc.get("do_unfold", True) + + self.stop_threshold = stop_threshold + self.stop_loss_weight = stop_loss_weight + self.stop_pos_weight = stop_pos_weight + # head0 goes through vLLM's SamplingParams; the other 31 heads use this + # in-model. There is no top_p: Gumbel-max has no nucleus step. + self.temperature = temperature + + # Carried for the cloning follow-up; zero-shot uses null_prefix only. + vc = voice_cloning or {} + self.voice_cloning_enabled = vc.get("enabled", True) + comp = vc.get("compressor", {}) or {} + self.num_speaker_prefix = comp.get("num_queries", 8) + self.ref_compressor_num_blocks = comp.get("num_layers", 2) + self.ref_compressor_num_heads = comp.get("num_heads", 8) + self.ref_compressor_d_model = comp.get("d_model", 1024) + self.ref_compressor_ffn_mult = comp.get("ffn_hidden_size_multiplier", 4) + + # Last on purpose — triggers transformers-5.x validators that call + # get_text_config(); every attribute they touch must already exist. + super().__init__(**kwargs) + + @classmethod + def from_checkpoint( + cls, + model: str, + backbone_config: dict | None = None, + revision: str | None = None, + ) -> GepardConfig: + """Build the full config for a checkpoint that self-identifies as the + bare backbone: audio fields from the sidecar, backbone fields from the + loaded config. + + ``revision`` must be the one the weights came from — a revision that + moves the audio-head cardinalities or the special tokens moves the + prompt layout and the STOP sentinel with them. + """ + sidecar: dict = {} + try: + from vllm.transformers_utils.config import get_hf_file_to_dict + + sidecar = dict(get_hf_file_to_dict("gepard_config.json", model, revision=revision) or {}) + except (OSError, ValueError) as e: + # A missing or malformed sidecar is expected; the defaults match the + # trained checkpoint. Anything else must surface. + logger.warning( + "GepardConfig: could not read gepard_config.json from %s (%s: %s); " + "audio fields fall back to trained-checkpoint defaults.", + model, + type(e).__name__, + e, + ) + if backbone_config: + bb = dict(backbone_config) + # The loaded config carries the hf_overrides-patched identity; the + # backbone must stay qwen3_5_text or get_text_config() recurses. + bb["model_type"] = _BACKBONE_MODEL_TYPE + bb.pop("architectures", None) + sidecar["backbone_config"] = bb + return cls(**sidecar) + + @staticmethod + def _parse_audio_heads(audio_heads: dict | None) -> list[int]: + """``{level_audio_0: 8, level_audio_1: 7, ...}`` -> ``[8, 7, ...]``. + + Sorted by the numeric suffix, so channel order does not depend on dict + insertion order. + """ + if not audio_heads: + return [8, 7, 6, 6] * 8 + return [audio_heads[k] for k in sorted(audio_heads, key=lambda s: int(str(s).rsplit("_", 1)[-1]))] + + @classmethod + def _normalize_backbone(cls, backbone_config: dict | None) -> dict: + """Return the backbone dict with vestigial mRoPE fields removed. + + They are template leftovers that the trained checkpoint's own + ``config.json`` does not carry. Left in, they flip vLLM's + ``uses_mrope()`` to True and wire ``positions`` wrongly. + """ + bb = dict(backbone_config or {}) + rope = bb.get("rope_parameters") + if isinstance(rope, dict) and any(k in rope for k in _MROPE_KEYS): + rope = {k: v for k, v in rope.items() if k not in _MROPE_KEYS} + bb["rope_parameters"] = rope + logger.info( + "GepardConfig: stripped vestigial mRoPE keys from backbone rope_parameters " + "(model uses standard 1D default RoPE)." + ) + return bb + + def get_text_config(self, **kwargs) -> PretrainedConfig: + """Return the Qwen3.5 backbone config, which vLLM runs the backbone on.""" + cached = getattr(self, "_text_config", None) + if cached is not None: + return cached + + bb = dict(self.backbone_config) + # pop, not setdefault: for_model()'s first positional is named + # model_type, so leaving the key in **bb passes it twice. + model_type = bb.pop("model_type", _BACKBONE_MODEL_TYPE) + try: + text_config = AutoConfig.for_model(model_type, **bb) + except (KeyError, ValueError) as e: + # Only an unknown model_type is a legitimate fallback. A TypeError + # means the kwargs are wrong and must not degrade silently. + text_config = PretrainedConfig(**bb) + logger.warning( + "GepardConfig: AutoConfig.for_model(%s) failed (%s: %s); using a plain " + "PretrainedConfig for the backbone. Verify vLLM still resolves " + "the Qwen3.5 model.", + model_type, + type(e).__name__, + e, + ) + self._text_config = text_config + return text_config + + +AutoConfig.register("gepard", GepardConfig) + +__all__ = ["GepardConfig"] diff --git a/vllm_omni/model_executor/models/gepard/gepard_talker.py b/vllm_omni/model_executor/models/gepard/gepard_talker.py new file mode 100644 index 000000000..02f5da533 --- /dev/null +++ b/vllm_omni/model_executor/models/gepard/gepard_talker.py @@ -0,0 +1,526 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gepard-1.0 native-AR talker. + +Single-stage AR TTS on a vLLM-native Qwen3.5 backbone. Each step samples one +32-channel FSQ frame; head0 is the vLLM-facing token and the other 31 are +side-channel. Generation ends on a learned binary stop head, not an EOS token. +The NeMo NanoCodec decodes committed frames to a waveform outside vLLM. + +Zero-shot only, and ``enforce_eager``; CUDA graph is a perf follow-up. +""" + +from __future__ import annotations + +import dataclasses +import os +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLM +from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig +from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + +# Decode cadence: a short first chunk for early audio, then larger ones to +# amortize the codec call. Does not affect the samples produced. +FIRST_CHUNK_FRAMES = int(os.environ.get("VLLM_GEPARD_FIRST_CHUNK_FRAMES", "10")) +CHUNK_FRAMES = int(os.environ.get("VLLM_GEPARD_CHUNK_FRAMES", "86")) + + +@dataclasses.dataclass(frozen=True) +class _ReqInfo: + """Scheduler state the runner threads into ``preprocess``. + + The hook key strings are read only here, so a rename fails in one place. + """ + + request_id: str + is_prefill: bool + prompt_len: int + num_computed_tokens: int + max_tokens: int | None + seed: int | None + + @classmethod + def from_hook_kwargs(cls, kwargs: dict[str, Any], *, span: int) -> _ReqInfo: + # Defaults cover direct calls, which thread no scheduler state. + max_tokens = kwargs.get("_omni_max_tokens") + seed = kwargs.get("_omni_seed") + return cls( + request_id=str(kwargs.get("request_id", "default")), + is_prefill=bool(kwargs.get("_omni_is_prefill", span > 1)), + prompt_len=int(kwargs.get("_omni_prompt_len") or 0), + num_computed_tokens=int(kwargs.get("_omni_num_computed_tokens") or 0), + max_tokens=None if max_tokens is None else int(max_tokens), + seed=None if seed is None else int(seed), + ) + + +@dataclasses.dataclass +class _GepardState: + """Per-request generation state, keyed by request_id.""" + + request_id: str + # Embedding of the last frame's 32 codes; the next step's input. + curr_embed_for_next: torch.Tensor | None = None + # Frame history (T, 32) handed to the codec. + frames: list[torch.Tensor] = dataclasses.field(default_factory=list) + emitted_frames: int = 0 + # None falls back to null_prefix; the cloning follow-up sets it. + speaker_prefix: torch.Tensor | None = None + # None means the request carried no seed and draws from the global RNG. + generator: torch.Generator | None = None + frame_count: int = 0 + past_first_step: bool = False + is_stopping: bool = False + + +class GepardTalkerForConditionalGeneration(nn.Module): + """Gepard native-AR TTS talker (Qwen3.5 backbone + 32 FSQ heads + stop).""" + + have_multimodal_outputs = True + has_preprocess = True + has_postprocess = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + cfg = vllm_config.model_config.hf_config + if not isinstance(cfg, GepardConfig): + # hf_overrides patches the loaded config instance rather than + # re-classing it, so what arrives here is the backbone config. + cfg = GepardConfig.from_checkpoint( + vllm_config.model_config.model, + backbone_config=cfg.to_dict(), + revision=vllm_config.model_config.revision, + ) + self.config = cfg + self._device = current_omni_platform.get_torch_device() + # Deterministic argmax sampling, for reproducible comparisons. + self._greedy = os.environ.get("VLLM_GEPARD_GREEDY", "0") == "1" + + self.vocab_sizes: list[int] = list(cfg.audio_head_levels) # [8,7,6,6]*8 + self.num_heads: int = cfg.num_audio_heads + self.head0_vocab: int = cfg.head0_vocab_size + self.stop_token: int = cfg.stop_token # sentinel, one past head0's range + self.total_vocab: int = sum(self.vocab_sizes) + self.stop_threshold: float = cfg.stop_threshold + self.temperature: float = cfg.temperature + hidden = self.config.get_text_config().hidden_size + + self.fused_codebook_head = nn.Linear(hidden, self.total_vocab, bias=True) + self.stop_head = nn.Linear(hidden, 1, bias=True) + self.audio_embeddings = nn.ModuleList([nn.Embedding(v, cfg.audio_embed_dim) for v in self.vocab_sizes]) + self.audio_embed_proj = nn.Sequential( + nn.Linear(self.num_heads * cfg.audio_embed_dim, hidden), + nn.GELU(), + nn.Linear(hidden, hidden), + nn.LayerNorm(hidden, elementwise_affine=False), + ) + self.register_buffer("audio_embed_scale", torch.tensor(1.0)) + # Learned default-voice fallback. + self.null_prefix = nn.Parameter(torch.zeros(cfg.num_speaker_prefix, hidden)) + + gather_idx, mask = self._build_gather_mask(self.vocab_sizes) + self.register_buffer("_cb_gather_idx", gather_idx, persistent=False) + self.register_buffer("_cb_mask", mask, persistent=False) + + self.model = Qwen3_5ForCausalLM(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + # Cheap and NeMo-free to construct; load() does the heavy import. + self._codec = NanoCodec(codec_id=cfg.codec_id, sample_rate=cfg.codec_sample_rate) + + self._max_model_len: int = int(vllm_config.model_config.max_model_len or 0) + + self._active_states: dict[str, _GepardState] = {} + # These queues stay index-aligned across preprocess -> forward -> + # compute_logits: preprocess appends in input_batch.req_ids order. + # (req_id, span, samples_frame, is_last_token) + self._pending_requests: list[tuple[str, int, bool, bool]] = [] + # (req_id, head0, do_stop); head0 < 0 marks a discarded placeholder row. + self._results_queue: list[tuple[str, int, bool]] = [] + self._audio_queue: list[tuple[str, torch.Tensor | None]] = [] + self._deferred_cleanup_ids: set[str] = set() + self._sample_rate: int = cfg.codec_sample_rate + + @staticmethod + def _build_gather_mask(vocab_sizes: list[int]) -> tuple[torch.Tensor, torch.Tensor]: + """Gather indices + -inf pad mask so one argmax samples all 31 heads.""" + num_heads = len(vocab_sizes) + max_cb_vocab = max(vocab_sizes[1:]) # == 8 for Gepard; do NOT hardcode 7 + gather_rows, mask_rows = [], [] + cum = vocab_sizes[0] + for i in range(1, num_heads): + v = vocab_sizes[i] + gather_rows.append(list(range(cum, cum + v)) + [cum] * (max_cb_vocab - v)) + mask_rows.append([0.0] * v + [float("-inf")] * (max_cb_vocab - v)) + cum += v + return ( + torch.tensor(gather_rows, dtype=torch.long), + torch.tensor(mask_rows, dtype=torch.float32), + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Fuse codebook_heads.{0..31} into fused_codebook_head; delegate the rest.""" + cb_w: dict[str, torch.Tensor] = {} + rest: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + if name.startswith(("aligner.", "supcon_head.")): + continue # training-only + if name.startswith("codebook_heads."): + cb_w[name] = tensor + continue + if name.startswith(("model.", "lm_head.")): + # We nest the backbone one level deeper than the checkpoint + # keys assume, so prepend the extra hop. + name = "model." + name + rest.append((name, tensor)) + + w_parts, b_parts = [], [] + for i in range(self.num_heads): + if (wk := f"codebook_heads.{i}.weight") in cb_w: + w_parts.append(cb_w[wk]) + if (bk := f"codebook_heads.{i}.bias") in cb_w: + b_parts.append(cb_w[bk]) + if w_parts: + rest.append(("fused_codebook_head.weight", torch.cat(w_parts, dim=0))) + if b_parts: + rest.append(("fused_codebook_head.bias", torch.cat(b_parts, dim=0))) + + # NanoCodec loads lazily on first decode, so load_format=dummy and + # NeMo-less environments never pay for it. + loader = AutoWeightsLoader(self, skip_prefixes=["mtp.", "ref_compressor."]) + return loader.load_weights(iter(rest)) + + def _get_or_create_state(self, request_id: str) -> _GepardState: + st = self._active_states.get(request_id) + if st is None: + st = _GepardState(request_id=request_id) + self._active_states[request_id] = st + return st + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + # Fires BEFORE forward -> defer the actual free until after forward runs, + # else we'd drop state the in-flight forward still reads. + self._deferred_cleanup_ids.update(str(r) for r in finished_req_ids) + + def _flush_deferred_cleanup(self) -> None: + # Cleanup only. Audio queued here is never delivered — the id has left + # the output batch by now — so the tail is emitted in forward instead, + # and anything still pending here is genuinely lost. Say so. + for req_id in self._deferred_cleanup_ids: + state = self._active_states.pop(req_id, None) + if state is not None and state.frame_count > state.emitted_frames: + logger.warning( + "Gepard request %s ended with %d undelivered frame(s): it finished " + "without the stop head firing and without reaching its token budget " + "(aborted or preempted). Its audio is truncated.", + req_id, + state.frame_count - state.emitted_frames, + ) + self._deferred_cleanup_ids.clear() + + def _is_last_output_token(self, info: _ReqInfo, *, span: int, samples_frame: bool) -> bool: + """Whether the token sampled this step is the request's budget-final one. + + That step is the last one on which the request is still in the output + batch, so it is also the last chance to ship its audio. Counted off the + engine's cursor rather than a local tally, so a re-prefill after + preemption cannot desync it. + """ + if not samples_frame: + # A partial prefill chunk's token is discarded, so it costs no budget. + return False + budget = info.max_tokens + if self._max_model_len: + len_cap = self._max_model_len - info.prompt_len + budget = len_cap if budget is None else min(budget, len_cap) + if budget is None or budget <= 0: + return False + produced = info.num_computed_tokens + span - info.prompt_len + 1 + return produced >= budget + + def _emit_audio(self, state: _GepardState, *, is_final: bool) -> None: + """Decode this request's un-emitted frames and queue the new samples. + + Queues deltas: each call contributes only samples not queued before. + """ + pending = state.frame_count - state.emitted_frames + if pending <= 0: + return + if not is_final: + threshold = FIRST_CHUNK_FRAMES if state.emitted_frames == 0 else CHUNK_FRAMES + if pending < threshold: + return + + frames = torch.stack(state.frames) # (T, 32) + if not self._codec.is_loaded: + # The first decode pays the codec load; warmup is a perf follow-up. + self._codec.load(frames.device) + audio = self._codec.decode_stream(frames, start_idx=state.emitted_frames, is_final=is_final) + if audio is None: + return + self._audio_queue.append((state.request_id, audio)) + state.emitted_frames = state.frame_count + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + @staticmethod + def _gumbel_noise(ref: torch.Tensor, generators: list[torch.Generator | None] | None) -> torch.Tensor: + """Gumbel noise shaped like ``ref``; row i uses ``generators[i]`` if set. + + A seeded row draws its own noise of a fixed shape, so the noise it sees + does not depend on who shared its batch. Only the noise: the logits come + from a bf16 matmul that is not row-wise reproducible, so a batched + request can still sample differently from a solo one. + """ + if not generators or all(g is None for g in generators): + u = torch.rand_like(ref) + else: + u = torch.stack( + [ + torch.rand(ref.shape[1:], generator=g, device=ref.device, dtype=ref.dtype) + if g is not None + else torch.rand(ref.shape[1:], device=ref.device, dtype=ref.dtype) + for g in generators + ] + ) + return -torch.log(-torch.log(u.clamp_min_(1e-20))) + + def _sample_frame( + self, + hidden: torch.Tensor, + generators: list[torch.Generator | None] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """(N, hidden) -> head0 (N,), heads_1_31 (N, 31), do_stop (N,). + + ``generators`` is row-aligned with ``hidden``. The 32 heads are sampled + here rather than by vLLM's sampler, so a request's seed only reaches + the audio through this. + """ + fused = self.fused_codebook_head(hidden) + + gathered = fused[:, self._cb_gather_idx].float() + self._cb_mask # (N,31,8) + h0_logits = fused[:, : self.head0_vocab].float() + if self._greedy: + heads = gathered.argmax(dim=-1) + head0 = h0_logits.argmax(dim=-1) + else: + heads = (gathered / self.temperature + self._gumbel_noise(gathered, generators)).argmax(dim=-1) + head0 = (h0_logits / self.temperature + self._gumbel_noise(h0_logits, generators)).argmax(dim=-1) + + p_stop = torch.sigmoid(self.stop_head(hidden)).squeeze(-1) + do_stop = p_stop > self.stop_threshold + return head0, heads, do_stop + + def _audio_frame_embed(self, head0: torch.Tensor, heads_1_31: torch.Tensor) -> torch.Tensor: + """Rebuild the next step's input embedding from the 32 sampled codes.""" + codes = torch.cat([head0.clamp(max=self.head0_vocab - 1).unsqueeze(-1), heads_1_31], dim=-1) + parts = [self.audio_embeddings[i](codes[:, i]) for i in range(self.num_heads)] + return self.audio_embed_proj(torch.cat(parts, dim=-1)) * self.audio_embed_scale + + def preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None = None, **req_infos: Any): + # Prefill injects the speaker prefix at the placeholder slots. The + # prompt ends on SOS, so forward samples the first frame there. + dev = input_ids.device + span = int(input_ids.shape[0]) + info = _ReqInfo.from_hook_kwargs(req_infos, span=span) + req_id = info.request_id + is_prefill = info.is_prefill + prompt_len = info.prompt_len + computed = info.num_computed_tokens + # A prefill chunk that does not finish the prompt commits no frame: + # vLLM discards its sampled token. + samples_frame = (not is_prefill) or prompt_len <= 0 or computed + span >= prompt_len + + if is_prefill: + state = self._get_or_create_state(req_id) + # An aborted id can be resubmitted, and the runner treats the two as + # distinct requests. Drop the pending free, or the cleanup at the + # end of this forward takes the new request's state with it. + resubmitted = req_id in self._deferred_cleanup_ids + self._deferred_cleanup_ids.discard(req_id) + if resubmitted: + # Stand in for the cleanup just cancelled. + undelivered = state.frame_count - state.emitted_frames + if undelivered > 0: + logger.warning( + "Gepard request %s ended with %d undelivered frame(s) and its id has been " + "resubmitted; the first request's audio is truncated.", + req_id, + undelivered, + ) + elif state.frame_count: + # Not a resubmit, so this is a preemption recompute, and it + # cannot be recovered: a frame's 31 side-channel codes are not + # in the token stream, so the replayed prompt cannot rebuild the + # KV that produced the audio already delivered. Generation + # restarts from frame 0 and the caller gets two unrelated halves. + logger.warning( + "Gepard request %s was preempted after %d frame(s) and is being recomputed; " + "its audio will not continue from where it stopped. Raise the stage's " + "gpu_memory_utilization or lower max_num_seqs to keep requests from being " + "preempted.", + req_id, + state.frame_count, + ) + state.curr_embed_for_next = None + state.frames.clear() + state.emitted_frames = 0 + state.frame_count = 0 + state.past_first_step = False + state.is_stopping = False + # SamplingParams.seed only reaches vLLM's sampler, so it never + # touches the audio unless it is threaded into a generator here. + state.generator = None + if info.seed is not None: + state.generator = torch.Generator(device=dev) + state.generator.manual_seed(info.seed) + + embeds = self.model.embed_input_ids(input_ids) + prefix = state.speaker_prefix if state.speaker_prefix is not None else self.null_prefix + base = self.config.speaker_token_base + spk = (input_ids >= base) & (input_ids < base + self.config.num_speaker_prefix) + if bool(spk.any()): + embeds[spk] = prefix[input_ids[spk] - base].to(dtype=embeds.dtype) + else: + state = self._active_states.get(req_id) + curr = state.curr_embed_for_next if state is not None else None + if curr is not None: + embeds = curr.to(device=dev).reshape(1, -1) + else: + # Defensive only: a decode step always follows a forward that + # stashed curr_embed_for_next. + hidden = self.config.get_text_config().hidden_size + embeds = torch.zeros(1, hidden, device=dev, dtype=self.null_prefix.dtype) + + is_last_token = self._is_last_output_token(info, span=span, samples_frame=samples_frame) + self._pending_requests.append((req_id, span, samples_frame, is_last_token)) + return input_ids, embeds, {} + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | OmniOutput: + # positions is plain 1D: Gepard uses default RoPE, not mRoPE. input_ids + # passes through only when no embeds were spliced (the profile run). + hidden = self.model( + input_ids=input_ids if inputs_embeds is None else None, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + # Empty _pending_requests means the engine-init profile run: skip. + if self._pending_requests: + # One sampling row per request: the last scheduled position of its + # span. Same row set compute_logits sees at logits_indices. + ends, generators, offset = [], [], 0 + for req_id, span, _samples_frame, _is_last in self._pending_requests: + offset += span + ends.append(offset - 1) + state = self._active_states.get(req_id) + generators.append(state.generator if state is not None else None) + rows = hidden[torch.tensor(ends, device=hidden.device)] + head0, heads_1_31, want_stop = self._sample_frame(rows, generators) + frame_embeds = self._audio_frame_embed(head0, heads_1_31) + head0_l = head0.tolist() + stop_l = want_stop.tolist() + for i, (req_id, _span, samples_frame, is_last_token) in enumerate(self._pending_requests): + if not samples_frame: + # Partial prefill chunk: vLLM discards this row's token. + self._results_queue.append((req_id, -1, False)) + continue + state = self._get_or_create_state(req_id) + if state.is_stopping: + # Async scheduling runs one more step after a stop token — + # only the max_tokens finish is guarded against it. This + # step's output is discarded, so a frame committed here + # would be stranded and reported as lost audio. + self._results_queue.append((req_id, int(head0_l[i]), True)) + continue + # Stop may only fire past the first audio frame, so the + # prefill-sampled frame and the one after it always commit. + do_stop = bool(stop_l[i]) and state.past_first_step + if do_stop: + state.is_stopping = True + else: + state.frames.append(torch.cat([head0[i : i + 1], heads_1_31[i]])) + state.frame_count += 1 + state.curr_embed_for_next = frame_embeds[i] + state.past_first_step = True + self._results_queue.append((req_id, int(head0_l[i]), do_stop)) + # Both ways a request ends have to flush here: this is the last + # step on which it is still in the output batch, and a payload + # queued after it is dropped by routing. + self._emit_audio(state, is_final=do_stop or is_last_token) + self._pending_requests.clear() + self._flush_deferred_cleanup() + return hidden + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None): + # Vocab-wide logits carrying head0, or STOP for a stopped row. Requires + # the pipeline's sampling_constraints, without which vLLM never acts on + # the sentinel. + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + bsz = hidden_states.shape[0] + vocab = self.config.get_text_config().vocab_size + logits = torch.full((bsz, vocab), float("-inf"), device=hidden_states.device, dtype=hidden_states.dtype) + if not self._results_queue: + # Engine-init dummy sampler run: neutral one-hot logits so KV-cache + # profiling can complete. + logits[:, 0] = 0.0 + return logits + # Row i belongs to _results_queue[i]. The pipeline pins temperature=0.0, + # so vLLM's sampler argmaxes the single finite entry. + for i, (_req_id, head0, do_stop) in enumerate(self._results_queue): + if i >= bsz: + break + if head0 < 0: + logits[i, 0] = 0.0 # partial-prefill placeholder; vLLM discards it + continue + logits[i, self.stop_token if do_stop else head0] = 0.0 + self._results_queue.clear() + return logits + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + by_req: dict[str, torch.Tensor] = {} + for req_id, audio in self._audio_queue: + if audio is None: + continue + by_req[req_id] = ( + torch.cat([by_req[req_id].reshape(-1), audio.reshape(-1)], dim=0) if req_id in by_req else audio + ) + self._audio_queue.clear() + + # The sparse marker is always carried, even on a step that decoded + # nothing: without it the runner ships the pooler payload instead, whose + # hidden states land under this stage's "audio" key. + ready_req_ids = list(by_req) + sr = torch.tensor(self._sample_rate, dtype=torch.int32) + mm: dict[str, Any] = { + "model_outputs": [by_req[r].reshape(-1) for r in ready_req_ids], + "sr": [sr for _ in ready_req_ids], + "meta": {"req_id": ready_req_ids, "sparse_audio": ["1"]}, + } + # Deltas, not the whole history: callers concatenate across steps. + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) diff --git a/vllm_omni/model_executor/models/gepard/nanocodec.py b/vllm_omni/model_executor/models/gepard/nanocodec.py new file mode 100644 index 000000000..2f93d310a --- /dev/null +++ b/vllm_omni/model_executor/models/gepard/nanocodec.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NeMo NanoCodec wrapper for Gepard — per-dimension FSQ decode. + +Gepard's 32 FSQ heads emit per-dimension codes (8 groups x 4 dims), while stock +NeMo decode wants composed per-group mixed-radix indices. ``decode_from_codes`` +dequantizes per-dim instead and hands the result to stock ``decode_audio`` — a +few lines on public NeMo API, no fork. + +NeMo is a heavy optional dependency, imported lazily inside ``load``. Decode +only; ``encode_to_dequantized`` belongs to the voice-cloning follow-up. +""" + +from __future__ import annotations + +import os + +import torch +import torch.nn as nn +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Left-context frames prepended to every streaming decode window. NanoCodec is +# convolutional, so decoding a bare chunk boundary yields a discontinuity; the +# extra frames are decoded for receptive-field context and trimmed afterwards. +# Mirrors the reference server's TTS_LOOKBACK_FRAMES. +LOOKBACK_FRAMES = int(os.environ.get("VLLM_GEPARD_LOOKBACK_FRAMES", "8")) + +# End-of-clip shaping. The codec leaves the waveform on a non-zero sample and +# the hard step to silence is an audible click; a linear fade ramps it out. +# Both default to 0.0 (off), matching the reference server's defaults. +END_OF_SPEECH_FADE_MS = float(os.environ.get("VLLM_GEPARD_END_FADE_MS", "0.0")) +END_OF_SPEECH_SILENCE_MS = float(os.environ.get("VLLM_GEPARD_END_SILENCE_MS", "0.0")) + + +def apply_end_of_speech_tail( + audio: torch.Tensor, + sample_rate: int, + fade_ms: float = END_OF_SPEECH_FADE_MS, + silence_ms: float = END_OF_SPEECH_SILENCE_MS, +) -> torch.Tensor: + """Fade the final ``fade_ms`` to zero, then append ``silence_ms`` of pad. + + The fade is what removes the end-of-clip click; padding alone does not. + """ + if audio.numel() == 0 or (fade_ms <= 0 and silence_ms <= 0): + return audio + out = audio + if fade_ms > 0: + n = min(out.shape[0], int(sample_rate * fade_ms / 1000.0)) + if n > 0: + # Copy: the caller's tensor is still referenced by the audio queue. + out = out.clone() + out[-n:] *= torch.linspace(1.0, 0.0, n, dtype=out.dtype, device=out.device) + if silence_ms > 0: + pad = torch.zeros(int(sample_rate * silence_ms / 1000.0), dtype=out.dtype, device=out.device) + out = torch.cat([out, pad], dim=0) + return out + + +def _make_unfolded_codec_cls(): + """Build the ``UnfoldedCodecModel`` subclass. Imports NeMo (optional dep). + + Defined lazily so importing this module never requires NeMo. The subclass + is OUR code layered on stock ``AudioCodecModel`` — not a NeMo fork. + """ + try: + from nemo.collections.tts.models import AudioCodecModel # noqa: PLC0415 + from omegaconf import open_dict # noqa: PLC0415 + except ImportError as e: # pragma: no cover - env-dependent + raise ImportError( + "Gepard TTS needs NeMo for the NanoCodec decoder. Install it with " + "`pip install nemo_toolkit[tts]` (heavy; NVIDIA Open Model License, " + "commercial use OK), or gate the model behind that extra." + ) from e + + class UnfoldedCodecModel(AudioCodecModel): + """AudioCodecModel + direct decode from per-dimension FSQ codes.""" + + def __init__(self, cfg, trainer=None): + # SLMDiscriminator pulls microsoft/wavlm-base-plus (~360 MB) and is + # training-only — strip it from the config before init. + with open_dict(cfg): + disc = cfg.get("discriminator", None) + if disc is not None and "discriminators" in disc: + disc.discriminators = [d for d in disc.discriminators if "SLM" not in d._target_] + super().__init__(cfg, trainer) + + def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor): + """Decode per-dimension discrete FSQ codes to a waveform. + + codes: (B, D, T) per-dim discrete values, D = num_groups * dims. + codes_len: (B,) valid frame count per batch element. + Returns (audio (B, T_audio), audio_len (B,)). + """ + num_levels = self.vector_quantizer.fsqs[0].num_levels.squeeze() + scale = (num_levels // 2).float().to(codes.device) # e.g. [4,3,3,3] + groups = codes.chunk(self.vector_quantizer.num_groups, dim=1) + dequantized = torch.cat( + [(g - scale[None, :, None]) / scale[None, :, None] for g in groups], + dim=1, + ) + return self.decode_audio(inputs=dequantized, input_len=codes_len) + + return UnfoldedCodecModel + + +class NanoCodec(nn.Module): + """Thin wrapper owning the (optional) NeMo NanoCodec decoder. + + Constructing this is cheap and NeMo-free; ``load`` does the heavy import + + weight load so ``load_format=dummy`` and NeMo-less environments don't pay + for it at model construction. + """ + + def __init__(self, codec_id: str, sample_rate: int) -> None: + super().__init__() + self.codec_id = codec_id + self.sample_rate = sample_rate + self._codec: nn.Module | None = None + + @property + def is_loaded(self) -> bool: + return self._codec is not None + + def load(self, device: torch.device) -> None: + """Import NeMo, build the UnfoldedCodecModel, move to device, eval().""" + cls = _make_unfolded_codec_cls() + codec = cls.from_pretrained(self.codec_id) + codec.to(device=device) + codec.eval() + self._codec = codec + logger.info("Gepard NanoCodec loaded from %s on %s", self.codec_id, device) + + @torch.inference_mode() + def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor) -> torch.Tensor: + """(B, 32, T) per-dim codes -> waveform (B, T_audio). Returns audio only.""" + if self._codec is None: + raise RuntimeError("NanoCodec.decode_from_codes called before load()") + audio, _audio_len = self._codec.decode_from_codes(codes, codes_len) + return audio + + @torch.inference_mode() + def decode_stream( + self, + frames: torch.Tensor, + *, + start_idx: int, + lookback: int = LOOKBACK_FRAMES, + is_final: bool = False, + ) -> torch.Tensor | None: + """Decode the tail of a ``(T, 32)`` code history past ``start_idx``. + + The window starts ``lookback`` frames earlier for receptive-field + context and those samples are trimmed after, so successive calls + concatenate to a single whole-history decode. None if nothing is new. + """ + count = int(frames.shape[0]) + if count <= start_idx: + return None + decode_start = max(0, start_idx - lookback) + window = frames[decode_start:count] # (T_win, 32) + codes = window.t().unsqueeze(0).contiguous().long() # (1, 32, T_win) + codes_len = torch.full((1,), codes.shape[-1], dtype=torch.long, device=codes.device) + audio = self.decode_from_codes(codes, codes_len).reshape(-1) + # Derive samples_per_frame from this decode rather than assuming it. + samples_per_frame = audio.shape[0] // (count - decode_start) + audio = audio[(start_idx - decode_start) * samples_per_frame :] + if is_final: + audio = apply_end_of_speech_tail(audio, self.sample_rate) + return audio diff --git a/vllm_omni/model_executor/models/gepard/pipeline.py b/vllm_omni/model_executor/models/gepard/pipeline.py new file mode 100644 index 000000000..e6ce4600d --- /dev/null +++ b/vllm_omni/model_executor/models/gepard/pipeline.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gepard-1.0 pipeline topology. + +Single-stage AR TTS. The Qwen3.5 backbone runs under paged attention, the 32 +codebook heads and the binary stop head are sampled inside the talker, and the +NeMo NanoCodec decodes the accumulated frames outside vLLM. + +``compute_logits`` scatters head0 into vLLM's vocab-wide logits; head0's valid +range is 0..7, so the synthetic STOP sentinel is 8. ``stop_token_ids=[8]`` +couples it to vLLM's stop machinery, with ``detokenize=False`` because head0 is +an audio code. Same idiom as VoxCPM2/MOSS. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +# head0 valid range is 0..7 (its FSQ cardinality is 8); STOP is the sentinel +# one past that range. Mirrors the reference config.py STOP_TOKEN = 8. +_GEPARD_STOP_TOKEN_ID = 8 + +GEPARD_PIPELINE = PipelineConfig( + model_type="gepard", + default_deploy_config_name="gepard.yaml", + model_arch="GepardTalkerForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="gepard", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="audio", + owns_tokenizer=True, + engine_output_type="audio", + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [_GEPARD_STOP_TOKEN_ID], + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/gepard/prompt.py b/vllm_omni/model_executor/models/gepard/prompt.py new file mode 100644 index 000000000..e6a29bc8f --- /dev/null +++ b/vllm_omni/model_executor/models/gepard/prompt.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gepard-1.0 prompt layout. + +``preprocess`` only consumes this layout — it injects the speaker prefix at the +placeholder slots and samples the first audio frame from the SOS position. The +layout itself is built here, by whichever entry point is assembling a request. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig + + +def build_gepard_prompt_ids(text_token_ids: list[int], *, config: GepardConfig) -> list[int]: + """Assemble ``[speaker_slots] + [SOT, *text, EOT]*(R-1) + [SOT, *text, EOT, SOS]``. + + Only the final copy carries SOS, the learned "audio starts now" gate, so + the earlier copies are read as context and never voiced. Short texts repeat + because a couple of text tokens carry too little mass against the speaker + prefix. This must match the training layout or WER collapses, which is why + the thresholds come from the checkpoint's ``text_repetition`` block rather + than from literals here. + + Raises ``ValueError`` on empty text: the layout would still be structurally + valid and the model would voice it as arbitrary audio, so an empty request + has to fail loudly. This is the only entry point in the offline path — there + is no adapter above it to check first. + """ + n = len(text_token_ids) + if n == 0: + raise ValueError("text_token_ids is empty; Gepard needs at least one text token to voice") + if ( + not config.text_repetition_enabled + or n >= config.text_repetition_apply_below + or n >= config.text_repetition_target_tokens + ): + repeats = 1 + else: + repeats = max( + 1, + min( + math.ceil(config.text_repetition_target_tokens / n), + config.text_repetition_max_repeats, + ), + ) + + block = [config.start_of_text, *text_token_ids, config.end_of_text] + text_region = block * (repeats - 1) + [*block, config.start_of_speech] + speaker_slots = [config.speaker_token_base + i for i in range(config.num_speaker_prefix)] + return speaker_slots + text_region diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index 579d0cd13..186e48714 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage-0 talker for higgs-audio v3 (Qwen3 backbone, fused multi-codebook). Architecture: @@ -344,6 +344,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): positions: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, omni_query_start_loc: torch.Tensor | None = None, + req_ids: list[str] | None = None, **_: Any, ) -> None: """Update per-step metadata before runner forward or CUDA graph replay.""" @@ -355,6 +356,91 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self._set_last_step_query_start_loc(omni_query_start_loc) self._decode_step_metadata_from_runner = True + if req_ids is not None: + self._sync_decode_state_with_batch(req_ids) + + def _ensure_state_pool_capacity(self, num_rows: int, device: torch.device) -> None: + if not hasattr(self, "_state_pool_indices"): + self._state_pool_indices: dict[str, int] = {} + self._free_state_pool_indices: set[int] = set() + self._state_pool_has_codes = torch.empty(0, dtype=torch.bool, device=device) + self._state_pool_last_codes = torch.empty((0, self.num_codebooks), dtype=torch.long, device=device) + self._state_pool_delay_count = torch.empty(0, dtype=torch.long, device=device) + self._state_pool_eoc_countdown = torch.empty(0, dtype=torch.long, device=device) + self._state_pool_generation_done = torch.empty(0, dtype=torch.bool, device=device) + + current_rows = self._state_pool_has_codes.shape[0] + if current_rows >= num_rows: + return + + new_rows = max(current_rows * 2, num_rows, 64) + + def resize(tensor: torch.Tensor, fill_value: int | bool) -> torch.Tensor: + result = torch.empty((new_rows, *tensor.shape[1:]), dtype=tensor.dtype, device=device) + result[:current_rows] = tensor + result[current_rows:] = fill_value + return result + + self._state_pool_has_codes = resize(self._state_pool_has_codes, False) + self._state_pool_last_codes = resize(self._state_pool_last_codes, 0) + self._state_pool_delay_count = resize(self._state_pool_delay_count, 0) + self._state_pool_eoc_countdown = resize(self._state_pool_eoc_countdown, -1) + self._state_pool_generation_done = resize(self._state_pool_generation_done, False) + + def _sync_decode_state_with_batch(self, req_ids: list[str]) -> None: + device = self._decode_has_codes.device + self._ensure_state_pool_capacity(len(req_ids), device) + previous_req_ids = getattr(self, "_last_batch_req_ids", []) + + if previous_req_ids != req_ids or any(req_id not in self._state_pool_indices for req_id in req_ids): + previous_rows = [row for row, req_id in enumerate(previous_req_ids) if req_id in self._state_pool_indices] + if previous_rows: + pool_rows = [self._state_pool_indices[previous_req_ids[row]] for row in previous_rows] + src = torch.tensor(previous_rows, dtype=torch.long, device=device) + dst = torch.tensor(pool_rows, dtype=torch.long, device=device) + self._state_pool_has_codes[dst] = self._decode_has_codes[src] + self._state_pool_last_codes[dst] = self._decode_last_codes[src] + self._state_pool_delay_count[dst] = self._decode_delay_count[src] + self._state_pool_eoc_countdown[dst] = self._decode_eoc_countdown[src] + self._state_pool_generation_done[dst] = self._decode_generation_done[src] + + for req_id in req_ids: + if req_id in self._state_pool_indices: + continue + if self._free_state_pool_indices: + pool_row = self._free_state_pool_indices.pop() + else: + pool_row = len(self._state_pool_indices) + self._ensure_state_pool_capacity(pool_row + 1, device) + self._state_pool_indices[req_id] = pool_row + self._state_pool_has_codes[pool_row] = False + self._state_pool_last_codes[pool_row] = 0 + self._state_pool_delay_count[pool_row] = 0 + self._state_pool_eoc_countdown[pool_row] = -1 + self._state_pool_generation_done[pool_row] = False + + self._ensure_decode_state_capacity(len(req_ids), device) + pool_rows = torch.tensor( + [self._state_pool_indices[req_id] for req_id in req_ids], + dtype=torch.long, + device=device, + ) + self._decode_has_codes[: len(req_ids)] = self._state_pool_has_codes[pool_rows] + self._decode_last_codes[: len(req_ids)] = self._state_pool_last_codes[pool_rows] + self._decode_delay_count[: len(req_ids)] = self._state_pool_delay_count[pool_rows] + self._decode_eoc_countdown[: len(req_ids)] = self._state_pool_eoc_countdown[pool_rows] + self._decode_generation_done[: len(req_ids)] = self._state_pool_generation_done[pool_rows] + + self._last_batch_req_ids = list(req_ids) + + def on_requests_finished(self, finished_req_ids: set[str]) -> None: + if not hasattr(self, "_state_pool_indices"): + return + for req_id in finished_req_ids: + pool_row = self._state_pool_indices.pop(req_id, None) + if pool_row is not None: + self._free_state_pool_indices.add(pool_row) + # ------------------------------------------------------------------ forward def forward( self, @@ -1386,7 +1472,19 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): fallback = x.argmax(dim=-1) safe_x = torch.where(all_masked.unsqueeze(-1), torch.zeros_like(x), x) probs = safe_x.softmax(dim=-1) - sampled = torch.multinomial(probs, num_samples=1).squeeze(-1) + if sampling_metadata.generators: + sampled = torch.cat( + [ + torch.multinomial( + probs[req_idx * num_codebooks : (req_idx + 1) * num_codebooks], + num_samples=1, + generator=sampling_metadata.generators.get(req_idx), + ) + for req_idx in range(int(probs.shape[0]) // num_codebooks) + ] + ).squeeze(-1) + else: + sampled = torch.multinomial(probs, num_samples=1).squeeze(-1) sampled = torch.where(all_masked, fallback, sampled) return torch.where(greedy, logits_2d.argmax(dim=-1), sampled) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py index 67eebbab3..6ba4ff00b 100644 --- a/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py +++ b/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import torch import torch.nn as nn @@ -175,7 +175,7 @@ class ConformerEncoderLayer(nn.Module): instance. `PositionwiseFeedForward` instance can be used as the argument. conv_module (torch.nn.Module): Convolution module instance. - `ConvlutionModule` instance can be used as the argument. + `ConvolutionModule` instance can be used as the argument. dropout_rate (float): Dropout rate. normalize_before (bool): True: use layer_norm before each sub-block. diff --git a/vllm_omni/model_executor/models/indextts2/pipeline.py b/vllm_omni/model_executor/models/indextts2/pipeline.py index e29d61563..6e8754daa 100644 --- a/vllm_omni/model_executor/models/indextts2/pipeline.py +++ b/vllm_omni/model_executor/models/indextts2/pipeline.py @@ -45,6 +45,7 @@ INDEXTTS2_PIPELINE = PipelineConfig( sync_process_input_func=f"{_PROC}.talker2s2mel_token_only", extras={"skip_tokenizer_init": True}, sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) @@ -81,6 +82,7 @@ INDEXTTS25_PIPELINE = PipelineConfig( sync_process_input_func=f"{_PROC}.talker2s2mel_token_only", extras={"skip_tokenizer_init": True}, sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 2af6abd42..4309d6177 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -284,14 +284,14 @@ class MiMoAudioDataParser(MultiModalDataParser): if not self.audio_tokenizer_path: raise ValueError( "Audio tokenizer path is not set. Provide " - "`model_config.audio_tokenizer_path` in the stage config " + "`model_config.audio_tokenizer_path` in the model configuration " "or export MIMO_AUDIO_TOKENIZER_PATH." ) if not os.path.exists(self.audio_tokenizer_path): raise ValueError( "Audio tokenizer not exists. Provide " - "`model_config.audio_tokenizer_path` in the stage config " + "`model_config.audio_tokenizer_path` in the model configuration " "or export MIMO_AUDIO_TOKENIZER_PATH." ) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index 26482c6bc..019670cf9 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py @@ -459,7 +459,7 @@ class MiMoAudioToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): if not self.audio_tokenizer_path: raise ValueError( "Audio tokenizer path is not set. Provide " - "`model_config.audio_tokenizer_path` in the stage config " + "`model_config.audio_tokenizer_path` in the model configuration " "or export MIMO_AUDIO_TOKENIZER_PATH." ) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 19a339056..c90adb467 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2025 Xiaomi Corporation. import logging import threading @@ -541,6 +544,11 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor architectures=["Qwen2ForCausalLM"], ) + # vLLM 0.28 turned SupportsPP.make_empty_intermediate_tensors from a + # method into a bare annotation, so declaring SupportsPP no longer + # supplies one. Delegate to the inner model, as the other omni talkers do. + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + self.device = current_omni_platform.get_torch_device() # global_sampler MUST stay greedy (do_sample=False) so its token decision # matches vLLM's external sampler (SamplingParams temperature=0.0). Both diff --git a/vllm_omni/model_executor/models/mimo_audio/pipeline.py b/vllm_omni/model_executor/models/mimo_audio/pipeline.py index 9e54979de..684004b7a 100644 --- a/vllm_omni/model_executor/models/mimo_audio/pipeline.py +++ b/vllm_omni/model_executor/models/mimo_audio/pipeline.py @@ -59,6 +59,7 @@ MIMO_AUDIO_PIPELINE = PipelineConfig( engine_output_type="audio", sync_process_input_func=f"{_PROC}.llm2code2wav_token_only", sampling_constraints={"detokenize": False}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index 8868e9d2a..d7e31beb3 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -25,14 +25,15 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin -from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig from vllm_omni.model_executor.models.ming_tts.constants import SPEAKER_EMBEDDING_DIM from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT -from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input +from .talker_module import CFM, DiT, MingAudioGenerator, build_tts_input from .text_processing import segment_and_normalize from .voice_presets import VoicePresetRegistry @@ -119,7 +120,11 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) DiT(llm_input_dim=self.hidden_size, **config.flowmodel), steps=config.steps, ) - self.aggregator = Aggregator(llm_input_dim=self.hidden_size, **config.aggregator) + # config.aggregator still wins `in_channels` if the checkpoint states it. + self.aggregator = Aggregator( + llm_input_dim=self.hidden_size, + **{"in_channels": self.latent_dim, **config.aggregator}, + ) self.stop_head = nn.Linear(self.hidden_size, 2, bias=True) self.spk_head = nn.Linear(SPEAKER_EMBEDDING_DIM, self.hidden_size, bias=True) @@ -525,27 +530,22 @@ class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin) def _iter_talker_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]: """Yield (name, tensor) pairs from talker/model*.safetensors.""" - model_path = self._model_path - # Try local path first - for candidate in (os.path.join(model_path, "talker"), model_path): + model_root = resolve_model_to_local_path( + self._model_path, + allow_download=True, + allow_patterns=["talker/model*.safetensors"], + cache_dir=self.vllm_config.load_config.download_dir, + ) + + # Nested talker/ layout first, then a flat checkpoint at the root. + for candidate in (os.path.join(model_root, "talker"), model_root): sf_files = sorted(glob_module.glob(os.path.join(candidate, "model*.safetensors"))) if sf_files: for sf_path in sf_files: yield from load_file(sf_path, device="cpu").items() return - # HF hub fallback: download only the talker checkpoint files - model_root = download_weights_from_hf_specific( - model_path, - self.vllm_config.load_config.download_dir, - allow_patterns=["talker/model*.safetensors"], - ) - talker_dir = os.path.join(model_root, "talker") - sf_files = sorted(glob_module.glob(os.path.join(talker_dir, "model*.safetensors"))) - if not sf_files: - raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.") - for sf_path in sf_files: - yield from load_file(sf_path, device="cpu").items() + raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.") def _load_vae_weights(self) -> set[str]: """Load AudioVAE weights from talker/vae/model.safetensors.""" diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index 8c527b8d9..03cf87abb 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py @@ -64,6 +64,7 @@ from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config from vllm_omni.transformers_utils.processors.ming import ( @@ -715,12 +716,8 @@ class MingFlashOmniThinkerForConditionalGeneration( p = Path(model_path) / "mlp" / "model.safetensors" if not p.exists() and not os.path.exists(model_path): - from vllm_omni.model_executor.model_loader.weight_utils import ( - download_weights_from_hf_specific, - ) - try: - local_path = download_weights_from_hf_specific(model_path, None, ["mlp/*"]) + local_path = resolve_model_to_local_path(model_path, allow_download=True, allow_patterns=["mlp/*"]) p = Path(local_path) / "mlp" / "model.safetensors" except Exception: logger.exception("[MingFlashOmniThinker] failed to download mlp/ from %s", model_path) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index cfb1090c2..8e3e55611 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -33,9 +33,14 @@ from vllm.platforms import current_platform from x_transformers.x_transformers import RotaryEmbedding from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE -from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps -from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps +from vllm_omni.model_executor.models.common.ming.cfm_solver import ( + apply_sway_sampling, + get_epss_timesteps, + integrate_cfm_steps, +) +from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer logger = init_logger(__name__) @@ -458,58 +463,6 @@ def silence_holder( return speech_out, sil_cache -######################################################################## -# Audio Postprocess -# Ported from: -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/aggregator.py -######################################################################## - - -class Aggregator(nn.Module): - """Maps generated audio latent patches back to LLM embedding space.""" - - def __init__( - self, - in_channels: int = 64, - hidden_size: int = 1152, - depth: int = 28, - num_heads: int = 16, - mlp_ratio: float = 4.0, - llm_input_dim: int = 896, - **kwargs, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = in_channels - self.num_heads = num_heads - - self.word_embedder = nn.Embedding(1, hidden_size) - self.x_embedder = nn.Linear(in_channels, hidden_size) - self.hidden_size = hidden_size - - self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) - - self.blocks = nn.ModuleList( - [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] - ) - self.final_layer = FinalLayer(hidden_size, llm_input_dim) - - def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: - x = self.x_embedder(x) - cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device)) - x = torch.cat([cls_embed, x], dim=1) - - rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) - if mask is not None: - mask_pad = mask.clone().detach()[:, :1] - mask = torch.cat([mask_pad, mask], dim=-1) - for block in self.blocks: - x = block(x, mask, rope) - x = self.final_layer(x) - x = x[:, :1, :] - return x - - ######################################################################## # Prompt Builder # Adapted from: diff --git a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py index 0a096df86..8f6db1afe 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py @@ -16,16 +16,15 @@ import torch from transformers.utils.hub import cached_file from vllm.logger import init_logger -from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path +from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor -from .spk_embedding import SpkembExtractor from .talker_module import resample if TYPE_CHECKING: + from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE - from .talker_module import Aggregator - logger = init_logger(__name__) @@ -273,18 +272,19 @@ class VoicePresetRegistry: if os.path.isfile(path): return path, candidate - if not os.path.isdir(self._model_path): - try: - hf_root = download_weights_from_hf_specific( - self._model_path, - self._download_dir, - allow_patterns=["talker/data/**"], - require_all=True, - ) - candidate = os.path.join(hf_root, "talker", "data", "voice_name.json") - if os.path.isfile(candidate): - return candidate, os.path.join(hf_root, "talker") - except Exception as e: # pragma: no cover - logger.info("Could not download voice presets from HF: %s", e) + try: + hf_root = resolve_model_to_local_path( + self._model_path, + allow_download=True, + allow_patterns=["talker/data/**"], + cache_dir=self._download_dir, + ) + except Exception as e: + logger.info("Could not download voice presets from HF: %s", e) + return None, None + + candidate = os.path.join(hf_root, "talker", "data", "voice_name.json") + if os.path.isfile(candidate): + return candidate, os.path.join(hf_root, "talker") return None, None diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py index c0ec5a513..b42c4e9cb 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py @@ -11,15 +11,21 @@ import torch import torch.nn as nn from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.model_executor.models.bailing_moe import BailingMoeModel from vllm.model_executor.models.qwen2 import Qwen2Model from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix from vllm.sequence import IntermediateTensors from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata +from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator +from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import ( + CFMGraphExecutor, + CFMSampler, +) +from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss from vllm_omni.model_executor.models.output_templates import OmniOutput -from .aggregator import Aggregator from .config_ming_tts import ( KEY_CFG, KEY_DECODE_STEP, @@ -34,7 +40,6 @@ from .config_ming_tts import ( MingTTSConfig, ) from .constants import SPEAKER_EMBEDDING_DIM -from .flowloss_head import FlowLoss from .patch_emission import ( MING_STOP_REASON_CODES, MING_STOP_REASON_CONTINUE, @@ -82,8 +87,6 @@ class MingLLMModel(nn.Module): self.prefix = prefix self.fm_dtype = _resolve_ming_runtime_dtype(vllm_config) if self.ming_config.model_variant == "moe": - from vllm.model_executor.models.bailing_moe import BailingMoeModel - # BailingMoeModel reads ``vllm_config.model_config.hf_config`` directly # (no get_text_config()), so re-wrap with the nested bailing_moe config. llm_config = vllm_config.model_config.hf_config.get_text_config() @@ -335,8 +338,6 @@ class MingLLMModel(nn.Module): if self._cfm_graph is not None: return self._cfm_graph try: - from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler - # TODO(perf): # (1) compile per-block (dit_model.blocks) instead of the # whole forward for tighter fusion / fewer graph breaks; diff --git a/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py b/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py index 251b2e0f9..355bd8873 100644 --- a/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py +++ b/vllm_omni/model_executor/models/ming_tts/prompt_encoder.py @@ -12,7 +12,7 @@ import torch from safetensors import safe_open from vllm.logger import init_logger -from vllm_omni.engine.stage_init_utils import _resolve_model_to_local_path +from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE from .audio_prep import ( @@ -78,7 +78,7 @@ def _load_prompt_encoder(wrapper: Any) -> AudioVAE: loaded_encoder_params = set() with torch.no_grad(): for shard_path in _iter_model_safetensors( - _resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model)) + resolve_model_to_local_path(str(wrapper.vllm_config.model_config.model)) ): with safe_open(str(shard_path), framework="pt", device="cpu") as handle: for key in handle.keys(): diff --git a/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py b/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py index 8e9f69068..6b6da71fe 100644 --- a/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py +++ b/vllm_omni/model_executor/models/ming_tts/speaker_extractor.py @@ -3,16 +3,13 @@ # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/spkemb_extractor.py from __future__ import annotations -import os import threading from typing import TYPE_CHECKING, Any import torch -import torchaudio from vllm.logger import init_logger -from vllm.multimodal.media.audio import load_audio -from vllm_omni.model_executor.models.ming_flash_omni.spk_embedding import SpkembExtractor +from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor from .audio_prep import coerce_prompt_waveform, coerce_speaker_embeddings from .constants import ( @@ -29,48 +26,9 @@ logger = init_logger(__name__) _EXTRACTOR_LOAD_LOCK = threading.Lock() -def resolve_model_to_local_path(model): - if os.path.isdir(model): - return model - - from huggingface_hub import snapshot_download - - return snapshot_download(model) - - -class MingSpeakerEmbeddingExtractor: - def __init__(self, model, target_sr=16000): - local_model_path = resolve_model_to_local_path(model) - campplus_path = os.path.join(local_model_path, "campplus.onnx") - if not os.path.exists(campplus_path): - raise RuntimeError(f"Missing Ming speaker extractor model: {campplus_path}") - - self.target_sr = int(target_sr) - self._core = SpkembExtractor(campplus_path, target_sr=self.target_sr) - - def extract_from_waveform(self, waveform, sample_rate): - if not isinstance(waveform, torch.Tensor): - waveform = torch.as_tensor(waveform) - - tensor = waveform.detach().to(torch.float32) - if tensor.ndim == 1: - tensor = tensor.unsqueeze(0) - if int(sample_rate) != self.target_sr: - tensor = torchaudio.transforms.Resample(orig_freq=int(sample_rate), new_freq=self.target_sr)(tensor) - - embedding = self._core._extract_spk_embedding(tensor) - return embedding.squeeze(0).to(dtype=torch.float32) - - def extract_from_file(self, audio_path): - audio, sample_rate = load_audio(audio_path, sr=None, mono=False) - waveform = torch.from_numpy(audio).contiguous() - return self.extract_from_waveform(waveform, sample_rate) - - def extract_many(self, audio_paths): - return [self.extract_from_file(path) for path in audio_paths] - - -def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> MingSpeakerEmbeddingExtractor: +# Ming's bundle ships the CAM++ artifact as `campplus.onnx` in the model +# directory, which is the shared extractor's default. +def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> SpeakerEmbeddingExtractor: extractor = getattr(wrapper, "_speaker_extractor", None) if extractor is not None: return extractor @@ -79,7 +37,7 @@ def _load_speaker_extractor(wrapper: MingTTSForConditionalGeneration) -> MingSpe extractor = getattr(wrapper, "_speaker_extractor", None) if extractor is None: model_path = wrapper.vllm_config.model_config.model - extractor = MingSpeakerEmbeddingExtractor(model_path, target_sr=16000) + extractor = SpeakerEmbeddingExtractor(model_path, target_sr=16000) wrapper._speaker_extractor = extractor return extractor @@ -174,8 +132,6 @@ def _resolve_speaker_embeddings( __all__ = [ - "MingSpeakerEmbeddingExtractor", "_load_speaker_extractor", "_resolve_speaker_embeddings", - "resolve_model_to_local_path", ] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py index fb6738a29..6e8677cf2 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py @@ -1,8 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # NOTE: Do not import model classes in this file. Importing any # submodule in this package triggers __init__.py execution, and # both the model registry and pipeline registry import submodules # directly — heavy imports here would be loaded as a side effect # even though nothing depends on these re-exports. + +# One MiniCPMTTS.generate_chunk emits 25 codec frames plus its terminator. +MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 41223e864..2bffeb49e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Strict, state-explicit batching for MiniCPM-o 4.5 Token2wav.""" from __future__ import annotations @@ -7,7 +7,7 @@ from __future__ import annotations from collections.abc import Mapping from contextlib import nullcontext from dataclasses import dataclass -from typing import Any +from typing import Any, cast import torch import torch.nn as nn @@ -19,6 +19,65 @@ from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper logger = init_logger(__name__) _SILENCE_TOKEN = 4218 +# CosyVoice2 RelPos PE is built for max_len=5000 and `forward_chunk` never +# calls ``extend_pe``. After the 2x upsample a single Code2Wav prefill longer +# than this many codec tokens overflows the PE slice (matrix_ac vs matrix_bd). +_DEFAULT_RELPOS_MAX_POS = 5000 +_MAX_ENCODE_TOKEN_CAP = 1024 + + +def relpos_encode_token_budget( + *, + max_pos: int, + stride: int, + cache_offset: int, + lookahead: int, + cap: int = _MAX_ENCODE_TOKEN_CAP, +) -> int: + """How many codec tokens ``forward_chunk`` can take before RelPos PE wraps. + + ``position_encoding(size)`` needs ``size <= max_pos``. After upsample, + ``size = cache_offset * stride + ~stride * token_frames`` (plus last-chunk + lookahead pad). + """ + stride = max(1, int(stride)) + lookahead = max(0, int(lookahead)) + room = int(max_pos) // stride - max(0, int(cache_offset)) - lookahead - 1 + return max(lookahead + 1, min(int(cap), room)) + + +def plan_token2wav_encode_slices( + num_frames: int, + *, + max_frames: int, + min_nonfinal: int, + last_chunk: bool, +) -> list[tuple[int, int]]: + """Split a Code2Wav token span so every non-final piece has a lookahead window.""" + if num_frames <= 0: + return [] + min_nonfinal = max(1, int(min_nonfinal)) + max_frames = max(int(max_frames), min_nonfinal) + last_min = 1 if last_chunk else min_nonfinal + if num_frames <= max_frames: + return [(0, num_frames)] + + slices: list[tuple[int, int]] = [] + start = 0 + while start < num_frames: + remaining = num_frames - start + if remaining <= max_frames: + slices.append((start, num_frames)) + break + take = max_frames + tail = remaining - take + if tail <= max_frames and tail < last_min: + take = remaining - last_min + if take < min_nonfinal: + take = remaining + slices.append((start, start + take)) + start += take + return slices def _autocast_disabled(device: torch.device): @@ -35,6 +94,17 @@ def _autocast_disabled(device: torch.device): return nullcontext() +def _token2wav_sdpa_context(device: torch.device): + if device.type != "npu": + return nullcontext() + + from vllm_omni.platforms.npu.models.step_audio2_token2wav import ( + npu_token2wav_sdpa_context, + ) + + return npu_token2wav_sdpa_context() + + def tensor_signature(value: torch.Tensor) -> tuple[tuple[int, ...], str, str]: return tuple(value.shape), str(value.dtype), value.device.type @@ -58,6 +128,25 @@ class BatchedToken2WavState: hift_cache: dict[str, torch.Tensor] +def _undecorate_dynamo(module: nn.Module, method: str) -> None: + """Restore ``method`` on ``module`` if TorchDynamo wrapped it. + + ``cosyvoice2`` decorates ``UpsampleConformerEncoderV2.forward_chunk`` with + ``torch.compile(backend="eager")``. That backend performs no Inductor + optimisation, so the wrapper only adds tracing and guard construction, and + duplex pays it again on every unseen chunk shape -- seconds inside a live + response. Dropping the wrapper leaves the original implementation, which is + what the eager backend was executing anyway. + """ + bound = getattr(module, method, None) + original = getattr(bound, "_torchdynamo_orig_callable", None) or getattr(bound, "__wrapped__", None) + if original is None: + return + function = getattr(original, "__func__", original) + module.__dict__[method] = function.__get__(module, type(module)) + logger.info("Bypassed TorchDynamo wrapper on %s.%s", type(module).__name__, method) + + class BatchedToken2Wav(nn.Module): """Drive Token2wav's modules with dynamically-sized, request-owned caches. @@ -74,6 +163,7 @@ class BatchedToken2Wav(nn.Module): connector_config: Mapping[str, int] | None = None, hift_graph_config: Mapping[str, Any] | None = None, cfm_graph_config: Mapping[str, Any] | None = None, + bfloat16_attention_cache: bool = False, ): super().__init__() self._token2wav = token2wav @@ -82,6 +172,9 @@ class BatchedToken2Wav(nn.Module): self._trt_stepper = trt_stepper self.flow = token2wav.flow self.hift = token2wav.hift + encoder = getattr(self.flow, "encoder", None) + if encoder is not None: + _undecorate_dynamo(encoder, "forward_chunk") hift_parameter = next(self.hift.parameters(), None) if hift_parameter is not None and hift_parameter.device.type == "cuda": # Prime the CUDA state used by HiFT during backend construction. @@ -104,6 +197,10 @@ class BatchedToken2Wav(nn.Module): del speech, source torch.accelerator.empty_cache() self.float16 = bool(token2wav.float16) + self._estimator_att_compute_dtype = torch.float16 if self.float16 else torch.float32 + self._estimator_att_cache_dtype = ( + torch.bfloat16 if bfloat16_attention_cache else self._estimator_att_compute_dtype + ) self.n_timesteps = int(token2wav.n_timesteps) self.mel_cache_len = int(token2wav.mel_cache_len) self.source_cache_len = int(token2wav.source_cache_len) @@ -217,6 +314,39 @@ class BatchedToken2Wav(nn.Module): width = getattr(layer, "pre_lookahead_len", None) return int(width) if width is not None else None + def _relpos_max_pos(self) -> int: + embed = getattr(self.flow.encoder, "embed", None) + pe = getattr(embed, "pe", None) + if isinstance(pe, torch.Tensor) and pe.numel() > 0: + return max(1, int(pe.size(1) // 2)) + return _DEFAULT_RELPOS_MAX_POS + + def _upsample_stride(self) -> int: + stride = getattr(getattr(self.flow.encoder, "up_layer", None), "stride", None) + return max(1, int(stride)) if stride is not None else 2 + + def _max_encode_token_frames(self, states: list[BatchedToken2WavState]) -> int: + att = states[0].flow_cache.get("conformer_att_cache") if states else None + offset1 = int(att.shape[3] // 2) if att is not None else 0 + lookahead = self._pre_lookahead_len() or 0 + return relpos_encode_token_budget( + max_pos=self._relpos_max_pos(), + stride=self._upsample_stride(), + cache_offset=offset1, + lookahead=lookahead, + ) + + def _ensure_relpos_pe(self, tokens: torch.Tensor, att_cache: torch.Tensor | None) -> None: + """Grow CosyVoice RelPos PE before ``forward_chunk``, which never calls extend_pe.""" + embed = getattr(self.flow.encoder, "embed", None) + extend_pe = getattr(embed, "extend_pe", None) + if not callable(extend_pe): + return + offset1 = int(att_cache.shape[3] // 2) if att_cache is not None else 0 + lookahead = self._pre_lookahead_len() or 0 + needed = offset1 * self._upsample_stride() + self._upsample_stride() * (int(tokens.shape[1]) + lookahead + 1) + extend_pe(tokens.new_zeros((1, max(needed, 1)))) + def _encode_chunk( self, tokens: torch.Tensor, @@ -225,6 +355,7 @@ class BatchedToken2Wav(nn.Module): cnn_cache: torch.Tensor | None, att_cache: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + self._ensure_relpos_pe(tokens, att_cache) embedded = self.flow.input_embedding(tokens) hidden, new_cnn, new_att = self.flow.encoder.forward_chunk( xs=embedded, @@ -265,8 +396,10 @@ class BatchedToken2Wav(nn.Module): cond: torch.Tensor, cnn_cache: torch.Tensor | None, att_cache: torch.Tensor | None, + attn_mask: torch.Tensor | None = None, + valid_lengths: list[int] | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if self._trt_stepper is not None: + if self._trt_stepper is not None and valid_lengths is None: out, new_cnn, new_att = self._trt_stepper.step( x=x, mu=mu, @@ -282,32 +415,139 @@ class BatchedToken2Wav(nn.Module): speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) - if self._cfm_graph_wrapper is not None: - if cnn_cache is None: - cnn_cache = torch.zeros_like(cnn_out) - if att_cache is None: - att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) + old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) + if isinstance(old_att, torch.Tensor) and old_att.dtype != estimator_input.dtype: + old_att = old_att.to(dtype=estimator_input.dtype) + if self._cfm_graph_wrapper is not None and valid_lengths is None: + graph_cnn = torch.zeros_like(cnn_out) if cnn_cache is None else cnn_cache + graph_att = ( + estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:]) + if att_cache is None + else old_att + ) return self._cfm_graph_wrapper.replay( estimator_input, time_embedding, - cnn_cache, - att_cache, + graph_cnn, + graph_att, + cnn_out, + att_out, + ) + if valid_lengths is not None: + if not hasattr(estimator, "in_proj"): + raise RuntimeError('MiniCPMO45Code2WavBatchError {"reason":"ragged_kernel_unavailable"}') + result = self._blocks_forward_chunk_ragged( + estimator, + estimator_input, + time_embedding, + attn_mask, + old_cnn, + old_att, + cnn_out, + att_out, + valid_lengths, + ) + else: + result = estimator.blocks_forward_chunk( + estimator_input, + time_embedding, + attn_mask, + old_cnn, + old_att, cnn_out, att_out, ) - old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) - old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks) - result = estimator.blocks_forward_chunk( - estimator_input, - time_embedding, - None, - old_cnn, - old_att, - cnn_out, - att_out, - ) return result, cnn_out, att_out + @staticmethod + def _gather_causal_cache( + history: torch.Tensor, + valid_lengths: torch.Tensor, + width: int, + ) -> torch.Tensor: + indices = valid_lengths[:, None] + torch.arange(width, device=history.device)[None, :] + return history.gather( + 2, + indices[:, None, :].expand(-1, int(history.shape[1]), -1), + ) + + def _blocks_forward_chunk_ragged( + self, + estimator: nn.Module, + estimator_input: torch.Tensor, + time_embedding: torch.Tensor, + attn_mask: torch.Tensor | None, + cnn_cache: Any, + att_cache: Any, + cnn_cache_buffer: torch.Tensor, + att_cache_buffer: torch.Tensor, + valid_lengths: list[int], + ) -> torch.Tensor: + """Run one padded DiT batch while capturing exact per-row CNN state.""" + lengths = torch.tensor( + (*valid_lengths, *valid_lengths), + device=estimator_input.device, + dtype=torch.long, + ) + x = estimator.in_proj(estimator_input.transpose(1, 2)) + for block_index, block in enumerate(estimator.blocks): + ( + shift_msa, + scale_msa, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + shift_conv, + scale_conv, + gate_conv, + ) = block.adaLN_modulation(time_embedding).chunk(9, dim=-1) + + normalized = block.norm1(x) * (1 + scale_msa) + shift_msa + x_att, new_att = block.attn.forward_chunk( + normalized, + att_cache[block_index], + attn_mask, + ) + x = x + gate_msa * x_att + + conv_input = block.norm3(x) * (1 + scale_conv) + shift_conv + old_cnn = cnn_cache[block_index] + if old_cnn is None: + width = int(block.conv.block[1].causal_padding[0]) + old_cnn = conv_input.new_zeros( + ( + int(conv_input.shape[0]), + int(block.conv.in_channels + block.conv.out_channels), + width, + ) + ) + old_cnn1, old_cnn2 = old_cnn.split((block.conv.in_channels, block.conv.out_channels), dim=1) + + conv1_input = block.conv.block[0](conv_input) + conv1_output, _ = block.conv.block[1].forward_chunk(conv1_input, old_cnn1) + new_cnn1 = self._gather_causal_cache( + torch.cat((old_cnn1, conv1_input), dim=2), + lengths, + int(old_cnn1.shape[2]), + ) + + conv2_input = block.conv.block[2:6](conv1_output) + conv2_output, _ = block.conv.block[6].forward_chunk(conv2_input, old_cnn2) + new_cnn2 = self._gather_causal_cache( + torch.cat((old_cnn2, conv2_input), dim=2), + lengths, + int(old_cnn2.shape[2]), + ) + x = x + gate_conv * block.conv.block[7](conv2_output) + x = x + gate_mlp * block.mlp(block.norm2(x) * (1 + scale_mlp) + shift_mlp) + + cnn_cache_buffer[block_index].copy_(torch.cat((new_cnn1, new_cnn2), dim=1)) + att_cache_buffer[block_index][:, :, : int(new_att.shape[2]), :].copy_(new_att) + + return estimator.final_layer(x, time_embedding).transpose(1, 2) + def _decode_cfm( self, mu: torch.Tensor, @@ -316,7 +556,12 @@ class BatchedToken2Wav(nn.Module): *, cnn_cache: torch.Tensor | None, att_cache: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + valid_lengths: list[int] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor | list[torch.Tensor], + ]: decoder = self.flow.decoder estimator = decoder.estimator batch_size = int(mu.shape[0]) @@ -341,36 +586,114 @@ class BatchedToken2Wav(nn.Module): mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0) speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0) cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0) - next_cnn: list[torch.Tensor] = [] - next_att: list[torch.Tensor] = [] - dt = timeline[1] - timeline[0] - for step in range(self.n_timesteps): - old_cnn = cnn_cache[step] if cnn_cache is not None else None - old_att = att_cache[step] if att_cache is not None else None - estimate, step_cnn, step_att = self._estimator_step( - estimator, - x=torch.cat((x, x), dim=0), - mu=mu_cfg, - time=torch.cat((time, time), dim=0), - speakers=speakers_cfg, - cond=cond_cfg, - cnn_cache=old_cnn, - att_cache=old_att, + attn_mask = None + if valid_lengths is not None: + if len(valid_lengths) != batch_size: + raise ValueError(f"valid length count {len(valid_lengths)} != batch {batch_size}") + cfg_lengths = torch.tensor( + (*valid_lengths, *valid_lengths), + device=mu.device, ) - conditional, unconditional = estimate.split(batch_size, dim=0) - velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional - x = x + dt * velocity - time = time + dt - if step + 1 < self.n_timesteps: - dt = timeline[step + 2] - time[0] - next_cnn.append(step_cnn) - next_att.append(step_att) - return x, torch.stack(next_cnn), torch.stack(next_att) + positions = torch.arange(int(mu.shape[2]), device=mu.device) + valid_queries = positions.unsqueeze(0) < cfg_lengths.unsqueeze(1) + current_keys = valid_queries.unsqueeze(1).expand(-1, int(mu.shape[2]), -1) + old_keys = torch.ones( + (2 * batch_size, int(mu.shape[2]), offset), + dtype=torch.bool, + device=mu.device, + ) + attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2) + next_cnn: list[torch.Tensor] = [] + next_att_cache: torch.Tensor | None = None + ragged_att_cache: list[torch.Tensor] | None = None + dt = timeline[1] - timeline[0] + with _token2wav_sdpa_context(mu.device): + for step in range(self.n_timesteps): + old_cnn = cnn_cache[step] if cnn_cache is not None else None + old_att = att_cache[step] if att_cache is not None else None + estimate, step_cnn, step_att = self._estimator_step( + estimator, + x=torch.cat((x, x), dim=0), + mu=mu_cfg, + time=torch.cat((time, time), dim=0), + speakers=speakers_cfg, + cond=cond_cfg, + cnn_cache=old_cnn, + att_cache=old_att, + attn_mask=attn_mask, + valid_lengths=valid_lengths, + ) + conditional, unconditional = estimate.split(batch_size, dim=0) + velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional + x = x + dt * velocity + time = time + dt + if step + 1 < self.n_timesteps: + dt = timeline[step + 2] - time[0] + next_cnn.append(step_cnn) + if valid_lengths is not None: + if ragged_att_cache is None: + ragged_att_cache = [ + torch.empty( + ( + self.n_timesteps, + int(step_att.shape[0]), + 2, + int(step_att.shape[2]), + valid_length + offset, + int(step_att.shape[4]), + ), + device=step_att.device, + dtype=self._estimator_att_cache_dtype, + ) + for valid_length in valid_lengths + ] + current_width = int(mu.shape[2]) + for row, (valid_length, row_cache) in enumerate(zip(valid_lengths, ragged_att_cache, strict=True)): + for cfg_row, source_row in enumerate((row, batch_size + row)): + row_cache[step, :, cfg_row, :, :valid_length].copy_( + step_att[:, source_row, :, :valid_length] + ) + if offset: + row_cache[step, :, cfg_row, :, valid_length:].copy_( + step_att[:, source_row, :, current_width:] + ) + del step_att + continue + if next_att_cache is None: + next_att_cache = torch.empty( + (self.n_timesteps, *step_att.shape), + device=step_att.device, + dtype=self._estimator_att_cache_dtype, + ) + next_att_cache[step].copy_(step_att) + del step_att + if ragged_att_cache is not None: + return x, torch.stack(next_cnn), ragged_att_cache + assert next_att_cache is not None + return x, torch.stack(next_cnn), next_att_cache - @staticmethod - def _split_flow_cache(cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]: + def _split_flow_cache(self, cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]: result: list[dict[str, torch.Tensor]] = [] for row in range(batch_size): + estimator_att = cache["estimator_att_cache"] + if estimator_att.dtype == self._estimator_att_cache_dtype: + request_att = torch.cat( + ( + estimator_att[:, :, row : row + 1], + estimator_att[:, :, batch_size + row : batch_size + row + 1], + ), + dim=2, + ).detach() + else: + request_shape = list(estimator_att.shape) + request_shape[2] = 2 + request_att = torch.empty( + request_shape, + device=estimator_att.device, + dtype=self._estimator_att_cache_dtype, + ) + request_att[:, :, 0:1].copy_(estimator_att[:, :, row : row + 1]) + request_att[:, :, 1:2].copy_(estimator_att[:, :, batch_size + row : batch_size + row + 1]) result.append( { "conformer_cnn_cache": cache["conformer_cnn_cache"][row : row + 1].detach().clone(), @@ -382,29 +705,23 @@ class BatchedToken2Wav(nn.Module): ), dim=2, ).detach(), - "estimator_att_cache": torch.cat( - ( - cache["estimator_att_cache"][:, :, row : row + 1], - cache["estimator_att_cache"][:, :, batch_size + row : batch_size + row + 1], - ), - dim=2, - ).detach(), + "estimator_att_cache": request_att, } ) return result - @staticmethod - def _stack_flow_cache(states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]: + def _stack_flow_cache(self, states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]: flows = [state.flow_cache for state in states] conditional_cnn = [flow["estimator_cnn_cache"][:, :, 0:1] for flow in flows] unconditional_cnn = [flow["estimator_cnn_cache"][:, :, 1:2] for flow in flows] conditional_att = [flow["estimator_att_cache"][:, :, 0:1] for flow in flows] unconditional_att = [flow["estimator_att_cache"][:, :, 1:2] for flow in flows] + estimator_att = torch.cat((*conditional_att, *unconditional_att), dim=2) return { "conformer_cnn_cache": torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0), "conformer_att_cache": torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1), "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2), - "estimator_att_cache": torch.cat((*conditional_att, *unconditional_att), dim=2), + "estimator_att_cache": estimator_att, } def setup_batch( @@ -488,14 +805,68 @@ class BatchedToken2Wav(nn.Module): # must carry at least one full kernel. Only the final chunk is allowed # to be shorter: ``forward_chunk`` zero-pads it by the lookahead width. lookahead = self._pre_lookahead_len() + num_frames = int(tokens.shape[1]) if lookahead is not None and not last_chunk: - num_frames = int(tokens.shape[1]) if num_frames <= lookahead: raise RuntimeError( "MiniCPMO45Code2WavBatchError " f'{{"reason":"chunk_below_lookahead_window","frames":{num_frames},' f'"minimum":{lookahead + 1}}}' ) + # A non-async Talker dump can land thousands of codec tokens in one + # Code2Wav prefill. CosyVoice RelPos PE (max_len=5000) plus 2x upsample + # cannot score that in one ``forward_chunk`` (6968 vs 985 on NPU). + max_frames = self._max_encode_token_frames(states) + slices = plan_token2wav_encode_slices( + num_frames, + max_frames=max_frames, + min_nonfinal=(lookahead + 1) if lookahead is not None else 1, + last_chunk=last_chunk, + ) + if len(slices) > 1: + logger.info( + "MiniCPM-o Code2Wav splitting %d codec tokens into %d encoder windows " + "(max_frames=%d) to stay inside RelPos PE.", + num_frames, + len(slices), + max_frames, + ) + parts: list[list[torch.Tensor]] = [[] for _ in range(batch_size)] + current = states + for index, (start, end) in enumerate(slices): + is_last_piece = index == len(slices) - 1 + audios, current = self._decode_batch_once( + tokens[:, start:end], + features, + current, + last_chunk=last_chunk and is_last_piece, + flush_encoder=flush_encoder and is_last_piece, + ) + for row, audio in enumerate(audios): + parts[row].append(audio) + merged = [ + torch.cat(row_parts) if row_parts else tokens.new_zeros((0,), dtype=torch.float32) + for row_parts in parts + ] + return merged, current + return self._decode_batch_once( + tokens, + features, + states, + last_chunk=last_chunk, + flush_encoder=flush_encoder, + ) + + def _decode_batch_once( + self, + tokens: torch.Tensor, + features: PromptFeatures, + states: list[BatchedToken2WavState], + *, + last_chunk: bool, + flush_encoder: bool = False, + ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: + batch_size = int(tokens.shape[0]) flow_cache = self._stack_flow_cache(states) speakers = features.speaker_embedding.expand(batch_size, -1) with self._autocast(tokens.device): @@ -558,3 +929,188 @@ class BatchedToken2Wav(nn.Module): ] audios = [emitted[row].reshape(-1).to(dtype=torch.float32) for row in range(batch_size)] return audios, next_states + + @staticmethod + def _require_complete_ragged_outputs( + audios: list[torch.Tensor | None], + next_states: list[BatchedToken2WavState | None], + ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: + missing_audio_rows = [row for row, audio in enumerate(audios) if audio is None] + missing_state_rows = [row for row, state in enumerate(next_states) if state is None] + if missing_audio_rows or missing_state_rows: + raise RuntimeError( + "MiniCPMO45Code2WavBatchError " + f'{{"reason":"incomplete_ragged_output","audio_rows":{missing_audio_rows},' + f'"state_rows":{missing_state_rows}}}' + ) + return ( + cast(list[torch.Tensor], audios), + cast(list[BatchedToken2WavState], next_states), + ) + + def decode_ragged_batch( + self, + tokens: list[torch.Tensor], + features: PromptFeatures, + states: list[BatchedToken2WavState], + *, + last_chunks: list[bool], + ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: + batch_size = len(tokens) + if batch_size != len(states) or batch_size != len(last_chunks): + raise ValueError("ragged token, state, and final-flag batches must have the same size") + if batch_size == 0: + return [], [] + + lookahead = self._pre_lookahead_len() + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + num_frames = int(row_tokens.numel()) + if lookahead is not None and not last_chunk and num_frames <= lookahead: + raise RuntimeError( + "MiniCPMO45Code2WavBatchError " + f'{{"reason":"chunk_below_lookahead_window","row":{row},' + f'"frames":{num_frames},"minimum":{lookahead + 1}}}' + ) + + max_frames = self._max_encode_token_frames(states) + if any(int(row_tokens.numel()) > max_frames for row_tokens in tokens): + # decode_batch owns the RelPos-safe encoder slicing used by + # non-async Talker dumps. Preserve that path for overlong rows; + # ordinary short duplex chunks continue through one padded DiT. + exact_groups: dict[tuple[int, bool], list[int]] = {} + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + exact_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row) + audios: list[torch.Tensor | None] = [None] * batch_size + next_states: list[BatchedToken2WavState | None] = [None] * batch_size + for (_, last_chunk), rows in exact_groups.items(): + group_audio, group_states = self.decode_batch( + torch.stack([tokens[row] for row in rows], dim=0), + features, + [states[row] for row in rows], + last_chunk=last_chunk, + ) + for group_row, row in enumerate(rows): + audios[row] = group_audio[group_row] + next_states[row] = group_states[group_row] + return self._require_complete_ragged_outputs(audios, next_states) + encoder_groups: dict[tuple[int, bool], list[int]] = {} + for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): + encoder_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row) + + hidden_rows: list[torch.Tensor | None] = [None] * batch_size + conformer_cnn_rows: list[torch.Tensor | None] = [None] * batch_size + conformer_att_rows: list[torch.Tensor | None] = [None] * batch_size + for (_, last_chunk), rows in encoder_groups.items(): + group_states = [states[row] for row in rows] + group_cache = self._stack_flow_cache(group_states) + group_tokens = torch.stack([tokens[row] for row in rows], dim=0) + with self._autocast(group_tokens.device): + hidden, conformer_cnn, conformer_att = self._encode_chunk( + group_tokens, + last_chunk=last_chunk, + cnn_cache=group_cache["conformer_cnn_cache"], + att_cache=group_cache["conformer_att_cache"], + ) + for group_row, row in enumerate(rows): + hidden_rows[row] = hidden[group_row : group_row + 1] + conformer_cnn_rows[row] = conformer_cnn[group_row : group_row + 1] + conformer_att_rows[row] = conformer_att[:, group_row : group_row + 1] + + missing_hidden_rows = [row for row, value in enumerate(hidden_rows) if value is None] + if missing_hidden_rows: + raise RuntimeError( + f'MiniCPMO45Code2WavBatchError {{"reason":"incomplete_encoder_output","rows":{missing_hidden_rows}}}' + ) + resolved_hidden = cast(list[torch.Tensor], hidden_rows) + hidden_lengths = [int(value.shape[1]) for value in resolved_hidden] + max_hidden_length = max(hidden_lengths) + padded_hidden = resolved_hidden[0].new_zeros((batch_size, max_hidden_length, int(resolved_hidden[0].shape[2]))) + for row, hidden in enumerate(resolved_hidden): + padded_hidden[row, : int(hidden.shape[1])].copy_(hidden[0]) + + flow_cache = self._stack_flow_cache(states) + speakers = features.speaker_embedding.expand(batch_size, -1) + with self._autocast(padded_hidden.device): + projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + cond = torch.zeros_like(padded_hidden).transpose(1, 2).contiguous() + chunk_mel, estimator_cnn, estimator_att = self._decode_cfm( + padded_hidden.transpose(1, 2).contiguous(), + projected_speakers, + cond, + cnn_cache=flow_cache["estimator_cnn_cache"], + att_cache=flow_cache["estimator_att_cache"], + valid_lengths=hidden_lengths, + ) + + prompt_len = int(features.mels.shape[1]) + assert isinstance(estimator_att, list) + new_flow: list[dict[str, torch.Tensor]] = [] + for row, row_estimator_att in enumerate(estimator_att): + conformer_cnn = conformer_cnn_rows[row] + conformer_att = conformer_att_rows[row] + assert conformer_cnn is not None and conformer_att is not None + if row_estimator_att.shape[4] > prompt_len + 100: + row_estimator_att = torch.cat( + ( + row_estimator_att[..., :prompt_len, :], + row_estimator_att[..., -100:, :], + ), + dim=4, + ) + if conformer_att.shape[3] > prompt_len + 100: + conformer_att = torch.cat( + ( + conformer_att[..., :prompt_len, :], + conformer_att[..., -100:, :], + ), + dim=3, + ) + new_flow.append( + { + "conformer_cnn_cache": conformer_cnn.detach().clone(), + "conformer_att_cache": conformer_att.detach().clone(), + "estimator_cnn_cache": torch.cat( + ( + estimator_cnn[:, :, row : row + 1], + estimator_cnn[:, :, batch_size + row : batch_size + row + 1], + ), + dim=2, + ).detach(), + "estimator_att_cache": row_estimator_att, + } + ) + + audios: list[torch.Tensor | None] = [None] * batch_size + next_states: list[BatchedToken2WavState | None] = [None] * batch_size + vocoder_groups: dict[tuple[int, bool], list[int]] = {} + for row, last_chunk in enumerate(last_chunks): + vocoder_groups.setdefault((hidden_lengths[row], last_chunk), []).append(row) + for (hidden_length, last_chunk), rows in vocoder_groups.items(): + old_mel = torch.cat([states[row].hift_cache["mel"] for row in rows], dim=0) + old_source = torch.cat([states[row].hift_cache["source"] for row in rows], dim=0) + old_speech = torch.cat([states[row].hift_cache["speech"] for row in rows], dim=0) + group_mel = torch.cat( + [chunk_mel[row : row + 1, :, :hidden_length] for row in rows], + dim=0, + ) + mel = torch.cat((old_mel, group_mel), dim=2) + speech, source = self._hift_inference(mel, old_source) + if old_speech.shape[-1] > 0: + window = self.speech_window.to(device=speech.device, dtype=speech.dtype) + speech = self._fade_in_out(speech, old_speech, window) + next_hift = { + "mel": mel[..., -self.mel_cache_len :].detach(), + "source": source[..., -self.source_cache_len :].detach(), + "speech": speech[..., -self.source_cache_len :].detach(), + } + emitted = speech if last_chunk else speech[..., : -self.source_cache_len] + for group_row, row in enumerate(rows): + audios[row] = emitted[group_row].reshape(-1).to(dtype=torch.float32) + next_states[row] = BatchedToken2WavState( + flow_cache=new_flow[row], + hift_cache={ + name: value[group_row : group_row + 1].detach().clone() for name, value in next_hift.items() + }, + ) + + return self._require_complete_ragged_outputs(audios, next_states) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index 97015855e..d737905fd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from functools import lru_cache +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import torch from torch.cuda import CUDAGraph @@ -141,19 +139,44 @@ def _tensor_signature(value: torch.Tensor) -> tuple: return tuple(value.shape), str(value.dtype), str(value.device) -_DTYPE_MAP = { - "torch.float32": torch.float32, - "torch.float16": torch.float16, - "torch.bfloat16": torch.bfloat16, - "torch.float64": torch.float64, -} +_DTYPE_MAP = {str(dtype): dtype for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.float64)} + + +def _memory_snapshot(device: torch.device) -> tuple[int, int] | None: + """(allocated, reserved) bytes, or None when the device cannot report them. + + Capture draws on the caching allocator, so these are the numbers that say + what a capture cost. Free device memory is not: the allocator serves a + capture out of memory it has already reserved, which is most of the device + on a normally configured worker. + """ + if device.type != "cuda": + return None + try: + return int(torch.accelerator.memory_allocated(device)), int(torch.accelerator.memory_reserved(device)) + except Exception: + return None + + +def _format_memory_delta(before: tuple[int, int] | None, after: tuple[int, int] | None) -> str: + if before is None or after is None: + return "" + mib = 1024 * 1024 + return ( + f" [allocated {after[0] / mib:.1f} MiB (+{(after[0] - before[0]) / mib:.1f}), " + f"reserved {after[1] / mib:.1f} MiB (+{(after[1] - before[1]) / mib:.1f})]" + ) def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]: - """Rebuild zero tensors from a cache key (shape, dtype, device tuples).""" + """Rebuild zero tensors from a cache key (shape, dtype, device tuples). + + Raises ``KeyError`` for a dtype the key cannot round-trip, so the caller + can fall back to eager instead of capturing wrong-dtype static buffers. + """ tensors = [] for shape, dtype_str, device_str in key[1:]: - dtype = _DTYPE_MAP.get(dtype_str, torch.float32) + dtype = _DTYPE_MAP[dtype_str] device = torch.device(device_str) tensors.append(torch.zeros(shape, dtype=dtype, device=device)) return tuple(tensors) @@ -166,56 +189,121 @@ class CFMGraphWrapper: as the graph target. The 10-step Euler loop stays in Python, replaying the graph 10 times per decode. - Uses functools.lru_cache for automatic LRU eviction. Cache misses - trigger capture; capture failures fall back to eager. Outputs are - cloned after replay to prevent streaming cache corruption. + Graphs are retired a whole generation at a time rather than one at a time. + Every capture shares one private memory pool, so a retired graph's blocks + return to that pool while its live peers still hold those addresses in + their recorded kernel arguments, and the next replay reads memory that now + belongs to something else. Retiring the whole generation bounds the cache + without ever leaving a live graph behind a freed one. + + Cache misses capture. A capture failure disables the wrapper for the rest + of the process, while a key whose static buffers cannot be rebuilt sends + only that one shape eager. Outputs are cloned after replay to prevent + streaming cache corruption. """ - def __init__(self, graph_fn, *, max_graphs: int = 32) -> None: + def __init__( + self, + graph_fn, + *, + max_graphs: int = 32, + ) -> None: self.graph_fn = graph_fn self.max_graphs = int(max_graphs) self.device = next(graph_fn.__self__.parameters()).device + # A non-positive budget means "no graphs", the same as + # `enable_cfm_graph: false`. Clamping to 1 would instead build a + # one-entry cache that flushes on every new shape. + self.enabled = self.max_graphs > 0 + self._cache: dict[tuple, tuple] = {} + # Shapes whose key cannot round-trip: eager for those, keep the rest. + self._unsupported: set[tuple] = set() + self._stats = { + "calls": 0, + "hits": 0, + "captures": 0, + "flushes": 0, + "eager": 0, + } - @lru_cache(maxsize=self.max_graphs) - def _capture_graph(key: tuple): - return self._capture(key) - - self._capture_graph = _capture_graph + def stats_snapshot(self) -> dict[str, int]: + """Bounded cumulative telemetry for the graph cache.""" + return {**self._stats, "cache_size": len(self._cache)} def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor: return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5]) + def _eager(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + self._stats["eager"] += 1 + with torch.no_grad(): + result = self._call_graph_fn(inputs) + return result, inputs[4], inputs[5] + + def _flush(self) -> None: + """Retire every captured graph at once. + + The sync keeps a replay from being in flight when the graphs go, and + the explicit ``reset()`` tears each one down here rather than whenever + Python drops the last reference. Nothing process-wide is touched: the + cuBLAS workspace is shared with graphs this wrapper does not own. + """ + if not self._cache: + return + torch.accelerator.synchronize(self.device) + for entry in self._cache.values(): + entry[2].reset() + self._cache.clear() + self._stats["flushes"] += 1 + logger.info("CFM graph cache flushed; stats=%s", self.stats_snapshot()) + + def _disable(self, reason: str, key: tuple) -> None: + logger.warning("Disabling CFM CUDA graphs (%s) for shape=%s; using eager", reason, key, exc_info=True) + self.enabled = False + self._flush() + def _capture(self, key: tuple) -> tuple | None: """Capture a CUDA graph for the given key. Returns None on failure.""" - static_inputs = _tensors_from_key(key) - - current_stream = torch.cuda.current_stream(self.device) - warmup_stream = torch.cuda.Stream(device=self.device) - warmup_stream.wait_stream(current_stream) - with torch.cuda.stream(warmup_stream), torch.no_grad(): - for _ in range(3): - self._call_graph_fn(static_inputs) - current_stream.wait_stream(warmup_stream) - try: + static_inputs = _tensors_from_key(key) + except KeyError: + # An unsupported dtype is a property of this shape alone, so run it + # eager and keep the other shapes on graphs. + logger.warning("CFM graph key carries an unsupported dtype: %s; using eager", key) + self._unsupported.add(key) + return None + + memory_before = _memory_snapshot(self.device) + try: + # Warmup runs the same kernels as the capture, so a fault here + # leaves the same dirty capture-stream and pool state, and must be + # handled the same way. + current_stream = torch.cuda.current_stream(self.device) + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(current_stream) + with torch.cuda.stream(warmup_stream), torch.no_grad(): + for _ in range(3): + warmup_output = self._call_graph_fn(static_inputs) + current_stream.wait_stream(warmup_stream) + del warmup_output + graph = CUDAGraph() with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): static_output = self._call_graph_fn(static_inputs) except Exception: - logger.warning( - "CFM graph capture failed for shape=%s; using eager", - key, - exc_info=True, - ) + # A failed capture can leave the capture stream current and the + # allocator still routing into the graph pool, so there is no safe + # way to keep capturing afterwards. + self._disable("capture failed", key) return None + self._stats["captures"] += 1 logger.info( - "Captured CFM CUDA Graph for shape %s (cache=%d/%d, hits=%d, misses=%d)", + "Captured CFM CUDA Graph for shape %s (cache=%d/%d, stats=%s)%s", key, - self._capture_graph.cache_info().currsize, + len(self._cache) + 1, self.max_graphs, - self._capture_graph.cache_info().hits, - self._capture_graph.cache_info().misses, + self.stats_snapshot(), + _format_memory_delta(memory_before, _memory_snapshot(self.device)), ) return (static_inputs, static_output, graph) @@ -229,25 +317,25 @@ class CFMGraphWrapper: att_out: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out) + self._stats["calls"] += 1 - if torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda": - with torch.no_grad(): - result = self._call_graph_fn(inputs) - return result, inputs[4], inputs[5] + if not self.enabled or torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda": + return self._eager(inputs) key = ("estimator_step",) + tuple(_tensor_signature(v) for v in inputs) - entry = self._capture_graph(key) + if key in self._unsupported: + return self._eager(inputs) + entry = self._cache.get(key) if entry is None: - logger.debug( - "CFM graph eager fallback for shape=%s (hits=%d, misses=%d)", - key, - self._capture_graph.cache_info().hits, - self._capture_graph.cache_info().misses, - ) - with torch.no_grad(): - result = self._call_graph_fn(inputs) - return result, inputs[4], inputs[5] + if len(self._cache) >= self.max_graphs: + self._flush() + entry = self._capture(key) + if entry is None: + return self._eager(inputs) + self._cache[key] = entry + else: + self._stats["hits"] += 1 static_inputs, static_output, graph = entry for static, current in zip(static_inputs, inputs, strict=True): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 686b39c5b..7017f703f 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Strict batched codec-to-waveform stage for MiniCPM-o 4.5.""" from __future__ import annotations @@ -69,6 +69,17 @@ def _codec_tensor(value: Any, fallback: torch.Tensor) -> torch.Tensor: # OmniGPUModelRunner._preprocess and the NPU _gather_runtime_additional_information # override). A step carrying only these has no producer payload at all. _RUNNER_STAMPED_KEYS = frozenset({"request_id", "req_id", "generated_len", "meta"}) +_PRODUCER_META_KEYS = frozenset( + { + "cache_epoch", + "chunk_seq", + "code_flat_numel", + "last_chunk", + "prompt_cache_id", + "prompt_wav", + "ref_audio_sr", + } +) def _carries_stage_payload(info: Mapping[str, Any], meta: Mapping[str, Any]) -> bool: @@ -77,9 +88,12 @@ def _carries_stage_payload(info: Mapping[str, Any], meta: Mapping[str, Any]) -> Any real async-chunk payload brings producer metadata along, whether the transport delivers it nested under ``meta`` or as flattened ``meta.*`` keys. """ + codes = info.get("codes") + if isinstance(codes, Mapping) and any(value is not None for value in codes.values()): + return True if any(key not in _RUNNER_STAMPED_KEYS for key in info): return True - return meta is not info and any(key not in _RUNNER_STAMPED_KEYS for key in meta) + return any(meta.get(key) is not None for key in _PRODUCER_META_KEYS) @dataclass(frozen=True) @@ -121,11 +135,12 @@ class _WorkItem: class MiniCPMO45Code2Wav(nn.Module): - """LLM_GENERATION model that admits only true exact-shape GPU batches.""" + """LLM_GENERATION model with request-owned state and compatible batching.""" input_modalities = "audio" have_multimodal_outputs = True enable_update_additional_information = True + replace_runtime_additional_information = True requires_raw_input_tokens = True requires_request_ids = True has_preprocess = False @@ -169,6 +184,11 @@ class MiniCPMO45Code2Wav(nn.Module): self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) if self._min_batch_size < 1: raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") + self._initial_batch_size = int(extra.get("code2wav_initial_batch_size", 0)) + if self._initial_batch_size < 0: + raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be >= 0") + if self._initial_batch_size and self._initial_batch_size < self._min_batch_size: + raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size") self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") @@ -252,6 +272,11 @@ class MiniCPMO45Code2Wav(nn.Module): ref_audio, meta.get("ref_audio_sr"), ) + logger.debug( + "MiniCPM-o Code2Wav selected runtime reference prompt_cache_id=%s prompt_wav=%s", + entry.cache_id, + entry.path, + ) return entry.cache_id, entry.path, cache_key if previous is not None: @@ -467,13 +492,63 @@ class MiniCPMO45Code2Wav(nn.Module): return ( item.prompt_cache_id, item.prompt_wav, - int(item.tokens.numel()), cache_signature, - item.last_chunk, - item.tts_is_last_chunk, item.cache_epoch, ) + def _iter_limited_batches( + self, + buckets: Iterable[list[_WorkItem]], + *, + maximum: int, + ) -> Iterable[list[_WorkItem]]: + for bucket in buckets: + if not bucket: + continue + if not maximum: + yield bucket + continue + batch_count = (len(bucket) + maximum - 1) // maximum + if len(bucket) < batch_count * self._min_batch_size: + raise _batch_error( + "initial_batch_partition_below_minimum", + size=len(bucket), + minimum=self._min_batch_size, + maximum=maximum, + ) + batch_size, larger_batches = divmod(len(bucket), batch_count) + start = 0 + for batch_index in range(batch_count): + stop = start + batch_size + (batch_index < larger_batches) + yield bucket[start:stop] + start = stop + + def _iter_decode_batches( + self, + buckets: Iterable[list[_WorkItem]], + ) -> Iterable[list[_WorkItem]]: + for bucket in buckets: + if not self._initial_batch_size: + yield bucket + continue + initial = [item for item in bucket if item.chunk_seq <= 1] + steady = [item for item in bucket if item.chunk_seq > 1] + + undersized = [ + {"wave": wave, "size": len(items)} + for wave, items in (("initial", initial), ("steady", steady)) + if items and len(items) < self._min_batch_size + ] + if undersized: + raise _batch_error( + "decode_wave_below_minimum", + minimum=self._min_batch_size, + waves=undersized, + ) + yield from self._iter_limited_batches([initial], maximum=self._initial_batch_size) + if steady: + yield steady + @torch.inference_mode() def forward( self, @@ -609,7 +684,10 @@ class MiniCPMO45Code2Wav(nn.Module): (item.prompt_cache_id, item.prompt_wav), [], ).append(item) - for bucket in initial_marker_buckets.values(): + for bucket in self._iter_limited_batches( + initial_marker_buckets.values(), + maximum=self._initial_batch_size, + ): try: features = self.backend.prepare_prompt( bucket[0].prompt_cache_id, @@ -641,7 +719,7 @@ class MiniCPMO45Code2Wav(nn.Module): prompt_wav=item.prompt_wav, token2wav=state, ) - for bucket in buckets.values(): + for bucket in self._iter_decode_batches(buckets.values()): batch_size = len(bucket) try: features = self.backend.prepare_prompt( @@ -652,13 +730,23 @@ class MiniCPMO45Code2Wav(nn.Module): states = self.backend.setup_batch(features, batch_size) else: states = [item.previous.token2wav for item in bucket if item.previous is not None] - tokens = torch.stack([item.tokens for item in bucket], dim=0) - audios, next_states = self.backend.decode_batch( - tokens, - features, - states, - last_chunk=bucket[0].last_chunk, - ) + token_lengths = {int(item.tokens.numel()) for item in bucket} + last_chunk_values = {item.last_chunk for item in bucket} + if len(token_lengths) > 1 or len(last_chunk_values) > 1: + audios, next_states = self.backend.decode_ragged_batch( + [item.tokens for item in bucket], + features, + states, + last_chunks=[item.last_chunk for item in bucket], + ) + else: + tokens = torch.stack([item.tokens for item in bucket], dim=0) + audios, next_states = self.backend.decode_batch( + tokens, + features, + states, + last_chunk=bucket[0].last_chunk, + ) except Exception as exc: self._prune_unowned_runtime_prompts() if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "): @@ -807,4 +895,5 @@ class MiniCPMO45Code2Wav(nn.Module): connector_config=self._connector_config, hift_graph_config=self._hift_graph_config, cfm_graph_config=self._cfm_graph_config, + bfloat16_attention_cache=bool(extra.get("code2wav_bfloat16_attention_cache", False)), ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index a182848cf..af9401872 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -17,6 +17,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect import math import os import warnings @@ -300,6 +301,10 @@ class ConditionalChatTTSConfig(PretrainedConfig): self.top_p = top_p self.top_k = top_k self.repetition_penalty = repetition_penalty + # Talker stage is a single-vocab codec LM: vLLM's Sampler and + # embed_tokens both address this table, not the Thinker text vocab. + self.vocab_size = int(num_audio_tokens) + self.eos_token_id = int(num_audio_tokens) - 1 class MiniCPMOConfig(Qwen2Config): @@ -2505,6 +2510,17 @@ def _in_projection( return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v) +def _get_audio_cache_length(past_key_values: Any) -> int: + cache = getattr(past_key_values, "self_attention_cache", past_key_values) + get_seq_length = getattr(cache, "get_seq_length", None) + if callable(get_seq_length): + try: + return int(get_seq_length()) + except TypeError: + return int(get_seq_length(0)) + return int(cache[0][0].shape[2]) + + # Copied from transformers.models.whisper.modeling_whisper.WhisperEncoderLayer and add use_cache for streaming inference class MiniCPMWhisperEncoderLayer(nn.Module): def __init__(self, config: WhisperConfig, layer_idx: int = None): @@ -2517,6 +2533,10 @@ class MiniCPMWhisperEncoderLayer(nn.Module): config=config, layer_idx=layer_idx, ) + attention_parameters = inspect.signature(self.self_attn.forward).parameters + self._past_key_values_kwarg = ( + "past_key_values" if "past_key_values" in attention_parameters else "past_key_value" + ) self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.dropout = config.dropout self.activation_fn = ACT2FN[config.activation_function] @@ -2554,16 +2574,18 @@ class MiniCPMWhisperEncoderLayer(nn.Module): """ residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) - attn_out = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - layer_head_mask=layer_head_mask, - output_attentions=output_attentions, - past_key_value=past_key_values, - ) + attention_kwargs = { + "hidden_states": hidden_states, + "attention_mask": attention_mask, + "layer_head_mask": layer_head_mask, + "output_attentions": output_attentions, + self._past_key_values_kwarg: past_key_values, + } + attn_out = self.self_attn(**attention_kwargs) hidden_states = attn_out[0] attn_weights = attn_out[1] if len(attn_out) > 1 else None - past_key_values = attn_out[2] if len(attn_out) > 2 else None + if len(attn_out) > 2: + past_key_values = attn_out[2] hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) hidden_states = residual + hidden_states @@ -2769,7 +2791,7 @@ class MiniCPMWhisperEncoder(WhisperEncoder): past_key_values = EncoderDecoderCache(past_key_values, DynamicCache()) else: pass - past_key_values_length = past_key_values.self_attention_cache.get_usable_length(inputs_embeds.shape[1]) + past_key_values_length = _get_audio_cache_length(past_key_values) if inputs_embeds.shape[1] + past_key_values_length > embed_pos.shape[0]: logger.warning("seems the audio is longer than 30s. repeating the last part of the audio") embed_pos_front = embed_pos[past_key_values_length:, :] @@ -4468,7 +4490,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S return [] if self.audio_past_key_values is not None: - cache_length = self.audio_past_key_values[0][0].shape[2] + cache_length = _get_audio_cache_length(self.audio_past_key_values) apm_max_len = self.apm.embed_positions.weight.shape[0] if cache_length + current_seq_len >= apm_max_len: logger.warning( @@ -4480,7 +4502,7 @@ class MiniCPMO45OmniLLMForConditionalGeneration(nn.Module, SupportsMultiModal, S past_len = 0 if self.audio_past_key_values is not None: - past_len = self.audio_past_key_values[0][0].shape[2] + past_len = _get_audio_cache_length(self.audio_past_key_values) total_seq_len = past_len + current_seq_len audio_attention_mask = torch.zeros( (batch_size, 1, current_seq_len, total_seq_len), diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 9e9dfe952..369219317 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from: # https://huggingface.co/openbmb/MiniCPM-o-4_5/blob/main/modeling_minicpmo.py """MiniCPM-o 4.5 native autoregressive Talker. @@ -8,11 +8,12 @@ Pipeline: 1. Receive thinker hidden_states + full token IDs via additional_information 2. Extract tts_bos..tts_eos region 3. Build condition: emb_text(tokens) + projector_semantic(hidden) (hidden_text_merge) - 4. Continuously generate request-aligned discrete audio-code deltas + 4. Project last hidden through head_code; vLLM Sampler picks the codec id + 5. Next decode embeds that id with emb_code and emits it to Code2Wav """ -from collections.abc import Iterable, Sequence -from dataclasses import dataclass +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from typing import Any import torch @@ -27,50 +28,48 @@ from vllm.model_executor.models.utils import maybe_prefix from vllm.v1.sample.sampler import Sampler from vllm_omni.experimental.fullduplex.engine.intermediate import get_tts_handoff +from vllm_omni.model_executor.models.minicpmo_4_5 import MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -_REPETITION_WINDOW = 16 _REPETITION_PENALTY_CHUNK_SIZE = 16 -_MIN_AUDIO_TOKENS = 64 -_MAX_AUDIO_TOKENS = 2048 -_AUDIO_TOKENS_PER_TEXT_TOKEN = 10 -# Codec-token sampling happens inside the model; vLLM sampling parameters -# only choose the Talker's binary continue/stop row. -_CODEC_SEED = 42 -_CODEC_TEMPERATURE = 0.8 -_CODEC_TOP_K = 25 -_CODEC_TOP_P = 0.85 -_CODEC_REPETITION_PENALTY = 1.05 -_CODEC_MIN_TOKENS = 50 -_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26 +# ``past_window`` of MiniCPMTTS's codec repetition penalty: both generate() and +# generate_chunk() build it through gen_logits(), which hardcodes +# CustomRepetitionPenaltyLogitsProcessorRepeat(penalty, num_code, 16). +_CODEC_PENALTY_WINDOW = 16 +# MiniCPMTTS.generate's max_new_token. The Talker context bounds this further; +# without it a request that never samples codec EOS keeps emitting frames for +# twice as long as upstream would, which is audible as a long silent tail. +_OFFLINE_CODEC_MAX_NEW_TOKENS = 2048 +# Native duplex Talker must finish after one MiniCPMTTS.generate_chunk: +# 25 codec frames (``codec_chunk_frames``) plus the terminating sample. +# Without this, the single-vocab Sampler keeps the stage-1 request alive +# until codec EOS / 4096 and Thinker never starts the next model turn. +_DUPLEX_CODEC_TOKENS_PER_CHUNK = MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK -@dataclass(slots=True) -class _PendingCodecSample: - output_index: int - hidden_row: torch.Tensor - codes: torch.Tensor - request_id: str - step: int - state: dict[str, Any] - info: dict[str, Any] +def _native_duplex_chunk_budget(meta: Mapping[str, Any] | None) -> tuple[int, int]: + """Return ``(max_tokens, min_tokens)`` for one native-duplex Talker request.""" + boundary = isinstance(meta, Mapping) and (bool(meta.get("turn_start")) or bool(meta.get("turn_end"))) + ceiling = _DUPLEX_CODEC_TOKENS_PER_CHUNK + return ceiling, 0 if boundary else ceiling -def _max_audio_tokens(condition_tokens: int) -> int: - """Bound codec generation with a conservative text-length estimate. +def blank_scheduler_prompt_for_penalties( + prompt_token_ids: torch.Tensor, + vocab_size: int, +) -> torch.Tensor: + """Return a penalty prompt whose every position is the pad id (``vocab_size``). - EOS is masked for the first 50 steps, so a direct ``text_tokens * 10`` - limit can terminate short responses before EOS is eligible. The 2048 - ceiling matches the checkpoint's native generation default and keeps the - sequence within the Talker's 4096-position context. + No Talker prompt position is codec history: prefill conditioning arrives as + embeddings from ``preprocess`` and decode embeds sampled ids with + ``emb_code``. The scheduler ids are placeholders (``llm2tts`` fills them + with ``0``, or with thinker token ids on the non-handoff path), so scoring + them would tax unrelated codec tokens. """ - return max( - _MIN_AUDIO_TOKENS, - min(_MAX_AUDIO_TOKENS, condition_tokens * _AUDIO_TOKENS_PER_TEXT_TOKEN), - ) + return torch.full_like(prompt_token_ids, int(vocab_size)) def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor: @@ -82,23 +81,31 @@ def _apply_batched_repetition_penalty( logits: torch.Tensor, histories: Sequence[torch.Tensor], *, - penalty: float, + penalty: float | torch.Tensor, window_size: int, ) -> torch.Tensor: - """Apply request-local frequency penalties to a batch of codec logits.""" + """Apply request-local frequency penalties to a batch of codec logits. + + ``penalty`` may be a scalar or one value per row, mirroring upstream's + per-request ``sampling_params.repetition_penalty``. + """ if logits.ndim != 2: raise ValueError(f"batched codec logits must be 2D, got shape {tuple(logits.shape)}") batch_size, vocab_size = logits.shape if len(histories) != batch_size: raise ValueError(f"expected {batch_size} codec histories, got {len(histories)}") - if penalty == 1.0: - return logits - if batch_size == 0: return logits + penalties = torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype).reshape(-1) + if penalties.numel() == 1: + penalties = penalties.expand(batch_size) + elif penalties.numel() != batch_size: + raise ValueError(f"expected 1 or {batch_size} codec repetition penalties, got {penalties.numel()}") + if not bool((penalties != 1.0).any()): + return logits + penalized = logits.clone() - penalty_tensor = torch.as_tensor(penalty, device=logits.device, dtype=logits.dtype) for start in range(0, batch_size, _REPETITION_PENALTY_CHUNK_SIZE): end = min(start + _REPETITION_PENALTY_CHUNK_SIZE, batch_size) chunk_logits = logits[start:end] @@ -116,38 +123,12 @@ def _apply_batched_repetition_penalty( encoded, minlength=(end - start) * vocab_size, ).reshape(end - start, vocab_size) - alpha = torch.pow(penalty_tensor, frequencies.to(dtype=logits.dtype)) + alpha = torch.pow(penalties[start:end].unsqueeze(1), frequencies.to(dtype=logits.dtype)) penalized[start:end] = torch.where(chunk_logits < 0, chunk_logits * alpha, chunk_logits / alpha) return penalized -def _apply_top_k_top_p( - logits: torch.Tensor, - *, - top_k: int | None, - top_p: float | None, - min_tokens_to_keep: int = 3, -) -> torch.Tensor: - """Apply the same candidate floors as the upstream Transformers warpers.""" - filtered = logits.clone() - vocab_size = filtered.shape[-1] - # MiniCPM-o's gen_logits() appends TopPLogitsWarper before - # TopKLogitsWarper. The order is observable for fixed-seed sampling. - if top_p is not None and 0.0 < top_p < 1.0: - sorted_logits, sorted_indices = torch.sort(filtered, descending=False, dim=-1) - cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) - remove = cumulative_probs <= (1.0 - float(top_p)) - remove[..., -min_tokens_to_keep:] = False - remove = remove.scatter(-1, sorted_indices, remove) - filtered.masked_fill_(remove, float("-inf")) - if top_k is not None and top_k > 0: - keep = min(vocab_size, max(int(top_k), min_tokens_to_keep)) - threshold = torch.topk(filtered, keep, dim=-1).values[..., -1, None] - filtered.masked_fill_(filtered < threshold, float("-inf")) - return filtered - - class _MiniCPMTTSProjector(nn.Module): """Checkpoint-compatible hidden-state projector used by MiniCPMTTS.""" @@ -173,9 +154,14 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): config: MiniCPMOConfig = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config - self._batch_stop_logits: torch.Tensor | None = None - self._request_generators: dict[str, torch.Generator] = {} + self._force_eos_rows: list[bool] | None = None + self._mask_eos_rows: list[bool] | None = None + self._pending_force_eos_rows: list[bool] | None = None + self._penalty_histories: list[torch.Tensor] | None = None self._request_audio_states: dict[str, dict[str, Any]] = {} + # Mirrors upstream TTSStreamingGenerator._chunk_info: one committed + # condition plus, during a rollover, one immutable recompute recipe. + self._request_condition_states: dict[str, dict[str, Any]] = {} self._deferred_cleanup_ids: set[str] = set() tts_config = getattr(config, "tts_config", None) @@ -186,23 +172,18 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self._tts_bos_id = getattr(tts_config, "audio_bos_token_id", 151687) self._text_eos_id = getattr(tts_config, "text_eos_token_id", 151692) self._num_audio_tokens = getattr(tts_config, "num_audio_tokens", 6562) + self._codec_eos_id = int(getattr(tts_config, "eos_token_id", self._num_audio_tokens - 1)) self._hidden_size = getattr(tts_config, "hidden_size", 768) self._normalize = getattr(tts_config, "normalize_projected_hidden", True) - self._codec_seed = int(getattr(tts_config, "seed", _CODEC_SEED)) - self._codec_temperature = float(getattr(tts_config, "temperature", _CODEC_TEMPERATURE)) - self._codec_top_k = int(getattr(tts_config, "top_k", _CODEC_TOP_K)) - self._codec_top_p = float(getattr(tts_config, "top_p", _CODEC_TOP_P)) - self._codec_repetition_penalty = float(getattr(tts_config, "repetition_penalty", _CODEC_REPETITION_PENALTY)) - self._codec_min_tokens = int(getattr(tts_config, "min_new_tokens", _CODEC_MIN_TOKENS)) else: self._tts_config = None + self._codec_eos_id = 0 self.has_preprocess = True self.has_postprocess = False - self.gpu_resident_buffer_keys: set[tuple[str, str]] = { - ("audio_codes", "current"), - ("audio_codes", "accumulated"), - } + # Same-step codes travel through make_omni_output from the previous + # sampled id (decode preprocess embeds that id via emb_code). + self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("codes", "audio")} self._init_native_talker(prefix) def _init_native_talker(self, prefix: str) -> None: @@ -285,6 +266,116 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): return torch.cat([condition, audio_bos], dim=0) return torch.cat([condition, self._boundary_embeddings()], dim=0) + def _build_streaming_recompute_embeddings( + self, + current_condition: torch.Tensor, + *, + request_id: str, + info_dict: Mapping[str, Any], + meta: Mapping[str, Any], + ) -> torch.Tensor: + """Return the official one-previous-chunk sliding-recompute window.""" + condition_seq = meta.get("streaming_condition_seq") + if not isinstance(condition_seq, int) or isinstance(condition_seq, bool): + if meta.get("streaming_prompt_recompute") is True: + raise ValueError("streaming prompt recompute is missing streaming_condition_seq") + # Direct model tests and non-connector callers do not participate in + # the persistent async-chunk lifecycle, so they need no window state. + return current_condition + + turn_start = bool(meta.get("turn_start")) + recompute = meta.get("streaming_prompt_recompute") is True + states = self._request_condition_states + state = states.get(request_id) + if turn_start: + if recompute: + raise ValueError("streaming prompt recompute cannot cross a native duplex turn boundary") + states[request_id] = { + "condition_seq": condition_seq, + "condition": current_condition.detach().clone(), + "base_recent_codes": (), + } + return current_condition + + if state is None: + if recompute: + raise ValueError("streaming prompt recompute is missing the previous Talker condition") + states[request_id] = { + "condition_seq": condition_seq, + "condition": current_condition.detach().clone(), + "base_recent_codes": (), + } + return current_condition + + previous_seq = state.get("condition_seq") + if not isinstance(previous_seq, int) or condition_seq < previous_seq: + raise ValueError( + f"stale native duplex Talker condition sequence: current={condition_seq}, previous={previous_seq}" + ) + if condition_seq > previous_seq + 1: + raise ValueError( + f"native duplex Talker skipped a condition sequence: current={condition_seq}, previous={previous_seq}" + ) + + if not recompute: + if condition_seq > previous_seq: + raise ValueError( + "a native duplex Talker condition advanced without its streaming recompute marker: " + f"current={condition_seq}, previous={previous_seq}" + ) + if "active_embeddings" in state: + raise ValueError("an active streaming recompute was replayed without its recompute marker") + return current_condition + + if condition_seq == previous_seq: + active_embeddings = state.get("active_embeddings") + if not isinstance(active_embeddings, torch.Tensor): + raise ValueError("streaming prompt window lost its cached recompute embeddings") + return active_embeddings + + if condition_seq != previous_seq + 1: + raise ValueError( + "streaming prompt recompute skipped a Talker condition: " + f"previous={previous_seq}, current={condition_seq}" + ) + + previous_condition = state.get("condition") + if not isinstance(previous_condition, torch.Tensor): + raise ValueError("streaming prompt recompute lost the previous Talker condition") + + ids = info_dict.get("ids") + previous_codes = ids.get("streaming_prompt_previous_codes") if isinstance(ids, Mapping) else None + if isinstance(previous_codes, torch.Tensor): + code_ids = previous_codes.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1) + elif isinstance(previous_codes, (list, tuple)): + code_ids = torch.as_tensor(previous_codes, device=self.emb_code[0].weight.device, dtype=torch.long) + else: + raise ValueError("streaming prompt recompute is missing confirmed codec ids") + if code_ids.numel() > _DUPLEX_CODEC_TOKENS_PER_CHUNK - 1: + raise ValueError(f"streaming prompt recompute has too many codec ids: {code_ids.numel()}") + if code_ids.numel() and bool(((code_ids < 0) | (code_ids >= self._codec_eos_id)).any()): + raise ValueError("streaming prompt recompute codec ids include an invalid or terminal token") + + parts = [previous_condition] + if code_ids.numel(): + parts.append(self.emb_code[0](code_ids)) + parts.append(current_condition) + full_embeddings = torch.cat(parts, dim=0) + previous_code_ids = tuple(int(code_id) for code_id in code_ids.tolist()) + previous_base_codes = state.get("base_recent_codes") + if not isinstance(previous_base_codes, tuple): + raise ValueError("streaming Talker condition lost its frozen codec history") + states[request_id] = { + "condition_seq": condition_seq, + "condition": current_condition.detach().clone(), + "active_embeddings": full_embeddings.detach().clone(), + # Official generate_with_buffer keeps all_generated_tokens across + # sliding recomputes, so the first sample in this chunk still sees + # the previous chunk's repetition-penalty window. + "base_recent_codes": (*previous_base_codes, *previous_code_ids)[-_CODEC_PENALTY_WINDOW:], + } + return full_embeddings + def preprocess( self, input_ids: torch.Tensor, @@ -297,6 +388,7 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): is_prefill = bool(info_dict.get("_omni_is_prefill", False)) state = info_dict.get("audio_state") first_call = not isinstance(state, dict) + request_id = str(info_dict.get("request_id", "0")) if is_prefill or first_call: token_ids, hidden_states = get_tts_handoff(info_dict) @@ -326,19 +418,41 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): info_dict.get("request_id"), ) native_duplex = bool(info_dict.get("native_duplex", False)) + meta = info_dict.get("meta") full_embeds = self._build_condition_embeddings( token_ids, hidden_states, native_duplex=native_duplex, ) + if native_duplex: + full_embeds = self._build_streaming_recompute_embeddings( + full_embeds, + request_id=request_id, + info_dict=info_dict, + meta=meta if isinstance(meta, Mapping) else {}, + ) + retained_codes: list[int] = [] + condition_seq = meta.get("streaming_condition_seq") if isinstance(meta, Mapping) else None + if native_duplex and isinstance(condition_seq, int) and not isinstance(condition_seq, bool): + condition_state = self._request_condition_states.get(request_id) + base_recent_codes = ( + condition_state.get("base_recent_codes") if isinstance(condition_state, dict) else None + ) + if not isinstance(base_recent_codes, tuple): + raise ValueError("streaming Talker condition lost its frozen codec history") + retained_codes = list(base_recent_codes) offset = int(info_dict.get("_omni_num_computed_tokens", 0)) - request_id = str(info_dict.get("request_id", "0")) - meta = info_dict.get("meta") # The handoff rebuilds only the tail-aligned Talker condition. # Materialize zero-token embeddings for any scheduler prompt # prefix so chunked prefill can slice from a non-zero offset. prompt_len = info_dict.get("_omni_prompt_len") target_len = int(prompt_len) if prompt_len is not None else offset + span_len + if native_duplex and isinstance(meta, Mapping) and meta.get("streaming_prompt_recompute") is True: + if target_len != full_embeds.shape[0]: + raise ValueError( + "streaming prompt recompute length mismatch: " + f"scheduler={target_len}, model={full_embeds.shape[0]}" + ) prefix_len = target_len - full_embeds.shape[0] if prefix_len > 0: placeholder_ids = torch.zeros( @@ -356,21 +470,23 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): f"tts_ids={token_ids.shape[0]} tts_hidden={hidden_states.shape[0]} " f"prompt_len={info_dict.get('_omni_prompt_len')}" ) - duplex_boundary = isinstance(meta, dict) and ( - bool(meta.get("turn_start", False)) or bool(meta.get("turn_end", False)) - ) if native_duplex: - max_tokens = _DUPLEX_CODEC_TOKENS_PER_CHUNK - min_tokens = 0 if duplex_boundary else _DUPLEX_CODEC_TOKENS_PER_CHUNK + max_tokens, min_tokens = _native_duplex_chunk_budget(meta if isinstance(meta, Mapping) else None) else: - max_tokens = _max_audio_tokens(int(token_ids.numel())) - min_tokens = self._codec_min_tokens - state = { + # MiniCPMTTS.generate()'s max_new_token, clamped to what the + # Talker context can still hold. Sampler min_tokens (upstream's + # min_new_token=50) comes from the deploy YAML. + remaining = int(self._tts_config.max_position_embeddings) - target_len + max_tokens = max(min(_OFFLINE_CODEC_MAX_NEW_TOKENS, remaining), 1) + min_tokens = None + state: dict[str, Any] = { + "finished": empty_condition, "step": 0, "max_tokens": max_tokens, "min_tokens": min_tokens, - "finished": empty_condition, } + if retained_codes: + state["recent_codes"] = retained_codes request_states = getattr(self, "_request_audio_states", None) if request_states is None: request_states = {} @@ -382,98 +498,38 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): embeds, { "audio_state": state, - "audio_codes": { - "current": empty_codes, - "accumulated": empty_codes, - }, + # Prefill has no previous codec id. vLLM samples the first + # code after this forward; the next decode emits it. + "codes": {"audio": empty_codes}, }, ) - current = (info_dict.get("audio_codes", {}) or {}).get("current") - if not isinstance(current, torch.Tensor) or current.numel() != 1: - if state.get("finished"): - # A request that finished before sampling any code can still be - # scheduled for decode steps while sampling min_tokens masks the - # stop token. make_omni_output ignores its hidden states, so any - # shape-correct embedding will do. - weight = self.emb_code[0].weight - return input_ids, weight.new_zeros((span_len, weight.shape[1])), {} - raise RuntimeError("MiniCPM-o Talker decode is missing the previous request-local audio code") - code = current.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(1) + stored = self._request_audio_states.get(request_id) + if isinstance(stored, dict): + state = stored + if isinstance(state, dict) and state.get("finished"): + # An empty speech segment can still be scheduled until EOS is + # eligible. The sampler is forced to EOS; any shape-correct + # embedding is enough for these leftover decode rows. + weight = self.emb_code[0].weight + empty_codes = torch.empty(0, dtype=torch.long, device=weight.device) + return input_ids, weight.new_zeros((span_len, weight.shape[1])), {"codes": {"audio": empty_codes}} + + # Decode: vLLM's previous sampled codec id is this step's input. + # Embed it with the codec table (not the 32k Llama embed_tokens) and + # hand the same id to make_omni_output so Code2Wav sees it this step. + code = input_ids.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1)[-1:] embeds = self.emb_code[0](code) - return input_ids, embeds, {} - - def _request_generator(self, request_id: str, device: torch.device) -> torch.Generator: - generator = self._request_generators.get(request_id) - if generator is None: - generator = torch.Generator(device=device) - generator.manual_seed(self._codec_seed) - self._request_generators[request_id] = generator - return generator - - def _sample_audio_codes( - self, - hidden_states: torch.Tensor, - histories: Sequence[torch.Tensor], - request_ids: Sequence[str], - steps: Sequence[int], - ) -> torch.Tensor: - batch_size = int(hidden_states.shape[0]) - if not (len(histories) == len(request_ids) == len(steps) == batch_size): - raise ValueError( - "MiniCPM-o batched codec sampling requires one history, request id, " - f"and step per hidden row, got batch={batch_size}, histories={len(histories)}, " - f"request_ids={len(request_ids)}, steps={len(steps)}" - ) - if batch_size == 0: - return torch.empty(0, dtype=torch.long, device=hidden_states.device) - - logits = self.head_code[0](hidden_states).float() / self._codec_temperature - eos_id = self._num_audio_tokens - 1 - logits = _apply_batched_repetition_penalty( - logits, - histories, - penalty=self._codec_repetition_penalty, - window_size=_REPETITION_WINDOW, - ) - request_states = getattr(self, "_request_audio_states", {}) - mask_eos_values: list[bool] = [] - for request_id, step in zip(request_ids, steps, strict=True): - state = request_states.get(request_id) - min_tokens = ( - int(state.get("min_tokens", self._codec_min_tokens)) - if isinstance(state, dict) - else self._codec_min_tokens - ) - mask_eos_values.append(step < min_tokens) - mask_eos = torch.tensor( - mask_eos_values, - dtype=torch.bool, - device=logits.device, - ) - logits[:, eos_id].masked_fill_(mask_eos, float("-inf")) - logits = _apply_top_k_top_p( - logits, - top_k=self._codec_top_k, - top_p=self._codec_top_p, - min_tokens_to_keep=3, - ) - probabilities = torch.softmax(logits, dim=-1) - # torch.multinomial accepts one generator for the entire call. Keep it - # per row so compaction, request ordering, or another request finishing - # cannot perturb request-local RNG streams; the expensive projection - # and logit transforms above remain fully batched. - return torch.cat( - [ - torch.multinomial( - probabilities[row : row + 1], - num_samples=1, - generator=self._request_generator(request_id, probabilities.device), - ) - for row, request_id in enumerate(request_ids) - ], - dim=0, - ).reshape(-1) + code_id = int(code.item()) + if code_id == int(self._codec_eos_id): + if isinstance(state, dict): + state["finished"] = True + elif stored is None: + self._request_audio_states[request_id] = {"finished": True, "step": 0} + delta = torch.empty(0, dtype=torch.long, device=code.device) + else: + delta = code.reshape(1, 1) + return input_ids, embeds, {"codes": {"audio": delta}} def make_omni_output( self, @@ -487,18 +543,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): spans = kwargs.get("request_token_spans") if spans is None or len(spans) != len(infos): raise RuntimeError("MiniCPM-o continuous Talker requires one request_token_span per request") - sample_eligible = kwargs.get("request_sample_eligible") - if sample_eligible is None: - sample_eligible = [True] * len(infos) - if len(sample_eligible) != len(infos): - raise RuntimeError( - f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests" - ) emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos) - # Rows default to continue. Only previously finished or newly terminal - # requests stop; prefill/ineligible rows stay aligned as False. - stop_flags = [False] * len(infos) native_duplex_flags: list[torch.Tensor] = [] duplex_epochs: list[torch.Tensor] = [] duplex_turn_ids: list[torch.Tensor] = [] @@ -507,7 +553,10 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): empty_delta = hidden.new_empty((0, 1), dtype=torch.long) codec_deltas = [empty_delta for _ in infos] terminal_flags = [torch.tensor(False, dtype=torch.bool) for _ in infos] - pending_samples: list[_PendingCodecSample] = [] + force_eos_rows = [False] * len(infos) + mask_eos_rows = [False] * len(infos) + empty_history = hidden.new_empty((0,), dtype=torch.long) + penalty_histories = [empty_history for _ in infos] for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} native_duplex = info_dict.get("native_duplex") is True @@ -554,100 +603,50 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): if not isinstance(info, dict): continue - start, end = spans[index] - end = min(int(end), int(hidden.shape[0])) - if int(start) >= end: - continue request_id = str(info.get("request_id", index)) - request_states = getattr(self, "_request_audio_states", None) - if request_states is None: - request_states = {} - self._request_audio_states = request_states - state = request_states.get(request_id) + state = self._request_audio_states.get(request_id) if not isinstance(state, dict): state = dict(info.get("audio_state", {}) or {}) - request_states[request_id] = state - if state.get("finished"): - stop_flags[index] = True - continue - if not sample_eligible[index]: - # vLLM computes a logit row for incomplete chunked prefills but - # discards its sampled token. Advancing codec/RNG state here - # would make output depend on prefill chunking and compaction. - continue - codes = state.get("codes") - if not isinstance(codes, torch.Tensor): - codes = (info.get("audio_codes", {}) or {}).get("accumulated") - if not isinstance(codes, torch.Tensor): - codes = torch.empty(0, dtype=torch.long, device=hidden.device) - else: - codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1) + self._request_audio_states[request_id] = state + empty_speech = bool(state.get("finished")) + codes = info.get("codes", {}) + audio = codes.get("audio") if isinstance(codes, Mapping) else None + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + codec_deltas[index] = audio.to(device=hidden.device, dtype=torch.long).reshape(-1, 1) + state["step"] = int(state.get("step", 0)) + 1 + # ``audio`` is the id sampled last step, i.e. exactly upstream's + # ``new_tokens[:, 0:t]`` history for the logits computed below. + recent = state.get("recent_codes") + recent = (recent if isinstance(recent, list) else []) + codec_deltas[index].reshape(-1).tolist() + state["recent_codes"] = recent[-_CODEC_PENALTY_WINDOW:] + recent_codes = state.get("recent_codes") + if recent_codes: + penalty_histories[index] = torch.tensor(recent_codes, dtype=torch.long, device=hidden.device) + max_tokens = state.get("max_tokens") + min_tokens = state.get("min_tokens") step = int(state.get("step", 0)) - pending_samples.append( - _PendingCodecSample( - output_index=index, - hidden_row=hidden[end - 1 : end], - codes=codes, - request_id=request_id, - step=step, - state=state, - info=info, - ) - ) + # Duplex: 26 samples include the terminating EOS, so force it after + # 25 forwarded frames — the same cadence as generate_chunk. + # Offline: force-stop at the remaining Talker context rather than + # waiting for a sampled EOS. + hit_chunk_limit = max_tokens is not None and step >= int(max_tokens) - 1 + chunk_done = empty_speech or hit_chunk_limit + if hit_chunk_limit: + # The next sampled id is forced EOS and will finish the + # request; mark the chunk done on this step so the data + # plane does not wait for a follow-up empty decode. + state["finished"] = True + force_eos_rows[index] = chunk_done + mask_eos_rows[index] = not force_eos_rows[index] and min_tokens is not None and step < int(min_tokens) + terminal_flags[index] = torch.tensor(chunk_done, dtype=torch.bool) - if pending_samples: - active_hidden_rows = [pending.hidden_row for pending in pending_samples] - active_hidden = ( - active_hidden_rows[0] if len(active_hidden_rows) == 1 else torch.cat(active_hidden_rows, dim=0) - ) - sampled_batch = self._sample_audio_codes( - active_hidden, - [pending.codes for pending in pending_samples], - [pending.request_id for pending in pending_samples], - [pending.step for pending in pending_samples], - ) - # One batched device-to-host synchronization replaces one .item() - # synchronization per request. - sampled_ids = sampled_batch.detach().to(device="cpu").tolist() - else: - sampled_batch = hidden.new_empty((0,), dtype=torch.long) - sampled_ids = [] - - for row, pending in enumerate(pending_samples): - sampled = sampled_batch[row].reshape(()) - sampled_id = int(sampled_ids[row]) - codes = pending.codes - state = pending.state - info = pending.info - is_eos = sampled_id == self._num_audio_tokens - 1 - state["step"] = int(state.get("step", 0)) + 1 - reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048)) - finished = is_eos or reached_limit - state["finished"] = finished - # MiniCPMTTS.generate_chunk consumes the boundary sample but - # returns only codes that were fed into the retained KV state. - if not is_eos and not reached_limit: - codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)]) - delta = sampled.reshape(1, 1) - else: - delta = empty_delta - state["codes"] = codes - info["audio_state"] = state - info["audio_codes"] = { - "current": sampled.reshape(1), - "accumulated": codes, - } - codec_deltas[pending.output_index] = delta - terminal_flags[pending.output_index] = torch.tensor(finished, dtype=torch.bool) - stop_flags[pending.output_index] = finished - - self._batch_stop_logits = ( - hidden.new_tensor([[float("-inf"), 0.0] if finished else [0.0, float("-inf")] for finished in stop_flags]) - if stop_flags - else hidden.new_empty((0, 2)) - ) - # Lists are deliberate: the runner routes element i to request i, - # preserving compaction alignment while emitting only this step's code. + # Empty-speech rows, finished duplex chunks, and offline requests that + # fill the remaining Talker context must sample codec EOS so the + # scheduler releases the request. Mid-chunk rows mask EOS until + # min_tokens. + self._force_eos_rows = force_eos_rows + self._mask_eos_rows = mask_eos_rows + self._penalty_histories = penalty_histories meta_outputs = {"finished": terminal_flags} if emit_duplex_metadata: meta_outputs.update( @@ -659,13 +658,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): "turn_end": turn_end_flags, } ) - multimodal_outputs: dict[str, Any] = { - "codes": {"audio": codec_deltas}, - "meta": meta_outputs, - } return OmniOutput( text_hidden_states=hidden, - multimodal_outputs=multimodal_outputs, + multimodal_outputs={ + "codes": {"audio": codec_deltas}, + "meta": meta_outputs, + }, ) def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: @@ -673,9 +671,10 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): def _flush_deferred_cleanup(self) -> None: request_audio_states = getattr(self, "_request_audio_states", {}) + request_condition_states = getattr(self, "_request_condition_states", {}) for request_id in self._deferred_cleanup_ids: - self._request_generators.pop(request_id, None) request_audio_states.pop(request_id, None) + request_condition_states.pop(request_id, None) self._deferred_cleanup_ids.clear() def _dummy_hidden_states( @@ -723,19 +722,101 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): def compute_logits(self, hidden_states, *args, **kwargs): if not isinstance(hidden_states, torch.Tensor): return None - if self._batch_stop_logits is None: - return torch.zeros( - hidden_states.shape[0], - 2, - device=hidden_states.device, - dtype=torch.float32, - ) - logits = self._batch_stop_logits - self._batch_stop_logits = None + if hidden_states.numel() == 0: + return hidden_states.new_empty((0, int(self._num_audio_tokens))) + logits = self.head_code[0](hidden_states).float() + force_eos = self._force_eos_rows + mask_eos = self._mask_eos_rows + self._force_eos_rows = None + self._mask_eos_rows = None + need_force = bool(force_eos and len(force_eos) == logits.shape[0] and any(force_eos)) + need_mask = bool(mask_eos and len(mask_eos) == logits.shape[0] and any(mask_eos)) + # sample() re-applies the decision on the sampled ids: vLLM's + # MinTokensLogitsProcessor runs after this and would blank the codec EOS + # we just forced (it is in the stage's ``stop_token_ids``), leaving an + # all -inf row and a request that never releases. + self._pending_force_eos_rows = force_eos if need_force else None + if need_force or need_mask: + logits = logits.clone() + eos_id = int(self._codec_eos_id) + if need_force: + assert force_eos is not None + forced = torch.tensor(force_eos, dtype=torch.bool, device=logits.device) + logits[forced] = float("-inf") + logits[forced, eos_id] = 0.0 + if need_mask: + assert mask_eos is not None + masked = torch.tensor(mask_eos, dtype=torch.bool, device=logits.device) + logits[masked, eos_id] = float("-inf") return logits def sample(self, logits, sampling_metadata): - return Sampler()(logits, sampling_metadata) + prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None) + if isinstance(logits, torch.Tensor) and isinstance(prompt_ids, torch.Tensor): + # Copy rather than mutate: the runner may hand us the input batch's + # own persistent SamplingMetadata. + sampling_metadata = replace( + sampling_metadata, + prompt_token_ids=blank_scheduler_prompt_for_penalties(prompt_ids, logits.shape[-1]), + ) + logits, sampling_metadata = self._apply_codec_repetition_penalty(logits, sampling_metadata) + force_eos = self._pending_force_eos_rows + self._pending_force_eos_rows = None + output = Sampler()(logits, sampling_metadata) + return self._force_eos_on_sampled_ids(output, force_eos) + + def _apply_codec_repetition_penalty(self, logits, sampling_metadata): + """Score MiniCPMTTS.generate's windowed codec penalty, not vLLM's. + + Upstream taxes a code by ``penalty ** frequency`` over the last + ``past_window`` frames only (``gen_logits`` builds + ``CustomRepetitionPenaltyLogitsProcessorRepeat(penalty, num_code, 16)``). + vLLM's is presence-based over the whole stream, so on a codec stream + thousands of frames long every code ever sampled ends up taxed by the + same flat factor while never-sampled codes stay untouched, and the tail + of a long answer drifts off the speech manifold into near-silence. + + Runs before ``Sampler`` so the penalty lands ahead of top-k/top-p, as + upstream does. Upstream scores it after dividing by temperature, but + the penalty only rescales and preserves sign, so the two orders agree. + """ + histories = self._penalty_histories + self._penalty_histories = None + penalties = getattr(sampling_metadata, "repetition_penalties", None) + if ( + not isinstance(logits, torch.Tensor) + or histories is None + or len(histories) != logits.shape[0] + or not isinstance(penalties, torch.Tensor) + or getattr(sampling_metadata, "no_penalties", False) + ): + return logits, sampling_metadata + logits = _apply_batched_repetition_penalty( + logits, + histories, + penalty=penalties.to(device=logits.device, dtype=logits.dtype), + window_size=_CODEC_PENALTY_WINDOW, + ) + # Neutralize the sampler's own pass so the penalty is scored once. + return logits, replace(sampling_metadata, repetition_penalties=torch.ones_like(penalties)) + + def _force_eos_on_sampled_ids(self, output: Any, force_eos: list[bool] | None) -> Any: + """Overwrite sampled ids for rows the model terminated this step. + + The codec EOS is a stage ``stop_token_ids`` entry, so vLLM's + ``min_tokens`` processor masks it for the first ``min_tokens`` steps. + A row the model forced to EOS therefore reaches the sampler as all + -inf and comes back as an arbitrary codec id, which keeps an + already-finished request decoding until its length cap. + """ + if not force_eos or not any(force_eos): + return output + sampled = getattr(output, "sampled_token_ids", None) + if not isinstance(sampled, torch.Tensor) or sampled.shape[0] != len(force_eos): + return output + rows = torch.tensor(force_eos, dtype=torch.bool, device=sampled.device) + sampled[rows] = int(self._codec_eos_id) + return output def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): return self._load_native_weights(weights) @@ -782,9 +863,10 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): return loaded def get_input_embeddings(self, input_ids, multimodal_embeddings=None, **kwargs): - if hasattr(self, "emb_text") and self.emb_text is not None: - return self.emb_text(input_ids) - return torch.zeros(input_ids.shape[0], 1) + del multimodal_embeddings + # Decode tokens live in the codec table. Prefill overwrites these + # embeddings in preprocess with emb_text + projected thinker hidden. + return self.emb_code[0](input_ids) def embed_input_ids(self, input_ids, **kwargs): return self.get_input_embeddings(input_ids, **kwargs) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index b8fed05f1..51898f27a 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniCPM-o 4.5 pipeline topology (frozen). Stage 0: Thinker — multimodal understanding + text generation. @@ -17,6 +17,7 @@ from vllm_omni.config.stage_config import ( ) _PROC = "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni" +MINICPMO45_REFERENCE_AUDIO_KEY = "_minicpmo45_reference_audio" MINICPMO_4_5_PIPELINE = PipelineConfig( @@ -62,7 +63,12 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( custom_process_input_func=f"{_PROC}.llm2tts", custom_process_next_stage_input_func=f"{_PROC}.tts2code2wav_full_payload", async_chunk_process_next_stage_input_func=f"{_PROC}.tts2code2wav_async_chunk", - sampling_constraints={"detokenize": False}, + sampling_constraints={ + "detokenize": False, + # MiniCPM-o 4.5 codec EOS is tts_config.num_audio_tokens - 1. + # Same pattern as Qwen3 talker's stop_token_ids: [2150]. + "stop_token_ids": [6561], + }, ), StagePipelineConfig( stage_id=2, @@ -75,6 +81,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( model_arch="MiniCPMO45Code2Wav", sync_process_input_func=f"{_PROC}.tts2code2wav_token_only", sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/minimax_h3/__init__.py b/vllm_omni/model_executor/models/minimax_h3/__init__.py new file mode 100644 index 000000000..bd977b2ec --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/__init__.py @@ -0,0 +1 @@ +"""MiniMax H3 model components.""" diff --git a/vllm_omni/model_executor/models/minimax_h3/checkpoint.py b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py new file mode 100644 index 000000000..e8350f41a --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/checkpoint.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 checkpoint resolution for the AR text-encoder stage.""" + +from pathlib import Path + +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) + +MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS = { + "fl2va": ["FL2VA/text_encoder/**"], + "ref2va": ["Ref2VA/text_encoder/**"], +} + + +def resolve_minimax_h3_partition( + model: str, + task_type: str | None, + *, + auto_partition: str, +) -> str: + task = str(task_type or "auto").lower() + if task not in {"auto", "combined", "t2va", "fl2va", "ref2va"}: + raise ValueError( + f"MiniMax-H3 task_type must be one of auto, combined, t2va, fl2va, or ref2va; got {task_type!r}" + ) + path = Path(model) + if task == "auto" and path.is_dir() and path.name in {"FL2VA", "Ref2VA"}: + return path.name.lower() + if task in {"auto", "combined"}: + return auto_partition + return "ref2va" if task == "ref2va" else "fl2va" + + +def resolve_minimax_h3_model_root( + model: str, + revision: str | None, + task_type: str | None, +) -> str: + path = Path(model) + partition = resolve_minimax_h3_partition(model, task_type, auto_partition="fl2va") + + if path.is_dir(): + if path.name == "text_encoder" and (path / "config.json").is_file(): + return str(path) + if path.name in {"FL2VA", "Ref2VA"}: + path = path.parent + subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" + return str(path / subdir / "text_encoder") + snapshot = download_weights_from_hf_specific( + model_name_or_path=model, + cache_dir=None, + allow_patterns=MINIMAX_H3_TEXT_ENCODER_DOWNLOAD_PATTERNS[partition], + revision=revision, + require_all=True, + ) + + subdir = "Ref2VA" if partition == "ref2va" else "FL2VA" + return str(Path(snapshot) / subdir / "text_encoder") + + +__all__ = ["resolve_minimax_h3_model_root", "resolve_minimax_h3_partition"] diff --git a/vllm_omni/model_executor/models/minimax_h3/conditioning.py b/vllm_omni/model_executor/models/minimax_h3/conditioning.py new file mode 100644 index 000000000..8a22c7737 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/conditioning.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 text-conditioning contract.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +MINIMAX_H3_TEXT_HIDDEN_SIZE = 5120 +MINIMAX_H3_PRESENTATION_TASK_KEY = "_minimax_h3_presentation_task" +MINIMAX_H3_CONDITION_LABELS_KEY = "_minimax_h3_condition_labels" + + +@dataclass(frozen=True) +class MiniMaxH3TextConditioning: + """Layer-50 Qwen3-VL hidden states and their H3 token roles.""" + + hidden_states: torch.Tensor + token_tags: torch.Tensor + + @classmethod + def from_payload( + cls, + payload: Mapping[str, Any], + ) -> MiniMaxH3TextConditioning: + hidden_states = payload.get("hidden_states") + token_tags = payload.get("token_tags") + if not isinstance(hidden_states, torch.Tensor) or not isinstance(token_tags, torch.Tensor): + raise ValueError("MiniMax H3 conditioning requires hidden_states and token_tags tensors") + if hidden_states.ndim != 2 or hidden_states.shape[-1] != MINIMAX_H3_TEXT_HIDDEN_SIZE: + raise ValueError( + "MiniMax H3 hidden_states must have shape " + f"[tokens, {MINIMAX_H3_TEXT_HIDDEN_SIZE}], got {tuple(hidden_states.shape)}" + ) + if token_tags.ndim != 1 or token_tags.shape[0] != hidden_states.shape[0]: + raise ValueError( + "MiniMax H3 token_tags must align with hidden_states, got " + f"token_tags={tuple(token_tags.shape)} and hidden_states={tuple(hidden_states.shape)}" + ) + if not torch.all((token_tags == 0) | (token_tags == 1)): + raise ValueError("MiniMax H3 text-encoder token_tags must contain only 0 and 1") + return cls(hidden_states=hidden_states, token_tags=token_tags) + + def to_payload(self) -> dict[str, torch.Tensor]: + return { + "hidden_states": self.hidden_states, + "token_tags": self.token_tags, + } diff --git a/vllm_omni/model_executor/models/minimax_h3/pipeline.py b/vllm_omni/model_executor/models/minimax_h3/pipeline.py new file mode 100644 index 000000000..5ac94b5ee --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/pipeline.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Two-stage MiniMax H3 topology.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROCESSOR = "vllm_omni.model_executor.stage_input_processors.minimax_h3" +_CHECKPOINT = "vllm_omni.model_executor.models.minimax_h3.checkpoint" +_DIFFUSION_PIPELINE = "vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3" + +MINIMAX_H3_PIPELINE = PipelineConfig( + model_type="minimax_h3_disaggregated", + default_deploy_config_name="minimax_h3_disaggregated.yaml", + stage_cli_aliases={"text_encoder_tp_size": (0, "tensor_parallel_size")}, + model_arch="MiniMaxH3TextEncoder", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="text_encoder", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + requires_multimodal_data=True, + model_arch="MiniMaxH3TextEncoder", + engine_output_type="latent", + prompt_transform_func=f"{_PROCESSOR}.prepare_text_encoder_prompt", + sampling_constraints={ + "max_tokens": 1, + "temperature": 0.0, + "detokenize": False, + }, + model_path_resolver=f"{_CHECKPOINT}.resolve_minimax_h3_model_root", + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + final_output=True, + final_output_type="video", + requires_multimodal_data=True, + model_arch="MiniMaxH3Pipeline", + custom_process_input_func=f"{_PROCESSOR}.text_encoder2diffusion", + omni_kv_config={"need_recv_cache": False}, + model_path_resolver=f"{_DIFFUSION_PIPELINE}.resolve_minimax_h3_diffusion_model_path", + inline_diffusion=True, + ), + ), +) diff --git a/vllm_omni/diffusion/models/minimax_h3/presentation.py b/vllm_omni/model_executor/models/minimax_h3/preprocessing.py similarity index 53% rename from vllm_omni/diffusion/models/minimax_h3/presentation.py rename to vllm_omni/model_executor/models/minimax_h3/preprocessing.py index a8e142189..36a1e20d6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/presentation.py +++ b/vllm_omni/model_executor/models/minimax_h3/preprocessing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""MiniMax H3 Qwen presentation building. +"""Shared MiniMax H3 media normalization and Qwen presentation building. Builds the positive presentation token stream: - fl2va: ': ' label + vision block (<|vision_start|> + @@ -13,10 +13,16 @@ accumulator so ids and AdaLN token tags cannot drift apart. from __future__ import annotations +import math +import os from collections.abc import Sequence from typing import Any +import numpy as np import torch +from PIL import Image + +from vllm_omni.errors import OmniClientError VISION_START = "<|vision_start|>" VISION_END = "<|vision_end|>" @@ -26,6 +32,142 @@ VIDEO_PAD = "<|video_pad|>" _TEXT_TAG = 1 _VIDEO_TAG = 0 +MINIMAX_H3_OUTPUT_SHORT_EDGE = 768 +MINIMAX_H3_OUTPUT_MAX_PIXELS = 768 * 1344 +MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048 +MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32 +MINIMAX_H3_SUPPORTED_ASPECT_RATIOS = { + "21:9": 21.0 / 9.0, + "16:9": 16.0 / 9.0, + "4:3": 4.0 / 3.0, + "1:1": 1.0, + "3:4": 3.0 / 4.0, + "9:16": 9.0 / 16.0, +} +MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024 +MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) + + +def _align_multiple(value: float, multiple: int = 32) -> int: + return max(multiple, int(round(float(value) / multiple)) * multiple) + + +def load_minimax_h3_images(value: Any) -> list[Image.Image]: + """Normalize one or more H3 image inputs to RGB PIL images.""" + if isinstance(value, (list, tuple)): + if not value: + raise OmniClientError("MiniMax H3 image input must not be empty") + images: list[Image.Image] = [] + for item in value: + loaded = load_minimax_h3_images(item) + if len(loaded) != 1: + raise OmniClientError(f"MiniMax H3 expected one image, got {len(loaded)}") + images.extend(loaded) + return images + if isinstance(value, (str, os.PathLike)): + if os.path.getsize(value) > MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES: + raise OmniClientError("MiniMax H3 reference image exceeds the 30 MiB size limit") + with Image.open(value) as image: + image_format = str(image.format or "").lower() + if image_format and image_format not in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: + raise OmniClientError( + f"MiniMax H3 reference image must use JPG, JPEG, PNG, WEBP, HEIC, or HEIF, got {image.format}" + ) + return [image.convert("RGB")] + if isinstance(value, Image.Image): + return [value.convert("RGB")] + if isinstance(value, torch.Tensor): + tensor = value.detach().float().cpu() + if tensor.ndim == 4 and tensor.shape[0] == 1: + tensor = tensor[0] + if tensor.ndim != 3: + raise OmniClientError(f"image tensor must be [C,H,W], got {tuple(tensor.shape)}") + if tensor.shape[0] in (1, 3, 4): + tensor = tensor.permute(1, 2, 0) + array = tensor.numpy() + if array.max(initial=0) <= 1.0: + array = array * 255.0 + return [Image.fromarray(array.clip(0, 255).astype(np.uint8)).convert("RGB")] + raise OmniClientError(f"unsupported MiniMax H3 image input {type(value)!r}") + + +def resolve_minimax_h3_aspect_ratio( + task: str, + value: Any, + image: Image.Image | None, +) -> float: + """Resolve H3's task-specific output aspect-ratio policy.""" + if task == "fl2va": + if image is None: + raise OmniClientError("fl2va requires an input image to resolve its aspect ratio") + return float(image.width) / float(image.height) + + if value is None: + if task == "t2va": + raise OmniClientError("t2va requires an explicit aspect_ratio") + return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"] + + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"adaptive", "auto"}: + if task == "t2va": + raise OmniClientError("t2va requires an explicit named aspect_ratio, not adaptive") + return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS["16:9"] + if normalized in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS: + return MINIMAX_H3_SUPPORTED_ASPECT_RATIOS[normalized] + try: + numeric_value = float(normalized) + except (TypeError, ValueError) as exc: + supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS) + raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}") from exc + elif isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool): + numeric_value = float(value) + else: + raise OmniClientError(f"MiniMax H3 aspect_ratio must be a string ratio, got {value!r}") + + if not math.isfinite(numeric_value) or not any( + math.isclose(numeric_value, ratio, rel_tol=0.0, abs_tol=1e-6) + for ratio in MINIMAX_H3_SUPPORTED_ASPECT_RATIOS.values() + ): + supported = ", ".join(MINIMAX_H3_SUPPORTED_ASPECT_RATIOS) + raise OmniClientError(f"MiniMax H3 aspect_ratio must be one of {supported}, got {value!r}") + return numeric_value + + +def resolve_minimax_h3_reference_image_shape(image: Image.Image) -> tuple[int, int]: + """Resize an H3 reference image to the official 2048-short-edge canvas.""" + width, height = image.size + ratio = width / height + if not 0.4 <= ratio <= 2.5: + raise OmniClientError(f"reference image aspect ratio must be in [0.4, 2.5], got {width}x{height}") + if min(width, height) < 256 or max(width, height) > 5760: + raise OmniClientError(f"reference image dimensions must be in [256, 5760] pixels, got {width}x{height}") + scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height) + return ( + _align_multiple(width * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE), + _align_multiple(height * scale, MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE), + ) + + +def resolve_minimax_h3_output_canvas(aspect_ratio: float, short_edge: int) -> tuple[int, int]: + """Resolve the official H3 ratio/area policy to a 32-pixel canvas.""" + if not math.isfinite(float(aspect_ratio)) or float(aspect_ratio) <= 0: + raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be positive, got {aspect_ratio!r}") + if short_edge != MINIMAX_H3_OUTPUT_SHORT_EDGE: + raise OmniClientError(f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {short_edge}") + if aspect_ratio >= 1.0: + width = float(short_edge) * aspect_ratio + height = float(short_edge) + else: + width = float(short_edge) + height = float(short_edge) / aspect_ratio + area = width * height + if area > MINIMAX_H3_OUTPUT_MAX_PIXELS: + scale = (MINIMAX_H3_OUTPUT_MAX_PIXELS / area) ** 0.5 + width *= scale + height *= scale + return _align_multiple(height, 32), _align_multiple(width, 32) + def _text_ids(tokenizer: Any, text: str) -> list[int]: return list(tokenizer(text, add_special_tokens=False)["input_ids"]) @@ -107,26 +249,18 @@ def _multi_image_presentation( return presentation.build() -def minimax_h3_multi_image_presentation_ids( +def minimax_h3_multi_image_presentation( tokenizer: Any, *, prompt: str, image_token_counts: list[int], -) -> torch.Tensor: - """Positive fl2va presentation with one or more labeled image slots.""" - ids, _ = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts) - return ids - - -def minimax_h3_multi_image_presentation_token_tags( - tokenizer: Any, - *, - prompt: str, - image_token_counts: list[int], -) -> torch.Tensor: - """AdaLN token tags aligned with minimax_h3_multi_image_presentation_ids.""" - _, tags = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts) - return tags +) -> tuple[torch.Tensor, torch.Tensor]: + """Return aligned FL2VA presentation token IDs and role IDs.""" + return _multi_image_presentation( + tokenizer, + prompt=prompt, + image_token_counts=image_token_counts, + ) def minimax_h3_ref2va_presentation( @@ -288,10 +422,72 @@ def minimax_h3_ref2va_video_presentation( return presentation.build() +def build_minimax_h3_presentation( + tokenizer: Any, + *, + prompt: str, + task: str, + condition_labels: list[tuple[str, int]], + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + video_timestamps: Sequence[Sequence[float]] | None, + merge_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build the token IDs and role IDs shared by fused and split H3.""" + if task == "t2va": + ids = minimax_h3_text_only_ids(tokenizer, prompt) + return ids, torch.ones_like(ids) + + merge_length = int(merge_size) ** 2 + image_counts = ( + [int(grid.prod().item()) // merge_length for grid in image_grid_thw] if image_grid_thw is not None else [] + ) + if task == "fl2va": + return minimax_h3_multi_image_presentation( + tokenizer, + prompt=prompt, + image_token_counts=image_counts, + ) + + video_counts: list[list[int]] = [] + if video_grid_thw is not None: + for grid in video_grid_thw: + block_count = int(grid[0].item()) + tokens_per_block = int(grid[1:].prod().item()) // merge_length + video_counts.append([tokens_per_block] * block_count) + timestamps = ( + [[float(value) for value in group] for group in video_timestamps] if video_timestamps is not None else [] + ) + if video_counts: + return minimax_h3_ref2va_video_presentation( + tokenizer, + prompt=prompt, + condition_labels=condition_labels, + image_token_count=image_counts or None, + video_block_token_counts=video_counts, + video_block_timestamps=timestamps, + ) + return minimax_h3_ref2va_presentation( + tokenizer, + prompt=prompt, + condition_labels=condition_labels, + image_token_count=image_counts or None, + ) + + __all__ = [ - "minimax_h3_multi_image_presentation_ids", - "minimax_h3_multi_image_presentation_token_tags", + "IMAGE_PAD", + "MINIMAX_H3_OUTPUT_SHORT_EDGE", + "VIDEO_PAD", + "VISION_END", + "VISION_START", + "build_minimax_h3_presentation", + "load_minimax_h3_images", + "minimax_h3_multi_image_presentation", "minimax_h3_ref2va_presentation", "minimax_h3_ref2va_video_presentation", "minimax_h3_text_only_ids", + "resolve_minimax_h3_aspect_ratio", + "resolve_minimax_h3_output_canvas", + "resolve_minimax_h3_reference_image_shape", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/reference_video.py b/vllm_omni/model_executor/models/minimax_h3/reference_video.py similarity index 79% rename from vllm_omni/diffusion/models/minimax_h3/reference_video.py rename to vllm_omni/model_executor/models/minimax_h3/reference_video.py index 0ecd69f31..4f98d9e13 100644 --- a/vllm_omni/diffusion/models/minimax_h3/reference_video.py +++ b/vllm_omni/model_executor/models/minimax_h3/reference_video.py @@ -15,6 +15,53 @@ import torch from vllm_omni.errors import OmniClientError +MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY = "minimax_h3_prepared_reference_videos" + + +def serialize_prepared_reference_videos(items: list[dict[str, Any]], artifact_dir: str) -> str: + fields = ( + "original_path", + "prepared_path", + "input_has_audio", + "width", + "height", + "start_time_seconds", + "duration_seconds", + "audio_duration_seconds", + ) + return json.dumps( + { + "artifact_dir": artifact_dir, + "videos": [{field: item[field] for field in fields} for item in items], + }, + separators=(",", ":"), + ) + + +def deserialize_prepared_reference_videos(value: str) -> tuple[str, list[dict[str, Any]]]: + try: + payload = json.loads(value) + artifact_dir = payload["artifact_dir"] + videos = payload["videos"] + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + raise ValueError("invalid MiniMax H3 prepared-reference-video descriptor") from exc + if not isinstance(artifact_dir, str) or not isinstance(videos, list): + raise ValueError("invalid MiniMax H3 prepared-reference-video descriptor") + required = { + "original_path", + "prepared_path", + "input_has_audio", + "width", + "height", + "start_time_seconds", + "duration_seconds", + "audio_duration_seconds", + } + if any(not isinstance(item, dict) or set(item) != required for item in videos): + raise ValueError("invalid MiniMax H3 prepared-reference-video descriptor") + return artifact_dir, videos + + MINIMAX_H3_FPS = 24.0 MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS = 2.0 MINIMAX_H3_QWEN_TEMPORAL_PATCH = 2 @@ -44,31 +91,33 @@ def _nearest_multiple(value: float, multiple: int) -> int: def _probe_video(path: str) -> dict[str, Any]: - result = subprocess.run( - [ - "ffprobe", - "-v", - "error", - "-count_frames", - "-show_entries", - ( - "stream=codec_type,codec_name,width,height,r_frame_rate,nb_read_frames,nb_frames,duration," - "sample_aspect_ratio:format=format_name,size,duration" - ), - "-of", - "json", - path, - ], - check=True, - capture_output=True, - text=True, + show_entries = ( + "stream=codec_type,codec_name,width,height,r_frame_rate,nb_read_frames,nb_frames,duration," + "sample_aspect_ratio:format=format_name,size,duration" ) + command = ["ffprobe", "-v", "error", "-show_entries", show_entries, "-of", "json", path] + result = subprocess.run(command, check=True, capture_output=True, text=True) probe = json.loads(result.stdout) streams = probe.get("streams") or [] video_streams = [stream for stream in streams if stream.get("codec_type") == "video"] if not video_streams: raise OmniClientError(f"media has no video stream: {path}") stream = video_streams[0] + raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") + if raw_count in (None, "", "N/A"): + result = subprocess.run( + ["ffprobe", "-v", "error", "-count_frames", *command[3:]], + check=True, + capture_output=True, + text=True, + ) + probe = json.loads(result.stdout) + streams = probe.get("streams") or [] + video_streams = [stream for stream in streams if stream.get("codec_type") == "video"] + if not video_streams: + raise OmniClientError(f"media has no video stream: {path}") + stream = video_streams[0] + raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") try: numerator, denominator = str(stream["r_frame_rate"]).split("/", 1) fps = float(numerator) / float(denominator) @@ -76,7 +125,6 @@ def _probe_video(path: str) -> dict[str, Any]: raise OmniClientError(f"cannot determine video FPS: {path}") from exc if not math.isfinite(fps) or fps <= 0: raise OmniClientError(f"cannot determine video FPS: {path}") - raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") if raw_count in (None, "", "N/A"): raise OmniClientError(f"cannot determine video frame count: {path}") raw_duration = stream.get("duration") @@ -279,6 +327,7 @@ def _transcode_reference_video( ) -> str: output = str(Path(workdir) / "prepared.mp4") duration_args = ["-t", f"{float(duration_seconds):.6f}"] if duration_seconds is not None else [] + frame_count_args = ["-frames:v", str(int(target_frame_count))] if target_frame_count > 0 else [] subprocess.run( [ "ffmpeg", @@ -295,8 +344,7 @@ def _transcode_reference_video( "-vf", (f"fps={MINIMAX_H3_FPS:g},scale={target_width}:{target_height}:flags=lanczos,setsar=1"), *duration_args, - "-frames:v", - str(int(target_frame_count)), + *frame_count_args, "-metadata:s:v:0", "rotate=0", # Keep the transformed RGB pixels lossless. The same prepared @@ -393,7 +441,7 @@ def prepare_reference_videos( # bounded duration separately for audio extraction. "audio_duration_seconds": min( duration, - float(target_frame_count) / MINIMAX_H3_FPS, + float(target_frame_count) / MINIMAX_H3_FPS if target_frame_count > 0 else duration, ), } ) @@ -404,13 +452,13 @@ def load_video_frames(path: str) -> np.ndarray: try: import decord except ImportError: - import av - - with av.open(path) as container: - frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)] - if not frames: - raise OmniClientError(f"video has no frames: {path}") - return np.stack(frames) + meta = _probe_video(path) + return _decode_video_frames_ffmpeg( + path, + frame_count=int(meta["frame_count"]), + width=int(meta["width"]), + height=int(meta["height"]), + ) reader = decord.VideoReader(path) if len(reader) <= 0: @@ -419,6 +467,70 @@ def load_video_frames(path: str) -> np.ndarray: return frames.asnumpy() if hasattr(frames, "asnumpy") else np.asarray(frames) +def _decode_video_frames_ffmpeg( + path: str, + *, + frame_count: int, + width: int, + height: int, + indices: list[int] | None = None, +) -> np.ndarray: + frame_count = int(frame_count) + width = int(width) + height = int(height) + if frame_count <= 0: + raise OmniClientError(f"video has no frames: {path}") + if width <= 0 or height <= 0: + raise OmniClientError(f"video has invalid dimensions {width}x{height}: {path}") + if indices is not None and (not indices or any(index < 0 or index >= frame_count for index in indices)): + raise OmniClientError(f"invalid frame indices for {path}: {indices}") + + output_frame_count = frame_count if indices is None else len(indices) + command = [ + "ffmpeg", + "-loglevel", + "error", + "-threads", + "0", + "-i", + path, + "-map", + "0:v:0", + "-an", + ] + if indices is not None: + select = "+".join(f"eq(n\\,{index})" for index in indices) + command.extend(["-vf", f"select={select}"]) + command.extend( + [ + "-vsync", + "0", + "-frames:v", + str(output_frame_count), + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "pipe:1", + ] + ) + result = subprocess.run( + command, + check=True, + capture_output=True, + ) + frame_size = width * height * 3 + expected_size = output_frame_count * frame_size + if len(result.stdout) != expected_size: + raise OmniClientError(f"decoded {len(result.stdout)} bytes from {path}, expected {expected_size}") + return np.frombuffer(result.stdout, dtype=np.uint8).reshape( + output_frame_count, + height, + width, + 3, + ) + + def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: meta = _probe_video(prepared_path) ratio = MINIMAX_H3_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS @@ -434,11 +546,17 @@ def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: if not indices: raise OmniClientError(f"no frames sampled from {prepared_path}") - # The preparation step emits a lossless RGB stream. Decode it once and - # sample from that array so Qwen3VL sees exactly the same pixels as the - # video VAE, without starting one full-video ffmpeg decode per sample. - decoded_frames = load_video_frames(prepared_path) - frames = [np.asarray(decoded_frames[source_index]) for source_index in indices] + # The preparation step emits a lossless RGB stream. Decode only the sampled + # frames in one ffmpeg process so Qwen3VL sees the exact prepared pixels + # without decoding the entire high-bitrate stream into host memory. + decoded_frames = _decode_video_frames_ffmpeg( + prepared_path, + frame_count=int(meta["frame_count"]), + indices=indices, + width=meta["width"], + height=meta["height"], + ) + frames = [np.asarray(frame) for frame in decoded_frames] timestamps = [index / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))] timestamps += [timestamps[-1]] * ((-len(timestamps)) % MINIMAX_H3_QWEN_TEMPORAL_PATCH) diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py new file mode 100644 index 000000000..85583f243 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 text encoder on the vLLM Qwen3-VL model runner.""" + +from __future__ import annotations + +import copy +from collections.abc import Iterable, Mapping, Sequence +from typing import Any + +import regex as re +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.model_executor.models.qwen3_vl import ( + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.multimodal import MULTIMODAL_REGISTRY + +from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MINIMAX_H3_CONDITION_LABELS_KEY, + MINIMAX_H3_PRESENTATION_TASK_KEY, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + IMAGE_PAD, + VIDEO_PAD, + VISION_END, + VISION_START, + build_minimax_h3_presentation, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + + +def _build_minimax_h3_presentation( + tokenizer: Any, + *, + prompt: str, + task: str, + condition_labels: list[tuple[str, int]], + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + video_timestamps: Sequence[Sequence[float]] | None, + merge_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build the same token stream consumed by the fused H3 encoder.""" + return build_minimax_h3_presentation( + tokenizer, + prompt=prompt, + task=task, + condition_labels=condition_labels, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + video_timestamps=video_timestamps, + merge_size=merge_size, + ) + + +class MiniMaxH3MultiModalProcessor(Qwen3VLMultiModalProcessor): + """Qwen3-VL media processing with H3's exact segmented presentation.""" + + @staticmethod + def _base_processor_kwargs( + kwargs: Mapping[str, object], + ) -> dict[str, object]: + return { + key: value + for key, value in kwargs.items() + if key + not in { + MINIMAX_H3_PRESENTATION_TASK_KEY, + MINIMAX_H3_CONDITION_LABELS_KEY, + } + } + + @staticmethod + def _condition_labels(value: object) -> list[tuple[str, int]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + return [] + labels: list[tuple[str, int]] = [] + for item in value: + if not isinstance(item, Sequence) or isinstance(item, (str, bytes)) or len(item) != 2: + raise ValueError("MiniMax H3 condition labels must be (type, index) pairs") + labels.append((str(item[0]), int(item[1]))) + return labels + + def _cached_apply_hf_processor(self, inputs: Any, timing_ctx: Any): + # H3 presentation IDs depend on the processed grids for every media + # item. On a sender-cache hit, the generic processor passes only cache + # misses to _apply_hf_processor_main, so the full presentation cannot + # be reconstructed there. Reprocess H3 media to preserve exact IDs. + return self._apply_hf_processor(inputs, timing_ctx) + + def _apply_hf_processor_main( + self, + prompt: str | list[int], + mm_items: Any, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + *, + enable_hf_prompt_update: bool, + ): + task_value = hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY) + if task_value is None: + return super()._apply_hf_processor_main( + prompt, + mm_items, + hf_processor_mm_kwargs, + tokenization_kwargs, + enable_hf_prompt_update=enable_hf_prompt_update, + ) + if not isinstance(prompt, str): + raise ValueError("MiniMax H3 presentation requires the original prompt text") + + task = str(task_value) + condition_labels = self._condition_labels(hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY)) + base_kwargs = self._base_processor_kwargs(hf_processor_mm_kwargs) + image_count = sum(kind == "image" for kind, _ in condition_labels) + video_count = sum(kind == "video" for kind, _ in condition_labels) + image_slot = f"{VISION_START}{IMAGE_PAD}{VISION_END}" + hf_processor = self.info.get_hf_processor(**base_kwargs) + video_slot = ( + VIDEO_PAD if self._expands_only_video_token(hf_processor) else f"{VISION_START}{VIDEO_PAD}{VISION_END}" + ) + processing_prompt = image_slot * image_count + video_slot * video_count + if not processing_prompt: + processing_prompt = prompt + + _, processed_data, _ = super()._apply_hf_processor_main( + processing_prompt, + mm_items, + base_kwargs, + tokenization_kwargs, + enable_hf_prompt_update=enable_hf_prompt_update, + ) + image_processor = self.info.get_image_processor(**base_kwargs) + ids, _ = _build_minimax_h3_presentation( + self.info.get_tokenizer(), + prompt=prompt, + task=task, + condition_labels=condition_labels, + image_grid_thw=processed_data.get("image_grid_thw"), + video_grid_thw=processed_data.get("video_grid_thw"), + video_timestamps=processed_data.get("timestamps"), + merge_size=int(image_processor.merge_size), + ) + return ids.tolist(), processed_data, True + + def _get_prompt_updates( + self, + mm_items: Any, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: Any, + ): + return super()._get_prompt_updates( + mm_items, + self._base_processor_kwargs(hf_processor_mm_kwargs), + out_mm_kwargs, + ) + + +class _ResidualMerge(nn.Module): + """Return the decoder residual stream without the final RMSNorm.""" + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, None]: + if residual is not None: + hidden_states = hidden_states + residual + return hidden_states, None + + +@MULTIMODAL_REGISTRY.register_processor( + MiniMaxH3MultiModalProcessor, + info=Qwen3VLProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class MiniMaxH3TextEncoder(Qwen3VLForConditionalGeneration): + """Qwen3-VL encoder used by MiniMax H3. + + MiniMax H3 consumes the residual stream after decoder layer 50. The + upstream vLLM Qwen3-VL implementation supplies the vision tower, + multimodal processor, M-RoPE, DeepStack injection, paged attention, and + tensor-parallel decoder layers. This class changes only the two parts of + the contract that differ from causal generation: decoder depth and output + normalization. + """ + + have_multimodal_outputs = True + omni_pooler_payload_include_hidden = False + requires_full_prefix_cached_hidden_states = False + num_encoder_layers = 50 + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): + if vllm_config.parallel_config.pipeline_parallel_size != 1: + raise ValueError("MiniMax H3 text encoding supports tensor parallelism, not pipeline parallelism") + + hf_config = copy.deepcopy(vllm_config.model_config.hf_config) + hf_config.text_config.num_hidden_layers = self.num_encoder_layers + # Reuse the token embedding as the unused LM head, avoiding a second + # vocab-sized allocation in the upstream constructor. Set the flag + # on both levels because VllmConfig.with_hf_config propagates the + # multimodal parent value into its nested text config. + hf_config.tie_word_embeddings = True + hf_config.text_config.tie_word_embeddings = True + encoder_config = vllm_config.with_hf_config(hf_config) + super().__init__(vllm_config=encoder_config, prefix=prefix) + + self.language_model.model.norm = _ResidualMerge() + self._token_tags: torch.Tensor | None = None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + # Qwen3-VL expands every visual block to exactly + # vision_start + image/video tokens + vision_end. Classifying each + # token directly makes tags independent of prefill chunk boundaries. + visual_token = ( + (input_ids == self.config.vision_start_token_id) + | (input_ids == self.config.vision_end_token_id) + | (input_ids == self.config.image_token_id) + | (input_ids == self.config.video_token_id) + ) + token_tags = (~visual_token).long() + self._token_tags = None + embeddings = super().embed_input_ids( + input_ids, + multimodal_embeddings, + is_multimodal=is_multimodal, + ) + self._token_tags = token_tags + return embeddings + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **_: object, + ) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + if self._token_tags is None: + raise RuntimeError("token tags were not prepared before MiniMax H3 text encoding") + token_tags = self._token_tags + self._token_tags = None + padding = model_outputs.shape[0] - token_tags.shape[0] + if padding < 0: + raise RuntimeError( + "MiniMax H3 token tags exceed the model output length: " + f"tags={token_tags.shape[0]}, outputs={model_outputs.shape[0]}" + ) + if padding: + # Multimodal embedding runs on scheduled tokens, while the AR + # forward may use a CUDA-graph-padded token buffer. Match that + # buffer so both conditioning tensors are recognized as + # per-token prefix-cache entries; runner slicing removes these + # padding rows before the stage payload is emitted. + token_tags = torch.cat((token_tags, token_tags.new_ones(padding))) + # Keep the stage-wire representation two-dimensional. The Omni + # runner treats tensors with shape [tokens, features] as per-token + # data, so they are concatenated across chunked-prefill steps and + # reconstructed on prefix-cache hits. The Stage 1 adapter restores + # H3's semantic [tokens] representation. + token_tags = token_tags.unsqueeze(-1) + return OmniOutput( + text_hidden_states=model_outputs, + multimodal_outputs={ + "hidden_states": {"output": model_outputs}, + "meta": {"token_role_ids": token_tags}, + }, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + **_: object, + ) -> torch.Tensor: + """Emit EOS immediately; the useful result is the prefill hidden state.""" + logits = hidden_states.new_full( + (hidden_states.shape[0], self.config.text_config.vocab_size), + torch.finfo(hidden_states.dtype).min, + ) + logits[:, int(self._tokenizer.eos_token_id)] = 0 + return logits + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + def encoder_weights(): + for name, weight in weights: + layer_match = re.match(r"model\.language_model\.layers\.(\d+)\.", name) + if layer_match and int(layer_match.group(1)) >= self.num_encoder_layers: + continue + if name == "model.language_model.norm.weight" or name.startswith("lm_head."): + continue + yield name, weight + + loader = AutoWeightsLoader(self) + return loader.load_weights( + encoder_weights(), + mapper=self.hf_to_vllm_mapper, + ) diff --git a/vllm_omni/model_executor/models/minimax_music3/__init__.py b/vllm_omni/model_executor/models/minimax_music3/__init__.py new file mode 100644 index 000000000..39e617d04 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax Music 3 text-to-music. + +Stage 0 is a Qwen3 backbone with an eight-codebook RVQ frame; stage 1 turns +the resulting conditioning frames into 32 kHz stereo audio through a +flow-matching transformer and a DAC decoder. +""" + +from .constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + AR_HIDDEN_SIZE, + MAX_AUDIO_FRAMES, + OUTPUT_SAMPLE_RATE, +) +from .prompt import ( + AUDIO_CODE_OFFSET, + SPECIAL_TOKEN_IDS, + build_cfg_null_token_ids, + build_prompt, +) + +__all__ = [ + "AR_CHUNK_FRAMES", + "AR_CHUNK_HOP_FRAMES", + "AR_HIDDEN_SIZE", + "AUDIO_CODE_OFFSET", + "MAX_AUDIO_FRAMES", + "OUTPUT_SAMPLE_RATE", + "SPECIAL_TOKEN_IDS", + "build_cfg_null_token_ids", + "build_prompt", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/acoustic.py b/vllm_omni/model_executor/models/minimax_music3/acoustic.py new file mode 100644 index 000000000..92fa332d0 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/acoustic.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stage-1 acoustic decoder for MiniMax Music 3. + +Consumes the AR stage's conditioning spans and returns 32 kHz stereo audio. +Each span is cut into 200-frame windows that overlap by half; every window is +solved by the flow-matching DiT with the previous window's tail latent pinned +as a prompt, decoded by the vocoder, and cropped so that concatenating the +survivors reconstructs the song exactly once. + +The window arithmetic lives in :mod:`.chunking` and is driven entirely by the +span's global frame range, so the streaming path (one window per engine step) +and the non-streaming path (the whole song in one payload) produce identical +audio for the same seed. + +The stage runs in float32. TF32 keeps that affordable, and bfloat16 measurably +degrades a 30-step solve. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Any + +import torch +from torch import Tensor, nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .chunking import ChunkWindow, chunk_windows, crop_sample_bounds, overlap_mel_length +from .constants import ( + AR_CHUNK_HOP_FRAMES, + AR_HIDDEN_SIZE, + DEFAULT_DIT_CFG_SCALE, + DEFAULT_DIT_STEPS, + OUTPUT_CHANNELS, + OUTPUT_SAMPLE_RATE, +) +from .dav import MiniMaxMusic3Vocoder, StreamingResampler, remove_weight_norm +from .dit import MiniMaxMusic3FlowMatchingDiT, remap_transformer_state +from .weights import ( + CONDITION_ENCODER_DIR, + TRANSFORMER_DIR, + VOCODER_DIR, + load_component_state, + resolve_repo_root, +) + +logger = init_logger(__name__) + +__all__ = ["MiniMaxMusic3AcousticForConditionalGeneration"] + +# Namespaced so a request seed reused by another component cannot collide with +# the noise draw, and so the draw is reproducible across processes. Python's +# hash() is randomized per interpreter and must not be used here. +_SEED_PERSON = b"minimax-ttm" +_SEED_BYTES = 8 + + +def _derive_seed(seed: int, *parts: object) -> int: + """Derive a stable 63-bit sub-seed from a request seed and some labels.""" + digest = hashlib.blake2b(digest_size=_SEED_BYTES, person=_SEED_PERSON) + digest.update(int(seed).to_bytes(8, "little", signed=False)) + for part in parts: + raw = str(part).encode("utf-8") + digest.update(len(raw).to_bytes(4, "little")) + digest.update(raw) + return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1) + + +def _meta_int(value: Any, default: int = 0) -> int: + """Read an int that transport may have wrapped in a tensor or a list.""" + if isinstance(value, torch.Tensor): + return int(value.reshape(-1)[0].item()) if value.numel() else default + if isinstance(value, list | tuple): + return _meta_int(value[0], default) if value else default + try: + return int(value) # type: ignore[arg-type] + except (TypeError, ValueError): + return default + + +def _meta_bool(value: Any) -> bool: + """Read a bool that transport may have wrapped in a tensor or a list.""" + if isinstance(value, torch.Tensor): + return bool(value.reshape(-1)[0].item()) if value.numel() else False + if isinstance(value, list | tuple): + return _meta_bool(value[0]) if value else False + return bool(value) + + +def _meta_str(value: Any) -> str | None: + """Read a string that transport may have wrapped in a list.""" + if isinstance(value, list | tuple): + return _meta_str(value[0]) if value else None + return value if isinstance(value, str) else None + + +@dataclass +class _RequestState: + """Cross-window streaming state, keyed by request id and never by row.""" + + last_latent: Tensor | None = None + last_condition: Tensor | None = None + next_window_index: int = 0 + decoded_windows: int = 0 + resampler: StreamingResampler = field(default_factory=StreamingResampler) + + def release(self) -> None: + self.last_latent = None + self.last_condition = None + self.resampler.reset() + + +class MiniMaxMusic3AcousticForConditionalGeneration(nn.Module): + """Flow-matching DiT plus DAC vocoder, driven by the generation runner.""" + + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + del prefix + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._configure_backends() + + # Skeleton only. vLLM's memory profiler walks the module tree at + # startup, before any weight loader has run. + self.dit = MiniMaxMusic3FlowMatchingDiT() + self.vocoder = MiniMaxMusic3Vocoder() + + extra = self._connector_extra_config() + self.dit_steps = max(1, _meta_int(extra.get("dit_steps"), DEFAULT_DIT_STEPS)) + cfg_scale = extra.get("dit_cfg_scale") + self.dit_cfg_scale = ( + float(cfg_scale) + if isinstance(cfg_scale, int | float) and not isinstance(cfg_scale, bool) + else DEFAULT_DIT_CFG_SCALE + ) + + self._states: dict[str, _RequestState] = {} + self._loaded = False + + @staticmethod + def _configure_backends() -> None: + """Pin the numeric backends this checkpoint was validated against. + + These are process-global switches, which is safe here only because a + stage owns its own engine-core process. Never call this from a module + shared with the AR stage: cuDNN's fused attention is not + bit-reproducible for grouped-query shapes, so flipping it underneath + the backbone makes the same seed produce a different song depending on + which stage loaded first. + + cuDNN itself is off because its 1-D convolution paths are slower than + the native ones for the vocoder's dilated stacks, and its SDPA backend + is not needed because the DiT runs the flash path. TF32 keeps + float32's range with a 10-bit mantissa, which is what makes a float32 + solve affordable. + """ + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + if torch.cuda.is_available(): + torch.backends.cudnn.enabled = False + torch.backends.cuda.enable_cudnn_sdp(False) + + def _connector_extra_config(self) -> dict[str, Any]: + """Return the stage connector's ``extra`` block, or an empty mapping.""" + connector = getattr(self.vllm_config.model_config, "stage_connector_config", None) + if isinstance(connector, dict): + extra = connector.get("extra", connector) + return extra if isinstance(extra, dict) else {} + extra = getattr(connector, "extra", None) + return extra if isinstance(extra, dict) else {} + + # ------------------------------------------------------------- vLLM hooks + + def embed_input_ids(self, input_ids: Tensor, **_: Any) -> Tensor: + """Stable dummy embedding: this stage never reads token embeddings.""" + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + del hidden_states, sampling_metadata + return None + + def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> set[str]: + """Load the three acoustic components from their own folders. + + The iterator vLLM hands over covers the stage's own checkpoint folder, + which holds nothing this stage needs. It is drained anyway: the loader + streams from a file handle the caller closes only once the iterator is + exhausted, so abandoning it wedges startup. + """ + try: + for _ in weights: + pass + except RuntimeError as exc: + # The iterator is lazy, so draining it is what actually asks the + # default loader for root-level weight files. This repository keeps + # every component in its own subfolder and has none at the root, so + # that lookup raises. Nothing here needs those weights; any other + # RuntimeError is a real problem and is re-raised. + if "Cannot find any model weights" not in str(exc): + raise + logger.debug("MiniMax Music 3 acoustic stage has no root-level weights, as expected") + return self._load_components() + + def _ensure_loaded(self) -> None: + """Load on first use when the weight loader skipped ``load_weights``. + + ``load_format=dummy`` (used by warmup and memory profiling) never calls + ``load_weights``, which would leave the modules holding their random + skeleton and make the dummy run produce noise or crash. + """ + if not self._loaded: + self._load_components() + + def _load_components(self) -> set[str]: + """Read ``condition_encoder/``, ``transformer/`` and ``vocoder/``. + + Every component loads with ``strict=True``. A silently partial load + would produce audio that sounds plausible but is wrong, which is far + worse than a startup failure. + """ + root = resolve_repo_root( + self.model_path, + revision=self.vllm_config.model_config.revision, + download_dir=self.vllm_config.load_config.download_dir, + ) + encoder = self.dit.condition_encoder + encoder.load_state_dict(load_component_state(root, CONDITION_ENCODER_DIR), strict=True) + transformer_state = remap_transformer_state(load_component_state(root, TRANSFORMER_DIR)) + self.dit.transformer.load_state_dict(transformer_state, strict=True) + self.vocoder.load_state_dict(load_component_state(root, VOCODER_DIR), strict=True) + + device = self.vllm_config.device_config.device + self.dit.to(device=device, dtype=torch.float32).eval() + self.vocoder.to(device=device, dtype=torch.float32).eval() + folded = remove_weight_norm(self.vocoder) + self._loaded = True + + loaded = {f"dit.{name}" for name, _ in self.dit.named_parameters()} + loaded |= {f"vocoder.{name}" for name, _ in self.vocoder.named_parameters()} + logger.info( + "MiniMax Music 3 acoustic loaded from %s device=%s dit_steps=%d " + "dit_cfg_scale=%.3f attention=%s folded_weight_norms=%d parameters=%d", + root, + device, + self.dit_steps, + self.dit_cfg_scale, + self.dit.transformer.attention_backend_name, + folded, + sum(p.numel() for p in self.parameters()), + ) + return loaded + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Drop the streaming state of finished or aborted requests.""" + for req_id in finished_req_ids: + state = self._states.pop(str(req_id), None) + if state is not None: + logger.info( + "MiniMax Music 3 acoustic done request=%s windows=%d", + req_id, + state.decoded_windows, + ) + state.release() + + def make_omni_output(self, model_outputs: Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput: + del kwargs + if isinstance(model_outputs, OmniOutput): + return model_outputs + # The runner may unpack the NamedTuple in transit; rebuild it. + if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields): + return OmniOutput(*model_outputs) + raise TypeError(f"MiniMaxMusic3AcousticForConditionalGeneration expected OmniOutput, got {type(model_outputs)}") + + # --------------------------------------------------------------- decoding + + @staticmethod + def _silence() -> Tensor: + """Placeholder for a step that produced no audio. + + Deliberately rank 2. A rank-1 empty tensor mixed with rank-2 chunks + makes the multimodal collector's ``torch.cat(..., dim=-1)`` fall back + to a flattening path that would planar-ize the stereo. + """ + return torch.zeros((OUTPUT_CHANNELS, 0), dtype=torch.float32) + + def _decode_window(self, state: _RequestState, hidden: Tensor, window: ChunkWindow, seed: int) -> Tensor: + """Solve and vocode one window, cropped to its unique contribution. + + Everything stays on the device: the crop is a slice and the result is + handed straight to the resampler, so a step costs one device-to-host + copy however many windows it decoded. + """ + device = self.vllm_config.device_config.device + hidden = hidden.unsqueeze(0).to(device=device, dtype=torch.float32) + generator = torch.Generator(device=device).manual_seed(_derive_seed(seed, "dit", window.index)) + align = self.dit.aligned_condition(hidden) + latent = self.dit( + align, + generator=generator, + initial_latent=state.last_latent, + initial_condition=state.last_condition, + num_steps=self.dit_steps, + cfg_scale=self.dit_cfg_scale, + ) + waveform = self.vocoder(latent) + + # Carry the middle of the overlap forward. The solver re-imposes it on + # the next window, so the two windows agree where they meet. + overlap = overlap_mel_length() + start = max(0, latent.shape[-1] - 2 * overlap) + end = max(start, latent.shape[-1] - overlap) + state.last_latent = latent[..., start:end].clone() + state.last_condition = align[..., start:end].clone() + + wave = waveform[0].clamp(-1.0, 1.0).float() + left, right = crop_sample_bounds(window) + stop = wave.shape[-1] - right if right else wave.shape[-1] + return wave[:, left:stop] + + def _decode_span(self, req_id: str, span: Tensor | None, meta: dict[str, Any]) -> Tensor: + """Decode every window a span covers and return this step's audio. + + A span is normally one window, but the terminal emission carries the + held-back complete window plus the truncated final one, and the + non-streaming path carries the whole song. Window boundaries and the + crop are re-derived from the span's global frame range so all three + cases go through one code path. + + Raises: + ValueError: If the span is malformed or lands on a window other + than the one this request is waiting for. + """ + span_start = _meta_int(meta.get("num_processed_tokens"), 0) + is_terminal = _meta_bool(meta.get("last_chunk")) or _meta_bool(meta.get("stream_finished")) + seed = _meta_int(meta.get("audio_seed"), 0) + state = self._states.setdefault(req_id, _RequestState()) + + # A metadata-only terminal marker carries no conditioning; it decodes + # zero windows and exists only to flush the resampler. + windows = chunk_windows(int(span.shape[0])) if span is not None else [] + pieces: list[Tensor] = [] + with torch.inference_mode(): + for local in windows: + start = span_start + local.start + index = start // AR_CHUNK_HOP_FRAMES + if index != state.next_window_index: + # Each window is prompted by its predecessor's tail, so a + # gap or a replay corrupts the join. Fail the request here + # rather than return audio with a seam in it. + raise ValueError( + f"MiniMax Music 3 acoustic request {req_id} expected window " + f"{state.next_window_index} but received window {index} " + f"(span starts at AR frame {span_start}); the overlap prompt " + f"for this window is stale and the audio would not join" + ) + window = ChunkWindow( + index=index, + start=start, + end=span_start + local.end, + is_first=start == 0, + is_last=local.is_last and is_terminal, + ) + pieces.append(self._decode_window(state, span[local.start : local.end], window, seed)) + state.next_window_index = index + 1 + state.decoded_windows += 1 + + emitted = state.resampler.push(pieces, flush=is_terminal) + return emitted.cpu() + + def _span_from_info(self, info: dict[str, Any], req_id: str) -> Tensor | None: + """Return the conditioning span as ``[T, 32768]``, or ``None``. + + ``None`` means this step carried no conditioning, which is normal: a + request can be scheduled before its span arrives, and the terminal + marker is metadata only. A span that is present but the wrong shape is + not normal and stops the request. + + Raises: + ValueError: If a span is present but is not ``[T, 32768]``. + """ + latent = info.get("latent") + if not isinstance(latent, Tensor) or latent.numel() == 0: + return None + if latent.ndim == 3 and latent.shape[0] == 1: + latent = latent[0] + if latent.ndim != 2 or latent.shape[1] != AR_HIDDEN_SIZE: + raise ValueError( + f"MiniMax Music 3 acoustic request {req_id} carries a conditioning " + f"span of shape {tuple(latent.shape)}; expected [T,{AR_HIDDEN_SIZE}]" + ) + return latent + + def forward( + self, + input_ids: Tensor | None = None, + positions: Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + """Decode this step's spans into one stereo waveform per request. + + Raises: + ValueError: If a payload is malformed. Failing the step is + deliberate: a request whose conditioning cannot be decoded has + no audio to return, and reporting an error is better than + leaving the caller waiting on a stream that will never end. + """ + del input_ids, positions, intermediate_tensors, inputs_embeds + self._ensure_loaded() + + infos = runtime_additional_information + if infos is None: + infos = kwargs.get("model_intermediate_buffer") + infos = infos or [] + + # One entry per scheduled request, including steps where a request's + # span has not arrived yet. + seq_token_counts = kwargs.get("seq_token_counts") + num_reqs = len(seq_token_counts) if seq_token_counts is not None else len(infos) + + audios: list[Tensor] = [] + sr_tensor = torch.tensor(OUTPUT_SAMPLE_RATE, dtype=torch.int32) + for index in range(num_reqs): + info = infos[index] if index < len(infos) else {} + info = info if isinstance(info, dict) else {} + meta = info.get("meta") + meta = meta if isinstance(meta, dict) else {} + req_id = _meta_str(meta.get("req_id")) or _meta_str(info.get("req_id")) + has_span = isinstance(info.get("latent"), Tensor) and info["latent"].numel() > 0 + is_terminal = _meta_bool(meta.get("last_chunk")) or _meta_bool(meta.get("stream_finished")) + if not has_span and not is_terminal: + audios.append(self._silence()) + continue + if req_id is None: + raise ValueError( + "MiniMax Music 3 acoustic payload carries no meta.req_id, so its " + "window cannot be joined to the rest of its request" + ) + audios.append(self._decode_span(req_id, self._span_from_info(info, req_id), meta)) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_reqs}, + ) diff --git a/vllm_omni/model_executor/models/minimax_music3/chunking.py b/vllm_omni/model_executor/models/minimax_music3/chunking.py new file mode 100644 index 000000000..ba917f39d --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/chunking.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Chunk boundary and overlap arithmetic for MiniMax Music 3. + +The AR stage emits fixed 200-frame windows that overlap by half. The acoustic +stage decodes each window with the previous window's latent as a prompt, then +crops the overlap back off so the concatenated waveform is seamless. All of +that is pure integer arithmetic and lives here so it can be unit tested +without a GPU. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + BLEND_MEL_FRAMES, + DAV_HOP_SAMPLES, + HOP_MEL_FRAMES, +) + + +@dataclass(frozen=True) +class ChunkWindow: + """One acoustic window, in AR frame coordinates.""" + + index: int + start: int + end: int + is_first: bool + is_last: bool + + @property + def length(self) -> int: + return self.end - self.start + + +def chunk_windows(frames: int) -> list[ChunkWindow]: + """Return the fixed-size windows covering ``frames``, overlapping by half. + + A song shorter than one window produces a single short window; otherwise + windows advance by ``AR_CHUNK_HOP_FRAMES`` and the final one is truncated + at the end of the song. + """ + frames = int(frames) + if frames <= 0: + return [] + if frames <= AR_CHUNK_FRAMES: + return [ChunkWindow(0, 0, frames, True, True)] + + windows: list[ChunkWindow] = [] + index = 0 + start = 0 + while start < frames: + end = min(start + AR_CHUNK_FRAMES, frames) + windows.append( + ChunkWindow( + index=index, + start=start, + end=end, + is_first=index == 0, + is_last=end >= frames, + ) + ) + if end >= frames: + break + index += 1 + start += AR_CHUNK_HOP_FRAMES + return windows + + +def overlap_mel_length() -> int: + """Latent frames carried forward as the next window's flow-matching prompt.""" + return HOP_MEL_FRAMES // 2 + + +def crop_sample_bounds(window: ChunkWindow) -> tuple[int, int]: + """Return ``(left, right)`` waveform samples to trim from a decoded window. + + The first window keeps its head and the last keeps its tail; every other + window is cropped on both sides so that concatenating the survivors + reconstructs the song exactly once. + """ + left = 0 if window.is_first else BLEND_MEL_FRAMES * DAV_HOP_SAMPLES + right = 0 if window.is_last else (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES + return left, right + + +__all__ = [ + "ChunkWindow", + "chunk_windows", + "crop_sample_bounds", + "overlap_mel_length", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/constants.py b/vllm_omni/model_executor/models/minimax_music3/constants.py new file mode 100644 index 000000000..4e93139c2 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/constants.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-contract constants for MiniMax Music 3. + +These values are fixed by the checkpoint, not tuning knobs. Changing any of +them changes the audio, so they are declared once here and imported +everywhere rather than re-derived per call site. +""" + +# --- Generation limits ----------------------------------------------------- + +# 25 audio frames per second, capped at six minutes of song. +AUDIO_FRAME_RATE = 25 +MAX_AUDIO_FRAMES = 9_000 +DEFAULT_MAX_AUDIO_FRAMES = 9_000 + +# The backbone was trained with a 10,240-token context; the prompt is capped +# well below it so there is room for the audio frames that follow. +MAX_PROMPT_TOKENS = 5_000 + +# --- Frame layout ---------------------------------------------------------- + +# Eight RVQ codebooks per frame: c0 lives in the backbone vocabulary, c1..c7 +# in the depth decoder's own 1024-entry books. +NUM_CODEBOOKS = 8 +C0_VOCAB_SIZE = 16_384 +AUDIO_VOCAB_SIZE = 1_024 + +# The acoustic stage is conditioned on the backbone hidden state concatenated +# with the seven depth-decoder hidden states. +BACKBONE_HIDDEN_SIZE = 4_096 +AR_HIDDEN_SIZE = BACKBONE_HIDDEN_SIZE * NUM_CODEBOOKS # 32768 + +# --- Sampling -------------------------------------------------------------- + +# Classifier-free guidance is not optional and has no request field: this is +# how the reference implementation samples this checkpoint. +AR_CFG_SCALE = 1.5 +AR_CFG_TOP_K = 50 +AR_TOP_K = 50 + +# --- Acoustic stage -------------------------------------------------------- + +DEFAULT_DIT_STEPS = 30 +DEFAULT_DIT_CFG_SCALE = 1.7 +DIT_LATENT_CHANNELS = 128 + +# Windowing: the AR stage hands the acoustic stage 200-frame windows that +# overlap by half, and the acoustic stage crops the overlap back off. +AR_CHUNK_FRAMES = 200 +AR_CHUNK_HOP_FRAMES = 100 + +# DAV runs at 44.1 kHz; the response is resampled to 32 kHz stereo. +DAV_SAMPLE_RATE = 44_100 +OUTPUT_SAMPLE_RATE = 32_000 +OUTPUT_CHANNELS = 2 + +# Condition resampling: AR frames land on a 24 kHz / 960-hop grid and are +# stretched onto the DAV latent's 44.1 kHz / 512-hop grid. +CONDITION_INPUT_SAMPLE_RATE = 24_000 +CONDITION_INPUT_HOP = 960 +CONDITION_OUTPUT_SAMPLE_RATE = 44_100 +CONDITION_OUTPUT_HOP = 512 + +# Mel frames per AR frame: 44100/24000 * 960/512 = 3.4453125. So a 100-frame +# hop is 344 mel frames and a 200-frame window is 689. +# +# HOP_MEL_FRAMES is the hop, not the window: the crop and carry-forward +# arithmetic is expressed in units of the hop, and a window is two of them +# plus one frame of rounding. Reading this as the window length is the easy +# mistake here, and it is off by a factor of two. +HOP_MEL_FRAMES = 344 +BLEND_MEL_FRAMES = HOP_MEL_FRAMES // 4 +DAV_HOP_SAMPLES = 512 + +# Only these two dtypes are supported for the acoustic stage. float32 is the +# default: TF32 keeps it affordable, and bfloat16 measurably degrades the +# solver. +SUPPORTED_ACOUSTIC_DTYPES = frozenset({"float32", "bfloat16"}) + + +__all__ = [ + "AR_CFG_SCALE", + "AR_CFG_TOP_K", + "AR_CHUNK_FRAMES", + "AR_CHUNK_HOP_FRAMES", + "AR_HIDDEN_SIZE", + "AR_TOP_K", + "AUDIO_FRAME_RATE", + "AUDIO_VOCAB_SIZE", + "BACKBONE_HIDDEN_SIZE", + "BLEND_MEL_FRAMES", + "C0_VOCAB_SIZE", + "CONDITION_INPUT_HOP", + "CONDITION_INPUT_SAMPLE_RATE", + "CONDITION_OUTPUT_HOP", + "CONDITION_OUTPUT_SAMPLE_RATE", + "DAV_HOP_SAMPLES", + "DAV_SAMPLE_RATE", + "DEFAULT_DIT_CFG_SCALE", + "DEFAULT_DIT_STEPS", + "DEFAULT_MAX_AUDIO_FRAMES", + "DIT_LATENT_CHANNELS", + "HOP_MEL_FRAMES", + "MAX_AUDIO_FRAMES", + "MAX_PROMPT_TOKENS", + "NUM_CODEBOOKS", + "OUTPUT_CHANNELS", + "OUTPUT_SAMPLE_RATE", + "SUPPORTED_ACOUSTIC_DTYPES", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/dav.py b/vllm_omni/model_executor/models/minimax_music3/dav.py new file mode 100644 index 000000000..bc0a478fd --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/dav.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DAC-style vocoder for MiniMax Music 3. + +Only the decoder half of the autoencoder ships in the checkpoint, which is all +inference needs: the flow-matching stage produces the latent directly, so the +encoder and the quantizer never run. + +Stereo is folded into the batch. A ``[B, 128, T]`` latent is read as two +64-channel mono latents per item, decoded as ``2B`` independent mono streams, +and unfolded again at the end. The four transposed-convolution blocks upsample +by 8, 8, 4 and 2, so one latent frame becomes 512 samples at 44.1 kHz. + +Module names match the ``vocoder/`` component folder key for key, so its state +dict loads with ``strict=True`` and no remapping. The convolutions keep the +checkpoint's weight-norm parameterization (``weight_g``/``weight_v``) through +loading; :func:`remove_weight_norm` folds it away afterwards, since the +factorization is fixed once the weights stop training. + +:class:`StreamingResampler` carries the 44.1 kHz output down to the 32 kHz +response rate across chunk boundaries, which is not something a plain +per-chunk resample can do. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence + +import torch +from torch import Tensor, nn +from vllm.logger import init_logger + +from .constants import DAV_SAMPLE_RATE, DIT_LATENT_CHANNELS, OUTPUT_CHANNELS, OUTPUT_SAMPLE_RATE + +logger = init_logger(__name__) + +__all__ = ["MiniMaxMusic3Vocoder", "StreamingResampler", "remove_weight_norm", "snake"] + +# Fixed by vocoder/config.json. +_DECODER_INPUT_DIM = 1_024 +_DECODER_HIDDEN_DIM = 1_536 +_UPSAMPLING_RATIOS = (8, 8, 4, 2) +_RESIDUAL_DILATIONS = (1, 3, 9) +_KERNEL_SIZE = 7 +# Each vocoder stream sees half the latent channels; the other half is the +# opposite stereo channel. +_STREAM_CHANNELS = DIT_LATENT_CHANNELS // OUTPUT_CHANNELS + + +def snake(x: Tensor, alpha: Tensor) -> Tensor: + """Periodic ``x + sin^2(alpha x) / alpha`` activation. + + The reshape to three dimensions is what the reference does and is a no-op + for the rank-3 activations used here; it is kept so the numerics are + identical for any rank. + """ + shape = x.shape + flat = x.reshape(shape[0], shape[1], -1) + flat = flat + (alpha + 1e-9).reciprocal() * torch.sin(alpha * flat).pow(2) + return flat.reshape(shape) + + +class Snake1d(nn.Module): + """Snake activation with one learned frequency per channel.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x: Tensor) -> Tensor: + return snake(x, self.alpha) + + +def _wn_conv(*args: object, **kwargs: object) -> nn.Module: + """A weight-normalized ``Conv1d``, named ``weight_g``/``weight_v``.""" + return nn.utils.weight_norm(nn.Conv1d(*args, **kwargs)) # type: ignore[arg-type] + + +def _wn_conv_transpose(*args: object, **kwargs: object) -> nn.Module: + """A weight-normalized ``ConvTranspose1d``, named ``weight_g``/``weight_v``.""" + return nn.utils.weight_norm(nn.ConvTranspose1d(*args, **kwargs)) # type: ignore[arg-type] + + +def remove_weight_norm(module: nn.Module) -> int: + """Fold weight normalization into every convolution weight in place. + + Returns: + How many submodules were folded. + """ + removed = 0 + for child in module.modules(): + try: + nn.utils.remove_weight_norm(child) + except (ValueError, RuntimeError): + continue + removed += 1 + return removed + + +class ResidualUnit(nn.Module): + """Dilated snake-gated residual block; preserves length and width.""" + + def __init__(self, dim: int, dilation: int) -> None: + super().__init__() + pad = (_KERNEL_SIZE - 1) * dilation // 2 + self.snake1 = Snake1d(dim) + self.conv1 = _wn_conv(dim, dim, kernel_size=_KERNEL_SIZE, dilation=dilation, padding=pad) + self.snake2 = Snake1d(dim) + self.conv2 = _wn_conv(dim, dim, kernel_size=1) + + def forward(self, x: Tensor) -> Tensor: + y = self.conv2(self.snake2(self.conv1(self.snake1(x)))) + if y.shape[-1] != x.shape[-1]: + # Defensive centre crop. The padding above makes this a no-op for + # every dilation the checkpoint ships. + pad = (x.shape[-1] - y.shape[-1]) // 2 + x = x[..., pad : x.shape[-1] - pad] + return x + y + + +class DecoderBlock(nn.Module): + """One upsampling stage: activate, stretch by ``stride``, then refine.""" + + def __init__(self, input_dim: int, output_dim: int, stride: int) -> None: + super().__init__() + self.snake1 = Snake1d(input_dim) + self.conv_t1 = _wn_conv_transpose( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + self.res_unit1 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[0]) + self.res_unit2 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[1]) + self.res_unit3 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[2]) + + def forward(self, x: Tensor) -> Tensor: + x = self.conv_t1(self.snake1(x)) + return self.res_unit3(self.res_unit2(self.res_unit1(x))) + + +class MiniMaxMusic3Vocoder(nn.Module): + """Latent to waveform: ``[B, 128, T] -> [B, 2, T * 512]`` at 44.1 kHz.""" + + def __init__(self) -> None: + super().__init__() + self.dec_in_proj = nn.Conv1d(_STREAM_CHANNELS, _DECODER_INPUT_DIM, kernel_size=1) + self.conv_in = _wn_conv( + _DECODER_INPUT_DIM, + _DECODER_HIDDEN_DIM, + kernel_size=_KERNEL_SIZE, + padding=_KERNEL_SIZE // 2, + ) + blocks: list[nn.Module] = [] + output_dim = _DECODER_HIDDEN_DIM + for index, stride in enumerate(_UPSAMPLING_RATIOS): + input_dim = _DECODER_HIDDEN_DIM // (2**index) + output_dim = _DECODER_HIDDEN_DIM // (2 ** (index + 1)) + blocks.append(DecoderBlock(input_dim, output_dim, stride)) + self.blocks = nn.ModuleList(blocks) + self.snake_out = Snake1d(output_dim) + self.conv_out = _wn_conv(output_dim, 1, kernel_size=_KERNEL_SIZE, padding=_KERNEL_SIZE // 2) + + @property + def upsampling_factor(self) -> int: + """Waveform samples produced per latent frame.""" + return math.prod(_UPSAMPLING_RATIOS) + + @torch.inference_mode() + def forward(self, latent: Tensor) -> Tensor: + """Decode a latent to interleaved stereo. + + Args: + latent: ``[B, 128, T]`` vocoder latent. + + Returns: + ``[B, 2, T * 512]`` waveform at 44.1 kHz, unclamped. + + Raises: + ValueError: If the latent is not ``[B, 128, T]``. + """ + if latent.ndim != 3 or latent.shape[1] != DIT_LATENT_CHANNELS: + raise ValueError( + f"MiniMax Music 3 vocoder latent must be [B,{DIT_LATENT_CHANNELS},T], got {tuple(latent.shape)}" + ) + bsz, _, frames = latent.shape + folded = latent.reshape(bsz * OUTPUT_CHANNELS, _STREAM_CHANNELS, frames) + h = self.conv_in(self.dec_in_proj(folded)) + for block in self.blocks: + h = block(h) + wave = torch.tanh(self.conv_out(self.snake_out(h))) + return wave.reshape(bsz, OUTPUT_CHANNELS, -1) + + +class StreamingResampler: + """Resample a chunked stream to the response rate without a seam. + + A polyphase resampler is time-invariant but not offset-invariant: it only + reproduces the same output grid when the input offset is a whole multiple + of ``orig // gcd(orig, new)`` samples, and it zero-pads both ends of every + call. Resampling each decoded chunk on its own therefore shifts the grid + by a fraction of a sample at every join and tapers the signal across it. + Measured against a single whole-stream resample on 20 s of broadband + noise, that costs a peak error of 2.2 (full scale) and drifts the length. + + This buffers instead. Only whole blocks are emitted, so every boundary + lands on an exact output sample; the previously emitted tail is prepended + as real left context and the not-yet-emitted head appended as real right + context, and both are dropped from the result. On the same measurement + the emitted stream is bit-identical to the whole-stream resample. + + One instance per request. It holds at most a couple of blocks of audio. + """ + + def __init__( + self, + *, + orig_freq: int = DAV_SAMPLE_RATE, + new_freq: int = OUTPUT_SAMPLE_RATE, + channels: int = OUTPUT_CHANNELS, + context_blocks: int = 1, + margin_blocks: int = 1, + ) -> None: + self.orig_freq = int(orig_freq) + self.new_freq = int(new_freq) + self.channels = int(channels) + # The output grid repeats every this many input samples. + self._block = self.orig_freq // math.gcd(self.orig_freq, self.new_freq) + self._context_len = self._block * context_blocks + self._margin = self._block * margin_blocks + self._pending: Tensor | None = None + self._context: Tensor | None = None + + def _empty(self, like: Tensor | None) -> Tensor: + if like is None: + return torch.zeros((self.channels, 0), dtype=torch.float32) + return like.new_zeros((self.channels, 0)) + + def reset(self) -> None: + """Drop all buffered audio.""" + self._pending = None + self._context = None + + @torch.inference_mode() + def push(self, chunks: Sequence[Tensor], *, flush: bool = False) -> Tensor: + """Append source-rate audio and return whatever is now emittable. + + Args: + chunks: New ``[channels, T]`` segments at ``orig_freq``, in order. + flush: Emit the remainder too, because the stream has ended. + + Returns: + ``[channels, T]`` at ``new_freq``; empty when nothing is ready. + """ + import torchaudio.functional as audio_functional + + parts = [chunk for chunk in chunks if chunk.numel()] + if self._pending is not None: + parts.insert(0, self._pending) + self._pending = None + if not parts: + return self._empty(self._context) + + buffered = torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0] + total = int(buffered.shape[-1]) + if flush: + emit_len, lookahead = total, 0 + else: + emit_len = max(0, ((total - self._margin) // self._block) * self._block) + lookahead = min(self._margin, total - emit_len) + if emit_len == 0: + self._pending = buffered + return self._empty(buffered) + + padded = buffered[..., : emit_len + lookahead] + start = 0 + if self._context is not None: + padded = torch.cat((self._context, padded), dim=-1) + # Exact because the context is a whole number of blocks. + start = int(self._context.shape[-1]) * self.new_freq // self.orig_freq + resampled = audio_functional.resample(padded.float(), self.orig_freq, self.new_freq) + # Round up, matching how a whole-stream resample sizes its output. This + # only bites on the flush, where the remainder is not a whole block; + # every other emission divides exactly. + stop = start + -(-emit_len * self.new_freq // self.orig_freq) + + if flush: + self.reset() + else: + self._context = buffered[..., max(0, emit_len - self._context_len) : emit_len].clone() + self._pending = buffered[..., emit_len:].clone() if emit_len < total else None + return resampled[..., start:stop] diff --git a/vllm_omni/model_executor/models/minimax_music3/depth_graph.py b/vllm_omni/model_executor/models/minimax_music3/depth_graph.py new file mode 100644 index 000000000..9c9a4f8e5 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/depth_graph.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CUDA graph capture for the RVQ depth loop. + +The depth loop is seven sequential forwards of a four-layer block over a +sequence that grows from two to eight positions, plus a guided draw after each +one. At the batch sizes this stage runs, none of those kernels is anywhere near +large enough to hide its own launch: measured in isolation on one L20X, the +loop costs 8.4 ms per frame when driven from Python and 4.0 ms when replayed +from a graph, and that figure barely moves between one and eight guided pairs. +It is launch-bound, so it is worth recording. + +Capture is lazy and keyed on the exact pair count, never padded. That keeps the +recorded kernels the same ones the eager loop selected for that shape, which is +what makes a replay bit-identical to the loop it replaces rather than merely +close to it. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +__all__ = ["DepthGraphCache"] + +# Inputs, in the order the eager loop takes them. +_DepthInputs = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +_DepthOutputs = tuple[torch.Tensor, torch.Tensor] + + +@dataclass +class _CapturedDepth: + """One recorded depth loop and the buffers it reads and writes.""" + + graph: torch.cuda.CUDAGraph + inputs: _DepthInputs + codes: torch.Tensor + depth_hidden: torch.Tensor + + +class DepthGraphCache: + """Replay the depth loop from a graph once its shape has settled. + + Args: + run_eager: The loop itself, taking ``(cond_hidden, uncond_hidden, c0, + seeds, positions)`` and returning ``(codes, depth_hidden)``. + warmup_steps: How many times a pair count runs eagerly before it is + recorded. + """ + + def __init__( + self, + run_eager: Callable[..., _DepthOutputs], + *, + warmup_steps: int = 2, + ) -> None: + self._run_eager = run_eager + self._warmup_steps = warmup_steps + self._captured: dict[int, _CapturedDepth] = {} + self._eager_steps: dict[int, int] = {} + self._disabled: set[int] = set() + + def run(self, *inputs: torch.Tensor) -> _DepthOutputs: + """Run one depth loop, recording or replaying it where possible.""" + pairs = int(inputs[2].shape[0]) + entry = self._captured.get(pairs) + if entry is None: + if not self._should_capture(pairs, inputs[0]): + return self._run_eager(*inputs) + entry = self._capture(pairs, inputs) + if entry is None: + return self._run_eager(*inputs) + + for static, live in zip(entry.inputs, inputs, strict=True): + static.copy_(live) + entry.graph.replay() + # The codes are fed back as the next frame's input and appended to the + # request's trajectory, so they outlive the step; the graph overwrites + # its output buffer on the next replay. The depth hidden states are + # consumed inside this step by the concatenation that builds the + # conditioning frame, which allocates, so that one can be handed back + # as is. + return entry.codes.clone(), entry.depth_hidden + + def _should_capture(self, pairs: int, sample: torch.Tensor) -> bool: + if pairs in self._disabled or not sample.is_cuda: + return False + if torch.cuda.is_current_stream_capturing(): + # Already inside somebody else's capture; nested graphs are not a + # thing and the outer recording covers this work anyway. + return False + count = self._eager_steps.get(pairs, 0) + 1 + self._eager_steps[pairs] = count + # Let a new pair count run eagerly a few times first. Capture records + # whatever the kernels do at that moment, and the first call into a + # cuBLAS or SDPA path for an unseen shape can allocate a workspace or + # pick an algorithm, neither of which belongs inside the recording. + return count > self._warmup_steps + + def _capture(self, pairs: int, inputs: _DepthInputs) -> _CapturedDepth | None: + try: + static = tuple(torch.empty_like(tensor) for tensor in inputs) + for target, live in zip(static, inputs, strict=True): + target.copy_(live) + graph = torch.cuda.CUDAGraph() + with ( + torch.inference_mode(), + torch.cuda.graph( + graph, + pool=current_platform.get_global_graph_pool(), + # The engine's connector threads keep their own streams + # busy while this runs. They never touch the depth loop's + # buffers, so restrict the legality check to this thread + # rather than failing the capture on their activity. + capture_error_mode="thread_local", + ), + ): + codes, depth_hidden = self._run_eager(*static) + entry = _CapturedDepth(graph=graph, inputs=static, codes=codes, depth_hidden=depth_hidden) + self._captured[pairs] = entry + logger.info("MiniMax Music 3 captured the RVQ depth CUDA graph for %d guided pair(s)", pairs) + return entry + except Exception as exc: # noqa: BLE001 - any capture failure must fall back, not crash + self._disabled.add(pairs) + logger.warning( + "MiniMax Music 3 RVQ depth CUDA graph disabled for %d guided pair(s), " + "falling back to the eager loop: %s", + pairs, + exc, + ) + return None diff --git a/vllm_omni/model_executor/models/minimax_music3/dit.py b/vllm_omni/model_executor/models/minimax_music3/dit.py new file mode 100644 index 000000000..cf180ac52 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/dit.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Condition encoder and flow-matching transformer for MiniMax Music 3. + +Stage 1 turns the AR stage's conditioning frames into a DAC latent in two +steps. The condition encoder collapses the eight 4096-wide layer slices of a +frame into one 2048-wide vector and stretches the result from the AR frame +grid onto the vocoder's latent grid. The transformer is then solved as a +flow-matching velocity field: a fixed-step Euler integrator walks Gaussian +noise to the latent the vocoder expects, under classifier-free guidance. + +Module names mirror the checkpoint's ``condition_encoder/`` and +``transformer/`` component folders key for key, with one deliberate +exception: the per-block ``to_q``/``to_k``/``to_v`` projections are fused into +a single ``to_qkv`` so each block runs one GEMM instead of three. +:func:`remap_transformer_state` performs that fusion, and is the only place +the checkpoint's layout is reinterpreted. + +Attention runs through the project's own diffusion attention layer, so this +model picks up the same backend selection, sequence-parallel plumbing and +KV-quant policy as every other transformer in the tree. That layer resolves +its backend from the diffusion config, which an ``LLM_GENERATION`` stage such +as this one does not set; it tolerates that and falls through to the platform +default. If it cannot be built at all the block runs plain SDPA instead, a +choice made once when the blocks are built and never inside ``forward``. +""" + +from __future__ import annotations + +import math + +import torch +from torch import Tensor, nn +from torch.nn import functional as F +from vllm.logger import init_logger + +from .constants import ( + AR_HIDDEN_SIZE, + BACKBONE_HIDDEN_SIZE, + CONDITION_INPUT_HOP, + CONDITION_INPUT_SAMPLE_RATE, + CONDITION_OUTPUT_HOP, + CONDITION_OUTPUT_SAMPLE_RATE, + DEFAULT_DIT_CFG_SCALE, + DEFAULT_DIT_STEPS, + DIT_LATENT_CHANNELS, + NUM_CODEBOOKS, +) + +logger = init_logger(__name__) + +__all__ = [ + "MiniMaxMusic3ConditionEncoder", + "MiniMaxMusic3FlowMatchingDiT", + "MiniMaxMusic3Transformer1DModel", + "remap_transformer_state", +] + +# Fixed by transformer/config.json. Declared rather than read from the file so +# the skeleton can be built before any checkpoint is on disk, which is what +# vLLM's memory profiler needs at startup. +_DIM = 2_048 +_NUM_LAYERS = 36 +_HEAD_DIM = 64 +_FF_INNER_DIM = 8_192 +_FOURIER_DIM = 256 +_ROTARY_DIM = 32 +_ROTARY_BASE = 10_000.0 +_CONDITION_DIM = 2_048 +# The solver feeds the transformer the current latent, a zero block of the +# same width, and the condition: 128 + 128 + 2048 = 2304. +_TRANSFORMER_IN_DIM = DIT_LATENT_CHANNELS * 2 + _CONDITION_DIM + +# The prompt interpolation floor from the reference sampler. Sigma never +# reaches exactly zero, so the noise term never fully vanishes. +_SIGMA_FLOOR = 1e-6 + + +class FourierFeatures(nn.Module): + """Random-Fourier timestep features: ``[B, 1] -> [B, out_features]``.""" + + def __init__(self, in_features: int, out_features: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(out_features // 2, in_features)) + + def forward(self, value: Tensor) -> Tensor: + f = 2.0 * math.pi * value @ self.weight.T + return torch.cat((f.cos(), f.sin()), dim=-1) + + +class TimestepEmbedding(nn.Module): + """The checkpoint's ``time_embed``: Linear, SiLU, Linear.""" + + def __init__(self, in_dim: int, dim: int) -> None: + super().__init__() + self.linear_1 = nn.Linear(in_dim, dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(dim, dim) + + def forward(self, x: Tensor) -> Tensor: + return self.linear_2(self.act(self.linear_1(x))) + + +def _rotate_half(x: Tensor) -> Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rope(x: Tensor, rope_cos: Tensor, rope_sin: Tensor) -> Tensor: + """Rotate the leading ``rope_cos.shape[-1]`` channels of ``[B, S, H, D]``. + + Only the first 32 of each head's 64 channels are rotated; the rest pass + through untouched, which is what this checkpoint was trained with. The + tables are cached at the exact sequence length, so there is nothing to + slice off the front. + """ + rot_dim = rope_cos.shape[-1] + rotated = x[..., :rot_dim] + rotated = rotated * rope_cos + _rotate_half(rotated) * rope_sin + return torch.cat((rotated, x[..., rot_dim:]), dim=-1) + + +# Set to the failure reason the first time the native attention layer cannot be +# built, so the 36 blocks agree on one path and the reason is logged once. +_NATIVE_ATTENTION_UNAVAILABLE: str | None = None + + +def _build_native_attention(*, num_heads: int, head_size: int, softmax_scale: float, prefix: str) -> nn.Module | None: + """Build the project's diffusion attention layer, or ``None`` to use SDPA. + + Anything that goes wrong here is a legitimate reason to run plain SDPA, + and is reported once rather than once per block. + """ + global _NATIVE_ATTENTION_UNAVAILABLE + if _NATIVE_ATTENTION_UNAVAILABLE is not None: + return None + try: + from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention + + return DiffusionAttention( + num_heads=num_heads, + head_size=head_size, + causal=False, + softmax_scale=softmax_scale, + prefix=prefix, + role="self", + role_category="self", + ) + except Exception as exc: # noqa: BLE001 - any failure means: use SDPA + _NATIVE_ATTENTION_UNAVAILABLE = repr(exc) + logger.warning( + "MiniMax Music 3 DiT could not build the diffusion attention layer " + "(%s); running torch SDPA instead. The stage decodes in float32, " + "for which the diffusion layer dispatches to SDPA anyway, so this " + "changes performance bookkeeping and not the audio.", + exc, + ) + return None + + +class Attention(nn.Module): + """Bidirectional self-attention with partial rotary position embeddings. + + Q, K and V are kept in ``[B, S, H, D]``, which is the layout the diffusion + attention backends consume and return. + """ + + def __init__(self, dim: int, *, head_dim: int, prefix: str = "") -> None: + super().__init__() + self.num_heads = dim // head_dim + self.head_dim = head_dim + self.softmax_scale = float(head_dim**-0.5) + # Fused from the checkpoint's separate to_q/to_k/to_v; see + # remap_transformer_state for the concatenation order. + self.to_qkv = nn.Linear(dim, dim * 3, bias=False) + self.to_out = nn.Linear(dim, dim, bias=False) + self.backend = _build_native_attention( + num_heads=self.num_heads, + head_size=head_dim, + softmax_scale=self.softmax_scale, + prefix=prefix, + ) + self.backend_name = self.backend.attn_backend.get_name() if self.backend is not None else "TORCH_SDPA" + + def _attend(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: + if self.backend is not None: + return self.backend(q, k, v) + q, k, v = (t.transpose(1, 2) for t in (q, k, v)) + out = F.scaled_dot_product_attention(q, k, v, is_causal=False, scale=self.softmax_scale) + return out.transpose(1, 2) + + def forward(self, x: Tensor, rope_cos: Tensor, rope_sin: Tensor) -> Tensor: + bsz, seq, dim = x.shape + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + q = q.reshape(bsz, seq, self.num_heads, self.head_dim) + k = k.reshape(bsz, seq, self.num_heads, self.head_dim) + v = v.reshape(bsz, seq, self.num_heads, self.head_dim) + q = _apply_rope(q, rope_cos, rope_sin) + k = _apply_rope(k, rope_cos, rope_sin) + out = self._attend(q, k, v) + return self.to_out(out.contiguous().reshape(bsz, seq, dim)) + + +class TransformerBlock(nn.Module): + """Pre-norm attention followed by a pre-norm gated feed-forward.""" + + def __init__(self, dim: int, *, head_dim: int, inner_dim: int, prefix: str = "") -> None: + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = Attention(dim, head_dim=head_dim, prefix=prefix) + self.norm2 = nn.LayerNorm(dim) + # One fused projection produces the GLU's value and gate halves. + self.ff_in = nn.Linear(dim, inner_dim * 2) + self.ff_out = nn.Linear(inner_dim, dim) + + def forward(self, x: Tensor, rope_cos: Tensor, rope_sin: Tensor) -> Tensor: + x = x + self.attn(self.norm1(x), rope_cos, rope_sin) + value, gate = self.ff_in(self.norm2(x)).chunk(2, dim=-1) + return x + self.ff_out(value * F.silu(gate)) + + +class MiniMaxMusic3Transformer1DModel(nn.Module): + """The 36-block velocity field: ``([B,128,T], [B], [B,2048,T]) -> [B,128,T]``.""" + + def __init__(self) -> None: + super().__init__() + self.preprocess_conv = nn.Conv1d(_TRANSFORMER_IN_DIM, _TRANSFORMER_IN_DIM, 1, bias=False) + self.postprocess_conv = nn.Conv1d(DIT_LATENT_CHANNELS, DIT_LATENT_CHANNELS, 1, bias=False) + self.proj_in = nn.Linear(_TRANSFORMER_IN_DIM, _DIM, bias=False) + self.proj_out = nn.Linear(_DIM, DIT_LATENT_CHANNELS, bias=False) + self.time_proj = FourierFeatures(1, _FOURIER_DIM) + self.time_embed = TimestepEmbedding(_FOURIER_DIM, _DIM) + self.transformer_blocks = nn.ModuleList( + TransformerBlock( + _DIM, + head_dim=_HEAD_DIM, + inner_dim=_FF_INNER_DIM, + prefix=f"transformer_blocks.{index}.attn", + ) + for index in range(_NUM_LAYERS) + ) + inv_freq = 1.0 / (_ROTARY_BASE ** (torch.arange(0, _ROTARY_DIM, 2).float() / _ROTARY_DIM)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self._rope_cache: dict[tuple[int, torch.dtype, torch.device], tuple[Tensor, Tensor]] = {} + + @property + def attention_backend_name(self) -> str: + """The attention path the blocks actually run, for startup logging.""" + return self.transformer_blocks[0].attn.backend_name + + def _rope(self, seq_len: int, *, dtype: torch.dtype, device: torch.device) -> tuple[Tensor, Tensor]: + """Return cached ``(cos, sin)`` shaped ``[seq_len, 1, 32]``. + + The head axis is left as a singleton so the tables broadcast over + ``[B, S, H, D]`` without transposing Q and K. + """ + key = (seq_len, dtype, device) + cached = self._rope_cache.get(key) + if cached is None: + t = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32) + freqs = torch.einsum("i,j->ij", t, self.inv_freq) + freqs = torch.cat((freqs, freqs), dim=-1).to(dtype=dtype, device=device).unsqueeze(1) + cached = (freqs.cos(), freqs.sin()) + self._rope_cache[key] = cached + return cached + + def forward(self, x: Tensor, t: Tensor, condition: Tensor) -> Tensor: + """Predict the flow velocity at latent ``x`` and time ``t``. + + Args: + x: Current latent, ``[B, 128, T_mel]``. + t: Flow time in ``[0, 1)``, ``[B]``. + condition: Aligned condition, ``[B, 2048, T_mel]``. + + Returns: + The velocity, ``[B, 128, T_mel]``. + """ + full = torch.cat((x, torch.zeros_like(x), condition), dim=1) + full = self.preprocess_conv(full) + full + timestep_embed = self.time_embed(self.time_proj(t[:, None])) + h = self.proj_in(full.transpose(1, 2)) + # The timestep rides as an extra leading position and is dropped again + # after the stack, so it shifts every rotary phase by one step. + h = torch.cat((timestep_embed.unsqueeze(1), h), dim=1) + rope_cos, rope_sin = self._rope(h.shape[1], dtype=h.dtype, device=h.device) + for block in self.transformer_blocks: + h = block(h, rope_cos, rope_sin) + out = self.proj_out(h[:, 1:]).transpose(1, 2) + return self.postprocess_conv(out) + out + + +def remap_transformer_state(state: dict[str, Tensor]) -> dict[str, Tensor]: + """Rewrite ``transformer/`` checkpoint keys onto this module tree. + + Two rewrites, both confined to the attention block: + + * ``attn.to_q/to_k/to_v.weight`` are concatenated along the output axis + into ``attn.to_qkv.weight``. ``nn.Linear`` output feature ``i`` reads + weight row ``i``, and :meth:`Attention.forward` splits the projection + with ``chunk(3, dim=-1)`` into ``q, k, v``, so query rows come first and + value rows last. + * ``attn.to_out.0.weight`` loses its index: the checkpoint models + ``to_out`` as a two-element list whose second element is a parameterless + dropout, and inference only needs the projection. + + Every other key is copied through untouched, including the fused ``ff_in`` + whose halves are consumed as ``value, gate`` in that order. + + Raises: + ValueError: If a block is missing one of its three projections. + """ + remapped: dict[str, Tensor] = {} + for key, value in state.items(): + if key.endswith(".attn.to_out.0.weight"): + remapped[key.replace(".to_out.0.weight", ".to_out.weight")] = value + continue + if ".attn.to_q.weight" in key: + prefix = key[: -len("to_q.weight")] + try: + fused = [state[f"{prefix}to_{name}.weight"] for name in "qkv"] + except KeyError as exc: + raise ValueError( + f"MiniMax Music 3 transformer block {prefix!r} is missing attention projection {exc}" + ) from exc + remapped[f"{prefix}to_qkv.weight"] = torch.cat(fused, dim=0) + continue + if ".attn.to_k.weight" in key or ".attn.to_v.weight" in key: + continue + remapped[key] = value + return remapped + + +class MiniMaxMusic3ConditionEncoder(nn.Module): + """Collapse an AR frame's eight layer slices and resample onto mel time.""" + + def __init__(self) -> None: + super().__init__() + self.layer_weight_logits = nn.Parameter(torch.zeros(NUM_CODEBOOKS)) + self.layer_scale = nn.Parameter(torch.ones(1)) + self.proj = nn.Conv1d(BACKBONE_HIDDEN_SIZE, _CONDITION_DIM, kernel_size=3, padding=1) + + def aligned_mel_length(self, frames: int) -> int: + """Return how many vocoder latent frames ``frames`` AR frames cover. + + AR frames sit on a 24 kHz / 960-hop grid (25 fps) and the vocoder + latent on a 44.1 kHz / 512-hop grid (about 86 fps), so one AR frame is + roughly 3.445 latent frames. The result is truncated, never rounded. + """ + return max( + 1, + int( + frames + * CONDITION_OUTPUT_SAMPLE_RATE + / CONDITION_INPUT_SAMPLE_RATE + * CONDITION_INPUT_HOP + / CONDITION_OUTPUT_HOP + ), + ) + + def condition(self, hidden: Tensor) -> Tensor: + """Project ``[B, T, 32768]`` AR frames to ``[B, 2048, T]``. + + Raises: + ValueError: If the conditioning is not ``[B, T, 32768]``. + """ + if hidden.ndim != 3 or hidden.shape[-1] != AR_HIDDEN_SIZE: + raise ValueError(f"MiniMax Music 3 condition must be [B,T,{AR_HIDDEN_SIZE}], got {tuple(hidden.shape)}") + hidden_cf = hidden.transpose(1, 2) + bsz, _, frames = hidden_cf.shape + hidden_cf = hidden_cf.reshape(bsz, NUM_CODEBOOKS, BACKBONE_HIDDEN_SIZE, frames) + weights = torch.softmax(self.layer_weight_logits, dim=0).to(hidden_cf.dtype) + hidden_cf = torch.einsum("blht,l->bht", hidden_cf, weights) + hidden_cf = self.layer_scale.to(hidden_cf.dtype) * hidden_cf + return self.proj(hidden_cf) + + def aligned_condition(self, hidden: Tensor) -> Tensor: + """Project AR hidden states and resample them onto the latent grid.""" + hidden = hidden.to(dtype=next(self.parameters()).dtype) + align = self.condition(hidden) + mel_len = self.aligned_mel_length(hidden.shape[1]) + return F.interpolate(align, size=mel_len, mode="nearest") + + +class MiniMaxMusic3FlowMatchingDiT(nn.Module): + """Condition encoder plus velocity field, with the Euler solver on top.""" + + def __init__(self) -> None: + super().__init__() + self.condition_encoder = MiniMaxMusic3ConditionEncoder() + self.transformer = MiniMaxMusic3Transformer1DModel() + + def aligned_mel_length(self, frames: int) -> int: + return self.condition_encoder.aligned_mel_length(frames) + + def condition(self, hidden: Tensor) -> Tensor: + return self.condition_encoder.condition(hidden) + + def aligned_condition(self, hidden: Tensor) -> Tensor: + return self.condition_encoder.aligned_condition(hidden) + + @torch.inference_mode() + def forward( + self, + align: Tensor, + *, + generator: torch.Generator, + initial_latent: Tensor | None = None, + initial_condition: Tensor | None = None, + num_steps: int = DEFAULT_DIT_STEPS, + cfg_scale: float = DEFAULT_DIT_CFG_SCALE, + ) -> Tensor: + """Solve an aligned condition into a vocoder latent ``[1, 128, T_mel]``. + + The previous window's latent is re-imposed at every step rather than + only at the start: the leading ``left`` frames are pinned to the exact + noise-to-latent interpolation the solver would have produced had it + generated them, so the two windows join without a seam. + + Args: + align: Aligned condition, ``[1, 2048, T_mel]``. Overwritten in + place over the prompt span when ``initial_condition`` is + given, which the caller relies on when it saves the condition + for the next window. + generator: Seeded generator for the initial noise draw. + initial_latent: Previous window's tail latent, or ``None``. + initial_condition: Previous window's tail condition, or ``None``. + num_steps: Euler steps. + cfg_scale: Classifier-free guidance weight. + + Returns: + The solved latent, ``[1, 128, T_mel]``. + + Raises: + ValueError: If ``num_steps`` is not positive. + """ + if num_steps < 1: + raise ValueError("MiniMax Music 3 DiT num_steps must be positive") + mel_len = align.shape[-1] + x = torch.randn( + (align.shape[0], DIT_LATENT_CHANNELS, mel_len), + device=align.device, + dtype=align.dtype, + generator=generator, + ) + + left = 0 + latent_prompt: Tensor | None = None + noise_prompt: Tensor | None = None + if initial_latent is not None: + left = min(initial_latent.shape[-1], mel_len) + if left <= 0: + left = 0 + else: + latent_prompt = initial_latent[..., :left].to(x) + noise_prompt = x[..., :left].clone() + if initial_condition is not None: + align[..., :left] = initial_condition[..., :left].to(align) + + dt = 1.0 / num_steps + # Row 0 carries the condition, row 1 stays zero: one batched forward + # evaluates both guidance branches. + cond_cfg = torch.zeros((2, *align.shape[1:]), device=align.device, dtype=align.dtype) + cond_cfg[0].copy_(align[0]) + # Every timestep in one host-to-device copy. Built from the same + # Python divisions the per-step construction would have used, so the + # values are identical, but the solver loop then touches the host only + # to launch kernels. + schedule = torch.tensor( + [step / num_steps for step in range(num_steps)], + device=align.device, + dtype=align.dtype, + ) + + for step in range(num_steps): + t = schedule[step].expand(x.shape[0]) + if left and latent_prompt is not None and noise_prompt is not None: + x[..., :left] = (1.0 - (1.0 - _SIGMA_FLOOR) * t[0]) * noise_prompt + t[0] * latent_prompt + d = self.transformer(x.expand(2, -1, -1), t.expand(2), cond_cfg) + d = cfg_scale * d[:1] + (1.0 - cfg_scale) * d[1:2] + x = x + dt * d + + if left and latent_prompt is not None: + x[..., :left] = latent_prompt + return x diff --git a/vllm_omni/model_executor/models/minimax_music3/frame_state.py b/vllm_omni/model_executor/models/minimax_music3/frame_state.py new file mode 100644 index 000000000..5191cebac --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/frame_state.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-request AR state for MiniMax Music 3. + +The talker accumulates one 32768-dimensional conditioning frame per decode +step and hands the acoustic stage overlapping 200-frame windows. Windows are +emitted one frame late on purpose: a window can only be declared non-final +once a further frame proves the song did not end at its boundary, and the +tail is flushed when the request finishes. + +State is keyed by request id rather than batch row, because a row's identity +is not stable across engine steps. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from .chunking import ChunkWindow +from .constants import AR_CHUNK_FRAMES, AR_CHUNK_HOP_FRAMES, NUM_CODEBOOKS + + +class FrameBuffer: + """Rolling buffer of conditioning frames addressed by absolute frame index. + + Frames before an emitted window's hop point are dropped, so peak memory is + bounded by the window size rather than by the length of the song. + """ + + def __init__(self) -> None: + self.base_frame = 0 + self._frames: list[torch.Tensor] = [] + self.peak_frames = 0 + + @property + def end_frame(self) -> int: + return self.base_frame + len(self._frames) + + def append(self, frame: torch.Tensor) -> None: + """Append one ``[1, hidden]`` frame.""" + self._frames.append(frame) + self.peak_frames = max(self.peak_frames, len(self._frames)) + + def slice(self, start: int, end: int) -> torch.Tensor: + """Return frames ``[start, end)`` stacked as ``[end - start, hidden]``. + + Raises: + RuntimeError: If the range has already been discarded or has not + been produced yet. + """ + rel_start = start - self.base_frame + rel_end = end - self.base_frame + if rel_start < 0 or rel_end > len(self._frames): + raise RuntimeError( + f"MiniMax Music 3 frame range [{start}, {end}) is unavailable; " + f"the buffer holds [{self.base_frame}, {self.end_frame})" + ) + return torch.cat(self._frames[rel_start:rel_end], dim=0) + + def discard_before(self, frame: int) -> None: + drop = min(max(frame - self.base_frame, 0), len(self._frames)) + if drop: + del self._frames[:drop] + self.base_frame += drop + + def clear(self) -> None: + self._frames.clear() + self.base_frame = 0 + + +@dataclass +class RequestARState: + """Everything the talker tracks for one conditioned request.""" + + seed: int + max_frames: int + # Prompt length in tokens, used to tell a decode row apart from a row that + # is still working through a chunked prefill. + prompt_tokens: int = 0 + sampling_position: int = 0 + generated_frames: int = 0 + emitted_windows: int = 0 + finished: bool = False + # The engine is told to stop one step after the song ends. Payloads are + # built from the buffer that ``postprocess`` filled on the previous step, + # so without that extra step the terminal span would be written and then + # discarded with the request. + pending_stop: bool = False + finish_reason: str = "length" + frames: FrameBuffer = field(default_factory=FrameBuffer) + codes: list[torch.Tensor] = field(default_factory=list) + codes_base_frame: int = 0 + # The previous frame's eight codes, fed back as the next step's input + # embedding. Held on the request rather than on a batch row, because the + # row a request occupies is not stable between steps. + last_codes: torch.Tensor | None = None + + def advance_position(self) -> int: + """Return this frame's base draw position and reserve the frame's slot. + + Each frame consumes ``NUM_CODEBOOKS`` positions so that the ``c0`` draw + and the seven depth draws never share an RNG stream. + """ + position = self.sampling_position + self.sampling_position += NUM_CODEBOOKS + return position + + def ready_window(self) -> ChunkWindow | None: + """Return the next window that a lookahead frame has proven non-final.""" + start = self.emitted_windows * AR_CHUNK_HOP_FRAMES + if self.frames.end_frame <= start + AR_CHUNK_FRAMES: + return None + return ChunkWindow( + index=self.emitted_windows, + start=start, + end=start + AR_CHUNK_FRAMES, + is_first=self.emitted_windows == 0, + is_last=False, + ) + + def take_window(self, window: ChunkWindow, *, is_final: bool) -> torch.Tensor: + """Materialize ``window`` and release the frames it retires.""" + chunk = self.frames.slice(window.start, window.end) + self.emitted_windows = max(self.emitted_windows, window.index + 1) + if not is_final: + self.frames.discard_before(window.start + AR_CHUNK_HOP_FRAMES) + return chunk + + def take_codes(self, start: int, end: int) -> torch.Tensor | None: + """Return the sampled codes for frames ``[start, end)``, if retained. + + Codes are carried alongside the conditioning frames so a downstream + stage or a debugging harness can recover the exact RVQ trajectory. + They are dropped rather than raising when a span predates what is + still buffered, since nothing in the audio path depends on them. + """ + if not self.codes or end <= start: + return None + offset = self.codes_base_frame + lo, hi = start - offset, end - offset + if lo < 0 or hi > len(self.codes): + return None + return torch.cat(self.codes[lo:hi], dim=0) + + def discard_codes_before(self, frame: int) -> None: + drop = min(max(frame - self.codes_base_frame, 0), len(self.codes)) + if drop: + del self.codes[:drop] + self.codes_base_frame += drop + + def release(self) -> None: + self.frames.clear() + self.codes.clear() + + +__all__ = ["FrameBuffer", "RequestARState"] diff --git a/vllm_omni/model_executor/models/minimax_music3/pipeline.py b/vllm_omni/model_executor/models/minimax_music3/pipeline.py new file mode 100644 index 000000000..1a8279f05 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/pipeline.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax Music 3 pipeline topology: AR talker -> acoustic decoder. + +Stage 0 ``minimax_music3_ar``: a Qwen3 backbone (the repo's ``language_model/`` +subfolder is a plain ``Qwen3ForCausalLM``) plus a depth decoder, wrapped by a +class that owns its sampler. It emits 200-frame windows of frame hidden states +and the matching 8-codebook RVQ codes. + +Stage 1 ``minimax_music3_acoustic``: flow-matching DiT plus DAV vocoder. Pure +compute, no KV cache, weights loaded from the repo's ``condition_encoder/``, +``transformer/`` and ``vocoder/`` subfolders. Output is 32 kHz stereo. + +Classifier-free guidance is mandatory for this checkpoint: every request +decodes as a cond/uncond pair, and the uncond row is produced by +``prompt_expand_func``. The deploy YAMLs therefore carry +``default_sampling_params.extra_args`` with ``cfg_role`` and ``cfg_scale`` so +that ``has_sampling_extra_args`` is on for the stage and per-row extra args +reach the model's forward. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.minimax_music3" + +# ``<|audio_end|>`` in the checkpoint tokenizer (pinned by +# ``models/minimax_music3/prompt.SPECIAL_TOKEN_IDS``). The talker's own sampler +# forces this at the frame budget; the stop id is the engine-side safety net. +MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID = 151670 + +# MiniMax captions carry the full style description and run to roughly 5000 +# characters, an order of magnitude past the 500-character serving default. +# ``tts_args`` is a pipeline-level stage extra, not a deploy-YAML field: it +# reaches the serving layer through ``StageConfig.yaml_extras``, which is where +# ``ServingSpeech._compute_max_instructions_length`` reads it from. +_MAX_INSTRUCTIONS_LENGTH = 20_000 + +MINIMAX_MUSIC3_PIPELINE = PipelineConfig( + model_type="minimax_music3", + model_arch="MiniMaxMusic3TalkerForConditionalGeneration", + hf_architectures=("MiniMaxMusic3ForConditionalGeneration",), + default_deploy_config_name="minimax_music3.yaml", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="minimax_music3_ar", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + # The repo root config.json only names the composite architecture; + # the AR weights and the tokenizer live in their own subfolders. + model_subdir="language_model", + tokenizer_subdir="tokenizer", + # detokenize: the sampled stream is audio codes, never text. + # stop_token_ids: <|audio_end|> closes the song. + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID], + }, + # CFG is mandatory here, so this expands on essentially every + # request rather than only on an opt-in flag. + prompt_expand_func=f"{_PROC}.expand_cfg_prompts", + async_chunk_process_next_stage_input_func=f"{_PROC}.ar2acoustic_async_chunk", + # retains_state_across_chunks stays False. It is read only by + # OmniGenerationScheduler, so it is a no-op on an LLM_AR stage; + # more to the point, stage 0 produces chunks rather than parking on + # them, so it never holds a runner slot while waiting for one. + extras={"tts_args": {"max_instructions_length": _MAX_INSTRUCTIONS_LENGTH}}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="minimax_music3_acoustic", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + # "audio" (not "latent"): the accumulator concatenates audio over + # time, whereas the latent path would stack the two stereo channels. + engine_output_type="audio", + model_arch="MiniMaxMusic3AcousticForConditionalGeneration", + # This stage reads the repo ROOT so it can reach condition_encoder/, + # transformer/ and vocoder/, and the root ships no tokenizer files. + tokenizer_subdir="tokenizer", + # The acoustic stage carries the previous window's latent forward as + # the next window's flow-matching prompt, so a request parked between + # chunks still owns its model state and must keep counting against + # max_num_seqs. + retains_state_across_chunks=True, + sync_process_input_func=f"{_PROC}.ar2acoustic", + sampling_constraints={"detokenize": True}, + ), + ), +) + +__all__ = [ + "MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID", + "MINIMAX_MUSIC3_PIPELINE", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/prompt.py b/vllm_omni/model_executor/models/minimax_music3/prompt.py new file mode 100644 index 000000000..25e18fd7e --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/prompt.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# +# Prompt normalization follows the MiniMaxAI/MiniMax-Music3 input contract. +# The transformations are reproduced byte-for-byte because any divergence +# shifts token positions and therefore changes the generated audio. +"""MiniMax Music 3 prompt construction. + +The prompt is a single flat string, not a chat template:: + + <|im_start|><|caption_start|>{caption}<|caption_end|> + <|lyrics_start|>{lyrics}<|lyrics_end|><|im_end|><|audio_start|> + +The classifier-free-guidance twin reuses the same token ids with the whole +caption-and-lyrics span (delimiters included) overwritten by +``<|audio_cfg|>``, so only the chat markers and ``<|audio_start|>`` survive. +Both rows therefore have identical length and identical audio-start position. +""" + +from __future__ import annotations + +from typing import Any + +import regex as re + +# Pinned by ``validate_tokenizer_ids``: a checkpoint whose tokenizer disagrees +# is not this model, and silently producing noise is worse than failing. +SPECIAL_TOKEN_IDS: dict[str, int] = { + "<|im_start|>": 151644, + "<|im_end|>": 151645, + "<|audio_cfg|>": 151654, + "<|audio_start|>": 151669, + "<|audio_end|>": 151670, + "<|caption_start|>": 151671, + "<|caption_end|>": 151672, + "<|lyrics_start|>": 151673, + "<|lyrics_end|>": 151674, +} + +# c0 codes occupy the backbone vocabulary immediately above the text tokens. +AUDIO_CODE_OFFSET = 151675 + +_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>") +_LEADING_TAGS_RE = re.compile(r"^[ \t]*((?:\[[^\]]+\][ \t]*)+)") + + +def _strip_markdown_residuals(text: str) -> str: + """Strip the Markdown syntax the model input contract does not accept.""" + if not text: + return text + lines_out: list[str] = [] + for raw_line in text.splitlines(): + line = raw_line + line = re.sub(r"^\s{0,3}#{1,6}\s+", "", line) + line = re.sub(r"^\s*[*+-]\s+", "", line) + line = re.sub(r"^\s*\*\s+", "", line) + while "**" in line: + updated = re.sub(r"\*\*([^*]+)\*\*", r"\1", line) + if updated == line: + break + line = updated + line = re.sub(r"(? str: + return _strip_markdown_residuals(text).replace("• ", "").replace(" ", "") + + +def clean_caption(caption: str) -> str: + """Apply the source order: special tokens, then Markdown, then newlines. + + ``<|tag value|>`` forms are rewritten to ``tag is value`` so a caption + cannot inject a real special token into the prompt. + """ + + def _special(match: re.Match[str]) -> str: + inner = match.group(1).strip() + parts = inner.split(None, 1) + return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else inner + + text = _SPECIAL_TAG_RE.sub(_special, caption) + text = _remove_markdown_format(text) + # The source collapses runs of two or more newlines, not three or more; + # this shifts blank-line token positions and therefore the audio. + return re.sub(r"\n{2,}", "\n", text) + + +def _strip_text_after_leading_tags(text: str) -> str: + """Keep only the consecutive structural tags that start a line. + + Lyrics written on the same line as a ``[Verse]`` tag are dropped, which is + the documented contract: a tag must sit on its own line. + """ + if not text: + return text + output: list[str] = [] + for line in text.split("\n"): + match = _LEADING_TAGS_RE.match(line) + output.append(match.group(1).strip() if match else line) + return "\n".join(output) + + +def normalize_lyrics(lyrics: str) -> str: + """Normalize lyrics to the form the checkpoint was trained on.""" + text = _strip_text_after_leading_tags(lyrics) + text = text.replace("] ", "]\n") + text = text.replace(" [", "\n[") + text = text.replace(" ^ ", "\n") + text = re.sub(r"\[([^\]]+)\]", lambda match: f"[{match.group(1).lower()}]", text) + return f"[start]\n{text}" + + +def build_prompt(caption: str, lyrics: str) -> str: + """Build the conditioned prompt string.""" + return ( + "<|im_start|><|caption_start|>" + f"{clean_caption(caption)}" + "<|caption_end|><|lyrics_start|>" + f"{normalize_lyrics(lyrics)}" + "<|lyrics_end|><|im_end|><|audio_start|>" + ) + + +def build_cfg_null_token_ids(prompt_token_ids: list[int]) -> list[int]: + """Return the unconditioned twin's token ids for a conditioned prompt. + + Everything between the leading ``<|im_start|>`` and the trailing + ``<|im_end|><|audio_start|>`` becomes ``<|audio_cfg|>``. Length and the + audio-start position are preserved, which is what lets the two rows decode + in lockstep. + + Raises: + ValueError: If the prompt is too short to carry a caption and lyrics, + or does not have the expected framing tokens. + """ + if len(prompt_token_ids) < 4: + raise ValueError(f"MiniMax Music 3 prompt is too short to build a CFG twin: {len(prompt_token_ids)} tokens") + expected = ( + SPECIAL_TOKEN_IDS["<|im_start|>"], + SPECIAL_TOKEN_IDS["<|im_end|>"], + SPECIAL_TOKEN_IDS["<|audio_start|>"], + ) + actual = (prompt_token_ids[0], prompt_token_ids[-2], prompt_token_ids[-1]) + if actual != expected: + raise ValueError( + "MiniMax Music 3 prompt framing is wrong: expected " + f"(im_start, im_end, audio_start)={expected}, got {actual}" + ) + null_ids = list(prompt_token_ids) + null_ids[1:-2] = [SPECIAL_TOKEN_IDS["<|audio_cfg|>"]] * (len(null_ids) - 3) + return null_ids + + +def validate_tokenizer_ids(tokenizer: Any) -> None: + """Fail fast when a tokenizer is not the supported music tokenizer.""" + for token, expected in SPECIAL_TOKEN_IDS.items(): + token_id = tokenizer.convert_tokens_to_ids(token) + if token_id != expected: + raise ValueError(f"MiniMax Music 3 tokenizer mismatch for {token}: expected {expected}, got {token_id}") + + +__all__ = [ + "AUDIO_CODE_OFFSET", + "SPECIAL_TOKEN_IDS", + "build_cfg_null_token_ids", + "build_prompt", + "clean_caption", + "normalize_lyrics", + "validate_tokenizer_ids", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py b/vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py new file mode 100644 index 000000000..64acc9b62 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RVQ depth decoder for MiniMax Music 3 (codebooks c1..c7). + +Each audio frame is eight codebooks deep. The backbone's own ``lm_head`` +emits ``c0``; this four-layer decoder walks ``c1..c7``, conditioning each code +on the previous ones. Its per-codebook hidden states are also the acoustic +stage's conditioning signal, so this module returns them alongside the codes. + +The sequence fed to the decoder is:: + + [proj(backbone_hidden), proj(embed(c0)), proj(embed(c1)), ...] + +and after step ``i`` the head ``audio_heads[i - 1]`` predicts ``c_i``. The +whole thing runs over at most eight positions, so it is written as a dense +causal block rather than anything with a KV cache. +""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn +from torch.nn import functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from .constants import AUDIO_VOCAB_SIZE, BACKBONE_HIDDEN_SIZE, NUM_CODEBOOKS + +# cuDNN's fused attention is not bit-reproducible on this hardware: two +# identical replays of the same input can differ by one bfloat16 ULP, which is +# enough to flip a sampled code and send the rest of the song somewhere else. +# The choice is pinned per call rather than through a process-global switch so +# it cannot be undone by whatever else shares the process. +_SDPA_BACKENDS = [SDPBackend.FLASH_ATTENTION, SDPBackend.MATH] + + +class DepthDecoderBlock(nn.Module): + """Pre-norm transformer block with fused QKV and a SwiGLU MLP.""" + + def __init__(self, hidden_size: int, num_heads: int, intermediate_size: int) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.input_layernorm = nn.RMSNorm(hidden_size, eps=1e-6) + # Stored fused so the three projections are one matmul; the checkpoint + # ships them split and ``load_checkpoint_state`` concatenates. + self.qkv_proj = nn.Linear(hidden_size, hidden_size * 3, bias=False) + self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.post_attention_layernorm = nn.RMSNorm(hidden_size, eps=1e-6) + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x: Tensor) -> Tensor: + h = self.input_layernorm(x) + bsz, seq, hidden = h.shape + q, k, v = self.qkv_proj(h).chunk(3, dim=-1) + q = q.reshape(bsz, seq, self.num_heads, self.head_dim).transpose(1, 2) + k = k.reshape(bsz, seq, self.num_heads, self.head_dim).transpose(1, 2) + v = v.reshape(bsz, seq, self.num_heads, self.head_dim).transpose(1, 2) + with sdpa_kernel(_SDPA_BACKENDS): + h = F.scaled_dot_product_attention(q, k, v, is_causal=True) + h = h.transpose(1, 2).contiguous().reshape(bsz, seq, hidden) + x = x + self.o_proj(h) + h = self.post_attention_layernorm(x) + h = self.down_proj(F.silu(self.gate_proj(h)) * self.up_proj(h)) + return x + h + + +class RVQDepthDecoder(nn.Module): + """Depth-autoregressive decoder over the seven residual codebooks.""" + + def __init__( + self, + *, + hidden_size: int = BACKBONE_HIDDEN_SIZE, + num_layers: int = 4, + num_heads: int = 16, + intermediate_size: int = 6144, + audio_vocab_size: int = AUDIO_VOCAB_SIZE, + num_codebooks: int = NUM_CODEBOOKS, + max_seq_len: int = 16, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.audio_vocab_size = audio_vocab_size + self.num_codebooks = num_codebooks + self.projection = nn.Linear(hidden_size, hidden_size, bias=False) + self.pos_embedding = nn.Embedding(max_seq_len, hidden_size) + self.layers = nn.ModuleList( + DepthDecoderBlock(hidden_size, num_heads, intermediate_size) for _ in range(num_layers) + ) + self.norm = nn.RMSNorm(hidden_size, eps=1e-6) + self.audio_heads = nn.ModuleList( + nn.Linear(hidden_size, audio_vocab_size, bias=False) for _ in range(num_codebooks - 1) + ) + + def forward(self, seq: Tensor) -> Tensor: + positions = torch.arange(seq.shape[1], device=seq.device) + x = seq + self.pos_embedding(positions).unsqueeze(0) + for layer in self.layers: + x = layer(x) + return self.norm(x) + + def load_checkpoint_state(self, state: dict[str, Tensor]) -> None: + """Load the component checkpoint, failing loudly on any mismatch. + + The checkpoint keeps ``q``/``k``/``v`` split and names attention + projections ``attn.to_*``; both are reshaped here. A silent partial + load would produce plausible-sounding but wrong audio, so every key is + accounted for and anything left over is an error. + + Args: + state: Tensors from the ``rvq_depth_decoder`` component, with the + ``audio_embeddings.weight`` entry already removed by the + caller. + + Raises: + ValueError: If any expected key is missing, any shape disagrees, + or the checkpoint carries keys this module does not consume. + """ + loaded: set[str] = set() + missing: list[str] = [] + mismatched: list[str] = [] + + def take(key: str) -> Tensor | None: + value = state.get(key) + if value is None: + missing.append(key) + return value + + def assign(key: str, target: Tensor, source: Tensor, keys: tuple[str, ...]) -> None: + if tuple(source.shape) != tuple(target.shape): + mismatched.append(f"{key}: checkpoint={tuple(source.shape)} model={tuple(target.shape)}") + return + target.data.copy_(source.to(dtype=target.dtype)) + loaded.update(keys) + + def copy_one(key: str, target: Tensor) -> None: + source = take(key) + if source is not None: + assign(key, target, source, (key,)) + + copy_one("projection.weight", self.projection.weight) + copy_one("pos_embedding.weight", self.pos_embedding.weight) + copy_one("norm.weight", self.norm.weight) + for index, head in enumerate(self.audio_heads): + copy_one(f"audio_heads.{index}.weight", head.weight) + + for index, layer in enumerate(self.layers): + prefix = f"layers.{index}." + copy_one(prefix + "input_layernorm.weight", layer.input_layernorm.weight) + copy_one( + prefix + "post_attention_layernorm.weight", + layer.post_attention_layernorm.weight, + ) + qkv_keys = tuple(prefix + f"attn.to_{name}.weight" for name in ("q", "k", "v")) + parts = [take(key) for key in qkv_keys] + if all(part is not None for part in parts): + assign( + prefix + "attn.qkv", + layer.qkv_proj.weight, + torch.cat(parts, dim=0), + qkv_keys, + ) + copy_one(prefix + "attn.to_out.weight", layer.o_proj.weight) + copy_one(prefix + "gate_proj.weight", layer.gate_proj.weight) + copy_one(prefix + "up_proj.weight", layer.up_proj.weight) + copy_one(prefix + "down_proj.weight", layer.down_proj.weight) + + unexpected = sorted(set(state) - loaded) + if missing or mismatched or unexpected: + raise ValueError( + "MiniMax Music 3 RVQ depth decoder failed to load: " + f"missing={missing[:8]}, mismatched={mismatched[:8]}, " + f"unexpected={unexpected[:8]}" + ) + + +__all__ = ["DepthDecoderBlock", "RVQDepthDecoder"] diff --git a/vllm_omni/model_executor/models/minimax_music3/sampling.py b/vllm_omni/model_executor/models/minimax_music3/sampling.py new file mode 100644 index 000000000..226addff5 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/sampling.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deterministic guided sampling for MiniMax Music 3. + +Two things happen per code, in this order: + +1. **Classifier-free guidance.** Every request decodes twice, once on the real + prompt and once on a prompt whose caption and lyrics have been replaced by + ``<|audio_cfg|>``. Codes are drawn from ``uncond + (cond - uncond) * 1.5``. + For ``c0`` the candidate set is additionally masked to the conditioned + row's own top 50, so the guided draw cannot wander outside what the + conditioned branch already considered plausible. +2. **Seeded top-k.** A batched, order-independent categorical draw: each row + gets its own RNG stream keyed by ``(seed, position)``, so a request's audio + depends only on its own seed and never on who else happens to share the + batch. + +The RNG is a MurmurHash3 x86_32 of ``(seed_lo, seed_hi, position, column)`` +turned into Gumbel noise, and the draw is an argmax over ``logits + gumbel``. +That is mathematically a categorical sample from ``softmax(logits)``, but +unlike ``torch.multinomial`` it needs no per-row generator state, so it is +CUDA-graph safe and reproducible across batch compositions. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + +from .constants import AR_CFG_SCALE, AR_CFG_TOP_K, AR_TOP_K + +_UINT32_MASK = 0xFFFFFFFF +_UINT32_MAX = 4_294_967_295 + +# MurmurHash3 x86_32 constants. +_C1 = 0xCC9E2D51 +_C2 = 0x1B873593 +_R1 = 15 +_R2 = 13 +_M = 5 +_N = 0xE6546B64 + + +def _rotl32(x: Tensor, r: int) -> Tensor: + return ((x << r) | (x >> (32 - r))) & _UINT32_MASK + + +def _murmur3_mix(h: Tensor, k: Tensor) -> Tensor: + """Mix one 32-bit key into the running hash state.""" + k = (k * _C1) & _UINT32_MASK + k = _rotl32(k, _R1) + k = (k * _C2) & _UINT32_MASK + h = h ^ k + h = _rotl32(h, _R2) + return (h * _M + _N) & _UINT32_MASK + + +def _fmix32(h: Tensor) -> Tensor: + """MurmurHash3 finalizer.""" + h = h ^ (h >> 16) + h = (h * 0x85EBCA6B) & _UINT32_MASK + h = h ^ (h >> 13) + h = (h * 0xC2B2AE35) & _UINT32_MASK + return h ^ (h >> 16) + + +# Memoized ``k`` term of the column block's mix. The column block is the same +# ``arange(width)`` on every draw, so its three ops are hoisted out of the per +# draw work and kept per (device, width). Small and bounded: two widths exist, +# the c0 candidate set and the residual codebook. +_COLUMN_MIX_CACHE: dict[tuple[torch.device, int], Tensor] = {} + + +def _column_mix(num_columns: int, device: torch.device) -> Tensor: + """Return the mixed column block, ``[cols]`` int64 holding uint32.""" + key = (device, num_columns) + cached = _COLUMN_MIX_CACHE.get(key) + if cached is None: + columns = torch.arange(num_columns, dtype=torch.int64, device=device) & _UINT32_MASK + k = (columns * _C1) & _UINT32_MASK + k = _rotl32(k, _R1) + cached = (k * _C2) & _UINT32_MASK + _COLUMN_MIX_CACHE[key] = cached + return cached + + +def _absorb_mixed_key(h: Tensor, k: Tensor) -> Tensor: + """Absorb a key whose own three-op transform is already applied.""" + h = h ^ k + h = _rotl32(h, _R2) + return (h * _M + _N) & _UINT32_MASK + + +def row_hash_state(seeds: Tensor, positions: Tensor) -> Tensor: + """Hash state after the seed and position blocks, before the columns. + + Args: + seeds: ``[rows]`` per-request seeds. + positions: ``[rows]`` or ``[rows, draws]`` draw identifiers. The extra + dimension lets every codebook of a frame be hashed in one batch + instead of one dispatch per codebook; the arithmetic is elementwise + integer, so a batched call is bit-identical to the loop it replaces. + + Returns: + A tensor shaped like ``positions``, int64 holding uint32. + """ + seeds = seeds.to(torch.int64) + while seeds.dim() < positions.dim(): + seeds = seeds.unsqueeze(-1) + h = torch.zeros((), dtype=torch.int64, device=seeds.device) + h = _murmur3_mix(h, seeds & _UINT32_MASK) + h = _murmur3_mix(h, (seeds >> 32) & _UINT32_MASK) + return _murmur3_mix(h, positions.to(torch.int64) & _UINT32_MASK) + + +def murmur_hash32(seeds: Tensor, positions: Tensor, columns: Tensor) -> Tensor: + """Hash ``(seed, position, column)`` to a uint32 held in int64. + + The 64-bit seed is consumed as two 32-bit blocks, then the position and + the column, for a hashed length of 16 bytes. + + Args: + seeds: ``[rows]`` per-request seeds. + positions: ``[rows]`` draw identifiers, unique per (frame, codebook). + columns: ``[cols]`` candidate indices, which must be ``arange(cols)``. + + Returns: + ``[rows, cols]`` int64 tensor holding uint32 values. + """ + h = row_hash_state(seeds, positions).unsqueeze(-1) + return _fmix32(_absorb_mixed_key(h, _column_mix(int(columns.shape[0]), h.device)) ^ 16) + + +def gumbel_noise(seeds: Tensor, positions: Tensor, num_columns: int) -> Tensor: + """Return the Gumbel noise for one draw per position, ``[..., cols]``. + + Args: + seeds: ``[rows]`` per-request seeds. + positions: ``[rows]`` or ``[rows, draws]`` draw identifiers. + num_columns: Candidate count, i.e. the width of the logits. + + Returns: + float64 noise shaped ``positions.shape + (num_columns,)``. + """ + h = row_hash_state(seeds, positions).unsqueeze(-1) + hashed = _fmix32(_absorb_mixed_key(h, _column_mix(num_columns, h.device)) ^ 16) + # float64 throughout: the double log is numerically unstable in float32 and + # a single flipped draw changes every frame that follows it. + uniform = hashed.to(torch.float64) / _UINT32_MAX + # Clamp both ends. uniform == 1 gives +inf Gumbel, which turns a -inf + # masked entry into NaN; the upper cap is the hash spacing, so a saturated + # bucket ties with its neighbour instead of always winning. + gumbel = uniform.log().clamp_(min=torch.finfo(torch.float64).min, max=-(2.0**-32)).neg_() + return gumbel.log().neg_() + + +def draw_from_gumbel(logits: Tensor, gumbel: Tensor, *, top_k: int = AR_TOP_K) -> Tensor: + """Argmax of ``top_k``-masked logits plus precomputed Gumbel noise. + + Args: + logits: ``[rows, vocab]`` unnormalized scores. + gumbel: ``[rows, vocab]`` float64 noise from :func:`gumbel_noise`. + top_k: Candidates retained per row before the draw. + + Returns: + ``[rows]`` sampled column indices. + """ + values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9) + k = min(int(top_k), values.shape[-1]) + threshold = torch.topk(values, k, dim=-1).values[..., -1, None] + values = values.masked_fill(values < threshold, -float("inf")) + return (values.to(torch.float64) + gumbel).argmax(dim=-1) + + +def sample_topk_seeded( + logits: Tensor, + seeds: Tensor, + positions: Tensor, + *, + top_k: int = AR_TOP_K, +) -> Tensor: + """Draw one column per row from its own top-k, under its own RNG stream. + + Args: + logits: ``[rows, vocab]`` unnormalized scores. + seeds: ``[rows]`` per-request sampling seeds. + positions: ``[rows]`` draw identifiers, unique per (frame, codebook). + top_k: Candidates retained per row before the draw. + + Returns: + ``[rows]`` sampled column indices. + """ + gumbel = gumbel_noise(seeds, positions, int(logits.shape[-1])) + return draw_from_gumbel(logits, gumbel, top_k=top_k) + + +def apply_c0_cfg(cond: Tensor, uncond: Tensor, *, scale: float = AR_CFG_SCALE) -> Tensor: + """Guide ``cond`` by ``uncond``, restricted to ``cond``'s own top candidates. + + The mask is a second, separate restriction from the top-k the sampler + applies afterwards. Dropping it widens the distribution the model draws + from and audibly loosens how closely the output follows the caption. + """ + guided = uncond + (cond - uncond) * scale + k = min(AR_CFG_TOP_K, cond.shape[-1]) + threshold = torch.topk(cond, k, dim=-1).values[..., -1, None] + return guided.masked_fill(cond < threshold, -float("inf")) + + +def apply_depth_cfg(cond: Tensor, uncond: Tensor, *, scale: float = AR_CFG_SCALE) -> Tensor: + """Guide the depth heads. Unlike ``c0`` these are not additionally masked.""" + return uncond + (cond - uncond) * scale + + +__all__ = [ + "apply_c0_cfg", + "apply_depth_cfg", + "draw_from_gumbel", + "gumbel_noise", + "murmur_hash32", + "row_hash_state", + "sample_topk_seeded", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/staging.py b/vllm_omni/model_executor/models/minimax_music3/staging.py new file mode 100644 index 000000000..ed3eb8a47 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/staging.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pinned host staging for the small per-step index vectors. + +Building a device tensor from a Python list (``torch.tensor(rows, +device="cuda")``) stages through pageable host memory, which makes the copy +synchronous: the host blocks until the copy retires, and because the copy is +queued behind the step's kernels that means blocking until the whole forward +has finished. On this model that turned four one-line conveniences into the +single largest entry in a CPU profile of the decode loop. + +Writing into a pinned block and pushing it with one non-blocking copy costs +one transfer for the whole step and leaves the host free to keep enqueueing. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +__all__ = ["PinnedIndexStage"] + + +class PinnedIndexStage: + """A fixed ``[rows, capacity]`` int64 block staged host-to-device. + + The device view returned by :meth:`push` is only populated once the stream + reaches the copy, so it may be consumed by kernels on that stream and by + nothing else. The host side is double-buffered: the next step writes the + slot the device is not reading, so refilling cannot corrupt a copy that is + still in flight. + """ + + def __init__(self, rows: int, capacity: int, device: torch.device) -> None: + self._capacity = capacity + self._host = [torch.zeros((rows, capacity), dtype=torch.int64, pin_memory=True) for _ in range(2)] + self._device = torch.zeros((rows, capacity), dtype=torch.int64, device=device) + self._slot = 0 + + @property + def capacity(self) -> int: + return self._capacity + + def push(self, values: Sequence[Sequence[int]], count: int) -> torch.Tensor: + """Stage ``count`` entries of each row and return the device view. + + Args: + values: One sequence of Python ints per row of the block. + count: How many entries of each row are meaningful this step. + + Returns: + ``[rows, count]`` int64 view of the device block. + + Raises: + ValueError: If ``count`` exceeds the reserved capacity. + """ + if count > self._capacity: + raise ValueError(f"staged {count} entries but only {self._capacity} were reserved") + self._slot ^= 1 + host = self._host[self._slot] + view = host.numpy() + for index, row in enumerate(values): + view[index, :count] = row + self._device.copy_(host, non_blocking=True) + return self._device[:, :count] diff --git a/vllm_omni/model_executor/models/minimax_music3/talker.py b/vllm_omni/model_executor/models/minimax_music3/talker.py new file mode 100644 index 000000000..c4fb876fa --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/talker.py @@ -0,0 +1,990 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-0 talker for MiniMax Music 3. + +A Qwen3 backbone predicts one audio frame per decode step. Each frame is +eight RVQ codebooks deep: the backbone's own ``lm_head`` emits ``c0`` and the +four-layer depth decoder walks ``c1..c7``. The conditioning signal handed to +the acoustic stage is the backbone hidden state concatenated with the seven +depth hidden states, so 8 x 4096 = 32768 dimensions per frame. + +Two things make this model unusual inside vLLM: + +* **Every request decodes as two rows.** Classifier-free guidance is not + optional here. The conditioned row sees the real prompt and the + unconditioned row sees the same prompt with the caption-and-lyrics span + replaced by ``<|audio_cfg|>``. Codes are drawn from the guided blend and + then fed back to *both* rows, so the pair stays token-identical. Rows are + paired by ``cfg_pair_id`` rather than by adjacency, because batch position + is not stable across engine steps. +* **The decode input is not a token embedding.** After prefill, each step's + input is the composed embedding of the previous frame's eight codes. The + sampled ``c0`` token id still flows through vLLM's normal machinery so that + stop handling and accounting work, but the embedding it would produce is + replaced before the backbone runs. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch +from torch import nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.v1.outputs import SamplerOutput + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .chunking import chunk_windows +from .constants import ( + AR_CHUNK_HOP_FRAMES, + AUDIO_VOCAB_SIZE, + BACKBONE_HIDDEN_SIZE, + C0_VOCAB_SIZE, + DEFAULT_MAX_AUDIO_FRAMES, + NUM_CODEBOOKS, +) +from .depth_graph import DepthGraphCache +from .frame_state import RequestARState +from .prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS +from .rvq_decoder import RVQDepthDecoder +from .sampling import ( + apply_c0_cfg, + apply_depth_cfg, + draw_from_gumbel, + gumbel_noise, + sample_topk_seeded, +) +from .staging import PinnedIndexStage +from .weights import load_depth_decoder_weights + +logger = init_logger(__name__) + +__all__ = ["MiniMaxMusic3TalkerForConditionalGeneration"] + +_COND = "cond" +_UNCOND = "uncond" +_STOP_TOKEN = SPECIAL_TOKEN_IDS["<|audio_end|>"] +# Emitted on the step a song ends, to keep the request alive for one more step +# while its terminal span is delivered. It is never fed back as audio: the +# state is already marked finished, so no frame is composed from it. +_HOLD_TOKEN = AUDIO_CODE_OFFSET + +# The unconditioned twin's request id is the conditioned one plus this suffix, +# which is what lets a pair be recognized without a serving-layer injection. +_CFG_UNCOND_SUFFIX = "__cfg_uncond" + + +@dataclass(frozen=True) +class _RowConstants: + """Per-request values that never change once the request is admitted.""" + + request_id: str + role: str + pair_id: str + seed: int + max_frames: int + + +def _request_key(value: Any) -> str: + """Normalize a request id to the form used to key per-request state. + + The engine injects ``global_request_id`` as a single-element list, while + ``on_requests_finished`` receives the bare string. Keying state on the + unwrapped value keeps allocation and release on the same key; getting this + wrong leaks a song's worth of frames per request. + """ + if isinstance(value, (list, tuple)): + return str(value[0]) if value else "" + return "" if value is None else str(value) + + +def _pair_id_from_request(request_id: str) -> str: + """Derive the CFG pair id when the request does not carry one.""" + if request_id.endswith(_CFG_UNCOND_SUFFIX): + return request_id[: -len(_CFG_UNCOND_SUFFIX)] + return request_id + + +def _role_from_request(request_id: str) -> str: + """Derive the CFG role from the companion suffix.""" + return _UNCOND if request_id.endswith(_CFG_UNCOND_SUFFIX) else _COND + + +def _first_int(value: Any, default: int) -> int: + """Read an int that the serving layer wraps in a single-element list.""" + if isinstance(value, (list, tuple)): + value = value[0] if value else None + if value is None: + return default + try: + parsed = int(value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +class MiniMaxMusic3TalkerForConditionalGeneration(nn.Module): + """Qwen3 backbone plus the audio embedding table and RVQ depth decoder.""" + + # Sampling, feedback and window emission are all model-owned. + prefer_model_sampler: bool = True + have_multimodal_outputs: bool = True + has_postprocess: bool = True + skips_model_sampler_output_token_history: bool = True + postprocess_uses_hidden_states: bool = False + postprocess_uses_multimodal_outputs: bool = False + postprocess_uses_req_infos: bool = True + supports_omni_query_start_loc: bool = True + # A row still completing a chunked prefill reaches the sampler but its + # token is discarded; drawing a frame for it would corrupt the song. + requires_request_sample_eligibility: bool = True + # The acoustic stage consumes the composed 32768-d frame, never the raw + # backbone hidden state, so suppress the runner's automatic "hidden" key. + omni_pooler_payload_include_hidden: bool = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + self.config = config + hidden_size = int(getattr(config, "hidden_size", BACKBONE_HIDDEN_SIZE)) + + self.model = Qwen3Model( + vllm_config=vllm_config, + prefix=f"{prefix}.model" if prefix else "model", + ) + if getattr(config, "tie_word_embeddings", False): + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + int(config.vocab_size), + hidden_size, + prefix=f"{prefix}.lm_head" if prefix else "lm_head", + ) + self.logits_processor = LogitsProcessor(int(config.vocab_size)) + + # c1..c7 share one table; codebook i occupies rows [i*1024, (i+1)*1024). + self.audio_embeddings = nn.Embedding(AUDIO_VOCAB_SIZE * (NUM_CODEBOOKS - 1), hidden_size) + self.rvq_decoder = RVQDepthDecoder(hidden_size=hidden_size) + + self.register_buffer( + "audio_embedding_offsets", + (torch.arange(NUM_CODEBOOKS - 1) * AUDIO_VOCAB_SIZE).unsqueeze(0), + persistent=False, + ) + # The only token ids c0 sampling can return, in vocabulary order: + # <|audio_end|> first, then the c0 code range. Column 0 therefore means + # stop and column code+1 means code. + self.register_buffer( + "c0_logit_ids", + torch.cat( + ( + torch.tensor([_STOP_TOKEN], dtype=torch.long), + torch.arange( + AUDIO_CODE_OFFSET, + AUDIO_CODE_OFFSET + C0_VOCAB_SIZE, + dtype=torch.long, + ), + ) + ), + persistent=False, + ) + self.frame_embedding_scale = float(NUM_CODEBOOKS) ** -0.5 + + scheduler_config = getattr(vllm_config, "scheduler_config", None) + # Guidance doubles the rows, so the batch is twice the admitted + # request count. + self._max_rows = max(64, 2 * int(getattr(scheduler_config, "max_num_seqs", 16))) + + # Per-step scratch, rebuilt every step from request-keyed state. + # Nothing here may survive a step: vLLM compacts and reorders the + # persistent batch, so a row index is only meaningful within the step + # that produced it. + self._last_logits_hidden: torch.Tensor | None = None + self._step_engine_ids: list[str] = [] + self._row_request_ids: list[str] = [] + self._row_roles: list[str] = [] + self._row_pair_ids: list[str] = [] + self._row_seeds: list[int] = [] + self._row_max_frames: list[int] = [] + self._row_prompt_tokens: list[int] = [] + self._row_sample_eligible: list[bool] = [] + self._feedback_positions: torch.Tensor | None = None + self._feedback_codes: torch.Tensor | None = None + + # Per-step index vectors are staged through pinned memory rather than + # built with ``torch.tensor(..., device="cuda")``; see + # ``staging.PinnedIndexStage`` for why that matters here. Allocated on + # first use because the device is not known at construction time. + # Rows are (cond_row, uncond_row, seed, position). + self._sample_stage: PinnedIndexStage | None = None + # One row, the flat token offsets of this step's decoding rows. + self._feedback_stage: PinnedIndexStage | None = None + + # Depth-loop draw offsets, ``[1..NUM_CODEBOOKS)``. Held as a buffer so + # the per-step position arithmetic is one broadcast add rather than one + # scalar add per codebook. + self.register_buffer( + "depth_draw_offsets", + torch.arange(1, NUM_CODEBOOKS, dtype=torch.int64), + persistent=False, + ) + + # Frame feedback slab. A captured decode graph replays fixed addresses + # and fixed shapes, so the rows to overwrite cannot be expressed as a + # variable-length ``index_copy`` against a freshly built index tensor. + # Instead every row of the reserved slab carries a codes entry and a + # flag, both written in place before the graph runs, and the forward + # selects with ``torch.where`` over the whole row set. Sized to the + # widest decode batch; wider prefill batches take the eager path in + # ``_apply_frame_feedback``, and prefill is never graph-replayed. + self.register_buffer( + "frame_feedback_codes", + torch.zeros((self._max_rows, NUM_CODEBOOKS), dtype=torch.long), + persistent=False, + ) + self.register_buffer( + "frame_feedback_active", + torch.zeros(self._max_rows, dtype=torch.bool), + persistent=False, + ) + + # Cross-step state. Keyed by CFG pair id, because a guided request is + # two engine rows sharing one song. The engine renames requests to + # internal UUIDs on admission, so the pair id is derived from the + # user-facing id and the engine ids are mapped onto it separately. + self._states: dict[str, RequestARState] = {} + self._pair_of_engine_id: dict[str, str] = {} + # Per-request constants, keyed by engine id. ``forward`` is the only + # hook the engine hands the sampling extra args to, but its Python body + # does not run on a step replayed from a CUDA graph, so what it learns + # is cached here and re-bound to rows in ``prepare_runner_inputs``. + self._row_constants: dict[str, _RowConstants] = {} + self._row_binding_complete = False + self._incomplete_steps = 0 + + # The depth loop is recorded per pair count and replayed after that. + self._depth_graphs = DepthGraphCache(self._depth_decode_eager) + + # ---------------------------------------------------------------- forward + + def prepare_runner_inputs( + self, + *, + req_ids: list[str], + num_computed_tokens: Any, + num_scheduled_tokens: Any, + input_ids: torch.Tensor, + positions: torch.Tensor, + **_: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Bind this step's batch rows to requests, before the forward runs. + + This is the only hook that sees the authoritative row-to-request + mapping for the step about to execute. vLLM compacts and reorders the + persistent batch whenever a request finishes or is preempted, so any + state indexed by row number from a previous step is stale by the time + it is read. Everything the model carries across steps is therefore + keyed by request id and gathered into row order here. + + It also resolves which rows are decoding. A row is decoding when its + prompt is already computed and it is advancing by a single token; a + row still working through a chunked prefill is not, and must keep its + real prompt embeddings. + + Python runs here on every step, including steps whose forward is + replayed from a CUDA graph rather than executed, so this is also where + the row-to-request binding that ``sample`` reads has to be rebuilt. + """ + self._step_engine_ids = [str(rid) for rid in req_ids] + computed = [int(value) for value in num_computed_tokens] + scheduled = [int(value) for value in num_scheduled_tokens] + # Prompt length for a row that is completing its prefill on this step. + # Only meaningful when the row is about to produce its first logit, + # which is the only moment state is created. + self._row_prompt_tokens = [computed[row] + scheduled[row] for row in range(len(req_ids))] + + # Flat token offset of each row's first scheduled token. + offsets: list[int] = [] + running = 0 + for count in scheduled: + offsets.append(running) + running += count + + feedback_rows: list[int] = [] + feedback_codes: list[torch.Tensor] = [] + for row, engine_id in enumerate(self._step_engine_ids): + # Both members of a guided pair are fed the same frame, so state + # is keyed by pair rather than by request. + pair_id = self._pair_of_engine_id.get(engine_id) + state = self._states.get(pair_id) if pair_id else None + if state is None or state.last_codes is None: + continue + # A decode row advances by exactly one token past a complete + # prompt. A row still working through a chunked prefill keeps its + # real prompt embeddings. + if scheduled[row] != 1 or computed[row] < state.prompt_tokens: + continue + feedback_rows.append(offsets[row]) + feedback_codes.append(state.last_codes) + + self._write_frame_feedback(feedback_rows, feedback_codes, input_ids.device) + self._bind_rows(computed, scheduled) + return input_ids, positions + + def _write_frame_feedback( + self, + feedback_rows: list[int], + feedback_codes: list[torch.Tensor], + device: torch.device, + ) -> None: + """Publish this step's frame feedback for the forward to consume. + + Two forms are written, because two forwards can consume it. The slab + is what a captured decode graph reads: fixed address, fixed shape, + selected row by row with a flag. The loose index and code tensors are + the eager fallback for a batch wider than the slab, which is only ever + a prefill batch and therefore never a replayed one. + """ + if self._feedback_stage is None: + self._feedback_stage = PinnedIndexStage(1, self._max_rows, device) + + active = self.frame_feedback_active + if not feedback_rows: + self._feedback_positions = None + self._feedback_codes = None + active.zero_() + return + + codes = torch.cat(feedback_codes, dim=0).to(device) + if max(feedback_rows) < self._max_rows: + positions = self._feedback_stage.push([feedback_rows], len(feedback_rows))[0] + self.frame_feedback_codes.index_copy_(0, positions, codes) + active.zero_() + active.index_fill_(0, positions, True) + else: + # Too wide for the slab. Nothing reads it on this step, but a stale + # flag would corrupt the next graph replay if the batch shrank. + positions = torch.tensor(feedback_rows, dtype=torch.long, device=device) + active.zero_() + self._feedback_positions = positions + self._feedback_codes = codes + + def _bind_rows(self, computed: list[int], scheduled: list[int]) -> None: + """Rebuild the row-order request metadata from the per-request cache. + + ``forward`` learns these values from the engine's own kwargs, but its + Python body is skipped whenever the step is replayed from a CUDA graph, + so the binding is re-established here. A row whose request has not been + seen by ``forward`` yet leaves the binding incomplete; that only + happens on the step a request is admitted, which is a prefill step, and + prefill is never replayed, so ``forward`` fills it in before the + sampler runs. + """ + constants = [self._row_constants.get(engine_id) for engine_id in self._step_engine_ids] + if not constants or any(entry is None for entry in constants): + self._row_binding_complete = False + return + self._row_request_ids = [entry.request_id for entry in constants] + self._row_roles = [entry.role for entry in constants] + self._row_pair_ids = [entry.pair_id for entry in constants] + self._row_seeds = [entry.seed for entry in constants] + self._row_max_frames = [entry.max_frames for entry in constants] + + # A row is eligible to draw a frame once it has computed its whole + # prompt. Mirrors the engine's own rule; the prompt length is taken + # from the pair's state, which exists for every row whose request has + # already produced a logit. + eligible: list[bool] = [] + for row, entry in enumerate(constants): + state = self._states.get(entry.pair_id) + if state is None: + self._row_binding_complete = False + return + eligible.append(computed[row] + scheduled[row] >= state.prompt_tokens) + self._row_sample_eligible = eligible + self._row_binding_complete = True + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + del intermediate_tensors + self._capture_row_metadata(kwargs) + + if inputs_embeds is None: + hidden_states = self.model.embed_tokens(input_ids) + hidden_states = self._apply_frame_feedback(hidden_states) + else: + hidden_states = inputs_embeds + + residual: torch.Tensor | None = None + for layer in self.model.layers: + hidden_states, residual = layer(positions, hidden_states, residual) + norm_out = self.model.norm(hidden_states, residual) + if isinstance(norm_out, tuple): + norm_out = norm_out[0] + return norm_out + + def _capture_row_metadata(self, kwargs: dict[str, Any]) -> None: + """Read this step's per-row CFG role, pair id and seed. + + Row order here matches ``prepare_runner_inputs``: both are built from + ``input_batch.req_ids``. The request ids resolved there are preferred, + and the info dicts are only a fallback for paths that do not route + through that hook. + """ + if self._row_binding_complete: + # Every row was resolved from the per-request cache before the + # forward, so there is nothing here the sampler still needs. Bail + # out rather than rebuild it: this runs on every eager decode step. + return + infos = kwargs.get("model_intermediate_buffer") + if infos is None: + infos = kwargs.get("runtime_additional_information") + extra_args = kwargs.get("sampling_extra_args") + eligible = kwargs.get("request_sample_eligible") + self._row_sample_eligible = [bool(flag) for flag in eligible] if eligible else [] + if infos is None and extra_args is None: + # Warm-up dummy run: no omni kwargs are supplied. + self._row_pair_ids = [] + return + + infos = infos or [] + extra_args = extra_args or [] + rows = max(len(infos), len(extra_args), len(self._step_engine_ids)) + request_ids: list[str] = [] + roles: list[str] = [] + pair_ids: list[str] = [] + seeds: list[int] = [] + max_frames: list[int] = [] + for index in range(rows): + info = infos[index] if index < len(infos) else {} + info = info if isinstance(info, dict) else {} + extra = extra_args[index] if index < len(extra_args) else {} + extra = extra if isinstance(extra, dict) else {} + # The user-facing id. The engine renames a request to an internal + # UUID on admission, so this is the only id that both the info + # dicts and postprocess agree on, and the only one that carries + # the companion suffix a pair is recognized by. + request_id = _request_key(info.get("global_request_id")) or f"row-{index}" + pair_id = str(extra.get("cfg_pair_id") or _pair_id_from_request(request_id)) + request_ids.append(request_id) + roles.append(str(extra.get("cfg_role") or _role_from_request(request_id))) + pair_ids.append(pair_id) + seeds.append(int(extra.get("tts_local_seed") or 0)) + max_frames.append(_first_int(info.get("max_audio_frames"), DEFAULT_MAX_AUDIO_FRAMES)) + # Remember which pair each engine-side row belongs to, so the next + # step's feedback can be gathered before the info dicts are seen, + # and cache the whole row constant so a graph-replayed step can be + # bound without the info dicts at all. + if index < len(self._step_engine_ids): + engine_id = self._step_engine_ids[index] + self._pair_of_engine_id[engine_id] = pair_id + self._row_constants[engine_id] = _RowConstants( + request_id=request_id, + role=roles[index], + pair_id=pair_id, + seed=seeds[index], + max_frames=max_frames[index], + ) + self._row_request_ids = request_ids + self._row_roles = roles + self._row_pair_ids = pair_ids + self._row_seeds = seeds + self._row_max_frames = max_frames + + def _apply_frame_feedback(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Replace each decoding row's token embedding with its frame embedding. + + The rows to replace and the codes to write were resolved in + ``prepare_runner_inputs`` against this step's own batch, so a mixed + prefill and decode step touches only the decoding rows and a prompt + token is never overwritten. + + Which of the two forms is read depends only on the row count, which is + a constant inside a captured graph: a decode batch fits the reserved + slab and is selected with ``torch.where`` over fixed buffers, while a + prefill batch wider than the slab falls back to a variable-length + ``index_copy``. Prefill is never replayed, so the fallback never has to + be graph-safe. + """ + rows = hidden_states.shape[0] + if rows <= self._max_rows: + frame_embeds = self.embed_frames(self.frame_feedback_codes[:rows]).to(dtype=hidden_states.dtype) + return torch.where( + self.frame_feedback_active[:rows].unsqueeze(1), + frame_embeds, + hidden_states, + ) + positions = self._feedback_positions + codes = self._feedback_codes + if positions is None or codes is None: + return hidden_states + frame_embeds = self.embed_frames(codes).to(dtype=hidden_states.dtype) + return hidden_states.index_copy(0, positions, frame_embeds) + + def embed_frames(self, codes: torch.Tensor) -> torch.Tensor: + """Compose one conditioning embedding per row from all eight codebooks. + + Args: + codes: ``[rows, 8]`` codebook indices, ``c0`` zero-based. + + Returns: + ``[rows, hidden]`` in the embedding table's dtype. + """ + c0 = self.model.embed_tokens(codes[:, 0] + AUDIO_CODE_OFFSET) + extra = self.audio_embeddings(codes[:, 1:] + self.audio_embedding_offsets) + return (c0 + extra.sum(dim=1).to(c0.dtype)) * self.frame_embedding_scale + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: + # sample() has no other route to the backbone hidden state. + self._last_logits_hidden = hidden_states + return self.logits_processor(self.lm_head, hidden_states) + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_tokens(input_ids) + + # ----------------------------------------------------------------- sample + + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: + """Draw one frame per conditioned request and feed it back to both rows.""" + del sampling_metadata + hidden = self._last_logits_hidden + self._last_logits_hidden = None + if hidden is None or not self._row_pair_ids: + return None + + rows = int(logits.shape[0]) + pairs = self._resolve_pairs(rows) + if not pairs: + # Nothing to draw this step. Emit a benign audio code so every row + # stays alive until its partner is admitted; returning None would + # hand the rows to the stock sampler, whose token would be read as + # a real code. + held = torch.full((rows, 1), _HOLD_TOKEN, dtype=torch.long, device=logits.device) + return SamplerOutput(sampled_token_ids=held, logprobs_tensors=None) + + states = [self._state_for(self._row_pair_ids[pair[0]], pair[0]) for pair in pairs] + cond_rows, uncond_rows, seeds, positions = self._stage_step_vectors( + [pair[0] for pair in pairs], + [pair[1] for pair in pairs], + [state.seed for state in states], + [state.advance_position() for state in states], + logits.device, + ) + + narrowed = logits.index_select(1, self.c0_logit_ids).float() + c0_logits = apply_c0_cfg(narrowed.index_select(0, cond_rows), narrowed.index_select(0, uncond_rows)) + drawn = sample_topk_seeded(c0_logits, seeds, positions) + stopped = drawn == 0 + c0 = (drawn - 1).clamp_min(0) + + codes, depth_hidden = self._depth_decode( + hidden.index_select(0, cond_rows), + hidden.index_select(0, uncond_rows), + c0, + seeds, + positions, + ) + frame_hidden = torch.cat((hidden.index_select(0, cond_rows), depth_hidden), dim=-1) + + stopped_flags = stopped.tolist() + drain_rows: list[int] = [] + for index, (cond_row, uncond_row) in enumerate(pairs): + state = states[index] + if state.finished: + # The drain step. The terminal span written last step is on + # its way out in this step's payload, so the pair can stop now. + state.pending_stop = False + drain_rows.extend((index, cond_row, uncond_row)) + continue + if stopped_flags[index]: + state.finished = True + state.pending_stop = True + state.finish_reason = "stop" + state.last_codes = None + continue + # Both rows of the pair are fed the same frame next step; the + # codes live on the request, not on a row, because the row this + # request occupies may change before then. + state.last_codes = codes[index : index + 1] + state.frames.append(frame_hidden[index : index + 1]) + state.codes.append(codes[index : index + 1]) + state.generated_frames += 1 + if state.generated_frames >= state.max_frames: + state.finished = True + state.pending_stop = True + state.finish_reason = "length" + del cond_row, uncond_row + + # Rows that resolved into no pair keep this value. A hold token rather + # than id 0, which is an arbitrary text token: a row waiting for its + # partner should wait, not emit something the vocabulary reads as + # prose. It is never fed back as audio, because feedback is gated on + # the pair's own state. + token_ids = torch.full((rows, 1), _HOLD_TOKEN, dtype=torch.long, device=logits.device) + sampled = self.c0_logit_ids[drawn] + token_ids[cond_rows, 0] = sampled + token_ids[uncond_rows, 0] = sampled + + # The stop token is deferred by exactly one step. A pair that ends this + # step must stay alive so that the terminal span ``postprocess`` is + # about to write can ride out on the next step's payload; a pair that + # already published its tail stops now. + # + # This covers the two ways a song ends on its own: the audio-end token, + # and the request's frame budget, which reaches the model as + # ``max_audio_frames``. It does not cover the engine ending a request + # underneath the model, because the payload for a step is fixed before + # the model is told. Hitting ``max_model_len`` mid-song therefore drops + # the final window. In practice the frame budget is reached first; a + # request can only hit the context limit with a prompt long enough that + # the two limits collide, which the deploy config documents. + for index, (cond_row, uncond_row) in enumerate(pairs): + state = states[index] + if state.pending_stop: + token_ids[cond_row, 0] = _HOLD_TOKEN + token_ids[uncond_row, 0] = _HOLD_TOKEN + for _, cond_row, uncond_row in zip(drain_rows[0::3], drain_rows[1::3], drain_rows[2::3], strict=True): + token_ids[cond_row, 0] = _STOP_TOKEN + token_ids[uncond_row, 0] = _STOP_TOKEN + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + def _stage_step_vectors( + self, + cond_rows: list[int], + uncond_rows: list[int], + seeds: list[int], + positions: list[int], + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Move this step's four small int64 vectors to the device in one copy. + + The copy is asynchronous, so the four returned views are only valid + once the stream reaches them. Every consumer is a kernel on that same + stream, which is what makes this safe. + """ + if self._sample_stage is None: + self._sample_stage = PinnedIndexStage(4, self._max_rows, device) + block = self._sample_stage.push([cond_rows, uncond_rows, seeds, positions], len(cond_rows)) + return block[0], block[1], block[2], block[3] + + def _resolve_pairs(self, rows: int) -> list[tuple[int, int]]: + """Map ``cfg_pair_id`` to its ``(cond_row, uncond_row)`` batch indices. + + Guidance is mandatory for this checkpoint, so an incomplete pair is a + hard error rather than something to route around. Skipping the step + would leave the conditioned request producing no frames and the client + waiting forever, which is the failure shape this module has been bitten + by before; raising surfaces it to the caller instead. + + Raises: + RuntimeError: If any pair is missing a member in this step. + """ + by_pair: dict[str, dict[str, int]] = {} + limit = min(rows, len(self._row_pair_ids)) + for row in range(limit): + # A row still finishing a chunked prefill is handed to the sampler + # but its token is discarded, so drawing a frame for it would add + # a bogus frame and desync the seeded RNG. + if row < len(self._row_sample_eligible) and not self._row_sample_eligible[row]: + continue + by_pair.setdefault(self._row_pair_ids[row], {})[self._row_roles[row]] = row + + pairs: list[tuple[int, int]] = [] + incomplete: list[str] = [] + for pair_id, roles in by_pair.items(): + if _COND in roles and _UNCOND in roles: + pairs.append((roles[_COND], roles[_UNCOND])) + else: + incomplete.append(f"{pair_id}({'+'.join(sorted(roles))})") + if incomplete: + # A member can be absent for a step while its partner is still + # being admitted. Holding the present row costs one wasted decode + # step; raising would abort the whole engine step for every other + # request in the batch, so the lone row simply does not sample this + # step and waits. A pair that never converges is bounded by the + # request's own token budget rather than hanging forever. + self._incomplete_steps += 1 + if self._incomplete_steps in (1, 64, 512): + logger.warning( + "MiniMax Music 3 waiting on incomplete CFG pair(s) %s " + "(step %d). If this does not resolve, the companion was " + "never enqueued: check that the stage declares " + "extra_args.cfg_role in its deploy config.", + incomplete[:4], + self._incomplete_steps, + ) + return pairs + + def _depth_decode( + self, + cond_hidden: torch.Tensor, + uncond_hidden: torch.Tensor, + c0: torch.Tensor, + seeds: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Walk ``c1..c7``, replaying the loop from a CUDA graph where possible. + + Seven four-layer forwards over at most eight positions is far too + little work per kernel to cover its own launch, so the loop is recorded + once per pair count and replayed after that. The cache falls back to + the eager loop for any shape it has not recorded or could not record. + """ + return self._depth_graphs.run(cond_hidden, uncond_hidden, c0, seeds, positions) + + def _depth_decode_eager( + self, + cond_hidden: torch.Tensor, + uncond_hidden: torch.Tensor, + c0: torch.Tensor, + seeds: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Walk ``c1..c7``, returning the codes and the stacked depth hiddens. + + Both branches run so the depth heads can be guided the same way the + ``c0`` head is; the drawn code is shared, so the two branches stay + aligned. + """ + decoder = self.rvq_decoder + pair_hidden = torch.cat((cond_hidden, uncond_hidden), dim=0) + c0_paired = c0.repeat(2) + c0_embed = self.model.embed_tokens(c0_paired + AUDIO_CODE_OFFSET) + seq = [ + decoder.projection(pair_hidden).unsqueeze(1), + decoder.projection(c0_embed).unsqueeze(1), + ] + n_pairs = int(c0.shape[0]) + codes = [c0] + depth_parts: list[torch.Tensor] = [] + # Every residual draw's noise comes from the same hash of + # ``(seed, position + codebook, candidate)``, so all seven are produced + # in one batch here instead of re-running the twenty-odd integer ops + # per codebook inside the loop. The arithmetic is elementwise, so the + # values are the same ones the per-codebook calls produced. + depth_gumbel = gumbel_noise( + seeds, + positions.unsqueeze(1) + self.depth_draw_offsets, + decoder.audio_vocab_size, + ) + for index in range(1, NUM_CODEBOOKS): + out = decoder(torch.cat(seq, dim=1))[:, -1] + depth_parts.append(out[:n_pairs].detach()) + head_logits = decoder.audio_heads[index - 1](out).float() + guided = apply_depth_cfg(head_logits[:n_pairs], head_logits[n_pairs:]) + code = draw_from_gumbel(guided, depth_gumbel[:, index - 1]) + codes.append(code) + if index < NUM_CODEBOOKS - 1: + embedding = self.audio_embeddings(code.repeat(2) + (index - 1) * AUDIO_VOCAB_SIZE) + seq.append(decoder.projection(embedding).unsqueeze(1)) + return torch.stack(codes, dim=1), torch.cat(depth_parts, dim=-1) + + def _state_for(self, pair_id: str, row: int) -> RequestARState: + """Fetch or create the state for a guided pair. + + Created on the step the pair first produces a logit, which is the step + its prefill completes, so the recorded prompt length is the real one. + """ + state = self._states.get(pair_id) + if state is None: + seed = self._row_seeds[row] if row < len(self._row_seeds) else 0 + state = RequestARState( + seed=seed, + max_frames=self._row_max_frames[row] if row < len(self._row_max_frames) else DEFAULT_MAX_AUDIO_FRAMES, + prompt_tokens=self._row_prompt_tokens[row] if row < len(self._row_prompt_tokens) else 0, + ) + self._states[pair_id] = state + logger.info( + "MiniMax Music 3 AR start pair=%s seed=%d max_frames=%d prompt_tokens=%d", + pair_id, + state.seed, + state.max_frames, + state.prompt_tokens, + ) + return state + + # ------------------------------------------------------------ emit / free + + def postprocess( + self, + hidden_states_slice: torch.Tensor, + multimodal_outputs: Any = None, + **req_infos: Any, + ) -> dict[str, Any]: + """Publish this request's next span of conditioning frames. + + Mid-stream this is one window, released a frame late so that a window + handed downstream is provably not the last. On the final step it is + instead a single contiguous span covering every window not yet + emitted, because one step produces exactly one payload and the tail + cannot be split across two. + """ + del hidden_states_slice, multimodal_outputs + request_id = _request_key(req_infos.get("global_request_id")) + # The pair's two rows share one song, so only the conditioned row + # publishes it. The companion would otherwise emit a duplicate span + # and the acoustic stage would decode the song twice. + if _role_from_request(request_id) == _UNCOND: + return {} + state = self._states.get(_pair_id_from_request(request_id)) + if state is None: + return {} + if state.finished: + return self._terminal_payload(state) + window = state.ready_window() + if window is None: + return {} + chunk = state.take_window(window, is_final=False) + return self._span_payload(state, chunk, start_frame=window.start, is_final=False) + + def _terminal_payload(self, state: RequestARState) -> dict[str, Any]: + """Emit every remaining frame as one span, once.""" + start = state.emitted_windows * AR_CHUNK_HOP_FRAMES + end = state.frames.end_frame + if end <= start: + return {} + chunk = state.frames.slice(start, end) + state.emitted_windows = len(chunk_windows(state.generated_frames)) + state.frames.clear() + return self._span_payload(state, chunk, start_frame=start, is_final=True) + + def _span_payload( + self, + state: RequestARState, + chunk: torch.Tensor, + *, + start_frame: int, + is_final: bool, + ) -> dict[str, Any]: + """Build the cross-stage payload for one contiguous frame span. + + ``num_processed_tokens`` is the span's first global AR frame index, so + the acoustic stage can re-derive window boundaries and therefore which + window is first and which is last. + """ + frames = int(chunk.shape[0]) + codes = state.take_codes(start_frame, start_frame + frames) + payload: dict[str, Any] = { + "latent": chunk.contiguous(), + "meta": { + "num_processed_tokens": start_frame, + "codec_chunk_frames": frames, + "last_chunk": is_final, + "audio_seed": state.seed, + }, + } + if codes is not None: + payload["codes"] = {"audio": codes} + return payload + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + infos = kwargs.get("model_intermediate_buffer") + if infos is None: + infos = kwargs.get("runtime_additional_information") + if not infos: + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=None) + + latents: list[torch.Tensor] = [] + codes: list[torch.Tensor] = [] + metas: list[dict[str, Any]] = [] + any_payload = False + for info in infos: + info = info if isinstance(info, dict) else {} + latent = info.get("latent") + if isinstance(latent, torch.Tensor) and latent.numel(): + latents.append(latent) + sub = info.get("codes") + codes.append(sub.get("audio") if isinstance(sub, dict) else torch.empty(0)) + meta = info.get("meta") + metas.append(dict(meta) if isinstance(meta, dict) else {}) + any_payload = True + # The runner merges postprocess output into this buffer and + # never clears it, so a window left behind would be re-emitted + # on every later step and the acoustic stage would decode the + # same conditioning again and again. Consume it here. + info.pop("latent", None) + info.pop("codes", None) + info.pop("meta", None) + else: + latents.append(torch.empty(0)) + codes.append(torch.empty(0)) + metas.append({}) + if not any_payload: + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=None) + payload: dict[str, Any] = {"latent": latents} + if any(isinstance(entry, torch.Tensor) and entry.numel() for entry in codes): + payload["codes"] = {"audio": codes} + # The span's start frame and terminal flag ride along with it: without + # them the acoustic stage cannot tell which window it is looking at, + # and therefore cannot crop it correctly. + # + # Every field is a list with one entry per batch row, in the same order + # as ``latent``. The payload is partitioned per request by indexing + # list values, so a scalar here would be broadcast to the whole batch + # and every request would decode with the last request's seed, window + # offset and terminal flag. + meta_keys = {key for meta in metas for key in meta} + if meta_keys: + payload["meta"] = {key: [meta.get(key) for meta in metas] for key in sorted(meta_keys)} + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=payload) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Release per-request state once the engine has retired the request. + + The key must be normalized the same way it was when the state was + created, or every finished song leaks its frame buffer. + """ + for engine_id in finished_req_ids: + # These are engine-side ids, so translate through the mapping the + # forward pass built. A pair's state is released when its second + # member retires. + self._row_constants.pop(str(engine_id), None) + pair_id = self._pair_of_engine_id.pop(str(engine_id), None) + if pair_id is None or any(pair_id == value for value in self._pair_of_engine_id.values()): + continue + request_id = pair_id + state = self._states.pop(pair_id, None) + if state is not None: + logger.info( + "MiniMax Music 3 AR done request=%s frames=%d reason=%s peak_buffer=%d", + request_id, + state.generated_frames, + state.finish_reason, + state.frames.peak_frames, + ) + state.release() + + # --------------------------------------------------------------- weights + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load the Qwen3 backbone, then the audio modules from their component. + + The backbone comes from the stage's own checkpoint folder and loads + through vLLM. The audio embedding table and the depth decoder live in + a separate component of the repo and are loaded explicitly. + """ + loader = AutoWeightsLoader(self, skip_prefixes=["audio_embeddings.", "rvq_decoder."]) + loaded = loader.load_weights(weights) + loaded |= load_depth_decoder_weights( + self.vllm_config, + audio_embeddings=self.audio_embeddings, + rvq_decoder=self.rvq_decoder, + ) + return loaded diff --git a/vllm_omni/model_executor/models/minimax_music3/weights.py b/vllm_omni/model_executor/models/minimax_music3/weights.py new file mode 100644 index 000000000..b019ff461 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/weights.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Component weight loading for MiniMax Music 3. + +The checkpoint is a multi-component repo. The Qwen3 backbone is a normal +``Qwen3ForCausalLM`` under ``language_model/`` and loads through vLLM. The +audio embedding table and the depth decoder live in ``rvq_depth_decoder/``, +and the acoustic stage's three components live in their own folders. None of +those are things vLLM's loader knows how to map, so they are read directly. + +Every loader here is strict. A silently partial load would produce audio that +sounds plausible but is wrong, which is far worse than a startup failure. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch +from torch import nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from .constants import AUDIO_VOCAB_SIZE, NUM_CODEBOOKS + +logger = init_logger(__name__) + +_AUDIO_EMBEDDING_KEY = "audio_embeddings.weight" +_COMPONENT_WEIGHT_NAME = "diffusion_pytorch_model.safetensors" + +# Component folders, relative to the repository root. +DEPTH_DECODER_DIR = "rvq_depth_decoder" +CONDITION_ENCODER_DIR = "condition_encoder" +TRANSFORMER_DIR = "transformer" +VOCODER_DIR = "vocoder" + +# Folder names that identify a repository root. This is deliberately a +# name-only probe and NOT the set a stage must be able to load: see +# ``_REQUIRED_COMPONENTS``. +_ROOT_MARKERS = (DEPTH_DECODER_DIR, TRANSFORMER_DIR, VOCODER_DIR) + +# Every component a stage actually reads. ``condition_encoder/`` is loaded by +# the acoustic stage, so a snapshot missing it is incomplete even though the +# root markers above are all present. +_REQUIRED_COMPONENTS = (DEPTH_DECODER_DIR, CONDITION_ENCODER_DIR, TRANSFORMER_DIR, VOCODER_DIR) + +# The subset a stage reading the repo ROOT needs on disk. The AR backbone under +# ``language_model/`` is pulled by vLLM's own loader, and the repo's top-level +# ``.pth`` files and ``qwen_7B/`` folder belong to the reference implementation, +# so neither is downloaded here. +_COMPONENT_PATTERNS = ( + "*.json", + f"{DEPTH_DECODER_DIR}/*", + f"{CONDITION_ENCODER_DIR}/*", + f"{TRANSFORMER_DIR}/*", + f"{VOCODER_DIR}/*", +) + + +def _component_is_complete(folder: Path) -> bool: + """True when *folder* holds a weight set ``load_component_state`` can read. + + A single shard satisfies a bare glob, so an interrupted multi-shard + download would otherwise pass as complete. Hub shards are named + ``--of-.safetensors``, which states how many pieces + to expect, so that run can be checked for gaps. + + Only that naming is provably incomplete. ``load_component_state`` loads + whatever the same glob returns, so any other shard-named file is accepted + here: judging it incomplete would force a needless re-download and fail + closed offline on weights that are present and readable. + """ + stem = Path(_COMPONENT_WEIGHT_NAME).stem + if (folder / _COMPONENT_WEIGHT_NAME).is_file(): + return True + + shards = list(folder.glob(f"{stem}-*.safetensors")) + if not shards: + return False + + totals: set[int] = set() + seen: set[int] = set() + for shard in shards: + marker = shard.name[len(stem) + 1 : -len(".safetensors")] + index_text, separator, total_text = marker.partition("-of-") + if not separator or not index_text.isdigit() or not total_text.isdigit(): + continue + totals.add(int(total_text)) + seen.add(int(index_text)) + if not totals: + # No piece announces a total, so there is nothing to be short of. + return True + if len(totals) != 1: + return False + return seen == set(range(1, totals.pop() + 1)) + + +def _snapshot_has_components(root: Path) -> bool: + """True when every required component folder holds its complete weights.""" + return all(_component_is_complete(root / component) for component in _REQUIRED_COMPONENTS) + + +def _download_component_snapshot( + model: str, + *, + revision: str | None = None, + download_dir: str | None = None, +) -> Path: + """Resolve an HF repo id to a snapshot holding the component folders. + + Cache-first, so a warm cache needs no Hub round trip. Depending on the + huggingface_hub version, ``local_files_only=True`` either raises on a + partial cache or returns the partial snapshot root unchanged; both + variants must fall through to the online call, or a cold cache fails + startup without ever consulting the Hub. + + ``revision`` and ``download_dir`` must be threaded through: resolving the + components against the default branch and the default cache while the AR + backbone was pinned elsewhere would silently mix weights from two commits. + """ + from vllm.transformers_utils.repo_utils import hf_api + + kwargs: dict[str, object] = {"allow_patterns": list(_COMPONENT_PATTERNS)} + if revision is not None: + kwargs["revision"] = revision + if download_dir is not None: + kwargs["cache_dir"] = download_dir + + try: + cached = Path(hf_api().snapshot_download(model, local_files_only=True, **kwargs)) + except Exception: + cached = None + if cached is not None and _snapshot_has_components(cached): + return cached + return Path(hf_api().snapshot_download(model, **kwargs)) + + +def _repo_id_and_revision_from_cache_path(path: Path) -> tuple[str | None, str | None]: + """Recover the Hub repo id and commit when *path* is inside an HF cache. + + Stage init hands the talker the snapshot's ``language_model`` folder. On a + cold cache the sibling component folders were never downloaded, so the + marker walk fails even though the missing pieces are one snapshot_download + away. The cache layout is ``models--{org}--{name}/snapshots/{revision}/...``, + so the exact commit is recoverable too and must be reused: re-resolving the + repo id alone would fetch the default branch instead of the snapshot the + caller is already sitting in. + """ + parts = path.parts + for position, part in enumerate(parts): + if not part.startswith("models--"): + continue + repo_id = part.removeprefix("models--").replace("--", "/") + revision = None + if position + 2 < len(parts) and parts[position + 1] == "snapshots": + revision = parts[position + 2] + return repo_id, revision + return None, None + + +def resolve_repo_root( + model_path: str | os.PathLike[str], + *, + revision: str | None = None, + download_dir: str | None = None, +) -> Path: + """Find the multi-component repository root from a stage's model path. + + A stage that declares ``model_subdir`` sees the subfolder as its model + path, so the root is usually the parent. Both are checked, and the search + walks up a couple of levels for snapshot layouts. + + The acoustic stage declares no ``model_subdir`` because it reads the root + itself, so on a hub-id deployment its model path is the repo id rather + than a directory. Resolve that to a local snapshot instead of treating it + as a relative path. + + A cold cache is recoverable too: stage init pre-downloads only + ``language_model/`` and ``tokenizer/``, so the talker's model path exists + while the component folders do not. Fetch them here instead of failing. + + A candidate is accepted only when its components are actually loadable. + Matching on folder names alone would accept a snapshot left weightless by + an interrupted download and defer the failure to ``load_component_state``, + past the point where the repair below could still run. + + Raises: + FileNotFoundError: If no ancestor looks like the repository root. + """ + path = Path(model_path).expanduser().resolve() + candidates = [path, *path.parents[:3]] + marker_roots = [ + candidate for candidate in candidates if all((candidate / marker).is_dir() for marker in _ROOT_MARKERS) + ] + for candidate in marker_roots: + if _snapshot_has_components(candidate): + return candidate + + # Either no ancestor carries the markers, or one does but is incomplete. + # Three recoverable cases: the path is a Hub repo id (the acoustic stage + # passes the id through), or it is a cache-snapshot subfolder whose sibling + # component folders were never downloaded, or a cached snapshot lost part + # of its weights. ``_download_component_snapshot`` is cache-first, so a + # warm cache stays offline either way. + cached_repo_id, cached_revision = _repo_id_and_revision_from_cache_path(path) + reference = str(model_path) if not path.exists() else cached_repo_id + if reference is not None: + try: + resolved = _download_component_snapshot( + reference, + revision=revision or cached_revision, + download_dir=download_dir, + ) + except Exception as exc: + raise FileNotFoundError( + f"MiniMax Music 3 could not resolve {model_path!r} to a local repository root; " + f"expected folders {_REQUIRED_COMPONENTS}" + ) from exc + if _snapshot_has_components(resolved): + return resolved + + if marker_roots: + raise FileNotFoundError( + f"MiniMax Music 3 found a repository root at {marker_roots[0]} but its components are incomplete; " + f"expected loadable weights under {_REQUIRED_COMPONENTS}" + ) + raise FileNotFoundError( + f"MiniMax Music 3 could not locate the repository root from {path}; " + f"expected sibling folders {_REQUIRED_COMPONENTS}" + ) + + +def load_component_state(root: Path, component: str, *, device: str = "cpu") -> dict[str, torch.Tensor]: + """Read one component's safetensors, sharded or not. + + Raises: + FileNotFoundError: If the component has no weight files. + """ + from safetensors.torch import load_file + + folder = root / component + single = folder / _COMPONENT_WEIGHT_NAME + if single.exists(): + return load_file(str(single), device=device) + + shards = sorted(folder.glob(f"{Path(_COMPONENT_WEIGHT_NAME).stem}-*.safetensors")) + if not shards: + raise FileNotFoundError(f"MiniMax Music 3 component {component!r} has no weights under {folder}") + state: dict[str, torch.Tensor] = {} + for shard in shards: + state.update(load_file(str(shard), device=device)) + return state + + +def load_depth_decoder_weights( + vllm_config: VllmConfig, + *, + audio_embeddings: nn.Embedding, + rvq_decoder: nn.Module, +) -> set[str]: + """Load the audio embedding table and the RVQ depth decoder in place. + + Args: + vllm_config: The stage's config, used to locate the checkpoint. + audio_embeddings: The ``c1..c7`` embedding table to populate. + rvq_decoder: The depth decoder to populate. + + Returns: + The model-side parameter names that were loaded. + + Raises: + ValueError: If the audio embedding is absent or the wrong shape. + """ + root = resolve_repo_root( + vllm_config.model_config.model, + revision=vllm_config.model_config.revision, + download_dir=vllm_config.load_config.download_dir, + ) + state = load_component_state(root, DEPTH_DECODER_DIR) + + weight = state.pop(_AUDIO_EMBEDDING_KEY, None) + expected = (AUDIO_VOCAB_SIZE * (NUM_CODEBOOKS - 1), audio_embeddings.embedding_dim) + if weight is None or tuple(weight.shape) != expected: + raise ValueError( + f"MiniMax Music 3 audio embedding mismatch: expected {expected}, got " + f"{None if weight is None else tuple(weight.shape)}" + ) + target = audio_embeddings.weight + target.data.copy_(weight.to(device=target.device, dtype=target.dtype)) + + rvq_decoder.load_checkpoint_state(state) + logger.info( + "MiniMax Music 3 loaded audio embedding and depth decoder from %s", + root / DEPTH_DECODER_DIR, + ) + loaded = {"audio_embeddings.weight"} + loaded.update(f"rvq_decoder.{name}" for name, _ in rvq_decoder.named_parameters()) + return loaded + + +__all__ = [ + "CONDITION_ENCODER_DIR", + "DEPTH_DECODER_DIR", + "TRANSFORMER_DIR", + "VOCODER_DIR", + "load_component_state", + "load_depth_decoder_weights", + "resolve_repo_root", +] diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py index 098f53792..2211e0fb2 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -2000,8 +2000,12 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): if codes_lengths is None: codes_lengths = torch.full((B,), T, device=device, dtype=torch.long) - # Keep eager execution and CUDA Graph capture on the same BF16 path. - with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + # Keep eager execution and CUDA/NPU Graph capture on the same BF16 path. + # Autocast is required on NPU because #5235's decode-LUT uses .float() + # for precision, producing float32 embeddings. Without autocast, these + # float32 tensors flow into bf16 LayerNorm weights and NPU's aclnnLayerNorm + # rejects the mixed dtype (CUDA auto-casts; NPU does not). + with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type != "cpu"): quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) zq = quantizer.decode_codes(codes) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py index 2fd0beefd..c9ce11a93 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py @@ -23,6 +23,8 @@ previous KV-cache loop -- causal attention, only the last position is read. from __future__ import annotations +from collections.abc import Iterable + import torch import torch.nn as nn import torch.nn.functional as F @@ -61,6 +63,22 @@ class MossTTSRealtimeLocalTransformer(nn.Module): prefix="model", ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Delegate ``model.*`` tensors to ``CodePredictorBaseModel.load_weights``. + + The shared body fuses ``q/k/v_proj`` into ``qkv_proj`` and + ``gate/up_proj`` into ``gate_up_proj``; its loader re-packs the HF + shards and raises on incomplete or missing fused parameters. Going + through it (instead of ``default_weight_loader`` per tensor) keeps + those guarantees for the MossTTSRealtime checkpoint too. + """ + model_weights: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + if name.startswith("model."): + model_weights.append((name[len("model.") :], tensor)) + loaded = self.model.load_weights(iter(model_weights)) + return {f"model.{n}" for n in loaded} + @torch.no_grad() def generate_frame( self, @@ -128,32 +146,45 @@ def _sample_token( do_sample: bool, generator: torch.Generator | None = None, ) -> torch.Tensor: - """Top-k + top-p sampling (matches upstream's ``sample_token`` for the - inference branch). + """Top-k + top-p sampling for the upstream inference branch. + + Nucleus filtering and multinomial sampling operate only on the retained + top-k candidates before the result is mapped back to the original + vocabulary. This preserves the categorical distribution when the top-k + boundary has no ties, but it is not seed/bit equivalent to multinomial over + a full-vocabulary tensor because random-number mapping depends on width. """ if not do_sample or temperature <= 0: return logits.argmax(dim=-1) logits = logits / max(temperature, 1e-6) + compact_indices = None if top_k and top_k > 0 and top_k < logits.shape[-1]: - top_vals, _ = torch.topk(logits, top_k, dim=-1) - thresh = top_vals[..., -1:].expand_as(logits) - logits = torch.where(logits < thresh, torch.full_like(logits, float("-inf")), logits) + top_vals, top_indices = torch.topk(logits, top_k, dim=-1, sorted=True) + # topk is already descending: keep the following nucleus filter, + # softmax, and multinomial at width k instead of the full vocab. + logits = top_vals + compact_indices = top_indices if 0.0 < top_p < 1.0: - sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) - probs = F.softmax(sorted_logits, dim=-1) + sorted_indices = None + if compact_indices is None: + logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) + probs = F.softmax(logits, dim=-1) cum = probs.cumsum(dim=-1) # Drop tail beyond top_p (keep at least one token). drop = cum > top_p drop[..., 1:] = drop[..., :-1].clone() drop[..., 0] = False - sorted_logits = sorted_logits.masked_fill(drop, float("-inf")) - logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_idx, sorted_logits) + logits = logits.masked_fill(drop, float("-inf")) + if sorted_indices is not None: + logits = torch.full_like(logits, float("-inf")).scatter_(-1, sorted_indices, logits) probs = F.softmax(logits, dim=-1) flat = probs.reshape(-1, probs.shape[-1]) sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1]) + if compact_indices is not None: + return compact_indices.gather(-1, sampled.unsqueeze(-1)).squeeze(-1) return sampled diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index 7ee551f29..e2d46ce2c 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -52,15 +52,48 @@ class _MossTTSLocalAttention(nn.Module): self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True) inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("_rope_cos_cache", torch.empty(0), persistent=False) + self.register_buffer("_rope_sin_cache", torch.empty(0), persistent=False) - def _rope_cos_sin( - self, seq_len: int, device: torch.device, dtype: torch.dtype - ) -> tuple[torch.Tensor, torch.Tensor]: - position_ids = torch.arange(seq_len, device=device, dtype=torch.float32) - freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device)) + def prepare_rope_cache( + self, + max_seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> None: + """Materialize frame-local RoPE values once per device/dtype. + + Local Depth positions always restart at zero and never exceed the + number of RVQ codebooks (12 for MOSS-TTS Local v1.5). Building + arange/einsum/cos/sin inside every one of the 12 depth steps creates + many tiny kernels and also bloats the enclosing talker CUDA graph. + Keep one non-persistent cache and slice it for prefix execution. + """ + if max_seq_len <= 0: + return + cache = self._rope_cos_cache + if cache.numel() > 0 and cache.device == device and cache.dtype == dtype and int(cache.shape[1]) >= max_seq_len: + return + + position_ids = torch.arange(max_seq_len, device=device, dtype=torch.float32) + inv_freq = self.inv_freq.to(device=device, dtype=torch.float32) + freqs = torch.einsum("s,d->sd", position_ids, inv_freq) cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype) sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype) - return cos.view(1, seq_len, 1, self.head_dim), sin.view(1, seq_len, 1, self.head_dim) + self._rope_cos_cache = cos.view(1, max_seq_len, 1, self.head_dim) + self._rope_sin_cache = sin.view(1, max_seq_len, 1, self.head_dim) + + def _rope_cos_sin( + self, + seq_len: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + self.prepare_rope_cache(seq_len, device, dtype) + return ( + self._rope_cos_cache[:, :seq_len], + self._rope_sin_cache[:, :seq_len], + ) @staticmethod def _rotate_half(x: torch.Tensor) -> torch.Tensor: @@ -200,9 +233,15 @@ class MossTTSLocalDepthTransformer(nn.Module): batch_size = backbone_last_hidden.shape[0] dtype = self.ln_f.weight.dtype + # Populate the complete [0, n_vq) table before torch.compile traces + # _forward_prefix or the runner captures talker_mtp. The compiled / + # captured body then contains only fixed-address cache slices rather + # than cache construction or buffer replacement. + for block in self.h: + block.attn.prepare_rope_cache(n_vq, backbone_last_hidden.device, dtype) + embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype) embeds[:, 0, :] = backbone_last_hidden.to(dtype) - hidden = self._run_prefix(embeds[:, :1, :]) local_hidden = hidden[:, 0, :] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index 171e57328..0495d79d8 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -1179,7 +1179,13 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): # ``gate_proj``/``up_proj`` → ``gate_up_proj``. ``default_weight_loader`` # alone leaves those fused params un-initialised, which silently turns # the backbone into a slightly-corrupted model that never emits EOS. + # The local depth transformer (shared ``CodePredictorBaseModel``) has + # the exact same hazard: its ``qkv_proj``/``gate_up_proj`` are fused, + # so its shards must go through ``CodePredictorBaseModel.load_weights`` + # (via ``MossTTSRealtimeLocalTransformer.load_weights``), which re-packs + # them and raises on incomplete/missing fused parameters. backbone_weights: list[tuple[str, torch.Tensor]] = [] + local_weights: list[tuple[str, torch.Tensor]] = [] skipped: list[str] = [] for name, tensor in weights: if name.startswith("language_model."): @@ -1194,8 +1200,9 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): continue # non-persistent buffer, recomputed at runtime if sub.startswith("embed_tokens."): sub = "codec_embedding." + sub[len("embed_tokens.") :] - tgt = "local_transformer.model." + sub - elif name.startswith("local_transformer.local_lm_heads."): + local_weights.append((f"model.{sub}", tensor)) + continue + if name.startswith("local_transformer.local_lm_heads."): tgt = "local_lm_heads." + name[len("local_transformer.local_lm_heads.") :] else: tgt = name @@ -1210,6 +1217,11 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module): for n in backbone_loaded: loaded.add(f"model.{n}") + # Delegate depth-transformer weights to the shared fused loader. + local_loaded = self.local_transformer.load_weights(iter(local_weights)) + for n in local_loaded: + loaded.add(f"local_transformer.{n}") + # Build stacked embedding cache for vectorised decode ops. # embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks. self._stacked_audio_emb_w = torch.stack( @@ -1323,7 +1335,11 @@ class MossTTSLocalTalkerForGeneration(nn.Module): self.mtp_hidden_size = hidden_size self.talker_mtp_graph_safe = True self.talker_mtp_output_key = ("audio_codes", "current") - self.use_async_omni_output = False + # ``make_omni_output`` keeps code rows fixed-shape and performs all + # state updates eagerly, so the runner can safely pack/snapshot them + # before CUDA-graph buffers are reused. This moves the pinned D2H and + # output materialization off the Stage-0 execution thread. + self.use_async_omni_output = True self.eager_omni_postprocess_before_async_output = True self.omni_pooler_payload_include_hidden = False self.postprocess_uses_multimodal_outputs = False @@ -1466,7 +1482,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module): ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: """Build per-step input embeddings (text + audio additive fusion). - Prefill: initialise the per-request ``{is_stopping, step}`` state. + Prefill: initialise the per-request ``{is_stopping}`` state. Decode: combine the forced text token (from ``compute_logits``) with the audio codes the local transformer produced last step. """ @@ -1474,8 +1490,9 @@ class MossTTSLocalTalkerForGeneration(nn.Module): span_len = int(input_ids.shape[0]) audio_state = info_dict.get("audio_state") is_first_call = not isinstance(audio_state, dict) + is_prefill = bool(info_dict.get("_omni_is_prefill", False)) - if span_len > 1 or is_first_call: + if span_len > 1 or is_first_call or is_prefill: embeds = self.model.embed_tokens(input_ids) ref_codes = (info_dict.get("codes", {}) or {}).get("ref") @@ -1668,9 +1685,14 @@ class MossTTSLocalTalkerForGeneration(nn.Module): should_continue = codes.ne(self.audio_pad_token_id).any(dim=-1) should_values.append(should_continue.reshape(-1)[:1]) have_mtp_output = True - emitted_codes = codes[should_continue] - per_req_codes.append(emitted_codes) - have_codes = have_codes or emitted_codes.numel() > 0 + # Do not boolean-index a CUDA tensor here. Even for the + # common single-row case, ``codes[should_continue]`` must + # discover a dynamic output size on the host and therefore + # synchronizes the CUDA stream once per request. Keep the + # fixed-shape row (including an all-pad stop row); the async + # chunk processor already filters all-pad rows on CPU. + per_req_codes.append(codes) + have_codes = True else: should_values.append(torch.ones((1,), device=hidden.device, dtype=torch.bool)) per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py index 701f673a5..ef1abfa86 100644 --- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -8,14 +8,13 @@ in ``serving_speech.py``. The serving layer just calls :func:`encode_reference_codes` with its generic helpers (the audio resolver and the process-wide speaker cache). -Kept import-light (only ``asyncio`` / ``hashlib`` / ``torch`` plus the logger) +Kept import-light (only ``asyncio`` / ``torch`` plus the logger) so importing it from the API-server process does not pull the talker/codec. """ from __future__ import annotations import asyncio -import hashlib from collections.abc import Awaitable, Callable from typing import Any @@ -43,7 +42,7 @@ async def encode_reference_codes( ref_str: str, *, processor: Any, - resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int]]], + resolve_ref_audio: Callable[[str], Awaitable[tuple[list, int, str]]], speaker_cache: Any, variant: str, n_vq: int, @@ -62,10 +61,15 @@ async def encode_reference_codes( ``asyncio.to_thread`` so cold/anonymous encodes from concurrent requests overlap instead of serializing on the event loop. + For named voices the speaker cache is checked *before* resolving the audio, + avoiding the decode cost when the cache is warm. For anonymous references + the audio is resolved first so the cache key incorporates mtime/size from + a single stat (no TOCTOU window). + Args: ref_str: The raw reference audio (URL / path / data URL) as received. processor: Upstream MOSS processor exposing ``encode_audios_from_wav``. - resolve_ref_audio: Async callable mapping ``ref_str`` to ``(wav_list, sr)``. + resolve_ref_audio: Async callable mapping ``ref_str`` to ``(wav_list, sr, cache_key)``. speaker_cache: Process-wide ``SpeakerEmbeddingCache``. variant: MOSS sub-variant (``tts`` / ``ttsd`` / ...), namespaces the cache key. n_vq: Number of RVQ codebooks (also namespaces the cache key). @@ -76,24 +80,45 @@ async def encode_reference_codes( Returns: The reference RVQ codes tensor (CPU), ready to pass to the processor. """ + model_type = f"moss_tts_{variant}_nq{n_vq}" + + # ---- Named-voice branch: check cache *before* resolving ---- + # When voice_name is set, the cache key is (voice_name, created_at) which + # does not depend on the resolved audio content. The serving layer only + # passes voice_name for uploaded speakers without an inline ref_audio; + # re-upload bumps created_at and clears the speaker-cache slot. if voice_name: speaker_name = voice_name created_at = int(voice_created_at) + cache_key = speaker_cache.make_cache_key( + speaker_name, + model_type=model_type, + created_at=created_at, + ) + cached = speaker_cache.get(cache_key) + if cached is not None: + logger.debug("Speaker cache HIT for MOSS-TTS reference '%s'", speaker_name) + return cached["codes"].clone() + # Cache miss — resolve the audio for encoding. + wav_list, sr, _ = await resolve_ref_audio(ref_str) else: - speaker_name = "ref:" + hashlib.sha1((ref_str or "").encode("utf-8")).hexdigest() + # ---- Anonymous branch: resolve first for content-addressed key ---- + # The resolve incorporates mtime/size from a single stat so the key + # and the waveform always come from the same filesystem snapshot. + wav_list, sr, resolve_cache_key = await resolve_ref_audio(ref_str) + speaker_name = "ref:" + resolve_cache_key created_at = 0 + cache_key = speaker_cache.make_cache_key( + speaker_name, + model_type=model_type, + created_at=created_at, + ) + cached = speaker_cache.get(cache_key) + if cached is not None: + logger.debug("Speaker cache HIT for MOSS-TTS reference '%s'", speaker_name) + return cached["codes"].clone() - cache_key = speaker_cache.make_cache_key( - speaker_name, - model_type=f"moss_tts_{variant}_nq{n_vq}", - created_at=created_at, - ) - cached = speaker_cache.get(cache_key) - if cached is not None: - logger.debug("Speaker cache HIT for MOSS-TTS reference '%s'", speaker_name) - return cached["codes"].clone() - - wav_list, sr = await resolve_ref_audio(ref_str) + # ---- Cold miss: encode in a worker thread and store ---- codes = await asyncio.to_thread(_encode_wav_sync, processor, wav_list, sr, sr_target, n_vq) speaker_cache.put(cache_key, {"codes": codes.detach().cpu()}) logger.debug("Speaker cache STORE for MOSS-TTS reference '%s'", speaker_name) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py index b6f8531a8..8400f350c 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/asr/conformer_encoder.py @@ -830,6 +830,75 @@ class ConformerEncoder(NeuralModule, StreamingEncoder, Exportable, AccessMixin): self.streaming_cfg = streaming_cfg + def streaming_post_process(self, rets, keep_all_outputs=True): + """Apply NeMo's cache-aware output trimming contract.""" + if len(rets) == 2: + return rets[0], rets[1], None, None, None + + encoded, encoded_len, cache_channel, cache_time, cache_channel_len = rets + if cache_channel is not None and self.streaming_cfg.last_channel_cache_size >= 0: + if self.streaming_cfg.last_channel_cache_size > 0: + cache_channel = cache_channel[:, :, -self.streaming_cfg.last_channel_cache_size :, :] + + if self.streaming_cfg.valid_out_len > 0 and (not keep_all_outputs or self.att_context_style == "regular"): + encoded = encoded[:, :, : self.streaming_cfg.valid_out_len] + encoded_len = torch.clamp(encoded_len, max=self.streaming_cfg.valid_out_len) + + return encoded, encoded_len, cache_channel, cache_time, cache_channel_len + + def get_initial_cache_state( + self, + batch_size=1, + dtype=torch.float32, + device=None, + max_dim=0, + ): + """Create the channel/time caches used by cache-aware streaming.""" + if self.streaming_cfg is None: + self.setup_streaming_params() + if device is None: + device = next(self.parameters()).device + create_tensor = torch.randn if max_dim > 0 else torch.zeros + cache_last_channel = create_tensor( + ( + len(self.layers), + batch_size, + self.streaming_cfg.last_channel_cache_size, + self.d_model, + ), + device=device, + dtype=dtype, + ) + cache_last_time = create_tensor( + ( + len(self.layers), + batch_size, + self.d_model, + self.conv_context_size[0], + ), + device=device, + dtype=dtype, + ) + if max_dim > 0: + cache_last_channel_len = torch.randint( + 0, + min(max_dim, self.streaming_cfg.last_channel_cache_size), + (batch_size,), + device=device, + dtype=torch.int64, + ) + for i in range(batch_size): + cache_last_channel[:, i, cache_last_channel_len[i] :, :] = 0 + if cache_last_channel_len[i] == 0: + cache_last_time[:, i, :, :] = 0 + else: + cache_last_channel_len = torch.zeros( + batch_size, + device=device, + dtype=torch.int64, + ) + return cache_last_channel, cache_last_time, cache_last_channel_len + def change_subsampling_conv_chunking_factor(self, subsampling_conv_chunking_factor: int): """ Update the conv_chunking_factor (int) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py index f1c648acd..e89b67107 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/compat.py @@ -259,6 +259,51 @@ class StreamingEncoder(ABC): """ pass + @abstractmethod + def get_initial_cache_state(self, batch_size, dtype, device, max_dim): + pass + + def cache_aware_stream_step( + self, + processed_signal, + processed_signal_length=None, + cache_last_channel=None, + cache_last_time=None, + cache_last_channel_len=None, + keep_all_outputs=True, + drop_extra_pre_encoded=None, + ): + """Run one incremental encoder step using NeMo's cache contract.""" + if self.streaming_cfg is None: + self.setup_streaming_params() + if drop_extra_pre_encoded is not None: + prev_drop_extra_pre_encoded = self.streaming_cfg.drop_extra_pre_encoded + self.streaming_cfg.drop_extra_pre_encoded = drop_extra_pre_encoded + else: + prev_drop_extra_pre_encoded = None + + if processed_signal_length is None: + processed_signal_length = processed_signal.new_full( + (processed_signal.size(0),), + processed_signal.size(-1), + ) + + try: + encoder_output = self( + audio_signal=processed_signal, + length=processed_signal_length, + cache_last_channel=cache_last_channel, + cache_last_time=cache_last_time, + cache_last_channel_len=cache_last_channel_len, + ) + return self.streaming_post_process( + encoder_output, + keep_all_outputs=keep_all_outputs, + ) + finally: + if prev_drop_extra_pre_encoded is not None: + self.streaming_cfg.drop_extra_pre_encoded = prev_drop_extra_pre_encoded + # ============================================================================== # CacheAwareStreamingConfig — exact copy of the dataclass from diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py index e83e78cb6..2d7d6277f 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_model.py @@ -935,8 +935,8 @@ class GatedProjectedSumRMSNorm(nn.Module): audio_emb = audio_emb / self.num_codebooks # projections run in model dtype (BF16) - audio_h = self.audio_proj(audio_emb) - text_h = self.text_proj(text_emb) + audio_h = self.audio_proj(audio_emb.to(dtype=self.audio_proj.weight.dtype)) + text_h = self.text_proj(text_emb.to(dtype=self.text_proj.weight.dtype)) dtype = audio_h.dtype @@ -1302,7 +1302,7 @@ class RVQEARTTSModel(nn.Module): pre_bos_mask = bos_mask.cumsum(dim=1) == 0 # [B, T, 1] # Apply projection to model size - code_embed = self.embed_code(code_embed) + code_embed = self.embed_code(code_embed.to(dtype=self.embed_code.weight.dtype)) # Choose projection if self.config.get("audio_prompt_encoder_config", None): @@ -1336,7 +1336,7 @@ class RVQEARTTSModel(nn.Module): code_embeds = code_embed + bos_mask * self.bos_emb else: # Inference - code_embeds = self.embed_code(self.depthsum_embedding(code)) + code_embeds = self.embed_code(self.depthsum_embedding(code).to(dtype=self.embed_code.weight.dtype)) uncond_dec_flag = torch.zeros(code.size(0), 1, 1, device=code.device, dtype=torch.bool) if guidance_enabled: @@ -1591,7 +1591,7 @@ class RVQEARTTSModel(nn.Module): top_p_or_k_i = top_p_or_k[i] if top_p_or_k is not None else 1.0 noise_scale_i = noise_scale[i] if noise_scale is not None else 1.0 - mog_input_embeds = self.embed_code(self.depthsum_embedding(code)) + mog_input_embeds = self.embed_code(self.depthsum_embedding(code).to(self.embed_code.weight.dtype)) if self.config.random_target_masking: mog_input_embeds += self.embed_target_mask(cnt + k - 1) if guidance_scale_i > 0.0: diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py index 3ceaf93de..d53af34e2 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_code2wav.py @@ -24,6 +24,7 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import scalar_bool from vllm_omni.model_executor.models.output_templates import OmniOutput logger = init_logger(__name__) @@ -83,6 +84,7 @@ class NemotronVoiceChatCode2Wav(nn.Module): # Constructed in load_weights (owns the tts_model.audio_codec.* subtree). self.audio_codec: nn.Module | None = None + self._duplex_codec_caches: dict[str, object] = {} # ------------------------------------------------------------------ # Runner-facing placeholder hooks. @@ -147,6 +149,12 @@ class NemotronVoiceChatCode2Wav(nn.Module): meta = runtime_info.get("meta") if isinstance(runtime_info, Mapping) else None if isinstance(meta, Mapping) and meta.get("nemotron_voicechat_dummy_profile") is True: continue + codec_streaming = scalar_bool(meta.get("codec_streaming")) if isinstance(meta, Mapping) else False + if isinstance(runtime_info, Mapping): + codec_streaming = codec_streaming or scalar_bool(runtime_info.get("meta.codec_streaming")) + cache_request_id = meta.get("request_id") if isinstance(meta, Mapping) else None + if cache_request_id is None and isinstance(runtime_info, Mapping): + cache_request_id = runtime_info.get("meta.request_id") codes = self._codes_from_runtime_info(runtime_info) if codes is None: # No input_ids fallback (unlike PersonaPlexCode2Wav): this @@ -159,6 +167,12 @@ class NemotronVoiceChatCode2Wav(nn.Module): "code stack (placeholder input_ids are never decoded)." ) codes = validate_code_stack(codes, self._num_quantizers, self._codebook_size).to(device=device) + logger.debug( + "Nemotron VoiceChat code2wav input: request=%s shape=%s codec_streaming=%s", + cache_request_id, + tuple(codes.shape), + codec_streaming, + ) frames = codes.shape[0] if frames == 0: continue @@ -176,19 +190,42 @@ class NemotronVoiceChatCode2Wav(nn.Module): if left_context_frames >= frames: continue # terminal empty chunk: all frames already emitted lens = torch.tensor([frames], device=device, dtype=torch.long) + codec_cache = None + if codec_streaming: + if not isinstance(cache_request_id, str) or not cache_request_id: + raise ValueError("Nemotron VoiceChat incremental codec input lacks meta.request_id") + from vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored.ear_tts_vae_codec import ( + CausalConv1dCache, + ) + + codec_cache = self._duplex_codec_caches.setdefault(cache_request_id, CausalConv1dCache()) # NeMo decodes the codec strictly in fp32. with torch.autocast(device_type=device.type, enabled=False): - wav, wav_len = self.audio_codec.decode(codes.unsqueeze(0), lens) + wav, wav_len = self.audio_codec.decode( + codes.unsqueeze(0), + lens, + cache=codec_cache, + ) wav = wav.reshape(-1)[: int(wav_len.reshape(-1)[0])] - if left_context_frames > 0: + if left_context_frames > 0 and not codec_streaming: wav = wav[left_context_frames * self._wav_per_frame :] audios[i] = wav.detach().to(dtype=torch.float32).cpu() + logger.debug( + "Nemotron VoiceChat code2wav output: request=%s frames=%s samples=%s", + cache_request_id, + frames, + int(audios[i].numel()), + ) return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, ) + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + for request_id in finished_req_ids: + self._duplex_codec_caches.pop(str(request_id), None) + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput: if isinstance(model_outputs, OmniOutput): return model_outputs diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py index f4f98ab00..515b29c6c 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_talker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage-1 talker for NemotronVoiceChat: frame-locked text -> 31-quantizer codes. ``LLM_AR`` stage driven per frame by the vLLM engine, with the vendored NeMo @@ -46,6 +46,7 @@ from vllm.logger import init_logger from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import ( merge_runtime_info, require_request_id, + scalar_bool, ) from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -164,6 +165,7 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] codes_list: list[torch.Tensor] = [] prompt_len: int | None = None + codec_streaming = False for info in info_dicts: if not isinstance(info, dict): continue @@ -177,16 +179,32 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): reported = info.get("meta.nvc_logical_prompt_len") if reported is not None: prompt_len = int(reported) + streaming = meta.get("codec_streaming") if isinstance(meta, dict) else None + if streaming is None: + streaming = info.get("meta.codec_streaming") + codec_streaming = codec_streaming or scalar_bool(streaming) if not codes_list: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) - outputs: dict[str, Any] = {"codes": {"audio": torch.cat(codes_list, dim=0)}} + codes = torch.cat(codes_list, dim=0) + logger.debug( + "Nemotron VoiceChat talker emitted codes: shapes=%s combined=%s prompt_len=%s", + [tuple(item.shape) for item in codes_list], + tuple(codes.shape), + prompt_len, + ) + outputs: dict[str, Any] = {"codes": {"audio": codes}} # The prompt length must ride the WIRE payload (multimodal_outputs): # the async-chunk dispatcher hands exactly this dict to the # talker2code2wav producer, while request.additional_information is # overwritten concurrently by arriving thinker chunks. + output_meta: dict[str, Any] = {} if prompt_len is not None: # 1-D on purpose: the wire-payload builder indexes element.shape[0]. - outputs["meta"] = {"nvc_logical_prompt_len": torch.tensor([prompt_len], dtype=torch.long)} + output_meta["nvc_logical_prompt_len"] = torch.tensor([prompt_len], dtype=torch.long) + if codec_streaming: + output_meta["codec_streaming"] = torch.tensor([True], dtype=torch.bool) + if output_meta: + outputs["meta"] = output_meta return OmniOutput( text_hidden_states=hidden, multimodal_outputs=outputs, @@ -269,8 +287,9 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): # offline_inference lines around get_init_inputs/first forward). self.tts.set_init_inputs(speaker_name=self._speaker_name) init_inputs = self.tts.get_init_inputs(B=1) - generation_config = self.tts._get_generation_config(self._guidance_enabled()) - init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": self._guidance_enabled()}) + guidance_enabled = self._guidance_enabled() + generation_config = self.tts._get_generation_config(guidance_enabled) + init_inputs.update({"use_cache": True, "past_key_values": None, "guidance_enabled": guidance_enabled}) with torch.inference_mode(): outputs = self.tts.tts_model(**init_inputs) prompt_len = self._logical_prompt_len(info) @@ -281,6 +300,10 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): "meta.num_processed_tokens); the code2wav producer cannot trim the " "prompt-region codes without it." ) + meta = info.get("meta") + codec_streaming = scalar_bool(meta.get("codec_streaming")) if isinstance(meta, dict) else False + codec_streaming = codec_streaming or scalar_bool(info.get("meta.codec_streaming")) + codec_request_id = meta.get("request_id") if isinstance(meta, dict) else None session = { "timeline": timeline, "prompt_len": prompt_len, @@ -288,9 +311,13 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): "past_key_values": outputs.past_key_values, "code": init_inputs["code"][:, -1:], "first_context_subword_id": init_inputs["subword_ids"][:, -1].unsqueeze(-1), + "current_subword_mask": torch.ones(1, 1, device=device, dtype=torch.bool), "tts_initialized": False, + "guidance_enabled": guidance_enabled, "generation_config": generation_config, "codes_rows": [], + "codec_streaming": codec_streaming, + "codec_request_id": codec_request_id, # Sync mode ships the whole timeline up front via # nvc_text_timeline; async-chunk mode grows it across chunks and # the last chunk sets meta.finished. @@ -333,25 +360,27 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): session["tts_initialized"] = True else: prev_subword_id = timeline[t - 1].reshape(1, 1) - current_subword_mask = torch.ones(1, 1, device=timeline.device, dtype=torch.bool) with torch.inference_mode(): code, past_key_values = self.tts.infer_codes_one_step( current_subword_id=current_subword_id, prev_subword_id=prev_subword_id, - current_subword_mask=current_subword_mask, + current_subword_mask=session["current_subword_mask"], prev_audio_tokens=session["code"], past_key_values=session["past_key_values"], - guidance_enabled=self._guidance_enabled(), + guidance_enabled=session["guidance_enabled"], generation_config=session["generation_config"], ignore_eos_flag_stop=True, ) session["code"] = code session["past_key_values"] = past_key_values session["step"] = t + 1 - # Async mode: exhausting the received timeline only ends the request - # once the upstream marked its final chunk; otherwise more entries are - # in flight and the scheduler parks us until they arrive. - finished = (t + 1) >= int(session["timeline"].numel()) and bool(session["upstream_finished"]) + # Native codec streaming uses one resumable scheduler segment per + # received timeline suffix. The pre-existing offline async-chunk path, + # however, must stay alive until the thinker sends its terminal marker; + # stopping at an intermediate cumulative snapshot starves Stage 2. + finished = (t + 1) >= int(session["timeline"].numel()) and ( + bool(session["codec_streaming"]) or bool(session["upstream_finished"]) + ) return code.reshape(1, -1).to(torch.long), finished def preprocess( @@ -363,14 +392,13 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): info = merge_runtime_info(info_dict) device = input_ids.device span = int(input_ids.shape[0]) - is_prefill_raw = info.get("_omni_is_prefill") - is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1 request_id = require_request_id(info, "talker") embeds = torch.zeros((span, self._hidden), device=device, dtype=self._dtype) - if is_prefill or request_id not in self._sessions: - # The 1-token vLLM prefill is the session boundary: run the vendored - # warmup here so the first decode step is a pure NeMo t=1 step. + if request_id not in self._sessions: + # The request id, not vLLM's per-segment prefill flag, is the + # session boundary. Resumable updates are prefill again in v0.28 + # and must preserve the existing EAR-TTS cache and step. self._init_session(request_id, info, device) return input_ids, embeds, {} @@ -388,21 +416,34 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): # chunk each. finished = False steps_run = 0 + new_code_rows: list[torch.Tensor] = [] + codec_streaming = bool(session["codec_streaming"]) while int(session["step"]) < int(session["timeline"].numel()): codes_row, finished = self._step_session(session) - session["codes_rows"].append(codes_row.cpu()) + if codec_streaming: + # One D2H copy per wake is enough; avoid synchronizing after every drained TTS step. + new_code_rows.append(codes_row) + else: + session["codes_rows"].append(codes_row.cpu()) steps_run += 1 if session["sync_mode"]: break if steps_run == 0: # Zero-progress wake (duplicate/terminal-marker chunk): no TTS step - # was run, so emit a plain CONTINUE placeholder — never a - # fabricated frame. The request finishes only when the upstream is - # exhausted AND the timeline is fully consumed. - finished = bool(session["upstream_finished"]) + # was run, so never fabricate a frame. Stop the current scheduler + # segment only for native codec streaming, or once the ordinary + # offline async producer has sent its terminal marker. + finished = int(session["step"]) >= int(session["timeline"].numel()) and ( + bool(session["codec_streaming"]) or bool(session["upstream_finished"]) + ) if finished: embeds[:, 0] = 1.0 return input_ids, embeds, {} + if codec_streaming: + # The code2wav producer owns suffix slicing and therefore requires + # cumulative history on every wake. Keep the history on CPU and do + # one D2H transfer for all rows produced by this wake. + session["codes_rows"].append(torch.cat(new_code_rows, dim=0).cpu()) cumulative = torch.cat(session["codes_rows"], dim=0) if finished: embeds[:, 0] = 1.0 # stop flag -> compute_logits emits the stop token @@ -413,7 +454,12 @@ class NemotronVoiceChatTalkerForConditionalGeneration(nn.Module): # chunks and by these model updates, so the async code2wav producer # reads the prompt length from the per-step multimodal_output # instead of racing the request-level dict. - "meta": {"nvc_tts_step": int(session["step"]), "nvc_logical_prompt_len": int(session["prompt_len"])}, + "meta": { + "nvc_tts_step": int(session["step"]), + "nvc_logical_prompt_len": int(session["prompt_len"]), + "codec_streaming": bool(session["codec_streaming"]), + "request_id": session["codec_request_id"], + }, } return input_ids, embeds, info_update diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py index 9f237a885..01fbed1f9 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/nemotron_voicechat_thinker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage-0 thinker for NemotronVoiceChat: 16 kHz speech -> frame-locked text. ``LLM_AR`` stage combining: @@ -12,7 +12,8 @@ * the NemotronH (hybrid Mamba) backbone on vLLM paged execution via ``init_vllm_registered_model`` — this is where vLLM earns its keep; * the separate ``lm_head`` (text channel, sampled by the vLLM engine) and - ``function_head`` (greedy, maintained internally per request). + ``function_head`` (greedy, maintained internally per request, with client + tool results injected into the function channel). Timeline contract (verified against NeMo ``_init_inference``/``_step_zero``/ ``_step_inference``; the golden shape test guards it): @@ -29,13 +30,14 @@ Timeline contract (verified against NeMo ``_init_inference``/``_step_zero``/ example sets ``max_tokens = acoustic_frame_count``. Never set ``min_tokens``: the tokenizer's eos_token is — the frame-locked PAD/silence token — and min_tokens masks "EOS" to -inf, forbidding silence entirely. -* The function channel exists for numerical correctness (fusion weight 2.0) - even though tool calling is out of scope; the greedy token timeline is - exposed for debugging via the request info. +* The function channel is load-bearing for numerical correctness (fusion + weight 2.0), emits tool calls, and accepts client tool results while keeping + the text channel padded during forced response injection. """ from __future__ import annotations +import copy import os from collections.abc import Iterable from typing import Any @@ -47,7 +49,10 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import HasInnerState, IsHybrid from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRow from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import ( merge_runtime_info, require_request_id, @@ -59,6 +64,49 @@ logger = init_logger(__name__) _LM_ARCHITECTURE = "NemotronHForCausalLM" +def slice_perception_streaming_mel( + processed_signal: torch.Tensor, + frame_idx: int, + streaming_cfg: Any, +) -> tuple[torch.Tensor, int]: + """Slice one 80 ms mel chunk using NeMo's cache-manager geometry.""" + + def _at(value: Any, index: int) -> int: + return int(value[index] if isinstance(value, list | tuple) else value) + + chunk_size_first = _at(streaming_cfg.chunk_size, 0) + chunk_size = _at(streaming_cfg.chunk_size, 1) + shift_size_first = _at(streaming_cfg.shift_size, 0) + shift_size = _at(streaming_cfg.shift_size, 1) + pre_encode_first = _at(streaming_cfg.pre_encode_cache_size, 0) + pre_encode_size = _at(streaming_cfg.pre_encode_cache_size, 1) + + if frame_idx == 0: + mel_chunk = processed_signal[:, :, :chunk_size_first] + if pre_encode_first > 0: + mel_chunk = torch.nn.functional.pad(mel_chunk, (pre_encode_first, 0)) + return mel_chunk, 0 + + chunk_start = shift_size_first + (frame_idx - 1) * shift_size + chunk_end = chunk_start + chunk_size + # Match NeMo's tail-boundary rule. With one 80 ms frame per call this + # normally leaves the expected trailing STFT columns unused. + offset = chunk_size - shift_size_first + if chunk_end > processed_signal.shape[-1] - offset: + chunk_end = processed_signal.shape[-1] - offset + chunk_start = chunk_end - chunk_size + main_chunk = processed_signal[:, :, chunk_start:chunk_end] + cache_start = max(0, chunk_start - pre_encode_size) + cache_mel = processed_signal[:, :, cache_start:chunk_start] + if cache_mel.shape[-1] < pre_encode_size: + cache_mel = torch.nn.functional.pad( + cache_mel, + (pre_encode_size - cache_mel.shape[-1], 0), + ) + mel_chunk = torch.cat([cache_mel, main_chunk], dim=-1) + return mel_chunk, int(streaming_cfg.drop_extra_pre_encoded) + + def validate_timeline_fits_max_model_len( logical_prompt_token_len: int, acoustic_frame_count: int, @@ -210,6 +258,14 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, from omegaconf import DictConfig self.perception = AudioPerceptionModule(DictConfig(perception_cfg)) + # NeMo's streaming wrapper deliberately uses a second, deterministic + # preprocessor. The training/offline frontend has dither enabled and + # pads the mel time axis; both are incompatible with exact incremental + # chunk boundaries. + streaming_preprocessor_cfg = copy.deepcopy(self.perception.cfg.preprocessor) + streaming_preprocessor_cfg.dither = 0.0 + streaming_preprocessor_cfg.pad_to = 0 + self._streaming_preprocessor = self.perception.from_config_dict(streaming_preprocessor_cfg) # AddFusion weights (the checkpoint uses fuse_method="add" defaults; # the same duplex_*_weight keys double as the fusion weights). @@ -235,6 +291,8 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, # Omni AR runner contract. self.have_multimodal_outputs = True self.has_preprocess = True + self.prefer_model_sampler = True + self.omni_client_multimodal_output_keys = ("nvc_function_token",) self.has_postprocess = True # per-step greedy function-channel token self.requires_full_prefix_cached_hidden_states = False # The per-step function token is a scalar; no GPU-resident buffers needed. @@ -251,6 +309,8 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, # request_id -> per-request state (audio frames, prefill embeds, counters). self._sessions: dict[str, dict[str, Any]] = {} + self._duplex_sampling_rows: dict[int, str] = {} + self._duplex_previous_text_tokens: dict[str, int] = {} # ------------------------------------------------------------------ # Special token ids (from the checkpoint's token strings). @@ -306,7 +366,57 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: if isinstance(model_outputs, OmniOutput): return model_outputs - return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs={}) + multimodal_outputs: dict[str, torch.Tensor] = {} + if self._use_function_head and model_outputs.numel(): + with torch.inference_mode(): + function_token = self.function_head(model_outputs[-1:, :].to(self._dtype)).argmax(dim=-1) + info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") + info = info_dicts[0] if isinstance(info_dicts, list) and len(info_dicts) == 1 else None + request_id = info.get("request_id") if isinstance(info, dict) else None + session = self._sessions.get(request_id) if isinstance(request_id, str) else None + forced_token = session.get("forced_function_token") if isinstance(session, dict) else None + if isinstance(forced_token, int): + function_token = torch.full_like(function_token, forced_token) + multimodal_outputs["nvc_function_token"] = function_token + return OmniOutput( + text_hidden_states=model_outputs, + multimodal_outputs=multimodal_outputs, + ) + + def prepare_duplex_sampling( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + rows: tuple[DuplexSamplingRow, ...], + ) -> None: + del logits, sampling_metadata + self._duplex_sampling_rows = {row.row_idx: row.request_id for row in rows} + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> SamplerOutput | None: + del sampling_metadata + rows = self._duplex_sampling_rows + if not rows: + return None + expected_rows = set(range(int(logits.shape[0]))) + if set(rows) != expected_rows: + raise RuntimeError( + "NemotronVoiceChat native-duplex sampling requires every active batch row " + "to be a duplex request; mixed duplex/non-duplex batches are unsupported." + ) + sampled = logits.argmax(dim=-1).to(dtype=torch.int32) + for row_idx, request_id in rows.items(): + session = self._sessions.get(request_id) + if isinstance(session, dict) and isinstance(session.get("forced_function_token"), int): + sampled[row_idx] = int(session["nvc_text_pad_id"]) + self._duplex_previous_text_tokens[request_id] = int(sampled[row_idx].item()) + return SamplerOutput( + sampled_token_ids=sampled.unsqueeze(-1), + logprobs_tensors=None, + ) # ------------------------------------------------------------------ # Fusion helpers. @@ -324,6 +434,263 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, fused = fused + self.embed_tokens(func_ids) * self._w_func return fused + @staticmethod + def _duplex_info(info: dict[str, Any]) -> dict[str, Any] | None: + duplex = info.get("duplex") + if isinstance(duplex, dict) and duplex.get("data_plane") is True: + return duplex + return None + + def _duplex_stable_frame( + self, + session: dict[str, Any], + duplex: dict[str, Any], + device: torch.device, + ) -> torch.Tensor: + try: + source_input_seq = int(duplex["source_input_seq"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Nemotron VoiceChat duplex input lacks source_input_seq") from exc + if source_input_seq == session.get("last_input_seq"): + return session["duplex_frame"] + last_input_seq = int(session.get("last_input_seq", 0)) + if source_input_seq < last_input_seq: + raise ValueError( + f"Nemotron VoiceChat duplex input sequence moved backwards: {source_input_seq} < {last_input_seq}" + ) + + from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import ( + decode_pcm_f32le, + ) + + raw = decode_pcm_f32le(duplex.get("payload"), exact_frame=True) + frame_audio = torch.from_numpy(np.frombuffer(raw, dtype=" tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + device = input_ids.device + runtime_config = duplex.get("runtime_config") + if not isinstance(runtime_config, dict): + raise ValueError("Nemotron VoiceChat duplex input lacks runtime_config") + try: + source_input_seq = int(duplex["source_input_seq"]) + pad_id = int(runtime_config["nvc_text_pad_id"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Nemotron VoiceChat duplex runtime token metadata is incomplete") from exc + + session = self._sessions.get(request_id) + if session is None: + prompt_ids_raw = runtime_config.get("nvc_prompt_token_ids") + if not isinstance(prompt_ids_raw, list | tuple) or not prompt_ids_raw: + raise ValueError("Nemotron VoiceChat duplex runtime requires nvc_prompt_token_ids") + prompt_ids = torch.as_tensor( + [int(token_id) for token_id in prompt_ids_raw], + dtype=torch.long, + device=device, + ) + session = { + "prompt_len": int(prompt_ids.numel()), + "func_token": pad_id, + "last_input_seq": 0, + "nvc_text_pad_id": pad_id, + "function_response_generation": 0, + "forced_function_tokens": [], + } + self._sessions[request_id] = session + frame = self._duplex_stable_frame(session, duplex, device) + timeline = torch.cat( + [self.embed_tokens(prompt_ids).to(self._dtype), frame], + dim=0, + ) + pad_ids = torch.full( + (timeline.shape[0],), + pad_id, + dtype=torch.long, + device=device, + ) + session["prefill_embeds"] = self._fuse(pad_ids, timeline, pad_ids) + else: + frame = self._duplex_stable_frame(session, duplex, device) + + self._sync_forced_function_response(session, runtime_config) + + if source_input_seq <= 1: + offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0)) + prefill = session["prefill_embeds"] + span = int(input_ids.shape[0]) + if offset + span > int(prefill.shape[0]): + raise ValueError("Nemotron VoiceChat duplex first append exceeds its fused prompt") + return ( + input_ids, + prefill[offset : offset + span].to(self._dtype), + { + "nvc_text_pad_id": pad_id, + }, + ) + + if int(input_ids.shape[0]) != 1: + raise ValueError("Nemotron VoiceChat duplex continuation must schedule exactly one frame token") + previous_function = info.get("nvc_prev_function_token") + if previous_function is not None: + if isinstance(previous_function, torch.Tensor): + previous_function = previous_function.reshape(-1)[-1].item() + session["func_token"] = int(previous_function) + previous_text = self._duplex_previous_text_tokens.get(request_id) + if previous_text is None: + raise RuntimeError("Nemotron VoiceChat duplex continuation has no previous sampled text token") + text_id = torch.tensor([previous_text], dtype=torch.long, device=device) + function_id = torch.tensor( + [int(session["func_token"])], + dtype=torch.long, + device=device, + ) + fused = self._fuse(text_id, frame, function_id) + return ( + input_ids, + fused, + { + "meta": {"nvc_frame": source_input_seq - 1}, + "nvc_text_pad_id": pad_id, + }, + ) + + @staticmethod + def _sync_forced_function_response( + session: dict[str, Any], + runtime_config: dict[str, Any], + ) -> None: + try: + generation = int(runtime_config.get("nvc_function_response_generation", 0)) + except (TypeError, ValueError) as exc: + raise ValueError("Nemotron VoiceChat function-response generation must be an integer") from exc + seen = int(session.get("function_response_generation", 0)) + if generation < seen: + raise ValueError("Nemotron VoiceChat function-response generation moved backwards") + if generation > seen: + raw_batches = runtime_config.get("nvc_function_response_batches") + batches = raw_batches if isinstance(raw_batches, list) else [] + unseen: list[tuple[int, list[int]]] = [] + for raw_batch in batches: + if not isinstance(raw_batch, dict): + raise ValueError("Nemotron VoiceChat function response batch must be an object") + try: + batch_generation = int(raw_batch["generation"]) + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Nemotron VoiceChat function response batch requires a generation") from exc + if batch_generation <= seen: + continue + raw_tokens = raw_batch.get("token_ids") + if not isinstance(raw_tokens, list | tuple) or not raw_tokens: + raise ValueError("Nemotron VoiceChat function response requires token ids") + unseen.append((batch_generation, [int(token_id) for token_id in raw_tokens])) + if not unseen: + # Backward-compatible single-response payload used by older + # serving adapters and serialized test fixtures. + raw_tokens = runtime_config.get("nvc_function_response_token_ids") + if not isinstance(raw_tokens, list | tuple) or not raw_tokens: + raise ValueError("Nemotron VoiceChat function response requires token ids") + unseen = [(generation, [int(token_id) for token_id in raw_tokens])] + expected = seen + 1 + queue = session.setdefault("forced_function_tokens", []) + if not isinstance(queue, list): + raise ValueError("Nemotron VoiceChat forced function-token queue is invalid") + for batch_generation, token_ids in unseen: + if batch_generation != expected: + raise ValueError("Nemotron VoiceChat function-response generations are not contiguous") + queue.extend(token_ids) + expected += 1 + if expected - 1 != generation: + raise ValueError("Nemotron VoiceChat function-response batches are incomplete") + session["function_response_generation"] = generation + queue = session.get("forced_function_tokens") + if session.get("forced_function_token") is None and isinstance(queue, list) and queue: + session["forced_function_token"] = int(queue[0]) + # ------------------------------------------------------------------ # Per-request session. # ------------------------------------------------------------------ @@ -418,6 +785,14 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, is_prefill = bool(is_prefill_raw) if isinstance(is_prefill_raw, bool) else span > 1 request_id = require_request_id(info, "thinker") + duplex = self._duplex_info(info) + if duplex is not None: + return self._preprocess_duplex( + request_id=request_id, + input_ids=input_ids, + info=info, + duplex=duplex, + ) if is_prefill or request_id not in self._sessions: session = self._sessions.get(request_id) or self._init_session(request_id, info, device) offset = max(0, int(info.get("_omni_num_computed_tokens", 0) or 0)) @@ -501,9 +876,25 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, prompt_len = int(kwargs.get("_omni_prompt_len", 0) or 0) if prompt_len and computed + int(hidden_states.shape[0]) < prompt_len: return {} # intermediate prefill chunk: frame 0 not reached yet - with torch.inference_mode(): - func_logits = self.function_head(hidden_states[-1, :].reshape(1, -1).to(self._dtype)) - token = func_logits.argmax(dim=-1).reshape(()).detach() + multimodal_outputs = kwargs.get("multimodal_outputs") + token = multimodal_outputs.get("nvc_function_token") if isinstance(multimodal_outputs, dict) else None + if not isinstance(token, torch.Tensor) or token.numel() == 0: + with torch.inference_mode(): + func_logits = self.function_head(hidden_states[-1, :].reshape(1, -1).to(self._dtype)) + token = func_logits.argmax(dim=-1) + token = token.reshape(-1)[-1].reshape(()).detach() + request_id = kwargs.get("request_id") + if isinstance(request_id, str): + session = self._sessions.get(request_id) + if session is not None: + # StreamingUpdate replaces the runner payload at each append; + # retain the function channel in model-owned session state. + session["func_token"] = int(token.item()) + forced = session.pop("forced_function_token", None) + if isinstance(forced, int): + queue = session.get("forced_function_tokens") + if not isinstance(queue, list) or not queue or queue.pop(0) != forced: + raise RuntimeError("Nemotron VoiceChat forced function-token queue lost alignment") update: dict[str, Any] = {"nvc_prev_function_token": token} if os.environ.get("NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE", "0") == "1": existing = kwargs.get("nvc_function_tokens") @@ -516,6 +907,7 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: for request_id in finished_req_ids: self._sessions.pop(str(request_id), None) + self._duplex_previous_text_tokens.pop(str(request_id), None) # ------------------------------------------------------------------ # Weights. @@ -581,6 +973,9 @@ class NemotronVoiceChatThinkerForConditionalGeneration(nn.Module, HasInnerState, ) device = self.vllm_config.device_config.device self.perception.to(device=device, dtype=self._dtype).eval() + # Keep the deterministic mel frontend in fp32, matching NeMo's + # PerceptionCacheManager. Its output is cast only at the encoder edge. + self._streaming_preprocessor.to(device=device, dtype=torch.float32).eval() self.embed_tokens.to(device=device, dtype=self._dtype) self.function_head.to(device=device, dtype=self._dtype) loaded |= {f"perception.{n}" for n in perception_state} diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py index 4f4834740..1bb1e19b8 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/pipeline.py @@ -31,6 +31,13 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.nemotron_voicechat" NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig( model_type="nemotron_voicechat", model_arch="NemotronVoiceChatThinkerForConditionalGeneration", + duplex_runtime_extension=( + "vllm_omni.experimental.fullduplex.nemotron_voicechat.runtime.NemotronVoiceChatDuplexRuntimeExtension" + ), + duplex_serving_adapter=( + "vllm_omni.experimental.fullduplex.nemotron_voicechat.serving_adapter.NemotronVoiceChatServingRuntimeAdapter" + ), + duplex_control_enabled=True, # Named after the alias key so the deploy-yaml stem substring-matches the # checkpoint directory name (NVIDIA-NemotronLabs-VoiceChat-11B) during # bare-path auto-detection. @@ -80,6 +87,7 @@ NEMOTRON_VOICECHAT_PIPELINE = PipelineConfig( engine_output_type="audio", sync_process_input_func=f"{_PROC}.talker2code2wav_token_only", sampling_constraints={"detokenize": False}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py index 7f5a897d2..9d8508329 100644 --- a/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py +++ b/vllm_omni/model_executor/models/nemotron_voicechat/runtime_info.py @@ -6,6 +6,27 @@ from __future__ import annotations from typing import Any +import torch + + +def scalar_bool(value: object) -> bool: + """Normalize a Python or tensor scalar to bool. + + Inter-stage scalar metadata is tensorized while it rides the engine output + channel, so model consumers must not rely on object identity with ``True``. + Empty tensors mean "unset"; non-scalar tensors are rejected instead of + silently choosing an element. + """ + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return False + if value.numel() != 1: + raise ValueError(f"Expected scalar boolean metadata, got tensor shape {tuple(value.shape)}.") + value = value.item() + elif hasattr(value, "item"): + value = value.item() + return bool(value) + def merge_runtime_info(info_dict: dict[str, Any]) -> dict[str, Any]: """Flatten ``additional_information`` into the info dict (top-level wins).""" diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index de79e1aae..2036c42aa 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py @@ -47,6 +47,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig( "detokenize": True, "stop_token_ids": [8294], }, + requires_full_payload_input=True, ), StagePipelineConfig( stage_id=2, @@ -59,6 +60,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig( engine_output_type="audio", sync_process_input_func=f"{_PROC}.talker2code2wav_token_only", sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index ce6617695..2f653b12f 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import glob import os from collections.abc import Iterable @@ -19,9 +22,6 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, ) from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix -from vllm.model_executor.models.vision import ( - get_llm_pos_ids_for_vision, -) # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav from vllm.multimodal import MULTIMODAL_REGISTRY @@ -33,6 +33,9 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.layers.rotary_embedding.mrope import ( + OmniMRotaryEmbedding, +) from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( @@ -43,6 +46,11 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges from vllm_omni.platforms import current_omni_platform +# `get_llm_pos_ids_for_vision` was removed from upstream vLLM (commit +# 0cf49a5d15); vllm-omni keeps an identical copy as a static method on +# OmniMRotaryEmbedding. Bind it at module level to keep call sites unchanged. +get_llm_pos_ids_for_vision = OmniMRotaryEmbedding._get_llm_pos_ids_for_vision + TALKER_CODEC_EOS_TOKEN_ID = 8294 TALKER_CODEC_BOS_TOKEN_ID = 8293 diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index c93e764ef..1e515d0dd 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -59,6 +59,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( "detokenize": False, "stop_token_ids": [2150], }, + requires_full_payload_input=True, ), StagePipelineConfig( stage_id=2, @@ -70,6 +71,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( hf_config_name="thinker_config", engine_output_type="audio", sampling_constraints={"detokenize": True}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/quantization.py b/vllm_omni/model_executor/models/qwen3_omni/quantization.py new file mode 100644 index 000000000..26fe46956 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/quantization.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Quantization mapping helpers for nested Qwen3-Omni stages.""" + +from vllm.model_executor.models.interfaces import SupportsQuant + +_OUTER_MAPPING_APPLIED = "_qwen3_omni_outer_mapping_applied" + + +def apply_outer_quant_config_mapping(model: SupportsQuant) -> None: + """Map checkpoint paths once at the unified model boundary.""" + SupportsQuant._maybe_apply_model_mapping(model) + if model.quant_config is not None: + setattr(model.quant_config, _OUTER_MAPPING_APPLIED, True) + + +def apply_nested_quant_config_mapping(model: SupportsQuant) -> None: + """Avoid applying a stage weight mapper twice to a mapped config.""" + quant_config = model.quant_config + if quant_config is None: + return + if not getattr(quant_config, _OUTER_MAPPING_APPLIED, False): + SupportsQuant._maybe_apply_model_mapping(model) + return + + # The outer model maps names, while the nested model still contributes + # fused-module metadata needed by quantization scheme matching. + if packed_modules_mapping := getattr(model, "packed_modules_mapping", None): + quant_config.packed_modules_mapping.update(packed_modules_mapping) + + +class Qwen3OmniNestedSupportsQuant(SupportsQuant): + """SupportsQuant variant for models initialized inside the stage wrapper.""" + + def _maybe_apply_model_mapping(self) -> None: + apply_nested_quant_config_mapping(self) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index ab524c578..71e2408d7 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -20,12 +20,22 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( from vllm.config import ModelConfig, VllmConfig from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger -from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP, SupportsRealtime +from vllm.model_executor.models.interfaces import ( + SupportsMRoPE, + SupportsMultiModal, + SupportsPP, + SupportsQuant, + SupportsRealtime, +) from vllm.model_executor.models.qwen3_asr_realtime import Qwen3ASRRealtimeBuffer from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeConditionalGenerationMixin, ) -from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.model_executor.models.utils import ( + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.sequence import IntermediateTensors @@ -39,6 +49,9 @@ from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload from vllm_omni.metrics import definitions as defs from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + apply_outer_quant_config_mapping, +) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerDummyInputsBuilder, Qwen3OmniMoeThinkerForConditionalGeneration, @@ -86,6 +99,7 @@ class Qwen3OmniMoeForConditionalGeneration( CustomProcessMixin, SupportsMRoPE, SupportsRealtime, + SupportsQuant, ): """ Unified Qwen3 Omni MoE model combining thinker, talker, and code2wav. @@ -99,6 +113,24 @@ class Qwen3OmniMoeForConditionalGeneration( Set `model_stage` in vllm_config to one of: "thinker", "talker", "code2wav" """ + # vLLM applies quantization-config name mapping before constructing this + # outer stage wrapper. Expose the final module paths here so checkpoint + # ignore lists (for example, the BF16 MoE routers in compressed-tensors + # checkpoints) match the nested thinker/talker modules at construction + # time. The keys intentionally omit a trailing dot so exact module names + # such as ``thinker.lm_head`` are mapped along with their parameters. + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "thinker.lm_head": "thinker.language_model.lm_head", + "thinker.model": "thinker.language_model.model", + "talker.model": "talker.language_model.model", + } + ) + packed_modules_mapping = Qwen3OmniMoeThinkerForConditionalGeneration.packed_modules_mapping + + def _maybe_apply_model_mapping(self) -> None: + apply_outer_quant_config_mapping(self) + realtime_max_tokens = 64 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 3a85599d6..5c3129be9 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -28,11 +28,14 @@ from vllm.model_executor.models.utils import ( # type: ignore ) from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + Qwen3OmniNestedSupportsQuant, +) logger = init_logger(__name__) -class Qwen3OmniMoeCode2Wav(nn.Module): +class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): """ Qwen3 Omni MoE Code2Wav - Converts num_quantizers-layer RVQ codec codes to audio waveform. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 4b27e4ba0..8337f8dad 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -18,6 +18,9 @@ from vllm.model_executor.models.utils import ( ) from vllm.sequence import IntermediateTensors +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + Qwen3OmniNestedSupportsQuant, +) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import ( Qwen3OmniMoeTalkerCodePredictor, ) @@ -32,6 +35,7 @@ logger = init_logger(__name__) class Qwen3OmniMoeTalkerForConditionalGeneration( nn.Module, SupportsPP, + Qwen3OmniNestedSupportsQuant, ): """ Qwen3 Omni MoE Talker - Converts text to audio codec codes. @@ -58,8 +62,10 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( "talker.model.": "language_model.model.", # Codec head remains separate (outputs audio codes, not text) "talker.codec_head.": "codec_head.", - # Code predictor: Now matches HF structure exactly (has .model sub-module) - # e.g., "talker.code_predictor.model.codec_embedding.0" → "code_predictor.model.codec_embedding.0" + # Code predictor: keep the subtree under ``code_predictor.``. + # ``CodePredictorWrapper.load_weights`` accepts both current + # ``model.*`` body names and legacy body-direct names such as + # ``layers.*`` / ``codec_embedding.*``. "talker.code_predictor.": "code_predictor.", # Projection layers "talker.text_projection.": "text_projection.", diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 623c799d9..bbfea498b 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -126,6 +126,9 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( _get_video_second_per_grid_t, _presampled_videos_hf_kwargs, ) +from vllm_omni.model_executor.models.qwen3_omni.quantization import ( + Qwen3OmniNestedSupportsQuant, +) from vllm_omni.quantization.component_config import ( PRE_QUANTIZED_METHODS, ComponentQuantizationConfig, @@ -266,7 +269,9 @@ class Qwen3OmniMoeAudioAttention(_Qwen3OmniMoeAudioAttention): self.num_heads = config.encoder_attention_heads self.head_dim = self.embed_dim // self.num_heads tp_size = get_tensor_model_parallel_world_size() - self.num_local_heads = self.num_heads // tp_size + disable_tp = self.num_heads % tp_size != 0 + effective_tp = 1 if disable_tp else tp_size + self.num_local_heads = self.num_heads // effective_tp if (self.head_dim * self.num_heads) != self.embed_dim: raise ValueError( @@ -282,6 +287,7 @@ class Qwen3OmniMoeAudioAttention(_Qwen3OmniMoeAudioAttention): total_num_heads=self.num_heads, total_num_kv_heads=self.num_heads, bias=True, + disable_tp=disable_tp, quant_config=quant_config, prefix=f"{prefix}.qkv", ) @@ -290,6 +296,7 @@ class Qwen3OmniMoeAudioAttention(_Qwen3OmniMoeAudioAttention): input_size=self.embed_dim, output_size=self.embed_dim, bias=True, + disable_tp=disable_tp, quant_config=quant_config, prefix=f"{prefix}.out_proj", ) @@ -316,10 +323,13 @@ class Qwen3OmniMoeAudioEncoderLayer(_Qwen3OmniMoeAudioEncoderLayer): self.self_attn = Qwen3OmniMoeAudioAttention(config, quant_config=quant_config, prefix=f"{prefix}.self_attn") self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function] + tp_size = get_tensor_model_parallel_world_size() + disable_tp = config.encoder_attention_heads % tp_size != 0 self.fc1 = ColumnParallelLinear( self.embed_dim, config.encoder_ffn_dim, bias=True, + disable_tp=disable_tp, quant_config=quant_config, prefix=f"{prefix}.fc1", ) @@ -327,6 +337,7 @@ class Qwen3OmniMoeAudioEncoderLayer(_Qwen3OmniMoeAudioEncoderLayer): config.encoder_ffn_dim, self.embed_dim, bias=True, + disable_tp=disable_tp, quant_config=quant_config, prefix=f"{prefix}.fc2", ) @@ -1129,6 +1140,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( SupportsMRoPE, Qwen3OmniMoeConditionalGenerationMixin, SupportsTranscription, + Qwen3OmniNestedSupportsQuant, ): # PEFT stores the expert LoRA matrices with an explicit expert dimension. is_3d_moe_weight: bool = True diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index 28ca4f362..36bf663dd 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -173,16 +173,13 @@ class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig): model_type = "qwen3_tts_talker_code_predictor" keys_to_ignore_at_inference = ["past_key_values"] - # Default tensor parallel plan for base model `Qwen3TTSTalkerCodePredictor` - base_model_tp_plan = { - "layers.*.self_attn.q_proj": "colwise", - "layers.*.self_attn.k_proj": "colwise", - "layers.*.self_attn.v_proj": "colwise", - "layers.*.self_attn.o_proj": "rowwise", - "layers.*.mlp.gate_proj": "colwise", - "layers.*.mlp.up_proj": "colwise", - "layers.*.mlp.down_proj": "rowwise", - } + # The code predictor currently runs TP=1. Do not advertise a tensor + # parallel plan for the packed qkv/gate_up projections until the fused + # loader and ``_split_qkv`` path are taught TP-aware packing/sharding: + # a generic colwise split of the plain ``nn.Linear`` would split the + # packed [q, k, v] layout across ranks instead of preserving each rank's + # local q/k/v groups. + base_model_tp_plan = {} base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 8a68e63cd..6dfe8736d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -52,6 +52,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig( sync_process_input_func=f"{_PROC}.talker2code2wav_token_only", sampling_constraints={"detokenize": True}, extras={"tts_args": {"max_instructions_length": 500}}, + requires_full_payload_input=True, ), ), ) diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index 0d8f03294..4053ddae4 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -85,6 +85,37 @@ def build_instruct_text(instruct: str) -> str: # --------------------------------------------------------------------------- +def resolve_x_vector_only(info_dict: dict) -> bool | None: + """Resolve whether a request runs Base voice-clone in x-vector-only mode. + + Mirrors the resolution in :meth:`Qwen3TTSPromptEmbedsBuilder.build_prompt_embeds` + for the ``task_type == "Base"`` branch: the ``x_vector_only_mode`` flag, then + the ``voice_clone_prompt.icl_mode`` override when present. + + Returns ``None`` when the mode does not apply (any non-Base task), so callers + can distinguish "in-context" from "not a voice-clone request at all". + """ + task_type = first_value(info_dict.get("task_type"), "CustomVoice") + if task_type != "Base": + return None + + xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0]) + + raw = info_dict.get("voice_clone_prompt") + if isinstance(raw, list): + raw = raw[0] if raw else None + if isinstance(raw, list) and raw and isinstance(raw[0], dict): + raw = raw[0] + if isinstance(raw, dict) and "icl_mode" in raw: + icl_flag = raw.get("icl_mode") + if isinstance(icl_flag, list): + icl_flag = icl_flag[0] if icl_flag else None + if isinstance(icl_flag, bool): + xvec_only = not icl_flag + + return xvec_only + + def first_value(value: object, default: object = None) -> object: if isinstance(value, list): return value[0] if value else default diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 3daeb4cf2..e65334ff5 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -27,10 +27,11 @@ from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, load_validated_profile_tensors, + validate_qwen3_tts_profile, ) from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig -from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder +from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder @@ -421,6 +422,27 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): codec_mask[self._codec_eos_token_id] = True self.register_buffer("_codec_disallowed_mask", ~codec_mask, persistent=False) + # Silence-region codec tokens (#4966). The talker has real probability mass on + # these in the opening decode frames; how many it draws is a sampling outcome, + # which is the source of both the leading silence and its seed-to-seed variance. + # The vocabulary is checkpoint-specific, so it is derived at load time in + # :meth:`_init_silence_mask`; this is only the placeholder. + self._silence_ban_frames = max(0, int(getattr(vllm_config.model_config, "silence_ban_frames", 0) or 0)) + # vLLM only fills sampling_metadata.output_token_ids when penalties, bad + # words, or a logits processor need history. compute_logits reads it for + # the per-request decode step, so request it when the ban is on. + self.logitsprocs_need_output_token_ids = self._silence_ban_frames > 0 + self.register_buffer("_silence_mask", torch.zeros((vocab,), dtype=torch.bool), persistent=False) + + # Per-request generation mode for the silence ban (#4966). The ban is a + # Base voice-clone (x-vector-only) fix; ICL onsets are legitimate audio + # and must not be masked. Mode is only knowable at prefill, so it is + # recorded in preprocess_batch and read back at decode. The batch order + # comes from set_batch_req_ids, which the runner calls each step with + # the same req_id ordering that indexes sampling_metadata. + self._req_x_vector_only: dict[str, bool | None] = {} + self._batch_req_ids: list[str] = [] + # Keys that should stay on GPU in model_intermediate_buffer to avoid # CPU-to-GPU round-trips on every decode step. self.gpu_resident_buffer_keys: set[tuple[str, str]] = { @@ -506,7 +528,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): tensors = load_validated_profile_tensors( profile, expected_model_type="qwen3_tts", - qwen3_embedding_dim=expected_dim, + validate_profile=lambda profile, tensors: validate_qwen3_tts_profile( + profile, + tensors, + expected_embedding_dim=expected_dim, + ), ) if tensors is None: continue @@ -560,6 +586,48 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # Mask out invalid codec ids using the pre-built constant buffer. logits = logits.masked_fill(self._codec_disallowed_mask, float("-inf")) + # Suppress silence-region tokens for the first N decode frames (#4966). The + # per-request decode step is the length of its generated-token history. + ban_n = self._silence_ban_frames + if ban_n > 0 and sampling_metadata is not None: + output_token_ids = getattr(sampling_metadata, "output_token_ids", None) + if output_token_ids is None or len(output_token_ids) != logits.shape[0]: + logger.warning_once( + "Silence ban is enabled but output_token_ids (%s) does not match the " + "logits batch (%d); the ban will not take effect.", + "None" if output_token_ids is None else len(output_token_ids), + logits.shape[0], + ) + else: + # Both gates are tiny [B] checks, so keep them on CPU: once no + # request is both in scope and inside the window, skip the device + # work rather than paying a host-to-device copy and a [B, vocab] + # mask on every decode frame. + # + # Scope the ban to x-vector-only requests (#4966). A request whose + # mode was never recorded is left alone: masking a request that + # cannot be classified is the risk this gate exists to remove. + batch_req_ids = self._batch_req_ids + if len(batch_req_ids) == logits.shape[0]: + in_scope = [self._req_x_vector_only.get(rid) is True for rid in batch_req_ids] + else: + logger.warning_once( + "Silence ban is enabled but the recorded batch req_ids (%d) do not match " + "the logits batch (%d); the ban will not take effect.", + len(batch_req_ids), + logits.shape[0], + ) + in_scope = [False] * logits.shape[0] + + steps = [len(t) for t in output_token_ids] + if any(scoped and step < ban_n for scoped, step in zip(in_scope, steps, strict=True)): + early = torch.tensor( + [scoped and step < ban_n for scoped, step in zip(in_scope, steps, strict=True)], + device=logits.device, + dtype=torch.bool, + ).unsqueeze(1) + logits = logits.masked_fill(early & self._silence_mask, float("-inf")) + return logits # -------------------- Omni multimodal output plumbing -------------------- @@ -944,12 +1012,30 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device: torch.device, ) -> None: """Delegate batched preprocess to :class:`Qwen3TTSPromptEmbedsBuilder`.""" + if self._silence_ban_frames > 0: + for req_id in req_ids: + info_dict = model_intermediate_buffer.get(req_id) + if isinstance(info_dict, dict): + self._req_x_vector_only[req_id] = resolve_x_vector_only(info_dict) self._prompt_builder.preprocess_batch( req_ids=req_ids, model_intermediate_buffer=model_intermediate_buffer, device=device, ) + def set_batch_req_ids(self, req_ids: Sequence[str]) -> None: + """Record the current batch's req_ids, in the order that indexes logits. + + Called by the runner after the forward and before ``compute_logits``, + with the same ``input_batch`` ordering that ``sampling_metadata`` + uses. Doubles as the eviction point: any request no longer in the batch + has finished or been aborted, so its recorded mode is dropped. + """ + self._batch_req_ids = list(req_ids) + if len(self._req_x_vector_only) > len(self._batch_req_ids): + live = set(self._batch_req_ids) + self._req_x_vector_only = {k: v for k, v in self._req_x_vector_only.items() if k in live} + def _encode_ref_audio_batch(self, wavs: list[np.ndarray], sr: int, *, device: torch.device) -> list[torch.Tensor]: fe = self._encoder_feature_extractor target_sr = int(fe.sampling_rate) @@ -1042,6 +1128,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self.encoder.to(device=device, dtype=torch.bfloat16) self._init_runtime_buffers() + self._init_silence_mask() logger.info("Loaded %d weights for Qwen3TTSTalkerForConditionalGeneration", len(loaded)) self._build_stacked_codec_embed() @@ -1056,6 +1143,84 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): device=w.device, dtype=w.dtype ) + @torch.no_grad() + def _init_silence_mask(self) -> None: + """Derive the silence-region codec vocabulary from the checkpoint. + + Encodes silence at several durations and amplitudes and collects the + codebook-0 ids the tokenizer emits. These are the tokens the talker draws + when it produces leading silence (#4966). Derived rather than hardcoded + because the codec vocabulary is checkpoint-specific. + """ + if self._silence_ban_frames <= 0: + return + device = next(self.parameters()).device + sr = 24000 + try: + rng = np.random.RandomState(0) + wavs = [ + (amp * rng.randn(int(dur * sr))).astype(np.float32) + for amp in (0.0, 1e-5, 1e-4, 1e-3) + for dur in (0.3, 0.5, 1.0) + ] + codes = self._encode_ref_audio_batch(wavs, sr, device=device) + ids: set[int] = set() + for code in codes: + code = code.detach().cpu() + cb0 = code[:, 0] if code.ndim == 2 else code + ids.update(int(t) for t in cb0.tolist()) + except Exception as exc: + logger.warning("Could not derive the silence codec vocabulary; disabling the ban: %s", exc) + self._silence_ban_frames = 0 + return + vocab = int(self.talker_config.vocab_size) + in_range = sorted(t for t in ids if 0 <= t < vocab) + num_out_of_range = len(ids) - len(in_range) + if num_out_of_range: + # Every derived id should be a codebook-0 index for this checkpoint. + # Ids outside the talker vocabulary mean the encoder and the talker + # disagree on vocabulary size, so the rest of the derivation is not + # trustworthy either. + logger.warning( + "%d of %d derived silence codec tokens fall outside the talker vocabulary " + "(size %d); disabling the ban.", + num_out_of_range, + len(ids), + vocab, + ) + self._silence_ban_frames = 0 + return + if not in_range: + # An empty derivation would leave an all-False mask: the ban would + # suppress nothing while still paying the per-frame check. + logger.warning("Derived no silence codec tokens; disabling the ban.") + self._silence_ban_frames = 0 + return + # The silence region is a small corner of codebook 0 (12 tokens on + # Qwen3-TTS-12Hz-1.7B-Base). A derivation this large means the encode + # returned speech-like codes, and masking that much of the vocabulary + # would distort generation rather than trim the onset. + max_silence_tokens = max(64, vocab // 20) + if len(in_range) > max_silence_tokens: + logger.warning( + "Derived %d silence codec tokens, above the %d sanity limit for a vocabulary of %d; disabling the ban.", + len(in_range), + max_silence_tokens, + vocab, + ) + self._silence_ban_frames = 0 + return + mask = torch.zeros((vocab,), dtype=torch.bool) + for token in in_range: + mask[token] = True + self._silence_mask.copy_(mask.to(self._silence_mask.device)) + logger.info( + "Derived %d silence codec tokens, suppressed for the first %d decode frames: %s", + len(in_range), + self._silence_ban_frames, + in_range, + ) + @torch.no_grad() def _init_runtime_buffers(self) -> None: """Populate request-independent runtime buffers from frozen weights. diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 498f9eb85..4438d2200 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -295,12 +295,24 @@ _OMNI_MODELS = { "fish_speech_dac_decoder", "FishSpeechDACDecoder", ), + ## Gepard-1.0 (nineninesix) — single-stage native-AR FSQ/NanoCodec TTS + "GepardTalkerForConditionalGeneration": ( + "gepard", + "gepard_talker", + "GepardTalkerForConditionalGeneration", + ), ## VoxCPM2 "VoxCPM2TalkerForConditionalGeneration": ( "voxcpm2", "voxcpm2_talker", "VoxCPM2TalkerForConditionalGeneration", ), + ## dots.tts + "DotsTTSForConditionalGeneration": ( + "dots_tts", + "dots_tts_talker", + "DotsTTSForConditionalGeneration", + ), ## Voxtral TTS "VoxtralTTSForConditionalGeneration": ( "voxtral_tts", @@ -445,6 +457,22 @@ _OMNI_MODELS = { "qwen3_vl", "AuraQwen3VLForConditionalGeneration", ), + "MiniMaxH3TextEncoder": ( + "minimax_h3", + "text_encoder", + "MiniMaxH3TextEncoder", + ), + ## MiniMax-Music3 (text-to-music; AR talker -> flow-matching acoustic decoder) + "MiniMaxMusic3TalkerForConditionalGeneration": ( + "minimax_music3", + "talker", + "MiniMaxMusic3TalkerForConditionalGeneration", + ), + "MiniMaxMusic3AcousticForConditionalGeneration": ( + "minimax_music3", + "acoustic", + "MiniMaxMusic3AcousticForConditionalGeneration", + ), } diff --git a/vllm_omni/model_executor/models/soulx_singer/__init__.py b/vllm_omni/model_executor/models/soulx_singer/__init__.py deleted file mode 100644 index 1ab85a8d3..000000000 --- a/vllm_omni/model_executor/models/soulx_singer/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Keep this package lightweight. The pipeline registry imports pipeline.py -# directly, while diffusion model classes are loaded lazily from their registry. diff --git a/vllm_omni/model_executor/models/soulx_singer/pipeline.py b/vllm_omni/model_executor/models/soulx_singer/pipeline.py deleted file mode 100644 index 4d2420528..000000000 --- a/vllm_omni/model_executor/models/soulx_singer/pipeline.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""SoulX-Singer pipeline topologies.""" - -from vllm_omni.config.stage_config import ( - PipelineConfig, - StageExecutionType, - StagePipelineConfig, -) - -SOULXSINGER_SVS_PIPELINE = PipelineConfig( - model_type="soulxsinger_svs", - default_deploy_config_name="soulxsinger_svs.yaml", - model_arch="SoulXSingerPipeline", - hf_architectures=("SoulXSingerPipeline",), - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="dit", - execution_type=StageExecutionType.DIFFUSION, - input_sources=(), - final_output=True, - final_output_type="audio", - model_arch="SoulXSingerPipeline", - ), - ), -) - -SOULXSINGER_SVC_PIPELINE = PipelineConfig( - model_type="soulxsinger_svc", - default_deploy_config_name="soulxsinger_svc.yaml", - model_arch="SoulXSingerSVCPipeline", - hf_architectures=("SoulXSingerSVCPipeline",), - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="dit", - execution_type=StageExecutionType.DIFFUSION, - input_sources=(), - final_output=True, - final_output_type="audio", - model_arch="SoulXSingerSVCPipeline", - ), - ), -) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py index b56229cf6..b00e24870 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py @@ -13,7 +13,6 @@ from __future__ import annotations from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, TypedDict -import librosa # noqa: TID251 import numpy as np import torch import torch.nn as nn @@ -52,6 +51,7 @@ from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( DEFAULT_TOKEN_CONFIG, STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID, ) +from vllm_omni.utils.audio import mel_filter_bank if TYPE_CHECKING: pass @@ -65,7 +65,7 @@ if TYPE_CHECKING: def _mel_filters(n_mels: int) -> torch.Tensor: """Generate mel filter banks""" assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}" - return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=n_mels)) + return mel_filter_bank(sr=16000, n_fft=400, n_mels=n_mels) def _normalize_audio(audio: Any) -> torch.Tensor: diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index 9666a4f5a..e9ca78b21 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -1,14 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import io import os from collections.abc import Iterable from dataclasses import dataclass, field from typing import Any -import librosa # noqa: TID251 import numpy as np import onnxruntime import s3tokenizer -import soundfile as sf import torch import torch.nn as nn import torchaudio @@ -23,6 +24,8 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import WeightsMapper +from vllm.multimodal.audio import AudioResampler +from vllm.multimodal.media.audio import load_audio from vllm.sequence import IntermediateTensors from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata @@ -41,6 +44,14 @@ from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( logger = init_logger(__name__) +def _get_left_context_size(info: dict[str, Any]) -> Any | None: + """Return the async-chunk marker from a nested metadata payload.""" + meta = info.get("meta") + if not isinstance(meta, dict): + return None + return meta.get("left_context_size") + + def fade_in_out( fade_in_mel: torch.Tensor, fade_out_mel: torch.Tensor, @@ -231,19 +242,11 @@ class StepAudio2Token2WavCore(nn.Module): def _prepare_prompt(self, prompt_wav: str): """Prepare prompt audio for conditioning""" - # Prefer soundfile/librosa path to avoid torchaudio->torchcodec runtime coupling. - try: - audio_np_raw, sample_rate = sf.read(prompt_wav, dtype="float32", always_2d=False) - except Exception: - audio_np_raw, sample_rate = librosa.load(prompt_wav, sr=None, mono=True) - if isinstance(audio_np_raw, np.ndarray) and audio_np_raw.ndim > 1: - audio_np_raw = np.mean(audio_np_raw, axis=-1) + audio_np_raw, sample_rate = load_audio(prompt_wav, sr=None, mono=True) sample_rate = int(sample_rate) # 16 kHz branch: s3tokenizer + speaker embedding - audio_16k = audio_np_raw.astype(np.float32) - if sample_rate != 16000: - audio_16k = librosa.resample(y=audio_16k, orig_sr=sample_rate, target_sr=16000) + audio_16k = AudioResampler(target_sr=16000).resample(audio_np_raw.astype(np.float32), orig_sr=sample_rate) audio = torch.from_numpy(audio_16k) mels = s3tokenizer.log_mel_spectrogram(audio) mels, mels_lens = s3tokenizer.padding([mels]) @@ -257,9 +260,7 @@ class StepAudio2Token2WavCore(nn.Module): # 24 kHz branch: mel spectrogram for flow model conditioning # Must resample from the ORIGINAL audio, not the 16 kHz version. - audio_24k = audio_np_raw.astype(np.float32) - if sample_rate != 24000: - audio_24k = librosa.resample(y=audio_24k, orig_sr=sample_rate, target_sr=24000) + audio_24k = AudioResampler(target_sr=24000).resample(audio_np_raw.astype(np.float32), orig_sr=sample_rate) audio_24k_t = torch.from_numpy(audio_24k).unsqueeze(0) # [1, T] prompt_mel = ( mel_spectrogram( @@ -547,7 +548,8 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): additional: Additional information dict (deprecated) runtime_additional_information: Per-request dicts from the async_chunk processor. Each dict contains ``audio_tokens`` - (list[int]) and ``last_chunk`` (bool). + (list[int]) and ``meta.left_context_size`` as the + last-chunk marker (0 = non-final, 1 = final). **kwargs: Other kwargs Returns: @@ -560,7 +562,7 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): # [{}]), so a bare truthiness check would incorrectly enter this # branch in synchronous mode. if runtime_additional_information and any( - "left_context_size" in info for info in runtime_additional_information + _get_left_context_size(info) is not None for info in runtime_additional_information ): return self._forward_async_chunk(input_ids, runtime_additional_information, **kwargs) @@ -629,15 +631,15 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): async_chunk payload to ``request.additional_information``). """ # --- batch=1 guard --- - batch_size = sum(1 for info in runtime_additional_information if "left_context_size" in info) + batch_size = sum(1 for info in runtime_additional_information if _get_left_context_size(info) is not None) if batch_size != 1: raise RuntimeError( f"Token2Wav async_chunk only supports batch=1, got {batch_size}. " "Batch>1 requires framework support for batch_req_ids." ) - info = next(info for info in runtime_additional_information if "left_context_size" in info) - last_chunk = info.get("left_context_size", 0) == 1 + info = next(info for info in runtime_additional_information if _get_left_context_size(info) is not None) + last_chunk = _get_left_context_size(info) == 1 # --- Manage single stream state --- # If the previous request was preempted (setup_done but not finished), diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 3a1677d49..f8dc962f3 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -42,6 +42,7 @@ from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, load_validated_profile_tensors, + validate_voxcpm2_profile, ) from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest @@ -963,7 +964,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): loaded = 0 for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type="voxcpm2"): - tensors = load_validated_profile_tensors(profile, expected_model_type="voxcpm2") + tensors = load_validated_profile_tensors( + profile, + expected_model_type="voxcpm2", + validate_profile=validate_voxcpm2_profile, + ) if tensors is None: continue @@ -3013,9 +3018,18 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): merged_cpu = merged.detach().cpu().float() mm["model_outputs"] = list(merged_cpu.split(sizes)) mm["sr"] = [sr for _ in ready_req_ids] + # Dense mode still yields a strict SUBSET of the batch when + # some requests emit no audio this step (e.g. prefill + # phase). Without the marker the runner indexes these + # per-request lists by batch position and misroutes audio + # across requests; the marker declares meta.req_id + # alignment. + mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} else: mm["model_outputs"] = list(audio_by_req.values()) mm["sr"] = [sr for _ in audio_by_req] + # Same subset hazard as the coalesce branch above. + mm["meta"] = {"req_id": list(audio_by_req), "sparse_audio": ["1"]} elif self._uses_sparse_audio_outputs(): mm["model_outputs"] = [] mm["sr"] = [] diff --git a/vllm_omni/model_executor/models/wan2_2/pipeline.py b/vllm_omni/model_executor/models/wan2_2/pipeline.py index 7059845fa..2d2bc3064 100644 --- a/vllm_omni/model_executor/models/wan2_2/pipeline.py +++ b/vllm_omni/model_executor/models/wan2_2/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Wan2.2 TI2V pipeline topology.""" from vllm_omni.config.stage_config import ( @@ -12,6 +12,7 @@ WAN2_2_TI2V_PIPELINE = PipelineConfig( model_type="wan2_2_ti2v", model_arch="WanPipeline", diffusers_class_name="WanPipeline", + diffusers_class_aliases=("WanDMDPipeline",), stages=( StagePipelineConfig( stage_id=0, diff --git a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py index ad35c2349..9b9462ad9 100644 --- a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py +++ b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging +from dataclasses import dataclass, field +from typing import ClassVar logger = logging.getLogger(__name__) @@ -127,3 +129,228 @@ def compute_ramp_emit( return True, 0 return True, length - prev_threshold + + +@dataclass +class AdaptiveChunkController: + """Buffer-feedback controller for adaptive chunk sizing. + + Replaces the fixed ramp table with a greedy buffer-feedback rule. + Chunk 0 uses dynamic IC; chunk 1+ sizes are decided per-chunk from + live buffer-depth signals. + + State is per-segment: the controller is discarded and rebuilt at + segment boundaries via ``_adaptive_states.pop()`` in the adapter. + """ + + frame_time_ewma_ms: float = 0.0 + first_emit_time: float = 0.0 + last_emit_time: float = 0.0 + frames_emitted_this_segment: int = 0 + steady_state_reached: bool = False + accumulated_at_last_emit: int = 0 + last_target: int = 0 + # Per-request underrun telemetry, mirroring stream_client.py's playback + # cursor model so server-side numbers are comparable to client-reported + # AUDIO_UNDERRUN. Collected in record_emit; logged once at request end via + # log_summary(). total_gap_s accumulates only gaps > 5ms (matching + # stream_client's silence-insert threshold); per_chunk_gap_ms records every + # chunk's gap (including 0) for the full trajectory. + chunk_sizes: list = field(default_factory=list) + per_chunk_gap_ms: list = field(default_factory=list) + playback_cursor_s: float = 0.0 + total_gap_s: float = 0.0 + # Ramp-rate tuning, sourced from connector config (codec_chunk_ramp_divisor + # / codec_chunk_ramp_delta_min) so deploys can tune the ramp without code + # changes. Defaults preserve the original 15/2 behavior. + ramp_divisor: int = 15 + ramp_delta_min: int = 2 + + EWMA_ALPHA: ClassVar[float] = 0.3 + FRAME_DURATION_MS: ClassVar[float] = 80.0 + + def compute_next_chunk_size( + self, + now: float, + min_frames: int, + max_frames: int, + safety_margin_ms: float, + ) -> int: + if self.steady_state_reached: + return max_frames + + emitted_ms = self.frames_emitted_this_segment * self.FRAME_DURATION_MS + elapsed_ms = (now - self.first_emit_time) * 1000.0 + buffer_ms = emitted_ms - elapsed_ms + + # Hysteresis: only check when EWMA has a valid measurement. + if self.frame_time_ewma_ms > 0: + steady_gen_ms = max_frames * self.frame_time_ewma_ms + if buffer_ms >= steady_gen_ms + safety_margin_ms: + self.steady_state_reached = True + return max_frames + + # Greedy: select the largest n where buffer can absorb + # n * frame_time + safety_margin. + available_ms = max(0.0, buffer_ms - safety_margin_ms) + if self.frame_time_ewma_ms > 0: + max_n = int(available_ms / self.frame_time_ewma_ms) + else: + max_n = min_frames + greedy_result = max(min_frames, min(max_n, max_frames)) + + # Ramp floor: chunk size grows by an EWMA-scaled delta per emit. + # Unconditional (never gated on buffer sign): in the saturated + # regime (per-stream RTF > 1) buffer is structurally negative and + # shrinking chunks would multiply Code2Wav decode cost -- each chunk + # re-decodes the left_context, cost ~ (72+n)/n, so n=2 costs ~10x + # n=25. Growing toward max minimizes decode passes, the only lever + # pre-W5. delta scales with per-frame time: slow hardware (NPU, + # ewma~150) ramps fast (delta=10, ~3 emits to max); fast hardware + # (NVIDIA, ewma~16) ramps gently (delta=2) for lower TTFP. This + # cross-hardware self-tuning is why the adaptive controller exists. + # last_target is set to the controller's target_size in record_emit + # (on actual emit), not here -- compute_next_chunk_size may be called + # multiple times per emit as frames accumulate one-by-one. + # WARNING: do NOT gate ramp_floor on buffer>0. That collapses to + # greedy=min_frames under load and triggers a vocoder-cost death + # spiral (see #4855 linyueqian 2026-07-19 comment + W5 coupling). + # Downward chunk-size adaptation is deferred until W5 (streaming + # decoder state) lands -- see PR description. + if self.frame_time_ewma_ms > 0: + ramp_delta = max(self.ramp_delta_min, int(self.frame_time_ewma_ms / self.ramp_divisor)) + else: + ramp_delta = self.ramp_delta_min + ramp_floor = self.last_target + ramp_delta + result = max(ramp_floor, greedy_result) + result = min(result, max_frames) + + return result + + def record_emit( + self, + now: float, + n_frames: int, + accumulated: int, + target_size: int, + ) -> None: + if self.frames_emitted_this_segment > 0: + dt_ms = (now - self.last_emit_time) * 1000.0 + measured_ms = dt_ms / n_frames if n_frames > 0 else self.frame_time_ewma_ms + if self.frame_time_ewma_ms == 0: + self.frame_time_ewma_ms = measured_ms + else: + self.frame_time_ewma_ms = ( + self.EWMA_ALPHA * measured_ms + (1 - self.EWMA_ALPHA) * self.frame_time_ewma_ms + ) + # Per-chunk underrun telemetry: mirror stream_client.py's playback + # cursor so server-side total_gap is comparable to the client's + # AUDIO_UNDERRUN. arrival_s = wall since first emit; gap_s = how late + # this chunk is vs where playback had reached. + arrival_s = now - self.first_emit_time + audio_dur_s = n_frames * self.FRAME_DURATION_MS / 1000.0 + gap_s = max(0.0, arrival_s - self.playback_cursor_s) + if gap_s > 0.005: + self.total_gap_s += gap_s + self.per_chunk_gap_ms.append(gap_s * 1000.0) + self.chunk_sizes.append(n_frames) + self.playback_cursor_s = max(arrival_s, self.playback_cursor_s) + audio_dur_s + self.last_emit_time = now + self.frames_emitted_this_segment += n_frames + self.accumulated_at_last_emit = accumulated + # last_target tracks the controller's intended target_size, NOT the + # actual n_frames shipped. With the B1 fix (compute_adaptive_emit + # returns new_frames), n_frames can overshoot target_size during a + # latency spike. Using n_frames here would ratchet ramp_floor to the + # spike size and lock the controller toward max_frames prematurely. + self.last_target = target_size + + def log_summary(self, request_id: str) -> None: + """Emit one INFO line at segment/request end. + + Mirrors stream_client.py's underrun model so total_gap_s is + comparable to the client-reported AUDIO_UNDERRUN. Fires on + ``finished`` which is ``is_segment_finished or is_finished`` (see + chunk_transfer_adapter), so a resumable /v1/realtime session emits + one line per segment. INFO level so it is visible by default + without any env var or dedicated logger. Only emits fields vLLM + stats do NOT already report: chunk-size trajectory, per-chunk + underrun, and total underrun. (pure/wall/ewma were dropped: vLLM + already reports audio_duration, e2e/stage wall, and + tpot/time_per_output_unit respectively; a trailing ewma also could + not reflect its intra-request evolution.) + """ + logger.info( + "[ADAPTIVE] req=%s done: chunks=%s gaps_ms=%s total_gap=%.3fs", + request_id, + self.chunk_sizes, + [round(g, 1) for g in self.per_chunk_gap_ms], + self.total_gap_s, + ) + + +def compute_adaptive_emit( + length: int, + accumulated_at_last_emit: int, + target_size: int, + finished: bool, +) -> tuple[bool, int]: + """Decide whether to emit under the adaptive controller. + + Same return contract as ``compute_ramp_emit``: + + - ``(False, 0)``: hold. + - ``(True, >0)``: emit with this many new frames. + - ``(True, 0)``: finished with no new frames (empty-finished sentinel). + """ + new_frames = length - accumulated_at_last_emit + + if finished and new_frames <= 0: + return True, 0 + + # When finished, flush ALL remaining frames to avoid data loss. + # Truncating to target_size would silently drop new_frames - target_size frames. + if finished: + return True, new_frames + + if not finished and new_frames < target_size: + return False, 0 + + # Ship ALL accumulated frames, not just target_size. When new_frames + # exceeds target_size (e.g. a scheduler stall dropped target below the + # already-queued frames), returning target_size would cause the caller's + # [-context_length:] slice to take the newest frames, silently skipping + # the surplus. Shipping new_frames keeps the slice contiguous with the + # previous emit and preserves Code2Wav's per-request decoder cache. + return True, new_frames + + +def parse_adaptive_config( + cfg: dict, +) -> tuple[bool, int, float, int, int]: + """Parse adaptive chunk settings from connector extra config. + + The adaptive max is ``codec_chunk_frames`` (applied at the call site + as ``max_frames``); there is no separate max config -- capping below + steady would undermine the controller's ramp-to-large design (more + Code2Wav re-decode tax) and is a saturated-regime death-spiral risk, + so it is intentionally not exposed. + + Returns: + ``(enabled, min_frames, safety_margin_ms, ramp_divisor, ramp_delta_min)`` + """ + enabled = bool(cfg.get("codec_chunk_adaptive", False)) + min_frames = int(cfg.get("codec_chunk_min_frames", 2)) + margin_ms = float(cfg.get("codec_chunk_safety_margin_ms", 50.0)) + ramp_divisor = int(cfg.get("codec_chunk_ramp_divisor", 15)) + ramp_delta_min = int(cfg.get("codec_chunk_ramp_delta_min", 2)) + if min_frames < 1: + logger.warning("codec_chunk_min_frames must be >= 1, clamping to 1.") + min_frames = 1 + if ramp_divisor < 1: + logger.warning("codec_chunk_ramp_divisor must be >= 1, clamping to 1.") + ramp_divisor = 1 + if ramp_delta_min < 1: + logger.warning("codec_chunk_ramp_delta_min must be >= 1, clamping to 1.") + ramp_delta_min = 1 + return enabled, min_frames, margin_ms, ramp_divisor, ramp_delta_min diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 2956be94f..07ea3af7f 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections import defaultdict from collections.abc import Mapping from contextlib import nullcontext from typing import Any -import numpy as np import torch from vllm.inputs import TextPrompt from vllm.logger import init_logger @@ -16,6 +15,7 @@ from vllm_omni.data_entry_keys import ( MetaStruct, OmniPayloadStruct, ) +from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.cosyvoice3.utils import unpad_prompt_conditioning @@ -93,32 +93,6 @@ def _to_cpu_tensor(x: Any) -> torch.Tensor | None: return None -def _decode_additional_information(raw_info: Any) -> dict[str, Any]: - if raw_info is None: - return {} - if isinstance(raw_info, dict): - return raw_info - - entries = getattr(raw_info, "entries", None) - if not isinstance(entries, dict): - return {} - - decoded: dict[str, Any] = {} - for key, entry in entries.items(): - tensor_data = getattr(entry, "tensor_data", None) - if tensor_data is not None: - dtype_name = getattr(entry, "tensor_dtype", "float32") - tensor_shape = getattr(entry, "tensor_shape", None) - if tensor_shape is None: - continue - dt = np.dtype(dtype_name) - arr = np.frombuffer(tensor_data, dtype=dt).reshape(tensor_shape) - decoded[key] = torch.from_numpy(arr.copy()) - else: - decoded[key] = getattr(entry, "list_data", None) - return decoded - - def talker2code2wav_async_chunk( transfer_manager: Any, multimodal_output: dict[str, Any] | None, @@ -149,7 +123,7 @@ def talker2code2wav_async_chunk( request_state = transfer_manager.request_payload.get(request_id) if not isinstance(request_state, dict) or "_cosyvoice3_async_state" not in request_state: with nullcontext(): - info = _decode_additional_information(getattr(request, "additional_information", None)) + info = deserialize_additional_information(getattr(request, "additional_information", None)) info_embed = info.get("embed", {}) if isinstance(info, dict) else {} prompt_payload = {} cond_keys = ("speech_token", "speech_feat", "embedding", "speech_token_len") diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index d12a3b382..b293a95c8 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniCPM-o 4.5 Thinker-to-Talker and Talker-to-Code2Wav bridges.""" import logging @@ -16,6 +16,8 @@ from vllm_omni.experimental.fullduplex.engine.intermediate import ( set_tts_handoff, ) from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.minicpmo_4_5 import MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY logger = logging.getLogger(__name__) _MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state" @@ -70,6 +72,23 @@ def _extract_first_audio_ref(multi_modal_data): return waveform.reshape(-1).cpu(), int(sample_rate) +def _extract_prompt_reference_audio( + prompt_item: OmniTokensPrompt | TextPrompt | None, +) -> tuple[torch.Tensor, int] | None: + if isinstance(prompt_item, Mapping): + multi_modal_data = prompt_item.get("multi_modal_data") + serving_reference_audio = prompt_item.get(MINICPMO45_REFERENCE_AUDIO_KEY) + else: + multi_modal_data = getattr(prompt_item, "multi_modal_data", None) + serving_reference_audio = getattr(prompt_item, MINICPMO45_REFERENCE_AUDIO_KEY, None) + + reference_audio = _extract_first_audio_ref(multi_modal_data) + if reference_audio is not None: + return reference_audio + + return _extract_first_audio_ref({"audio": serving_reference_audio}) + + def _extract_native_runtime_ref_audio(data_plane_metadata): if not isinstance(data_plane_metadata, dict): return None @@ -141,6 +160,17 @@ def _codec_config(transfer_manager: Any) -> tuple[int, int]: return chunk_frames, left_context_frames +def _request_intermediate_section(request: object, section: str) -> dict[str, object]: + merged: dict[str, object] = {} + for attribute in ("additional_information", "model_intermediate_buffer"): + information = getattr(request, attribute, None) + if isinstance(information, Mapping): + values = information.get(section) + if isinstance(values, Mapping): + merged.update(values) + return merged + + def _codec_scalars(value: Any) -> list[int]: """Normalize one request-routed codec delta to CPU scalar token IDs.""" if value is None: @@ -260,7 +290,15 @@ def tts2code2wav_async_chunk( return None if native_duplex and turn_end and record.get("last_terminal_turn") == duplex_turn_key: - return None + # Emit an empty replacement snapshot so Code2Wav cannot replay the + # prior terminal audio when this control-only boundary arrives. + return OmniPayloadStruct( + meta=_MiniCPMO45MetaStruct( + is_segment_finished=torch.tensor(True, dtype=torch.bool), + replace_runtime_additional_information=True, + ), + request_id=request_id, + ) state = container.get(_MINICPMO45_ASYNC_STATE) if not isinstance(state, dict): @@ -340,15 +378,13 @@ def tts2code2wav_async_chunk( ref_audio = None ref_audio_sr = None if int(record["cache_epoch"]) == 0 and chunk_seq == 0: - request_info = getattr(request, "additional_information", None) - if isinstance(request_info, Mapping): - codes_info = request_info.get("codes") - meta_info = request_info.get("meta") - raw_ref_audio = codes_info.get("ref") if isinstance(codes_info, Mapping) else None - raw_ref_audio_sr = meta_info.get("ref_audio_sr") if isinstance(meta_info, Mapping) else None - ref_audio_sr = _coerce_int(raw_ref_audio_sr) - if raw_ref_audio is not None: - ref_audio = torch.as_tensor(raw_ref_audio, dtype=torch.float32).reshape(-1).cpu() + codes_info = _request_intermediate_section(request, "codes") + meta_info = _request_intermediate_section(request, "meta") + raw_ref_audio = codes_info.get("ref") + raw_ref_audio_sr = meta_info.get("ref_audio_sr") + ref_audio_sr = _coerce_int(raw_ref_audio_sr) + if raw_ref_audio is not None: + ref_audio = torch.as_tensor(raw_ref_audio, dtype=torch.float32).reshape(-1).cpu() finished_tensor = torch.tensor(last_chunk, dtype=torch.bool) payload = OmniPayloadStruct( codes=CodesStruct( @@ -373,6 +409,7 @@ def tts2code2wav_async_chunk( llm_output_text_utf8=segment_text_utf8, tts_is_last_chunk=flush_pending, turn_end=turn_end and last_chunk, + replace_runtime_additional_information=True, ref_audio_sr=ref_audio_sr, ), request_id=request_id, @@ -399,18 +436,9 @@ def tts2code2wav_full_payload( context = [_MINICPMO45_SILENCE_CODE] * left_context_frames if codes else [] output_codes = [*context, *codes] - request_info = getattr(request, "additional_information", None) - if not isinstance(request_info, Mapping): - request_info = {} - codes_info = request_info.get("codes") - if not isinstance(codes_info, Mapping): - codes_info = {} - meta_info = request_info.get("meta") - if not isinstance(meta_info, Mapping): - meta_info = {} - duplex_info = request_info.get("duplex") - if not isinstance(duplex_info, Mapping): - duplex_info = {} + codes_info = _request_intermediate_section(request, "codes") + meta_info = _request_intermediate_section(request, "meta") + duplex_info = _request_intermediate_section(request, "duplex") ref_audio = codes_info.get("ref") finished = torch.tensor(True, dtype=torch.bool) @@ -432,6 +460,7 @@ def tts2code2wav_full_payload( finished=finished, req_id=[request_id], ref_audio_sr=_coerce_int(meta_info.get("ref_audio_sr")), + replace_runtime_additional_information=True, native_duplex_segment_text=( str(meta_info["native_duplex_segment_text"]) if isinstance(meta_info.get("native_duplex_segment_text"), str) @@ -540,6 +569,29 @@ def _native_tts_boundary_token_ids(special_token_ids): } +def _native_duplex_output_is_control_only( + output_ids: Sequence[int], + special_token_ids: Mapping[str, int], +) -> bool: + if not output_ids: + return True + control_ids = { + token_id + for token_id in ( + special_token_ids.get("tts_eos_token_id"), + special_token_ids.get("tts_pad_token_id"), + special_token_ids.get("listen_token_id"), + special_token_ids.get("chunk_eos_token_id"), + special_token_ids.get("chunk_tts_eos_token_id"), + special_token_ids.get("turn_eos_token_id"), + special_token_ids.get("unit_token_id"), + special_token_ids.get("unit_end_token_id"), + ) + if token_id is not None + } + return bool(control_ids) and all(token_id in control_ids for token_id in output_ids) + + def _decode_native_duplex_token_ids( token_ids: list[int], streaming_context, @@ -688,11 +740,13 @@ def llm2tts( prompt = [prompt] multi_modal_data = {} + reference_audio_by_request_id = {} for llm_output, p in zip(llm_outputs, prompt): if isinstance(p, dict): multi_modal_data[llm_output.request_id] = p.get("multi_modal_data", None) else: multi_modal_data[llm_output.request_id] = getattr(p, "multi_modal_data", None) + reference_audio_by_request_id[llm_output.request_id] = _extract_prompt_reference_audio(p) for llm_output in llm_outputs: output = llm_output.outputs[0] @@ -747,10 +801,16 @@ def llm2tts( ) prompt_token_ids_len = len(prompt_token_ids) + is_native_duplex_handoff = _has_native_duplex_prompt_metadata(mm_output) + latent = mm_output.get("latent", None) if latent is None: latent = output.hidden_states if hasattr(output, "hidden_states") else None if latent is None: + if is_native_duplex_handoff and _native_duplex_output_is_control_only( + llm_output_ids, special_token_ids + ): + continue raise ValueError("No latent or hidden_states found in thinker output") thinker_hidden_states = latent.detach() @@ -761,7 +821,6 @@ def llm2tts( full_token_ids = prompt_token_ids + llm_output_ids tts_bos_id = special_token_ids.get("tts_bos_token_id") - is_native_duplex_handoff = _has_native_duplex_prompt_metadata(mm_output) if is_native_duplex_handoff: _require_native_tts_boundary_metadata(special_token_ids) tts_end_ids = _native_tts_boundary_token_ids(special_token_ids) @@ -922,8 +981,7 @@ def llm2tts( meta["turn_start"] = native_turn_start if native_segment_end: meta["segment_end"] = True - req_mm_data = multi_modal_data.get(llm_output.request_id) - ref_audio = _extract_first_audio_ref(req_mm_data) + ref_audio = reference_audio_by_request_id.get(llm_output.request_id) if ref_audio is None: ref_audio = _extract_native_runtime_ref_audio( model_intermediate_buffer.get("duplex"), @@ -945,9 +1003,14 @@ def llm2tts( if handoff_ids is not None and handoff_hidden is not None: condition_suffix_length = 1 if is_native_duplex_handoff else 2 condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length + # Dummy ids only reserve scheduler slots; prefill overwrites the + # embeddings. Talker.sample() blanks the whole prompt out of the + # repetition penalty, so codec id 0 is not taxed from step one. scheduler_prompt_token_ids = [0] * condition_length handoff_meta = model_intermediate_buffer.setdefault("meta", {}) handoff_meta["next_stage_prompt_len"] = condition_length + if is_native_duplex_handoff: + handoff_meta["next_stage_generation_tokens"] = MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK # Native duplex resumes one Talker request within a turn, but a new # assistant turn must discard the previous turn's prompt and KV. if not is_native_duplex_handoff or native_turn_start: diff --git a/vllm_omni/model_executor/stage_input_processors/minimax_h3.py b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py new file mode 100644 index 000000000..372d675b6 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/minimax_h3.py @@ -0,0 +1,364 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniMax H3 text-encoder stage input and output adapters.""" + +from __future__ import annotations + +import copy +import shutil +import tempfile +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from vllm_omni.data_entry_keys import REQUEST_ARTIFACT_DIRS_KEY +from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_align_frame_count +from vllm_omni.errors import OmniClientError +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.models.minimax_h3.conditioning import ( + MINIMAX_H3_CONDITION_LABELS_KEY, + MINIMAX_H3_PRESENTATION_TASK_KEY, + MiniMaxH3TextConditioning, +) +from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( + MINIMAX_H3_OUTPUT_SHORT_EDGE, + load_minimax_h3_images, + resolve_minimax_h3_aspect_ratio, + resolve_minimax_h3_output_canvas, + resolve_minimax_h3_reference_image_shape, +) +from vllm_omni.model_executor.models.minimax_h3.reference_video import ( + MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, + MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, + prepare_reference_videos, + sample_reference_video_frames, + serialize_prepared_reference_videos, +) + + +def _items(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, list): + return value + if isinstance(value, tuple) and not (len(value) == 2 and isinstance(value[1], Mapping)): + return list(value) + return [value] + + +def _audio_items(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)): + return [value] + return list(value) if isinstance(value, (list, tuple)) else [value] + + +def _resolve_task( + extra_args: Mapping[str, Any], + multi_modal_data: Mapping[str, Any], +) -> str: + requested = extra_args.get("task") + if requested is not None: + return str(requested).lower() + if multi_modal_data.get("video") is not None or multi_modal_data.get("audio") is not None: + return "ref2va" + if multi_modal_data.get("image") is not None: + return "fl2va" + return "t2va" + + +def _diffusion_sampling_params(sampling_params_list: Sequence[Any]) -> Any: + diffusion_params = [ + sampling_params + for sampling_params in sampling_params_list + if isinstance(sampling_params, OmniDiffusionSamplingParams) + ] + if len(diffusion_params) != 1: + raise RuntimeError( + "MiniMax H3 text encoding requires exactly one OmniDiffusionSamplingParams stage parameter, " + f"got {len(diffusion_params)}" + ) + return diffusion_params[0] + + +def _ref2va_target_frame_count(sampling_params_list: Sequence[Any]) -> int: + sampling = _diffusion_sampling_params(sampling_params_list) + extra_args = sampling.extra_args or {} + target = extra_args.get("target") + target = target if isinstance(target, Mapping) else {} + duration = target.get("duration_seconds", extra_args.get("duration_seconds", extra_args.get("duration"))) + if duration is not None: + requested = int(round(float(duration) * 24)) + elif int(getattr(sampling, "num_frames", None) or 1) > 1: + requested = int(sampling.num_frames) + else: + requested = 124 + return minimax_h3_align_frame_count(requested) + + +def _prepare_qwen_images( + task: str, + values: list[Any], + sampling_params_list: Sequence[Any], +) -> list[Any]: + if not values: + return [] + images = load_minimax_h3_images(values) + if task == "ref2va": + return [ + image.resize( + resolve_minimax_h3_reference_image_shape(image), + Image.Resampling.LANCZOS, + ) + for image in images + ] + if task != "fl2va": + return images + + sampling = _diffusion_sampling_params(sampling_params_list) + extra_args = sampling.extra_args or {} + target = extra_args.get("target") + if target is not None and not isinstance(target, Mapping): + raise OmniClientError("MiniMax H3 extra_args['target'] must be an object") + target = target if isinstance(target, Mapping) else {} + aspect_ratio = resolve_minimax_h3_aspect_ratio( + task, + target.get("aspect_ratio", extra_args.get("aspect_ratio")), + images[0], + ) + if not 0.25 <= aspect_ratio <= 4.0: + raise OmniClientError(f"MiniMax H3 canvas aspect ratio must be in [1:4, 4:1], got {aspect_ratio}") + height = sampling.height + width = sampling.width + if height is None or width is None: + short_edge = target.get( + "short_edge", + extra_args.get("short_edge", MINIMAX_H3_OUTPUT_SHORT_EDGE), + ) + if isinstance(short_edge, bool) or not isinstance(short_edge, (int, np.integer)): + raise OmniClientError( + f"MiniMax H3 target.short_edge must be {MINIMAX_H3_OUTPUT_SHORT_EDGE}, got {short_edge!r}" + ) + height, width = resolve_minimax_h3_output_canvas(aspect_ratio, int(short_edge)) + height = int(height) // 32 * 32 + width = int(width) // 32 * 32 + if min(height, width) <= 0: + raise OmniClientError(f"invalid MiniMax H3 canvas {width}x{height}") + if width > 4 * height or height > 4 * width: + raise OmniClientError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]") + return [image.resize((width, height), Image.Resampling.LANCZOS) for image in images] + + +def prepare_text_encoder_prompt( + prompt: Any, + sampling_params_list: Sequence[Any], +) -> Any: + """Build H3's labeled Qwen3-VL presentation for Stage 0. + + The upstream Qwen3-VL multimodal processor expands each image/video + placeholder into the exact number of vision tokens and adds timestamped + video blocks. Audio is represented only by its H3 text label and is not + sent to Qwen3-VL. + """ + if isinstance(prompt, str): + return prompt + if not isinstance(prompt, dict): + raise TypeError(f"MiniMax H3 expects a string or dict prompt, got {type(prompt)!r}") + + prompt = copy.copy(prompt) + additional_information = dict(prompt.get("additional_information") or {}) + meta = dict(additional_information.get("meta") or {}) + meta.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY, None) + additional_information["meta"] = meta + prompt["additional_information"] = additional_information + + text = str(prompt.get("prompt") or "") + if not text: + raise OmniClientError("MiniMax H3 requires a non-empty prompt") + multi_modal_data = prompt.get("multi_modal_data") or {} + if not isinstance(multi_modal_data, Mapping): + raise TypeError("multi_modal_data must be a mapping") + + image_values = _items(multi_modal_data.get("image")) + videos = _items(multi_modal_data.get("video")) + audios = _audio_items(multi_modal_data.get("audio")) + diffusion_sampling = _diffusion_sampling_params(sampling_params_list) + extra_args = diffusion_sampling.extra_args or {} + task = _resolve_task(extra_args, multi_modal_data) + images = _prepare_qwen_images(task, image_values, sampling_params_list) + qwen_video_inputs: list[tuple[np.ndarray, dict[str, Any]]] = [] + condition_labels: list[tuple[str, int]] = [] + + if task == "t2va": + if images or videos or audios: + raise OmniClientError("t2va does not accept image, video, or audio conditions") + elif task == "fl2va": + if not images or videos or audios: + raise OmniClientError("fl2va requires image conditions only") + condition_labels.extend(("image", index) for index in range(1, len(images) + 1)) + elif task == "ref2va": + if not images and not videos: + raise OmniClientError("ref2va requires an image or video condition") + condition_labels.extend(("image", index) for index in range(1, len(images) + 1)) + prepared_videos: list[dict[str, Any]] = [] + artifact_dir: str | None = None + if videos: + artifact_dir = tempfile.mkdtemp(prefix="minimax_h3_ref2va_") + try: + prepared_videos = prepare_reference_videos( + videos, + target_frame_count=_ref2va_target_frame_count(sampling_params_list), + workdir=artifact_dir, + start_time_seconds=extra_args.get("start_time_seconds"), + ) + for item in prepared_videos: + sampled = sample_reference_video_frames(item["prepared_path"]) + frames = np.stack(sampled["frames"]) + frame_count = int(frames.shape[0]) + qwen_video_inputs.append( + ( + frames, + { + "total_num_frames": frame_count, + "fps": MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, + "duration": frame_count / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS, + "video_backend": "minimax_h3", + "frames_indices": list(range(frame_count)), + "do_sample_frames": False, + }, + ) + ) + except BaseException: + shutil.rmtree(artifact_dir, ignore_errors=True) + raise + audio_index = 0 + for video_index, item in enumerate(prepared_videos, start=1): + if item["input_has_audio"]: + audio_index += 1 + condition_labels.append(("audio", audio_index)) + condition_labels.append(("video", video_index)) + for _ in audios: + audio_index += 1 + condition_labels.append(("audio", audio_index)) + else: + raise OmniClientError(f"unsupported MiniMax H3 task {task!r}") + + transformed = copy.copy(prompt) + if isinstance(prompt.get("additional_information"), Mapping): + transformed["additional_information"] = dict(prompt["additional_information"]) + transformed["prompt"] = text + if task == "ref2va" and prepared_videos and artifact_dir is not None: + additional_information = dict(transformed.get("additional_information") or {}) + meta = dict(additional_information.get("meta") or {}) + meta[MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY] = serialize_prepared_reference_videos( + prepared_videos, + artifact_dir, + ) + additional_information["meta"] = meta + transformed["additional_information"] = additional_information + transformed[REQUEST_ARTIFACT_DIRS_KEY] = [artifact_dir] + qwen_mm_data = dict(multi_modal_data) + qwen_mm_data.pop("audio", None) + if images: + qwen_mm_data["image"] = images + if qwen_video_inputs: + qwen_mm_data["video"] = qwen_video_inputs + transformed["multi_modal_data"] = qwen_mm_data or None + + mm_processor_kwargs = dict(prompt.get("mm_processor_kwargs") or {}) + mm_processor_kwargs[MINIMAX_H3_PRESENTATION_TASK_KEY] = task + mm_processor_kwargs[MINIMAX_H3_CONDITION_LABELS_KEY] = condition_labels + transformed["mm_processor_kwargs"] = mm_processor_kwargs + return transformed + + +def _original_prompt(prompt: Any) -> dict[str, Any]: + if isinstance(prompt, list): + prompt = prompt[0] if prompt else {} + if isinstance(prompt, dict): + return copy.copy(prompt) + if isinstance(prompt, str): + return {"prompt": prompt} + raise TypeError(f"invalid MiniMax H3 prompt type {type(prompt)!r}") + + +def _global_request_id(prompt: Mapping[str, Any]) -> str | None: + additional_information = prompt.get("additional_information") + if not isinstance(additional_information, Mapping): + return None + value = additional_information.get("global_request_id") + if isinstance(value, (list, tuple)): + value = value[0] if value else None + return str(value) if value is not None else None + + +def text_encoder2diffusion( + source_outputs: list[Any], + prompt: Any = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, +) -> dict[str, Any] | None: + """Attach Stage 0 hidden states and token tags to the original request.""" + del requires_multimodal_data, streaming_context + if not source_outputs: + return None + if len(source_outputs) != 1: + raise RuntimeError(f"MiniMax H3 diffusion requires exactly one text-encoder source, got {len(source_outputs)}") + + diffusion_prompt = _original_prompt(prompt) + source_output = source_outputs[0] + source_request_id = getattr(source_output, "request_id", None) + expected_request_id = _global_request_id(diffusion_prompt) + if ( + source_request_id is not None + and expected_request_id is not None + and str(source_request_id) != expected_request_id + ): + raise RuntimeError( + "MiniMax H3 text-encoder request ID does not match the diffusion request: " + f"source={source_request_id!r}, expected={expected_request_id!r}" + ) + + outputs = getattr(source_output, "outputs", None) + if not isinstance(outputs, list) or len(outputs) != 1: + output_count = len(outputs) if isinstance(outputs, list) else 0 + raise RuntimeError(f"MiniMax H3 text encoder must return exactly one completion, got {output_count}") + + completion = outputs[0] + payload = completion.multimodal_output + if not isinstance(payload, Mapping): + raise RuntimeError("MiniMax H3 text encoder returned no conditioning payload") + hidden_states = payload.get("hidden_states") + if not isinstance(hidden_states, Mapping): + raise RuntimeError("MiniMax H3 text encoder returned no hidden_states payload") + hidden = hidden_states.get("output") + if not isinstance(hidden, torch.Tensor): + raise RuntimeError("MiniMax H3 text encoder returned no hidden_states.output tensor") + meta = payload.get("meta") + if not isinstance(meta, Mapping): + raise RuntimeError("MiniMax H3 text encoder returned no conditioning metadata") + token_tags = meta.get("token_role_ids") + if not isinstance(token_tags, torch.Tensor): + raise RuntimeError("MiniMax H3 text encoder returned no token_role_ids tensor") + if token_tags.ndim != 2 or token_tags.shape[-1] != 1: + raise RuntimeError( + f"MiniMax H3 stage-wire token_tags must have shape [tokens, 1], got {tuple(token_tags.shape)}" + ) + try: + conditioning = MiniMaxH3TextConditioning.from_payload( + { + "hidden_states": hidden, + "token_tags": token_tags.squeeze(-1), + } + ) + except ValueError as exc: + raise RuntimeError(str(exc)) from exc + + additional_information = dict(diffusion_prompt.get("additional_information") or {}) + additional_information["text_encoder_output"] = conditioning.to_payload() + diffusion_prompt["additional_information"] = additional_information + return diffusion_prompt diff --git a/vllm_omni/model_executor/stage_input_processors/minimax_music3.py b/vllm_omni/model_executor/stage_input_processors/minimax_music3.py new file mode 100644 index 000000000..44037e4fc --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/minimax_music3.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-input processors for MiniMax Music 3: AR talker -> acoustic decoder. + +Three entry points, all referenced from +``vllm_omni/model_executor/models/minimax_music3/pipeline.py``: + +* ``expand_cfg_prompts`` (stage 0 ``prompt_expand_func``) builds the + unconditioned twin of a conditioned request. Classifier-free guidance is + mandatory for this checkpoint, so in practice every request expands. +* ``ar2acoustic_async_chunk`` (stage 0 ``async_chunk_process_next_stage_input_func``) + forwards one AR frame span per call to stage 1, and is also the flush point + for the request's tail. +* ``ar2acoustic`` (stage 1 ``sync_process_input_func``) builds the stage-1 + engine input for the non-streaming path, where the whole song arrives at + once after stage 0 finishes. + +Stage-0 emit contract +--------------------- +Per emission the talker publishes a multimodal output shaped like:: + + { + "latent": FloatTensor[1, T, 32768] or FloatTensor[T, 32768], + "codes": {"audio": LongTensor[T, 8]}, + "meta": { + "num_processed_tokens": int, # global AR frame index of frame 0 + "last_chunk": bool, # True only on the terminal emission + "audio_seed": int, # optional + }, + } + +The span is normally one window: ``T == AR_CHUNK_FRAMES``, windows advance by +``AR_CHUNK_HOP_FRAMES``, and window ``k`` starts at AR frame +``k * AR_CHUNK_HOP_FRAMES``. The terminal emission is the exception. A window +cannot be proven non-final until a further frame exists, so the talker holds +the most recently completed window back and, on the step where it knows the +request is ending, publishes the whole remaining frame span at once. That span +covers up to two windows (the held-back complete one plus the truncated final +one), so ``T`` can reach ``AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES``. Stage 1 +re-derives the window boundaries inside a span from +``meta.num_processed_tokens`` and the span length through +``models/minimax_music3/chunking.chunk_windows``. + +``meta.num_processed_tokens`` is required rather than derived: once a terminal +span carries more than one window, a running ``chunk_seq * hop`` count no +longer identifies the span's first frame. The derived value is kept only as a +fallback for pre-terminal emissions. + +``hidden_states.output`` is accepted as an alias for ``latent`` because both +are legal homes for a float payload; ``latent`` is canonical and matches the +stage's ``engine_output_type``. + +Stage-0 -> stage-1 payload +-------------------------- +Every key below is a declared field of ``OmniPayloadStruct`` / +``MetaStruct``, which are ``forbid_unknown_fields=True`` msgspec structs, so +an undeclared key is a construction error rather than silent data loss:: + + latent FloatTensor[T, 32768] span conditioning, bf16 + codes.audio LongTensor[T, 8] RVQ codes for the span + meta.chunk_seq int emission index; is_first == (chunk_seq == 0) + meta.num_processed_tokens int absolute AR frame index of the span start + meta.codec_chunk_frames int T; span end (exclusive) = start + T + meta.last_chunk bool terminal span + meta.stream_finished BoolTensor terminal marker that survives transport + meta.audio_seed int seed for the acoustic stage's noise draw + meta.req_id list[str] single-element external request id + meta.finished BoolTensor overwritten by the chunk adapter + +``codes.audio`` is rank 2 on purpose: the chunk transfer adapter reroutes rank-1 +code tensors onto the token path (``.tolist()``), while rank-2 tensors ride the +payload channel and leave a single ``[0]`` placeholder prompt token per chunk. +``meta.finished`` is stripped before the payload reaches the stage-1 model, so +``meta.stream_finished`` and ``meta.last_chunk`` are what stage 1 reads. + +The overlap crop is deliberately NOT in the payload. Stage 1 recovers it from +the span's global frame range and ``last_chunk`` through +``models/minimax_music3/chunking.crop_sample_bounds``, which is the single +owner of that arithmetic; duplicating it here would let the two copies drift. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct +from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AR_CFG_SCALE, + AR_CHUNK_HOP_FRAMES, + AR_HIDDEN_SIZE, + NUM_CODEBOOKS, +) +from vllm_omni.model_executor.models.minimax_music3.prompt import build_cfg_null_token_ids +from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt + +__all__ = [ + "CFG_UNCOND_SUFFIX", + "ar2acoustic", + "ar2acoustic_async_chunk", + "expand_cfg_prompts", +] + +logger = init_logger(__name__) + +CFG_UNCOND_SUFFIX = "__cfg_uncond" + +_COND = "cond" +_UNCOND = "uncond" + +# Per-request forwarder state, parked on the transfer manager's request map. +_STATE_KEY = "_minimax_music3_async_state" + +# The window conditioning is emitted in bfloat16. Each chunk is its own +# /dev/shm segment and consecutive windows overlap by half, so a float32 +# 200-frame window would move 26 MB of mostly-redundant bytes per hop against +# 13 MB in bfloat16. The backbone runs in bfloat16 anyway, so this discards no +# information the AR stage actually produced. +_LATENT_TRANSPORT_DTYPE = torch.bfloat16 + + +# --------------------------------------------------------------------------- +# Stage 0: classifier-free-guidance prompt expansion +# --------------------------------------------------------------------------- + + +def _prompt_token_ids(prompt: dict[str, Any] | str | Any) -> list[int] | None: + """Return the prompt's token ids, or ``None`` when only text is available.""" + token_ids: Any = None + if isinstance(prompt, dict): + token_ids = prompt.get("prompt_token_ids") + elif isinstance(prompt, (list, tuple)): + token_ids = prompt + if token_ids is None: + return None + if isinstance(token_ids, torch.Tensor): + token_ids = token_ids.reshape(-1).tolist() + if not isinstance(token_ids, (list, tuple)) or not token_ids: + return None + return [int(t) for t in token_ids] + + +def _uncond_extra_args(extra_args: dict[str, Any], cfg_scale: float) -> dict[str, Any]: + """Clone the conditioned row's extra_args for the unconditioned twin. + + Every key is carried through, not only the ``cfg_*`` ones: the stage runs + with ``has_sampling_extra_args`` on, so extra_args reach the model's + ``forward`` per row, and dropping a key here would make the two rows + execute under different settings and decode out of lockstep. + ``cfg_null_prompt`` is the exception: the companion IS the null prompt. + """ + companion = dict(extra_args) + companion.pop("cfg_null_prompt", None) + companion["cfg_role"] = _UNCOND + companion["cfg_scale"] = cfg_scale + return companion + + +def expand_cfg_prompts( + prompt: dict[str, Any] | str, + sampling_params: Any, +) -> list[ExpandedPrompt]: + """Emit the unconditioned companion for a guided MiniMax Music 3 request. + + Classifier-free guidance is not optional for this checkpoint, and the + deploy YAMLs set ``cfg_role: cond`` / ``cfg_scale: 1.5`` as stage-0 + defaults, so an ordinary request expands into exactly one companion and + the pair decodes in lockstep. + + The companion's token ids come from ``build_cfg_null_token_ids``, which + overwrites the caption-and-lyrics span with ``<|audio_cfg|>`` while keeping + the framing tokens. Token ids are preferred over text because the two rows + are only blendable when they have identical length and an identical + ``<|audio_start|>`` position, and only the token-id form guarantees that + without a tokenizer round trip. + + Args: + prompt: The raw user prompt. A dict carrying ``prompt_token_ids`` is + the preferred form; a text-only prompt must supply a pre-built, + length-matched ``cfg_null_prompt`` in ``extra_args``. + sampling_params: Stage-0 sampling params for the conditioned row. + + Returns: + A one-element list with the unconditioned companion, or an empty list + when the request is not CFG-guided. + + Raises: + ValueError: If the request is guided but neither prompt token ids nor + a pre-built ``cfg_null_prompt`` is available. + """ + extra_args = getattr(sampling_params, "extra_args", None) or {} + if extra_args.get("cfg_role") != _COND: + # Not guided, or this IS the companion. Either way, nothing to expand. + return [] + try: + cfg_scale = float(extra_args.get("cfg_scale", AR_CFG_SCALE)) + except (TypeError, ValueError): + logger.warning( + "MiniMax Music 3: ignoring non-numeric cfg_scale %r; request decodes unguided", + extra_args.get("cfg_scale"), + ) + return [] + if cfg_scale <= 1.0: + return [] + + token_ids = _prompt_token_ids(prompt) + companion_prompt: dict[str, Any] | str + if token_ids is not None: + null_ids = build_cfg_null_token_ids(token_ids) + if isinstance(prompt, dict): + companion_prompt = {**prompt, "prompt_token_ids": null_ids} + # Drop the text form so the input processor does not re-tokenize + # and hand back a row of a different length. + companion_prompt.pop("prompt", None) + companion_prompt.pop("prompt_embeds", None) + else: + companion_prompt = {"prompt_token_ids": null_ids} + else: + null_prompt = extra_args.get("cfg_null_prompt") + if not isinstance(null_prompt, str) or not null_prompt: + raise ValueError( + "MiniMax Music 3 CFG request carries neither prompt_token_ids nor a " + "pre-built cfg_null_prompt in extra_args. Submit the prompt as token " + "ids (preferred; the null twin is then built by " + "models.minimax_music3.prompt.build_cfg_null_token_ids) or attach a " + "length-matched cfg_null_prompt string." + ) + if isinstance(prompt, dict): + companion_prompt = {**prompt, "prompt": null_prompt} + companion_prompt.pop("prompt_token_ids", None) + companion_prompt.pop("prompt_embeds", None) + else: + companion_prompt = null_prompt + + return [ + ExpandedPrompt( + prompt=companion_prompt, + role=_UNCOND, + request_id_suffix=CFG_UNCOND_SUFFIX, + sampling_params_override={"extra_args": _uncond_extra_args(extra_args, cfg_scale)}, + ) + ] + + +# --------------------------------------------------------------------------- +# Stage 0 -> stage 1: async-chunk window forwarder +# --------------------------------------------------------------------------- + + +def _payload_section(payload: Any, key: str) -> dict[str, Any]: + """Return a nested payload section as a dict (empty when absent).""" + if not isinstance(payload, dict): + return {} + section = payload.get(key) + return section if isinstance(section, dict) else {} + + +def _window_latent(payload: Any) -> torch.Tensor | None: + """Return this window's conditioning as ``[T, 32768]``, or ``None``. + + ``latent`` is the canonical key; ``hidden_states.output`` is accepted for + talkers that publish through the hidden-state channel instead. A leading + singleton batch dimension is squeezed so stage 1 sees a stable rank. + """ + if not isinstance(payload, dict): + return None + tensor = payload.get("latent") + if not isinstance(tensor, torch.Tensor): + tensor = _payload_section(payload, "hidden_states").get("output") + if not isinstance(tensor, torch.Tensor) or tensor.numel() == 0: + return None + if tensor.ndim == 3: + if tensor.shape[0] != 1: + raise ValueError( + f"MiniMax Music 3 window latent must have a singleton batch dim; got {tuple(tensor.shape)}" + ) + tensor = tensor[0] + if tensor.ndim != 2: + raise ValueError(f"MiniMax Music 3 window latent must be rank 2 or 3; got {tuple(tensor.shape)}") + if tensor.shape[1] != AR_HIDDEN_SIZE: + raise ValueError(f"MiniMax Music 3 window latent has {tensor.shape[1]} features; expected {AR_HIDDEN_SIZE}") + return tensor + + +def _window_codes(payload: Any) -> torch.Tensor | None: + """Return this window's RVQ codes as ``[T, 8]``, or ``None``.""" + codes = _payload_section(payload, "codes").get("audio") + if not isinstance(codes, torch.Tensor) or codes.numel() == 0: + return None + if codes.ndim == 1: + if codes.numel() % NUM_CODEBOOKS != 0: + raise ValueError( + f"MiniMax Music 3 flat code payload of length {codes.numel()} is not a " + f"multiple of {NUM_CODEBOOKS} codebooks" + ) + codes = codes.reshape(-1, NUM_CODEBOOKS) + if codes.ndim != 2 or codes.shape[1] != NUM_CODEBOOKS: + raise ValueError(f"MiniMax Music 3 codes must be [T, {NUM_CODEBOOKS}]; got {tuple(codes.shape)}") + return codes.to(torch.long) + + +def _request_audio_seed(request: Any, payload: Any) -> int | None: + """Resolve the seed the acoustic stage draws its noise from. + + A seed published by the talker wins, because a talker that already + consumed the request seed knows which value it used. Otherwise fall back + to the stage-0 request's own sampling seed. + """ + published = _payload_section(payload, "meta").get("audio_seed") + if isinstance(published, int): + return published + seed = getattr(getattr(request, "sampling_params", None), "seed", None) + return int(seed) if isinstance(seed, int) else None + + +def _forwarder_state(transfer_manager: Any, request_id: str) -> dict[str, Any]: + """Return (creating if needed) the per-request forwarder state.""" + request_state = transfer_manager.request_payload.get(request_id) + if not isinstance(request_state, dict) or _STATE_KEY not in request_state: + request_state = {_STATE_KEY: {"chunk_seq": 0, "terminal_sent": False}} + transfer_manager.request_payload[request_id] = request_state + return request_state[_STATE_KEY] + + +def ar2acoustic_async_chunk( + transfer_manager: Any, + multimodal_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Forward one AR frame span to the acoustic stage. + + The talker publishes whole spans, so this is a relay plus bookkeeping + rather than an accumulator: it stamps the emission index, the span's global + start frame, its length and the terminal flags, then hands the tensors to + the connector. + + This hook is the flush point for the request's tail. The AR scheduler calls + ``save_async`` on the step that stops the request, including when that + step produced no inter-stage output, so this function is guaranteed one + call with ``is_finished=True`` (omni_ar_scheduler.py:551, base + chunk_transfer_adapter.py:341). Whatever the talker published on that step + therefore reaches stage 1 in the same hop. + + Returns ``None`` on steps where the talker released no span and the request + is still running. + """ + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + state = _forwarder_state(transfer_manager, request_id) + + latent = _window_latent(multimodal_output) + codes = _window_codes(multimodal_output) + published_meta = _payload_section(multimodal_output, "meta") + published_last = published_meta.get("last_chunk") is True + + if state["terminal_sent"]: + if latent is not None: + logger.error( + "MiniMax Music 3 talker emitted %d more frames for request %s after the " + "terminal span; the downstream stream is already closed and they are dropped", + int(latent.shape[0]), + request_id, + ) + return None + + if latent is None: + if not finished: + return None + # Terminal step with nothing left to send: still deliver the marker so + # stage 1 can flush and complete instead of polling until timeout. + state["terminal_sent"] = True + # The unconditioned twin never publishes a span: the pair shares one + # song and the conditioned row emits it. Its silence is the contract, + # not a failure. + if state["chunk_seq"] == 0 and not request_id.endswith(CFG_UNCOND_SUFFIX): + logger.error("MiniMax Music 3 talker produced no audio for request %s", request_id) + return OmniPayloadStruct( + meta=MetaStruct( + chunk_seq=state["chunk_seq"], + last_chunk=True, + stream_finished=torch.tensor(True, dtype=torch.bool), + finished=torch.tensor(True, dtype=torch.bool), + req_id=[request_id], + ) + ) + + span_frames = int(latent.shape[0]) + if codes is not None and codes.shape[0] != span_frames: + raise ValueError( + f"MiniMax Music 3 span mismatch for request {request_id}: latent has " + f"{span_frames} frames but codes have {codes.shape[0]}" + ) + + chunk_seq = int(state["chunk_seq"]) + state["chunk_seq"] = chunk_seq + 1 + + # A span is terminal when the engine stopped the request or the talker said + # so. Only the engine's verdict closes the stream, so a talker that flags + # a span last while generation continues is reported rather than obeyed. + terminal = bool(finished or published_last) + if published_last and not finished: + logger.warning( + "MiniMax Music 3 talker flagged a terminal span for request %s while the " + "engine still has the request running", + request_id, + ) + state["terminal_sent"] = terminal + + span_start = published_meta.get("num_processed_tokens") + if not isinstance(span_start, int): + # Fallback for talkers that publish one hop-aligned window per step. + # It is only sound before the terminal span, which may cover more than + # one window and therefore break the running count. + span_start = chunk_seq * AR_CHUNK_HOP_FRAMES + if terminal and span_frames > AR_CHUNK_HOP_FRAMES: + logger.warning( + "MiniMax Music 3 terminal span for request %s carries %d frames but no " + "meta.num_processed_tokens; the derived start frame %d assumes a single " + "hop-aligned window and may misplace the tail", + request_id, + span_frames, + span_start, + ) + + finished_flag = torch.tensor(terminal, dtype=torch.bool) + meta = MetaStruct( + chunk_seq=chunk_seq, + num_processed_tokens=int(span_start), + codec_chunk_frames=span_frames, + last_chunk=terminal, + # meta.finished is stripped before the payload reaches the stage-1 + # model; stream_finished is what survives and drives the final flush. + stream_finished=finished_flag, + finished=finished_flag, + req_id=[request_id], + ) + audio_seed = _request_audio_seed(request, multimodal_output) + if audio_seed is not None: + meta.audio_seed = audio_seed + + payload = OmniPayloadStruct( + latent=latent.detach().to(_LATENT_TRANSPORT_DTYPE).contiguous(), + meta=meta, + ) + if codes is not None: + payload.codes = CodesStruct(audio=codes.detach().contiguous()) + return payload + + +# --------------------------------------------------------------------------- +# Stage 1: non-streaming engine-input builder +# --------------------------------------------------------------------------- + + +def _empty_acoustic_prompt() -> OmniTokensPrompt: + """Zero-length stage-1 request. + + Stage 1 decodes it to empty audio, which the serving layer's no-audio + guard turns into a request-level error. That is preferable to raising + here, where the failure would surface as a stalled request. + """ + return OmniTokensPrompt( + prompt_token_ids=[], + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=None, + ) + + +def ar2acoustic( + source_outputs: list[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + """Build stage-1 engine inputs for the non-streaming path. + + With ``async_chunk: false`` the whole song arrives once stage 0 has + finished, so the song is handed over as a single logical window and stage + 1 does its own windowing through ``chunking.chunk_windows``. The payload + layout matches the streaming one key for key, so the acoustic stage reads + one contract in both modes. + + The prompt is ``[0] * num_windows``: the generation scheduler sizes the + request from ``prompt_token_ids``, and one placeholder per acoustic window + keeps that budget proportional to the work without shipping the + conditioning through the token path. + """ + del prompt, _requires_multimodal_data + + acoustic_inputs: list[OmniTokensPrompt] = [] + for ar_output in source_outputs: + if not ar_output.finished: + continue + output = ar_output.outputs[0] + multimodal_output = getattr(output, "multimodal_output", None) + + try: + latent = _window_latent(multimodal_output) + codes = _window_codes(multimodal_output) + except ValueError as exc: + logger.warning("Skipping malformed MiniMax Music 3 stage-0 output: %s", exc) + acoustic_inputs.append(_empty_acoustic_prompt()) + continue + + if latent is None: + logger.error("MiniMax Music 3 talker produced no audio for request %s", ar_output.request_id) + acoustic_inputs.append(_empty_acoustic_prompt()) + continue + + total_frames = int(latent.shape[0]) + meta: dict[str, Any] = { + "chunk_seq": 0, + "num_processed_tokens": 0, + "codec_chunk_frames": total_frames, + "last_chunk": True, + "stream_finished": torch.tensor(True, dtype=torch.bool), + "req_id": [ar_output.request_id], + } + published_seed = _payload_section(multimodal_output, "meta").get("audio_seed") + if isinstance(published_seed, int): + meta["audio_seed"] = published_seed + + additional_information: dict[str, Any] = { + "latent": latent.detach().to(_LATENT_TRANSPORT_DTYPE).contiguous(), + "meta": meta, + } + if codes is not None: + additional_information["codes"] = {"audio": codes.detach().contiguous()} + + num_windows = max(1, -(-total_frames // AR_CHUNK_HOP_FRAMES)) + acoustic_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * num_windows, + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=additional_information, + ) + ) + return acoustic_inputs diff --git a/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py index 48de5a380..918db4ad5 100644 --- a/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py +++ b/vllm_omni/model_executor/stage_input_processors/nemotron_voicechat.py @@ -11,9 +11,10 @@ loop), so the prompt slots are reconstructed here, and the timeline metadata rides ``additional_information``. talker (1) -> code2wav (2): PersonaPlex-style full payload. The talker -accumulates one 31-quantizer code stack per frame under ``("codes","audio")``; -the producer ships the prompt-trimmed ``[frames, 31]`` stack once the request -finishes, and the sync placeholder builder sizes stage-2's dummy prompt. +accumulates one 31-quantizer code stack per frame under ``("codes","audio")``. +The synchronous producer ships the prompt-trimmed ``[frames, 31]`` stack once +the request finishes. Native duplex wakes expose that same cumulative stack, +so the streaming producer must delta it before feeding code2wav's causal cache. """ from __future__ import annotations @@ -28,6 +29,7 @@ from vllm_omni.data_entry_keys import ( MetaStruct, OmniPayloadStruct, ) +from vllm_omni.model_executor.models.nemotron_voicechat.runtime_info import scalar_bool # Full-sequence tensors are re-shipped whole; the full-payload accumulator must # replace (not concatenate) them if the producer fires more than once. The AR @@ -39,6 +41,7 @@ _FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"codes", "codes.audio", # stt pad_token. The thinker also reports it in its latent metadata, which # takes precedence when present. _DEFAULT_TEXT_PAD_ID = 12 +_STEP_TOKEN_SNAPSHOT_MISSING = object() def _info_get(info: Any, key: str) -> Any: @@ -51,6 +54,27 @@ def _info_get(info: Any, key: str) -> Any: return None +def _request_info(request: Any) -> dict[str, Any]: + """Return the live runner payload, accepting both transport generations.""" + for name in ( + "model_intermediate_buffer", + "additional_information", + "additional_information_cpu", + ): + info = getattr(request, name, None) + if isinstance(info, dict): + return info + if info is not None and hasattr(info, "entries"): + from vllm_omni.engine.serialization import ( + deserialize_additional_information, + ) + + decoded = deserialize_additional_information(info) + if isinstance(decoded, dict): + return decoded + return {} + + def thinker2talker_token_only( source_outputs: list, prompt: Any = None, @@ -118,32 +142,60 @@ def thinker2talker_async_chunk( the payload self-contained. ``meta.num_processed_tokens`` carries the logical prompt length for the code2wav producer's prompt-region trim. """ - del multimodal_output, kwargs + step_snapshot = kwargs.pop("new_token_ids", _STEP_TOKEN_SNAPSHOT_MISSING) + has_step_snapshot = step_snapshot is not _STEP_TOKEN_SNAPSHOT_MISSING + generated_delta = [int(token_id) for token_id in ((step_snapshot or ()) if has_step_snapshot else ())] request_id = request.external_req_id + request_info = _request_info(request) + duplex_info = request_info.get("duplex") if isinstance(request_info, dict) else None + is_duplex = isinstance(duplex_info, dict) and duplex_info.get("data_plane") is True + + state = transfer_manager.request_payload.get(request_id) + generated_now = [int(token_id) for token_id in (getattr(request, "output_token_ids", None) or [])] prompt_ids = list(getattr(request, "prompt_token_ids", None) or []) - # vLLM prompt = logical prompt + 1 placeholder (acoustic frame 0). - logical_prompt_len = max(len(prompt_ids) - 1, 0) - generated = list(getattr(request, "output_token_ids", None) or []) + if is_duplex: + runtime_config = duplex_info.get("runtime_config") + runtime_prompt = runtime_config.get("nvc_prompt_token_ids") if isinstance(runtime_config, dict) else None + if not isinstance(runtime_prompt, list | tuple) or not runtime_prompt: + raise ValueError("Nemotron VoiceChat duplex stage transfer lacks nvc_prompt_token_ids") + logical_prompt_len = len(runtime_prompt) + previous = list(state.get("nvc_duplex_text_tokens", ())) if isinstance(state, dict) else [] + if has_step_snapshot: + generated = previous + generated_delta + elif generated_now[: len(previous)] == previous and len(generated_now) > len(previous): + generated = generated_now + elif generated_now: + generated = previous + [generated_now[-1]] + else: + generated = previous + else: + # vLLM prompt = logical prompt + 1 placeholder (acoustic frame 0). + logical_prompt_len = max(len(prompt_ids) - 1, 0) + generated = generated_now pad_id = _DEFAULT_TEXT_PAD_ID - reported_pad = _info_get(getattr(request, "additional_information", None), "nvc_text_pad_id") + reported_pad = _info_get(request_info, "nvc_text_pad_id") if reported_pad is not None: pad_id = int(reported_pad) - timeline = [pad_id] * logical_prompt_len + [int(t) for t in generated] + timeline = [pad_id] * logical_prompt_len + generated # Skip no-progress wakeups (same length as the last emitted chunk). - state = transfer_manager.request_payload.get(request_id) last_len = int(state.get("nvc_timeline_len", -1)) if isinstance(state, dict) else -1 if len(timeline) <= last_len and not is_finished: return None - transfer_manager.request_payload[request_id] = {"nvc_timeline_len": len(timeline)} + transfer_manager.request_payload[request_id] = { + "nvc_timeline_len": len(timeline), + "nvc_duplex_text_tokens": generated if is_duplex else (), + } return OmniPayloadStruct( - ids=IdsStruct(all=timeline), + ids=IdsStruct(all=timeline, prompt=[0]), meta=MetaStruct( - finished=torch.tensor(bool(is_finished), dtype=torch.bool), + finished=torch.tensor(bool(is_finished and not is_duplex), dtype=torch.bool), num_processed_tokens=logical_prompt_len, # The talker's vLLM prompt stays a single placeholder token. next_stage_prompt_len=1, + codec_streaming=is_duplex, + request_id=request_id if is_duplex else None, ), ) @@ -155,18 +207,30 @@ def talker2code2wav_async_chunk( is_finished: bool = False, **kwargs: Any, ) -> OmniPayloadStruct | None: - """Streaming producer: ship prompt-trimmed CUMULATIVE code stacks in chunks. + """Streaming producer: ship prompt-trimmed code stacks in chunks. - Follows the NeMo incremental-decode recipe (``decode_one_audio_step``): - each chunk carries the FULL trimmed code history and - ``meta.left_context_size`` = frames the codec already emitted, so the - code2wav stage re-decodes the whole prefix and slices off only the new - samples — no seams, no crossfade, and the concatenated stream equals the - prefix-decode of the final stack. Chunk cadence comes from the connector - config key ``codec_chunk_frames`` (default 13 frames ~= 1 s of audio). + The regular path follows NeMo's prefix-decode recipe: each chunk carries + the full trimmed history and ``meta.left_context_size`` identifies frames + already emitted. Native duplex talker wakes also carry cumulative history, + but code2wav keeps causal codec state for that path and therefore must see + only rows not observed on an earlier wake. Chunk cadence comes from + ``codec_chunk_frames`` (default 13 frames ~= 1 s of audio). """ - del kwargs request_id = request.external_req_id + del kwargs + request_info = _request_info(request) + codec_streaming = scalar_bool(_info_get(request_info, "codec_streaming")) or scalar_bool( + _info_get(multimodal_output, "codec_streaming") + ) + state = transfer_manager.request_payload.get(request_id) + if isinstance(state, dict): + codec_streaming = codec_streaming or bool(state.get("nvc_codec_stream_started", False)) + # ``is_finished`` describes the current resumable Stage-1 scheduler + # segment, not the lifetime of a native-duplex codec stream. Marking each + # segment terminal makes the connector clean sender state after its first + # chunk; every later one-row wake is then mistaken for a fresh stream and + # trimmed as prompt context. Duplex lifetime is closed by session abort. + stream_finished = bool(is_finished and not codec_streaming) codes = None if isinstance(multimodal_output, dict): @@ -175,9 +239,19 @@ def talker2code2wav_async_chunk( if codes is None: codes = multimodal_output.get("codes.audio") if not isinstance(codes, torch.Tensor) or codes.numel() == 0: - if not is_finished: + if not stream_finished: + # Returning ``None`` at a resumable segment boundary makes the + # generic chunk adapter synthesize ``is_segment_finished=True``. + # Stage 2 then stops polling after the first 80 ms wake. An + # explicit non-terminal empty payload consumes this connector + # sequence number without closing the long-lived codec stream. + if codec_streaming: + return _empty_streaming_payload(request_id=request_id) return None - return _empty_finished_payload() + return _empty_finished_payload( + codec_streaming=codec_streaming, + request_id=request_id, + ) if codes.ndim == 1: codes = codes.reshape(1, -1) @@ -186,15 +260,13 @@ def talker2code2wav_async_chunk( # request-level additional_information is overwritten both by arriving # thinker chunks and by talker info updates, so it races. Cache the first # sighting in the transfer-manager state as a fallback. - state = transfer_manager.request_payload.get(request_id) prompt_len = _info_get(multimodal_output if isinstance(multimodal_output, dict) else None, "nvc_logical_prompt_len") if prompt_len is None and isinstance(multimodal_output, dict): prompt_len = multimodal_output.get("meta.nvc_logical_prompt_len") if prompt_len is None: - info = getattr(request, "additional_information", None) - prompt_len = _info_get(info, "nvc_logical_prompt_len") + prompt_len = _info_get(request_info, "nvc_logical_prompt_len") if prompt_len is None: - prompt_len = _info_get(info, "num_processed_tokens") + prompt_len = _info_get(request_info, "num_processed_tokens") if prompt_len is None and isinstance(state, dict): prompt_len = state.get("nvc_prompt_len") if prompt_len is None: @@ -203,41 +275,111 @@ def talker2code2wav_async_chunk( "('nvc_logical_prompt_len' or async meta.num_processed_tokens); cannot " "trim the prompt-region codes for streaming code2wav." ) - trimmed = codes[max(int(prompt_len) - 1, 0) :] - frames_sent = int(state.get("nvc_frames_sent", 0)) if isinstance(state, dict) else 0 - new_frames = int(trimmed.shape[0]) - frames_sent - if new_frames <= 0 and not is_finished: + if codec_streaming: + # The AR runner exposes its cumulative model_intermediate_buffer on + # every wake. Feeding that whole prefix into a stateful causal codec + # replays history into its cache (live traces showed lengths + # 19,20,20,21,21,...) and corrupts the waveform. Track rows OBSERVED, + # separately from rows EMITTED at the chunk cadence, and ship only the + # suffix first seen on this wake. + full_trimmed = codes[max(int(prompt_len) - 1, 0) :].detach().to(dtype=torch.long, device="cpu") + frames_seen = int(state.get("nvc_frames_seen", 0)) if isinstance(state, dict) else 0 + if int(full_trimmed.shape[0]) < frames_seen: + raise ValueError( + "NemotronVoiceChat duplex talker code history shrank from " + f"{frames_seen} to {int(full_trimmed.shape[0])} frames; cannot " + "safely continue the stateful code2wav stream." + ) + trimmed = full_trimmed[frames_seen:] + frames_seen = int(full_trimmed.shape[0]) + pending = state.get("nvc_pending_codes") if isinstance(state, dict) else None + if isinstance(pending, torch.Tensor) and pending.numel() > 0: + trimmed = torch.cat([pending.reshape(-1, codes.shape[-1]), trimmed], dim=0) + new_frames = int(trimmed.shape[0]) + else: + trimmed = codes[max(int(prompt_len) - 1, 0) :] + new_frames = int(trimmed.shape[0]) - frames_sent + if new_frames <= 0 and not stream_finished: + if codec_streaming: + transfer_manager.request_payload[request_id] = { + "nvc_frames_sent": frames_sent, + "nvc_frames_seen": frames_seen, + "nvc_prompt_len": int(prompt_len), + "nvc_codec_stream_started": True, + "nvc_pending_codes": trimmed, + } + return _empty_streaming_payload(request_id=request_id) return None connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} chunk_frames = int(cfg.get("codec_chunk_frames", 13)) - if new_frames < chunk_frames and not is_finished: + if new_frames < chunk_frames and not stream_finished: + if codec_streaming: + transfer_manager.request_payload[request_id] = { + "nvc_frames_sent": frames_sent, + "nvc_frames_seen": frames_seen, + "nvc_prompt_len": int(prompt_len), + "nvc_codec_stream_started": True, + "nvc_pending_codes": trimmed, + } + return _empty_streaming_payload(request_id=request_id) return None - if new_frames <= 0 and is_finished and frames_sent > 0: + if new_frames <= 0 and stream_finished and frames_sent > 0: # Everything already shipped; emit a terminal empty chunk so the # code2wav request observes meta.finished. - return _empty_finished_payload() + return _empty_finished_payload( + codec_streaming=codec_streaming, + request_id=request_id, + ) transfer_manager.request_payload[request_id] = { - "nvc_frames_sent": int(trimmed.shape[0]), + "nvc_frames_sent": frames_sent + new_frames if codec_streaming else int(trimmed.shape[0]), + **({"nvc_frames_seen": frames_seen} if codec_streaming else {}), "nvc_prompt_len": int(prompt_len), + "nvc_codec_stream_started": codec_streaming, } + code_payload = trimmed return OmniPayloadStruct( - codes=CodesStruct(audio=trimmed.detach().to(dtype=torch.long, device="cpu")), + codes=CodesStruct(audio=code_payload.detach().to(dtype=torch.long, device="cpu")), meta=MetaStruct( - left_context_size=frames_sent, - finished=torch.tensor(bool(is_finished), dtype=torch.bool), + left_context_size=0 if codec_streaming else frames_sent, + codec_streaming=codec_streaming, + request_id=request_id if codec_streaming else None, + finished=torch.tensor(stream_finished, dtype=torch.bool), + # Native duplex scheduler segments are transport wakes, not codec + # stream boundaries. Preserve polling across successive wakes. + is_segment_finished=(torch.tensor(False, dtype=torch.bool) if codec_streaming else None), ), ) -def _empty_finished_payload() -> OmniPayloadStruct: +def _empty_streaming_payload(*, request_id: str) -> OmniPayloadStruct: return OmniPayloadStruct( codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), - meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + meta=MetaStruct( + codec_streaming=True, + request_id=request_id, + finished=torch.tensor(False, dtype=torch.bool), + is_segment_finished=torch.tensor(False, dtype=torch.bool), + ), + ) + + +def _empty_finished_payload( + *, + codec_streaming: bool = False, + request_id: str | None = None, +) -> OmniPayloadStruct: + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + finished=torch.tensor(True, dtype=torch.bool), + codec_streaming=codec_streaming, + request_id=request_id if codec_streaming else None, + ), ) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 42377faa8..42a9b19e3 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -23,8 +23,7 @@ TALKER_CODEC_END_TOKEN_ID = 8294 # ============================================================================ # Worker-connector data plane (non-async-chunk path). # Both transitions ship payloads via the worker connector -# (registered in ``_FULL_PAYLOAD_INPUT_STAGES`` in -# omni_scheduling_coordinator): +# (requires_full_payload_input=True in the pipeline topology): # - thinker->talker reads accumulated ``pooling_output["hidden"]`` and # packs an OmniPayload-shaped dict (embed.prefill / # hidden_states.output / ids.prompt / ids.output) for the talker, which @@ -155,9 +154,9 @@ def thinker2talker_token_only( ``thinker2talker_full_payload`` via the worker connector and lands in ``model_intermediate_buffer`` before the talker's forward() runs. - Consumer-wait gating is whitelist-driven via - ``_FULL_PAYLOAD_INPUT_STAGES`` (see the mixin - ``should_accumulate_full_payload_output`` docstring). + Consumer-wait gating is driven by + ``requires_full_payload_input=True`` in the pipeline topology + (see the mixin ``should_accumulate_full_payload_output`` docstring). """ thinker_outputs = source_outputs talker_inputs = [] diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index da5fa840d..423d288fe 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -1,5 +1,6 @@ """Stage input processor for Qwen3-TTS: Talker -> Code2Wav.""" +import time from collections.abc import Mapping from typing import Any @@ -14,9 +15,12 @@ from vllm_omni.data_entry_keys import ( to_dict, ) from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( + AdaptiveChunkController, + compute_adaptive_emit, compute_dynamic_initial_chunk_size, compute_ramp_emit, max_ic_for_chunk_size, + parse_adaptive_config, parse_chunk_ramp, ) from vllm_omni.model_executor.stage_input_processors.tts_utils import ( @@ -108,6 +112,18 @@ def talker2code2wav_async_chunk( transfer_manager._ramp_parsed = parse_chunk_ramp(cfg, steady=chunk_size) ramp = transfer_manager._ramp_parsed + # Adaptive: parse once per transfer_manager. + # Takes precedence over fixed ramp when both are configured. + if not hasattr(transfer_manager, "_adaptive_parsed"): + transfer_manager._adaptive_parsed = parse_adaptive_config(cfg) + ( + adaptive_enabled, + adaptive_min, + adaptive_margin, + adaptive_divisor, + adaptive_delta_min, + ) = transfer_manager._adaptive_parsed + # Per-request override takes priority over dynamic IC. fixed_initial_chunk_size = configured_initial_chunk_size > 0 initial_chunk_size = configured_initial_chunk_size @@ -124,8 +140,12 @@ def talker2code2wav_async_chunk( fixed_initial_chunk_size = True # Dynamic IC: cache per request so boundaries stay stable for its lifetime. - # Skipped when ramp is active (ramp replaces IC/steady entirely). - if ramp is None and not fixed_initial_chunk_size: + # Skipped when fixed ramp is active (ramp replaces IC/steady entirely) or + # a fixed IC is configured — unless adaptive is active, in which case + # dynamic IC always runs (used for chunk 0), even if + # initial_codec_chunk_frames is configured. + skip_dynamic_ic = (ramp is not None or fixed_initial_chunk_size) and not adaptive_enabled + if not skip_dynamic_ic: _ic_cache = getattr(transfer_manager, "_cached_ic", None) if _ic_cache is None: _ic_cache = {} @@ -172,7 +192,71 @@ def talker2code2wav_async_chunk( ) return None - if ramp is not None: + if adaptive_enabled: + _adaptive_states = getattr(transfer_manager, "_adaptive_states", None) + if _adaptive_states is None: + _adaptive_states = {} + transfer_manager._adaptive_states = _adaptive_states + + ctrl = _adaptive_states.get(request_id) + chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) + first_chunk = chunk_index <= 0 + + if ctrl is None or chunk_index == 0: + use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < chunk_size + + if use_first_chunk and length <= initial_chunk_size: + if not finished and length < initial_chunk_size: + return None + context_length = length if finished and length < initial_chunk_size else initial_chunk_size + else: + initial_coverage = initial_chunk_size if use_first_chunk else 0 + adjusted = length - initial_coverage + if not finished and adjusted % chunk_size != 0: + return None + chunk_length = adjusted % chunk_size + context_length = chunk_length if chunk_length != 0 else chunk_size + + # Chunk 0: target_size equals context_length (no overshoot possible + # — IC logic emits exactly what's accumulated up to the IC boundary). + target_size = context_length + + now = time.monotonic() + ctrl = AdaptiveChunkController( + first_emit_time=now, + last_emit_time=now, + ramp_divisor=adaptive_divisor, + ramp_delta_min=adaptive_delta_min, + ) + _adaptive_states[request_id] = ctrl + else: + now = time.monotonic() + target_size = ctrl.compute_next_chunk_size( + now, + adaptive_min, + chunk_size, + adaptive_margin, + ) + emit, context_length = compute_adaptive_emit( + length, + ctrl.accumulated_at_last_emit, + target_size, + finished, + ) + if not emit: + return None + if context_length == 0: + ctrl.log_summary(request_id) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + request_id=request_id, + left_context_size=0, + finished=torch.tensor(True, dtype=torch.bool), + ), + ) + + elif ramp is not None: chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) first_chunk = chunk_index <= 0 emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size, finished) @@ -206,6 +290,13 @@ def talker2code2wav_async_chunk( if finished and first_chunk: context_length = length + if adaptive_enabled: + ctrl = transfer_manager._adaptive_states.get(request_id) + if ctrl is not None: + ctrl.record_emit(time.monotonic(), context_length, length, target_size) + if finished: + ctrl.log_summary(request_id) + # Code2Wav keeps the quantizer/conv/Transformer context per request. Ship # only the newly completed codec frames after the first chunk. window_frames = transfer_manager.code_prompt_token_ids[request_id][-context_length:] diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index b11ceadec..f10ca9806 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -10,8 +10,10 @@ from vllm_omni.model_extras.registry import ( get_extra_output_params, get_model_class_name, get_output_tensor_range, + get_transformer_config_subfolder, get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, + should_preserve_reference_image_size, ) __all__ = [ @@ -23,6 +25,8 @@ __all__ = [ "get_extra_output_params", "get_model_class_name", "get_output_tensor_range", + "get_transformer_config_subfolder", "get_x_to_text_model_family", "should_init_extra_args_for_non_diffusion_stages", + "should_preserve_reference_image_size", ] diff --git a/vllm_omni/model_extras/audiox.py b/vllm_omni/model_extras/audiox.py deleted file mode 100644 index 220972148..000000000 --- a/vllm_omni/model_extras/audiox.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -# Request params the AudioX pipeline reads from ``sampling_params.extra_args`` -# (see ``pipeline_audiox.py`` ``forward``). Declaring them here routes -# ``extra_body`` (online) / ``--extra-body`` (offline) keys into ``extra_args`` -# instead of silently dropping them. -AUDIOX_EXTRA_BODY_PARAMS = frozenset( - { - "audiox_task", # t2a/t2m/v2a/v2m/tv2a/tv2m - "seconds_start", - "seconds_total", - "sigma_min", - "sigma_max", - "cfg_rescale", - "video_path", # v2*/tv2* conditioning - "audio_path", # optional reference-audio conditioning - } -) -# Echoed back in diffusion output metadata (pipeline_audiox.py). -AUDIOX_EXTRA_OUTPUT_PARAMS = frozenset( - { - "audiox_task", - } -) diff --git a/vllm_omni/model_extras/cosmos3.py b/vllm_omni/model_extras/cosmos3.py index d0ffa5182..01cecf94f 100644 --- a/vllm_omni/model_extras/cosmos3.py +++ b/vllm_omni/model_extras/cosmos3.py @@ -9,6 +9,7 @@ COSMOS3_EXTRA_BODY_PARAMS = frozenset( "max_sequence_length", "use_resolution_template", "use_duration_template", + "negative_metadata_mode", "use_system_prompt", "system_prompt", "negative_prompt", diff --git a/vllm_omni/model_extras/ltx2.py b/vllm_omni/model_extras/ltx2.py index 5bac58469..7516a55f2 100644 --- a/vllm_omni/model_extras/ltx2.py +++ b/vllm_omni/model_extras/ltx2.py @@ -21,7 +21,35 @@ LTX_EXTRA_BODY_PARAMS = frozenset( "audio_rescale_scale", "video_stg_blocks", "audio_stg_blocks", + "sigmas", + "stage_1_sigmas", + "stage_2_sigmas", + "image_crf", } ) LTX_EXTRA_OUTPUT_PARAMS = frozenset() + + +def ltx_preserves_reference_image_size(*, model: str | None, revision: str | None = None) -> bool: + """Keep source geometry until LTX-2.5 applies conditioning compression.""" + if model is None: + return False + from vllm_omni.diffusion.models.ltx2.ltx2_components import ( + preserves_reference_image_size as ltx_preserves_reference_image_size, + ) + + return ltx_preserves_reference_image_size(model=model, revision=revision) + + +def ltx_transformer_config_subfolder( + *, + model: str | None, + revision: str | None = None, +) -> str: + """Select the transformer config used by the requested LTX checkpoint.""" + if model is None: + return "transformer" + from vllm_omni.diffusion.models.ltx2.ltx2_components import detect_ltx_model_version + + return "transformer_full" if detect_ltx_model_version(model, revision=revision) == "2.5" else "transformer" diff --git a/vllm_omni/model_extras/magi_human.py b/vllm_omni/model_extras/magi_human.py deleted file mode 100644 index 470a1455b..000000000 --- a/vllm_omni/model_extras/magi_human.py +++ /dev/null @@ -1,20 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -# MagiHuman is a text -> video+audio model. Its model-specific knobs are read -# from sampling_params.extra_args in pipeline_magi_human.py; declaring them here -# routes request `extra_body` fields into OmniDiffusionSamplingParams.extra_args -# so the model can be driven through the standard task example. -MAGI_HUMAN_EXTRA_BODY_PARAMS = frozenset( - { - "seconds", - "audio_path", - "image_path", - "sr_height", - "sr_width", - "sr_num_inference_steps", - } -) -MAGI_HUMAN_EXTRA_OUTPUT_PARAMS = frozenset() diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 8d7ee1e29..a23dd016c 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -4,14 +4,10 @@ from __future__ import annotations from collections.abc import Callable, Mapping -from typing import Any, Literal +from typing import Any, Literal, Protocol from PIL import Image -from vllm_omni.model_extras.audiox import ( - AUDIOX_EXTRA_BODY_PARAMS, - AUDIOX_EXTRA_OUTPUT_PARAMS, -) from vllm_omni.model_extras.bagel import ( BAGEL_EXTRA_BODY_PARAMS, BAGEL_EXTRA_OUTPUT_PARAMS, @@ -34,10 +30,11 @@ from vllm_omni.model_extras.helios import ( ) from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS -from vllm_omni.model_extras.ltx2 import LTX_EXTRA_BODY_PARAMS, LTX_EXTRA_OUTPUT_PARAMS -from vllm_omni.model_extras.magi_human import ( - MAGI_HUMAN_EXTRA_BODY_PARAMS, - MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, +from vllm_omni.model_extras.ltx2 import ( + LTX_EXTRA_BODY_PARAMS, + LTX_EXTRA_OUTPUT_PARAMS, + ltx_preserves_reference_image_size, + ltx_transformer_config_subfolder, ) from vllm_omni.model_extras.mammothmodal2_preview import ( MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS, @@ -61,6 +58,7 @@ from vllm_omni.model_extras.ming_flash_omni import ( from vllm_omni.model_extras.ming_flash_omni import ( build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt, ) +from vllm_omni.model_extras.sana_video import SANA_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.sensenova_u1 import ( SENSENOVA_U1_EXTRA_BODY_PARAMS, SENSENOVA_U1_EXTRA_OUTPUT_PARAMS, @@ -96,6 +94,24 @@ XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"] +class ReferenceImageSizeResolver(Protocol): + def __call__( + self, + *, + model: str | None, + revision: str | None = None, + ) -> bool: ... + + +class TransformerConfigSubfolderResolver(Protocol): + def __call__( + self, + *, + model: str | None, + revision: str | None = None, + ) -> str: ... + + def default_x_to_text_prompt( model: str, prompt: str, @@ -156,10 +172,6 @@ def default_image_to_image_prompt( _EXTRA_SPECS: dict[str, dict[str, Any]] = { - "AudioXPipeline": { - "extra_body_params": AUDIOX_EXTRA_BODY_PARAMS, - "extra_output_params": AUDIOX_EXTRA_OUTPUT_PARAMS, - }, "BagelPipeline": { "extra_body_params": BAGEL_EXTRA_BODY_PARAMS, "extra_output_params": BAGEL_EXTRA_OUTPUT_PARAMS, @@ -180,10 +192,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS, "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS, }, - "MagiHumanPipeline": { - "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS, - "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, - }, "HeliosPipeline": { "extra_body_params": HELIOS_EXTRA_BODY_PARAMS, "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS, @@ -202,12 +210,22 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { model_class_name: { "extra_body_params": LTX_EXTRA_BODY_PARAMS, "extra_output_params": LTX_EXTRA_OUTPUT_PARAMS, + "reference_image_size_resolver": ltx_preserves_reference_image_size, } for model_class_name in ( "LTX2Pipeline", + "LTX2TwoStagePipeline", + "LTX2DistilledOneStagePipeline", "LTX2DistilledPipeline", + "LTX2DistilledTwoStagePipeline", ) }, + "SanaVideoPipeline": { + "extra_body_params": SANA_VIDEO_EXTRA_BODY_PARAMS, + }, + "SanaImageToVideoPipeline": { + "extra_body_params": SANA_VIDEO_EXTRA_BODY_PARAMS, + }, "WanVACEPipeline": { "extra_body_params": VACE_EXTRA_BODY_PARAMS, "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS, @@ -228,6 +246,10 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { }, } +for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"): + _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder + + # Multi-stage discovery reports the top-level wrapper rather than its DiT # submodule, so both names must resolve to the same request builders. _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] @@ -276,6 +298,36 @@ def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange: return spec.get("output_tensor_range", "negative_one_to_one") +def get_transformer_config_subfolder( + model_class_name: str | None, + *, + model: str | None, + revision: str | None = None, +) -> str: + """Return the model-declared DiT config subfolder, or the standard default.""" + spec = _get_spec(model_class_name) + resolver: TransformerConfigSubfolderResolver | None = ( + spec.get("transformer_config_subfolder_resolver") if spec else None + ) + return resolver(model=model, revision=revision) if resolver else "transformer" + + +def should_preserve_reference_image_size( + model_class_name: str | None, + *, + model: str | None, + revision: str | None = None, +) -> bool: + """Return whether the selected pipeline owns reference-image resizing.""" + if model_class_name is None and model is not None: + from vllm_omni.diffusion.data import resolve_model_class_name + + model_class_name = resolve_model_class_name(model, revision=revision) + spec = _get_spec(model_class_name) + resolver: ReferenceImageSizeResolver | None = spec.get("reference_image_size_resolver") if spec else None + return bool(resolver and resolver(model=model, revision=revision)) + + def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool: spec = _get_spec(model_class_name) return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) diff --git a/vllm_omni/model_extras/sana_video.py b/vllm_omni/model_extras/sana_video.py new file mode 100644 index 000000000..67234de48 --- /dev/null +++ b/vllm_omni/model_extras/sana_video.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +SANA_VIDEO_EXTRA_BODY_PARAMS = frozenset( + { + "clean_caption", + "motion_score", + "use_resolution_binning", + } +) diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index 105d0222d..90197f6e6 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -85,7 +85,9 @@ _REQUEST_OUTPUT_CONTENT_ATTRS = ( "encoder_prompt", "encoder_prompt_token_ids", "num_cached_tokens", + "num_cache_creation_tokens", "kv_transfer_params", + "ec_transfer_params", ) # Omni-specific content copied when the wrapped stage output is itself an @@ -147,7 +149,9 @@ class OmniRequestOutput(RequestOutput): encoder_prompt: str | None = None encoder_prompt_token_ids: list[int] | None = None num_cached_tokens: int | None = None + num_cache_creation_tokens: int | None = None kv_transfer_params: dict[str, Any] | None = None + ec_transfer_params: dict[str, Any] | None = None # --- Pipeline stage fields --- stage_id: int | None = None diff --git a/vllm_omni/outputs/output_metadata.py b/vllm_omni/outputs/output_metadata.py index f3e364a34..82faed499 100644 --- a/vllm_omni/outputs/output_metadata.py +++ b/vllm_omni/outputs/output_metadata.py @@ -130,7 +130,6 @@ def _validate_audio_metadata(metadata: DiffusionMetadataMapping | None) -> None: if metadata is None: return _check_optional_int(metadata, "sample_rate", group="audio", min_value=1) - _check_optional_non_empty_str(metadata, "audiox_task", group="audio") def _validate_actions_metadata(metadata: DiffusionMetadataMapping | None) -> None: diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index 9558d66e5..13844518e 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -1,3 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections.abc import Iterable from dataclasses import fields as dataclass_fields from typing import Any @@ -148,6 +152,8 @@ class OmniRequestState(RequestState): finish_reason: FinishReason | None, stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, + ec_transfer_params: dict[str, Any] | None = None, + *, routed_experts: Any = None, ) -> OmniRequestOutput | PoolingRequestOutput | None: """Create a request output from generation results. @@ -162,8 +168,13 @@ class OmniRequestState(RequestState): finish_reason: Optional finish reason indicating why generation stopped stop_reason: Optional stop reason (token ID or stop string) kv_transfer_params: Optional KV cache transfer parameters + ec_transfer_params: Optional encoder-cache transfer parameters + (6th positional, matching upstream RequestState so that + super().process_outputs() calls do not misroute it). routed_experts: Optional MoE routed-expert ids for this step, attached to the completion output for generation stages + (omni-specific keyword; upstream moved this accumulation + into RequestState.routed_experts_chunks). Returns: OmniRequestOutput or PoolingRequestOutput if output should be @@ -177,6 +188,7 @@ class OmniRequestState(RequestState): finish_reason, stop_reason, kv_transfer_params, + ec_transfer_params, ) is_delta = self.output_kind == RequestOutputKind.DELTA @@ -230,6 +242,7 @@ class OmniRequestState(RequestState): outputs, finished, kv_transfer_params, + ec_transfer_params, ) def _new_completion_output( @@ -309,6 +322,7 @@ class OmniRequestState(RequestState): outputs: list, finished: bool, kv_transfer_params: dict[str, Any] | None = None, + ec_transfer_params: dict[str, Any] | None = None, ) -> RequestOutput | PoolingRequestOutput: """Create request output, handling no-detokenizer generation stages. @@ -323,6 +337,7 @@ class OmniRequestState(RequestState): outputs, finished, kv_transfer_params, + ec_transfer_params, ) # No-detokenizer path: build RequestOutput directly. @@ -341,7 +356,9 @@ class OmniRequestState(RequestState): outputs=outputs, finished=finished, kv_transfer_params=kv_transfer_params, + ec_transfer_params=ec_transfer_params, num_cached_tokens=self.num_cached_tokens, + num_cache_creation_tokens=self.num_cache_creation_tokens, metrics=self.stats, ) @@ -416,6 +433,116 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): return self._native_text_metrics_by_request.pop(request_id, {}) def abort_requests(self, request_ids, internal: bool) -> list[str]: + aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal) + return aborted_ids + + def abort_requests_collecting_outputs( + self, + request_ids: Iterable[str], + *, + internal: bool, + commit_state: bool = True, + ) -> tuple[list[str], list[RequestOutput | PoolingRequestOutput]]: + """Abort requests and return terminal abort outputs with partial tokens. + + Mirrors upstream ``OutputProcessor.abort_requests``, but also returns + abort ``RequestOutput`` objects when ``queue`` is ``None`` (Omni stage + pools register OP state without a collector queue). + + Terminal abort outputs use ``RequestOutputKind.CUMULATIVE`` so DELTA + streaming requests still surface the full prefix generated so far. + ``new_token_ids`` is the detokenizer prefix: Omni's + ``make_request_output`` writes that list onto ``CompletionOutput.token_ids``. + + When ``commit_state`` is False, processor mappings stay in place so a + failed physical EngineCore abort can retry with the same prefix. + """ + request_ids = list(request_ids) + if commit_state: + for request_id in request_ids: + if internal: + req_state = self.request_states.get(request_id) + if req_state is not None: + self._native_text_metrics_by_request.pop(req_state.external_req_id, None) + else: + self._native_text_metrics_by_request.pop(request_id, None) + + internal_req_ids: list[str] = [] + for request_id in request_ids: + if internal: + internal_req_ids.append(request_id) + if commit_state: + if req_state := self.request_states.get(request_id): + external_req_id = req_state.external_req_id + internal_ids = self.external_req_ids.get(external_req_id) + if internal_ids is not None: + try: + internal_ids.remove(request_id) + except ValueError: + pass + if not internal_ids: + self.external_req_ids.pop(external_req_id, None) + elif commit_state: + if internal_ids := self.external_req_ids.pop(request_id, []): + internal_req_ids.extend(internal_ids) + else: + internal_req_ids.extend(self.external_req_ids.get(request_id, [])) + + request_ids_to_abort: list[str] = [] + abort_outputs: list[RequestOutput | PoolingRequestOutput] = [] + for request_id in internal_req_ids: + req_state = ( + self.request_states.pop(request_id, None) if commit_state else self.request_states.get(request_id) + ) + if req_state is not None: + if commit_state: + self.lora_states.request_finished(request_id, req_state.lora_name) + request_ids_to_abort.append(request_id) + original_kind = req_state.output_kind + req_state.output_kind = RequestOutputKind.CUMULATIVE + detok = req_state.detokenizer + new_token_ids = list(detok.output_token_ids) if detok is not None else [] + saved_children = None + if not commit_state and req_state.parent_req is not None: + saved_children = set(req_state.parent_req.child_requests) + try: + request_output = req_state.make_request_output( + new_token_ids=new_token_ids, + pooling_output=None, + finish_reason=FinishReason.ABORT, + stop_reason=None, + kv_transfer_params=None, + ) + finally: + req_state.output_kind = original_kind + if saved_children is not None and req_state.parent_req is not None: + req_state.parent_req.child_requests = saved_children + if request_output is not None: + abort_outputs.append(request_output) + if commit_state and req_state.queue is not None: + req_state.queue.put(request_output) + if commit_state: + parent_req = req_state.parent_req + if parent_req is not None: + parent_req.child_requests.discard(request_id) + if not parent_req.child_requests: + self.parent_requests.pop(parent_req.request_id, None) + elif parent := self.parent_requests.get(request_id): + if parent.child_requests: + child_reqs = list(parent.child_requests) + child_aborted, child_outputs = self.abort_requests_collecting_outputs( + child_reqs, + internal=True, + commit_state=commit_state, + ) + request_ids_to_abort.extend(child_aborted) + abort_outputs.extend(child_outputs) + if commit_state: + self.parent_requests.pop(request_id, None) + return request_ids_to_abort, abort_outputs + + def commit_aborted_request_state(self, request_ids: Iterable[str], *, internal: bool) -> None: + """Drop processor mappings after a successful physical EngineCore abort.""" request_ids = list(request_ids) for request_id in request_ids: if internal: @@ -424,7 +551,37 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): self._native_text_metrics_by_request.pop(req_state.external_req_id, None) else: self._native_text_metrics_by_request.pop(request_id, None) - return super().abort_requests(request_ids, internal) + + internal_req_ids: list[str] = [] + for request_id in request_ids: + if internal: + internal_req_ids.append(request_id) + if req_state := self.request_states.get(request_id): + external_req_id = req_state.external_req_id + internal_ids = self.external_req_ids.get(external_req_id) + if internal_ids is not None: + try: + internal_ids.remove(request_id) + except ValueError: + pass + if not internal_ids: + self.external_req_ids.pop(external_req_id, None) + elif internal_ids := self.external_req_ids.pop(request_id, []): + internal_req_ids.extend(internal_ids) + + for request_id in internal_req_ids: + req_state = self.request_states.pop(request_id, None) + if req_state is not None: + self.lora_states.request_finished(request_id, req_state.lora_name) + parent_req = req_state.parent_req + if parent_req is not None: + parent_req.child_requests.discard(request_id) + if not parent_req.child_requests: + self.parent_requests.pop(parent_req.request_id, None) + elif parent := self.parent_requests.get(request_id): + if parent.child_requests: + self.commit_aborted_request_state(list(parent.child_requests), internal=True) + self.parent_requests.pop(request_id, None) def add_request( self, @@ -560,6 +717,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): finish_reason = eco.finish_reason stop_reason = eco.stop_reason kv_transfer_params = eco.kv_transfer_params + ec_transfer_params = getattr(eco, "ec_transfer_params", None) routed_experts = eco.routed_experts self._update_stats_from_output( req_state, @@ -570,6 +728,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): prefill_stats = getattr(eco, "prefill_stats", None) if prefill_stats is not None: req_state.num_cached_tokens = prefill_stats.num_cached_tokens + req_state.num_cache_creation_tokens = prefill_stats.num_cache_creation_tokens req_state.is_prefilling = False is_non_final_audio_chunk = ( @@ -583,7 +742,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): finish_reason, stop_reason, kv_transfer_params, - routed_experts, + ec_transfer_params, + routed_experts=routed_experts, ): if req_state.queue is not None: req_state.queue.put(request_output) diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 7759cb6d1..7ab4dd084 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import nullcontext from enum import Enum @@ -150,6 +150,12 @@ class OmniPlatform(Platform): ) backend_cls.validate_available() + @classmethod + def supports_diffusion_dense_flash_attention(cls) -> bool: + """Whether the platform's dense ``FLASH_ATTN`` dependencies exist.""" + + return True + @classmethod def supports_torch_inductor(cls) -> bool: """Check if the platform supports torch.compile with inductor backend.""" @@ -184,6 +190,47 @@ class OmniPlatform(Platform): """ return "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner" + @classmethod + def get_diffusion_kv_block_tables_cls(cls) -> type: + """Return the platform's native paged-KV BlockTables implementation.""" + from vllm.v1.worker.gpu.block_table import BlockTables + + return BlockTables + + @classmethod + def build_diffusion_kv_attn_metadata(cls, **kwargs: Any) -> dict[str, Any]: + """Build native attention metadata for the diffusion paged path. + + This default is the GPU/common path and uses vLLM's builder. NPU + overrides it to build Ascend attention metadata without making the + shared diffusion adapter import ``vllm_ascend``. ``seq_lens_cpu`` is an + adapter-only convenience value and is removed before calling the + upstream builder. + """ + from vllm.v1.worker.gpu.attn_utils import build_attn_metadata + + kwargs.pop("seq_lens_cpu", None) + return build_attn_metadata(**kwargs) + + @classmethod + def get_diffusion_paged_kv_attn_backend(cls, attn_backend: type, *, ulysses_degree: int) -> type: + """Specialize a native paged backend for diffusion execution.""" + + del ulysses_degree + return attn_backend + + @classmethod + def requires_diffusion_paged_kv_prewrite(cls) -> bool: + """Whether paged attention must write K/V before native execution. + + The default GPU path keeps cache-update ownership in its native + attention call. Ascend overrides this because piecewise FIA should + write the complete K/V span once, then read it from cache for every + segment. + """ + + return False + @classmethod def init_diffusion_worker_vllm_config( cls, @@ -192,6 +239,15 @@ class OmniPlatform(Platform): """Initialize platform-specific state for diffusion worker VllmConfig.""" return None + @classmethod + def configure_diffusion_vllm_config( + cls, + vllm_config: Any, + od_config: Any, + ) -> None: + """Apply platform-specific native cache geometry for diffusion.""" + return None + @classmethod def init_diffusion_model_runner_runtime( cls, diff --git a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py index 713b67f81..dd03ba7d2 100644 --- a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py +++ b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py @@ -148,6 +148,15 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) + config = args[0] if args else kwargs.get("config") + if config is not None: + self.num_heads = getattr(self, "num_heads", config.num_attention_heads) + self.num_kv_heads = getattr(self, "num_kv_heads", config.num_key_value_heads) + self.head_dim = getattr( + self, + "head_dim", + getattr(config, "head_dim", config.hidden_size // config.num_attention_heads), + ) self._buffers.pop("_fusion_causal_mask", None) self._q_size = self.num_heads * self.head_dim self._kv_size = self.num_kv_heads * self.head_dim @@ -157,11 +166,9 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti def prepare_qkv_weights(self) -> None: # Pack QKV once so each graph replay uses one matmul and consumes the # weight directly in the 310P matmul layout. - self._fused_qkv_weight = maybe_trans_nz( - torch.cat((self.q_proj.weight, self.k_proj.weight, self.v_proj.weight), dim=0).contiguous() - ) - if self.q_proj.bias is not None: - self._fused_qkv_bias = torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias), dim=0) + self._fused_qkv_weight = maybe_trans_nz(self.qkv_proj.weight.contiguous()) + if self.qkv_proj.bias is not None: + self._fused_qkv_bias = self.qkv_proj.bias def forward( self, @@ -170,11 +177,17 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: bsz, seq_len, _ = hidden_states.shape + hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) + hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) + + # ``prepare_qkv_weights`` packs qkv_proj into the 310P matmul layout; + # split it back out the same way before re-applying the 310P + # flash-attention path. qkv = F.linear(hidden_states, self._fused_qkv_weight, self._fused_qkv_bias) - q, k, v = qkv.split((self._q_size, self._kv_size, self._kv_size), dim=-1) - q = self.q_norm(q.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2) - k = self.k_norm(k.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2) - v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + q_raw, k_raw, v_raw = self._split_qkv(qkv) + q = self.q_norm(q_raw.view(hidden_shape_q)).transpose(1, 2) + k = self.k_norm(k_raw.view(hidden_shape_kv)).transpose(1, 2) + v = v_raw.view(hidden_shape_kv).transpose(1, 2) cos, sin = position_embeddings cos = cos.unsqueeze(1) @@ -304,7 +317,7 @@ class _Qwen3TTSTalkerCodePredictor310P( for layer in self.model.layers: attention = layer.self_attn attention.prepare_qkv_weights() - qkv_projections.update((attention.q_proj, attention.k_proj, attention.v_proj)) + qkv_projections.add(attention.qkv_proj) for module in self.modules(): if isinstance(module, nn.Linear) and module not in qkv_projections: diff --git a/vllm_omni/platforms/npu/__init__.py b/vllm_omni/platforms/npu/__init__.py index fb95f7b29..74dfc419c 100644 --- a/vllm_omni/platforms/npu/__init__.py +++ b/vllm_omni/platforms/npu/__init__.py @@ -1,6 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm_omni.platforms.npu.platform import NPUOmniPlatform +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm_omni.platforms.npu.platform import NPUOmniPlatform __all__ = ["NPUOmniPlatform"] + + +def __getattr__(name: str): + if name != "NPUOmniPlatform": + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + from vllm_omni.platforms.npu.platform import NPUOmniPlatform + + globals()[name] = NPUOmniPlatform + return NPUOmniPlatform diff --git a/vllm_omni/platforms/npu/graph_tools.py b/vllm_omni/platforms/npu/graph_tools.py new file mode 100644 index 000000000..d943d1584 --- /dev/null +++ b/vllm_omni/platforms/npu/graph_tools.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reusable exact-signature NPUGraph capture and replay helpers.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Protocol + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class _ReplayGraph(Protocol): + def replay(self) -> None: ... + + +@dataclass +class CapturedDeviceGraph: + graph: _ReplayGraph + static_inputs: tuple[torch.Tensor, ...] + static_outputs: tuple[torch.Tensor, ...] + + def replay(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + with torch.inference_mode(): + for static, current in zip(self.static_inputs, inputs, strict=True): + static.copy_(current) + self.graph.replay() + # Graph outputs are persistent and overwritten by the next replay. + # Clone before they become request-owned streaming cache entries. + return tuple(output.detach().clone() for output in self.static_outputs) + + +def _tensor_signature(value: torch.Tensor) -> tuple[tuple[int, ...], str, str]: + return tuple(value.shape), str(value.dtype), str(value.device) + + +class NPUExactGraphRunner: + """Capture and replay tensor-only functions for exact NPU signatures.""" + + def __init__( + self, + *, + max_graphs: int = 32, + component_name: str = "device graph", + disable_config_hint: str = "disable graph capture", + ) -> None: + self.max_graphs = max(0, int(max_graphs)) + self.component_name = component_name + self.disable_config_hint = disable_config_hint + self._enabled = self.max_graphs > 0 + self._graphs: dict[tuple[object, ...], CapturedDeviceGraph] = {} + self._failed_keys: set[tuple[object, ...]] = set() + self._hits = 0 + self._graph_pool: object | None = None + + @staticmethod + def is_supported() -> bool: + npu = getattr(torch, "npu", None) + return npu is not None and all( + hasattr(npu, name) + for name in ( + "NPUGraph", + "graph", + "is_current_stream_capturing", + "synchronize", + ) + ) + + @staticmethod + def _stream_is_capturing() -> bool: + npu = getattr(torch, "npu", None) + is_capturing = getattr(npu, "is_current_stream_capturing", None) + if not callable(is_capturing): + return False + try: + return bool(is_capturing()) + except (RuntimeError, TypeError): + return False + + def _eligible(self, inputs: tuple[torch.Tensor, ...]) -> bool: + return ( + self._enabled + and bool(inputs) + and all(value.device.type == "npu" for value in inputs) + and self.is_supported() + and not self._stream_is_capturing() + ) + + @property + def stats(self) -> dict[str, int]: + return { + "captures": len(self._graphs), + "failed": len(self._failed_keys), + "hits": self._hits, + } + + def capture( + self, + inputs: tuple[torch.Tensor, ...], + compute: Callable[..., tuple[torch.Tensor, ...]], + ) -> CapturedDeviceGraph: + npu = torch.npu + static_inputs = tuple(value.detach().clone() for value in inputs) + npu.synchronize() + graph = npu.NPUGraph() + if self._graph_pool is None: + from vllm.platforms import current_platform + + self._graph_pool = current_platform.get_global_graph_pool() + with torch.inference_mode(), npu.graph(graph, pool=self._graph_pool): + static_outputs = compute(*static_inputs) + npu.synchronize() + return CapturedDeviceGraph( + graph=graph, + static_inputs=static_inputs, + static_outputs=static_outputs, + ) + + def run( + self, + operation: str, + inputs: tuple[torch.Tensor, ...], + constants: tuple[object, ...], + compute: Callable[..., tuple[torch.Tensor, ...]], + ) -> tuple[torch.Tensor, ...]: + if self._failed_keys: + raise RuntimeError( + f"{self.component_name} cannot continue after a failed NPUGraph capture; " + f"restart the stage process and {self.disable_config_hint} before retrying." + ) + if not self._eligible(inputs): + return compute(*inputs) + + key = ( + operation, + constants, + tuple(_tensor_signature(value) for value in inputs), + ) + graph = self._graphs.get(key) + if graph is not None: + self._hits += 1 + if self._hits == 1: + logger.info("%s started NPUGraph replay", self.component_name) + return graph.replay(inputs) + + # Prime lazy kernels and allocator state before capture. The next call + # with the same exact tensor signature replays this graph. + eager_outputs = compute(*inputs) + if len(self._graphs) >= self.max_graphs: + logger.warning_once( + "%s reached the %d-entry NPUGraph limit; new tensor shapes will use eager execution.", + self.component_name, + self.max_graphs, + ) + return eager_outputs + try: + self._graphs[key] = self.capture(inputs, compute) + except Exception as exc: + self._failed_keys.add(key) + self._enabled = False + logger.exception( + "%s failed to capture NPUGraph for %s; the torch-npu " + "allocator/RNG capture state may be invalid and this stage " + "process must be restarted.", + self.component_name, + operation, + ) + raise RuntimeError( + f"{self.component_name} NPUGraph capture failed for {operation}; " + f"restart the stage process. To run eagerly, {self.disable_config_hint}." + ) from exc + logger.info( + "%s captured NPUGraph %d/%d for %s", + self.component_name, + len(self._graphs), + self.max_graphs, + operation, + ) + return eager_outputs diff --git a/vllm_omni/platforms/npu/layers/fused_moe.py b/vllm_omni/platforms/npu/layers/fused_moe.py index e17610701..4af0e9e33 100644 --- a/vllm_omni/platforms/npu/layers/fused_moe.py +++ b/vllm_omni/platforms/npu/layers/fused_moe.py @@ -29,6 +29,40 @@ def _ensure_forward_context_attr(name: str, annotation: Any, default: Any) -> No setattr(_vllm_fc.ForwardContext, name, default) +def _ensure_ascend_moe_customop_registered() -> None: + """Register vllm-ascend's out-of-tree MoE layer implementations here. + + vllm-ascend registers its OOT layers from ``register_ascend_customop``, + which only ever runs inside their own ``NPUWorker.__init__``. The diffusion + path runs on omni's DiffusionWorker, so the OOT registry stays empty and + ``PluggableLayer.__new__`` falls back to the upstream classes. FusedMoE then + breaks outright: vllm-ascend's ``patch_fused_moe`` passes + ``n_shared_experts`` through ``routed_experts_args``, which only + ``AscendRoutedExperts`` accepts. + + Only the MoE pair is registered, not the full ``REGISTERED_ASCEND_OPS`` + table. Most of that table assumes a text-generation ``ModelConfig``: + ``AscendRotaryEmbedding.__init__`` reads ``model_config.use_mla`` and + ``hf_text_config.qk_rope_head_dim``, neither of which exists on the + diffusion config shim. ``AscendRoutedExperts``/``AscendMoERunner`` only need + ``parallel_config``, which the shim does provide. + """ + from vllm.model_executor.custom_op import CustomOp, op_registry_oot + from vllm_ascend.ops.fused_moe.fused_moe import AscendMoERunner + from vllm_ascend.ops.fused_moe.routed_experts import AscendRoutedExperts + + ascend_moe_ops: dict[str, type] = { + "MoERunner": AscendMoERunner, + "RoutedExperts": AscendRoutedExperts, + } + for name, op_cls in ascend_moe_ops.items(): + # register_oot asserts on duplicates, so skip names vllm-ascend already + # claimed. Its own registration is authoritative when it has run. + if name in op_registry_oot: + continue + CustomOp.register_oot(_decorated_op_cls=op_cls, name=name) + + def _init_mc2_group_for_diffusion( backend: str, local_rank: int, @@ -66,6 +100,8 @@ def _select_moe_comm_method(vllm_config: VllmConfig) -> MoECommType | None: def prepare_fused_moe_runtime() -> None: vllm_config = omni_get_ctx().vllm_config + _ensure_ascend_moe_customop_registered() + backend = torch.distributed.get_backend(get_world_group().device_group) local_rank = get_world_group().local_rank _init_mc2_group_for_diffusion( diff --git a/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py b/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py index 7e1d5d540..7cf280a00 100644 --- a/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py +++ b/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py @@ -18,7 +18,7 @@ This module: from __future__ import annotations from collections.abc import Iterator -from contextlib import contextmanager, nullcontext +from contextlib import ExitStack, contextmanager import torch import torch.nn.functional as F @@ -107,17 +107,20 @@ def _patched_attention_forward_chunk( @contextmanager -def npu_math_sdpa_context() -> Iterator[None]: +def npu_math_sdpa_context(*, require_available: bool = False) -> Iterator[None]: """Force SDPA MATH backend so Ascend does not call fused FA.""" - try: - from torch.nn.attention import SDPBackend, sdpa_kernel + with ExitStack() as stack: + try: + from torch.nn.attention import SDPBackend, sdpa_kernel - with sdpa_kernel(SDPBackend.MATH): - yield - except Exception: - # Older torch / missing backend enum — just run as-is. - with nullcontext(): - yield + stack.enter_context(sdpa_kernel(SDPBackend.MATH)) + except Exception as exc: + if require_available: + raise RuntimeError( + "MiniCPM-o Code2Wav NPU graph capture requires torch.nn.attention.sdpa_kernel(SDPBackend.MATH)" + ) from exc + logger.warning_once("MATH SDPA selection is unavailable; using the default Ascend SDPA backend.") + yield def _disable_upsample_encoder_compile() -> None: diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py new file mode 100644 index 000000000..0d34c93a5 --- /dev/null +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Inject MiniCPM-o Code2Wav NPUGraph acceleration on Ascend.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from contextlib import nullcontext +from typing import cast +from weakref import WeakKeyDictionary + +import torch +from vllm.logger import init_logger + +from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner + +logger = init_logger(__name__) + +_PATCHED = False +_original_build_backend = None +_original_estimator_step = None +_original_setup_batch = None +_original_decode_batch = None +_backend_graph_runners: WeakKeyDictionary[object, NPUExactGraphRunner] = WeakKeyDictionary() +_ENABLE_KEY = "code2wav_enable_npu_graph" +_MAX_GRAPHS_KEY = "code2wav_max_npu_graphs" +_BF16_ATTENTION_CACHE_KEY = "code2wav_bfloat16_attention_cache" + + +def _config_bool(value: object, default: bool) -> bool: + if value is None: + return default + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + + +def _graph_config(model: object) -> dict[str, object]: + config = getattr(getattr(model, "vllm_config", None), "additional_config", None) + return dict(config) if isinstance(config, Mapping) else {} + + +def prepare_code2wav_graph_runtime() -> None: + """Select graph-capturable ACLNN kernels before Token2Wav is loaded.""" + if os.environ.get("ASCEND_LAUNCH_BLOCKING") == "1": + raise RuntimeError( + "MiniCPM-o Code2Wav NPUGraph capture is incompatible with " + "ASCEND_LAUNCH_BLOCKING=1; unset it or set it to 0 before startup." + ) + npu = torch.npu + npu.config.allow_internal_format = False + npu.set_compile_mode(jit_compile=False) + logger.info("Configured MiniCPM-o Code2Wav NPUGraph runtime (allow_internal_format=False, jit_compile=False)") + + +def _flow_execution_context(device: torch.device, *, require_math: bool): + if device.type != "npu": + return nullcontext() + from vllm_omni.platforms.npu.models.step_audio2_token2wav import ( + npu_token2wav_sdpa_context, + ) + + return npu_token2wav_sdpa_context(require_math=require_math) + + +def _graphable_estimator_step( + backend, + estimator, + *, + x, + mu, + time_embedding, + speakers, + cond, + cnn_cache, + att_cache, +): + """Run the CFM estimator body after host-backed timestep embedding.""" + width = int(x.shape[-1]) + speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) + estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) + cnn_out, att_out = backend._estimator_buffers(estimator, estimator_input, att_cache) + old_cnn = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) + old_att = att_cache if att_cache is not None else [None] * len(estimator.blocks) + result = estimator.blocks_forward_chunk( + estimator_input, + time_embedding, + None, + old_cnn, + old_att, + cnn_out, + att_out, + ) + return result, cnn_out, att_out + + +def _patched_estimator_step( + self, + estimator, + *, + x, + mu, + time, + speakers, + cond, + cnn_cache, + att_cache, + attn_mask=None, + valid_lengths=None, +): + assert _original_estimator_step is not None + graph_runner = _backend_graph_runners.get(self) + if ( + graph_runner is None + or self._trt_stepper is not None + or self._cfm_graph_wrapper is not None + or attn_mask is not None + or valid_lengths is not None + ): + return _original_estimator_step( + self, + estimator, + x=x, + mu=mu, + time=time, + speakers=speakers, + cond=cond, + cnn_cache=cnn_cache, + att_cache=att_cache, + attn_mask=attn_mask, + valid_lengths=valid_lengths, + ) + if (cnn_cache is None) != (att_cache is None): + raise ValueError("estimator CNN and attention caches must both be present or absent") + + # The upstream embedder creates a frequency tensor on the host. Keep it + # outside capture while retaining the tensor-only estimator body in graph. + time_embedding = estimator.t_embedder(time).unsqueeze(1) + if cnn_cache is None: + return graph_runner.run( + "cfm_estimator", + (x, mu, time_embedding, speakers, cond), + (False,), + lambda step_x, step_mu, step_time, step_speakers, step_cond: _graphable_estimator_step( + self, + estimator, + x=step_x, + mu=step_mu, + time_embedding=step_time, + speakers=step_speakers, + cond=step_cond, + cnn_cache=None, + att_cache=None, + ), + ) + + return graph_runner.run( + "cfm_estimator", + (x, mu, time_embedding, speakers, cond, cnn_cache, att_cache), + (True,), + lambda step_x, step_mu, step_time, step_speakers, step_cond, step_cnn, step_att: _graphable_estimator_step( + self, + estimator, + x=step_x, + mu=step_mu, + time_embedding=step_time, + speakers=step_speakers, + cond=step_cond, + cnn_cache=step_cnn, + att_cache=step_att, + ), + ) + + +def _patched_setup_batch(self, features, batch_size): + assert _original_setup_batch is not None + with _flow_execution_context( + features.speech_tokens.device, + require_math=self in _backend_graph_runners, + ): + return _original_setup_batch(self, features, batch_size) + + +def _patched_decode_batch( + self, + tokens, + features, + states, + *, + last_chunk, + flush_encoder=False, +): + assert _original_decode_batch is not None + with _flow_execution_context( + tokens.device, + require_math=self in _backend_graph_runners, + ): + return _original_decode_batch( + self, + tokens, + features, + states, + last_chunk=last_chunk, + flush_encoder=flush_encoder, + ) + + +def _patched_build_backend(self) -> None: + if self.backend is not None: + return + + extra = self._extra_config() + if bool(extra.get(_BF16_ATTENTION_CACHE_KEY, False)): + raise ValueError( + "MiniCPM-o Code2Wav code2wav_bfloat16_attention_cache is " + "currently supported only on CUDA; leave it unset or false on NPU." + ) + + config = _graph_config(self) + max_graphs = max(0, int(cast(int | str, config.get(_MAX_GRAPHS_KEY, 32)))) + graph_enabled = max_graphs > 0 and _config_bool(config.get(_ENABLE_KEY), False) + if graph_enabled: + # NPUOmniPlatform enables internal format for quantized LLM kernels. + # Code2Wav uses regular convolution kernels that must remain in the + # graph-capturable ACLNN path. + prepare_code2wav_graph_runtime() + + assert _original_build_backend is not None + _original_build_backend(self) + + graph_runner = None + if graph_enabled: + graph_runner = NPUExactGraphRunner( + max_graphs=max_graphs, + component_name="MiniCPM-o Code2Wav", + disable_config_hint=( + "set platforms.npu.stages[stage_id=2].additional_config.code2wav_enable_npu_graph=false" + ), + ) + if self.backend.speech_window.device.type == "npu" and not graph_runner.is_supported(): + raise RuntimeError( + "MiniCPM-o Code2Wav NPUGraph capture requires torch.npu " + "NPUGraph, graph, is_current_stream_capturing, and synchronize APIs." + ) + if self.backend.flow.training: + raise ValueError("MiniCPM-o Code2Wav NPUGraph capture requires flow.eval()") + _backend_graph_runners[self.backend] = graph_runner + + if graph_enabled: + logger.info( + "MiniCPM-o Code2Wav NPUGraph replay enabled (max_graphs=%d)", + max_graphs, + ) + + +def apply_minicpmo_4_5_code2wav_patch() -> None: + """Patch the generic Code2Wav backend builder with Ascend acceleration.""" + global _PATCHED, _original_build_backend + global _original_decode_batch, _original_estimator_step, _original_setup_batch + if _PATCHED: + return + + from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( + BatchedToken2Wav, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( + MiniCPMO45Code2Wav, + ) + + _original_build_backend = MiniCPMO45Code2Wav._build_backend + _original_estimator_step = BatchedToken2Wav._estimator_step + _original_setup_batch = BatchedToken2Wav.setup_batch + _original_decode_batch = BatchedToken2Wav.decode_batch + + MiniCPMO45Code2Wav._build_backend = _patched_build_backend # type: ignore[method-assign] + BatchedToken2Wav._estimator_step = _patched_estimator_step # type: ignore[method-assign] + BatchedToken2Wav.setup_batch = _patched_setup_batch # type: ignore[method-assign] + BatchedToken2Wav.decode_batch = _patched_decode_batch # type: ignore[method-assign] + _PATCHED = True + logger.debug("Applied NPU patch for MiniCPM-o 4.5 Code2Wav") diff --git a/vllm_omni/platforms/npu/models/minimax_h3.py b/vllm_omni/platforms/npu/models/minimax_h3.py new file mode 100644 index 000000000..c8889f18f --- /dev/null +++ b/vllm_omni/platforms/npu/models/minimax_h3.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""NPU patches for the MiniMax H3 Qwen3-VL text encoder.""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn.functional as F +import torch_npu +from vllm.logger import init_logger + +from vllm_omni.platforms.npu.layers.rotary_embedding import ( + npu_rotary_mul_with_bsnd_fallback, +) + +logger = init_logger(__name__) + +_ROPE_PATCHED = False +_SDPA_PATCHED = False +_SWIGLU_PATCHED = False + + +def _apply_rotary_pos_emb_npu( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply Qwen3-VL text RoPE with fused BNSD rotary multiplication.""" + return ( + npu_rotary_mul_with_bsnd_fallback(q, cos, sin, unsqueeze_dim=1), + npu_rotary_mul_with_bsnd_fallback(k, cos, sin, unsqueeze_dim=1), + ) + + +def apply_minimax_h3_qwen3vl_patch() -> None: + """Route MiniMax H3 Qwen3-VL text RoPE to the Ascend fused operator.""" + global _ROPE_PATCHED + if _ROPE_PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder._apply_rotary_pos_emb = _apply_rotary_pos_emb_npu + _ROPE_PATCHED = True + logger.debug("Applied NPU fused RoPE patch for MiniMax H3 Qwen3-VL text encoder") + + +def _scaled_dot_product_attention_npu( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, +) -> torch.Tensor: + """Run causal SDPA with compressed K/V heads through NPU native GQA.""" + num_heads = query.shape[1] + num_key_value_heads = key.shape[1] + num_value_heads = value.shape[1] + if num_key_value_heads != num_value_heads: + raise ValueError( + "GQA requires key and value to have the same number of heads, " + f"got k_heads={num_key_value_heads} and v_heads={num_value_heads}." + ) + if num_key_value_heads == 0: + raise ValueError("GQA requires at least one KV head.") + if num_heads % num_key_value_heads != 0: + raise ValueError( + "GQA requires query heads to be a multiple of KV heads, " + f"got q_heads={num_heads} and kv_heads={num_key_value_heads}." + ) + return F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.0, + is_causal=True, + enable_gqa=num_heads != num_key_value_heads, + ) + + +def apply_minimax_h3_qwen3vl_sdpa_patch() -> None: + """Route MiniMax H3 Qwen3-VL text attention to NPU native GQA.""" + global _SDPA_PATCHED + if _SDPA_PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder._scaled_dot_product_attention = _scaled_dot_product_attention_npu + _SDPA_PATCHED = True + logger.debug("Applied NPU SDPA patch for MiniMax H3 Qwen3-VL text encoder") + + +def npu_swiglu_from_packed(gate_up: torch.Tensor) -> torch.Tensor: + """Apply fused SwiGLU to a packed gate/up projection tensor.""" + return torch_npu.npu_swiglu(gate_up, dim=-1) + + +def _forward_minimax_h3_qwen3vl_text_mlp_npu(self: Any, x: torch.Tensor) -> torch.Tensor: + """Run the Qwen3-VL MLP with one packed GEMM and fused SwiGLU.""" + gate_up = F.linear(x, self.gate_up_proj.weight) + return self.down_proj(npu_swiglu_from_packed(gate_up)) + + +def apply_minimax_h3_qwen3vl_swiglu_patch() -> None: + """Route MiniMax H3 Qwen3-VL text MLP to the Ascend fused SwiGLU path.""" + global _SWIGLU_PATCHED + if _SWIGLU_PATCHED: + return + + from vllm_omni.diffusion.models.minimax_h3 import encoder + + encoder.MiniMaxH3Qwen3VLTextMLP.forward = _forward_minimax_h3_qwen3vl_text_mlp_npu # type: ignore[method-assign] + _SWIGLU_PATCHED = True diff --git a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py index 1641fee27..1999a08a0 100644 --- a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py +++ b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py @@ -13,7 +13,7 @@ Ascend-specific workarounds that must not live in the shared GPU model file: from __future__ import annotations from collections.abc import Iterator -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from types import MethodType import numpy as np @@ -113,20 +113,16 @@ def patch_step_audio2_hift_for_npu(hift: torch.nn.Module) -> None: @contextmanager -def npu_token2wav_sdpa_context() -> Iterator[None]: +def npu_token2wav_sdpa_context(*, require_math: bool = False) -> Iterator[None]: """Expand CosyVoice masks + force MATH SDPA to avoid FA 161001.""" - try: - from vllm_omni.platforms.npu.models.cosyvoice2_dit_attn import ( - apply_cosyvoice2_dit_attn_npu_patch, - npu_math_sdpa_context, - ) + from vllm_omni.platforms.npu.models.cosyvoice2_dit_attn import ( + apply_cosyvoice2_dit_attn_npu_patch, + npu_math_sdpa_context, + ) - apply_cosyvoice2_dit_attn_npu_patch() - with npu_math_sdpa_context(): - yield - except Exception: - with nullcontext(): - yield + apply_cosyvoice2_dit_attn_npu_patch() + with npu_math_sdpa_context(require_available=require_math): + yield def _patched_ensure_models_loaded(self) -> None: diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 525725fcf..cf65da414 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import nullcontext +from functools import cache from typing import Any import torch @@ -21,6 +22,28 @@ _DIFFUSION_PACKED_MODULES_MAPPING = { } +@cache +def _get_strict_ulysses_paged_backend() -> type: + """Return an Ascend backend that bypasses vLLM PCP dispatch.""" + + from vllm_ascend.attention.attention_v1 import ( + AscendAttentionBackend, + AscendAttentionBackendImpl, + AscendAttentionMetadataBuilder, + ) + + class AscendStrictUlyssesPagedBackend(AscendAttentionBackend): + @staticmethod + def get_impl_cls() -> type: + return AscendAttentionBackendImpl + + @staticmethod + def get_builder_cls() -> type: + return AscendAttentionMetadataBuilder + + return AscendStrictUlyssesPagedBackend + + class NPUOmniPlatform(OmniPlatform, NPUPlatform): """NPU/Ascend implementation of OmniPlatform. @@ -36,6 +59,9 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): from vllm_ascend.utils import adapt_patch from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches + from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import ( + apply_minicpmo_4_5_code2wav_patch, + ) from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import ( apply_qwen3_tts_code2wav_patch, ) @@ -44,6 +70,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): ) adapt_patch(is_global_patch=True) + apply_minicpmo_4_5_code2wav_patch() apply_qwen3_tts_code2wav_patch() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() @@ -73,13 +100,95 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): @classmethod def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None: from vllm_ascend.ascend_config import init_ascend_config + from vllm_ascend.utils import adapt_patch + # Omni's custom DiffusionWorker does not pass through vLLM-Ascend's + # NPUWorker constructor, where worker-local patches are normally + # installed. In particular, AscendBlockTables needs the patched + # non-UVA buffer implementation on NPU. + adapt_patch() init_ascend_config(vllm_config) + @classmethod + def configure_diffusion_vllm_config(cls, vllm_config: Any, od_config: Any) -> None: + """Use the block geometry required by Ascend's native paged kernel.""" + if getattr(od_config, "diffusion_kv_mode", None) is None: + return + from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode + + if od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER: + return + from vllm_ascend.attention.attention_v1 import AscendAttentionBackend + + supported_sizes = [ + size for size in AscendAttentionBackend.get_supported_kernel_block_sizes() if type(size) is int and size > 0 + ] + if not supported_sizes: + raise RuntimeError("Ascend paged attention did not expose an integer kernel block size") + # vLLM's generic default is 16, while the Ascend FIA backend stores + # cache pages as 128-token blocks. Set the Manager geometry + # before KV specs are collected so Scheduler and Worker agree. + vllm_config.cache_config.block_size = supported_sizes[0] + + @classmethod + def requires_diffusion_paged_kv_prewrite(cls) -> bool: + """Write the full K/V span once before piecewise FIA segments.""" + + return True + + @classmethod + def get_diffusion_paged_kv_attn_backend(cls, attn_backend: type, *, ulysses_degree: int) -> type: + """Keep strict Ulysses paged FIA out of vLLM's PCP implementation.""" + + del cls + if ulysses_degree <= 1: + return attn_backend + from vllm_ascend.attention.attention_v1 import AscendAttentionBackend + + if not isinstance(attn_backend, type) or not issubclass(attn_backend, AscendAttentionBackend): + return attn_backend + return _get_strict_ulysses_paged_backend() + + @classmethod + def get_diffusion_kv_block_tables_cls(cls) -> type: + from vllm_ascend.worker.v2.block_table import AscendBlockTables + + return AscendBlockTables + + @classmethod + def build_diffusion_kv_attn_metadata(cls, **kwargs: Any) -> dict[str, Any]: + """Build the Ascend metadata required by the native NPU backend.""" + from vllm_ascend.attention.attention_v1 import AscendAttentionState + from vllm_ascend.worker.v2.attn_utils import build_attn_metadata + + kwargs = dict(kwargs) + seq_lens_cpu = kwargs.pop("seq_lens_cpu") + kwargs["seq_lens_np"] = seq_lens_cpu.detach().cpu().numpy() + # The diffusion adapter always supplies a paged cache and the current + # K/V write span. ChunkedPrefill is Ascend's cache-backed FIA state for + # both multi-token updates and single-token updates in this path. + kwargs["attn_state"] = AscendAttentionState.ChunkedPrefill + return build_attn_metadata(**kwargs) + @classmethod def init_diffusion_model_runner_runtime(cls, vllm_config: Any, od_config: Any, device: torch.device) -> None: from vllm_ascend.ascend_forward_context import set_mc2_mask, set_mc2_tokens_capacity + from vllm_omni.platforms.npu.models.minimax_h3 import ( + apply_minimax_h3_qwen3vl_patch, + apply_minimax_h3_qwen3vl_sdpa_patch, + apply_minimax_h3_qwen3vl_swiglu_patch, + ) + + # These patches import the MiniMax encoder package, whose __init__ loads + # pipeline_minimax_h3 → diffusion.data. Doing that during platform + # construction races vllm_omni/__init__.py (patch before config) and + # closes a cycle through pipeline_registry → PI0_PIPELINE → + # DiffusionOutput. Apply them only after the platform exists, before + # the diffusion pipeline is loaded. + apply_minimax_h3_qwen3vl_patch() + apply_minimax_h3_qwen3vl_sdpa_patch() + apply_minimax_h3_qwen3vl_swiglu_patch() set_mc2_tokens_capacity(vllm_config, od_config.max_num_seqs, 1) set_mc2_mask(vllm_config, device) @@ -139,10 +248,15 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): ) backend_upper = "FLASH_ATTN" - if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"): - # The NPU FLASH_ATTN backend imports mindiesd lazily at first - # forward, but CANN snapshots the custom-op registry at the - # first custom-op regInfo lookup in the process (e.g. a + if backend_upper in ("FLASH_ATTN", "RAINFUSION_ATTN") and find_spec("mindiesd"): + # Eager-import mindiesd only for backends that actually reach + # mindiesd kernels: FLASH_ATTN directly, and RAINFUSION_ATTN + # via its dense FlashAttention fallback (used before + # start_step and on any layer without a sparsifiable video + # segment). Other backends (e.g. TORCH_SDPA) never touch + # mindiesd, so a broken optional install must not block them. + # CANN snapshots the custom-op registry at the first + # custom-op regInfo lookup in the process (e.g. a # vllm-ascend custom op during model load/warmup). Import # mindiesd here so its env.py prepends the mindiesd vendor # dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before @@ -166,6 +280,14 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): logger.debug("Falling back to diffusion attention backend SDPA") return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + @classmethod + def supports_diffusion_dense_flash_attention(cls) -> bool: + """Return whether MindIE-SD is installed for dense NPU FlashAttention.""" + + from importlib.util import find_spec + + return find_spec("mindiesd") is not None + @classmethod def supports_torch_inductor(cls) -> bool: return False @@ -200,7 +322,12 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): """ try: torch.npu.current_stream().synchronize() - event = torch.npu.Event() + # The async output worker uses the public ``torch.Stream`` API. + # With torch_npu 2.10, a native ``torch.npu.Event`` cannot be + # consumed by that wrapper stream (the reverse direction is + # supported), while ``torch.Event`` is dispatched to the active + # NPU backend and remains cross-stream compatible. + event = torch.Event() event.record() return event except Exception: diff --git a/vllm_omni/platforms/npu/profiler.py b/vllm_omni/platforms/npu/profiler.py index 65c9c92d9..27f5a647a 100644 --- a/vllm_omni/platforms/npu/profiler.py +++ b/vllm_omni/platforms/npu/profiler.py @@ -53,7 +53,7 @@ class NPUTorchProfilerWrapper(OmniTorchProfilerWrapper): export_type=torch_npu.profiler.ExportType.Text, profiler_level=torch_npu.profiler.ProfilerLevel.Level1, msprof_tx=False, - aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone, + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, l2_cache=False, op_attr=False, data_simplification=True, diff --git a/vllm_omni/platforms/npu/worker/base.py b/vllm_omni/platforms/npu/worker/base.py index ff1a860b0..5dd929ea1 100644 --- a/vllm_omni/platforms/npu/worker/base.py +++ b/vllm_omni/platforms/npu/worker/base.py @@ -45,7 +45,7 @@ class OmniNPUWorkerBase(NPUWorker): if self.profiler is None: raise RuntimeError( "Profiling is not enabled. For diffusion models, set --profiler-config via CLI. " - "For omni models, add profiler_config to your stage config file." + "For omni models, add profiler_config to the relevant deploy config stage." ) if is_start: from vllm_omni.profiler import OmniTorchProfilerWrapper diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 63f3912f7..36a3499cc 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import logging import time from collections.abc import Mapping -from copy import copy, deepcopy +from copy import copy +from dataclasses import replace from typing import Any, NamedTuple import numpy as np @@ -44,14 +45,15 @@ from vllm_ascend.worker.model_runner_v1 import graph_capture from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager -from vllm_omni.distributed.omni_connectors.utils.config import ( - get_stage_connector_role, - stage_sends_async_output, -) +from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output +from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element -from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin +from vllm_omni.worker.omni_connector_model_runner_mixin import ( + OmniConnectorModelRunnerMixin, + needs_omni_connector, +) from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids @@ -106,7 +108,7 @@ class ExecuteModelState(NamedTuple): batch_desc: BatchDescriptor multimodal_outputs: Any # Omni-Specific -class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): +class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin): """Autoregressive NPU model runner that returns hidden states per request.""" def __init__(self, *args, **kwargs): @@ -118,30 +120,24 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) self._async_chunk = getattr(self.model_config, "async_chunk", False) - # -------------------------------------- Omni-new ------------------------------------------------- - _OMNI_CONNECTOR_INIT_ARCHS = { - "Qwen3OmniMoeForConditionalGeneration", - "Qwen2_5OmniForConditionalGeneration", - "CovoAudioForConditionalGeneration", - "MiMoAudioModel", - "Qwen3TTSTalkerForConditionalGeneration", - "Qwen3TTSCode2Wav", - "CosyVoice3Model", - "DyninOmniForConditionalGeneration", - "IndexTTS2TalkerForConditionalGeneration", - } - stage_archs = set(getattr(self.model_config, "architectures", None) or ()) - model_arch_override = getattr(self.model_config, "model_arch", None) - if model_arch_override: - stage_archs.add(model_arch_override) - if stage_archs & _OMNI_CONNECTOR_INIT_ARCHS or get_stage_connector_role(self.model_config) is not None: + if needs_omni_connector(self.model_config): self.init_omni_connectors( model_config=self.model_config, kv_transfer_manager=self.kv_transfer_manager, ) self._downstream_payload_cache: dict[str, bool] = {} + self._init_duplex_sampling_state() # -------------------------------------- Omni-new ------------------------------------------------- + def load_model(self, *args, **kwargs) -> None: + super().load_model(*args, **kwargs) + self._resolve_duplex_sampling_hook(force=True) + + def _update_states(self, scheduler_output: SchedulerOutput): + deferred_state_corrections_fn = super()._update_states(scheduler_output) + self._update_duplex_sampling_states(scheduler_output) + return deferred_state_corrections_fn + def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size from vllm_omni.distributed.ray_utils.utils import ( @@ -399,9 +395,10 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config + if profiling_chunk_config.enabled and profiling_chunk_config.need_timing: if getattr(scheduler_output, "disable_profiling_timing", False): - self.ascend_config.scheduler_config.profiling_chunk_config.need_timing = False + profiling_chunk_config.need_timing = False else: self._sync_device() self._execution_start_time = time.perf_counter() @@ -457,21 +454,17 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): set(scheduler_output.finished_req_ids), self.input_batch, ) - # self._draft_token_ids is None when `input_fits_in_drafter=False` - # and there is no draft tokens scheduled. so it need to update the - # spec_decoding info in scheduler_output with async_scheduling. - # use deepcopy to avoid the modification has influence on the - # scheduler_output in engine core process. - # TODO(Ronald1995): deepcopy is expensive when there is a large - # number of requests, optimize it later. - if ( - self.use_async_scheduling - and self.num_spec_tokens - and self._draft_token_ids is None # type: ignore[has-type] - ): - scheduler_output = deepcopy(scheduler_output) # -------------------------------------- Omni-new ------------------------------------------------- + if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy() + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) + if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata if kv_connector_metadata is not None: @@ -490,7 +483,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): # -------------------------------------- Omni-new ------------------------------------------------- if has_ec_transfer() and get_ec_transfer().is_producer: - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -541,7 +534,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): "it when the requests need prompt logprobs" ) - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) num_reqs = self.input_batch.num_reqs req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] @@ -717,13 +710,6 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): with record_function_or_nullcontext("EPLB weight D2D"): self.eplb_updator.forward_before() - # Set cudagraph mode to none if calc_kv_scales is true. - # KV scales calculation involves dynamic operations that are incompatible - # with CUDA graph capture. - if self.calculate_kv_scales: # type: ignore[has-type] - cudagraph_mode = CUDAGraphMode.NONE - # Mark KV scales as calculated after the first forward pass - self.calculate_kv_scales = False # type: ignore[has-type] # Encoder-decoder models can only compile the pure decode steps where no # encoder inputs are present. Use eager for the first pass. num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) @@ -767,6 +753,15 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): if callable(flush_pending_metadata): flush_pending_metadata(req_ids[:num_reqs]) + # [Omni] Hand the model the batch's req_ids in logits order, for + # models that gate logits per request. Mirrors gpu_ar_model_runner. + # Only valid without spec decode: there logits_indices carries several + # rows per request, so row i no longer corresponds to req_ids[i]. + if spec_decode_metadata is None: + set_batch_req_ids = getattr(self.model, "set_batch_req_ids", None) + if callable(set_batch_req_ids): + set_batch_req_ids(req_ids[:num_reqs]) + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) if multimodal_outputs is not None: @@ -899,10 +894,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.input_batch.idx_mapping_np, self.input_batch.positions[self.input_batch.logits_indices], ) - sampler_output = model_sample( - logits, - self._sampling_metadata_for_model_sampler(sampling_metadata), - ) + prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata) + self._apply_duplex_sampling(logits, prepared_sampling_metadata) + sampler_output = model_sample(logits, prepared_sampling_metadata) if sampler_output is not None: return sampler_output return self.sampler( @@ -916,6 +910,7 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): def sample_tokens( self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config kv_connector_output = self.kv_connector_output self.kv_connector_output = None pp = get_pp_group() @@ -1291,8 +1286,10 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): model_runner_output.omni_connector_output = self.get_omni_connector_output() # -------------------------------------- Omni-new ------------------------------------------------- - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing and hasattr( - self, "_execution_start_time" + if ( + profiling_chunk_config.enabled + and profiling_chunk_config.need_timing + and hasattr(self, "_execution_start_time") ): self._sync_device() model_runner_output.execution_time_ms = (time.perf_counter() - self._execution_start_time) * 1000.0 diff --git a/vllm_omni/platforms/npu/worker/npu_ar_worker.py b/vllm_omni/platforms/npu/worker/npu_ar_worker.py index d34d83ef9..19f3f89c8 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_worker.py @@ -11,9 +11,11 @@ from vllm_omni.worker.mixins import OmniWorkerMixin class NPUARWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU AR worker for thinker/talker stages in Omni model.""" + model_runner_cls = NPUARModelRunner + def init_device(self): self.device = self._init_device() num_ubatches = 1 init_workspace_manager(self.device, num_ubatches) - self.model_runner = NPUARModelRunner(self.vllm_config, self.device) + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 94c6f8916..eeac7d059 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -7,7 +7,8 @@ import gc import logging import time from collections.abc import Mapping -from copy import copy, deepcopy +from copy import copy +from dataclasses import replace import numpy as np import torch @@ -33,12 +34,14 @@ from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.utils import enable_sp, lmhead_tp_enable from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE -from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_role from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import partition_payload_list -from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin +from vllm_omni.worker.omni_connector_model_runner_mixin import ( + OmniConnectorModelRunnerMixin, + needs_omni_connector, +) class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): @@ -47,31 +50,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._async_chunk = getattr(self.model_config, "async_chunk", False) - # -------------------------------------- Omni-new ------------------------------------------------- - # Mirrors the init allowlist in gpu_generation_model_runner.py. The - # connector-role fallback keeps consumer stages whose arch is not on the - # allowlist (e.g. MiniCPM-o 4.5 stage 2) from starving on input. - _OMNI_CONNECTOR_INIT_ARCHS = { - "Qwen3OmniMoeForConditionalGeneration", - "Qwen2_5OmniForConditionalGeneration", - "CovoAudioForConditionalGeneration", - "MiMoAudioModel", - "Qwen3TTSTalkerForConditionalGeneration", - "Qwen3TTSCode2Wav", - "AudexCode2Wav", - "AudexXCodec1", - "CosyVoice3Model", - "DyninOmniForConditionalGeneration", - "IndexTTS2S2MelDecoder", - } - if ( - getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS - or get_stage_connector_role(self.model_config) is not None - ): + if needs_omni_connector(self.model_config): self.init_omni_connectors( model_config=self.model_config, ) - # -------------------------------------- Omni-new ------------------------------------------------- def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests @@ -110,25 +92,24 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin capturer = self.routed_experts_capturer if capturer is not None and hasattr(capturer, "finalize_pending_copy"): capturer.finalize_pending_copy() - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing: + profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config + if profiling_chunk_config.enabled and profiling_chunk_config.need_timing: if getattr(scheduler_output, "disable_profiling_timing", False): - self.ascend_config.scheduler_config.profiling_chunk_config.need_timing = False + profiling_chunk_config.need_timing = False else: self._sync_device() self._execution_start_time = time.perf_counter() if self.execute_model_state is not None: raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") - # self._draft_token_ids is None when `input_fits_in_drafter=False` - # and there is no draft tokens scheduled. so it need to update the - # spec_decoding info in scheduler_output with async_scheduling. - # use deepcopy to avoid the modification has influence on the - # scheduler_output in engine core process. - # TODO(Ronald1995): deepcopy is expensive when there is a large - # number of requests, optimize it later. - if ( - self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None # type: ignore[has-type] - ): - scheduler_output = deepcopy(scheduler_output) + + if self.speculative_config is not None and self.speculative_config.use_ngram_gpu(): + num_scheduled_tokens_copy = scheduler_output.num_scheduled_tokens.copy() + spec_decode_tokens_copy = scheduler_output.scheduled_spec_decode_tokens.copy() + scheduler_output = replace( + scheduler_output, + num_scheduled_tokens=num_scheduled_tokens_copy, + scheduled_spec_decode_tokens=spec_decode_tokens_copy, + ) if has_kv_transfer_group(): kv_connector_metadata = scheduler_output.kv_connector_metadata @@ -161,7 +142,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin # -------------------------------------- Omni-new ------------------------------------------------- if has_ec_transfer() and get_ec_transfer().is_producer: - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) with self.maybe_get_ec_connector_output( scheduler_output, encoder_cache=self.encoder_cache, @@ -198,7 +179,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin "it when the requests need prompt logprobs" ) - self._start_dump_data() + self._start_dump_data(scheduled_tokens=scheduler_output.num_scheduled_tokens) num_reqs = self.input_batch.num_reqs req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] @@ -368,13 +349,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin with record_function_or_nullcontext("EPLB weight D2D"): self.eplb_updator.forward_before() - # Set cudagraph mode to none if calc_kv_scales is true. - # KV scales calculation involves dynamic operations that are incompatible - # with CUDA graph capture. - if self.calculate_kv_scales: # type: ignore[has-type] - cudagraph_mode = CUDAGraphMode.NONE - # Mark KV scales as calculated after the first forward pass - self.calculate_kv_scales = False # type: ignore[has-type] # Encoder-decoder models can only compile the pure decode steps where no # encoder inputs are present. Use eager for the first pass. num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs) @@ -450,6 +424,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin def sample_tokens( self, grammar_output: GrammarOutput | None ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + profiling_chunk_config = self.ascend_config.scheduler_config.profiling_chunk_config kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -565,8 +540,10 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if self.speculative_config is not None: self.finalize_kv_connector() - if self.ascend_config.scheduler_config.profiling_chunk_config.need_timing and hasattr( - self, "_execution_start_time" + if ( + profiling_chunk_config.enabled + and profiling_chunk_config.need_timing + and hasattr(self, "_execution_start_time") ): self._sync_device() output.execution_time_ms = (time.perf_counter() - self._execution_start_time) * 1000.0 @@ -875,9 +852,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin if get_pp_group().is_first_rank: intermediate_tensors = None else: - # When PP and flashcomm1 are enabled, during dummy_run the estimated space should divide num_tokens by - # tp_size; otherwise, on non-first PP ranks it would effectively perform an extra all-gather, leading - # to incorrect memory estimation and potentially causing OOM. + # When PP and sequence parallelism are enabled, during dummy_run the estimated space should divide + # num_tokens by tp_size; otherwise, on non-first PP ranks it would effectively perform an extra + # all-gather, leading to incorrect memory estimation and potentially causing OOM. intermediate_tokens = num_tokens_padded if enable_sp(): tp_size = get_tensor_model_parallel_world_size() diff --git a/vllm_omni/platforms/npu/worker/npu_generation_worker.py b/vllm_omni/platforms/npu/worker/npu_generation_worker.py index f77975a5d..ac66c6a54 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_worker.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_worker.py @@ -11,9 +11,11 @@ from vllm_omni.worker.mixins import OmniWorkerMixin class NPUGenerationWorker(OmniWorkerMixin, OmniNPUWorkerBase): """NPU generation worker for code2wav stage in Omni model.""" + model_runner_cls = NPUGenerationModelRunner + def init_device(self): self.device = self._init_device() num_ubatches = 1 init_workspace_manager(self.device, num_ubatches) - self.model_runner = NPUGenerationModelRunner(self.vllm_config, self.device) + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 375d36047..8f0369808 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -335,9 +335,9 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): if get_pp_group().is_first_rank: intermediate_tensors = None else: - # When PP and flashcomm1 are enabled, during dummy_run the estimated space should divide num_tokens by - # tp_size; otherwise, on non-first PP ranks it would effectively perform an extra all-gather, leading - # to incorrect memory estimation and potentially causing OOM. + # When PP and sequence parallelism are enabled, during dummy_run the estimated space should divide + # num_tokens by tp_size; otherwise, on non-first PP ranks it would effectively perform an extra + # all-gather, leading to incorrect memory estimation and potentially causing OOM. intermediate_tokens = num_tokens_padded if enable_sp(): tp_size = get_tensor_model_parallel_world_size() @@ -448,7 +448,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): 1. Accepts num_tokens_padded as required by NPUModelRunner 2. Injects omni-specific kwargs (runtime_additional_information) 3. Caches model output for multimodal results - 4. Handles NPU-specific post-forward logic (graph params update, SP all-gather) + 4. Handles NPU-specific post-forward logic (graph params update) """ # Omni-specific: build and inject extra model kwargs model_kwargs_extra = self._build_model_kwargs_extra() @@ -482,10 +482,6 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # Omni-specific: cache model output for later sample_tokens self._omni_last_model_output = model_output - # NPU-specific: all-gather for sequence parallelism - if forward_context.flash_comm_v1_enabled and not isinstance(model_output, IntermediateTensors): - model_output = self._all_gather_hidden_states_and_aux(model_output) - return model_output def _gather_runtime_additional_information(self): diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 0d9fe3639..3d2cd2d62 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -102,8 +102,20 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): @classmethod def record_device_event(cls) -> torch.Event | None: + """Record an XPU event on the current stream to mark tensor readiness. + + Deliberately a device-agnostic ``torch.Event`` rather than a + ``torch.xpu.Event``. The consumer (the async diffusion output thread) + waits with ``torch.Stream.wait_event`` on a generic ``torch.Stream``, + and that C-level binding silently no-ops for a ``torch.xpu.Event`` + instead of enqueuing the dependency — the side stream then starts its + D2H copy while the compute stream is still writing the tensor, so the + host reads a partially-written image (garbage rows at the bottom of the + output). ``torch.Event`` dispatches through the accelerator hooks and + the wait is honored, which is the actual fix. + """ try: - event = torch.xpu.Event() + event = torch.Event() event.record() return event except Exception: diff --git a/vllm_omni/platforms/xpu/utils.py b/vllm_omni/platforms/xpu/utils.py index 18437cae3..6eb314745 100644 --- a/vllm_omni/platforms/xpu/utils.py +++ b/vllm_omni/platforms/xpu/utils.py @@ -1,16 +1,27 @@ from contextlib import contextmanager +from functools import partial import torch +from vllm.utils.torch_utils import supports_xpu_graph @contextmanager def torch_cuda_wrapper(): - try: - # replace cuda APIs with xpu APIs, this should work by default - torch.cuda.Stream = torch.xpu.Stream - torch.cuda.default_stream = torch.xpu.current_stream - torch.cuda.current_stream = torch.xpu.current_stream - torch.cuda.stream = torch.xpu.stream - yield - finally: - pass + # Use partial() so callables differ from torch.xpu.* (Dynamo registration). + torch.cuda.Stream = torch.xpu.Stream + torch.cuda.default_stream = partial(torch.xpu.current_stream) + torch.cuda.current_stream = partial(torch.xpu.current_stream) + torch.cuda.stream = partial(torch.xpu.stream) + torch.cuda.set_stream = partial(torch.xpu.set_stream) + torch.cuda.set_device = partial(torch.xpu.set_device) + + def _xpu_event(*args, blocking=None, **kwargs): + return torch.xpu.Event(*args, **kwargs) + + torch.cuda.Event = _xpu_event + if supports_xpu_graph(): + torch.cuda.graph = partial(torch.xpu.graph) + torch.cuda.CUDAGraph = torch.xpu.XPUGraph + torch.cuda.graph_pool_handle = partial(torch.xpu.graph_pool_handle) + torch.cuda.is_current_stream_capturing = torch.xpu.is_current_stream_capturing + yield diff --git a/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py b/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py index 9acf869f7..e25ea4cc6 100644 --- a/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py +++ b/vllm_omni/platforms/xpu/worker/xpu_ar_worker.py @@ -13,10 +13,12 @@ logger = init_logger(__name__) class XPUARWorker(OmniWorkerMixin, XPUWorker): """XPU AR worker for thinker/talker stages in Omni model.""" + model_runner_cls = XPUARModelRunner + def init_device(self): super().init_device() if self.use_v2_model_runner: # OMNI: v2 model runner does not yet include omni hooks. logger.warning("OMNI XPUARWorker forces v1 model runner for omni hooks.") self.use_v2_model_runner = False - self.model_runner: XPUARModelRunner = XPUARModelRunner(self.vllm_config, self.device) + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) diff --git a/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py b/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py index d9bb790d3..3036f5a5e 100644 --- a/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py +++ b/vllm_omni/platforms/xpu/worker/xpu_generation_worker.py @@ -13,10 +13,12 @@ logger = init_logger(__name__) class XPUGenerationWorker(OmniWorkerMixin, XPUWorker): """XPU generation worker for the code2wav (non-AR waveform generation) stage in the Omni model.""" + model_runner_cls = XPUGenerationModelRunner + def init_device(self): super().init_device() if self.use_v2_model_runner: # OMNI: v2 model runner does not yet include omni hooks. logger.warning("OMNI XPUGenerationWorker forces v1 model runner for omni hooks.") self.use_v2_model_runner = False - self.model_runner: XPUGenerationModelRunner = XPUGenerationModelRunner(self.vllm_config, self.device) + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) diff --git a/vllm_omni/quantization/__init__.py b/vllm_omni/quantization/__init__.py index 8896903e8..89d445d61 100644 --- a/vllm_omni/quantization/__init__.py +++ b/vllm_omni/quantization/__init__.py @@ -11,7 +11,7 @@ Adds per-component quantization for multi-stage models. config = build_quant_config({"transformer": {"method": "fp8"}, "vae": None}) """ -from .component_config import ComponentQuantizationConfig +from .component_config import ComponentQuantizationConfig, resolve_component_quant_config from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override from .inc_config import OmniINCConfig @@ -23,6 +23,7 @@ from .inc_config import OmniINCConfig __all__ = [ "build_quant_config", "ComponentQuantizationConfig", + "resolve_component_quant_config", "OmniINCConfig", "SUPPORTED_QUANTIZATION_METHODS", "register_quantization_override", diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py index 28c33d637..bf6f3b6b2 100644 --- a/vllm_omni/quantization/component_config.py +++ b/vllm_omni/quantization/component_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Per-component quantization routing for multi-stage models. Routes get_quant_method() to different configs based on longest-prefix match: @@ -21,12 +21,31 @@ if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizeMethodBase, ) + from vllm.model_executor.models.utils import ( + WeightsMapper, + ) -# Pre-quantized checkpoints (modelopt FP8/FP4/MXFP8/mixed) only quantize the -# Thinker LM. Vision and audio encoder weights remain in BF16 with no -# corresponding scale tensors in the checkpoint. -PRE_QUANTIZED_METHODS: frozenset[str] = frozenset({"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"}) +# These pre-quantized formats require serialized scale or correction tensors +# that the vision and audio encoder checkpoints do not provide. +PRE_QUANTIZED_METHODS: frozenset[str] = frozenset( + {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed", "svdquant"} +) + + +def resolve_component_quant_config( + quant_config: QuantizationConfig | None, + component: str, +) -> QuantizationConfig | None: + """Resolve one pipeline component from a global or component config. + + A plain config is global and therefore applies unchanged to every + quantization-aware component. Only ``ComponentQuantizationConfig`` narrows + the scope through its explicit prefix map. + """ + if isinstance(quant_config, ComponentQuantizationConfig): + return quant_config.resolve(component) + return quant_config def resolve_encoder_quant_config( @@ -97,6 +116,14 @@ class ComponentQuantizationConfig(QuantizationConfig): return self._components[comp_prefix] return self._default + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: + """Apply a weight mapper to every routed quantization config.""" + for quant_config in self._components.values(): + if quant_config is not None: + quant_config.apply_vllm_mapper(hf_to_vllm_mapper) + if self._default is not None: + self._default.apply_vllm_mapper(hf_to_vllm_mapper) + def get_name(self) -> str: return "component" diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index b5d0a8ad7..f3b6e8cd2 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Factory for building quantization configs. build_quant_config() delegates to vLLM's quantization registry. @@ -55,12 +55,12 @@ def _register_humming_stubs() -> None: registry[name] = mod # wire parent references - registry["humming"].config = registry["humming.config"] - registry["humming"].dtypes = registry["humming.dtypes"] - registry["humming"].layer = registry["humming.layer"] - registry["humming"].schema = registry["humming.schema"] - registry["humming"].utils = registry["humming.utils"] - registry["humming.utils"].weight = registry["humming.utils.weight"] + setattr(registry["humming"], "config", registry["humming.config"]) + setattr(registry["humming"], "dtypes", registry["humming.dtypes"]) + setattr(registry["humming"], "layer", registry["humming.layer"]) + setattr(registry["humming"], "schema", registry["humming.schema"]) + setattr(registry["humming"], "utils", registry["humming.utils"]) + setattr(registry["humming.utils"], "weight", registry["humming.utils.weight"]) for name, mod in registry.items(): sys.modules[name] = mod @@ -125,6 +125,13 @@ def _build_mxfp4_dualscale(**kw: Any) -> QuantizationConfig: return DiffusionMXFP4DualScaleMixedConfig(**kw) +def _build_svdquant(**kw: Any) -> QuantizationConfig: + """Build the serialized SVDQuant diffusion checkpoint loader.""" + from .svdquant_config import DiffusionSVDQuantConfig + + return DiffusionSVDQuantConfig.from_config(kw) + + def _build_inc(**kw: Any) -> QuantizationConfig: """Lazy import for INC/AutoRound config with checkpoint kwarg normalization.""" from .inc_config import OmniINCConfig @@ -145,6 +152,7 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "mxfp8": _build_mxfp8, "mxfp4": _build_mxfp4, "mxfp4_dualscale": _build_mxfp4_dualscale, + "svdquant": _build_svdquant, "inc": _build_inc, "auto-round": _build_inc, "auto_round": _build_inc, @@ -470,10 +478,14 @@ def resolve_quant_config_from_disk( ) return build_quant_config(qc_method, **qc_kwargs) - # AutoRound MXFP8 checkpoints use data_type="mx_fp" instead of - # is_checkpoint_*_serialized; rebuild so the offline path is selected. + # AutoRound MXFP checkpoints use data_type="mx_fp" instead of + # is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is + # selected according to the checkpoint's bit width. if qc_kwargs.get("data_type") == "mx_fp": - logger.info("config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP8.") + logger.info( + "config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP%d.", + qc_kwargs.get("bits", getattr(quant_config, "weight_bits", 0)), + ) return build_quant_config(qc_method, **qc_kwargs) if ( diff --git a/vllm_omni/quantization/int8_config.py b/vllm_omni/quantization/int8_config.py index 01e57e708..de31fbc3a 100644 --- a/vllm_omni/quantization/int8_config.py +++ b/vllm_omni/quantization/int8_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """INT8 quantization config for diffusion transformers. Supports both online (dynamic) and offline (checkpoint) INT8 quantization @@ -26,10 +26,12 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.layers.quantization.fp8 import CopyNumelCounter from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) +from vllm.model_executor.model_loader.reload.meta import ( + CopyCounter as CopyNumelCounter, +) from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight from vllm.model_executor.parameter import ( ChannelQuantScaleParameter, diff --git a/vllm_omni/quantization/mxfp8_config.py b/vllm_omni/quantization/mxfp8_config.py index 433531d56..69e51dc6b 100644 --- a/vllm_omni/quantization/mxfp8_config.py +++ b/vllm_omni/quantization/mxfp8_config.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """W8A8 MXFP8 (Microscaling FP8) online/offline quantization for diffusion transformers. Architecture: @@ -49,8 +49,10 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) -from vllm.model_executor.layers.quantization.fp8 import CopyNumelCounter from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped +from vllm.model_executor.model_loader.reload.meta import ( + CopyCounter as CopyNumelCounter, +) from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight from vllm.model_executor.parameter import ModelWeightParameter from vllm.model_executor.utils import replace_parameter diff --git a/vllm_omni/quantization/svdquant_config.py b/vllm_omni/quantization/svdquant_config.py new file mode 100644 index 000000000..ad22592af --- /dev/null +++ b/vllm_omni/quantization/svdquant_config.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Serialized SVDQuant NVFP4 support for diffusion transformers. + +The checkpoint stores NVFP4 weights plus a rank-R correction for each +quantized linear. The four-bit GEMM uses vLLM's existing NVFP4 kernel +registry, while the rank correction uses ordinary BF16 matrix multiplication. +Native SVDQuant fusion is a separate optimization and is not required to load +or run the checkpoint. +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, Any + +import torch +from torch.nn import Parameter +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import ( + CutlassNvFp4LinearKernel, + FbgemmNvFp4LinearKernel, + FlashInferB12xNvFp4LinearKernel, + FlashInferCudnnNvFp4LinearKernel, + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + FlashInferTrtllmNvFp4LinearKernel, + NvFp4LinearKernel, + init_nvfp4_linear_kernel, +) +from vllm.model_executor.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization import QuantizationMethods + +logger = init_logger(__name__) + +_SUPPORTED_CAPABILITIES = {(10, 3)} +_COMPATIBLE_NVFP4_KERNELS = ( + CutlassNvFp4LinearKernel, + FbgemmNvFp4LinearKernel, + FlashInferB12xNvFp4LinearKernel, + FlashInferCudnnNvFp4LinearKernel, + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + FlashInferTrtllmNvFp4LinearKernel, +) + + +def _supports_capability(capability: DeviceCapability | None) -> bool: + return ( + capability is not None + and ( + capability.major, + capability.minor, + ) + in _SUPPORTED_CAPABILITIES + ) + + +def _assert_supported() -> None: + if not current_platform.is_cuda(): + raise RuntimeError("SVDQuant NVFP4 requires a CUDA device") + capability = current_platform.get_device_capability() + if not _supports_capability(capability): + device = current_platform.device_name + sm = capability.to_int() if capability is not None else "unknown" + raise RuntimeError(f"SVDQuant NVFP4 is validated on SM103 only; got {device!r} (SM{sm})") + + +@functools.cache +def _nvfp4_kernel() -> NvFp4LinearKernel: + kernel = init_nvfp4_linear_kernel() + if not isinstance(kernel, _COMPATIBLE_NVFP4_KERNELS): + raise RuntimeError( + "The selected vLLM NVFP4 backend is incompatible with the " + f"SVDQuant checkpoint layout: {type(kernel).__name__}. Use one of " + "the FlashInfer, CUTLASS, or FBGEMM NVFP4 backends." + ) + return kernel + + +class DiffusionSVDQuantConfig(QuantizationConfig): + """Configuration for serialized NVFP4 W4A4 plus low-rank correction.""" + + def __init__( + self, + rank: int = 32, + precision: str = "nvfp4", + act_unsigned: bool = False, + modules_to_not_convert: list[str] | None = None, + ) -> None: + super().__init__() + if rank <= 0: + raise ValueError(f"SVDQuant rank must be positive, got {rank}") + if precision != "nvfp4": + raise ValueError( + f"Phase 1 SVDQuant supports serialized NVFP4 checkpoints only; got precision={precision!r}" + ) + if act_unsigned: + raise ValueError("Phase 1 SVDQuant does not support unsigned activations") + self.rank = rank + self.precision = precision + self.modules_to_not_convert = modules_to_not_convert or [] + + def __repr__(self) -> str: + return f"DiffusionSVDQuantConfig(rank={self.rank}, precision={self.precision!r})" + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "svdquant" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 103 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return ["quantization_config.json"] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> DiffusionSVDQuantConfig: + return cls( + rank=config.get("rank", 32), + precision=config.get("precision", "nvfp4"), + act_unsigned=config.get("act_unsigned", False), + modules_to_not_convert=config.get("modules_to_not_convert"), + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> QuantizeMethodBase | None: + if not isinstance(layer, LinearBase): + return None + if is_layer_skipped( + prefix, + self.modules_to_not_convert, + self.packed_modules_mapping, + # vLLM 0.28: skip_with_substr was replaced by match_mode + # ("substring" preserves the v0.27 skip_with_substr=True + # behavior verbatim). + match_mode="substring", + ): + return UnquantizedLinearMethod() + return DiffusionSVDQuantLinearMethod(self) + + +class DiffusionSVDQuantLinearMethod(LinearMethodBase): + """Load and execute a serialized SVDQuant linear layer.""" + + def __init__(self, quant_config: DiffusionSVDQuantConfig) -> None: + _assert_supported() + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs: Any, + ) -> None: + del params_dtype + weight_loader = extra_weight_attrs.pop( + "weight_loader", + default_weight_loader, + ) + if input_size_per_partition % 16 != 0: + raise ValueError( + "SVDQuant NVFP4 requires each input partition to be divisible " + f"by the block size 16; got {input_size_per_partition}" + ) + output_size_per_partition = sum(output_partition_sizes) + rank = self.quant_config.rank + + qweight = Parameter( + torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + set_weight_attrs( + qweight, + { + "input_dim": 1, + "output_dim": 0, + "weight_loader": weight_loader, + }, + ) + + wscales = Parameter( + torch.empty( + input_size_per_partition // 16, + output_size_per_partition, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + set_weight_attrs( + wscales, + { + "input_dim": 0, + "output_dim": 1, + "weight_loader": weight_loader, + }, + ) + + proj_down = Parameter( + torch.empty( + input_size_per_partition, + rank, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + set_weight_attrs( + proj_down, + { + "input_dim": 0, + "weight_loader": weight_loader, + }, + ) + + proj_up = Parameter( + torch.empty( + output_size_per_partition, + rank, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + set_weight_attrs( + proj_up, + { + "output_dim": 0, + "weight_loader": weight_loader, + }, + ) + + smooth_factor = Parameter( + torch.empty( + input_size_per_partition, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + set_weight_attrs( + smooth_factor, + { + "input_dim": 0, + "weight_loader": weight_loader, + }, + ) + + wcscales = Parameter( + torch.ones( + output_size_per_partition, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + set_weight_attrs( + wcscales, + { + "output_dim": 0, + "weight_loader": weight_loader, + }, + ) + + wtscale = Parameter( + torch.ones(1, dtype=torch.bfloat16), + requires_grad=False, + ) + set_weight_attrs(wtscale, {"weight_loader": default_weight_loader}) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("wscales", wscales) + layer.register_parameter("proj_down", proj_down) + layer.register_parameter("proj_up", proj_up) + layer.register_parameter("smooth_factor", smooth_factor) + layer.register_parameter("wcscales", wcscales) + layer.register_parameter("wtscale", wtscale) + + del input_size, output_size + layer.output_size_per_partition = output_size_per_partition + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Adapt the canonical row-major checkpoint to vLLM's NVFP4 ABI.""" + qweight = layer.qweight + wscales = layer.wscales + del layer.qweight + del layer.wscales + layer.register_parameter( + "weight", + Parameter( + qweight.detach().view(torch.uint8), + requires_grad=False, + ), + ) + layer.register_parameter( + "weight_scale", + Parameter( + wscales.detach().transpose(0, 1).contiguous(), + requires_grad=False, + ), + ) + + layer.register_parameter( + "input_global_scale_inv", + Parameter( + torch.ones( + 1, + dtype=torch.float32, + device=layer.weight.device, + ), + requires_grad=False, + ), + ) + + wtscale = layer.wtscale.detach().to(dtype=torch.float32) + del layer.wtscale + layer.register_parameter( + "alpha", + Parameter(wtscale, requires_grad=False), + ) + + channel_scale = layer.wcscales.detach() + del layer.wcscales + if torch.all(channel_scale == 1).item(): + layer.output_channel_scale = None + else: + layer.register_parameter( + "output_channel_scale", + Parameter(channel_scale, requires_grad=False), + ) + + _nvfp4_kernel().process_weights_after_loading(layer) + logger.info_once( + "SVDQuant NVFP4 is using vLLM's compatibility path; the rank correction is not fused into the GEMM." + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Compute the base NVFP4 GEMM plus the BF16 rank correction.""" + if x.dtype != torch.bfloat16: + raise ValueError(f"SVDQuant NVFP4 requires BF16 activations; got {x.dtype}") + + original_shape = x.shape + x_2d = x.reshape(-1, original_shape[-1]).contiguous() + + # The residual branch consumes the original activation. Only the + # four-bit base GEMM consumes the smoothed activation. + smoothed = x_2d / layer.smooth_factor + out = _nvfp4_kernel().apply_weights( + layer=layer, + x=smoothed, + bias=None, + ) + + channel_scale = getattr(layer, "output_channel_scale", None) + if channel_scale is not None: + # Fused QKV can store independent Q/K/V outer scales. Apply them + # explicitly until a vector-alpha GEMM epilogue is available. + out.mul_(channel_scale) + + correction_input = torch.mm(x_2d, layer.proj_down) + out = torch.addmm( + out, + correction_input, + layer.proj_up.transpose(0, 1), + ) + if bias is not None: + out.add_(bias) + return out.reshape( + *original_shape[:-1], + layer.output_size_per_partition, + ) + + +__all__ = [ + "DiffusionSVDQuantConfig", + "DiffusionSVDQuantLinearMethod", +] diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 0606bcfc8..a6da16696 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING @@ -103,6 +106,7 @@ class OmniRequest(Request): additional_information=request.additional_information, model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None), resumable=request.resumable, + session_id=request.session_id, reasoning_ended=request.reasoning_ended, reasoning_parser_kwargs=request.reasoning_parser_kwargs, abort_immediately=request.abort_immediately, diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 88a23dc0c..42d6d8d45 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -20,8 +20,10 @@ _CLASS_TO_MODULE: dict[str, str] = { "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "GLMTTSConfig": "vllm_omni.transformers_utils.configs.glm_tts", "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2", + "DotsTTSConfig": "vllm_omni.transformers_utils.configs.dots_tts", "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3", + "MiniMaxMusic3Config": "vllm_omni.transformers_utils.configs.minimax_music3", "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice", "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", @@ -29,6 +31,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "Qwen3VLMoeVisionConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", + "SenseNovaU1Config": "vllm_omni.transformers_utils.configs.sensenova_u1", } __all__ = [ @@ -42,8 +45,10 @@ __all__ = [ "FishSpeechFastARConfig", "GLMTTSConfig", "VoxCPM2Config", + "DotsTTSConfig", "VoxtralTTSConfig", "CosyVoice3Config", + "MiniMaxMusic3Config", "OmniVoiceConfig", "BailingMoeV2Config", "BailingMM2Config", @@ -51,6 +56,7 @@ __all__ = [ "MingFlashOmniConfig", "Qwen3VLMoeVisionConfig", "WhisperEncoderConfig", + "SenseNovaU1Config", ] @@ -70,11 +76,14 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import minimax_music3 as _minimax_music3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import sensenova_u1 as _sensenova_u1 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/cosyvoice3.py b/vllm_omni/transformers_utils/configs/cosyvoice3.py index b08cec063..dddcb4659 100644 --- a/vllm_omni/transformers_utils/configs/cosyvoice3.py +++ b/vllm_omni/transformers_utils/configs/cosyvoice3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from transformers import AutoConfig from transformers.configuration_utils import PretrainedConfig @@ -18,6 +18,7 @@ class CosyVoice3Config(PretrainedConfig): self.llm_output_size = 896 self.hidden_size = self.llm_output_size self.num_attention_heads = 14 + self.num_key_value_heads = 2 self.num_hidden_layers = 24 self.spk_embed_dim = 192 self.token_frame_rate = 25 diff --git a/vllm_omni/transformers_utils/configs/dots_tts.py b/vllm_omni/transformers_utils/configs/dots_tts.py new file mode 100644 index 000000000..1950a2b3a --- /dev/null +++ b/vllm_omni/transformers_utils/configs/dots_tts.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HF config for dots.tts. + +dots.tts splits its config across two files in the HF checkpoint: + * ``config.json`` — top-level dots.tts fields (latent_dim / patch_size + / PatchEncoder / DiT / vocoder / etc.). + * ``llm_config.json`` — full Qwen2 LM config (vocab_size / hidden_size / + num_hidden_layers / rope_theta / etc.). + +``AutoConfig.from_pretrained`` only consumes ``config.json``: its dots.tts +sub-blocks (``vocoder`` / ``DiT`` / ``PatchEncoder``) arrive through +``**kwargs`` and land as attributes, which the talker's config factories +consume directly. The Qwen2 LM fields live in ``llm_config.json``, which +AutoConfig never reads — this class hoists them to top-level attributes +with defaults matching the released checkpoints (soar / base / mf all ship +the same Qwen2.5-1.5B backbone). vLLM's ``Qwen2Model.__init__`` then reads +``config.hidden_size`` / ``config.vocab_size`` / etc. through +``get_text_config()``, the same pattern voxcpm2 follows. + +Future work: read ``llm_config.json`` during ``from_pretrained`` so a +checkpoint with a different LM shape works without editing this file. + +Upstream reference: + https://huggingface.co/rednote-hilab/dots.tts-soar/blob/main/llm_config.json +""" + +from __future__ import annotations + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig + + +class DotsTTSConfig(PretrainedConfig): + """Configuration for dots.tts integration (Qwen2.5-1.5B-Base AR backbone).""" + + model_type = "dots_tts" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + # -- top-level dots.tts params -- + lm_config: dict | None = None, + # -- LM defaults (overridden by lm_config dict if present) -- + # All values match dots.tts-soar's llm_config.json. + vocab_size: int = 151672, + hidden_size: int = 1536, + intermediate_size: int = 8960, + num_hidden_layers: int = 28, + num_attention_heads: int = 12, + num_key_value_heads: int = 2, + max_position_embeddings: int = 131072, + rms_norm_eps: float = 1e-6, + rope_theta: float = 1_000_000.0, + tie_word_embeddings: bool = True, + hidden_act: str = "silu", + bos_token_id: int = 151643, + eos_token_id: int = 151643, + **kwargs, + ): + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + # The dots.tts sub-blocks (vocoder / DiT / PatchEncoder / latent_dim + # / patch_size / ...) arrive via **kwargs and are stored as + # attributes by PretrainedConfig; the talker's config factories + # read them from there. + self.lm_config = lm_config or {} + + # -- Hoist LM parameters to top-level for vLLM's Qwen2Model -- + lm = self.lm_config + self.vocab_size = lm.get("vocab_size", vocab_size) + self.hidden_size = lm.get("hidden_size", hidden_size) + self.intermediate_size = lm.get("intermediate_size", intermediate_size) + self.num_hidden_layers = lm.get("num_hidden_layers", num_hidden_layers) + self.num_attention_heads = lm.get("num_attention_heads", num_attention_heads) + self.num_key_value_heads = lm.get("num_key_value_heads", num_key_value_heads) + self.max_position_embeddings = lm.get("max_position_embeddings", max_position_embeddings) + self.rms_norm_eps = lm.get("rms_norm_eps", rms_norm_eps) + self.rope_theta = lm.get("rope_theta", rope_theta) + self.hidden_act = lm.get("hidden_act", hidden_act) + self.head_dim = self.hidden_size // self.num_attention_heads + + # vLLM's Qwen2Model reads ``max_window_layers`` and checks against + # ``num_hidden_layers`` only when ``use_sliding_window`` is true. + # soar disables sliding window — set both to the same value so the + # assertion (if executed) stays vacuous. + self.use_sliding_window = lm.get("use_sliding_window", False) + self.sliding_window = lm.get("sliding_window", None) + self.max_window_layers = lm.get("max_window_layers", self.num_hidden_layers) + + def get_text_config(self, **kwargs): + """Return self — LM fields are hoisted to top-level attributes.""" + return self + + +AutoConfig.register("dots_tts", DotsTTSConfig) + +__all__ = ["DotsTTSConfig"] diff --git a/vllm_omni/transformers_utils/configs/minimax_music3.py b/vllm_omni/transformers_utils/configs/minimax_music3.py new file mode 100644 index 000000000..be951055d --- /dev/null +++ b/vllm_omni/transformers_utils/configs/minimax_music3.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HuggingFace config for the MiniMax Music 3 composite repository. + +The checkpoint root carries only:: + + {"architectures": ["MiniMaxMusic3ForConditionalGeneration"], + "model_type": "minimax_music3"} + +and ships no modelling code, so ``trust_remote_code`` has nothing to load and +``AutoConfig`` rejects the repository outright. Stage 0 sidesteps this by +declaring ``model_subdir="language_model"`` and reading the plain +``Qwen3ForCausalLM`` config, but the acoustic stage must open the root to reach +``condition_encoder/``, ``transformer/`` and ``vocoder/``. Registering this +class is what lets that stage build a ``ModelConfig`` at all. + +The generic transformer fields describe the AR backbone, because the composite +root has no transformer of its own: they are metadata for vLLM's config +plumbing, not a description of the acoustic stage. The acoustic stage holds no +vLLM attention layers, so no KV cache is derived from them; it reads its real +geometry from ``models/minimax_music3/constants.py`` and its weights from the +component folders. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig + + +class MiniMaxMusic3Config(PretrainedConfig): + """Root config for the multi-component MiniMax Music 3 repository.""" + + model_type = "minimax_music3" + + def __init__(self, **kwargs: Any) -> None: + # AR backbone geometry (language_model/config.json). Kept in sync with + # the checkpoint so max_model_len validation sees the real context. + kwargs.setdefault("hidden_size", 4096) + kwargs.setdefault("num_hidden_layers", 36) + kwargs.setdefault("num_attention_heads", 32) + kwargs.setdefault("num_key_value_heads", 8) + kwargs.setdefault("head_dim", 128) + kwargs.setdefault("intermediate_size", 12288) + kwargs.setdefault("vocab_size", 200000) + kwargs.setdefault("max_position_embeddings", 10240) + kwargs.setdefault("tie_word_embeddings", False) + super().__init__(**kwargs) + + # Audio contract, mirroring models/minimax_music3/constants.py. These + # are fixed by the checkpoint; they are exposed here only so a caller + # inspecting the config sees the same numbers the model uses. + self.audio_frame_rate = 25 + self.num_codebooks = 8 + self.audio_vocab_size = 1024 + self.c0_vocab_size = 16384 + self.sample_rate = 32000 + self.audio_channels = 2 + + +AutoConfig.register("minimax_music3", MiniMaxMusic3Config) + +__all__ = ["MiniMaxMusic3Config"] diff --git a/vllm_omni/transformers_utils/configs/sensenova_u1.py b/vllm_omni/transformers_utils/configs/sensenova_u1.py new file mode 100644 index 000000000..81e8b3406 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/sensenova_u1.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HuggingFace-style configuration classes for SenseNova-U1.""" + +from transformers import AutoConfig, PretrainedConfig, Qwen3Config + +_SENSENOVA_U1_LLM_HIDDEN_DEFAULT = 4096 + + +# Adapted from: https://github.com/OpenSenseNova/SenseNova-U1/blob/main/src/sensenova_u1/models/neo_unify/configuration_neo_chat.py +class SenseNovaU1LLMConfig(Qwen3Config): + """Qwen3-based LLM backbone config with 3D RoPE extensions.""" + + model_type = "sensenova_u1_llm" + + def __init__( + self, + rope_theta: float = 10000.0, + rope_theta_hw: float = 10000.0, + max_position_embeddings_hw: int = 10000, + **kwargs, + ): + self.rope_theta = rope_theta + self.rope_theta_hw = rope_theta_hw + self.max_position_embeddings_hw = max_position_embeddings_hw + super().__init__(**kwargs) + + +# Adapted from https://github.com/OpenSenseNova/SenseNova-U1/blob/main/src/sensenova_u1/models/neo_unify/configuration_neo_vit.py#L10 +class SenseNovaU1VisionConfig(PretrainedConfig): + """Vision embedding config (2D RoPE + conv patch embed, no transformer).""" + + model_type = "sensenova_u1_vision" + + def __init__( + self, + num_channels: int = 3, + patch_size: int = 16, + hidden_size: int = 1024, + llm_hidden_size: list[int] | None = None, + downsample_ratio: list[float] | None = None, + rope_theta_vision: float = 10000.0, + max_position_embeddings_vision: int = 10000, + **kwargs, + ): + self.num_channels = num_channels + self.patch_size = patch_size + self.hidden_size = hidden_size + self.llm_hidden_size = llm_hidden_size if llm_hidden_size is not None else [_SENSENOVA_U1_LLM_HIDDEN_DEFAULT] + self.downsample_ratio = downsample_ratio if downsample_ratio is not None else [0.5] + self.rope_theta_vision = rope_theta_vision + self.max_position_embeddings_vision = max_position_embeddings_vision + super().__init__(**kwargs) + + +# Adapted from https://github.com/OpenSenseNova/SenseNova-U1/blob/main/training/sensenovavl/model/sensenovavl_moe_chat/configuration_sensenovavl_chat.py#L17 +class SenseNovaU1Config(PretrainedConfig): + """Top-level composite config for SenseNova-U1. + + Nests ``llm_config`` and ``vision_config`` sub-configs alongside the + flow-matching / diffusion parameters. When constructed from a dict + (e.g. via ``from_pretrained``), sub-dicts are automatically promoted + to their typed config objects. + """ + + model_type = "sensenova_u1" + + def __init__( + self, + llm_config: dict | SenseNovaU1LLMConfig | None = None, + vision_config: dict | SenseNovaU1VisionConfig | None = None, + downsample_ratio: float = 0.5, + template: str = "neo1_0", + fm_head_layers: int = 2, + fm_head_dim: int = 4096, + fm_head_mlp_ratio: float = 1.0, + use_pixel_head: bool = False, + noise_scale: float = 1.0, + noise_scale_mode: str = "none", + noise_scale_base_image_seq_len: int = 256, + noise_scale_max_value: float = 10.0, + add_noise_scale_embedding: bool = False, + time_schedule: str = "standard", + time_shift_type: str = "exponential", + base_shift: float = 0.5, + max_shift: float = 1.15, + base_image_seq_len: int = 256, + max_image_seq_len: int = 4096, + concat_time_token_num: int = 0, + t_eps: float = 0.02, + **kwargs, + ): + if isinstance(llm_config, dict): + llm_config = SenseNovaU1LLMConfig(**llm_config) + self.llm_config = llm_config or SenseNovaU1LLMConfig() + + if isinstance(vision_config, dict): + vision_config = SenseNovaU1VisionConfig(**vision_config) + self.vision_config = vision_config or SenseNovaU1VisionConfig() + + self.downsample_ratio = downsample_ratio + self.template = template + self.fm_head_layers = fm_head_layers + self.fm_head_dim = fm_head_dim + self.fm_head_mlp_ratio = fm_head_mlp_ratio + self.use_pixel_head = use_pixel_head + self.noise_scale = noise_scale + self.noise_scale_mode = noise_scale_mode + self.noise_scale_base_image_seq_len = noise_scale_base_image_seq_len + self.noise_scale_max_value = noise_scale_max_value + self.add_noise_scale_embedding = add_noise_scale_embedding + self.time_schedule = time_schedule + self.time_shift_type = time_shift_type + self.base_shift = base_shift + self.max_shift = max_shift + self.base_image_seq_len = base_image_seq_len + self.max_image_seq_len = max_image_seq_len + self.concat_time_token_num = concat_time_token_num + self.t_eps = t_eps + super().__init__(**kwargs) + + +AutoConfig.register("sensenova_u1", SenseNovaU1Config) + +__all__ = [ + "SenseNovaU1Config", + "SenseNovaU1LLMConfig", + "SenseNovaU1VisionConfig", +] diff --git a/vllm_omni/transformers_utils/processors/audiox.py b/vllm_omni/transformers_utils/processors/audiox.py deleted file mode 100644 index b0e53cfc7..000000000 --- a/vllm_omni/transformers_utils/processors/audiox.py +++ /dev/null @@ -1,240 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Input transform utilities for the AudioX diffusion pipeline. - -Loads and normalizes the raw audio/video conditioning signals (file path / URL / -``data:`` URI / ``np.ndarray`` / ``torch.Tensor``) into the (channels, samples) and -[T, C, H, W] tensors the pipeline needs. The pipeline itself stays focused on model -forward + sampling logic. -""" - -from __future__ import annotations - -import base64 -import os -import tempfile -from typing import Any -from urllib.parse import urlparse -from urllib.request import urlopen - -import av -import numpy as np -import soundfile -import torch -import torch.nn.functional as F -import torchaudio.functional as taF -from einops import rearrange -from torchvision.io import read_image - -# AudioX task taxonomy. Tasks beginning with "v" require a video input; tasks containing -# "t" carry a text prompt. tv2*/v2* share the same conditioning pathways. -VIDEO_ONLY_TASKS = frozenset({"v2a", "v2m"}) -TEXT_VIDEO_TASKS = frozenset({"tv2a", "tv2m"}) -VIDEO_CONDITIONED_TASKS = VIDEO_ONLY_TASKS | TEXT_VIDEO_TASKS - -_IMAGE_EXTS = frozenset({".jpg", ".jpeg", ".png"}) - - -def normalize_prompts(prompts: list[Any]) -> list[dict[str, Any]]: - """Coerce raw prompt entries into ``{"prompt": str, ...}`` dicts (preserves extras).""" - out: list[dict[str, Any]] = [] - for i, raw in enumerate(prompts): - if isinstance(raw, str): - out.append({"prompt": raw.strip()}) - elif isinstance(raw, dict): - p = dict(raw) - p["prompt"] = str(p.get("prompt") or "").strip() - out.append(p) - else: - raise TypeError(f"AudioX prompt {i} must be str or dict, got {type(raw)!r}") - return out - - -def materialize_media_source(source: str) -> str: - """Return a local filesystem path for ``source``. - - Accepts a local path, a ``data:;base64,...`` URI, or an ``http(s)://`` URL. - Anything non-local is fetched into a NamedTemporaryFile and that path is returned; - callers don't need to clean the tempfile up (the OS does on exit). - """ - if source.startswith("data:"): - _, _, payload = source.partition(",") - raw = base64.b64decode(payload) - f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False) - f.write(raw) - f.close() - return f.name - parsed = urlparse(source) - if parsed.scheme in ("http", "https"): - with urlopen(source) as resp: - data = resp.read() - f = tempfile.NamedTemporaryFile(prefix="audiox_media_", suffix=".bin", delete=False) - f.write(data) - f.close() - return f.name - return source - - -def _load_video_path_pyav( - path: str, - *, - target_fps: int, - duration: float, - seek_time: float, -) -> torch.Tensor: - path = materialize_media_source(path) - seek_time = float(seek_time) - duration = float(duration) - end_time = seek_time + duration if duration > 0 else None - - with av.open(path) as container: - stream = container.streams.video[0] - time_base = float(stream.time_base) - src_fps = float(stream.average_rate or target_fps) - - if seek_time > 0: - container.seek(int(seek_time / time_base), stream=stream) - - frames: list[np.ndarray] = [] - for frame in container.decode(stream): - t = float(frame.pts) * time_base - if t < seek_time: - continue - if end_time is not None and t >= end_time: - break - frames.append(frame.to_ndarray(format="rgb24")) - - if not frames: - raise ValueError(f"No frames in range seek_time={seek_time!r}, duration={duration!r} for {path!r}") - - # PyAV gives [H, W, C] uint8 RGB per frame; AudioX expects [T, C, H, W]. - video = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).contiguous() - n_target = max(1, int(round(video.shape[0] * float(target_fps) / src_fps))) - if n_target >= video.shape[0]: - return video - indices = torch.linspace(0, video.shape[0] - 1, n_target).round().long() - return video[indices] - - -def load_video_source( - source: Any, - *, - target_fps: int, - duration: float, - seek_time: float = 0.0, -) -> torch.Tensor: - if isinstance(source, str): - ext = os.path.splitext(source)[1].lower() - if ext in _IMAGE_EXTS: - return read_image(materialize_media_source(source)).unsqueeze(0) - return _load_video_path_pyav( - source, - target_fps=target_fps, - duration=duration, - seek_time=seek_time, - ) - - if isinstance(source, torch.Tensor): - return source - if isinstance(source, np.ndarray): - return torch.from_numpy(source) - raise TypeError(f"Unsupported video source type: {type(source)!r}") - - -def normalize_video_tensor(frames: torch.Tensor, size: int = 224) -> torch.Tensor: - if frames.dim() != 4: - raise ValueError(f"Expected [T, C, H, W], got {tuple(frames.shape)}") - - frames = frames.float() - if frames.max() > 1.5: - frames = frames / 255.0 - - if frames.shape[-2:] != (size, size): - frames = F.interpolate(frames, size=(size, size), mode="bicubic", align_corners=False) - return frames - - -def adjust_video_duration(frames: torch.Tensor, duration: float, target_fps: int) -> torch.Tensor: - target_t = int(duration * target_fps) - cur_t = frames.shape[0] - - if cur_t > target_t: - return frames[:target_t] - if cur_t < target_t: - last = frames[-1:].repeat(target_t - cur_t, 1, 1, 1) - return torch.cat([frames, last], dim=0) - return frames - - -def prepare_video_reference( - source: Any, - *, - duration: float, - target_fps: int, - seek_time: float = 0.0, -) -> torch.Tensor: - """Decode a video clip (or single image) into the AudioX [T, 3, 224, 224] form.""" - frames = load_video_source( - source, - target_fps=target_fps, - duration=duration, - seek_time=seek_time, - ) - - if frames.dim() == 4 and frames.shape[-1] == 3: - frames = rearrange(frames, "t h w c -> t c h w") - - frames = normalize_video_tensor(frames, size=224) - if duration > 0: - frames = adjust_video_duration(frames, duration, target_fps) - return frames - - -def prepare_audio_reference( - source: Any, - *, - model_sample_rate: int, - seconds_start: float, - seconds_total: float, - device: torch.device, -) -> torch.Tensor: - """Decode an audio source into a stereo (2, samples) tensor at the model's rate.""" - target_len = int(model_sample_rate * seconds_total) - start = int(model_sample_rate * seconds_start) - if isinstance(source, str): - data, sr = soundfile.read(materialize_media_source(source), dtype="float32", always_2d=True) - # soundfile returns channels-last (T, C); project convention is (C, T). - wav = torch.from_numpy(data).transpose(0, 1).contiguous() - if sr != model_sample_rate: - wav = taF.resample(wav, sr, model_sample_rate) - elif isinstance(source, torch.Tensor): - wav = source.float() - elif isinstance(source, np.ndarray): - wav = torch.from_numpy(source).float() - else: - raise TypeError(f"Unsupported audio source type: {type(source)!r}") - if wav.dim() == 1: - wav = wav.unsqueeze(0) - if wav.shape[0] == 1: - wav = wav.repeat(2, 1) - elif wav.shape[0] > 2: - wav = wav[:2] - wav = wav[:, start : start + target_len] - if wav.shape[1] < target_len: - wav = F.pad(wav, (target_len - wav.shape[1], 0)) - return wav.to(device=device, dtype=torch.float32) - - -__all__ = [ - "VIDEO_ONLY_TASKS", - "TEXT_VIDEO_TASKS", - "VIDEO_CONDITIONED_TASKS", - "normalize_prompts", - "materialize_media_source", - "load_video_source", - "normalize_video_tensor", - "adjust_video_duration", - "prepare_video_reference", - "prepare_audio_reference", -] diff --git a/vllm_omni/utils/speaker_cache.py b/vllm_omni/utils/speaker_cache.py index d29141489..d571a36d4 100644 --- a/vllm_omni/utils/speaker_cache.py +++ b/vllm_omni/utils/speaker_cache.py @@ -10,6 +10,7 @@ import json import os import threading from collections import OrderedDict +from collections.abc import Callable from pathlib import Path from typing import Any @@ -131,7 +132,7 @@ def _first_dim(tensor: torch.Tensor | None) -> int: return int(tensor.shape[0]) -def _validate_qwen3_tts_profile( +def validate_qwen3_tts_profile( profile: dict[str, Any], tensors: dict[str, torch.Tensor], *, @@ -164,7 +165,7 @@ def _validate_qwen3_tts_profile( return None -def _validate_voxcpm2_profile(profile: dict[str, Any], tensors: dict[str, torch.Tensor]) -> str | None: +def validate_voxcpm2_profile(profile: dict[str, Any], tensors: dict[str, torch.Tensor]) -> str | None: # Successful validation also normalizes/augments `profile` for serving metadata. ref_audio_feat = tensors.get("ref_audio_feat") audio_feat = tensors.get("audio_feat") @@ -207,7 +208,7 @@ def load_validated_profile_tensors( profile: dict[str, Any], *, expected_model_type: str, - qwen3_embedding_dim: int | None = None, + validate_profile: Callable[[dict[str, Any], dict[str, torch.Tensor]], str | None], ) -> dict[str, torch.Tensor] | None: tensors = _load_profile_tensors(profile) error = None @@ -215,12 +216,8 @@ def load_validated_profile_tensors( error = "safetensors file could not be loaded" elif str(profile.get("model_type") or expected_model_type or "") != expected_model_type: error = f"model_type={profile.get('model_type')!r}, expected={expected_model_type!r}" - elif expected_model_type == "qwen3_tts": - error = _validate_qwen3_tts_profile(profile, tensors, expected_embedding_dim=qwen3_embedding_dim) - elif expected_model_type == "voxcpm2": - error = _validate_voxcpm2_profile(profile, tensors) else: - error = f"unsupported custom voice model_type={expected_model_type!r}" + error = validate_profile(profile, tensors) if error is not None: logger.warning("Skipping custom %s voice %s: %s", expected_model_type, profile.get("name"), error) diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index c1e229118..d616347d8 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -74,7 +74,7 @@ class OmniGPUWorkerBase(GPUWorker): if self.profiler is None: raise RuntimeError( "Profiling is not enabled. For diffusion models, set --profiler-config via CLI. " - "For omni models, add profiler_config to your stage config file." + "For omni models, add profiler_config to the relevant deploy config stage." ) if is_start: from vllm_omni.profiler import OmniTorchProfilerWrapper @@ -201,11 +201,19 @@ class OmniGPUWorkerBase(GPUWorker): """ from vllm.device_allocator.cumem import CuMemAllocator + # Drain in-flight kernels before CuMem unmap. Abort/decode can still + # have GPU work after the scheduler reports idle; unmapping then + # returns CUDA_ERROR_INVALID_VALUE (cudaErrorInvalidValue). + current_omni_platform.synchronize() + gc.collect() + mem_before = current_omni_platform.get_current_memory_usage(self.device) offload_tags = ("weights",) if level == 1 else tuple() allocator = CuMemAllocator.get_instance() allocator.sleep(offload_tags=offload_tags) - current_omni_platform.empty_cache() + # allocator.sleep already gc.collects and empty_cache()s. A second + # empty_cache on the pluggable allocator is a known CUDA invalid + # argument (pytorch/pytorch#145168). Only sync for the readout. current_omni_platform.synchronize() mem_after = current_omni_platform.get_current_memory_usage(self.device) freed = max(0, mem_before - mem_after) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b3c388e3a..cd03cf3fc 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """AR GPU Model Runner for vLLM-Omni. Exposes per-request hidden representations via ModelRunnerOutput.pooler_output @@ -41,21 +44,23 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager -from vllm_omni.distributed.omni_connectors.utils.config import ( - get_stage_connector_role, - stage_sends_async_output, -) +from vllm_omni.distributed.omni_connectors.utils.config import stage_sends_async_output +from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import ( build_mm_cpu, partition_payload_list, snapshot_mm_payload, - to_payload_element, ) from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner -from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin +from vllm_omni.worker.omni_connector_model_runner_mixin import ( + OmniConnectorModelRunnerMixin, + needs_omni_connector, +) +from vllm_omni.worker.output.payload_build import build_omni_mm_payload from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest -from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids +from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding, sanitize_min_tokens_stop_ids +from vllm_omni.worker.sparse_audio import resolve_sparse_mm_routing logger = init_logger(__name__) @@ -167,6 +172,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): async_output_copy_stream = kwargs.pop("async_output_copy_stream") vocab_size = kwargs.pop("vocab_size") routed_experts = kwargs.pop("routed_experts", None) + num_nans = kwargs.pop("num_nans", None) # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault # for EP all2all fault tolerance (PR #43637). Omni doesn't use this # feature but must consume the kwarg to prevent TypeError from stray @@ -184,6 +190,11 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): self._logprobs_tensors = logprobs_tensors self._routed_experts = routed_experts self._has_fault: torch.Tensor | None = None + # Upstream b1e12d142d (PR #51304) added device-side NaN-in-logits + # counts (num_nans) to AsyncGPUModelRunnerOutput. Omni keeps the + # counts on the async copy stream and lets super().get_output() + # populate num_nans_in_logits from the CPU copy. + self._num_nans = num_nans default_stream = torch.cuda.current_stream() with torch.cuda.stream(async_output_copy_stream): @@ -197,6 +208,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): self._routed_experts_cpu = ( self._routed_experts.to_cpu_nonblocking() if self._routed_experts is not None else None ) + self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None self.async_copy_ready_event.record() self._model_runner_output_builder = model_runner_output_builder @@ -240,6 +252,12 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): # (e.g. unit tests using object.__new__). if not hasattr(self, "_has_fault"): self._has_fault = None + # Upstream b1e12d142d (PR #51304) also touches _num_nans/_num_nans_cpu + # in get_output(). Guard them the same way for object.__new__ tests. + if not hasattr(self, "_num_nans"): + self._num_nans = None + if not hasattr(self, "_num_nans_cpu"): + self._num_nans_cpu = None with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"): return super().get_output() @@ -294,7 +312,7 @@ class ExecuteModelState(NamedTuple): slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None -class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): +class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin): """Autoregressive GPU model runner that returns hidden states per request. Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and @@ -321,66 +339,24 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) self._async_chunk = getattr(self.model_config, "async_chunk", False) - # Worker-connector init is gated by a per-`model_arch` allowlist - # (covers both producer-side and consumer-side runners for the - # arches below). Consumer-wait stages must be registered - # separately as `(model_arch, model_stage)` tuples in - # `omni_scheduling_coordinator._FULL_PAYLOAD_INPUT_STAGES`; - # forgetting that produces a Stage-1 hang on the consumer. - _OMNI_CONNECTOR_INIT_ARCHS = { - "Qwen3OmniMoeForConditionalGeneration", - "Qwen2_5OmniForConditionalGeneration", - "CovoAudioForConditionalGeneration", - "MiMoAudioModel", - "Qwen3TTSTalkerForConditionalGeneration", - "Qwen3TTSCode2Wav", - "CosyVoice3Model", - "NemotronDenseForCausalLM", - "NemotronHForCausalLM", - "DyninOmniForConditionalGeneration", - "IndexTTS2TalkerForConditionalGeneration", - # nemotron_voicechat: the talker (stage 1) is the full-payload - # producer for code2wav (stage 2); the thinker (stage 0) only - # produces over the connector in async-chunk (streaming) mode. - "NemotronVoiceChatTalkerForConditionalGeneration", - "NemotronVoiceChatThinkerForConditionalGeneration", - } - # The stage-level ``model_arch`` override may be blank so the class - # resolves from the checkpoint's own ``architectures`` (e.g. the Audex - # TTA/TTS thinker stages, which bind dense on the 2B and NemotronH on - # the 30B-A3B). Gate on the RESOLVED architectures, not only the raw - # override — otherwise a blank-override producer stage never creates - # its worker connector and the sync full-payload flush silently never - # runs (downstream stage starves on connector input). - stage_archs = set(getattr(self.model_config, "architectures", None) or ()) - model_arch_override = getattr(self.model_config, "model_arch", None) - if model_arch_override: - stage_archs.add(model_arch_override) - if stage_archs & _OMNI_CONNECTOR_INIT_ARCHS or get_stage_connector_role(self.model_config) is not None: + # Connector init is derived from pipeline capabilities and active + # producer/KV/connector configuration. A stage needs it when it is: + # - a full-payload consumer, or + # - a producer (custom_process_next_stage_input_func is set, meaning + # it sends inter-stage data to a downstream consumer via connector), or + # - explicitly configured with a connector role via stage_connector_config. + if needs_omni_connector(self.model_config): self.init_omni_connectors( model_config=self.model_config, kv_transfer_manager=self.kv_transfer_manager, ) self._downstream_payload_cache: dict[str, bool] = {} - self._duplex_sampling_hook = None - self._duplex_sampling_hook_resolved = False + self._init_duplex_sampling_state() def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) self._resolve_duplex_sampling_hook(force=True) - def _resolve_duplex_sampling_hook(self, *, force: bool = False): - if not force and getattr(self, "_duplex_sampling_hook_resolved", False): - return self._duplex_sampling_hook - candidate = getattr(getattr(self, "model", None), "prepare_duplex_sampling", None) - self._duplex_sampling_hook = candidate if callable(candidate) else None - self._duplex_sampling_hook_resolved = True - if self._duplex_sampling_hook is not None and not hasattr(self, "_duplex_sampling_helper"): - from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingHelper - - self._duplex_sampling_helper = DuplexSamplingHelper() - return self._duplex_sampling_hook - def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size from vllm_omni.distributed.ray_utils.utils import ( @@ -449,11 +425,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): def _update_states(self, scheduler_output: SchedulerOutput) -> Callable | None: deferred_state_corrections_fn = super()._update_states(scheduler_output) - if self._resolve_duplex_sampling_hook() is None: - return deferred_state_corrections_fn - helper = getattr(self, "_duplex_sampling_helper", None) - if helper is not None: - helper.update_states(self, scheduler_output) + self._update_duplex_sampling_states(scheduler_output) return deferred_state_corrections_fn def _request_final_stage_id(self, req_id: str) -> int | None: @@ -473,65 +445,29 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): cached = self._downstream_payload_cache.get(req_id) if cached is not None: return cached - # Conservative default: keep payload if marker is missing. final_stage_id = self._request_final_stage_id(req_id) - needs_payload = final_stage_id is None or final_stage_id > 0 + if final_stage_id is None: + # Conservative default while the marker is missing: keep the + # payload, but do NOT memoize - the marker arrives via + # `model_intermediate_buffer`, which may be unpopulated on the + # first call (memoizing here pinned the request to True forever, + # never refreshing once the marker landed). + return True + needs_payload = final_stage_id > 0 self._downstream_payload_cache[req_id] = needs_payload return needs_payload def _resolve_pooler_payload_req_ids(self, req_ids_output_copy: list[str]) -> tuple[str, list[str]]: downstream_req_ids = [rid for rid in req_ids_output_copy if self._request_needs_downstream_stage_payload(rid)] engine_output_type = (self.vllm_config.model_config.engine_output_type or "").lower() + if self._client_multimodal_output_keys(): + downstream_req_ids = req_ids_output_copy # Single-stage AR TTS models (e.g. VoxCPM2) finish on this stage but still # need multimodal payloads for final audio postprocess/output. - if engine_output_type == "audio" and not downstream_req_ids: + elif engine_output_type == "audio" and not downstream_req_ids: downstream_req_ids = req_ids_output_copy return engine_output_type, downstream_req_ids - @staticmethod - def _sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None: - if not isinstance(multimodal_outputs, dict): - return None - meta = multimodal_outputs.get("meta") - req_ids = None - sparse_audio = False - if isinstance(meta, dict): - req_ids = meta.get("req_id") - sparse_audio = GPUARModelRunner._is_sparse_audio_marker(meta.get("sparse_audio")) - if req_ids is None: - req_ids = multimodal_outputs.get("meta.req_id") - sparse_audio = GPUARModelRunner._is_sparse_audio_marker(multimodal_outputs.get("meta.sparse_audio")) - if not sparse_audio: - return None - if not isinstance(req_ids, list): - return None - return [rid for rid in req_ids if isinstance(rid, str)] - - @staticmethod - def _resolve_sparse_mm_routing( - *, - engine_output_type: str, - req_ids_output_copy: list[str], - downstream_req_ids: list[str], - multimodal_outputs: Any, - ) -> tuple[list[str], dict[str, int], bool]: - sparse_mm_req_ids = GPUARModelRunner._sparse_mm_req_ids(multimodal_outputs) - sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])} - if engine_output_type != "audio" or sparse_mm_req_ids is None: - return downstream_req_ids, sparse_mm_index, False - - sparse_req_id_set = set(sparse_mm_req_ids) - sparse_downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set] - return sparse_downstream_req_ids, sparse_mm_index, True - - @staticmethod - def _is_sparse_audio_marker(value: Any) -> bool: - if isinstance(value, list): - return any(str(item).lower() in ("1", "true", "yes", "on") for item in value) - if isinstance(value, str): - return value.lower() in ("1", "true", "yes", "on") - return bool(value) - def capture_model(self) -> int: result = super().capture_model() self._capture_talker_mtp_graphs() @@ -581,9 +517,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self._downstream_payload_cache.clear() if hasattr(self, "model_intermediate_buffer"): self.model_intermediate_buffer.clear() - duplex_helper = getattr(self, "_duplex_sampling_helper", None) - if duplex_helper is not None: - duplex_helper.clear() + self._clear_duplex_sampling() # 5. Release all CUDA graphs unconditionally (upstream only does this # on ROCm; on CUDA the graphs are only freed by Python GC during @@ -880,79 +814,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): ) return hidden_states_cpu, combined_hidden_states, combined_multimodal_outputs - @staticmethod - def _build_combined_prefix_cache_mm_payload( - combined_multimodal_outputs: dict, - *, - rid: str, - idx: int, - ) -> dict[str, object]: - def _unwrap_lists(v): - if isinstance(v, list): - return v[idx] if idx < len(v) else v[0] - if isinstance(v, dict): - return {k: _unwrap_lists(sv) for k, sv in v.items()} - return v - - return { - mm_key: _unwrap_lists(combined_multimodal_outputs[mm_key][rid]) - for mm_key in combined_multimodal_outputs.keys() - } - - def _build_omni_mm_payload( - self, - *, - combined_multimodal_outputs: dict | None, - mm_cpu: dict[str, object] | None, - rid: str, - idx: int, - start: int, - end: int, - audio_sparse_output: bool, - sparse_mm_index: dict[str, int], - hidden_seq_len: int, - scheduled_seq_len: int, - ) -> dict[str, object]: - if combined_multimodal_outputs: - return self._build_combined_prefix_cache_mm_payload( - combined_multimodal_outputs, - rid=rid, - idx=idx, - ) - - mm_payload: dict[str, object] = {} - if not mm_cpu: - return mm_payload - - for mm_key, mm_val in mm_cpu.items(): - if mm_key in {"meta.req_id", "meta.sparse_audio"}: - continue - if audio_sparse_output and isinstance(mm_val, list): - sparse_idx = sparse_mm_index.get(rid) - if sparse_idx is None: - continue - if sparse_idx >= len(mm_val): - logger.warning( - "Sparse multimodal payload mismatch for request %s: index %d >= %d.", - rid, - sparse_idx, - len(mm_val), - ) - continue - sparse_val = mm_val[sparse_idx] - mm_payload[mm_key] = sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val - continue - mm_payload[mm_key] = to_payload_element( - element=mm_val, - idx=idx, - start=start, - end=end, - pass_lists_through=False, - seq_len=hidden_seq_len, - scheduled_seq_len=scheduled_seq_len, - ) - return mm_payload - def _build_omni_pooler_payload( self, *, @@ -985,7 +846,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if req_hidden_states is not None: payload["hidden"] = req_hidden_states - mm_payload = self._build_omni_mm_payload( + mm_payload = build_omni_mm_payload( combined_multimodal_outputs=combined_multimodal_outputs, mm_cpu=mm_cpu, rid=rid, @@ -1000,6 +861,47 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): payload.update(mm_payload) return payload + def _merge_model_kv_transfer_metadata( + self, + finished_reqs: dict[str, dict[str, Any]], + ) -> dict[str, dict[str, Any]]: + """Merge model-provided KV-transfer metadata copy-on-write. + + The outer mapping is always new; a per-request dict is shallow-copied + (with a fresh merged ``custom_metadata``) only when the model supplies + metadata for it, and passes through unchanged otherwise. + `scheduler_output` (and its per-request dicts) can be shared with the + engine-core process, so nothing here mutates them — the ngram block + in `execute_model` `replace()`-copies for exactly the same reason. The + merged mapping is only consumed by + `kv_transfer_manager.handle_finished_requests_kv_transfer`. + """ + merged: dict[str, dict[str, Any]] = {} + for req_id, data in finished_reqs.items(): + try: + # NOTE: seq_len is the same as num_computed_tokens_cpu in current + # async scheduling, since both exclude async placeholders. We use + # seq_len since we control it, just in case upstream async scheduler + # semantics change in the future. + model_meta = self.model.get_kv_transfer_metadata( + req_id, + num_computed_tokens=data.get("seq_len"), + ) + except Exception: + # A swallowed failure here used to ship the KV transfer WITHOUT + # the model's custom metadata and surface much later as a + # corrupt/incomplete transfer on the consumer stage. Fail at + # the point of the defect instead. + logger.exception("Failed to get custom KV-transfer metadata from model for %s", req_id) + raise + if model_meta: + data = { + **data, + "custom_metadata": {**(data.get("custom_metadata") or {}), **model_meta}, + } + merged[req_id] = data + return merged + @torch.inference_mode() def execute_model( self, @@ -1030,24 +932,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests) finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"): - for req_id, data in finished_reqs.items(): - try: - # NOTE: seq_len is the same as num_computed_tokens_cpu in current - # async scheduling, since both exclude async placeholders. We use - # seq_len since we control it, just in case upstream async scheduler - # semantics change in the future. - num_computed = data.get("seq_len") - - model_meta = self.model.get_kv_transfer_metadata( - req_id, - num_computed_tokens=num_computed, - ) - if model_meta: - existing = data.get("custom_metadata") or {} - existing.update(model_meta) - data["custom_metadata"] = existing - except Exception as e: - logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}") + finished_reqs = self._merge_model_kv_transfer_metadata(finished_reqs) self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer( finished_reqs=finished_reqs, kv_caches=self.kv_caches, @@ -1161,9 +1046,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens - logits_indices, spec_decode_metadata = self._prepare_inputs( - scheduler_output, - num_scheduled_tokens_np, + logits_indices, spec_decode_metadata, max_num_sampled_tokens = self._prepare_inputs( + scheduler_output, num_scheduled_tokens_np ) cascade_attn_prefix_lens = None @@ -1257,6 +1141,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): max_query_len=max_num_scheduled_tokens, ubatch_slices=ubatch_slices_attn, logits_indices=logits_indices, + max_num_sampled_tokens=max_num_sampled_tokens, use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, @@ -1296,15 +1181,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): input_ids_buffer=self.input_ids.gpu[:num_tokens_padded], ) - # Set cudagraph mode to none if calc_kv_scales is true. - # KV scales calculation involves dynamic operations that are incompatible - # with CUDA graph capture. - if self.calculate_kv_scales: - cudagraph_mode = CUDAGraphMode.NONE - runner_assisted_full_attn = False - # Mark KV scales as calculated after the first forward pass - self.calculate_kv_scales = False - runner_assisted_context_enabled = False if runner_assisted_full_attn: runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context( @@ -1352,6 +1228,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None) if callable(flush_pending_metadata): flush_pending_metadata(req_ids[:num_reqs]) + + # [Omni] Hand the model the batch's req_ids in logits order, for + # models that gate logits per request. Only valid without spec + # decode: there logits_indices carries several rows per request, + # so row i no longer corresponds to req_ids[i]. + if spec_decode_metadata is None: + set_batch_req_ids = getattr(self.model, "set_batch_req_ids", None) + if callable(set_batch_req_ids): + set_batch_req_ids(req_ids[:num_reqs]) finally: if runner_assisted_context_enabled: self._set_runner_assisted_full_attention_metadata_context(enabled=False) @@ -1485,6 +1370,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): logits: torch.Tensor | None, spec_decode_metadata: Any, ): + """Sample tokens, preferring the model's own sampler when declared. + + Model-interface contract for `prefer_model_sampler` models: a custom + `model.sample(logits, sampling_metadata)` may return ``None`` to + DECLINE the step, and the runner then falls back to the default + sampler. Declarers rely on this deliberately for empty-logits or + inapplicable-stage steps (e.g. cosyvoice3, glm_tts), so the + fallthrough is load-bearing, not an error path; it is logged once per + model for visibility. + """ sampling_metadata = self.input_batch.sampling_metadata if spec_decode_metadata is None: model_sample = getattr(self.model, "sample", None) @@ -1501,17 +1396,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.input_batch.positions[self.input_batch.logits_indices], ) prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata) - prepare_duplex_sampling = self._resolve_duplex_sampling_hook() - if prepare_duplex_sampling is not None: - helper = getattr(self, "_duplex_sampling_helper", None) - rows = helper.rows(self) if helper is not None and helper.active_request_ids else () - if rows or (helper is not None and helper.hook_active): - prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) - if helper is not None: - helper.hook_active = bool(rows) + self._apply_duplex_sampling(logits, prepared_sampling_metadata) sampler_output = model_sample(logits, prepared_sampling_metadata) if sampler_output is not None: return sampler_output + # Contract: None => fall back to the default sampler (see + # docstring). If a custom sampler returns None by accident, + # its tokens silently come from the default sampler - the + # one-time log is the visibility hook for that case. + logger.warning_once( + "prefer_model_sampler model %s returned None from sample(); falling back to the default sampler.", + type(self.model).__name__, + ) return self.sampler( logits=logits, sampling_metadata=sampling_metadata, @@ -1668,7 +1564,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): lists are wrapped into tensors, and anything that cannot be safely converted is dropped. """ - if self.vllm_config.model_config.engine_output_type == "text": + if self.vllm_config.model_config.engine_output_type == "text" and not self._client_multimodal_output_keys(): return None if per_req_payloads is None: return None @@ -1727,6 +1623,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): def _runner_model_omni_flag(self, name: str, default: bool = False) -> bool: return self._model_omni_flag(getattr(self, "model", None), name, default) + def _client_multimodal_output_keys(self) -> tuple[str, ...]: + raw = getattr( + getattr(self, "model", None), + "omni_client_multimodal_output_keys", + (), + ) + if not isinstance(raw, tuple) or any(not isinstance(key, str) or not key for key in raw): + raise TypeError("omni_client_multimodal_output_keys must be a tuple of non-empty strings") + if len(raw) != len(set(raw)): + raise ValueError("omni_client_multimodal_output_keys must not contain duplicates") + return raw + def _model_omni_pooler_payload_include_hidden(self) -> bool: return self._runner_model_omni_flag("omni_pooler_payload_include_hidden", default=True) @@ -1914,7 +1822,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): combined_multimodal_outputs = None engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) - downstream_req_ids, sparse_mm_index, audio_sparse_output = self._resolve_sparse_mm_routing( + downstream_req_ids, sparse_mm_index, audio_sparse_output = resolve_sparse_mm_routing( engine_output_type=engine_output_type, req_ids_output_copy=req_ids_output_copy, downstream_req_ids=downstream_req_ids, @@ -2034,6 +1942,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # Connector-less stages expose the same payload through the # orchestrator bridge; non-async stages preserve legacy behavior. pooler_inter, pooler_client = pooler_output, pooler_output + client_output_keys = self._client_multimodal_output_keys() + if client_output_keys: + allowed = frozenset(client_output_keys) + pooler_client = [ + {key: value for key, value in payload.items() if key in allowed} for payload in pooler_output + ] if pooler_inter and self._should_accumulate_full_payload_output(): with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"): @@ -2114,13 +2028,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): if grammar_output is not None: apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits) - # Correct padding values of prompt_token_ids to match the logits vocabulary size + # Correct padding values of prompt_token_ids to match the logits vocabulary size. if logits is not None and not self.input_batch.sampling_metadata.no_penalties: smd = self.input_batch.sampling_metadata if smd.prompt_token_ids is not None: logits_vocab = logits.shape[-1] if self.input_batch.vocab_size > logits_vocab: - smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab) + smd.prompt_token_ids = clamp_prompt_ids_to_penalty_padding(smd.prompt_token_ids, logits_vocab) # Drop min-tokens stop ids the head cannot emit (e.g. the text # tokenizer EOS folded into all_stop_token_ids on a narrow codec @@ -2194,6 +2108,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): with record_function_or_nullcontext("gpu_model_runner: bookkeep"): ( num_nans_in_logits, + num_nans, logprobs_lists, valid_sampled_token_ids, prompt_logprobs_dict, @@ -2311,6 +2226,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, vocab_size=self.input_batch.vocab_size, + num_nans=num_nans, ) if use_async_omni_output: async_output = async_output_cls( diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 2afa59576..654cd4dc5 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -28,6 +28,8 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): model runners for text generation stages (e.g., thinker stages). """ + model_runner_cls = GPUARModelRunner + @instrument(span_name="Init device") def init_device(self): if self.device_config.device_type in ("cuda", "musa"): @@ -117,7 +119,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): self.use_v2_model_runner = False # Construct the model runner - self.model_runner = GPUARModelRunner(self.vllm_config, self.device) + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) if self.rank == 0: # If usage stat is enabled, collect relevant info. diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 7cf26d005..307a5a428 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Code2Wav GPU Model Runner for vLLM-Omni. Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav. @@ -35,12 +38,14 @@ from vllm.v1.worker.gpu_model_runner import ( from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs -from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_role from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import partition_payload_list from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner -from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin +from vllm_omni.worker.omni_connector_model_runner_mixin import ( + OmniConnectorModelRunnerMixin, + needs_omni_connector, +) logger = logging.getLogger(__name__) @@ -56,27 +61,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._async_chunk = getattr(self.model_config, "async_chunk", False) - # Mirrors the init allowlist in gpu_ar_model_runner.py. - _OMNI_CONNECTOR_INIT_ARCHS = { - "Qwen3OmniMoeForConditionalGeneration", - "Qwen2_5OmniForConditionalGeneration", - "CovoAudioForConditionalGeneration", - "MiMoAudioModel", - "Qwen3TTSTalkerForConditionalGeneration", - "Qwen3TTSCode2Wav", - "AudexCode2Wav", - "AudexXCodec1", - "CosyVoice3Model", - "DyninOmniForConditionalGeneration", - "IndexTTS2S2MelDecoder", - # nemotron_voicechat: code2wav (stage 2) consumes the talker's - # full-payload code stacks. - "NemotronVoiceChatCode2Wav", - } - if ( - getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS - or get_stage_connector_role(self.model_config) is not None - ): + if needs_omni_connector(self.model_config): self.init_omni_connectors( model_config=self.model_config, ) @@ -167,10 +152,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): self.model.on_requests_finished(scheduler_output.finished_req_ids) - # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). - if scheduler_output.total_num_scheduled_tokens <= 0: - return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) - if has_ec_transfer() and not get_ec_transfer().is_consumer: with self.maybe_get_ec_connector_output( scheduler_output, @@ -179,7 +160,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self._execute_mm_encoder(scheduler_output) return self.attach_omni_connector_output(make_empty_encoder_model_runner_output(scheduler_output)) - if not num_scheduled_tokens: + # OMNI: keep this block in lock-step with the same block in + # GPUARModelRunner.execute_model. + # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). + if num_scheduled_tokens <= 0: if ( self.parallel_config.distributed_executor_backend == "external_launcher" and self.parallel_config.data_parallel_size > 1 @@ -211,9 +195,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens - logits_indices, spec_decode_metadata = self._prepare_inputs( - scheduler_output, - num_scheduled_tokens_np, + logits_indices, spec_decode_metadata, max_num_sampled_tokens = self._prepare_inputs( + scheduler_output, num_scheduled_tokens_np ) cascade_attn_prefix_lens = None @@ -297,6 +280,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin max_query_len=max_num_scheduled_tokens, ubatch_slices=ubatch_slices_attn, logits_indices=logits_indices, + max_num_sampled_tokens=max_num_sampled_tokens, use_spec_decode=use_spec_decode, num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, @@ -328,14 +312,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if getattr(self.model, "requires_request_ids", False): model_kwargs["request_ids"] = list(req_ids) - # Set cudagraph mode to none if calc_kv_scales is true. - # KV scales calculation involves dynamic operations that are incompatible - # with CUDA graph capture. - if self.calculate_kv_scales: - cudagraph_mode = CUDAGraphMode.NONE - # Mark KV scales as calculated after the first forward pass - self.calculate_kv_scales = False - # Run the model. # Use persistent buffers for CUDA graphs. # When spec decode is enabled, defer connector finalization @@ -439,16 +415,29 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # pooler_output is no longer used for multimodal data. per_req_payloads: list[dict[str, object]] = [] if isinstance(multimodal_outputs_raw, torch.Tensor): - assert multimodal_outputs_raw.shape[0] == 1, ( - "model should return a single tensor, to return multiple tensors, use a dict" - ) - assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs - for i in range(self.input_batch.num_reqs): + # One row per request. The old asserts (`shape[0] == 1` AND + # `shape[0] == num_reqs`) jointly forced num_reqs == 1, silently + # rejecting batched steps the generation scheduler does admit. + num_reqs = self.input_batch.num_reqs + if multimodal_outputs_raw.shape[0] != num_reqs: + raise ValueError( + f"Multimodal output tensor has leading dim {multimodal_outputs_raw.shape[0]} " + f"but the batch has {num_reqs} requests (one row per request; " + "to return multiple tensors per request, use a dict)." + ) + for i in range(num_reqs): per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) elif isinstance(multimodal_outputs_raw, list): - assert len(multimodal_outputs_raw) == 1, ( - "model should return a single list, to return multiple lists, use a dict" - ) + # One entry per request. The old `len == 1` assert did not check + # num_reqs, so a batched step built a length-1 payload list that + # misaligned with `req_ids` downstream. + num_reqs = self.input_batch.num_reqs + if len(multimodal_outputs_raw) != num_reqs: + raise ValueError( + f"Multimodal output list has length {len(multimodal_outputs_raw)} " + f"but the batch has {num_reqs} requests (one entry per request; " + "to return multiple lists per request, use a dict)." + ) for out in multimodal_outputs_raw: per_req_payloads.append( {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} @@ -794,6 +783,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] model_kwargs = self._init_model_kwargs() + elif getattr(getattr(self, "model", None), "has_preprocess", False): + # Capture CUDA graph with inputs_embeds path so replay reads + # from the same buffer that _preprocess writes into. + input_ids = self.input_ids.gpu[:num_tokens_padded] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 38a88c84c..d6a8d6c8e 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -23,10 +23,12 @@ logger = init_logger(__name__) class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): """GPU Worker for Generation model (non-autoregressive waveform generation). - Usage in stage config: - worker_cls: "vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner" + Selected for stages whose pipeline topology uses + ``execution_type=StageExecutionType.LLM_GENERATION``. """ + model_runner_cls = GPUGenerationModelRunner + @instrument(span_name="Init device") def init_device(self): if self.device_config.device_type in ("cuda", "musa"): @@ -99,7 +101,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.") self.use_v2_model_runner = False - self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device) + self.model_runner = self.model_runner_cls(self.vllm_config, self.device) if self.rank == 0: # If usage stat is enabled, collect relevant info. diff --git a/vllm_omni/worker/gpu_memory_utils.py b/vllm_omni/worker/gpu_memory_utils.py index d005c5e58..1a8072c83 100644 --- a/vllm_omni/worker/gpu_memory_utils.py +++ b/vllm_omni/worker/gpu_memory_utils.py @@ -96,7 +96,7 @@ def get_process_gpu_memory(local_rank: int) -> int | None: except Exception as e: raise RuntimeError( f"Failed to get NVML handle for device '{device_id}' (local_rank={local_rank}). " - f"Check CUDA_VISIBLE_DEVICES or stage config 'devices' setting." + f"Check CUDA_VISIBLE_DEVICES or the deploy config `devices` setting." ) from e else: # No CUDA_VISIBLE_DEVICES or local_rank out of range: use index directly @@ -104,7 +104,7 @@ def get_process_gpu_memory(local_rank: int) -> int | None: if local_rank >= device_count: raise RuntimeError( f"Invalid GPU device {local_rank}. Only {device_count} GPU(s) available. " - f"Check CUDA_VISIBLE_DEVICES or stage config 'devices' setting." + f"Check CUDA_VISIBLE_DEVICES or the deploy config `devices` setting." ) handle = nvmlDeviceGetHandleByIndex(local_rank) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 2edea7fff..eca15c16e 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import contextlib import inspect from collections.abc import Callable @@ -1494,14 +1497,18 @@ class OmniGPUModelRunner(GPUModelRunner): self.inputs_embeds.gpu[start_offset : start_offset + overlay_len].copy_(src) def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None: + replace = getattr(self.model, "replace_runtime_additional_information", False) + update_buffer = self._replace_intermediate_buffer if replace else self._update_intermediate_buffer for new_req in scheduler_output.scheduled_new_reqs: model_buffer = getattr(new_req, "model_intermediate_buffer", None) if isinstance(model_buffer, dict) and model_buffer: - self._update_intermediate_buffer(new_req.req_id, model_buffer) + update_buffer(new_req.req_id, model_buffer) + if replace: + continue payload_info = getattr(new_req, "additional_information", None) decoded_info = deserialize_additional_information(payload_info) if decoded_info: - self._update_intermediate_buffer(new_req.req_id, decoded_info) + update_buffer(new_req.req_id, decoded_info) if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"): cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {}) @@ -1509,7 +1516,7 @@ class OmniGPUModelRunner(GPUModelRunner): for req_id, req_infos in cached_infos.items(): decoded_info = deserialize_additional_information(req_infos) if decoded_info: - self._update_intermediate_buffer(req_id, decoded_info) + update_buffer(req_id, decoded_info) def _maybe_attach_mimo_audio_req_infos( self, @@ -1772,6 +1779,16 @@ class OmniGPUModelRunner(GPUModelRunner): req_infos["_omni_prompt_len"] = prompt_len req_infos["_omni_num_computed_tokens"] = num_computed_tokens req_infos["_omni_is_prefill"] = is_prefill + # Output-token cap, so a model that must ship a payload on the + # request's final step can tell which step that is. A finished + # request drops out of req_ids_output_copy, and downstream + # routing keys off that list, so a payload produced after its + # last step never reaches the caller. + sampling_params = getattr(req_state, "sampling_params", None) + req_infos["_omni_max_tokens"] = getattr(sampling_params, "max_tokens", None) + # Seed, so a model that samples inside forward() can be + # reproducible: vLLM's own sampler seeding does not reach it. + req_infos["_omni_seed"] = getattr(sampling_params, "seed", None) if callable(batch_decode_preprocess) and self.has_talker_mtp and span_len == 1 and not is_prefill: decode_batch_items.append((req_id, s, req_infos)) continue @@ -1974,6 +1991,7 @@ class OmniGPUModelRunner(GPUModelRunner): positions=positions, inputs_embeds=inputs_embeds, omni_query_start_loc=model_kwargs_extra.get("omni_query_start_loc"), + req_ids=self.input_batch.req_ids, ) model_output = super()._model_forward( @@ -2029,6 +2047,27 @@ class OmniGPUModelRunner(GPUModelRunner): # Backward compatible: mirror to old setattr location setattr(req_state, "additional_information_cpu", existing) + def _replace_intermediate_buffer(self, req_id: str, upd: dict) -> None: + """Replace one request's runtime payload with the current chunk.""" + if not isinstance(upd, dict): + return + req_state = self.requests.get(req_id) + if req_state is None: + return + snapshot = dict(upd) + previous = self.model_intermediate_buffer.get(req_id) + previous_meta = previous.get("meta") if isinstance(previous, dict) else None + incoming_meta = snapshot.get("meta") + meta = dict(incoming_meta) if isinstance(incoming_meta, dict) else {} + if isinstance(previous_meta, dict): + for key in ("num_processed_tokens", "resumable"): + if key not in meta and key in previous_meta: + meta[key] = previous_meta[key] + if meta or "meta" in snapshot: + snapshot["meta"] = meta + self.model_intermediate_buffer[req_id] = snapshot + setattr(req_state, "additional_information_cpu", snapshot) + def _update_streaming_input_additional_info(self, new_req_data, req_id): # For streaming input prefill case only. Update buffer from last segment input. cached_additional_info = self.model_intermediate_buffer.get(req_id, {}) @@ -2036,6 +2075,15 @@ class OmniGPUModelRunner(GPUModelRunner): if not isinstance(inc_info, dict) or not inc_info: payload_info = getattr(new_req_data, "additional_information", None) inc_info = deserialize_additional_information(payload_info) + if getattr(self.model, "replace_runtime_additional_information", False): + snapshot = dict(inc_info) if isinstance(inc_info, dict) else {} + meta = snapshot.get("meta") + meta = dict(meta) if isinstance(meta, dict) else {} + meta["num_processed_tokens"] = 0 + meta["resumable"] = True + snapshot["meta"] = meta + self._replace_intermediate_buffer(req_id, snapshot) + return if not isinstance(inc_info, dict) or not inc_info: return merged_info = dict(cached_additional_info) if isinstance(cached_additional_info, dict) else {} diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 07a0327a6..1c8fe5b9f 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -49,6 +49,15 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def needs_omni_connector(model_config: Any) -> bool: + """Whether a runner owns an input, output, or explicitly routed connector.""" + return ( + bool(getattr(model_config, "requires_full_payload_input", False)) + or bool(getattr(model_config, "custom_process_next_stage_input_func", None)) + or get_stage_connector_role(model_config) is not None + ) + + def _should_create_payload_connector(model_config: Any) -> bool: """Whether this stage owns runner payload transport for its edge. @@ -170,6 +179,7 @@ class OmniConnectorModelRunnerMixin: # does not have segment boundary infrastructure; multi-segment support # is only available via chunk_transfer_adapter (distributed path). self._ramp_chunk_count: dict[str, int] = defaultdict(int) + self._adaptive_states: dict[str, Any] = {} # Send-side async accumulation / staging buffer. Receive-side payload # ownership lives in ``_local_stage_payload_cache``. self._send_side_request_payload: dict[str, dict[str, Any]] = {} @@ -320,6 +330,7 @@ class OmniConnectorModelRunnerMixin: self._code_prompt_token_ids.pop(k, None) self._cached_ic.pop(k, None) self._ramp_chunk_count.pop(k, None) + self._adaptive_states.pop(k, None) self._kv_pending_transfers.pop(req_id, None) self._kv_active_transfers.discard(req_id) self._kv_completed_transfers.discard(req_id) @@ -406,11 +417,13 @@ class OmniConnectorModelRunnerMixin: # thread by one execute_model() cycle, especially when the request was # added after the current scheduler_output snapshot. # - # Orphaned pending recv entries (e.g. from upstream stage crash) - # are handled by OmniSchedulingCoordinator.collect_timed_out_request_ids() - # which detects wait-time violations. The scheduler then removes the - # request from its queues, sets FINISHED_ERROR, and calls _free_request() - # which ultimately triggers cleanup_finished_request() here. + # Orphaned pending recv entries (e.g. from upstream stage crash) are + # handled by collect_timed_out_request_ids() -- on + # OmniSchedulingCoordinator for full-payload requests, and on + # OmniChunkTransferAdapter for async-chunk ones -- which detect + # wait-time violations. The scheduler then removes the request from + # its queues, sets FINISHED_ERROR, and calls _free_request() which + # ultimately triggers cleanup_finished_request() here. for attr_name in ( "_request_ids_mapping", "_get_req_chunk", @@ -2016,6 +2029,7 @@ class OmniConnectorModelRunnerMixin: self._code_prompt_token_ids.pop(cleanup_req_id, None) self._cached_ic.pop(cleanup_req_id, None) self._ramp_chunk_count.pop(cleanup_req_id, None) + self._adaptive_states.pop(cleanup_req_id, None) # ------------------------------------------------------------------ # # Payload accumulation (ported from OmniChunkTransferAdapter) diff --git a/vllm_omni/worker/omni_connector_validation.py b/vllm_omni/worker/omni_connector_validation.py new file mode 100644 index 000000000..867a6d95f --- /dev/null +++ b/vllm_omni/worker/omni_connector_validation.py @@ -0,0 +1,42 @@ +"""Validation for Omni connector support on selected LLM workers.""" + +import importlib +from typing import Any + + +def validate_worker_omni_connector( + worker_cls_path: str | type[Any] | None, + required: bool, +) -> None: + """Require the selected worker's runner to implement the connector mixin.""" + if not required: + return + if worker_cls_path is None: + raise ValueError("Omni connector support requires a resolved worker_cls.") + + worker_cls = _resolve_worker_cls(worker_cls_path) + model_runner_cls = getattr(worker_cls, "model_runner_cls", None) + + from vllm_omni.worker.omni_connector_model_runner_mixin import ( + OmniConnectorModelRunnerMixin, + ) + + if not isinstance(model_runner_cls, type) or not issubclass( + model_runner_cls, + OmniConnectorModelRunnerMixin, + ): + worker_name = getattr(worker_cls, "__name__", str(worker_cls_path)) + raise ValueError(f"Worker {worker_name!r} does not provide an Omni connector model runner.") + + +def _resolve_worker_cls(worker_cls_path: str | type[Any]) -> type[Any]: + if not isinstance(worker_cls_path, str): + return worker_cls_path + module_path, separator, class_name = worker_cls_path.rpartition(".") + if not separator: + raise ValueError(f"worker_cls must be a fully qualified class path, got {worker_cls_path!r}.") + try: + worker_module = importlib.import_module(module_path) + return getattr(worker_module, class_name) + except (ImportError, AttributeError) as exc: + raise ValueError(f"Unable to resolve worker_cls {worker_cls_path!r} for connector validation.") from exc diff --git a/vllm_omni/worker/output/__init__.py b/vllm_omni/worker/output/__init__.py new file mode 100644 index 000000000..33ea418a3 --- /dev/null +++ b/vllm_omni/worker/output/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worker output helpers. + +Package seeded ahead of the async-output refactor (RFC #5450 series, G5-C2), +which relocates the remaining payload-copy helpers, snapshot types, and +`ExecuteModelState` here. Currently hosts the per-request multimodal payload +builders. +""" + +from vllm_omni.worker.output.payload_build import ( + build_combined_prefix_cache_mm_payload, + build_omni_mm_payload, + unwrap_combined_payload_value, +) + +__all__ = [ + "build_combined_prefix_cache_mm_payload", + "build_omni_mm_payload", + "unwrap_combined_payload_value", +] diff --git a/vllm_omni/worker/output/payload_build.py b/vllm_omni/worker/output/payload_build.py new file mode 100644 index 000000000..0f65ac293 --- /dev/null +++ b/vllm_omni/worker/output/payload_build.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-request multimodal payload builders. + +State-free extraction helpers that turn a step's `multimodal_outputs` +(combined prefix-cache-merged form, or the direct `mm_cpu` form) into one +request's payload dict. Sparse-routing context (`audio_sparse_output`, +`sparse_mm_index`) is resolved upstream by +`vllm_omni.worker.sparse_audio.resolve_sparse_mm_routing` and arrives here +as plain arguments. +""" + +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.utils.mm_outputs import to_payload_element + +logger = init_logger(__name__) + +_SPARSE_AUDIO_PROTOCOL_KEYS = frozenset({"meta.req_id", "meta.sparse_audio"}) + + +def unwrap_combined_payload_value( + value: Any, + *, + rid: str, + list_idx: int, + mm_key: str, +) -> tuple[bool, Any]: + """Unwrap one value from the prefix-cache-merged payload. + + Returns ``(keep, unwrapped)``: ``keep`` is False when a non-singleton + per-request list is misaligned and must be dropped. A singleton list is + request-invariant passthrough data and is shared across the batch. + """ + if isinstance(value, list): + if len(value) == 1: + # Prefix-cache passthrough intentionally preserves list + # containers and copies the whole list under every request id. + # Qwen3-Omni uses this shape for its request-invariant TTS + # BOS/EOS/PAD embeddings, so every batch index unwraps the same + # singleton. Sparse one-request lists are safe here too: routing + # limits them to the declared request at sparse index zero. + return True, value[0] + if list_idx < len(value): + return True, value[list_idx] + # A non-singleton list shorter than the request's index is per-request + # data that lost alignment. Never substitute another request's + # element (`v[0]` used to ship request 0's payload as `rid`'s); drop + # the key loudly instead — the SAME protocol-break policy as the + # sparse mismatch in + # `build_omni_mm_payload`. Not a raise: v1 runners do not catch + # per-request exceptions, so raising would turn a per-request + # protocol break into a whole-stage outage. + logger.error( + "Combined prefix-cache mm payload misaligned for request %s: index %d out of " + "range for per-request list of length %d; dropping key %s.", + rid, + list_idx, + len(value), + mm_key, + ) + return False, None + if isinstance(value, dict): + nested = {} + for k, sv in value.items(): + keep, unwrapped = unwrap_combined_payload_value(sv, rid=rid, list_idx=list_idx, mm_key=mm_key) + if keep: + nested[k] = unwrapped + return True, nested + return True, value + + +def build_combined_prefix_cache_mm_payload( + combined_multimodal_outputs: dict, + *, + rid: str, + list_idx: int, +) -> dict[str, object]: + """Unwrap one request's entry from the prefix-cache-merged payload. + + Sparse-agnostic by design: ``list_idx`` arrives pre-resolved by the + caller (`build_omni_mm_payload`, which owns the sparse-routing context) + — batch index normally, sparse index under sparse routing. + """ + payload: dict[str, object] = {} + for mm_key in combined_multimodal_outputs.keys(): + if mm_key in _SPARSE_AUDIO_PROTOCOL_KEYS: + continue + keep, unwrapped = unwrap_combined_payload_value( + combined_multimodal_outputs[mm_key][rid], + rid=rid, + list_idx=list_idx, + mm_key=mm_key, + ) + if keep: + payload[mm_key] = unwrapped + return payload + + +def build_omni_mm_payload( + *, + combined_multimodal_outputs: dict | None, + mm_cpu: dict[str, object] | None, + rid: str, + idx: int, + start: int, + end: int, + audio_sparse_output: bool, + sparse_mm_index: dict[str, int], + hidden_seq_len: int, + scheduled_seq_len: int, +) -> dict[str, object]: + """Build one request's multimodal payload for the step.""" + if combined_multimodal_outputs: + # Sparse-audio producers emit per-request lists aligned to the + # sparse request order (`meta.req_id`), not the batch order; under + # sparse routing the request's sparse index selects its element. + list_idx = sparse_mm_index.get(rid, idx) if audio_sparse_output else idx + return build_combined_prefix_cache_mm_payload( + combined_multimodal_outputs, + rid=rid, + list_idx=list_idx, + ) + + mm_payload: dict[str, object] = {} + if not mm_cpu: + return mm_payload + + for mm_key, mm_val in mm_cpu.items(): + if mm_key in _SPARSE_AUDIO_PROTOCOL_KEYS: + continue + if audio_sparse_output and isinstance(mm_val, list): + sparse_idx = sparse_mm_index.get(rid) + if sparse_idx is None: + continue + if sparse_idx >= len(mm_val): + # Misalignment means the sparse-marker protocol broke; the + # request's output for this key is dropped, so be loud. + logger.error( + "Sparse multimodal payload mismatch for request %s: index %d >= %d; dropping key %s.", + rid, + sparse_idx, + len(mm_val), + mm_key, + ) + continue + sparse_val = mm_val[sparse_idx] + mm_payload[mm_key] = sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val + continue + mm_payload[mm_key] = to_payload_element( + element=mm_val, + idx=idx, + start=start, + end=end, + pass_lists_through=False, + seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, + ) + return mm_payload diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py index c38a9bbdc..845881138 100644 --- a/vllm_omni/worker/sampling_utils.py +++ b/vllm_omni/worker/sampling_utils.py @@ -8,7 +8,21 @@ from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsPro logger = init_logger(__name__) -__all__ = ["sanitize_min_tokens_stop_ids"] +__all__ = ["clamp_prompt_ids_to_penalty_padding", "sanitize_min_tokens_stop_ids"] + + +def clamp_prompt_ids_to_penalty_padding(prompt_token_ids: torch.Tensor, logits_vocab: int) -> torch.Tensor: + """Clamp batch-level pad ids down to ``logits_vocab`` — upstream's + designed penalty padding value. + + ``max=logits_vocab`` (NOT ``logits_vocab - 1``) is deliberate: upstream + penalty computation allocates ``vocab_size + 1`` bins and drops the last + column, so ``vocab_size`` is the padding value that never affects + penalties (vllm/model_executor/layers/utils.py:: + get_token_bin_counts_and_mask). Clamping one lower would count padding + as real occurrences of the last vocab token. + """ + return prompt_token_ids.clamp(max=logits_vocab) def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: int) -> None: diff --git a/vllm_omni/worker/sparse_audio.py b/vllm_omni/worker/sparse_audio.py new file mode 100644 index 000000000..f82c7a59e --- /dev/null +++ b/vllm_omni/worker/sparse_audio.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse-audio output protocol: marker classification and payload routing. + +A step's per-request `multimodal_outputs` lists may cover only a SUBSET of +the batch (requests that produced audio this step). The producer declares +this with two metadata keys (nested ``meta`` dict or flattened ``meta.*``): + +* ``meta.req_id`` — the request ids that DO have output, in list order; +* ``meta.sparse_audio`` — the marker saying "per-request lists are aligned + to ``meta.req_id`` order, not batch order". + +The marker's carrier form is ENFORCED, not normalized. The design contract +is a list of strings (e.g. ``["1"]`` — what every in-tree producer emits); a +bare string is also accepted. Interpreting deviant carriers would need a +second truth predicate, and predicate divergence is exactly how +``np.array(["0"])`` came to route as sparse while ``["0"]`` routed as dense; +per the refactor series' Policy A the runner<->model interface owes no +tolerance shim. + +A present-but-invalid declaration (illegal marker carrier, or marker set +with an unusable ``meta.req_id``) FAILS CLOSED: it most likely belongs to a +producer that INTENDED sparse output, and falling back to dense would +batch-index-assign its payloads to the WRONG requests. `InvalidSparseDeclarationError` +is raised by the classifiers and caught INSIDE this module by +`resolve_sparse_mm_routing`, which routes no payload for the affected step — +the exception never crosses the module boundary, so it can never escape a +v1 runner (which does not catch per-request exceptions). + +This module is the single home of the contract until the model-state +refactor series (G9-C2.6) turns the marker into a typed bool resolved at +one adapter site and unifies the nested/flattened encodings. The NPU runner +still carries a tolerant pre-enforcement copy — NPU-track handoff item. +""" + +import threading +import time +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# The DESIGNED marker carrier is a list of strings (bare str also accepted); +# truthiness of a string item: +_SPARSE_MARKER_TRUTHY = ("1", "true", "yes", "on") + +# Per-process dedup so a broken producer logs once per offending shape, not +# once per step: +_logged_offenders: set[str] = set() +_error_state_lock = threading.Lock() + +# A persistent invalid producer must remain observable without flooding the +# per-step output path. Classifier errors above stay deduplicated per carrier +# type; the routing consequence is reported at most once per interval, with a +# count of repeats suppressed since the previous report. +_FAIL_CLOSED_LOG_INTERVAL_S = 60.0 +_fail_closed_last_log_at: float | None = None +_fail_closed_suppressed = 0 + + +class InvalidSparseDeclarationError(ValueError): + """Sparse-audio metadata is PRESENT but out of contract. + + Internal control signal: raised by `is_sparse_audio_marker` / + `resolve_sparse_mm_req_ids` and caught by `resolve_sparse_mm_routing` + in this same module, which fails closed. Callers outside this module + only ever see the routing result, never the exception. + """ + + +def _error_once(key: str, msg: str, *args: Any) -> None: + with _error_state_lock: + if key in _logged_offenders: + return + _logged_offenders.add(key) + logger.error(msg, *args) + + +def _error_fail_closed_rate_limited(cause: str) -> None: + global _fail_closed_last_log_at, _fail_closed_suppressed + + now = time.monotonic() + with _error_state_lock: + if _fail_closed_last_log_at is not None and now - _fail_closed_last_log_at < _FAIL_CLOSED_LOG_INTERVAL_S: + _fail_closed_suppressed += 1 + return + suppressed = _fail_closed_suppressed + _fail_closed_last_log_at = now + _fail_closed_suppressed = 0 + + logger.error( + "Invalid sparse-audio declaration: failing closed - no multimodal payload is routed " + "for the affected step(s). Fix the emitting model. Suppressed repeats before this " + "report: %d. Cause: %s", + suppressed, + cause, + ) + + +def is_sparse_audio_marker(value: Any) -> bool: + """Classify a ``meta.sparse_audio`` marker. + + Legal carriers (list of strings, bare string) evaluate against the + literal truthy set; ``None`` means "marker absent" and returns ``False``. + An ILLEGAL carrier (tensors, arrays, numbers, lists with non-string + items) is a producer bug: logged at error level once per offending type, + then raised as `InvalidSparseDeclarationError`. Deviant carriers are never + read element-wise, so a device tensor can never introduce a D2H sync + here. + """ + if value is None: + return False + if isinstance(value, str): + return value.lower() in _SPARSE_MARKER_TRUTHY + if isinstance(value, list) and all(isinstance(item, str) for item in value): + return any(item.lower() in _SPARSE_MARKER_TRUTHY for item in value) + vtype = type(value).__name__ + _error_once( + vtype, + 'Illegal sparse-audio marker of type %s: the marker MUST be a list of strings (e.g. ["1"]). ' + "The emitting model is out of contract; routing fails closed for the affected step(s).", + vtype, + ) + raise InvalidSparseDeclarationError(f"illegal sparse-audio marker carrier: {vtype}") + + +def resolve_sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None: + """Resolve the sparse request-id list. + + Returns the request-id list when a legal sparse declaration is present, + and ``None`` when there is no sparse declaration (dense). Raises + `InvalidSparseDeclarationError` when the declaration is present but out of + contract (illegal marker carrier, or marker set with an unusable + ``meta.req_id``). + """ + if not isinstance(multimodal_outputs, dict): + return None + meta = multimodal_outputs.get("meta") + nested_req_ids = meta.get("req_id") if isinstance(meta, dict) else None + flat_req_ids = multimodal_outputs.get("meta.req_id") + + # Resolve both marker encodings independently, even when the request ids + # came from the other encoding. Previously a nested `req_id` prevented the + # flattened marker from being read, while the mirror combination worked. + # Evaluating both also preserves the strict contract: an illegal carrier + # on either side wins over an otherwise-valid declaration and fails closed. + nested_marker_value = meta.get("sparse_audio") if isinstance(meta, dict) else None + flat_marker_value = multimodal_outputs.get("meta.sparse_audio") + nested_marker = is_sparse_audio_marker(nested_marker_value) + flat_marker = is_sparse_audio_marker(flat_marker_value) + if nested_marker_value is not None and flat_marker_value is not None and nested_marker != flat_marker: + _error_once( + "marker:conflict", + "Nested and flattened sparse-audio markers disagree. Failing closed.", + ) + raise InvalidSparseDeclarationError("nested and flattened sparse markers disagree") + sparse_audio = nested_marker or flat_marker + if not sparse_audio: + return None + + req_id_carriers = [ + (name, value) for name, value in (("nested", nested_req_ids), ("flattened", flat_req_ids)) if value is not None + ] + if not req_id_carriers: + req_id_carriers = [("missing", None)] + + validated_req_ids: list[list[str]] = [] + for carrier_name, req_ids in req_id_carriers: + if isinstance(req_ids, list) and all(isinstance(rid, str) for rid in req_ids): + validated_req_ids.append(req_ids) + continue + rtype = type(req_ids).__name__ + _error_once( + f"req_id:{carrier_name}:{rtype}", + "Sparse-audio marker is set but the %s meta.req_id carrier is unusable (%s; " + "must be a list of request-id strings). Failing closed.", + carrier_name, + rtype, + ) + raise InvalidSparseDeclarationError(f"sparse marker set but {carrier_name} meta.req_id is unusable: {rtype}") + + req_ids = validated_req_ids[0] + if len(validated_req_ids) == 2 and validated_req_ids[1] != req_ids: + _error_once( + "req_id:conflict", + "Nested and flattened sparse-audio request-id lists disagree. Failing closed.", + ) + raise InvalidSparseDeclarationError("nested and flattened sparse request-id lists disagree") + if len(req_ids) != len(set(req_ids)): + _error_once( + "req_id:duplicate", + "Sparse-audio request-id list contains duplicates. Failing closed.", + ) + raise InvalidSparseDeclarationError("sparse request-id list contains duplicates") + return list(req_ids) + + +def resolve_sparse_mm_routing( + *, + engine_output_type: str, + req_ids_output_copy: list[str], + downstream_req_ids: list[str], + multimodal_outputs: Any, +) -> tuple[list[str], dict[str, int], bool]: + """Resolve `(downstream_req_ids, sparse_mm_index, audio_sparse_output)`. + + The only entry point runners use; all protocol errors are absorbed here. + """ + try: + sparse_mm_req_ids = resolve_sparse_mm_req_ids(multimodal_outputs) + except InvalidSparseDeclarationError as exc: + if engine_output_type != "audio": + # Non-audio pipelines never consult sparse routing; the carrier + # error is already logged by the classifier. + return downstream_req_ids, {}, False + # FAIL CLOSED: a present-but-invalid sparse declaration most likely + # belongs to a producer that INTENDED sparse output; routing it as + # dense would batch-index-assign its payloads to the WRONG requests. + # Route no payload for this step instead - a visible, bounded gap + # rather than silent cross-request corruption. (Empty-sparse is an + # already-exercised path: the zero-audio sparse step ships + # req_id=[] through the same shape.) + _error_fail_closed_rate_limited(str(exc)) + return [], {}, True + sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])} + if engine_output_type != "audio" or sparse_mm_req_ids is None: + return downstream_req_ids, sparse_mm_index, False + + sparse_req_id_set = set(sparse_mm_req_ids) + sparse_downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set] + return sparse_downstream_req_ids, sparse_mm_index, True